diff --git a/.github/PULL_REQUEST_TEMPLATE b/.github/PULL_REQUEST_TEMPLATE index 4949331a2110d..c8f6b9a65abea 100644 --- a/.github/PULL_REQUEST_TEMPLATE +++ b/.github/PULL_REQUEST_TEMPLATE @@ -7,7 +7,7 @@ Thanks for sending a pull request! Here are some tips for you: 5. Please write your PR title to summarize what this PR proposes. 6. If possible, provide a concise example to reproduce the issue for a faster review. 7. If you want to add a new configuration, please read the guideline first for naming configurations in - 'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'. + 'common/utils/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'. 8. If you want to add or modify an error type or message, please read the guideline first in 'common/utils/src/main/resources/error/README.md'. --> diff --git a/.github/actions/checkout-and-sync/action.yml b/.github/actions/checkout-and-sync/action.yml new file mode 100644 index 0000000000000..8481678b51152 --- /dev/null +++ b/.github/actions/checkout-and-sync/action.yml @@ -0,0 +1,68 @@ +# +# 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: 'Checkout and Sync' +description: >- + Check out apache/spark at a pinned ref and, on forks, squash-merge the fork + branch on top so the build tests the combined result. Callers must add a + bare 'actions/checkout' step before invoking this action so that the latest + action definition is present in the workspace. +inputs: + ref: + description: 'Git ref to check out from apache/spark' + required: true + set-safe-directory: + description: 'Add GITHUB_WORKSPACE to git safe.directory (needed inside containers)' + required: false + default: 'false' +outputs: + head_sha: + description: 'The apache/spark HEAD SHA checked out (before the fork merge)' + value: ${{ steps.resolve-sha.outputs.head_sha }} +runs: + using: 'composite' + steps: + - name: Checkout Spark repository + uses: actions/checkout@v6 + with: + fetch-depth: 0 + repository: apache/spark + ref: ${{ inputs.ref }} + - name: Add GITHUB_WORKSPACE to git trust safe.directory + if: inputs.set-safe-directory == 'true' + shell: bash + run: git config --global --add safe.directory ${GITHUB_WORKSPACE} + - name: Resolve apache/spark HEAD SHA + id: resolve-sha + shell: bash + run: echo "head_sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT + - name: Sync the current branch with the latest in Apache Spark + if: github.repository != 'apache/spark' + shell: bash + run: | + echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV + git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} + git \ + -c user.name='Apache Spark Test Account' \ + -c user.email='sparktestacc@gmail.com' \ + merge --no-commit --progress --squash FETCH_HEAD + git \ + -c user.name='Apache Spark Test Account' \ + -c user.email='sparktestacc@gmail.com' \ + commit -m "Merged commit" --allow-empty diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index fd3ab715d6085..4f0bc158c7311 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -214,7 +214,7 @@ jobs: key: tpcds-${{ hashFiles('.github/workflows/benchmark.yml', 'sql/core/src/test/scala/org/apache/spark/sql/TPCDSSchema.scala') }} - name: Run benchmarks run: | - ./build/sbt -Pscala-${{ inputs.scala }} -Pyarn -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Pkinesis-asl -Pspark-ganglia-lgpl Test/package + ./build/sbt -Pscala-${{ inputs.scala }} -Pyarn -Pkubernetes -Phive -Phive-thriftserver -Phadoop-cloud -Pkinesis-asl -Pcredential-aws -Pspark-ganglia-lgpl Test/package # Make less noisy cp conf/log4j2.properties.template conf/log4j2.properties sed -i 's/rootLogger.level = info/rootLogger.level = warn/g' conf/log4j2.properties diff --git a/.github/workflows/branch43_scheduler.yml b/.github/workflows/branch43_scheduler.yml new file mode 100644 index 0000000000000..85c0ab74a4216 --- /dev/null +++ b/.github/workflows/branch43_scheduler.yml @@ -0,0 +1,114 @@ +# +# 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: "Branch-4.3 CI Scheduler" + +on: + schedule: + - cron: '0 0 */2 * *' + - cron: '0 1 */2 * *' + - cron: '0 6 */2 * *' + - cron: '0 8 */2 * *' + - cron: '0 10 */2 * *' + - cron: '0 14 */2 * *' + - cron: '0 20 */2 * *' + - cron: '0 23 */2 * *' + workflow_dispatch: + inputs: + target: + description: Target workflow to run + required: true + type: choice + default: all + options: + - all + - maven + - maven_java21 + - java17 + - java21 + - java25 + - non_ansi + - python_3.11 + - python_3.14 + +jobs: + schedule: + if: github.repository == 'apache/spark' + runs-on: ubuntu-latest + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + permissions: + actions: write + steps: + - name: Maven Build + if: >- + github.event.schedule == '0 10 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'maven' || inputs.target == 'all')) + run: | + gh workflow run build_maven.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Maven Java 21 Build + if: >- + github.event.schedule == '0 14 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'maven_java21' || inputs.target == 'all')) + run: | + gh workflow run build_maven_java21.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Java 17 Build + if: >- + github.event.schedule == '0 8 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'java17' || inputs.target == 'all')) + run: | + gh workflow run build_java17.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Java 21 Build + if: >- + github.event.schedule == '0 1 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'java21' || inputs.target == 'all')) + run: | + gh workflow run build_java21.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Java 25 Build + if: >- + github.event.schedule == '0 6 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'java25' || inputs.target == 'all')) + run: | + gh workflow run build_java25.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Non-ANSI Build + if: >- + github.event.schedule == '0 0 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'non_ansi' || inputs.target == 'all')) + run: | + gh workflow run build_non_ansi.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Python 3.11 Build + if: >- + github.event.schedule == '0 20 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'python_3.11' || inputs.target == 'all')) + run: | + gh workflow run build_python_3.11.yml --repo ${{ github.repository }} --ref branch-4.3 + - name: Python 3.14 Build + if: >- + github.event.schedule == '0 23 */2 * *' || + (github.event_name == 'workflow_dispatch' && + (inputs.target == 'python_3.14' || inputs.target == 'all')) + run: | + gh workflow run build_python_3.14.yml --repo ${{ github.repository }} --ref branch-4.3 diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 6bcd93d5340e4..fd6aa5114add0 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -51,6 +51,23 @@ on: required: false type: string default: '' + build_timeout_minutes: + description: >- + Timeout for each `build` matrix job. Raise it for a caller whose configuration makes the + tests slower, e.g. a non-default codegen compiler. + required: false + type: number + default: 150 + pyspark_timeout_minutes: + description: Timeout for each `pyspark` matrix job. Same reason as above. + required: false + type: number + default: 120 + docker_integration_tests_timeout_minutes: + description: Timeout for the `docker-integration-tests` job. Same reason as above. + required: false + type: number + default: 120 secrets: codecov_token: description: The upload token of codecov. @@ -69,7 +86,7 @@ jobs: outputs: required: ${{ steps.set-outputs.outputs.required }} # Pinned so every downstream job checks out the same snapshot, even if `master` advances mid-run. - head_sha: ${{ steps.resolve-sha.outputs.head_sha }} + head_sha: ${{ steps.checkout.outputs.head_sha }} image_docs_url: ${{ steps.infra-image-docs-outputs.outputs.image_docs_url }} image_docs_url_link: ${{ steps.infra-image-link.outputs.image_docs_url_link }} image_lint_url: ${{ steps.infra-image-lint-outputs.outputs.image_lint_url }} @@ -79,22 +96,13 @@ jobs: image_pyspark_url: ${{ steps.infra-image-pyspark-outputs.outputs.image_pyspark_url }} image_pyspark_url_link: ${{ steps.infra-image-link.outputs.image_pyspark_url_link }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + id: checkout + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ inputs.branch }} - - name: Resolve apache/spark HEAD SHA - id: resolve-sha - run: echo "head_sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Check all modules id: set-outputs run: | @@ -113,6 +121,9 @@ jobs: docker=`./dev/is-changed.py -m docker-integration-tests` # Skip PySpark, SparkR, TPC-DS when only static UI resources (JS/CSS/HTML) changed. # These tests are unaffected by UI static resource modifications. + # Also we only run the very slow transpile changes when that logic changes + transpile=false + static_only=false changed_files=$(git diff --name-only "$APACHE_SPARK_REF" HEAD 2>/dev/null) if [ -n "$changed_files" ]; then static_only=true @@ -122,8 +133,11 @@ jobs: *) static_only=false; break ;; esac done - else - static_only=false + for f in $changed_files; do + case "$f" in + *transpile*) transpile=true; break ;; + esac + done fi if [ "$static_only" = "true" ]; then pyspark=false @@ -150,7 +164,7 @@ jobs: docs=false java25=false fi - build=`./dev/is-changed.py -m "core,unsafe,kvstore,avro,utils,utils-java,network-common,network-shuffle,repl,launcher,examples,sketch,variant,api,catalyst,hive-thriftserver,mllib-local,mllib,graphx,streaming,sql-kafka-0-10,streaming-kafka-0-10,streaming-kinesis-asl,kubernetes,hadoop-cloud,spark-ganglia-lgpl,profiler,protobuf,yarn,connect,sql,hive,pipelines"` + build=`./dev/is-changed.py -m "core,unsafe,kvstore,avro,utils,utils-java,network-common,network-shuffle,repl,launcher,examples,sketch,variant,api,catalyst,hive-thriftserver,mllib-local,mllib,graphx,streaming,sql-kafka-0-10,streaming-kafka-0-10,streaming-kinesis-asl,credential-aws,kubernetes,hadoop-cloud,spark-ganglia-lgpl,profiler,protobuf,yarn,connect,sql,hive,pipelines"` build_core_utils=`./dev/is-changed.py -m "core,unsafe,kvstore,utils,utils-java,network-common,network-shuffle,sketch,variant,launcher"` precondition=" { @@ -170,6 +184,7 @@ jobs: \"k8s-integration-tests\" : \"$kubernetes\", \"buf\" : \"$buf\", \"ui\" : \"$ui\", + \"transpile\": \"$transpile\", }" echo $precondition # For debugging # Remove `\n` to avoid "Invalid format" error @@ -252,9 +267,9 @@ jobs: build: name: "Build modules: ${{ matrix.modules }} ${{ matrix.comment }}" needs: [precondition, precompile] - if: (!cancelled()) && fromJson(needs.precondition.outputs.required).build == 'true' + if: fromJson(needs.precondition.outputs.required).build == 'true' runs-on: ubuntu-latest - timeout-minutes: 150 + timeout-minutes: ${{ inputs.build_timeout_minutes }} strategy: fail-fast: false max-parallel: 20 @@ -276,7 +291,7 @@ jobs: mllib-local, mllib, graphx, profiler, pipelines, repl, examples - >- streaming, sql-kafka-0-10, streaming-kafka-0-10, streaming-kinesis-asl, - kubernetes, hadoop-cloud, spark-ganglia-lgpl, protobuf, connect, avro + credential-aws, kubernetes, hadoop-cloud, spark-ganglia-lgpl, protobuf, connect, avro - yarn # Here, we split Hive and SQL tests into some of slow ones and the rest of them. included-tags: [""] @@ -333,20 +348,12 @@ jobs: SKIP_MIMA: true SKIP_PACKAGING: true steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven uses: actions/cache@v5 @@ -367,6 +374,8 @@ jobs: coursier-${{ runner.os }}-${{ matrix.java }}-${{ matrix.hadoop }}- coursier-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true run: | if [ -f ./dev/free_disk_space ]; then ./dev/free_disk_space @@ -391,16 +400,10 @@ jobs: python3.12 -m pip install 'numpy>=1.23.2' pyarrow 'pandas==2.3.3' pyyaml scipy unittest-xml-reporting 'lxml==4.9.4' 'grpcio==1.76.0' 'grpcio-status==1.76.0' 'protobuf==6.33.5' 'zstandard==0.25.0' python3.12 -m pip list - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst @@ -413,10 +416,8 @@ jobs: export TERM=vt100 # Hive "other tests" test needs larger metaspace size based on experiment. if [[ "$MODULES_TO_TEST" == "hive" ]] && [[ "$EXCLUDED_TAGS" == "org.apache.spark.tags.SlowHiveTest" ]]; then export METASPACE_SIZE=2g; fi - if [ "${{ steps.extract-precompiled.outcome }}" = "success" ]; then - export SKIP_SCALA_BUILD=true - echo "Reusing precompiled artifact, skipping local SBT build." - fi + export SKIP_SCALA_BUILD=true + echo "Reusing precompiled artifact, skipping local SBT build." export SERIAL_SBT_TESTS=1 ./dev/run-tests --parallelism 1 --modules "$MODULES_TO_TEST" --included-tags "$INCLUDED_TAGS" --excluded-tags "$EXCLUDED_TAGS" - name: Upload test results to report @@ -461,33 +462,25 @@ jobs: packages: write steps: - name: Login to GitHub Container Registry - uses: docker/login-action@4907a6ddec9925e35a0a9e82d7399ccc52663121 + uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f with: registry: ghcr.io username: ${{ github.actor }} password: ${{ secrets.GITHUB_TOKEN }} - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Set up QEMU - uses: docker/setup-qemu-action@ce360397dd3f832beb865e1373c09c0e9f86d70a + uses: docker/setup-qemu-action@96fe6ef7f33517b61c61be40b68a1882f3264fb8 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c - name: Build and push (Documentation) if: ${{ fromJson(needs.precondition.outputs.required).docs == 'true' && hashFiles('dev/spark-test-image/docs/Dockerfile') != '' }} id: docker_build_docs - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/docs/ build-contexts: | @@ -500,7 +493,7 @@ jobs: - name: Build and push (Linter) if: ${{ fromJson(needs.precondition.outputs.required).lint == 'true' && hashFiles('dev/spark-test-image/lint/Dockerfile') != '' }} id: docker_build_lint - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/lint/ build-contexts: | @@ -513,7 +506,7 @@ jobs: - name: Build and push (SparkR) if: ${{ fromJson(needs.precondition.outputs.required).sparkr == 'true' && hashFiles('dev/spark-test-image/sparkr/Dockerfile') != '' }} id: docker_build_sparkr - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/sparkr/ push: true @@ -525,7 +518,7 @@ jobs: if: ${{ (fromJson(needs.precondition.outputs.required).pyspark == 'true' || fromJson(needs.precondition.outputs.required).pyspark-pandas == 'true') && env.PYSPARK_IMAGE_TO_TEST != '' }} id: docker_build_pyspark env: ${{ fromJSON(inputs.envs) }} - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/${{ env.PYSPARK_IMAGE_TO_TEST }}/ build-contexts: | @@ -552,27 +545,17 @@ jobs: name: "Precompile Spark" runs-on: ubuntu-latest timeout-minutes: 60 - # Optional optimization: if this job fails or is cancelled, the pyspark - # matrix entries fall back to running the SBT build locally as before. - continue-on-error: true env: HADOOP_PROFILE: ${{ inputs.hadoop }} HIVE_PROFILE: hive2.3 GITHUB_PREV_SHA: ${{ github.event.before }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven uses: actions/cache@v5 with: @@ -598,7 +581,7 @@ jobs: - name: Build Spark run: | ./build/sbt -Phadoop-3 -Pyarn -Pspark-ganglia-lgpl -Phadoop-cloud -Phive \ - -Pkubernetes -Pjvm-profiler -Pkinesis-asl -Phive-thriftserver \ + -Pkubernetes -Pjvm-profiler -Pkinesis-asl -Pcredential-aws -Phive-thriftserver \ -Pdocker-integration-tests -Pkubernetes-integration-tests -Pvolcano \ Test/package streaming-kinesis-asl-assembly/assembly connect/assembly assembly/package - name: Package compile output @@ -621,7 +604,7 @@ jobs: fromJson(needs.precondition.outputs.required).pyspark-pandas == 'true' name: "Build modules: ${{ matrix.modules }}" runs-on: ubuntu-latest - timeout-minutes: 120 + timeout-minutes: ${{ inputs.pyspark_timeout_minutes }} container: image: ${{ needs.precondition.outputs.image_pyspark_url_link }} options: >- @@ -678,24 +661,19 @@ jobs: METASPACE_SIZE: 1g BRANCH: ${{ inputs.branch }} PYSPARK_TEST_TIMEOUT: 450 + RUN_HYPOTHESIS: ${{ fromJson(needs.precondition.outputs.required).transpile }} + # Each generated example runs two full Spark jobs (transpiled vs. interpreted + # differential), so the local default of 1000 blows past PYSPARK_TEST_TIMEOUT. + # Cap CI fuzz volume; the explicit @example edge seeds still always run. + RUN_HYPOTHESIS_MAX_EXAMPLES: 50 steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Add GITHUB_WORKSPACE to git trust safe.directory - run: | - git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty + set-safe-directory: 'true' # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven uses: actions/cache@v5 @@ -715,6 +693,8 @@ jobs: restore-keys: | coursier-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true shell: 'script -q -e -c "bash {0}"' run: ./dev/free_disk_space_container - name: Install Java ${{ matrix.java }} @@ -735,16 +715,10 @@ jobs: echo "" done - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst @@ -754,10 +728,8 @@ jobs: if: ${{ matrix.modules != 'pyspark-connect-old-client' }} shell: 'script -q -e -c "bash {0}"' run: | - if [ "${{ steps.extract-precompiled.outcome }}" = "success" ]; then - export SKIP_SCALA_BUILD=true - echo "Reusing precompiled artifact, skipping local SBT build." - fi + export SKIP_SCALA_BUILD=true + echo "Reusing precompiled artifact, skipping local SBT build." if [[ "$MODULES_TO_TEST" == *"pyspark-install"* ]]; then export SKIP_PACKAGING=false echo "Python Packaging Tests Enabled!" @@ -770,13 +742,6 @@ jobs: SPARK_CONNECT_TESTING_REMOTE: sc://localhost if: ${{ matrix.modules == 'pyspark-connect-old-client' && inputs.branch == 'master' }} run: | - # Build Spark - if [ "${{ steps.extract-precompiled.outcome }}" = "success" ]; then - echo "Reusing precompiled artifact, skipping local SBT build." - else - ./build/sbt -Phive Test/package - fi - # Make less noisy cp conf/log4j2.properties.template conf/log4j2.properties sed -i 's/rootLogger.level = info/rootLogger.level = warn/g' conf/log4j2.properties @@ -859,23 +824,13 @@ jobs: SKIP_MIMA: true SKIP_PACKAGING: true steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Add GITHUB_WORKSPACE to git trust safe.directory - run: | - git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty + set-safe-directory: 'true' # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven uses: actions/cache@v5 @@ -895,6 +850,8 @@ jobs: restore-keys: | coursier-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true run: ./dev/free_disk_space_container - name: Install Java ${{ inputs.java }} uses: actions/setup-java@v5 @@ -902,16 +859,10 @@ jobs: distribution: zulu java-version: ${{ inputs.java }} - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst @@ -922,10 +873,8 @@ jobs: # R issues at docker environment export TZ=UTC export _R_CHECK_SYSTEM_CLOCK_=FALSE - if [ "${{ steps.extract-precompiled.outcome }}" = "success" ]; then - export SKIP_SCALA_BUILD=true - echo "Reusing precompiled artifact, skipping local SBT build." - fi + export SKIP_SCALA_BUILD=true + echo "Reusing precompiled artifact, skipping local SBT build." ./dev/run-tests --parallelism 1 --modules sparkr - name: Upload test results to report if: always() @@ -949,18 +898,12 @@ jobs: name: Protobuf breaking change detection and Python CodeGen check runs-on: ubuntu-latest steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Install Buf uses: bufbuild/buf-setup-action@a47c93e0b1648d5651a065437926377d060baa99 # v1 with: @@ -1000,22 +943,13 @@ jobs: container: image: ${{ needs.precondition.outputs.image_lint_url_link }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Add GITHUB_WORKSPACE to git trust safe.directory - run: | - git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty + set-safe-directory: 'true' # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven uses: actions/cache@v5 @@ -1042,6 +976,8 @@ jobs: restore-keys: | docs-maven-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true run: ./dev/free_disk_space_container - name: Install Java ${{ inputs.java }} uses: actions/setup-java@v5 @@ -1125,7 +1061,7 @@ jobs: run: | export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" export MAVEN_CLI_OPTS="--no-transfer-progress" - ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl clean install + ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws clean install # Documentation build docs: @@ -1143,22 +1079,13 @@ jobs: container: image: ${{ needs.precondition.outputs.image_docs_url_link }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Add GITHUB_WORKSPACE to git trust safe.directory - run: | - git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty + set-safe-directory: 'true' # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven uses: actions/cache@v5 @@ -1185,6 +1112,8 @@ jobs: restore-keys: | docs-maven-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true run: ./dev/free_disk_space_container - name: Install Java ${{ inputs.java }} uses: actions/setup-java@v5 @@ -1240,25 +1169,19 @@ jobs: # Any TPC-DS related updates on this job need to be applied to tpcds-1g-gen job of benchmark.yml as well tpcds-1g: needs: [precondition, precompile] - if: (!cancelled()) && fromJson(needs.precondition.outputs.required).tpcds-1g == 'true' + if: fromJson(needs.precondition.outputs.required).tpcds-1g == 'true' name: Run TPC-DS queries with SF=1 runs-on: ubuntu-latest timeout-minutes: 120 env: SPARK_LOCAL_IP: localhost steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven uses: actions/cache@v5 with: @@ -1282,16 +1205,10 @@ jobs: distribution: zulu java-version: ${{ inputs.java }} - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst @@ -1319,6 +1236,7 @@ jobs: SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 build/sbt "sql/testOnly org.apache.spark.sql.TPCDSQueryTestSuite" env: SPARK_ANSI_SQL_MODE: ${{ fromJSON(inputs.envs).SPARK_ANSI_SQL_MODE }} + SPARK_CODEGEN_COMPILER: ${{ fromJSON(inputs.envs).SPARK_CODEGEN_COMPILER }} SPARK_TPCDS_JOIN_CONF: | spark.sql.autoBroadcastJoinThreshold=-1 spark.sql.join.preferSortMergeJoin=true @@ -1327,6 +1245,7 @@ jobs: SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 build/sbt "sql/testOnly org.apache.spark.sql.TPCDSQueryTestSuite" env: SPARK_ANSI_SQL_MODE: ${{ fromJSON(inputs.envs).SPARK_ANSI_SQL_MODE }} + SPARK_CODEGEN_COMPILER: ${{ fromJSON(inputs.envs).SPARK_CODEGEN_COMPILER }} SPARK_TPCDS_JOIN_CONF: | spark.sql.autoBroadcastJoinThreshold=10485760 - name: Run TPC-DS queries (Shuffled hash join) @@ -1334,12 +1253,15 @@ jobs: SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 build/sbt "sql/testOnly org.apache.spark.sql.TPCDSQueryTestSuite" env: SPARK_ANSI_SQL_MODE: ${{ fromJSON(inputs.envs).SPARK_ANSI_SQL_MODE }} + SPARK_CODEGEN_COMPILER: ${{ fromJSON(inputs.envs).SPARK_CODEGEN_COMPILER }} SPARK_TPCDS_JOIN_CONF: | spark.sql.autoBroadcastJoinThreshold=-1 spark.sql.join.forceApplyShuffledHashJoin=true - name: Run TPC-DS queries on collated data run: | SPARK_TPCDS_DATA=`pwd`/tpcds-sf-1 build/sbt "sql/testOnly org.apache.spark.sql.TPCDSCollationQueryTestSuite" + env: + SPARK_CODEGEN_COMPILER: ${{ fromJSON(inputs.envs).SPARK_CODEGEN_COMPILER }} - name: Upload test results to report if: always() uses: actions/upload-artifact@v7 @@ -1364,10 +1286,10 @@ jobs: docker-integration-tests: needs: [precondition, precompile] - if: (!cancelled()) && fromJson(needs.precondition.outputs.required).docker-integration-tests == 'true' + if: fromJson(needs.precondition.outputs.required).docker-integration-tests == 'true' name: Run Docker integration tests runs-on: ubuntu-latest - timeout-minutes: 120 + timeout-minutes: ${{ inputs.docker_integration_tests_timeout_minutes }} env: HADOOP_PROFILE: ${{ inputs.hadoop }} HIVE_PROFILE: hive2.3 @@ -1377,19 +1299,12 @@ jobs: SKIP_MIMA: true SKIP_PACKAGING: true steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven uses: actions/cache@v5 with: @@ -1413,26 +1328,18 @@ jobs: distribution: zulu java-version: ${{ inputs.java }} - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst - name: Run tests env: ${{ fromJSON(inputs.envs) }} run: | - if [ "${{ steps.extract-precompiled.outcome }}" = "success" ]; then - export SKIP_SCALA_BUILD=true - echo "Reusing precompiled artifact, skipping local SBT build." - fi + export SKIP_SCALA_BUILD=true + echo "Reusing precompiled artifact, skipping local SBT build." ./dev/run-tests --parallelism 1 --modules docker-integration-tests --included-tags org.apache.spark.tags.DockerTest - name: Upload test results to report if: always() @@ -1458,24 +1365,17 @@ jobs: k8s-integration-tests: needs: [precondition, precompile] - if: (!cancelled()) && fromJson(needs.precondition.outputs.required).k8s-integration-tests == 'true' + if: fromJson(needs.precondition.outputs.required).k8s-integration-tests == 'true' name: Run Spark on Kubernetes Integration test runs-on: ubuntu-latest timeout-minutes: 120 steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ needs.precondition.outputs.head_sha }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven uses: actions/cache@v5 with: @@ -1494,6 +1394,8 @@ jobs: restore-keys: | coursier-${{ runner.os }}- - name: Free up disk space + timeout-minutes: 10 + continue-on-error: true run: | if [ -f ./dev/free_disk_space ]; then ./dev/free_disk_space @@ -1504,16 +1406,10 @@ jobs: distribution: zulu java-version: ${{ inputs.java }} - name: Download precompiled artifact - id: download-precompiled - if: needs.precompile.result == 'success' - continue-on-error: true uses: actions/download-artifact@v8 with: name: spark-compile-${{ inputs.branch }}-${{ github.run_id }} - name: Extract precompiled artifact - id: extract-precompiled - if: steps.download-precompiled.outcome == 'success' - continue-on-error: true run: | zstd -dc compile-artifact.tar.zst | tar -xf - rm compile-artifact.tar.zst diff --git a/.github/workflows/build_codegen_jdk.yml b/.github/workflows/build_codegen_jdk.yml new file mode 100644 index 0000000000000..3df5f94ce4845 --- /dev/null +++ b/.github/workflows/build_codegen_jdk.yml @@ -0,0 +1,65 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Runs the test suite with `spark.sql.codegen.compiler=jdk`, so that codegen +# reaching the JDK backend is covered by more than `CodeCompilerSuite`. The +# default backend is Janino, which leaves the javac path (SPARK-57403) exercised +# only by its own suite even though every generated unit could take it. +# +# javac is slower than Janino, so this runs on its own schedule rather than on +# every push, skips the jobs that do not compile generated code (docs, lint, buf, +# ui), and raises the timeout of the three jobs that no longer fit their default. +name: "Build / Codegen JDK backend (master, Scala 2.13, JDK 17)" + +on: + schedule: + - cron: '37 8 */2 * *' + workflow_dispatch: + +jobs: + run-build: + permissions: + packages: write + name: Run + uses: ./.github/workflows/build_and_test.yml + if: github.repository == 'apache/spark' + with: + java: 17 + branch: master + hadoop: hadoop3 + # See SPARK-57411 for the measurements behind these three. + build_timeout_minutes: 200 + pyspark_timeout_minutes: 130 + docker_integration_tests_timeout_minutes: 140 + envs: >- + { + "PYSPARK_IMAGE_TO_TEST": "python-312", + "PYTHON_TO_TEST": "python3.12", + "SPARK_CODEGEN_COMPILER": "jdk" + } + jobs: >- + { + "build": "true", + "pyspark": "true", + "pyspark-pandas": "true", + "sparkr": "true", + "tpcds-1g": "true", + "docker-integration-tests": "true", + "yarn": "true" + } diff --git a/.github/workflows/build_infra_images_cache.yml b/.github/workflows/build_infra_images_cache.yml index d5063f8db9d2b..69af217f45343 100644 --- a/.github/workflows/build_infra_images_cache.yml +++ b/.github/workflows/build_infra_images_cache.yml @@ -59,11 +59,11 @@ jobs: - name: Checkout Spark repository uses: actions/checkout@v6 - name: Set up QEMU - uses: docker/setup-qemu-action@ce360397dd3f832beb865e1373c09c0e9f86d70a + uses: docker/setup-qemu-action@96fe6ef7f33517b61c61be40b68a1882f3264fb8 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@4d04d5d9486b7bd6fa91e7baf45bbb4f8b9deedd + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c - name: Login to DockerHub - uses: docker/login-action@4907a6ddec9925e35a0a9e82d7399ccc52663121 + uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f with: registry: ghcr.io username: ${{ github.actor }} @@ -72,7 +72,7 @@ jobs: if: hashFiles('dev/spark-test-image/docs/Dockerfile') != '' id: docker_build_docs continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/docs/ build-contexts: | @@ -88,7 +88,7 @@ jobs: if: hashFiles('dev/spark-test-image/lint/Dockerfile') != '' id: docker_build_lint continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/lint/ build-contexts: | @@ -104,7 +104,7 @@ jobs: if: hashFiles('dev/spark-test-image/sparkr/Dockerfile') != '' id: docker_build_sparkr continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/sparkr/ push: true @@ -118,9 +118,11 @@ jobs: if: hashFiles('dev/spark-test-image/python-minimum/Dockerfile') != '' id: docker_build_pyspark_python_minimum continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-minimum/ + build-contexts: | + root=./ push: true tags: ghcr.io/apache/spark/apache-spark-github-action-image-pyspark-python-minimum-cache:${{ github.ref_name }}-static cache-from: type=registry,ref=ghcr.io/apache/spark/apache-spark-github-action-image-pyspark-python-minimum-cache:${{ github.ref_name }} @@ -132,7 +134,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-311/Dockerfile') != '' id: docker_build_pyspark_python_311 continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-311/ build-contexts: | @@ -148,9 +150,11 @@ jobs: if: hashFiles('dev/spark-test-image/python-312-classic-only/Dockerfile') != '' id: docker_build_pyspark_python_312_classic_only continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-312-classic-only/ + build-contexts: | + root=./ push: true tags: ghcr.io/apache/spark/apache-spark-github-action-image-pyspark-python-312-classic-only-cache:${{ github.ref_name }}-static cache-from: type=registry,ref=ghcr.io/apache/spark/apache-spark-github-action-image-pyspark-python-312-classic-only-cache:${{ github.ref_name }} @@ -162,7 +166,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-312/Dockerfile') != '' id: docker_build_pyspark_python_312 continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-312/ build-contexts: | @@ -178,7 +182,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-312-pandas-3/Dockerfile') != '' id: docker_build_pyspark_python_312_pandas_3 continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-312-pandas-3/ push: true @@ -192,7 +196,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-313/Dockerfile') != '' id: docker_build_pyspark_python_313 continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-313/ build-contexts: | @@ -208,7 +212,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-314/Dockerfile') != '' id: docker_build_pyspark_python_314 continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-314/ build-contexts: | @@ -224,7 +228,7 @@ jobs: if: hashFiles('dev/spark-test-image/python-314-nogil/Dockerfile') != '' id: docker_build_pyspark_python_314_nogil continue-on-error: true - uses: docker/build-push-action@bcafcacb16a39f128d818304e6c9c0c18556b85f + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a with: context: ./dev/spark-test-image/python-314-nogil/ push: true diff --git a/.github/workflows/build_main.yml b/.github/workflows/build_main.yml index e8f7054b2c32f..bdc7c21173bbc 100644 --- a/.github/workflows/build_main.yml +++ b/.github/workflows/build_main.yml @@ -29,12 +29,9 @@ jobs: permissions: packages: write name: Run - # Skip: - # - pushes to `branch-4.x` on apache/spark: post-merge CI on this - # integration/staging branch is disabled to save resources. - # - pushes to `master` on forks: the "Sync fork" button mirrors - # apache/spark and would otherwise re-run the full build on every sync. + # Skip pushes to `master` on forks: the "Sync fork" button mirrors + # apache/spark and would otherwise re-run the full build on every sync. if: >- - (github.repository == 'apache/spark' && github.ref != 'refs/heads/branch-4.x') - || (github.repository != 'apache/spark' && github.ref != 'refs/heads/master') + github.repository == 'apache/spark' + || github.ref != 'refs/heads/master' uses: ./.github/workflows/build_and_test.yml diff --git a/.github/workflows/maven_test.yml b/.github/workflows/maven_test.yml index f8edd0b723361..a3ea4c14af223 100644 --- a/.github/workflows/maven_test.yml +++ b/.github/workflows/maven_test.yml @@ -63,29 +63,20 @@ jobs: continue-on-error: true outputs: # Pinned so the build job checks out the same snapshot, even if the branch advances mid-run. - head_sha: ${{ steps.resolve-sha.outputs.head_sha }} + head_sha: ${{ steps.checkout.outputs.head_sha }} env: HADOOP_PROFILE: ${{ inputs.hadoop }} HIVE_PROFILE: hive2.3 SPARK_LOCAL_IP: localhost GITHUB_PREV_SHA: ${{ github.event.before }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + id: checkout + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ inputs.branch }} - - name: Resolve apache/spark HEAD SHA - id: resolve-sha - run: echo "head_sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven # TODO(SPARK-54466): https://github.com/actions/runner-images/issues/13341 if: ${{ runner.os != 'macOS' }} @@ -120,7 +111,7 @@ jobs: export MAVEN_OPTS="-Xss64m -Xmx4g -Xms4g -XX:ReservedCodeCacheSize=128m -Dorg.slf4j.simpleLogger.defaultLogLevel=WARN" export MAVEN_CLI_OPTS="--no-transfer-progress" export JAVA_VERSION=${{ inputs.java }} - ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} clean install + ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} clean install - name: Package compile output run: | # Exclude assembly/ from the artifact: 11 of 12 matrix entries wipe it @@ -171,7 +162,7 @@ jobs: - >- repl,sql#hive-thriftserver - >- - connector#kafka-0-10,connector#kafka-0-10-sql,connector#kafka-0-10-token-provider,connector#spark-ganglia-lgpl,connector#protobuf,connector#avro,connector#kinesis-asl + connector#kafka-0-10,connector#kafka-0-10-sql,connector#kafka-0-10-token-provider,connector#spark-ganglia-lgpl,connector#protobuf,connector#avro,connector#kinesis-asl,connector#credential-aws - >- sql#api,sql#catalyst,resource-managers#yarn,resource-managers#kubernetes#core - >- @@ -222,21 +213,13 @@ jobs: SPARK_LOCAL_IP: localhost GITHUB_PREV_SHA: ${{ github.event.before }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark # Fall back to the branch when the precompile job failed before resolving the SHA. ref: ${{ needs.precompile-maven.outputs.head_sha || inputs.branch }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven # TODO(SPARK-54466): https://github.com/actions/runner-images/issues/13341 @@ -325,10 +308,10 @@ jobs: # it here. if [ "$MODULES_TO_TEST" = "connect" ]; then echo "Building assembly module for connect tests." - ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} -pl assembly install + ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} -pl assembly install fi else - ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} clean install + ./build/mvn $MAVEN_CLI_OPTS -DskipTests -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} clean install # SPARK-51628: wipe the assembly module so tests exercise the # SPARK-51600 prepend fallback path. Connect tests strongly depend # on a built assembly module, so they are excluded. @@ -339,27 +322,27 @@ jobs: fi if [[ "$INCLUDED_TAGS" != "" ]]; then - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} -Dtest.include.tags="$INCLUDED_TAGS" test -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} -Dtest.include.tags="$INCLUDED_TAGS" test -fae elif [[ "$MODULES_TO_TEST" == "connect" && "$INPUT_BRANCH" == "branch-4.0" ]]; then # SPARK-53914: Remove sql/connect/client/jdbc from `-pl` for branch-4.0, this branch can be deleted after the EOL of branch-4.0. ./build/mvn $MAVEN_CLI_OPTS -Djava.version=${JAVA_VERSION/-ea} -pl sql/connect/client/jvm,sql/connect/common,sql/connect/server test -fae elif [[ "$MODULES_TO_TEST" == "connect" ]]; then ./build/mvn $MAVEN_CLI_OPTS -Djava.version=${JAVA_VERSION/-ea} -pl sql/connect/client/jdbc,sql/connect/client/jvm,sql/connect/common,sql/connect/server test -fae elif [[ "$EXCLUDED_TAGS" != "" ]]; then - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} -Dtest.exclude.tags="$EXCLUDED_TAGS" test -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} -Dtest.exclude.tags="$EXCLUDED_TAGS" test -fae elif [[ "$MODULES_TO_TEST" == *"sql#hive-thriftserver"* ]]; then # To avoid a compilation loop, for the `sql/hive-thriftserver` module, run `clean install` instead - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} clean install -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Phadoop-cloud -Pjvm-profiler -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} clean install -fae elif [[ "$MODULES_TO_TEST" == *"sql#pipelines"* && "$INPUT_BRANCH" == "branch-4.0" ]]; then # SPARK-52441: Remove sql/pipelines from TEST_MODULES for branch-4.0, this branch can be deleted after the EOL of branch-4.0. TEST_MODULES=${TEST_MODULES/,sql\/pipelines/} - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} test -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} test -fae elif [[ "$MODULES_TO_TEST" == *"common#utils-java"* && "$INPUT_BRANCH" == "branch-4.0" ]]; then # SPARK-53138: Remove common/utils-java from TEST_MODULES for branch-4.0, this branch can be deleted after the EOL of branch-4.0. TEST_MODULES=${TEST_MODULES/,common\/utils-java/} - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} test -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} test -fae else - ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Djava.version=${JAVA_VERSION/-ea} test -fae + ./build/mvn $MAVEN_CLI_OPTS -pl "$TEST_MODULES" -Pyarn -Pkubernetes -Pvolcano -Phive -Phive-thriftserver -Pspark-ganglia-lgpl -Phadoop-cloud -Pjvm-profiler -Pkinesis-asl -Pcredential-aws -Djava.version=${JAVA_VERSION/-ea} test -fae fi - name: Clean up local Maven repository run: | diff --git a/.github/workflows/publish_snapshot.yml b/.github/workflows/publish_snapshot.yml index 224f3c3207c07..80f450ba7995f 100644 --- a/.github/workflows/publish_snapshot.yml +++ b/.github/workflows/publish_snapshot.yml @@ -28,7 +28,7 @@ on: description: 'list of branches to publish (JSON)' required: true # keep in sync with default value of strategy matrix 'branch' - default: '["master", "branch-4.x", "branch-4.2", "branch-4.1", "branch-4.0", "branch-3.5"]' + default: '["master", "branch-4.x", "branch-4.3", "branch-4.2", "branch-4.1", "branch-4.0", "branch-3.5"]' jobs: publish-snapshot: @@ -39,7 +39,7 @@ jobs: max-parallel: 20 matrix: # keep in sync with default value of workflow_dispatch input 'branch' - branch: ${{ fromJSON( inputs.branch || '["master", "branch-4.x", "branch-4.2", "branch-4.1", "branch-4.0", "branch-3.5"]' ) }} + branch: ${{ fromJSON( inputs.branch || '["master", "branch-4.x", "branch-4.3", "branch-4.2", "branch-4.1", "branch-4.0", "branch-3.5"]' ) }} steps: - name: Checkout Spark repository uses: actions/checkout@v6 diff --git a/.github/workflows/python_hosted_runner_test.yml b/.github/workflows/python_hosted_runner_test.yml index 7488a0b7b9d15..b44098130c61f 100644 --- a/.github/workflows/python_hosted_runner_test.yml +++ b/.github/workflows/python_hosted_runner_test.yml @@ -67,29 +67,20 @@ jobs: continue-on-error: true outputs: # Pinned so the build job checks out the same snapshot, even if the branch advances mid-run. - head_sha: ${{ steps.resolve-sha.outputs.head_sha }} + head_sha: ${{ steps.checkout.outputs.head_sha }} env: HADOOP_PROFILE: ${{ inputs.hadoop }} HIVE_PROFILE: hive2.3 SPARK_LOCAL_IP: localhost GITHUB_PREV_SHA: ${{ github.event.before }} steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 + - name: Checkout and sync Spark repository + id: checkout + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark ref: ${{ inputs.branch }} - - name: Resolve apache/spark HEAD SHA - id: resolve-sha - run: echo "head_sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty - name: Cache SBT and Maven # TODO(SPARK-54466): https://github.com/actions/runner-images/issues/13341 if: ${{ runner.os != 'macOS' }} @@ -119,7 +110,7 @@ jobs: - name: Build Spark run: | ./build/sbt -Phadoop-3 -Pyarn -Pspark-ganglia-lgpl -Phadoop-cloud -Phive \ - -Pkubernetes -Pjvm-profiler -Pkinesis-asl -Phive-thriftserver \ + -Pkubernetes -Pjvm-profiler -Pkinesis-asl -Pcredential-aws -Phive-thriftserver \ -Pdocker-integration-tests -Pvolcano \ Test/package streaming-kinesis-asl-assembly/assembly connect/assembly assembly/package - name: Package compile output @@ -185,21 +176,13 @@ jobs: BRANCH: ${{ inputs.branch }} PYSPARK_TEST_TIMEOUT: 450 steps: - - name: Checkout Spark repository + - name: Bootstrap composite actions uses: actions/checkout@v6 - # In order to fetch changed files + - name: Checkout and sync Spark repository + uses: ./.github/actions/checkout-and-sync with: - fetch-depth: 0 - repository: apache/spark # Fall back to the branch when the precompile job failed before resolving the SHA. ref: ${{ needs.precompile.outputs.head_sha || inputs.branch }} - - name: Sync the current branch with the latest in Apache Spark - if: github.repository != 'apache/spark' - run: | - echo "APACHE_SPARK_REF=$(git rev-parse HEAD)" >> $GITHUB_ENV - git fetch https://github.com/$GITHUB_REPOSITORY.git ${GITHUB_REF#refs/heads/} - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' merge --no-commit --progress --squash FETCH_HEAD - git -c user.name='Apache Spark Test Account' -c user.email='sparktestacc@gmail.com' commit -m "Merged commit" --allow-empty # Cache local repositories. Note that GitHub Actions cache has a 10G limit. - name: Cache SBT and Maven # TODO(SPARK-54466): https://github.com/actions/runner-images/issues/13341 diff --git a/.github/workflows/test_report.yml b/.github/workflows/test_report.yml index 80e57b38d2758..8cfa152da7c67 100644 --- a/.github/workflows/test_report.yml +++ b/.github/workflows/test_report.yml @@ -41,20 +41,15 @@ jobs: github-token: ${{ secrets.GITHUB_TOKEN }} run-id: ${{ github.event.workflow_run.id }} pattern: "test-*" - - name: Check if test results exist - id: check - run: | - if find . -path '*/target/test-reports/*.xml' -print -quit | grep -q .; then - echo "has_results=true" >> $GITHUB_OUTPUT - else - echo "No test result files found. Skipping report." - echo "has_results=false" >> $GITHUB_OUTPUT - fi - name: Publish test report - if: steps.check.outputs.has_results == 'true' - uses: scacap/action-surefire-report@5609ce4db72c09db044803b344a8968fd1f315da + uses: EnricoMi/publish-unit-test-result-action@d0a4676d0e0b938bc201470d88276b7c74c712b3 # v2.24.0 with: check_name: Report test results - github_token: ${{ secrets.GITHUB_TOKEN }} - report_paths: "**/target/test-reports/*.xml" + files: "**/target/test-reports/*.xml" + comment_mode: off + large_files: true commit: ${{ github.event.workflow_run.head_commit.id }} + github_token: ${{ secrets.GITHUB_TOKEN }} + compare_to_earlier_commit: false + test_changes_limit: 0 + check_run_annotations: none diff --git a/.github/workflows/update_build_status.yml b/.github/workflows/update_build_status.yml index 6b16c59ce6ac5..1f1e8fafd5d92 100644 --- a/.github/workflows/update_build_status.yml +++ b/.github/workflows/update_build_status.yml @@ -46,6 +46,72 @@ jobs: // See https://docs.github.com/en/graphql/reference/enums#mergestatestatus const maybeReady = ['behind', 'clean', 'draft', 'has_hooks', 'unknown', 'unstable']; + // A fork workflow-run lookup can fail transiently (server errors, network drops, or + // REST rate limiting, which GitHub returns as 403 with rate-limit headers or 429). + // These should be retried on a later scheduled pass rather than reported to the + // contributor as a broken fork. Any other failure (e.g. a 404 for a missing + // build_main.yml) is treated as permanent. + // https://docs.github.com/en/rest/using-the-rest-api/rate-limits-for-the-rest-api + const isTransientError = (e) => { + if (!e.status) return true; + if (e.status >= 500 || e.status == 429) return true; + if (e.status == 403) { + const headers = (e.response && e.response.headers) || {}; + return !!headers['retry-after'] || headers['x-ratelimit-remaining'] === '0' + || /rate limit/i.test(e.message || ''); + } + return false; + }; + + // List all check-runs for a commit. per_page=100 (not the default 30) matches + // notify_test_workflow.yml: a SHA can accumulate more check-runs than one page + // (CI matrix, external checks, duplicate Build checks from reopened PRs), which + // could otherwise push the target Build check off the first page and leave the PR + // stuck in 'queued' forever. + const listCheckRuns = (ref) => github.paginate( + 'GET /repos/{owner}/{repo}/commits/{ref}/check-runs', + { + owner: context.repo.owner, + repo: context.repo.repo, + ref: ref, + per_page: 100 + } + ); + + // Parse a Build check's output text into the {owner, repo, run_id} the run fetch + // needs, or return null if it is absent, malformed, or missing a field. JSON.parse + // succeeding is not enough: a check from an older version, a manual run, or another + // app can carry null, {}, or unrelated JSON that parses but lacks these fields, and + // the run fetch would then fail. + const parseRunParams = (cr) => { + let params; + try { + params = JSON.parse(cr.output.text); + } catch (error) { + return null; + } + if (!params || !params.owner || !params.repo || !params.run_id) { + return null; + } + return params; + }; + + // A Build check this updater can actually sync: not action_required (notify writes + // that when it missed the fork run, and it carries no run params) and with output + // text carrying the {owner, repo, run_id} needed to fetch the fork run. A Build check + // with empty/malformed/fieldless output can never be synced, so it must not count as + // present - otherwise it would suppress the backfill and leave the PR stuck with an + // unsyncable check. + const isSyncableBuildCheck = (cr) => + cr.name == 'Build' && cr.conclusion != 'action_required' + && parseRunParams(cr) != null; + + // An action_required Build check: the contributor-facing "enable Actions / rebase" + // status notify writes when it found no fork run. It is not syncable, but it already + // carries the guidance the no-run backfill branch would create, so it is "useful". + const isActionRequiredBuildCheck = (cr) => + cr.name == 'Build' && cr.conclusion == 'action_required'; + // Iterate open PRs for await (const prs of github.paginate.iterator(endpoint,params)) { // Each page @@ -53,35 +119,38 @@ jobs: console.log('SHA: ' + pr.head.sha) console.log(' Mergeable status: ' + pr.mergeable_state) if (pr.mergeable_state == null || maybeReady.includes(pr.mergeable_state)) { - // Paginate with per_page=100 to match notify_test_workflow.yml. The default - // page size is 30, and a SHA can accumulate more check-runs than that (CI - // matrix, external checks, duplicate Build checks from reopened PRs), which - // could push the target Build check off the first page and leave the PR - // stuck in 'queued' forever. - const checkRuns = await github.paginate( - 'GET /repos/{owner}/{repo}/commits/{ref}/check-runs', - { - owner: context.repo.owner, - repo: context.repo.repo, - ref: pr.head.sha, - per_page: 100 - } - ) + const checkRuns = await listCheckRuns(pr.head.sha) - // Iterator GitHub Checks in the PR + // Does this SHA already carry an action_required Build check? notify (or an + // earlier pass of this updater) writes one when no fork run was found. It is + // not syncable, so it never suppresses the backfill below - which means a PR + // that permanently lacks a fork run (Actions disabled, old master) would + // otherwise re-poll on every 15-minute pass forever. When one is already + // present, the backfill's re-poll is unnecessary (see below). + const hasActionRequiredBuildCheck = checkRuns.some(isActionRequiredBuildCheck) + + // Track whether a syncable Build check exists (see isSyncableBuildCheck). + // notify_test_workflow.yml creates one per push; if that job never completed + // (e.g. cancelled while starved of an ASF runner) the check is missing and the + // backfill after this loop recreates it. Sync every match (no early break): a + // SHA can carry more than one Build check (reopened PRs, or a backfill that + // raced notify). Branch protection evaluates the newest check-run of a given + // name, so syncing only the first would leave a newer duplicate stuck in + // 'queued' and block the PR. + let syncableBuildCheck = false + // Build checks whose referenced fork run is permanently gone (404). They parse + // as syncable (their output still carries valid run params) but can never be + // synced, so the recheck below must not let them suppress the backfill. + const staleCheckIds = new Set() for await (const cr of checkRuns) { if (cr.name == 'Build' && cr.conclusion != "action_required") { - // text contains parameters to make request in JSON. A Build check - // created by something other than notify_test_workflow.yml (an older - // version, a manual run, or another app) may have empty or malformed - // output text; skip it instead of aborting the whole scheduled run, - // which would block updates for every PR queued behind it. - let params - try { - params = JSON.parse(cr.output.text) - } catch (error) { - console.error('Skipping Build check ' + cr.id + ' with unparseable output text') - console.error(error) + // Skip a check with unusable output (see parseRunParams) instead of + // aborting the whole scheduled run, which would block every PR queued + // behind it. Leaving syncableBuildCheck false lets the backfill below + // replace it rather than stranding the PR. + const params = parseRunParams(cr) + if (!params) { + console.error('Skipping Build check ' + cr.id + ' with unusable output') continue } @@ -91,10 +160,22 @@ jobs: run = await github.request('GET /repos/{owner}/{repo}/actions/runs/{run_id}', params) } catch (error) { console.error(error) - // Run not found. This can happen when the PR author removes GitHub Actions runs or - // disables GitHub Actions. + // The referenced fork run could not be fetched. A transient failure (5xx, + // rate limit) should be retried on a later pass, so count the check as + // syncable to suppress the backfill this pass. A permanent not-found (404: + // the PR author deleted the fork run or disabled GitHub Actions) means this + // check can never be synced, so leave syncableBuildCheck false and let the + // backfill below recreate a useful status - otherwise the flag would stay + // set and leave Build 'queued' indefinitely on every pass. + if (isTransientError(error)) { + syncableBuildCheck = true + } else { + staleCheckIds.add(cr.id) + } continue } + // Only now that the run is fetched is this a check we can actually sync. + syncableBuildCheck = true // Keep syncing the status of the checks if (run.data.status == 'completed') { @@ -122,8 +203,183 @@ jobs: details_url: run.data.details_url }) } + } + } - break + // No syncable Build check: notify_test_workflow.yml never created one, or only + // a stale action_required one exists. Recreate it here, mirroring notify. Skip + // if the head repo was deleted. + if (!syncableBuildCheck && pr.head.repo) { + const forkOwner = pr.head.repo.owner.login + const forkRepo = pr.head.repo.name + // Look up the fork's build_main.yml runs for the PR branch, matching + // notify_test_workflow.yml. Filter by branch (head.ref), not just head_sha: + // build_main.yml skips fork pushes to master (the "Sync fork" case), so the + // same SHA can carry a skipped/unrelated run on another branch, and a + // head_sha-only lookup could attach the Build check to that run instead of + // the PR branch's. Then pick the run whose head_sha equals the settled head + // SHA so stale runs from earlier commits on the branch cannot drive the + // decision. Re-poll a few times: the run is often not yet registered right + // after a push, and we must not create the sticky action_required check below + // over registration lag. A transient lookup error (5xx, network, rate limit) + // aborts to a later scheduled pass; any other error is treated as "no runs". + // + // When an action_required check already exists that lag risk is already + // realized, so a single sleepless lookup suffices - this stops a PR that + // permanently lacks a fork run (Actions disabled, old master) from re-polling + // (3 lookups plus two 3s sleeps) on every 15-minute pass forever, while still + // promoting it to a queued check on whichever pass first sees a run appear. + const attempts = hasActionRequiredBuildCheck ? 1 : 3 + let matched_run + let transient = false + for (let attempt = 0; attempt < attempts; attempt++) { + let forkRuns + try { + const runs = await github.request( + 'GET /repos/{owner}/{repo}/actions/workflows/{id}/runs', + { + owner: forkOwner, + repo: forkRepo, + id: 'build_main.yml', + branch: pr.head.ref + } + ) + forkRuns = runs.data.workflow_runs + } catch (error) { + console.error(error) + // Transient -> retry on a later pass. Permanent (e.g. a 404 for a + // missing build_main.yml) -> there is no run to find and it will not + // appear in 3s, so stop polling with matched_run undefined and let the + // action_required branch below handle it. Either way, do not burn the + // remaining retries and sleeps on an error that cannot resolve here. + if (isTransientError(error)) { + transient = true + } + break + } + matched_run = forkRuns.find(r => r.head_sha == pr.head.sha) + if (matched_run) { + break + } + if (attempt < attempts - 1) { + await new Promise(resolve => setTimeout(resolve, 3000)) + } + } + if (transient) { + console.log(' Fork run lookup failed transiently; will retry next pass') + continue + } + + // An action_required Build check already exists and we still found no run: + // the guidance check the no-run branch below would create is already present, + // so there is nothing to backfill. Skip the recheck listing and the create + // outright. (This is only reachable via the single, sleepless lookup above, + // so there was no ~6s window for notify to have created a syncable check in + // the meantime; even if it had, skipping is safe - with no run there is + // nothing to promote to 'queued' this pass anyway.) + if (hasActionRequiredBuildCheck && !matched_run) { + console.log(' action_required Build check present; nothing to backfill') + continue + } + + // Recheck for an existing Build check immediately before creating one. The + // poll above can span up to ~6 seconds, during which notify_test_workflow.yml + // may create the check; this recheck skips the backfill in that common case. + // It narrows but cannot fully close the window (there is no atomic + // create-if-absent, and the list endpoint is eventually consistent) - the + // sync loop above stays tolerant of duplicates so any that slip through still + // converge. + // + // Skip the backfill only if a check that is already at least as useful as the + // one we would create exists. When we found a fork run, that means a syncable + // check (the queued check we would create). When we found no run, that means a + // syncable OR an action_required check (the guidance check we would create). + // In both cases a malformed/fieldless check does NOT count - the sync loop + // cannot sync it, so refusing here would strand the PR on an unsyncable check; + // we fall through and create a usable one instead. A check whose fork run the + // sync loop just found permanently gone (staleCheckIds) is likewise not useful, + // even though its output still parses as syncable, so it too is excluded. + // + // The recheck listing and the create below are wrapped so a failure on one + // PR degrades to "skip this PR, retry next pass" rather than aborting the + // whole scheduled run and starving every PR later in the iteration - the same + // reason the sync loop above guards its calls. + try { + const recheck = await listCheckRuns(pr.head.sha) + const existingBuild = recheck.some(cr => + (isSyncableBuildCheck(cr) && !staleCheckIds.has(cr.id)) + || (!matched_run && isActionRequiredBuildCheck(cr))) + if (existingBuild) { + console.log(' Build check appeared during polling; skipping backfill') + continue + } + + if (matched_run) { + // A run exists for this head commit; point the check at it and let the + // next pass sync its status. + const actions_url = 'https://github.com/' + forkOwner + '/' + forkRepo + + '/actions/runs/' + matched_run.id + console.log(' Backfilling missing Build check -> ' + actions_url) + await github.rest.checks.create({ + owner: context.repo.owner, + repo: context.repo.repo, + name: 'Build', + head_sha: pr.head.sha, + status: 'queued', + output: { + title: 'Test results', + summary: '[See test results](' + actions_url + ')\n\n' + + 'If the tests fail for reasons unrelated to this pull request, ' + + 'please rerun the workflow in your forked repository.\n' + + 'If the failures are related to this pull request, ' + + 'please investigate them and push follow-up changes.', + text: JSON.stringify({ + owner: forkOwner, + repo: forkRepo, + run_id: matched_run.id + }) + }, + details_url: actions_url + }) + } else { + // No run for this head commit after re-polling (empty result, permanent + // lookup failure, or build_main.yml missing): Actions is disabled, the + // branch is on an old master, or the commit never triggered a run. Mirror + // notify's action_required check so the PR carries a required status + // telling the contributor how to fix it. + console.log(' No forked run for head SHA; creating action_required check') + await github.rest.checks.create({ + owner: context.repo.owner, + repo: context.repo.repo, + name: 'Build', + head_sha: pr.head.sha, + status: 'completed', + conclusion: 'action_required', + output: { + title: 'Workflow run detection failed', + summary: ` + Unable to detect the workflow run for testing the changes in your PR. + + 1. If you did not enable GitHub Actions in your forked repository, please enable it by clicking the button as shown in the image below. See also [Managing Github Actions Settings for a repository](https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/enabling-features-for-your-repository/managing-github-actions-settings-for-a-repository) for more details. + 2. It is possible your branch is based on the old \`master\` branch in Apache Spark, please sync your branch to the latest master branch. For example as below: + \`\`\`bash + git fetch upstream + git rebase upstream/master + git push origin YOUR_BRANCH --force + \`\`\``, + images: [ + { + alt: 'enabling workflows button', + image_url: 'https://raw.githubusercontent.com/apache/spark/master/.github/workflows/images/workflow-enable-button.png' + } + ] + } + }) + } + } catch (error) { + console.error(' Backfill failed for this PR; will retry next pass') + console.error(error) + continue } } } diff --git a/AGENTS.md b/AGENTS.md index 36f1eccb02dd0..8ec93d418ba92 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,7 +9,7 @@ Before the first code edit or running test in a session, ensure a clean working 3. If there are uncommitted changes (check with `git status`), ask the user to stash them before proceeding. 4. Switch to the appropriate branch: - **Existing PR**: resolve the PR branch name via `gh api repos/apache/spark/pulls/ --jq '.head.ref'`, then look for a local branch matching that name. If found, switch to it and inform the user. If not found, ask whether to fetch it or if there is a local branch under a different name. - - **New edits**: ask the user to choose: create a new git worktree from `/master` and work from there (recommended), or create and switch to a new branch from `/master`. + - **New edits**: ask the user to choose: create a new git worktree from `/master` with `--no-track` and work from there (recommended), or create and switch to a new branch from `/master` with `--no-track`. - **Running tests**: use `/master`. ## Development Notes @@ -18,6 +18,8 @@ SQL golden file tests are managed by `SQLQueryTestSuite` and its variants. Read Spark Connect protocol is defined in proto files under `sql/connect/common/src/main/protobuf/`. Read the README there before modifying proto definitions. +When adding a member to an existing class or object, follow the sectioning the file already uses and put the new member with the code it belongs with. The common failure mode is dropping it wherever it is first used without checking how the file is organized, splitting a section of unrelated code in the process. Beyond grouping related code there is no prescribed order -- the Databricks Scala guide, which Spark follows, asks only that a long class group its members into logical sections with comment headers. Do not reorganize existing members unless the change requires it. + Avoid introducing non-ASCII characters in code or comments. String literals may contain non-ASCII when the content requires it (error messages, test data, etc.). Identifiers are ASCII by convention. The common failure mode is typographic characters (em-dash, smart quotes, ellipsis, non-breaking space) sneaking into comments; scalastyle flags some of these. Spot-check before committing: `grep -rn -P "[^\x00-\x7F]" `. Keep source lines within 100 characters — the linters enforce this for Scala, Java, and Python, and LLMs commonly overrun it in comments and long expressions. A quick scan of just the changed files catches most cases in seconds, far cheaper than a CI round trip: @@ -101,6 +103,9 @@ These are combined with a base above rather than used on their own: Build and tests can take a long time. If the user explicitly asked to run tests, run them. Otherwise (you are running tests on your own to verify a change), first ask the user if they have more changes to make. +For build and test setup, including how to run tests and troubleshoot common +local failures, see `docs/building-spark.md`. + Prefer SBT over Maven for faster incremental compilation. Module names are defined in `project/SparkBuild.scala`. Compile a single module: @@ -202,6 +207,8 @@ It lists `master` and the latest major's release branches the commit reached (e. PR title format is `[SPARK-xxxx][COMPONENT] Title`. Draft, WIP, MINOR, and TRIVIAL PRs may omit the JIRA ID. The component tag is derived from the JIRA component name: take the last word and uppercase it (e.g. `Project Infra` → `[INFRA]`, `Spark Core` → `[CORE]`, `Structured Streaming` → `[STREAMING]`, `SQL` → `[SQL]`). +Use `[FOLLOWUP]` only for small PRs that directly modify or correct unreleased earlier PRs. For separately planned work, non-trivial changes, or work outside the earlier JIRA's scope, create a separate JIRA ticket for each PR and use the normal title format without `[FOLLOWUP]`. + Infer the PR title from the changes. If no ticket ID is given and the PR is not draft, WIP, MINOR, or TRIVIAL, create one using `dev/create_spark_jira.py`, using the PR title (without the JIRA ID and component tag) as the ticket title. python3 dev/create_spark_jira.py "" -c <component> { -t <type> | -p <parent-jira-id> } @@ -241,6 +248,8 @@ When exploring or working in any directory, always check for nested `AGENTS.md` directory and its ancestors. Read and follow every applicable file; instructions in a more specific directory take precedence for that directory's scope. +A directory may carry these instructions as `CLAUDE.md` rather than `AGENTS.md`, so check **both** names — a directory with only a `CLAUDE.md` has project instructions you must read too. Where both files exist, one is typically the real file and the other is a symlink to it, although either name may be the real file. Reading either one is sufficient. When adding instructions to a directory that has neither file, create AGENTS.md as the real file and add a CLAUDE.md symlink beside it (ln -s AGENTS.md CLAUDE.md). This keeps a single source of truth accessible under both names. + ## Versioning and Branch Policy When a change needs a version — `@since` annotations, config `.version("...")` (`SQLConf` / `*Conf`), new `MimaExcludes` sections, etc. — use the version of the branch it first ships in, with `-SNAPSHOT` stripped. Determine that branch: diff --git a/LICENSE-binary b/LICENSE-binary index 9d2a6a1acc6a5..7051e09a6383a 100644 --- a/LICENSE-binary +++ b/LICENSE-binary @@ -218,10 +218,7 @@ com.google.code.gson:gson com.google.crypto.tink:tink com.google.flatbuffers:flatbuffers-java com.google.guava:guava -com.jamesmurty.utils:java-xmlbuilder com.ning:compress-lzf -com.squareup.okhttp3:okhttp -com.squareup.okio:okio com.tdunning:json com.twitter:chill-java com.twitter:chill_2.13 diff --git a/README.md b/README.md index 717822ef42035..b57b53688f542 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,7 @@ This README file only contains basic setup instructions. | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java21.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_java21.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java25.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_java25.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_codegen_jdk.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_codegen_jdk.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_uds.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_uds.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_rockdb_as_ui_backend.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_rockdb_as_ui_backend.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_maven.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_maven.yml) | @@ -54,7 +55,8 @@ This README file only contains basic setup instructions. | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_connect40.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_python_connect40.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_connect.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_python_connect.yml) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_sparkr_window.yml/badge.svg)](https://github.com/apache/spark/actions/workflows/build_sparkr_window.yml) | -| branch-4.x | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java17.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_java17.yml?query=branch%3Abranch-4.x) | +| branch-4.x | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_main.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_main.yml?query=branch%3Abranch-4.x) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java17.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_java17.yml?query=branch%3Abranch-4.x) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java21.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_java21.yml?query=branch%3Abranch-4.x) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java25.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_java25.yml?query=branch%3Abranch-4.x) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml?query=branch%3Abranch-4.x) | @@ -62,6 +64,15 @@ This README file only contains basic setup instructions. | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_maven_java21.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_maven_java21.yml?query=branch%3Abranch-4.x) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_3.11.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_python_3.11.yml?query=branch%3Abranch-4.x) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_3.14.yml/badge.svg?branch=branch-4.x)](https://github.com/apache/spark/actions/workflows/build_python_3.14.yml?query=branch%3Abranch-4.x) | +| branch-4.3 | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_main.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_main.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java17.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_java17.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java21.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_java21.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java25.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_java25.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_non_ansi.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_maven.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_maven.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_maven_java21.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_maven_java21.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_3.11.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_python_3.11.yml?query=branch%3Abranch-4.3) | +| | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_python_3.14.yml/badge.svg?branch=branch-4.3)](https://github.com/apache/spark/actions/workflows/build_python_3.14.yml?query=branch%3Abranch-4.3) | | branch-4.2 | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_main.yml/badge.svg?branch=branch-4.2)](https://github.com/apache/spark/actions/workflows/build_main.yml?query=branch%3Abranch-4.2) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java17.yml/badge.svg?branch=branch-4.2)](https://github.com/apache/spark/actions/workflows/build_java17.yml?query=branch%3Abranch-4.2) | | | [![GitHub Actions Build](https://github.com/apache/spark/actions/workflows/build_java21.yml/badge.svg?branch=branch-4.2)](https://github.com/apache/spark/actions/workflows/build_java21.yml?query=branch%3Abranch-4.2) | diff --git a/assembly/pom.xml b/assembly/pom.xml index ade100e3bcd46..b1d00776bdf68 100644 --- a/assembly/pom.xml +++ b/assembly/pom.xml @@ -371,5 +371,19 @@ </dependency> </dependencies> </profile> + + <!-- + Pull in spark-credential-aws and its associated JARs, + --> + <profile> + <id>credential-aws</id> + <dependencies> + <dependency> + <groupId>org.apache.spark</groupId> + <artifactId>spark-credential-aws_${scala.binary.version}</artifactId> + <version>${project.version}</version> + </dependency> + </dependencies> + </profile> </profiles> </project> diff --git a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java index c4e1c5b96b1d2..b47bab42d6fa9 100644 --- a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java +++ b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java @@ -127,6 +127,7 @@ public enum LogKeys implements LogKey { CONFIG3, CONFIG4, CONFIG5, + CONFIGS, CONFIG_DEPRECATION_MESSAGE, CONFIG_KEY_UPDATED, CONFIG_VERSION, @@ -140,6 +141,7 @@ public enum LogKeys implements LogKey { CREATED_POOL_NAME, CREATION_SITE, CREDENTIALS_RENEWAL_INTERVAL_RATIO, + CREDENTIAL_VERSION, CROSS_VALIDATION_METRIC, CROSS_VALIDATION_METRICS, CSV_HEADER_COLUMN_NAME, @@ -366,6 +368,7 @@ public enum LogKeys implements LogKey { LOG_TYPE, LOSSES, LOWER_BOUND, + MAINTENANCE_TASK_TYPE, MALFORMATTED_STRING, MAP_ID, MASTER_URL, @@ -718,6 +721,7 @@ public enum LogKeys implements LogKey { SCHEMA, SCHEMA2, SERVER_NAME, + SERVICE_ACCOUNT_NAME, SERVICE_NAME, SERVLET_CONTEXT_HANDLER_PATH, SESSION_HANDLE, diff --git a/common/utils-java/src/main/java/org/apache/spark/internal/SparkLogger.java b/common/utils-java/src/main/java/org/apache/spark/internal/SparkLogger.java index 84d6d7cf4238c..ca7da3abff830 100644 --- a/common/utils-java/src/main/java/org/apache/spark/internal/SparkLogger.java +++ b/common/utils-java/src/main/java/org/apache/spark/internal/SparkLogger.java @@ -96,7 +96,7 @@ public void error(String msg, Throwable throwable) { } public void error(String msg, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.error(msg); } else if (slf4jLogger.isErrorEnabled()) { withLogContext(msg, mdcs, null, mt -> slf4jLogger.error(mt.message)); @@ -104,7 +104,7 @@ public void error(String msg, MDC... mdcs) { } public void error(String msg, Throwable throwable, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.error(msg, throwable); } else if (slf4jLogger.isErrorEnabled()) { withLogContext(msg, mdcs, throwable, mt -> slf4jLogger.error(mt.message, mt.throwable)); @@ -124,7 +124,7 @@ public void warn(String msg, Throwable throwable) { } public void warn(String msg, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.warn(msg); } else if (slf4jLogger.isWarnEnabled()) { withLogContext(msg, mdcs, null, mt -> slf4jLogger.warn(mt.message)); @@ -132,7 +132,7 @@ public void warn(String msg, MDC... mdcs) { } public void warn(String msg, Throwable throwable, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.warn(msg, throwable); } else if (slf4jLogger.isWarnEnabled()) { withLogContext(msg, mdcs, throwable, mt -> slf4jLogger.warn(mt.message, mt.throwable)); @@ -152,7 +152,7 @@ public void info(String msg, Throwable throwable) { } public void info(String msg, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.info(msg); } else if (slf4jLogger.isInfoEnabled()) { withLogContext(msg, mdcs, null, mt -> slf4jLogger.info(mt.message)); @@ -160,7 +160,7 @@ public void info(String msg, MDC... mdcs) { } public void info(String msg, Throwable throwable, MDC... mdcs) { - if (mdcs == null || mdcs.length == 0) { + if (isNullOrEmpty(mdcs)) { slf4jLogger.info(msg, throwable); } else if (slf4jLogger.isInfoEnabled()) { withLogContext(msg, mdcs, throwable, mt -> slf4jLogger.info(mt.message, mt.throwable)); @@ -215,6 +215,10 @@ public void trace(String msg, Throwable throwable) { slf4jLogger.trace(msg, throwable); } + private boolean isNullOrEmpty(MDC[] mdcs) { + return mdcs == null || mdcs.length == 0; + } + private void withLogContext( String pattern, MDC[] mdcs, diff --git a/common/utils/src/main/resources/error/README.md b/common/utils/src/main/resources/error/README.md index 575e2ebad35a3..7a628a418c7cb 100644 --- a/common/utils/src/main/resources/error/README.md +++ b/common/utils/src/main/resources/error/README.md @@ -157,6 +157,14 @@ Spark prefers to re-use existing SQLSTATEs, preferably used by multiple vendors. For extension Spark claims the `K**` sub-class range. If a new class is needed it will also claim the `K0` class. +Every error condition and its sub-conditions normally belong to a single error state: the +condition declares the SQLSTATE and the sub-conditions inherit it. As a documented +exception, a sub-condition may declare its own `sqlState` when regrouping would break +released clients that match the condition name on the wire. The only such exception is +`INVALID_HANDLE.SESSION_*`, pinned by a test in `SparkThrowableSuite`. Do not add +overrides, even within the same error class; a future compatibility layer rewriting +condition names per client version may remove the existing exception. + Internal errors should use the `XX` class. You can subdivide internal errors by component. For example: The existing `XXKD0` is used for an internal analyzer error. diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index bff7ea9bc264a..5af823a3f24b2 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -256,7 +256,7 @@ }, "AUXILIARY_TABLE_PROPERTY_MISSING" : { "message" : [ - "The internal auxiliary table is missing the required <propertyName> table property; cannot validate AutoCDC key columns. The auxiliary table metadata may be corrupted or have been modified externally. Perform a full refresh of the target table to recreate the auxiliary table." + "The internal auxiliary table is missing the required <propertyName> table property; cannot validate the AutoCDC configuration. The auxiliary table metadata may be corrupted or have been modified externally. Perform a full refresh of the target table to recreate the auxiliary table." ] }, "KEY_SCHEMA_DRIFT" : { @@ -268,6 +268,16 @@ "message" : [ "One or more AutoCDC flows writing to the target use SCD type <expectedScdType>, which is inconsistent with the SCD type recorded for it (recorded <recordedScdType>). AutoCDC does not support changing a target's SCD type across incremental pipeline runs. Correct the conflicting flow(s) or perform a full refresh of the target table." ] + }, + "SEQUENCING_TYPE_DRIFT" : { + "message" : [ + "One or more AutoCDC flows writing to the target use a sequencing expression of type <expectedSequencingType>, which is inconsistent with the sequencing type recorded for it (recorded <recordedSequencingType>). The sequencing expression may change across incremental pipeline runs, but its result type must not, so recorded sequencing values remain comparable. Correct the conflicting flow(s) or perform a full refresh of the target table." + ] + }, + "TRACK_HISTORY_DRIFT" : { + "message" : [ + "One or more AutoCDC flows writing to the target track history on columns [<expectedTrackHistoryColumns>], which are inconsistent with the track-history columns recorded for it (recorded [<recordedTrackHistoryColumns>]). AutoCDC does not support changing the SCD Type 2 track-history columns across incremental pipeline runs. Correct the conflicting flow(s) or perform a full refresh of the target table." + ] } }, "sqlState" : "42000" @@ -355,6 +365,13 @@ ], "sqlState" : "42613" }, + "AVRO_CANNOT_READ_NULL_FIELD" : { + "message" : [ + "Cannot read null value for <name> into a non-nullable SQL type.", + "To allow null values, declare the corresponding type as nullable in the read schema." + ], + "sqlState" : "22004" + }, "AVRO_CANNOT_WRITE_NULL_FIELD" : { "message" : [ "Cannot write null value for field <name> defined as non-null Avro data type <dataType>.", @@ -472,6 +489,12 @@ ], "sqlState" : "0A000" }, + "BITMAP_INPUT_TOO_LARGE" : { + "message" : [ + "The input bitmap has <inputNumBytes> bytes, which exceeds the maximum supported size of <maxNumBytes> bytes." + ], + "sqlState" : "22001" + }, "CALL_ON_STREAMING_DATASET_UNSUPPORTED" : { "message" : [ "The method <methodName> can not be called on streaming Dataset/DataFrame." @@ -809,6 +832,12 @@ ], "sqlState" : "22007" }, + "CANNOT_READ_ZIP_ENTRY" : { + "message" : [ + "Cannot read zip entry <entry> in archive <path>: encrypted or an unsupported compression method." + ], + "sqlState" : "KD003" + }, "CANNOT_RECOGNIZE_HIVE_TYPE" : { "message" : [ "Cannot recognize hive type string: <fieldType>, column: <fieldName>. The specified data type for the field cannot be recognized by Spark SQL. Please check the data type of the specified field and ensure that it is a valid Spark SQL data type. Refer to the Spark SQL documentation for a list of valid data types and their format. If the data type is correct, please ensure that you are using a supported version of Spark SQL." @@ -970,6 +999,13 @@ }, "sqlState" : "XX000" }, + "CHECKPOINT_DIRECTORY_NOT_SET" : { + "message" : [ + "Cannot checkpoint because no checkpoint directory is configured.", + "Set one with `SparkContext.setCheckpointDir` or the \"spark.checkpoint.dir\" configuration." + ], + "sqlState" : "55019" + }, "CHECKPOINT_FILE_CHECKSUM_VERIFICATION_FAILED" : { "message" : [ "Checksum verification failed, the file may be corrupted. File: <fileName>", @@ -985,6 +1021,13 @@ ], "sqlState" : "56000" }, + "CHECKPOINT_RDD_PARTITION_COUNT_MISMATCH" : { + "message" : [ + "The checkpoint of RDD <originalRDDId> has <newRDDLength> partition(s), but the RDD itself has <originalRDDLength>. The checkpoint RDD is <newRDDId>.", + "This usually means the checkpoint directory is not on storage that both the driver and the executors can read and write, or that some checkpoint files did not survive until the checkpoint was read back." + ], + "sqlState" : "58030" + }, "CHECK_CONSTRAINT_VIOLATION" : { "message" : [ "CHECK constraint <constraintName> <expression> violated by row with values:", @@ -1025,6 +1068,12 @@ ], "sqlState" : "42000" }, + "CLUSTER_MANAGER_APPLICATION_FAILURE" : { + "message" : [ + "Exiting due to an error reported by the cluster manager: <message>" + ], + "sqlState" : "56000" + }, "CODEC_NOT_AVAILABLE" : { "message" : [ "The codec <codecName> is not available." @@ -1224,6 +1273,14 @@ ], "sqlState" : "KD009" }, + "CONFLICTING_PIPELINE_FLOW_CASE_SENSITIVITY" : { + "message" : [ + "The flows writing to table <tableName> do not agree on '<configKey>': <flowConfigurations>.", + "The effective value determines whether column names that differ only in case identify the same column, so the table's schema would depend on the order the flows are evaluated in.", + "Set '<configKey>' to the same value for every flow writing to this table." + ], + "sqlState" : "42KD9" + }, "CONNECT" : { "message" : [ "Generic Spark Connect error." @@ -1921,11 +1978,36 @@ "Input schema <schema> can only contain STRING as a key type for a MAP." ] }, + "INVALID_JSON_PATH" : { + "message" : [ + "<functionName> has an invalid or unsupported JSON path <path>." + ] + }, + "INVALID_JSON_QUERY_RETURNING_TYPE" : { + "message" : [ + "<functionName> cannot return a value of type <returningType>. The RETURNING type must be a string type." + ] + }, + "INVALID_JSON_QUERY_WRAPPER_AND_QUOTES" : { + "message" : [ + "<functionName> cannot combine an OMIT QUOTES clause with a WITH ARRAY WRAPPER clause. OMIT QUOTES applies only to an unwrapped result." + ] + }, + "INVALID_JSON_SCALAR_RETURNING_TYPE" : { + "message" : [ + "<functionName> cannot return a value of type <returningType>. The RETURNING type must be a scalar (string, numeric, boolean, or datetime) type." + ] + }, "INVALID_JSON_SCHEMA" : { "message" : [ "Input schema <schema> must be a struct, an array, a map or a variant." ] }, + "INVALID_JSON_TABLE_PATH" : { + "message" : [ + "The <location> of JSON_TABLE has an invalid or unsupported JSON path <path>. Only simple, wildcard-free paths are supported (a trailing '[*]' is allowed on the row path)." + ] + }, "INVALID_MAP_KEY_TYPE" : { "message" : [ "The key of map cannot be/contain <keyType>." @@ -2253,6 +2335,12 @@ ], "sqlState" : "22003" }, + "DECIMAL_SCALE_EXCEEDS_PRECISION" : { + "message" : [ + "Decimal scale (<scale>) cannot be greater than precision (<precision>)." + ], + "sqlState" : "22003" + }, "DEFAULT_DATABASE_NOT_EXISTS" : { "message" : [ "Default database <defaultDatabase> does not exist, please create it first or change default database to `<defaultDatabase>`." @@ -2433,6 +2521,12 @@ ], "sqlState" : "42815" }, + "EMPTY_COLLECTION_NOT_ALLOWED" : { + "message" : [ + "empty collection" + ], + "sqlState" : "0A000" + }, "EMPTY_JSON_FIELD_VALUE" : { "message" : [ "Failed to parse an empty string for data type <dataType>." @@ -2598,6 +2692,12 @@ ], "sqlState" : "42822" }, + "FAILED_CREATE_CHECKPOINT_DIRECTORY" : { + "message" : [ + "Failed to create the checkpoint directory <path> as FileSystem.mkdirs returned false." + ], + "sqlState" : "58030" + }, "FAILED_EXECUTE_UDF" : { "message" : [ "User defined function (<functionName>: (<signature>) => <result>) failed due to: <reason>." @@ -3694,6 +3794,12 @@ }, "sqlState" : "42K03" }, + "INVALID_CHECKPOINT_DIRECTORY" : { + "message" : [ + "Cannot read the checkpoint directory <path>: expected the partition file <expectedFileName> but found <fileName>. The partition files must be numbered contiguously from part-00000." + ], + "sqlState" : "58030" + }, "INVALID_CLONE_SESSION_REQUEST" : { "message" : [ "Invalid session clone request." @@ -3899,9 +4005,16 @@ }, "INVALID_DRIVER_MEMORY" : { "message" : [ - "System memory <systemMemory> must be at least <minSystemMemory>.", - "Please increase heap size using the --driver-memory option or \"<config>\" in Spark configuration." + "Insufficient driver memory:" ], + "subClass" : { + "SYSTEM_MEMORY" : { + "message" : [ + "System memory <systemMemory> must be at least <minSystemMemory>.", + "Please increase heap size using the --driver-memory option or \"<config>\" in Spark configuration." + ] + } + }, "sqlState" : "F0000" }, "INVALID_EMPTY_LOCATION" : { @@ -3947,9 +4060,22 @@ }, "INVALID_EXECUTOR_MEMORY" : { "message" : [ - "Executor memory <executorMemory> must be at least <minSystemMemory>.", - "Please increase executor memory using the --executor-memory option or \"<config>\" in Spark configuration." + "Insufficient executor memory:" ], + "subClass" : { + "CONFIG_MEMORY" : { + "message" : [ + "Executor memory <executorMemory> must be at least <minSystemMemory>.", + "Please increase executor memory using the --executor-memory option or \"<config>\" in Spark configuration." + ] + }, + "SYSTEM_MEMORY" : { + "message" : [ + "System memory <systemMemory> must be at least <minSystemMemory>.", + "Please increase heap size using the --executor-memory option or \"<config>\" in Spark configuration." + ] + } + }, "sqlState" : "F0000" }, "INVALID_EXPLODE_EMBEDDED_ARRAY_SCHEMA" : { @@ -4137,17 +4263,20 @@ "SESSION_CHANGED" : { "message" : [ "The existing Spark server driver instance has restarted. Please reconnect." - ] + ], + "sqlState" : "08003" }, "SESSION_CLOSED" : { "message" : [ "Session was closed." - ] + ], + "sqlState" : "08003" }, "SESSION_NOT_FOUND" : { "message" : [ "Session not found." - ] + ], + "sqlState" : "08003" } }, "sqlState" : "HY000" @@ -4704,6 +4833,11 @@ "expects a long literal, but got <invalidValue>." ] }, + "NORMALIZE_FORM" : { + "message" : [ + "expects one of the normalization forms 'NFC', 'NFD', 'NFKC', 'NFKD', but got <form>." + ] + }, "NULL" : { "message" : [ "expects a non-NULL value." @@ -4744,6 +4878,11 @@ "expects one of the units 'HOUR', 'MINUTE', 'SECOND', 'MILLISECOND', 'MICROSECOND', but got '<invalidValue>'." ] }, + "TRIM_ARRAY_LENGTH" : { + "message" : [ + "Expects a value between 0 and the array cardinality (<numElements>), but got <length>." + ] + }, "ZERO_INDEX" : { "message" : [ "expects %1$, %2$ and so on, but got %0$." @@ -4800,6 +4939,18 @@ ], "sqlState" : "42602" }, + "INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA" : { + "message" : [ + "The incremental Python aggregate function <functionName> requires a struct (StructType) intermediate buffer schema, but got <bufferType>." + ], + "sqlState" : "42000" + }, + "INVALID_PYTHON_UDF_PLACEMENT" : { + "message" : [ + "The Python user-defined aggregate function(s) <functionList> cannot be invoked together with other kinds of aggregate function in the same aggregation. Multiple such functions may be used together, but not mixed with SQL or differently-typed aggregate functions." + ], + "sqlState" : "0A000" + }, "INVALID_QUERY_MIXED_QUERY_PARAMETERS" : { "message" : [ "Parameterized query must either use positional, or named parameters, but not both." @@ -5017,6 +5168,11 @@ "CREATE TEMPORARY TABLE ... USING ... is a deprecated syntax. To overcome the issue, please use CREATE TEMPORARY VIEW instead." ] }, + "DUPLICATE_JSON_TABLE_COLUMN" : { + "message" : [ + "JSON_TABLE has duplicate column name <columnName>. Column names in the COLUMNS clause must be unique." + ] + }, "EMPTY_IN_PREDICATE" : { "message" : [ "IN predicate requires at least one value. Empty IN clauses like 'IN ()' are not allowed. Consider using 'WHERE FALSE' if you need an always-false condition, or provide at least one value in the IN list." @@ -5271,6 +5427,18 @@ ], "sqlState" : "38000" }, + "INVALID_UDF_PARAMETER_PLACEHOLDER" : { + "message" : [ + "Invalid Python UDF transpiled-expression parameter placeholder: <placeholder>. Placeholders must be of the form `_udf_param_N` where N is a non-negative integer index into the UDF's positional arguments. This is an internal error in the Python UDF transpiler." + ], + "sqlState" : "42000" + }, + "INVALID_UDF_PARAMETER_PLACEHOLDER_INDEX" : { + "message" : [ + "Python UDF transpiled-expression parameter placeholder `_udf_param_<index>` references position <index> but the UDF only has <numParams> argument(s). This is an internal error in the Python UDF transpiler." + ], + "sqlState" : "42000" + }, "INVALID_URL" : { "message" : [ "The url is invalid: <url>. Use `try_parse_url` to tolerate invalid URL and return NULL instead." @@ -5489,6 +5657,48 @@ ], "sqlState" : "42K0E" }, + "JSON_EXISTS_ON_ERROR" : { + "message" : [ + "<functionName> could not evaluate the path <path>: the input is not valid JSON. This error was requested by the ERROR ON ERROR clause." + ], + "sqlState" : "2203G" + }, + "JSON_QUERY_ON_ERROR" : { + "message" : [ + "<functionName> could not extract a value at path <path>." + ], + "subClass" : { + "EMPTY" : { + "message" : [ + "The path matched no value. This error was requested by the ERROR ON EMPTY clause." + ] + }, + "ERROR" : { + "message" : [ + "The input is not valid JSON. This error was requested by the ERROR ON ERROR clause." + ] + } + }, + "sqlState" : "2203G" + }, + "JSON_VALUE_ON_ERROR" : { + "message" : [ + "<functionName> could not extract a scalar value at path <path>." + ], + "subClass" : { + "EMPTY" : { + "message" : [ + "The path matched no value. This error was requested by the ERROR ON EMPTY clause." + ] + }, + "ERROR" : { + "message" : [ + "The input is not valid JSON, the matched value is not a scalar, or it could not be cast to the RETURNING type. This error was requested by the ERROR ON ERROR clause." + ] + } + }, + "sqlState" : "2203G" + }, "KAFKA_DATA_SOURCE_NOT_ENABLED" : { "message" : [ "Failed to find data source: <provider>. Please deploy the application as per the deployment section of Structured Streaming + Kafka Integration Guide." @@ -5556,6 +5766,12 @@ ], "sqlState" : "42K03" }, + "LOCAL_BLOCK_DATA_NOT_FOUND" : { + "message" : [ + "Block <blockId> had block metadata but its data was missing from both the memory and disk stores. The block has been removed, so a later read must fetch it from another replica or recompute it." + ], + "sqlState" : "58030" + }, "LOCAL_MUST_WITH_SCHEMA_FILE" : { "message" : [ "LOCAL must be used together with the schema of `file`, but got: `<actualSchema>`." @@ -5586,6 +5802,39 @@ ], "sqlState" : "KD000" }, + "MALFORMED_EXPRESSION_INFO" : { + "message" : [ + "'<fieldName>' is malformed in the expression [<exprName>]:" + ], + "subClass" : { + "DEPRECATED" : { + "message" : [ + "it should start with a newline and 4 leading spaces; end with a newline and two spaces; however, got [<deprecated>]." + ] + }, + "GROUP" : { + "message" : [ + "it should be a value in <validGroups>; however, got [<group>]." + ] + }, + "NOTE" : { + "message" : [ + "it should start with a newline and 4 leading spaces; end with a newline and two spaces; however, got [<note>]." + ] + }, + "SINCE" : { + "message" : [ + "it should not start with a negative number; however, got [<since>]." + ] + }, + "SOURCE" : { + "message" : [ + "it should be a value in <validSources>; however, got [<source>]." + ] + } + }, + "sqlState" : "22023" + }, "MALFORMED_LOG_FILE" : { "message" : [ "Log file was malformed: failed to read correct log version from <text>." @@ -6405,6 +6654,18 @@ ], "sqlState" : "42K03" }, + "PIPELINED_SHUFFLE_CROSS_JOB_REUSE" : { + "message" : [ + "The pipelined shuffle <shuffleId> is being reused across jobs. A pipelined (incrementally-readable) shuffle is transient and has no retained output for another job to read, so its producer stage cannot be shared across jobs." + ], + "sqlState" : "0A000" + }, + "PIPELINED_SHUFFLE_UNSUPPORTED" : { + "message" : [ + "A pipelined shuffle stage group cannot be scheduled because it uses an unsupported feature: <reason>. Pipelined (incrementally-readable) shuffles run their producer and consumer stages concurrently over a transient, once-through stream, which is incompatible with this feature in this version." + ], + "sqlState" : "0A000" + }, "PIPELINE_DATASET_WITHOUT_FLOW" : { "message" : [ "Pipeline dataset <identifier> does not have any defined flows. Please attach a query with the dataset's definition, or explicitly define at least one flow that writes to the dataset." @@ -6737,6 +6998,24 @@ ], "sqlState" : "21000" }, + "SCHEDULER_BACKEND_SHUTDOWN_FAILED" : { + "message" : [ + "Failed to shut down the scheduler backend while the application was terminating:" + ], + "subClass" : { + "DRIVER_ENDPOINT" : { + "message" : [ + "the driver endpoint did not stop cleanly." + ] + }, + "EXECUTORS" : { + "message" : [ + "the driver could not confirm that the executors were asked to shut down, so the later shutdown steps were skipped." + ] + } + }, + "sqlState" : "58030" + }, "SCHEMA_ALREADY_EXISTS" : { "message" : [ "Cannot create schema <schemaName> because it already exists.", @@ -6777,6 +7056,12 @@ ], "sqlState" : "42K05" }, + "SHUFFLE_BLOCK_MIGRATION_NOT_SUPPORTED" : { + "message" : [ + "Shuffle block <blockId> cannot be migrated to the receiving executor: its shuffle block resolver <resolverClass> must implement MigratableResolver. See the cause for details." + ], + "sqlState" : "0A000" + }, "SKETCH_INVALID_LG_NOM_ENTRIES" : { "message" : [ "Invalid call to <function>; the `lgNomEntries` value must be between <min> and <max>, inclusive: <value>." @@ -7571,6 +7856,12 @@ "The checkpointing interval for async progress tracking must be set to 0, which means each progress update is checkpointed. Set option asyncProgressTrackingCheckpointIntervalMs to 0 in DataStreamWriter options and retry your query." ] }, + "CHECKPOINT_FORMAT_V1_NOT_SUPPORTED" : { + "message" : [ + "Real-time mode does not support state store checkpoint format v1, because a re-executed batch can reuse the state file names of a partially-written failed batch and lose data. If you encountered this while switching an existing query to real-time mode, use a fresh checkpoint location.", + "If you must continue with your existing checkpoint and accept that doing so exposes you to the possibility of data loss on failure, set <config> to true." + ] + }, "IDENTICAL_SOURCES_IN_UNION_NOT_SUPPORTED" : { "message" : [ "Real-time mode does not support union on two or more identical streaming data sources in a single query. This includes scenarios such as referencing the same source DataFrame more than once, or using two data sources with identical configurations for some sources. For Kafka, avoid reusing the same DataFrame and create different ones. Sources provided in the query: <sources>" @@ -7595,6 +7886,16 @@ "message" : [ "The <className> sink is currently not supported. See the Real-Time Mode User Guide for a list of supported sinks." ] + }, + "SQL_CONFIGURATION_NOT_SUPPORTED" : { + "message" : [ + "The following session configuration(s) are incompatible with Real-Time Mode: <invalidReasons>. Update or remove them and restart the query." + ] + }, + "STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED" : { + "message" : [ + "Streaming queries in real-time mode cannot include stateful operators (e.g. aggregate, deduplicate, transformWithState) before a union. Please restructure your query to apply the union before any stateful operations." + ] } }, "sqlState" : "0A000" @@ -7617,6 +7918,14 @@ ], "sqlState" : "XXKST" }, + "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT" : { + "message" : [ + "Streaming shuffle <shuffleId> writer <writerId> timed out after <timeoutMs> ms waiting for reader <readerId> to connect.", + "This can occur when the reader task is delayed because cluster resources are unavailable or when network connectivity prevents the reader from reaching the writer.", + "Check cluster resources and network connectivity. If the reader needs more time to start, increase 'spark.shuffle.streaming.writerConnectionTimeout', or set it to -1 to wait indefinitely." + ], + "sqlState" : "XXKST" + }, "STREAMING_STATEFUL_OPERATOR_MISSING_STATE_DIRECTORY" : { "message" : [ "Cannot restart streaming query with stateful operators because the state directory is empty or missing.", @@ -7717,6 +8026,12 @@ }, "sqlState" : "42601" }, + "TABLE_LOCATION_URI_NOT_SPECIFIED" : { + "message" : [ + "Table <identifier> did not specify locationUri." + ], + "sqlState" : "42601" + }, "TABLE_OR_VIEW_ALREADY_EXISTS" : { "message" : [ "Cannot create table or view <relationName> because it already exists.", @@ -7935,7 +8250,7 @@ }, "UNABLE_TO_ACQUIRE_MEMORY" : { "message" : [ - "Unable to acquire <requestedBytes> bytes of memory, got <receivedBytes>." + "Unable to acquire <requestedBytes> bytes of memory, got <receivedBytes>.<consumerBreakdown>" ], "sqlState" : "53200" }, @@ -7970,6 +8285,12 @@ ], "sqlState" : "42KD9" }, + "UNABLE_TO_REGISTER_WITH_EXTERNAL_SHUFFLE_SERVICE" : { + "message" : [ + "Unable to register with the external shuffle service: <message>" + ], + "sqlState" : "58030" + }, "UNBOUND_SQL_PARAMETER" : { "message" : [ "Found the unbound parameter: <name>. Please, fix `args` and provide a mapping of the parameter to either a SQL literal or collection constructor functions such as `map()`, `array()`, `struct()`." @@ -8069,7 +8390,8 @@ }, "UNRECOGNIZED_SQL_TYPE" : { "message" : [ - "Unrecognized SQL type - name: <typeName>, id: <jdbcType>." + "Unrecognized SQL type - name: <typeName>, id: <jdbcType>.", + "To read this column, map it explicitly with the `customSchema` option, or register a custom `JdbcDialect` that handles this type." ], "sqlState" : "42704" }, @@ -8196,6 +8518,29 @@ }, "sqlState" : "0A000" }, + "UNSUPPORTED_ARRAY_KEY" : { + "message" : [ + "Array keys are not supported by:" + ], + "subClass" : { + "HASH_PARTITIONER" : { + "message" : [ + "HashPartitioner." + ] + }, + "MAP_SIDE_COMBINE" : { + "message" : [ + "map-side combining." + ] + }, + "REDUCE_BY_KEY_LOCALLY" : { + "message" : [ + "reduceByKeyLocally()." + ] + } + }, + "sqlState" : "0A000" + }, "UNSUPPORTED_ARROWTYPE" : { "message" : [ "Unsupported arrow type <typeName>." @@ -8241,6 +8586,11 @@ "The row shall have a schema to get an index of the field <fieldName>." ] }, + "TASK_NOT_FINISHED" : { + "message" : [ + "The task has not finished yet, so its duration is not available." + ] + }, "WITHOUT_SUGGESTION" : { "message" : [ "" @@ -8574,7 +8924,8 @@ }, "LAMBDA_FUNCTION_WITH_PYTHON_UDF" : { "message" : [ - "Lambda function with Python UDF <funcName> in a higher order function." + "Cannot evaluate the Python UDF <funcName> inside the lambda of a higher-order function.", + "This placement is not supported. Rewrite the query so the UDF is applied outside the lambda." ] }, "LAMBDA_FUNCTION_WITH_SQL_UDF" : { @@ -8627,6 +8978,11 @@ "Multiple bucket TRANSFORMs." ] }, + "MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW" : { + "message" : [ + "Cannot use Python user-defined functions of different types together over a single window: <functionList>. Use a separate window specification for each." + ] + }, "MULTI_ACTION_ALTER" : { "message" : [ "The target JDBC server hosting table <tableName> does not support ALTER TABLE with multiple actions. Split the ALTER TABLE up into individual actions to avoid this error." @@ -9513,6 +9869,12 @@ ], "sqlState" : "42601" }, + "WRITING_JOB_FAILED" : { + "message" : [ + "Writing job failed." + ], + "sqlState" : "58030" + }, "WRONG_COMMAND_FOR_OBJECT_TYPE" : { "message" : [ "The operation <operation> requires a <requiredType>. But <objectName> is a <foundType>. Use <alternative> instead." @@ -9781,11 +10143,6 @@ "ADD COLUMN with v1 tables cannot specify NOT NULL." ] }, - "_LEGACY_ERROR_TEMP_1058" : { - "message" : [ - "Cannot create table with both USING <provider> and <serDeInfo>." - ] - }, "_LEGACY_ERROR_TEMP_1059" : { "message" : [ "STORED AS with file format '<serdeInfo>' is invalid." @@ -9821,11 +10178,6 @@ "Table <identifier> did not specify database." ] }, - "_LEGACY_ERROR_TEMP_1081" : { - "message" : [ - "Table <identifier> did not specify locationUri." - ] - }, "_LEGACY_ERROR_TEMP_1082" : { "message" : [ "Partition [<specString>] did not specify locationUri." @@ -10197,11 +10549,6 @@ "The SQL config '<configName>' was removed in the version <version>. <comment>" ] }, - "_LEGACY_ERROR_TEMP_1228" : { - "message" : [ - "Decimal scale (<scale>) cannot be greater than precision (<precision>)." - ] - }, "_LEGACY_ERROR_TEMP_1232" : { "message" : [ "Partition spec is invalid. The spec (<specKeys>) must match the partition spec (<partitionColumnNames>) defined in table '<tableName>'." @@ -10609,11 +10956,6 @@ "Missing database location." ] }, - "_LEGACY_ERROR_TEMP_2070" : { - "message" : [ - "Writing job failed." - ] - }, "_LEGACY_ERROR_TEMP_2071" : { "message" : [ "Commit denied for partition <partId> (task <taskId>, attempt <attemptId>, stage <stageId>.<stageAttempt>)." @@ -11282,21 +11624,6 @@ "empty RDD" ] }, - "_LEGACY_ERROR_TEMP_3008" : { - "message" : [ - "Cannot use map-side combining with array keys." - ] - }, - "_LEGACY_ERROR_TEMP_3009" : { - "message" : [ - "HashPartitioner cannot partition array keys." - ] - }, - "_LEGACY_ERROR_TEMP_3010" : { - "message" : [ - "reduceByKeyLocally() does not support array keys" - ] - }, "_LEGACY_ERROR_TEMP_3011" : { "message" : [ "This RDD lacks a SparkContext. It could happen in the following cases:", @@ -11314,148 +11641,16 @@ "Can only zip RDDs with same number of elements in each partition" ] }, - "_LEGACY_ERROR_TEMP_3014" : { - "message" : [ - "empty collection" - ] - }, "_LEGACY_ERROR_TEMP_3015" : { "message" : [ "countByValueApprox() does not support arrays" ] }, - "_LEGACY_ERROR_TEMP_3016" : { - "message" : [ - "Checkpoint directory has not been set in the SparkContext" - ] - }, - "_LEGACY_ERROR_TEMP_3017" : { - "message" : [ - "Invalid checkpoint file: <path>" - ] - }, - "_LEGACY_ERROR_TEMP_3018" : { - "message" : [ - "Failed to create checkpoint path <checkpointDirPath>" - ] - }, - "_LEGACY_ERROR_TEMP_3019" : { - "message" : [ - "Checkpoint RDD has a different number of partitions from original RDD. Original", - "RDD [ID: <originalRDDId>, num of partitions: <originalRDDLength>];", - "Checkpoint RDD [ID: <newRDDId>, num of partitions: <newRDDLength>]." - ] - }, - "_LEGACY_ERROR_TEMP_3020" : { - "message" : [ - "Checkpoint dir must be specified." - ] - }, - "_LEGACY_ERROR_TEMP_3021" : { - "message" : [ - "Error asking standalone scheduler to shut down executors" - ] - }, - "_LEGACY_ERROR_TEMP_3022" : { - "message" : [ - "Error stopping standalone scheduler's driver endpoint" - ] - }, - "_LEGACY_ERROR_TEMP_3023" : { - "message" : [ - "Can't run submitMapStage on RDD with 0 partitions" - ] - }, - "_LEGACY_ERROR_TEMP_3024" : { - "message" : [ - "attempted to access non-existent accumulator <id>" - ] - }, - "_LEGACY_ERROR_TEMP_3025" : { - "message" : [ - "TaskSetManagers should only send Resubmitted task statuses for tasks in ShuffleMapStages." - ] - }, - "_LEGACY_ERROR_TEMP_3026" : { - "message" : [ - "duration() called on unfinished task" - ] - }, "_LEGACY_ERROR_TEMP_3028" : { "message" : [ "<errorMsg>" ] }, - "_LEGACY_ERROR_TEMP_3029" : { - "message" : [ - "Exiting due to error from cluster scheduler: <message>" - ] - }, - "_LEGACY_ERROR_TEMP_3030" : { - "message" : [ - "Task <currentTaskAttemptId> has not locked block <blockId> for writing" - ] - }, - "_LEGACY_ERROR_TEMP_3031" : { - "message" : [ - "Block <blockId> does not exist" - ] - }, - "_LEGACY_ERROR_TEMP_3032" : { - "message" : [ - "Error occurred while waiting for replication to finish" - ] - }, - "_LEGACY_ERROR_TEMP_3033" : { - "message" : [ - "Unable to register with external shuffle server due to : <message>" - ] - }, - "_LEGACY_ERROR_TEMP_3034" : { - "message" : [ - "Error occurred while waiting for async. reregistration" - ] - }, - "_LEGACY_ERROR_TEMP_3035" : { - "message" : [ - "Unexpected shuffle block <blockId> with unsupported shuffle resolver <shuffleBlockResolver>" - ] - }, - "_LEGACY_ERROR_TEMP_3036" : { - "message" : [ - "Failure while trying to store block <blockId> on <blockManagerId>." - ] - }, - "_LEGACY_ERROR_TEMP_3037" : { - "message" : [ - "Block <blockId> was not found even though it's read-locked" - ] - }, - "_LEGACY_ERROR_TEMP_3038" : { - "message" : [ - "get() failed for block <blockId> even though we held a lock" - ] - }, - "_LEGACY_ERROR_TEMP_3039" : { - "message" : [ - "BlockManager returned null for BlockStatus query: <blockId>" - ] - }, - "_LEGACY_ERROR_TEMP_3040" : { - "message" : [ - "BlockManagerMasterEndpoint returned false, expected true." - ] - }, - "_LEGACY_ERROR_TEMP_3041" : { - "message" : [ - "" - ] - }, - "_LEGACY_ERROR_TEMP_3042" : { - "message" : [ - "Failed to get block <blockId>, which is not a shuffle block" - ] - }, "_LEGACY_ERROR_TEMP_3052" : { "message" : [ "Unexpected resolved action: <other>" @@ -11869,31 +12064,6 @@ "Read-ahead limit < 0" ] }, - "_LEGACY_ERROR_TEMP_3201" : { - "message" : [ - "'note' is malformed in the expression [<exprName>]. It should start with a newline and 4 leading spaces; end with a newline and two spaces; however, got [<note>]." - ] - }, - "_LEGACY_ERROR_TEMP_3202" : { - "message" : [ - "'group' is malformed in the expression [<exprName>]. It should be a value in <validGroups>; however, got <group>." - ] - }, - "_LEGACY_ERROR_TEMP_3203" : { - "message" : [ - "'source' is malformed in the expression [<exprName>]. It should be a value in <validSources>; however, got [<source>]." - ] - }, - "_LEGACY_ERROR_TEMP_3204" : { - "message" : [ - "'since' is malformed in the expression [<exprName>]. It should not start with a negative number; however, got [<since>]." - ] - }, - "_LEGACY_ERROR_TEMP_3205" : { - "message" : [ - "'deprecated' is malformed in the expression [<exprName>]. It should start with a newline and 4 leading spaces; end with a newline and two spaces; however, got [<deprecated>]." - ] - }, "_LEGACY_ERROR_TEMP_3206" : { "message" : [ "<value> is not a boolean string." diff --git a/common/utils/src/main/scala/org/apache/spark/ErrorClassesJSONReader.scala b/common/utils/src/main/scala/org/apache/spark/ErrorClassesJSONReader.scala index 18a47e7ee37a7..55153425ef134 100644 --- a/common/utils/src/main/scala/org/apache/spark/ErrorClassesJSONReader.scala +++ b/common/utils/src/main/scala/org/apache/spark/ErrorClassesJSONReader.scala @@ -124,10 +124,15 @@ class ErrorClassesJsonReader(jsonFileURLs: Seq[URL]) { } def getSqlState(errorClass: String): String = { - Option(errorClass) - .flatMap(_.split('.').headOption) - .flatMap(errorInfoMap.get) - .flatMap(_.sqlState) + val errorClasses = Option(errorClass).map(_.split('.')).getOrElse(Array.empty[String]) + val errorInfo = errorClasses.headOption.flatMap(errorInfoMap.get) + val subClassSqlState = errorClasses match { + case Array(_, subClass) => + errorInfo.flatMap(_.subClass).flatMap(_.get(subClass)).flatMap(_.sqlState) + case _ => None + } + subClassSqlState + .orElse(errorInfo.flatMap(_.sqlState)) .orNull } @@ -192,10 +197,13 @@ private case class ErrorInfo( * * @param message Message format with optional placeholders (e.g. <parm>). * The error message is constructed by concatenating the lines with newlines. + * @param sqlState SQLSTATE associated with this subclass. If absent, the subclass inherits + * the SQLSTATE of its main error class. * @param breakingChangeInfo Additional metadata if the error is due to a breaking change. */ private case class ErrorSubInfo( message: Seq[String], + sqlState: Option[String] = None, breakingChangeInfo: Option[BreakingChangeInfo] = None) { // For compatibility with multi-line error messages @JsonIgnore diff --git a/common/utils/src/main/scala/org/apache/spark/util/SparkClassUtils.scala b/common/utils/src/main/scala/org/apache/spark/util/SparkClassUtils.scala index 3f22719240c06..8b50c3da53333 100644 --- a/common/utils/src/main/scala/org/apache/spark/util/SparkClassUtils.scala +++ b/common/utils/src/main/scala/org/apache/spark/util/SparkClassUtils.scala @@ -164,6 +164,16 @@ private[spark] trait SparkClassUtils { currentClass = currentClass.getSuperclass } } + + /** + * Returns the companion object for the given class. + */ + def getCompanionObject(name: String): Any = { + val companionCls = classForName(name + "$") + // The companion object instance is held in a static MODULE$ field + val moduleField = companionCls.getField("MODULE$") + moduleField.get(null) + } } private[spark] object SparkClassUtils extends SparkClassUtils diff --git a/common/utils/src/main/scala/org/apache/spark/util/SparkCollectionUtils.scala b/common/utils/src/main/scala/org/apache/spark/util/SparkCollectionUtils.scala index 9c255daa0522d..956d9c8bc9e34 100644 --- a/common/utils/src/main/scala/org/apache/spark/util/SparkCollectionUtils.scala +++ b/common/utils/src/main/scala/org/apache/spark/util/SparkCollectionUtils.scala @@ -52,6 +52,8 @@ private[spark] trait SparkCollectionUtils { Arrays.fill(arr.asInstanceOf[Array[Byte]], defaultValue.asInstanceOf[Byte]) case c if c == classOf[Short] => Arrays.fill(arr.asInstanceOf[Array[Short]], defaultValue.asInstanceOf[Short]) + case c if c == classOf[Char] => + Arrays.fill(arr.asInstanceOf[Array[Char]], defaultValue.asInstanceOf[Char]) case c if c == classOf[Int] => Arrays.fill(arr.asInstanceOf[Array[Int]], defaultValue.asInstanceOf[Int]) case c if c == classOf[Long] => diff --git a/common/utils/src/test/scala/org/apache/spark/util/SparkCollectionUtilsSuite.scala b/common/utils/src/test/scala/org/apache/spark/util/SparkCollectionUtilsSuite.scala new file mode 100644 index 0000000000000..8040587e873a1 --- /dev/null +++ b/common/utils/src/test/scala/org/apache/spark/util/SparkCollectionUtilsSuite.scala @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.util + +import org.scalatest.funsuite.AnyFunSuite // scalastyle:ignore funsuite + +class SparkCollectionUtilsSuite extends AnyFunSuite { // scalastyle:ignore funsuite + + test("toMapWithIndex maps each key to its position, matching zipWithIndex.toMap") { + assert(SparkCollectionUtils.toMapWithIndex(Seq("a", "b", "c")) === + Map("a" -> 0, "b" -> 1, "c" -> 2)) + assert(SparkCollectionUtils.toMapWithIndex(Seq.empty[String]) === Map.empty[String, Int]) + // Duplicate keys keep the last index, matching zipWithIndex.toMap. + assert(SparkCollectionUtils.toMapWithIndex(Seq("a", "a")) === Map("a" -> 1)) + val keys = Seq(10, 20, 30, 40) + assert(SparkCollectionUtils.toMapWithIndex(keys) === keys.zipWithIndex.toMap) + } + + test("isEmpty and isNotEmpty handle null, empty and non-empty maps") { + val nullMap: java.util.Map[String, Int] = null + assert(SparkCollectionUtils.isEmpty(nullMap)) + assert(!SparkCollectionUtils.isNotEmpty(nullMap)) + + val empty = new java.util.HashMap[String, Int]() + assert(SparkCollectionUtils.isEmpty(empty)) + assert(!SparkCollectionUtils.isNotEmpty(empty)) + + val nonEmpty = new java.util.HashMap[String, Int]() + nonEmpty.put("a", 1) + assert(!SparkCollectionUtils.isEmpty(nonEmpty)) + assert(SparkCollectionUtils.isNotEmpty(nonEmpty)) + } + + test("createArray fills primitive-typed arrays with the default value") { + assert(SparkCollectionUtils.createArray(3, 7) === Array(7, 7, 7)) + assert(SparkCollectionUtils.createArray(2, true) === Array(true, true)) + assert(SparkCollectionUtils.createArray(2, 1L) === Array(1L, 1L)) + assert(SparkCollectionUtils.createArray(2, 2.5d) === Array(2.5d, 2.5d)) + assert(SparkCollectionUtils.createArray(2, 1.5f) === Array(1.5f, 1.5f)) + assert(SparkCollectionUtils.createArray(2, 3.toByte) === Array(3.toByte, 3.toByte)) + assert(SparkCollectionUtils.createArray(2, 4.toShort) === Array(4.toShort, 4.toShort)) + assert(SparkCollectionUtils.createArray(2, 'a') === Array('a', 'a')) + assert(SparkCollectionUtils.createArray(0, 7).isEmpty) + } + + test("createArray fills reference-typed arrays and returns empty for size 0") { + assert(SparkCollectionUtils.createArray(2, "x") === Array("x", "x")) + assert(SparkCollectionUtils.createArray(0, "x").isEmpty) + } +} diff --git a/common/utils/src/test/scala/org/apache/spark/util/SparkStringUtilsSuite.scala b/common/utils/src/test/scala/org/apache/spark/util/SparkStringUtilsSuite.scala new file mode 100644 index 0000000000000..93fb4a745d440 --- /dev/null +++ b/common/utils/src/test/scala/org/apache/spark/util/SparkStringUtilsSuite.scala @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.util + +import org.scalatest.funsuite.AnyFunSuite // scalastyle:ignore funsuite + +class SparkStringUtilsSuite extends AnyFunSuite { // scalastyle:ignore funsuite + + test("isBlank treats null and whitespace-only strings as blank") { + assert(SparkStringUtils.isBlank(null)) + assert(SparkStringUtils.isBlank("")) + assert(SparkStringUtils.isBlank(" ")) + assert(!SparkStringUtils.isBlank("a")) + assert(SparkStringUtils.isNotBlank("a")) + assert(!SparkStringUtils.isNotBlank(" ")) + } + + test("leftPad and rightPad pad with spaces up to the requested width") { + assert(SparkStringUtils.leftPad("hi", 5) === " hi") + assert(SparkStringUtils.rightPad("hi", 5) === "hi ") + // A width that is not larger than the input is a no-op, and null is passed through. + assert(SparkStringUtils.leftPad("hello", 5) === "hello") + assert(SparkStringUtils.rightPad("hello", 3) === "hello") + assert(SparkStringUtils.leftPad(null, 5) === null) + assert(SparkStringUtils.rightPad(null, 5) === null) + } + + test("rightPad repeats the pad string and truncates the last repetition") { + assert(SparkStringUtils.rightPad("x", 6, "*") === "x*****") + assert(SparkStringUtils.rightPad("a", 5, "xy") === "axyxy") + // Only the first character of the final "xy" fits within the width. + assert(SparkStringUtils.rightPad("a", 4, "xy") === "axyx") + assert(SparkStringUtils.rightPad("hello", 3, "*") === "hello") + assert(SparkStringUtils.rightPad(null, 5, "*") === null) + } + + test("rightPad with an empty pad string fails when padding is required") { + // The pad length is used as a divisor, so an empty pad string divides by zero. + // Documented here rather than guarded, since callers are expected to pass a + // non-empty pad string. + intercept[ArithmeticException] { + SparkStringUtils.rightPad("a", 5, "") + } + // No padding is needed in these cases, so the divisor is never reached. + assert(SparkStringUtils.rightPad("hello", 5, "") === "hello") + assert(SparkStringUtils.rightPad("hello", 3, "") === "hello") + assert(SparkStringUtils.rightPad(null, 5, "") === null) + } + + test("abbreviate truncates with the marker and leaves short inputs untouched") { + assert(SparkStringUtils.abbreviate("hello world", 8) === "hello...") + assert(SparkStringUtils.abbreviate("abc", 8) === "abc") + assert(SparkStringUtils.abbreviate("abcdefgh", "..", 4) === "ab..") + // An empty marker makes abbreviate a plain truncation. + assert(SparkStringUtils.abbreviate("abcdefgh", "", 3) === "abc") + assert(SparkStringUtils.abbreviate(null, 5) === null) + assert(SparkStringUtils.abbreviate("abc", null, 2) === null) + } + + test("strip removes the given prefix and suffix when present") { + assert(SparkStringUtils.strip("\"path\"", "\"") === "path") + assert(SparkStringUtils.strip("path", "\"") === "path") + assert(SparkStringUtils.strip(null, "\"") === null) + assert(SparkStringUtils.strip("path", null) === "path") + } + + test("stringToSeq trims entries and drops empty ones") { + assert(SparkStringUtils.stringToSeq(" a, b ,,c ") === Seq("a", "b", "c")) + assert(SparkStringUtils.stringToSeq("") === Seq.empty) + } +} diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java b/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java index 8996caeb364cf..cefb957f9b976 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java +++ b/common/variant/src/main/java/org/apache/spark/types/variant/Variant.java @@ -144,18 +144,17 @@ public int objectSize() { // It is only legal to call it when `getType()` is `Type.OBJECT`. public Variant getFieldByKey(String key) { return handleObject(value, pos, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> { - // Use linear search for a short list. Switch to binary search when the length reaches - // `BINARY_SEARCH_THRESHOLD`. - final int BINARY_SEARCH_THRESHOLD = 32; - if (size < BINARY_SEARCH_THRESHOLD) { - for (int i = 0; i < size; ++i) { - int id = readUnsigned(value, idStart + idSize * i, idSize); - if (key.equals(getMetadataKey(metadata, id))) { - int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize); - return new Variant(value, metadata, dataStart + offset); - } + byte[] keyBytes = encodeKey(key); + int numAttempts = 1; + // UTF-8 and UTF-16 order can differ only for keys with a code unit at or above U+D800. + for (int i = 0; i < key.length(); ++i) { + if (key.charAt(i) >= Character.MIN_SURROGATE) { + numAttempts = 2; + break; } - } else { + } + // Search the spec's UTF-8 order first, then the UTF-16 order written by older Spark versions. + for (int attempt = 0; attempt < numAttempts; ++attempt) { int low = 0; int high = size - 1; while (low <= high) { @@ -164,7 +163,10 @@ public Variant getFieldByKey(String key) { // overflows int. int mid = (low + high) >>> 1; int id = readUnsigned(value, idStart + idSize * mid, idSize); - int cmp = getMetadataKey(metadata, id).compareTo(key); + String midKey = getMetadataKey(metadata, id); + int cmp = attempt == 0 + ? compareKeys(encodeKey(midKey), keyBytes) + : midKey.compareTo(key); if (cmp < 0) { low = mid + 1; } else if (cmp > 0) { diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java index 296bc339ee084..ef00c8379b89c 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java +++ b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java @@ -171,6 +171,17 @@ public static Variant arrayAppendAtPath(Variant v, PathSegment[] segments, Varia return builder.result(); } + // Return a new variant with null-valued object fields removed, recursing into nested objects + // and arrays. When `includeArrays` is true, null array elements are removed too; when false, + // arrays keep their nulls but objects inside them are still cleaned. A container emptied by + // stripping is preserved as {} / [] rather than collapsed to a variant null, and a top-level + // variant null is returned unchanged. The result is always rebuilt with fresh metadata. + public static Variant stripNulls(Variant v, boolean includeArrays) { + VariantBuilder builder = new VariantBuilder(false); + builder.appendWithNullStrippingImpl(v.value, v.metadata, v.pos, includeArrays); + return builder.result(); + } + // Build the variant metadata from `dictionaryKeys` and return the variant result. public Variant result() { int numKeys = dictionaryKeys.size(); @@ -365,7 +376,7 @@ public int addKey(String key) { } else { id = dictionaryKeys.size(); dictionary.put(key, id); - dictionaryKeys.add(key.getBytes(StandardCharsets.UTF_8)); + dictionaryKeys.add(encodeKey(key)); } return id; } @@ -899,6 +910,54 @@ private void appendNewPath(PathSegment[] segments, int depth, Variant val) { } } + private void appendWithNullStrippingImpl( + byte[] value, byte[] metadata, int pos, boolean includeArrays) { + checkIndex(pos, value.length); + int basicType = value[pos] & BASIC_TYPE_MASK; + if (basicType == OBJECT) { + handleObject(value, pos, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> { + ArrayList<FieldEntry> fields = new ArrayList<>(size); + int start = writePos; + for (int i = 0; i < size; ++i) { + int id = readUnsigned(value, idStart + idSize * i, idSize); + int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize); + int elementPos = dataStart + offset; + // Drop the whole field when its value is a variant null. + if (getType(value, elementPos) == Type.NULL) { + continue; + } + String key = getMetadataKey(metadata, id); + int newId = addKey(key); + fields.add(new FieldEntry(key, newId, writePos - start)); + appendWithNullStrippingImpl(value, metadata, elementPos, includeArrays); + } + finishWritingObject(start, fields); + return null; + }); + } else if (basicType == ARRAY) { + handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> { + ArrayList<Integer> offsets = new ArrayList<>(size); + int start = writePos; + for (int i = 0; i < size; ++i) { + int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize); + int elementPos = dataStart + offset; + // Drop variant-null elements only when stripping arrays; otherwise keep them but still + // recurse into nested containers. + if (includeArrays && getType(value, elementPos) == Type.NULL) { + continue; + } + offsets.add(writePos - start); + appendWithNullStrippingImpl(value, metadata, elementPos, includeArrays); + } + finishWritingArray(start, offsets); + return null; + }); + } else { + // Scalars and standalone variant nulls are appended unchanged. + appendVariantImpl(value, metadata, pos); + } + } + // Append the variant value without rewriting or creating any metadata. This is used when // building an object during shredding, where there is a fixed pre-existing metadata that // all shredded values will refer to. @@ -935,6 +994,7 @@ public static final class FieldEntry implements Comparable<FieldEntry> { final String key; final int id; final int offset; + private byte[] keyBytes; public FieldEntry(String key, int id, int offset) { this.key = key; @@ -946,9 +1006,16 @@ FieldEntry withNewOffset(int newOffset) { return new FieldEntry(key, id, newOffset); } + private byte[] keyBytes() { + if (keyBytes == null) { + keyBytes = encodeKey(key); + } + return keyBytes; + } + @Override public int compareTo(FieldEntry other) { - return key.compareTo(other.key); + return compareKeys(keyBytes(), other.keyBytes()); } } diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java index 681c4038a9a98..184a7b29a9781 100644 --- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java +++ b/common/variant/src/main/java/org/apache/spark/types/variant/VariantUtil.java @@ -688,6 +688,16 @@ private static void validateImpl(byte[] value, byte[] metadata, int pos) { } } + // Encode an object field key for comparison in the order required by the Variant spec. + public static byte[] encodeKey(String key) { + return key.getBytes(StandardCharsets.UTF_8); + } + + // Compare UTF-8-encoded object field keys using unsigned lexicographic byte ordering. + public static int compareKeys(byte[] left, byte[] right) { + return Arrays.compareUnsigned(left, right); + } + // Get a key at `id` in the variant metadata. // Throw `MALFORMED_VARIANT` if the variant is malformed. An out-of-bound `id` is also considered // a malformed variant because it is read from the corresponding variant value. diff --git a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala index 49d597bfc8a77..7bffcc7da9bab 100644 --- a/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala +++ b/connector/avro/src/test/scala/org/apache/spark/sql/avro/AvroCatalystDataConversionSuite.scala @@ -24,7 +24,7 @@ import org.apache.avro.Schema import org.apache.avro.generic.{GenericData, GenericRecordBuilder} import org.apache.avro.message.{BinaryMessageDecoder, BinaryMessageEncoder} -import org.apache.spark.SparkException +import org.apache.spark.{SparkException, SparkRuntimeException} import org.apache.spark.sql.{RandomDataGenerator, Row} import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow, NoopFilters, OrderedFilters, StructFilters} import org.apache.spark.sql.catalyst.expressions.{ExpressionEvalHelper, GenericInternalRow, Literal} @@ -300,6 +300,72 @@ class AvroCatalystDataConversionSuite extends SharedSparkSession } } + private def deserializerFor(schema: Schema, dataType: DataType): AvroDeserializer = { + new AvroDeserializer( + schema, + dataType, + false, + RebaseSpec(LegacyBehaviorPolicy.CORRECTED), + new NoopFilters, + false, + "", + -1) + } + + test("SPARK-58218: null array element for non-null Catalyst type reports a typed error") { + val avroSchema = new Schema.Parser().parse( + """ + |{ "type": "record", + | "name": "record", + | "fields": [{ + | "name": "array", + | "type": { "type": "array", "items": ["null", "int"] } + | }] + |} + """.stripMargin) + // The Avro element is nullable, but the target Catalyst element type is not, so a null + // element must surface as a typed error rather than a generic RuntimeException. + val catalystType = new StructType() + .add("array", ArrayType(IntegerType, containsNull = false), nullable = false) + val data = new GenericRecordBuilder(avroSchema) + .set("array", util.Arrays.asList(1, null, 3)) + .build() + + checkError( + exception = intercept[SparkRuntimeException] { + deserializerFor(avroSchema, catalystType).deserialize(data) + }, + condition = "AVRO_CANNOT_READ_NULL_FIELD", + parameters = Map("name" -> "field 'array.element'")) + } + + test("SPARK-58218: null map value for non-null Catalyst type reports a typed error") { + val avroSchema = new Schema.Parser().parse( + """ + |{ "type": "record", + | "name": "record", + | "fields": [{ + | "name": "map", + | "type": { "type": "map", "values": ["null", "int"] } + | }] + |} + """.stripMargin) + val catalystType = new StructType() + .add("map", MapType(StringType, IntegerType, valueContainsNull = false), nullable = false) + val values = new util.HashMap[String, Integer]() + values.put("k", null) + val data = new GenericRecordBuilder(avroSchema) + .set("map", values) + .build() + + checkError( + exception = intercept[SparkRuntimeException] { + deserializerFor(avroSchema, catalystType).deserialize(data) + }, + condition = "AVRO_CANNOT_READ_NULL_FIELD", + parameters = Map("name" -> "field 'map.value'")) + } + test("avro array can be generic java collection") { val jsonFormatSchema = """ diff --git a/connector/credential-aws/pom.xml b/connector/credential-aws/pom.xml new file mode 100644 index 0000000000000..552a481605c2f --- /dev/null +++ b/connector/credential-aws/pom.xml @@ -0,0 +1,93 @@ +<?xml version="1.0" encoding="UTF-8"?> +<!-- + ~ 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. + --> + +<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> + <modelVersion>4.0.0</modelVersion> + <parent> + <groupId>org.apache.spark</groupId> + <artifactId>spark-parent_2.13</artifactId> + <version>5.0.0-SNAPSHOT</version> + <relativePath>../../pom.xml</relativePath> + </parent> + + <artifactId>spark-credential-aws_2.13</artifactId> + <properties> + <sbt.project.name>credential-aws</sbt.project.name> + </properties> + <packaging>jar</packaging> + <name>Spark AWS Credential Provider</name> + <description> + OIDC-to-AWS credential provider that exchanges identity tokens for + temporary AWS credentials via STS AssumeRoleWithWebIdentity. + </description> + <url>https://spark.apache.org/</url> + + <dependencies> + <dependency> + <groupId>org.apache.spark</groupId> + <artifactId>spark-core_${scala.binary.version}</artifactId> + <version>${project.version}</version> + <scope>provided</scope> + </dependency> + <dependency> + <groupId>org.apache.spark</groupId> + <artifactId>spark-tags_${scala.binary.version}</artifactId> + </dependency> + <!-- + Keep pinned to ${aws.java.sdk.v2.version}. When this module is built into the + assembly alongside the hadoop-cloud profile, this `sts` jar co-exists with + hadoop-cloud's `bundle` jar and they share the same software.amazon.awssdk.* + classes (e.g. StsClient). This is safe only because both resolve to the same + version, making the duplicated classes byte-identical. A divergent version here + would place duplicate classes at DIFFERENT versions on the classpath and risk + NoSuchMethodError / LinkageError at runtime. + --> + <dependency> + <groupId>software.amazon.awssdk</groupId> + <artifactId>sts</artifactId> + <version>${aws.java.sdk.v2.version}</version> + </dependency> + + <!-- Test dependencies --> + <!-- + This spark-tags test-dep is needed even though it isn't used in this module, otherwise testing-cmds that exclude + them will yield errors. + --> + <dependency> + <groupId>org.apache.spark</groupId> + <artifactId>spark-tags_${scala.binary.version}</artifactId> + <type>test-jar</type> + <scope>test</scope> + </dependency> + <dependency> + <groupId>org.junit.jupiter</groupId> + <artifactId>junit-jupiter</artifactId> + <scope>test</scope> + </dependency> + <dependency> + <groupId>org.mockito</groupId> + <artifactId>mockito-core</artifactId> + <scope>test</scope> + </dependency> + </dependencies> + + <build> + <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory> + <testOutputDirectory>target/scala-${scala.binary.version}/test-classes</testOutputDirectory> + </build> +</project> diff --git a/connector/credential-aws/src/main/java/org/apache/spark/security/aws/AwsStsCredentialProvider.java b/connector/credential-aws/src/main/java/org/apache/spark/security/aws/AwsStsCredentialProvider.java new file mode 100644 index 0000000000000..d694c897bcee1 --- /dev/null +++ b/connector/credential-aws/src/main/java/org/apache/spark/security/aws/AwsStsCredentialProvider.java @@ -0,0 +1,467 @@ +/* + * 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.security.aws; + +import java.net.URI; +import java.time.Duration; +import java.time.Instant; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Set; +import java.util.regex.Pattern; +import java.util.stream.Stream; + +import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; +import software.amazon.awssdk.core.exception.SdkException; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.StsClientBuilder; +import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityResponse; +import software.amazon.awssdk.services.sts.model.Credentials; + +import org.apache.spark.security.CredentialProvider; +import org.apache.spark.security.CredentialResolutionException; +import org.apache.spark.security.ServiceCredential; +import org.apache.spark.security.UserContext; + +/** + * A {@link CredentialProvider} that exchanges an OIDC identity token for temporary + * AWS credentials via the STS {@code AssumeRoleWithWebIdentity} API. + * <p> + * The returned {@link ServiceCredential} contains S3A-compatible Hadoop configuration + * properties ({@code fs.s3a.access.key}, {@code fs.s3a.secret.key}, + * {@code fs.s3a.session.token}) that can be propagated to executors for accessing + * S3-compatible storage. + * <p> + * <b>Configuration keys</b> (passed via {@code spark.security.oidc.*}): + * <ul> + * <li>{@code spark.security.oidc.aws.roleArn} (required) -- the ARN of the IAM role + * to assume</li> + * <li>{@code spark.security.oidc.aws.sessionName} (optional) -- the role session name; + * defaults to a value derived from the user's principal or "spark-oidc"</li> + * <li>{@code spark.security.oidc.aws.durationSeconds} (optional) -- credential duration + * in seconds (900-43200); if unset, STS uses the role's default maximum</li> + * <li>{@code spark.security.oidc.aws.region} (optional) -- the AWS region for the STS + * endpoint; defaults to us-east-1 when only {@code stsEndpoint} is set. When neither + * {@code region} nor {@code stsEndpoint} is configured, the STS client falls back to + * the AWS SDK default region resolution (AWS_REGION / AWS_DEFAULT_REGION environment + * variables, then the ~/.aws/config profile).</li> + * <li>{@code spark.security.oidc.aws.stsEndpoint} (optional) -- a custom STS endpoint URL + * for non-AWS environments (MinIO, Ceph, LocalStack, etc.)</li> + * </ul> + * <p> + * <b>Security note:</b> The OIDC raw token is never included in log messages or + * exception messages. It is passed directly to the STS API call and discarded. + * + * @since 4.4.0 + */ +public class AwsStsCredentialProvider implements CredentialProvider { + + // Configuration key constants + static final String CONF_ROLE_ARN = "spark.security.oidc.aws.roleArn"; + static final String CONF_SESSION_NAME = "spark.security.oidc.aws.sessionName"; + static final String CONF_DURATION_SECONDS = "spark.security.oidc.aws.durationSeconds"; + static final String CONF_REGION = "spark.security.oidc.aws.region"; + static final String CONF_STS_ENDPOINT = "spark.security.oidc.aws.stsEndpoint"; + + /** Minimum duration allowed by STS AssumeRoleWithWebIdentity (15 minutes). */ + static final int MIN_DURATION_SECONDS = 900; + /** Maximum duration allowed by STS AssumeRoleWithWebIdentity (12 hours). */ + static final int MAX_DURATION_SECONDS = 43200; + + private static final String DEFAULT_REGION = "us-east-1"; + private static final String DEFAULT_SESSION_NAME = "spark-oidc"; + + /** + * Precompiled pattern matching characters that are NOT valid in STS session names. + * Valid characters are: alphanumeric, underscore, plus, equals, comma, period, at, hyphen. + */ + private static final Pattern SESSION_NAME_INVALID_CHARS = + Pattern.compile("[^a-zA-Z0-9_+=,.@\\-]"); + + /** + * Precompiled pattern for validating a configured STS session name. + * Must match {@code [a-zA-Z0-9_+=,.@-]{2,64}} per AWS STS documentation. + * Uses an explicit ASCII character class rather than {@code \w} to avoid + * accepting non-ASCII characters (e.g. accented letters, CJK) that AWS STS + * would reject. + */ + private static final Pattern SESSION_NAME_VALID_PATTERN = + Pattern.compile("[a-zA-Z0-9_+=,.@\\-]{2,64}"); + + /** + * Immutable configuration holder that is safely published via the volatile + * {@link #config} field. All fields are set once during construction and are + * final, ensuring correct visibility across threads after init() completes. + */ + static final class ResolvedConfig { + final String roleArn; + final String roleSessionName; + final Integer durationSeconds; + final Region resolvedRegion; + final URI endpointOverride; + final StsClient stsClient; + + ResolvedConfig(String roleArn, String roleSessionName, Integer durationSeconds, + Region resolvedRegion, URI endpointOverride, StsClient stsClient) { + this.roleArn = roleArn; + this.roleSessionName = roleSessionName; + this.durationSeconds = durationSeconds; + this.resolvedRegion = resolvedRegion; + this.endpointOverride = endpointOverride; + this.stsClient = stsClient; + } + } + + /** Safely published via volatile write in init(); read in resolve()/suggestedTtl(). */ + private volatile ResolvedConfig config; + + /** Guards against double-close and allows resolve() to fail fast after close(). */ + private volatile boolean closed = false; + + /** + * Default no-arg constructor used by {@link java.util.ServiceLoader}. + */ + public AwsStsCredentialProvider() { + // ServiceLoader requires a public no-arg constructor + } + + /** + * Package-private constructor for testing with an injected STS client. + * <p> + * This constructor is visible for testing only; production code must use the + * no-arg constructor followed by {@link #init(Map)}. + * + * @param stsClient the STS client to use (must not be null) + * @param roleArn the IAM role ARN (must not be null or blank) + * @param roleSessionName the session name (may be null for default) + * @param durationSeconds the credential duration in seconds (may be null) + */ + AwsStsCredentialProvider(StsClient stsClient, String roleArn, String roleSessionName, + Integer durationSeconds) { + this.config = new ResolvedConfig(roleArn, roleSessionName, durationSeconds, + null, null, stsClient); + } + + @Override + public void init(Map<String, String> conf) { + if (this.config != null) { + throw new IllegalStateException("AwsStsCredentialProvider is already initialized"); + } + + String roleArn = conf.get(CONF_ROLE_ARN); + if (roleArn != null) { + roleArn = roleArn.trim(); + } + if (roleArn == null || roleArn.isBlank()) { + throw new IllegalArgumentException( + "Configuration key '" + CONF_ROLE_ARN + "' is required but was not set. " + + "Specify the ARN of the IAM role to assume via AssumeRoleWithWebIdentity."); + } + + String roleSessionName = conf.get(CONF_SESSION_NAME); + if (roleSessionName != null) { + roleSessionName = roleSessionName.trim(); + } + if (roleSessionName != null && roleSessionName.isBlank()) { + roleSessionName = null; + } + if (roleSessionName != null && !SESSION_NAME_VALID_PATTERN.matcher(roleSessionName).matches()) { + throw new IllegalArgumentException( + "Configuration key '" + CONF_SESSION_NAME + + "' must match [a-zA-Z0-9_+=,.@-]{2,64}, got: " + + roleSessionName); + } + String regionStr = conf.get(CONF_REGION); + if (regionStr != null) { + regionStr = regionStr.trim(); + } + String stsEndpoint = conf.get(CONF_STS_ENDPOINT); + if (stsEndpoint != null) { + stsEndpoint = stsEndpoint.trim(); + } + + Integer durationSeconds = null; + String durationStr = conf.get(CONF_DURATION_SECONDS); + if (durationStr != null && !durationStr.isBlank()) { + try { + durationSeconds = Integer.parseInt(durationStr.trim()); + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + "Configuration key '" + CONF_DURATION_SECONDS + + "' must be a valid integer, got: " + durationStr, e); + } + if (durationSeconds < MIN_DURATION_SECONDS || durationSeconds > MAX_DURATION_SECONDS) { + throw new IllegalArgumentException( + "Configuration key '" + CONF_DURATION_SECONDS + "' must be between " + + MIN_DURATION_SECONDS + " and " + MAX_DURATION_SECONDS + + " seconds, got: " + durationSeconds); + } + } + + // Resolve the region and endpoint before building the client + Region resolvedRegion = resolveRegion(regionStr, stsEndpoint); + URI endpointOverride = resolveEndpoint(stsEndpoint); + + StsClient stsClient = buildStsClient(resolvedRegion, endpointOverride); + + // Single volatile write publishes all configuration atomically + this.config = new ResolvedConfig(roleArn, roleSessionName, durationSeconds, + resolvedRegion, endpointOverride, stsClient); + } + + @Override + public void close() { + if (closed) { + return; + } + closed = true; + ResolvedConfig cfg = this.config; + if (cfg != null && cfg.stsClient != null) { + cfg.stsClient.close(); + } + } + + /** + * Resolves the AWS region based on explicit configuration and endpoint presence. + * When a custom endpoint is provided without an explicit region, defaults to us-east-1. + * When neither region nor stsEndpoint is configured, returns null so the STS client + * falls back to the AWS SDK default region resolution (AWS_REGION / AWS_DEFAULT_REGION + * environment variables, then the ~/.aws/config profile). + */ + private static Region resolveRegion(String regionStr, String stsEndpoint) { + if (regionStr != null && !regionStr.isBlank()) { + return Region.of(regionStr); + } else if (stsEndpoint != null && !stsEndpoint.isBlank()) { + // When a custom endpoint is set but no explicit region, use a default region. + // The region is required by the SDK but not meaningful for non-AWS endpoints. + return Region.of(DEFAULT_REGION); + } + return null; + } + + /** + * Resolves the endpoint override URI from configuration. + */ + private static URI resolveEndpoint(String stsEndpoint) { + if (stsEndpoint != null && !stsEndpoint.isBlank()) { + try { + return URI.create(stsEndpoint); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + "Configuration key '" + CONF_STS_ENDPOINT + "' contains a malformed URI: " + + stsEndpoint, e); + } + } + return null; + } + + /** + * Builds the STS client with the resolved region and endpoint. + */ + private static StsClient buildStsClient(Region resolvedRegion, URI endpointOverride) { + StsClientBuilder builder = StsClient.builder() + // AssumeRoleWithWebIdentity does not require AWS credentials; + // the OIDC token itself serves as the authentication mechanism. + .credentialsProvider(AnonymousCredentialsProvider.create()); + + if (resolvedRegion != null) { + builder.region(resolvedRegion); + } + + if (endpointOverride != null) { + builder.endpointOverride(endpointOverride); + } + + return builder.build(); + } + + /** + * Returns the resolved configuration for testing purposes. + * Package-private visibility allows test assertions on resolved region/endpoint. + */ + ResolvedConfig resolvedConfig() { + return config; + } + + @Override + public Set<String> supportedSchemes() { + return Set.of("s3a"); + } + + @Override + public ServiceCredential resolve(UserContext user, URI target) + throws CredentialResolutionException { + if (closed) { + throw new CredentialResolutionException("provider is closed"); + } + if (user == null) { + throw new CredentialResolutionException( + "UserContext must not be null when resolving AWS credentials"); + } + if (target == null) { + throw new CredentialResolutionException( + "Target URI must not be null when resolving AWS credentials"); + } + String rawToken = user.getRawToken(); + if (rawToken == null || rawToken.isBlank()) { + throw new CredentialResolutionException( + "UserContext raw token must not be null or blank; cannot perform " + + "AssumeRoleWithWebIdentity without an identity token"); + } + + ResolvedConfig cfg = this.config; + if (cfg == null) { + throw new CredentialResolutionException("resolve() called before init()"); + } + String sessionName = cfg.roleSessionName; + if (sessionName == null || sessionName.isBlank()) { + // Derive from principal, sanitizing for STS session name constraints + // (alphanumeric, =,.@- only, max 64 chars) + String principal = user.getPrincipal(); + if (principal != null && !principal.isBlank()) { + sessionName = sanitizeSessionName(principal); + } else { + sessionName = DEFAULT_SESSION_NAME; + } + } + + try { + AssumeRoleWithWebIdentityRequest.Builder reqBuilder = + AssumeRoleWithWebIdentityRequest.builder() + .roleArn(cfg.roleArn) + .roleSessionName(sessionName) + .webIdentityToken(rawToken); + + if (cfg.durationSeconds != null) { + reqBuilder.durationSeconds(cfg.durationSeconds); + } + + AssumeRoleWithWebIdentityResponse response = + cfg.stsClient.assumeRoleWithWebIdentity(reqBuilder.build()); + + Credentials creds = response.credentials(); + if (creds == null || creds.accessKeyId() == null + || creds.secretAccessKey() == null || creds.sessionToken() == null) { + throw new CredentialResolutionException( + "STS returned incomplete credentials for role '" + cfg.roleArn + "'"); + } + + Map<String, String> properties = Map.of( + "fs.s3a.access.key", creds.accessKeyId(), + "fs.s3a.secret.key", creds.secretAccessKey(), + "fs.s3a.session.token", creds.sessionToken() + ); + + Instant expiration = creds.expiration(); + return new ServiceCredential(properties, expiration); + } catch (SdkException e) { + // SECURITY: Never include the token in exception messages. + // Defensively strip any occurrence of the raw token from the STS error message + // in case the service accidentally echoed it. + String errorMsg = e.getMessage(); + String redactedMsg = errorMsg; + if (redactedMsg != null && rawToken != null && redactedMsg.contains(rawToken)) { + redactedMsg = redactedMsg.replace(rawToken, "[REDACTED]"); + } + // Walk the entire cause chain to check if the token leaked into any layer. + boolean tokenInAnyMessage = causeChainContainsToken(e, rawToken); + Throwable cause; + if (tokenInAnyMessage || (errorMsg != null && errorMsg.contains(rawToken))) { + // Drop the original cause chain entirely; replace with a sanitized wrapper. + cause = SdkException.builder() + .message(redactedMsg) + .build(); + } else { + cause = e; + } + throw new CredentialResolutionException( + "Failed to assume role '" + cfg.roleArn + "' via AssumeRoleWithWebIdentity: " + + redactedMsg, cause); + } catch (IllegalStateException e) { + // The SDK throws IllegalStateException when the client has been closed. + throw new CredentialResolutionException( + "Failed to assume role '" + cfg.roleArn + "' via AssumeRoleWithWebIdentity: " + + "the STS client has been closed", e); + } + } + + @Override + public Duration suggestedTtl() { + ResolvedConfig cfg = this.config; + if (cfg != null && cfg.durationSeconds != null) { + return Duration.ofSeconds(cfg.durationSeconds); + } + return Duration.ofMinutes(15); + } + + @Override + public Map<String, String> additionalSparkProperties() { + return Map.of( + "spark.hadoop.fs.s3a.aws.credentials.provider", + "org.apache.spark.security.aws.SparkOidcAwsCredentialsProvider"); + } + + /** + * Walks the cause chain of the given throwable and checks whether the raw token + * appears in any layer's message. Used to decide whether the original exception + * chain is safe to preserve as a cause. + */ + private static boolean causeChainContainsToken(Throwable root, String rawToken) { + if (rawToken == null) { + return false; + } + return causeChainStream(root) + .anyMatch(t -> t.getMessage() != null && t.getMessage().contains(rawToken)); + } + + /** + * Returns a sequential stream over the cause chain starting from {@code root}. + * Uses identity-based cycle detection to guard against circular cause chains. + */ + private static Stream<Throwable> causeChainStream(Throwable root) { + Stream.Builder<Throwable> builder = Stream.builder(); + Set<Throwable> visited = Collections.newSetFromMap(new IdentityHashMap<>()); + Throwable current = root; + while (current != null && visited.add(current)) { + builder.accept(current); + current = current.getCause(); + } + return builder.build(); + } + + /** + * Sanitizes a principal string to be valid as an STS role session name. + * STS session names must match [a-zA-Z0-9_+=,.@-]{2,64}. Both the validation + * pattern and the sanitization replacement use an explicit ASCII character class + * rather than {@code \w} to avoid locale-dependent behavior and to reject + * non-ASCII characters that AWS STS would not accept. + */ + static String sanitizeSessionName(String principal) { + String sanitized = SESSION_NAME_INVALID_CHARS.matcher(principal).replaceAll("-"); + if (sanitized.length() > 64) { + sanitized = sanitized.substring(0, 64); + } + if (sanitized.length() < 2) { + return DEFAULT_SESSION_NAME; + } + return sanitized; + } +} diff --git a/connector/credential-aws/src/main/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProvider.java b/connector/credential-aws/src/main/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProvider.java new file mode 100644 index 0000000000000..0248dbb34cbdb --- /dev/null +++ b/connector/credential-aws/src/main/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProvider.java @@ -0,0 +1,126 @@ +/* + * 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.security.aws; + +import java.util.Map; + +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; + +import org.apache.spark.SparkEnv; +import org.apache.spark.VersionedCredentials; +import org.apache.spark.deploy.security.UserCredentialManager; +import org.apache.spark.security.ServiceCredential; +import org.apache.spark.security.UserCredentials; + +/** + * A dynamic AWS credentials provider for executor-side S3A access that reads from + * Spark's credential store. + * <p> + * This provider uses version-based caching: credentials are deserialized only when + * the credential store version changes (i.e., after a driver-initiated refresh). + * Since the store version is monotonically increasing and only changes on renewal + * (minutes-scale), the cache hit rate is {@literal >}99.99% for I/O-heavy workloads. + * <p> + * This implementation is thread-safe. Multiple threads may call + * {@code resolveCredentials()} concurrently without external synchronization. + * Each invocation reads the credential version atomically and returns either + * the cached result or a freshly deserialized one. + * <p> + * Configure via: + * {@code fs.s3a.aws.credentials.provider= + * org.apache.spark.security.aws.SparkOidcAwsCredentialsProvider} + * + * @since 4.4.0 + */ +public class SparkOidcAwsCredentialsProvider implements AwsCredentialsProvider { + + /** S3A credential property keys (same as produced by AwsStsCredentialProvider). */ + private static final String ACCESS_KEY = "fs.s3a.access.key"; + private static final String SECRET_KEY = "fs.s3a.secret.key"; + private static final String SESSION_TOKEN = "fs.s3a.session.token"; + + /** The S3A scheme used to look up credentials in the UserCredentials bundle. */ + private static final String S3A_SCHEME = "s3a"; + + /** Version-keyed cache to avoid repeated deserialization on every S3A API call. */ + private volatile CachedResult cached; + + private record CachedResult(long version, AwsSessionCredentials credentials) {} + + @Override + public AwsCredentials resolveCredentials() { + SparkEnv env = SparkEnv.get(); + if (env == null) { + throw new IllegalStateException( + "SparkEnv is not available. SparkOidcAwsCredentialsProvider can only be used " + + "within an active Spark executor."); + } + + VersionedCredentials versioned = env.userCredentials().get(); + if (versioned == null) { + throw new IllegalStateException( + "No credentials available in the executor credential store. " + + "Ensure spark.security.oidc.enabled=true and the driver has acquired " + + "credentials before executor tasks run."); + } + + // Fast path: return cached credentials if version hasn't changed. + CachedResult current = cached; + if (current != null && current.version() == versioned.version()) { + return current.credentials(); + } + + UserCredentials credentials; + try { + credentials = UserCredentialManager.deserializeUserCredentials(versioned.bytes()); + } catch (Exception e) { + throw new IllegalStateException( + "Failed to deserialize credentials from executor store (version=" + + versioned.version() + "). The credential bytes may be corrupted or " + + "incompatible with this Spark version.", e); + } + + ServiceCredential s3aCred = credentials.forScheme(S3A_SCHEME).orElse(null); + if (s3aCred == null) { + throw new IllegalStateException( + "No credential found for scheme '" + S3A_SCHEME + "' in the executor " + + "credential store. Ensure an S3A-compatible CredentialProvider " + + "(e.g., AwsStsCredentialProvider) is configured on the driver."); + } + + Map<String, String> props = s3aCred.getProperties(); + String accessKey = props.get(ACCESS_KEY); + String secretKey = props.get(SECRET_KEY); + String sessionToken = props.get(SESSION_TOKEN); + + if (accessKey == null || accessKey.isEmpty() + || secretKey == null || secretKey.isEmpty() + || sessionToken == null || sessionToken.isEmpty()) { + throw new IllegalStateException( + "ServiceCredential for scheme '" + S3A_SCHEME + "' is missing required " + + "properties. Expected non-empty values for: " + ACCESS_KEY + ", " + + SECRET_KEY + ", " + SESSION_TOKEN); + } + + AwsSessionCredentials result = AwsSessionCredentials.create(accessKey, secretKey, sessionToken); + cached = new CachedResult(versioned.version(), result); + return result; + } +} diff --git a/connector/credential-aws/src/main/resources/META-INF/services/org.apache.spark.security.CredentialProvider b/connector/credential-aws/src/main/resources/META-INF/services/org.apache.spark.security.CredentialProvider new file mode 100644 index 0000000000000..84b31f53ec263 --- /dev/null +++ b/connector/credential-aws/src/main/resources/META-INF/services/org.apache.spark.security.CredentialProvider @@ -0,0 +1,18 @@ +# +# 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.security.aws.AwsStsCredentialProvider diff --git a/connector/credential-aws/src/test/java/org/apache/spark/security/aws/AwsStsCredentialProviderSuite.java b/connector/credential-aws/src/test/java/org/apache/spark/security/aws/AwsStsCredentialProviderSuite.java new file mode 100644 index 0000000000000..0dae6227a002c --- /dev/null +++ b/connector/credential-aws/src/test/java/org/apache/spark/security/aws/AwsStsCredentialProviderSuite.java @@ -0,0 +1,953 @@ +/* + * 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.security.aws; + +import java.net.URI; +import java.time.Duration; +import java.time.Instant; +import java.util.HashMap; +import java.util.Map; +import java.util.ServiceLoader; +import java.util.Set; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.sts.StsClient; +import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest; +import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityResponse; +import software.amazon.awssdk.services.sts.model.Credentials; +import software.amazon.awssdk.services.sts.model.StsException; + +import org.apache.spark.security.CredentialProvider; +import org.apache.spark.security.CredentialResolutionException; +import org.apache.spark.security.ServiceCredential; +import org.apache.spark.security.UserContext; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link AwsStsCredentialProvider}. + */ +public class AwsStsCredentialProviderSuite { + + private static String previousAwsRegion; + + /** + * Set the aws.region system property so that the AWS SDK's + * DefaultAwsRegionProviderChain can resolve a region on any environment + * (including CI runners with no AWS configuration). This does NOT affect + * AwsStsCredentialProvider.resolveRegion() which reads only from the conf Map. + */ + @BeforeAll + static void setUpClass() { + previousAwsRegion = System.getProperty("aws.region"); + System.setProperty("aws.region", "us-east-1"); + } + + @AfterAll + static void tearDownClass() { + if (previousAwsRegion == null) { + System.clearProperty("aws.region"); + } else { + System.setProperty("aws.region", previousAwsRegion); + } + } + + /** Suite-level field closed by tearDown to avoid resource leaks from real StsClients. */ + private AwsStsCredentialProvider provider; + + @AfterEach + void tearDown() { + if (provider != null) { + provider.close(); + provider = null; + } + } + + private static final String TEST_ROLE_ARN = "arn:aws:iam::123456789012:role/test-role"; + private static final String TEST_TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.test-payload"; + private static final String TEST_PRINCIPAL = "user@example.com"; + private static final String TEST_ISSUER = "https://idp.example.com"; + private static final URI TEST_TARGET = URI.create("s3a://my-bucket/data/file.parquet"); + + // ========== ServiceLoader Discovery ========== + + @Test + public void testServiceLoaderDiscovery() { + ServiceLoader<CredentialProvider> loader = ServiceLoader.load(CredentialProvider.class); + boolean found = false; + for (CredentialProvider provider : loader) { + if (provider instanceof AwsStsCredentialProvider) { + found = true; + break; + } + } + assertTrue(found, "AwsStsCredentialProvider should be discoverable via ServiceLoader"); + } + + // ========== init() ========== + + @Test + public void testMissingRoleArnThrowsIllegalArgumentException() { + Map<String, String> conf = new HashMap<>(); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, + () -> provider.init(conf)); + assertTrue(ex.getMessage().contains(AwsStsCredentialProvider.CONF_ROLE_ARN)); + } + + @Test + public void testBlankRoleArnThrowsIllegalArgumentException() { + Map<String, String> conf = new HashMap<>(); + conf.put(AwsStsCredentialProvider.CONF_ROLE_ARN, " "); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, + () -> provider.init(conf)); + assertTrue(ex.getMessage().contains(AwsStsCredentialProvider.CONF_ROLE_ARN)); + } + + @Test + public void testInitWithInvalidDurationSeconds() { + assertInitThrowsForConfig(AwsStsCredentialProvider.CONF_DURATION_SECONDS, "not-a-number"); + } + + @Test + public void testInitWithZeroDurationSecondsThrowsIllegalArgumentException() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_DURATION_SECONDS, "0"); + assertTrue(ex.getMessage().contains("900")); + assertTrue(ex.getMessage().contains("43200")); + } + + @Test + public void testInitWithNegativeDurationSecondsThrowsIllegalArgumentException() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_DURATION_SECONDS, "-100"); + assertTrue(ex.getMessage().contains("900")); + } + + @Test + public void testInitWithDurationBelowMinimumThrowsIllegalArgumentException() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_DURATION_SECONDS, "899"); + assertTrue(ex.getMessage().contains("900")); + assertTrue(ex.getMessage().contains("43200")); + } + + @Test + public void testInitWithDurationAboveMaximumThrowsIllegalArgumentException() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_DURATION_SECONDS, "43201"); + assertTrue(ex.getMessage().contains("900")); + assertTrue(ex.getMessage().contains("43200")); + } + + @Test + public void testInitWithMinimumValidDurationSucceeds() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_DURATION_SECONDS, "900"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + assertNotNull(provider.resolvedConfig()); + assertEquals(900, provider.resolvedConfig().durationSeconds); + } + + @Test + public void testInitWithMaximumValidDurationSucceeds() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_DURATION_SECONDS, "43200"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + assertNotNull(provider.resolvedConfig()); + assertEquals(43200, provider.resolvedConfig().durationSeconds); + } + + @Test + public void testReInitializationThrowsIllegalStateException() { + Map<String, String> conf = confWithRoleArn(); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + IllegalStateException ex = assertThrows(IllegalStateException.class, + () -> provider.init(conf)); + assertTrue(ex.getMessage().contains("already initialized")); + } + + @Test + public void testInitWithEndpointAndRegionResolvesCorrectly() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_STS_ENDPOINT, "http://localhost:9000"); + conf.put(AwsStsCredentialProvider.CONF_REGION, "us-west-2"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + AwsStsCredentialProvider.ResolvedConfig cfg = provider.resolvedConfig(); + assertNotNull(cfg); + assertEquals(Region.of("us-west-2"), cfg.resolvedRegion); + assertEquals(URI.create("http://localhost:9000"), cfg.endpointOverride); + assertNotNull(cfg.stsClient); + } + + @Test + public void testInitWithEndpointNoRegionUsesDefault() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_STS_ENDPOINT, "http://localhost:9000"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + AwsStsCredentialProvider.ResolvedConfig cfg = provider.resolvedConfig(); + assertNotNull(cfg); + assertEquals(Region.of("us-east-1"), cfg.resolvedRegion); + assertEquals(URI.create("http://localhost:9000"), cfg.endpointOverride); + } + + @Test + public void testInitWithNeitherEndpointNorRegion() { + Map<String, String> conf = confWithRoleArn(); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + AwsStsCredentialProvider.ResolvedConfig cfg = provider.resolvedConfig(); + assertNotNull(cfg); + assertNull(cfg.resolvedRegion); + assertNull(cfg.endpointOverride); + } + + @Test + public void testInitWithMalformedEndpointThrowsIllegalArgumentException() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_STS_ENDPOINT, "not a valid uri^[]"); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, + () -> provider.init(conf)); + assertTrue(ex.getMessage().contains(AwsStsCredentialProvider.CONF_STS_ENDPOINT), + "Error should mention the config key"); + assertTrue(ex.getMessage().contains("not a valid uri^[]"), + "Error should mention the bad value"); + assertNotNull(ex.getCause(), "Original IllegalArgumentException should be preserved"); + } + + @Test + public void testInitTrimsConfigValues() { + Map<String, String> conf = new HashMap<>(); + conf.put(AwsStsCredentialProvider.CONF_ROLE_ARN, " " + TEST_ROLE_ARN + " "); + conf.put(AwsStsCredentialProvider.CONF_STS_ENDPOINT, " http://localhost:9000 "); + conf.put(AwsStsCredentialProvider.CONF_REGION, " us-west-2 "); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + AwsStsCredentialProvider.ResolvedConfig cfg = provider.resolvedConfig(); + assertNotNull(cfg); + assertEquals(TEST_ROLE_ARN, cfg.roleArn); + assertEquals(Region.of("us-west-2"), cfg.resolvedRegion); + assertEquals(URI.create("http://localhost:9000"), cfg.endpointOverride); + } + + @Test + public void testInitTrimsRoleSessionName() throws CredentialResolutionException { + Instant expiration = Instant.now().plusSeconds(3600); + StsClient mockSts = createMockStsClient("AK", "SK", "ST", expiration); + + // Use init() with a padded sessionName to exercise the trim path + Map<String, String> conf = new HashMap<>(); + conf.put(AwsStsCredentialProvider.CONF_ROLE_ARN, TEST_ROLE_ARN); + conf.put(AwsStsCredentialProvider.CONF_SESSION_NAME, " my-session "); + conf.put(AwsStsCredentialProvider.CONF_STS_ENDPOINT, "http://localhost:9000"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + // Verify the stored config has the trimmed value + assertEquals("my-session", provider.resolvedConfig().roleSessionName); + + // Now resolve() with a separate provider that uses the test constructor + // so we can capture the STS request via mock + String sessionName = provider.resolvedConfig().roleSessionName; + provider.close(); + provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, sessionName, null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + provider.resolve(user, TEST_TARGET); + + ArgumentCaptor<AssumeRoleWithWebIdentityRequest> captor = + ArgumentCaptor.forClass(AssumeRoleWithWebIdentityRequest.class); + verify(mockSts).assumeRoleWithWebIdentity(captor.capture()); + + // The session name in the request must be trimmed, not " my-session " + assertEquals("my-session", captor.getValue().roleSessionName()); + } + + // ========== resolve() ========== + + @Test + public void testSuccessfulResolve() throws CredentialResolutionException { + Instant expiration = Instant.now().plusSeconds(3600); + StsClient mockSts = createMockStsClient("AKIAIOSFODNN7EXAMPLE", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", "FwoGZXIvY...token", expiration); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", 3600); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + ServiceCredential credential = provider.resolve(user, TEST_TARGET); + + assertNotNull(credential); + assertEquals("AKIAIOSFODNN7EXAMPLE", credential.getProperties().get("fs.s3a.access.key")); + assertEquals("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + credential.getProperties().get("fs.s3a.secret.key")); + assertEquals("FwoGZXIvY...token", credential.getProperties().get("fs.s3a.session.token")); + assertEquals(expiration, credential.getExpiresAt()); + assertEquals(3, credential.getProperties().size()); + + // Verify the STS request was built correctly + ArgumentCaptor<AssumeRoleWithWebIdentityRequest> captor = + ArgumentCaptor.forClass(AssumeRoleWithWebIdentityRequest.class); + verify(mockSts).assumeRoleWithWebIdentity(captor.capture()); + + AssumeRoleWithWebIdentityRequest request = captor.getValue(); + assertEquals(TEST_ROLE_ARN, request.roleArn()); + assertEquals(TEST_TOKEN, request.webIdentityToken()); + assertEquals("test-session", request.roleSessionName()); + assertEquals(3600, request.durationSeconds()); + } + + @Test + public void testStsFailureWrappedInCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + StsException stsException = (StsException) StsException.builder() + .message("Access denied for role") + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenThrow(stsException); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + + // Verify the exception wraps the STS exception + assertEquals(stsException, ex.getCause()); + assertTrue(ex.getMessage().contains(TEST_ROLE_ARN)); + assertTrue(ex.getMessage().contains("AssumeRoleWithWebIdentity")); + + // SECURITY: Verify that the raw token is NOT in any exception message + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must never appear in exception messages"); + assertFalse(ex.getCause().getMessage() != null + && ex.getCause().getMessage().contains(TEST_TOKEN), + "Raw token must never appear in cause exception messages"); + } + + @Test + public void testTokenRedactedFromStsErrorMessage() { + StsClient mockSts = mock(StsClient.class); + // Simulate STS echoing the token back in an error message + StsException stsException = (StsException) StsException.builder() + .message("Invalid identity token: " + TEST_TOKEN) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenThrow(stsException); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + + // SECURITY: The wrapped message must NOT contain the raw token + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must be redacted from exception messages even if STS echoed it"); + assertTrue(ex.getMessage().contains("[REDACTED]"), + "Token should be replaced with [REDACTED]"); + + // SECURITY: The cause exception message must also NOT contain the raw token + assertNotNull(ex.getCause(), "Cause should be a sanitized wrapper exception"); + assertFalse(ex.getCause().getMessage().contains(TEST_TOKEN), + "Raw token must be redacted from cause exception message"); + } + + @Test + public void testCustomSessionNameAndDurationPropagateToRequest() + throws CredentialResolutionException { + Instant expiration = Instant.now().plusSeconds(900); + StsClient mockSts = createMockStsClient("AK", "SK", "ST", expiration); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "custom-session-name", 900); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + provider.resolve(user, TEST_TARGET); + + ArgumentCaptor<AssumeRoleWithWebIdentityRequest> captor = + ArgumentCaptor.forClass(AssumeRoleWithWebIdentityRequest.class); + verify(mockSts).assumeRoleWithWebIdentity(captor.capture()); + + AssumeRoleWithWebIdentityRequest request = captor.getValue(); + assertEquals("custom-session-name", request.roleSessionName()); + assertEquals(900, request.durationSeconds()); + } + + @Test + public void testNoDurationSecondsInRequestWhenNotConfigured() + throws CredentialResolutionException { + Instant expiration = Instant.now().plusSeconds(3600); + StsClient mockSts = createMockStsClient("AK", "SK", "ST", expiration); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + provider.resolve(user, TEST_TARGET); + + ArgumentCaptor<AssumeRoleWithWebIdentityRequest> captor = + ArgumentCaptor.forClass(AssumeRoleWithWebIdentityRequest.class); + verify(mockSts).assumeRoleWithWebIdentity(captor.capture()); + + // durationSeconds should be null when not configured + assertNull(captor.getValue().durationSeconds()); + } + + @Test + public void testNullUserContextThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + assertThrows(CredentialResolutionException.class, + () -> provider.resolve(null, TEST_TARGET)); + } + + @Test + public void testNullTargetThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, null)); + assertTrue(ex.getMessage().contains("Target URI must not be null")); + } + + @Test + public void testNullRawTokenThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = mock(UserContext.class); + when(user.getRawToken()).thenReturn(null); + + assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + } + + @Test + public void testBlankRawTokenThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, " ", + Instant.now(), Instant.now().plusSeconds(300)); + + assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + } + + @Test + public void testResolveBeforeInitThrowsCredentialResolutionException() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + + assertTrue(ex.getMessage().contains("resolve() called before init()")); + } + + @Test + public void testNullCredentialsResponseThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + AssumeRoleWithWebIdentityResponse response = AssumeRoleWithWebIdentityResponse.builder() + .credentials((Credentials) null) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenReturn(response); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + + assertTrue(ex.getMessage().contains("incomplete credentials")); + assertTrue(ex.getMessage().contains(TEST_ROLE_ARN)); + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must never appear in exception messages"); + } + + @Test + public void testMissingSessionTokenThrowsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + Credentials creds = Credentials.builder() + .accessKeyId("AKID") + .secretAccessKey("SECRET") + .sessionToken(null) + .expiration(Instant.now().plusSeconds(3600)) + .build(); + AssumeRoleWithWebIdentityResponse response = AssumeRoleWithWebIdentityResponse.builder() + .credentials(creds) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenReturn(response); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> provider.resolve(user, TEST_TARGET)); + + assertTrue(ex.getMessage().contains("incomplete credentials")); + assertTrue(ex.getMessage().contains(TEST_ROLE_ARN)); + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must never appear in exception messages"); + } + + // ========== supportedSchemes() ========== + + @Test + public void testSupportedSchemesContainsS3a() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + Set<String> schemes = provider.supportedSchemes(); + assertTrue(schemes.contains("s3a")); + assertEquals(1, schemes.size()); + } + + // ========== suggestedTtl() ========== + + @Test + public void testSuggestedTtlWithDurationSeconds() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", 1800); + + assertEquals(Duration.ofSeconds(1800), provider.suggestedTtl()); + } + + @Test + public void testSuggestedTtlDefaultsTo15Minutes() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + assertEquals(Duration.ofMinutes(15), provider.suggestedTtl()); + } + + // ========== Session name derivation ========== + + @Test + public void testDefaultSessionNameDerivedFromPrincipal() + throws CredentialResolutionException { + Instant expiration = Instant.now().plusSeconds(3600); + StsClient mockSts = createMockStsClient("AK", "SK", "ST", expiration); + + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, null, null); + + UserContext user = new UserContext("alice@corp.example.com", TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + provider.resolve(user, TEST_TARGET); + + ArgumentCaptor<AssumeRoleWithWebIdentityRequest> captor = + ArgumentCaptor.forClass(AssumeRoleWithWebIdentityRequest.class); + verify(mockSts).assumeRoleWithWebIdentity(captor.capture()); + + String sessionName = captor.getValue().roleSessionName(); + assertNotNull(sessionName); + assertFalse(sessionName.isBlank()); + assertTrue(sessionName.contains("alice")); + // @ is valid in STS session names and should be preserved + assertTrue(sessionName.contains("@")); + } + + @Test + public void testSanitizeSessionName() { + // {input, expected} + String[][] cases = { + // truncates to 64 + {"a".repeat(70), "a".repeat(64)}, + // replaces invalid chars + {"user name!#$%", "user-name----"}, + // preserves valid chars + {"user_+=,.@-test", "user_+=,.@-test"}, + // falls back for short result + {"x", "spark-oidc"}, + // falls back when all invalid + {"!", "spark-oidc"}, + // replaces backslash + {"DOMAIN\\user", "DOMAIN-user"}, + }; + + for (String[] c : cases) { + assertEquals(c[1], AwsStsCredentialProvider.sanitizeSessionName(c[0]), + "input: " + c[0]); + } + } + + // ========== close() ========== + + @Test + public void testCloseShutsStsClient() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider provider = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "session", null); + + provider.close(); + + verify(mockSts).close(); + } + + @Test + public void testCloseBeforeInitDoesNotThrow() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + // Should not throw even when config is null + provider.close(); + } + + // ========== Deep cause-chain token redaction (Item 1) ========== + + @Test + public void testTokenInCauseChainIsRedacted() { + StsClient mockSts = mock(StsClient.class); + // Build a cause chain where the INNER cause contains the raw token, + // but the top-level message does NOT. + RuntimeException innerCause = new RuntimeException( + "Token validation failed: " + TEST_TOKEN); + StsException stsException = (StsException) StsException.builder() + .message("Access denied for role") + .cause(innerCause) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenThrow(stsException); + + AwsStsCredentialProvider p = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> p.resolve(user, TEST_TARGET)); + + // Top-level message must not contain the raw token + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must not appear in top-level exception message"); + + // Walk the entire cause chain and assert the token is absent everywhere + Throwable current = ex.getCause(); + while (current != null) { + if (current.getMessage() != null) { + assertFalse(current.getMessage().contains(TEST_TOKEN), + "Raw token must not appear in any cause message, but found in: " + + current.getClass().getSimpleName()); + } + current = current.getCause(); + } + } + + // ========== Session name validation (Item 3) ========== + + @Test + public void testInitWithValidCustomSessionName() { + Map<String, String> conf = confWithRoleArn(); + conf.put(AwsStsCredentialProvider.CONF_SESSION_NAME, "valid_session+=,.@-name"); + + provider = new AwsStsCredentialProvider(); + provider.init(conf); + + assertEquals("valid_session+=,.@-name", provider.resolvedConfig().roleSessionName); + } + + @Test + public void testInitWithInvalidSessionNameContainingSpace() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_SESSION_NAME, "bad session"); + assertTrue(ex.getMessage().contains("bad session"), + "Error must echo the bad value"); + } + + @Test + public void testInitWithInvalidSessionNameContainingQuestionMark() { + assertInitThrowsForConfig(AwsStsCredentialProvider.CONF_SESSION_NAME, "bad?name"); + } + + @Test + public void testInitWithSessionNameTooShort() { + IllegalArgumentException ex = assertInitThrowsForConfig( + AwsStsCredentialProvider.CONF_SESSION_NAME, "x"); + assertTrue(ex.getMessage().contains("x")); + } + + @Test + public void testInitWithSessionNameTooLong() { + assertInitThrowsForConfig(AwsStsCredentialProvider.CONF_SESSION_NAME, "a".repeat(65)); + } + + // ========== Session name: non-ASCII rejection (regression) ========== + + @Test + public void testInitRejectsSessionNameWithAccentedChar() { + // U+00E9 (accented e) is valid under \w but NOT valid in STS session names. + // This test would PASS (incorrectly) under the buggy \w pattern and must + // FAIL (correctly) under the explicit ASCII pattern. + assertInitThrowsForConfig(AwsStsCredentialProvider.CONF_SESSION_NAME, "café"); + } + + @Test + public void testInitRejectsSessionNameWithCjkChar() { + // U+4E16 (CJK ideograph) is valid under \w but NOT valid in STS session names. + // Regression test for the ASCII-only fix. + assertInitThrowsForConfig(AwsStsCredentialProvider.CONF_SESSION_NAME, "session世"); + } + + // ========== close() idempotency ========== + + @Test + public void testDoubleCloseIsIdempotent() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider p = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + + // First close should succeed normally + p.close(); + // Second close must be a no-op (no exception, no double-close on the client) + p.close(); + + // Verify the underlying client was closed exactly once + verify(mockSts).close(); + } + + @Test + public void testResolveAfterCloseThrowsWithClosedMessage() { + StsClient mockSts = mock(StsClient.class); + AwsStsCredentialProvider p = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + p.close(); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> p.resolve(user, TEST_TARGET)); + assertTrue(ex.getMessage().contains("closed"), + "Message should indicate the provider is closed"); + } + + // ========== causeChainStream cycle protection ========== + + @Test + public void testCauseChainCycleDoesNotLoop() throws Exception { + // Construct a circular cause chain via reflection: nodeA -> nodeB -> nodeA. + // Java's Throwable.initCause() prevents self-causation and re-initialization, + // so we use Field.set to create the cycle. + RuntimeException nodeA = new RuntimeException("nodeA"); + RuntimeException nodeB = new RuntimeException("nodeB", nodeA); + // nodeA.cause is currently the sentinel (this), which means "not yet set" in + // OpenJDK's implementation. Force it to nodeB to create the cycle. + java.lang.reflect.Field causeField = Throwable.class.getDeclaredField("cause"); + causeField.setAccessible(true); + causeField.set(nodeA, nodeB); + + // Now we have nodeA -> nodeB -> nodeA (a cycle). + // Simulate what resolve() does: call causeChainContainsToken and verify termination. + StsClient mockSts = mock(StsClient.class); + StsException stsException = (StsException) StsException.builder() + .message("some error") + .cause(nodeA) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenThrow(stsException); + + AwsStsCredentialProvider p = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + // This must terminate (no infinite loop) and throw CredentialResolutionException + assertThrows(CredentialResolutionException.class, + () -> p.resolve(user, TEST_TARGET)); + } + + // ========== resolve() wraps SDK IllegalStateException (Item 5) ========== + + @Test + public void testResolveWrapsSdkIllegalStateExceptionAsCredentialResolutionException() { + StsClient mockSts = mock(StsClient.class); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenThrow(new IllegalStateException("client has been closed")); + + AwsStsCredentialProvider p = new AwsStsCredentialProvider( + mockSts, TEST_ROLE_ARN, "test-session", null); + // Do NOT call p.close() -- keep the provider's closed flag false so that + // resolve() reaches the STS call, where the mock throws IllegalStateException. + // This exercises the catch (IllegalStateException e) branch in resolve(). + + UserContext user = new UserContext(TEST_PRINCIPAL, TEST_ISSUER, TEST_TOKEN, + Instant.now(), Instant.now().plusSeconds(300)); + + CredentialResolutionException ex = assertThrows(CredentialResolutionException.class, + () -> p.resolve(user, TEST_TARGET)); + + assertTrue(ex.getMessage().contains("closed"), + "Message should indicate the STS client has been closed"); + assertFalse(ex.getMessage().contains(TEST_TOKEN), + "Raw token must never appear in exception messages"); + + // Verify the cause is the original IllegalStateException from the SDK client + assertTrue(ex.getCause() instanceof IllegalStateException, + "Cause should be the IllegalStateException thrown by the STS client"); + assertEquals("client has been closed", ex.getCause().getMessage()); + + // Verify the mock was actually invoked (proving the closed-flag short-circuit + // was NOT hit and the test truly exercises the ISE catch branch) + verify(mockSts).assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class)); + } + + // ========== Helpers ========== + + /** Creates a config map with {@link #TEST_ROLE_ARN} pre-set. */ + private Map<String, String> confWithRoleArn() { + Map<String, String> conf = new HashMap<>(); + conf.put(AwsStsCredentialProvider.CONF_ROLE_ARN, TEST_ROLE_ARN); + return conf; + } + + /** + * Asserts that {@code init()} throws {@link IllegalArgumentException} whose message + * contains the given config key. Returns the exception for additional assertions. + */ + private IllegalArgumentException assertInitThrowsForConfig(String key, String value) { + Map<String, String> conf = confWithRoleArn(); + conf.put(key, value); + AwsStsCredentialProvider p = new AwsStsCredentialProvider(); + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, + () -> p.init(conf)); + assertTrue(ex.getMessage().contains(key), + "Error message must contain the config key '" + key + "' but was: " + ex.getMessage()); + return ex; + } + + private StsClient createMockStsClient(String accessKeyId, String secretAccessKey, + String sessionToken, Instant expiration) { + StsClient mockSts = mock(StsClient.class); + Credentials creds = Credentials.builder() + .accessKeyId(accessKeyId) + .secretAccessKey(secretAccessKey) + .sessionToken(sessionToken) + .expiration(expiration) + .build(); + AssumeRoleWithWebIdentityResponse response = AssumeRoleWithWebIdentityResponse.builder() + .credentials(creds) + .build(); + when(mockSts.assumeRoleWithWebIdentity(any(AssumeRoleWithWebIdentityRequest.class))) + .thenReturn(response); + return mockSts; + } + + // ========== additionalSparkProperties() ========== + + @Test + public void testAdditionalSparkPropertiesReturnsS3aProviderMapping() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + Map<String, String> props = provider.additionalSparkProperties(); + assertEquals(1, props.size()); + assertEquals( + "org.apache.spark.security.aws.SparkOidcAwsCredentialsProvider", + props.get("spark.hadoop.fs.s3a.aws.credentials.provider")); + } + + @Test + public void testAdditionalSparkPropertiesKeyIncludesSparkHadoopPrefix() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + Map<String, String> props = provider.additionalSparkProperties(); + props.keySet().forEach(key -> + assertTrue(key.startsWith("spark.hadoop."), + "Key must include spark.hadoop. prefix: " + key)); + } + + @Test + public void testAdditionalSparkPropertiesIsUnmodifiable() { + AwsStsCredentialProvider provider = new AwsStsCredentialProvider(); + Map<String, String> props = provider.additionalSparkProperties(); + assertThrows(UnsupportedOperationException.class, + () -> props.put("foo", "bar")); + } +} diff --git a/connector/credential-aws/src/test/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProviderSuite.java b/connector/credential-aws/src/test/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProviderSuite.java new file mode 100644 index 0000000000000..e4c7b54ac3a3f --- /dev/null +++ b/connector/credential-aws/src/test/java/org/apache/spark/security/aws/SparkOidcAwsCredentialsProviderSuite.java @@ -0,0 +1,468 @@ +/* + * 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.security.aws; + +import java.io.ByteArrayOutputStream; +import java.io.ObjectOutputStream; +import java.time.Instant; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import software.amazon.awssdk.auth.credentials.AwsCredentials; +import software.amazon.awssdk.auth.credentials.AwsSessionCredentials; + +import org.apache.spark.SparkEnv; +import org.apache.spark.VersionedCredentials; +import org.apache.spark.security.ServiceCredential; +import org.apache.spark.security.UserCredentials; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link SparkOidcAwsCredentialsProvider}. + * + * <p>Unit tests mock SparkEnv.get() to inject controlled credential store state. + * End-to-end tests use real serialization/deserialization to verify the full path. + */ +public class SparkOidcAwsCredentialsProviderSuite { + + private MockedStatic<SparkEnv> sparkEnvMock; + private SparkEnv mockEnv; + private AtomicReference<VersionedCredentials> credentialStore; + + @BeforeEach + void setUp() { + mockEnv = mock(SparkEnv.class); + credentialStore = new AtomicReference<>(); + when(mockEnv.userCredentials()).thenReturn(credentialStore); + + sparkEnvMock = mockStatic(SparkEnv.class); + sparkEnvMock.when(SparkEnv::get).thenReturn(mockEnv); + } + + @AfterEach + void tearDown() { + sparkEnvMock.close(); + } + + // ========================================================================= + // Happy path + // ========================================================================= + + @Test + void testResolveCredentialsReturnsValidSessionCredentials() { + populateStore("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "FwoGZXIvYXdzEBYaDH...", 1L); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsCredentials creds = provider.resolveCredentials(); + + assertInstanceOf(AwsSessionCredentials.class, creds); + AwsSessionCredentials session = (AwsSessionCredentials) creds; + assertEquals("AKIAIOSFODNN7EXAMPLE", session.accessKeyId()); + assertEquals("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", session.secretAccessKey()); + assertEquals("FwoGZXIvYXdzEBYaDH...", session.sessionToken()); + } + + @Test + void testResolveCredentialsAlwaysReadsFreshFromStore() { + // Populate v1 + populateStore("key-v1", "secret-v1", "token-v1", 1L); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + + // First call returns v1 + AwsSessionCredentials creds1 = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-v1", creds1.accessKeyId()); + + // Update store to v2 (simulates credential refresh) + populateStore("key-v2", "secret-v2", "token-v2", 2L); + + // Second call returns v2 -- no caching + AwsSessionCredentials creds2 = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-v2", creds2.accessKeyId()); + assertEquals("secret-v2", creds2.secretAccessKey()); + assertEquals("token-v2", creds2.sessionToken()); + } + + @Test + void testResolveCredentialsNeverReturnsStaleAfterRefresh() { + // Start with v1 + populateStore("key-old", "secret-old", "token-old", 1L); + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + provider.resolveCredentials(); // consume v1 + + // Refresh to v2 (simulates RPC UpdateUserCredentials) + populateStore("key-new", "secret-new", "token-new", 2L); + + // Multiple subsequent calls all return v2 + for (int i = 0; i < 10; i++) { + AwsSessionCredentials creds = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-new", creds.accessKeyId(), + "Call " + i + " returned stale credentials"); + } + } + + // ========================================================================= + // Error cases + // ========================================================================= + + @Test + void testThrowsWhenSparkEnvIsNull() { + sparkEnvMock.when(SparkEnv::get).thenReturn(null); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + + assertTrue(ex.getMessage().contains("SparkEnv is not available")); + } + + @Test + void testThrowsWhenCredentialStoreIsEmpty() { + // Store is null (no credentials delivered yet) + credentialStore.set(null); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + + assertTrue(ex.getMessage().contains("No credentials available")); + assertTrue(ex.getMessage().contains("spark.security.oidc.enabled=true")); + } + + @Test + void testThrowsWhenS3aSchemeNotPresent() { + // Populate with a credential that has a different scheme (e.g., "hdfs") + ServiceCredential hdfsCred = new ServiceCredential( + Map.of("some.key", "some.value"), Instant.now().plusSeconds(3600)); + UserCredentials credentials = new UserCredentials(Map.of("hdfs", hdfsCred)); + byte[] bytes = serializeCredentials(credentials); + credentialStore.set(new VersionedCredentials(1L, bytes)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + + assertTrue(ex.getMessage().contains("No credential found for scheme 's3a'")); + assertTrue(ex.getMessage().contains("AwsStsCredentialProvider")); + } + + @Test + void testThrowsWhenAccessKeyMissing() { + // Credential with secret and token but no access key + ServiceCredential incompleteCred = new ServiceCredential( + Map.of("fs.s3a.secret.key", "secret", "fs.s3a.session.token", "token"), + Instant.now().plusSeconds(3600)); + UserCredentials credentials = new UserCredentials(Map.of("s3a", incompleteCred)); + byte[] bytes = serializeCredentials(credentials); + credentialStore.set(new VersionedCredentials(1L, bytes)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + + assertTrue(ex.getMessage().contains("missing required properties")); + } + + @Test + void testThrowsWhenSecretKeyMissing() { + ServiceCredential incompleteCred = new ServiceCredential( + Map.of("fs.s3a.access.key", "access", "fs.s3a.session.token", "token"), + Instant.now().plusSeconds(3600)); + UserCredentials credentials = new UserCredentials(Map.of("s3a", incompleteCred)); + byte[] bytes = serializeCredentials(credentials); + credentialStore.set(new VersionedCredentials(1L, bytes)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + assertTrue(ex.getMessage().contains("missing required properties")); + } + + @Test + void testThrowsWhenSessionTokenMissing() { + ServiceCredential incompleteCred = new ServiceCredential( + Map.of("fs.s3a.access.key", "access", "fs.s3a.secret.key", "secret"), + Instant.now().plusSeconds(3600)); + UserCredentials credentials = new UserCredentials(Map.of("s3a", incompleteCred)); + byte[] bytes = serializeCredentials(credentials); + credentialStore.set(new VersionedCredentials(1L, bytes)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + IllegalStateException ex = assertThrows(IllegalStateException.class, + provider::resolveCredentials); + assertTrue(ex.getMessage().contains("missing required properties")); + } + + // ========================================================================= + // End-to-end: real serialization roundtrip + // ========================================================================= + + @Test + void testEndToEndSerializationRoundtrip() { + // Simulate the full driver->executor path: + // 1. Driver creates ServiceCredential with S3A properties + // 2. Driver wraps in UserCredentials + // 3. Driver serializes via UserCredentialManager.serializeUserCredentials + // (Java ObjectOutputStream) + // 4. Bytes stored in VersionedCredentials + // 5. Executor reads via SparkOidcAwsCredentialsProvider.resolveCredentials() + + String expectedAccessKey = "AKIAI44QH8DHBEXAMPLE"; + String expectedSecretKey = "je7MtGbClwBF/2Zp9Utk/h3yCo8nvbEXAMPLEKEY"; + String expectedToken = "AQoDYXdzEJr...<very-long-session-token>..."; + + ServiceCredential s3aCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", expectedAccessKey, + "fs.s3a.secret.key", expectedSecretKey, + "fs.s3a.session.token", expectedToken + ), Instant.now().plusSeconds(3600)); + + UserCredentials userCreds = new UserCredentials(Map.of("s3a", s3aCred)); + byte[] serialized = serializeCredentials(userCreds); + + // Place in store (as executor would receive via TaskDescription) + credentialStore.set(new VersionedCredentials(42L, serialized)); + + // Resolve -- full path + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsSessionCredentials result = (AwsSessionCredentials) provider.resolveCredentials(); + + assertEquals(expectedAccessKey, result.accessKeyId()); + assertEquals(expectedSecretKey, result.secretAccessKey()); + assertEquals(expectedToken, result.sessionToken()); + } + + @Test + void testEndToEndVersionGuardWithMultipleUpdates() { + // Simulate credential refresh cycle: + // v1 arrives via TaskDescription, v2 arrives via RPC, stale v1 arrives again + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + + // v1 arrives + populateStore("key-v1", "secret-v1", "token-v1", 1L); + AwsSessionCredentials r1 = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-v1", r1.accessKeyId()); + + // v2 arrives (credential refresh from driver) + populateStoreWithVersionGuard("key-v2", "secret-v2", "token-v2", 2L); + AwsSessionCredentials r2 = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-v2", r2.accessKeyId()); + + // Stale v1 arrives (delayed TaskDescription) -- version guard rejects it + populateStoreWithVersionGuard("key-v1-stale", "secret-v1-stale", "token-v1-stale", 1L); + AwsSessionCredentials r3 = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-v2", r3.accessKeyId(), "Stale v1 should not overwrite v2"); + } + + @Test + void testEndToEndSchemeNormalizationByUserCredentials() { + // UserCredentials constructor normalizes scheme keys to lowercase. + // This test verifies our provider works correctly with that normalized store. + ServiceCredential s3aCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", "key-normalized", + "fs.s3a.secret.key", "secret-normalized", + "fs.s3a.session.token", "token-normalized" + ), Instant.now().plusSeconds(3600)); + + // UserCredentials normalizes "s3a" to lowercase internally + UserCredentials userCreds = new UserCredentials(Map.of("s3a", s3aCred)); + byte[] serialized = serializeCredentials(userCreds); + credentialStore.set(new VersionedCredentials(1L, serialized)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsSessionCredentials result = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-normalized", result.accessKeyId()); + } + + @Test + void testEndToEndWithExpiredCredentialStillReturns() { + // Expired credentials should still be returned -- the provider does NOT check expiry. + // S3A will get a 403 and retry, triggering another resolveCredentials() which + // should by then have fresh creds from the renewal loop. + ServiceCredential expiredCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", "key-expired", + "fs.s3a.secret.key", "secret-expired", + "fs.s3a.session.token", "token-expired" + ), Instant.now().minusSeconds(3600)); // expired 1 hour ago + + UserCredentials userCreds = new UserCredentials(Map.of("s3a", expiredCred)); + byte[] serialized = serializeCredentials(userCreds); + credentialStore.set(new VersionedCredentials(1L, serialized)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsSessionCredentials result = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-expired", result.accessKeyId()); + } + + @Test + void testEndToEndMultipleSchemesBundled() { + // UserCredentials can have multiple schemes -- we only read s3a + ServiceCredential s3aCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", "s3a-key", + "fs.s3a.secret.key", "s3a-secret", + "fs.s3a.session.token", "s3a-token" + ), Instant.now().plusSeconds(3600)); + ServiceCredential abfsCred = new ServiceCredential( + Map.of("fs.azure.account.key", "azure-key"), + Instant.now().plusSeconds(3600)); + + UserCredentials userCreds = new UserCredentials(Map.of( + "s3a", s3aCred, + "abfs", abfsCred + )); + byte[] serialized = serializeCredentials(userCreds); + credentialStore.set(new VersionedCredentials(1L, serialized)); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsSessionCredentials result = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("s3a-key", result.accessKeyId()); + assertEquals("s3a-secret", result.secretAccessKey()); + assertEquals("s3a-token", result.sessionToken()); + } + + @Test + void testImplementsAwsCredentialsProviderInterface() { + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + assertInstanceOf(software.amazon.awssdk.auth.credentials.AwsCredentialsProvider.class, + provider); + } + + // ========================================================================= + // Concurrency: thread-safety of resolveCredentials + // ========================================================================= + + @Test + void testRapidStoreUpdatesReturnLatestVersion() throws InterruptedException { + // Note: Mockito mockStatic is scoped to the declaring thread by default. + // We test concurrency by rapidly updating the store between sequential calls, + // verifying atomic read consistency from a single thread (which is what the + // real AtomicReference guarantees for the multi-threaded executor case). + populateStore("key-initial", "secret-initial", "token-initial", 1L); + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + + // Rapidly alternate between reading and updating + for (int i = 0; i < 100; i++) { + AwsSessionCredentials creds = (AwsSessionCredentials) provider.resolveCredentials(); + // Every read must return a complete, non-null credential set + assertNotNull(creds.accessKeyId()); + assertNotNull(creds.secretAccessKey()); + assertNotNull(creds.sessionToken()); + + // Update store mid-loop (simulates concurrent RPC updates) + populateStore("key-" + i, "secret-" + i, "token-" + i, (long) (i + 2)); + } + + // Final read must return the latest version + AwsSessionCredentials finalCreds = (AwsSessionCredentials) provider.resolveCredentials(); + assertEquals("key-99", finalCreds.accessKeyId()); + } + + // ========================================================================= + // Auto-config integration (verifies contract with AwsStsCredentialProvider) + // ========================================================================= + + @Test + void testClassNameMatchesAutoConfigValue() { + AwsStsCredentialProvider stsProvider = new AwsStsCredentialProvider(); + stsProvider.init(Map.of( + "spark.security.oidc.aws.roleArn", "arn:aws:iam::123456789012:role/test", + "spark.security.oidc.aws.region", "us-east-1")); + Map<String, String> props = stsProvider.additionalSparkProperties(); + assertEquals( + SparkOidcAwsCredentialsProvider.class.getName(), + props.get("spark.hadoop.fs.s3a.aws.credentials.provider")); + stsProvider.close(); + } + + // ========================================================================= + // Cache behavior: same version returns same instance (no re-deserialization) + // ========================================================================= + + @Test + void testCacheHitReturnsSameInstanceWhenVersionUnchanged() { + populateStore("key-1", "secret-1", "token-1", 1L); + + SparkOidcAwsCredentialsProvider provider = new SparkOidcAwsCredentialsProvider(); + AwsCredentials result1 = provider.resolveCredentials(); + AwsCredentials result2 = provider.resolveCredentials(); + + // Same cached instance proves deserialization was skipped on second call + assertSame(result1, result2); + } + + // ========================================================================= + // Helpers + // ========================================================================= + + private void populateStore(String accessKey, String secretKey, String sessionToken, + long version) { + ServiceCredential s3aCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", accessKey, + "fs.s3a.secret.key", secretKey, + "fs.s3a.session.token", sessionToken + ), Instant.now().plusSeconds(3600)); + UserCredentials userCreds = new UserCredentials(Map.of("s3a", s3aCred)); + byte[] bytes = serializeCredentials(userCreds); + credentialStore.set(new VersionedCredentials(version, bytes)); + } + + private void populateStoreWithVersionGuard(String accessKey, String secretKey, + String sessionToken, long version) { + ServiceCredential s3aCred = new ServiceCredential(Map.of( + "fs.s3a.access.key", accessKey, + "fs.s3a.secret.key", secretKey, + "fs.s3a.session.token", sessionToken + ), Instant.now().plusSeconds(3600)); + UserCredentials userCreds = new UserCredentials(Map.of("s3a", s3aCred)); + byte[] bytes = serializeCredentials(userCreds); + VersionedCredentials.updateIfNewer(credentialStore, version, bytes); + } + + /** + * Serialize UserCredentials to bytes using Java ObjectOutputStream. + * This mirrors UserCredentialManager.serializeUserCredentials which is + * package-private to org.apache.spark.deploy.security. + */ + private static byte[] serializeCredentials(UserCredentials credentials) { + try { + ByteArrayOutputStream bos = new ByteArrayOutputStream(); + try (ObjectOutputStream oos = new ObjectOutputStream(bos)) { + oos.writeObject(credentials); + oos.flush(); + } + return bos.toByteArray(); + } catch (java.io.IOException e) { + throw new java.io.UncheckedIOException("Failed to serialize UserCredentials", e); + } + } +} diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala index c2714587e2d24..b18e9e4dac824 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/MySQLIntegrationSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.jdbc.v2 import java.sql.{Connection, SQLFeatureNotSupportedException} import org.apache.spark.{SparkConf, SparkSQLFeatureNotSupportedException} -import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.{AnalysisException, Row} import org.apache.spark.sql.execution.datasources.v2.jdbc.JDBCTableCatalog import org.apache.spark.sql.jdbc.MySQLDatabaseOnDocker import org.apache.spark.sql.types._ @@ -89,6 +89,7 @@ class MySQLIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTest new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", dataType.isInstanceOf[NumericType]) .putString("jdbcClientType", jdbcClientType) .build() @@ -313,6 +314,15 @@ class MySQLIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTest assert(rows10(0).getString(0) === "amy") assert(rows10(1).getString(0) === "alex") } + + test("do not push down casts to double") { + val df = sql( + s"SELECT name FROM $catalogName.employee " + + "WHERE CAST(salary AS DOUBLE) > 10000.5") + + checkFilterPushed(df, pushed = false) + checkAnswer(df, Seq(Row("alex"), Row("jen"))) + } } /** @@ -334,6 +344,7 @@ class MySQLOverMariaConnectorIntegrationSuite extends MySQLIntegrationSuite { new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", true) .putString("jdbcClientType", jdbcClientType) .build() diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/OracleIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/OracleIntegrationSuite.scala index f7ba1e1e0dbdf..68e74ef731fed 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/OracleIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/OracleIntegrationSuite.scala @@ -101,6 +101,7 @@ class OracleIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCTes new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean( "isSigned", dataType.isInstanceOf[NumericType] || dataType.isInstanceOf[StringType]) diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala index d57d3aa5ea03e..9b349b0484ff3 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/PostgresIntegrationSuite.scala @@ -50,6 +50,7 @@ class PostgresIntegrationSuite extends DockerJDBCIntegrationV2Suite with V2JDBCT new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", dataType.isInstanceOf[NumericType]) .putString("jdbcClientType", jdbcClientType) .build() diff --git a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala index 6f2a9e97ff005..3b1724877241f 100644 --- a/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala +++ b/connector/docker-integration-tests/src/test/scala/org/apache/spark/sql/jdbc/v2/V2JDBCTest.scala @@ -56,6 +56,7 @@ private[v2] trait V2JDBCTest jdbcClientType: String = "STRING"): Metadata = new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", dataType.isInstanceOf[NumericType]) .putString("jdbcClientType", jdbcClientType) .build() diff --git a/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json b/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json index c256b8cbebb2d..6aaf067cd3f4c 100644 --- a/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json +++ b/connector/kafka-0-10-sql/src/main/resources/error/kafka-error-conditions.json @@ -36,6 +36,13 @@ "Specified: <specifiedPartitions> Assigned: <assignedPartitions>" ] }, + "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED" : { + "message" : [ + "Topics specified with a topic-level offset for Kafka offsets don't have any assigned partition. Maybe the topics are misspelled, ", + "or they are not part of the topics being subscribed.", + "Specified: <specifiedTopics> Assigned: <assignedTopics>" + ] + }, "KAFKA_TIMESTAMP_OFFSET_DOES_NOT_MATCH_ASSIGNED" : { "message" : [ "Partitions specified for Kafka timestamp based <position> offsets don't match what are assigned. Maybe topic partitions are created ", diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala index 7f3eb0370a078..618064d883145 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/JsonUtils.scala @@ -17,11 +17,14 @@ package org.apache.spark.sql.kafka010 +import java.util.Locale + import scala.collection.mutable.HashMap import scala.util.control.NonFatal import org.apache.kafka.common.TopicPartition -import org.json4s.{Formats, NoTypeHints} +import org.json4s.{Formats, JObject, JString, NoTypeHints} +import org.json4s.jackson.JsonMethods.parse import org.json4s.jackson.Serialization /** @@ -76,6 +79,50 @@ private object JsonUtils { } } + /** + * Read the offsets of the `startingOffsets` / `endingOffsets` json string. On top of the + * per-TopicPartition form read by [[partitionOffsets]], a topic may bind to "earliest" or + * "latest" as a whole, e.g. {"topicA":"earliest","topicB":{"0":23,"1":-1}}. Such topic-level + * values are returned unexpanded in `SpecificOffsetRangeLimit.topicOffsets`, since the + * partitions of the topic are only known once the offsets get resolved against Kafka. + */ + def specificOffsets(str: String): SpecificOffsetRangeLimit = { + def fail(): Nothing = throw new IllegalArgumentException( + s"""Expected e.g. {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} or + |{"topicA":"earliest","topicB":"latest"}, got $str""".stripMargin) + + val partitionOffsets = new HashMap[TopicPartition, Long] + val topicOffsets = new HashMap[String, Long] + try { + parse(str) match { + case JObject(topics) => + topics.foreach { + case (topic, JString(value)) => + topicOffsets += topic -> (value.toLowerCase(Locale.ROOT) match { + case "earliest" => KafkaOffsetRangeLimit.EARLIEST + case "latest" => KafkaOffsetRangeLimit.LATEST + case _ => fail() + }) + case (topic, partOffsets) => + partOffsets.extract[Map[Int, Long]].foreach { case (part, offset) => + partitionOffsets += new TopicPartition(topic, part) -> offset + } + } + case _ => fail() + } + } catch { + case NonFatal(_) => fail() + } + val bothForms = topicOffsets.keySet.intersect(partitionOffsets.keySet.map(_.topic)) + if (bothForms.nonEmpty) { + throw new IllegalArgumentException( + s"""Topic(s) ${bothForms.toSeq.sorted.mkString(", ")} are given both a topic-level offset + |and per-partition offsets, only one of the two forms is allowed per topic, + |got $str""".stripMargin) + } + SpecificOffsetRangeLimit(partitionOffsets.toMap, topicOffsets.toMap) + } + def partitionTimestamps(str: String): Map[TopicPartition, Long] = { try { Serialization.read[Map[String, Map[Int, Long]]](str).flatMap { case (topic, partTimestamps) => diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala index 041fe074f7e40..e0be21dadc2cd 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaContinuousStream.scala @@ -70,7 +70,7 @@ class KafkaContinuousStream( val offsets = initialOffsets match { case EarliestOffsetRangeLimit => KafkaSourceOffset(offsetReader.fetchEarliestOffsets()) case LatestOffsetRangeLimit => KafkaSourceOffset(offsetReader.fetchLatestOffsets(None)) - case SpecificOffsetRangeLimit(p) => offsetReader.fetchSpecificOffsets(p, reportDataLoss) + case o: SpecificOffsetRangeLimit => offsetReader.fetchSpecificOffsets(o, reportDataLoss) case SpecificTimestampRangeLimit(p, strategy) => offsetReader.fetchSpecificTimestampBasedOffsets(p, isStartingOffsets = true, strategy) case GlobalTimestampRangeLimit(ts, strategy) => diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala index 17472fb920f19..4b7e60205a115 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaExceptions.scala @@ -166,6 +166,16 @@ object KafkaExceptions { "assignedPartitions" -> assignedPartitions.toString)) } + def topicOffsetDoesNotMatchAssigned( + specifiedTopics: Set[String], + assignedTopics: Set[String]): KafkaIllegalStateException = { + new KafkaIllegalStateException( + errorClass = "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED", + messageParameters = Map( + "specifiedTopics" -> specifiedTopics.toString, + "assignedTopics" -> assignedTopics.toString)) + } + def timestampOffsetDoesNotMatchAssigned( isStartingOffsets: Boolean, specifiedPartitions: Set[TopicPartition], diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala index 17323165b451f..2144cd29b758d 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala @@ -379,8 +379,8 @@ private[kafka010] class KafkaMicroBatchStream( KafkaSourceOffset(kafkaOffsetReader.fetchEarliestOffsets()) case LatestOffsetRangeLimit => KafkaSourceOffset(kafkaOffsetReader.fetchLatestOffsets(None)) - case SpecificOffsetRangeLimit(p) => - kafkaOffsetReader.fetchSpecificOffsets(p, reportDataLoss) + case o: SpecificOffsetRangeLimit => + kafkaOffsetReader.fetchSpecificOffsets(o, reportDataLoss) case SpecificTimestampRangeLimit(p, strategy) => kafkaOffsetReader.fetchSpecificTimestampBasedOffsets(p, isStartingOffsets = true, strategy) @@ -514,7 +514,7 @@ object KafkaMicroBatchStream extends Logging { latestAvailablePartitionOffsets: Option[PartitionOffsetMap]): ju.Map[String, String] = { val offset = Option(latestConsumedOffset.orElse(null)) - if (offset.nonEmpty && latestAvailablePartitionOffsets.isDefined) { + if (offset.nonEmpty && latestAvailablePartitionOffsets.exists(_ != null)) { val consumedPartitionOffsets = offset.map(KafkaSourceOffset(_)).get.partitionToOffsets val offsetsBehindLatest = latestAvailablePartitionOffsets.get .map(partitionOffset => partitionOffset._2 - diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala index ef2feb679afa4..365e8cb434b90 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimit.scala @@ -40,9 +40,37 @@ private[kafka010] case object LatestOffsetRangeLimit extends KafkaOffsetRangeLim /** * Represents the desire to bind to specific offsets. A offset == -1 binds to the * latest offset, and offset == -2 binds to the earliest offset. + * + * `topicOffsets` holds topic-level bindings, written in the JSON as "earliest" or "latest" in + * place of a topic's per-partition object. They are expanded against the partitions discovered + * for the topic when the offsets are resolved, so they keep working across repartitioning. */ private[kafka010] case class SpecificOffsetRangeLimit( - partitionOffsets: Map[TopicPartition, Long]) extends KafkaOffsetRangeLimit + partitionOffsets: Map[TopicPartition, Long], + topicOffsets: Map[String, Long] = Map.empty) extends KafkaOffsetRangeLimit { + + /** + * Expands the topic-level bindings against `assignedPartitions` and returns the resulting + * offset per topic-partition. `assignedPartitions` is by-name so that a fully enumerated + * limit never pays for discovering the partitions. + */ + def resolve(assignedPartitions: => Set[TopicPartition]): Map[TopicPartition, Long] = { + if (topicOffsets.isEmpty) { + partitionOffsets + } else { + val partitions = assignedPartitions + val expanded = partitions.collect { + case tp if topicOffsets.contains(tp.topic) => tp -> topicOffsets(tp.topic) + }.toMap + val unmatchedTopics = topicOffsets.keySet.diff(expanded.keySet.map(_.topic)) + if (unmatchedTopics.nonEmpty) { + throw KafkaExceptions.topicOffsetDoesNotMatchAssigned( + unmatchedTopics, partitions.map(_.topic)) + } + expanded ++ partitionOffsets + } + } +} /** * Represents the desire to bind to earliest offset which timestamp for the offset is equal or diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala index b8624c5ae6376..9233205cbdcea 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReader.scala @@ -62,11 +62,14 @@ private[kafka010] trait KafkaOffsetReader { * This method resolves offset value -1 to the latest and -2 to the * earliest Kafka seek position. * - * @param partitionOffsets the specific offsets to resolve + * Offsets bound to a topic as a whole are expanded against the partitions currently assigned + * to the query. + * + * @param offsets the specific offsets to resolve * @param reportDataLoss callback to either report or log data loss depending on setting */ def fetchSpecificOffsets( - partitionOffsets: Map[TopicPartition, Long], + offsets: SpecificOffsetRangeLimit, reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset /** @@ -176,6 +179,11 @@ private[kafka010] object KafkaOffsetReader extends Logging { private[kafka010] abstract class KafkaOffsetReaderBase extends KafkaOffsetReader with Logging { protected val rangeCalculator: KafkaOffsetRangeCalculator + /** + * @return The set of TopicPartitions currently assigned to the query. + */ + protected def fetchTopicPartitions(): Set[TopicPartition] + private def getSortedExecutorList: Array[String] = { def compare(a: ExecutorCacheTaskLocation, b: ExecutorCacheTaskLocation): Boolean = { if (a.host == b.host) { diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala index f182e156c76e9..f6e54c49cc14e 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderAdmin.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.kafka010 import java.{util => ju} import java.util.Locale +import java.util.concurrent.TimeUnit import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -62,7 +63,7 @@ private[kafka010] class KafkaOffsetReaderAdmin( // Protected by this.synchronized (always accessed inside withRetries, which holds the lock). private var cachedPartitions: Set[TopicPartition] = Set.empty - private var cacheTimestampMs: Long = 0L + private var cacheTimestampNanos: Option[Long] = None /** * An AdminClient used in the driver to query the latest Kafka offsets. @@ -123,6 +124,9 @@ private[kafka010] class KafkaOffsetReaderAdmin( stopAdmin() } + override protected def fetchTopicPartitions(): Set[TopicPartition] = + withRetries { resolvePartitions() } + override def fetchPartitionOffsets( offsetRangeLimit: KafkaOffsetRangeLimit, isStartingOffsets: Boolean): Map[TopicPartition, Long] = { @@ -134,7 +138,7 @@ private[kafka010] class KafkaOffsetReaderAdmin( logDebug(s"Assigned partitions: $partitions. Seeking to $partitionOffsets") partitionOffsets } - val partitions = withRetries { resolvePartitions() } + val partitions = fetchTopicPartitions() // Obtain TopicPartition offsets with late binding support offsetRangeLimit match { case EarliestOffsetRangeLimit => partitions.map { @@ -143,8 +147,8 @@ private[kafka010] class KafkaOffsetReaderAdmin( case LatestOffsetRangeLimit => partitions.map { case tp => tp -> KafkaOffsetRangeLimit.LATEST }.toMap - case SpecificOffsetRangeLimit(partitionOffsets) => - validateTopicPartitions(partitions, partitionOffsets) + case offsets: SpecificOffsetRangeLimit => + validateTopicPartitions(partitions, offsets.resolve(partitions)) case SpecificTimestampRangeLimit(partitionTimestamps, strategyOnNoMatchingStartingOffset) => fetchSpecificTimestampBasedOffsets(partitionTimestamps, isStartingOffsets, strategyOnNoMatchingStartingOffset).partitionToOffsets @@ -155,9 +159,12 @@ private[kafka010] class KafkaOffsetReaderAdmin( } override def fetchSpecificOffsets( - partitionOffsets: Map[TopicPartition, Long], + offsets: SpecificOffsetRangeLimit, reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset = { + // Topic-level offsets are expanded against the partitions the fetch is actually run with, + // so that metadata changing in between cannot make valid offsets fail the assertion below val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions => + val partitionOffsets = offsets.resolve(partitions.asScala.toSet) assert(partitions.asScala == partitionOffsets.keySet, "If startingOffsets contains specific offsets, you must specify all TopicPartitions.\n" + "Use -1 for latest, -2 for earliest, if you don't care.\n" + @@ -165,8 +172,8 @@ private[kafka010] class KafkaOffsetReaderAdmin( logDebug(s"Assigned partitions: $partitions. Seeking to $partitionOffsets") } - val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { _ => - partitionOffsets + val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { + partitions => offsets.resolve(partitions.asScala.toSet) } fetchSpecificOffsets0(fnAssertParametersWithPartitions, fnRetrievePartitionOffsets) @@ -424,9 +431,11 @@ private[kafka010] class KafkaOffsetReaderAdmin( // No need to report data loss here val resolvedFromOffsets = - fetchSpecificOffsets(fromOffsetsMap, (_, _) => ()).partitionToOffsets + fetchSpecificOffsets(SpecificOffsetRangeLimit(fromOffsetsMap), (_, _) => ()) + .partitionToOffsets val resolvedUntilOffsets = - fetchSpecificOffsets(untilOffsetsMap, (_, _) => ()).partitionToOffsets + fetchSpecificOffsets(SpecificOffsetRangeLimit(untilOffsetsMap), (_, _) => ()) + .partitionToOffsets val ranges = offsetRangesBase.map(_.topicPartition).map { tp => KafkaOffsetRange(tp, resolvedFromOffsets(tp), resolvedUntilOffsets(tp), preferredLoc = None) } @@ -451,16 +460,20 @@ private[kafka010] class KafkaOffsetReaderAdmin( if (partitionMetadataCacheTtlMs <= 0) { return consumerStrategy.assignedTopicPartitions(admin) } - val now = System.currentTimeMillis() - if (cacheTimestampMs > 0 && (now - cacheTimestampMs) < partitionMetadataCacheTtlMs) { - logDebug(s"Reusing cached partitions (age ${now - cacheTimestampMs}ms < " + - s"${partitionMetadataCacheTtlMs}ms TTL): $cachedPartitions") - cachedPartitions - } else { - val fresh = consumerStrategy.assignedTopicPartitions(admin) - cachedPartitions = fresh - cacheTimestampMs = now - fresh + val nowNanos = System.nanoTime() + cacheTimestampNanos match { + case Some(timestampNanos) + if nowNanos - timestampNanos < + TimeUnit.MILLISECONDS.toNanos(partitionMetadataCacheTtlMs) => + val cacheAgeMs = TimeUnit.NANOSECONDS.toMillis(nowNanos - timestampNanos) + logDebug(s"Reusing cached partitions (age ${cacheAgeMs}ms < " + + s"${partitionMetadataCacheTtlMs}ms TTL): $cachedPartitions") + cachedPartitions + case _ => + val fresh = consumerStrategy.assignedTopicPartitions(admin) + cachedPartitions = fresh + cacheTimestampNanos = Some(nowNanos) + fresh } } @@ -519,6 +532,6 @@ private[kafka010] class KafkaOffsetReaderAdmin( stopAdmin() _admin = null // will automatically get reinitialized again cachedPartitions = Set.empty - cacheTimestampMs = 0L + cacheTimestampNanos = None } } diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala index 5196e967399cf..3666dfe7541e5 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaOffsetReaderConsumer.scala @@ -125,10 +125,7 @@ private[kafka010] class KafkaOffsetReaderConsumer( uninterruptibleThreadRunner.shutdown() } - /** - * @return The Set of TopicPartitions for a given topic - */ - private def fetchTopicPartitions(): Set[TopicPartition] = + override protected def fetchTopicPartitions(): Set[TopicPartition] = uninterruptibleThreadRunner.runUninterruptibly { assert(Thread.currentThread().isInstanceOf[UninterruptibleThread]) // Poll to get the latest assigned partitions @@ -158,8 +155,8 @@ private[kafka010] class KafkaOffsetReaderConsumer( case LatestOffsetRangeLimit => partitions.map { case tp => tp -> KafkaOffsetRangeLimit.LATEST }.toMap - case SpecificOffsetRangeLimit(partitionOffsets) => - validateTopicPartitions(partitions, partitionOffsets) + case offsets: SpecificOffsetRangeLimit => + validateTopicPartitions(partitions, offsets.resolve(partitions)) case SpecificTimestampRangeLimit(partitionTimestamps, strategy) => fetchSpecificTimestampBasedOffsets(partitionTimestamps, isStartingOffsets, strategy).partitionToOffsets @@ -170,9 +167,12 @@ private[kafka010] class KafkaOffsetReaderConsumer( } override def fetchSpecificOffsets( - partitionOffsets: Map[TopicPartition, Long], + offsets: SpecificOffsetRangeLimit, reportDataLoss: (String, () => Throwable) => Unit): KafkaSourceOffset = { + // Topic-level offsets are expanded against the partitions the fetch is actually run with, + // so that metadata changing in between cannot make valid offsets fail the assertion below val fnAssertParametersWithPartitions: ju.Set[TopicPartition] => Unit = { partitions => + val partitionOffsets = offsets.resolve(partitions.asScala.toSet) assert(partitions.asScala == partitionOffsets.keySet, "If startingOffsets contains specific offsets, you must specify all TopicPartitions.\n" + "Use -1 for latest, -2 for earliest, if you don't care.\n" + @@ -180,12 +180,14 @@ private[kafka010] class KafkaOffsetReaderConsumer( logDebug(s"Partitions assigned to consumer: $partitions. Seeking to $partitionOffsets") } - val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { _ => - partitionOffsets + val fnRetrievePartitionOffsets: ju.Set[TopicPartition] => Map[TopicPartition, Long] = { + partitions => offsets.resolve(partitions.asScala.toSet) } val fnAssertFetchedOffsets: Map[TopicPartition, Long] => Unit = { fetched => - partitionOffsets.foreach { + // Only the explicitly specified offsets can be checked here. Topic-level ones always + // expand to the earliest/latest sentinels, which have no expected value to compare to. + offsets.partitionOffsets.foreach { case (tp, off) if off != KafkaOffsetRangeLimit.LATEST && off != KafkaOffsetRangeLimit.EARLIEST => if (fetched(tp) != off) { @@ -466,9 +468,11 @@ private[kafka010] class KafkaOffsetReaderConsumer( // No need to report data loss here val resolvedFromOffsets = - fetchSpecificOffsets(fromOffsetsMap, (_, _) => ()).partitionToOffsets + fetchSpecificOffsets(SpecificOffsetRangeLimit(fromOffsetsMap), (_, _) => ()) + .partitionToOffsets val resolvedUntilOffsets = - fetchSpecificOffsets(untilOffsetsMap, (_, _) => ()).partitionToOffsets + fetchSpecificOffsets(SpecificOffsetRangeLimit(untilOffsetsMap), (_, _) => ()) + .partitionToOffsets val ranges = offsetRangesBase.map(_.topicPartition).map { tp => KafkaOffsetRange(tp, resolvedFromOffsets(tp), resolvedUntilOffsets(tp), preferredLoc = None) } diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala index 62cd20448453b..3ff43cef1fb5a 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala @@ -125,7 +125,7 @@ private[kafka010] class KafkaSource( val offsets = startingOffsets match { case EarliestOffsetRangeLimit => KafkaSourceOffset(kafkaReader.fetchEarliestOffsets()) case LatestOffsetRangeLimit => KafkaSourceOffset(kafkaReader.fetchLatestOffsets(None)) - case SpecificOffsetRangeLimit(p) => kafkaReader.fetchSpecificOffsets(p, reportDataLoss) + case o: SpecificOffsetRangeLimit => kafkaReader.fetchSpecificOffsets(o, reportDataLoss) case SpecificTimestampRangeLimit(p, strategy) => kafkaReader.fetchSpecificTimestampBasedOffsets(p, isStartingOffsets = true, strategy) case GlobalTimestampRangeLimit(ts, strategy) => diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala index b5a7bcb98fba2..e950cdefc35da 100644 --- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala +++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSourceProvider.scala @@ -375,10 +375,10 @@ private[kafka010] class KafkaSourceProvider extends DataSourceRegister case LatestOffsetRangeLimit => throw new IllegalArgumentException("starting offset can't be latest " + "for batch queries on Kafka") - case SpecificOffsetRangeLimit(partitionOffsets) => - partitionOffsets.foreach { - case (tp, off) if off == KafkaOffsetRangeLimit.LATEST => - throw new IllegalArgumentException(s"startingOffsets for $tp can't " + + case SpecificOffsetRangeLimit(partitionOffsets, topicOffsets) => + (partitionOffsets.map { case (tp, off) => tp.toString -> off } ++ topicOffsets).foreach { + case (name, off) if off == KafkaOffsetRangeLimit.LATEST => + throw new IllegalArgumentException(s"startingOffsets for $name can't " + "be latest for batch queries on Kafka") case _ => // ignore } @@ -393,10 +393,10 @@ private[kafka010] class KafkaSourceProvider extends DataSourceRegister throw new IllegalArgumentException("ending offset can't be earliest " + "for batch queries on Kafka") case LatestOffsetRangeLimit => // good to go - case SpecificOffsetRangeLimit(partitionOffsets) => - partitionOffsets.foreach { - case (tp, off) if off == KafkaOffsetRangeLimit.EARLIEST => - throw new IllegalArgumentException(s"ending offset for $tp can't be " + + case SpecificOffsetRangeLimit(partitionOffsets, topicOffsets) => + (partitionOffsets.map { case (tp, off) => tp.toString -> off } ++ topicOffsets).foreach { + case (name, off) if off == KafkaOffsetRangeLimit.EARLIEST => + throw new IllegalArgumentException(s"ending offset for $name can't be " + "earliest for batch queries on Kafka") case _ => // ignore } @@ -640,19 +640,25 @@ private[kafka010] object KafkaSourceProvider extends Logging { startOffset match { case start: SpecificOffsetRangeLimit if endOffset.isInstanceOf[SpecificOffsetRangeLimit] => val end = endOffset.asInstanceOf[SpecificOffsetRangeLimit] - if (start.partitionOffsets.keySet != end.partitionOffsets.keySet) { + // Topic-level offsets are only expanded once the partitions are discovered, so with them + // the two sides can legitimately enumerate different topic-partitions here. Matching them + // up is then left to the offset readers, which see the assigned partitions. + val enumeratesAllPartitions = start.topicOffsets.isEmpty && end.topicOffsets.isEmpty + if (enumeratesAllPartitions && + start.partitionOffsets.keySet != end.partitionOffsets.keySet) { throw KafkaExceptions.unmatchedTopicPartitionsBetweenOffsets( start.partitionOffsets.keySet, end.partitionOffsets.keySet ) } start.partitionOffsets.foreach { - case (tp, startOffset) => + case (tp, startOffset) if end.partitionOffsets.contains(tp) => checkStartOffsetNotGreaterThanEndOffset( startOffset, end.partitionOffsets(tp), tp, KafkaExceptions.unresolvedStartOffsetGreaterThanEndOffset ) + case _ => // the offsets of this partition are not comparable before resolution } case start: SpecificTimestampRangeLimit @@ -711,7 +717,7 @@ private[kafka010] object KafkaSourceProvider extends Logging { LatestOffsetRangeLimit case Some(offset) if offset.toLowerCase(Locale.ROOT) == "earliest" => EarliestOffsetRangeLimit - case Some(json) => SpecificOffsetRangeLimit(JsonUtils.partitionOffsets(json)) + case Some(json) => JsonUtils.specificOffsets(json) case None => defaultOffsets } } diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala index 54b980049d1a2..e7b08c2575561 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/JsonUtilsSuite.scala @@ -42,4 +42,59 @@ class JsonUtilsSuite extends SparkFunSuite { assert(parsed(new TopicPartition("topicA", 1)) === -1) assert(parsed(new TopicPartition("topicB", 0)) === -2) } + + test("parsing specificOffsets") { + val parsed = JsonUtils.specificOffsets( + """{"topicA":{"0":23,"1":-1},"topicB":{"0":-2}}""") + + assert(parsed.partitionOffsets === Map( + new TopicPartition("topicA", 0) -> 23L, + new TopicPartition("topicA", 1) -> -1L, + new TopicPartition("topicB", 0) -> -2L)) + assert(parsed.topicOffsets.isEmpty) + } + + test("parsing specificOffsets with topic-level earliest/latest") { + val parsed = JsonUtils.specificOffsets( + """{"topicA":"earliest","topicB":"latest"}""") + + assert(parsed.partitionOffsets.isEmpty) + assert(parsed.topicOffsets === Map( + "topicA" -> KafkaOffsetRangeLimit.EARLIEST, + "topicB" -> KafkaOffsetRangeLimit.LATEST)) + } + + test("parsing specificOffsets mixing topic-level and partition-level values") { + val parsed = JsonUtils.specificOffsets( + """{"topicA":"earliest","topicB":{"0":23,"1":-1},"topicC":"LATEST"}""") + + assert(parsed.partitionOffsets === Map( + new TopicPartition("topicB", 0) -> 23L, + new TopicPartition("topicB", 1) -> -1L)) + assert(parsed.topicOffsets === Map( + "topicA" -> KafkaOffsetRangeLimit.EARLIEST, + "topicC" -> KafkaOffsetRangeLimit.LATEST)) + } + + test("parsing specificOffsets rejects a topic given in both forms") { + val ex = intercept[IllegalArgumentException] { + JsonUtils.specificOffsets("""{"topicA":"earliest","topicA":{"0":23}}""") + } + assert(ex.getMessage.contains("topicA")) + assert(ex.getMessage.contains("both a topic-level offset")) + } + + test("parsing specificOffsets rejects malformed json") { + Seq( + """{"topicA":"first"}""", + """{"topicA":[0,1]}""", + """{"topicA":{"0":"earliest"}}""", + """"earliest"""", + "not json").foreach { json => + val ex = intercept[IllegalArgumentException] { + JsonUtils.specificOffsets(json) + } + assert(ex.getMessage.contains(json)) + } + } } diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaDelegationTokenSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaDelegationTokenSuite.scala index a55b6e0068519..5d6356cb0dad7 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaDelegationTokenSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaDelegationTokenSuite.scala @@ -17,15 +17,23 @@ package org.apache.spark.sql.kafka010 +import java.security.PrivilegedExceptionAction import java.util.UUID +import java.util.concurrent.ExecutionException import org.apache.hadoop.conf.Configuration import org.apache.hadoop.security.{Credentials, UserGroupInformation} +import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclOperation, AclPermissionType} +import org.apache.kafka.common.errors.DelegationTokenAuthorizationException +import org.apache.kafka.common.resource.{PatternType, ResourcePattern, ResourceType} import org.apache.kafka.common.security.auth.SecurityProtocol.SASL_PLAINTEXT +import org.apache.kafka.common.security.token.delegation.TokenInformation +import org.apache.spark.SparkException import org.apache.spark.deploy.SparkHadoopUtil import org.apache.spark.deploy.security.HadoopDelegationTokenManager import org.apache.spark.internal.config.{KEYTAB, PRINCIPAL} +import org.apache.spark.kafka010.{KafkaTokenSparkConf, KafkaTokenUtil} import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.streaming.{OutputMode, StreamTest} import org.apache.spark.sql.test.SharedSparkSession @@ -34,6 +42,13 @@ class KafkaDelegationTokenSuite extends StreamTest with SharedSparkSession with import testImplicits._ + private val clusterIdentifier = "cluster1" + + // The client principal is allowed to obtain tokens for `proxyUser` but not for + // `deniedProxyUser`. + private val proxyUser = "proxyUser" + private val deniedProxyUser = "deniedProxyUser" + protected var testUtils: KafkaTestUtils = _ protected override def sparkConf = super.sparkConf @@ -41,12 +56,26 @@ class KafkaDelegationTokenSuite extends StreamTest with SharedSparkSession with .set("spark.security.credentials.hbase.enabled", "false") .set(KEYTAB, testUtils.clientKeytab) .set(PRINCIPAL, testUtils.clientPrincipal) - .set("spark.kafka.clusters.cluster1.auth.bootstrap.servers", testUtils.brokerAddress) - .set("spark.kafka.clusters.cluster1.security.protocol", SASL_PLAINTEXT.name) + .set(s"spark.kafka.clusters.$clusterIdentifier.auth.bootstrap.servers", testUtils.brokerAddress) + .set(s"spark.kafka.clusters.$clusterIdentifier.security.protocol", SASL_PLAINTEXT.name) override def beforeAll(): Unit = { testUtils = new KafkaTestUtils(Map.empty, true) - testUtils.setup() + try { + testUtils.setup( + testUtils.allowAllAcls(testUtils.kafkaPrincipal(proxyUser).toString) :+ + createTokensAcl(proxyUser)) + } catch { + case e: Throwable => + // ScalaTest skips afterAll when beforeAll throws, so tear down here to avoid leaking + // the KDC, broker, and the global JAAS system property into later suites. + try { + testUtils.teardown() + } finally { + testUtils = null + } + throw e + } super.beforeAll() } @@ -62,17 +91,53 @@ class KafkaDelegationTokenSuite extends StreamTest with SharedSparkSession with } } - testRetry("Roundtrip", 3) { - val hadoopConf = new Configuration() - val manager = new HadoopDelegationTokenManager(spark.sparkContext.conf, hadoopConf, null) + /** + * Allow the client principal to obtain tokens owned by `owner`. The broker looks the ACL up + * by the `KafkaPrincipal.toString` of the owner, hence the `User:` prefixed resource name. + */ + private def createTokensAcl(owner: String): AclBinding = new AclBinding( + new ResourcePattern( + ResourceType.USER, testUtils.kafkaPrincipal(owner).toString, PatternType.LITERAL), + new AccessControlEntry( + testUtils.clientKafkaPrincipal, "*", AclOperation.CREATE_TOKENS, AclPermissionType.ALLOW)) + + private def proxyUgi(user: String): UserGroupInformation = + UserGroupInformation.createProxyUser(user, UserGroupInformation.getCurrentUser()) + + private def obtainDelegationTokens(ugi: UserGroupInformation): Credentials = { val credentials = new Credentials() - manager.obtainDelegationTokens(credentials) + ugi.doAs(new PrivilegedExceptionAction[Unit]() { + override def run(): Unit = { + val manager = new HadoopDelegationTokenManager( + spark.sparkContext.conf, new Configuration(), null) + manager.obtainDelegationTokens(credentials) + } + }) + credentials + } + + /** Ask the broker what it recorded for the token which was just obtained. */ + private def describeToken(credentials: Credentials): TokenInformation = { + val token = credentials.getToken(KafkaTokenUtil.getTokenService(clusterIdentifier)) + assert(token != null, s"No delegation token was obtained for cluster $clusterIdentifier") + val tokenId = new String(token.getIdentifier) + testUtils.describeDelegationTokens().find(_.tokenId() == tokenId) + .getOrElse(fail(s"Token $tokenId is unknown to the broker")) + } + + private def distributeTokens(credentials: Credentials): Unit = { val serializedCredentials = SparkHadoopUtil.get.serialize(credentials) SparkHadoopUtil.get.addDelegationTokens(serializedCredentials, spark.sparkContext.conf) + } + private def createTopic(): String = { val topic = "topic-" + UUID.randomUUID().toString testUtils.createTopic(topic, partitions = 5) + topic + } + /** Write to and read back from `topic`, authenticating with the distributed token. */ + private def roundtrip(topic: String): Unit = { withTempDir { checkpointDir => val input = MemoryStream[String] @@ -115,4 +180,55 @@ class KafkaDelegationTokenSuite extends StreamTest with SharedSparkSession with StopStream ) } + + testRetry("Roundtrip", 3) { + val credentials = obtainDelegationTokens(UserGroupInformation.getCurrentUser()) + + // Without impersonation no owner is sent, so the broker assigns ownership to the requester. + val tokenInfo = describeToken(credentials) + assert(tokenInfo.owner().toString === testUtils.clientKafkaPrincipal) + assert(tokenInfo.tokenRequester().toString === testUtils.clientKafkaPrincipal) + + distributeTokens(credentials) + roundtrip(createTopic()) + } + + testRetry("SPARK-28173: Roundtrip with proxy user", 3) { + // The manager preserves the proxy UGI here despite KEYTAB/PRINCIPAL being set only because + // Hadoop security is SIMPLE in this JVM (a keytab re-login is then a no-op). The production + // proxy flow (ticket cache or direct providers, no principal) preserves it by design. + val credentials = obtainDelegationTokens(proxyUgi(proxyUser)) + + // The token is requested with the client's credentials but owned by the proxy user, so + // connectors authenticate to Kafka as the proxy user. + val tokenInfo = describeToken(credentials) + assert(tokenInfo.owner() === testUtils.kafkaPrincipal(proxyUser)) + assert(tokenInfo.tokenRequester().toString === testUtils.clientKafkaPrincipal) + + distributeTokens(credentials) + val topic = createTopic() + // Deny the client principal on the topic so the roundtrip only succeeds when the connectors + // authenticate as the proxy user (deny overrides the wildcard allow). + testUtils.createAcls(Seq(AclOperation.READ, AclOperation.WRITE).map { op => + new AclBinding( + new ResourcePattern(ResourceType.TOPIC, topic, PatternType.LITERAL), + new AccessControlEntry(testUtils.clientKafkaPrincipal, "*", op, AclPermissionType.DENY)) + }) + roundtrip(topic) + } + + test("SPARK-28173: Obtaining token for proxy user without CreateTokens permission fails") { + val conf = spark.sparkContext.conf + val clusterConf = KafkaTokenSparkConf.getClusterConfig(conf, clusterIdentifier) + proxyUgi(deniedProxyUser).doAs(new PrivilegedExceptionAction[Unit]() { + override def run(): Unit = { + val e = intercept[SparkException] { + KafkaTokenUtil.obtainToken(conf, clusterConf) + } + assert(e.getMessage.contains("CreateTokens")) + assert(e.getCause.asInstanceOf[ExecutionException] + .getCause.isInstanceOf[DelegationTokenAuthorizationException]) + } + }) + } } diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala index 274be623da6fa..b148af71412e4 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala @@ -911,6 +911,54 @@ abstract class KafkaMicroBatchSourceSuiteBase extends KafkaSourceSuiteBase with ) } + test("topic-level offsets tolerate a new topic matching subscribePattern after the start") { + val topicPrefix = newTopic() + val topic = s"$topicPrefix-a" + val topic2 = s"$topicPrefix-b" + testUtils.createTopic(topic, partitions = 1) + testUtils.sendMessages(topic, Array("1"), Some(0)) + + // The starting offsets only name the topic that exists when the query starts + val ds = spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("kafka.metadata.max.age.ms", "1") + .option("subscribePattern", s"$topicPrefix-.*") + .option("startingOffsets", s"""{"$topic":"earliest"}""") + .load() + .selectExpr("CAST(value AS STRING)") + .as[String] + .map(_.toInt) + + testStream(ds)( + StartStream(), + AssertOnQuery { q => + q.processAllAvailable() + true + }, + CheckAnswer(1), + // A topic created after the initial offsets were resolved is picked up as a new partition, + // starting at earliest, instead of tripping the strict topic check + WithOffsetSync(new TopicPartition(topic2, 0), expectedOffset = 1) { () => + testUtils.createTopic(topic2, partitions = 1) + testUtils.sendMessages(topic2, Array("2"), Some(0)) + }, + AssertOnQuery { q => + // The consumer based reader only sees a newly created topic on its next metadata refresh, + // so keep triggering batches until the topic shows up in the query's offsets + eventually(timeout(streamingTimeout)) { + q.processAllAvailable() + val progress = q.lastProgress + assert(progress != null && progress.sources.exists(_.endOffset.contains(topic2)), + s"$topic2 has not been discovered yet") + } + true + }, + CheckAnswer(1, 2) + ) + } + test("ensure that initial offset are written with an extra byte in the beginning (SPARK-19517)") { withTempDir { metadataPath => val topic = "kafka-initial-offset-current" @@ -1882,6 +1930,8 @@ abstract class KafkaMicroBatchV2SourceSuite extends KafkaMicroBatchSourceSuiteBa // test null latestAvailablePartitionOffsets assert(KafkaMicroBatchStream.metrics(Optional.ofNullable(offset), None).isEmpty) + assert(KafkaMicroBatchStream.metrics( + Optional.ofNullable(offset), Some(null).asInstanceOf[Option[PartitionOffsetMap]]).isEmpty) } test("SPARK-57438: metrics should not NPE when latestPartitionOffsets is null") { @@ -2100,6 +2150,42 @@ abstract class KafkaSourceSuiteBase extends KafkaSourceTest { } } + test("subscribing topics from topic-level offsets") { + val topic1 = newTopic() + val topic2 = newTopic() + testUtils.createTopic(topic1, partitions = 3) + testUtils.createTopic(topic2, partitions = 2) + testUtils.sendMessages(topic1, (1 to 3).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic2, Array("11"), Some(0)) + + // topic1 is read from its beginning while topic2 only contributes the records added after + // the query started, without either topic having its partitions enumerated in the option. + val kafka = spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("kafka.metadata.max.age.ms", "1") + .option("subscribe", s"$topic1,$topic2") + .option("startingOffsets", s"""{"$topic1":"earliest","$topic2":"latest"}""") + .load() + .selectExpr("CAST(value AS STRING)") + .as[String] + val mapped = kafka.map(_.toInt) + + testStream(mapped)( + makeSureGetOffsetCalled, + // Records written to topic2 after the query started are read, unlike "11" which predates it + WithOffsetSync(new TopicPartition(topic2, 0), expectedOffset = 2) { () => + testUtils.sendMessages(topic2, Array("12"), Some(0)) + }, + AddKafkaData(Set(topic1, topic2), 4), + CheckAnswer(1, 2, 3, 12, 4), + StopStream, + StartStream(), + CheckAnswer(1, 2, 3, 12, 4) // Should get the data back on recovery + ) + } + test("subscribing topic by name from specific timestamps with non-matching starting offset") { val topic = newTopic() testFromSpecificTimestampsWithNoMatchingStartingOffset(topic, "subscribe" -> topic) @@ -2339,7 +2425,11 @@ abstract class KafkaSourceSuiteBase extends KafkaSourceTest { (STARTING_OFFSETS_OPTION_KEY, "earLiEst", EarliestOffsetRangeLimit), (ENDING_OFFSETS_OPTION_KEY, "laTest", LatestOffsetRangeLimit), (STARTING_OFFSETS_OPTION_KEY, """{"topic-A":{"0":23}}""", - SpecificOffsetRangeLimit(Map(new TopicPartition("topic-A", 0) -> 23))))) { + SpecificOffsetRangeLimit(Map(new TopicPartition("topic-A", 0) -> 23))), + (STARTING_OFFSETS_OPTION_KEY, """{"topic-A":"eArLiEst","topic-B":{"0":23}}""", + SpecificOffsetRangeLimit( + Map(new TopicPartition("topic-B", 0) -> 23), + Map("topic-A" -> KafkaOffsetRangeLimit.EARLIEST))))) { val offset = getKafkaOffsetRangeLimit( CaseInsensitiveMap[String](Map(optionKey -> optionValue)), "dummy", "dummy", optionKey, answer) diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala new file mode 100644 index 0000000000000..8041d7e56d920 --- /dev/null +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaOffsetRangeLimitSuite.scala @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.kafka010 + +import org.apache.kafka.common.TopicPartition + +import org.apache.spark.SparkFunSuite + +class KafkaOffsetRangeLimitSuite extends SparkFunSuite { + + test("resolving topic-level offsets against the assigned partitions") { + val limit = SpecificOffsetRangeLimit( + Map(new TopicPartition("topicB", 0) -> 23L), + Map("topicA" -> KafkaOffsetRangeLimit.EARLIEST)) + val partitions = Set( + new TopicPartition("topicA", 0), + new TopicPartition("topicA", 1), + new TopicPartition("topicB", 0)) + + assert(limit.resolve(partitions) === Map( + new TopicPartition("topicA", 0) -> KafkaOffsetRangeLimit.EARLIEST, + new TopicPartition("topicA", 1) -> KafkaOffsetRangeLimit.EARLIEST, + new TopicPartition("topicB", 0) -> 23L)) + } + + test("resolving fully enumerated offsets doesn't need the assigned partitions") { + val limit = SpecificOffsetRangeLimit(Map(new TopicPartition("topicB", 0) -> 23L)) + assert(limit.resolve(fail("the partitions should not have been fetched")) === + Map(new TopicPartition("topicB", 0) -> 23L)) + } + + test("resolving a topic-level offset of a topic without assigned partitions fails") { + val limit = SpecificOffsetRangeLimit( + Map.empty, + Map("topicA" -> KafkaOffsetRangeLimit.EARLIEST, + "topicC" -> KafkaOffsetRangeLimit.EARLIEST)) + val ex = intercept[KafkaIllegalStateException] { + limit.resolve(Set(new TopicPartition("topicA", 0))) + } + assert(ex.getCondition === "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED") + assert(ex.getMessage.contains("topicC")) + } +} diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeAggregationSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeAggregationSuite.scala new file mode 100644 index 0000000000000..28b397f13e446 --- /dev/null +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeAggregationSuite.scala @@ -0,0 +1,271 @@ +/* + * 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.kafka010 + +import scala.collection.mutable + +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema +import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryWrapper +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.{StringType, StructField, StructType} + +class KafkaRealTimeModeAggregationSuite extends KafkaRealTimeModeBaseSuite { + + test("tumbling window max") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val tumblingWindowDuration = 10 + val numRows = 10 + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy(window(column("timestamp"), s"${tumblingWindowDuration} seconds")) + .max() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("max(value)").cast("STRING") + ).as("value") + ) + + params.query = + writeToKafka("tumbling_window_max_low_latency", outputTopic, checkpointDir, readPart) + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + + for (i <- 0 until 3) { + for (k <- (1 to numRows).reverse) { + val data = ((i * 10).toLong, k) + read.addData(0, Seq(data)) + + val windowDurationMs = tumblingWindowDuration * 1000 + + val startTime = getDateTimeString(((i + 1) * windowDurationMs) - windowDurationMs) + val endTime = getDateTimeString((i + 1) * windowDurationMs) + + expectedResults += new GenericRowWithSchema( + Array(s"{$startTime, $endTime}-${numRows}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + // advance to next batch + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } + + test("tumbling window min") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val tumblingWindowDuration = 10 + val numRows = 10 + + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy(window(column("timestamp"), s"${tumblingWindowDuration} seconds")) + .min() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("min(value)").cast("STRING") + ).as("value") + ) + + params.query = + writeToKafka("tumbling_window_min_low_latency", outputTopic, checkpointDir, readPart) + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + + for (i <- 0 until 3) { + for (k <- (1 to numRows)) { + val data = ((i * 10).toLong, k) + read.addData(0, Seq(data)) + + val windowDurationMs = tumblingWindowDuration * 1000 + + val startTime = getDateTimeString(((i + 1) * windowDurationMs) - windowDurationMs) + val endTime = getDateTimeString((i + 1) * windowDurationMs) + + expectedResults += new GenericRowWithSchema( + Array(s"{$startTime, $endTime}-${1}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + // advance to next batch + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } + + test("tumbling window sum") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val tumblingWindowDuration = 10 + val numRows = 10 + + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy(window(column("timestamp"), s"${tumblingWindowDuration} seconds")) + .sum() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("sum(value)").cast("STRING") + ).as("value") + ) + + params.query = + writeToKafka("tumbling_window_sum_low_latency", outputTopic, checkpointDir, readPart) + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + + for (i <- 0 until 3) { + var sum = 0 + for (k <- (1 to numRows)) { + val data = ((i * 10).toLong, k) + read.addData(0, Seq(data)) + + val windowDurationMs = tumblingWindowDuration * 1000 + + val startTime = getDateTimeString(((i + 1) * windowDurationMs) - windowDurationMs) + val endTime = getDateTimeString((i + 1) * windowDurationMs) + + sum += k + expectedResults += new GenericRowWithSchema( + Array(s"{$startTime, $endTime}-${sum}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + // advance to next batch + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } + + test("tumbling window avg") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val tumblingWindowDuration = 10 + val numRows = 10 + + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy(window(column("timestamp"), s"${tumblingWindowDuration} seconds")) + .avg() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("avg(value)").cast("INT").cast("STRING") + ).as("value") + ) + + params.query = + writeToKafka("tumbling_window_avg_low_latency", outputTopic, checkpointDir, readPart) + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + + for (i <- 0 until 3) { + var sum = 0 + for (k <- (1 to numRows)) { + // Feed varying values (the row index k) so the assertion below actually exercises the + // average rather than passing for any single-value aggregate: a constant input would + // make avg == first == last == that value. + val data = ((i * 10).toLong, k) + read.addData(0, Seq(data)) + + val windowDurationMs = tumblingWindowDuration * 1000 + + val startTime = getDateTimeString(((i + 1) * windowDurationMs) - windowDurationMs) + val endTime = getDateTimeString((i + 1) * windowDurationMs) + + // Update mode emits the running average after each record: sum(1..k) / k. Cast to INT + // to match the query's avg(value) cast. + sum += k + val runningAvg = sum / k + expectedResults += new GenericRowWithSchema( + Array(s"{$startTime, $endTime}-$runningAvg"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + // advance to next batch + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } +} diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeBaseSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeBaseSuite.scala new file mode 100644 index 0000000000000..8eefa70682daa --- /dev/null +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeBaseSuite.scala @@ -0,0 +1,184 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.kafka010 + +import java.io.File +import java.time.{Instant, ZoneId} +import java.time.format.DateTimeFormatter + +import org.apache.kafka.clients.producer.ProducerRecord +import org.scalatest.BeforeAndAfterEach +import org.scalatest.matchers.should.Matchers +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.{SparkContext, ThreadAudit} +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock +import org.apache.spark.sql.execution.streaming.RealTimeTrigger +import org.apache.spark.sql.execution.streaming.sources.LowLatencyMemoryStream +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.{OutputMode, StreamingQuery} +import org.apache.spark.sql.streaming.util.{GlobalSingletonManualClock, StreamManualClock} +import org.apache.spark.sql.test.TestSparkSession +import org.apache.spark.util.Utils + +abstract class KafkaRealTimeModeBaseSuite + extends KafkaSourceTest + with ThreadAudit + with BeforeAndAfterEach + with Matchers { + + import testImplicits._ + + private def defaultTriggerBatchDurationMs: Long = 1000L + + override def beforeAll(): Unit = { + super.beforeAll() + // testing to make sure the cluster is usable + testUtils.createTopic("_test") + testUtils.sendMessage(new ProducerRecord[String, String]("_test", "", "")) + testUtils.deleteTopic("_test") + logInfo("Kafka cluster setup complete....") + + spark.conf.set(SQLConf.SHUFFLE_PARTITIONS.key, 5) + spark.conf.set( + SQLConf.STATE_STORE_PROVIDER_CLASS.key, + "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider" + ) + spark.conf.set("spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled", "true") + spark.conf.set("spark.sql.streaming.stateStore.rocksdb.trackTotalNumberOfRows", "false") + spark.conf.set("spark.sql.streaming.stateStore.checkpointFormatVersion", "2") + spark.conf.set( + SQLConf.STREAMING_REAL_TIME_MODE_MIN_BATCH_DURATION, + defaultTriggerBatchDurationMs + ) + } + + override protected def createSparkSession = + new TestSparkSession( + new SparkContext( + // Ensure we have enough for both stages. 5 source partitions and 5 shuffle partitions + "local[15]", + "microbatch-context", + sparkConf + .set("spark.sql.testkey", "true") + .set("spark.sql.shuffle.partitions", "5") + .set("spark.sql.adaptive.enabled", "false") + .set( + "spark.executor.extraJavaOptions", + "-Dio.netty.leakDetection.level=paranoid" + ) + ) + ) + + override def beforeEach(): Unit = { + super.beforeEach() + GlobalSingletonManualClock.reset() + } + + protected def writeToKafka( + queryName: String, + outputTopic: String, + checkpointDir: File, + df: DataFrame): StreamingQuery = { + df.writeStream + .outputMode(OutputMode.Update()) + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("topic", outputTopic) + .option("checkpointLocation", checkpointDir.getAbsolutePath) + .queryName(queryName) + // The batch duration set here doesn't matter because we manually control batch durations + // via the manual clock. + .trigger(RealTimeTrigger(defaultTriggerBatchDurationMs)) + .start() + } + + protected def readKafkaTopic(topic: String): DataFrame = { + spark.read + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribe", topic) + .option("startingOffsets", "earliest") + .load() + .select(col("value").cast("STRING")) + } + + + protected def runTest(test: (TestParams) => Unit): Unit = { + withTempDir { checkpointDir => + val outputTopic = newTopic() + testUtils.createTopic(outputTopic, partitions = 5) + + val clock = new GlobalSingletonManualClock() + + LowLatencyClock.setClock(clock) + val read = LowLatencyMemoryStream[(Long, Int)](5) + val param = TestParams(null, clock, read, outputTopic, checkpointDir) + try { + test(param) + } finally { + if (param.query != null) { + param.query.stop() + } + + try { + eventually(timeout(60.seconds)) { + + val currentRunningTasks = sparkContext.statusTracker.getExecutorInfos + .map( + i => + s"[host: ${i.host()}" + + s" port: ${i.port} tasks:${i.numRunningTasks()}]" + ) + .toList + + logInfo(s"Current tasks: ${currentRunningTasks}") + + assert( + spark.sparkContext.statusTracker.getExecutorInfos.map(_.numRunningTasks()).sum <= 0, + currentRunningTasks + ) + } + } catch { + case t: Throwable => + // Best-effort diagnostic for a task that never wound down. + logWarning(s"Tasks still running after the query stopped", t) + logWarning(Utils.getThreadDump().map(_.toString).mkString("\n")) + throw t + } + } + } + } + + protected def getDateTimeString(millis: Long): String = { + val instant = Instant.ofEpochMilli(millis) + val formatter = DateTimeFormatter + .ofPattern("yyyy-MM-dd HH:mm:ss") + .withZone(ZoneId.systemDefault()) + formatter.format(instant) + } +} + +case class TestParams( + var query: StreamingQuery, + clock: StreamManualClock, + read: LowLatencyMemoryStream[(Long, Int)], + outputTopic: String, + checkpointDir: File) diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeSuite.scala index ae23b53fc35ad..606f5bce35e30 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeSuite.scala @@ -18,23 +18,51 @@ package org.apache.spark.sql.kafka010 import java.util.UUID +import java.util.regex.Pattern import org.scalatest.matchers.should.Matchers import org.scalatest.time.SpanSugar._ import org.apache.spark.{SparkConf, SparkContext, SparkIllegalStateException} -import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock +import org.apache.spark.sql.Encoders +import org.apache.spark.sql.execution.datasources.v2.{LowLatencyClock, RealTimeStreamScanExec} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming._ +import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateExec import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider +import org.apache.spark.sql.functions.{count, timestamp_seconds, window} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.kafka010.consumer.KafkaDataConsumer -import org.apache.spark.sql.streaming.{StreamingQuery, Trigger} +import org.apache.spark.sql.streaming.{OutputMode, StatefulProcessor, StreamingQuery, TimeMode, + TimerValues, Trigger, TTLConfig, ValueState} import org.apache.spark.sql.streaming.OutputMode.Update import org.apache.spark.sql.streaming.util.GlobalSingletonManualClock import org.apache.spark.sql.test.TestSparkSession import org.apache.spark.util.SystemClock +private class KafkaRunningCountStatefulProcessor + extends StatefulProcessor[String, String, (String, Long)] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState( + "count", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[String], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.map { _ => + val count = Option(countState.get()).getOrElse(0L) + 1L + countState.update(count) + (key, count) + } + } +} + class KafkaRealTimeModeSuite extends KafkaSourceTest with Matchers { @@ -164,6 +192,97 @@ class KafkaRealTimeModeSuite WaitUntilCurrentBatchProcessed) } + test("transformWithState uses a pipelined shuffle and recovers RocksDB changelog state") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + "spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled" -> "true") { + val topic = newTopic() + testUtils.createTopic(topic, partitions = 2) + testUtils.sendMessages(topic, Array("a", "a"), Some(0)) + testUtils.sendMessages(topic, Array("b"), Some(1)) + + val counts = spark.readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribe", topic) + .option("startingOffsets", "earliest") + .load() + .selectExpr("CAST(value AS STRING)") + .as[String] + .groupByKey(identity) + .transformWithState( + new KafkaRunningCountStatefulProcessor, + TimeMode.None(), + Update) + + testStream(counts, Update, sink = new ContinuousMemorySink())( + StartStream(), + CheckAnswerWithTimeout(60000, ("a", 1L), ("a", 2L), ("b", 1L)), + Execute { q => + val plan = q.lastExecution.executedPlan + val statefulOperators = plan.collect { case t: TransformWithStateExec => t } + assert(statefulOperators.size == 1, plan) + assert(statefulOperators.head.isRealTimeMode, plan) + + val exchanges = plan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, plan) + assert(exchanges.forall(_.pipelined), plan) + }, + WaitUntilCurrentBatchProcessed, + StopStream, + new ExternalAction() { + override def runAction(): Unit = { + testUtils.sendMessages(topic, Array("a"), Some(0)) + testUtils.sendMessages(topic, Array("b", "b"), Some(1)) + } + }, + StartStream(), + CheckAnswerWithTimeout( + 60000, + ("a", 1L), ("a", 2L), ("b", 1L), + ("a", 3L), ("b", 2L), ("b", 3L)), + WaitUntilCurrentBatchProcessed) + } + } + + test("transformWithState remains in real-time mode when latestOffset returns null") { + val topic = newTopic() + + val counts = spark.readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("kafka.metadata.max.age.ms", "1") + .option("subscribePattern", "^" + Pattern.quote(topic) + "$") + .option("startingOffsets", "earliest") + .load() + .selectExpr("CAST(value AS STRING)") + .as[String] + .groupByKey(identity) + .transformWithState( + new KafkaRunningCountStatefulProcessor, + TimeMode.None(), + Update) + + testStream(counts, Update, sink = new ContinuousMemorySink())( + StartStream(), + WaitUntilBatchProcessed(0), + Execute { q => + val plan = q.lastExecution.executedPlan + assert(plan.collect { case _: RealTimeStreamScanExec => true }.isEmpty, plan) + val statefulOperators = plan.collect { case t: TransformWithStateExec => t } + assert(statefulOperators.size == 1, plan) + assert(statefulOperators.head.isRealTimeMode, plan) + }, + new ExternalAction() { + override def runAction(): Unit = { + testUtils.createTopic(topic, partitions = 1) + testUtils.sendMessages(topic, Array("a")) + } + }, + CheckAnswerWithTimeout(60000, ("a", 1L)), + StopStream) + } + // A simple unit test that reads from Kakfa source, does a simple map and writes to memory // sink. Make sure there is no data for a whole batch. Also, after restart the first batch // has no data. @@ -679,6 +798,85 @@ class KafkaRealTimeModeSuite } ) } + + // Aggregation over a Kafka source in Real-Time Mode. The aggregate is planned as the streamline + // operator, which merges each record against state and emits as it goes, so in update mode a key + // seen twice produces two outputs -- the running value after each record. + test("aggregation over a Kafka source") { + val topic = newTopic() + testUtils.createTopic(topic, partitions = 2) + + testUtils.sendMessages(topic, Array("a", "b"), Some(0)) + testUtils.sendMessages(topic, Array("a"), Some(1)) + + val aggregated = spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribe", topic) + .option("startingOffsets", "earliest") + .load() + .selectExpr("CAST(value AS STRING) AS key") + .groupBy($"key") + .agg(count("*").as("count")) + .as[(String, Long)] + + testStream(aggregated, Update, sink = new ContinuousMemorySink())( + StartStream(), + // "a" arrives twice, so it is emitted at count 1 and again at count 2. + CheckAnswerWithTimeout(60000, ("a", 1L), ("a", 2L), ("b", 1L)), + WaitUntilCurrentBatchProcessed, + new ExternalAction() { + override def runAction(): Unit = { + testUtils.sendMessages(topic, Array("b", "c"), Some(0)) + } + }, + CheckAnswerWithTimeout(30000, + ("a", 1L), ("a", 2L), ("b", 1L), ("b", 2L), ("c", 1L)), + WaitUntilCurrentBatchProcessed, + StopStream, + new ExternalAction() { + override def runAction(): Unit = { + testUtils.sendMessages(topic, Array("a"), Some(1)) + } + }, + StartStream(), + // The counts continue from the committed state across the restart rather than restarting. + CheckAnswerWithTimeout(30000, + ("a", 1L), ("a", 2L), ("b", 1L), ("b", 2L), ("c", 1L), ("a", 3L)), + WaitUntilCurrentBatchProcessed) + } + + // A tumbling window aggregation over a Kafka source, where the grouping key is derived from an + // event time column rather than being a bare field. + test("tumbling window aggregation over a Kafka source") { + val topic = newTopic() + testUtils.createTopic(topic, partitions = 2) + + // Values are seconds since the epoch; a 10 second window buckets 1-9 together and 11-19 next. + testUtils.sendMessages(topic, Array("1", "2"), Some(0)) + testUtils.sendMessages(topic, Array("11"), Some(1)) + + val windowed = spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribe", topic) + .option("startingOffsets", "earliest") + .load() + .selectExpr("CAST(value AS STRING) AS value") + .select(timestamp_seconds($"value".cast("long")).as("eventTime")) + .groupBy(window($"eventTime", "10 seconds").as("window")) + .agg(count("*").as("count")) + .select($"window".getField("end").cast("long").as[Long], $"count".as[Long]) + + testStream(windowed, Update, sink = new ContinuousMemorySink())( + StartStream(), + // window ending at 10 holds 1 and 2, emitted at count 1 then 2; window ending at 20 holds 11. + CheckAnswerWithTimeout(60000, (10L, 1L), (10L, 2L), (20L, 1L)), + WaitUntilCurrentBatchProcessed) + } + } class KafkaConsumerPoolRealTimeModeSuite diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeWindowSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeWindowSuite.scala new file mode 100644 index 0000000000000..98d3077a97849 --- /dev/null +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRealTimeModeWindowSuite.scala @@ -0,0 +1,247 @@ +/* + * 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.kafka010 + +import scala.collection.mutable + +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema +import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryWrapper +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.{StringType, StructField, StructType} + +class KafkaRealTimeModeWindowSuite extends KafkaRealTimeModeBaseSuite { + + test("tumbling window count") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val tumblingWindowDuration = 10 + val numRows = 10 + + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy( + window(column("timestamp"), s"${tumblingWindowDuration} seconds"), + column("value") + ) + .count() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("value").cast("STRING"), + lit("-"), + col("count").cast("STRING") + ).as("value") + ) + + params.query = writeToKafka("tumbling_window_count_low_latency", + outputTopic, checkpointDir, readPart) + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + + for (i <- 0 until 3) { + for (k <- 0 until numRows) { + val value = k % 2 + val data = ((i * 10).toLong, value) + read.addData({ + data + }) + + /** + * results should be something like this + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-0-1 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-1-1 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-0-2 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-1-2 + */ + val windowDurationMs = tumblingWindowDuration * 1000 + + val startTime = getDateTimeString(((i + 1) * windowDurationMs) - windowDurationMs) + val endTime = getDateTimeString((i + 1) * windowDurationMs) + + expectedResults += new GenericRowWithSchema( + Array(s"{$startTime, $endTime}-$value-${Math.ceil((k + 1) / 2.0).toInt}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + // advance to next batch + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } + + test("sliding window count") { + runTest { + case params @ TestParams(query, clock, read, outputTopic, checkpointDir) => + val numRows = 10 + val slideWindowDuration = 5 + val windowDuration = 10 + + val readPart = read + .toDF() + .select(col("_1").as("timestamp").cast("TIMESTAMP"), col("_2").as("value")) + .groupBy( + window( + column("timestamp"), + s"$windowDuration seconds", + s"$slideWindowDuration seconds" + ), + column("value") + ) + .count() + .select( + concat( + col("window").cast("STRING"), + lit("-"), + col("value").cast("STRING"), + lit("-"), + col("count").cast("STRING") + ).as("value") + ) + + params.query = + writeToKafka("sliding_window_count_low_latency", outputTopic, checkpointDir, readPart) + + // -5 -> 5 + val bucket0 = mutable.HashMap[Int, Int]() + // 0 -> 10 + val bucket1 = mutable.HashMap[Int, Int]() + // 5 -> 15 + val bucket2 = mutable.HashMap[Int, Int]() + // 10 -> 20 + val bucket3 = mutable.HashMap[Int, Int]() + + val expectedResults = mutable.ListBuffer[GenericRowWithSchema]() + for (i <- 0 until 3) { + for (k <- 0 until numRows) { + val value = k % 2 + val data = ((i * 5).toLong, value) + read.addData({ + data + }) + + /** + * Results should be something like + * + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-0-1 + * {1969-12-31 15:59:55, 1969-12-31 16:00:05}-0-1 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-1-1 + * {1969-12-31 15:59:55, 1969-12-31 16:00:05}-1-1 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-0-2 + * {1969-12-31 15:59:55, 1969-12-31 16:00:05}-0-2 + * {1969-12-31 16:00:00, 1969-12-31 16:00:10}-1-2 + * {1969-12-31 15:59:55, 1969-12-31 16:00:05}-1-2 + * ... + */ + val ts = data._1 + if (ts >= -5 && ts < 5) { + bucket0(value) = bucket0.getOrElse(value, 0) + 1 + } + + if (ts >= 0 && ts < 10) { + bucket1(value) = bucket1.getOrElse(value, 0) + 1 + } + + if (ts >= 5 && ts < 15) { + bucket2(value) = bucket2.getOrElse(value, 0) + 1 + } + + if (ts >= 10 && ts < 20) { + bucket3(value) = bucket3.getOrElse(value, 0) + 1 + } + } + + bucket0.foreach(pair => { + val k = pair._1 + val count = pair._2 + for (i <- 1 to count) { + expectedResults += new GenericRowWithSchema( + Array(s"{${getDateTimeString(-5000)}, ${getDateTimeString(5000)}}-$k-${i}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + }) + + bucket1.foreach(pair => { + val k = pair._1 + val count = pair._2 + for (i <- 1 to count) { + expectedResults += new GenericRowWithSchema( + Array(s"{${getDateTimeString(0)}, ${getDateTimeString(10000)}}-$k-${i}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + }) + + bucket2.foreach(pair => { + val k = pair._1 + val count = pair._2 + for (i <- 1 to count) { + expectedResults += new GenericRowWithSchema( + Array(s"{${getDateTimeString(5000)}, ${getDateTimeString(15000)}}-$k-${i}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + }) + + bucket3.foreach(pair => { + val k = pair._1 + val count = pair._2 + for (i <- 1 to count) { + expectedResults += new GenericRowWithSchema( + Array(s"{${getDateTimeString(10000)}, ${getDateTimeString(20000)}}-$k-${i}"), + schema = new StructType().add(StructField("value", StringType)) + ) + } + }) + + eventually(timeout(60.seconds)) { + checkAnswer(readKafkaTopic(outputTopic), expectedResults.toSeq) + } + + expectedResults.clear() + clock.advance(1000) + + eventually(timeout(60.seconds)) { + params.query + .asInstanceOf[StreamingQueryWrapper] + .streamingQuery + .getLatestExecutionContext() + .batchId should be(i + 1) + params.query.lastProgress.sources(0).numInputRows should be(numRows) + } + } + } + } +} diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala index ab515e2a8e494..3b71abf026d2e 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaRelationSuite.scala @@ -25,7 +25,7 @@ import org.apache.kafka.clients.producer.ProducerRecord import org.apache.kafka.common.TopicPartition import org.apache.spark.{SparkConf, TestUtils} -import org.apache.spark.sql.DataFrameReader +import org.apache.spark.sql.{DataFrame, DataFrameReader} import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation import org.apache.spark.sql.internal.SQLConf @@ -151,6 +151,105 @@ abstract class KafkaRelationSuiteBase extends SharedSparkSession with KafkaTest checkAnswer(df, (0 to 30).map(_.toString).toDF()) } + test("topic-level offsets") { + val topic1 = newTopic() + val topic2 = newTopic() + testUtils.createTopic(topic1, partitions = 3) + testUtils.createTopic(topic2, partitions = 2) + testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1)) + testUtils.sendMessages(topic1, Array("20"), Some(2)) + testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray, Some(1)) + + // A topic bound as a whole covers all of its partitions, without enumerating them. This is a + // def so that every query plans against the partitions discovered at that point. + def df: DataFrame = createDF(s"$topic1,$topic2", withOptions = Map( + "kafka.metadata.max.age.ms" -> "1", + "startingOffsets" -> s"""{"$topic1":"earliest","$topic2":"earliest"}""", + "endingOffsets" -> s"""{"$topic1":"latest","$topic2":"latest"}""")) + checkAnswer(df, ((0 to 20) ++ (100 to 119)).map(_.toString).toDF()) + + // Topic-level offsets are expanded against the partitions discovered when the query is + // planned, so a partition added afterwards is picked up without touching the options + testUtils.addPartitions(topic2, 3) + testUtils.sendMessages(topic2, Array("120"), Some(2)) + checkAnswer(df, ((0 to 20) ++ (100 to 120)).map(_.toString).toDF()) + } + + test("mixed topic-level and partition-level offsets") { + val topic1 = newTopic() + val topic2 = newTopic() + testUtils.createTopic(topic1, partitions = 2) + testUtils.createTopic(topic2, partitions = 2) + testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1)) + testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray, Some(1)) + + // topic1 starts from its earliest offsets, while topic2 is enumerated per partition + val df = createDF(s"$topic1,$topic2", withOptions = Map( + "startingOffsets" -> s"""{"$topic1":"earliest","$topic2":{"0":5,"1":-2}}""", + "endingOffsets" -> s"""{"$topic1":"latest","$topic2":{"0":-1,"1":8}}""")) + checkAnswer(df, ((0 to 19) ++ (105 to 109) ++ (110 to 117)).map(_.toString).toDF()) + } + + test("topic-level offsets with subscribePattern") { + val prefix = newTopic() + val topic1 = s"$prefix-a" + val topic2 = s"$prefix-b" + testUtils.createTopic(topic1, partitions = 2) + testUtils.createTopic(topic2, partitions = 2) + testUtils.sendMessages(topic1, (0 to 9).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic1, (10 to 19).map(_.toString).toArray, Some(1)) + testUtils.sendMessages(topic2, (100 to 109).map(_.toString).toArray, Some(0)) + testUtils.sendMessages(topic2, (110 to 119).map(_.toString).toArray, Some(1)) + + val df = spark + .read + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribePattern", s"$prefix-.*") + .option("startingOffsets", s"""{"$topic1":"earliest","$topic2":{"0":5,"1":-2}}""") + .option("endingOffsets", s"""{"$topic1":"latest","$topic2":{"0":-1,"1":8}}""") + .load() + .selectExpr("CAST(value AS STRING)") + checkAnswer(df, ((0 to 19) ++ (105 to 109) ++ (110 to 117)).map(_.toString).toDF()) + } + + test("topic-level offsets must line up with the topics matched by subscribePattern") { + val prefix = newTopic() + val topic1 = s"$prefix-a" + val topic2 = s"$prefix-b" + testUtils.createTopic(topic1, partitions = 1) + testUtils.createTopic(topic2, partitions = 1) + testUtils.sendMessages(topic1, Array("1"), Some(0)) + testUtils.sendMessages(topic2, Array("2"), Some(0)) + + def readWith(startingOffsets: String): Unit = { + spark + .read + .format("kafka") + .option("kafka.bootstrap.servers", testUtils.brokerAddress) + .option("subscribePattern", s"$prefix-.*") + .option("startingOffsets", startingOffsets) + .load() + .collect() + } + + // A topic matched by the pattern must still be covered by the offsets + val uncovered = intercept[KafkaIllegalStateException] { + readWith(s"""{"$topic1":"earliest"}""") + } + assert(uncovered.getCondition === "KAFKA_START_OFFSET_DOES_NOT_MATCH_ASSIGNED") + + // ... and a topic-level offset for a topic the pattern doesn't match is rejected + val unknown = intercept[KafkaIllegalStateException] { + readWith(s"""{"$topic1":"earliest","$topic2":"earliest","$prefix-ghost":"earliest"}""") + } + assert(unknown.getCondition === "KAFKA_TOPIC_OFFSET_DOES_NOT_MATCH_ASSIGNED") + } + test("default starting and ending offsets with headers") { val topic = newTopic() testUtils.createTopic(topic, partitions = 3) @@ -470,6 +569,9 @@ abstract class KafkaRelationSuiteBase extends SharedSparkSession with KafkaTest testBadOptions("subscribe" -> "t", "startingOffsets" -> startingOffsets)( "startingOffsets for t-0 can't be latest for batch queries on Kafka") + // Now do it with a topic-level start offset indicating latest + testBadOptions("subscribe" -> "t", "startingOffsets" -> """{"t":"latest"}""")( + "startingOffsets for t can't be latest for batch queries on Kafka") // Make sure we catch ending offsets that indicate earliest testBadOptions("endingOffsets" -> "earliest")("ending offset can't be earliest " + @@ -481,6 +583,10 @@ abstract class KafkaRelationSuiteBase extends SharedSparkSession with KafkaTest testBadOptions("subscribe" -> "t", "endingOffsets" -> endingOffsets)( "ending offset for t-0 can't be earliest for batch queries on Kafka") + // Make sure we catch a topic-level ending offset indicating earliest + testBadOptions("subscribe" -> "t", "endingOffsets" -> """{"t":"earliest"}""")( + "ending offset for t can't be earliest for batch queries on Kafka") + // No strategy specified testBadOptions()("options must be specified", "subscribe", "subscribePattern") diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala index b31f6af1c794c..14ba16667ac87 100644 --- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala +++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaTestUtils.scala @@ -28,6 +28,7 @@ import scala.io.Source import scala.jdk.CollectionConverters._ import kafka.log.LogManager +import kafka.security.authorizer.AclAuthorizer import kafka.server.{HostedPartition, KafkaConfig, KafkaServer} import kafka.server.checkpoints.OffsetCheckpointFile import kafka.zk.KafkaZkClient @@ -37,10 +38,14 @@ import org.apache.kafka.clients.CommonClientConfigs import org.apache.kafka.clients.admin._ import org.apache.kafka.clients.producer._ import org.apache.kafka.common.TopicPartition +import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding, AclOperation, AclPermissionType} import org.apache.kafka.common.config.SaslConfigs import org.apache.kafka.common.network.ListenerName import org.apache.kafka.common.requests.FetchRequest +import org.apache.kafka.common.resource.{PatternType, Resource, ResourcePattern, ResourceType} +import org.apache.kafka.common.security.auth.KafkaPrincipal import org.apache.kafka.common.security.auth.SecurityProtocol.{PLAINTEXT, SASL_PLAINTEXT} +import org.apache.kafka.common.security.token.delegation.TokenInformation import org.apache.kafka.common.serialization.StringSerializer import org.apache.kafka.common.utils.Time import org.apache.zookeeper.client.ZKClientConfig @@ -95,7 +100,10 @@ class KafkaTestUtils( private var brokerConf: KafkaConfig = _ private val brokerServiceName = "kafka" - private val clientUser = s"client/$localCanonicalHostName" + private val brokerUser = s"$brokerServiceName/$localCanonicalHostName" + private var brokerKeytabFile: File = _ + private val clientShortName = "client" + private val clientUser = s"$clientShortName/$localCanonicalHostName" private var clientKeytabFile: File = _ // Kafka broker server @@ -137,6 +145,16 @@ class KafkaTestUtils( clientKeytabFile.getAbsolutePath() } + /** The Kafka principal for `user`, i.e. `User:<user>`. */ + def kafkaPrincipal(user: String): KafkaPrincipal = + new KafkaPrincipal(KafkaPrincipal.USER_TYPE, user) + + /** + * The Kafka principal the client keytab authenticates as. The broker derives it from + * `client/<host>@<realm>` with the default `sasl.kerberos.principal.to.local.rules`. + */ + def clientKafkaPrincipal: String = kafkaPrincipal(clientShortName).toString + private def setUpMiniKdc(): Unit = { val kdcDir = Utils.createTempDir() val kdcConf = MiniKdc.createConf() @@ -197,10 +215,9 @@ class KafkaTestUtils( kdc.createPrincipal(zkClientKeytabFile, zkClientUser) logDebug(s"Created keytab file: ${zkClientKeytabFile.getAbsolutePath()}") - val kafkaServerUser = s"kafka/$localCanonicalHostName" - val kafkaServerKeytabFile = new File(baseDir, "kafka.keytab") - kdc.createPrincipal(kafkaServerKeytabFile, kafkaServerUser) - logDebug(s"Created keytab file: ${kafkaServerKeytabFile.getAbsolutePath()}") + brokerKeytabFile = new File(baseDir, "kafka.keytab") + kdc.createPrincipal(brokerKeytabFile, brokerUser) + logDebug(s"Created keytab file: ${brokerKeytabFile.getAbsolutePath()}") clientKeytabFile = new File(baseDir, "client.keytab") kdc.createPrincipal(clientKeytabFile, clientUser) @@ -235,8 +252,8 @@ class KafkaTestUtils( | serviceName="$brokerServiceName" | useKeyTab=true | storeKey=true - | keyTab="${kafkaServerKeytabFile.getAbsolutePath()}" - | principal="$kafkaServerUser@$realm"; + | keyTab="${brokerKeytabFile.getAbsolutePath()}" + | principal="$brokerUser@$realm"; |}; """.stripMargin.trim Files.writeString(file.toPath, content) @@ -276,8 +293,12 @@ class KafkaTestUtils( brokerReady = true } - /** setup the whole embedded servers, including Zookeeper and Kafka brokers */ - def setup(): Unit = { + /** + * Setup the whole embedded servers, including Zookeeper and Kafka brokers. In secure mode, + * `extraAcls` are created together with the client principal's allow-all ACLs. + */ + def setup(extraAcls: Seq[AclBinding] = Nil): Unit = { + assert(secure || extraAcls.isEmpty, "ACLs require the cluster to be set up in secure mode") // Set up a KafkaTestUtils leak detector so that we can see where the leak KafkaTestUtils is // created. val exception = new SparkException("It was created at: ") @@ -299,8 +320,45 @@ class KafkaTestUtils( eventually(timeout(1.minute)) { assert(zkClient.getAllBrokersInCluster.nonEmpty, "Broker was not up in 60 seconds") } + if (secure) { + createAcls(allowAllAcls(clientKafkaPrincipal) ++ extraAcls) + } } + /** All the ACLs granting `principal` unrestricted access to the cluster. */ + def allowAllAcls(principal: String): Seq[AclBinding] = { + val entry = new AccessControlEntry(principal, "*", AclOperation.ALL, AclPermissionType.ALLOW) + def wildcard(resourceType: ResourceType): ResourcePattern = + new ResourcePattern(resourceType, ResourcePattern.WILDCARD_RESOURCE, PatternType.LITERAL) + Seq( + wildcard(ResourceType.TOPIC), + wildcard(ResourceType.GROUP), + wildcard(ResourceType.TRANSACTIONAL_ID), + new ResourcePattern(ResourceType.CLUSTER, Resource.CLUSTER_NAME, PatternType.LITERAL) + ).map(new AclBinding(_, entry)) + } + + /** + * Add ACLs as a super user. The single zk-mode broker updates its authorizer cache before + * completing the request, so the ACLs are enforced once this returns. + */ + def createAcls(bindings: Seq[AclBinding]): Unit = { + assert(secure, "ACLs are only enforced when the cluster is set up in secure mode") + val props = new Properties() + props.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, brokerAddress) + addSaslClientProps(props, brokerKeytabFile, s"$brokerUser@${kdc.getRealm()}") + val superuserAdminClient = AdminClient.create(props) + try { + superuserAdminClient.createAcls(bindings.asJava).all().get() + } finally { + superuserAdminClient.close() + } + } + + /** All delegation tokens the broker knows, as seen by the client principal. */ + def describeDelegationTokens(): Seq[TokenInformation] = + adminClient.describeDelegationToken().delegationTokens().get().asScala.toSeq.map(_.tokenInfo()) + /** Teardown the whole servers, including Kafka broker and Zookeeper */ def teardown(): Unit = { if (leakDetector != null) { @@ -501,6 +559,13 @@ class KafkaTestUtils( props.put("inter.broker.listener.name", "SASL_PLAINTEXT") props.put("delegation.token.master.key", UUID.randomUUID().toString) props.put("sasl.enabled.mechanisms", "GSSAPI,SCRAM-SHA-512") + // Enforce ACLs so that delegation token authorization can be tested. The broker + // authenticates to itself as `kafka/<host>@<realm>`, which the default + // principal-to-local rules map to `User:kafka`, and it must stay unrestricted. + // `setup` grants the client principal full access, so tests which do not care about + // ACLs are unaffected. + props.put("authorizer.class.name", classOf[AclAuthorizer].getName) + props.put("super.users", kafkaPrincipal(brokerServiceName).toString) } props.putAll(withBrokerProps.asJava) @@ -542,13 +607,16 @@ class KafkaTestUtils( private def setAuthenticationConfigIfNeeded(props: Properties): Unit = { if (secure) { - val jaasParams = KafkaTokenUtil.getKeytabJaasParams( - clientKeytabFile.getAbsolutePath, clientPrincipal, brokerServiceName) - props.put(SaslConfigs.SASL_JAAS_CONFIG, jaasParams) - props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, SASL_PLAINTEXT.name) + addSaslClientProps(props, clientKeytabFile, clientPrincipal) } } + private def addSaslClientProps(props: Properties, keytabFile: File, principal: String): Unit = { + props.put(SaslConfigs.SASL_JAAS_CONFIG, KafkaTokenUtil.getKeytabJaasParams( + keytabFile.getAbsolutePath, principal, brokerServiceName)) + props.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, SASL_PLAINTEXT.name) + } + /** Verify topic is deleted in all places, e.g, brokers, zookeeper. */ private def verifyTopicDeletion( topic: String, diff --git a/connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaTokenUtil.scala b/connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaTokenUtil.scala index 54014859bc5e5..9a2325cce3b2b 100644 --- a/connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaTokenUtil.scala +++ b/connector/kafka-0-10-token-provider/src/main/scala/org/apache/spark/kafka010/KafkaTokenUtil.scala @@ -20,6 +20,7 @@ package org.apache.spark.kafka010 import java.{util => ju} import java.time.{Instant, ZoneId} import java.time.format.DateTimeFormatter +import java.util.concurrent.ExecutionException import java.util.regex.Pattern import scala.jdk.CollectionConverters._ @@ -32,14 +33,17 @@ import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenIdenti import org.apache.kafka.clients.CommonClientConfigs import org.apache.kafka.clients.admin.{AdminClient, CreateDelegationTokenOptions} import org.apache.kafka.common.config.{SaslConfigs, SslConfigs} +import org.apache.kafka.common.errors.{DelegationTokenAuthorizationException, UnsupportedVersionException} import org.apache.kafka.common.security.JaasContext +import org.apache.kafka.common.security.auth.KafkaPrincipal import org.apache.kafka.common.security.auth.SecurityProtocol.{SASL_PLAINTEXT, SASL_SSL, SSL} import org.apache.kafka.common.security.scram.ScramLoginModule import org.apache.kafka.common.security.token.delegation.DelegationToken -import org.apache.spark.SparkConf +import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.deploy.SparkHadoopUtil import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.USER_NAME import org.apache.spark.internal.config._ import org.apache.spark.util.{SecurityUtils, Utils} import org.apache.spark.util.Utils.REDACTION_REPLACEMENT_TEXT @@ -65,12 +69,24 @@ object KafkaTokenUtil extends Logging { def obtainToken( sparkConf: SparkConf, clusterConf: KafkaTokenClusterConf): (Token[KafkaDelegationTokenIdentifier], Long) = { - checkProxyUser() - + val options = createDelegationTokenOptions() val adminClient = AdminClient.create(createAdminClientProperties(sparkConf, clusterConf)) - val createDelegationTokenOptions = new CreateDelegationTokenOptions() - val createResult = adminClient.createDelegationToken(createDelegationTokenOptions) - val token = createResult.delegationToken().get() + val token = try { + adminClient.createDelegationToken(options).delegationToken().get() + } catch { + case e: ExecutionException if options.owner().isPresent => + e.getCause match { + case _: UnsupportedVersionException => + throw new SparkException("Obtaining delegation token for proxy user requires " + + "Kafka 3.3.0 or later brokers (KAFKA-6945).", e) + case _: DelegationTokenAuthorizationException => + throw new SparkException("The real user must be granted the CreateTokens operation " + + s"on the ${options.owner().get()} resource.", e) + case _ => throw e + } + } finally { + Utils.closeQuietly(adminClient) + } printToken(token) (new Token[KafkaDelegationTokenIdentifier]( @@ -81,12 +97,26 @@ object KafkaTokenUtil extends Logging { ), token.tokenInfo.expiryTimestamp) } - def checkProxyUser(): Unit = { + /** + * When impersonating, the token is requested with the real user's credentials but owned by the + * proxy user, so connectors authenticate to Kafka as the proxy user. This needs Kafka 3.3.0 or + * later brokers (KAFKA-6945), and the real user must be granted the `CreateTokens` operation on + * the `User:<proxy user>` resource. + * + * Must be invoked under the UGI that should own the token. HadoopDelegationTokenManager + * preserves a proxy UGI only when no principal/keytab is configured (spark-submit forbids + * combining --proxy-user with --principal) or when direct credential providers are used; + * a keytab re-login drops the proxy identity and the token is then owned by the real user. + */ + private[kafka010] def createDelegationTokenOptions(): CreateDelegationTokenOptions = { + val options = new CreateDelegationTokenOptions() val currentUser = UserGroupInformation.getCurrentUser() - // Obtaining delegation token for proxy user is planned but not yet implemented - // See https://issues.apache.org/jira/browse/KAFKA-6945 - require(!SparkHadoopUtil.get.isProxyUser(currentUser), "Obtaining delegation token for proxy " + - "user is not yet supported.") + if (SparkHadoopUtil.get.isProxyUser(currentUser)) { + val owner = currentUser.getUserName + logInfo(log"Obtaining kafka delegation token for proxy user ${MDC(USER_NAME, owner)}.") + options.owner(new KafkaPrincipal(KafkaPrincipal.USER_TYPE, owner)) + } + options } def createAdminClientProperties( @@ -225,13 +255,14 @@ object KafkaTokenUtil extends Logging { private def printToken(token: DelegationToken): Unit = { if (log.isDebugEnabled) { - logDebug("%-15s %-30s %-15s %-25s %-15s %-15s %-15s".format( - "TOKENID", "HMAC", "OWNER", "RENEWERS", "ISSUEDATE", "EXPIRYDATE", "MAXDATE")) + logDebug("%-15s %-30s %-15s %-15s %-25s %-15s %-15s %-15s".format( + "TOKENID", "HMAC", "OWNER", "REQUESTER", "RENEWERS", "ISSUEDATE", "EXPIRYDATE", "MAXDATE")) val tokenInfo = token.tokenInfo - logDebug("%-15s %-15s %-15s %-25s %-15s %-15s %-15s".format( + logDebug("%-15s %-30s %-15s %-15s %-25s %-15s %-15s %-15s".format( tokenInfo.tokenId, REDACTION_REPLACEMENT_TEXT, tokenInfo.owner, + tokenInfo.tokenRequester, tokenInfo.renewersAsString, DATE_TIME_FORMATTER.format(Instant.ofEpochMilli(tokenInfo.issueTimestamp)), DATE_TIME_FORMATTER.format(Instant.ofEpochMilli(tokenInfo.expiryTimestamp)), diff --git a/connector/kafka-0-10-token-provider/src/test/scala/org/apache/spark/kafka010/KafkaTokenUtilSuite.scala b/connector/kafka-0-10-token-provider/src/test/scala/org/apache/spark/kafka010/KafkaTokenUtilSuite.scala index 8a606a1adc767..c7d6fa647db99 100644 --- a/connector/kafka-0-10-token-provider/src/test/scala/org/apache/spark/kafka010/KafkaTokenUtilSuite.scala +++ b/connector/kafka-0-10-token-provider/src/test/scala/org/apache/spark/kafka010/KafkaTokenUtilSuite.scala @@ -26,6 +26,7 @@ import org.apache.hadoop.io.Text import org.apache.hadoop.security.UserGroupInformation import org.apache.kafka.clients.CommonClientConfigs import org.apache.kafka.common.config.{SaslConfigs, SslConfigs} +import org.apache.kafka.common.security.auth.KafkaPrincipal import org.apache.kafka.common.security.auth.SecurityProtocol.{SASL_PLAINTEXT, SASL_SSL, SSL} import org.apache.spark.{SparkConf, SparkFunSuite} @@ -39,16 +40,24 @@ class KafkaTokenUtilSuite extends SparkFunSuite with KafkaDelegationTokenTest { sparkConf = new SparkConf() } - test("checkProxyUser with proxy current user should throw exception") { + test("SPARK-28173: createDelegationTokenOptions without proxy user should not set owner") { + UserGroupInformation.createUserForTesting("realUser", Array()).doAs( + new PrivilegedExceptionAction[Unit]() { + override def run(): Unit = { + assert(!KafkaTokenUtil.createDelegationTokenOptions().owner().isPresent) + } + } + ) + } + + test("SPARK-28173: createDelegationTokenOptions with proxy user should set owner") { val realUser = UserGroupInformation.createUserForTesting("realUser", Array()) UserGroupInformation.createProxyUserForTesting("proxyUser", realUser, Array()).doAs( new PrivilegedExceptionAction[Unit]() { override def run(): Unit = { - val thrown = intercept[IllegalArgumentException] { - KafkaTokenUtil.checkProxyUser() - } - assert(thrown.getMessage contains - "Obtaining delegation token for proxy user is not yet supported.") + val owner = KafkaTokenUtil.createDelegationTokenOptions().owner() + assert(owner.isPresent) + assert(owner.get() === new KafkaPrincipal(KafkaPrincipal.USER_TYPE, "proxyUser")) } } ) diff --git a/connector/kinesis-asl/src/main/python/examples/streaming/kinesis_wordcount_asl.py b/connector/kinesis-asl/src/main/python/examples/streaming/kinesis_wordcount_asl.py index 53a6b69dc93a8..73b942ece024e 100644 --- a/connector/kinesis-asl/src/main/python/examples/streaming/kinesis_wordcount_asl.py +++ b/connector/kinesis-asl/src/main/python/examples/streaming/kinesis_wordcount_asl.py @@ -59,7 +59,7 @@ from pyspark import SparkContext from pyspark.streaming import StreamingContext -from pyspark.streaming.kinesis import KinesisUtils, InitialPositionInStream +from pyspark.streaming.kinesis import InitialPositionInStream, KinesisUtils if __name__ == "__main__": if len(sys.argv) != 5: diff --git a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufUtils.scala b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufUtils.scala index 47898f73d0165..533311e311ea7 100644 --- a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufUtils.scala +++ b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/utils/ProtobufUtils.scala @@ -18,11 +18,14 @@ package org.apache.spark.sql.protobuf.utils import java.util.Locale +import java.util.concurrent.{ExecutionException, TimeUnit} import scala.collection.mutable import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ +import com.google.common.hash.Hashing +import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException} import com.google.protobuf.{DescriptorProtos, Descriptors, DynamicMessage, ExtensionRegistry, InvalidProtocolBufferException, Message} import com.google.protobuf.DescriptorProtos.{FileDescriptorProto, FileDescriptorSet} import com.google.protobuf.Descriptors.{Descriptor, FieldDescriptor} @@ -33,7 +36,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.spark.util.Utils +import org.apache.spark.util.{NonFateSharingCache, Utils} private[sql] object ProtobufUtils extends Logging { @@ -238,10 +241,37 @@ private[sql] object ProtobufUtils extends Logging { DescriptorWithExtensions(descriptor, ExtensionRegistry.getEmptyRegistry, Map.empty) } + // Parsing a FileDescriptorSet materializes a potentially-large FileDescriptor graph, and each + // task instance would otherwise parse the same bytes independently. + + // Read fresh each call so size 0 is an immediate kill-switch; the built cache's bound is fixed. + private def isDescriptorCacheEnabled: Boolean = + SQLConf.get.getConf(SQLConf.PROTOBUF_DESCRIPTOR_CACHE_SIZE) > 0 + + private def descriptorCacheMaxSize: Long = + SQLConf.get.getConf(SQLConf.PROTOBUF_DESCRIPTOR_CACHE_SIZE).toLong + + private type FileDescriptors = List[Descriptors.FileDescriptor] + + // Keyed on a content hash rather than the bytes, to avoid pinning the descriptor arrays. + // NonFateSharingCache (SPARK-43300) prevents a cancelled task's failed load from failing the + // other tasks blocked on the same key. The loader is passed per get() rather than as a + // CacheLoader because the key is the hash, not the bytes needed to parse; this overload also + // keeps the shaded Guava Cache type out of the signature (SPARK-44064). + private lazy val fileDescriptorCache: NonFateSharingCache[String, FileDescriptors] = + NonFateSharingCache(descriptorCacheMaxSize, 0, TimeUnit.SECONDS) + + private[protobuf] def clearDescriptorCacheForTesting(): Unit = + fileDescriptorCache.invalidateAll() + + private[protobuf] def fileDescriptorCacheSizeForTesting(): Long = fileDescriptorCache.size() + + private def contentHash(bytes: Array[Byte]): String = + Hashing.sha256().hashBytes(bytes).toString + def buildDescriptorFromFDS( messageName: String, binaryFileDescriptorSet: Array[Byte]): DescriptorWithExtensions = { - // Find the first message descriptor that matches the name. val fileDescriptors = parseFileDescriptorSet(binaryFileDescriptorSet) val descriptor = fileDescriptors .flatMap { fileDesc => @@ -262,7 +292,20 @@ private[sql] object ProtobufUtils extends Logging { } } - private def parseFileDescriptorSet(bytes: Array[Byte]): List[Descriptors.FileDescriptor] = { + private def parseFileDescriptorSet(bytes: Array[Byte]): FileDescriptors = { + if (!isDescriptorCacheEnabled) { + return doParseFileDescriptorSet(bytes) + } + try { + fileDescriptorCache.get(contentHash(bytes), () => doParseFileDescriptorSet(bytes)) + } catch { + case e: UncheckedExecutionException if e.getCause != null => throw e.getCause + case e: ExecutionException if e.getCause != null => throw e.getCause + case e: ExecutionError if e.getCause != null => throw e.getCause + } + } + + private def doParseFileDescriptorSet(bytes: Array[Byte]): FileDescriptors = { var fileDescriptorSet: DescriptorProtos.FileDescriptorSet = null try { fileDescriptorSet = DescriptorProtos.FileDescriptorSet.parseFrom(bytes) diff --git a/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala b/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala index 5c875989e7144..c25b4b264c65c 100644 --- a/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala +++ b/connector/protobuf/src/test/scala/org/apache/spark/sql/protobuf/ProtobufFunctionsSuite.scala @@ -27,6 +27,7 @@ import org.json4s.jackson.JsonMethods import org.apache.spark.sql.{AnalysisException, Column, DataFrame, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.functions.{array, lit, map, struct, typedLit} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.protobuf.protos.Proto2Messages.Proto2AllTypes import org.apache.spark.sql.protobuf.protos.SimpleMessageProtos._ import org.apache.spark.sql.protobuf.protos.SimpleMessageProtos.SimpleMessageRepeated.NestedEnum @@ -2368,4 +2369,84 @@ class ProtobufFunctionsSuite extends SharedSparkSession with ProtobufTestBase assert(expectedDf.schema === fromProtoDf.schema) checkAnswer(fromProtoDf, expectedDf) } + + test("descriptor cache: repeated builds on the same bytes share the cached parse") { + ProtobufUtils.clearDescriptorCacheForTesting() + val first = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + val second = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + // Same descriptor instance (from the shared parse), not a rebuilt copy. + assert(first.descriptor eq second.descriptor) + val repeated = ProtobufUtils.buildDescriptor("RepeatedMessage", Some(testFileDesc)) + assert(repeated.descriptor ne first.descriptor) + // Distinct message names on the same bytes still share one parse. + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 1) + } + + test("descriptor cache: distinct bytes are parsed and cached separately") { + ProtobufUtils.clearDescriptorCacheForTesting() + // Derive a second, genuinely distinct descriptor set from the same source rather than relying + // on testFileDesc vs proto2FileDesc being different bytes: under SBT both point at one combined + // descriptor file, so they hash to the same cache key. proto2_messages.proto has no imports, so + // its single-file set parses standalone and is distinct from the full combined set either way. + val proto2Standalone = descriptorSetWithoutImports(proto2FileDesc, "FoobarWithRequiredFieldBar") + val basic = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + val proto2 = ProtobufUtils.buildDescriptor("FoobarWithRequiredFieldBar", Some(proto2Standalone)) + assert(proto2.descriptor ne basic.descriptor) + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 2) + } + + test("descriptor cache: buildTypeRegistry shares the parse with buildDescriptor") { + ProtobufUtils.clearDescriptorCacheForTesting() + ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 1) + // The other entry point reuses the same parse rather than adding an entry. + ProtobufUtils.buildTypeRegistry(testFileDesc) + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 1) + } + + test("descriptor cache: the extensions-enabled flag is honored per call") { + ProtobufUtils.clearDescriptorCacheForTesting() + val extConf = SQLConf.PROTOBUF_EXTENSIONS_SUPPORT_ENABLED + // The flag is read per build, so a shared cached parse still yields the flag-correct result. + withSQLConf(extConf.key -> "false") { + val disabled = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + assert(disabled.extensionRegistry eq com.google.protobuf.ExtensionRegistry.getEmptyRegistry) + } + withSQLConf(extConf.key -> "true") { + val enabled = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + assert(enabled.extensionRegistry ne com.google.protobuf.ExtensionRegistry.getEmptyRegistry) + } + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 1) + } + + test("descriptor cache: disabled (size 0) reparses every time and caches nothing") { + ProtobufUtils.clearDescriptorCacheForTesting() + withSQLConf(SQLConf.PROTOBUF_DESCRIPTOR_CACHE_SIZE.key -> "0") { + val first = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + val second = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + // Each call reparses, so the descriptors come from different graphs. + assert(first.descriptor ne second.descriptor) + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 0) + } + } + + test("descriptor cache: an unknown message name surfaces the domain exception") { + ProtobufUtils.clearDescriptorCacheForTesting() + // The parse succeeds; only the message lookup fails, so the parse stays cached below. + intercept[AnalysisException] { + ProtobufUtils.buildDescriptor("NoSuchMessage", Some(testFileDesc)) + } + val good = ProtobufUtils.buildDescriptor("BasicMessage", Some(testFileDesc)) + assert(good.descriptor != null) + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 1) + } + + test("descriptor cache: a failed parse is not cached") { + ProtobufUtils.clearDescriptorCacheForTesting() + intercept[AnalysisException] { + ProtobufUtils.buildTypeRegistry(Array[Byte](1, 2, 3, 4)) + } + // A failed parse leaves the cache empty. + assert(ProtobufUtils.fileDescriptorCacheSizeForTesting() == 0) + } } diff --git a/core/benchmarks/OpenHashMapBenchmark-jdk21-results.txt b/core/benchmarks/OpenHashMapBenchmark-jdk21-results.txt new file mode 100644 index 0000000000000..c98d2b88560e3 --- /dev/null +++ b/core/benchmarks/OpenHashMapBenchmark-jdk21-results.txt @@ -0,0 +1,26 @@ +================================================================================================ +OpenHashMap vs java.util.HashMap +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Insert 1000000 distinct String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +OpenHashMap 262 280 15 3.8 262.0 1.0X +java.util.HashMap 48 84 57 21.0 47.7 5.5X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Aggregate 5000000 ops on 1000000 String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +---------------------------------------------------------------------------------------------------------------------------- +OpenHashMap changeValue 858 885 28 5.8 171.7 1.0X +java.util.HashMap merge 772 811 60 6.5 154.5 1.1X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Look up 1000000 String keys in random order: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +OpenHashMap 101 106 3 9.9 100.9 1.0X +java.util.HashMap 48 53 3 20.7 48.3 2.1X + + diff --git a/core/benchmarks/OpenHashMapBenchmark-jdk25-results.txt b/core/benchmarks/OpenHashMapBenchmark-jdk25-results.txt new file mode 100644 index 0000000000000..f21fd323a84f4 --- /dev/null +++ b/core/benchmarks/OpenHashMapBenchmark-jdk25-results.txt @@ -0,0 +1,26 @@ +================================================================================================ +OpenHashMap vs java.util.HashMap +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Insert 1000000 distinct String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +OpenHashMap 326 345 19 3.1 326.3 1.0X +java.util.HashMap 50 96 67 19.8 50.5 6.5X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Aggregate 5000000 ops on 1000000 String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +---------------------------------------------------------------------------------------------------------------------------- +OpenHashMap changeValue 1103 1208 148 4.5 220.6 1.0X +java.util.HashMap merge 1273 1286 19 3.9 254.5 0.9X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Look up 1000000 String keys in random order: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +OpenHashMap 120 136 15 8.3 120.0 1.0X +java.util.HashMap 58 64 3 17.1 58.4 2.1X + + diff --git a/core/benchmarks/OpenHashMapBenchmark-results.txt b/core/benchmarks/OpenHashMapBenchmark-results.txt new file mode 100644 index 0000000000000..6bdd08f4ca197 --- /dev/null +++ b/core/benchmarks/OpenHashMapBenchmark-results.txt @@ -0,0 +1,26 @@ +================================================================================================ +OpenHashMap vs java.util.HashMap +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Insert 1000000 distinct String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +OpenHashMap 228 243 13 4.4 227.9 1.0X +java.util.HashMap 50 88 53 20.1 49.8 4.6X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Aggregate 5000000 ops on 1000000 String keys: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +---------------------------------------------------------------------------------------------------------------------------- +OpenHashMap changeValue 772 778 7 6.5 154.3 1.0X +java.util.HashMap merge 740 763 37 6.8 148.1 1.0X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +Look up 1000000 String keys in random order: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +OpenHashMap 109 111 2 9.2 108.5 1.0X +java.util.HashMap 52 54 1 19.1 52.3 2.1X + + diff --git a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java index d2ee6269c134e..63ee137561249 100644 --- a/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java +++ b/core/src/main/java/org/apache/spark/memory/MemoryConsumer.java @@ -162,7 +162,9 @@ private void throwOom(final MemoryBlock page, final long required) { got = page.size(); taskMemoryManager.freePage(page, this); } - taskMemoryManager.showMemoryUsage(); - throw SparkCoreErrors.outOfMemoryError(required, got); + // Log the full breakdown and attach the bounded one to the error from a single snapshot, so the + // executor logs and the driver/UI error message describe the same instant and cannot disagree. + String consumerBreakdown = taskMemoryManager.logMemoryUsageAndGetBreakdown(); + throw SparkCoreErrors.outOfMemoryError(required, got, consumerBreakdown); } } diff --git a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java index 099d95c6f9de5..cefebcffa2fc8 100644 --- a/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java +++ b/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java @@ -449,36 +449,162 @@ public void releaseExecutionMemory(long size, MemoryConsumer consumer) { } /** - * Dump the memory usage of all consumers. + * A point-in-time snapshot of this task's execution-memory usage: each consumer that is holding + * memory (largest first) paired with its used bytes, plus the bytes not attributable to any + * specific consumer. Rendering the executor-log dump and the error-message breakdown from a + * single snapshot is what lets the two describe the same instant; see + * {@link #logMemoryUsageAndGetBreakdown()}. */ - public void showMemoryUsage() { - logger.info("Memory used in task {}", - MDC.of(LogKeys.TASK_ATTEMPT_ID, taskAttemptId)); + private static final class MemoryUsageSnapshot { + private final List<Map.Entry<MemoryConsumer, Long>> consumerUsages; + private final long memoryNotAccountedFor; + + MemoryUsageSnapshot( + List<Map.Entry<MemoryConsumer, Long>> consumerUsages, long memoryNotAccountedFor) { + this.consumerUsages = consumerUsages; + this.memoryNotAccountedFor = memoryNotAccountedFor; + } + } + + /** + * Snapshot the per-consumer memory usage once, holding the monitor for the whole read. + * <p> + * {@link MemoryConsumer#getUsed()} reads an {@code AtomicLong} that is not guarded by this + * monitor, so callers must not re-read it while sorting or rendering: doing so could observe + * changing values and trip {@code TimSort}'s "Comparison method violates its general contract!" + * check, masking the OOM we are about to report with an unrelated failure. Consumers are returned + * largest first, since the biggest consumers are the most likely culprits of an OOM. + */ + private MemoryUsageSnapshot snapshotMemoryUsage() { + List<Map.Entry<MemoryConsumer, Long>> consumerUsages = new ArrayList<>(); + long memoryAccountedForByConsumers = 0; + long memoryNotAccountedFor; synchronized (this) { - long memoryAccountedForByConsumers = 0; - for (MemoryConsumer c: consumers) { + for (MemoryConsumer c : consumers) { long totalMemUsage = c.getUsed(); - memoryAccountedForByConsumers += totalMemUsage; if (totalMemUsage > 0) { - logger.info("Acquired by {}: {}", - MDC.of(LogKeys.MEMORY_CONSUMER, c), - MDC.of(LogKeys.MEMORY_SIZE, Utils.bytesToString(totalMemUsage))); + memoryAccountedForByConsumers += totalMemUsage; + consumerUsages.add(new AbstractMap.SimpleEntry<>(c, totalMemUsage)); } } - long memoryNotAccountedFor = + memoryNotAccountedFor = memoryManager.getExecutionMemoryUsageForTask(taskAttemptId) - memoryAccountedForByConsumers; - logger.info( - "{} bytes of memory were used by task {} but are not associated with specific consumers", - MDC.of(LogKeys.MEMORY_SIZE, memoryNotAccountedFor), - MDC.of(LogKeys.TASK_ATTEMPT_ID, taskAttemptId)); - logger.info( - "{} bytes of memory are used for execution " + - "and {} bytes of memory are used for storage " + - "and {} bytes of unmanaged memory are used", - MDC.of(LogKeys.EXECUTION_MEMORY_SIZE, memoryManager.executionMemoryUsed()), - MDC.of(LogKeys.STORAGE_MEMORY_SIZE, memoryManager.storageMemoryUsed()), - MDC.of(LogKeys.MEMORY_SIZE, UnifiedMemoryManager$.MODULE$.getUnmanagedMemoryUsed())); } + consumerUsages.sort(Map.Entry.<MemoryConsumer, Long>comparingByValue().reversed()); + return new MemoryUsageSnapshot(consumerUsages, memoryNotAccountedFor); + } + + /** + * Dump the given memory-usage snapshot to the executor logs, one line per consumer (uncapped). + */ + private void logMemoryUsage(MemoryUsageSnapshot snapshot) { + logger.info("Memory used in task {}", + MDC.of(LogKeys.TASK_ATTEMPT_ID, taskAttemptId)); + for (Map.Entry<MemoryConsumer, Long> usage : snapshot.consumerUsages) { + logger.info("Acquired by {}: {}", + MDC.of(LogKeys.MEMORY_CONSUMER, usage.getKey()), + MDC.of(LogKeys.MEMORY_SIZE, Utils.bytesToString(usage.getValue()))); + } + logger.info( + "{} bytes of memory were used by task {} but are not associated with specific consumers", + MDC.of(LogKeys.MEMORY_SIZE, snapshot.memoryNotAccountedFor), + MDC.of(LogKeys.TASK_ATTEMPT_ID, taskAttemptId)); + logger.info( + "{} bytes of memory are used for execution " + + "and {} bytes of memory are used for storage " + + "and {} bytes of unmanaged memory are used", + MDC.of(LogKeys.EXECUTION_MEMORY_SIZE, memoryManager.executionMemoryUsed()), + MDC.of(LogKeys.STORAGE_MEMORY_SIZE, memoryManager.storageMemoryUsed()), + MDC.of(LogKeys.MEMORY_SIZE, UnifiedMemoryManager$.MODULE$.getUnmanagedMemoryUsed())); + } + + /** + * Render the given snapshot as the compact, bounded breakdown embedded in the + * {@code UNABLE_TO_ACQUIRE_MEMORY} error. Returns an empty string when there is nothing to + * show -- that is, when the snapshot has neither attributed nor unattributed memory to report + * (or when the breakdown is disabled by a limit of 0) -- so callers can append it + * unconditionally. + */ + private String renderConsumerBreakdown(MemoryUsageSnapshot snapshot) { + // Bound the message that travels to the driver and the UI. The largest consumers -- the most + // likely culprits -- are listed individually up to this limit; the rest are collapsed into a + // single summary line so total byte accounting is preserved without unbounded noise. The full, + // uncapped breakdown is still available in the executor logs via logMemoryUsage(). A limit of + // 0 omits the breakdown from the error message entirely. + int limit = memoryManager.oomErrorConsumerBreakdownLimit(); + if (limit == 0) { + return ""; + } + List<Map.Entry<MemoryConsumer, Long>> usages = snapshot.consumerUsages; + StringBuilder sb = new StringBuilder(); + int shown = Math.min(limit, usages.size()); + for (int i = 0; i < shown; i++) { + Map.Entry<MemoryConsumer, Long> usage = usages.get(i); + sb.append("\n ").append(usage.getKey()).append(": ") + .append(Utils.bytesToString(usage.getValue())); + } + if (usages.size() > shown) { + long remainingBytes = 0; + for (int i = shown; i < usages.size(); i++) { + remainingBytes = Math.addExact(remainingBytes, usages.get(i).getValue()); + } + sb.append("\n (").append(usages.size() - shown).append(" more consumers): ") + .append(Utils.bytesToString(remainingBytes)); + } + if (snapshot.memoryNotAccountedFor > 0) { + sb.append("\n (not attributed to a specific consumer): ") + .append(Utils.bytesToString(snapshot.memoryNotAccountedFor)); + } + if (sb.length() == 0) { + return ""; + } + return "\nMemory used by task " + taskAttemptId + " grouped by consumer:" + sb; + } + + /** + * Dump the memory usage of all consumers to the executor logs. + */ + public void showMemoryUsage() { + logMemoryUsage(snapshotMemoryUsage()); + } + + /** + * Build a compact, human-readable breakdown of this task's execution-memory usage grouped by + * {@link MemoryConsumer}, with the biggest consumers first, followed by the bytes that are not + * attributable to any specific consumer. + * <p> + * The returned string is meant to be embedded in the {@code UNABLE_TO_ACQUIRE_MEMORY} error so + * that the consumers competing for memory at the moment of failure travel with the task failure + * reason all the way to the driver and the Spark UI. Returns an empty string when there is + * nothing to report -- no consumer is holding memory <i>and</i> there is no unattributed + * memory, or the breakdown is disabled by a limit of 0 -- so the caller can append it + * unconditionally. Otherwise, a task holding only unattributed memory still gets a breakdown, + * consisting of the unattributed line alone. + */ + public String getMemoryConsumptionBreakdown() { + // Skip the snapshot entirely when the breakdown is disabled: renderConsumerBreakdown would + // discard it anyway. The combined logging path (logMemoryUsageAndGetBreakdown) still needs the + // snapshot to write the executor-log dump, so it keeps snapshotting and relies on the + // renderer's own limit-0 check. + if (memoryManager.oomErrorConsumerBreakdownLimit() == 0) { + return ""; + } + return renderConsumerBreakdown(snapshotMemoryUsage()); + } + + /** + * Snapshot this task's memory usage once, write the full (uncapped) breakdown to the executor + * logs, and return the bounded breakdown to embed in the {@code UNABLE_TO_ACQUIRE_MEMORY} error. + * <p> + * Taking a single snapshot for both outputs is what guarantees the log dump and the error message + * describe the same instant and cannot disagree; this is the method the OOM path should call + * rather than invoking {@link #showMemoryUsage()} and {@link #getMemoryConsumptionBreakdown()} + * separately. + */ + public String logMemoryUsageAndGetBreakdown() { + MemoryUsageSnapshot snapshot = snapshotMemoryUsage(); + logMemoryUsage(snapshot); + return renderConsumerBreakdown(snapshot); } /** diff --git a/core/src/main/java/org/apache/spark/security/CredentialProvider.java b/core/src/main/java/org/apache/spark/security/CredentialProvider.java index 610962a594945..2d3e51f8ca101 100644 --- a/core/src/main/java/org/apache/spark/security/CredentialProvider.java +++ b/core/src/main/java/org/apache/spark/security/CredentialProvider.java @@ -36,10 +36,10 @@ * Implementations must be thread-safe: {@code resolve()} may be called concurrently from * multiple threads after {@code init()} completes. * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi -public interface CredentialProvider { +public interface CredentialProvider extends AutoCloseable { /** * Initializes this provider with configuration properties. @@ -58,7 +58,7 @@ public interface CredentialProvider { * * @param conf Spark configuration properties scoped to {@code spark.security.oidc.*} * keys (must not be null) - * @since 4.3.0 + * @since 4.4.0 */ void init(Map<String, String> conf); @@ -69,7 +69,7 @@ public interface CredentialProvider { * set must be non-empty and stable across calls. * * @return a non-empty set of supported scheme names - * @since 4.3.0 + * @since 4.4.0 */ Set<String> supportedSchemes(); @@ -85,7 +85,7 @@ public interface CredentialProvider { * @param target the target URI for which credentials are requested (must not be null) * @return a short-lived service credential for the target * @throws CredentialResolutionException if the credential exchange fails - * @since 4.3.0 + * @since 4.4.0 */ ServiceCredential resolve(UserContext user, URI target) throws CredentialResolutionException; @@ -96,9 +96,60 @@ public interface CredentialProvider { * The default is 15 minutes. * * @return the suggested credential TTL (never null) - * @since 4.3.0 + * @since 4.4.0 */ default Duration suggestedTtl() { return Duration.ofMinutes(15); } + + /** + * Returns additional Spark configuration properties that should be set when this + * provider is active. + * <p> + * This method is called after {@link #init(Map)} and a successful + * {@link #resolve(UserContext, URI)} invocation. Implementations may + * assume that provider state is fully initialized when this is called. + * <p> + * The credential management layer applies these entries to {@code SparkConf} after + * successful startup, only if the user has not already set them explicitly. This + * allows provider modules to declare executor-side wiring (e.g., the Hadoop + * credentials provider class for a particular filesystem scheme) without requiring + * core to have vendor-specific knowledge. + * <p> + * Keys must use the {@code spark.} prefix to be effective (SparkConf convention). + * Keys with the {@code spark.hadoop.} prefix are propagated to executor-side + * Hadoop {@code Configuration} with the prefix stripped. Other {@code spark.*} + * keys are applied as Spark-internal configuration. + * <p> + * The default implementation returns an empty map (no additional properties). + * + * @return an unmodifiable map of property key-value pairs (never null). + * Keys and values within the map must not be {@code null}. + * @since 4.4.0 + */ + default Map<String, String> additionalSparkProperties() { + return Map.of(); + } + + /** + * Releases any resources held by this provider (e.g., HTTP clients, connection pools). + * <p> + * Called by the credential management layer during shutdown. The default implementation + * is a no-op; providers that allocate long-lived resources in {@link #init(Map)} should + * override this method to clean them up. + * <p> + * Shutdown interrupts the renewal thread and waits a bounded time for in-flight calls to + * complete. If the wait times out or the shutdown thread is interrupted, {@code close()} may + * be invoked while another thread is still executing {@link #resolve(UserContext, URI)}. + * Implementations must tolerate a concurrent or subsequent {@code resolve()} failing after + * resources have been released, and {@code close()} itself must not block indefinitely. + * <p> + * Implementations that do not throw checked exceptions may narrow the {@code throws} + * clause in their override (e.g., declare {@code close()} with no {@code throws} or + * with a more specific exception type). + * + * @since 4.4.0 + */ + @Override + default void close() throws Exception {} } diff --git a/core/src/main/java/org/apache/spark/security/CredentialProviderLoader.java b/core/src/main/java/org/apache/spark/security/CredentialProviderLoader.java index 9d6e9e6fedf08..918d11f5b7b7a 100644 --- a/core/src/main/java/org/apache/spark/security/CredentialProviderLoader.java +++ b/core/src/main/java/org/apache/spark/security/CredentialProviderLoader.java @@ -40,8 +40,9 @@ * Discovers {@link CredentialProvider} implementations via {@link ServiceLoader} and selects * the appropriate provider for a given URI scheme using Binding Policy A (explicit selection). * <p> - * Provider discovery happens once (lazily on first call) and the list is cached. Each provider - * is initialized exactly once per provider instance via {@link CredentialProvider#init(Map)} + * Provider discovery happens once per loader (lazily on first call) and the list is cached. + * Each provider is initialized exactly once per provider instance via + * {@link CredentialProvider#init(Map)} * with the configuration from the first call that selects it (first-conf-wins semantics); * subsequent resolutions reuse the already-initialized instance without re-calling {@code init}. * <p> @@ -53,13 +54,13 @@ * is unset (or empty) do the count-based rules apply: a single candidate is auto-selected; * multiple candidates produce an ambiguity error; no candidates produce {@code Optional.empty()}. * <p> - * <b>Thread-safety:</b> This class uses synchronized access to the cached provider list and + * <b>Thread-safety:</b> Each loader uses synchronized access to its cached provider list and * initialization tracking, and callers may invoke {@link #providerFor(String, Map)} from * multiple threads. A provider instance is cached and shared across callers; per the * {@link CredentialProvider} contract, implementations must be thread-safe, so a returned * instance may be used concurrently. * - * @since 4.3.0 + * @since 4.4.0 */ @Private public final class CredentialProviderLoader { @@ -78,18 +79,19 @@ public final class CredentialProviderLoader { */ private static final String OIDC_CONF_PREFIX = "spark.security.oidc."; - private static volatile List<CredentialProvider> cachedProviders; + private volatile List<CredentialProvider> cachedProviders; /** - * Tracks which provider instances have already been initialized. Guarded by the class lock. + * Tracks which provider instances have already been initialized. Guarded by this loader's lock. * Uses identity semantics (reference equality) to handle multiple provider instances correctly. */ - private static final Set<CredentialProvider> initializedProviders = + private final Set<CredentialProvider> initializedProviders = Collections.newSetFromMap(new IdentityHashMap<>()); - private CredentialProviderLoader() { - // utility class - } + /** Guarded by this loader's lock. Prevents providers from being reinitialized after shutdown. */ + private boolean providersClosed; + + public CredentialProviderLoader() {} /** * Returns the {@link CredentialProvider} for the given URI scheme, applying Binding Policy A: @@ -118,14 +120,20 @@ private CredentialProviderLoader() { * @return the selected provider, or empty if no provider supports the scheme * @throws IllegalArgumentException if explicit selection names an unknown or non-supporting * class, or if multiple candidates exist without explicit selection - * @throws IllegalStateException if a provider returns null from {@code supportedSchemes()} + * @throws IllegalStateException if providers have already been closed or if a provider returns + * null from {@code supportedSchemes()} */ - public static Optional<CredentialProvider> providerFor(String scheme, Map<String, String> conf) { + public Optional<CredentialProvider> providerFor(String scheme, Map<String, String> conf) { Objects.requireNonNull(scheme, "scheme must not be null"); Objects.requireNonNull(conf, "conf must not be null"); if (scheme.isEmpty()) { throw new IllegalArgumentException("scheme must not be empty"); } + synchronized (this) { + if (providersClosed) { + throw new IllegalStateException("Credential providers have already been closed"); + } + } String normalizedScheme = scheme.toLowerCase(Locale.ROOT); List<CredentialProvider> providers = getProviders(); @@ -179,7 +187,11 @@ public static Optional<CredentialProvider> providerFor(String scheme, Map<String // precedent of DataSourceV2Utils.extractSessionConfigs() which scopes configuration // to a specific prefix. We keep the full key (unlike extractSessionConfigs which // strips the prefix) so providers can distinguish sub-keys unambiguously. - synchronized (CredentialProviderLoader.class) { + synchronized (this) { + // Re-check under the initialization lock in case closeAll() ran during selection. + if (providersClosed) { + throw new IllegalStateException("Credential providers have already been closed"); + } if (!initializedProviders.contains(selected)) { Map<String, String> filteredConf = new HashMap<>(); for (Map.Entry<String, String> entry : conf.entrySet()) { @@ -197,10 +209,10 @@ public static Optional<CredentialProvider> providerFor(String scheme, Map<String /** * Returns the cached list of discovered providers, loading them on first access. */ - private static List<CredentialProvider> getProviders() { + private List<CredentialProvider> getProviders() { List<CredentialProvider> providers = cachedProviders; if (providers == null) { - synchronized (CredentialProviderLoader.class) { + synchronized (this) { providers = cachedProviders; if (providers == null) { providers = loadProviders(); @@ -211,7 +223,7 @@ private static List<CredentialProvider> getProviders() { return providers; } - private static List<CredentialProvider> loadProviders() { + private List<CredentialProvider> loadProviders() { ClassLoader cl = Thread.currentThread().getContextClassLoader(); if (cl == null) { cl = CredentialProvider.class.getClassLoader(); @@ -236,8 +248,14 @@ private static List<CredentialProvider> loadProviders() { * (e.g., {@code spark.security.oidc.provider.<scheme>}) is provided. * * @return a set of all supported scheme names (lowercased), possibly empty + * @throws IllegalStateException if providers have already been closed */ - public static Set<String> discoverAllSchemes() { + public Set<String> discoverAllSchemes() { + synchronized (this) { + if (providersClosed) { + throw new IllegalStateException("Credential providers have already been closed"); + } + } List<CredentialProvider> providers = getProviders(); Set<String> schemes = new HashSet<>(); for (CredentialProvider provider : providers) { @@ -251,24 +269,71 @@ public static Set<String> discoverAllSchemes() { return schemes; } + /** + * Closes all initialized providers, suppressing individual close exceptions. + * <p> + * This method iterates over all providers that have been initialized via + * {@link CredentialProvider#init(Map)} and calls {@link CredentialProvider#close()} + * on each. The first exception is retained, later exceptions are suppressed onto it, and the + * first exception is rethrown after all providers have been attempted. + * <p> + * After shutdown begins, subsequent {@link #providerFor} calls fail rather than + * re-initializing a cached provider whose resources have already been released. + * <p> + * <b>Contract:</b> {@code close()} implementations must not call back into + * {@code CredentialProviderLoader} methods (e.g., {@code providerFor}). + * + * @throws Exception if one or more providers threw during close + */ + public void closeAll() throws Exception { + List<CredentialProvider> toClose; + synchronized (this) { + // Copy and clear under the lock to prevent double-close if closeAll() is called + // again concurrently, and to avoid ConcurrentModificationException. + providersClosed = true; + toClose = new ArrayList<>(initializedProviders); + initializedProviders.clear(); + } + // Close outside the lock so a slow or blocking close() cannot stall + // providerFor() callers or deadlock against them. + Exception firstException = null; + for (CredentialProvider provider : toClose) { + try { + provider.close(); + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } else { + firstException.addSuppressed(e); + } + } + } + if (firstException != null) { + throw firstException; + } + } + /** * Resets the cached provider list and initialization tracking. Intended for testing only. */ @VisibleForTesting - public static void resetForTesting() { - synchronized (CredentialProviderLoader.class) { + public void resetForTesting() { + synchronized (this) { cachedProviders = null; initializedProviders.clear(); + providersClosed = false; } } /** * Overrides the cached provider list for testing. Intended for testing only. */ - static void setProvidersForTesting(List<CredentialProvider> providers) { - synchronized (CredentialProviderLoader.class) { + @VisibleForTesting + void setProvidersForTesting(List<CredentialProvider> providers) { + synchronized (this) { cachedProviders = providers; initializedProviders.clear(); + providersClosed = false; } } } diff --git a/core/src/main/java/org/apache/spark/security/CredentialResolutionException.java b/core/src/main/java/org/apache/spark/security/CredentialResolutionException.java index eb5aaf2372984..86db39835168e 100644 --- a/core/src/main/java/org/apache/spark/security/CredentialResolutionException.java +++ b/core/src/main/java/org/apache/spark/security/CredentialResolutionException.java @@ -26,7 +26,7 @@ * This is a checked exception to ensure callers handle credential resolution failures * explicitly (e.g., retry, fail the job, or fall back to another mechanism). * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi public class CredentialResolutionException extends Exception { diff --git a/core/src/main/java/org/apache/spark/security/FileTokenIngestor.java b/core/src/main/java/org/apache/spark/security/FileTokenIngestor.java index f99e957ad6e73..c43fd5dc601ca 100644 --- a/core/src/main/java/org/apache/spark/security/FileTokenIngestor.java +++ b/core/src/main/java/org/apache/spark/security/FileTokenIngestor.java @@ -36,13 +36,13 @@ /** * A {@link TokenIngestor} that reads an OIDC identity token from a file. * <p> - * The file path is typically a Kubernetes projected service account token + * The file path typically points to a Kubernetes projected service account token * (e.g., {@code /var/run/secrets/tokens/spark-identity}) or a path configured via * {@code spark.security.oidc.identityToken.file}. * <p> * This implementation: * <ul> - * <li>Detects file rotation via mtime change (only re-parses when the file changes)</li> + * <li>Detects file rotation by comparing file content (only re-parses when it changes)</li> * <li>Parses JWT claims by Base64-decoding the payload segment directly, without * signature verification, since the token is trusted from the local filesystem * and works with both signed (RS256, ES256) and unsigned tokens</li> @@ -50,7 +50,7 @@ * returns empty rather than throwing</li> * </ul> * - * @since 4.3.0 + * @since 4.4.0 */ @Private public class FileTokenIngestor implements TokenIngestor { @@ -60,11 +60,17 @@ public class FileTokenIngestor implements TokenIngestor { private final Path tokenPath; - // Cached state for rotation detection. - // Write order matters for thread-safety: cachedContext must be visible before - // lastMtime, so a concurrent reader never sees a new mtime with a stale context. - private volatile UserContext cachedContext = null; - private volatile long lastMtime = -1L; + private volatile CachedToken cachedToken; + + private static final class CachedToken { + private final String content; + private final UserContext context; + + private CachedToken(String content, UserContext context) { + this.content = content; + this.context = context; + } + } /** * Construct a new FileTokenIngestor. @@ -83,24 +89,22 @@ public Optional<UserContext> load() { return Optional.empty(); } - // File did not change since last successful parse - long currentMtime = Files.getLastModifiedTime(tokenPath).toMillis(); - if (currentMtime == lastMtime && cachedContext != null) { - return Optional.of(cachedContext); - } - String content = new String(Files.readAllBytes(tokenPath), StandardCharsets.UTF_8).trim(); if (content.isEmpty()) { LOG.warn("Token file is empty: {}", MDC.of(LogKeys.PATH, tokenPath)); return Optional.empty(); } + CachedToken currentCache = cachedToken; + if (currentCache != null && content.equals(currentCache.content)) { + return Optional.of(currentCache.context); + } + Optional<UserContext> userContext = parseJwt(content); if (userContext.isPresent()) { // If the new file has invalid content, return empty and do NOT // fall back to the previously cached context. The caller will retry on next poll. - cachedContext = userContext.get(); - lastMtime = currentMtime; + cachedToken = new CachedToken(content, userContext.get()); } return userContext; } catch (Exception e) { @@ -114,8 +118,7 @@ public Optional<UserContext> load() { /** * Parse a JWT token string into a UserContext by Base64-decoding the payload segment. * This works with both signed (RS256, ES256) and unsigned (alg:none) tokens since - * we never verify the signature - the token is trusted from the local filesystem and - * will be re-verified downstream at the STS token exchange. + * we never verify the signature - the token is trusted from the local filesystem. */ private Optional<UserContext> parseJwt(String token) { try { diff --git a/core/src/main/java/org/apache/spark/security/ServiceCredential.java b/core/src/main/java/org/apache/spark/security/ServiceCredential.java index b0502cf922a45..4e988821a3983 100644 --- a/core/src/main/java/org/apache/spark/security/ServiceCredential.java +++ b/core/src/main/java/org/apache/spark/security/ServiceCredential.java @@ -36,7 +36,7 @@ * <p> * This class is immutable and {@link Serializable}. * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi public final class ServiceCredential implements Serializable { diff --git a/core/src/main/java/org/apache/spark/security/TokenIngestor.java b/core/src/main/java/org/apache/spark/security/TokenIngestor.java index ad4a208e663f0..fd15b5f734387 100644 --- a/core/src/main/java/org/apache/spark/security/TokenIngestor.java +++ b/core/src/main/java/org/apache/spark/security/TokenIngestor.java @@ -23,18 +23,21 @@ /** * :: DeveloperApi :: - * Read an OIDC identity token and produces a {@link UserContext}. + * Reads an OIDC identity token and produces a {@link UserContext}. * <p> * Implementation should be stateless with respect to Spark configuration; * configuration is passed at construction time. + * Implementations must be thread-safe because {@link #load()} may be called concurrently. * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi public interface TokenIngestor { /** * Attempt to load the current identity token and parse it into a UserContext. + * This method may be called repeatedly. Implementations may cache parsed tokens, but must + * detect changes to the underlying token source and return the current identity. * * @return a present Optional containing the UserContext if a valid token is available, * or empty if unavailable (e.g. empty content / missing file). diff --git a/core/src/main/java/org/apache/spark/security/UserContext.java b/core/src/main/java/org/apache/spark/security/UserContext.java index 9d9f372dad34b..a063c6c2eae1c 100644 --- a/core/src/main/java/org/apache/spark/security/UserContext.java +++ b/core/src/main/java/org/apache/spark/security/UserContext.java @@ -31,7 +31,7 @@ * <b>not</b> {@link java.io.Serializable} and must never be transmitted to executors. * The {@code rawToken} field is always redacted in {@link #toString()}. * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi public final class UserContext { diff --git a/core/src/main/java/org/apache/spark/security/UserCredentials.java b/core/src/main/java/org/apache/spark/security/UserCredentials.java index 3f5e1d54b9239..3d4427651a0e9 100644 --- a/core/src/main/java/org/apache/spark/security/UserCredentials.java +++ b/core/src/main/java/org/apache/spark/security/UserCredentials.java @@ -38,7 +38,7 @@ * This class is transmitted to executors and does <b>not</b> contain any reference * to {@link UserContext} or raw identity tokens. It is immutable and {@link Serializable}. * - * @since 4.3.0 + * @since 4.4.0 */ @DeveloperApi public final class UserCredentials implements Serializable { diff --git a/core/src/main/resources/org/apache/spark/ui/static/webui.css b/core/src/main/resources/org/apache/spark/ui/static/webui.css index c407c8cb22bd0..94b4ec0971e9d 100755 --- a/core/src/main/resources/org/apache/spark/ui/static/webui.css +++ b/core/src/main/resources/org/apache/spark/ui/static/webui.css @@ -139,6 +139,11 @@ a.kill-link { color: var(--bs-secondary-color); } +a.confirm-link { + margin-left: 4px; + color: var(--bs-secondary-color); +} + a.name-link { word-wrap: break-word; } diff --git a/core/src/main/resources/org/apache/spark/ui/static/webui.js b/core/src/main/resources/org/apache/spark/ui/static/webui.js index 436b987a17139..aab41a6a6717a 100644 --- a/core/src/main/resources/org/apache/spark/ui/static/webui.js +++ b/core/src/main/resources/org/apache/spark/ui/static/webui.js @@ -131,6 +131,13 @@ $(function() { } }); + // generic links guarded by a confirmation prompt before navigating + $(document).on("click", "a.confirm-link[data-confirm-message]", function(e) { + if (!window.confirm($(this).data("confirm-message"))) { + e.preventDefault(); + } + }); + // loadMore / loadNew buttons $(document).on("click", ".log-more-btn", function() { loadMore(); }); $(document).on("click", ".log-new-btn", function() { loadNew(); }); diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala index 0b3c22a22cb46..6c1b49157cc01 100644 --- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala +++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala @@ -236,6 +236,10 @@ private[spark] class ContextCleaner( /** Perform shuffle cleanup. */ def doCleanupShuffle(shuffleId: Int, blocking: Boolean): Unit = { try { + // A shuffle lives in exactly one tracker, split by dependency type: a regular shuffle in the + // MapOutputTracker, a pipelined shuffle in the driver-only StreamingShuffleOutputTracker (see + // DAGScheduler.createShuffleMapStage). Clean up whichever holds it -- the two branches are + // independent, each keyed on its own tracker's membership, so neither depends on the other. if (mapOutputTrackerMaster.containsShuffle(shuffleId)) { logDebug("Cleaning shuffle " + shuffleId) // Shuffle must be removed before it's unregistered from the output tracker @@ -244,6 +248,19 @@ private[spark] class ContextCleaner( mapOutputTrackerMaster.unregisterShuffle(shuffleId) listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) logDebug("Cleaned shuffle " + shuffleId) + } else if (streamingShuffleOutputTrackerMaster.exists(_.containsShuffle(shuffleId))) { + // A pipelined shuffle's output state is driver-only (the worker tracker caches nothing), so + // it is unregistered directly from the StreamingShuffleOutputTracker rather than the + // MapOutputTracker. Still call shuffleDriverComponents.removeShuffle to balance the + // registerShuffle that the ShuffleDependency constructor issues for EVERY shuffle (incl. a + // pipelined one): a custom components impl may track per-shuffle driver state that would + // otherwise leak. For the default impl this is just a RemoveShuffle RPC that finds no + // blocks (a pipelined shuffle has none), matching the regular branch's ordering. + logDebug("Cleaning pipelined shuffle " + shuffleId) + shuffleDriverComponents.removeShuffle(shuffleId, blocking) + streamingShuffleOutputTrackerMaster.foreach(_.unregisterShuffle(shuffleId)) + listeners.asScala.foreach(_.shuffleCleaned(shuffleId)) + logDebug("Cleaned pipelined shuffle " + shuffleId) } else { logDebug("Asked to cleanup non-existent shuffle (maybe it was already removed)") } @@ -308,6 +325,8 @@ private[spark] class ContextCleaner( private def broadcastManager = sc.env.broadcastManager private def mapOutputTrackerMaster = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + private def streamingShuffleOutputTrackerMaster: Option[StreamingShuffleOutputTrackerMaster] = + sc.env.streamingShuffleOutputTracker.map(_.asInstanceOf[StreamingShuffleOutputTrackerMaster]) } private object ContextCleaner { diff --git a/core/src/main/scala/org/apache/spark/ExecutorAllocationClient.scala b/core/src/main/scala/org/apache/spark/ExecutorAllocationClient.scala index c533f5cbc0b0b..870a24a9b3b4f 100644 --- a/core/src/main/scala/org/apache/spark/ExecutorAllocationClient.scala +++ b/core/src/main/scala/org/apache/spark/ExecutorAllocationClient.scala @@ -104,6 +104,19 @@ private[spark] trait ExecutorAllocationClient { countFailures = false) } + /** + * Decommission only executors which are still idle when the request is accepted. + * Implementations must check for assigned tasks and stop further task assignment atomically. + * The default declines all requests. Clients that support graceful dynamic-allocation + * scale-down must override this method and provide the atomicity guarantee. + * + * @param executorsAndDecomInfo identifiers of executors and decommission information + * @param adjustTargetNumExecutors whether to reduce the target number of executors + * @return the ids of the executors accepted for decommissioning + */ + def decommissionExecutorsIfIdle( + executorsAndDecomInfo: Array[(String, ExecutorDecommissionInfo)], + adjustTargetNumExecutors: Boolean): Seq[String] = Seq.empty /** * Request that the cluster manager decommission the specified executor. diff --git a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala index 57f2d21bc7b70..fe6a3c0d1bf63 100644 --- a/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala +++ b/core/src/main/scala/org/apache/spark/ExecutorAllocationManager.scala @@ -69,7 +69,9 @@ import org.apache.spark.util.{Clock, SystemClock, ThreadUtils, Utils} * blocks will be removed if it has been idle for more than L seconds. * * There is no retry logic in either case because we make the assumption that the cluster manager - * will eventually fulfill all requests it receives asynchronously. + * will eventually fulfill all requests it receives asynchronously. The one exception is the + * deferred target push used while the executors are held (see `targetSyncPending`), which + * retries until the cluster manager acknowledges the current targets. * * The relevant Spark properties are below. Each of these properties applies separately to * every ResourceProfile. So if you set a minimum number of executors, that is a minimum @@ -175,6 +177,28 @@ private[spark] class ExecutorAllocationManager( // (2) an executor idle timeout has elapsed. @volatile private var initializing: Boolean = true + // Whether allocation is suspended because the executors are held. While this is true, + // `schedule()` is a no-op so that pending tasks do not bring up new executors. + // See `SparkContext.holdExecutors()`. + private var suspended: Boolean = false + + // Whether the current executor targets still have to be pushed to the cluster manager. Set + // when a push from `suspend()`/`resume()` fails or is rejected (e.g. before the YARN AM has + // registered), and by `reset()`, which may run inside a cluster manager RPC handler where a + // synchronous request would self-deadlock (e.g. YARN's RegisterClusterManager). The push is + // performed from the allocation thread in `schedule()` and retried until acknowledged, with + // an exponential backoff: the conditions under which it retries (an AM restart, an + // unreachable cluster manager) resolve on a timescale far above the allocation tick. + private var targetSyncPending: Boolean = false + + // Ticks to wait before the next deferred push attempt, and the delay to arm after another + // failure. Zero keeps the first attempt after arming `targetSyncPending` immediate. + private var ticksUntilTargetSync: Int = 0 + private var targetSyncBackoffTicks: Int = 0 + + // Upper bound of the deferred push backoff, in ticks (10 seconds with the 100ms interval). + private val maxTargetSyncBackoffTicks = 100 + // Number of locality aware tasks for each ResourceProfile, used for executor placement. private var numLocalityAwareTasksPerResourceProfileId = new mutable.HashMap[Int, Int] numLocalityAwareTasksPerResourceProfileId(defaultProfileId) = 0 @@ -282,12 +306,107 @@ private[spark] class ExecutorAllocationManager( def reset(): Unit = synchronized { addTime = 0L numExecutorsTargetPerResourceProfileId.keys.foreach { rpId => - numExecutorsTargetPerResourceProfileId(rpId) = initialNumExecutors + numExecutorsTargetPerResourceProfileId(rpId) = if (suspended) 0 else initialNumExecutors } numExecutorsToAddPerResourceProfileId.keys.foreach { rpId => numExecutorsToAddPerResourceProfileId(rpId) = 1 } executorMonitor.reset() + if (suspended) { + // A restarted cluster manager AM may have allocated executors on its own, so the zero + // targets have to be pushed again. Leave that to `schedule()`: this method may run + // inside the cluster manager's RPC handler, where a synchronous request would + // self-deadlock. + targetSyncPending = true + ticksUntilTargetSync = 0 + targetSyncBackoffTicks = 0 + } + } + + /** + * Push the current executor targets to the cluster manager and report whether the request + * was acknowledged, arming the `schedule()` retry otherwise. + */ + private def syncTargetsWithClient(): Boolean = { + val acknowledged = try { + client.requestTotalExecutors( + numExecutorsTargetPerResourceProfileId.toMap, + numLocalityAwareTasksPerResourceProfileId.toMap, + rpIdToHostToLocalTaskCount) + } catch { + case NonFatal(e) => + // Use INFO level to be consistent with `doUpdateRequest`: errors here are more + // commonly caused by YARN AM restarts, which is a recoverable issue. + logInfo("Error reaching cluster manager.", e) + false + } + if (acknowledged) { + targetSyncPending = false + ticksUntilTargetSync = 0 + targetSyncBackoffTicks = 0 + } else { + targetSyncPending = true + ticksUntilTargetSync = targetSyncBackoffTicks + targetSyncBackoffTicks = + math.min(maxTargetSyncBackoffTicks, math.max(1, targetSyncBackoffTicks * 2)) + } + acknowledged + } + + /** + * Suspend allocation and lower the executor targets of all resource profiles to zero, so that + * pending tasks do not bring up new executors while the executors are held. Target updates + * are a no-op in `schedule()` until [[resume()]] is called. + * + * @return whether the zero targets were acknowledged by the cluster manager. When false, the + * push is retried from `schedule()` until acknowledged. + */ + def suspend(): Boolean = synchronized { + if (!suspended) { + suspended = true + numExecutorsTargetPerResourceProfileId.keys.foreach { rpId => + numExecutorsTargetPerResourceProfileId(rpId) = 0 + } + numExecutorsToAddPerResourceProfileId.keys.foreach { rpId => + numExecutorsToAddPerResourceProfileId(rpId) = 1 + } + syncTargetsWithClient() + } else { + true + } + } + + /** + * Resume allocation suspended by [[suspend()]]. The executor targets are recomputed from the + * current load on the next `schedule()` run. + * + * @return whether the restored targets were acknowledged by the cluster manager. When false, + * the push is retried from `schedule()` until acknowledged. + */ + def resume(): Boolean = synchronized { + if (suspended) { + suspended = false + // Restore at least the lower bound of the target: the initial warm-up requested by + // `start()` when no stage has been submitted yet (`updateAndSyncNumExecutorsTarget` + // skips it while initializing), and `minNumExecutors` otherwise, so that an idle + // application does not stay below the minimum until the next backlog. + val floor = if (initializing) initialNumExecutors else minNumExecutors + numExecutorsTargetPerResourceProfileId.keys.foreach { rpId => + numExecutorsTargetPerResourceProfileId(rpId) = + math.max(numExecutorsTargetPerResourceProfileId(rpId), floor) + } + val acknowledged = syncTargetsWithClient() + if (!initializing && numExecutorsTargetPerResourceProfileId.keys + .exists(maxNumExecutorsNeededPerResourceProfile(_) > 0)) { + // Trigger an immediate ramp-up for the load that built up while suspended. Do not + // touch `addTime` otherwise: nothing resets it to NOT_SET while idle, and a stale + // value would let a much later backlog skip the scheduler backlog timeout. + addTime = clock.nanoTime() + } + acknowledged + } else { + true + } } /** @@ -339,6 +458,21 @@ private[spark] class ExecutorAllocationManager( * This is factored out into its own method for testing. */ private def schedule(): Unit = synchronized { + if (targetSyncPending) { + if (ticksUntilTargetSync <= 0) { + // Deferred target push, retried with a backoff until acknowledged: an earlier push + // was rejected or failed (e.g. before the YARN AM registered), or `reset()` ran + // inside a cluster manager RPC handler where a synchronous request would + // self-deadlock. No `testing` short-circuit here: tests assert on this call through + // the mocked client. + syncTargetsWithClient() + } else { + ticksUntilTargetSync -= 1 + } + } + if (suspended) { + return + } val executorIdsToBeRemoved = executorMonitor.timedOutExecutors() if (executorIdsToBeRemoved.nonEmpty) { initializing = false @@ -584,10 +718,9 @@ private[spark] class ExecutorAllocationManager( if (decommissionEnabled) { val executorIdsWithoutHostLoss = executorIdsToBeRemoved.map( id => (id, ExecutorDecommissionInfo("spark scale down"))).toArray - client.decommissionExecutors( + client.decommissionExecutorsIfIdle( executorIdsWithoutHostLoss, - adjustTargetNumExecutors = false, - triggeredByExecutor = false) + adjustTargetNumExecutors = false) } else { client.killExecutors(executorIdsToBeRemoved.toSeq, adjustTargetNumExecutors = false, countFailures = false, force = false) @@ -718,8 +851,11 @@ private[spark] class ExecutorAllocationManager( updateExecutorPlacementHints() if (!numExecutorsTargetPerResourceProfileId.contains(profId)) { - numExecutorsTargetPerResourceProfileId.put(profId, initialNumExecutors) - if (initialNumExecutors > 0) { + // While suspended, a new resource profile must start at a zero target so that the + // hold is not bypassed; the target is recomputed on resume. + numExecutorsTargetPerResourceProfileId.put( + profId, if (suspended) 0 else initialNumExecutors) + if (!suspended && initialNumExecutors > 0) { logDebug(s"requesting executors, rpId: $profId, initial number is $initialNumExecutors") // we need to trigger a schedule since we add an initial number here. client.requestTotalExecutors( diff --git a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala index 742baf40ecbf0..933da94772725 100644 --- a/core/src/main/scala/org/apache/spark/MapOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/MapOutputTracker.scala @@ -799,7 +799,7 @@ private[spark] class MapOutputTrackerMaster( conf: SparkConf, private[spark] val broadcastManager: BroadcastManager, private[spark] val isLocal: Boolean) - extends MapOutputTracker(conf) { + extends MapOutputTracker(conf) with ShuffleOutputTrackerMaster { // The size at which we use Broadcast to send the map output statuses to the executors private val minSizeForBroadcast = conf.get(SHUFFLE_MAPOUTPUT_MIN_SIZE_FOR_BROADCAST).toInt @@ -943,6 +943,10 @@ private[spark] class MapOutputTrackerMaster( } } + // ShuffleOutputTrackerMaster: a regular shuffle has no per-job registration, so jobId is ignored. + override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = + registerShuffle(shuffleId, numMaps, numReduces) + def updateMapOutput(shuffleId: Int, mapId: Long, bmAddress: BlockManagerId): Unit = { shuffleStatuses.get(shuffleId) match { case Some(shuffleStatus) => @@ -1030,7 +1034,7 @@ private[spark] class MapOutputTrackerMaster( } /** Unregister shuffle data */ - def unregisterShuffle(shuffleId: Int): Unit = { + override def unregisterShuffle(shuffleId: Int): Unit = { shuffleStatuses.remove(shuffleId).foreach { shuffleStatus => shuffleStatus.invalidateSerializedMapOutputStatusCache() shuffleStatus.invalidateSerializedMergeOutputStatusCache() @@ -1077,7 +1081,7 @@ private[spark] class MapOutputTrackerMaster( } /** Check if the given shuffle is being tracked */ - def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId) + override def containsShuffle(shuffleId: Int): Boolean = shuffleStatuses.contains(shuffleId) def getNumAvailableOutputs(shuffleId: Int): Int = { shuffleStatuses.get(shuffleId).map(_.numAvailableMapOutputs).getOrElse(0) diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala index 8cb5eef770a07..04d04a1fd3828 100644 --- a/core/src/main/scala/org/apache/spark/SparkContext.scala +++ b/core/src/main/scala/org/apache/spark/SparkContext.scala @@ -61,7 +61,7 @@ import org.apache.spark.resource._ import org.apache.spark.resource.ResourceUtils._ import org.apache.spark.rpc.RpcEndpointRef import org.apache.spark.scheduler._ -import org.apache.spark.scheduler.cluster.StandaloneSchedulerBackend +import org.apache.spark.scheduler.cluster.{CoarseGrainedSchedulerBackend, SchedulerBackendUtils, StandaloneSchedulerBackend} import org.apache.spark.scheduler.local.LocalSchedulerBackend import org.apache.spark.shuffle.ShuffleDataIOUtils import org.apache.spark.shuffle.api.ShuffleDriverComponents @@ -245,6 +245,13 @@ class SparkContext(config: SparkConf) extends Logging { private var _plugins: Option[PluginContainer] = None private var _resourceProfileManager: ResourceProfileManager = _ + // Whether the executors are held via `holdExecutors()`, and, when dynamic allocation is + // disabled, the number of executors to restore on `resumeExecutors()`. Declared here so + // that the initializers run before the constructor exposes the context through the UI: + // a hold served while `postStartHook()` is still waiting must not be erased. + @volatile private var _executorsHeld: Boolean = false + private var heldNumExecutors: Int = 0 + /* ------------------------------------------------------------------------------------- * | Accessors and public fields. These provide access to the internal state of the | | context. | @@ -452,9 +459,18 @@ class SparkContext(config: SparkConf) extends Logging { // instead of relying on the default value of the config constant. if (SparkMasterRegex.isK8s(master) && _conf.getBoolean("spark.kubernetes.executor.useDriverPodIP", true)) { - logInfo("Use DRIVER_BIND_ADDRESS instead of DRIVER_HOST_ADDRESS as driver address " + - "because spark.kubernetes.executor.useDriverPodIP is true in K8s mode.") - _conf.set(DRIVER_HOST_ADDRESS, _conf.get(DRIVER_BIND_ADDRESS)) + val bindAddress = _conf.get(DRIVER_BIND_ADDRESS) + if (Utils.isAnyLocalAddress(bindAddress)) { + val driverHost = _conf.get(DRIVER_HOST_ADDRESS) + logInfo(log"spark.kubernetes.executor.useDriverPodIP is true but bind address " + + log"${MDC(LogKeys.BIND_ADDRESS, bindAddress)} is a wildcard; " + + log"preserving advertised driver host ${MDC(LogKeys.HOST, driverHost)}") + _conf.set(DRIVER_HOST_ADDRESS, driverHost) + } else { + logInfo("Use DRIVER_BIND_ADDRESS instead of DRIVER_HOST_ADDRESS as driver address " + + "because spark.kubernetes.executor.useDriverPodIP is true in K8s mode.") + _conf.set(DRIVER_HOST_ADDRESS, Utils.normalizeIpIfNeeded(bindAddress)) + } } else { _conf.set(DRIVER_HOST_ADDRESS, _conf.get(DRIVER_HOST_ADDRESS)) } @@ -699,6 +715,10 @@ class SparkContext(config: SparkConf) extends Logging { postEnvironmentUpdate() postApplicationStart() + // Advertise whether this application can be held, now that the shuffle driver components, + // which decide it, are up. + reportExecutorHoldStatus() + // After application started, attach handlers to started server and start handler. _ui.foreach(_.attachAllHandlers()) // Attach the driver metrics servlet handler to the web ui after the metrics system is started. @@ -2066,6 +2086,253 @@ class SparkContext(config: SparkConf) extends Logging { } } + /** + * Whether `holdExecutors()` is supported in the current deployment. It requires a scheduler + * backend that can adjust the number of executors and can hold them, decommission support, + * and shuffle data kept outside the executors: either an external shuffle service or a + * `ShuffleDataIO` with reliable storage. + */ + private[spark] def executorHoldSupported: Boolean = { + (schedulerBackend match { + case cg: CoarseGrainedSchedulerBackend => cg.supportsExecutorHold + case _ => false + }) && + (conf.get(SHUFFLE_SERVICE_ENABLED) || shuffleDriverComponents.supportsReliableStorage()) && + conf.get(DECOMMISSION_ENABLED) + } + + /** Whether the executors are currently held by `holdExecutors()`. */ + private[spark] def executorsHeld: Boolean = _executorsHeld + + /** + * :: DeveloperApi :: + * Hold the whole application by declining to allocate new executors and gracefully + * decommissioning all existing ones. Each executor finishes its running tasks and then + * exits -- unless `spark.executor.decommission.forceKillTimeout` is set, in which case an + * executor still running tasks is killed after that timeout. The shuffle data already + * written remains available outside the executors, so the application can later pick up + * where it left off via `resumeExecutors()`. Cached blocks are not preserved and are + * recomputed after resuming. While held the application has no executors, so with + * `spark.default.parallelism` unset the default parallelism falls back to 2, and an RDD + * created during the hold keeps that partition count after resuming. + * + * This requires decommission support (`spark.decommission.enabled`), shuffle data kept + * outside the executors -- either an external shuffle service + * (`spark.shuffle.service.enabled`) or a `ShuffleDataIO` with reliable storage -- and a + * scheduler backend that can hold executors: Standalone, YARN, and Kubernetes with the + * `direct` pods allocator. Fallback storage + * (`spark.storage.decommission.fallbackStorage.path`) deliberately does not qualify: + * shuffle blocks not yet migrated when an executor exits are dropped. + * + * Executor requirements requested while held, through `requestExecutors` or + * `requestTotalExecutors`, are recorded but nothing is allocated until `resumeExecutors()` + * restores them. + * + * Pipelined-shuffle jobs are outside the hold's scope. A hold is rejected, on a + * best-effort check, while a pipelined job is running: its transient shuffle data lives + * only on the executors and would not survive the drain, and a group that slips past the + * check is aborted rather than drained, since a pipelined task set tolerates no task + * failure. A pipelined job submitted while held fails its gang admission immediately + * instead of waiting; resubmit it after the resume (with + * `spark.scheduler.pipelinedGroup.slotCheck.enabled=false` there is no admission check, + * so it waits for the resume instead). Workloads with long-running tasks (a + * streaming receiver, continuous processing) are outside the scope too: their tasks never + * finish, so the drain cannot complete. + * + * @throws IllegalArgumentException when the decommission or shuffle-storage precondition is + * not met; an unsupported scheduler backend instead returns false with a warning. + * @return whether the lowered executor requirement was acknowledged by the cluster manager. + * With dynamic allocation a rejected request is retried in the background; the + * executors are drained in either case. + */ + @DeveloperApi + def holdExecutors(): Boolean = { + val acknowledged = schedulerBackend match { + case cg: CoarseGrainedSchedulerBackend if cg.supportsExecutorHold => + require(executorHoldSupported, + s"holdExecutors() requires ${DECOMMISSION_ENABLED.key} and either " + + s"${SHUFFLE_SERVICE_ENABLED.key} or a ShuffleDataIO with reliable storage") + val pipelinedRunning = taskScheduler match { + case ts: TaskSchedulerImpl => ts.hasPipelinedTaskSets + case _ => false + } + if (pipelinedRunning) { + // A pipelined group reads and writes transient shuffle data that lives only on + // its executors: a partially launched group would deadlock the drain, and a + // force-killed member aborts the whole group. + logWarning(log"Cannot hold the executors while a pipelined job is running.") + false + } else synchronized { + if (_executorsHeld) { + // A repeated hold re-asserts the zero requirement: the earlier publish may not + // have been acknowledged, and with dynamic allocation off nothing retries it. + if (executorAllocationManager.isDefined) true else zeroExecutorRequirementAndDrain(cg) + } else { + if (executorAllocationManager.isEmpty) { + // The requirement to restore on resume when none was explicitly requested + // (explicitly requested totals, made before or during the hold, are + // republished from the backend directly). Only killExecutors' bookkeeping + // zero is kill-seeded (read atomically against a concurrent reset): restore + // the count of executors not already being removed, so that resume neither + // parks the application at zero nor undoes the downscale. Otherwise + // Standalone has no explicit requirement by default (and ignores + // spark.executor.instances, even a leftover value), so restore an unbounded + // one; elsewhere follow the conf, or fall back to the cluster manager's + // default when no executor has registered yet. + heldNumExecutors = if (cg.hasKillSeededTotalsOnly) { + cg.activeExecutorCount + } else { + schedulerBackend match { + case _: StandaloneSchedulerBackend => Int.MaxValue + case _ => + conf.get(EXECUTOR_INSTANCES).getOrElse(math.max(cg.getExecutorIds().size, + SchedulerBackendUtils.DEFAULT_NUMBER_EXECUTORS)) + } + } + } + // Mark the hold before talking to the cluster manager, so that a partial failure + // below leaves the executors held, and thus resumable, instead of half-held. + _executorsHeld = true + cg.setExecutorsHeld(true) + executorAllocationManager match { + case Some(manager) => + val acknowledged = manager.suspend() + drainHeldExecutors(cg) + acknowledged + case None => zeroExecutorRequirementAndDrain(cg) + } + } + } + case _ => + logWarning("Holding executors is not supported by current scheduler.") + false + } + reportExecutorHoldStatus() + acknowledged + } + + // Gracefully decommission all the current executors of a held application and let the + // executor monitor know, so that it does not try to remove the draining executors again and + // reports them in the decommissioning metrics. + private def drainHeldExecutors(b: ExecutorAllocationClient): Unit = { + val executors = b.getExecutorIds() + if (executors.nonEmpty) { + val decommissioned = b.decommissionExecutors( + executors.map(id => (id, ExecutorDecommissionInfo("Executors are held"))).toArray, + adjustTargetNumExecutors = false, + triggeredByExecutor = false) + executorAllocationManager.foreach( + _.executorMonitor.executorsDecommissioned(decommissioned)) + } + } + + // Restore the hold invariant without dynamic allocation: publish the zero requirement + // (the requested totals are kept and republished on resume) and drain the current + // executors. The publish must not abort the drain, which has to run even when the cluster + // manager is temporarily unreachable; the registration guard remains the backstop when the + // publish fails. + private[spark] def zeroExecutorRequirementAndDrain( + cg: CoarseGrainedSchedulerBackend): Boolean = { + val acknowledged = try { + cg.republishRequestedTotals() + } catch { + case NonFatal(e) => + logWarning(log"Failed to lower the executor requirement while holding the " + + log"executors.", e) + false + } + drainHeldExecutors(cg) + acknowledged + } + + /** + * :: DeveloperApi :: + * Resume an application held by `holdExecutors()` by restoring its executor requirements. + * + * @return whether the restored executor requirement was acknowledged by the cluster manager. + * With dynamic allocation a rejected request is retried in the background and the + * hold is lifted; otherwise the executors stay held so the call can be retried. + */ + @DeveloperApi + def resumeExecutors(): Boolean = { + val acknowledged = schedulerBackend match { + case cg: CoarseGrainedSchedulerBackend => + synchronized { + if (!_executorsHeld) { + true + } else { + // Lift the backend guard before restoring the requirement, so that an executor + // granted by the restored requirement cannot race with its own registration and + // be drained. + cg.setExecutorsHeld(false) + val acknowledged = executorAllocationManager match { + case Some(manager) => + // resume() retries a rejected push in the background, so the hold can be + // lifted regardless of the acknowledgment. + manager.resume() + case None => + try { + // Totals requested before or during the hold are republished as-is; the + // check and the publish are atomic, so a concurrent cluster manager reset + // cannot turn this into publishing an empty map. When none are recorded + // (never requested, or cleared by such a reset), restore the requirement + // captured at hold. + cg.republishExplicitTotals().getOrElse { + // Publish without recording: the pre-hold state had no explicitly + // requested totals, and recording the restore (in particular + // Standalone's unbounded sentinel) would flip killExecutors' empty-map + // seeding and mark the totals explicit for good. + cg.publishTotalsWithoutRecording( + immutable.Map(resourceProfileManager.defaultResourceProfile -> + heldNumExecutors)) + } + } catch { + case NonFatal(e) => + logWarning(log"Failed to restore the executor requirement while resuming " + + log"the executors.", e) + false + } + } + if (executorAllocationManager.isDefined || acknowledged) { + _executorsHeld = false + } else { + // The requirement could not be restored: stay held and re-arm the guard. The + // recorded totals are still the non-zero pre-hold ones, and an executor may + // have registered while the guard was down, so restore the hold invariant: + // push the zero requirement again and drain any current executors. The call + // can be retried. + cg.setExecutorsHeld(true) + zeroExecutorRequirementAndDrain(cg) + logWarning(log"The cluster manager did not acknowledge the restored executor " + + log"requirement; the executors remain held and resumeExecutors() can be " + + log"retried.") + } + acknowledged + } + } + case _ => + logWarning("Resuming executors is not supported by current scheduler.") + false + } + reportExecutorHoldStatus() + acknowledged + } + + /** + * Tell the cluster manager whether this application can be held and whether it currently is, + * so that it can show the hold status on its own UI. Called once the context is fully started + * -- `executorHoldSupported` reads the shuffle driver components, which are initialized late + * -- and again after every transition. Synchronized so that concurrent transitions cannot + * deliver their reports inverted: whichever call serializes last reports the current state. + */ + private def reportExecutorHoldStatus(): Unit = synchronized { + schedulerBackend match { + case cg: CoarseGrainedSchedulerBackend => + cg.reportExecutorHoldStatus(executorHoldSupported, _executorsHeld) + case _ => + } + } + /** The version of Spark on which this application is running. */ def version: String = SPARK_VERSION @@ -2730,16 +2997,17 @@ class SparkContext(config: SparkConf) extends Logging { * * @param tag The tag to be cancelled. Cannot contain ',' (comma) character. * @param reason reason for cancellation. - * @return A future with [[ActiveJob]]s, allowing extraction of information such as Job ID and - * tags. + * @return A future with the cancelled jobs' [[CancelledJobInfo]], allowing extraction of + * information such as Job ID and tags. Covers active jobs and barrier jobs cancelled while + * deferred for their slot-check retry (which have no [[ActiveJob]]). */ private[spark] def cancelJobsWithTagWithFuture( tag: String, - reason: String): Future[Seq[ActiveJob]] = { + reason: String): Future[Seq[CancelledJobInfo]] = { SparkContext.throwIfInvalidTag(tag) assertNotStopped() - val cancelledJobs = Promise[Seq[ActiveJob]]() + val cancelledJobs = Promise[Seq[CancelledJobInfo]]() dagScheduler.cancelJobsWithTag(tag, Some(reason), Some(cancelledJobs)) cancelledJobs.future } @@ -2777,6 +3045,20 @@ class SparkContext(config: SparkConf) extends Logging { dagScheduler.cancelAllJobs() } + /** + * Cancel all jobs that have been scheduled or are running. + * + * @param reason reason for cancellation. It is surfaced in the error of every cancelled job, so + * that a job aborted as collateral of a context-wide cancellation can be told + * apart from one that failed on its own. + * + * @since 4.4.0 + */ + def cancelAllJobs(reason: String): Unit = { + assertNotStopped() + dagScheduler.cancelAllJobs(Option(reason)) + } + /** * Cancel a given job if it's scheduled or running. * @@ -2903,8 +3185,34 @@ class SparkContext(config: SparkConf) extends Logging { private val nextRddId = new AtomicInteger(0) - /** Register a new RDD, returning its RDD ID */ - private[spark] def newRddId(): Int = nextRddId.getAndIncrement() + /** + * Testing helper: set the next value that [[newRddId]] will return. + */ + private[spark] def setNextRddIdForTesting(value: Int): Unit = nextRddId.set(value) + + /** + * Register a new RDD, returning its RDD ID. + * + * Fails if the 32-bit counter would wrap to a negative value. Continuing with + * wrapped ids can break BlockManager, UI, and other id-keyed state. + * [[org.apache.spark.storage.BlockId]] still parses negative RDD names as a + * safety net for any in-flight or pre-upgrade cached blocks. + */ + private[spark] def newRddId(): Int = { + val id = nextRddId.getAndIncrement() + if (id < 0) { + // Stay pegged so subsequent allocations keep failing clearly. + nextRddId.set(Int.MinValue) + throw new SparkException( + "RDD id counter overflowed Int.MaxValue (" + Int.MaxValue + + "). This application has created too many RDDs; restart it.") + } + if (id == Int.MaxValue) { + logWarning("Allocated the last valid RDD id (Int.MaxValue). " + + "Further RDD creation will fail.") + } + id + } /** * Registers listeners specified in spark.extraListeners, then starts the listener bus. @@ -3157,6 +3465,7 @@ object SparkContext extends Logging { private[spark] val SQL_EXECUTION_ID_KEY = "spark.sql.execution.id" private[spark] val DATASET_QUERY_EXECUTION_ID_KEY = "spark.sql.dataset.queryExecution.id" + private[spark] val SPARK_CONNECT_OPERATION_ID_PROPERTY = "spark.connect.operation_id" /** * Executor id for the driver. In earlier versions of Spark, this was `<driver>`, but this was diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala b/core/src/main/scala/org/apache/spark/SparkEnv.scala index ca48ee473eb0f..d48640d9469b4 100644 --- a/core/src/main/scala/org/apache/spark/SparkEnv.scala +++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala @@ -19,6 +19,7 @@ package org.apache.spark import java.io.File import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference import scala.collection.concurrent import scala.collection.mutable @@ -268,6 +269,19 @@ class SparkEnv ( private[spark] var executorBackend: Option[ExecutorBackend] = None + /** + * Versioned credential store for OIDC-based user credentials on both driver and executors. + * Updated via `UpdateUserCredentials` RPC and `TaskDescription` credential delivery. + * Read by connector-specific credential providers (e.g., SparkOidcAwsCredentialsProvider). + * Contains serialized `UserCredentials` (no raw identity token). + * + * The version field is a monotonically increasing counter assigned by `UserCredentialManager` + * on each credential renewal. It is used to guard against stale credentials from delayed + * `TaskDescription` delivery overwriting fresher credentials delivered via RPC broadcast. + */ + private[spark] val userCredentials: AtomicReference[VersionedCredentials] = + new AtomicReference[VersionedCredentials]() + private[spark] def stop(): Unit = { if (!isStopped) { @@ -498,7 +512,10 @@ class SparkEnv ( } else { conf.clone.set(MEMORY_OFFHEAP_ENABLED, false).set(MEMORY_OFFHEAP_SIZE, 0L) } - _memoryManager = UnifiedMemoryManager(memoryManagerConf, numUsableCores) + _memoryManager = UnifiedMemoryManager( + memoryManagerConf, + numUsableCores, + isDriver = SparkContext.isDriver(executorId)) } } @@ -825,3 +842,33 @@ object SparkEnv extends Logging { "Metrics Properties" -> metricsProperties.toSeq.sorted) } } + +/** + * Container for versioned OIDC user credentials. + * + * @param version Monotonically increasing counter assigned by `UserCredentialManager` on each + * credential renewal. Used to prevent stale credentials from overwriting fresher + * ones on executors. + * @param bytes Serialized `UserCredentials` payload (no raw identity token). + */ +private[spark] case class VersionedCredentials(version: Long, bytes: Array[Byte]) + +private[spark] object VersionedCredentials { + /** + * Atomically update a credential store only if the given version is strictly newer + * than what is currently stored. This prevents stale credentials (e.g., from a delayed + * `TaskDescription`) from overwriting fresher credentials delivered via RPC broadcast. + * + * Uses `AtomicReference.updateAndGet` to ensure the check-and-set is atomic even + * when called concurrently from multiple task threads and the RPC dispatcher thread. + */ + def updateIfNewer( + store: AtomicReference[VersionedCredentials], + version: Long, + bytes: Array[Byte]): Unit = { + val newValue = VersionedCredentials(version, bytes) + store.updateAndGet { current => + if (current == null || version > current.version) newValue else current + } + } +} diff --git a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala index 090be7f691d4b..df4e76b017dbc 100644 --- a/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala +++ b/core/src/main/scala/org/apache/spark/StreamingShuffleOutputTracker.scala @@ -28,6 +28,23 @@ import org.apache.spark.internal.config.SHUFFLE_MAPOUTPUT_DISPATCHER_NUM_THREADS import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv} import org.apache.spark.util.ThreadUtils +/** + * The driver-side registry for a shuffle's output. A regular shuffle is served by the + * `MapOutputTrackerMaster`, a pipelined (streaming) shuffle by the + * `StreamingShuffleOutputTrackerMaster` -- split by dependency type, with no overlap. This is the + * small common surface the DAGScheduler and ContextCleaner drive polymorphically, so they select + * the right tracker by shuffle-dependency type (see `DAGScheduler.outputTrackerMaster`) rather than + * special-casing pipelined shuffles at each call site. + */ +private[spark] trait ShuffleOutputTrackerMaster { + /** Register a shuffle so its outputs can be tracked. `jobId` is used by the streaming tracker. */ + def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit + /** Whether the given shuffle is registered with this tracker. */ + def containsShuffle(shuffleId: Int): Boolean + /** Unregister a shuffle and release its tracked state. */ + def unregisterShuffle(shuffleId: Int): Unit +} + private[spark] sealed trait StreamingShuffleTaskLocationTrackerMessage private[spark] case class UpdateStreamingShuffleTaskLocation( @@ -196,7 +213,7 @@ private[spark] abstract class StreamingShuffleOutputTracker(conf: SparkConf) ext private[spark] case class StreamingShuffleInfo(numMaps: Int, numReduces: Int, jobId: Int) private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) - extends StreamingShuffleOutputTracker(conf) { + extends StreamingShuffleOutputTracker(conf) with ShuffleOutputTrackerMaster { // map that stores task location information organized in the following fashion // shuffle id -> {mapId -> location} @@ -220,7 +237,7 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) pool } - def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = { + override def registerShuffle(shuffleId: Int, numMaps: Int, numReduces: Int, jobId: Int): Unit = { logInfo(log"Registering shuffleId ${MDC(LogKeys.SHUFFLE_ID, shuffleId)} with ${ MDC(LogKeys.NUM_MAPPERS, numMaps)} mappers and ${ MDC(LogKeys.NUM_REDUCERS, numReduces)} reducers") @@ -230,6 +247,8 @@ private[spark] class StreamingShuffleOutputTrackerMaster(conf: SparkConf) } } + override def containsShuffle(shuffleId: Int): Boolean = shuffleInfos.containsKey(shuffleId) + // for testing purposes private[spark] def getShuffleInfo(shuffleId: Int): Option[StreamingShuffleInfo] = { Option(shuffleInfos.get(shuffleId)) diff --git a/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala b/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala index bb8b02616c93f..35de8f5312506 100644 --- a/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala +++ b/core/src/main/scala/org/apache/spark/api/java/JavaSparkContext.scala @@ -819,6 +819,15 @@ class JavaSparkContext(val sc: SparkContext) extends Closeable { /** Cancel all jobs that have been scheduled or are running. */ def cancelAllJobs(): Unit = sc.cancelAllJobs() + /** + * Cancel all jobs that have been scheduled or are running. + * + * @param reason reason for cancellation + * + * @since 4.4.0 + */ + def cancelAllJobs(reason: String): Unit = sc.cancelAllJobs(reason) + /** * Returns a Java map of JavaRDDs that have marked themselves as persistent via cache() call. * diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala index 20424e4f37dff..e643045134f45 100644 --- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala +++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala @@ -51,6 +51,15 @@ private[spark] object PythonEvalType { val SQL_BATCHED_UDF = 100 val SQL_ARROW_BATCHED_UDF = 101 + // A scalar Python UDF applied element-wise over the elements of an array column, used to + // support Python UDFs inside higher-order function lambdas. See ExtractPythonUDFFromLambda. + // 102 lifts a row-at-a-time UDF (SQL_BATCHED_UDF / SQL_ARROW_BATCHED_UDF); 103-106 lift the + // vectorized scalar UDFs, preserving pandas- vs. Arrow-shaped batches and the iterator contract. + val SQL_ARROW_ELEMENTWISE_UDF = 102 + val SQL_SCALAR_PANDAS_ELEMENTWISE_UDF = 103 + val SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF = 104 + val SQL_SCALAR_ARROW_ELEMENTWISE_UDF = 105 + val SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF = 106 val SQL_SCALAR_PANDAS_UDF = 200 val SQL_GROUPED_MAP_PANDAS_UDF = 201 @@ -78,6 +87,19 @@ private[spark] object PythonEvalType { val SQL_WINDOW_AGG_ARROW_UDF = 253 val SQL_GROUPED_AGG_ARROW_ITER_UDF = 254 + // Incremental (partial + final) Arrow aggregator. Unlike the whole-group grouped-agg UDFs + // above, these support true partial aggregation: the PARTIAL eval type folds input rows into a + // per-group buffer (via the aggregator's `reduce`) on the map side, and the FINAL eval type + // merges partial buffers across the shuffle (via `merge`) and produces the output (via `finish`). + // See PythonIncrementalAggregateExec and the Python `Aggregator` API. + val SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF = 255 + val SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF = 256 + + // Window aggregation with an incremental Arrow aggregator. A window has no shuffle, so it needs + // neither the PARTIAL nor the FINAL eval type above: the operator sends each frame's rows to the + // worker, which folds them with `reduce` (from `zero`) and produces the value with `finish`. + val SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF = 257 + val SQL_TABLE_UDF = 300 val SQL_ARROW_TABLE_UDF = 301 val SQL_ARROW_UDTF = 302 @@ -86,6 +108,11 @@ private[spark] object PythonEvalType { case NON_UDF => "NON_UDF" case SQL_BATCHED_UDF => "SQL_BATCHED_UDF" case SQL_ARROW_BATCHED_UDF => "SQL_ARROW_BATCHED_UDF" + case SQL_ARROW_ELEMENTWISE_UDF => "SQL_ARROW_ELEMENTWISE_UDF" + case SQL_SCALAR_PANDAS_ELEMENTWISE_UDF => "SQL_SCALAR_PANDAS_ELEMENTWISE_UDF" + case SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF => "SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF" + case SQL_SCALAR_ARROW_ELEMENTWISE_UDF => "SQL_SCALAR_ARROW_ELEMENTWISE_UDF" + case SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF => "SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF" case SQL_SCALAR_PANDAS_UDF => "SQL_SCALAR_PANDAS_UDF" case SQL_GROUPED_MAP_PANDAS_UDF => "SQL_GROUPED_MAP_PANDAS_UDF" case SQL_GROUPED_AGG_PANDAS_UDF => "SQL_GROUPED_AGG_PANDAS_UDF" @@ -116,6 +143,23 @@ private[spark] object PythonEvalType { case SQL_GROUPED_AGG_ARROW_UDF => "SQL_GROUPED_AGG_ARROW_UDF" case SQL_WINDOW_AGG_ARROW_UDF => "SQL_WINDOW_AGG_ARROW_UDF" case SQL_GROUPED_AGG_ARROW_ITER_UDF => "SQL_GROUPED_AGG_ARROW_ITER_UDF" + case SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF => + "SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF" + case SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF => + "SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF" + case SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF => "SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF" + } + + // The eval types produced by ExtractPythonUDFFromLambda: a scalar UDF lifted out of a + // higher-order function's lambda, which receives each argument as an `array<T>` column and is + // applied element-wise inside the Python worker. See ExtractPythonUDFFromLambda. + def isElementwiseUDF(evalType: Int): Boolean = evalType match { + case SQL_ARROW_ELEMENTWISE_UDF | + SQL_SCALAR_PANDAS_ELEMENTWISE_UDF | + SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF | + SQL_SCALAR_ARROW_ELEMENTWISE_UDF | + SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF => true + case _ => false } } diff --git a/core/src/main/scala/org/apache/spark/deploy/DeployMessage.scala b/core/src/main/scala/org/apache/spark/deploy/DeployMessage.scala index 12e031711aa2a..8ad8efff4acda 100644 --- a/core/src/main/scala/org/apache/spark/deploy/DeployMessage.scala +++ b/core/src/main/scala/org/apache/spark/deploy/DeployMessage.scala @@ -201,6 +201,12 @@ private[deploy] object DeployMessages { case class KillExecutors(appId: String, executorIds: Seq[String]) + // Whether this application can be held and whether it currently is. Pushed on every + // transition and re-sent on failover, since a new Master starts without it. The number of + // executors still draining is not carried here: the Master already tracks the executors of + // the application and derives it from them. + case class ApplicationHoldUpdated(appId: String, supported: Boolean, held: Boolean) + // Master to AppClient case class RegisteredApplication(appId: String, master: RpcEndpointRef) extends DeployMessage @@ -251,6 +257,10 @@ private[deploy] object DeployMessages { case object StopAppClient + // Reports the hold status of the application to the Master, and caches it in the endpoint so + // that it can be re-sent on failover. + case class ReportApplicationHold(supported: Boolean, held: Boolean) + // Master to Worker & AppClient case class MasterChanged(master: RpcEndpointRef, masterWebUiUrl: String) diff --git a/core/src/main/scala/org/apache/spark/deploy/JsonProtocol.scala b/core/src/main/scala/org/apache/spark/deploy/JsonProtocol.scala index 2a3fd0d004e11..1119f47094d6d 100644 --- a/core/src/main/scala/org/apache/spark/deploy/JsonProtocol.scala +++ b/core/src/main/scala/org/apache/spark/deploy/JsonProtocol.scala @@ -95,6 +95,9 @@ private[deploy] object JsonProtocol { * `resourcesperexecutor` minimal resources required to each executor * `submitdate` time in Date that the application is submitted * `state` state of the application, see [[ApplicationState]] + * `holdsupported` whether the driver reported that it can be held + * `held` whether the application is currently held; always false once it finishes + * `draining` the number of executors still draining while the application is held * `duration` time in milliseconds that the application has been running * For compatibility also returns the deprecated `memoryperslave` & `resourcesperslave` fields. */ @@ -112,6 +115,9 @@ private[deploy] object JsonProtocol { .toList.map(writeResourceRequirement)) ~ ("submitdate" -> obj.submitDate.toString) ~ ("state" -> obj.state.toString) ~ + ("holdsupported" -> obj.holdSupported) ~ + ("held" -> obj.isHeld) ~ + ("draining" -> obj.numDrainingExecutors) ~ ("duration" -> obj.duration) } diff --git a/core/src/main/scala/org/apache/spark/deploy/client/StandaloneAppClient.scala b/core/src/main/scala/org/apache/spark/deploy/client/StandaloneAppClient.scala index 2d742b31f99c2..1b85ce06eced8 100644 --- a/core/src/main/scala/org/apache/spark/deploy/client/StandaloneAppClient.scala +++ b/core/src/main/scala/org/apache/spark/deploy/client/StandaloneAppClient.scala @@ -71,6 +71,9 @@ private[spark] class StandaloneAppClient( private val alreadyDead = new AtomicBoolean(false) private val registerMasterFutures = new AtomicReference[Array[JFuture[_]]] private val registrationRetryTimer = new AtomicReference[JScheduledFuture[_]] + // The hold status last reported to the Master, re-sent on failover since a new Master + // starts without it. None until the driver has reported one. + private var holdStatus: Option[(Boolean, Boolean)] = None // A thread pool for registering with masters. Because registering with a master is a blocking // action, this thread pool must be able to create "masterRpcAddresses.size" threads at the same @@ -169,6 +172,7 @@ private[spark] class StandaloneAppClient( registered.set(true) master = Some(masterRef) listener.connected(appId.get) + sendHoldStatus() case ApplicationRemoved(message) => markDead("Master removed our application: %s".format(message)) @@ -204,6 +208,18 @@ private[spark] class StandaloneAppClient( master = Some(masterRef) alreadyDisconnected = false masterRef.send(MasterChangeAcknowledged(appId.get)) + // The new Master recovered the application without its hold status, so report it again. + sendHoldStatus() + + case ReportApplicationHold(supported, held) => + holdStatus = Some((supported, held)) + sendHoldStatus() + } + + /** Report the last known hold status, if any, to the current Master. */ + private def sendHoldStatus(): Unit = holdStatus.foreach { case (supported, held) => + // Dropped while no Master is known; re-sent from `RegisteredApplication`/`MasterChanged`. + sendToMaster(ApplicationHoldUpdated(appId.get, supported, held)) } override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { @@ -330,6 +346,19 @@ private[spark] class StandaloneAppClient( } } + /** + * Report to the Master whether this application can be held and whether it currently is, so + * that the Master UI can show the hold status of the application. The status is cached and + * re-sent on failover, so a report made before the registration completes is not lost. + */ + def reportHoldStatus(supported: Boolean, held: Boolean): Unit = { + if (endpoint.get != null) { + endpoint.get.send(ReportApplicationHold(supported, held)) + } else { + logWarning("Attempted to report the hold status before driver fully initialized.") + } + } + /** * Kill the given list of executors through the Master. * @return whether the kill request is acknowledged. diff --git a/core/src/main/scala/org/apache/spark/deploy/history/EventLogFileWriters.scala b/core/src/main/scala/org/apache/spark/deploy/history/EventLogFileWriters.scala index 601515e57dc82..a725da2f3a44c 100644 --- a/core/src/main/scala/org/apache/spark/deploy/history/EventLogFileWriters.scala +++ b/core/src/main/scala/org/apache/spark/deploy/history/EventLogFileWriters.scala @@ -257,6 +257,15 @@ class SingleEventLogFileWriter( } object SingleEventLogFileWriter { + /** Returns names for completed and in-progress single event logs. */ + private[history] def getLogFileNames(appId: String, appAttemptId: Option[String]): Seq[String] = { + val logBaseName = EventLogFileWriter.nameForAppAndAttempt(appId, appAttemptId) + val names = logBaseName +: CompressionCodec.shortCompressionCodecNames.keys.toSeq.map { + codec => s"$logBaseName.$codec" + } + names ++ names.map(_ + EventLogFileWriter.IN_PROGRESS) + } + /** * Return a file-system-safe path to the log file for the given application. * diff --git a/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala b/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala index a87db35575536..748aec37570d7 100644 --- a/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala +++ b/core/src/main/scala/org/apache/spark/deploy/history/FsHistoryProvider.scala @@ -412,16 +412,11 @@ private[history] class FsHistoryProvider(conf: SparkConf, clock: Clock) override def getLastUpdatedTime(): Long = lastScanTime.get() override def getAppUI(appId: String, attemptId: Option[String]): Option[LoadedAppUI] = { - val logPath = RollingEventLogFilesWriter.EVENT_LOG_DIR_NAME_PREFIX + - EventLogFileWriter.nameForAppAndAttempt(appId, attemptId) val app = try { load(appId) } catch { - case _: NoSuchElementException if this.conf.get(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED) => - loadFromFallbackLocation(appId, attemptId, logPath) match { - case Some(wrapper) => wrapper - case None => return None - } + case _: NoSuchElementException if isOnDemandLogLoadEnabled => + loadFromFallbackLocations(appId, attemptId).getOrElse(return None) case _: NoSuchElementException => return None } @@ -462,6 +457,27 @@ private[history] class FsHistoryProvider(conf: SparkConf, clock: Clock) Some(loadedUI) } + private def loadFromFallbackLocations( + appId: String, + attemptId: Option[String]): Option[ApplicationInfoWrapper] = { + val logPaths = mutable.ArrayBuffer.empty[String] + if (conf.get(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED)) { + logPaths += RollingEventLogFilesWriter.EVENT_LOG_DIR_NAME_PREFIX + + EventLogFileWriter.nameForAppAndAttempt(appId, attemptId) + } + if (conf.get(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED)) { + logPaths ++= SingleEventLogFileWriter.getLogFileNames(appId, attemptId) + } + logPaths.iterator.map(loadFromFallbackLocation(appId, attemptId, _)).collectFirst { + case Some(app) => app + } + } + + private def isOnDemandLogLoadEnabled: Boolean = { + conf.get(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED) || + conf.get(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED) + } + private def loadFromFallbackLocation(appId: String, attemptId: Option[String], logPath: String) : Option[ApplicationInfoWrapper] = { // Call mergeApplicationListing to populate accurate metadata immediately. diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala b/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala index e66933b84af55..0c329e96328c1 100644 --- a/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala +++ b/core/src/main/scala/org/apache/spark/deploy/master/ApplicationInfo.scala @@ -44,6 +44,11 @@ private[spark] class ApplicationInfo( @transient var endTime: Long = _ @transient var appSource: ApplicationSource = _ + // Hold state reported by the driver. Transient, because it belongs to the running driver and + // not to the recovered application: a new Master learns it again when the driver re-registers. + @transient var holdSupported: Boolean = _ + @transient var held: Boolean = _ + @transient private var executorsPerResourceProfileId: mutable.HashMap[Int, mutable.Set[Int]] = _ @transient private var targetNumExecutorsPerResourceProfileId: mutable.HashMap[Int, Int] = _ @transient private var rpIdToResourceProfile: mutable.HashMap[Int, ResourceProfile] = _ @@ -63,6 +68,8 @@ private[spark] class ApplicationInfo( executors = new mutable.HashMap[Int, ExecutorDesc] coresGranted = 0 endTime = -1L + holdSupported = false + held = false appSource = new ApplicationSource(this) nextExecutorId = 0 removedExecutors = new ArrayBuffer[ExecutorDesc] @@ -197,6 +204,40 @@ private[spark] class ApplicationInfo( targetNumExecutorsPerResourceProfileId.values.sum } + /** + * Whether the application is held right now, per the last report from its running driver. An + * application that did not report the hold as supported is not considered held -- its driver + * UI offers no control to change it -- and neither is a finished application, whose driver is + * gone, so the last reported hold is stale. + */ + private[deploy] def isHeld: Boolean = held && holdSupported && !isFinished + + /** + * The number of executors that have not exited yet while the application is held. They are + * still draining their running tasks; the hold is complete once this reaches zero. Counting + * `executors` relies on `DECOMMISSIONED` not being a finished `ExecutorState`: the Master + * keeps a decommissioning executor in the map until it actually exits, which is what makes + * this count agree with the driver's own `draining` from its `/holdstatus` endpoint. + */ + private[deploy] def numDrainingExecutors: Int = if (isHeld) executors.size else 0 + + /** + * The application state, annotated with the hold reported by its driver, for example + * `RUNNING (held, draining 2 executors)`. + */ + private[deploy] def stateText: String = { + if (!isHeld) { + state.toString + } else { + val draining = numDrainingExecutors + if (draining == 0) { + s"$state (held)" + } else { + s"$state (held, draining $draining executor${if (draining > 1) "s" else ""})" + } + } + } + def duration: Long = { if (endTime != -1) { endTime - startTime diff --git a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala index e95c4bd8c6222..ad42aa5a5e1ba 100644 --- a/core/src/main/scala/org/apache/spark/deploy/master/Master.scala +++ b/core/src/main/scala/org/apache/spark/deploy/master/Master.scala @@ -393,6 +393,16 @@ private[deploy] class Master( log" ${MDC(LogKeys.APP_ID, applicationId)}") idToApp.get(applicationId).foreach(finishApplication) + case ApplicationHoldUpdated(appId, supported, held) => + idToApp.get(appId) match { + case Some(app) => + app.holdSupported = supported + app.held = held + case None => + logWarning(log"Got hold status for unknown application " + + log"${MDC(LogKeys.APP_ID, appId)}") + } + case CheckForWorkerTimeOut => timeOutDeadWorkers() diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala b/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala index 3de530f1d252a..095eaa3afe68e 100644 --- a/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala +++ b/core/src/main/scala/org/apache/spark/deploy/master/ui/ApplicationPage.scala @@ -89,7 +89,7 @@ private[ui] class ApplicationPage(parent: MasterWebUI) extends WebUIPage("app") </li> <li><strong>Submit Date:</strong> {UIUtils.formatDate(app.submitDate)}</li> <li><strong>Duration:</strong> {UIUtils.formatDuration(app.duration)}</li> - <li><strong>State:</strong> {app.state}</li> + <li><strong>State:</strong> {app.stateText}</li> { if (!app.isFinished) { if (app.desc.appUiUrl.isBlank()) { diff --git a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala index 2f38b6fcfb4f7..af195adcae54a 100644 --- a/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala +++ b/core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala @@ -348,7 +348,7 @@ private[ui] class MasterPage(parent: MasterWebUI) extends WebUIPage("") { </td> <td>{UIUtils.formatDate(app.submitDate)}</td> <td>{app.desc.user}</td> - <td>{app.state.toString}</td> + <td>{app.stateText}</td> <td sorttable_customkey={app.duration.toString}> {UIUtils.formatDuration(app.duration)} </td> diff --git a/core/src/main/scala/org/apache/spark/deploy/security/HadoopDelegationTokenManager.scala b/core/src/main/scala/org/apache/spark/deploy/security/HadoopDelegationTokenManager.scala index ca202bd9bbfb7..b98a71b13cbad 100644 --- a/core/src/main/scala/org/apache/spark/deploy/security/HadoopDelegationTokenManager.scala +++ b/core/src/main/scala/org/apache/spark/deploy/security/HadoopDelegationTokenManager.scala @@ -24,6 +24,7 @@ import java.util.ServiceLoader import java.util.concurrent.{ScheduledExecutorService, TimeUnit} import scala.collection.mutable +import scala.util.control.NonFatal import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.FileSystem @@ -111,7 +112,16 @@ private[spark] class HadoopDelegationTokenManager( def renewalEnabled: Boolean = { hasKerberosCredentials || (sparkConf.get(DIRECT_CREDENTIAL_PROVIDERS_ENABLED) && - delegationTokenProviders.values.exists(_.delegationTokensRequired(sparkConf, hadoopConf))) + delegationTokenProviders.values.exists { provider => + try { + provider.delegationTokensRequired(sparkConf, hadoopConf) + } catch { + case NonFatal(e) => + logWarning(log"Failed to determine whether credentials are required from " + + log"${MDC(LogKeys.SERVICE_NAME, provider.serviceName)}.", e) + false + } + }) } /** @@ -198,24 +208,20 @@ private[spark] class HadoopDelegationTokenManager( val creds = new Credentials() var failureCount = 0 val nextRenewal = delegationTokenProviders.values.flatMap { provider => - if (provider.delegationTokensRequired(sparkConf, hadoopConf)) { - if (isolateFailures) { - try { - provider.obtainDelegationTokens(hadoopConf, sparkConf, creds) - } catch { - case e: Exception => - logWarning(log"Failed to obtain credentials from " + - log"${MDC(LogKeys.SERVICE_NAME, provider.serviceName)}.", e) - failureCount += 1 - None - } - } else { + try { + if (provider.delegationTokensRequired(sparkConf, hadoopConf)) { provider.obtainDelegationTokens(hadoopConf, sparkConf, creds) + } else { + logDebug(s"Service ${provider.serviceName} does not require a token." + + s" Check your configuration to see if security is disabled or not.") + None } - } else { - logDebug(s"Service ${provider.serviceName} does not require a token." + - s" Check your configuration to see if security is disabled or not.") - None + } catch { + case NonFatal(e) if isolateFailures => + logWarning(log"Failed to obtain credentials from " + + log"${MDC(LogKeys.SERVICE_NAME, provider.serviceName)}.", e) + failureCount += 1 + None } }.foldLeft(Long.MaxValue)(math.min) (creds, nextRenewal, failureCount) diff --git a/core/src/main/scala/org/apache/spark/deploy/security/UserCredentialManager.scala b/core/src/main/scala/org/apache/spark/deploy/security/UserCredentialManager.scala index c98771ee02659..a5e395e547b97 100644 --- a/core/src/main/scala/org/apache/spark/deploy/security/UserCredentialManager.scala +++ b/core/src/main/scala/org/apache/spark/deploy/security/UserCredentialManager.scala @@ -22,6 +22,7 @@ import java.net.URI import java.nio.file.Paths import java.time.Instant import java.util.concurrent.{RejectedExecutionException, ScheduledExecutorService, TimeUnit} +import java.util.concurrent.atomic.AtomicLong import scala.collection.mutable import scala.jdk.CollectionConverters._ @@ -60,11 +61,26 @@ import org.apache.spark.util.{ThreadUtils, Utils} private[spark] class UserCredentialManager( sparkConf: SparkConf, tokenIngestor: TokenIngestor, - onCredentialsUpdate: Array[Byte] => Unit) extends Logging { + onCredentialsUpdate: (Long, Array[Byte]) => Unit, + credentialProviderLoader: CredentialProviderLoader) + extends Logging { + + def this( + sparkConf: SparkConf, + tokenIngestor: TokenIngestor, + onCredentialsUpdate: (Long, Array[Byte]) => Unit) = { + this(sparkConf, tokenIngestor, onCredentialsUpdate, new CredentialProviderLoader()) + } private val safetyMargin = sparkConf.get(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN) private val minInterval = sparkConf.get(SECURITY_OIDC_RENEWAL_MIN_INTERVAL) + // Monotonically increasing version counter for credential updates. + // Incremented on each successful credential acquisition (initial + renewals). + // Used by executors to guard against stale TaskDescription credentials overwriting + // fresher credentials delivered via RPC broadcast. + private val credentialVersion = new AtomicLong(0) + // Counter for exponential backoff calculation. // Only accessed from the single-thread renewal executor. private var consecutiveFailures: Int = 0 @@ -84,10 +100,10 @@ private[spark] class UserCredentialManager( * no credentials can be resolved, an exception is thrown. Subsequent renewal failures * are handled with exponential backoff. * - * @return The serialized initial [[UserCredentials]]. + * @return The version and serialized initial [[UserCredentials]]. * @throws IllegalStateException if the initial credential acquisition fails. */ - def start(): Array[Byte] = { + def start(): (Long, Array[Byte]) = { require(renewalExecutor == null, "start() must not be called more than once") // Initial acquisition is fail-fast (no retry/backoff). @@ -105,11 +121,12 @@ private[spark] class UserCredentialManager( log"${MDC(LogKeys.PRINCIPAL, ctx.getPrincipal)} " + log"(issuer: ${MDC(LogKeys.URI, ctx.getIssuer)})") - val (credentials, earliestExpiry) = resolveCredentials(ctx) + val (credentials, earliestExpiry, activeProviders) = resolveCredentials(ctx) val serialized = UserCredentialManager.serializeUserCredentials(credentials) + val version = credentialVersion.incrementAndGet() // Propagate initial credentials - onCredentialsUpdate(serialized) + onCredentialsUpdate(version, serialized) // Create the renewal executor only after successful initial acquisition. // This avoids leaking a daemon thread if the fail-fast path throws, and @@ -123,12 +140,70 @@ private[spark] class UserCredentialManager( logInfo(log"Credential acquisition successful. Next renewal in " + log"${MDC(LogKeys.TIME_UNITS, UIUtils.formatDuration(renewalDelay))}.") - serialized + + // Apply additional Spark properties declared by active providers. + // Only providers that successfully resolved credentials contribute properties. + // This allows provider modules to wire executor-side configuration + // (e.g., fs.s3a.aws.credentials.provider) without core having + // vendor-specific knowledge. Properties are only set if the user + // has not already configured them explicitly. + for (provider <- activeProviders) { + try { + val props = provider.additionalSparkProperties() + if (props != null) { + props.forEach { (key, value) => + if (!sparkConf.contains(key)) { + sparkConf.set(key, value) + logInfo(log"Auto-configured ${MDC(LogKeys.CONFIG, key)} from " + + log"${MDC(LogKeys.CLASS_NAME, provider.getClass.getName)}") + } else { + logDebug(log"Skipped ${MDC(LogKeys.CONFIG, key)} from " + + log"${MDC(LogKeys.CLASS_NAME, provider.getClass.getName)} " + + log"(already configured)") + } + } + } + } catch { + case scala.util.control.NonFatal(e) => + logWarning(log"Failed to apply additionalSparkProperties from " + + log"${MDC(LogKeys.CLASS_NAME, provider.getClass.getName)}. " + + log"Skipping.", e) + } + } + + (version, serialized) } def stop(): Unit = { + var interrupted = false if (renewalExecutor != null) { renewalExecutor.shutdownNow() + try { + if (!renewalExecutor.awaitTermination( + UserCredentialManager.RENEWAL_SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { + logWarning(log"Timed out waiting for credential renewal to stop; " + + log"closing credential providers while renewal may still be running.") + } + } catch { + case e: InterruptedException => + interrupted = true + logWarning(log"Interrupted while waiting for credential renewal to stop.", e) + } + } + // Close all initialized credential providers to release resources (e.g., HTTP clients). + // This loader belongs to this manager, so closing it cannot affect a later SparkContext. + try { + credentialProviderLoader.closeAll() + } catch { + case e: InterruptedException => + interrupted = true + logWarning(log"Interrupted while closing credential providers during shutdown.", e) + case NonFatal(e) => + logWarning(log"Error closing credential providers during shutdown.", e) + } finally { + if (interrupted) { + Thread.currentThread().interrupt() + } } } @@ -149,13 +224,14 @@ private[spark] class UserCredentialManager( log"${MDC(LogKeys.PRINCIPAL, ctx.getPrincipal)} " + log"(issuer: ${MDC(LogKeys.URI, ctx.getIssuer)})") - val (credentials, earliestExpiry) = resolveCredentials(ctx) + val (credentials, earliestExpiry, _) = resolveCredentials(ctx) val serialized = UserCredentialManager.serializeUserCredentials(credentials) + val version = credentialVersion.incrementAndGet() // Propagate credentials to executors. Errors here are logged separately // so that credential-fetch success is not conflated with distribution failure. try { - onCredentialsUpdate(serialized) + onCredentialsUpdate(version, serialized) } catch { case e: Exception => logWarning(log"Credentials were resolved successfully but failed to propagate " + @@ -200,15 +276,16 @@ private[spark] class UserCredentialManager( * @return Tuple of (UserCredentials, earliest expiry across all service credentials) */ private def resolveCredentials( - ctx: UserContext): (UserCredentials, Option[Instant]) = { + ctx: UserContext): (UserCredentials, Option[Instant], Seq[CredentialProvider]) = { val schemes = discoverSchemes() val credentialMap = new mutable.HashMap[String, ServiceCredential]() + val activeProviders = new mutable.ArrayBuffer[CredentialProvider]() var earliestExpiry: Option[Instant] = None for (scheme <- schemes) { try { - val providerOpt = CredentialProviderLoader.providerFor(scheme, credentialConfMap) + val providerOpt = credentialProviderLoader.providerFor(scheme, credentialConfMap) if (providerOpt.isPresent) { val provider = providerOpt.get() // Use a synthetic target URI with just the scheme for initial resolution. @@ -223,6 +300,7 @@ private[spark] class UserCredentialManager( log"returned null; skipping.") } else { credentialMap.put(scheme, credential) + activeProviders += provider val expiry = credential.getExpiresAt if (expiry != null) { @@ -249,7 +327,7 @@ private[spark] class UserCredentialManager( "Check that providers are on the classpath and configured correctly.") } - (new UserCredentials(credentialMap.asJava), earliestExpiry) + (new UserCredentials(credentialMap.asJava), earliestExpiry, activeProviders.toSeq) } /** @@ -284,7 +362,7 @@ private[spark] class UserCredentialManager( // available on the classpath by probing CredentialProviderLoader. // This covers both built-in providers (e.g., connector/credential-aws for "s3a") // and third-party providers registered via ServiceLoader. - CredentialProviderLoader.discoverAllSchemes().asScala.toSet + credentialProviderLoader.discoverAllSchemes().asScala.toSet } } @@ -360,6 +438,8 @@ private[spark] class UserCredentialManager( private[spark] object UserCredentialManager { + private val RENEWAL_SHUTDOWN_TIMEOUT_SECONDS = 10L + /** * Synthetic authority used in target URIs for scheme-based provider resolution. * Providers should not rely on this value; it signals that no specific endpoint @@ -404,7 +484,7 @@ private[spark] object UserCredentialManager { */ def create( sparkConf: SparkConf, - onCredentialsUpdate: Array[Byte] => Unit): Option[UserCredentialManager] = { + onCredentialsUpdate: (Long, Array[Byte]) => Unit): Option[UserCredentialManager] = { if (!sparkConf.get(SECURITY_OIDC_ENABLED)) { None } else { diff --git a/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala b/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala index 4813f3f94107e..d23469da5b4e0 100644 --- a/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala +++ b/core/src/main/scala/org/apache/spark/errors/SparkCoreErrors.scala @@ -107,19 +107,25 @@ private[spark] object SparkCoreErrors { def cannotUseMapSideCombiningWithArrayKeyError(): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3008", messageParameters = Map.empty, cause = null + errorClass = "UNSUPPORTED_ARRAY_KEY.MAP_SIDE_COMBINE", + messageParameters = Map.empty, + cause = null ) } def hashPartitionerCannotPartitionArrayKeyError(): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3009", messageParameters = Map.empty, cause = null + errorClass = "UNSUPPORTED_ARRAY_KEY.HASH_PARTITIONER", + messageParameters = Map.empty, + cause = null ) } def reduceByKeyLocallyNotSupportArrayKeysError(): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3010", messageParameters = Map.empty, cause = null + errorClass = "UNSUPPORTED_ARRAY_KEY.REDUCE_BY_KEY_LOCALLY", + messageParameters = Map.empty, + cause = null ) } @@ -140,7 +146,7 @@ private[spark] object SparkCoreErrors { } def emptyCollectionError(): Throwable = { - new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3014") + new SparkUnsupportedOperationException("EMPTY_COLLECTION_NOT_ALLOWED") } def countByValueApproxNotSupportArraysError(): Throwable = { @@ -151,22 +157,30 @@ private[spark] object SparkCoreErrors { def checkpointDirectoryHasNotBeenSetInSparkContextError(): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3016", messageParameters = Map.empty, cause = null + errorClass = "CHECKPOINT_DIRECTORY_NOT_SET", + messageParameters = Map.empty, + cause = null ) } - def invalidCheckpointFileError(path: Path): Throwable = { + def invalidCheckpointDirectoryError( + partitionFilePath: Path, + expectedFileName: String): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3017", - messageParameters = Map("path" -> s"$path"), + errorClass = "INVALID_CHECKPOINT_DIRECTORY", + messageParameters = Map( + "path" -> s"${partitionFilePath.getParent}", + "expectedFileName" -> expectedFileName, + "fileName" -> partitionFilePath.getName + ), cause = null ) } def failToCreateCheckpointPathError(checkpointDirPath: Path): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3018", - messageParameters = Map("checkpointDirPath" -> s"$checkpointDirPath"), + errorClass = "FAILED_CREATE_CHECKPOINT_DIRECTORY", + messageParameters = Map("path" -> s"$checkpointDirPath"), cause = null ) } @@ -177,7 +191,7 @@ private[spark] object SparkCoreErrors { newRDDId: Int, newRDDLength: Int): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3019", + errorClass = "CHECKPOINT_RDD_PARTITION_COUNT_MISMATCH", messageParameters = Map( "originalRDDId" -> s"$originalRDDId", "originalRDDLength" -> s"$originalRDDLength", @@ -194,20 +208,21 @@ private[spark] object SparkCoreErrors { } def mustSpecifyCheckpointDirError(): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3020", messageParameters = Map.empty, cause = null - ) + SparkException.internalError( + "SparkContext.checkpointDir is unset when creating ReliableRDDCheckpointData.") } def askStandaloneSchedulerToShutDownExecutorsError(e: Exception): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3021", messageParameters = Map.empty, cause = e + errorClass = "SCHEDULER_BACKEND_SHUTDOWN_FAILED.EXECUTORS", + messageParameters = Map.empty, cause = e ) } def stopStandaloneSchedulerDriverEndpointError(e: Exception): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3022", messageParameters = Map.empty, cause = e + errorClass = "SCHEDULER_BACKEND_SHUTDOWN_FAILED.DRIVER_ENDPOINT", + messageParameters = Map.empty, cause = e ) } @@ -242,29 +257,26 @@ private[spark] object SparkCoreErrors { } def cannotRunSubmitMapStageOnZeroPartitionRDDError(): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3023", messageParameters = Map.empty, cause = null - ) + SparkException.internalError("Can't run submitMapStage on RDD with 0 partitions.") } def accessNonExistentAccumulatorError(id: Long): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3024", messageParameters = Map("id" -> s"$id"), cause = null - ) + SparkException.internalError(s"Attempted to access non-existent accumulator $id.") } def sendResubmittedTaskStatusForShuffleMapStagesOnlyError(): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3025", messageParameters = Map.empty, cause = null - ) + SparkException.internalError( + "TaskSetManagers should only send Resubmitted task statuses for tasks in ShuffleMapStages.") } def nonEmptyEventQueueAfterTimeoutError(timeoutMillis: Long): Throwable = { new TimeoutException(s"The event queue is not empty after $timeoutMillis ms.") } - def durationCalledOnUnfinishedTaskError(): Throwable = { - new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3026") + def durationCalledOnUnfinishedTaskError(className: String, methodName: String): Throwable = { + new SparkUnsupportedOperationException( + errorClass = "UNSUPPORTED_CALL.TASK_NOT_FINISHED", + messageParameters = Map("className" -> className, "methodName" -> methodName)) } def sparkError(errorMsg: String): Throwable = { @@ -277,7 +289,7 @@ private[spark] object SparkCoreErrors { def clusterSchedulerError(message: String): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3029", + errorClass = "CLUSTER_MANAGER_APPLICATION_FAILURE", messageParameters = Map("message" -> message), cause = null ) @@ -292,22 +304,13 @@ private[spark] object SparkCoreErrors { } def taskHasNotLockedBlockError(currentTaskAttemptId: Long, blockId: BlockId): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3030", - messageParameters = Map( - "currentTaskAttemptId" -> s"$currentTaskAttemptId", - "blockId" -> s"$blockId" - ), - cause = null - ) + SparkException.internalError( + s"Task $currentTaskAttemptId has not locked block $blockId for writing.", + category = "STORAGE") } def blockDoesNotExistError(blockId: BlockId): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3031", - messageParameters = Map("blockId" -> s"$blockId"), - cause = null - ) + SparkException.internalError(s"Block $blockId does not exist.", category = "STORAGE") } def cannotSaveBlockOnDecommissionedExecutorError(blockId: BlockId): Throwable = { @@ -315,54 +318,47 @@ private[spark] object SparkCoreErrors { } def waitingForReplicationToFinishError(e: Throwable): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3032", messageParameters = Map.empty, cause = e - ) + SparkException.internalError("Error occurred while waiting for replication to finish.", e) } def unableToRegisterWithExternalShuffleServerError(e: Throwable): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3033", - messageParameters = Map("message" -> e.getMessage), + errorClass = "UNABLE_TO_REGISTER_WITH_EXTERNAL_SHUFFLE_SERVICE", + messageParameters = Map("message" -> Option(e.getMessage).getOrElse(e.toString)), cause = e ) } def waitingForAsyncReregistrationError(e: Throwable): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3034", messageParameters = Map.empty, cause = e - ) + SparkException.internalError("Error occurred while waiting for async. reregistration.", e) } - def unexpectedShuffleBlockWithUnsupportedResolverError( + def shuffleBlockMigrationNotSupportedError( + blockId: BlockId, shuffleBlockResolver: ShuffleBlockResolver, - blockId: BlockId): Throwable = { + e: Throwable): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3035", + errorClass = "SHUFFLE_BLOCK_MIGRATION_NOT_SUPPORTED", messageParameters = Map( "blockId" -> s"$blockId", - "shuffleBlockResolver" -> s"$shuffleBlockResolver" + "resolverClass" -> shuffleBlockResolver.getClass.getName ), - cause = null + cause = e ) } def failToStoreBlockOnBlockManagerError( blockManagerId: BlockManagerId, blockId: BlockId): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3036", - messageParameters = Map( - "blockId" -> s"$blockId", - "blockManagerId" -> s"$blockManagerId" - ), - cause = null - ) + SparkException.internalError( + s"Failed to store block $blockId on $blockManagerId. This mostly happens when there is " + + "not enough storage memory for the block and its storage level has no disk fallback.", + category = "STORAGE") } - def readLockedBlockNotFoundError(blockId: BlockId): Throwable = { + def localBlockDataNotFoundError(blockId: BlockId): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3037", + errorClass = "LOCAL_BLOCK_DATA_NOT_FOUND", messageParameters = Map( "blockId" -> s"$blockId" ), @@ -371,13 +367,8 @@ private[spark] object SparkCoreErrors { } def failToGetBlockWithLockError(blockId: BlockId): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3038", - messageParameters = Map( - "blockId" -> s"$blockId" - ), - cause = null - ) + SparkException.internalError( + s"get() failed for block $blockId even though we held a lock.", category = "STORAGE") } def blockNotFoundError(blockId: BlockId): Throwable = { @@ -389,17 +380,14 @@ private[spark] object SparkCoreErrors { } def blockStatusQueryReturnedNullError(blockId: BlockId): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3039", - messageParameters = Map("blockId" -> s"$blockId"), - cause = null - ) + SparkException.internalError( + s"BlockManager returned null for BlockStatus query: $blockId.", category = "STORAGE") } - def unexpectedBlockManagerMasterEndpointResultError(): Throwable = { - new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3040", messageParameters = Map.empty, cause = null - ) + def unexpectedBlockManagerMasterEndpointResultError(message: Any): Throwable = { + SparkException.internalError( + s"BlockManagerMasterEndpoint returned false for message $message, expected true.", + category = "STORAGE") } def failToCreateDirectoryError(path: String, maxAttempts: Int): Throwable = { @@ -407,10 +395,6 @@ private[spark] object SparkCoreErrors { s"Failed to create directory ${path} with permission 770 after $maxAttempts attempts!") } - def unsupportedOperationError(): Throwable = { - new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3041") - } - def noSuchElementError(): Throwable = { new NoSuchElementException() } @@ -428,10 +412,10 @@ private[spark] object SparkCoreErrors { def failToGetNonShuffleBlockError(blockId: BlockId, e: Throwable): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_3042", - messageParameters = Map("blockId" -> s"$blockId"), - cause = e - ) + errorClass = "INTERNAL_ERROR_STORAGE", + messageParameters = Map( + "message" -> s"Failed to get block $blockId, which is not a shuffle block."), + cause = e) } def graphiteSinkInvalidProtocolError(invalidProtocol: String): Throwable = { @@ -448,12 +432,16 @@ private[spark] object SparkCoreErrors { cause = null) } - def outOfMemoryError(requestedBytes: Long, receivedBytes: Long): OutOfMemoryError = { + def outOfMemoryError( + requestedBytes: Long, + receivedBytes: Long, + consumerBreakdown: String): OutOfMemoryError = { new SparkOutOfMemoryError( "UNABLE_TO_ACQUIRE_MEMORY", Map( "requestedBytes" -> requestedBytes.toString, - "receivedBytes" -> receivedBytes.toString).asJava) + "receivedBytes" -> receivedBytes.toString, + "consumerBreakdown" -> consumerBreakdown).asJava) } def failedRenameTempFileError(srcFile: File, dstFile: File): Throwable = { diff --git a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala index 206a6a0fe385c..375682383b1f8 100644 --- a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala +++ b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala @@ -223,6 +223,12 @@ private[spark] class CoarseGrainedExecutorBackend( logInfo(log"Received tokens of ${MDC(LogKeys.NUM_BYTES, tokenBytes.length)} bytes") SparkHadoopUtil.get.addDelegationTokens(tokenBytes, env.conf) + case UpdateUserCredentials(version, credentials) => + logInfo(log"Received user credentials of " + + log"${MDC(LogKeys.NUM_BYTES, credentials.length)} bytes " + + log"(version ${MDC(LogKeys.CREDENTIAL_VERSION, version)})") + VersionedCredentials.updateIfNewer(env.userCredentials, version, credentials) + case DecommissionExecutor => decommissionSelf() } @@ -505,6 +511,18 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging { } val env = SparkEnv.createExecutorEnv(driverConf, arguments.executorId, arguments.bindAddress, arguments.hostname, arguments.cores, cfg.ioEncryptionKey, isLocal = false) + + // Apply initial user credentials to the executor credential store. + // Uses unconditional set() rather than updateIfNewer() because the store is guaranteed + // null at this point (executor startup, before any RPC or task is received). + // Note: there is a narrow window where a renewal broadcast (vN+1) could arrive between + // the SparkAppConfig reply (vN) and executor registration in executorDataMap, leaving + // this executor on vN until vN+2. The TaskDescription path covers this case since + // every dispatched task carries the latest credentials. + cfg.userCredentials.foreach { case (version, credentials) => + env.userCredentials.set(VersionedCredentials(version, credentials)) + } + // Set the application attemptId in the BlockStoreClient if available. val appAttemptId = env.conf.get(APP_ATTEMPT_ID) appAttemptId.foreach(attemptId => diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala index f288ceacef7e6..3056c3b38ad69 100644 --- a/core/src/main/scala/org/apache/spark/executor/Executor.scala +++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala @@ -847,6 +847,15 @@ private[spark] class Executor( // requires access to properties contained within (e.g. for access control). Executor.taskDeserializationProps.set(taskDescription.properties) + // Apply user credentials from TaskDescription to the executor credential store. + // This ensures credentials are available before any task code runs, avoiding + // the race between RPC broadcast and task dispatch. + // Only apply if the version is newer than what the store already has, preventing + // a delayed TaskDescription from overwriting fresher credentials delivered via RPC. + taskDescription.userCredentials.foreach { case (version, creds) => + VersionedCredentials.updateIfNewer(env.userCredentials, version, creds) + } + updateDependencies( taskDescription.artifacts.files, taskDescription.artifacts.jars, @@ -1602,7 +1611,7 @@ private[spark] object Executor extends Logging { * minus executorRunTime -- would misreport the un-consumed share of the core as scheduler * delay. Relative comparisons for speculation are unaffected because all tasks of a stage * share the same cpu amount. Trailing zeros are stripped from the scale-9 cpus so the - * product stays in BigDecimal's compact (long-backed) form instead of inflating a long + * product stays in BigDecimal's compact (long-backed) form instead of inflating a long * duration into a BigInteger. */ private[spark] def cpuWeightedNanos(intervalNs: Long, cpus: BigDecimal): Long = { diff --git a/core/src/main/scala/org/apache/spark/internal/config/History.scala b/core/src/main/scala/org/apache/spark/internal/config/History.scala index 1936ecf68103d..570113d145a14 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/History.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/History.scala @@ -193,6 +193,14 @@ private[spark] object History { .booleanConf .createWithDefault(true) + val EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED = + ConfigBuilder("spark.history.fs.eventLog.onDemandLoadEnabled") + .doc("Whether to look up single event log locations on demand manner before listing files.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val DRIVER_LOG_CLEANER_ENABLED = ConfigBuilder("spark.history.fs.driverlog.cleaner.enabled") .version("3.0.0") .doc("Specifies whether the History Server should periodically clean up driver logs from " + diff --git a/core/src/main/scala/org/apache/spark/internal/config/UI.scala b/core/src/main/scala/org/apache/spark/internal/config/UI.scala index 65f26e9ca2799..e70743fa9df0e 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/UI.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/UI.scala @@ -92,6 +92,19 @@ private[spark] object UI { .booleanConf .createWithDefault(true) + val UI_HOLD_ENABLED = ConfigBuilder("spark.ui.holdEnabled") + .doc("Allows the whole application to be held and resumed from the web UI. Holding " + + "gracefully decommissions all executors and stops requesting new ones. Cached blocks " + + "are not preserved and are recomputed after resuming. This takes effect only when " + + "spark.decommission.enabled is true, the shuffle data is kept outside the executors " + + "(through either spark.shuffle.service.enabled or a ShuffleDataIO with reliable " + + "storage), and the cluster manager can hold executors: Standalone, YARN, and " + + "Kubernetes with spark.kubernetes.allocation.pods.allocator=direct.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val UI_THREAD_DUMPS_ENABLED = ConfigBuilder("spark.ui.threadDumpsEnabled") .doc("Whether to show a link for executor thread dumps in Stages and Executor pages.") .version("1.2.0") diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala b/core/src/main/scala/org/apache/spark/internal/config/package.scala index 8038307afc65a..6303f751c9c14 100644 --- a/core/src/main/scala/org/apache/spark/internal/config/package.scala +++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala @@ -503,6 +503,21 @@ package object config { .doubleConf .createWithDefault(0.6) + private[spark] val MEMORY_OOM_ERROR_CONSUMER_BREAKDOWN_LIMIT = + ConfigBuilder("spark.memory.oomErrorConsumerBreakdownLimit") + .internal() + .doc("The maximum number of memory consumers listed individually in the per-consumer " + + "memory breakdown attached to an UNABLE_TO_ACQUIRE_MEMORY error. The largest consumers " + + "are listed first; any beyond this limit are collapsed into a single summary line. This " + + "bounds the size of the error message that is sent to the driver and shown in the UI. " + + "It does not affect the full breakdown written to the executor logs. Set to 0 to omit " + + "the breakdown from the error message entirely.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .intConf + .checkValue(_ >= 0, "The consumer breakdown limit must not be negative") + .createWithDefault(5) + private[spark] val UNMANAGED_MEMORY_POLLING_INTERVAL = ConfigBuilder("spark.memory.unmanagedMemoryPollingInterval") .doc("Interval for polling unmanaged memory users to track their memory usage. " + @@ -1695,7 +1710,7 @@ package object config { .doc("Whether to enable OIDC credential propagation. When enabled, the driver reads an " + "identity token from a file, exchanges it for short-lived service credentials via " + "CredentialProvider implementations, and propagates those credentials to executors.") - .version("4.3.0") + .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .booleanConf .createWithDefault(false) @@ -1705,7 +1720,7 @@ package object config { .doc("Path to the OIDC identity token file on the driver. Required when " + "spark.security.oidc.enabled is true. The file should contain a JWT token " + "(e.g., a Kubernetes projected service account token).") - .version("4.3.0") + .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .stringConf .createOptional @@ -1714,7 +1729,7 @@ package object config { ConfigBuilder("spark.security.oidc.renewal.safetyMargin") .doc("How long before credential expiry to trigger renewal. Credentials are refreshed " + "at min(identity token expiry, service credential expiry) minus this margin.") - .version("4.3.0") + .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .timeConf(TimeUnit.MILLISECONDS) .checkValue(_ > 0, "The safety margin must be a positive time value.") @@ -1724,7 +1739,7 @@ package object config { ConfigBuilder("spark.security.oidc.renewal.minInterval") .doc("Minimum interval between credential renewal attempts. This prevents tight renewal " + "loops when credentials have very short TTLs or when failures cause rapid retries.") - .version("4.3.0") + .version("4.4.0") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) .timeConf(TimeUnit.MILLISECONDS) .checkValue(_ > 0, "The minimum renewal interval must be a positive time value.") @@ -1920,6 +1935,18 @@ package object config { .longConf .createWithDefault(50) + private[spark] val STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS = + ConfigBuilder("spark.shuffle.streaming.writerConnectionTimeout") + .doc("Maximum time a streaming shuffle writer waits for each reader to connect. " + + "Set to -1 to wait indefinitely.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .timeConf(TimeUnit.MILLISECONDS) + .checkValue( + timeoutMs => timeoutMs == -1 || timeoutMs > 0, + "The reader connection timeout must be positive or -1 to wait indefinitely.") + .createWithDefaultString("1h") + private[spark] val STREAMING_SHUFFLE_WRITER_MAX_MEMORY = ConfigBuilder("spark.shuffle.streaming.writerMaxMemory") .doc("Best-effort memory limit in bytes for in-flight data buffers in a streaming " + diff --git a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala index 639b82b6080b3..dac82cc763577 100644 --- a/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala +++ b/core/src/main/scala/org/apache/spark/memory/MemoryManager.scala @@ -62,6 +62,14 @@ private[spark] abstract class MemoryManager( protected[this] val offHeapStorageMemory = (maxOffHeapMemory * conf.get(MEMORY_STORAGE_FRACTION)).toLong + /** + * The maximum number of consumers listed individually in the per-consumer memory breakdown + * attached to an UNABLE_TO_ACQUIRE_MEMORY error. Read by [[TaskMemoryManager]]. See + * `spark.memory.oomErrorConsumerBreakdownLimit`. + */ + private[memory] val oomErrorConsumerBreakdownLimit: Int = + conf.get(MEMORY_OOM_ERROR_CONSUMER_BREAKDOWN_LIMIT) + offHeapExecutionMemoryPool.incrementPoolSize(maxOffHeapMemory - offHeapStorageMemory) offHeapStorageMemoryPool.incrementPoolSize(offHeapStorageMemory) diff --git a/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala b/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala index 6b278c47f32f1..954dc0ae06c44 100644 --- a/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala +++ b/core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala @@ -447,7 +447,11 @@ object UnifiedMemoryManager extends Logging { } def apply(conf: SparkConf, numCores: Int): UnifiedMemoryManager = { - val maxMemory = getMaxMemory(conf) + apply(conf, numCores, isDriver = true) + } + + def apply(conf: SparkConf, numCores: Int, isDriver: Boolean): UnifiedMemoryManager = { + val maxMemory = getMaxMemory(conf, isDriver) new UnifiedMemoryManager( conf, maxHeapMemory = maxMemory, @@ -459,25 +463,34 @@ object UnifiedMemoryManager extends Logging { /** * Return the total amount of memory shared between execution and storage, in bytes. */ - private def getMaxMemory(conf: SparkConf): Long = { + private def getMaxMemory(conf: SparkConf, isDriver: Boolean): Long = { val systemMemory = conf.get(TEST_MEMORY) val reservedMemory = conf.getLong(TEST_RESERVED_MEMORY.key, if (conf.contains(IS_TESTING)) 0 else RESERVED_SYSTEM_MEMORY_BYTES) val minSystemMemory = (reservedMemory * 1.5).ceil.toLong if (systemMemory < minSystemMemory) { - throw new SparkIllegalArgumentException( - errorClass = "INVALID_DRIVER_MEMORY", - messageParameters = Map( - "systemMemory" -> systemMemory.toString, - "minSystemMemory" -> minSystemMemory.toString, - "config" -> config.DRIVER_MEMORY.key)) + if (isDriver) { + throw new SparkIllegalArgumentException( + errorClass = "INVALID_DRIVER_MEMORY.SYSTEM_MEMORY", + messageParameters = Map( + "systemMemory" -> systemMemory.toString, + "minSystemMemory" -> minSystemMemory.toString, + "config" -> config.DRIVER_MEMORY.key)) + } else { + throw new SparkIllegalArgumentException( + errorClass = "INVALID_EXECUTOR_MEMORY.SYSTEM_MEMORY", + messageParameters = Map( + "systemMemory" -> systemMemory.toString, + "minSystemMemory" -> minSystemMemory.toString, + "config" -> config.EXECUTOR_MEMORY.key)) + } } // SPARK-12759 Check executor memory to fail fast if memory is insufficient if (conf.contains(config.EXECUTOR_MEMORY)) { val executorMemory = conf.getSizeAsBytes(config.EXECUTOR_MEMORY.key) if (executorMemory < minSystemMemory) { throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXECUTOR_MEMORY", + errorClass = "INVALID_EXECUTOR_MEMORY.CONFIG_MEMORY", messageParameters = Map( "executorMemory" -> executorMemory.toString, "minSystemMemory" -> minSystemMemory.toString, diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala index fd42cea795d60..244d79aa217a4 100644 --- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala +++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala @@ -24,7 +24,7 @@ import scala.reflect.ClassTag import scala.util.control.NonFatal import com.google.common.cache.{CacheBuilder, CacheLoader} -import org.apache.hadoop.fs.Path +import org.apache.hadoop.fs.{FileAlreadyExistsException, Path} import org.apache.spark._ import org.apache.spark.broadcast.Broadcast @@ -78,8 +78,9 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag]( .sortBy(_.getName.stripPrefix("part-").toInt) // Fail fast if input files are invalid inputFiles.zipWithIndex.foreach { case (path, i) => - if (path.getName != ReliableCheckpointRDD.checkpointFileName(i)) { - throw SparkCoreErrors.invalidCheckpointFileError(path) + val expectedFileName = ReliableCheckpointRDD.checkpointFileName(i) + if (path.getName != expectedFileName) { + throw SparkCoreErrors.invalidCheckpointDirectoryError(path, expectedFileName) } } Array.tabulate(inputFiles.length)(i => new CheckpointRDDPartition(i)) @@ -225,7 +226,19 @@ private[spark] object ReliableCheckpointRDD extends Logging { serializeStream.close() }) - if (!fs.rename(tempOutputPath, finalOutputPath)) { + // On HDFS, renaming onto an existing destination reports failure by returning false, which + // is handled below. Some FileSystem implementations instead raise FileAlreadyExistsException + // (e.g. S3A since HADOOP-16721, ABFS); treat it the same way, as it means another attempt of + // this task has already committed the final output (SPARK-58750). + val renamed = try { + fs.rename(tempOutputPath, finalOutputPath) + } catch { + case e: FileAlreadyExistsException => + logDebug(log"Rename from ${MDC(TEMP_OUTPUT_PATH, tempOutputPath)} to" + + log" ${MDC(FINAL_OUTPUT_PATH, finalOutputPath)} failed", e) + false + } + if (!renamed) { if (!fs.exists(finalOutputPath)) { logInfo(log"Deleting tempOutputPath ${MDC(TEMP_OUTPUT_PATH, tempOutputPath)}") fs.delete(tempOutputPath, false) diff --git a/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala b/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala index cdf1a1ec7bce2..8341a723508d8 100644 --- a/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala +++ b/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala @@ -32,9 +32,9 @@ import org.apache.spark.annotation.{Since, Stable} * numbers, since a task's amount must map onto discrete resource addresses - * ie amount equals 0.5 translates into 2 tasks per resource address. CPUs * (resource name "cpus") are a plain quantity drawn from the executor's core - * pool rather than an addressable resource, so any amount of at least 1e-9 is - * valid, e.g. 1.5; the cpus amount is rounded to the nearest 1e-9, so precision - * beyond 9 decimal places is not preserved. + * pool rather than an addressable resource, so any amount from 1e-9 through + * Int.MaxValue after rounding is valid, e.g. 1.5; the cpus amount is rounded to + * the nearest 1e-9, so precision beyond 9 decimal places is not preserved. */ @Stable @Since("3.1.0") diff --git a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala index 9876668194a84..0ea568693aebc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ActiveJob.scala @@ -48,7 +48,14 @@ private[spark] class ActiveJob( val callSite: CallSite, val listener: JobListener, val artifacts: JobArtifactSet, - val properties: Properties) { + val properties: Properties, + // Whether this job's RDD graph uses a `PipelinedShuffleDependency`. Every pipelined-group + // scheduling path -- co-scheduling, deferral, and the per-submit `TaskSet.isPipelined` tag -- + // is inert for a job without one, so this flag lets those paths short-circuit the + // group-membership graph walk for the common regular job at no cost. Defaults to false; the + // result-job path (handleJobSubmitted) passes the value it computed up front, and the + // map-stage-job path (which rejects pipelined dependencies) leaves it false. + val hasPipelinedDependency: Boolean = false) { /** * Number of partitions we need to compute for this job. Note that result stages may not need diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index bfd24ceb2f895..fa2ab84c64e78 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -45,7 +45,7 @@ import org.apache.spark.network.shuffle.{BlockStoreClient, MergeFinalizerListene import org.apache.spark.network.shuffle.protocol.MergeStatuses import org.apache.spark.network.util.JavaUtils import org.apache.spark.partial.{ApproximateActionListener, ApproximateEvaluator, PartialResult} -import org.apache.spark.rdd.{RDD, RDDCheckpointData} +import org.apache.spark.rdd.{DeterministicLevel, RDD, RDDCheckpointData, ReliableRDDCheckpointData} import org.apache.spark.resource.{CpuAmount, ResourceProfile, TaskResourceProfile} import org.apache.spark.resource.ResourceProfile.{CPUS, DEFAULT_RESOURCE_PROFILE_ID, EXECUTOR_CORES_LOCAL_PROPERTY, MAX_TASKS_PER_EXECUTOR_LOCAL_PROPERTY, PYSPARK_MEMORY_LOCAL_PROPERTY} import org.apache.spark.rpc.RpcTimeout @@ -187,6 +187,11 @@ private[spark] class DAGScheduler( delayedTaskCompletionEvents: ListBuffer[CompletionEvent] = new ListBuffer[CompletionEvent]) private[scheduler] val dependentStageMap = new HashMap[Stage, DependentStageInfo] + // Whether we have already logged that the pipelined-group slot check is disabled. Only the + // (single-threaded) event loop touches this, so a plain var is safe. Limits the warning to once + // per scheduler rather than once per submitted batch job. + private var warnedPipelinedSlotCheckDisabled = false + private[scheduler] val activeJobs = new HashSet[ActiveJob] // Track all the jobs submitted by the same query execution, will clean up after @@ -401,6 +406,42 @@ private[spark] class DAGScheduler( */ private[scheduler] val barrierJobIdToNumTasksCheckFailures = new ConcurrentHashMap[Int, Int] + /** + * The barrier jobs deferred while their max concurrent tasks check is being retried (the + * submission is re-posted on a timer), keyed by job id. A deferred job is registered nowhere + * else, so this is what lets a cancellation fail it: the cancellation swaps in + * `deferredJobCancelledMarker` and `handleJobSubmitted` drops the re-posted submission when + * it finds the marker. + */ + private[scheduler] val deferredBarrierJobs = + new ConcurrentHashMap[Int, DAGScheduler.DeferredBarrierJob] + + /** + * Marker left in `deferredBarrierJobs` when a deferred job is cancelled. The entry is + * replaced rather than removed so that the decision to drop the pending re-post is made on + * the event loop (in `handleJobSubmitted`), atomically with respect to the cancellation; the + * re-post itself always fires and is what finally clears the entry. + */ + private val deferredJobCancelledMarker = DAGScheduler.DeferredBarrierJob(null, null) + + /** (id, properties) of the deferred barrier jobs whose submission-time properties match `p`. */ + private def deferredJobsMatching(p: Properties => Boolean): Seq[(Int, Properties)] = { + val matched = mutable.ArrayBuffer[(Int, Properties)]() + deferredBarrierJobs.forEach { (jobId, deferred) => + if ((deferred ne deferredJobCancelledMarker) && p(deferred.properties)) { + matched += ((jobId, deferred.properties)) + } + } + matched.toSeq + } + + /** + * Whether the executors are held (see `SparkContext.holdExecutors()`). While held, the + * barrier slot check reads zero capacity, so its retry budget must not be consumed: the job + * should wait for the resume. Extracted as a seam so tests can control the hold state. + */ + protected def executorsHeld: Boolean = sc.executorsHeld + /** * Time in seconds to wait between a max concurrent tasks check failure and the next check. */ @@ -622,6 +663,21 @@ private[spark] class DAGScheduler( firstJobId: Int): ShuffleMapStage = { shuffleIdToMapStage.get(shuffleDep.shuffleId) match { case Some(stage) => + // A pipelined shuffle is transient: it is a once-through live stream with no retained, + // addressable output, so reusing its producer stage across jobs is unsound (there is no + // durable output for a second job to read). Reuse must be prevented explicitly -- from the + // scheduler's view a shuffle-map stage can be reused unless something forbids it. If a + // pipelined dependency's shuffleId is already bound to a stage from a different job, that + // is the forbidden cross-job reuse; fail fast. (Within the same job the cached stage is the + // one we just created, so returning it is correct and not reuse.) + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] && + !stage.jobIds.contains(firstJobId)) { + throw new SparkException( + errorClass = "PIPELINED_SHUFFLE_CROSS_JOB_REUSE", + messageParameters = scala.collection.immutable.Map( + "shuffleId" -> shuffleDep.shuffleId.toString), + cause = null) + } stage case None => @@ -671,6 +727,12 @@ private[spark] class DAGScheduler( checkBarrierStageWithDynamicAllocation(rdd) checkBarrierStageWithNumSlots(rdd, resourceProfile) checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions) + checkPipelinedProducerSupported(shuffleDep) + // Resolve the tracker that owns this shuffle's output, by dependency type (see + // outputTrackerMaster). Resolve it up front, BEFORE any stage-map mutation below, so that a + // fail-loud on a misconfigured pipelined shuffle (no StreamingShuffleOutputTracker) leaves no + // partial scheduler state and a re-submit re-throws the same error. + val outputTracker = outputTrackerMaster(shuffleDep) val numTasks = rdd.partitions.length val parents = getOrCreateParentStages(shuffleDeps, jobId) val id = nextStageId.getAndIncrement() @@ -682,18 +744,108 @@ private[spark] class DAGScheduler( shuffleIdToMapStage(shuffleDep.shuffleId) = stage updateJobIdStageIdMaps(jobId, stage) - if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) { - // Kind of ugly: need to register RDDs with the cache and map output tracker here - // since we can't do it in the RDD constructor because # of partitions is unknown + // Register the shuffle with its own tracker (a pipelined shuffle in the + // StreamingShuffleOutputTracker, a regular one in the MapOutputTracker -- split by dependency + // type, no overlap). Self-guarded on the tracker's own membership: createShuffleMapStage runs + // once per shuffleId via getOrCreateShuffleMapStage, but guard defensively against re-entry. + // (A pipelined shuffle is never registered with the MapOutputTracker; its availability is + // tracked on the stage via pipelinedCompletedPartitions and its writers are located through the + // StreamingShuffleOutputTracker. jobId is used only by the streaming tracker.) The + // MapOutputTracker.getStatistics paths, which WOULD throw ShuffleStatusNotFoundException on the + // absent entry, are unreachable for a pipelined dependency: markMapStageJobsAsFinished calls it + // only when mapStageJobs is non-empty, but handleMapStageSubmitted rejects a pipelined dep + // before addActiveJob (its sole populator) runs, so a pipelined stage's mapStageJobs is always + // empty; and checkAndScheduleShuffleMergeFinalize's getStatistics is on the push-based-merge + // path, which a pipelined dependency rejects up front (checkPipelinedProducerSupported). + if (!outputTracker.containsShuffle(shuffleDep.shuffleId)) { logInfo(log"Registering RDD ${MDC(RDD_ID, rdd.id)} " + log"(${MDC(CREATION_SITE, rdd.getCreationSite)}) as input to " + log"shuffle ${MDC(SHUFFLE_ID, shuffleDep.shuffleId)}") - mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, - shuffleDep.partitioner.numPartitions) + outputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length, + shuffleDep.partitioner.numPartitions, jobId) } stage } + /** + * The driver-side output tracker that owns a shuffle's outputs, selected by dependency type: the + * StreamingShuffleOutputTracker for a pipelined shuffle, the MapOutputTracker otherwise. The two + * are split with no overlap. A pipelined shuffle REQUIRES a StreamingShuffleOutputTracker + * (created with a streaming-capable shuffle manager, see + * SparkEnv.initializeStreamingShuffleOutputTracker), so fail loud if one is not configured rather + * than silently register it nowhere (a consumer would then find no writer locations; the reader + * enforces the same invariant, see StreamingShuffleReader). + */ + private def outputTrackerMaster( + shuffleDep: ShuffleDependency[_, _, _]): ShuffleOutputTrackerMaster = { + if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + sc.env.streamingShuffleOutputTracker + .getOrElse(throw new IllegalStateException( + s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " + + "StreamingShuffleOutputTracker, but none is configured")) + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + } else { + mapOutputTracker + } + } + + private def pipelinedUnsupportedError(reason: String): PipelinedShuffleUnsupportedException = + new PipelinedShuffleUnsupportedException(reason) + + /** + * Fail-fast on producer-side idioms a pipelined shuffle cannot support, checked when the producer + * stage is created. A pipelined shuffle runs its producer and consumer stages concurrently over a + * transient, once-through stream that a group never recomputes in isolation (any failure aborts + * the whole group), so mechanisms that recompute/roll back a single stage are moot, and features + * that expose output only after a global barrier are incompatible with + * incremental reads. Rejecting here (before the stage is used) keeps a misuse from silently + * mis-scheduling. Inert for a regular ShuffleDependency. + * + * Group-level idioms are handled elsewhere, since they are properties of the group rather than a + * single producer stage: fan-out (a producer with more than one consumer) and a group with a + * non-default resource profile are rejected up front at job submission by + * `checkPipelinedGroupsSupportedInRDDGraph` (before any stage is created). A regular shuffle + * internal to a group does not arise for the all-pipelined job shape (groups are split at + * regular-shuffle boundaries) and so is not checked. + */ + private def checkPipelinedProducerSupported(shuffleDep: ShuffleDependency[_, _, _]): Unit = { + if (!shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) { + return + } + val rdd: RDD[_] = shuffleDep.rdd + // Barrier: exposes output only after a global sync, contradicting concurrent partial reads. + if (rdd.isBarrier()) { + throw pipelinedUnsupportedError("barrier execution in a pipelined-group member stage") + } + // Dynamic resource allocation: gang admission needs a stable slot set; reclaiming executors + // from a pinned-open group can deadlock it. + if (Utils.isDynamicAllocationEnabled(sc.conf)) { + throw pipelinedUnsupportedError("dynamic resource allocation with a pipelined shuffle") + } + // Statically-indeterminate producer: its recovery is stage rollback-and-recompute, which a + // group never performs (any failure aborts the whole group); reject rather than carry dead + // machinery. + if (rdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) { + throw pipelinedUnsupportedError("a statically-indeterminate pipelined producer") + } + // Checksum-mismatch full retry: the runtime counterpart to static indeterminism; it rolls back + // and re-runs succeeding stages on a cross-attempt mismatch, which a group never keeps (moot). + // A PipelinedShuffleDependency does not enable it (see its definition), so this is defensive. + if (shuffleDep.checksumMismatchFullRetryEnabled) { + throw pipelinedUnsupportedError("checksum-mismatch full retry with a pipelined shuffle") + } + // Push-based shuffle merge on a pipelined shuffle: exposes output only after a + // post-completion finalize step, the opposite of incremental reads. A + // PipelinedShuffleDependency disables merge in its constructor, so this is a defensive + // backstop against that being bypassed. + if (shuffleDep.shuffleMergeEnabled) { + throw pipelinedUnsupportedError("push-based shuffle merge as a pipelined shuffle") + } + // A reliable RDD checkpoint in a member's within-stage chain (producer OR consumer side) is + // rejected in checkPipelinedGroupsSupportedInRDDGraph, at job submission before any stage is + // created -- so a reject leaves no partial stage state and both chain sides are covered. + } + /** * We don't support run a barrier stage with dynamic resource allocation enabled, it shall lead * to some confusing behaviors (e.g. with dynamic resource allocation enabled, it may happen that @@ -1060,6 +1212,126 @@ private[spark] class DAGScheduler( false } + /** + * Reject group-level idioms a pipelined group cannot support, checked against the RDD graph + * BEFORE any stage is created -- so a rejection fails the job up front (via handleJobSubmitted's + * listener.jobFailed) without leaving partial scheduler state behind, exactly like the + * speculation check. + * + * Call only for a job that has a pipelined dependency (handleJobSubmitted gates on + * hasPipelined): the resource-profile check below is not keyed on a pipelined dependency, so on a + * regular job it would reject an ordinary RDD.withResources(...) use. Throws + * PIPELINED_SHUFFLE_UNSUPPORTED on violation. Enforces: + * - Fan-out: a pipelined producer feeding more than one consumer. 1:N is a supported model not + * yet built (it needs multicast to N live readers), so it is rejected for now. A + * PipelinedShuffleDependency's producer is `dep.rdd`; a "consumer" is any RDD that lists that + * dependency. More than one distinct consumer RDD for the same pipelined shuffle is fan-out. + * - Reliable RDD checkpoint in a group member's within-stage chain (producer OR consumer side): + * a reliable `checkpoint()` writes a durable, lineage-truncated snapshot, which both + * reintroduces cross-time reuse of a transient edge and requires a post-success recompute + * of the member's transient input -- for a consumer, that input is the vanished pipelined + * shuffle. Checked here (not at stage creation) so a reject leaves no partial stage state, and + * so BOTH the producer chain (rooted at pd.rdd) and each consumer chain (rooted at a consuming + * RDD) are covered from the whole-graph view. + * - A non-default resource profile on any member. The gang slot check compares one demand + * against one profile's capacity and measures it against the default profile, so the whole + * group is required to run on the default profile; any member with an explicit non-default + * profile is rejected. Per-profile accounting is a follow-up. + * + * (The remaining group-level case -- a regular shuffle internal to a group -- is a structural + * invariant that does not arise for the prefix -> pipelined-group -> suffix shapes targeted here: + * groups are split at regular-shuffle boundaries. The producer-side idioms -- barrier, DRA, + * indeterminate, checksum, push-merge -- are rejected in checkPipelinedProducerSupported at stage + * creation, where a producer-only throw leaves no partial state.) + */ + private def checkPipelinedGroupsSupportedInRDDGraph(finalRDD: RDD[_]): Unit = { + // Walk the whole RDD graph once, collecting for each pipelined shuffleId the distinct consumer + // RDDs that read it (for the fan-out check), the producer RDDs that write it (roots of producer + // member stages), and every reliably-checkpointed RDD (to locate ones inside a member stage). + val consumersByShuffleId = new HashMap[Int, HashSet[Int]] + val producerRoots = new HashSet[RDD[_]] // RDDs that WRITE a pipelined shuffle + val reliablyCheckpointed = new HashSet[RDD[_]] // RDDs with a reliable checkpoint pending + var hasNonDefaultResourceProfile = false // any member RDD with a non-default RP + traverseRDDGraph(finalRDD) { (rdd, enqueue) => + if (rdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) { + reliablyCheckpointed += rdd + } + // An RDD's EFFECTIVE profile is its explicit one, or the default when unset. The slot check + // measures capacity against the default profile (see rejectUnadmittablePipelinedGroup), so + // the whole group must run on the default profile; any explicit non-default profile on a + // member makes the group span profiles (against the default the rest use) and is rejected. + val rp = rdd.getResourceProfile() + if (rp != null && rp.id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) { + hasNonDefaultResourceProfile = true + } + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => + consumersByShuffleId.getOrElseUpdate(pd.shuffleId, new HashSet[Int]) += rdd.id + producerRoots += pd.rdd + enqueue(pd.rdd) + case dep => + enqueue(dep.rdd) + } + } + if (consumersByShuffleId.values.exists(_.size > 1)) { + throw pipelinedUnsupportedError( + "a pipelined producer with more than one consumer (fan-out / branching)") + } + // Resource profile. The gang slot check compares one demand against one profile's capacity + // (maxNumConcurrentTasks is defined per profile) and measures it against the DEFAULT profile, + // so the whole group is required to run on the default profile and fails fast otherwise; + // per-profile accounting is a follow-up. Reject if ANY member carries an explicit non-default + // profile -- that both spans profiles (against the default the other members use) and would be + // admitted against the wrong (default) capacity pool. This is a real check, not just a + // documented assumption: nothing else enforces it. + if (hasNonDefaultResourceProfile) { + throw pipelinedUnsupportedError( + "a pipelined group member with a non-default resource profile (the whole group must run " + + "on the default profile)") + } + // Reject a reliable checkpoint anywhere in a pipelined-group MEMBER's within-stage chain. + // Keyed on checkpointData being ReliableRDDCheckpointData, not isCheckpointed, since the + // write has not happened yet. Cache / .persist() / local checkpoint are whole-partition and + // ephemeral and are not rejected. A reliably-checkpointed RDD `cp` is inside a member stage + // iff: + // - PRODUCER side: cp is within some producer root's own within-stage chain (walk parents from + // the producer root, stopping at shuffle boundaries), OR + // - CONSUMER side: cp's OWN within-stage chain reads a pipelined shuffle (walk parents from + // cp, stopping at shuffle boundaries, and check whether any stopped-at boundary is + // pipelined). + // Rooting the consumer check at each checkpointed RDD (rather than at the PSD-reading RDD) is + // what makes it cover a checkpoint anywhere DOWNSTREAM in the consumer stage, not just on the + // reading RDD itself. + def chainHasReliableCheckpoint(root: RDD[_]): Boolean = + !traverseParentRDDsWithinStage(root, (r: RDD[_]) => + !r.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val offending = + // CONSUMER side: a checkpointed RDD whose own within-stage chain reads a pipelined shuffle is + // inside a consumer member stage (covers a checkpoint anywhere in that stage, not just on the + // reading RDD). PRODUCER side: a producer root's within-stage chain carries a checkpoint. + reliablyCheckpointed.exists(rddChainReadsPipelinedShuffle) || + producerRoots.exists(chainHasReliableCheckpoint) + if (offending) { + throw pipelinedUnsupportedError( + "a reliable RDD checkpoint in a pipelined-group member's within-stage chain") + } + } + + /** Whether `rdd`'s within-stage chain (parents, stopping at shuffle boundaries) reads through a + * [[PipelinedShuffleDependency]] -- i.e. `rdd` is inside a pipelined CONSUMER member stage. */ + private def rddChainReadsPipelinedShuffle(rdd: RDD[_]): Boolean = { + !traverseRDDGraphUntil(rdd) { (r, enqueue) => + val readsPipelined = r.dependencies.exists { + case _: PipelinedShuffleDependency[_, _, _] => true + case _: ShuffleDependency[_, _, _] => false // regular boundary: not within this stage + case narrowDep => + enqueue(narrowDep.rdd) + false + } + !readsPipelined + } + } + /** Invoke `.partitions` on the given RDD and all of its ancestors */ private def eagerlyComputePartitionsForRddAndAncestors(rdd: RDD[_]): Unit = { val startTime = System.nanoTime @@ -1093,24 +1365,25 @@ private[spark] class DAGScheduler( } /** - * Removes state for job and any stages that are not needed by any other job. Does not - * handle cancelling tasks or notifying the SparkListener about finished jobs/stages/tasks. - * - * @param job The job whose state to cleanup. + * Removes the given job from every stage registered for it, unregistering each stage that no + * other job needs. Unlike [[cleanupStateForJobAndIndependentStages]] this does not require an + * `ActiveJob`: stage creation may register ancestor stages before a barrier slot check throws + * (see SPARK-58887), and a job cancelled while deferred for the retry must still drop those + * registrations. */ - private def cleanupStateForJobAndIndependentStages(job: ActiveJob): Unit = { - val registeredStages = jobIdToStageIds.get(job.jobId) + private def cleanupStagesForJob(jobId: Int): Unit = { + val registeredStages = jobIdToStageIds.get(jobId) if (registeredStages.isEmpty || registeredStages.get.isEmpty) { - logError(log"No stages registered for job ${MDC(JOB_ID, job.jobId)}") + logError(log"No stages registered for job ${MDC(JOB_ID, jobId)}") } else { stageIdToStage.filter { case (stageId, _) => registeredStages.get.contains(stageId) }.foreach { case (stageId, stage) => val jobSet = stage.jobIds - if (!jobSet.contains(job.jobId)) { + if (!jobSet.contains(jobId)) { // scalastyle:off line.size.limit - logError(log"Job ${MDC(JOB_ID, job.jobId)} not registered for stage ${MDC(STAGE_ID, stageId)} even though that stage was registered for the job") + logError(log"Job ${MDC(JOB_ID, jobId)} not registered for stage ${MDC(STAGE_ID, stageId)} even though that stage was registered for the job") // scalastyle:on } else { def removeStage(stageId: Int): Unit = { @@ -1155,14 +1428,24 @@ private[spark] class DAGScheduler( .format(stageId, stageIdToStage.size)) } - jobSet -= job.jobId + jobSet -= jobId if (jobSet.isEmpty) { // no other job needs this stage removeStage(stageId) } } } } - jobIdToStageIds -= job.jobId + jobIdToStageIds -= jobId + } + + /** + * Removes state for job and any stages that are not needed by any other job. Does not + * handle cancelling tasks or notifying the SparkListener about finished jobs/stages/tasks. + * + * @param job The job whose state to cleanup. + */ + private def cleanupStateForJobAndIndependentStages(job: ActiveJob): Unit = { + cleanupStagesForJob(job.jobId) jobIdToActiveJob -= job.jobId activeJobs -= job job.finalStage match { @@ -1384,7 +1667,7 @@ private[spark] class DAGScheduler( def cancelJobsWithTag( tag: String, reason: Option[String], - cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = { + cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]): Unit = { SparkContext.throwIfInvalidTag(tag) logInfo(log"Asked to cancel jobs with tag ${MDC(TAG, tag)}") eventProcessLoop.post(JobTagCancelled(tag, reason, cancelledJobs)) @@ -1392,9 +1675,14 @@ private[spark] class DAGScheduler( /** * Cancel all jobs that are running or waiting in the queue. + * + * @param reason reason for cancellation. It is surfaced in the error of every cancelled job, so + * that a job aborted as collateral of a context-wide cancellation can be told + * apart from one that failed on its own. */ - def cancelAllJobs(): Unit = { - eventProcessLoop.post(AllJobsCancelled) + def cancelAllJobs(reason: Option[String] = None): Unit = { + logInfo(log"Asked to cancel all jobs${MDC(REASON, reason.map(" " + _).getOrElse(""))}") + eventProcessLoop.post(AllJobsCancelled(reason)) } /** @@ -1405,10 +1693,17 @@ private[spark] class DAGScheduler( eventProcessLoop.post(CleanupQueryJobs(executionId)) } - private[scheduler] def doCancelAllJobs(): Unit = { - // Cancel all running jobs. - runningStages.map(_.firstJobId).foreach(handleJobCancellation(_, - Option("as part of cancellation of all jobs"))) + private[scheduler] def doCancelAllJobs(reason: Option[String] = None): Unit = { + // Cancel all running jobs. A job caught here was healthy and is being aborted as collateral of + // a context-wide cancellation, not because it failed on its own, so attribute the reason when + // the caller supplied one. + val updatedReason = reason.getOrElse(DAGScheduler.DEFAULT_CANCEL_ALL_JOBS_REASON) + runningStages.map(_.firstJobId).foreach(handleJobCancellation(_, Option(updatedReason))) + // Also fail the barrier jobs deferred for a slot-check retry: they are registered nowhere + // else, and their pending re-posts are dropped once the entries are marked cancelled. + deferredBarrierJobs.keySet().forEach { jobId => + handleJobCancellation(jobId, Option(updatedReason)) + } activeJobs.clear() // These should already be empty by this point, jobIdToActiveJob.clear() // but just in case we lost track of some jobs... } @@ -1574,49 +1869,17 @@ private[spark] class DAGScheduler( */ private def rejectUnadmittablePipelinedGroup( jobId: Int, finalRDD: RDD[_], partitions: Array[Int], listener: JobListener): Boolean = { - // Reject up-front (before any stage is created) two group members the admission model cannot - // support, walking the group's RDD graph once: - // - A barrier member. A barrier stage exposes its output only after a global sync, which - // contradicts a pipelined consumer reading the producer's output incrementally as it runs; - // and its failure recovery resubmits the stage, which the group's atomic completion cannot - // accommodate (a resubmitted producer would drop a co-scheduled consumer's buffered - // completions). Reject it here rather than let it be co-scheduled. - // - A member on a non-default resource profile. Admission below measures capacity and - // occupancy against the default profile, but each stage derives its profile from its RDDs - // (see createShuffleMapStage/createResultStage). A member carrying a non-default profile would be - // admitted against the default profile's free slots yet run in a different, often smaller - // pool -- and could then queue or deadlock there. - var offendingBarrier = false - var offendingRp: Option[ResourceProfile] = None - traverseRDDGraph(finalRDD) { (rdd, enqueue) => - if (rdd.isBarrier()) { - offendingBarrier = true - } - val rp = rdd.getResourceProfile() - if (offendingRp.isEmpty && rp != null && rp.id != DEFAULT_RESOURCE_PROFILE_ID) { - offendingRp = Some(rp) - } - rdd.dependencies.foreach(dep => enqueue(dep.rdd)) - } - if (offendingBarrier) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group contains a " + - log"barrier member") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a barrier stage in the pipelined group is not supported: a " + - "barrier stage exposes its output only after a global sync, which is incompatible with " + - "a pipelined consumer reading its output incrementally.")) - return true - } - if (offendingRp.nonEmpty) { - logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: a pipelined stage group member uses a " + - log"non-default resource profile") - listener.jobFailed(new SparkException( - "A pipelined shuffle job with a member on a non-default resource profile is not " + - s"supported (resource profile id ${offendingRp.get.id}): the whole pipelined stage " + - "group must run on the default resource profile.")) - return true - } if (!sc.conf.get(config.PIPELINED_GROUP_SLOT_CHECK_ENABLED)) { + // The only deadlock-prevention check for gang admission is off. Legitimate only when the + // deployment admits capacity out-of-band (e.g. a slot reservation); otherwise a pipelined + // group that cannot co-fit will be gang-scheduled and can deadlock. Warn once so this is + // never a silent state. + if (!warnedPipelinedSlotCheckDisabled) { + warnedPipelinedSlotCheckDisabled = true + logWarning(log"${MDC(CONFIG, config.PIPELINED_GROUP_SLOT_CHECK_ENABLED.key)}=false: " + + log"pipelined-group gang admission is NOT checking free slots. This is safe only if " + + log"capacity is reserved out-of-band; otherwise a group that cannot co-fit may deadlock.") + } return false } val rp = sc.resourceProfileManager.defaultResourceProfile @@ -1680,24 +1943,12 @@ private[spark] class DAGScheduler( * member's `TaskSet.isPipelined` at submission and routing a member FetchFailed to a whole-group * abort. */ - private def isPipelinedGroupMember(stage: Stage): Boolean = { - if (isPipelinedProducer(stage)) { - return true - } - // Consumer check: does this stage read through a PipelinedShuffleDependency at one of its - // shuffle boundaries? Walk the stage's own RDD graph (descending narrow deps, stopping at every - // shuffle boundary) -- the same edges that define the stage -- and look for a pipelined one. - !traverseRDDGraphUntil(stage.rdd) { (rdd, enqueue) => - val hasPipelinedBoundary = rdd.dependencies.exists { - case _: PipelinedShuffleDependency[_, _, _] => true - case _: ShuffleDependency[_, _, _] => false // regular shuffle boundary: do not descend - case narrowDep => - enqueue(narrowDep.rdd) - false - } - !hasPipelinedBoundary // keep walking until a pipelined boundary is found - } - } + private def isPipelinedGroupMember(stage: Stage): Boolean = + // Producer side: it writes a pipelined shuffle. Consumer side: its within-stage chain reads + // one. rddChainReadsPipelinedShuffle is the single source of truth for that within-stage walk + // (descend narrow deps, stop at every shuffle boundary, look for a pipelined one) -- do not + // re-inline it; the consumer walk here and that method used to be byte-identical copies. + isPipelinedProducer(stage) || rddChainReadsPipelinedShuffle(stage.rdd) /** Finds the earliest-created active job that needs the stage */ // TODO: Probably should actually find among the active jobs that need this @@ -1727,11 +1978,16 @@ private[spark] class DAGScheduler( _.getProperty(SparkContext.SPARK_JOB_GROUP_ID) == groupId } } - if (activeInGroup.isEmpty && !cancelFutureJobs) { + // A barrier job deferred for a slot-check retry is not in `activeJobs` yet, so match it by + // the properties captured at submission. + val deferredInGroup = deferredJobsMatching { properties => + Option(properties).exists(_.getProperty(SparkContext.SPARK_JOB_GROUP_ID) == groupId) + } + if (activeInGroup.isEmpty && deferredInGroup.isEmpty && !cancelFutureJobs) { logWarning(log"Failed to cancel job group ${MDC(GROUP_ID, groupId)}. " + log"Cannot find active jobs for it.") } - val jobIds = activeInGroup.map(_.jobId) + val jobIds = activeInGroup.map(_.jobId) ++ deferredInGroup.map(_._1) val updatedReason = reason.getOrElse("part of cancelled job group %s".format(groupId)) jobIds.foreach(handleJobCancellation(_, Option(updatedReason))) } @@ -1739,19 +1995,32 @@ private[spark] class DAGScheduler( private[scheduler] def handleJobTagCancelled( tag: String, reason: Option[String], - cancelledJobs: Option[Promise[Seq[ActiveJob]]]): Unit = { + cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]): Unit = { // Cancel all jobs that have all provided tags. // First finds all active jobs with this group id, and then kill stages for them. val jobsToBeCancelled = activeJobs.filter { activeJob => - Option(activeJob.properties).exists { properties => - Option(properties.getProperty(SparkContext.SPARK_JOB_TAGS)).getOrElse("") - .split(SparkContext.SPARK_JOB_TAGS_SEP).filter(!_.isEmpty).toSet.contains(tag) - } + hasJobTag(activeJob.properties, tag) } + // A barrier job deferred for a slot-check retry is not in `activeJobs` yet, so match it by + // the properties captured at submission. + val deferredTagged = deferredJobsMatching(hasJobTag(_, tag)) val updatedReason = reason.getOrElse("part of cancelled job tags %s".format(tag)) - jobsToBeCancelled.map(_.jobId).foreach(handleJobCancellation(_, Option(updatedReason))) - cancelledJobs.map(_.success(jobsToBeCancelled.toSeq)) + (jobsToBeCancelled.map(_.jobId) ++ deferredTagged.map(_._1)) + .foreach(handleJobCancellation(_, Option(updatedReason))) + // Report the deferred jobs too: they have no ActiveJob, but consumers (e.g. classic + // SparkSession.interruptTag) read the SQL execution id from the properties. + cancelledJobs.map(_.success( + jobsToBeCancelled.toSeq.map(job => CancelledJobInfo(job.jobId, job.properties)) ++ + deferredTagged.map { case (jobId, properties) => CancelledJobInfo(jobId, properties) })) + } + + /** Whether the job properties carry the given job tag. */ + private def hasJobTag(properties: Properties, tag: String): Boolean = { + Option(properties).exists { props => + Option(props.getProperty(SparkContext.SPARK_JOB_TAGS)).getOrElse("") + .split(SparkContext.SPARK_JOB_TAGS_SEP).filter(!_.isEmpty).toSet.contains(tag) + } } private[scheduler] def handleBeginEvent(task: Task[_], taskInfo: TaskInfo): Unit = { @@ -1806,6 +2075,15 @@ private[spark] class DAGScheduler( } listenerBus.post(SparkListenerJobEnd(job.jobId, clock.getTimeMillis(), JobFailed(error))) } + // Also complete the waiters of the barrier jobs deferred for a slot-check retry: they are + // in `activeJobs` above only once re-processed, which will never happen now. + deferredBarrierJobs.forEach { (jobId, deferred) => + if (deferred ne deferredJobCancelledMarker) { + deferred.listener.jobFailed( + new SparkException(s"Job $jobId cancelled because SparkContext was shut down")) + } + } + deferredBarrierJobs.clear() } private[scheduler] def handleGetTaskResult(taskInfo: TaskInfo): Unit = { @@ -1833,6 +2111,11 @@ private[spark] class DAGScheduler( listener: JobListener, artifacts: JobArtifactSet, properties: Properties): Unit = { + // The job is being (re-)processed, so it is no longer merely deferred. The marker means it + // was cancelled while deferred: its listener has already been failed, so drop this re-post. + if (deferredBarrierJobs.remove(jobId) eq deferredJobCancelledMarker) { + return + } // If this job belongs to a cancelled job group, skip running it val jobGroupIdOpt = Option(properties).map(_.getProperty(SparkContext.SPARK_JOB_GROUP_ID)) if (jobGroupIdOpt.exists(cancelledJobGroups.contains(_))) { @@ -1875,17 +2158,34 @@ private[spark] class DAGScheduler( if (hasPipelined && rejectUnadmittablePipelinedGroup(jobId, finalRDD, partitions, listener)) { return } - var finalStage: ResultStage = null try { + // Reject group-level unsupported pipelined idioms (e.g. fan-out, a non-default resource + // profile, a reliable checkpoint in a member stage) from the RDD graph, up front -- before + // any stage is created, so a rejection leaves no partial scheduler state. Gated on + // hasPipelined: every idiom this checks concerns a pipelined group, so it must not run for a + // job with no pipelined dependency (the resource-profile check in particular is not keyed on + // a pipelined dependency and would otherwise reject an ordinary job that merely uses a + // non-default profile via RDD.withResources). Inside this try so any incidental exception + // from the graph walk is handled by the same listener.jobFailed path as stage creation. + if (hasPipelined) { + checkPipelinedGroupsSupportedInRDDGraph(finalRDD) + } // New stage creation may throw an exception if, for example, jobs are run on a // HadoopRDD whose underlying HDFS files have been deleted. finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite) } catch { case e: BarrierJobSlotsNumberCheckFailed => // If jobId doesn't exist in the map, Scala coverts its value null to 0: Int automatically. - val numCheckFailures = barrierJobIdToNumTasksCheckFailures.compute(jobId, - (_: Int, value: Int) => value + 1) + // Do not consume the retry budget while the executors are held: the slot check sees + // zero slots for the whole hold, and the job should wait for the resume like any + // other job instead of failing when the retries run out. + val numCheckFailures = if (executorsHeld) { + barrierJobIdToNumTasksCheckFailures.getOrDefault(jobId, 0) + } else { + barrierJobIdToNumTasksCheckFailures.compute(jobId, + (_: Int, value: Int) => value + 1) + } logWarning(log"Barrier stage in job ${MDC(JOB_ID, jobId)} " + log"requires ${MDC(NUM_SLOTS, e.requiredConcurrentTasks)} slots, " + @@ -1894,6 +2194,7 @@ private[spark] class DAGScheduler( log"more times") if (numCheckFailures <= maxFailureNumTasksCheck) { + deferredBarrierJobs.put(jobId, DAGScheduler.DeferredBarrierJob(listener, properties)) messageScheduler.schedule( new Runnable { override def run(): Unit = eventProcessLoop.post(JobSubmitted(jobId, finalRDD, func, @@ -1910,6 +2211,15 @@ private[spark] class DAGScheduler( return } + case e: PipelinedShuffleUnsupportedException => + // An up-front idiom rejection (checkPipelinedGroupsSupportedInRDDGraph / a producer-side + // check in createShuffleMapStage), not a stage-creation failure. Log it as such (the + // generic "Creating new stage failed" message below would be misleading). Matched by TYPE, + // not by the error-condition string, so a rename or a wrapped cause cannot misroute it. + logWarning(log"Rejecting job ${MDC(JOB_ID, jobId)}: unsupported pipelined-shuffle idiom", e) + listener.jobFailed(e) + return + case e: Exception => logWarning(log"Creating new stage failed due to exception - job: ${MDC(JOB_ID, jobId)}", e) listener.jobFailed(e) @@ -1918,7 +2228,9 @@ private[spark] class DAGScheduler( // Job submitted, clear internal data. barrierJobIdToNumTasksCheckFailures.remove(jobId) - val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties) + // Pass hasPipelined (computed above) into the job; see ActiveJob.hasPipelinedDependency. + val job = new ActiveJob(jobId, finalStage, callSite, listener, artifacts, properties, + hasPipelinedDependency = hasPipelined) clearCacheLocs() logInfo( log"Got job ${MDC(JOB_ID, job.jobId)} (${MDC(CALL_SITE_SHORT_FORM, callSite.shortForm)}) " + @@ -2077,7 +2389,9 @@ private[spark] class DAGScheduler( // rejectUnadmittablePipelinedGroup) before any member was submitted, so the group is // known to fit; just co-schedule this consumer with its running producer(s). No slot // check here -- that would re-measure capacity against a mid-flight snapshot and is - // unnecessary once admission is decided up front (gang admission). + // unnecessary once admission is decided up front (gang admission). Group-level + // idiom rejection (fan-out, internal regular shuffle) already happened at job + // submission (checkPipelinedGroupsSupportedInRDDGraph + the all-pipelined check). logInfo(log"Submitting ${MDC(STAGE, stage)} concurrently with its running " + log"pipelined producer(s) ${MDC(MISSING_PARENT_STAGES, pipelinedMissing)}") // Record that this stage is co-scheduled with still-running pipelined producers, @@ -2538,9 +2852,13 @@ private[spark] class DAGScheduler( case _: ResultStage => None } + // Only a job that uses a pipelined shuffle can have a pipelined-group member; gate the + // group-membership graph walk on that cheap per-job flag so a regular job pays nothing here. + val isPipelined = jobIdToActiveJob.get(jobId).exists(_.hasPipelinedDependency) && + isPipelinedGroupMember(stage) taskScheduler.submitTasks(new TaskSet( tasks.toArray, stage.id, stage.latestInfo.attemptNumber(), jobId, properties, - stage.resourceProfileId, shuffleId)) + stage.resourceProfileId, shuffleId, isPipelined = isPipelined)) } else { // Because we posted SparkListenerStageSubmitted earlier, we should mark // the stage as completed here in case there are no tasks to run @@ -2865,8 +3183,17 @@ private[spark] class DAGScheduler( // will be re-executed. if (clearShuffle) { logInfo(log"Cleaning up shuffle for stage ${MDC(STAGE, sms)} to ensure re-execution") - mapOutputTracker.unregisterAllMapAndMergeOutput(sms.shuffleDep.shuffleId) - sms.shuffleDep.newShuffleMergeState() + // A pipelined shuffle is not registered with the MapOutputTracker, so unregistering there + // would throw ShuffleStatusNotFoundException. Not reachable today -- an indeterminate + // pipelined producer is rejected up front (checkPipelinedProducerSupported), and a job is + // all-regular or all-pipelined (mixed rejected), so a pipelined stage is never a succeeding + // stage of a regular indeterminate producer that rolls back -- but guard defensively, like + // the pipelined branch on the FetchFailed base path. (A transient pipelined producer cannot + // be rolled back and recomputed anyway; a genuine member failure fails the group, not this.) + if (!sms.isPipelined) { + mapOutputTracker.unregisterAllMapAndMergeOutput(sms.shuffleDep.shuffleId) + sms.shuffleDep.newShuffleMergeState() + } } } @@ -3081,7 +3408,29 @@ private[spark] class DAGScheduler( case smt: ShuffleMapTask => val shuffleStage = stage.asInstanceOf[ShuffleMapStage] - if (!ignoreOldTaskAttempts) { + if (shuffleStage.isPipelined) { + // A pipelined shuffle's completed partitions are tracked locally and monotonically on + // the stage, not in the MapOutputTracker (the reader finds the producer via the + // streaming transport, not the tracker). See ShuffleMapStage's + // pipelinedCompletedPartitions scaladoc for why -- the crux of avoiding the + // streaming-writer resubmit hang. Checksum-mismatch detection does not apply (a + // dependency never enables checksum retry -- see PipelinedShuffleDependency). + // + // Record the partition and decrement pendingPartitions UNCONDITIONALLY -- outside the + // `!ignoreOldTaskAttempts` and bogus-epoch guards that gate a regular shuffle. Both + // guards exist only to avoid trusting a MapOutputTracker registration that a later + // rollback (ignoreOldTaskAttempts, from an indeterminate/rolled-back stage) or an + // executor-loss strip (bogus epoch) could invalidate. A pipelined stage never + // registers there, and its completed set is monotonic and never rolled back (a + // transient shuffle cannot be recomputed; any real group failure aborts the whole + // group). Skipping the record for an "old" or "bogus" straggler would be actively + // harmful: with pendingPartitions decremented but the partition unrecorded, a dropped + // last partition leaves the stage "done but not available" -> processShuffleMapStage- + // Completion resubmits the transient producer, reopening the streaming-writer hang. + // An already-successful straggler is not a failure, so recording it is correct. + shuffleStage.pendingPartitions -= task.partitionId + shuffleStage.addPipelinedCompletedPartition(smt.partitionId) + } else if (!ignoreOldTaskAttempts) { shuffleStage.pendingPartitions -= task.partitionId val status = event.result.asInstanceOf[MapStatus] val execId = status.location.executorId @@ -3149,6 +3498,33 @@ private[spark] class DAGScheduler( log"${MDC(STAGE_ATTEMPT_ID, task.stageAttemptId)} and there is a more recent attempt for " + log"that stage (attempt " + log"${MDC(NUM_ATTEMPT, failedStage.latestInfo.attemptNumber())}) running") + } else if (activeJobForStage(failedStage).flatMap(jobIdToActiveJob.get) + .exists(_.hasPipelinedDependency) && + (isPipelinedGroupMember(failedStage) || isPipelinedGroupMember(mapStage))) { + // Failure is group-atomic for a pipelined group. The base scheduler handles a + // FetchFailed by resubmitting just the map stage in isolation and recomputing serially, + // but a transient pipelined shuffle cannot be re-read and its members are co-scheduled, + // so a lone-stage resubmit is never valid and would deadlock the group. Abort the + // whole group instead: aborting the failed stage tears down its running co-scheduled + // members and fails the job, and the caller (e.g. the streaming batch loop) reruns the + // batch from scratch. This is distinct from the maxTaskFailures=1 lever (which handles + // task failures the TaskSetManager counts): a FetchFailed is NOT counted there (the base + // TaskSetManager marks the task successful and zombies the set), so the routing to group + // failure must be enforced here. + logInfo(log"Failing pipelined group containing ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) atomically due to a fetch failure " + + log"from ${MDC(STAGE, mapStage)} (${MDC(STAGE_NAME, mapStage.name)})") + failedStage.failedAttemptIds.add(task.stageAttemptId) + // Still unregister the failed executor's outputs, exactly as the base FetchFailed path + // does -- aborting the group tears down only THIS job's stages, but the FetchFailed is + // authoritative evidence that the executor's shuffle data is gone, and other/concurrent + // jobs sharing that executor must not keep stale MapOutputTracker entries (with an + // external shuffle service, an ExecutorLost would NOT clean these, so this is the only + // proactive channel). Safe for the pipelined shuffle itself: it registers no map outputs + // in the tracker, so this can only strip regular/durable outputs. + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { val ignoreStageFailure = ignoreDecommissionFetchFailure && isExecutorDecommissioningOrDecommissioned(taskScheduler, bmAddress) @@ -3181,123 +3557,106 @@ private[spark] class DAGScheduler( "longer running") } - if (mapStage.rdd.isBarrier()) { - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and - // TODO: finalized as we are clearing all the merge results. - mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) - } else if (mapIndex != -1) { - // Mark the map whose fetch failed as broken in the map stage - mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) - if (pushBasedShuffleEnabled) { - // Possibly unregister the merge result <shuffleId, reduceId>, if the FetchFailed - // mapIndex is part of the merge result of <shuffleId, reduceId> - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) - } + if (mapStage.isPipelined) { + // Defense-in-depth for a pipelined-shuffle FetchFailed that reaches this base path + // rather than the group-atomic branch above (e.g. a job whose hasPipelinedDependency + // flag was not propagated through the group check). The base path is invalid for a + // pipelined shuffle in two ways: (a) the MapOutputTracker invalidation below would + // throw ShuffleStatusNotFoundException (a pipelined shuffle is never registered there, + // see createShuffleMapStage); (b) the resubmit branch would enqueue a lone-stage + // resubmit of the transient producer, which -- as the group-atomic branch's comment + // explains -- is never valid and would deadlock the group. So abort the group here + // instead, matching the group-atomic branch's outcome, and skip both. + abortStage(failedStage, + s"A pipelined group member failed with a fetch failure: $failureMessage", None) } else { - // Unregister the merge result of <shuffleId, reduceId> if there is a FetchFailed event - // and is not a MetaDataFetchException which is signified by bmAddress being null - if (bmAddress != null && - bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { - assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + - "be enabled when handling merge block fetch failure.") - mapOutputTracker. - unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + if (mapStage.rdd.isBarrier()) { + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + // TODO: SPARK-35547: Clean all push-based shuffle metadata like merge enabled and + // TODO: finalized as we are clearing all the merge results. + mapOutputTracker.unregisterAllMapAndMergeOutput(shuffleId) + } else if (mapIndex != -1) { + // Mark the map whose fetch failed as broken in the map stage + mapOutputTracker.unregisterMapOutput(shuffleId, mapIndex, bmAddress) + if (pushBasedShuffleEnabled) { + // Possibly unregister the merge result <shuffleId, reduceId>, if the FetchFailed + // mapIndex is part of the merge result of <shuffleId, reduceId> + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, Option(mapIndex)) + } + } else { + // Unregister the merge result of <shuffleId, reduceId> if there is a FetchFailed + // event and is not a MetaDataFetchException (signified by bmAddress being null) + if (bmAddress != null && + bmAddress.executorId.equals(BlockManagerId.SHUFFLE_MERGER_IDENTIFIER)) { + assert(pushBasedShuffleEnabled, "Push based shuffle expected to " + + "be enabled when handling merge block fetch failure.") + mapOutputTracker. + unregisterMergeResult(shuffleId, reduceId, bmAddress, None) + } } - } - - if (failedStage.rdd.isBarrier()) { - failedStage match { - case failedMapStage: ShuffleMapStage => - // Mark all the map as broken in the map stage, to ensure retry all the tasks on - // resubmitted stage attempt. - mapOutputTracker.unregisterAllMapAndMergeOutput(failedMapStage.shuffleDep.shuffleId) - case failedResultStage: ResultStage => - // Abort the failed result stage since we may have committed output for some - // partitions. - val reason = "Could not recover from a failed barrier ResultStage. Most recent " + - s"failure reason: $failureMessage" - abortStage(failedResultStage, reason, None) + if (failedStage.rdd.isBarrier()) { + failedStage match { + case failedMapStage: ShuffleMapStage => + // Mark all the map as broken in the map stage, to ensure retry all the tasks on + // resubmitted stage attempt. + mapOutputTracker.unregisterAllMapAndMergeOutput( + failedMapStage.shuffleDep.shuffleId) + + case failedResultStage: ResultStage => + // Abort the failed result stage since we may have committed output for some + // partitions. + val reason = "Could not recover from a failed barrier ResultStage. Most recent " + + s"failure reason: $failureMessage" + abortStage(failedResultStage, reason, None) + } } - } - if (shouldAbortStage) { - abortStage(failedStage, abortReason.get, None) - } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued - // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 - val noResubmitEnqueued = !failedStages.contains(failedStage) - failedStages += failedStage - failedStages += mapStage - if (noResubmitEnqueued) { - // For statically indeterminate stages, trigger rollback early (here and in - // submitMissingTasks) rather than deferring to task completion. This is more - // efficient because it clears shuffle outputs before the retry is submitted, - // ensuring findMissingPartitions() returns all partitions. - // - // For runtime detection (checksum mismatch), rollback is triggered at task - // completion when the mismatch is discovered. - // - // The `rollbackCurrentStage = true` parameter ensures the failed map stage is - // included in the cleanup: clearing its shuffle outputs, marking old task results - // to be ignored, and creating a new shuffle merge state for the upcoming retry. - if (mapStage.isStaticallyIndeterminate && - !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { - rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) - } + if (shouldAbortStage) { + abortStage(failedStage, abortReason.get, None) + } else { // update failedStages and make sure a ResubmitFailedStages event is enqueued + // TODO: Cancel running tasks in the failed stage -- cf. SPARK-17064 + val noResubmitEnqueued = !failedStages.contains(failedStage) + failedStages += failedStage + failedStages += mapStage + if (noResubmitEnqueued) { + // For statically indeterminate stages, trigger rollback early (here and in + // submitMissingTasks) rather than deferring to task completion. This is more + // efficient because it clears shuffle outputs before the retry is submitted, + // ensuring findMissingPartitions() returns all partitions. + // + // For runtime detection (checksum mismatch), rollback is triggered at task + // completion when the mismatch is discovered. + // + // The `rollbackCurrentStage = true` parameter ensures the failed map stage is + // included in the cleanup: clearing its shuffle outputs, marking old task results + // to be ignored, and creating a new shuffle merge state for the upcoming retry. + if (mapStage.isStaticallyIndeterminate && + !mapStage.shuffleDep.checksumMismatchFullRetryEnabled) { + rollbackSucceedingStages(mapStage, rollbackCurrentStage = true) + } - // We expect one executor failure to trigger many FetchFailures in rapid succession, - // but all of those task failures can typically be handled by a single resubmission of - // the failed stage. We avoid flooding the scheduler's event queue with resubmit - // messages by checking whether a resubmit is already in the event queue for the - // failed stage. If there is already a resubmit enqueued for a different failed - // stage, that event would also be sufficient to handle the current failed stage, but - // producing a resubmit for each failed stage makes debugging and logging a little - // simpler while not producing an overwhelming number of scheduler events. - logInfo( - log"Resubmitting ${MDC(STAGE, mapStage)} " + - log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + - log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") - scheduleResubmit() + // We expect one executor failure to trigger many FetchFailures in rapid succession, + // but all of those task failures can typically be handled by a single resubmission + // of the failed stage. We avoid flooding the scheduler's event queue with resubmit + // messages by checking whether a resubmit is already in the event queue for the + // failed stage. If there is already a resubmit enqueued for a different failed + // stage, that event would also be sufficient to handle the current failed stage, + // but producing a resubmit for each failed stage makes debugging and logging a + // little simpler while not producing an overwhelming number of scheduler events. + logInfo( + log"Resubmitting ${MDC(STAGE, mapStage)} " + + log"(${MDC(STAGE_NAME, mapStage.name)}) and ${MDC(FAILED_STAGE, failedStage)} " + + log"(${MDC(FAILED_STAGE_NAME, failedStage.name)}) due to fetch failure") + scheduleResubmit() + } } } // TODO: mark the executor as failed only if there were lots of fetch failures on it - if (bmAddress != null) { - val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled - val isHostDecommissioned = taskScheduler - .getExecutorDecommissionState(bmAddress.executorId) - .exists(_.workerHost.isDefined) - - // Shuffle output of all executors on host `bmAddress.host` may be lost if: - // - External shuffle service is enabled, so we assume that all shuffle data on node is - // bad. - // - Host is decommissioned, thus all executors on that host will die. - val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || - isHostDecommissioned - val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost - && unRegisterOutputOnHostOnFetchFailure) { - Some(bmAddress.host) - } else { - // Unregister shuffle data just for one executor (we don't have any - // reason to believe shuffle data has been lost for the entire host). - None - } - removeExecutorAndUnregisterOutputs( - execId = bmAddress.executorId, - fileLost = true, - hostToUnregisterOutputs = hostToUnregisterOutputs, - maybeEpoch = Some(task.epoch), - // shuffleFileLostEpoch is ignored when a host is decommissioned because some - // decommissioned executors on that host might have been removed before this fetch - // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and - // proceed with unconditional removal of shuffle outputs from all executors on that - // host, including from those that we still haven't confirmed as lost due to heartbeat - // delays. - ignoreShuffleFileLostEpoch = isHostDecommissioned) - } + unregisterOutputsOnFetchFailedExecutor(bmAddress, task) } case failure: TaskFailedReason if task.isBarrier => @@ -3837,6 +4196,52 @@ private[spark] class DAGScheduler( maybeEpoch = None) } + /** + * On a FetchFailed, unregister the shuffle outputs of the executor (or its whole host) whose + * fetch failed, treating the FetchFailed as authoritative evidence that its shuffle data is gone. + * Extracted from the base FetchFailed handler so the pipelined-group-abort branch can also run + * it: aborting the group fails only this job's stages, but a dead executor's REGULAR outputs must + * still be cleaned up for other/concurrent jobs (with an external shuffle service, an + * ExecutorLost does not clean them, so FetchFailed is the only proactive channel). No-op when + * `bmAddress` is null. Safe for a pipelined shuffle: it registers no map outputs in the tracker, + * so this can only strip regular/durable outputs. + */ + private def unregisterOutputsOnFetchFailedExecutor( + bmAddress: BlockManagerId, task: Task[_]): Unit = { + // TODO: mark the executor as failed only if there were lots of fetch failures on it + if (bmAddress != null) { + val externalShuffleServiceEnabled = env.blockManager.externalShuffleServiceEnabled + val isHostDecommissioned = taskScheduler + .getExecutorDecommissionState(bmAddress.executorId) + .exists(_.workerHost.isDefined) + + // Shuffle output of all executors on host `bmAddress.host` may be lost if: + // - External shuffle service is enabled, so we assume that all shuffle data on node is bad. + // - Host is decommissioned, thus all executors on that host will die. + val shuffleOutputOfEntireHostLost = externalShuffleServiceEnabled || isHostDecommissioned + val hostToUnregisterOutputs = if (shuffleOutputOfEntireHostLost + && unRegisterOutputOnHostOnFetchFailure) { + Some(bmAddress.host) + } else { + // Unregister shuffle data just for one executor (we don't have any + // reason to believe shuffle data has been lost for the entire host). + None + } + removeExecutorAndUnregisterOutputs( + execId = bmAddress.executorId, + fileLost = true, + hostToUnregisterOutputs = hostToUnregisterOutputs, + maybeEpoch = Some(task.epoch), + // shuffleFileLostEpoch is ignored when a host is decommissioned because some + // decommissioned executors on that host might have been removed before this fetch + // failure and might have bumped up the shuffleFileLostEpoch. We ignore that, and + // proceed with unconditional removal of shuffle outputs from all executors on that + // host, including from those that we still haven't confirmed as lost due to heartbeat + // delays. + ignoreShuffleFileLostEpoch = isHostDecommissioned) + } + } + /** * Handles removing an executor from the BlockManagerMaster as well as unregistering shuffle * outputs for the executor or optionally its host. @@ -3971,7 +4376,25 @@ private[spark] class DAGScheduler( } private[scheduler] def handleJobCancellation(jobId: Int, reason: Option[String]): Unit = { - if (!jobIdToStageIds.contains(jobId)) { + val deferred = deferredBarrierJobs.get(jobId) + if (deferred != null) { + if (deferred ne deferredJobCancelledMarker) { + // A barrier job deferred for a slot-check retry is registered nowhere else, so fail its + // listener directly. Leave the marker in place instead of removing the entry: the + // pending re-post always fires, and handleJobSubmitted drops it on finding the marker. + deferredBarrierJobs.put(jobId, deferredJobCancelledMarker) + barrierJobIdToNumTasksCheckFailures.remove(jobId) + // Stage creation may have registered ancestor stages (e.g. an ordinary shuffle upstream + // of the barrier stage) before the slot check threw. Drop this job from them, and + // unregister the stages no other job needs, so the abandoned registrations cannot break + // a later cancellation of this job id or pin the stages. + if (jobIdToStageIds.contains(jobId)) { + cleanupStagesForJob(jobId) + } + deferred.listener.jobFailed( + SparkCoreErrors.sparkJobCancelled(jobId, reason.getOrElse(""), null)) + } + } else if (!jobIdToStageIds.contains(jobId)) { logDebug("Trying to cancel unregistered job " + jobId) } else { failJobAndIndependentStages( @@ -4369,8 +4792,8 @@ private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler case JobTagCancelled(tag, reason, cancelledJobs) => dagScheduler.handleJobTagCancelled(tag, reason, cancelledJobs) - case AllJobsCancelled => - dagScheduler.doCancelAllJobs() + case AllJobsCancelled(reason) => + dagScheduler.doCancelAllJobs(reason) case CleanupQueryJobs(executionId) => dagScheduler.doCleanupQueryJobs(executionId) @@ -4429,7 +4852,8 @@ private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler override def onError(e: Throwable): Unit = { logError("DAGSchedulerEventProcessLoop failed; shutting down SparkContext", e) try { - dagScheduler.doCancelAllJobs() + dagScheduler.doCancelAllJobs( + Option("because the DAGScheduler event loop failed and the SparkContext is shutting down")) } catch { case t: Throwable => logError("DAGScheduler failed to cancel all jobs.", t) } @@ -4447,8 +4871,42 @@ private[spark] object DAGScheduler { // this is a simplistic way to avoid resubmitting tasks in the non-fetchable map stage one by one // as more failure events come in val RESUBMIT_TIMEOUT = 200 + + // Fallback reason used when a context-wide cancellation does not supply a more specific one. + // Kept as the historical wording so existing log and error consumers are unaffected. + val DEFAULT_CANCEL_ALL_JOBS_REASON = "as part of cancellation of all jobs" + + /** + * A barrier job deferred while its max concurrent tasks check is being retried, tracked in + * `deferredBarrierJobs`. The submission-time properties are kept so group/tag cancellations + * can match the job. + */ + private[scheduler] case class DeferredBarrierJob(listener: JobListener, properties: Properties) } +/** + * Metadata of a job cancelled by a tag cancellation, reported through the promise of + * `SparkContext.cancelJobsWithTagWithFuture`. Not restricted to jobs with an `ActiveJob`: a + * barrier job cancelled while deferred for its slot-check retry is reported too, and consumers + * (e.g. classic `SparkSession.interruptTag`) read the SQL execution id from the submission-time + * properties. + */ +private[spark] case class CancelledJobInfo(jobId: Int, properties: Properties) + +/** + * Thrown when a job uses a pipelined-shuffle idiom that is not supported (fan-out, a barrier / + * indeterminate / checksum-retry / push-merge producer, a reliable checkpoint in a member's chain, + * or a non-default resource profile on a member). `handleJobSubmitted` matches on this + * TYPE to distinguish an up-front idiom rejection from an ordinary stage-creation failure, not on + * the error-condition string (which a rename or a wrapped cause would silently break). Carries the + * `PIPELINED_SHUFFLE_UNSUPPORTED` error class so the user-facing message is unchanged. + */ +private[scheduler] class PipelinedShuffleUnsupportedException(reason: String) + extends SparkException( + errorClass = "PIPELINED_SHUFFLE_UNSUPPORTED", + messageParameters = scala.collection.immutable.Map("reason" -> reason), + cause = null) + /** * A NOT thread-safe set that only keeps the last `capacity` elements added to it. */ diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala index cc788e5c65bcc..20905ba58ba1f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGSchedulerEvent.scala @@ -74,9 +74,10 @@ private[scheduler] case class JobGroupCancelled( private[scheduler] case class JobTagCancelled( tagName: String, reason: Option[String], - cancelledJobs: Option[Promise[Seq[ActiveJob]]]) extends DAGSchedulerEvent + cancelledJobs: Option[Promise[Seq[CancelledJobInfo]]]) extends DAGSchedulerEvent -private[scheduler] case object AllJobsCancelled extends DAGSchedulerEvent +private[scheduler] case class AllJobsCancelled(reason: Option[String] = None) + extends DAGSchedulerEvent private[scheduler] case class CleanupQueryJobs(executionId: Long) extends DAGSchedulerEvent diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala index 79f7af48f102a..11cec94f2b311 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapStage.scala @@ -19,7 +19,7 @@ package org.apache.spark.scheduler import scala.collection.mutable.HashSet -import org.apache.spark.{MapOutputTrackerMaster, ShuffleDependency} +import org.apache.spark.{MapOutputTrackerMaster, PipelinedShuffleDependency, ShuffleDependency} import org.apache.spark.rdd.{DeterministicLevel, RDD} import org.apache.spark.util.CallSite @@ -59,6 +59,33 @@ private[spark] class ShuffleMapStage( */ val pendingPartitions = new HashSet[Int] + /** Whether this stage produces a pipelined (incrementally-readable) shuffle. */ + val isPipelined: Boolean = shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]] + + /** + * Availability tracking for a pipelined shuffle. A pipelined shuffle produces no durable, + * addressable map output and is NOT registered with the `MapOutputTracker` at all -- + * `createShuffleMapStage` registers it only with the `StreamingShuffleOutputTracker` (the two + * trackers are split by dependency type, with no overlap). Its map-stage availability therefore + * cannot be read from the `MapOutputTracker`; instead we track completed partitions here, + * monotonically: a partition is added when its map task succeeds and is NEVER removed on + * executor/host loss. + * + * This is the crux of avoiding the streaming-writer resubmit hang: if a pipelined shuffle's + * availability were read from the `MapOutputTracker`, losing an executor that held a completed + * (already-consumed) pipelined output would strip it there, flip `isAvailable` to false, and make + * the DAGScheduler resubmit the producer -- whose streaming writer then blocks forever waiting + * for termination acks from reducers that already finished. Keeping availability local and + * monotonic means executor loss never triggers such a resubmit; a genuine mid-group failure is + * handled group-atomically instead. Unused (and empty) for a non-pipelined stage. + */ + private[this] val pipelinedCompletedPartitions = new HashSet[Int] + + /** Record a successful map task's partition as completed (pipelined stages only). */ + private[scheduler] def addPipelinedCompletedPartition(partitionId: Int): Unit = { + pipelinedCompletedPartitions += partitionId + } + override def toString: String = "ShuffleMapStage " + id /** @@ -81,7 +108,12 @@ private[spark] class ShuffleMapStage( * Number of partitions that have shuffle outputs. * When this reaches [[numPartitions]], this map stage is ready. */ - def numAvailableOutputs: Int = mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + def numAvailableOutputs: Int = { + // A pipelined shuffle is not registered with the MapOutputTracker at all (see + // pipelinedCompletedPartitions); read its locally-tracked, monotonic completed set instead. + if (isPipelined) pipelinedCompletedPartitions.size + else mapOutputTrackerMaster.getNumAvailableOutputs(shuffleDep.shuffleId) + } /** * Returns true if the map stage is ready, i.e. all partitions have shuffle outputs. @@ -90,9 +122,13 @@ private[spark] class ShuffleMapStage( /** Returns the sequence of partition ids that are missing (i.e. needs to be computed). */ override def findMissingPartitions(): Seq[Int] = { - mapOutputTrackerMaster - .findMissingPartitions(shuffleDep.shuffleId) - .getOrElse(0 until numPartitions) + if (isPipelined) { + (0 until numPartitions).filterNot(pipelinedCompletedPartitions.contains) + } else { + mapOutputTrackerMaster + .findMissingPartitions(shuffleDep.shuffleId) + .getOrElse(0 until numPartitions) + } } /** diff --git a/core/src/main/scala/org/apache/spark/scheduler/SupportsDelegationToken.scala b/core/src/main/scala/org/apache/spark/scheduler/SupportsDelegationToken.scala index 3fdbe4c25017a..6edcc7806e01c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/SupportsDelegationToken.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/SupportsDelegationToken.scala @@ -43,8 +43,9 @@ private[spark] trait SupportsDelegationToken { protected def updateDelegationTokens(tokens: Array[Byte]): Unit /** - * Whether the token manager should be started. Returns true if Hadoop security is enabled - * or if direct credential providers are configured. + * Whether the token manager should be started. The default implementation returns true when + * Hadoop security is enabled. Backends that support direct credential providers override this + * method to also check whether those providers are configured. */ protected def tokenManagerRequired(): Boolean = UserGroupInformation.isSecurityEnabled diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala index 7eba116c56800..db67b69fd9d7c 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskDescription.scala @@ -61,6 +61,16 @@ private[spark] class TaskDescription( // Eg, Map("gpu" -> Map("0" -> ResourceAmountUtils.toInternalResource(0.7))): // assign 0.7 of the gpu address "0" to this task val resources: immutable.Map[String, immutable.Map[String, Long]], + // OIDC credentials with version for stale-update prevention. + // The version is a monotonic counter from UserCredentialManager; executors only apply + // credentials if the version is newer than what they already have. This prevents a + // delayed TaskDescription from overwriting fresher credentials delivered via RPC. + // Trade-off: every TaskDescription carries the full serialized UserCredentials (a few KB). + // This is acceptable because OIDC credentials are short-lived (minutes) unlike Hadoop + // delegation tokens (hours/days), making the race between RPC broadcast and task dispatch + // a practical concern. For short-task-heavy workloads, the overhead is bounded by + // credential size x tasks-in-flight at any instant (not total task count). + val userCredentials: Option[(Long, Array[Byte])], val serializedTask: ByteBuffer) { assert(cpus > 0, "CPUs per task should be > 0") @@ -121,6 +131,17 @@ private[spark] object TaskDescription { // Write resources. serializeResources(taskDescription.resources, dataOut) + // Write user credentials (OIDC). + taskDescription.userCredentials match { + case Some((version, creds)) => + dataOut.writeBoolean(true) + dataOut.writeLong(version) + dataOut.writeInt(creds.length) + dataOut.write(creds) + case None => + dataOut.writeBoolean(false) + } + // Write the task. The task is already serialized, so write it directly to the byte buffer. Utils.writeByteBuffer(taskDescription.serializedTask, bytesOut) @@ -228,10 +249,21 @@ private[spark] object TaskDescription { // Read resources. val resources = deserializeResources(dataIn) + // Read user credentials (OIDC). + val userCredentials = if (dataIn.readBoolean()) { + val version = dataIn.readLong() + val length = dataIn.readInt() + val creds = new Array[Byte](length) + dataIn.readFully(creds) + Some((version, creds)) + } else { + None + } + // Create a sub-buffer for the serialized task into its own buffer (to be deserialized later). val serializedTask = byteBuffer.slice() new TaskDescription(taskId, attemptNumber, executorId, name, index, partitionId, artifacts, - properties, cpus, resources, serializedTask) + properties, cpus, resources, userCredentials, serializedTask) } } diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskInfo.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskInfo.scala index 9ed95870d2406..1eca39c26fccc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskInfo.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskInfo.scala @@ -149,7 +149,8 @@ class TaskInfo( def duration: Long = { if (!finished) { - throw SparkCoreErrors.durationCalledOnUnfinishedTaskError() + throw SparkCoreErrors.durationCalledOnUnfinishedTaskError( + classOf[TaskInfo].getName, "duration") } else { finishTime - launchTime } diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala index e5fdfc68701f0..4600993d2664e 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala @@ -193,6 +193,29 @@ private[spark] class TaskSchedulerImpl( }.sum } + /** + * Whether any live task set belongs to a pipelined group (see `TaskSet.isPipelined`). Such a + * group's shuffle data is transient and lives only on its executors, so a caller about to + * disturb the executors needs to know that a group is running. + * + * Zombie (superseded) attempts are skipped: their tasks are no longer scheduled, so a stale + * pipelined attempt left in the map must not make the scheduler look like it still has a + * pipelined group in flight. Matches the !isZombie filtering used elsewhere on this map. + * + * Best-effort: this reports what the TASK scheduler currently holds, which is narrower than "a + * pipelined job is active". It is false before the group's first task set is submitted, and + * false again once the last member's TaskSetManager has gone zombie but the DAGScheduler has + * not yet processed the final completion event. + * + * Synchronized like every other access to this map, so it is safe to call from any thread -- + * including off the DAGScheduler event loop (e.g. SparkContext, on a user thread). + */ + private[spark] def hasPipelinedTaskSets: Boolean = synchronized { + taskSetsByStageIdAndAttempt.values.exists(_.values.exists { tsm => + !tsm.isZombie && tsm.taskSet.isPipelined + }) + } + // The set of executors we have on each host; this is used to compute hostsAlive, which // in turn is used to decide when we can attain data locality on a given host protected val hostToExecutors = new HashMap[String, HashSet[String]] @@ -520,7 +543,7 @@ private[spark] class TaskSchedulerImpl( availWorkerResources: ExecutorResourcesAmounts): Option[Map[String, Map[String, Long]]] = { val rpId = taskSet.taskSet.resourceProfileId val taskSetProf = sc.resourceProfileManager.resourceProfileFromId(rpId) - // check if the ResourceProfile has cpus first since that is common case. Both values are in + // check if the ResourceProfile has cpus first since that is the common case. Both values are in // the internal exact BigDecimal representation, so this comparison is exact regardless of // whether spark.task.cpus is fractional (e.g. 0.2). if (availCpus < taskCpus) return None @@ -585,7 +608,7 @@ private[spark] class TaskSchedulerImpl( val shuffledOffers = shuffleOffers(filteredOffers) // Build a list of tasks to assign to each worker. // Note the size estimate here might be off with different ResourceProfiles but should be - // close estimate. It is only a capacity hint, so cap it: with a tiny fractional + // a close estimate. It is only a capacity hint, so cap it: with a tiny fractional // spark.task.cpus the exact slot count can be huge and would preallocate a giant buffer. val tasks = shuffledOffers.map { o => val sizeHint = diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala index 3513cb1f93764..407ebc53fcffc 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSet.scala @@ -33,7 +33,13 @@ private[spark] class TaskSet( val priority: Int, val properties: Properties, val resourceProfileId: Int, - val shuffleId: Option[Int]) { + val shuffleId: Option[Int], + // True if this stage is a member of a pipelined group (connected to another stage by a + // PipelinedShuffleDependency). Such a stage's transient shuffle output cannot be re-read in + // isolation, so any task failure must fail the whole group rather than be retried per-task; + // the TaskSetManager uses this to fail fast (maxTaskFailures = 1) and to count every failure, + // including otherwise-uncounted ones like executor loss. Defaults to false. + val isPipelined: Boolean = false) { val id: String = s"$stageId.$stageAttemptId" override def toString: String = "TaskSet " + id diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala index 3a5f60ae19358..9920d7ad971fb 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala @@ -70,6 +70,15 @@ private[spark] class TaskSetManager( val ser = env.closureSerializer.newInstance() val tasks = taskSet.tasks + + // A pipelined-group member reads/writes a transient shuffle that cannot be re-read in isolation, + // so a single task failure must fail the whole group (which fails the job, triggering a rerun) + // rather than be retried per-task. For such a task set we cap attempts at 1 and count every + // failure, including reasons that are normally not counted (e.g. executor loss). This is the + // native equivalent of the prototype's per-batch "any failure counts" behavior, keyed on the + // pipelined dependency (via TaskSet.isPipelined) rather than a job property. + private val effectiveMaxTaskFailures = if (taskSet.isPipelined) 1 else maxTaskFailures + private val isShuffleMapTasks = tasks(0).isInstanceOf[ShuffleMapTask] // shuffleId is only available when isShuffleMapTasks=true private val shuffleId = taskSet.shuffleId @@ -595,6 +604,7 @@ private[spark] class TaskSetManager( task.localProperties, taskCpus, taskResourceAssignments, + Option(env.userCredentials.get()).map(vc => (vc.version, vc.bytes)), serializedTask) } @@ -1106,16 +1116,34 @@ private[spark] class TaskSetManager( emptyTaskInfoAccumulablesAndNotifyDagScheduler(tid, tasks(index), reason, null, accumUpdates, metricPeaks) - if (!isZombie && reason.countTowardsTaskFailures) { + // A pipelined-group member's transient shuffle cannot be recovered per-task, so a GENUINE task + // failure must fail the whole group even when the reason is normally not counted -- most + // importantly executor loss not caused by the app (ExecutorLostFailure with + // exitCausedByApp=false), which strands the transient output. But we must NOT force-count + // reasons that are benign: TaskKilled (a deliberate kill) and TaskCommitDenied. These are not + // normally even reachable for a pipelined group -- speculation is rejected up front for a + // pipelined job, and a group never does an in-place stage/task retry (a member failure aborts + // the whole group; the caller reruns as a fresh job), so the usual sources of these reasons do + // not arise. The exclusion is kept as a defensive guard: if one ever did reach a live member + // (e.g. a manual killTaskAttempt), it must be treated as benign rather than spuriously aborting + // the group. So for a pipelined set we count everything EXCEPT those two benign reasons. + // Non-pipelined sets are unchanged. + val forceCountForPipelined = taskSet.isPipelined && (reason match { + case _: TaskKilled | _: TaskCommitDenied => false + case _ => true + }) + val countTowardsTaskFailures = reason.countTowardsTaskFailures || forceCountForPipelined + if (!isZombie && countTowardsTaskFailures) { assert (null != failureReason) taskSetExcludelistHelperOpt.foreach(_.updateExcludedForFailedTask( info.host, info.executorId, index, failureReasonString)) numFailures(index) += 1 - if (numFailures(index) >= maxTaskFailures) { + if (numFailures(index) >= effectiveMaxTaskFailures) { logError(log"Task ${MDC(TASK_INDEX, index)} in stage " + taskSet.logId + - log" failed ${MDC(MAX_ATTEMPTS, maxTaskFailures)} times; aborting job") + log" failed ${MDC(MAX_ATTEMPTS, effectiveMaxTaskFailures)} times; aborting job") abort("Task %d in stage %s failed %d times, most recent failure: %s\nDriver stacktrace:" - .format(index, taskSet.id, maxTaskFailures, failureReasonString), failureException) + .format(index, taskSet.id, effectiveMaxTaskFailures, failureReasonString), + failureException) return } } @@ -1191,7 +1219,16 @@ private[spark] class TaskSetManager( // could serve the shuffle outputs or the executor lost is caused by decommission (which // can destroy the whole host). The reason is the next stage wouldn't be able to fetch the // data from this dead executor so we would need to rerun these tasks on other executors. - val maybeShuffleMapOutputLoss = isShuffleMapTasks && + // A pipelined-group member must never single-resubmit an already-successful map task: its + // transient shuffle output is not addressable and a lone producer rerun hangs the streaming + // writer in awaitTerminationAcks. This "Resubmitted" re-enqueue loop bypasses handleFailedTask, + // so the group-atomic abort (driven from handleFailedTask via maxTaskFailures=1) would never + // see it; exclude pipelined sets here. A genuine executor loss still flows through the + // running-task loop below (iter2 -> handleFailedTask(ExecutorLostFailure)), force-counted for a + // pipelined set and aborts the whole group. Note isZombie already skips a fully-complete + // producer's set; this guard also covers a PARTIALLY-complete producer losing an executor on + // decommission. + val maybeShuffleMapOutputLoss = isShuffleMapTasks && !taskSet.isPipelined && !sched.sc.shuffleDriverComponents.supportsReliableStorage() && (reason.isInstanceOf[ExecutorDecommission] || !env.blockManager.externalShuffleServiceEnabled) if (maybeShuffleMapOutputLoss && !isZombie) { diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala index 84a6f9b3be2a1..e1be2af9294b3 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala @@ -35,6 +35,7 @@ private[spark] object CoarseGrainedClusterMessages { sparkProperties: Seq[(String, String)], ioEncryptionKey: Option[Array[Byte]], hadoopDelegationCreds: Option[Array[Byte]], + userCredentials: Option[(Long, Array[Byte])], resourceProfile: ResourceProfile, logLevel: Option[String]) extends CoarseGrainedClusterMessage @@ -60,6 +61,9 @@ private[spark] object CoarseGrainedClusterMessages { case class UpdateDelegationTokens(tokens: Array[Byte]) extends CoarseGrainedClusterMessage + case class UpdateUserCredentials(version: Long, credentials: Array[Byte]) + extends CoarseGrainedClusterMessage + // Executors to driver case class RegisterExecutor( executorId: String, diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala index 99fdc534a6fb9..1b5cd70cef908 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala @@ -26,8 +26,9 @@ import scala.concurrent.Future import com.google.common.cache.CacheBuilder -import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskState} +import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskState, VersionedCredentials} import org.apache.spark.deploy.SparkHadoopUtil +import org.apache.spark.deploy.security.UserCredentialManager import org.apache.spark.errors.SparkCoreErrors import org.apache.spark.executor.ExecutorLogUrlHandler import org.apache.spark.internal.{config, Logging} @@ -106,6 +107,17 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp // Executors which are being decommissioned. Maps from executorId to ExecutorDecommissionInfo. protected val executorsPendingDecommission = new HashMap[String, ExecutorDecommissionInfo] + // Whether the executors are held (see `SparkContext.holdExecutors()`). While true, newly + // registered executors are decommissioned immediately, so that executors the cluster manager + // granted before the hold cannot outlive it. + @volatile private var executorsHeld = false + + // Whether an executor total was ever explicitly requested (through requestExecutors or + // requestTotalExecutors), as opposed to the bookkeeping seed that `adjustExecutors` writes + // when killing an executor that was started by default. + @GuardedBy("CoarseGrainedSchedulerBackend.this") + private var explicitExecutorRequest = false + // Unknown Executors which are being decommissioned. This could be caused by unregistered executor // This executor should be decommissioned after registration. // Maps from executorId to (ExecutorDecommissionInfo, adjustTargetNumExecutors, @@ -132,6 +144,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp // Current set of delegation tokens to send to executors. private val delegationTokens = new AtomicReference[Array[Byte]]() + // UserCredentialManager for OIDC credential propagation (if enabled). + private var userCredentialManager: Option[UserCredentialManager] = None + private val reviveThread = ThreadUtils.newDaemonSingleThreadScheduledExecutor("driver-revive-thread") @@ -220,6 +235,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp case UpdateDelegationTokens(newDelegationTokens) => updateDelegationTokens(newDelegationTokens) + case UpdateUserCredentials(version, newCredentials) => + updateUserCredentials(version, newCredentials) + case RemoveExecutor(executorId, reason) => // We will remove the executor's state and cannot restore it. However, the connection // between the driver and the executor may be still alive so that the executor won't exit @@ -314,6 +332,22 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp decommissionExecutors(Array((executorId, v._1)), v._2, v._3) unknownExecutorsPendingDecommission.invalidate(executorId) }) + if (executorsHeld) { + // The executors are held; drain this late-registered executor immediately. Note + // that ExecutorMonitor cannot be told here: it builds its state asynchronously + // from the SparkListenerExecutorAdded event posted above, so this executor is not + // tracked there yet and its decommissioning is under-reported in the metrics. + // The monitor's own executor-removal handling covers the eventual exit. + decommissionExecutors( + Array((executorId, ExecutorDecommissionInfo("Executors are held"))), + adjustTargetNumExecutors = false, + triggeredByExecutor = false) + // The cluster manager granted this executor against a stale requirement (e.g. a + // restarted AM using its own initial target, or a lost response to an earlier + // request), so re-assert the zero requirement, or it would keep granting + // replacements that churn through this drain. + reassertHeldRequirement() + } context.reply(true) } @@ -358,6 +392,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp sparkProperties, SparkEnv.get.securityManager.getIOEncryptionKey(), Option(delegationTokens.get()), + Option(SparkEnv.get.userCredentials.get()).map(vc => (vc.version, vc.bytes)), rp, currentLogLevel) context.reply(reply) @@ -616,8 +651,24 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp executorsToDecommission.toImmutableArraySeq } + override def decommissionExecutorsIfIdle( + executorsAndDecomInfo: Array[(String, ExecutorDecommissionInfo)], + adjustTargetNumExecutors: Boolean): Seq[String] = withLock { + val idleExecutors = executorsAndDecomInfo.distinctBy(_._1).filter { case (executorId, _) => + isExecutorActive(executorId) && !scheduler.isExecutorBusy(executorId) + } + if (idleExecutors.isEmpty) { + Seq.empty + } else { + // Keep both locks until the existing path marks these executors pending decommission. + // Use virtual dispatch so cluster-manager overrides receive only the filtered IDs. + decommissionExecutors(idleExecutors, adjustTargetNumExecutors, triggeredByExecutor = false) + } + } + override def start(): Unit = { setupTokenManager() + setupUserCredentialManager() if (conf.get(DIRECT_CREDENTIAL_PROVIDERS_ENABLED) && delegationTokenManager.isEmpty) { logWarning("spark.security.directCredentialProviders.enabled is set but " + "this cluster manager does not support credential distribution. " + @@ -648,6 +699,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp cleanupService.foreach(_.shutdownNow()) stopExecutors() stopTokenManager() + stopUserCredentialManager() try { if (driverEndpoint != null) { driverEndpoint.askSync[Boolean](StopDriver) @@ -671,7 +723,19 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp * */ protected[scheduler] def reset(): Unit = { val executors: Set[String] = synchronized { - requestedTotalExecutorsPerResourceProfile.clear() + if (executorsHeld) { + // Keep the requested totals so that resume can restore them, and re-assert the zero + // requirement so the restarted cluster manager AM does not allocate executors + // against its own initial target while the executors are held. Do not await the + // response: this may run inside the cluster manager's RPC handler. + publishTotals() + } else { + requestedTotalExecutorsPerResourceProfile.clear() + // Clear the explicit-request record with the totals it describes, so that a later + // resume does not republish an empty map, which YARN would treat as cancelling + // every target. + explicitExecutorRequest = false + } executorDataMap.keys.toSet } @@ -729,6 +793,117 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp executorDataMap.keySet.toSeq } + /** + * Whether this backend can hold its executors gracefully: publish an all-zero executor + * requirement to the cluster manager, and let the executors that are already running finish + * their tasks and exit on their own instead of being terminated. + * + * False by default so that a backend has to opt in: the base `doRequestTotalExecutors` does + * not acknowledge the published requirement at all, and a cluster manager that terminates + * running executors once the requirement drops to zero cannot drain them gracefully. + */ + private[spark] def supportsExecutorHold: Boolean = false + + /** See `SparkContext.holdExecutors()`. */ + private[spark] def setExecutorsHeld(held: Boolean): Unit = { + executorsHeld = held + } + + /** + * Report the hold status of the application to the cluster manager, so that it can show the + * status on its own UI. Called once the application is fully started and again on every + * transition. Ignored by default: only Standalone renders it today. + */ + private[spark] def reportExecutorHoldStatus(supported: Boolean, held: Boolean): Unit = {} + + /** Whether an executor total was ever explicitly requested. Visible for testing only. */ + private[spark] def hasExplicitExecutorRequests: Boolean = synchronized { + explicitExecutorRequest + } + + /** The number of executors that are neither being removed nor decommissioned. */ + private[spark] def activeExecutorCount: Int = synchronized { + executorDataMap.keys.count(isExecutorActive) + } + + /** + * Whether the tracked totals are only killExecutors' bookkeeping seed: present, but never + * explicitly requested. Read atomically, so that a concurrent `reset()` cannot split the + * decision. + */ + private[spark] def hasKillSeededTotalsOnly: Boolean = synchronized { + !explicitExecutorRequest && requestedTotalExecutorsPerResourceProfile.nonEmpty + } + + /** + * Republish the explicitly requested totals and await the acknowledgment, or None when + * there are none: never explicitly requested, or cleared by a concurrent `reset()`. The + * check and the publish are atomic, and an empty map is never published, since some + * cluster managers treat it as cancelling every target. + */ + private[spark] def republishExplicitTotals(): Option[Boolean] = { + val response = synchronized { + if (explicitExecutorRequest && requestedTotalExecutorsPerResourceProfile.nonEmpty) { + Some(publishTotals()) + } else { + None + } + } + response.map(defaultAskTimeout.awaitResult(_)) + } + + /** + * Publish the given totals once, without recording them, and await the acknowledgment. + * Used to restore a requirement the application never explicitly requested (in particular + * Standalone's unbounded default on resume): recording it would flip `adjustExecutors`' + * empty-map seeding and mark the totals explicit, changing how `killExecutors` and a later + * hold behave for the rest of the application's life. + */ + private[spark] def publishTotalsWithoutRecording( + totals: Map[ResourceProfile, Int]): Boolean = { + val response = synchronized { + doRequestTotalExecutors(totals) + } + defaultAskTimeout.awaitResult(response) + } + + /** + * Re-assert the zero executor requirement of a held application without awaiting the + * response and without touching the requested totals. The held flag is re-checked under + * the lock: a concurrent `resumeExecutors()` lifts the flag before restoring the + * requirement, so a stale re-assertion cannot overwrite a restored one. + */ + private[spark] def reassertHeldRequirement(): Unit = synchronized { + if (executorsHeld) { + publishTotals() + } + } + + /** + * Publish the current executor totals (all-zero while the executors are held) and await + * the acknowledgment. + */ + private[spark] def republishRequestedTotals(): Boolean = { + defaultAskTimeout.awaitResult(publishTotals()) + } + + /** + * Publish the executor totals to the cluster manager: the requested totals normally, or + * all-zero totals while the executors are held, so that requirements recorded during a + * hold are retained but nothing is allocated until resume. + * + * @return a future whose evaluation indicates whether the request is acknowledged. + */ + private def publishTotals(): Future[Boolean] = synchronized { + val totals = if (executorsHeld) { + requestedTotalExecutorsPerResourceProfile.map { case (rp, _) => (rp, 0) }.toMap + + (scheduler.sc.resourceProfileManager.defaultResourceProfile -> 0) + } else { + requestedTotalExecutorsPerResourceProfile.toMap + } + doRequestTotalExecutors(totals) + } + def getExecutorsWithRegistrationTs(): Map[String, Long] = synchronized { executorDataMap.toMap.transform((_, v) => v.registrationTs) } @@ -784,6 +959,11 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp execDataOption.map(_.resourceProfileId).getOrElse(ResourceProfile.UNKNOWN_RESOURCE_PROFILE_ID) } + // this function is for testing only + private[spark] def getRequestedTotalExecutors(): Map[ResourceProfile, Int] = synchronized { + requestedTotalExecutorsPerResourceProfile.toMap + } + /** * Request an additional number of executors from the cluster manager. This is * requesting against the default ResourceProfile, we will need an API change to @@ -800,12 +980,17 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp log"executor(s) from the cluster manager") val response = synchronized { + explicitExecutorRequest = true val defaultProf = scheduler.sc.resourceProfileManager.defaultResourceProfile val numExisting = requestedTotalExecutorsPerResourceProfile.getOrElse(defaultProf, 0) - requestedTotalExecutorsPerResourceProfile(defaultProf) = numExisting + numAdditionalExecutors + // Saturate instead of overflowing when the current requirement is already huge (e.g. + // Int.MaxValue after an unbounded `requestTotalExecutors`). + val newTotal = + math.min(numExisting.toLong + numAdditionalExecutors, Int.MaxValue.toLong).toInt + requestedTotalExecutorsPerResourceProfile(defaultProf) = newTotal // Account for executors pending to be added or removed - updateExecRequestTime(defaultProf.id, numAdditionalExecutors) - doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap) + updateExecRequestTime(defaultProf.id, newTotal - numExisting) + publishTotals() } defaultAskTimeout.awaitResult(response) @@ -843,6 +1028,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp (scheduler.sc.resourceProfileManager.resourceProfileFromId(rpid), num) } val response = synchronized { + explicitExecutorRequest = true val oldResourceProfileToNumExecutors = requestedTotalExecutorsPerResourceProfile.map { case (rp, num) => (rp.id, num) @@ -852,7 +1038,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp this.numLocalityAwareTasksPerResourceProfileId = numLocalityAwareTasksPerResourceProfileId this.rpHostToLocalTaskCount = hostToLocalTaskCount updateExecRequestTimes(oldResourceProfileToNumExecutors, resourceProfileIdToNumExecutors) - doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap) + publishTotals() } defaultAskTimeout.awaitResult(response) } @@ -925,7 +1111,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp } } } - doRequestTotalExecutors(requestedTotalExecutorsPerResourceProfile.toMap) + publishTotals() } else { Future.successful(true) } @@ -1047,6 +1233,45 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp protected def currentDelegationTokens: Array[Byte] = delegationTokens.get() + /** + * Update user credentials and broadcast to all registered executors. + * Called from the DriverEndpoint receive loop (thread-safe access to executorDataMap). + */ + private def updateUserCredentials(version: Long, credentials: Array[Byte]): Unit = { + VersionedCredentials.updateIfNewer(SparkEnv.get.userCredentials, version, credentials) + executorDataMap.values.foreach { ed => + ed.executorEndpoint.send(UpdateUserCredentials(version, credentials)) + } + } + + /** + * Start the UserCredentialManager if OIDC credential propagation is enabled. + * Called from start(), independently of Kerberos/HadoopDelegationTokenManager. + */ + private def setupUserCredentialManager(): Unit = { + userCredentialManager = UserCredentialManager.create(conf, { (version, credentials) => + // Send to DriverEndpoint to ensure thread-safe access to executorDataMap. + // This mirrors HadoopDelegationTokenManager's pattern of sending + // UpdateDelegationTokens via schedulerRef. + driverEndpoint.send(UpdateUserCredentials(version, credentials)) + }) + userCredentialManager.foreach { manager => + val (version, initialCredentials) = manager.start() + // Store initial credentials synchronously so they are available for SparkAppConfig + // (late-registering executors) and TaskDescription (task dispatch) immediately. + // Note: the onCredentialsUpdate callback above also triggers an async + // UpdateUserCredentials message that will redundantly call updateIfNewer. + // The synchronous set here ensures no null window before the async message + // is processed by DriverEndpoint. + VersionedCredentials.updateIfNewer( + SparkEnv.get.userCredentials, version, initialCredentials) + } + } + + private def stopUserCredentialManager(): Unit = { + userCredentialManager.foreach(_.stop()) + } + /** * Checks whether the executor is excluded due to failure(s). This is called when the executor * tries to register with the scheduler, and will deny registration if this method returns true. diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala index 061b54914c839..aebe364c92939 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala @@ -248,6 +248,14 @@ private[spark] class StandaloneSchedulerBackend( } } + // The Master honors a zero executor limit without killing the running executors, so the + // executors can be held gracefully. + private[spark] override def supportsExecutorHold: Boolean = true + + private[spark] override def reportExecutorHoldStatus(supported: Boolean, held: Boolean): Unit = { + Option(client).foreach(_.reportHoldStatus(supported, held)) + } + /** * Kill the given list of executors through the Master. * @return whether the kill request is acknowledged. diff --git a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala index 65f0ffcbe18ea..9b120d2f5bdb3 100644 --- a/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala +++ b/core/src/main/scala/org/apache/spark/serializer/KryoSerializer.scala @@ -555,7 +555,8 @@ private[serializer] object KryoSerializer { classOf[SparkConf], classOf[TaskCommitMessage], classOf[SerializedLambda], - classOf[BitSet] + classOf[BitSet], + classOf[java.util.HashMap[_, _]] ) private val toRegisterSerializer = Map[Class[_], KryoClassSerializer[_]]( diff --git a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala index 03ca257353c59..5ff9f63394258 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala +++ b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala @@ -74,6 +74,20 @@ object StreamingShuffleManager extends Logging { errorClass = "STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE", messageParameters = Map("messageType" -> messageType.toString)) } + + def streamingShuffleWriterConnectionTimeout( + shuffleId: Int, + writerId: Int, + readerId: Int, + timeoutMs: Long): RuntimeException = { + new SparkRuntimeException( + errorClass = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT", + messageParameters = Map( + "shuffleId" -> shuffleId.toString, + "writerId" -> writerId.toString, + "readerId" -> readerId.toString, + "timeoutMs" -> timeoutMs.toString)) + } } private[spark] class StreamingShuffleManager extends PipelinedShuffleManager with Logging { @@ -113,8 +127,8 @@ private[spark] class StreamingShuffleManager extends PipelinedShuffleManager wit override def unregisterShuffle(shuffleId: Int): Boolean = { // No manager-side state to release here: the driver's StreamingShuffleOutputTracker is - // unregistered in BlockManagerStorageEndpoint's RemoveShuffle handler, and per-task writer - // and reader resources are released via task completion listeners. + // unregistered in ContextCleaner.doCleanupShuffle (its state is driver-only), and per-task + // writer and reader resources are released via task completion listeners. true } diff --git a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala index b0cc249f73422..0b6719127fd18 100644 --- a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala +++ b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala @@ -17,7 +17,7 @@ package org.apache.spark.shuffle.streaming -import java.util.concurrent.{CancellationException, CompletableFuture, CountDownLatch, LinkedBlockingDeque, Semaphore, TimeUnit} +import java.util.concurrent.{CancellationException, CompletableFuture, CompletionException, CountDownLatch, LinkedBlockingDeque, Semaphore, TimeoutException, TimeUnit} import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} import javax.annotation.concurrent.NotThreadSafe @@ -29,7 +29,13 @@ import io.netty.channel.{ChannelFuture, ChannelOption} import org.apache.spark.{SparkContext, SparkEnv, StreamingShuffleTaskLocation, TaskContext} import org.apache.spark.internal.LogKeys -import org.apache.spark.internal.config.{EXECUTOR_ID, STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, STREAMING_SHUFFLE_WRITER_MAX_MEMORY} +import org.apache.spark.internal.config.{ + EXECUTOR_ID, + STREAMING_SHUFFLE_CHECKSUM_ENABLED, + STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, + STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, + STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS, + STREAMING_SHUFFLE_WRITER_MAX_MEMORY} import org.apache.spark.internal.config.Network.RPC_IO_THREADS import org.apache.spark.memory.{MemoryConsumer, MemoryMode} import org.apache.spark.network.TransportContext @@ -58,6 +64,7 @@ class StreamingShuffleWriter[K, V]( private val BUFFER_SIZE: Integer = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE) // The interval at which we flush pending messages. private val MAX_BUFFERING_TIME_MS = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS) + private val CONNECTION_TIMEOUT_MS = conf.get(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS) // Shuffle details. private val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K, V, _]] @@ -205,11 +212,27 @@ class StreamingShuffleWriter[K, V]( private[streaming] case class ShardState(id: Int) { // client may be accessed from other threads via cancel(); @volatile to be safe. @volatile private var client: Either[TransportClient, CompletableFuture[TransportClient]] = - Right(transportServerHandler.futureClients(id).thenApply(c => { - c.getChannel.config.setOption(ChannelOption.SO_SNDBUF, SEND_BUFFER_SIZE) - c.getChannel.config.setOption(ChannelOption.SO_RCVBUF, RECV_BUFFER_SIZE) - c - })) + Right( + (if (CONNECTION_TIMEOUT_MS == -1) { + transportServerHandler.futureClients(id) + } else { + transportServerHandler.futureClients(id) + .orTimeout(CONNECTION_TIMEOUT_MS, TimeUnit.MILLISECONDS) + .exceptionally { + case _: TimeoutException => + throw StreamingShuffleManager.streamingShuffleWriterConnectionTimeout( + streamingShuffleHandle.shuffleId, + shuffleWriterId, + id, + CONNECTION_TIMEOUT_MS) + case error => + throw error + } + }).thenApply(c => { + c.getChannel.config.setOption(ChannelOption.SO_SNDBUF, SEND_BUFFER_SIZE) + c.getChannel.config.setOption(ChannelOption.SO_RCVBUF, RECV_BUFFER_SIZE) + c + })) val buffer: AtomicReference[TimestampedBuffer] = new AtomicReference(null) val lastSentSequenceNum: AtomicLong = new AtomicLong(-1) val terminationAckReceived: AtomicBoolean = new AtomicBoolean(false) @@ -257,7 +280,10 @@ class StreamingShuffleWriter[K, V]( val newFuture = future.whenComplete { (client, ex) => ex match { case null => sendToClient(client) - case _ => buf.release(); done() + case error => + buf.release() + errorNotifier.markError(error) + done() } } // Once the future is completed, stop accumulating CompletionStages. @@ -318,7 +344,7 @@ class StreamingShuffleWriter[K, V]( // For testing only. def hasClient: Boolean = client match { case Left(_) => true - case Right(future) => future.isDone + case Right(future) => future.isDone && !future.isCompletedExceptionally } } @@ -395,6 +421,11 @@ class StreamingShuffleWriter[K, V]( System.currentTimeMillis() - cleanupStartTime)} ms") } + private def unwrapCompletionException(error: Throwable): Throwable = error match { + case e: CompletionException if e.getCause != null => e.getCause + case _ => error + } + private def throwErrorIfExists(): Unit = { context.getTaskFailure.foreach { throw _ } errorNotifier.throwErrorIfExists() @@ -430,7 +461,7 @@ class StreamingShuffleWriter[K, V]( Try { while (!isWriteFinished.await(MAX_BUFFERING_TIME_MS, TimeUnit.MILLISECONDS)) shards.foreach(_.send()) - }.recover { case e => errorNotifier.markError(e) } + }.recover { case e => errorNotifier.markError(unwrapCompletionException(e)) } , "time-based-flush-for-shuffle-writer-" + s"${streamingShuffleHandle.shuffleId}-${shuffleWriterId}") try { @@ -508,6 +539,9 @@ class StreamingShuffleWriter[K, V]( logInfo(log"Received all termination acks for shuffle writer ${MDC( LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}. Closing server channel.") throwErrorIfExists() + } catch { + case e: CompletionException => + throw unwrapCompletionException(e) } finally { isWriteFinished.countDown() // Duplicate countDowns are a no-op. flushThread.join() diff --git a/core/src/main/scala/org/apache/spark/status/api/v1/OneApplicationResource.scala b/core/src/main/scala/org/apache/spark/status/api/v1/OneApplicationResource.scala index de25e7c524ead..c0f73559caecb 100644 --- a/core/src/main/scala/org/apache/spark/status/api/v1/OneApplicationResource.scala +++ b/core/src/main/scala/org/apache/spark/status/api/v1/OneApplicationResource.scala @@ -57,7 +57,7 @@ private[v1] class AbstractApplicationResource extends BaseAppResource { @Path("executors/{executorId}/threads") def threadDump(@PathParam("executorId") execId: String): Array[ThreadStackTrace] = withUI { ui => checkExecutorId(execId) - val safeSparkContext = checkAndGetSparkContext() + val safeSparkContext = checkAndGetSparkContext("Thread dumps") ui.store.asOption(ui.store.executorSummary(execId)) match { case Some(executorSummary) if executorSummary.isActive => val safeThreadDump = safeSparkContext.getExecutorThreadDump(execId).getOrElse { @@ -75,7 +75,7 @@ private[v1] class AbstractApplicationResource extends BaseAppResource { @QueryParam("taskId") taskId: Long, @QueryParam("executorId") execId: String): ThreadStackTrace = { checkExecutorId(execId) - val safeSparkContext = checkAndGetSparkContext() + val safeSparkContext = checkAndGetSparkContext("Thread dumps") safeSparkContext .getTaskThreadDump(taskId, execId) .getOrElse { @@ -92,6 +92,18 @@ private[v1] class AbstractApplicationResource extends BaseAppResource { @Path("allmiscellaneousprocess") def allProcessList(): Seq[ProcessSummary] = withUI(_.store.miscellaneousProcessList(false)) + @GET + @Path("holdstatus") + def holdStatus(): ApplicationHoldStatus = { + val safeSparkContext = checkAndGetSparkContext("Hold status") + val held = safeSparkContext.executorsHeld + new ApplicationHoldStatus( + supported = safeSparkContext.executorHoldSupported, + held = held, + // The executors that have not exited yet are still draining their running tasks. + draining = if (held) safeSparkContext.getExecutorIds().size else 0) + } + @Path("stages") def stages(): Class[StagesResource] = classOf[StagesResource] @@ -188,9 +200,13 @@ private[v1] class AbstractApplicationResource extends BaseAppResource { } } - private def checkAndGetSparkContext(): SparkContext = withUI { ui => + /** + * Returns the live `SparkContext` for live-only endpoints. `feature` is the subject of the + * error message, e.g. "Thread dumps" -> "Thread dumps not available through the history server." + */ + private def checkAndGetSparkContext(feature: String): SparkContext = withUI { ui => ui.sc.getOrElse { - throw new ServiceUnavailable("Thread dumps not available through the history server.") + throw new ServiceUnavailable(s"$feature not available through the history server.") } } } diff --git a/core/src/main/scala/org/apache/spark/status/api/v1/api.scala b/core/src/main/scala/org/apache/spark/status/api/v1/api.scala index 3e7f5e6d2ec3e..1eb25f14e3348 100644 --- a/core/src/main/scala/org/apache/spark/status/api/v1/api.scala +++ b/core/src/main/scala/org/apache/spark/status/api/v1/api.scala @@ -66,6 +66,11 @@ case class ApplicationAttemptInfo private[spark]( } +class ApplicationHoldStatus private[spark]( + val supported: Boolean, + val held: Boolean, + val draining: Int) + class ResourceProfileInfo private[spark]( val id: Int, val executorResources: Map[String, ExecutorResourceRequest], diff --git a/core/src/main/scala/org/apache/spark/storage/BlockId.scala b/core/src/main/scala/org/apache/spark/storage/BlockId.scala index 3e46a53ee082c..8d44e79a223af 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockId.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockId.scala @@ -263,7 +263,11 @@ case class CacheId(sessionUUID: String, hash: String) extends BlockId { @DeveloperApi object BlockId { - val RDD = "rdd_([0-9]+)_([0-9]+)".r + // Safety net: newRddId() fails fast before minting negative ids, but names like + // rdd_-1330910599_36 may still appear from blocks cached before upgrade or from + // tests. Accept an optional minus so BlockId.apply does not throw + // UnrecognizedBlockId for those names. + val RDD = "rdd_(-?[0-9]+)_([0-9]+)".r val SHUFFLE = "shuffle_([0-9]+)_([0-9]+)_([0-9]+)".r val SHUFFLE_BATCH = "shuffle_([0-9]+)_([0-9]+)_([0-9]+)_([0-9]+)".r val SHUFFLE_DATA = "shuffle_([0-9]+)_([0-9]+)_([0-9]+).data".r diff --git a/core/src/main/scala/org/apache/spark/storage/BlockInfoManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockInfoManager.scala index 4dc8d1f1b50c7..2c0f82283e1bc 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockInfoManager.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockInfoManager.scala @@ -288,7 +288,7 @@ private[storage] class BlockInfoManager(trackingCacheVisibility: Boolean = false } /** - * Helper for lock acquisistion. + * Helper for lock acquisition. */ private def acquireLock( blockId: BlockId, diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala index 801f4c0b74720..7267348031ed5 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala @@ -948,9 +948,9 @@ private[spark] class BlockManager( try { return migratableResolver.putShuffleBlockAsStream(blockId, serializerManager) } catch { - case _: ClassCastException => - throw SparkCoreErrors.unexpectedShuffleBlockWithUnsupportedResolverError( - shuffleBlockResolver, blockId) + case e: ClassCastException => + throw SparkCoreErrors.shuffleBlockMigrationNotSupportedError( + blockId, shuffleBlockResolver, e) } } logDebug(s"Putting regular block ${blockId}") @@ -1128,7 +1128,7 @@ private[spark] class BlockManager( releaseLock(blockId) // Remove the missing block so that its unavailability is reported to the driver removeBlock(blockId) - throw SparkCoreErrors.readLockedBlockNotFoundError(blockId) + throw SparkCoreErrors.localBlockDataNotFoundError(blockId) } private def isIORelatedException(t: Throwable): Boolean = diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala index 84a70d27047ca..47b6e2a50525f 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMaster.scala @@ -342,7 +342,7 @@ class BlockManagerMaster( /** Send a one-way message to the master endpoint, to which we expect it to reply with true. */ private def tell(message: Any): Unit = { if (!driverEndpoint.askSync[Boolean](message)) { - throw SparkCoreErrors.unexpectedBlockManagerMasterEndpointResultError() + throw SparkCoreErrors.unexpectedBlockManagerMasterEndpointResultError(message) } } diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala index 08d8027406188..beca0d11d1c5e 100644 --- a/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala +++ b/core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala @@ -804,7 +804,7 @@ class BlockManagerMasterEndpoint( case ShuffleIndexBlockId(shuffleId, mapId, _) => // SPARK-36782: Invoke `MapOutputTracker.updateMapOutput` within the thread // `dispatcher-BlockManagerMaster` could lead to the deadlock when - // `MapOutputTracker.serializeOutputStatuses` broadcasts the serialized mapstatues under + // `MapOutputTracker.serializeOutputStatuses` broadcasts the serialized map statuses under // the acquired write lock. The broadcast block would report its status to // `BlockManagerMasterEndpoint`, while the `BlockManagerMasterEndpoint` is occupied by // `updateMapOutput` since it's waiting for the write lock. Thus, we use `Future` to call diff --git a/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala b/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala index 9964e64724f6b..3ede609743879 100644 --- a/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala +++ b/core/src/main/scala/org/apache/spark/storage/DiskBlockObjectWriter.scala @@ -22,8 +22,7 @@ import java.nio.channels.{ClosedByInterruptException, FileChannel} import java.nio.file.Files import java.util.zip.Checksum -import org.apache.spark.SparkException -import org.apache.spark.errors.SparkCoreErrors +import org.apache.spark.{SparkException, SparkUnsupportedOperationException} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys._ import org.apache.spark.io.MutableCheckedOutputStream @@ -338,7 +337,11 @@ private[spark] class DiskBlockObjectWriter( recordWritten() } - override def write(b: Int): Unit = throw SparkCoreErrors.unsupportedOperationError() + override def write(b: Int): Unit = throw new SparkUnsupportedOperationException( + errorClass = "UNSUPPORTED_CALL.WITHOUT_SUGGESTION", + messageParameters = Map( + "className" -> classOf[DiskBlockObjectWriter].getName, + "methodName" -> "write")) override def write(kvBytes: Array[Byte], offs: Int, len: Int): Unit = { if (!streamOpen) { diff --git a/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala b/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala index cb15e954bb38a..84707424d7b0e 100644 --- a/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala +++ b/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala @@ -885,7 +885,7 @@ final class ShuffleBlockFetcherIterator( bufIn } } catch { - // The exception could only be throwed by local shuffle block + // The exception could only be thrown by local shuffle block case e: IOException => assert(buf.isInstanceOf[FileSegmentManagedBuffer]) e match { @@ -1668,7 +1668,7 @@ object ShuffleBlockFetcherIterator { * of shuffle by an indeterminate stage attempt. * @param reduceId reduce id. * @param bitmaps bitmaps for every chunk. - * @param localDirs local directories where the push-merged shuffle files are storedl + * @param localDirs local directories where the push-merged shuffle files are stored */ private[storage] case class PushMergedLocalMetaFetchResult( shuffleId: Int, diff --git a/core/src/main/scala/org/apache/spark/ui/HttpSecurityFilter.scala b/core/src/main/scala/org/apache/spark/ui/HttpSecurityFilter.scala index edfa3581d9ab2..86c0df0e50b9c 100644 --- a/core/src/main/scala/org/apache/spark/ui/HttpSecurityFilter.scala +++ b/core/src/main/scala/org/apache/spark/ui/HttpSecurityFilter.scala @@ -49,9 +49,11 @@ private class HttpSecurityFilter( val hres = res.asInstanceOf[HttpServletResponse] hres.setHeader("Cache-Control", "no-cache, no-store, must-revalidate") + val isProxyRequest = hreq.getContextPath == "/proxy" + val cspNonce = CspNonce.generate() try { - if (conf.get(UI_CONTENT_SECURITY_POLICY_ENABLED)) { + if (conf.get(UI_CONTENT_SECURITY_POLICY_ENABLED) && !isProxyRequest) { // Use CSP frame-ancestors as the primary clickjacking protection mechanism. // X-Frame-Options ALLOW-FROM is deprecated and ignored by modern browsers // (Chrome, Firefox, Edge, Safari), so frame-ancestors is used instead. diff --git a/core/src/main/scala/org/apache/spark/ui/JettyUtils.scala b/core/src/main/scala/org/apache/spark/ui/JettyUtils.scala index 3c21fd6d2a003..164a3bccf2f39 100644 --- a/core/src/main/scala/org/apache/spark/ui/JettyUtils.scala +++ b/core/src/main/scala/org/apache/spark/ui/JettyUtils.scala @@ -27,7 +27,7 @@ import scala.xml.Node import jakarta.servlet.{DispatcherType, Filter, FilterChain, ServletRequest, ServletResponse} import jakarta.servlet.http._ -import org.eclipse.jetty.client.{Response => CResponse} +import org.eclipse.jetty.client.{Request => CRequest, Response => CResponse} import org.eclipse.jetty.client.HttpClient import org.eclipse.jetty.client.transport.HttpClientTransportOverHTTP import org.eclipse.jetty.compression.server.CompressionHandler @@ -209,6 +209,25 @@ private[spark] object JettyUtils extends Logging { .orNull } + override def addProxyHeaders( + clientRequest: HttpServletRequest, + proxyRequest: CRequest): Unit = { + super.addProxyHeaders(clientRequest, proxyRequest) + val path = clientRequest.getPathInfo + if (path != null) { + val prefixTrailingSlashIndex = path.indexOf('/', 1) + val prefix = if (prefixTrailingSlashIndex == -1) { + path + } else { + path.substring(0, prefixTrailingSlashIndex) + } + val existingContext = Option(clientRequest.getHeader("X-Forwarded-Context")).getOrElse("") + val contextPath = Option(clientRequest.getContextPath).getOrElse("") + val proxyContext = existingContext + contextPath + prefix + proxyRequest.headers(headers => headers.put("X-Forwarded-Context", proxyContext)) + } + } + override def newHttpClient(): HttpClient = { // SPARK-21176: Use the Jetty logic to calculate the number of selector threads (#CPUs/2), // but limit it to 8 max. diff --git a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala index 862e150acd441..97c5d16873445 100644 --- a/core/src/main/scala/org/apache/spark/ui/SparkUI.scala +++ b/core/src/main/scala/org/apache/spark/ui/SparkUI.scala @@ -55,6 +55,8 @@ private[spark] class SparkUI private ( val killEnabled = sc.map(_.conf.get(UI_KILL_ENABLED)).getOrElse(false) + val holdEnabled = sc.map(_.conf.get(UI_HOLD_ENABLED)).getOrElse(false) + var appId: String = _ private var streamingJobProgressListener: Option[SparkListener] = None @@ -94,9 +96,11 @@ private[spark] class SparkUI private ( } } + private var jobsTab: JobsTab = _ + /** Initialize all components of the server. */ def initialize(): Unit = { - val jobsTab = new JobsTab(this, store) + jobsTab = new JobsTab(this, store) attachTab(jobsTab) val stagesTab = new StagesTab(this, store) attachTab(stagesTab) @@ -123,6 +127,10 @@ private[spark] class SparkUI private ( attachHandler(createRedirectHandler( "/stages/stage/kill", "/stages/", stagesTab.handleKillRequest, httpMethods = Set("GET", "POST"))) + attachHandler(createRedirectHandler( + "/jobs/hold", "/jobs/", jobsTab.handleHoldRequest, httpMethods = Set("GET", "POST"))) + attachHandler(createRedirectHandler( + "/jobs/resume", "/jobs/", jobsTab.handleResumeRequest, httpMethods = Set("GET", "POST"))) } initialize() @@ -164,6 +172,7 @@ private[spark] class SparkUI private ( /** Stop the server behind this web interface. Only valid after bind(). */ override def stop(): Unit = { super.stop() + Option(jobsTab).foreach(_.stop()) logInfo(log"Stopped Spark web UI at ${MDC(WEB_URL, webUrl)}") } diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/AllJobsPage.scala b/core/src/main/scala/org/apache/spark/ui/jobs/AllJobsPage.scala index 70d0a9faf4de2..a0f9e52f2fb83 100644 --- a/core/src/main/scala/org/apache/spark/ui/jobs/AllJobsPage.scala +++ b/core/src/main/scala/org/apache/spark/ui/jobs/AllJobsPage.scala @@ -354,6 +354,31 @@ private[ui] class AllJobsPage(parent: JobsTab, store: AppStatusStore) extends We <strong>Scheduling Mode: </strong> {schedulingMode} </li> + { + if (parent.holdEnabled && parent.sc.exists(_.executorHoldSupported)) { + val basePathUri = UIUtils.prependBaseUri(request, parent.basePath) + val (status, action, confirm) = if (parent.sc.get.executorsHeld) { + val numDraining = parent.sc.get.getExecutorIds().size + val status = if (numDraining > 0) { + s"Held (draining $numDraining executor${if (numDraining > 1) "s" else ""})" + } else { + "Held" + } + (status, "resume", "Are you sure you want to resume this application?") + } else { + ("Running", "hold", "Are you sure you want to hold this application? All " + + "executors will be decommissioned after finishing their running tasks.") + } + <li> + <strong>Application:</strong> + {status} + <a href={s"$basePathUri/jobs/$action/"} + data-confirm-message={confirm} + class="confirm-link">{s"($action)"}</a> + {parent.lastHoldRequestStatus.getOrElse("")} + </li> + } + } { if (shouldShowActiveJobs) { <li> diff --git a/core/src/main/scala/org/apache/spark/ui/jobs/JobsTab.scala b/core/src/main/scala/org/apache/spark/ui/jobs/JobsTab.scala index a3c2d05414a88..98ae0de22fb3f 100644 --- a/core/src/main/scala/org/apache/spark/ui/jobs/JobsTab.scala +++ b/core/src/main/scala/org/apache/spark/ui/jobs/JobsTab.scala @@ -17,6 +17,8 @@ package org.apache.spark.ui.jobs +import java.util.concurrent.{ExecutorService, RejectedExecutionException} + import jakarta.servlet.http.HttpServletRequest import org.apache.spark.JobExecutionStatus @@ -24,6 +26,7 @@ import org.apache.spark.internal.config.SCHEDULER_MODE import org.apache.spark.scheduler.SchedulingMode import org.apache.spark.status.AppStatusStore import org.apache.spark.ui._ +import org.apache.spark.util.{ThreadUtils, Utils} /** Web UI showing progress status of all jobs in the given SparkContext. */ private[ui] class JobsTab(parent: SparkUI, store: AppStatusStore) @@ -32,6 +35,7 @@ private[ui] class JobsTab(parent: SparkUI, store: AppStatusStore) val sc = parent.sc val conf = parent.conf val killEnabled = parent.killEnabled + val holdEnabled = parent.holdEnabled // Show pool information for only live UI. def isFairScheduler: Boolean = { @@ -62,4 +66,90 @@ private[ui] class JobsTab(parent: SparkUI, store: AppStatusStore) } } } + + // Serves the hold/resume requests off the Jetty serving thread: they talk to the cluster + // manager and may block up to the RPC ask timeout. A single thread also serializes + // concurrent requests. Created on first use and shut down by `stop()`, so that the thread + // does not outlive the SparkContext. + private var holdRequestExecutor: Option[ExecutorService] = None + private var stopped = false + + // None once stopped, so that a request served during teardown neither hits a rejected + // execution on the shut-down pool nor recreates it and leaks the thread. + private def holdRequestExecutorPool: Option[ExecutorService] = synchronized { + if (stopped) { + None + } else { + Some(holdRequestExecutor.getOrElse { + val pool = ThreadUtils.newDaemonSingleThreadExecutor("spark-ui-hold-resume") + holdRequestExecutor = Some(pool) + pool + }) + } + } + + def stop(): Unit = synchronized { + stopped = true + holdRequestExecutor.foreach(_.shutdownNow()) + } + + // Outcome of the last hold/resume request served by this tab, as (isHold, message): Some + // while a request is running or after it did not take effect, None when idle or after a + // success. Rendered with the operation it belongs to, so a stale message never reads as + // the opposite operation's, and a failed hold stays visible even though the page already + // shows the (resume) control (the hold is marked before the cluster manager is asked). + @volatile private var holdRequestStatus: Option[(Boolean, String)] = None + + private[jobs] def lastHoldRequestStatus: Option[String] = + holdRequestStatus.map { case (isHold, message) => + s"(${if (isHold) "hold" else "resume"}: $message)" + } + + def handleHoldRequest(request: HttpServletRequest): Unit = { + if (holdEnabled && parent.securityManager.checkModifyPermissions(request.getRemoteUser)) { + sc.filter(_.executorHoldSupported).foreach { ctx => + holdRequestExecutorPool.foreach { pool => + holdRequestStatus = Some((true, "requested")) + try { + pool.execute { () => + var acknowledged = false + Utils.tryLogNonFatalError { acknowledged = ctx.holdExecutors() } + holdRequestStatus = if (acknowledged) { + None + } else { + Some((true, "the last request did not take effect, see the driver logs")) + } + } + } catch { + // stop() may have shut the pool down after the accessor returned it + case _: RejectedExecutionException => + } + } + } + } + } + + def handleResumeRequest(request: HttpServletRequest): Unit = { + if (holdEnabled && parent.securityManager.checkModifyPermissions(request.getRemoteUser)) { + sc.filter(_.executorHoldSupported).foreach { ctx => + holdRequestExecutorPool.foreach { pool => + holdRequestStatus = Some((false, "requested")) + try { + pool.execute { () => + var acknowledged = false + Utils.tryLogNonFatalError { acknowledged = ctx.resumeExecutors() } + holdRequestStatus = if (acknowledged) { + None + } else { + Some((false, "the last request did not take effect, see the driver logs")) + } + } + } catch { + // stop() may have shut the pool down after the accessor returned it + case _: RejectedExecutionException => + } + } + } + } + } } diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala index c8030e5a4c7c8..8507c889c64ba 100644 --- a/core/src/main/scala/org/apache/spark/util/Utils.scala +++ b/core/src/main/scala/org/apache/spark/util/Utils.scala @@ -946,11 +946,12 @@ private[spark] object Utils */ private[spark] def normalizeIpIfNeeded(host: String): String = { // Is this a v6 address. We ask users to add [] around v6 addresses as strs but - // there not always there. If it's just 0-9 and : and [] we treat it as a v6 address. + // they're not always there. If it's just hexadecimal digits, :, and [] we treat it as a + // v6 address. // This means some invalid addresses are treated as v6 addresses, but since they are // not valid hostnames it doesn't matter. // See https://www.rfc-editor.org/rfc/rfc1123#page-13 for context around valid hostnames. - val addressRe = """^\[{0,1}([0-9:]+?:[0-9]*)\]{0,1}$""".r + val addressRe = """^\[{0,1}([0-9a-fA-F:]+?:[0-9a-fA-F]*)\]{0,1}$""".r host match { case addressRe(unbracketed) => addBracketsIfNeeded(InetAddresses.toAddrString(InetAddresses.forString(unbracketed))) @@ -959,6 +960,12 @@ private[spark] object Utils } } + /** Returns whether a literal IPv4 or IPv6 address binds to every local interface. */ + private[spark] def isAnyLocalAddress(host: String): Boolean = { + val address = host.stripPrefix("[").stripSuffix("]") + InetAddresses.isInetAddress(address) && InetAddresses.forString(address).isAnyLocalAddress + } + /** * Checks if the host contains only valid hostname/ip without port * NOTE: Incase of IPV6 ip it should be enclosed inside [] diff --git a/core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala b/core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala index d892fa0e47060..06dcd136faacc 100644 --- a/core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala +++ b/core/src/main/scala/org/apache/spark/util/collection/ExternalAppendOnlyMap.scala @@ -309,7 +309,7 @@ class ExternalAppendOnlyMap[K, V, C]( inputStreams.foreach { it => val kcPairs = new ArrayBuffer[(K, C)] readNextHashCode(it, kcPairs) - if (kcPairs.length > 0) { + if (kcPairs.nonEmpty) { mergeHeap.enqueue(new StreamBuffer(it, kcPairs)) } } @@ -426,11 +426,11 @@ class ExternalAppendOnlyMap[K, V, C]( val pairs: ArrayBuffer[(K, C)]) extends Comparable[StreamBuffer] { - def isEmpty: Boolean = pairs.length == 0 + def isEmpty: Boolean = pairs.isEmpty // Invalid if there are no more pairs in this stream def minKeyHash: Int = { - assert(pairs.length > 0) + assert(pairs.nonEmpty) hashKey(pairs.head) } diff --git a/core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala b/core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala index e421a1f4746ea..3c7f083381c3c 100644 --- a/core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala +++ b/core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala @@ -21,8 +21,9 @@ import scala.reflect.ClassTag /** * A fast hash map implementation for nullable keys. This hash map supports insertions and updates, - * but not deletions. This map is about 5X faster than java.util.HashMap, while using much less - * space overhead. + * but not deletions. This map uses much less space than java.util.HashMap and is competitive with + * it for aggregation workloads (`changeValue`), while java.util.HashMap is faster for pure + * insertions and lookups on modern JDKs. See `OpenHashMapBenchmark` for details. * * Under the hood, it uses our OpenHashSet implementation. * diff --git a/core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala b/core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala index 3d1eb5788c707..c56ae0224740a 100644 --- a/core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala +++ b/core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala @@ -28,9 +28,10 @@ import org.apache.spark.annotation.Private * removed. * * The underlying implementation uses Scala compiler's specialization to generate optimized - * storage for four primitive types (Long, Int, Double, and Float). It is much faster than Java's - * standard HashSet while incurring much less memory overhead. This can serve as building blocks - * for higher level data structures such as an optimized HashMap. + * storage for four primitive types (Long, Int, Double, and Float). It incurs much less memory + * overhead than Java's standard HashSet, and the specialized versions avoid boxing of primitive + * keys. This can serve as building blocks for higher level data structures such as an optimized + * HashMap. * * This OpenHashSet is designed to serve as building blocks for higher level data structures * such as an optimized hash map. Compared with standard hash set implementations, this class diff --git a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java index 6c408c0c928b5..44e17d040ea32 100644 --- a/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java +++ b/core/src/test/java/org/apache/spark/memory/TaskMemoryManagerSuite.java @@ -151,6 +151,26 @@ public long spill(long size, MemoryConsumer trigger) throws IOException { } } + private static final class NonSpillingConsumer extends MemoryConsumer { + NonSpillingConsumer(TaskMemoryManager manager) { + super(manager, 1024L, MemoryMode.ON_HEAP); + } + + void use(long size) { + used.getAndAdd(taskMemoryManager.acquireExecutionMemory(size, this)); + } + + void free(long size) { + used.getAndAdd(-size); + taskMemoryManager.releaseExecutionMemory(size, this); + } + + @Override + public long spill(long size, MemoryConsumer trigger) { + return 0; + } + } + private static final class NonSpillingAllocatingConsumer extends TestMemoryConsumer { private int spillAttempts; private MemoryBlock nestedPage; @@ -841,6 +861,155 @@ public void shouldNotForceSpillingInDifferentModes() { Assertions.assertEquals(80, c1.getUsed()); // not spilled } + @Test + public void memoryConsumptionBreakdownIsEmptyWhenThereIsNoMemoryToReport() { + final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf()); + memoryManager.limit(100); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + // The breakdown is empty only when the task holds neither attributed nor unattributed memory. + // A consumer that has never acquired memory must not appear in the breakdown. + new TestMemoryConsumer(manager); + Assertions.assertEquals(0, memoryManager.getExecutionMemoryUsageForTask(0)); + Assertions.assertEquals("", manager.getMemoryConsumptionBreakdown()); + } + + @Test + public void memoryConsumptionBreakdownListsConsumersLargestFirst() { + final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf()); + memoryManager.limit(100); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + + TestMemoryConsumer small = new TestMemoryConsumer(manager); + TestMemoryConsumer large = new TestMemoryConsumer(manager); + small.use(20); + large.use(60); + + String breakdown = manager.getMemoryConsumptionBreakdown(); + Assertions.assertTrue(breakdown.contains("grouped by consumer"), breakdown); + // The larger consumer is listed before the smaller one so the likely culprit surfaces first. + int largeIdx = breakdown.indexOf(large.toString()); + int smallIdx = breakdown.indexOf(small.toString()); + Assertions.assertTrue(largeIdx >= 0 && smallIdx >= 0, breakdown); + Assertions.assertTrue(largeIdx < smallIdx, breakdown); + + small.free(20); + large.free(60); + Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory()); + } + + @Test + public void memoryConsumptionBreakdownCapsConsumersAndSummarizesTheRest() { + final SparkConf conf = new SparkConf() + .set(package$.MODULE$.MEMORY_OOM_ERROR_CONSUMER_BREAKDOWN_LIMIT(), 2); + final TestMemoryManager memoryManager = new TestMemoryManager(conf); + memoryManager.limit(1000); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + + TestMemoryConsumer c1 = new TestMemoryConsumer(manager); + TestMemoryConsumer c2 = new TestMemoryConsumer(manager); + TestMemoryConsumer c3 = new TestMemoryConsumer(manager); + TestMemoryConsumer c4 = new TestMemoryConsumer(manager); + c1.use(40); + c2.use(30); + c3.use(20); + c4.use(10); + + String breakdown = manager.getMemoryConsumptionBreakdown(); + // Only the two largest consumers are listed individually. + Assertions.assertTrue(breakdown.contains(c1.toString()), breakdown); + Assertions.assertTrue(breakdown.contains(c2.toString()), breakdown); + Assertions.assertFalse(breakdown.contains(c3.toString()), breakdown); + Assertions.assertFalse(breakdown.contains(c4.toString()), breakdown); + // The remaining two are collapsed into a summary line covering their combined 30 bytes. + Assertions.assertTrue(breakdown.contains("(2 more consumers): 30.0 B"), breakdown); + + c1.free(40); + c2.free(30); + c3.free(20); + c4.free(10); + Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory()); + } + + @Test + public void memoryConsumptionBreakdownIsOmittedWhenLimitIsZero() { + final SparkConf conf = new SparkConf() + .set(package$.MODULE$.MEMORY_OOM_ERROR_CONSUMER_BREAKDOWN_LIMIT(), 0); + final TestMemoryManager memoryManager = new TestMemoryManager(conf); + memoryManager.limit(100); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + final TestMemoryConsumer c = new TestMemoryConsumer(manager); + c.use(50); + + Assertions.assertEquals("", manager.getMemoryConsumptionBreakdown()); + + c.free(50); + Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory()); + } + + @Test + public void memoryConsumptionBreakdownReportsMemoryNotAttributedToAConsumer() { + final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf()); + memoryManager.limit(4096); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + // Acquire task memory that is not tracked by any registered consumer's getUsed(). The + // NonSpillingConsumer intentionally does not increment its `used` counter, so from the + // manager's perspective this memory is used by the task but unattributed. + final NonSpillingConsumer c = new NonSpillingConsumer(manager); + manager.acquireExecutionMemory(2048, c); + + String breakdown = manager.getMemoryConsumptionBreakdown(); + Assertions.assertTrue( + breakdown.contains("(not attributed to a specific consumer)"), breakdown); + // The consumer never bumped its own counter, so it must not appear as an attributed line. + Assertions.assertFalse(breakdown.contains(c.toString() + ": "), breakdown); + + manager.releaseExecutionMemory(2048, c); + Assertions.assertEquals("", manager.getMemoryConsumptionBreakdown()); + } + + @Test + public void outOfMemoryErrorCarriesConsumerBreakdown() { + final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf()); + memoryManager.limit(1024); + final TestAllocator allocator = new TestAllocator(Integer.MAX_VALUE); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0, allocator); + // A non-spilling consumer holds all the memory so the page allocation below cannot be + // satisfied and MemoryConsumer#throwOom fires. + final NonSpillingConsumer hog = new NonSpillingConsumer(manager); + hog.use(1024); + final PageAllocatingConsumer requestingConsumer = new PageAllocatingConsumer(manager, 4096); + + SparkOutOfMemoryError e = Assertions.assertThrows( + SparkOutOfMemoryError.class, () -> requestingConsumer.allocate(4096)); + Assertions.assertEquals("UNABLE_TO_ACQUIRE_MEMORY", e.getCondition()); + String breakdown = e.getMessageParameters().get("consumerBreakdown"); + Assertions.assertNotNull(breakdown); + // The breakdown captured at failure time names the consumer that was hogging memory. + Assertions.assertTrue(breakdown.contains(hog.toString()), breakdown); + Assertions.assertTrue(e.getMessage().contains("grouped by consumer"), e.getMessage()); + + hog.free(1024); + Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory()); + } + + @Test + public void logMemoryUsageAndGetBreakdownReturnsSameBreakdownAsRenderer() { + final TestMemoryManager memoryManager = new TestMemoryManager(new SparkConf()); + memoryManager.limit(100); + final TaskMemoryManager manager = new TaskMemoryManager(memoryManager, 0); + TestMemoryConsumer c = new TestMemoryConsumer(manager); + c.use(50); + + // The combined OOM-path method logs and returns the breakdown from a single snapshot; on a + // quiescent manager the returned breakdown matches a standalone render of the same state. + String combined = manager.logMemoryUsageAndGetBreakdown(); + Assertions.assertEquals(manager.getMemoryConsumptionBreakdown(), combined); + Assertions.assertTrue(combined.contains(c.toString()), combined); + + c.free(50); + Assertions.assertEquals(0, manager.cleanUpAllAllocatedMemory()); + } + @Test public void offHeapConfigurationBackwardsCompatibility() { // Tests backwards-compatibility with the old `spark.unsafe.offHeap` configuration, which diff --git a/core/src/test/java/org/apache/spark/security/AnotherFakeCredentialProvider.java b/core/src/test/java/org/apache/spark/security/AnotherFakeCredentialProvider.java index b4f3ae40700b0..ccc987e7588c5 100644 --- a/core/src/test/java/org/apache/spark/security/AnotherFakeCredentialProvider.java +++ b/core/src/test/java/org/apache/spark/security/AnotherFakeCredentialProvider.java @@ -31,6 +31,12 @@ public class AnotherFakeCredentialProvider implements CredentialProvider { /** Sentinel URI host that triggers a CredentialResolutionException. */ public static final String ERROR_HOST = "error.example.com"; + /** + * When set to true, {@link #additionalSparkProperties()} throws a RuntimeException. + * Used to test exception isolation in UserCredentialManager. + */ + public static volatile boolean throwOnProperties = false; + private Map<String, String> initConf; @Override @@ -58,4 +64,12 @@ public ServiceCredential resolve(UserContext user, URI target) public Map<String, String> getInitConf() { return initConf; } + + @Override + public Map<String, String> additionalSparkProperties() { + if (throwOnProperties) { + throw new RuntimeException("Simulated failure in additionalSparkProperties"); + } + return Map.of(); + } } diff --git a/core/src/test/java/org/apache/spark/security/CredentialProviderLoaderSuite.java b/core/src/test/java/org/apache/spark/security/CredentialProviderLoaderSuite.java index 54c64db893338..b325fba1ff5d8 100644 --- a/core/src/test/java/org/apache/spark/security/CredentialProviderLoaderSuite.java +++ b/core/src/test/java/org/apache/spark/security/CredentialProviderLoaderSuite.java @@ -25,6 +25,8 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -43,9 +45,12 @@ */ public class CredentialProviderLoaderSuite { + private CredentialProviderLoader loader; + @BeforeEach public void setUp() { - CredentialProviderLoader.resetForTesting(); + loader = new CredentialProviderLoader(); + loader.resetForTesting(); } @Test @@ -53,7 +58,7 @@ public void testServiceLoaderDiscoversFakeProviders() { // The "fake" scheme is supported only by FakeCredentialProvider (single candidate). // If discovery works, providerFor should find it. Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent(), "ServiceLoader should discover FakeCredentialProvider"); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -62,7 +67,7 @@ public void testServiceLoaderDiscoversFakeProviders() { public void testSingleCandidateSchemeResolvesWithNoConf() { // "fake" is supported only by FakeCredentialProvider Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -72,7 +77,7 @@ public void testSharedSchemeWithNoConfThrowsAmbiguity() { // "shared" is supported by both FakeCredentialProvider and AnotherFakeCredentialProvider Map<String, String> conf = Map.of(); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("shared", conf)); + () -> loader.providerFor("shared", conf)); assertTrue(e.getMessage().contains("Multiple credential providers"), "Should mention multiple providers: " + e.getMessage()); assertTrue(e.getMessage().contains("shared"), @@ -92,7 +97,7 @@ public void testEmptyStringConfTreatedAsUnsetThrowsAmbiguity() { Map<String, String> conf = new HashMap<>(); conf.put("spark.security.oidc.provider.shared", ""); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("shared", conf)); + () -> loader.providerFor("shared", conf)); assertTrue(e.getMessage().contains("Multiple credential providers"), "Empty conf value should behave as unset: " + e.getMessage()); } @@ -102,7 +107,7 @@ public void testSharedSchemeWithExplicitConfSelectsFake() { Map<String, String> conf = Map.of( "spark.security.oidc.provider.shared", FakeCredentialProvider.class.getName()); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("shared", conf); + Optional<CredentialProvider> result = loader.providerFor("shared", conf); assertTrue(result.isPresent()); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -112,7 +117,7 @@ public void testSharedSchemeWithExplicitConfSelectsAnother() { Map<String, String> conf = Map.of( "spark.security.oidc.provider.shared", AnotherFakeCredentialProvider.class.getName()); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("shared", conf); + Optional<CredentialProvider> result = loader.providerFor("shared", conf); assertTrue(result.isPresent()); assertInstanceOf(AnotherFakeCredentialProvider.class, result.get()); } @@ -123,7 +128,7 @@ public void testConfNamingUnknownClassThrowsClearError() { "spark.security.oidc.provider.fake", "com.example.NonExistentProvider"); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("fake", conf)); + () -> loader.providerFor("fake", conf)); assertTrue(e.getMessage().contains("com.example.NonExistentProvider"), "Should mention the configured class: " + e.getMessage()); assertTrue(e.getMessage().contains("fake"), @@ -139,7 +144,7 @@ public void testConfNamingNonSupportingClassThrowsClearError() { "spark.security.oidc.provider.fake", AnotherFakeCredentialProvider.class.getName()); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("fake", conf)); + () -> loader.providerFor("fake", conf)); assertTrue(e.getMessage().contains(AnotherFakeCredentialProvider.class.getName()), "Should mention the configured class: " + e.getMessage()); assertTrue(e.getMessage().contains("fake"), @@ -154,7 +159,7 @@ public void testSingleCandidateWithCorrectExplicitConfSelectsIt() { Map<String, String> conf = Map.of( "spark.security.oidc.provider.fake", FakeCredentialProvider.class.getName()); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -167,7 +172,7 @@ public void testSingleCandidateWithWrongExplicitConfThrowsClearError() { "spark.security.oidc.provider.fake", "org.apache.spark.security.SomeOtherProvider"); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("fake", conf)); + () -> loader.providerFor("fake", conf)); assertTrue(e.getMessage().contains("fake"), "Should mention the scheme: " + e.getMessage()); assertTrue(e.getMessage().contains("org.apache.spark.security.SomeOtherProvider"), @@ -180,7 +185,7 @@ public void testSingleCandidateWithWrongExplicitConfThrowsClearError() { public void testUnknownSchemeReturnsEmpty() { Map<String, String> conf = Map.of(); Optional<CredentialProvider> result = - CredentialProviderLoader.providerFor("nonexistent", conf); + loader.providerFor("nonexistent", conf); assertFalse(result.isPresent(), "Unknown scheme should return empty"); } @@ -190,7 +195,7 @@ public void testInitConfIsInvokedOnSelectedProvider() { conf.put("spark.security.oidc.endpoint", "https://sts.example.com"); conf.put("spark.security.oidc.roleArn", "arn:aws:iam::123456:role/test"); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); FakeCredentialProvider fake = (FakeCredentialProvider) result.get(); assertNotNull(fake.getInitConf(), "init() should have been called"); @@ -210,8 +215,8 @@ public void testProviderInitializedExactlyOnce() { Map<String, String> conf2 = new HashMap<>(); conf2.put("spark.security.oidc.tag", "second-call"); - Optional<CredentialProvider> result1 = CredentialProviderLoader.providerFor("fake", conf1); - Optional<CredentialProvider> result2 = CredentialProviderLoader.providerFor("fake", conf2); + Optional<CredentialProvider> result1 = loader.providerFor("fake", conf1); + Optional<CredentialProvider> result2 = loader.providerFor("fake", conf2); assertTrue(result1.isPresent()); assertTrue(result2.isPresent()); @@ -242,11 +247,11 @@ public ServiceCredential resolve(UserContext user, URI target) { return null; } }; - CredentialProviderLoader.setProvidersForTesting( + loader.setProvidersForTesting( List.of(nullSchemesProvider)); IllegalStateException e = assertThrows(IllegalStateException.class, - () -> CredentialProviderLoader.providerFor("anything", Map.of())); + () -> loader.providerFor("anything", Map.of())); assertTrue(e.getMessage().contains("returned null from supportedSchemes()"), "Should have a clear null-schemes message: " + e.getMessage()); } @@ -254,7 +259,7 @@ public ServiceCredential resolve(UserContext user, URI target) { @Test public void testResolveReturnsExpectedServiceCredential() throws Exception { Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); UserContext user = new UserContext( @@ -270,7 +275,7 @@ public void testResolveReturnsExpectedServiceCredential() throws Exception { @Test public void testResolveSentinelThrowsCredentialResolutionException() { Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); UserContext user = new UserContext( @@ -287,7 +292,7 @@ public void testResolveSentinelThrowsCredentialResolutionException() { public void testSchemeNormalizationIsCaseInsensitive() { // "FAKE" should resolve the same as "fake" Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("FAKE", conf); + Optional<CredentialProvider> result = loader.providerFor("FAKE", conf); assertTrue(result.isPresent()); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -298,7 +303,7 @@ public void testExplicitSelectionWithUppercaseSchemeNormalizesConfKey() { Map<String, String> conf = Map.of( "spark.security.oidc.provider.shared", FakeCredentialProvider.class.getName()); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("SHARED", conf); + Optional<CredentialProvider> result = loader.providerFor("SHARED", conf); assertTrue(result.isPresent()); assertInstanceOf(FakeCredentialProvider.class, result.get()); } @@ -306,7 +311,7 @@ public void testExplicitSelectionWithUppercaseSchemeNormalizesConfKey() { @Test public void testNullSchemeThrowsNPE() { NullPointerException e = assertThrows(NullPointerException.class, - () -> CredentialProviderLoader.providerFor(null, Map.of())); + () -> loader.providerFor(null, Map.of())); assertTrue(e.getMessage().contains("scheme must not be null"), "Should have a clear message: " + e.getMessage()); } @@ -314,7 +319,7 @@ public void testNullSchemeThrowsNPE() { @Test public void testNullConfThrowsNPE() { NullPointerException e = assertThrows(NullPointerException.class, - () -> CredentialProviderLoader.providerFor("fake", null)); + () -> loader.providerFor("fake", null)); assertTrue(e.getMessage().contains("conf must not be null"), "Should have a clear message: " + e.getMessage()); } @@ -322,7 +327,7 @@ public void testNullConfThrowsNPE() { @Test public void testSuggestedTtlDefaultValue() { Map<String, String> conf = Map.of(); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); assertEquals(Duration.ofMinutes(15), result.get().suggestedTtl()); } @@ -338,7 +343,7 @@ public void testInitConfScopedToOidcKeysOnly() { conf.put("spark.authenticate.secret", "TOPSECRET"); conf.put("spark.ssl.keyPassword", "keypass"); - Optional<CredentialProvider> result = CredentialProviderLoader.providerFor("fake", conf); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); assertTrue(result.isPresent()); FakeCredentialProvider fake = (FakeCredentialProvider) result.get(); Map<String, String> initConf = fake.getInitConf(); @@ -362,10 +367,138 @@ public void testInitConfScopedToOidcKeysOnly() { @Test public void testEmptySchemeThrowsIllegalArgument() { IllegalArgumentException e = assertThrows(IllegalArgumentException.class, - () -> CredentialProviderLoader.providerFor("", Map.of())); + () -> loader.providerFor("", Map.of())); assertEquals("scheme must not be empty", e.getMessage()); } + @Test + public void testCloseAllClosesInitializedProviders() throws Exception { + // Initialize a provider by calling providerFor + Map<String, String> conf = Map.of(); + Optional<CredentialProvider> result = loader.providerFor("fake", conf); + assertTrue(result.isPresent()); + FakeCredentialProvider fake = (FakeCredentialProvider) result.get(); + assertEquals(0, fake.getCloseCount(), "close() not yet called"); + + // Call closeAll + loader.closeAll(); + + assertEquals(1, fake.getCloseCount(), "close() should be called exactly once"); + } + + @Test + public void testCloseAllSuppressesExceptionsAndClosesAll() throws Exception { + // Two providers: first throws on close, second should still be closed + CredentialProvider throwingProvider = new CredentialProvider() { + @Override + public void init(Map<String, String> conf) {} + + @Override + public Set<String> supportedSchemes() { + return Set.of("throwing"); + } + + @Override + public ServiceCredential resolve(UserContext user, URI target) { + return new ServiceCredential(Map.of(), Instant.now().plusSeconds(60)); + } + + @Override + public void close() throws Exception { + throw new RuntimeException("Simulated close failure"); + } + }; + + FakeCredentialProvider fakeProvider = new FakeCredentialProvider(); + + loader.setProvidersForTesting( + List.of(throwingProvider, fakeProvider)); + + // Initialize both by selecting them + Map<String, String> conf = new HashMap<>(); + conf.put("spark.security.oidc.provider.throwing", throwingProvider.getClass().getName()); + loader.providerFor("throwing", conf); + loader.providerFor("fake", conf); + + // closeAll should throw (from throwingProvider) but still close fakeProvider + Exception e = assertThrows(Exception.class, + () -> loader.closeAll()); + assertTrue(e.getMessage().contains("Simulated close failure")); + assertEquals(1, fakeProvider.getCloseCount(), + "Second provider should still be closed even when first throws"); + } + + @Test + public void testCloseAllWithNoInitializedProvidersIsNoOp() throws Exception { + // No providers initialized; closeAll should not throw + loader.closeAll(); + // If we reach here, no exception was thrown: success + } + + @Test + public void testProviderCannotBeReinitializedAfterCloseAll() throws Exception { + Map<String, String> conf = Map.of(); + CredentialProvider first = loader.providerFor("fake", conf).orElseThrow(); + + loader.closeAll(); + + IllegalStateException e = assertThrows(IllegalStateException.class, + () -> loader.providerFor("fake", conf)); + assertEquals("Credential providers have already been closed", e.getMessage()); + assertThrows(IllegalStateException.class, + () -> loader.providerFor("nonexistent", conf)); + IllegalStateException discoverError = assertThrows( + IllegalStateException.class, loader::discoverAllSchemes); + assertEquals("Credential providers have already been closed", discoverError.getMessage()); + + CredentialProviderLoader nextLoader = new CredentialProviderLoader(); + CredentialProvider second = nextLoader.providerFor("fake", conf).orElseThrow(); + assertInstanceOf(FakeCredentialProvider.class, second); + assertTrue(first != second, "A new loader should discover a fresh provider instance"); + } + + @Test + public void testStaleLifecycleCannotUseNextLifecycleProviders() throws Exception { + Map<String, String> conf = Map.of(); + CredentialProvider retiredProvider = loader.providerFor("fake", conf).orElseThrow(); + CountDownLatch releaseStaleCaller = new CountDownLatch(1); + AtomicReference<Throwable> staleFailure = new AtomicReference<>(); + Thread staleCaller = new Thread(() -> { + try { + releaseStaleCaller.await(); + loader.providerFor("fake", conf); + } catch (Throwable t) { + staleFailure.set(t); + } + }); + staleCaller.start(); + + try { + loader.closeAll(); + + CredentialProviderLoader nextLoader = new CredentialProviderLoader(); + try { + CredentialProvider nextProvider = nextLoader.providerFor("fake", conf).orElseThrow(); + assertTrue(retiredProvider != nextProvider, + "The next lifecycle should discover a fresh provider instance"); + + releaseStaleCaller.countDown(); + staleCaller.join(10000); + + assertFalse(staleCaller.isAlive(), "The stale caller should have completed"); + assertInstanceOf(IllegalStateException.class, staleFailure.get()); + assertEquals("Credential providers have already been closed", + staleFailure.get().getMessage()); + } finally { + nextLoader.closeAll(); + } + } finally { + releaseStaleCaller.countDown(); + staleCaller.interrupt(); + staleCaller.join(10000); + } + } + @Test public void testInitRetryAfterFailure() { // A provider whose init() throws on the first call then succeeds on the second. @@ -400,17 +533,17 @@ public String toString() { } }; - CredentialProviderLoader.setProvidersForTesting(List.of(failOnceThenSucceed)); + loader.setProvidersForTesting(List.of(failOnceThenSucceed)); // First call: init() throws, providerFor should propagate Map<String, String> conf = Map.of(); RuntimeException e = assertThrows(RuntimeException.class, - () -> CredentialProviderLoader.providerFor("retryscheme", conf)); + () -> loader.providerFor("retryscheme", conf)); assertTrue(e.getMessage().contains("Simulated transient init failure")); // Second call: init() should be retried and succeed Optional<CredentialProvider> result = - CredentialProviderLoader.providerFor("retryscheme", conf); + loader.providerFor("retryscheme", conf); assertTrue(result.isPresent(), "Second providerFor should succeed after init retry"); // Verify init was called exactly twice (proving the retry) diff --git a/core/src/test/java/org/apache/spark/security/FakeCredentialProvider.java b/core/src/test/java/org/apache/spark/security/FakeCredentialProvider.java index 9eb9be01446aa..b9108980a4738 100644 --- a/core/src/test/java/org/apache/spark/security/FakeCredentialProvider.java +++ b/core/src/test/java/org/apache/spark/security/FakeCredentialProvider.java @@ -21,6 +21,7 @@ import java.time.Instant; import java.util.Map; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; /** * A fake credential provider for testing. Supports schemes "fake" and "shared". @@ -30,13 +31,14 @@ public class FakeCredentialProvider implements CredentialProvider { /** Sentinel URI host that triggers a CredentialResolutionException. */ public static final String ERROR_HOST = "error.example.com"; - private Map<String, String> initConf; - private int initCount; + private volatile Map<String, String> initConf; + private final AtomicInteger initCount = new AtomicInteger(); + private final AtomicInteger closeCount = new AtomicInteger(); @Override public void init(Map<String, String> conf) { this.initConf = conf; - this.initCount++; + this.initCount.incrementAndGet(); } @Override @@ -55,6 +57,11 @@ public ServiceCredential resolve(UserContext user, URI target) return new ServiceCredential(Map.of("provider", "fake"), expiresAt); } + @Override + public void close() { + this.closeCount.incrementAndGet(); + } + /** Returns the configuration map passed to {@link #init(Map)}, or null if not yet called. */ public Map<String, String> getInitConf() { return initConf; @@ -62,6 +69,17 @@ public Map<String, String> getInitConf() { /** Returns the number of times {@link #init(Map)} has been called. */ public int getInitCount() { - return initCount; + return initCount.get(); + } + + /** Returns the number of times {@link #close()} has been called. */ + public int getCloseCount() { + return closeCount.get(); + } + + @Override + public Map<String, String> additionalSparkProperties() { + return Map.of("spark.hadoop.fs.fake.credentials.provider", + "org.apache.spark.security.FakeExecutorCredentialProvider"); } } diff --git a/core/src/test/java/org/apache/spark/security/FileTokenIngestorSuite.java b/core/src/test/java/org/apache/spark/security/FileTokenIngestorSuite.java index f6f8ea2c3636a..9929e23964b32 100644 --- a/core/src/test/java/org/apache/spark/security/FileTokenIngestorSuite.java +++ b/core/src/test/java/org/apache/spark/security/FileTokenIngestorSuite.java @@ -19,6 +19,7 @@ import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.attribute.FileTime; import java.security.KeyPair; import java.security.KeyPairGenerator; import java.util.Comparator; @@ -143,6 +144,25 @@ public void loadDetectsFileRotation() throws Exception { assertEquals(token2, result2.get().getRawToken()); } + @Test + public void loadDetectsFileRotationWhenMtimeIsUnchanged() throws Exception { + Path tokenFile = tempDir.resolve("token"); + String token1 = createUnsignedJwt("user1", "https://issuer.example.com"); + writeToken(tokenFile, token1); + + FileTokenIngestor ingestor = new FileTokenIngestor(tokenFile); + assertEquals("user1", ingestor.load().get().getPrincipal()); + FileTime originalMtime = Files.getLastModifiedTime(tokenFile); + + String token2 = createUnsignedJwt("user2", "https://issuer.example.com"); + writeToken(tokenFile, token2); + Files.setLastModifiedTime(tokenFile, originalMtime); + + Optional<UserContext> result = ingestor.load(); + assertEquals("user2", result.get().getPrincipal()); + assertEquals(token2, result.get().getRawToken()); + } + @Test public void loadReturnsEmptyForMissingFile() { Path tokenFile = tempDir.resolve("nonexistent"); diff --git a/core/src/test/resources/META-INF/services/org.apache.spark.security.HadoopDelegationTokenProvider b/core/src/test/resources/META-INF/services/org.apache.spark.security.HadoopDelegationTokenProvider index 8621d66256e04..add25b0c26fe1 100644 --- a/core/src/test/resources/META-INF/services/org.apache.spark.security.HadoopDelegationTokenProvider +++ b/core/src/test/resources/META-INF/services/org.apache.spark.security.HadoopDelegationTokenProvider @@ -19,4 +19,5 @@ org.apache.spark.deploy.security.ExceptionThrowingDelegationTokenProvider org.apache.spark.deploy.security.TestNonKerberosTokenProvider org.apache.spark.deploy.security.TestDisabledProvider org.apache.spark.deploy.security.TestFailingProvider +org.apache.spark.deploy.security.TestRequirementFailingProvider org.apache.spark.deploy.security.TestNoExpiryProvider diff --git a/core/src/test/scala/org/apache/spark/CheckErrorHelper.scala b/core/src/test/scala/org/apache/spark/CheckErrorHelper.scala index d01600bb439f1..d0fbb8ead436e 100644 --- a/core/src/test/scala/org/apache/spark/CheckErrorHelper.scala +++ b/core/src/test/scala/org/apache/spark/CheckErrorHelper.scala @@ -67,7 +67,10 @@ trait CheckErrorHelper { self: Suite => * Test suites may override this to add or change ignorable parameters per condition. */ protected def checkErrorIgnorableParameters: Map[String, Set[String]] = Map( - "TABLE_OR_VIEW_NOT_FOUND" -> Set("searchPath") + "TABLE_OR_VIEW_NOT_FOUND" -> Set("searchPath"), + // The per-consumer memory breakdown is a best-effort diagnostic whose content (live memory + // sizes and consumer identities) is inherently non-deterministic, so tests need not pin it. + "UNABLE_TO_ACQUIRE_MEMORY" -> Set("consumerBreakdown") ) /** diff --git a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala index 58512a2282ac2..af06e8f29f3b7 100644 --- a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala +++ b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala @@ -18,18 +18,22 @@ package org.apache.spark import java.io.File +import java.net.URI +import java.util.Properties import scala.reflect.ClassTag -import org.apache.hadoop.fs.Path +import org.apache.hadoop.fs.{FileAlreadyExistsException, LocalFileSystem, Path, RawLocalFileSystem} import org.apache.spark.internal.config.CACHE_CHECKPOINT_PREFERRED_LOCS_EXPIRE_TIME import org.apache.spark.internal.config.UI._ import org.apache.spark.io.CompressionCodec +import org.apache.spark.memory.TaskMemoryManager import org.apache.spark.rdd._ import org.apache.spark.shuffle.FetchFailedException import org.apache.spark.storage.{BlockId, StorageLevel, TestBlockId} import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.util.SerializableConfiguration import org.apache.spark.util.Utils trait RDDCheckpointTester { self: SparkFunSuite => @@ -669,6 +673,34 @@ class CheckpointStorageSuite extends SparkFunSuite with LocalSparkContext { } } + test("SPARK-58750: checkpointing tolerates FileAlreadyExistsException on part file rename") { + withTempDir { checkpointDir => + val conf = new SparkConf().set(UI_ENABLED.key, "false") + sc = new SparkContext("local", "test", conf) + sc.hadoopConfiguration.set( + "fs.faee.impl", classOf[FileAlreadyExistsRenameFileSystem].getName) + val broadcastedConf = SerializableConfiguration.broadcast(sc, sc.hadoopConfiguration) + val outputDir = s"faee://${checkpointDir.getAbsolutePath}" + + def writePartition(taskAttemptId: Long, attemptNumber: Int): Unit = { + val ctx = new TaskContextImpl(0, 0, 0, taskAttemptId, attemptNumber, 1, + new TaskMemoryManager(sc.env.memoryManager, 0L), new Properties, sc.env.metricsSystem) + ReliableCheckpointRDD.writePartitionToCheckpointFile[Int]( + outputDir, broadcastedConf)(ctx, Iterator(1, 2, 3)) + } + + writePartition(taskAttemptId = 0L, attemptNumber = 0) + // A speculative or retried attempt of the same partition finds the part file already + // committed by the first attempt. On filesystems that raise FileAlreadyExistsException + // from rename (S3A, ABFS), this must be treated as success rather than fail the task. + writePartition(taskAttemptId = 1L, attemptNumber = 1) + + val fs = new Path(outputDir).getFileSystem(sc.hadoopConfiguration) + val fileNames = fs.listStatus(new Path(outputDir)).map(_.getPath.getName) + assert(fileNames === Array("part-00000")) + } + } + test("SPARK-48268: checkpoint directory via configuration") { withTempDir { checkpointDir => val conf = new SparkConf() @@ -684,4 +716,88 @@ class CheckpointStorageSuite extends SparkFunSuite with LocalSparkContext { assert(flatMappedRDD.collect() === result) } } + + test("checkpoint() without a checkpoint directory") { + sc = new SparkContext("local", "test", new SparkConf().set(UI_ENABLED.key, "false")) + checkError( + exception = intercept[SparkException](sc.makeRDD(1 to 4).checkpoint()), + condition = "CHECKPOINT_DIRECTORY_NOT_SET", + sqlState = Some("55019")) + } + + test("reading a checkpoint directory with a missing partition file") { + withTempDir { checkpointDir => + sc = new SparkContext("local", "test", new SparkConf().set(UI_ENABLED.key, "false")) + sc.setCheckpointDir(checkpointDir.toString) + val rdd = sc.makeRDD(1 to 20, numSlices = 4) + rdd.checkpoint() + assert(rdd.collect().toSeq === (1 to 20)) + + // Drop a middle partition file so the remaining names are no longer part-00000..part-0000N. + val checkpointPath = new Path(rdd.getCheckpointFile.get) + val fs = checkpointPath.getFileSystem(sc.hadoopConfiguration) + assert(fs.delete(new Path(checkpointPath, "part-00001"), false)) + + val recovered = sc.checkpointFile[Int](checkpointPath.toString) + checkError( + exception = intercept[SparkException](recovered.partitions), + condition = "INVALID_CHECKPOINT_DIRECTORY", + sqlState = Some("58030"), + parameters = Map( + "path" -> checkpointPath.toString, + "expectedFileName" -> "part-00001", + "fileName" -> "part-00002")) + } + } + + test("checkpoint path that cannot be created") { + withTempDir { checkpointDir => + // MkdirsFailingFilesystem refuses to create the per-RDD directory and reports it the way + // HDFS and S3A do, by returning false rather than throwing. + val conf = new SparkConf() + .set("spark.hadoop.fs.file.impl", classOf[MkdirsFailingFilesystem].getName) + .set("spark.hadoop.fs.file.impl.disable.cache", "true") + .set(UI_ENABLED.key, "false") + sc = new SparkContext("local", "test", conf) + sc.setCheckpointDir(checkpointDir.toString) + val rdd = sc.makeRDD(1 to 4) + val rddPath = ReliableRDDCheckpointData.checkpointPath(sc, rdd.id).get + + rdd.checkpoint() + checkError( + exception = intercept[SparkException](rdd.collect()), + condition = "FAILED_CREATE_CHECKPOINT_DIRECTORY", + sqlState = Some("58030"), + parameters = Map("path" -> rddPath.toString)) + } + } +} + +/** + * A filesystem that refuses to create a per-RDD checkpoint directory, returning false from + * `mkdirs` instead of throwing. `LocalFileSystem` throws `FileAlreadyExistsException` when the + * path is taken, while HDFS and S3A can report the failure through the return value, which is + * what `ReliableCheckpointRDD` checks. + */ +class MkdirsFailingFilesystem extends LocalFileSystem { + override def mkdirs(f: Path): Boolean = { + if (f.getName.startsWith("rdd-")) false else super.mkdirs(f) + } +} + +/** + * A local filesystem mimicking how some Hadoop FileSystem implementations report a rename onto + * an existing file: by raising FileAlreadyExistsException (e.g. S3A since HADOOP-16721, ABFS) + * rather than returning false as HDFS does. + */ +class FileAlreadyExistsRenameFileSystem extends RawLocalFileSystem { + override def getUri: URI = URI.create("faee:///") + + override def rename(src: Path, dst: Path): Boolean = { + if (exists(dst)) { + throw new FileAlreadyExistsException( + s"Failed to rename $src to $dst; destination file exists") + } + super.rename(src, dst) + } } diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala index bb2d7d5c4d8e4..0bd0ee103a190 100644 --- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark import scala.collection.mutable.HashSet import scala.util.Random +import org.apache.logging.log4j.Level import org.scalatest.BeforeAndAfter import org.scalatest.concurrent.Eventually._ import org.scalatest.concurrent.PatienceConfiguration @@ -127,6 +128,52 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { assert(rdd.collect().toList.equals(collected)) } + test("cleanup shuffle unregisters a pipelined shuffle from the streaming output tracker") { + // A pipelined shuffle lives ONLY in the driver-only StreamingShuffleOutputTracker -- it is + // never registered with the MapOutputTracker (see DAGScheduler.createShuffleMapStage). So the + // MapOutputTracker.containsShuffle guard is false for it, and doCleanupShuffle must still clean + // it up via its own streaming branch; otherwise the tracker grows without bound across + // micro-batches in Real-Time Mode. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Register a shuffle id in the streaming tracker ONLY, exactly as the scheduler now does for a + // pipelined shuffle (nothing is put in the MapOutputTracker). + val shuffleId = 1000 + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") + + cleaner.doCleanupShuffle(shuffleId, blocking = true) + + // The streaming tracker entry is gone -- cleanup fired even though the shuffle was never in the + // MapOutputTracker, proving the streaming cleanup branch is independent of MapOutputTracker. + assert(!streamingTracker.containsShuffle(shuffleId)) + } + + test("cleanup shuffle leaves a regular shuffle's streaming tracker untouched and quiet") { + // A regular (non-pipelined) shuffle is registered only with the MapOutputTracker, never the + // streaming tracker. doCleanupShuffle must take the MapOutputTracker branch and never touch the + // streaming tracker -- in particular it must not emit the streaming tracker's "attempting to + // unregister a shuffle that hasn't been registered" warning. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + + val (rdd, shuffleDeps) = newRDDWithShuffleDependencies() + rdd.collect() + shuffleDeps.foreach(s => assert(!streamingTracker.containsShuffle(s.shuffleId))) + + val logAppender = new LogAppender("unregister streaming shuffle") + logAppender.setThreshold(Level.WARN) + withLogAppender(logAppender, level = Some(Level.WARN)) { + shuffleDeps.foreach(s => cleaner.doCleanupShuffle(s.shuffleId, blocking = true)) + } + assert(!logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("hasn't been registered"))) + } + test("cleanup broadcast") { val broadcast = newBroadcast() val tester = new CleanerTester(sc, broadcastIds = Seq(broadcast.id)) @@ -174,6 +221,49 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase { postGCTester.assertCleanup() } + test("automatically cleanup pipelined shuffle from the streaming output tracker") { + // The load-bearing leak guarantee for a pipelined (streaming) shuffle: it is registered ONLY + // with the driver-only StreamingShuffleOutputTracker (never the MapOutputTracker), and its only + // cleanup channel is the ContextCleaner weak-reference path fired when the + // PipelinedShuffleDependency is garbage-collected (registerShuffleForCleanup in the + // ShuffleDependency constructor). If that GC path did not reach the streaming tracker, + // the tracker would grow without bound across Real-Time Mode micro-batches. The sibling + // doCleanupShuffle tests call cleanup directly; this one proves the end-to-end + // GC -> ContextCleaner -> streaming-tracker unregister link. + val streamingTracker = sc.env.streamingShuffleOutputTracker.get + .asInstanceOf[StreamingShuffleOutputTrackerMaster] + val mapOutputTracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + + // Build a real PipelinedShuffleDependency (registers itself for cleanup on construction) and + // register its shuffleId in the streaming tracker exactly as DAGScheduler.createShuffleMapStage + // does for a pipelined shuffle (nothing is put in the MapOutputTracker). + var pipelinedDep: PipelinedShuffleDependency[Int, Int, Int] = + new PipelinedShuffleDependency(newPairRDD(), new HashPartitioner(2)) + val shuffleId = pipelinedDep.shuffleId + streamingTracker.registerShuffle(shuffleId, numMaps = 2, numReduces = 2, jobId = 0) + assert(streamingTracker.containsShuffle(shuffleId)) + assert(!mapOutputTracker.containsShuffle(shuffleId), + "a pipelined shuffle must not be registered with the MapOutputTracker") + + // A strong reference to the dependency must prevent GC-triggered cleanup. + runGC() + intercept[Exception] { + eventually(timeout(1.second), interval(100.milliseconds)) { + assert(!streamingTracker.containsShuffle(shuffleId), + "cleanup must NOT fire while the dependency is strongly referenced") + } + } + + // Dereference the dependency; GC must then drive ContextCleaner to unregister it from the + // streaming tracker (and only there -- it was never in the MapOutputTracker). + pipelinedDep = null + runGC() + eventually(timeout(10.seconds), interval(100.milliseconds)) { + assert(!streamingTracker.containsShuffle(shuffleId), + "GC of the PipelinedShuffleDependency must unregister it from the streaming tracker") + } + } + test("automatically cleanup broadcast") { var broadcast = newBroadcast() diff --git a/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala b/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala index 0872d39df5356..32289e74ccacd 100644 --- a/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/ExecutorAllocationManagerSuite.scala @@ -21,7 +21,8 @@ import java.util.concurrent.TimeUnit import scala.collection.mutable -import org.mockito.ArgumentMatchers.any +import org.mockito.ArgumentCaptor +import org.mockito.ArgumentMatchers.{any, anyBoolean, eq => mockitoEq} import org.mockito.Mockito._ import org.scalatest.PrivateMethodTester @@ -1070,6 +1071,53 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite { assert(manager.executorAllocationManagerSource.exitedUnexpectedly.getCount() === 1) } + test("SPARK-58879: dynamic allocation accounts only for accepted idle decommissions") { + val conf = createConf(0, 5, 0, decommissioningEnabled = true) + .set(config.DYN_ALLOCATION_TESTING, false) + when(client.requestTotalExecutors(any(), any(), any())).thenReturn(true) + when(client.decommissionExecutorsIfIdle(any(), anyBoolean())) + .thenReturn(Seq("executor-2")) + val manager = createManager(conf) + onExecutorAddedDefaultProfile(manager, "executor-1") + onExecutorAddedDefaultProfile(manager, "executor-2") + + assert(removeExecutorsDefaultProfile(manager, Seq("executor-1", "executor-2")) === + Seq("executor-2")) + assert(executorsDecommissioning(manager) === Set("executor-2")) + assert(executorsPendingToRemove(manager).isEmpty) + + val requests = ArgumentCaptor.forClass( + classOf[Array[(String, ExecutorDecommissionInfo)]]) + verify(client).decommissionExecutorsIfIdle(requests.capture(), mockitoEq(false)) + assert(requests.getValue.toSeq === Seq( + "executor-1" -> ExecutorDecommissionInfo("spark scale down"), + "executor-2" -> ExecutorDecommissionInfo("spark scale down"))) + verify(client, never()).decommissionExecutors(any(), anyBoolean(), anyBoolean()) + verify(client, never()).killExecutors(any(), anyBoolean(), anyBoolean(), anyBoolean()) + + // A stale timeout can be rejected again without accounting the executor as removed. + when(client.decommissionExecutorsIfIdle(any(), anyBoolean())) + .thenReturn(Seq.empty[String]) + assert(removeExecutorsDefaultProfile(manager, Seq("executor-1")).isEmpty) + assert(executorsDecommissioning(manager) === Set("executor-2")) + + // The same executor remains eligible for a later request once it is actually idle. + when(client.decommissionExecutorsIfIdle(any(), anyBoolean())) + .thenReturn(Seq("executor-1")) + assert(removeExecutorsDefaultProfile(manager, Seq("executor-1")) === Seq("executor-1")) + assert(executorsDecommissioning(manager) === Set("executor-1", "executor-2")) + } + + test("SPARK-58879: unsupported idle decommission does not force removal") { + val unsupportedClient = mock(classOf[ExecutorAllocationClient], CALLS_REAL_METHODS) + assert(unsupportedClient.decommissionExecutorsIfIdle( + Array("executor-1" -> ExecutorDecommissionInfo("spark scale down")), + adjustTargetNumExecutors = false).isEmpty) + verify(unsupportedClient, never()).decommissionExecutors(any(), anyBoolean(), anyBoolean()) + verify(unsupportedClient, never()) + .killExecutors(any(), anyBoolean(), anyBoolean(), anyBoolean()) + } + test("remove multiple executors") { val manager = createManager(createConf(5, 10, 5)) (1 to 10).map(_.toString).foreach { id => onExecutorAddedDefaultProfile(manager, id) } @@ -1348,6 +1396,147 @@ class ExecutorAllocationManagerSuite extends SparkFunSuite { assert(numExecutorsTargetForDefaultProfileId(manager) === 20) // limit reached } + test("SPARK-58828: suspend and resume executor allocation") { + val clock = new ManualClock(2020L) + val manager = createManager(createConf(0, 20, 0), clock = clock) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000))) + + // Ramp up the target normally first + onSchedulerBacklogged(manager) + clock.advance(schedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + clock.advance(sustainedSchedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1 + 2) + + // Suspending lowers the target to zero and stops the ramp-up + manager.suspend() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + clock.advance(sustainedSchedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + clock.advance(sustainedSchedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + + // Resuming ramps the target up again for the still-pending tasks + manager.resume() + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + clock.advance(sustainedSchedulerBacklogTimeout * 1000) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 1 + 2) + } + + test("SPARK-58828: suspend keeps a zero target across reset()") { + val manager = createManager(createConf(1, 10, 3)) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000))) + manager.suspend() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + manager.reset() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + manager.resume() + manager.reset() + assert(numExecutorsTargetForDefaultProfileId(manager) === 3) + } + + test("SPARK-58828: resume before the first stage restores the initial target") { + val manager = createManager(createConf(1, 10, 3)) + manager.suspend() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + manager.resume() + assert(numExecutorsTargetForDefaultProfileId(manager) === 3) + } + + test("SPARK-58828: a new resource profile submitted while suspended gets a zero target") { + val manager = createManager(createConf(1, 10, 1)) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000, rp = defaultProfile))) + manager.suspend() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + val rp1 = new ResourceProfileBuilder() + val execReqs = new ExecutorResourceRequests().cores(4).resource("gpu", 4) + val taskReqs = new TaskResourceRequests().cpus(1).resource("gpu", 1) + rp1.require(execReqs).require(taskReqs) + val rprof1 = rp1.build() + rpManager.addResourceProfile(rprof1) + post(SparkListenerStageSubmitted(createStageInfo(1, 1000, rp = rprof1))) + assert(numExecutorsTarget(manager, rprof1.id) === 0) + // After resume, the new profile is restored to the minimum and ramps up like any other + manager.resume() + assert(numExecutorsTargetForDefaultProfileId(manager) === 1) + assert(numExecutorsTarget(manager, rprof1.id) === 1) + schedule(manager) + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + assert(numExecutorsTarget(manager, rprof1.id) === 2) + } + + test("SPARK-58828: resume restores at least the minimum target when idle") { + val manager = createManager(createConf(2, 10, 2)) + post(SparkListenerStageSubmitted(createStageInfo(0, 8))) + post(SparkListenerStageCompleted(createStageInfo(0, 8))) + manager.suspend() + assert(numExecutorsTargetForDefaultProfileId(manager) === 0) + manager.resume() + assert(numExecutorsTargetForDefaultProfileId(manager) === 2) + } + + test("SPARK-58828: a rejected suspend push is retried from the schedule loop") { + val manager = createManager(createConf(1, 10, 3)) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000))) + when(client.requestTotalExecutors(any(), any(), any())).thenReturn(false) + clearInvocations(client) + assert(!manager.suspend()) + verify(client).requestTotalExecutors(mockitoEq(Map(defaultProfile.id -> 0)), any(), any()) + // Rejected: re-pushed on every tick until acknowledged, then no more + when(client.requestTotalExecutors(any(), any(), any())).thenReturn(true) + clearInvocations(client) + schedule(manager) + verify(client).requestTotalExecutors(mockitoEq(Map(defaultProfile.id -> 0)), any(), any()) + clearInvocations(client) + schedule(manager) + verify(client, never()).requestTotalExecutors(any(), any(), any()) + } + + test("SPARK-58828: deferred target pushes back off exponentially") { + val manager = createManager(createConf(1, 10, 3)) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000))) + when(client.requestTotalExecutors(any(), any(), any())).thenReturn(false) + clearInvocations(client) + assert(!manager.suspend()) // first push fails; the next retry is immediate + verify(client).requestTotalExecutors(any(), any(), any()) + clearInvocations(client) + schedule(manager) // immediate retry fails; the next retry waits a tick + verify(client).requestTotalExecutors(any(), any(), any()) + clearInvocations(client) + schedule(manager) // backing off + verify(client, never()).requestTotalExecutors(any(), any(), any()) + schedule(manager) // retried after the backoff + verify(client).requestTotalExecutors(any(), any(), any()) + } + + test("SPARK-58828: reset while suspended defers the zero-target push to the schedule loop") { + // Keep DYN_ALLOCATION_TESTING on: turning it off would start the real polling thread + // (`start()` ignores TEST_DYNAMIC_ALLOCATION_SCHEDULE_ENABLED when not testing) and its + // background ticks would race with the assertions below. + val manager = createManager(createConf(1, 10, 3)) + when(client.requestTotalExecutors(any(), any(), any())).thenReturn(true) + post(SparkListenerStageSubmitted(createStageInfo(0, 1000))) + manager.suspend() + clearInvocations(client) + // reset() itself must not talk to the cluster manager: it may run inside a cluster + // manager RPC handler (e.g. YARN's RegisterClusterManager) where a synchronous request + // would self-deadlock. + manager.reset() + verify(client, never()).requestTotalExecutors(any(), any(), any()) + // The zero target is pushed from the schedule loop instead, exactly once + schedule(manager) + verify(client).requestTotalExecutors(mockitoEq(Map(defaultProfile.id -> 0)), any(), any()) + clearInvocations(client) + schedule(manager) + verify(client, never()).requestTotalExecutors(any(), any(), any()) + } + test("mock polling loop remove behavior") { val clock = new ManualClock(2020L) val manager = createManager(createConf(1, 20, 1), clock = clock) diff --git a/core/src/test/scala/org/apache/spark/PartitioningSuite.scala b/core/src/test/scala/org/apache/spark/PartitioningSuite.scala index d0423e267baf3..1dca5c6c11d51 100644 --- a/core/src/test/scala/org/apache/spark/PartitioningSuite.scala +++ b/core/src/test/scala/org/apache/spark/PartitioningSuite.scala @@ -215,24 +215,34 @@ class PartitioningSuite extends SparkFunSuite with SharedSparkContext with Priva val arrPairs: RDD[(Array[Int], Int)] = sc.parallelize(Array(1, 2, 3, 4).toImmutableArraySeq, 2).map(x => (Array(x), x)) - def verify(testFun: => Unit): Unit = { - intercept[SparkException](testFun).getMessage.contains("array") + def verify(subCondition: String)(testFun: => Unit): Unit = { + checkError( + exception = intercept[SparkException](testFun), + condition = s"UNSUPPORTED_ARRAY_KEY.$subCondition", + sqlState = Some("0A000")) } - verify(arrs.distinct()) + // combineByKeyWithClassTag checks mapSideCombine before it looks at the partitioner, so the + // calls below that leave mapSideCombine at its default report MAP_SIDE_COMBINE even though + // they also end up with a HashPartitioner. + verify("MAP_SIDE_COMBINE")(arrs.distinct()) // We can't catch all usages of arrays, since they might occur inside other collections: // assert(fails { arrPairs.distinct() }) - verify(arrPairs.partitionBy(new HashPartitioner(2))) - verify(arrPairs.join(arrPairs)) - verify(arrPairs.leftOuterJoin(arrPairs)) - verify(arrPairs.rightOuterJoin(arrPairs)) - verify(arrPairs.fullOuterJoin(arrPairs)) - verify(arrPairs.groupByKey()) - verify(arrPairs.countByKey()) - verify(arrPairs.countByKeyApprox(1)) - verify(arrPairs.cogroup(arrPairs)) - verify(arrPairs.reduceByKeyLocally(_ + _)) - verify(arrPairs.reduceByKey(_ + _)) + verify("HASH_PARTITIONER")(arrPairs.partitionBy(new HashPartitioner(2))) + verify("HASH_PARTITIONER")(arrPairs.join(arrPairs)) + verify("HASH_PARTITIONER")(arrPairs.leftOuterJoin(arrPairs)) + verify("HASH_PARTITIONER")(arrPairs.rightOuterJoin(arrPairs)) + verify("HASH_PARTITIONER")(arrPairs.fullOuterJoin(arrPairs)) + verify("HASH_PARTITIONER")(arrPairs.groupByKey()) + verify("MAP_SIDE_COMBINE")(arrPairs.countByKey()) + // countByKeyApprox() is rejected by RDD.countByValueApprox, whose own array check is a + // separate condition that is still legacy. + checkError( + exception = intercept[SparkException](arrPairs.countByKeyApprox(1)), + condition = "_LEGACY_ERROR_TEMP_3015") + verify("HASH_PARTITIONER")(arrPairs.cogroup(arrPairs)) + verify("REDUCE_BY_KEY_LOCALLY")(arrPairs.reduceByKeyLocally(_ + _)) + verify("MAP_SIDE_COMBINE")(arrPairs.reduceByKey(_ + _)) } test("zero-length partitions should be correctly handled") { diff --git a/core/src/test/scala/org/apache/spark/SecurityManagerSuite.scala b/core/src/test/scala/org/apache/spark/SecurityManagerSuite.scala index 9db1afe0853a6..f4ecc540155df 100644 --- a/core/src/test/scala/org/apache/spark/SecurityManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/SecurityManagerSuite.scala @@ -53,7 +53,7 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { assert(securityManager.aclsEnabled()) assert(securityManager.checkUIViewPermissions("user1")) assert(securityManager.checkUIViewPermissions("user2")) - assert(securityManager.checkUIViewPermissions("user3") === false) + assert(!securityManager.checkUIViewPermissions("user3")) } test("set security with conf for groups") { @@ -65,8 +65,8 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { // default ShellBasedGroupsMappingProvider is used to resolve user groups val securityManager = new SecurityManager(conf); // assuming executing user does not belong to group1,group2 - assert(securityManager.checkUIViewPermissions("user1") === false) - assert(securityManager.checkUIViewPermissions("user2") === false) + assert(!securityManager.checkUIViewPermissions("user1")) + assert(!securityManager.checkUIViewPermissions("user2")) val conf2 = new SparkConf conf2.set(NETWORK_AUTH_ENABLED, true) @@ -91,8 +91,8 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { val securityManager3 = new SecurityManager(conf3) // BogusServiceProvider cannot be loaded and an error is logged returning an empty group set - assert(securityManager3.checkUIViewPermissions("user1") === false) - assert(securityManager3.checkUIViewPermissions("user2") === false) + assert(!securityManager3.checkUIViewPermissions("user1")) + assert(!securityManager3.checkUIViewPermissions("user2")) } test("set security with api") { @@ -102,7 +102,7 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { securityManager.setAcls(true) assert(securityManager.aclsEnabled()) securityManager.setAcls(false) - assert(securityManager.aclsEnabled() === false) + assert(!securityManager.aclsEnabled()) // acls are off so doesn't matter what view acls set to assert(securityManager.checkUIViewPermissions("user4")) @@ -110,11 +110,11 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { securityManager.setAcls(true) assert(securityManager.aclsEnabled()) securityManager.setViewAcls(Set[String]("user5"), Seq("user6", "user7")) - assert(securityManager.checkUIViewPermissions("user1") === false) + assert(!securityManager.checkUIViewPermissions("user1")) assert(securityManager.checkUIViewPermissions("user5")) assert(securityManager.checkUIViewPermissions("user6")) assert(securityManager.checkUIViewPermissions("user7")) - assert(securityManager.checkUIViewPermissions("user8") === false) + assert(!securityManager.checkUIViewPermissions("user8")) assert(securityManager.checkUIViewPermissions(null)) } @@ -132,8 +132,8 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { // change groups so they do not match securityManager.setViewAclsGroups(Seq("group4", "group5")) - assert(securityManager.checkUIViewPermissions("user1") === false) - assert(securityManager.checkUIViewPermissions("user2") === false) + assert(!securityManager.checkUIViewPermissions("user1")) + assert(!securityManager.checkUIViewPermissions("user2")) val conf2 = new SparkConf conf.set(USER_GROUPS_MAPPING, "BogusServiceProvider") @@ -143,13 +143,13 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { securityManager2.setViewAclsGroups(Seq("group1", "group2")) // group1,group2 do not match because of BogusServiceProvider - assert(securityManager.checkUIViewPermissions("user1") === false) - assert(securityManager.checkUIViewPermissions("user2") === false) + assert(!securityManager.checkUIViewPermissions("user1")) + assert(!securityManager.checkUIViewPermissions("user2")) // setting viewAclsGroups to empty should still not match because of BogusServiceProvider securityManager2.setViewAclsGroups(Nil) - assert(securityManager.checkUIViewPermissions("user1") === false) - assert(securityManager.checkUIViewPermissions("user2") === false) + assert(!securityManager.checkUIViewPermissions("user1")) + assert(!securityManager.checkUIViewPermissions("user2")) } test("set security modify acls") { @@ -160,7 +160,7 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { securityManager.setAcls(true) assert(securityManager.aclsEnabled()) securityManager.setAcls(false) - assert(securityManager.aclsEnabled() === false) + assert(!securityManager.aclsEnabled()) // acls are off so doesn't matter what view acls set to assert(securityManager.checkModifyPermissions("user4")) @@ -168,11 +168,11 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { securityManager.setAcls(true) assert(securityManager.aclsEnabled()) securityManager.setModifyAcls(Set("user5"), Seq("user6", "user7")) - assert(securityManager.checkModifyPermissions("user1") === false) + assert(!securityManager.checkModifyPermissions("user1")) assert(securityManager.checkModifyPermissions("user5")) assert(securityManager.checkModifyPermissions("user6")) assert(securityManager.checkModifyPermissions("user7")) - assert(securityManager.checkModifyPermissions("user8") === false) + assert(!securityManager.checkModifyPermissions("user8")) assert(securityManager.checkModifyPermissions(null)) } @@ -190,8 +190,8 @@ class SecurityManagerSuite extends SparkFunSuite with ResetSystemProperties { // change groups so they do not match securityManager.setModifyAclsGroups(Seq("group4", "group5")) - assert(securityManager.checkModifyPermissions("user1") === false) - assert(securityManager.checkModifyPermissions("user2") === false) + assert(!securityManager.checkModifyPermissions("user1")) + assert(!securityManager.checkModifyPermissions("user2")) // change so they match again securityManager.setModifyAclsGroups(Seq("group2", "group3")) diff --git a/core/src/test/scala/org/apache/spark/SparkContextSuite.scala b/core/src/test/scala/org/apache/spark/SparkContextSuite.scala index 709f267de155a..f7d668377c1cd 100644 --- a/core/src/test/scala/org/apache/spark/SparkContextSuite.scala +++ b/core/src/test/scala/org/apache/spark/SparkContextSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark import java.io.File import java.net.{MalformedURLException, URI} import java.nio.file.Files -import java.util.concurrent.{CountDownLatch, Semaphore, TimeUnit} +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, Semaphore, TimeUnit} import scala.concurrent.duration._ import scala.io.Source @@ -33,6 +33,8 @@ import org.apache.hadoop.mapred.TextInputFormat import org.apache.hadoop.mapreduce.lib.input.{TextInputFormat => NewTextInputFormat} import org.apache.logging.log4j.{Level, LogManager} import org.json4s.{DefaultFormats, Extraction} +import org.mockito.ArgumentMatchers.{any, eq => meq} +import org.mockito.Mockito.{mock, verify, when} import org.scalatest.concurrent.Eventually import org.scalatest.matchers.must.Matchers._ @@ -42,10 +44,14 @@ import org.apache.spark.internal.config._ import org.apache.spark.internal.config.Tests._ import org.apache.spark.internal.config.UI._ import org.apache.spark.launcher.SparkLauncher +import org.apache.spark.network.TransportContext +import org.apache.spark.network.netty.SparkTransportConf +import org.apache.spark.network.shuffle.ExternalBlockHandler import org.apache.spark.resource.ResourceAllocation import org.apache.spark.resource.ResourceUtils._ import org.apache.spark.resource.TestResourceIDs._ -import org.apache.spark.scheduler.{SparkListener, SparkListenerExecutorMetricsUpdate, SparkListenerJobStart, SparkListenerTaskEnd, SparkListenerTaskStart} +import org.apache.spark.scheduler.{LiveListenerBus, SparkListener, SparkListenerExecutorMetricsUpdate, SparkListenerJobStart, SparkListenerStageSubmitted, SparkListenerTaskEnd, SparkListenerTaskStart} +import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend import org.apache.spark.shuffle.FetchFailedException import org.apache.spark.util.{ThreadUtils, Utils} import org.apache.spark.util.ArrayImplicits._ @@ -713,6 +719,50 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu } } + test("SPARK-58748: Kubernetes drivers do not advertise wildcard bind addresses") { + Seq( + ("0.0.0.0", "10.129.36.37", "10.129.36.37"), + ("::", "10.129.36.37", "10.129.36.37"), + ("10.138.148.230", "driver-service", "10.138.148.230"), + ("2001:DB8:0:0::BEEF", "driver-service", "[2001:db8::beef]")).foreach { + case (bindAddress, advertisedAddress, expectedAddress) => + var observedAddress: Option[String] = None + val logAppender = new LogAppender("wildcard driver bind address") + val conf = new SparkConf(false) + .setMaster("k8s://https://localhost:6443") + .setAppName("driver-bind-address") + .set(DRIVER_BIND_ADDRESS, bindAddress) + .set(DRIVER_HOST_ADDRESS, advertisedAddress) + + withLogAppender(logAppender) { + val error = intercept[SparkException] { + new SparkContext(conf) { + override private[spark] def createSparkEnv( + conf: SparkConf, + isLocal: Boolean, + listenerBus: LiveListenerBus): SparkEnv = { + observedAddress = Some(conf.get(DRIVER_HOST_ADDRESS)) + throw new SparkException("stop after resolving the driver address") + } + } + } + assert(error.getMessage === "stop after resolving the driver address") + } + + assert(observedAddress.contains(expectedAddress)) + val wildcardMessages = logAppender.loggingEvents + .map(_.getMessage.getFormattedMessage) + .filter(_.contains("is a wildcard; preserving advertised driver host")) + if (Utils.isAnyLocalAddress(bindAddress)) { + assert(wildcardMessages.size === 1) + assert(wildcardMessages.head.contains(bindAddress)) + assert(wildcardMessages.head.contains(advertisedAddress)) + } else { + assert(wildcardMessages.isEmpty) + } + } + } + testCancellingTasks("that raise interrupted exception on cancel") { Thread.sleep(9999999) } @@ -1521,6 +1571,162 @@ class SparkContextSuite extends SparkFunSuite with LocalSparkContext with Eventu sc = new SparkContext(conf) assert(sc.env.memoryManager.maxOffHeapStorageMemory > 0) } + + test("SPARK-41246: fail-fast on RDD id overflow") { + val conf = new SparkConf().setAppName("test").setMaster("local[1]") + sc = new SparkContext(conf) + sc.setNextRddIdForTesting(Int.MaxValue) + val last = sc.parallelize(Seq(1), 1) + assert(last.id === Int.MaxValue) + val err = intercept[SparkException] { + sc.parallelize(Seq(2), 1) + } + assert(err.getMessage.contains("Int.MaxValue")) + assert(err.getMessage.contains("overflowed")) + } + + test("SPARK-58828: holdExecutors and resumeExecutors are unsupported by the local scheduler") { + sc = new SparkContext(new SparkConf().setAppName("test").setMaster("local")) + assert(!sc.executorHoldSupported) + assert(!sc.holdExecutors()) + assert(!sc.resumeExecutors()) + } + + test("SPARK-58828: holdExecutors requires external shuffle service and decommission support") { + sc = new SparkContext( + new SparkConf().setAppName("test").setMaster("local-cluster[1,1,1024]")) + assert(!sc.executorHoldSupported) + val err = intercept[IllegalArgumentException] { + sc.holdExecutors() + } + assert(err.getMessage.contains(SHUFFLE_SERVICE_ENABLED.key)) + assert(err.getMessage.contains(DECOMMISSION_ENABLED.key)) + } + + private def withExternalShuffleServer(conf: SparkConf)(body: => Unit): Unit = { + // The executors register with the external shuffle service on startup, so run one + val transportConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numUsableCores = 2) + val rpcHandler = new ExternalBlockHandler(transportConf, null) + val transportContext = new TransportContext(transportConf, rpcHandler) + val server = transportContext.createServer() + try { + conf.set(SHUFFLE_SERVICE_PORT, server.getPort) + body + } finally { + Utils.tryLogNonFatalError(server.close()) + Utils.tryLogNonFatalError(rpcHandler.close()) + Utils.tryLogNonFatalError(transportContext.close()) + } + } + + private def verifyHoldAndResumeExecutors(conf: SparkConf): Unit = { + withExternalShuffleServer(conf) { + sc = new SparkContext(conf) + TestUtils.waitUntilExecutorsUp(sc, 1, 60000) + assert(sc.executorHoldSupported) + assert(!sc.executorsHeld) + + // Shuffle output written before the hold must survive the drain + val shuffled = sc.parallelize(1 to 100, 4).map(i => (i % 8, i)).reduceByKey(_ + _) + assert(shuffled.count() === 8) + val shuffleId = + shuffled.dependencies.head.asInstanceOf[ShuffleDependency[_, _, _]].shuffleId + val tracker = sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster] + assert(tracker.getNumAvailableOutputs(shuffleId) === 4) + + assert(sc.holdExecutors()) + assert(sc.executorsHeld) + // The executor finishes decommissioning and exits, and no new one replaces it + eventually(timeout(60.seconds)) { + assert(sc.getExecutorIds().isEmpty) + } + // And stays drained: a transiently empty poll would also pass the check above under a + // register-and-drain churn, so verify the zero requirement actually settled + Thread.sleep(2000) + assert(sc.getExecutorIds().isEmpty) + + assert(sc.resumeExecutors()) + assert(!sc.executorsHeld) + // The restored requirement brings an executor back + eventually(timeout(60.seconds)) { + assert(sc.getExecutorIds().nonEmpty) + } + + // The map output survived the drain: only the reduce stage re-runs when the shuffle is + // read again. Asserted after the job, since the executor removal is processed + // asynchronously on the DAGScheduler event loop. + val submitted = new ConcurrentLinkedQueue[Int]() + sc.addSparkListener(new SparkListener { + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = + submitted.add(e.stageInfo.stageId) + }) + assert(shuffled.collect().length === 8) + sc.listenerBus.waitUntilEmpty() + assert(submitted.size() === 1, s"expected only the reduce stage, got $submitted") + assert(tracker.getNumAvailableOutputs(shuffleId) === 4) + } + } + + test("SPARK-58828: holdExecutors drains the executors and resumeExecutors brings them back") { + verifyHoldAndResumeExecutors( + new SparkConf().setAppName("test").setMaster("local-cluster[1,1,1024]") + .set(SHUFFLE_SERVICE_ENABLED, true) + .set(DECOMMISSION_ENABLED, true)) + } + + test("SPARK-58828: hold and resume the executors with dynamic allocation") { + verifyHoldAndResumeExecutors( + new SparkConf().setAppName("test").setMaster("local-cluster[1,1,1024]") + .set(SHUFFLE_SERVICE_ENABLED, true) + .set(DECOMMISSION_ENABLED, true) + .set(DYN_ALLOCATION_ENABLED, true) + .set(DYN_ALLOCATION_INITIAL_EXECUTORS, 1) + .set(DYN_ALLOCATION_MIN_EXECUTORS, 1)) + } + + test("SPARK-58828: a task running at the hold finishes and a pending one runs after resume") { + val conf = new SparkConf().setAppName("test").setMaster("local-cluster[1,1,1024]") + .set(SHUFFLE_SERVICE_ENABLED, true) + .set(DECOMMISSION_ENABLED, true) + withExternalShuffleServer(conf) { + sc = new SparkContext(conf) + TestUtils.waitUntilExecutorsUp(sc, 1, 60000) + val taskStarted = new Semaphore(0) + sc.addSparkListener(new SparkListener { + override def onTaskStart(taskStart: SparkListenerTaskStart): Unit = taskStarted.release() + }) + // Two tasks on one core: the second is still pending when the hold starts + val result = sc.parallelize(1 to 2, 2).map { i => Thread.sleep(2000); i * 10 }.collectAsync() + assert(taskStarted.tryAcquire(1, 60, TimeUnit.SECONDS)) + assert(sc.holdExecutors()) + // The running task finishes before the executor exits + eventually(timeout(60.seconds)) { + assert(sc.getExecutorIds().isEmpty) + } + assert(sc.resumeExecutors()) + // The pending task runs after the resume and the job completes with no lost work + assert(ThreadUtils.awaitResult(result, 2.minutes).sorted === Seq(10, 20)) + } + } + + test("SPARK-58828: restoring the hold invariant pushes a zero requirement and drains") { + sc = new SparkContext(new SparkConf().setAppName("test").setMaster("local")) + val backend = mock(classOf[CoarseGrainedSchedulerBackend]) + when(backend.republishRequestedTotals()).thenReturn(true) + when(backend.getExecutorIds()).thenReturn(Seq("1", "2")) + when(backend.decommissionExecutors(any(), any(), any())).thenReturn(Seq("1", "2")) + assert(sc.zeroExecutorRequirementAndDrain(backend)) + verify(backend).republishRequestedTotals() + verify(backend).decommissionExecutors(any(), meq(false), meq(false)) + + // A failed publish must not abort the drain + val failing = mock(classOf[CoarseGrainedSchedulerBackend]) + when(failing.republishRequestedTotals()).thenThrow(new RuntimeException("boom")) + when(failing.getExecutorIds()).thenReturn(Seq("3")) + when(failing.decommissionExecutors(any(), any(), any())).thenReturn(Seq("3")) + assert(!sc.zeroExecutorRequirementAndDrain(failing)) + verify(failing).decommissionExecutors(any(), meq(false), meq(false)) + } } object SparkContextSuite { diff --git a/core/src/test/scala/org/apache/spark/SparkThrowableSuite.scala b/core/src/test/scala/org/apache/spark/SparkThrowableSuite.scala index 5fb6924383f0e..785325001be05 100644 --- a/core/src/test/scala/org/apache/spark/SparkThrowableSuite.scala +++ b/core/src/test/scala/org/apache/spark/SparkThrowableSuite.scala @@ -125,7 +125,9 @@ class SparkThrowableSuite extends SparkFunSuite { errorClassesJson.openStream(), new TypeReference[Map[String, String]]() {}) val errorStates = mapper.readValue( errorStatesJson.openStream(), new TypeReference[Map[String, ErrorStateInfo]]() {}) - val errorConditionStates = errorReader.errorInfoMap.values.toSeq.flatMap(_.sqlState).toSet + val errorConditionStates = errorReader.errorInfoMap.values.toSeq.flatMap { i => + i.sqlState ++ i.subClass.getOrElse(Map.empty).values.flatMap(_.sqlState) + }.toSet assert(Set("22012", "22003", "42601").subsetOf(errorStates.keySet)) assert(errorClasses.keySet.filter(!_.matches("[A-Z0-9]{2}")).isEmpty) assert(errorStates.keySet.filter(!_.matches("[A-Z0-9]{5}")).isEmpty) @@ -133,6 +135,27 @@ class SparkThrowableSuite extends SparkFunSuite { assert(errorConditionStates.diff(errorStates.keySet).isEmpty) } + test("Sub-condition SQLSTATE overrides are limited to the documented exceptions") { + // Sub-conditions inherit their condition's SQLSTATE. The only permitted overrides are + // the wire-compatibility exceptions documented in the error README's SQLSTATE section. + val allowedOverrides = Set( + "INVALID_HANDLE.SESSION_CHANGED", + "INVALID_HANDLE.SESSION_CLOSED", + "INVALID_HANDLE.SESSION_NOT_FOUND") + errorReader.errorInfoMap.foreach { case (condition, info) => + info.subClass.getOrElse(Map.empty).foreach { case (sub, subInfo) => + subInfo.sqlState.foreach { subState => + val name = s"$condition.$sub" + assert( + allowedOverrides(name), + s"$name declares its own SQLSTATE ($subState). Sub-conditions inherit their " + + "condition's SQLSTATE; do not add new overrides. See the SQLSTATE section " + + "of the error README.") + } + } + } + } + test("Message invariants") { val messageSeq = errorReader.errorInfoMap.values.toSeq.flatMap { i => Seq(i.message) ++ i.subClass.getOrElse(Map.empty).values.toSeq.map(_.message) @@ -631,6 +654,70 @@ class SparkThrowableSuite extends SparkFunSuite { } } + test("sub-condition SQLSTATE overrides the main condition's SQLSTATE") { + withTempDir { dir => + val json = new File(dir, "errors.json") + Files.writeString( + json.toPath, + """ + |{ + | "TEST_MAIN_STATE": { + | "message": [ + | "Main message." + | ], + | "sqlState": "42000", + | "subClass": { + | "SUB_WITHOUT_STATE": { + | "message": [ + | "Sub-condition without its own SQLSTATE." + | ] + | }, + | "SUB_WITH_STATE": { + | "message": [ + | "Sub-condition with its own SQLSTATE." + | ], + | "sqlState": "08003" + | } + | } + | } + |} + |""".stripMargin, + StandardCharsets.UTF_8) + + val reader = + new ErrorClassesJsonReader(Seq(errorJsonFilePath.toUri.toURL, json.toURI.toURL)) + // A sub-condition with its own SQLSTATE overrides the main condition's. + assert(reader.getSqlState("TEST_MAIN_STATE.SUB_WITH_STATE") == "08003") + // A sub-condition without its own SQLSTATE inherits the main condition's. + assert(reader.getSqlState("TEST_MAIN_STATE.SUB_WITHOUT_STATE") == "42000") + assert(reader.getSqlState("TEST_MAIN_STATE") == "42000") + // Degenerate inputs keep the pre-existing non-throwing behavior: anything that is not + // a known "MAIN.SUB" pair resolves to the main condition's SQLSTATE, or null. + assert(reader.getSqlState("TEST_MAIN_STATE.NON_EXISTENT_SUB") == "42000") + assert(reader.getSqlState("TEST_MAIN_STATE.SUB_WITH_STATE.EXTRA") == "42000") + assert(reader.getSqlState("NON_EXISTENT") == null) + assert(reader.getSqlState(null) == null) + } + } + + test("INVALID_HANDLE session sub-conditions carry SQLSTATE 08003") { + // The session sub-conditions mean the server-side session backing a Connect client is + // gone (08003, connection does not exist); the wire-visible names are unchanged. + val sessionSubConditions = Seq("SESSION_CHANGED", "SESSION_CLOSED", "SESSION_NOT_FOUND") + sessionSubConditions.foreach { sub => + assert(errorReader.getSqlState(s"INVALID_HANDLE.$sub") == "08003", sub) + } + // Every other sub-condition concerns a single operation on a healthy session and keeps + // the condition's SQLSTATE. + val otherSubConditions = errorReader + .errorInfoMap("INVALID_HANDLE").subClass.get.keys.toSeq.diff(sessionSubConditions) + assert(otherSubConditions.nonEmpty) + otherSubConditions.foreach { sub => + assert(errorReader.getSqlState(s"INVALID_HANDLE.$sub") == "HY000", sub) + } + assert(errorReader.getSqlState("INVALID_HANDLE") == "HY000") + } + test("detect unused message parameters") { checkError( exception = intercept[SparkException] { diff --git a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala index 3c8d9b2f14e21..a9351630ad8db 100644 --- a/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala +++ b/core/src/test/scala/org/apache/spark/StreamingShuffleOutputTrackerSuite.scala @@ -270,6 +270,20 @@ class StreamingShuffleOutputTrackerSuite assert(tracker.getAllShuffleWriterTaskLocations(0).isEmpty) } + test("StreamingShuffleOutputTrackerMaster - containsShuffle reflects register/unregister") { + val conf = new SparkConf(false) + val tracker = newTrackerMaster(conf) + + // A shuffle that was never registered is absent. + assert(!tracker.containsShuffle(0)) + + tracker.registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 1) + assert(tracker.containsShuffle(0)) + + tracker.unregisterShuffle(0) + assert(!tracker.containsShuffle(0)) + } + test("StreamingShuffleOutputTrackerMaster - register writer before shuffle fails") { val conf = new SparkConf(false) val tracker = newTrackerMaster(conf) diff --git a/core/src/test/scala/org/apache/spark/deploy/JsonProtocolSuite.scala b/core/src/test/scala/org/apache/spark/deploy/JsonProtocolSuite.scala index 6d2c663a2588e..ac839103aef58 100644 --- a/core/src/test/scala/org/apache/spark/deploy/JsonProtocolSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/JsonProtocolSuite.scala @@ -39,6 +39,18 @@ class JsonProtocolSuite extends SparkFunSuite with JsonTestUtils { assertValidDataInJson(output, JsonMethods.parse(JsonConstants.appInfoJsonStr)) } + test("SPARK-59055: writeApplicationInfo with the hold status") { + val appInfo = createAppInfo() + appInfo.holdSupported = true + appInfo.held = true + val output = JsonProtocol.writeApplicationInfo(appInfo) + assertValidJson(output) + assert(output \ "holdsupported" === JBool(true)) + assert(output \ "held" === JBool(true)) + // The application has no executor left, so its hold is complete. + assert(output \ "draining" === JInt(0)) + } + test("writeWorkerInfo") { val output = JsonProtocol.writeWorkerInfo(createWorkerInfo()) assertValidJson(output) @@ -191,7 +203,9 @@ object JsonConstants { |"resourcesperslave":[{"name":"fpga", |"amount":3},{"name":"gpu","amount":3}], |"submitdate":"%s", - |"state":"WAITING","duration":%d} + |"state":"WAITING", + |"holdsupported":false,"held":false,"draining":0, + |"duration":%d} """.format(System.getProperty("user.name", "<unknown>"), submitDate.toString, currTimeInMillis - appInfoStartTime).stripMargin diff --git a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala index 90ef0aa510c24..4614435af44cf 100644 --- a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala @@ -496,6 +496,51 @@ class StandaloneDynamicAllocationSuite } } + test("SPARK-59055: SparkContext reports the hold status of the application to the Master") { + sc = new SparkContext(appConf + .set(config.SHUFFLE_SERVICE_ENABLED, true) + .set(config.DECOMMISSION_ENABLED, true)) + // The report from the end of the SparkContext constructor. + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().length === 1) + assert(getApplications().head.holdSupported) + } + assert(!getApplications().head.held) + + assert(sc.holdExecutors()) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().head.held) + } + // The faked executors never exit, so all of them are still counted as draining. + assert(getApplications().head.numDrainingExecutors === 2) + + assert(sc.resumeExecutors()) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(!getApplications().head.held) + } + } + + test("SPARK-59055: spark.ui.holdEnabled=false does not suppress the hold status") { + sc = new SparkContext(appConf + .set(config.SHUFFLE_SERVICE_ENABLED, true) + .set(config.DECOMMISSION_ENABLED, true) + .set(config.UI.UI_HOLD_ENABLED, false)) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().length === 1) + } + + // spark.ui.holdEnabled gates the controls on the driver UI, not the capability itself: + // holdExecutors() is a developer API that the config does not touch, so a hold made with + // the config off must still be visible on the Master. + assert(sc.holdExecutors()) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().head.held) + } + assert(getApplications().head.holdSupported) + assert(getApplications().head.isHeld) + assert(getApplications().head.numDrainingExecutors === 2) + } + test("executor registration on a excluded host must fail") { // The context isn't really used by the test, but it helps with creating a test scheduler, // since CoarseGrainedSchedulerBackend makes a lot of calls to the context instance. diff --git a/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala b/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala index 877aee47cd658..4147535eb0db0 100644 --- a/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/client/AppClientSuite.scala @@ -217,6 +217,43 @@ class AppClientSuite } } + test("SPARK-59055: report the hold status of the application to the Master") { + Utils.tryWithResource(new AppClientInst(masterRpcEnv.address.toSparkURL)) { ci => + ci.client.start() + + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().length === 1, "master should have 1 registered app") + } + + // The Master knows nothing about the hold until the driver reports it. + val app = getApplications().head + assert(!app.holdSupported && !app.held && app.numDrainingExecutors === 0) + + ci.client.reportHoldStatus(supported = true, held = false) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().head.holdSupported) + } + assert(!getApplications().head.held) + + ci.client.reportHoldStatus(supported = true, held = true) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().head.held) + } + + ci.client.reportHoldStatus(supported = true, held = false) + eventually(timeout(10.seconds), interval(10.millis)) { + assert(!getApplications().head.held) + } + + // Issue stop command for Client to disconnect from Master + ci.client.stop() + + eventually(timeout(10.seconds), interval(10.millis)) { + assert(getApplications().isEmpty, "master should have 0 registered apps") + } + } + } + test("request from AppClient before initialized with master") { Utils.tryWithResource(new AppClientInst(masterRpcEnv.address.toSparkURL)) { ci => diff --git a/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileWritersSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileWritersSuite.scala index 00a92c503be4e..36c5b1915edd7 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileWritersSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/EventLogFileWritersSuite.scala @@ -352,6 +352,20 @@ class SingleEventLogFileWriterSuite extends EventLogFileWritersSuite { "a fine:mind$dollar{bills}.1", None, Some(CompressionCodec.LZ4))) } + test("Event log file names") { + Seq(None, Some("attempt1")).foreach { attemptId => + val baseName = attemptId.map(id => s"app1_$id").getOrElse("app1") + val names = SingleEventLogFileWriter.getLogFileNames("app1", attemptId) + assert(names.head === baseName) + assert(names.contains(s"$baseName.inprogress")) + assert(names.size === (1 + CompressionCodec.shortCompressionCodecNames.size) * 2) + CompressionCodec.shortCompressionCodecNames.keys.foreach { codec => + assert(names.contains(s"$baseName.$codec")) + assert(names.contains(s"$baseName.$codec.inprogress")) + } + } + } + override protected def createWriter( appId: String, appAttemptId: Option[String], diff --git a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala index 873431d7942e3..d932265791dcb 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/FsHistoryProviderSuite.scala @@ -83,9 +83,10 @@ abstract class FsHistoryProviderSuite extends SparkFunSuite with Matchers with P appId: String, appAttemptId: Option[String], inProgress: Boolean, - codec: Option[String] = None): File = { + codec: Option[String] = None, + logDir: File = testDir): File = { val ip = if (inProgress) EventLogFileWriter.IN_PROGRESS else "" - val logUri = SingleEventLogFileWriter.getLogPath(testDir.toURI, appId, appAttemptId, codec) + val logUri = SingleEventLogFileWriter.getLogPath(logDir.toURI, appId, appAttemptId, codec) val logPath = new Path(logUri).toUri.getPath + ip new File(logPath) } @@ -1731,6 +1732,7 @@ abstract class FsHistoryProviderSuite extends SparkFunSuite with Matchers with P val conf = createTestConf(true) conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, onDemandEnabled) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, false) val hadoopConf = SparkHadoopUtil.newConfiguration(conf) val provider = new FsHistoryProvider(conf) @@ -1759,6 +1761,186 @@ abstract class FsHistoryProviderSuite extends SparkFunSuite with Matchers with P } } + test("Support spark.history.fs.eventLog.onDemandLoadEnabled") { + Seq(true, false).foreach { onDemandEnabled => + Seq(None, Some(CompressionCodec.LZF)).foreach { codecName => + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, onDemandEnabled) + val provider = new FsHistoryProvider(conf) + val appId = s"app1-${codecName.getOrElse("none")}" + val logFile = newLogFile(appId, None, inProgress = false, codecName, dir) + val codec = codecName.map(CompressionCodec.createCodec(conf, _)) + writeFile(logFile, codec, + SparkListenerApplicationStart(appId, Some(appId), 1000, "testuser", None), + SparkListenerApplicationEnd(5000)) + + assert(provider.getListing().isEmpty) + assert(provider.getAppUI(appId, None).isDefined == onDemandEnabled) + assert(provider.getListing().length === (if (onDemandEnabled) 1 else 0)) + + if (onDemandEnabled) { + val appInfo = provider.getListing().next() + assert(appInfo.name === appId) + assert(appInfo.attempts.head.sparkUser === "testuser") + assert(appInfo.attempts.head.completed) + assert(appInfo.attempts.head.duration === 4000) + } + + assert(provider.getAppUI("nonexist", None).isEmpty) + assert(provider.getListing().length === (if (onDemandEnabled) 1 else 0)) + + provider.stop() + } + } + } + } + + test("Support on-demand loading for in-progress single event logs") { + Seq(None, Some(CompressionCodec.LZF)).foreach { codecName => + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val provider = new FsHistoryProvider(conf) + val appId = s"app1-${codecName.getOrElse("none")}" + val logFile = newLogFile(appId, None, inProgress = true, codecName, dir) + val codec = codecName.map(CompressionCodec.createCodec(conf, _)) + writeFile(logFile, codec, + SparkListenerApplicationStart(appId, Some(appId), 1000, "testuser", None)) + + assert(provider.getAppUI(appId, None).isDefined) + assert(!provider.getListing().next().attempts.head.completed) + + provider.stop() + } + } + } + + test("On-demand loading preserves in-progress single event logs during cleanup") { + withTempDir { dir => + val clock = new ManualClock(0) + val conf = createTestConf(true) + .set(HISTORY_LOG_DIR, dir.getAbsolutePath) + .set(CLEANER_ENABLED, true) + .set(MAX_LOG_AGE_S, 0L) + .set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + .set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val provider = new FsHistoryProvider(conf, clock) + val logFile = newLogFile("app1", None, inProgress = true, logDir = dir) + writeFile(logFile, None, + SparkListenerApplicationStart("app1", Some("app1"), 0, "testuser", None)) + + assert(provider.getAppUI("app1", None).isDefined) + assert(provider.getListing().next().attempts.head.appSparkVersion === SPARK_VERSION) + assert(logFile.exists()) + + provider.cleanLogs() + assert(logFile.exists()) + + provider.stop() + } + } + + test("On-demand loading respects single event log attempts") { + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val provider = new FsHistoryProvider(conf) + val logFile = newLogFile("app1", Some("attempt1"), inProgress = false, logDir = dir) + writeFile(logFile, None, + SparkListenerApplicationStart("app1", Some("app1"), 1000, "testuser", Some("attempt1")), + SparkListenerApplicationEnd(5000)) + + assert(provider.getAppUI("app1", None).isEmpty) + assert(provider.getListing().isEmpty) + assert(provider.getAppUI("app1", Some("attempt1")).isDefined) + + provider.stop() + } + } + + test("On-demand loading finds single event logs in all directories") { + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, s"${testDir.getAbsolutePath},${dir.getAbsolutePath}") + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val provider = new FsHistoryProvider(conf) + val logFile = newLogFile("app1", None, inProgress = false, logDir = dir) + writeFile(logFile, None, + SparkListenerApplicationStart("app1", Some("app1"), 1000, "testuser", None), + SparkListenerApplicationEnd(5000)) + + assert(provider.getAppUI("app1", None).isDefined) + assert(provider.getListing().next().attempts.head.logSourceFullPath === + Some(dir.getAbsolutePath)) + + provider.stop() + } + } + + test("On-demand loading supports single event logs when scanning is disabled") { + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) + conf.set(SCAN_DISABLED_PATH_PATTERNS, Seq("file:.*")) + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, false) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val provider = new FsHistoryProvider(conf) + val logFile = newLogFile("app1", None, inProgress = false, logDir = dir) + writeFile(logFile, None, + SparkListenerApplicationStart("app1", Some("app1"), 1000, "testuser", None), + SparkListenerApplicationEnd(5000)) + + provider.checkForLogs() + assert(provider.getListing().isEmpty) + assert(provider.getAppUI("app1", None).isDefined) + + provider.checkForLogs() + assert(provider.getListing().length === 1) + + provider.stop() + } + } + + test("On-demand loading supports rolling and single event logs together") { + withTempDir { dir => + val conf = createTestConf(true) + conf.set(HISTORY_LOG_DIR, dir.getAbsolutePath) + conf.set(EVENT_LOG_ROLLING_ON_DEMAND_LOAD_ENABLED, true) + conf.set(EVENT_LOG_SINGLE_ON_DEMAND_LOAD_ENABLED, true) + val hadoopConf = SparkHadoopUtil.newConfiguration(conf) + val provider = new FsHistoryProvider(conf) + + val rollingWriter = new RollingEventLogFilesWriter( + "app-rolling", None, dir.toURI, conf, hadoopConf) + rollingWriter.start() + writeEventsToRollingWriter(rollingWriter, Seq( + SparkListenerApplicationStart("app-rolling", Some("app-rolling"), 0, "user", None), + SparkListenerJobStart(1, 0, Seq.empty)), rollFile = false) + rollingWriter.stop() + + val singleLog = newLogFile("app-single", None, inProgress = false, logDir = dir) + writeFile(singleLog, None, + SparkListenerApplicationStart("app-single", Some("app-single"), 0, "user", None), + SparkListenerApplicationEnd(1)) + + assert(provider.getAppUI("app-rolling", None).isDefined) + assert(provider.getAppUI("app-single", None).isDefined) + assert(provider.getListing().length === 2) + assert(provider.getAppUI("nonexist", None).isEmpty) + assert(provider.getListing().length === 2) + + provider.stop() + } + } + test("SPARK-56278: On-demand loading populates accurate metadata via mergeApplicationListing") { withTempDir { dir => val conf = createTestConf(true) diff --git a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala index 79e0001fa22fb..f252ac22c4a27 100644 --- a/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/history/HistoryServerSuite.scala @@ -348,6 +348,12 @@ abstract class HistoryServerSuite extends SparkFunSuite with BeforeAndAfter with getContentAndCode("foobar")._1 should be (HttpServletResponse.SC_NOT_FOUND) } + test("SPARK-59010: hold status is not available through the history server") { + val holdStatus = getContentAndCode("applications/local-1422981780767/holdstatus") + holdStatus._1 should be (HttpServletResponse.SC_SERVICE_UNAVAILABLE) + holdStatus._3 should be (Some("Hold status not available through the history server.")) + } + test("automatically retrieve uiRoot from request through Knox") { assert(sys.props.get("spark.ui.proxyBase").isEmpty, "spark.ui.proxyBase is defined but it should not for this UT") diff --git a/core/src/test/scala/org/apache/spark/deploy/master/MasterSuite.scala b/core/src/test/scala/org/apache/spark/deploy/master/MasterSuite.scala index c138d1142a474..2b68f3db1b449 100644 --- a/core/src/test/scala/org/apache/spark/deploy/master/MasterSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/master/MasterSuite.scala @@ -250,6 +250,40 @@ class MasterSuite extends MasterSuiteBase { assert(master.invokePrivate(_createApplication(desc, null)).id === "spark-45756") } + test("SPARK-59055: The executors of a held application are counted as draining") { + val appInfo = makeAppInfo(1024) + val worker = DeployTestUtils.createWorkerInfo() + appInfo.addExecutor(worker, 1, 1024, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) + appInfo.addExecutor(worker, 1, 1024, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) + + // The executors of an application that is not held are not draining. + assert(appInfo.numDrainingExecutors === 0) + + // A hold that the driver did not report as supported is not treated as held. + appInfo.held = true + assert(appInfo.numDrainingExecutors === 0) + + // While held, an executor that has not exited yet is still draining its running tasks. + appInfo.holdSupported = true + assert(appInfo.numDrainingExecutors === 2) + + // The hold is complete once the last executor is gone. + appInfo.executors.values.toSeq.foreach(appInfo.removeExecutor) + assert(appInfo.numDrainingExecutors === 0) + assert(appInfo.isHeld) + assert(appInfo.stateText === "WAITING (held)") + + // A single draining executor is reported in the singular. + appInfo.addExecutor(worker, 1, 1024, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) + assert(appInfo.stateText === "WAITING (held, draining 1 executor)") + + // A finished application is never held, even though the Master keeps its executors for the + // UI: its driver is gone, so the last reported hold is stale. + appInfo.markFinished(ApplicationState.FINISHED) + assert(!appInfo.isHeld) + assert(appInfo.numDrainingExecutors === 0) + } + test("SPARK-57451: Allows REST server and spark.authenticate.secret to be enabled together") { val conf = new SparkConf() .set(MASTER_REST_SERVER_ENABLED, true) diff --git a/core/src/test/scala/org/apache/spark/deploy/master/ui/ReadOnlyMasterWebUISuite.scala b/core/src/test/scala/org/apache/spark/deploy/master/ui/ReadOnlyMasterWebUISuite.scala index 2679349bfe028..66661aa1cd6c4 100644 --- a/core/src/test/scala/org/apache/spark/deploy/master/ui/ReadOnlyMasterWebUISuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/master/ui/ReadOnlyMasterWebUISuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.deploy.master.ui.MasterWebUISuite._ import org.apache.spark.internal.config.DECOMMISSION_ENABLED import org.apache.spark.internal.config.UI.MASTER_UI_VISIBLE_ENV_VAR_PREFIXES import org.apache.spark.internal.config.UI.UI_KILL_ENABLED +import org.apache.spark.resource.ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID import org.apache.spark.rpc.{RpcEndpointRef, RpcEnv} import org.apache.spark.util.Utils @@ -73,6 +74,27 @@ class ReadOnlyMasterWebUISuite extends SparkFunSuite { } } + test("SPARK-59055: annotate the state of a held application") { + val url = s"http://${Utils.localHostNameForURI()}:${masterWebUI.boundPort}/" + app1.holdSupported = true + app1.held = true + app1.addExecutor(createWorkerInfo(), 1, 1234, Map.empty, DEFAULT_RESOURCE_PROFILE_ID) + app1.addExecutor(createWorkerInfo(), 1, 1234, Map.empty, DEFAULT_RESOURCE_PROFILE_ID) + try { + var result = Source.fromInputStream(sendHttpRequest(url, "GET", "").getInputStream).mkString + assert(result.contains("WAITING (held, draining 2 executors)")) + + // An application that did not report the hold as supported is not annotated. + app1.holdSupported = false + result = Source.fromInputStream(sendHttpRequest(url, "GET", "").getInputStream).mkString + assert(!result.contains("(held")) + } finally { + app1.holdSupported = false + app1.held = false + app1.executors.values.toSeq.foreach(app1.removeExecutor) + } + } + test("/app/kill POST method is not allowed") { val url = s"http://${Utils.localHostNameForURI()}:${masterWebUI.boundPort}/app/kill/" val body = convPostDataToString(Map(("id", "1"), ("terminate", "true"))) diff --git a/core/src/test/scala/org/apache/spark/deploy/security/NonKerberosCredentialsSuite.scala b/core/src/test/scala/org/apache/spark/deploy/security/NonKerberosCredentialsSuite.scala index d6ecc954cfe0b..f9f77b4d27985 100644 --- a/core/src/test/scala/org/apache/spark/deploy/security/NonKerberosCredentialsSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/security/NonKerberosCredentialsSuite.scala @@ -79,6 +79,24 @@ private class TestFailingProvider extends HadoopDelegationTokenProvider { } } +private class TestRequirementFailingProvider extends HadoopDelegationTokenProvider { + override def serviceName: String = "test-requirement-failing" + + override def delegationTokensRequired( + sparkConf: SparkConf, hadoopConf: Configuration): Boolean = { + if (sparkConf.getBoolean( + "spark.test.requirement-failing.throw", false)) { + throw new RuntimeException("Simulated provider requirement failure") + } + false + } + + override def obtainDelegationTokens( + hadoopConf: Configuration, + sparkConf: SparkConf, + creds: Credentials): Option[Long] = None +} + // Adds a credential but reports no expiry (returns None), mimicking providers such as // HBaseDelegationTokenProvider. This exercises the nextRenewal == Long.MaxValue case where // credentials were nevertheless obtained. @@ -149,6 +167,32 @@ class NonKerberosCredentialsSuite extends SparkFunSuite { assert(new String(creds.getSecretKey(new Text("test.direct.credential"))) === "test-token") } + test("provider requirement failure does not prevent other providers from running") { + val sparkConf = baseConf + .set("spark.security.credentials.test-requirement-failing.enabled", "true") + .set("spark.test.requirement-failing.throw", "true") + val manager = new HadoopDelegationTokenManager(sparkConf, hadoopConf, null) + + assert(manager.renewalEnabled) + + val creds = new Credentials() + manager.obtainDelegationTokens(creds) + + assert(creds.getSecretKey(new Text("test.direct.credential")) != null) + } + + test("provider requirement failure does not abort renewalEnabled") { + val sparkConf = baseConf + .set("spark.security.credentials.test-direct.enabled", "false") + .set("spark.security.credentials.test-failing.enabled", "false") + .set("spark.security.credentials.test-noexpiry.enabled", "false") + .set("spark.security.credentials.test-requirement-failing.enabled", "true") + .set("spark.test.requirement-failing.throw", "true") + val manager = new HadoopDelegationTokenManager(sparkConf, hadoopConf, null) + + assert(!manager.renewalEnabled) + } + test("individual provider can be disabled via per-service config") { val sparkConf = baseConf .set("spark.security.credentials.test-direct.enabled", "false") diff --git a/core/src/test/scala/org/apache/spark/deploy/security/OidcCredentialIntegrationSuite.scala b/core/src/test/scala/org/apache/spark/deploy/security/OidcCredentialIntegrationSuite.scala new file mode 100644 index 0000000000000..47d19e05a2a5e --- /dev/null +++ b/core/src/test/scala/org/apache/spark/deploy/security/OidcCredentialIntegrationSuite.scala @@ -0,0 +1,593 @@ +/* + * 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.deploy.security + +import java.io.File +import java.nio.file.Files +import java.time.Instant +import java.util.Optional +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference} + +import scala.concurrent.duration._ + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.io.Text +import org.mockito.ArgumentCaptor +import org.mockito.Mockito.{mock, verify} +import org.scalatest.concurrent.Eventually.{eventually, timeout} + +import org.apache.spark.{SparkConf, SparkFunSuite, VersionedCredentials} +import org.apache.spark.deploy.SparkHadoopUtil +import org.apache.spark.internal.config._ +import org.apache.spark.internal.config.Network.NETWORK_CRYPTO_ENABLED +import org.apache.spark.rpc.RpcEndpointRef +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.UpdateDelegationTokens +import org.apache.spark.security._ + +/** + * Integration tests for SPARK-57896: Kerberos coexistence and per-user token tests. + * + * Verifies that: + * 1. UserCredentialManager (OIDC) delivers credentials via the update callback + * 2. Credential refresh works end-to-end with expiring tokens + * 3. Per-user identity tokens produce valid credentials identically to workload tokens + * 4. Both UserCredentialManager and HadoopDelegationTokenManager can run simultaneously + * without interfering with each other + * 5. Failure in one credential system does not affect the other + * 6. TaskDescription credentials are applied to the executor store with version guard + */ +class OidcCredentialIntegrationSuite extends SparkFunSuite { + + private var tokenFile: File = _ + + override def beforeEach(): Unit = { + super.beforeEach() + tokenFile = File.createTempFile("oidc-token-", ".jwt") + tokenFile.deleteOnExit() + writeTokenFile("fake.jwt.token.workload") + } + + override def afterEach(): Unit = { + try { + if (tokenFile != null) tokenFile.delete() + } finally { + super.afterEach() + } + } + + private def writeTokenFile(content: String): Unit = { + Files.writeString(tokenFile.toPath, content) + } + + private def createOidcConf(): SparkConf = { + new SparkConf(loadDefaults = false) + .set(SECURITY_OIDC_ENABLED, true) + .set(SECURITY_OIDC_IDENTITY_TOKEN_FILE, tokenFile.getAbsolutePath) + .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 5000L) + .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 1000L) + } + + private def createUserContext( + principal: String = "test-user", + expiresInSeconds: Long = 300): UserContext = { + val now = Instant.now() + new UserContext( + principal, + "https://issuer.example.com", + "fake.jwt.token", + now, + now.plusSeconds(expiresInSeconds)) + } + + private def createIngestor(ctx: UserContext): TokenIngestor = { + new TokenIngestor { + override def load(): Optional[UserContext] = Optional.of(ctx) + } + } + + private def createFreshExpiryIngestor(expiresInSeconds: Long): TokenIngestor = { + new TokenIngestor { + override def load(): Optional[UserContext] = + Optional.of(createUserContext(expiresInSeconds = expiresInSeconds)) + } + } + + private def createFailingIngestor(): TokenIngestor = { + new TokenIngestor { + override def load(): Optional[UserContext] = Optional.empty() + } + } + + test("OIDC credential delivery via update callback") { + val conf = createOidcConf() + val ctx = createUserContext() + val callbackRef = new AtomicReference[Array[Byte]]() + val callbackVersion = new AtomicLong(0L) + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (version, bytes) => { + callbackVersion.set(version) + callbackRef.set(bytes) + }) + + try { + val (version, initialBytes) = manager.start() + + assert(version == 1L, "Initial version should be 1") + assert(initialBytes != null, "Initial credentials should not be null") + + val credentials = UserCredentialManager.deserializeUserCredentials(initialBytes) + assert(credentials != null, "Deserialized credentials should not be null") + + val fakeCred = credentials.forScheme("fake") + assert(fakeCred.isPresent, "Should have credential for scheme 'fake'") + assert(fakeCred.get().getProperties.get("provider") == "fake", + "Credential should come from FakeCredentialProvider") + assert(!fakeCred.get().isExpired(Instant.now()), + "Credential should not be expired immediately after resolution") + } finally { + manager.stop() + } + } + + test("credential refresh works end-to-end on expiry") { + val conf = createOidcConf() + .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 2000L) + .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 500L) + + val updateCount = new AtomicInteger(0) + val latestVersion = new AtomicLong(0L) + + // Return a fresh UserContext on each load() so each renewal gets a genuinely + // new expiry rather than spinning on an already-expired token. + val manager = new UserCredentialManager( + conf, + createFreshExpiryIngestor(expiresInSeconds = 3), + (version, _) => { + latestVersion.set(version) + updateCount.incrementAndGet() + }) + + try { + manager.start() + assert(updateCount.get() == 1, "Should have exactly 1 update after start()") + + eventually(timeout(15.seconds)) { + assert(updateCount.get() >= 2, + s"Expected at least 2 updates (got ${updateCount.get()}), " + + "indicating credential renewal occurred") + } + + assert(latestVersion.get() >= 2L, + "Version should be at least 2 after renewal") + } finally { + manager.stop() + } + } + + test("per-user identity token produces valid credentials") { + val conf = createOidcConf() + + val userCtx = createUserContext(principal = "alice@corp.example.com") + val callbackRef = new AtomicReference[Array[Byte]]() + + val manager = new UserCredentialManager( + conf, + createIngestor(userCtx), + (_, bytes) => callbackRef.set(bytes)) + + try { + val (_, initialBytes) = manager.start() + + val credentials = UserCredentialManager.deserializeUserCredentials(initialBytes) + val fakeCred = credentials.forScheme("fake") + assert(fakeCred.isPresent, + "Per-user token should produce credentials for scheme 'fake'") + assert(fakeCred.get().getProperties.get("provider") == "fake", + "Per-user credential should come from FakeCredentialProvider") + assert(!fakeCred.get().isExpired(Instant.now()), + "Per-user credential should not be expired") + + // Verify the credential is identical in structure to workload token output + val workloadCtx = createUserContext(principal = "workload-identity") + val workloadManager = new UserCredentialManager( + conf, + createIngestor(workloadCtx), + (_, _) => ()) + try { + val (_, workloadBytes) = workloadManager.start() + val workloadCreds = UserCredentialManager.deserializeUserCredentials(workloadBytes) + val workloadFake = workloadCreds.forScheme("fake") + assert(workloadFake.isPresent) + assert(workloadFake.get().getProperties == fakeCred.get().getProperties, + "Per-user and workload tokens should produce identical credential properties") + } finally { + workloadManager.stop() + } + } finally { + manager.stop() + } + } + + test("UserCredentialManager and HadoopDelegationTokenManager coexist") { + val hadoopConf = new Configuration() + val mockRef = mock(classOf[RpcEndpointRef]) + + val conf = createOidcConf() + .set(DIRECT_CREDENTIAL_PROVIDERS_ENABLED, true) + .set(NETWORK_AUTH_ENABLED, true) + .set(NETWORK_CRYPTO_ENABLED, true) + + val ctx = createUserContext() + val oidcCallbackRef = new AtomicReference[Array[Byte]]() + val oidcVersion = new AtomicLong(0L) + + val oidcManager = new UserCredentialManager( + conf, + createIngestor(ctx), + (version, bytes) => { + oidcVersion.set(version) + oidcCallbackRef.set(bytes) + }) + + val dtManager = new HadoopDelegationTokenManager(conf, hadoopConf, mockRef) + + try { + val (oidcVer, oidcBytes) = oidcManager.start() + assert(oidcVer == 1L) + assert(oidcBytes != null) + + val dtTokens = dtManager.start() + assert(dtTokens != null, "DT manager should produce tokens") + + val oidcCreds = UserCredentialManager.deserializeUserCredentials(oidcBytes) + assert(oidcCreds.forScheme("fake").isPresent, + "OIDC credentials should contain 'fake' scheme") + + // HadoopDelegationTokenManager.start() is synchronous -- verify directly + val captor = ArgumentCaptor.forClass(classOf[Any]) + verify(mockRef).send(captor.capture()) + val msg = captor.getValue.asInstanceOf[UpdateDelegationTokens] + val dtCreds = SparkHadoopUtil.get.deserialize(msg.tokens) + assert(dtCreds.getSecretKey(new Text("test.direct.credential")) != null, + "DT credentials should contain test.direct.credential") + assert(new String(dtCreds.getSecretKey(new Text("test.direct.credential"))) === "test-token", + "DT credential value should match") + + assert(oidcVersion.get() == 1L, "OIDC version should remain at 1") + } finally { + oidcManager.stop() + dtManager.stop() + } + } + + test("OIDC failure does not affect HadoopDelegationTokenManager") { + val hadoopConf = new Configuration() + val mockRef = mock(classOf[RpcEndpointRef]) + + val conf = createOidcConf() + .set(DIRECT_CREDENTIAL_PROVIDERS_ENABLED, true) + .set(NETWORK_AUTH_ENABLED, true) + .set(NETWORK_CRYPTO_ENABLED, true) + + val failingOidcManager = new UserCredentialManager( + conf, + createFailingIngestor(), + (_, _) => ()) + + val dtManager = new HadoopDelegationTokenManager(conf, hadoopConf, mockRef) + + try { + // OIDC start should fail with IllegalStateException (missing token) + val oidcException = intercept[IllegalStateException] { + failingOidcManager.start() + } + assert(oidcException.getMessage.contains( + "identity token file is missing or malformed")) + + // DT manager should still work perfectly despite OIDC failure + val dtTokens = dtManager.start() + assert(dtTokens != null, "DT manager should succeed despite OIDC failure") + + val captor = ArgumentCaptor.forClass(classOf[Any]) + verify(mockRef).send(captor.capture()) + val msg = captor.getValue.asInstanceOf[UpdateDelegationTokens] + val dtCreds = SparkHadoopUtil.get.deserialize(msg.tokens) + assert(dtCreds.getSecretKey(new Text("test.direct.credential")) != null, + "DT credentials should be unaffected by OIDC failure") + } finally { + failingOidcManager.stop() + dtManager.stop() + } + } + + test("DT provider failure does not affect UserCredentialManager") { + val hadoopConf = new Configuration() + val mockRef = mock(classOf[RpcEndpointRef]) + + val conf = createOidcConf() + .set(DIRECT_CREDENTIAL_PROVIDERS_ENABLED, true) + .set(NETWORK_AUTH_ENABLED, true) + .set(NETWORK_CRYPTO_ENABLED, true) + .set("spark.security.credentials.test-direct.enabled", "false") + .set("spark.security.credentials.test-noexpiry.enabled", "false") + + val ctx = createUserContext() + val oidcCallbackRef = new AtomicReference[Array[Byte]]() + + val oidcManager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, bytes) => oidcCallbackRef.set(bytes)) + + val dtManager = new HadoopDelegationTokenManager(conf, hadoopConf, mockRef) + + try { + val dtTokens = dtManager.start() + assert(dtTokens == null, "DT manager should return null when all providers fail") + + val (oidcVer, oidcBytes) = oidcManager.start() + assert(oidcVer == 1L, "OIDC version should be 1") + assert(oidcBytes != null, "OIDC should produce credentials despite DT failure") + + val oidcCreds = UserCredentialManager.deserializeUserCredentials(oidcBytes) + assert(oidcCreds.forScheme("fake").isPresent, + "OIDC credentials should be unaffected by DT failure") + } finally { + oidcManager.stop() + dtManager.stop() + } + } + + test("deserialization idempotency preserves credential content") { + val conf = createOidcConf() + val ctx = createUserContext() + val serializedRef = new AtomicReference[Array[Byte]]() + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, bytes) => serializedRef.set(bytes)) + + try { + manager.start() + + val bytes = serializedRef.get() + assert(bytes != null && bytes.length > 0, "Serialized credentials should be non-empty") + + val creds1 = UserCredentialManager.deserializeUserCredentials(bytes) + val creds2 = UserCredentialManager.deserializeUserCredentials(bytes) + + assert(creds1.forScheme("fake").isPresent) + assert(creds2.forScheme("fake").isPresent) + assert(creds1.forScheme("fake").get().getProperties == + creds2.forScheme("fake").get().getProperties, + "Multiple deserializations of same bytes should produce identical credentials") + + val cred = creds1.forScheme("fake").get() + assert(cred.getProperties.containsKey("provider")) + assert(cred.getExpiresAt != null, "Credential should have an expiry set") + assert(!cred.isExpired(Instant.now()), "Freshly resolved credential should not be expired") + } finally { + manager.stop() + } + } + + test("case-insensitive scheme lookup in credential bundle") { + val conf = createOidcConf() + val ctx = createUserContext() + val serializedRef = new AtomicReference[Array[Byte]]() + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, bytes) => serializedRef.set(bytes)) + + try { + manager.start() + + val creds = UserCredentialManager.deserializeUserCredentials(serializedRef.get()) + + // FakeCredentialProvider declares supportedSchemes = Set("fake", "shared") + // but "shared" is ambiguous (AnotherFakeCredentialProvider also claims it), + // so only "fake" auto-resolves without explicit config. + assert(creds.forScheme("fake").isPresent, "Should resolve 'fake' scheme") + assert(creds.forScheme("FAKE").isPresent, + "Scheme lookup should be case-insensitive") + assert(creds.forScheme("Fake").isPresent, + "Scheme lookup should be case-insensitive") + } finally { + manager.stop() + } + } + + test("stop() after start() completes cleanly without exceptions") { + val conf = createOidcConf() + val ctx = createUserContext(expiresInSeconds = 60) + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, _) => ()) + + manager.start() + manager.stop() + + // Double stop should also be safe + manager.stop() + } + + test("credential version is monotonically increasing across renewals") { + val conf = createOidcConf() + .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 2000L) + .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 500L) + + val versions = new java.util.concurrent.CopyOnWriteArrayList[Long]() + + val manager = new UserCredentialManager( + conf, + createFreshExpiryIngestor(expiresInSeconds = 3), + (version, _) => versions.add(version)) + + try { + manager.start() + + eventually(timeout(15.seconds)) { + assert(versions.size() >= 3, + s"Expected at least 3 callbacks (got ${versions.size()})") + } + + val versionList = new java.util.ArrayList(versions) + for (i <- 1 until versionList.size()) { + assert(versionList.get(i) > versionList.get(i - 1), + s"Version ${versionList.get(i)} should be > ${versionList.get(i - 1)} " + + s"at index $i (full list: $versionList)") + } + } finally { + manager.stop() + } + } + + test("every renewal callback provides non-null non-empty credentials") { + val conf = createOidcConf() + .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 2000L) + .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 500L) + + val allBytes = new java.util.concurrent.CopyOnWriteArrayList[Array[Byte]]() + + val manager = new UserCredentialManager( + conf, + createFreshExpiryIngestor(expiresInSeconds = 3), + (_, bytes) => allBytes.add(bytes)) + + try { + manager.start() + + eventually(timeout(15.seconds)) { + assert(allBytes.size() >= 2, + s"Expected at least 2 callbacks (got ${allBytes.size()})") + } + + val it = allBytes.iterator() + while (it.hasNext) { + val bytes = it.next() + assert(bytes != null, "Callback bytes should never be null") + assert(bytes.length > 0, "Callback bytes should never be empty") + val creds = UserCredentialManager.deserializeUserCredentials(bytes) + assert(creds.forScheme("fake").isPresent, + "Every renewed credential bundle should contain 'fake' scheme") + } + } finally { + manager.stop() + } + } + + test("OIDC disabled does not interfere with DT manager") { + val hadoopConf = new Configuration() + val mockRef = mock(classOf[RpcEndpointRef]) + + val conf = new SparkConf(loadDefaults = false) + .set(SECURITY_OIDC_ENABLED, false) + .set(DIRECT_CREDENTIAL_PROVIDERS_ENABLED, true) + .set(NETWORK_AUTH_ENABLED, true) + .set(NETWORK_CRYPTO_ENABLED, true) + + val oidcManager = UserCredentialManager.create(conf, (_, _) => ()) + assert(oidcManager.isEmpty, "OIDC manager should not be created when disabled") + + val dtManager = new HadoopDelegationTokenManager(conf, hadoopConf, mockRef) + try { + val dtTokens = dtManager.start() + assert(dtTokens != null, "DT manager should work when OIDC is disabled") + + val captor = ArgumentCaptor.forClass(classOf[Any]) + verify(mockRef).send(captor.capture()) + val msg = captor.getValue.asInstanceOf[UpdateDelegationTokens] + val dtCreds = SparkHadoopUtil.get.deserialize(msg.tokens) + assert(dtCreds.getSecretKey(new Text("test.direct.credential")) != null) + } finally { + dtManager.stop() + } + } + + test("credential serialization roundtrip through VersionedCredentials store") { + val conf = createOidcConf() + val ctx = createUserContext() + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, _) => ()) + + try { + val (version, credentialBytes) = manager.start() + assert(version == 1L) + + // Set up the credential store exactly as CoarseGrainedSchedulerBackend does + val store = new AtomicReference[VersionedCredentials]() + VersionedCredentials.updateIfNewer(store, version, credentialBytes) + + // Exercise the EXACT expression from TaskSetManager.scala line 607: + // Option(env.userCredentials.get()).map(vc => (vc.version, vc.bytes)) + // This is what TaskSetManager reads when constructing TaskDescription. + val credentialTuple: Option[(Long, Array[Byte])] = + Option(store.get()).map(vc => (vc.version, vc.bytes)) + + assert(credentialTuple.isDefined, + "Credential store should produce Some when credentials are set") + assert(credentialTuple.get._1 == 1L, + "TaskDescription should carry version 1") + assert(credentialTuple.get._2 === credentialBytes, + "TaskDescription should carry the credential bytes from the store") + + // Verify the bytes are valid credentials end-to-end + val creds = UserCredentialManager.deserializeUserCredentials(credentialTuple.get._2) + assert(creds.forScheme("fake").isPresent, + "Credentials from TaskDescription should contain 'fake' scheme") + + // Update the store to version 2 (simulating a renewal) + val v2Bytes = UserCredentialManager.serializeUserCredentials( + new UserCredentials(java.util.Map.of("fake", + new ServiceCredential(java.util.Map.of("provider", "fake-v2"), + Instant.now().plusSeconds(300))))) + VersionedCredentials.updateIfNewer(store, 2L, v2Bytes) + + // Read again -- same expression as TaskSetManager line 607 + val credentialTupleV2: Option[(Long, Array[Byte])] = + Option(store.get()).map(vc => (vc.version, vc.bytes)) + + assert(credentialTupleV2.get._1 == 2L, + "After renewal, TaskDescription should carry version 2") + + val credsV2 = UserCredentialManager.deserializeUserCredentials(credentialTupleV2.get._2) + assert(credsV2.forScheme("fake").get().getProperties.get("provider") == "fake-v2", + "TaskDescription should carry renewed credentials") + + // Verify empty store produces None (no credentials yet scenario) + val emptyStore = new AtomicReference[VersionedCredentials]() + val emptyTuple: Option[(Long, Array[Byte])] = + Option(emptyStore.get()).map(vc => (vc.version, vc.bytes)) + assert(emptyTuple.isEmpty, + "Empty store should produce None for TaskDescription.userCredentials") + } finally { + manager.stop() + } + } +} diff --git a/core/src/test/scala/org/apache/spark/deploy/security/UserCredentialManagerSuite.scala b/core/src/test/scala/org/apache/spark/deploy/security/UserCredentialManagerSuite.scala index b5ee471013b38..af8e02c11e51f 100644 --- a/core/src/test/scala/org/apache/spark/deploy/security/UserCredentialManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/security/UserCredentialManagerSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark.deploy.security import java.time.Instant import java.util import java.util.Optional +import java.util.concurrent.{CountDownLatch, TimeUnit} import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} import scala.concurrent.duration._ @@ -32,11 +33,6 @@ import org.apache.spark.security._ class UserCredentialManagerSuite extends SparkFunSuite { - override def beforeEach(): Unit = { - super.beforeEach() - CredentialProviderLoader.resetForTesting() - } - private def createSparkConf(): SparkConf = { new SparkConf(loadDefaults = false) .set(SECURITY_OIDC_ENABLED, true) @@ -80,10 +76,10 @@ class UserCredentialManagerSuite extends SparkFunSuite { val manager = new UserCredentialManager( conf, createIngestor(ctx), - bytes => callbackRef.set(bytes)) + (_, bytes) => callbackRef.set(bytes)) try { - val result = manager.start() + val (_, result) = manager.start() assert(result != null, "start() should return serialized credentials") assert(callbackRef.get() != null, "callback should have been invoked") @@ -105,7 +101,7 @@ class UserCredentialManagerSuite extends SparkFunSuite { val manager = new UserCredentialManager( conf, createFailingIngestor(), - _ => ()) + (_: Long, _: Array[Byte]) => ()) try { val ex = intercept[IllegalStateException] { @@ -129,7 +125,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { val original = new UserCredentials(credsMap) val conf = createSparkConf() - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { val serialized = UserCredentialManager.serializeUserCredentials(original) @@ -168,7 +165,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 10000L) // 10s .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 5000L) // 5s - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { // Token expires in 60s, credential expires in 30s // Expected: min(60s, 30s) - 10s = 20s @@ -189,7 +187,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 10000L) .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 5000L) - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { // Token expires in 5s, safetyMargin is 10s -> computed delay would be negative // Should be bounded by minInterval (5s) @@ -208,7 +207,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { .set(SECURITY_OIDC_RENEWAL_SAFETY_MARGIN, 10000L) // 10s .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 5000L) // 5s - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { // Token expires in 60s, no credential expiry // Expected: 60s - 10s = 50s @@ -224,7 +224,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { test("computeRenewalDelay returns default when no expiry information") { val conf = createSparkConf() - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { // UserContext with null expiresAt val ctx = new UserContext( @@ -243,7 +244,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { val conf = createSparkConf() .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 1000L) - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { val failuresField = classOf[UserCredentialManager].getDeclaredField("consecutiveFailures") failuresField.setAccessible(true) @@ -274,7 +276,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { val conf = createSparkConf() .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 1000L) - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { val failuresField = classOf[UserCredentialManager].getDeclaredField("consecutiveFailures") failuresField.setAccessible(true) @@ -293,7 +296,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { val conf = createSparkConf() .set(SECURITY_OIDC_RENEWAL_MIN_INTERVAL, 1000L) - val manager = new UserCredentialManager(conf, createFailingIngestor(), _ => ()) + val manager = new UserCredentialManager( + conf, createFailingIngestor(), (_: Long, _: Array[Byte]) => ()) try { val failuresField = classOf[UserCredentialManager].getDeclaredField("consecutiveFailures") failuresField.setAccessible(true) @@ -314,13 +318,13 @@ class UserCredentialManagerSuite extends SparkFunSuite { val conf = new SparkConf(loadDefaults = false) .set(SECURITY_OIDC_ENABLED, false) - val result = UserCredentialManager.create(conf, _ => ()) + val result = UserCredentialManager.create(conf, (_: Long, _: Array[Byte]) => ()) assert(result.isEmpty) } test("UserCredentialManager.create returns Some when enabled with valid config") { val conf = createSparkConf() - val result = UserCredentialManager.create(conf, _ => ()) + val result = UserCredentialManager.create(conf, (_: Long, _: Array[Byte]) => ()) assert(result.isDefined) } @@ -330,7 +334,7 @@ class UserCredentialManagerSuite extends SparkFunSuite { // Deliberately not setting SECURITY_OIDC_IDENTITY_TOKEN_FILE val ex = intercept[IllegalArgumentException] { - UserCredentialManager.create(conf, _ => ()) + UserCredentialManager.create(conf, (_: Long, _: Array[Byte]) => ()) } assert(ex.getMessage.contains("spark.security.oidc.identityToken.file")) } @@ -345,10 +349,10 @@ class UserCredentialManagerSuite extends SparkFunSuite { val manager = new UserCredentialManager( conf, createIngestor(ctx), - _ => { callbackCount += 1 }) + (_: Long, _: Array[Byte]) => { callbackCount += 1 }) try { - val result = manager.start() + val (_, result) = manager.start() assert(result != null) assert(callbackCount === 1, "callback should be invoked once on start") } finally { @@ -362,7 +366,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { "org.apache.spark.security.FakeCredentialProvider") val ctx = createUserContext() - val manager = new UserCredentialManager(conf, createIngestor(ctx), _ => ()) + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_: Long, _: Array[Byte]) => ()) manager.start() // Should not throw manager.stop() @@ -386,10 +391,10 @@ class UserCredentialManagerSuite extends SparkFunSuite { val manager = new UserCredentialManager( conf, createIngestor(ctx), - bytes => callbackRef.set(bytes)) + (_, bytes) => callbackRef.set(bytes)) try { - val result = manager.start() + val (_, result) = manager.start() assert(result != null, "start() should return serialized credentials") // Verify that "fake" credentials were resolved despite "shared" failing @@ -412,7 +417,8 @@ class UserCredentialManagerSuite extends SparkFunSuite { val ctx = createUserContext() - val manager = new UserCredentialManager(conf, createIngestor(ctx), _ => ()) + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_: Long, _: Array[Byte]) => ()) try { val ex = intercept[IllegalStateException] { @@ -450,14 +456,14 @@ class UserCredentialManagerSuite extends SparkFunSuite { } } - val callbacks = new java.util.concurrent.CopyOnWriteArrayList[Array[Byte]]() + val callbacks = new java.util.concurrent.CopyOnWriteArrayList[(Long, Array[Byte])]() val manager = new UserCredentialManager( conf, rotatingIngestor, - bytes => callbacks.add(bytes)) + (version, bytes) => callbacks.add((version, bytes))) try { - val initial = manager.start() + val (_, initial) = manager.start() assert(initial != null) assert(callbacks.size() === 1, "Should have one callback from start()") @@ -471,8 +477,226 @@ class UserCredentialManagerSuite extends SparkFunSuite { // Verify that the ingestor was called more than once (rotation detected) assert(callCount.get() >= 2, s"TokenIngestor should have been called at least twice, got ${callCount.get()}") + + // Verify version monotonicity: each callback receives a strictly increasing version + val versions = (0 until callbacks.size()).map(i => callbacks.get(i)._1) + assert(versions === versions.sorted, + s"Versions should be monotonically increasing: $versions") + assert(versions.head === 1L, "First version should be 1") + assert(versions(1) === 2L, "Second version should be 2") } finally { manager.stop() } } + + test("stop() closes initialized credential providers") { + val conf = createSparkConf() + val ctx = createUserContext() + val callbackRef = new AtomicReference[Array[Byte]]() + val loader = new CredentialProviderLoader() + + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, bytes) => callbackRef.set(bytes), + loader) + + manager.start() + + // Get the FakeCredentialProvider instance to verify close was called + val providerOpt = loader.providerFor("fake", + new util.HashMap[String, String]()) + assert(providerOpt.isPresent) + val fakeProvider = providerOpt.get().asInstanceOf[FakeCredentialProvider] + assert(fakeProvider.getCloseCount === 0, "close() not yet called before stop()") + + manager.stop() + + assert(fakeProvider.getCloseCount === 1, + "stop() should close initialized providers exactly once") + } + + test("a later manager uses fresh credential providers after stop") { + val conf = createSparkConf() + val ctx = createUserContext() + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + + val firstLoader = new CredentialProviderLoader() + val firstManager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, _) => (), + firstLoader) + firstManager.start() + val firstProvider = firstLoader.providerFor("fake", + new util.HashMap[String, String]()).get().asInstanceOf[FakeCredentialProvider] + firstManager.stop() + assert(firstProvider.getCloseCount === 1) + + val secondLoader = new CredentialProviderLoader() + val secondManager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, _) => (), + secondLoader) + try { + secondManager.start() + val secondProvider = secondLoader.providerFor("fake", + new util.HashMap[String, String]()).get().asInstanceOf[FakeCredentialProvider] + assert(secondProvider ne firstProvider) + assert(secondProvider.getCloseCount === 0) + } finally { + secondManager.stop() + } + } + + test("stop() waits for credential renewal before closing providers") { + val conf = createSparkConf() + val ctx = createUserContext(expiresInSeconds = 6) + val callbackCount = new AtomicInteger() + val renewalStarted = new CountDownLatch(1) + val renewalInterrupted = new CountDownLatch(1) + val releaseRenewal = new CountDownLatch(1) + val renewalCompleted = new CountDownLatch(1) + val loader = new CredentialProviderLoader() + + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + + val manager = new UserCredentialManager( + conf, + createIngestor(ctx), + (_, _) => { + if (callbackCount.incrementAndGet() > 1) { + renewalStarted.countDown() + var released = false + while (!released) { + try { + releaseRenewal.await() + released = true + } catch { + case _: InterruptedException => renewalInterrupted.countDown() + } + } + renewalCompleted.countDown() + } + }, + loader) + + manager.start() + val provider = loader.providerFor("fake", + new util.HashMap[String, String]()).get().asInstanceOf[FakeCredentialProvider] + assert(renewalStarted.await(10, TimeUnit.SECONDS)) + + val stopThread = new Thread(() => manager.stop()) + stopThread.start() + + try { + assert(renewalInterrupted.await(10, TimeUnit.SECONDS)) + assert(stopThread.isAlive, "stop() should wait for credential renewal to finish") + assert(provider.getCloseCount === 0, + "stop() should not close providers while credential renewal is still running") + } finally { + releaseRenewal.countDown() + stopThread.join(10000) + } + + assert(renewalCompleted.await(10, TimeUnit.SECONDS)) + assert(!stopThread.isAlive, "stop() should finish after credential renewal exits") + assert(provider.getCloseCount === 1, + "stop() should close providers after credential renewal exits") + } + + // ========== additionalSparkProperties application ========== + + test("start() applies additionalSparkProperties from active providers") { + val conf = createSparkConf() + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + val ctx = createUserContext() + + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_, _) => ()) + + try { + manager.start() + assert(conf.get("spark.hadoop.fs.fake.credentials.provider") === + "org.apache.spark.security.FakeExecutorCredentialProvider") + } finally { + manager.stop() + } + } + + test("start() does not overwrite user-set properties") { + val conf = createSparkConf() + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + // User explicitly sets the property before start() + conf.set("spark.hadoop.fs.fake.credentials.provider", "user.Custom") + val ctx = createUserContext() + + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_, _) => ()) + + try { + manager.start() + // User-set value must NOT be overwritten + assert(conf.get("spark.hadoop.fs.fake.credentials.provider") === "user.Custom") + } finally { + manager.stop() + } + } + + test("start() handles provider returning null from additionalSparkProperties") { + // AnotherFakeCredentialProvider uses default (empty map), not null. + // This test verifies the defensive null check doesn't crash + // with a provider that inherits the default empty map. + val conf = createSparkConf() + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + val ctx = createUserContext() + + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_, _) => ()) + + try { + // Should not throw + manager.start() + assert(conf.contains("spark.hadoop.fs.fake.credentials.provider")) + } finally { + manager.stop() + } + } + + test("start() succeeds when a provider throws from additionalSparkProperties") { + // AnotherFakeCredentialProvider is configured to throw; FakeCredentialProvider should + // still have its properties applied (exception isolation via NonFatal catch). + val conf = createSparkConf() + conf.set("spark.security.oidc.provider.fake", + "org.apache.spark.security.FakeCredentialProvider") + conf.set("spark.security.oidc.provider.shared", + "org.apache.spark.security.AnotherFakeCredentialProvider") + val ctx = createUserContext() + + AnotherFakeCredentialProvider.throwOnProperties = true + try { + val manager = new UserCredentialManager( + conf, createIngestor(ctx), (_, _) => ()) + try { + // start() must not fail even though AnotherFakeCredentialProvider throws + manager.start() + // FakeCredentialProvider's property must still be applied + assert(conf.get("spark.hadoop.fs.fake.credentials.provider") === + "org.apache.spark.security.FakeExecutorCredentialProvider") + } finally { + manager.stop() + } + } finally { + AnotherFakeCredentialProvider.throwOnProperties = false + } + } } diff --git a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala index a038d8e8613ef..c7307066c78d3 100644 --- a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala @@ -312,7 +312,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // We don't really verify the data, just pass it around. val data = ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)) val taskDescription = new TaskDescription(taskId, 2, "1", "TASK 1000000", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) val serializedTaskDescription = TaskDescription.encode(taskDescription) backend.rpcEnv.setupEndpoint("Executor 1", backend) backend.executor = mock[Executor](CALLS_REAL_METHODS) @@ -437,7 +437,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // Fake tasks with different taskIds. val taskDescriptions = (1 to numTasks).map { taskId => new TaskDescription(taskId, 2, "1", s"TASK $taskId", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) } assert(taskDescriptions.length == numTasks) @@ -531,7 +531,7 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite // Fake tasks with different taskIds. val taskDescriptions = (1 to numTasks).map { taskId => new TaskDescription(taskId, 2, "1", s"TASK $taskId", 19, - 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, data) + 1, JobArtifactSet.emptyJobArtifactSet, new Properties, 1, resourcesAmounts, None, data) } assert(taskDescriptions.length == numTasks) diff --git a/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala b/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala index 4362eb2f816f1..80b0bf69d45cf 100644 --- a/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala +++ b/core/src/test/scala/org/apache/spark/executor/ExecutorSuite.scala @@ -870,6 +870,7 @@ class ExecutorSuite extends SparkFunSuite properties = new Properties, cpus = 1, resources = Map.empty, + None, serializedTask) } diff --git a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala index d87e5c0245421..bc6a9ec213333 100644 --- a/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala +++ b/core/src/test/scala/org/apache/spark/internal/plugin/PluginContainerSuite.scala @@ -290,6 +290,7 @@ class PluginContainerSuite extends SparkFunSuite with LocalSparkContext { case _: TestSparkPluginEvent => // Count down upon receiving the event sent from the plugin during shutdown. countDownLatch.countDown() + case _ => } } }) diff --git a/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala b/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala index 9f0e622b1d515..9a21fa25b35b7 100644 --- a/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.memory import org.scalatest.PrivateMethodTester -import org.apache.spark.SparkConf +import org.apache.spark.{SparkConf, SparkIllegalArgumentException} import org.apache.spark.internal.config._ import org.apache.spark.internal.config.Tests._ import org.apache.spark.storage.TestBlockId @@ -234,10 +234,16 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes // Try using a system memory that's too small val conf2 = conf.clone().set(TEST_MEMORY, reservedMemory / 2) - val exception = intercept[IllegalArgumentException] { + val exception = intercept[SparkIllegalArgumentException] { UnifiedMemoryManager(conf2, numCores = 1) } - assert(exception.getMessage.contains("increase heap size")) + checkError( + exception, + condition = "INVALID_DRIVER_MEMORY.SYSTEM_MEMORY", + parameters = Map( + "systemMemory" -> (reservedMemory / 2).toString, + "minSystemMemory" -> (reservedMemory * 1.5).ceil.toLong.toString, + "config" -> DRIVER_MEMORY.key)) } test("insufficient executor memory") { @@ -253,10 +259,38 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes // Try using an executor memory that's too small val conf2 = conf.clone().set(EXECUTOR_MEMORY.key, (reservedMemory / 2).toString) - val exception = intercept[IllegalArgumentException] { + val exception = intercept[SparkIllegalArgumentException] { UnifiedMemoryManager(conf2, numCores = 1) } - assert(exception.getMessage.contains("increase executor memory")) + checkError( + exception, + condition = "INVALID_EXECUTOR_MEMORY.CONFIG_MEMORY", + parameters = Map( + "executorMemory" -> (reservedMemory / 2).toString, + "minSystemMemory" -> (reservedMemory * 1.5).ceil.toLong.toString, + "config" -> EXECUTOR_MEMORY.key)) + } + + test("SPARK-58513: executor validates executor heap") { + val systemMemory = 400L * 1024 + val reservedMemory = 300L * 1024 + val memoryFraction = 0.8 + val conf = new SparkConf() + .set(MEMORY_FRACTION, memoryFraction) + .set(TEST_MEMORY, systemMemory) + .set(TEST_RESERVED_MEMORY, reservedMemory) + .set(EXECUTOR_MEMORY.key, (500L * 1024).toString) + + val exception = intercept[SparkIllegalArgumentException] { + UnifiedMemoryManager(conf, numCores = 1, isDriver = false) + } + checkError( + exception, + condition = "INVALID_EXECUTOR_MEMORY.SYSTEM_MEMORY", + parameters = Map( + "systemMemory" -> systemMemory.toString, + "minSystemMemory" -> (reservedMemory * 1.5).ceil.toLong.toString, + "config" -> EXECUTOR_MEMORY.key)) } test("execution can evict cached blocks when there are multiple active tasks (SPARK-12155)") { diff --git a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala index 3c295aecf2748..5ef10b6d7e2e0 100644 --- a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala +++ b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala @@ -325,10 +325,12 @@ class RDDSuite extends SparkFunSuite with SharedSparkContext with Eventually { assert(empty.count() === 0) assert(empty.collect().length === 0) - val thrown = intercept[UnsupportedOperationException]{ - empty.reduce(_ + _) - } - assert(thrown.getMessage.contains("empty")) + checkError( + exception = intercept[SparkUnsupportedOperationException] { + empty.reduce(_ + _) + }, + condition = "EMPTY_COLLECTION_NOT_ALLOWED", + parameters = Map.empty[String, String]) val emptyKv = new EmptyRDD[(Int, Int)](sc) val rdd = sc.parallelize(1 to 2, 2).map(x => (x, x)) diff --git a/core/src/test/scala/org/apache/spark/rpc/RpcAddressSuite.scala b/core/src/test/scala/org/apache/spark/rpc/RpcAddressSuite.scala index 9fb08c79420cb..203810b7ca4a8 100644 --- a/core/src/test/scala/org/apache/spark/rpc/RpcAddressSuite.scala +++ b/core/src/test/scala/org/apache/spark/rpc/RpcAddressSuite.scala @@ -80,4 +80,20 @@ class RpcAddressSuite extends SparkFunSuite { val address = RpcAddress("2600::", 1234) assert(address.toSparkURL == "spark://[2600::]:1234") } + + test("SPARK-58719: Normalize hexadecimal IPv6 addresses") { + val expected = RpcAddress("2001:db8::dead:beef", 1234) + Seq( + "2001:db8::dead:beef", + "[2001:db8::dead:beef]", + "2001:0DB8:0000::DEAD:BEEF", + "[2001:0DB8:0000::DEAD:BEEF]").foreach { host => + val address = RpcAddress(host, 1234) + assert(address.host == "[2001:db8::dead:beef]") + assert(address.hostPort == "[2001:db8::dead:beef]:1234") + assert(address.toSparkURL == "spark://[2001:db8::dead:beef]:1234") + assert(address == expected) + assert(RpcAddress.fromSparkURL(address.toSparkURL) == address) + } + } } diff --git a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala index 99b757496379d..ecdd02585a123 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala @@ -18,7 +18,8 @@ package org.apache.spark.scheduler import java.util.Properties -import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.{Callable, CountDownLatch, FutureTask, LinkedBlockingQueue, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} import scala.collection.mutable import scala.concurrent.Future @@ -348,7 +349,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 1 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, taskResources, bytebuffer))) + taskCpus, taskResources, None, bytebuffer))) val ts = backend.getTaskSchedulerImpl() when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) @@ -455,7 +456,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 1 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, taskResources, bytebuffer))) + taskCpus, taskResources, None, bytebuffer))) val ts = backend.getTaskSchedulerImpl() when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) @@ -548,7 +549,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val taskCpus = 2 val taskDescs: Seq[Seq[TaskDescription]] = Seq(Seq(new TaskDescription(1, 0, "1", "t1", 0, 1, JobArtifactSet.emptyJobArtifactSet, new Properties(), - taskCpus, Map.empty, bytebuffer))) + taskCpus, Map.empty, None, bytebuffer))) when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(taskDescs) backend.driverEndpoint.send(ReviveOffers) @@ -580,6 +581,153 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo } } + Seq(false, true).foreach { isBarrier => + test("SPARK-58879: idle decommission rejects tasks assigned before LaunchTask " + + s"(barrier=$isBarrier)") { + val numTasks = if (isBarrier) 2 else 1 + val conf = new SparkConf().set(EXECUTOR_CORES, 2) + val backend = createDecommissionBackend(conf) + val scheduler = backend.taskScheduler + val executor = registerDecommissionExecutor(backend, "1", 2) + val decommissionCalled = new AtomicBoolean(false) + backend.beforeDecommission = (_, _, _) => decommissionCalled.set(true) + val launchEntered = new CountDownLatch(1) + val allowLaunch = new CountDownLatch(1) + executor.beforeLaunch = _ => { + launchEntered.countDown() + require(allowLaunch.await(30, TimeUnit.SECONDS), "LaunchTask was not released") + } + val taskSet = if (isBarrier) { + FakeTask.createBarrierTaskSet(numTasks) + } else { + FakeTask.createTaskSet(numTasks) + } + scheduler.submitTasks(taskSet) + backend.driverEndpoint.send(ReviveOffers) + + var requestThread: Thread = null + try { + assert(launchEntered.await(10, TimeUnit.SECONDS)) + assert(scheduler.runningTasksByExecutors("1") === numTasks) + val (thread, request) = startDecommissionRequest { + backend.decommissionExecutorsIfIdle( + Array("1" -> ExecutorDecommissionInfo("idle timeout")), false) + } + requestThread = thread + assert(request.get(10, TimeUnit.SECONDS).isEmpty) + assert(backend.isExecutorActive("1")) + assert(!executor.decommissionReceived) + assert(!decommissionCalled.get()) + } finally { + allowLaunch.countDown() + if (requestThread != null) { + requestThread.join(TimeUnit.SECONDS.toMillis(10)) + assert(!requestThread.isAlive) + } + } + + val tasks = (0 until numTasks).map(_ => executor.nextTask()) + flushDecommissionBackend(backend) + tasks.foreach(completeDecommissionTestTask(backend, _)) + assert(!scheduler.isExecutorBusy("1")) + assert(backend.getExecutorAvailableCpus("1").contains(BigDecimal(2))) + assert(backend.decommissionExecutorsIfIdle( + Array("1" -> ExecutorDecommissionInfo("idle timeout")), false) === Seq("1")) + assert(executor.decommissionReceived) + } + } + + test("SPARK-58879: idle decommission fences an executor before concurrent resource offers") { + val backend = createDecommissionBackend() + val retired = registerDecommissionExecutor(backend, "1") + val survivor = registerDecommissionExecutor(backend, "2") + backend.taskScheduler.submitTasks(FakeTask.createTaskSet(1)) + val admissionEntered = new CountDownLatch(1) + val allowAdmission = new CountDownLatch(1) + val offerEntered = new CountDownLatch(1) + val locksHeld = new AtomicBoolean(false) + val decommissionReleased = new AtomicBoolean(false) + backend.beforeDecommission = (_, _, _) => { + locksHeld.set(Thread.holdsLock(backend.taskScheduler) && Thread.holdsLock(backend)) + admissionEntered.countDown() + decommissionReleased.set(allowAdmission.await(30, TimeUnit.SECONDS)) + } + backend.beforeOffers = () => offerEntered.countDown() + + val (requestThread, request) = startDecommissionRequest { + backend.decommissionExecutorsIfIdle( + Array("1" -> ExecutorDecommissionInfo("idle timeout")), false) + } + try { + assert(admissionEntered.await(10, TimeUnit.SECONDS)) + backend.driverEndpoint.send(ReviveOffers) + assert(offerEntered.await(10, TimeUnit.SECONDS)) + } finally { + allowAdmission.countDown() + requestThread.join(TimeUnit.SECONDS.toMillis(10)) + assert(!requestThread.isAlive) + } + assert(request.get(10, TimeUnit.SECONDS) === Seq("1")) + assert(locksHeld.get()) + assert(decommissionReleased.get()) + flushDecommissionBackend(backend) + assert(retired.launchedTasks.isEmpty) + assert(retired.decommissionReceived) + val task = survivor.nextTask() + assert(task.executorId === "2") + completeDecommissionTestTask(backend, task) + } + + test("SPARK-58879: idle decommission skips duplicates and does not replay rejected requests") { + // Retention defaults to zero. Enable it so an incorrectly queued request is observable. + val conf = new SparkConf().set(SCHEDULER_MAX_RETAINED_UNKNOWN_EXECUTORS, 1) + val backend = createDecommissionBackend(conf) + val first = registerDecommissionExecutor(backend, "1") + val second = registerDecommissionExecutor(backend, "2") + val info = ExecutorDecommissionInfo("idle timeout") + val requests = mutable.ArrayBuffer.empty[(Seq[String], Boolean, Boolean)] + backend.beforeDecommission = (ids, adjustTarget, triggeredByExecutor) => { + requests += ((ids, adjustTarget, triggeredByExecutor)) + } + + assert(backend.decommissionExecutorsIfIdle( + Array("1" -> info, "1" -> info, "3" -> info), false) === Seq("1")) + assert(backend.decommissionExecutorsIfIdle(Array("1" -> info), false).isEmpty) + assert(requests.toSeq === Seq((Seq("1"), false, false))) + assert(first.decommissionReceived) + assert(!second.decommissionReceived) + assert(!backend.hasUnknownDecommission("1")) + assert(!backend.hasUnknownDecommission("3")) + + val third = registerDecommissionExecutor(backend, "3") + assert(!third.decommissionReceived) + assert(backend.isExecutorActive("3")) + + assert(backend.decommissionExecutors(Array("4" -> info), false, false).isEmpty) + assert(backend.hasUnknownDecommission("4")) + val fourth = registerDecommissionExecutor(backend, "4") + assert(fourth.decommissionReceived) + assert(!backend.isExecutorActive("4")) + assert(!backend.hasUnknownDecommission("4")) + } + + test("SPARK-58879: forced decommission still accepts a busy executor") { + val backend = createDecommissionBackend() + val executor = registerDecommissionExecutor(backend, "1") + backend.taskScheduler.submitTasks(FakeTask.createTaskSet(1)) + backend.driverEndpoint.send(ReviveOffers) + val task = executor.nextTask() + assert(backend.taskScheduler.isExecutorBusy("1")) + + assert(backend.decommissionExecutors( + Array("1" -> ExecutorDecommissionInfo("host drain")), false, false) === Seq("1")) + assert(executor.decommissionReceived) + assert(!backend.isExecutorActive("1")) + completeDecommissionTestTask(backend, task) + assert(!backend.taskScheduler.isExecutorBusy("1")) + assert(executor.launchedTasks.isEmpty) + } + test("SPARK-41766: New registered executor should receive decommission request" + " sent before registration") { val conf = new SparkConf() @@ -606,6 +754,325 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo assert(mockEndpointRef.decommissionReceived) } + test("SPARK-58886: requestExecutors should saturate instead of overflowing a huge" + + " requested total") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + sc.requestTotalExecutors(Int.MaxValue - 1, 0, Map.empty) + backend.requestExecutors(2) + + val defaultProf = sc.resourceProfileManager.defaultResourceProfile + assert(backend.getRequestedTotalExecutors()(defaultProf) === Int.MaxValue) + + // Only the applied increase (1, not the requested 2) may be recorded as a pending + // request time. Shrink the total to 1 to consume the huge seed entry, leaving just + // that increment: exactly one of the two executors registered below should get a + // request time. + sc.requestTotalExecutors(1, 0, Map.empty) + + val infos = mutable.ArrayBuffer[ExecutorInfo]() + sc.addSparkListener(new SparkListener() { + override def onExecutorAdded(executorAdded: SparkListenerExecutorAdded): Unit = { + infos += executorAdded.executorInfo + } + }) + val mockAddress = mock[RpcAddress] + Seq("1", "2").foreach { id => + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor(id, new MockExecutorRpcEndpointRef(conf), mockAddress.host, 1, + Map(), Map(), Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + } + sc.listenerBus.waitUntilEmpty(executorUpTimeout.toMillis) + assert(infos.size === 2) + assert(infos.head.requestTime.isDefined) + assert(infos.last.requestTime.isEmpty) + } + + test("SPARK-58828: New registered executor should be decommissioned while held") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + val mockEndpointRef = new MockExecutorRpcEndpointRef(conf) + val mockAddress = mock[RpcAddress] + + backend.setExecutorsHeld(true) + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map(), Map(), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + sc.listenerBus.waitUntilEmpty(executorUpTimeout.toMillis) + assert(mockEndpointRef.decommissionReceived) + // The zero requirement is re-published without touching the requested totals + assert(backend.getRequestedTotalExecutors().isEmpty) + } + + test("SPARK-58828: reset clears the explicit request record with the requested totals") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + val defaultProf = sc.resourceProfileManager.defaultResourceProfile + sc.requestTotalExecutors(2, 0, Map.empty) + assert(backend.hasExplicitExecutorRequests) + + // While held, reset() preserves the requested totals and the explicit record + backend.setExecutorsHeld(true) + backend.reset() + assert(backend.hasExplicitExecutorRequests) + assert(backend.getRequestedTotalExecutors().getOrElse(defaultProf, -1) === 2) + + // Not held, reset() clears both + backend.setExecutorsHeld(false) + backend.reset() + assert(!backend.hasExplicitExecutorRequests) + assert(backend.getRequestedTotalExecutors().isEmpty) + } + + test("SPARK-58828: the restore sentinel is published without being recorded") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + val defaultProf = sc.resourceProfileManager.defaultResourceProfile + assert(backend.publishTotalsWithoutRecording(Map(defaultProf -> Int.MaxValue))) + // Neither the totals nor the explicit-request record change, so killExecutors' empty-map + // seeding and a later hold's kill-seeded restore keep their pre-hold behavior + assert(backend.getRequestedTotalExecutors().isEmpty) + assert(!backend.hasExplicitExecutorRequests) + } + + test("SPARK-58828: executor requests made while held are retained until resume") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + val defaultProf = sc.resourceProfileManager.defaultResourceProfile + assert(!backend.hasExplicitExecutorRequests) + + backend.setExecutorsHeld(true) + // Both public request APIs keep recording the requested totals while held + sc.requestTotalExecutors(3, 0, Map.empty) + assert(backend.getRequestedTotalExecutors().getOrElse(defaultProf, -1) === 3) + sc.requestExecutors(2) + assert(backend.getRequestedTotalExecutors().getOrElse(defaultProf, -1) === 5) + assert(backend.hasExplicitExecutorRequests) + + // A held reassertion publishes zero but does not overwrite the requested totals + backend.reassertHeldRequirement() + assert(backend.getRequestedTotalExecutors().getOrElse(defaultProf, -1) === 5) + + // Resume republishes the requested totals as-is, and a stale reassertion after the hold + // is lifted leaves them alone + backend.setExecutorsHeld(false) + assert(backend.republishRequestedTotals()) + backend.reassertHeldRequirement() + assert(backend.getRequestedTotalExecutors().getOrElse(defaultProf, -1) === 5) + } + + test("UpdateUserCredentials is broadcast to all registered executors") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + // Register two mock executors + val mockEndpointRef1 = new MockExecutorRpcEndpointRef(conf) + val mockEndpointRef2 = new MockExecutorRpcEndpointRef(conf) + val mockAddress = mock[RpcAddress] + + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1", mockEndpointRef1, mockAddress.host, 1, Map(), Map(), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("2", mockEndpointRef2, mockAddress.host, 1, Map(), Map(), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + sc.listenerBus.waitUntilEmpty(executorUpTimeout.toMillis) + + // Neither executor should have received credentials yet + assert(mockEndpointRef1.receivedUserCredentials.isEmpty) + assert(mockEndpointRef2.receivedUserCredentials.isEmpty) + + // Send UpdateUserCredentials via DriverEndpoint + val testCredentials = Array[Byte](1, 2, 3, 4, 5) + backend.driverEndpoint.send(UpdateUserCredentials(1L, testCredentials)) + + // Wait for the message to be processed + eventually(timeout(5 seconds)) { + assert(mockEndpointRef1.receivedUserCredentials.isDefined) + assert(mockEndpointRef2.receivedUserCredentials.isDefined) + } + + assert(mockEndpointRef1.receivedUserCredentials.get._2 === testCredentials) + assert(mockEndpointRef2.receivedUserCredentials.get._2 === testCredentials) + + // Verify SparkEnv credential store is also updated + assert(SparkEnv.get.userCredentials.get().bytes === testCredentials) + } + + test("SparkAppConfig includes current user credentials for late-registering executors") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + // Simulate credential acquisition by setting credentials before any executor registers + val testCredentials = Array[Byte](10, 20, 30, 40, 50) + backend.driverEndpoint.send(UpdateUserCredentials(1L, testCredentials)) + + // Wait for the message to be processed + eventually(timeout(5 seconds)) { + assert(SparkEnv.get.userCredentials.get() != null) + } + + // Now retrieve SparkAppConfig as a late-registering executor would + val appConfig = backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfig(ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + // Verify that userCredentials is present in the response + assert(appConfig.userCredentials.isDefined, + "SparkAppConfig should include user credentials for late-registering executors") + assert(appConfig.userCredentials.get._2 === testCredentials) + } + + test("version guard prevents stale credentials from overwriting newer ones") { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + + sc = new SparkContext(conf) + val backend = sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend] + + // Send version 3 credentials first + val credsV3 = Array[Byte](30, 30, 30) + backend.driverEndpoint.send(UpdateUserCredentials(3L, credsV3)) + + eventually(timeout(5 seconds)) { + assert(SparkEnv.get.userCredentials.get() != null) + assert(SparkEnv.get.userCredentials.get().version === 3L) + } + + // Now send version 1 (stale) -- should be rejected + val credsV1 = Array[Byte](10, 10, 10) + backend.driverEndpoint.send(UpdateUserCredentials(1L, credsV1)) + + // Flush the DriverEndpoint mailbox by sending a synchronous request. + // Since DriverEndpoint is single-threaded, when this returns, v1 has been processed. + backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfig(ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + + // Store should still hold v3 (stale v1 was rejected by version guard) + assert(SparkEnv.get.userCredentials.get().version === 3L, + "Stale version 1 should not overwrite newer version 3") + + // Send version 5 (newer) -- should be accepted + val credsV5 = Array[Byte](50, 50, 50) + backend.driverEndpoint.send(UpdateUserCredentials(5L, credsV5)) + + eventually(timeout(5 seconds)) { + assert(SparkEnv.get.userCredentials.get().version === 5L) + } + + // Verify version 5 credentials are in the store (not v1 or v3) + assert(SparkEnv.get.userCredentials.get().bytes === credsV5) + } + + test("executor-side credential store version guard rejects stale and accepts newer") { + // This tests VersionedCredentials.updateIfNewer -- the same method used in + // CoarseGrainedExecutorBackend.receive and Executor.TaskRunner. + val store = new AtomicReference[VersionedCredentials]() + + // Initial write to null store should succeed + VersionedCredentials.updateIfNewer(store, 2L, Array[Byte](20, 20)) + assert(store.get().version === 2L) + assert(store.get().bytes === Array[Byte](20, 20)) + + // Stale version (1) should be rejected + VersionedCredentials.updateIfNewer(store, 1L, Array[Byte](10, 10)) + assert(store.get().version === 2L, "Stale version should not overwrite newer") + + // Same version (2) should also be rejected (strict >) + VersionedCredentials.updateIfNewer(store, 2L, Array[Byte](22, 22)) + assert(store.get().version === 2L) + assert(store.get().bytes === Array[Byte](20, 20), "Same version should not overwrite") + + // Newer version (5) should be accepted + VersionedCredentials.updateIfNewer(store, 5L, Array[Byte](50, 50)) + assert(store.get().version === 5L) + assert(store.get().bytes === Array[Byte](50, 50)) + } + + private def createDecommissionBackend( + conf: SparkConf = new SparkConf()): DecommissionTestSchedulerBackend = { + conf.setMaster(s"coarseclustermanager[${classOf[DecommissionTestSchedulerBackend].getName}]") + .setAppName("idle decommission test") + .set(EXECUTOR_INSTANCES, 0) + .set(DYN_ALLOCATION_ENABLED, false) + .set(DECOMMISSION_ENABLED, true) + sc = new SparkContext(conf) + sc.schedulerBackend.asInstanceOf[DecommissionTestSchedulerBackend] + } + + private def registerDecommissionExecutor( + backend: DecommissionTestSchedulerBackend, + executorId: String, + cores: Int = 1) + : DecommissionTestExecutorRpcEndpointRef = { + val executor = new DecommissionTestExecutorRpcEndpointRef(sc.conf, executorId) + assert(backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor(executorId, executor, "localhost", cores, Map.empty, Map.empty, + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID))) + backend.driverEndpoint.send(LaunchedExecutor(executorId)) + flushDecommissionBackend(backend) + executor + } + + private def flushDecommissionBackend(backend: DecommissionTestSchedulerBackend): Unit = { + // Any synchronous request flushes earlier driver-endpoint messages. Ignore its result: + // executor "1" may not be registered or may already be retired. + backend.driverEndpoint.askSync[Boolean](IsExecutorAlive("1")) + } + + private def completeDecommissionTestTask( + backend: DecommissionTestSchedulerBackend, + task: TaskDescription): Unit = { + val result = new DirectTaskResult[Int]( + sc.env.serializer.newInstance().serialize(0), Seq.empty, Array.emptyLongArray) + val serializedResult = sc.env.closureSerializer.newInstance().serialize(result) + backend.driverEndpoint.send(StatusUpdate( + task.executorId, task.taskId, TaskState.FINISHED, new SerializableBuffer(serializedResult), + task.cpus, task.resources)) + flushDecommissionBackend(backend) + } + + private def startDecommissionRequest[T](body: => T): (Thread, FutureTask[T]) = { + val request = new FutureTask[T](new Callable[T] { + override def call(): T = body + }) + val thread = new Thread(request, "idle-decommission-test") + thread.setDaemon(true) + thread.start() + (thread, request) + } + private def testSubmitJob(sc: SparkContext, rdd: RDD[Int]): Unit = { sc.submitJob( rdd, @@ -617,7 +1084,7 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo } } -/** Simple cluster manager that wires up our mock backend for the resource tests. */ +/** Cluster manager for the mock resource tests and real-scheduler decommission tests. */ private class CSMockExternalClusterManager extends ExternalClusterManager { private var ts: TaskSchedulerImpl = _ @@ -628,12 +1095,18 @@ private class CSMockExternalClusterManager extends ExternalClusterManager { override def createTaskScheduler( sc: SparkContext, masterURL: String): TaskScheduler = { - ts = mock[TaskSchedulerImpl] - when(ts.sc).thenReturn(sc) - when(ts.applicationId()).thenReturn("appid1") - when(ts.applicationAttemptId()).thenReturn(Some("attempt1")) - when(ts.schedulingMode).thenReturn(SchedulingMode.FIFO) - when(ts.excludedNodes()).thenReturn(Set.empty[String]) + masterURL match { + case MOCK_REGEX(backendClassName) + if backendClassName == classOf[DecommissionTestSchedulerBackend].getName => + ts = new TaskSchedulerImpl(sc, sc.conf.get(TASK_MAX_FAILURES)) + case _ => + ts = mock[TaskSchedulerImpl] + when(ts.sc).thenReturn(sc) + when(ts.applicationId()).thenReturn("appid1") + when(ts.applicationAttemptId()).thenReturn(Some("attempt1")) + when(ts.schedulingMode).thenReturn(SchedulingMode.FIFO) + when(ts.excludedNodes()).thenReturn(Set.empty[String]) + } ts } @@ -661,18 +1134,90 @@ class TestCoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, override v def getTaskSchedulerImpl(): TaskSchedulerImpl = scheduler } +private[spark] class DecommissionTestSchedulerBackend( + scheduler: TaskSchedulerImpl, + override val rpcEnv: RpcEnv) + extends CoarseGrainedSchedulerBackend(scheduler, rpcEnv) { + + val taskScheduler = scheduler + @volatile var beforeOffers: () => Unit = () => () + @volatile var beforeDecommission: (Seq[String], Boolean, Boolean) => Unit = (_, _, _) => () + + // Tests drive the real offer paths explicitly, without periodic or scheduler-triggered offers. + override protected def createDriverEndpoint(): DriverEndpoint = new DriverEndpoint { + override def onStart(): Unit = {} + + override def receive: PartialFunction[Any, Unit] = { + case ReviveOffers => + beforeOffers() + super.receive(ReviveOffers) + case message => super.receive(message) + } + } + + override def reviveOffers(): Unit = {} + + override def decommissionExecutors( + executorsAndDecomInfo: Array[(String, ExecutorDecommissionInfo)], + adjustTargetNumExecutors: Boolean, + triggeredByExecutor: Boolean): Seq[String] = { + beforeDecommission( + executorsAndDecomInfo.map(_._1).toSeq, adjustTargetNumExecutors, triggeredByExecutor) + super.decommissionExecutors( + executorsAndDecomInfo, adjustTargetNumExecutors, triggeredByExecutor) + } + + def hasUnknownDecommission(executorId: String): Boolean = synchronized { + unknownExecutorsPendingDecommission.getIfPresent(executorId) != null + } +} + +private[spark] class DecommissionTestExecutorRpcEndpointRef( + conf: SparkConf, + executorId: String) extends RpcEndpointRef(conf) { + + val launchedTasks = new LinkedBlockingQueue[TaskDescription]() + @volatile var decommissionReceived = false + @volatile var beforeLaunch: TaskDescription => Unit = (_: TaskDescription) => () + + override def address: RpcAddress = RpcAddress("localhost", 10000 + executorId.toInt) + override def name: String = s"executor-$executorId" + + override def send(message: Any): Unit = message match { + case LaunchTask(data) => + val task = TaskDescription.decode(data.value) + beforeLaunch(task) + launchedTasks.add(task) + case DecommissionExecutor => decommissionReceived = true + case _ => + } + + override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = { + Future.successful(true.asInstanceOf[T]) + } + + def nextTask(): TaskDescription = { + val task = launchedTasks.poll(10, TimeUnit.SECONDS) + require(task != null, s"No task was launched on executor $executorId") + task + } +} + private[spark] class MockExecutorRpcEndpointRef(conf: SparkConf) extends RpcEndpointRef(conf) { // scalastyle:off executioncontextglobal import scala.concurrent.ExecutionContext.Implicits.global // scalastyle:on executioncontextglobal - var decommissionReceived = false + @volatile var decommissionReceived = false + @volatile var receivedUserCredentials: Option[(Long, Array[Byte])] = None override def address: RpcAddress = null override def name: String = "executor" override def send(message: Any): Unit = message match { case DecommissionExecutor => decommissionReceived = true + case UpdateUserCredentials(version, creds) => receivedUserCredentials = Some((version, creds)) + case _ => } override def ask[T: ClassTag](message: Any, timeout: RpcTimeout): Future[T] = { Future{true.asInstanceOf[T]} diff --git a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala index f43354c808652..a44ed77b65113 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala @@ -23,6 +23,7 @@ import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, At import scala.annotation.meta.param import scala.collection.mutable.{ArrayBuffer, HashMap, HashSet, Map} +import scala.concurrent.Promise import scala.jdk.CollectionConverters._ import scala.language.reflectiveCalls import scala.util.control.NonFatal @@ -41,7 +42,7 @@ import org.apache.spark.executor.ExecutorMetrics import org.apache.spark.internal.config import org.apache.spark.internal.config.{LEGACY_ABORT_STAGE_AFTER_KILL_TASKS, Tests} import org.apache.spark.network.shuffle.ExternalBlockStoreClient -import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.rdd.{DeterministicLevel, RDD, ReliableRDDCheckpointData} import org.apache.spark.resource.{ExecutorResourceRequests, ResourceProfile, ResourceProfileBuilder, TaskResourceProfile, TaskResourceRequests} import org.apache.spark.resource.ResourceUtils.{FPGA, GPU} import org.apache.spark.rpc.RpcTimeoutException @@ -394,6 +395,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti override protected def outstandingTasksForOtherWork(rpId: Int, excludeStageIds: Set[Int]): Int = outstandingTasksForOtherWorkForTest(rpId, excludeStageIds) + // Seam for the hold state (SPARK-58828) read by the barrier retry-budget freeze. Default + // false so existing tests are unaffected. + @volatile var executorsHeldForTest: Boolean = false + override protected def executorsHeld: Boolean = executorsHeldForTest + /** * Schedules shuffle merge finalize. */ @@ -1559,6 +1565,149 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("SPARK-58887: a barrier job can be cancelled while its slot check is being retried") { + // 3 barrier tasks on the local[2] backend fail the max concurrent tasks check, so the + // submission enters the retry window during which the job is registered nowhere but + // `deferredBarrierJobs`. A cancellation in that window must fail the job immediately + // instead of silently no-oping until the retries run out. + val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener) + assert(failure.get() === null, "a barrier job in its slot-check retry window must wait") + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + assert(scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId)) + cancel(jobId) + assert(failure.get() !== null, "cancelling a deferred barrier job must fail its listener") + assert(failure.get().getMessage.contains(s"Job $jobId cancelled"), + "the listener must fail with the cancellation error, not the exhausted-retries one") + assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId), + "a cancelled deferral must not leak its slot-check failure count") + // The pending re-post fires regardless of the cancellation; simulate its arrival and check + // it is dropped instead of resurrecting the cancelled job. + runEvent(JobSubmitted(jobId, barrierRdd, jobComputeFunc, Array(0, 1, 2), CallSite("", ""), + failListener, JobArtifactSet.getActiveOrDefault(sc), null)) + assert(scheduler.deferredBarrierJobs.isEmpty, + "a re-post arriving after the cancellation must be dropped, not re-deferred") + assert(scheduler.barrierJobIdToNumTasksCheckFailures.isEmpty) + assertDataStructuresEmpty() + } + + test("SPARK-58887: cancelAllJobs also fails a barrier job deferred for a slot-check retry") { + val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener) + assert(failure.get() === null) + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + runEvent(AllJobsCancelled()) + assert(failure.get() !== null, "cancelAllJobs must fail a deferred barrier job's listener") + assert(failure.get().getMessage.contains(s"Job $jobId cancelled")) + assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId)) + assertDataStructuresEmpty() + } + + test("SPARK-58887: cancelJobGroup also fails a barrier job deferred for a slot-check retry") { + val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val props = new Properties() + props.setProperty(SparkContext.SPARK_JOB_GROUP_ID, "deferredGroup") + val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props) + assert(failure.get() === null) + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + runEvent(JobGroupCancelled("deferredGroup", cancelFutureJobs = false, None)) + assert(failure.get() !== null, "cancelJobGroup must fail a deferred barrier job's listener") + assert(failure.get().getMessage.contains(s"Job $jobId cancelled")) + assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId)) + assertDataStructuresEmpty() + } + + test("SPARK-58887: cancelJobsWithTag also fails a deferred barrier job") { + val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val props = new Properties() + props.setProperty(SparkContext.SPARK_JOB_TAGS, "deferredTag") + val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props) + assert(failure.get() === null) + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + runEvent(JobTagCancelled("deferredTag", None, None)) + assert(failure.get() !== null, + "cancelJobsWithTag must fail a deferred barrier job's listener") + assert(failure.get().getMessage.contains(s"Job $jobId cancelled")) + assert(!scheduler.barrierJobIdToNumTasksCheckFailures.containsKey(jobId)) + assertDataStructuresEmpty() + } + + test("SPARK-58887: tag cancellation reports a deferred barrier job in its promise") { + val barrierRdd = new MyRDD(sc, 3, Nil).barrier().mapPartitions(iter => iter) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val props = new Properties() + props.setProperty(SparkContext.SPARK_JOB_TAGS, "reportedTag") + val jobId = submit(barrierRdd, Array(0, 1, 2), listener = failListener, properties = props) + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + val cancelledJobs = Promise[Seq[CancelledJobInfo]]() + runEvent(JobTagCancelled("reportedTag", None, Some(cancelledJobs))) + assert(failure.get() !== null) + // The handler completes the promise synchronously; the deferred job must be reported with + // its submission-time properties (e.g. for SQL execution id extraction), even though it + // never had an ActiveJob. + val reported = cancelledJobs.future.value.get.get + assert(reported.map(_.jobId) === Seq(jobId)) + assert(reported.head.properties.getProperty(SparkContext.SPARK_JOB_TAGS) === "reportedTag") + assertDataStructuresEmpty() + } + + test("SPARK-58887: cancelling a deferred barrier job drops its partial stage registrations") { + // An ordinary shuffle upstream of the barrier stage is created and registered before the + // barrier slot check throws, so the deferred job is NOT registered nowhere. Cancelling it + // must drop those registrations too, or a later cancellation of the same job id finds + // jobIdToStageIds populated without an ActiveJob and crashes the event loop. + val ordinaryRdd = new MyRDD(sc, 2, Nil) + val ordinaryDep = new ShuffleDependency(ordinaryRdd, new HashPartitioner(3)) + val barrierRdd = new MyRDD(sc, 3, List(ordinaryDep), tracker = mapOutputTracker) + .barrier().mapPartitions(iter => iter) + val barrierDep = new ShuffleDependency(barrierRdd, new HashPartitioner(2)) + val resultRdd = new MyRDD(sc, 2, List(barrierDep), tracker = mapOutputTracker) + val failure = new AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + val jobId = submit(resultRdd, Array(0, 1), listener = failListener) + assert(failure.get() === null) + assert(scheduler.deferredBarrierJobs.containsKey(jobId)) + assert(scheduler.jobIdToStageIds.contains(jobId), + "the ordinary ancestor stage must have been registered before the slot check threw") + cancel(jobId) + assert(failure.get() !== null) + assert(!scheduler.jobIdToStageIds.contains(jobId), + "cancelling a deferred job must drop its partial stage registrations") + // The pending re-post arrives and is dropped; cancelling once more must then be a harmless + // no-op instead of tripping over leftover registrations without an ActiveJob. + runEvent(JobSubmitted(jobId, resultRdd, jobComputeFunc, Array(0, 1), CallSite("", ""), + failListener, JobArtifactSet.getActiveOrDefault(sc), null)) + cancel(jobId) + assertDataStructuresEmpty() + } + test("Fail the job if a barrier ResultTask failed") { val shuffleMapRdd = new MyRDD(sc, 2, Nil) val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(2)) @@ -2334,6 +2483,34 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assert(scheduler.waitingStages.isEmpty) } + test("SPARK-58616: cancelAllJobs surfaces the supplied reason in the job failure") { + val jobId = submit(new MyRDD(sc, 1, Nil), Array(0)) + assert(scheduler.runningStages.size === 1) + runEvent(AllJobsCancelled(Some("because the user requested cancellation"))) + checkError( + exception = failure.asInstanceOf[SparkException], + condition = "SPARK_JOB_CANCELLED", + sqlState = "XXKDA", + parameters = scala.collection.immutable.Map( + "jobId" -> jobId.toString, + "reason" -> "because the user requested cancellation")) + assertDataStructuresEmpty() + } + + test("SPARK-58616: cancelAllJobs falls back to the default reason when none is supplied") { + val jobId = submit(new MyRDD(sc, 1, Nil), Array(0)) + assert(scheduler.runningStages.size === 1) + runEvent(AllJobsCancelled()) + checkError( + exception = failure.asInstanceOf[SparkException], + condition = "SPARK_JOB_CANCELLED", + sqlState = "XXKDA", + parameters = scala.collection.immutable.Map( + "jobId" -> jobId.toString, + "reason" -> "as part of cancellation of all jobs")) + assertDataStructuresEmpty() + } + test("misbehaved accumulator should not crash DAGScheduler and SparkContext") { val acc = new LongAccumulator { override def add(v: java.lang.Long): Unit = throw new DAGSchedulerSuiteDummyException @@ -6432,6 +6609,31 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: a regular-shuffle prefix feeding a pipelined producer is rejected") { + // Also mixed: a regular shuffle in the PREFIX that feeds a pipelined producer + // (regularRoot --regular--> producer(pipelined) --pipelined--> consumer). Rather than treat + // the regular edge as an ordinary external input to the group and support a mid-DAG regular + // prefix, a pipelined job rejects ANY regular shuffle up front (the supported shape is + // scan-of-files --pipelined--> stateful, with no upstream shuffle). + val regularRoot = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(regularRoot, new HashPartitioner(2)) + val producerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + + assert(failure.get() != null, "a pipelined job with a regular-shuffle prefix must fail") + assert(failure.get().getMessage.contains("all-regular or all-pipelined"), + s"expected a mixed-job rejection, got: ${failure.get().getMessage}") + assert(taskSets.isEmpty, "no stage should be submitted for a rejected mixed job") + assertDataStructuresEmpty() + } + test("pipelined shuffle: deep chain A->B->C is submitted fully concurrently") { // A --pipelined--> B --pipelined--> C : all three co-scheduled (each edge is non-sequencing). val rddA = new MyRDD(sc, 2, Nil) @@ -6459,6 +6661,35 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti assertDataStructuresEmpty() } + test("pipelined shuffle: an explicit job cancellation cleans up a buffered consumer deferral") { + // A buffered deferral is the only mutable state this feature adds; it must never outlive its + // job. Job cancellation goes through failJobAndIndependentStages -> + // cleanupStateForJobAndIndependentStages, which drops the entry (as a consumer key) and removes + // the stage from every other consumer's pending-producer set. Verify a consumer whose + // completion is buffered while its producer runs leaves NO deferral behind when the job is + // cancelled, and no buffered success is later applied as a result. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val jobId = submit(consumerRdd, Array(0, 1)) + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + + // Consumer finishes ahead of its producer -> its completion is buffered (a deferral exists). + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(scheduler.dependentStageMap.nonEmpty, "a deferral must exist while the producer runs") + assert(results.isEmpty, "the consumer result must be deferred, not applied yet") + + // Cancel the job. The deferral must be torn down with everything else -- no stale entry, and + // the buffered success must not be replayed as a result. + cancel(jobId) + assert(scheduler.dependentStageMap.isEmpty, + "job cancellation must clean up the buffered consumer deferral (no state outlives the job)") + assert(results.isEmpty, "a cancelled job's buffered consumer success must not be applied") + assertDataStructuresEmpty() + } + test("regular shuffle job with speculation enabled is NOT rejected (rejection path is inert)") { // The speculation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with speculation on runs normally. @@ -6532,65 +6763,6 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a member on a non-default resource profile is rejected up front") { - // Admission measures capacity/occupancy against the DEFAULT resource profile, but a stage - // derives its profile from its RDDs. A pipelined group member on a custom profile would be - // admitted against the default pool's free slots yet run in a different (often smaller) pool - // and could deadlock there. Reject such a job before any stage is created. - // Ensure the default profile exists (id 0) before building a custom one, so the custom profile - // gets a distinct, non-default id regardless of test-execution order (profile ids come from a - // process-wide counter, and the default profile occupies id 0). - sc.resourceProfileManager.defaultResourceProfile - val ereqs = new ExecutorResourceRequests().cores(4) - val treqs = new TaskResourceRequests().cpus(2) - val customRp = new ResourceProfileBuilder().require(ereqs).require(treqs).build() - val producerRdd = new MyRDD(sc, 2, Nil) - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - .withResources(customRp) - // Sanity: the consumer really carries a non-default profile (otherwise the test is vacuous). - assert(consumerRdd.getResourceProfile() != null && - consumerRdd.getResourceProfile().id != ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, - s"test setup: expected a non-default profile, got id " + - s"${Option(consumerRdd.getResourceProfile()).map(_.id)}") - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a non-default-resource-profile member should be rejected") - assert(failure.get().getMessage.contains("non-default resource profile")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - - test("pipelined shuffle: a barrier group member is rejected up front") { - // A barrier stage exposes its output only after a global sync, which is incompatible with a - // pipelined consumer reading the producer's output incrementally. It is also a recovery hazard: - // a barrier task failure fails the stage and RESUBMITS it (markStageAsFinished without - // willRetry), and if such a producer were co-scheduled, the resubmit would drop a fan-in - // consumer's buffered completions and the job could never complete. Reject any job whose - // pipelined group contains a barrier member before any stage is created, so that path is - // unreachable. - val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) - assert(producerRdd.isBarrier(), "test setup: the producer must be a barrier RDD") - val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) - val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) - val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() - val failListener = new JobListener { - override def taskSucceeded(index: Int, result: Any): Unit = {} - override def jobFailed(exception: Exception): Unit = failure.set(exception) - } - submit(consumerRdd, Array(0, 1), listener = failListener) - assert(failure.get() != null, - "a pipelined job with a barrier group member should be rejected") - assert(failure.get().getMessage.contains("barrier")) - assert(taskSets.isEmpty, "no stage should be created for a rejected job") - assertDataStructuresEmpty() - } - test("regular shuffle job with dynamic allocation enabled is NOT rejected (path is inert)") { // The dynamic-allocation fail-fast must apply only to jobs with a pipelined dependency; a plain // regular-shuffle job with dynamic allocation on runs normally. @@ -6701,6 +6873,27 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("SPARK-58828: a barrier job does not consume its retry budget while held") { + // 4 barrier tasks on a local[2] backend fail the slot check; while held that must not + // count against spark.scheduler.barrier.maxConcurrentTasksCheck.maxFailures. + val barrierRdd = new MyRDD(sc, 4, Nil).barrier().mapPartitions(iter => iter) + val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] + myScheduler.executorsHeldForTest = true + try { + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(barrierRdd, Array(0, 1, 2, 3), listener = failListener) + assert(failure.get() === null, "a barrier job submitted while held must wait, not fail") + assert(scheduler.barrierJobIdToNumTasksCheckFailures.isEmpty, + "the retry budget must not be consumed while held") + } finally { + myScheduler.executorsHeldForTest = false + } + } + test("pipelined shuffle: an over-capacity 3-stage all-pipelined chain is rejected up front") { // A 3-stage all-pipelined chain A->B->C has whole-group demand 2+2+2 = 6. With capacity 3 it // cannot co-fit, so the up-front gang admission check (handleJobSubmitted) fails the job before @@ -6832,13 +7025,19 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a diamond reusing one producer counts its tasks once for admission") { + test("pipelined shuffle: admission counts a reused producer once (diamond rejected for " + + "fan-out, not for slots)") { // Fan-out/diamond: ONE PipelinedShuffleDependency (one producer, one shuffle id) is read by TWO - // consumers that are then narrow-joined into the result. Execution creates ONE producer stage - // (getOrCreateShuffleMapStage keys on shuffle id), so the group's real concurrent demand is - // producer(2) + result(2) = 4 -- NOT 6. Counting the producer once per consumer EDGE would make - // demand 6 and wrongly reject this group at capacity 4. Pinning capacity to exactly 4 admits - // the correctly-deduped group and would fail if the producer were double-counted. + // consumers that are then narrow-joined into the result. Fan-out is unsupported, so the group + // is ultimately rejected -- but WHICH rejection it gets proves the admission demand is deduped. + // The up-front slot admission (rejectUnadmittablePipelinedGroup) runs BEFORE the fan-out idiom + // check (checkPipelinedGroupsSupportedInRDDGraph). Execution creates ONE producer stage + // (getOrCreateShuffleMapStage keys on shuffle id), so the real concurrent demand is + // producer(2) + result(2) = 4. Counting the producer once per consumer EDGE would inflate it + // to 6. Pinning capacity to exactly 4: the deduped group PASSES the slot check and is then + // rejected for fan-out (PIPELINED_SHUFFLE_UNSUPPORTED); a double-counted group (6 > 4) would + // instead be rejected earlier for CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. So a dedup regression + // flips the error class -- which this test detects. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6852,13 +7051,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti val resultRdd = new MyRDD( sc, 2, List(new OneToOneDependency(consumer1), new OneToOneDependency(consumer2)), tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"diamond group is producer + result (both consumers); got ${taskSets.size} task sets") - // Drain to completion: producer stage first, then the result stage. - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + // Rejected for fan-out (demand fit within capacity 4), NOT for slots -- proving the reused + // producer was counted once. + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 @@ -6866,10 +7065,13 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } - test("pipelined shuffle: a wide fan-out reusing one producer counts its tasks once") { + test("pipelined shuffle: admission counts a reused producer once across a wide fan-out " + + "(rejected for fan-out, not for slots)") { // Wider fan-out: ONE producer feeds THREE consumers, all narrow-joined into the result. Real - // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. Capacity - // 4 admits the deduped group and would reject the double-counted one. + // demand is still producer(2) + result(2) = 4; a per-edge count would be 2 + 3*2 = 8. At + // capacity 4 the deduped group passes the slot check and is rejected for fan-out; a + // double-counted group (8 > 4) would be rejected first for insufficient slots. See the diamond + // test above for the ordering rationale. val producerRdd = new MyRDD(sc, 2, Nil) val myScheduler = scheduler.asInstanceOf[MyDAGScheduler] myScheduler.maxConcurrentTasksForTest = 4 @@ -6880,12 +7082,11 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti new OneToOneDependency(new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker)) }.toList val resultRdd = new MyRDD(sc, 2, consumers, tracker = mapOutputTracker) - submit(resultRdd, Array(0, 1)) - assert(taskSets.size === 2, - s"wide fan-out group is producer + result; got ${taskSets.size} task sets") - completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) - complete(taskSets(1), Seq((Success, 42), (Success, 43))) - assert(results === Map(0 -> 42, 1 -> 43)) + val failure = submitAndCaptureFailure(resultRdd, Array(0, 1)) + assertPipelinedUnsupported(failure, "more than one consumer") + assert(!failure.getMessage.contains("CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT"), + s"a reused producer must be counted once (fit capacity 4), got a slot rejection: " + + s"${failure.getMessage}") assertDataStructuresEmpty() } finally { myScheduler.maxConcurrentTasksForTest = 1000 @@ -7075,6 +7276,85 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti s"a dropped fan-in consumer must not replay when a surviving producer succeeds; got $results") } + test("pipelined shuffle: a fan-in group leaks no scheduler state when a producer fails and the " + + "job is torn down") { + // Leak check for the multi-producer (fan-in) deferral: a consumer deferred against TWO + // producers buffers its completions; when one producer fails and the job is torn down, ALL + // scheduler state -- the consumer's dependentStageMap deferral keyed on both producers -- + // must be cleaned up. The sibling "dropped as soon as one producer fails" test stops at the + // drop-vs-replay outcome; this one drives the failure to full job teardown and asserts no leak. + val producerA = new MyRDD(sc, 2, Nil) + val psdA = new PipelinedShuffleDependency(producerA, new HashPartitioner(2)) + val producerB = new MyRDD(sc, 2, Nil) + val psdB = new PipelinedShuffleDependency(producerB, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdA, psdB), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + assert(taskSets.size === 3, s"A, B and the consumer must be co-scheduled; got ${taskSets.size}") + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val consumerStage = scheduler.stageIdToStage(consumerTaskSet.stageId) + val producerATaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerA).get + + // Consumer finishes early -> buffered (deferred) against BOTH producers. + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while its producers run") + assert(scheduler.dependentStageMap.get(consumerStage).exists(_.parents.size == 2), + "the consumer must be deferred against both producers") + + // Producer A fails its whole task set -> the job fails and the group is torn down. + failed(producerATaskSet, "producer A blew up") + assert(failure.get() != null, "the job must fail when a fan-in producer fails") + assert(results.isEmpty, "buffered consumer successes must be dropped, not applied, on failure") + sc.listenerBus.waitUntilEmpty(10000) + // The whole point: no leak. Every scheduler map (including the two-parent deferral) empties. + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a multi-producer consumer's buffered successes are dropped when one " + + "producer fails (no stale replay via a sibling)") { + // Group-atomic drop across MULTIPLE producers. Consumer C reads TWO pipelined producers P1 and + // P2 (a join; purely all-pipelined, no regular shuffle). C finishes early -> its successes are + // buffered against BOTH producers. If P1 then fails, the whole group must be torn down and C's + // successes DROPPED -- they depended on P1's (now-invalid) output. Note the release path + // evaluates producerFailed per finishing producer, so a naive impl could remove P1 + // (parents still has P2, no drop), then later see P2 "succeed" and REPLAY. This asserts the + // shipped stack drops instead: a member failure aborts the whole group, which tears C down. + val rddP1 = new MyRDD(sc, 2, Nil) + val psdP1 = new PipelinedShuffleDependency(rddP1, new HashPartitioner(2)) + val rddP2 = new MyRDD(sc, 2, Nil) + val psdP2 = new PipelinedShuffleDependency(rddP2, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(psdP1, psdP2), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // P1, P2, and C are all co-scheduled (purely all-pipelined join, admitted up front). + assert(taskSets.size === 3, s"expected P1, P2, consumer co-scheduled, got ${taskSets.size}") + val tsP1 = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq rddP1).get + val tsC = taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + + // Consumer finishes early -> buffered against both P1 and P2. + complete(tsC, Seq((Success, 42), (Success, 43))) + assert(results.isEmpty, "consumer completions should be buffered while producers run") + assert(scheduler.dependentStageMap.keys.exists(_.rdd eq consumerRdd), + "consumer should be deferred against its two producers") + + // P1 fails. The group is torn down; C's buffered successes must be DROPPED, never replayed -- + // even though P2 has not (and now will not) complete. + failed(tsP1, "producer P1 blew up") + assert(failure.get() != null, "the job must fail when a pipelined producer fails") + assert(results.isEmpty, + "a multi-producer consumer's buffered successes must be dropped when any producer fails") + assertDataStructuresEmpty() + } test("pipelined shuffle: a deferred consumer task fires its TaskEnd exactly once (at replay)") { // A deferred CompletionEvent must have its side effects (task-end listener event, accumulator @@ -7219,6 +7499,788 @@ class DAGSchedulerSuite extends SparkFunSuite with TempLocalSparkContext with Ti } } + test("pipelined shuffle: both producer and consumer task sets are marked isPipelined") { + // The TaskSet.isPipelined flag drives group-atomic failure in the task scheduler; verify the + // DAGScheduler sets it for both members of a pipelined group (producer and consumer). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.size === 2) + val producerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTs = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + assert(producerTs.isPipelined, "the pipelined producer's task set must be marked isPipelined") + assert(consumerTs.isPipelined, "the pipelined consumer's task set must be marked isPipelined") + + completeShuffleMapStageSuccessfully(producerTs.stageId, 0, 2) + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("regular shuffle: task sets are NOT marked isPipelined (inertness)") { + // A regular producer/consumer must not be marked isPipelined. + val producerRdd = new MyRDD(sc, 2, Nil) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + assert(taskSets.head.isPipelined === false, "a regular producer must not be marked isPipelined") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + assert(taskSets(1).isPipelined === false, "a regular consumer must not be marked isPipelined") + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + // ========================================================================================== + // Cross-job / cross-time reuse prevention at both layers: a consumed pipelined + // producer whose executor is lost must not be resubmitted + // ========================================================================================== + + test("pipelined shuffle: producer availability is tracked on the stage, not the" + + "MapOutputTracker") { + // A pipelined producer's completed partitions are tracked on ShuffleMapStage (monotonic), not + // registered as durable outputs in the MapOutputTracker. Verify: after the producer's map tasks + // succeed, the stage is available WITHOUT the shuffle having map outputs in the tracker. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val shuffleId = pipelinedDep.shuffleId + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + + // Complete the producer's two map tasks. + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + + // The stage is available via its local completed-partition set... + assert(producerStage.isAvailable, + "pipelined producer should be available after its tasks finish") + assert(producerStage.isPipelined) + // ...but the pipelined shuffle has NO map outputs registered in the MapOutputTracker. + assert(mapOutputTracker.getNumAvailableOutputs(shuffleId) === 0, + "a pipelined shuffle must not register durable map outputs in the MapOutputTracker") + + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: numAvailableOutputs/findMissingPartitions reflect the exact " + + "incomplete partition set") { + // Pin the availability CONTRACT for a partially-complete pipelined producer, not just the + // all-done / none-done endpoints: numAvailableOutputs must be the count of completed partitions + // and findMissingPartitions must return the exact ids still missing (identity, not just size). + // Without this, a plausible refactor (revert the findMissingPartitions isPipelined branch to + // the + // tracker, use a bare counter, drop/take by size, or invert the filter) would silently break + // the + // contract yet pass every full-completion test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Complete ONLY partition 1, leaving partition 0 missing. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + assert(producerStage.numAvailableOutputs === 1) + assert(!producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq(0), + "findMissingPartitions must name the exact missing partition, not just a right-sized set") + + // Complete partition 0; now nothing is missing. + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 2) + assert(producerStage.isAvailable) + assert(producerStage.findMissingPartitions() === Seq.empty) + + // Drain the consumer cleanly. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a post-executor-loss straggler success must not resubmit the producer " + + "(bogus-epoch race)") { + // A pipelined producer's map task can succeed on an executor whose loss is already recorded + // (its StatusUpdate raced the executor-loss event). That completion hits the "possibly bogus + // epoch" branch, which for a regular shuffle simply ignores it (a healthy reattempt will + // re-register the output). But the same branch also runs `pendingPartitions -= partitionId` + // first, so if it is the last pending partition and we do NOT record it in the pipelined + // completed set, the stage looks "done but not available" and processShuffleMapStageCompletion + // resubmits the transient producer -- the exact streaming-writer hang. A pipelined stage must + // record the partition as completed even on the bogus-epoch path (its output is monotonic and + // MapOutputTracker's executor-loss stripping does not apply to it). + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Partition 0 completes normally on executor "hostA-exec" (recorded). Note makeMapStatus(host) + // builds executorId = host + "-exec", so pass host "hostA" to get executorId "hostA-exec". + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Executor "hostA-exec" is lost: this records executorFailureEpoch("hostA-exec") = current + // epoch. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + // Now a delayed straggler Success for partition 1 arrives FROM the lost executor "hostA-exec". + // Its task epoch is the default (-1) <= the recorded failure epoch, so it is treated as + // possibly-bogus. It must still be recorded for a pipelined stage, so the producer becomes + // available and is NOT resubmitted. + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostA", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer must be available after all partitions succeed, even via a " + + "bogus-epoch " + + "straggler") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted by a post-loss straggler success") + + // Consumer drains normally; no hang. + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a producer success with ignoreOldTaskAttempts set is still recorded " + + "(no resubmit)") { + // Sibling of the bogus-epoch case, for the OTHER guard that gates a regular shuffle's + // recording: + // ignoreOldTaskAttempts (set when a stage is rolled back, e.g. as a succeeding stage of an + // indeterminate ancestor). A pipelined producer's completed set is monotonic and never rolled + // back, so its success must be recorded even when ignoreOldTaskAttempts is true -- otherwise + // the + // last partition is dropped (pendingPartitions decremented, not recorded) -> "done but not + // available" -> processShuffleMapStageCompletion resubmits the transient producer. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + runEvent(makeCompletionEvent(producerTs.tasks(0), Success, makeMapStatus("hostA", 2))) + assert(producerStage.numAvailableOutputs === 1) + val taskSetsBefore = taskSets.size + + // Force ignoreOldTaskAttempts=true for the next completion: maxAttemptIdToIgnore >= the task's + // stageAttemptId (0). For a regular stage this would drop the completion; a pipelined stage + // must + // still record it. + producerStage.maxAttemptIdToIgnore = Some(0) + runEvent(makeCompletionEvent(producerTs.tasks(1), Success, makeMapStatus("hostB", 2))) + + assert(producerStage.isAvailable, + "a pipelined producer success must be recorded even when ignoreOldTaskAttempts is set") + assert(producerStage.findMissingPartitions() === Seq.empty) + assert(taskSets.size === taskSetsBefore, + "the pipelined producer must NOT be resubmitted when a success arrives under " + + "ignoreOldTaskAttempts") + + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTs, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: losing an executor does NOT flip a completed producer to unavailable " + + "or resubmit it") { + // A completed, consumed pipelined producer whose executor is lost must NOT be + // resubmitted (which would hang the streaming writer in awaitTerminationAcks). Because + // pipelined + // availability is tracked on the stage (monotonic) and not the MapOutputTracker, executor loss + // cannot flip isAvailable, so the producer is not resubmitted. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerStageId = taskSets.head.stageId + val producerStage = + scheduler.stageIdToStage(producerStageId).asInstanceOf[ShuffleMapStage] + // Producer completes on hostA (both partitions). makeMapStatus("hostA") registers the output + // under executor id "hostA-exec" (makeBlockManagerId appends "-exec"), so the ExecutorLost + // below -- which loses executor id "hostA-exec" -- actually matches the registered output. + complete(taskSets.head, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostA", 2)))) + assert(producerStage.isAvailable) + val taskSetsAfterProducer = taskSets.size + + // Lose the executor that ran the producer ("hostA-exec"). For a regular shuffle this would + // strip the outputs and flip isAvailable -> resubmit; for a pipelined shuffle it must be inert. + runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed"))) + + assert(producerStage.isAvailable, + "a completed pipelined producer must remain available after executor loss (no tracker strip)") + // Pin the underlying completed set directly (not just the derived isAvailable boolean): + // executor + // loss must not remove any completed partition, so nothing is missing. + assert(producerStage.findMissingPartitions() === Seq.empty, + "executor loss must not strip a pipelined producer's completed partitions (monotonic)") + assert(taskSets.size === taskSetsAfterProducer, + "the pipelined producer must NOT be resubmitted on executor loss") + + // The consumer completes normally; no hang, no extra producer attempt. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + complete(consumerTaskSet, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: binding a live producer stage to a second concurrent job fails fast") { + // A pipelined shuffle is a once-through live stream with no retained output, so two concurrent + // jobs cannot share one producer stage. While job 0 is still active its producer + // stage stays cached in shuffleIdToMapStage bound to job 0; a second concurrent job that reuses + // the SAME PipelinedShuffleDependency would bind that live stage to a second jobId, which is + // the + // forbidden cross-job reuse. Fail fast rather than let job 1 attach to job 0's live stream. + // (A sequential re-run is NOT this case: after job 0 finishes, cleanup drops the stage from + // shuffleIdToMapStage, so a later job gets a fresh producer bound to only its own job -- + // exactly how each streaming micro-batch reruns its producer.) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + // Job 0: submit but leave it active (do not complete the producer or the result stage), so the + // producer stage remains cached in shuffleIdToMapStage bound to job 0. + val firstConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(firstConsumer, Array(0, 1)) + val producerStage = + scheduler.stageIdToStage(taskSets.head.stageId).asInstanceOf[ShuffleMapStage] + assert(producerStage.jobIds === Set(0)) + val shuffleStagesBeforeJob1 = scheduler.shuffleIdToMapStage.size + val stagesBeforeJob1 = scheduler.stageIdToStage.size + + // Job 1: a second concurrent job reusing the SAME pipelined dependency -> cross-job reuse of a + // live producer stage -> fail fast. + val secondConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(secondConsumer, Array(0, 1), listener = failListener) + assert(failure.get() != null, + "binding a live pipelined producer to a second concurrent job must fail the job") + assert(failure.get().getMessage.contains("PIPELINED_SHUFFLE_CROSS_JOB_REUSE") || + failure.get().getMessage.contains("reused across jobs"), + s"expected a cross-job-reuse error, got: ${failure.get().getMessage}") + // Job 0's producer stage was never bound to job 1, and the failed job 1 left no scheduler + // residue (the throw happened during stage creation, before any job-1 state was registered). + assert(producerStage.jobIds === Set(0)) + assert(scheduler.shuffleIdToMapStage.size === shuffleStagesBeforeJob1, + "the failed cross-job submission must not add a shuffle->stage mapping") + assert(scheduler.stageIdToStage.size === stagesBeforeJob1, + "the failed cross-job submission must not leave orphaned stages") + + // Job 0 can still complete normally -- the rejected job 1 did not corrupt its shared producer. + completeShuffleMapStageSuccessfully(producerStage.id, 0, 2) + val job0Consumer = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq firstConsumer + }.get + complete(job0Consumer, Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: sequential re-run of the same producer is sound (fresh stage per job)") { + // Contrast with the concurrent case above: after a job using a pipelined dependency finishes, + // cleanup removes its producer from shuffleIdToMapStage, so re-submitting a job on the SAME + // dependency creates a fresh producer stage bound to only the new job. This is sound (it is how + // a streaming micro-batch reruns) and must NOT trip the cross-job-reuse fail-fast. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + + val consumer0 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer0, Array(0, 1)) + complete(taskSets(0), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + results.clear() + + // Second, sequential job on the same dependency: a fresh producer stage, no fail-fast, no hang. + // taskSets is a growing buffer across both jobs, so the second job's task sets are at index 2+. + val consumer1 = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumer1, Array(0, 1)) + complete(taskSets(2), Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(taskSets(3), Seq((Success, 4), (Success, 5))) + assert(results === Map(0 -> 4, 1 -> 5)) + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a FetchFailed on a group member fails the group, not a single-stage " + + "resubmit") { + // A FetchFailed must fail a pipelined (streaming) query promptly rather than trigger the + // base scheduler's single-stage resubmit -> serial recompute -> deadlock. The transient + // pipelined shuffle cannot be re-read, and members are co-scheduled, so a lone-stage resubmit + // is never valid. Any member's FetchFailed must abort the whole group (-> job abort -> + // the caller reruns the batch). Note the base TaskSetManager does NOT count a FetchFailed (it + // marks the task successful and zombies the set), so the group-atomic maxTaskFailures=1 lever + // does not apply to FetchFailed; the routing must be enforced in the DAGScheduler. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val producerTs = taskSets.head + val producerStage = + scheduler.stageIdToStage(producerTs.stageId).asInstanceOf[ShuffleMapStage] + + // Producer completes (its outputs are tracked locally on the pipelined stage). + complete(producerTs, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val taskSetsBeforeFetchFailure = taskSets.size + + // A consumer task hits a FetchFailed reading the pipelined shuffle. + runEvent(makeCompletionEvent( + consumerTs.tasks(0), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + + // The job must be failed (group-atomic), and no single stage may be resubmitted: no new task + // set is created, and the scheduler is not left waiting to recompute the producer in isolation. + // scheduleResubmit posts ResubmitFailedStages on a timer, so drive any pending resubmit and + // confirm nothing new is launched. + scheduler.resubmitFailedStages() + assert(failure != null, "a FetchFailed on a pipelined group member must fail the job") + assert(taskSets.size === taskSetsBeforeFetchFailure, + "a pipelined group member's FetchFailed must NOT resubmit a single stage") + assert(!scheduler.runningStages.exists(_.isInstanceOf[ShuffleMapStage]), + "the pipelined producer must not be left running/resubmitted after the group fails") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a consumer result task throwing aborts the whole group") { + // Defense-in-depth for the group-atomic model at the DAGScheduler layer: a CONSUMER (result) + // task failing must tear down the whole group, not just its own stage. Result and map tasks + // share handleFailedTask / effectiveMaxTaskFailures=1, so the first consumer-task exception + // makes the TaskSetManager abort the set (delivered here as TaskSetFailed, exactly as a + // maxTaskFailures=1 abort surfaces to the DAGScheduler), which must abort the group and tear + // down the still-running producer -- the caller then reruns the batch. Complements the + // TaskSetManager-layer maxTaskFailures=1 tests and the producer-failure drop test. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + val consumerTs = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + val producerStage = scheduler.stageIdToStage(taskSets.head.stageId) + // The producer is STILL RUNNING (co-scheduled, not yet completed) when the consumer fails -- + // this is what gives the assertion teeth: group-atomic teardown must reach a running producer. + assert(scheduler.runningStages.contains(producerStage), + "the pipelined producer must be co-scheduled and running alongside the consumer") + + // The consumer's task set aborts on the first task exception (maxTaskFailures=1 for a pipelined + // member). This must fail the whole group and tear down the still-running producer. + failed(consumerTs, "consumer result task threw") + assert(failure != null, "a consumer task failure must fail the group's job") + assert(!scheduler.runningStages.contains(producerStage), + "the still-running pipelined producer must be torn down when the consumer fails the group") + sc.listenerBus.waitUntilEmpty() + assertDataStructuresEmpty() + } + + // ========================================================================================== + // Group-atomic rerun resets per-partition commit authorization + // ========================================================================================== + + test("pipelined shuffle: a group rerun resets per-partition commit authorization") { + // A pipelined group is atomic, so a failure reruns the WHOLE group -- including a result + // stage whose tasks already succeeded and committed. Those committed partitions are rerun and + // must be allowed to commit again. OutputCommitCoordinator permanently denies re-commit for a + // committed partition (a Success clears nothing; keyed by stage id). The rerun stays correct + // two ways, both asserted here: (b) the group teardown runs the committed result stage through + // markStageAsFinished -> stageEnd, clearing its committer state (so no stale authorization + // survives); and (a) the caller's rerun is a NEW job whose stages get FRESH stage ids, so the + // coordinator has no prior committer for them regardless. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Locate producer/consumer defensively by RDD identity (not task-set order). + val producerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq producerRdd).get + val consumerTaskSet = + taskSets.find(ts => scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd).get + val firstConsumerStageId = consumerTaskSet.stageId + + // Both stages register OutputCommitCoordinator state at submission (stageStart), so the + // coordinator holds per-stage commit-authorization state that a rerun must not inherit. (This + // asserts stage state exists to be reset; it does not claim a committer was authorized -- that + // needs a live canCommit RPC, which the mock backend does not drive.) + assert(!scheduler.outputCommitCoordinator.isEmpty, + "the coordinator must hold commit-authorization state for the submitted group's stages") + + // The producer now fails -> group-atomic failure -> the whole group is torn down and the job + // fails; the caller will rerun the batch as a new job. + failed(producerTaskSet, "producer blew up") + assert(failure.get() != null, "the group must fail atomically when the producer fails") + // Teardown ran the group's stages through markStageAsFinished -> stageEnd, clearing their + // coordinator state: the commit-authorization state does not survive the failed attempt. + // (isEmpty here proves stageEnd reached every stage the teardown covered.) + assert(scheduler.outputCommitCoordinator.isEmpty, + "group teardown must reset per-partition commit authorization; none may survive") + assertDataStructuresEmpty() + + // The caller reruns the batch as a NEW job on the same dependency. Its stages get fresh ids, so + // the coordinator has no prior committer, and the rerun's partition 0 can commit again. Only + // the + // task sets submitted from here on belong to the rerun (earlier ones' stages were cleaned up, + // so + // look them up defensively). + val taskSetsBeforeRerun = taskSets.size + val rerunConsumer = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + submit(rerunConsumer, Array(0, 1)) + val rerunTaskSets = taskSets.drop(taskSetsBeforeRerun) + val rerunConsumerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq rerunConsumer) + }.get + val rerunProducerTaskSet = rerunTaskSets.find { ts => + scheduler.stageIdToStage.get(ts.stageId).exists(_.rdd eq producerRdd) + }.get + assert(rerunConsumerTaskSet.stageId != firstConsumerStageId, + "the rerun's result stage must get a fresh stage id (fresh coordinator state)") + complete(rerunProducerTaskSet, Seq( + (Success, makeMapStatus("hostA", 2)), + (Success, makeMapStatus("hostB", 2)))) + complete(rerunConsumerTaskSet, Seq((Success, 7), (Success, 8))) + assert(results === Map(0 -> 7, 1 -> 8), "the rerun must complete, re-committing its partitions") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a buffered consumer success is dropped on group abort, and its " + + "TaskEnd is flushed so the stage is not leaked as running") { + // A consumer's successful task has its whole completion event buffered while its producer runs + // (coarse model) -- no TaskEnd yet. If the group is then torn down before the producer finishes + // (here: the consumer's OTHER task hits a FetchFailed, which aborts the group), the buffered + // success's result must NOT be applied, but its TaskEnd IS flushed on the drop path. A listener + // that tracks active tasks (e.g. AppStatusListener, which removes a stage once activeTasks hits + // 0) therefore does not leak the consumer stage as perpetually running: the success's TaskStart + // was delivered, and its matching TaskEnd arrives at teardown even though the result is + // dropped. + val endedTaskIds = new java.util.concurrent.ConcurrentHashMap[Long, Boolean]() + val recordingListener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = + endedTaskIds.put(taskEnd.taskInfo.taskId, true) + } + sc.addSparkListener(recordingListener) + try { + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = results.put(index, result) + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(consumerRdd, Array(0, 1), listener = failListener) + // Identify the co-scheduled consumer (ResultStage over consumerRdd) and its still-running + // pipelined producer by RDD identity rather than task-set order. + val consumerTaskSet = taskSets.find { ts => + scheduler.stageIdToStage(ts.stageId).rdd eq consumerRdd + }.get + assert(scheduler.dependentStageMap.size === 1, + "the consumer must be co-scheduled with a running producer (a deferral must exist)") + + // Consumer partition 0 succeeds. Its whole completion event is buffered (no job result and no + // TaskEnd yet). Give it a known taskId so we can assert the drop path flushes it later. + val bufferedTaskId = 7007L + runEvent(makeCompletionEvent(consumerTaskSet.tasks(0), Success, 42, + taskInfo = createFakeTaskInfoWithId(bufferedTaskId))) + sc.listenerBus.waitUntilEmpty() + assert(!endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success must not emit its TaskEnd while the producer still runs") + assert(results.isEmpty, "the consumer's job result must still be deferred (producer running)") + assert(scheduler.dependentStageMap.get(scheduler.stageIdToStage(consumerTaskSet.stageId)) + .exists(_.delayedTaskCompletionEvents.nonEmpty), + "the consumer's completion event must be buffered while its producer runs") + + // The consumer's OTHER task now hits a FetchFailed. For a pipelined group member this aborts + // the whole group (group-atomic failure) WITHOUT the producer ever finishing. The buffered + // success is dropped (its result must not be applied), but its TaskEnd is flushed on the drop + // path so active-task-tracking listeners see the task finish. + runEvent(makeCompletionEvent( + consumerTaskSet.tasks(1), + FetchFailed(makeBlockManagerId("hostA"), pipelinedDep.shuffleId, 0L, 0, 0, "ignored"), + null)) + scheduler.resubmitFailedStages() + assert(failure.get() != null, "the job must fail") + + sc.listenerBus.waitUntilEmpty() + assert(results.isEmpty, + "the deferred consumer success must be dropped on abort, not applied as a result") + assert(endedTaskIds.containsKey(bufferedTaskId), + "the buffered consumer success's TaskEnd must be flushed on the drop path, so the stage " + + "is not leaked as perpetually running") + assertDataStructuresEmpty() + } finally { + sc.removeSparkListener(recordingListener) + } + } + + test("pipelined shuffle: a barrier producer stage is rejected") { + // A barrier stage exposes output only after a global sync, incompatible with incremental reads. + val producerRdd = new MyRDD(sc, 2, Nil).barrier().mapPartitions(iter => iter) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "barrier") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a statically-indeterminate producer is rejected") { + // Indeterminate output's recovery is stage rollback-and-recompute, which a group never performs + // (so it never applies); reject rather than carry dead machinery. + val producerRdd = new MyRDD(sc, 2, Nil, indeterminate = true) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported(submitAndCaptureFailure(consumerRdd, Array(0, 1)), "indeterminate") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a reliable RDD checkpoint in a PRODUCER's chain is rejected") { + // A reliable checkpoint writes a durable, lineage-truncated snapshot -> reintroduces cross-time + // reuse of a transient edge and needs a post-success recompute of the vanished input. Rejected + // by walking the producer's within-stage chain for a ReliableRDDCheckpointData. Keyed on + // checkpointData (not isCheckpointed): the write has not happened yet at group-creation time. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val checkpointableRdd = new MyCheckpointRDD(sc, 2, Nil) + checkpointableRdd.checkpoint() // sets checkpointData to a ReliableRDDCheckpointData + assert(checkpointableRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + val pipelinedDep = new PipelinedShuffleDependency(checkpointableRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable RDD checkpoint in a CONSUMER's chain is rejected") { + // The rejection must cover a consumer member's chain too, not just the producer's: a consumer's + // transient input IS the pipelined shuffle, so a reliable checkpoint there would re-read the + // vanished stream on recompute. Purely all-pipelined shape (no regular shuffle -- those are + // rejected separately): producer --pipelined--> consumer(checkpointed, result). The consumer + // reads the pipelined shuffle and its within-stage chain carries the reliable checkpoint. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + // The consumer RDD reads the pipelined shuffle AND is reliably checkpointed. + val consumerRdd = new MyCheckpointRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + consumerRdd.checkpoint() + assert(consumerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a reliable checkpoint DOWNSTREAM in the consumer stage (not on the " + + "reading RDD) is still rejected") { + // Coverage for a checkpoint that is NOT on the RDD reading the pipelined shuffle, but on a + // narrow-dep child of it WITHIN the same consumer stage. Purely all-pipelined shape: + // producer --pipelined--> reads --(narrow)--> checkpointed (result). + // `reads` and `checkpointed` are one stage. Rooting the check at the pipelined-reading RDD and + // walking parents would MISS this (the checkpoint is downstream of the read); the check must + // instead recognize that `checkpointed`'s own within-stage chain reads a pipelined shuffle. + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val readsRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + // A narrow-dep child of the reading RDD, in the SAME stage, that is reliably checkpointed, + // and is the result RDD (no regular shuffle after -- that would be separately rejected). + val checkpointedRdd = new MyCheckpointRDD(sc, 2, List(new OneToOneDependency(readsRdd))) + checkpointedRdd.checkpoint() + assert(checkpointedRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assertPipelinedUnsupported( + submitAndCaptureFailure(checkpointedRdd, Array(0, 1)), "reliable RDD checkpoint") + assertDataStructuresEmpty() + } + } + + test("pipelined shuffle: a producer feeding more than one consumer (fan-out) is rejected") { + // 1:N fan-out is deferred to a later version and rejected up front here. Fan-out is + // detected at the RDD level -- two DISTINCT RDDs listing the same pipelined shuffle as a + // dependency -- so it is expressible in an all-pipelined job without any regular shuffle: two + // consumer RDDs both read the same pipelined producer, unioned by a narrow dependency into the + // result. checkPipelinedGroupsSupportedInRDDGraph counts 2 distinct consumers for the shuffle. + val producerRdd = new MyRDD(sc, 2, Nil) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerA = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val consumerB = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + val union = new MyRDD(sc, 2, + List(new OneToOneDependency(consumerA), new OneToOneDependency(consumerB)), + tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(union, Array(0, 1)), "more than one consumer") + assertDataStructuresEmpty() + } + + // Resource-profile rejection. The gang slot check measures capacity against the DEFAULT + // resource profile, so the whole group must run on the default profile; any member with an + // explicit non-default profile is rejected. The three shapes below must all be rejected; the + // uniform-non-default and non-default-plus-default cases in particular would each pass a + // "more than one DISTINCT profile" check yet still be admitted against the wrong (default) pool. + private def rpWithCores(cores: Int, cpus: Int): ResourceProfile = + new ResourceProfileBuilder() + .require(new ExecutorResourceRequests().cores(cores)) + .require(new TaskResourceRequests().cpus(cpus)).build() + + test("pipelined shuffle: a group with two distinct non-default resource profiles is rejected") { + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rpWithCores(8, 2)) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group uniformly on one non-default resource profile is rejected") { + // Both members share ONE non-default profile -- a "distinct profile count" of 1, which a + // more-than-one-distinct-profile check would wrongly admit, then measure against the default + // profile's capacity (the wrong pool). Must be rejected: the whole group is off the default RP. + val rp = rpWithCores(4, 1) + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rp) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + .withResources(rp) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("pipelined shuffle: a group mixing a non-default profile with the default is rejected") { + // The producer carries an explicit non-default profile; the consumer is left on the default. + // The set of EXPLICIT profiles is again size 1, so a distinct-explicit-profile check would miss + // it, but the group genuinely spans the non-default and default profiles. Must be rejected. + val producerRdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 1)) + val pipelinedDep = new PipelinedShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(pipelinedDep), tracker = mapOutputTracker) + assertPipelinedUnsupported( + submitAndCaptureFailure(consumerRdd, Array(0, 1)), "non-default resource profile") + assertDataStructuresEmpty() + } + + test("regular shuffle idioms are NOT rejected (inertness of the pipelined fail-fast)") { + // The pipelined fail-fast checks (indeterminate producer, reliable-checkpoint-in-chain) must be + // inert for a job with NO pipelined dependency: a regular shuffle whose producer is BOTH + // indeterminate AND reliably checkpointed -- the two idioms most likely to false-fire -- must + // run exactly as before. (Actually exercise both, not just the flag.) + withTempDir { dir => + sc.setCheckpointDir(dir.getCanonicalPath) + val producerRdd = new MyCheckpointRDD(sc, 2, Nil, indeterminate = true) + producerRdd.checkpoint() + assert(producerRdd.checkpointData.exists(_.isInstanceOf[ReliableRDDCheckpointData[_]])) + assert(producerRdd.outputDeterministicLevel == DeterministicLevel.INDETERMINATE) + val regularDep = new ShuffleDependency(producerRdd, new HashPartitioner(2)) + val consumerRdd = new MyRDD(sc, 2, List(regularDep), tracker = mapOutputTracker) + submit(consumerRdd, Array(0, 1)) + // The job proceeds normally (producer stage submitted), i.e. NOT failed by a pipelined check. + assert(failure === null, "a regular shuffle must not be rejected by the pipelined fail-fast") + completeShuffleMapStageSuccessfully(taskSets.head.stageId, 0, 2) + complete(taskSets(1), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + } + + test("regular job on a non-default resource profile is NOT rejected (RP check is pipelined)") { + // The resource-profile rejection is not keyed on a pipelined dependency, so it must run ONLY + // for a job that has one (handleJobSubmitted gates checkPipelinedGroupsSupportedInRDDGraph on + // hasPipelined). A perfectly ordinary job that merely attaches a non-default profile via + // RDD.withResources -- a GA stage-level-scheduling feature -- has NO pipelined dependency and + // must run untouched. Without the gate the whole graph walk fires and rejects it with + // PIPELINED_SHUFFLE_UNSUPPORTED, a regression on plain withResources jobs. + val rdd = new MyRDD(sc, 2, Nil).withResources(rpWithCores(4, 2)) + // A non-default profile drives submitMissingTasks through addPySparkConfigsToProperties, which + // needs a non-null Properties; pass one (the default submit() overload leaves it null). + submit(rdd, Array(0, 1), properties = new Properties()) + assert(failure === null, + "a regular job with a non-default resource profile must not be rejected by the pipelined " + + "fail-fast") + assert(taskSets.nonEmpty, "the job must proceed to task submission, not be failed up front") + complete(taskSets(0), Seq((Success, 42), (Success, 43))) + assert(results === Map(0 -> 42, 1 -> 43)) + assertDataStructuresEmpty() + } + + private def submitAndCaptureFailure(finalRdd: RDD[_], partitions: Array[Int]): Exception = { + val failure = new java.util.concurrent.atomic.AtomicReference[Exception]() + val failListener = new JobListener { + override def taskSucceeded(index: Int, result: Any): Unit = {} + override def jobFailed(exception: Exception): Unit = failure.set(exception) + } + submit(finalRdd, partitions, listener = failListener) + failure.get() + } + + private def assertPipelinedUnsupported(failure: Exception, reasonSubstring: String): Unit = { + assert(failure != null, "the job must fail fast on the unsupported pipelined idiom") + val msg = failure.getMessage + assert(msg.contains("PIPELINED_SHUFFLE_UNSUPPORTED") || msg.contains("unsupported feature"), + s"expected a PIPELINED_SHUFFLE_UNSUPPORTED error, got: $msg") + assert(msg.contains(reasonSubstring), + s"expected reason to mention '$reasonSubstring', got: $msg") + } + + test("submitMapStage on an RDD with 0 partitions reports an internal error") { + val shuffleMapRdd = new MyRDD(sc, 0, Nil) + val shuffleDep = new ShuffleDependency(shuffleMapRdd, new HashPartitioner(1)) + // `ShuffleExchangeExec` guards this case, so reaching it means an internal invariant broke. + checkError( + exception = intercept[SparkException] { + scheduler.submitMapStage(shuffleDep, (_: MapOutputStatistics) => (), CallSite("", ""), + new Properties()) + }, + condition = "INTERNAL_ERROR", + sqlState = Some("XX000"), + parameters = scala.collection.immutable.Map( + "message" -> "Can't run submitMapStage on RDD with 0 partitions.")) + assertDataStructuresEmpty() + } + } class DAGSchedulerAbortStageOffSuite extends DAGSchedulerSuite { diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala index dcb22f8b0d7d8..4744b41809bb0 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskDescriptionSuite.scala @@ -85,6 +85,7 @@ class TaskDescriptionSuite extends SparkFunSuite { originalProperties, cpus = 2, originalResources, + None, taskBuffer ) @@ -102,6 +103,7 @@ class TaskDescriptionSuite extends SparkFunSuite { assert(decodedTaskDescription.properties.equals(originalTaskDescription.properties)) assert(decodedTaskDescription.cpus.equals(originalTaskDescription.cpus)) assert(decodedTaskDescription.resources === originalTaskDescription.resources) + assert(decodedTaskDescription.userCredentials === None) assert(decodedTaskDescription.serializedTask.equals(taskBuffer)) } @@ -119,6 +121,7 @@ class TaskDescriptionSuite extends SparkFunSuite { new Properties(), cpus, Map.empty, + None, ByteBuffer.wrap(Array[Byte](1, 2, 3, 4))) val decoded = TaskDescription.decode(TaskDescription.encode(taskDescription)) // The decoded amount is what the executor echoes back in its status update and what the @@ -129,4 +132,66 @@ class TaskDescriptionSuite extends SparkFunSuite { } } + test("encoding and decoding TaskDescription with userCredentials preserves credentials") { + val taskBuffer = ByteBuffer.wrap(Array[Byte](1, 2, 3, 4)) + val properties = new Properties() + properties.put("key", "value") + + val credentialBytes = Array[Byte](10, 20, 30, 40, 50) + + val taskDescription = new TaskDescription( + taskId = 42, + attemptNumber = 0, + executorId = "exec-1", + name = "task with credentials", + index = 0, + partitionId = 0, + JobArtifactSet.emptyJobArtifactSet, + properties, + cpus = 1, + Map.empty, + Some((1L, credentialBytes)), + taskBuffer + ) + + val serialized = TaskDescription.encode(taskDescription) + val decoded = TaskDescription.decode(serialized) + + assert(decoded.taskId === 42) + assert(decoded.name === "task with credentials") + assert(decoded.userCredentials.isDefined) + assert(decoded.userCredentials.get._1 === 1L) + assert(decoded.userCredentials.get._2 === credentialBytes) + assert(decoded.serializedTask.equals(taskBuffer)) + } + + test("encoding and decoding TaskDescription without userCredentials round-trips correctly") { + val taskBuffer = ByteBuffer.wrap(Array[Byte](5, 6, 7, 8)) + val properties = new Properties() + + val taskDescription = new TaskDescription( + taskId = 99, + attemptNumber = 1, + executorId = "exec-2", + name = "task without credentials", + index = 3, + partitionId = 2, + JobArtifactSet.emptyJobArtifactSet, + properties, + cpus = 2, + Map.empty, + None, + taskBuffer + ) + + val serialized = TaskDescription.encode(taskDescription) + val decoded = TaskDescription.decode(serialized) + + assert(decoded.taskId === 99) + assert(decoded.attemptNumber === 1) + assert(decoded.name === "task without credentials") + assert(decoded.userCredentials === None) + assert(decoded.serializedTask.equals(taskBuffer)) + } + } diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskInfoSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskInfoSuite.scala new file mode 100644 index 0000000000000..7467f25f933d5 --- /dev/null +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskInfoSuite.scala @@ -0,0 +1,53 @@ +/* + * 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.scheduler + +import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException, TaskState} + +class TaskInfoSuite extends SparkFunSuite { + + private def newTaskInfo(): TaskInfo = new TaskInfo( + taskId = 42L, + index = 0, + attemptNumber = 0, + partitionId = 0, + launchTime = 100L, + executorId = "exec-1", + host = "host-1", + taskLocality = TaskLocality.PROCESS_LOCAL, + speculative = false) + + test("duration is not available before the task finishes") { + val info = newTaskInfo() + checkError( + exception = intercept[SparkUnsupportedOperationException] { + info.duration + }, + condition = "UNSUPPORTED_CALL.TASK_NOT_FINISHED", + sqlState = Some("0A000"), + parameters = Map( + "className" -> "org.apache.spark.scheduler.TaskInfo", + "methodName" -> "duration")) + } + + test("duration is available once the task has finished") { + val info = newTaskInfo() + info.markFinished(TaskState.FINISHED, 300L) + assert(info.duration === 200L) + } +} diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala index 81782df283eec..810a0b7c483b5 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala @@ -2866,4 +2866,84 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 0) } + /** + * A TaskSet marked as a pipelined-group member. FakeTask.createTaskSet always leaves + * `isPipelined` at its default (false), so build the TaskSet directly here. + */ + private def pipelinedTaskSet(numTasks: Int, stageId: Int, stageAttemptId: Int = 0): TaskSet = { + val tasks = Array.tabulate[Task[_]](numTasks)(i => new FakeTask(stageId, i, Nil)) + new TaskSet(tasks, stageId, stageAttemptId, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + } + + test("SPARK-58913: hasPipelinedTaskSets is true only while a pipelined-group member is live") { + val taskScheduler = setupScheduler() + assert(!taskScheduler.hasPipelinedTaskSets, "no task set has been submitted yet") + + // A regular task set leaves isPipelined at false and must not register as a group member. + taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0)) + assert(!taskScheduler.hasPipelinedTaskSets, + "a regular task set is not a pipelined-group member") + + // One pipelined member anywhere in the map is enough, even next to regular task sets. + taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 1)) + assert(taskScheduler.hasPipelinedTaskSets, + "a live pipelined task set must be reported even alongside regular ones") + } + + test("SPARK-58913: hasPipelinedTaskSets ignores zombie attempts") { + // A zombie (superseded by a retry/kill) attempt no longer schedules tasks, so a stale pipelined + // attempt left in the map must not make the scheduler look like a group is still in flight. + val taskScheduler = setupScheduler() + taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 0, stageAttemptId = 0)) + assert(taskScheduler.hasPipelinedTaskSets) + + taskScheduler.taskSetManagerForAttempt(0, 0).get.isZombie = true + assert(!taskScheduler.hasPipelinedTaskSets, + "a zombie pipelined attempt must not count as a live group member") + + // The live retry of the same stage is a group member again. + taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 0, stageAttemptId = 1)) + assert(taskScheduler.hasPipelinedTaskSets, + "the live retry attempt must be reported even though the zombie attempt is skipped") + } + + test("SPARK-58913: hasPipelinedTaskSets is false once every member's task set has finished") { + // The task scheduler drops a task set once it finishes, so the flag goes false at that point + // even though the job that owns the group is not done yet (the DAGScheduler has still to + // process the final completion). This is the documented narrowing of the task-scheduler view. + val taskScheduler = setupScheduler() + taskScheduler.submitTasks(pipelinedTaskSet(1, stageId = 0, stageAttemptId = 0)) + assert(taskScheduler.hasPipelinedTaskSets) + + val tsm = taskScheduler.taskSetManagerForAttempt(0, 0).get + taskScheduler.taskSetFinished(tsm) + assert(!taskScheduler.hasPipelinedTaskSets, + "a finished task set is no longer held by the task scheduler") + } + + test("error() with no active task sets reports a cluster manager failure") { + val taskScheduler = setupScheduler() + // No task set has been submitted, so `error` cannot abort anything and throws instead. + checkError( + exception = intercept[SparkException] { + taskScheduler.error("Master removed our application: FAILED") + }, + condition = "CLUSTER_MANAGER_APPLICATION_FAILURE", + sqlState = Some("56000"), + parameters = Map("message" -> "Master removed our application: FAILED")) + } + + test("error() with an active task set aborts it instead of throwing") { + val taskScheduler = setupScheduler() + taskScheduler.submitTasks(FakeTask.createTaskSet(1)) + val tsm = taskScheduler.taskSetManagerForAttempt(0, 0).get + taskScheduler.error("Master removed our application: FAILED") + assert(tsm.isZombie) + // The zombie flag cannot show whether the reason survived the trip, so assert that the + // DAGScheduler was notified with the message verbatim. + assert(failedTaskSet) + assert(failedTaskSetReason === "Master removed our application: FAILED") + } + } diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala index 28d774eb27f7e..94f4ca251c6b4 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSetManagerSuite.scala @@ -3046,6 +3046,250 @@ class TaskSetManagerSuite "ExceptionFailure must count towards task failures (contrast)") } + // ========================================================================================== + // Pipelined-group member: group-atomic failure via job-abort + // ========================================================================================== + + /** A single-task TaskSet marked as a pipelined-group member. */ + private def pipelinedTaskSet(): TaskSet = { + val tasks = Array.tabulate[Task[_]](1)(i => new FakeTask(0, i, Nil)) + new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + } + + test("pipelined task set aborts after a single task failure (maxTaskFailures capped at 1)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + // Even though MAX_TASK_FAILURES (4) is passed, a pipelined task set caps attempts at 1. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // A single ordinary task failure must abort the whole task set (group -> job), no retry. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, TaskResultLost) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must abort after the first task failure") + } + + test("pipelined task set counts an executor-loss failure that would normally be uncounted") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = pipelinedTaskSet() + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offerResult = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offerResult.isDefined) + // ExecutorLostFailure with exitCausedByApp=false is normally NOT counted toward task failures, + // so a non-pipelined task set would not abort. For a pipelined task set it counts, so the + // single failure aborts the group. + manager.handleFailedTask(offerResult.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(sched.taskSetsFailed.contains(taskSet.id), + "a pipelined task set must count executor loss and abort") + } + + test("pipelined task set does NOT abort on a benign TaskKilled (losing duplicate attempt)") { + // A pipelined set with 2 tasks. One task gets a duplicate attempt; when one wins, the loser is + // killed with TaskKilled. That benign kill must NOT count as a failure and must NOT abort the + // group (the task actually succeeded). Regression for the speculative-loser abort bug. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // Launch task index 0 and fail it with a benign TaskKilled (as if a duplicate attempt won). + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined && offer.get.index === 0) + manager.handleFailedTask(offer.get.taskId, TaskState.KILLED, + TaskKilled("another attempt succeeded")) + // The group must NOT be aborted: TaskKilled is benign even for a pipelined set. + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a benign TaskKilled must not abort a pipelined group") + } + + test("pipelined task set does NOT abort on TaskCommitDenied (speculation commit race)") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val tasks = Array.tabulate[Task[_]](2)(i => new FakeTask(0, i, Nil)) + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true) + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined) + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, + TaskCommitDenied(jobID = 0, partitionID = 0, attemptNumber = 0)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "TaskCommitDenied must not abort a pipelined group") + } + + test("pipelined task set does NOT single-resubmit a completed map task on executor " + + "decommission") { + // The TaskSetManager.executorLost "Resubmitted" loop re-enqueues an already-successful + // ShuffleMapTask when its executor is lost and the map output looks gone -- for a decommission, + // it looks gone whenever MapOutputTracker.getMapOutputLocation returns None. A pipelined + // shuffle is registered only with the StreamingShuffleOutputTracker, never the + // MapOutputTracker, so getMapOutputLocation is always None and this loop would resubmit the + // lone producer task -- the exact streaming-writer hang. This + // loop bypasses handleFailedTask, so the group-atomic abort would never see it. The guard + // (!taskSet.isPipelined on maybeShuffleMapOutputLoss) must keep a pipelined set out of the + // loop. + // + // Scenario: a 2-task pipelined producer, one task succeeds on the to-be-lost executor while the + // other is still running (so the set is NOT a zombie and the loop is entered for a + // non-pipelined set). Decommission that executor. Assert NO Resubmitted is emitted for the + // completed task. + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("execA", "host1"), ("execB", "host2")) + sched.initialize(new FakeSchedulerBackend()) + + var resubmittedTasks = 0 + val dagScheduler = new FakeDAGScheduler(sc, sched) { + override def taskEnded( + task: Task[_], + reason: TaskEndReason, + result: Any, + accumUpdates: Seq[AccumulatorV2[_, _]], + metricPeaks: Array[Long], + taskInfo: TaskInfo): Unit = { + super.taskEnded(task, reason, result, accumUpdates, metricPeaks, taskInfo) + reason match { + case Resubmitted => resubmittedTasks += 1 + case _ => + } + } + } + sched.dagScheduler.stop() + sched.setDAGScheduler(dagScheduler) + + // Two real ShuffleMapTasks (so isShuffleMapTasks = true), marked as a pipelined-group member. + // Give each task a distinct preferred location so one lands on execA and the other on execB. + val prefs = Array( + Seq[TaskLocation](TaskLocation("host1", "execA")), + Seq[TaskLocation](TaskLocation("host2", "execB"))) + val tasks = Array.tabulate[Task[_]](2) { i => + new ShuffleMapTask(0, 0, null, new Partition { override def index: Int = i }, 1, + prefs(i), JobArtifactSet.getActiveOrDefault(sc), new Properties, null) + } + val taskSet = new TaskSet(tasks, stageId = 0, stageAttemptId = 0, priority = 0, null, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, Some(0), isPipelined = true) + // A pipelined shuffle is deliberately NOT registered in the MapOutputTracker, so a decommission + // lookup of its output returns None (looks lost) -- the trigger for channel 3. + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES) + + // Launch task 0 on execA (PROCESS_LOCAL) and task 1 on execB; complete task 0 so the set is not + // yet a zombie (task 1 still running). + val t0 = manager.resourceOffer("execA", "host1", TaskLocality.PROCESS_LOCAL)._1.get + val t1 = manager.resourceOffer("execB", "host2", TaskLocality.PROCESS_LOCAL)._1.get + assert(manager.runningTasks === 2) + assert(t0.index === 0 && t1.index === 1, "each task should land on its preferred executor") + val result0 = new DirectTaskResult[String]() { + override def value(resultSer: SerializerInstance): String = "" + } + manager.handleSuccessfulTask(t0.taskId, result0) + assert(!manager.isZombie, "the set must not be a zombie (task 1 still running)") + assert(manager.successful(t0.index)) + assert(resubmittedTasks === 0) + + // Decommission execA (which ran the completed task 0) and lose it. For a NON-pipelined set this + // would re-enqueue task 0 via Resubmitted (its output looks lost); for a pipelined set the + // guard must prevent that. + manager.executorDecommission("execA") + manager.executorLost("execA", "host1", ExecutorDecommission()) + + assert(resubmittedTasks === 0, + "a pipelined producer's completed task must NOT be single-resubmitted on decommission") + assert(manager.successful(t0.index), + "the completed task must remain successful (no Resubmitted re-enqueue)") + } + + test("non-pipelined task set is unaffected: executor loss is not counted, retries still apply") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = FakeTask.createTaskSet(1) // isPipelined defaults to false + val clock = new ManualClock + clock.advance(1) + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + // An uncounted executor-loss failure must NOT abort a normal task set. + val offer1 = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer1.isDefined) + manager.handleFailedTask(offer1.get.taskId, TaskState.FINISHED, + ExecutorLostFailure("exec1", exitCausedByApp = false, reason = None)) + assert(!sched.taskSetsFailed.contains(taskSet.id), + "a normal task set must not count uncaused executor loss") + + // And a normal task set still retries a counted failure up to MAX_TASK_FAILURES before abort. + (1 to MAX_TASK_FAILURES).foreach { index => + val offer = manager.resourceOffer("exec1", "host1", ANY)._1 + assert(offer.isDefined, s"expected an offer on iteration $index") + manager.handleFailedTask(offer.get.taskId, TaskState.FINISHED, TaskResultLost) + if (index < MAX_TASK_FAILURES) { + assert(!sched.taskSetsFailed.contains(taskSet.id), + s"must not abort before $MAX_TASK_FAILURES failures (iteration $index)") + } else { + assert(sched.taskSetsFailed.contains(taskSet.id), + "must abort after MAX_TASK_FAILURES counted failures") + } + } + } + + test("resourceOffer attaches current userCredentials to TaskDescription") { + sc = new SparkContext("local", "test") + sched = new FakeTaskScheduler(sc, ("exec1", "host1")) + val taskSet = FakeTask.createTaskSet(3) + val clock = new ManualClock() + val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, clock = clock) + + try { + // Initially no credentials in SparkEnv store + assert(SparkEnv.get.userCredentials.get() == null) + val taskOpt1 = manager.resourceOffer("exec1", "host1", TaskLocality.ANY)._1 + assert(taskOpt1.isDefined) + assert(taskOpt1.get.userCredentials.isEmpty, + "TaskDescription should have None when credential store is empty") + + // Set credentials to version 1 + val v1Bytes = Array[Byte](10, 20, 30, 40, 50) + VersionedCredentials.updateIfNewer(SparkEnv.get.userCredentials, 1L, v1Bytes) + + // Offer another task from the same set -- should carry v1 + val taskOpt2 = manager.resourceOffer("exec1", "host1", TaskLocality.ANY)._1 + assert(taskOpt2.isDefined) + assert(taskOpt2.get.userCredentials.isDefined, + "TaskDescription should carry credentials after store is populated") + assert(taskOpt2.get.userCredentials.get._1 === 1L, + "TaskDescription should carry version 1") + assert(taskOpt2.get.userCredentials.get._2 === v1Bytes, + "TaskDescription should carry the v1 credential bytes") + + // Update store to version 2 + val v2Bytes = Array[Byte](50, 60, 70, 80, 90) + VersionedCredentials.updateIfNewer(SparkEnv.get.userCredentials, 2L, v2Bytes) + + // Offer another task -- should carry v2 + val taskOpt3 = manager.resourceOffer("exec1", "host1", TaskLocality.ANY)._1 + assert(taskOpt3.isDefined) + assert(taskOpt3.get.userCredentials.get._1 === 2L, + "After renewal, TaskDescription should carry version 2") + assert(taskOpt3.get.userCredentials.get._2 === v2Bytes, + "TaskDescription should carry the v2 credential bytes") + } finally { + SparkEnv.get.userCredentials.set(null) + } + } + } class FakeLongTasks(stageId: Int, partitionId: Int) extends FakeTask(stageId, partitionId) { diff --git a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala index 2368900233648..f963d2a815e73 100644 --- a/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala +++ b/core/src/test/scala/org/apache/spark/serializer/KryoSerializerSuite.scala @@ -117,6 +117,9 @@ class KryoSerializerSuite extends SparkFunSuite with SharedSparkContext { check(Array.empty[Int]) check(Array(Array("1", "2"), Array("1", "2", "3", "4"))) check(Array(Array(1.toByte))) + val javaMap = new java.util.HashMap[String, Int]() + javaMap.put("one", 1) + check(javaMap) } test("pairs") { diff --git a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala index da74803a10ebc..8e0f9e04c1bd8 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala @@ -240,6 +240,46 @@ class PipelinedShuffleRoutingSuite extends SparkFunSuite with LocalSparkContext assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) } + test("createShuffleMapStage fails loud for a pipelined dependency when no streaming tracker") { + // A pipelined shuffle is served solely by the StreamingShuffleOutputTracker. With a non- + // streaming incremental manager configured, that tracker is absent (verified above), so a job + // whose stage graph contains a PipelinedShuffleDependency cannot be scheduled. + // createShuffleMapStage must fail loud rather than silently register the shuffle in no tracker + // (which would strand a consumer with no writer locations). This guards the decoupling's + // invariant that a pipelined dependency implies a streaming tracker. + // local[4]: the pipelined group's up-front slot admission (producer 2 + result 2 = 4) must pass + // so that stage creation is actually reached -- otherwise the job is rejected for slots first. + sc = new SparkContext("local[4]", "test", newConf()) + assert(SparkEnv.get.streamingShuffleOutputTracker.isEmpty) + val producer = sc.parallelize(1 to 4, 2).map(x => (x, x)) + val pipelined = new PipelinedShuffleDependency[Int, Int, Int](producer, new HashPartitioner(2)) + // A minimal reduce-side RDD whose single dependency is the pipelined shuffle, so submitting an + // action forces createShuffleMapStage for the pipelined producer. + val consumer = new RDD[(Int, Int)](sc, Seq(pipelined)) { + override def compute(split: Partition, ctx: TaskContext): Iterator[(Int, Int)] = + Iterator.empty + override protected def getPartitions: Array[Partition] = + Array.tabulate(2)(i => new Partition { override def index: Int = i }) + } + val ex = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"expected a fail-loud IllegalStateException about the missing tracker, got: $ex") + + // The fail-loud throw happens BEFORE createShuffleMapStage mutates stageIdToStage / + // shuffleIdToMapStage, so it leaves no partial scheduler state behind. If it left a + // half-created stage cached in shuffleIdToMapStage, a re-submission would reuse that stale + // stage instead of re-throwing. Re-submitting the same job must therefore re-throw the error. + val ex2 = intercept[Exception] { + consumer.count() + } + assert(findCause[IllegalStateException](ex2).exists( + _.getMessage.contains("requires a StreamingShuffleOutputTracker")), + s"a re-submission must re-throw the fail-loud error (no leaked partial stage), got: $ex2") + } + test("spark.shuffle.manager.incremental resolves the same short aliases as the default manager") { // "sort" is a short alias the incremental slot must resolve (to SortShuffleManager) rather than // treat as a class name. SortShuffleManager is blocking, so it is rejected from the pipelined diff --git a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala index e421272cc1199..4cbbf44daa364 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala @@ -25,12 +25,20 @@ import io.netty.channel.{Channel, ChannelConfig, ChannelFuture} import io.netty.util.concurrent.GenericFutureListener import org.mockito.ArgumentMatchers.any import org.mockito.Mockito.when +import org.scalatest.concurrent.Eventually.eventually +import org.scalatest.concurrent.PatienceConfiguration.Timeout import org.scalatest.matchers.should.Matchers +import org.scalatest.time.SpanSugar._ import org.scalatestplus.mockito.MockitoSugar import org.apache.spark._ import org.apache.spark.LocalSparkContext.withSpark -import org.apache.spark.internal.config.{SHUFFLE_MANAGER_INCREMENTAL, STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE} +import org.apache.spark.internal.config.{ + SHUFFLE_MANAGER_INCREMENTAL, + STREAMING_SHUFFLE_CHECKSUM_ENABLED, + STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, + STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, + STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS} import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} import org.apache.spark.metrics.MetricsSystem import org.apache.spark.network.client.TransportClient @@ -229,6 +237,69 @@ class StreamingShuffleWriterSuite } } + test("writer fails when a reader does not connect before timeout") { + val conf = newConf() + .set(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS, 100L) + withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf)) { sc => + val context = createTaskContext(sc.conf, 0) + try { + val writer = newWriter(sc, context) + val error = intercept[SparkRuntimeException] { + writer.write(Iterator.empty) + } + checkError( + exception = error, + condition = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT", + sqlState = "XXKST", + parameters = Map( + "shuffleId" -> "0", + "writerId" -> "0", + "readerId" -> "0", + "timeoutMs" -> "100")) + } finally { + context.markTaskCompleted(None) + } + } + } + + test("writer surfaces a reader connection timeout from the flush thread") { + val timeoutMs = 100L + val flushIntervalMs = 200L + val conf = newConf() + .set(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS, timeoutMs) + .set(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, flushIntervalMs) + withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf)) { sc => + val context = createTaskContext(sc.conf, 0) + try { + val writer = newWriter(sc, context) + val records = Iterator((1, 1), (2, 2)).map { record => + if (record._1 == 2) { + // Block before yielding the second record until the flush thread observes the + // connection timeout while sending the first record. + eventually(Timeout(10.seconds)) { + writer.errorNotifier.getError() shouldBe defined + } + } + record + } + val error = intercept[SparkRuntimeException] { + writer.write(records) + } + checkError( + exception = error, + condition = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT", + sqlState = "XXKST", + parameters = Map( + "shuffleId" -> "0", + "writerId" -> "0", + "readerId" -> "0", + "timeoutMs" -> timeoutMs.toString)) + } finally { + context.markTaskCompleted(None) + } + } + } + test("cleanupResources releases queued pooled buffers") { withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", newConf())) { sc => val context = createTaskContext(sc.conf, 0) diff --git a/core/src/test/scala/org/apache/spark/storage/BlockIdSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockIdSuite.scala index 4b7f1fafb211c..259b38c7f83d2 100644 --- a/core/src/test/scala/org/apache/spark/storage/BlockIdSuite.scala +++ b/core/src/test/scala/org/apache/spark/storage/BlockIdSuite.scala @@ -51,6 +51,16 @@ class BlockIdSuite extends SparkFunSuite { assertSame(id, BlockId(id.toString)) } + test("SPARK-41246: rdd with a negative id") { + // Safety net: fail-fast prevents minting negative ids, but BlockId must still + // parse names from pre-upgrade caches or tests. + val id = RDDBlockId(-1330910599, 36) + assert(id.name === "rdd_-1330910599_36") + assertSame(id, BlockId(id.name)) + assertDifferent(id, RDDBlockId(1330910599, 36)) + assertSame(RDDBlockId(Int.MinValue, 0), BlockId("rdd_-2147483648_0")) + } + test("shuffle") { val id = ShuffleBlockId(1, 2, 3) assertSame(id, ShuffleBlockId(1, 2, 3)) diff --git a/core/src/test/scala/org/apache/spark/storage/BlockInfoManagerSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockInfoManagerSuite.scala index 4b34c13706ad8..016c49c2d1a4b 100644 --- a/core/src/test/scala/org/apache/spark/storage/BlockInfoManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/storage/BlockInfoManagerSuite.scala @@ -313,9 +313,12 @@ class BlockInfoManagerSuite extends SparkFunSuite { test("removing a non-existent block throws SparkException") { withTaskId(0) { - intercept[SparkException] { - blockInfoManager.removeBlock("non-existent-block") - } + checkError( + exception = intercept[SparkException] { + blockInfoManager.removeBlock("non-existent-block") + }, + condition = "INTERNAL_ERROR_STORAGE", + parameters = Map("message" -> "Block test_non-existent-block does not exist.")) } } diff --git a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala index e2e2552270f77..f60e7800ab49d 100644 --- a/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala +++ b/core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala @@ -1781,9 +1781,13 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe // Because the BlockManager's metadata claims that the block exists (i.e. that it's present // in at least one store), the read attempts to read it and fails when the on-disk file is // missing. - intercept[SparkException] { - readMethod(store) - } + checkError( + exception = intercept[SparkException] { + readMethod(store) + }, + condition = "LOCAL_BLOCK_DATA_NOT_FOUND", + sqlState = Some("58030"), + parameters = Map("blockId" -> "test_blockId")) // Subsequent read attempts will succeed; the block isn't present but we return an expected // "block not found" response rather than a fatal error: assert(readMethod(store).isEmpty) @@ -2065,7 +2069,16 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe val exception = intercept[SparkException] { bm.putBlockDataAsStream(shuffleBlockId, StorageLevel.DISK_ONLY, ClassTag(message.getClass)) } - assert(exception.getMessage.contains("unsupported shuffle resolver")) + checkError( + exception = exception, + condition = "SHUFFLE_BLOCK_MIGRATION_NOT_SUPPORTED", + sqlState = Some("0A000"), + parameters = Map( + "blockId" -> shuffleBlockId.toString, + "resolverClass" -> badShuffleResolver.getClass.getName)) + // The `ClassCastException` is kept as the cause. The `try` only covers building the callback, + // so this pins the cast failure itself rather than anything the resolver does while writing. + assert(exception.getCause.isInstanceOf[ClassCastException]) } test("SPARK-54796: putBlockDataAsStream throws ShuffleManagerNotInitializedException " + @@ -2401,10 +2414,18 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe conf.set(SHUFFLE_SERVICE_PORT.key, shufflePort.toString) conf.set(SHUFFLE_REGISTRATION_TIMEOUT.key, "40") conf.set(SHUFFLE_REGISTRATION_MAX_ATTEMPTS.key, "1") - val e = intercept[SparkException] { - makeBlockManager(8000, "timeoutExec") - }.getMessage - assert(e.contains("TimeoutException")) + // `matchPVals` because the parameter is the text of the `RuntimeException` that + // `TransportClient.sendRpcSync` wraps Guava's `TimeoutException` in (`cause.toString`). + // Only the class-name prefix of that text is stable; the rest is a Guava internal that + // varies run to run (e.g. the future's identity hash and a scheduling-delay clause). + checkError( + exception = intercept[SparkException] { + makeBlockManager(8000, "timeoutExec") + }, + condition = "UNABLE_TO_REGISTER_WITH_EXTERNAL_SHUFFLE_SERVICE", + sqlState = Some("58030"), + parameters = Map("message" -> "(?s)java\\.util\\.concurrent\\.TimeoutException: .*"), + matchPVals = true) verify(master, times(0)) .registerBlockManager(mc.any(), mc.any(), mc.any(), mc.any(), mc.any(), mc.any()) server.close() diff --git a/core/src/test/scala/org/apache/spark/storage/DiskBlockObjectWriterSuite.scala b/core/src/test/scala/org/apache/spark/storage/DiskBlockObjectWriterSuite.scala index 4352436c872fe..f103f3c305cf3 100644 --- a/core/src/test/scala/org/apache/spark/storage/DiskBlockObjectWriterSuite.scala +++ b/core/src/test/scala/org/apache/spark/storage/DiskBlockObjectWriterSuite.scala @@ -21,7 +21,7 @@ import java.nio.ByteBuffer import scala.reflect.ClassTag -import org.apache.spark.{SparkConf, SparkException, SparkFunSuite} +import org.apache.spark.{SparkConf, SparkException, SparkFunSuite, SparkUnsupportedOperationException} import org.apache.spark.executor.ShuffleWriteMetrics import org.apache.spark.serializer.{DeserializationStream, JavaSerializer, SerializationStream, Serializer, SerializerInstance, SerializerManager} import org.apache.spark.util.Utils @@ -214,6 +214,18 @@ class DiskBlockObjectWriterSuite extends SparkFunSuite { assert(bs.isClosed) assert(objOut.isClosed) } + + test("write(b: Int) is not supported") { + val (writer, _, _) = createWriter() + checkError( + exception = intercept[SparkUnsupportedOperationException] { + writer.write(1) + }, + condition = "UNSUPPORTED_CALL.WITHOUT_SUGGESTION", + parameters = Map( + "className" -> "org.apache.spark.storage.DiskBlockObjectWriter", + "methodName" -> "write")) + } } trait CloseDetecting { diff --git a/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala b/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala index 9096076f8212a..fde52ed57a3d0 100644 --- a/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala +++ b/core/src/test/scala/org/apache/spark/ui/UISeleniumSuite.scala @@ -302,6 +302,59 @@ class UISeleniumSuite extends SparkFunSuite with WebBrowser with Matchers { } } + test("SPARK-58828: spark.ui.holdEnabled should properly control hold button display") { + def hasHoldLink: Boolean = find(className("confirm-link")).isDefined + // The control needs a coarse-grained backend (so local-cluster, not local) and the hold + // preconditions; the external shuffle service itself is not needed to render the page. + val holdConfs = Map( + SHUFFLE_SERVICE_ENABLED.key -> "true", + DECOMMISSION_ENABLED.key -> "true") + + withSpark(newSparkContext(master = "local-cluster[1,1,1024]", + additionalConfs = holdConfs)) { sc => + eventually(timeout(10.seconds), interval(50.milliseconds)) { + goToUi(sc, "/jobs") + assert(hasHoldLink) + } + } + + withSpark(newSparkContext(master = "local-cluster[1,1,1024]", + additionalConfs = holdConfs + (UI_HOLD_ENABLED.key -> "false"))) { sc => + goToUi(sc, "/jobs") + assert(!hasHoldLink) + } + } + + test("SPARK-59010: hold status api") { + def holdStatus(sc: SparkContext): (Boolean, Boolean, Int) = { + val json = getJson(sc.ui.get, "holdstatus") + ((json \ "supported").extract[Boolean], (json \ "held").extract[Boolean], + (json \ "draining").extract[Int]) + } + + // A local backend cannot hold its executors. + withSpark(newSparkContext()) { sc => + assert(holdStatus(sc) === (false, false, 0)) + } + + withSpark(newSparkContext(master = "local-cluster[1,1,1024]", + additionalConfs = Map( + SHUFFLE_SERVICE_ENABLED.key -> "true", + DECOMMISSION_ENABLED.key -> "true"))) { sc => + assert(holdStatus(sc) === (true, false, 0)) + // Wait for the executor to register, so that the hold has something to drain. + eventually(timeout(30.seconds), interval(200.milliseconds)) { + assert(sc.getExecutorIds().nonEmpty) + } + assert(sc.holdExecutors()) + assert(holdStatus(sc)._2) + // The held executors exit once they are done, and the drain is then complete. + eventually(timeout(30.seconds), interval(200.milliseconds)) { + assert(holdStatus(sc) === (true, true, 0)) + } + } + } + test("jobs page should not display job group name unless some job was submitted in a job group") { withSpark(newSparkContext()) { sc => // If no job has been run in a job group, then "(Job Group)" should not appear in the header diff --git a/core/src/test/scala/org/apache/spark/ui/UISuite.scala b/core/src/test/scala/org/apache/spark/ui/UISuite.scala index 9bd649994258a..7a6f409c04efc 100644 --- a/core/src/test/scala/org/apache/spark/ui/UISuite.scala +++ b/core/src/test/scala/org/apache/spark/ui/UISuite.scala @@ -534,9 +534,39 @@ class UISuite extends SparkFunSuite { lastRequest = req res.sendError(HttpServletResponse.SC_OK) } - } + test("SPARK-58521: createProxyHandler forwards X-Forwarded-Context header") { + val (conf, securityMgr, sslOptions) = sslDisabledConf() + val targetServer = JettyUtils.startJettyServer("0.0.0.0", 0, sslOptions, conf) + val proxyServer = JettyUtils.startJettyServer("0.0.0.0", 0, sslOptions, conf) + + @volatile var forwardedContext: String = null + val targetHandler = new ServletContextHandler() + targetHandler.setContextPath("/") + targetHandler.addServlet(new ServletHolder(new HttpServlet { + override def doGet(req: HttpServletRequest, resp: HttpServletResponse): Unit = { + forwardedContext = req.getHeader("X-Forwarded-Context") + resp.setStatus(HttpServletResponse.SC_OK) + } + }), "/*") + targetServer.addHandler(targetHandler, securityMgr) + + val targetAddr = s"http://$localhost:${targetServer.boundPort}" + val proxyHandler = JettyUtils.createProxyHandler(_ => Some(targetAddr)) + proxyServer.addHandler(proxyHandler, securityMgr) + + try { + val proxyUrl = s"http://$localhost:${proxyServer.boundPort}/proxy/app-123/stages/" + TestUtils.withHttpConnection(new URI(proxyUrl).toURL) { conn => + assert(conn.getResponseCode === HttpServletResponse.SC_OK) + assert(forwardedContext === "/proxy/app-123") + } + } finally { + stopServer(proxyServer) + stopServer(targetServer) + } + } } // Filter for testing; returns a configurable code for every request. diff --git a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala index e1e49a226735c..c467b7f0189c9 100644 --- a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala @@ -1486,6 +1486,15 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties { assert(Utils.buildLocationMetadata(paths, 18) == "(5 paths)[path0, path1, ...]") } + test("SPARK-58748: wildcard IPv4 and IPv6 bind addresses are not advertised driver addresses") { + Seq("0.0.0.0", "::", "[::]", "0:0:0:0:0:0:0:0").foreach { address => + assert(Utils.isAnyLocalAddress(address), s"$address should be a wildcard address") + } + Seq("10.129.36.37", "::1", "[::1]", "2001:db8::1", "localhost").foreach { address => + assert(!Utils.isAnyLocalAddress(address), s"$address should not be a wildcard address") + } + } + test("checkHost supports both IPV4 and IPV6") { // IPV4 ips Utils.checkHost("0.0.0.0") diff --git a/core/src/test/scala/org/apache/spark/util/collection/OpenHashMapBenchmark.scala b/core/src/test/scala/org/apache/spark/util/collection/OpenHashMapBenchmark.scala new file mode 100644 index 0000000000000..4aa5ee8b4b84d --- /dev/null +++ b/core/src/test/scala/org/apache/spark/util/collection/OpenHashMapBenchmark.scala @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.util.collection + +import java.util.{HashMap => JHashMap, Random} + +import org.apache.spark.benchmark.{Benchmark, BenchmarkBase} + +/** + * Benchmark for OpenHashMap vs java.util.HashMap. + * Measures insert, aggregate (changeValue/merge), and random-order lookup performance + * with String keys and Long values. + * {{{ + * To run this benchmark: + * 1. without sbt: bin/spark-submit --class <this class> <spark core test jar> + * 2. build/sbt "core/Test/runMain <this class>" + * 3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "core/Test/runMain <this class>" + * Results will be written to "benchmarks/OpenHashMapBenchmark-results.txt". + * }}} + */ +object OpenHashMapBenchmark extends BenchmarkBase { + + private val numKeys = 1000000 + private val numAggOps = 5000000 + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runBenchmark("OpenHashMap vs java.util.HashMap") { + insertBenchmark() + aggregateBenchmark() + lookupBenchmark() + } + } + + private def stringKeys: Array[String] = Array.tabulate(numKeys)(i => "key_" + i) + + private def insertBenchmark(): Unit = { + val keys = stringKeys + val benchmark = new Benchmark(s"Insert $numKeys distinct String keys", numKeys, + output = output) + benchmark.addCase("OpenHashMap") { _ => + val map = new OpenHashMap[String, Long](64) + var i = 0 + while (i < numKeys) { + map.update(keys(i), i.toLong) + i += 1 + } + } + benchmark.addCase("java.util.HashMap") { _ => + val map = new JHashMap[String, java.lang.Long](16) + var i = 0 + while (i < numKeys) { + map.put(keys(i), i.toLong) + i += 1 + } + } + benchmark.run() + } + + private def aggregateBenchmark(): Unit = { + val keys = stringKeys + val random = new Random(42) + val aggIndices = Array.fill(numAggOps)(random.nextInt(numKeys)) + val benchmark = new Benchmark(s"Aggregate $numAggOps ops on $numKeys String keys", numAggOps, + output = output) + benchmark.addCase("OpenHashMap changeValue") { _ => + val map = new OpenHashMap[String, Long](64) + var i = 0 + while (i < numAggOps) { + map.changeValue(keys(aggIndices(i)), 1L, _ + 1L) + i += 1 + } + } + benchmark.addCase("java.util.HashMap merge") { _ => + val map = new JHashMap[String, java.lang.Long](16) + var i = 0 + while (i < numAggOps) { + map.merge(keys(aggIndices(i)), 1L, (a, b) => a + b) + i += 1 + } + } + benchmark.run() + } + + private def lookupBenchmark(): Unit = { + val keys = stringKeys + val random = new Random(42) + // Look up in random order so that neither map benefits from the memory layout produced by + // sequential insertion. + val lookupIndices = Array.fill(numKeys)(random.nextInt(numKeys)) + val openHashMap = new OpenHashMap[String, Long](64) + val jHashMap = new JHashMap[String, java.lang.Long](16) + var i = 0 + while (i < numKeys) { + openHashMap.update(keys(i), i.toLong) + jHashMap.put(keys(i), i.toLong) + i += 1 + } + val benchmark = new Benchmark(s"Look up $numKeys String keys in random order", numKeys, + output = output) + benchmark.addCase("OpenHashMap") { _ => + var sum = 0L + var i = 0 + while (i < numKeys) { + sum += openHashMap(keys(lookupIndices(i))) + i += 1 + } + } + benchmark.addCase("java.util.HashMap") { _ => + var sum = 0L + var i = 0 + while (i < numKeys) { + sum += jHashMap.get(keys(lookupIndices(i))) + i += 1 + } + } + benchmark.run() + } +} diff --git a/dev/check-protos.py b/dev/check-protos.py index 4ddd1f1058820..42ad3e005ed14 100755 --- a/dev/check-protos.py +++ b/dev/check-protos.py @@ -20,11 +20,11 @@ # Utility for checking whether generated codes in PySpark are out of sync. # usage: ./dev/check-protos.py +import filecmp import os +import subprocess import sys -import filecmp import tempfile -import subprocess # Location of your Spark git development area SPARK_HOME = os.environ.get("SPARK_HOME", os.getcwd()) diff --git a/dev/create-release/generate-contributors.py b/dev/create-release/generate-contributors.py index e32d81838e83d..69ae19b859be3 100755 --- a/dev/create-release/generate-contributors.py +++ b/dev/create-release/generate-contributors.py @@ -22,13 +22,12 @@ import sys from github import Github - from releaseutils import ( - tag_exists, - get_commits, - yesOrNoPrompt, contributors_file_name, + get_commits, get_github_name, + tag_exists, + yesOrNoPrompt, ) # You must set the following before use! diff --git a/dev/create-release/generate-llms-txt.py b/dev/create-release/generate-llms-txt.py index 6175a2e3f3985..593c48924cf33 100755 --- a/dev/create-release/generate-llms-txt.py +++ b/dev/create-release/generate-llms-txt.py @@ -18,8 +18,8 @@ # # This script generates llms.txt file for Apache Spark documentation -import sys import argparse +import sys from pathlib import Path diff --git a/dev/create-release/release-build.sh b/dev/create-release/release-build.sh index 161dca83738b3..cdeec37c4eddf 100755 --- a/dev/create-release/release-build.sh +++ b/dev/create-release/release-build.sh @@ -40,7 +40,7 @@ SPARK_VERSION - (optional) Version of Spark being built (e.g. 2.1.2) ASF_USERNAME - Username of ASF committer account ASF_PASSWORD - Password of ASF committer account -ASF_NEXUS_TOKEN - API token in ASF Nexus reposiotry +ASF_NEXUS_TOKEN - API token in ASF Nexus repository GPG_KEY - GPG key used to sign release artifacts GPG_PASSPHRASE - Passphrase for GPG key @@ -683,7 +683,7 @@ SCALA_2_12_PROFILES="-Pscala-2.12" HIVE_PROFILES="-Phive -Phive-thriftserver" # Profiles for publishing snapshots and release to Maven Central # We use Apache Hive 2.3 for publishing -PUBLISH_PROFILES="$BASE_PROFILES $HIVE_PROFILES -Pspark-ganglia-lgpl -Pkinesis-asl -Phadoop-cloud -Pjvm-profiler" +PUBLISH_PROFILES="$BASE_PROFILES $HIVE_PROFILES -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Phadoop-cloud -Pjvm-profiler" # Profiles for building binary releases BASE_RELEASE_PROFILES="$BASE_PROFILES -Psparkr" @@ -697,14 +697,6 @@ elif [[ $JAVA_VERSION < "17.0." ]] && [[ $SPARK_VERSION > "3.5.99" ]]; then exit 1 fi -# This is a band-aid fix to avoid the failure of Maven nightly snapshot in some Jenkins -# machines by explicitly calling /usr/sbin/lsof. Please see SPARK-22377 and the discussion -# in its pull request. -LSOF=lsof -if ! hash $LSOF 2>/dev/null; then - LSOF=/usr/sbin/lsof -fi - if [ -z "$SPARK_PACKAGE_VERSION" ]; then SPARK_PACKAGE_VERSION="${SPARK_VERSION}-$(date +%Y_%m_%d_%H_%M)-${git_hash}" fi diff --git a/dev/create-release/releaseutils.py b/dev/create-release/releaseutils.py index 728e58426b3d2..383a574e47aa7 100755 --- a/dev/create-release/releaseutils.py +++ b/dev/create-release/releaseutils.py @@ -20,11 +20,13 @@ import re import sys -from subprocess import Popen, PIPE +from subprocess import PIPE, Popen try: - from github import Github # noqa: F401 - from github import GithubException + from github import ( + Github, # noqa: F401 + GithubException, + ) except ImportError: print("This tool requires the PyGithub library") print("Install using 'pip install PyGithub'") diff --git a/dev/deps/spark-deps-hadoop-3-hive-2.3 b/dev/deps/spark-deps-hadoop-3-hive-2.3 index 17c7f3f2db378..482262ad3793b 100644 --- a/dev/deps/spark-deps-hadoop-3-hive-2.3 +++ b/dev/deps/spark-deps-hadoop-3-hive-2.3 @@ -1,3 +1,7 @@ +# +# Generated by `dev/test-dependencies.sh --replace-manifest`. +# Do not edit manually. +# HdrHistogram/2.1.12//HdrHistogram-2.1.12.jar HikariCP/2.5.1//HikariCP-2.5.1.jar JLargeArrays/1.5//JLargeArrays-1.5.jar @@ -12,9 +16,11 @@ aliyun-java-sdk-kms/2.11.0//aliyun-java-sdk-kms-2.11.0.jar aliyun-java-sdk-ram/3.1.0//aliyun-java-sdk-ram-3.1.0.jar aliyun-sdk-oss/3.18.1//aliyun-sdk-oss-3.18.1.jar analyticsaccelerator-s3/1.3.1//analyticsaccelerator-s3-1.3.1.jar +annotations/2.35.4//annotations-2.35.4.jar antlr-runtime/3.5.2//antlr-runtime-3.5.2.jar antlr4-runtime/4.13.1//antlr4-runtime-4.13.1.jar aopalliance-repackaged/3.0.6//aopalliance-repackaged-3.0.6.jar +apache-client/2.35.4//apache-client-2.35.4.jar arpack/3.2.0//arpack-3.2.0.jar arpack_combined_all/0.1//arpack_combined_all-0.1.jar arrow-compression/19.0.0//arrow-compression-19.0.0.jar @@ -24,9 +30,12 @@ arrow-memory-netty-buffer-patch/19.0.0//arrow-memory-netty-buffer-patch-19.0.0.j arrow-memory-netty/19.0.0//arrow-memory-netty-19.0.0.jar arrow-vector/19.0.0//arrow-vector-19.0.0.jar audience-annotations/0.12.0//audience-annotations-0.12.0.jar +auth/2.35.4//auth-2.35.4.jar avro-ipc/1.12.1//avro-ipc-1.12.1.jar avro-mapred/1.12.1//avro-mapred-1.12.1.jar avro/1.12.1//avro-1.12.1.jar +aws-core/2.35.4//aws-core-2.35.4.jar +aws-query-protocol/2.35.4//aws-query-protocol-2.35.4.jar azure-data-lake-store-sdk/2.3.9//azure-data-lake-store-sdk-2.3.9.jar azure-keyvault-core/1.0.0//azure-keyvault-core-1.0.0.jar azure-storage/7.0.1//azure-storage-7.0.1.jar @@ -35,12 +44,14 @@ breeze-macros_2.13/2.1.0//breeze-macros_2.13-2.1.0.jar breeze_2.13/2.1.0//breeze_2.13-2.1.0.jar bundle/2.35.4//bundle-2.35.4.jar cats-kernel_2.13/2.8.0//cats-kernel_2.13-2.8.0.jar +checksums-spi/2.35.4//checksums-spi-2.35.4.jar +checksums/2.35.4//checksums-2.35.4.jar chill-java/0.10.0//chill-java-0.10.0.jar chill_2.13/0.10.0//chill_2.13-0.10.0.jar commons-cli/1.11.0//commons-cli-1.11.0.jar -commons-codec/1.22.0//commons-codec-1.22.0.jar -commons-collections4/4.5.0//commons-collections4-4.5.0.jar -commons-compiler/3.1.9//commons-compiler-3.1.9.jar +commons-codec/1.22.1//commons-codec-1.22.1.jar +commons-collections4/4.6.0//commons-collections4-4.6.0.jar +commons-compiler/3.1.12//commons-compiler-3.1.12.jar commons-compress/1.28.0//commons-compress-1.28.0.jar commons-crypto/1.1.0//commons-crypto-1.1.0.jar commons-dbcp/1.4//commons-dbcp-1.4.jar @@ -64,7 +75,8 @@ derbyshared/10.16.1.1//derbyshared-10.16.1.1.jar derbytools/10.16.1.1//derbytools-10.16.1.1.jar dom4j/2.1.4//dom4j-2.1.4.jar dropwizard-metrics-hadoop-metrics2-reporter/0.1.2//dropwizard-metrics-hadoop-metrics2-reporter-0.1.2.jar -esdk-obs-java/3.20.4.2//esdk-obs-java-3.20.4.2.jar +endpoints-spi/2.35.4//endpoints-spi-2.35.4.jar +eventstream/1.0.1//eventstream-1.0.1.jar failureaccess/1.0.3//failureaccess-1.0.3.jar flatbuffers-java/25.2.10//flatbuffers-java-25.2.10.jar gmetric4j/1.0.10//gmetric4j-1.0.10.jar @@ -79,7 +91,6 @@ hadoop-client-api/3.5.0//hadoop-client-api-3.5.0.jar hadoop-client-runtime/3.5.0//hadoop-client-runtime-3.5.0.jar hadoop-cloud-storage/3.5.0//hadoop-cloud-storage-3.5.0.jar hadoop-gcp/3.5.0//hadoop-gcp-3.5.0.jar -hadoop-huaweicloud/3.5.0//hadoop-huaweicloud-3.5.0.jar hadoop-shaded-guava/1.5.0//hadoop-shaded-guava-1.5.0.jar hive-beeline/2.3.10//hive-beeline-2.3.10.jar hive-cli/2.3.10//hive-cli-2.3.10.jar @@ -97,19 +108,25 @@ hive-storage-api/2.8.1//hive-storage-api-2.8.1.jar hk2-api/3.0.6//hk2-api-3.0.6.jar hk2-locator/3.0.6//hk2-locator-3.0.6.jar hk2-utils/3.0.6//hk2-utils-3.0.6.jar +http-auth-aws-eventstream/2.35.4//http-auth-aws-eventstream-2.35.4.jar +http-auth-aws/2.35.4//http-auth-aws-2.35.4.jar +http-auth-spi/2.35.4//http-auth-spi-2.35.4.jar +http-auth/2.35.4//http-auth-2.35.4.jar +http-client-spi/2.35.4//http-client-spi-2.35.4.jar httpclient/4.5.14//httpclient-4.5.14.jar httpcore/4.4.16//httpcore-4.4.16.jar icu4j/78.3//icu4j-78.3.jar +identity-spi/2.35.4//identity-spi-2.35.4.jar ini4j/0.5.4//ini4j-0.5.4.jar istack-commons-runtime/4.1.2//istack-commons-runtime-4.1.2.jar ivy/2.5.3//ivy-2.5.3.jar jackson-annotations/2.22//jackson-annotations-2.22.jar -jackson-core/2.22.0//jackson-core-2.22.0.jar -jackson-databind/2.22.0//jackson-databind-2.22.0.jar -jackson-dataformat-cbor/2.22.0//jackson-dataformat-cbor-2.22.0.jar -jackson-dataformat-yaml/2.22.0//jackson-dataformat-yaml-2.22.0.jar -jackson-datatype-jsr310/2.22.0//jackson-datatype-jsr310-2.22.0.jar -jackson-module-scala_2.13/2.22.0//jackson-module-scala_2.13-2.22.0.jar +jackson-core/2.22.1//jackson-core-2.22.1.jar +jackson-databind/2.22.1//jackson-databind-2.22.1.jar +jackson-dataformat-cbor/2.22.1//jackson-dataformat-cbor-2.22.1.jar +jackson-dataformat-yaml/2.22.1//jackson-dataformat-yaml-2.22.1.jar +jackson-datatype-jsr310/2.22.1//jackson-datatype-jsr310-2.22.1.jar +jackson-module-scala_2.13/2.22.1//jackson-module-scala_2.13-2.22.1.jar jakarta.activation-api/2.1.4//jakarta.activation-api-2.1.4.jar jakarta.annotation-api/2.1.1//jakarta.annotation-api-2.1.1.jar jakarta.inject-api/2.0.1//jakarta.inject-api-2.0.1.jar @@ -117,10 +134,9 @@ jakarta.servlet-api/6.0.0//jakarta.servlet-api-6.0.0.jar jakarta.validation-api/3.0.2//jakarta.validation-api-3.0.2.jar jakarta.ws.rs-api/3.1.0//jakarta.ws.rs-api-3.1.0.jar jakarta.xml.bind-api/4.0.5//jakarta.xml.bind-api-4.0.5.jar -janino/3.1.9//janino-3.1.9.jar +janino/3.1.12//janino-3.1.12.jar java-diff-utils/4.16//java-diff-utils-4.16.jar java-trace-api/0.2.11-beta//java-trace-api-0.2.11-beta.jar -java-xmlbuilder/1.2//java-xmlbuilder-1.2.jar javassist/3.30.2-GA//javassist-3.30.2-GA.jar javax.jdo/3.2.0-m3//javax.jdo-3.2.0-m3.jar javax.servlet-api/4.0.1//javax.servlet-api-4.0.1.jar @@ -144,6 +160,7 @@ jline/2.14.6//jline-2.14.6.jar jline/3.29.0/jdk8/jline-3.29.0-jdk8.jar joda-time/2.14.3//joda-time-2.14.3.jar jpam/1.1//jpam-1.1.jar +json-utils/2.35.4//json-utils-2.35.4.jar json/1.8//json-1.8.jar json4s-ast_2.13/4.0.7//json4s-ast_2.13-4.0.7.jar json4s-core_2.13/4.0.7//json4s-core_2.13-4.0.7.jar @@ -184,50 +201,50 @@ lapack/3.2.0//lapack-3.2.0.jar leveldbjni-all/1.8//leveldbjni-all-1.8.jar libfb303/0.9.3//libfb303-0.9.3.jar libthrift/0.16.0//libthrift-0.16.0.jar -log4j-1.2-api/2.26.0//log4j-1.2-api-2.26.0.jar -log4j-api/2.26.0//log4j-api-2.26.0.jar -log4j-core/2.26.0//log4j-core-2.26.0.jar -log4j-layout-template-json/2.26.0//log4j-layout-template-json-2.26.0.jar -log4j-slf4j2-impl/2.26.0//log4j-slf4j2-impl-2.26.0.jar -lz4-java/1.11.1//lz4-java-1.11.1.jar +log4j-1.2-api/2.26.1//log4j-1.2-api-2.26.1.jar +log4j-api/2.26.1//log4j-api-2.26.1.jar +log4j-core/2.26.1//log4j-core-2.26.1.jar +log4j-layout-template-json/2.26.1//log4j-layout-template-json-2.26.1.jar +log4j-slf4j2-impl/2.26.1//log4j-slf4j2-impl-2.26.1.jar +lz4-java/1.11.2//lz4-java-1.11.2.jar metrics-core/4.2.37//metrics-core-4.2.37.jar metrics-graphite/4.2.37//metrics-graphite-4.2.37.jar metrics-jmx/4.2.37//metrics-jmx-4.2.37.jar metrics-json/4.2.37//metrics-json-4.2.37.jar metrics-jvm/4.2.37//metrics-jvm-4.2.37.jar +metrics-spi/2.35.4//metrics-spi-2.35.4.jar minlog/1.3.0//minlog-1.3.0.jar -netty-all/4.2.16.Final//netty-all-4.2.16.Final.jar -netty-buffer/4.2.16.Final//netty-buffer-4.2.16.Final.jar -netty-codec-base/4.2.16.Final//netty-codec-base-4.2.16.Final.jar -netty-codec-compression/4.2.16.Final//netty-codec-compression-4.2.16.Final.jar -netty-codec-dns/4.2.16.Final//netty-codec-dns-4.2.16.Final.jar -netty-codec-http/4.2.16.Final//netty-codec-http-4.2.16.Final.jar -netty-codec-http2/4.2.16.Final//netty-codec-http2-4.2.16.Final.jar -netty-codec-socks/4.2.16.Final//netty-codec-socks-4.2.16.Final.jar -netty-codec/4.2.16.Final//netty-codec-4.2.16.Final.jar -netty-common/4.2.16.Final//netty-common-4.2.16.Final.jar -netty-handler-proxy/4.2.16.Final//netty-handler-proxy-4.2.16.Final.jar -netty-handler/4.2.16.Final//netty-handler-4.2.16.Final.jar -netty-resolver-dns/4.2.16.Final//netty-resolver-dns-4.2.16.Final.jar -netty-resolver/4.2.16.Final//netty-resolver-4.2.16.Final.jar +netty-all/4.2.17.Final//netty-all-4.2.17.Final.jar +netty-buffer/4.2.17.Final//netty-buffer-4.2.17.Final.jar +netty-codec-base/4.2.17.Final//netty-codec-base-4.2.17.Final.jar +netty-codec-compression/4.2.17.Final//netty-codec-compression-4.2.17.Final.jar +netty-codec-dns/4.2.17.Final//netty-codec-dns-4.2.17.Final.jar +netty-codec-http/4.2.17.Final//netty-codec-http-4.2.17.Final.jar +netty-codec-http2/4.2.17.Final//netty-codec-http2-4.2.17.Final.jar +netty-codec-socks/4.2.17.Final//netty-codec-socks-4.2.17.Final.jar +netty-codec/4.2.17.Final//netty-codec-4.2.17.Final.jar +netty-common/4.2.17.Final//netty-common-4.2.17.Final.jar +netty-handler-proxy/4.2.17.Final//netty-handler-proxy-4.2.17.Final.jar +netty-handler/4.2.17.Final//netty-handler-4.2.17.Final.jar +netty-nio-client/2.35.4//netty-nio-client-2.35.4.jar +netty-resolver-dns/4.2.17.Final//netty-resolver-dns-4.2.17.Final.jar +netty-resolver/4.2.17.Final//netty-resolver-4.2.17.Final.jar netty-tcnative-boringssl-static/2.0.81.Final/linux-aarch_64/netty-tcnative-boringssl-static-2.0.81.Final-linux-aarch_64.jar netty-tcnative-boringssl-static/2.0.81.Final/linux-x86_64/netty-tcnative-boringssl-static-2.0.81.Final-linux-x86_64.jar netty-tcnative-boringssl-static/2.0.81.Final/osx-aarch_64/netty-tcnative-boringssl-static-2.0.81.Final-osx-aarch_64.jar netty-tcnative-boringssl-static/2.0.81.Final/osx-x86_64/netty-tcnative-boringssl-static-2.0.81.Final-osx-x86_64.jar netty-tcnative-boringssl-static/2.0.81.Final/windows-x86_64/netty-tcnative-boringssl-static-2.0.81.Final-windows-x86_64.jar netty-tcnative-classes/2.0.81.Final//netty-tcnative-classes-2.0.81.Final.jar -netty-transport-classes-epoll/4.2.16.Final//netty-transport-classes-epoll-4.2.16.Final.jar -netty-transport-classes-kqueue/4.2.16.Final//netty-transport-classes-kqueue-4.2.16.Final.jar -netty-transport-native-epoll/4.2.16.Final/linux-aarch_64/netty-transport-native-epoll-4.2.16.Final-linux-aarch_64.jar -netty-transport-native-epoll/4.2.16.Final/linux-riscv64/netty-transport-native-epoll-4.2.16.Final-linux-riscv64.jar -netty-transport-native-epoll/4.2.16.Final/linux-x86_64/netty-transport-native-epoll-4.2.16.Final-linux-x86_64.jar -netty-transport-native-kqueue/4.2.16.Final/osx-aarch_64/netty-transport-native-kqueue-4.2.16.Final-osx-aarch_64.jar -netty-transport-native-kqueue/4.2.16.Final/osx-x86_64/netty-transport-native-kqueue-4.2.16.Final-osx-x86_64.jar -netty-transport-native-unix-common/4.2.16.Final//netty-transport-native-unix-common-4.2.16.Final.jar -netty-transport/4.2.16.Final//netty-transport-4.2.16.Final.jar +netty-transport-classes-epoll/4.2.17.Final//netty-transport-classes-epoll-4.2.17.Final.jar +netty-transport-classes-kqueue/4.2.17.Final//netty-transport-classes-kqueue-4.2.17.Final.jar +netty-transport-native-epoll/4.2.17.Final/linux-aarch_64/netty-transport-native-epoll-4.2.17.Final-linux-aarch_64.jar +netty-transport-native-epoll/4.2.17.Final/linux-riscv64/netty-transport-native-epoll-4.2.17.Final-linux-riscv64.jar +netty-transport-native-epoll/4.2.17.Final/linux-x86_64/netty-transport-native-epoll-4.2.17.Final-linux-x86_64.jar +netty-transport-native-kqueue/4.2.17.Final/osx-aarch_64/netty-transport-native-kqueue-4.2.17.Final-osx-aarch_64.jar +netty-transport-native-kqueue/4.2.17.Final/osx-x86_64/netty-transport-native-kqueue-4.2.17.Final-osx-x86_64.jar +netty-transport-native-unix-common/4.2.17.Final//netty-transport-native-unix-common-4.2.17.Final.jar +netty-transport/4.2.17.Final//netty-transport-4.2.17.Final.jar objenesis/3.5//objenesis-3.5.jar -okhttp/3.12.12//okhttp-3.12.12.jar -okio/1.17.6//okio-1.17.6.jar opencsv/2.3//opencsv-2.3.jar opentelemetry-api/1.49.0//opentelemetry-api-1.49.0.jar opentelemetry-context/1.49.0//opentelemetry-context-1.49.0.jar @@ -248,9 +265,14 @@ parquet-format-structures/1.17.1//parquet-format-structures-1.17.1.jar parquet-hadoop/1.17.1//parquet-hadoop-1.17.1.jar parquet-jackson/1.17.1//parquet-jackson-1.17.1.jar pickle/1.5//pickle-1.5.jar +profiles/2.35.4//profiles-2.35.4.jar +protocol-core/2.35.4//protocol-core-2.35.4.jar py4j/0.10.9.9//py4j-0.10.9.9.jar -reactive-streams/1.0.3//reactive-streams-1.0.3.jar +reactive-streams/1.0.4//reactive-streams-1.0.4.jar +regions/2.35.4//regions-2.35.4.jar remotetea-oncrpc/1.1.2//remotetea-oncrpc-1.1.2.jar +retries-spi/2.35.4//retries-spi-2.35.4.jar +retries/2.35.4//retries-2.35.4.jar rocksdbjni/10.10.1.1//rocksdbjni-10.10.1.1.jar scala-compiler/2.13.18//scala-compiler-2.13.18.jar scala-library/2.13.18//scala-library-2.13.18.jar @@ -258,6 +280,7 @@ scala-parallel-collections_2.13/1.2.0//scala-parallel-collections_2.13-1.2.0.jar scala-parser-combinators_2.13/2.4.0//scala-parser-combinators_2.13-2.4.0.jar scala-reflect/2.13.18//scala-reflect-2.13.18.jar scala-xml_2.13/2.4.0//scala-xml_2.13-2.4.0.jar +sdk-core/2.35.4//sdk-core-2.35.4.jar slf4j-api/2.0.17//slf4j-api-2.0.17.jar snakeyaml-engine/3.0.1//snakeyaml-engine-3.0.1.jar snakeyaml/2.5//snakeyaml-2.5.jar @@ -267,11 +290,15 @@ spire-platform_2.13/0.18.0//spire-platform_2.13-0.18.0.jar spire-util_2.13/0.18.0//spire-util_2.13-0.18.0.jar spire_2.13/0.18.0//spire_2.13-0.18.0.jar stream/2.9.8//stream-2.9.8.jar +sts/2.35.4//sts-2.35.4.jar super-csv/2.2.0//super-csv-2.2.0.jar +third-party-jackson-core/2.35.4//third-party-jackson-core-2.35.4.jar threeten-extra/1.9.0//threeten-extra-1.9.0.jar tink/1.23.0//tink-1.23.0.jar transaction-api/1.1//transaction-api-1.1.jar univocity-parsers/2.9.1//univocity-parsers-2.9.1.jar +utils-lite/2.35.4//utils-lite-2.35.4.jar +utils/2.35.4//utils-2.35.4.jar vertx-auth-common/4.5.28//vertx-auth-common-4.5.28.jar vertx-core/4.5.28//vertx-core-4.5.28.jar vertx-uri-template/4.5.28//vertx-uri-template-4.5.28.jar @@ -287,4 +314,4 @@ xz/1.12//xz-1.12.jar zjsonpatch/7.8.0//zjsonpatch-7.8.0.jar zookeeper-jute/3.9.5//zookeeper-jute-3.9.5.jar zookeeper/3.9.5//zookeeper-3.9.5.jar -zstd-jni/1.5.7-9//zstd-jni-1.5.7-9.jar +zstd-jni/1.5.7-13//zstd-jni-1.5.7-13.jar diff --git a/dev/free_disk_space b/dev/free_disk_space index d0916a32f301a..bcb70ce62b1b1 100755 --- a/dev/free_disk_space +++ b/dev/free_disk_space @@ -44,15 +44,26 @@ sudo rm -rf /opt/hostedtoolcache/go sudo rm -rf /opt/hostedtoolcache/node du -sh /opt/* -sudo apt-get update --fix-missing -sudo apt-get remove --purge -y '^aspnet.*' -sudo apt-get remove --purge -y '^dotnet-.*' -sudo apt-get remove --purge -y '^llvm-.*' -sudo apt-get remove --purge -y '^temurin-.*' -sudo apt-get remove --purge -y 'php.*' -sudo apt-get remove --purge -y '^mongodb-.*' -sudo apt-get remove --purge -y snapd google-chrome-stable microsoft-edge-stable firefox -sudo apt-get remove --purge -y azure-cli google-cloud-sdk mono-devel powershell libgl1-mesa-dri +package_patterns=( + '^aspnet.*' + '^dotnet-.*' + '^llvm-.*' + '^mongodb-.*' + '^temurin-.*' + 'php.*' + azure-cli + firefox + google-chrome-stable + google-cloud-sdk + libgl1-mesa-dri + microsoft-edge-stable + mono-devel + powershell + snapd +) +for package_pattern in "${package_patterns[@]}"; do + sudo apt-get remove --purge -y "$package_pattern" || true +done sudo apt-get autoremove --purge -y sudo apt-get clean diff --git a/dev/is-changed.py b/dev/is-changed.py index 92f61c8f07cd7..2e34a268f742d 100755 --- a/dev/is-changed.py +++ b/dev/is-changed.py @@ -17,17 +17,18 @@ # limitations under the License. # -import warnings -import traceback import os import sys +import traceback +import warnings from argparse import ArgumentParser + +import sparktestsupport.modules as modules from sparktestsupport.utils import ( determine_modules_for_files, determine_modules_to_test, identify_changed_files_from_git_commits, ) -import sparktestsupport.modules as modules def parse_opts(): diff --git a/dev/lint-java b/dev/lint-java index ff431301773f3..8095de68fb2a3 100755 --- a/dev/lint-java +++ b/dev/lint-java @@ -20,7 +20,7 @@ SCRIPT_DIR="$( cd "$( dirname "$0" )" && pwd )" SPARK_ROOT_DIR="$(dirname $SCRIPT_DIR)" -ERRORS=$($SCRIPT_DIR/../build/mvn -Pkinesis-asl -Pspark-ganglia-lgpl -Pkubernetes -Pyarn -Phive -Phive-thriftserver checkstyle:check | grep ERROR) +ERRORS=$($SCRIPT_DIR/../build/mvn -Pkinesis-asl -Pcredential-aws -Pspark-ganglia-lgpl -Pkubernetes -Pyarn -Phive -Phive-thriftserver checkstyle:check | grep ERROR) if test ! -z "$ERRORS"; then echo -e "Checkstyle checks failed at following occurrences:\n$ERRORS" diff --git a/dev/lint-scala b/dev/lint-scala index 2b242838d6b35..c718ec0850b06 100755 --- a/dev/lint-scala +++ b/dev/lint-scala @@ -39,7 +39,7 @@ ERRORS=$(./build/mvn \ ) if test ! -z "$ERRORS"; then - echo -e "The scalafmt check failed on sql/connect or sql/connect at following occurrences:\n\n$ERRORS\n" + echo -e "The scalafmt check failed on sql/api or sql/connect at following occurrences:\n\n$ERRORS\n" echo "Before submitting your change, please make sure to format your code using the following command:" echo "./build/mvn scalafmt:format -Dscalafmt.skip=false -Dscalafmt.validateOnly=false -Dscalafmt.changedOnly=false -pl sql/api -pl sql/connect/common -pl sql/connect/server -pl sql/connect/shims -pl sql/connect/client/jvm" exit 1 diff --git a/dev/merge_spark_pr.py b/dev/merge_spark_pr.py index 09598adc86110..bfa633c29cb89 100755 --- a/dev/merge_spark_pr.py +++ b/dev/merge_spark_pr.py @@ -37,6 +37,7 @@ # have added remotes corresponding to both (i) the github apache Spark # mirror and (ii) the apache git repo. +import argparse import json import os import re @@ -44,9 +45,13 @@ import sys import traceback from typing import List -from urllib.request import urlopen -from urllib.request import Request from urllib.error import HTTPError +from urllib.request import Request, urlopen + +# Shared with dev/pr_merge_status.py so the two committer tools agree on where a PR landed. +# Importable because Python puts this script's own directory first on sys.path. +from spark_merge_footer import branches_with_merge_footer as _branches_with_merge_footer +from spark_merge_footer import has_merge_footer try: import jira.client @@ -78,6 +83,16 @@ # exceeding your IP's unauthenticated request rate limit. You can create an OAuth key at # https://github.com/settings/tokens. This script only requires the "public_repo" scope. GITHUB_OAUTH_KEY = os.environ.get("GITHUB_OAUTH_KEY") +# Setting the DRY_RUN env var to any non-empty value makes the --dry-run/-n flag default to on +# (see build_arg_parser); the flag is the primary interface and can also turn it on explicitly. +# Consistent with the SKIP_VERSION_CHECK env toggle above. In a dry run every read-only step and +# local git op still runs -- fetching the PR, JIRA lookup, project_versions, JIRA/GitHub token +# validation, and the local squash-merge and cherry-picks on the throwaway PR_TOOL_* branches so +# conflicts and the computed merge hash stay realistic -- but every outbound effect is routed +# through a DryRun* client (see Git/GitHub/Jira below) that logs a "DRY-RUN: would ..." line +# instead of executing it: the git push to PUSH_REMOTE_NAME, the GitHub PR close/comment, and all +# JIRA writes (component and fixVersion updates, assignment, and the resolve transition). +DRY_RUN_ENV = bool(os.environ.get("DRY_RUN")) GITHUB_BASE = "https://github.com/apache/spark/pull" @@ -175,7 +190,7 @@ def compute_merge_default_fix_versions(merge_branches, unreleased_version_names) contain the commit, leveraging the Upstream-First backporting policy (cherry-picks flow master -> branch-M.x -> branch-M.N): - master contributes the greatest unreleased N.0.0; - - branch-M.x with master contributes that major's greatest unreleased minor.0; + - branch-M.x contributes that major's greatest unreleased minor.0; - branch-M.N contributes its greatest unreleased M.N.patch. Redundant entries are then suppressed: master's N.0.0 is dropped when any branch-M.x is in the merge set (a cherry-pick to branch-M.x has already landed on master); branch-M.x's @@ -212,7 +227,7 @@ def compute_merge_default_fix_versions(merge_branches, unreleased_version_names) ([], 1, True) >>> compute_merge_default_fix_versions(["branch-4.x"], ["4.3.0"]) - ([], []) + (['4.3.0'], []) >>> d, w = compute_merge_default_fix_versions(["branch-4.99"], ["4.3.0"]) >>> d == [] and len(w) == 1 and "branch-4.99" in w[0] @@ -272,8 +287,6 @@ def compute_merge_default_fix_versions(merge_branches, unreleased_version_names) continue line_major = _integration_major_from_branch(b) if line_major is not None: - if "master" not in merge_branches: - continue line_versions = [n for n in names if re.match(r"^%s\.\d+\.\d+$" % line_major, n)] chosen = _semver_max_version(line_versions) if chosen: @@ -281,7 +294,7 @@ def compute_merge_default_fix_versions(merge_branches, unreleased_version_names) else: warnings.append( "Could not infer an unreleased Spark %s (minor.maintenance) fix version " - "for branch-%s.x + master merge; enter version(s) manually when prompted." + "for branch-%s.x merge; enter version(s) manually when prompted." % (line_major, line_major) ) continue @@ -320,10 +333,74 @@ def keep(item): return list(dict.fromkeys(v for _, v in filtered)), warnings +def additional_fix_versions(inferred_versions, existing_versions): + """Return inferred Fix Versions not already present on a JIRA issue. + + Existing versions are preserved separately when the issue is updated, so this only + identifies the additions needed after a later backport. + + >>> additional_fix_versions(["4.4.0"], ["5.0.0"]) + ['4.4.0'] + >>> additional_fix_versions(["4.4.0"], ["5.0.0", "4.4.0"]) + [] + >>> additional_fix_versions(["4.4.0", "4.3.4"], ["5.0.0", "4.4.0"]) + ['4.3.4'] + """ + existing = set(existing_versions) + return [version for version in inferred_versions if version not in existing] + + +def fix_version_additions(inferred_versions, existing_versions): + """Return (additions, all_inferred_present). + + An empty ``additions`` list is ambiguous on its own: either the issue already carries + every inferred version, or nothing was inferred. Only the first means there is nothing + to do; the second still needs the committer prompted. + + >>> fix_version_additions(["4.3.0"], ["5.0.0"]) + (['4.3.0'], False) + >>> fix_version_additions(["5.0.0"], ["5.0.0"]) + ([], True) + >>> fix_version_additions([], ["5.0.0"]) + ([], False) + """ + additions = additional_fix_versions(inferred_versions, existing_versions) + return additions, bool(inferred_versions) and not additions + + +def fix_versions_from_input(raw_input, default_fix_versions): + """Resolve the Fix Version prompt's raw input into a list of version names. + + Blank falls back to the inferred default. With no default to fall back on, blank means + skip: the empty list, not [""], which no known version can match. + + >>> fix_versions_from_input("4.2.2", "") + ['4.2.2'] + >>> fix_versions_from_input("", "5.0.0") + ['5.0.0'] + >>> fix_versions_from_input("5.0.0, 4.3.0", "") + ['5.0.0', '4.3.0'] + >>> fix_versions_from_input("", "") + [] + >>> fix_versions_from_input(" ", "") + [] + """ + if raw_input == "": + raw_input = default_fix_versions + stripped = raw_input.replace(" ", "") + if stripped == "": + return [] + return stripped.split(",") + + def red(text): return "\033[91m%s\033[0m" % text +def bold(text): + return "\033[1m%s\033[0m" % text + + def print_error(msg): print(red(msg)) @@ -411,36 +488,112 @@ def get_json(url): sys.exit(-1) -def close_pr(pr_num): - url = "%s/pulls/%s" % (GITHUB_API_BASE, pr_num) - data = json.dumps({"state": "closed"}).encode("utf-8") - request = Request(url, data=data, method="PATCH") - request.add_header("Content-Type", "application/json") - request.add_header("Accept", "application/vnd.github+json") - if GITHUB_OAUTH_KEY: - request.add_header("Authorization", "token %s" % GITHUB_OAUTH_KEY) - try: - return json.load(urlopen(request)) - except HTTPError as e: - print_error("Failed to close PR #%s: HTTP %s %s" % (pr_num, e.code, e.reason)) - return None +def merge_commit_candidates(pr_events): + """Split `pr_events` into (closed_commits, referenced_commits), each oldest-first. + + Ordered by time so that a PR reopened and merged again yields its latest merge last. + + >>> merge_commit_candidates([{"event": "closed", "commit_id": "a", "created_at": "t2"}, + ... {"event": "referenced", "commit_id": "b", "created_at": "t1"}]) + (['a'], ['b']) + >>> merge_commit_candidates([{"event": "closed", "commit_id": None, "created_at": "t1"}]) + ([], []) + >>> merge_commit_candidates([{"event": "referenced", "commit_id": "c", "created_at": "t2"}, + ... {"event": "referenced", "commit_id": "b", "created_at": "t1"}]) + ([], ['b', 'c']) + """ + def commits_of(event_name): + matched = [e for e in pr_events if e["event"] == event_name and e["commit_id"] is not None] + return [e["commit_id"] for e in sorted(matched, key=lambda x: x["created_at"])] -def comment_pr(pr_num, body): - url = "%s/issues/%s/comments" % (GITHUB_API_BASE, pr_num) - data = json.dumps({"body": body}).encode("utf-8") - request = Request(url, data=data, method="POST") - request.add_header("Content-Type", "application/json") - request.add_header("Accept", "application/vnd.github+json") - if GITHUB_OAUTH_KEY: + return commits_of("closed"), commits_of("referenced") + + +def find_merge_commit(pr_num, pr_events): + """Return (hash, message) of the commit that merged `pr_num`, or (None, None). + + GitHub attributes the merge commit to the `closed` event only when that commit lands + on the default branch (master), because the "Closes #N" keyword in the commit message + is what closes the PR and the keyword is honored only there. A PR merged into any + other branch -- e.g. one opened against a rolling branch-M.x -- is instead closed by + this script through the API, and that `closed` event carries no commit, so the merge + survives only as a `referenced` event. Prefer the `closed` commit, which GitHub itself + linked; otherwise fall back to `referenced` events, which are also raised by any commit + merely mentioning the PR, so confirm each against the merge footer `merge_pr` generates. + """ + + def message_of(commit_hash): + return get_json("%s/commits/%s" % (GITHUB_API_BASE, commit_hash))["commit"]["message"] + + closed_commits, referenced_commits = merge_commit_candidates(pr_events) + if closed_commits: + return closed_commits[-1], message_of(closed_commits[-1]) + + for commit_hash in reversed(referenced_commits): + message = message_of(commit_hash) + if has_merge_footer(message, pr_num): + return commit_hash, message + return None, None + + +class GitHub: + """GitHub REST writes used by the merge script -- the single seam for PR mutations. + + Reads go through the module-level get_json(); only the two mutations (closing the PR and + posting a comment) live here, so DryRunGitHub can override them without touching any read + path. main() constructs GitHub() normally and DryRunGitHub() for a dry run. + """ + + def close_pr(self, pr_num): + url = "%s/pulls/%s" % (GITHUB_API_BASE, pr_num) + data = json.dumps({"state": "closed"}).encode("utf-8") + request = Request(url, data=data, method="PATCH") + request.add_header("Content-Type", "application/json") + request.add_header("Accept", "application/vnd.github+json") + if GITHUB_OAUTH_KEY: + request.add_header("Authorization", "token %s" % GITHUB_OAUTH_KEY) + try: + return json.load(urlopen(request)) + except HTTPError as e: + print_error("Failed to close PR #%s: HTTP %s %s" % (pr_num, e.code, e.reason)) + return None + + def comment(self, pr_num, body): + # Posting a comment is a write and needs auth; without a token, skip rather than 401. + if not GITHUB_OAUTH_KEY: + print_error("GITHUB_OAUTH_KEY is not set; skipping the comment.") + return None + url = "%s/issues/%s/comments" % (GITHUB_API_BASE, pr_num) + data = json.dumps({"body": body}).encode("utf-8") + request = Request(url, data=data, method="POST") + request.add_header("Content-Type", "application/json") + request.add_header("Accept", "application/vnd.github+json") request.add_header("Authorization", "token %s" % GITHUB_OAUTH_KEY) - try: - return json.load(urlopen(request)) - except HTTPError as e: - print_error("Failed to comment on PR #%s: HTTP %s %s" % (pr_num, e.code, e.reason)) + try: + return json.load(urlopen(request)) + except HTTPError as e: + print_error("Failed to comment on PR #%s: HTTP %s %s" % (pr_num, e.code, e.reason)) + return None + + +class DryRunGitHub(GitHub): + """Logs the intended PR close/comment instead of calling the GitHub API.""" + + def close_pr(self, pr_num): + print("DRY-RUN: would close PR #%s via the GitHub API." % pr_num) + return None + + def comment(self, pr_num, body): + print("DRY-RUN: would post the following comment on PR #%s:\n%s" % (pr_num, body)) return None +# Set in main(): GitHub() normally, DryRunGitHub() for a dry run. Module-level default so the +# top-level except handler and any early failure path have a usable client. +github = GitHub() + + def post_merge_comment(pr_num, merged_commits): """Post a comment on the PR recording every branch the change landed on and a link to the resulting commit, so the merge is traceable from the PR page. @@ -454,12 +607,14 @@ def post_merge_comment(pr_num, merged_commits): "- merged into %s %s/%s" % (ref, GITHUB_COMMIT_BASE, commit_hash) for ref, commit_hash in merged_commits ] - body = "**Merge Summary:**\n" + "\n".join(lines) + "\n\n*Posted by `merge_spark_pr.py`*" - print("Posting merge comment on PR #%s:\n%s" % (pr_num, body)) - if not GITHUB_OAUTH_KEY: - print_error("GITHUB_OAUTH_KEY is not set; skipping the merge comment.") - return - comment_pr(pr_num, body) + summary = "**Merge Summary:**\n" + "\n".join(lines) + attribution = "*Posted by `merge_spark_pr.py`*" + body = "%s\n\n%s" % (summary, attribution) + print( + "\n%s\n\n%s\n%s" + % (bold("Posting merge comment on PR #%s:" % pr_num), bold(summary), attribution) + ) + github.comment(pr_num, body) def fail(msg): @@ -468,47 +623,89 @@ def fail(msg): sys.exit(-1) -def run_cmd(cmd): - print(cmd) - if isinstance(cmd, list): - return subprocess.check_output(cmd).decode("utf-8") - else: - return subprocess.check_output(cmd.split(" ")).decode("utf-8") +class Git: + """Runs the merge-flow git commands for the merge script. + + This is the seam for the merge/backport flow, not every git invocation: + check_script_up_to_date() shells out to git merge-base directly. ``run`` executes read-only + and local commands (fetch, checkout, merge, commit, cherry-pick, rev-parse, config, branch + bookkeeping); those are safe even in a dry run because the merge and cherry-picks happen on + the throwaway PR_TOOL_* branches that clean_up removes. ``push`` is the only command that + mutates the shared repo at PUSH_REMOTE_NAME, so it is the one method DryRunGit overrides; + everything else it inherits and runs for real, keeping conflict detection and the computed + merge hash realistic. main() constructs Git() normally and DryRunGit() for a dry run. + """ + + def run(self, cmd): + print(cmd) + if isinstance(cmd, list): + return subprocess.check_output(cmd).decode("utf-8") + else: + return subprocess.check_output(cmd.split(" ")).decode("utf-8") + + def push(self, remote, local_ref, remote_ref): + return self.run("git push %s %s:%s" % (remote, local_ref, remote_ref)) + + +class DryRunGit(Git): + """Logs the intended push instead of updating the remote; all other git runs for real.""" + + def push(self, remote, local_ref, remote_ref): + print("DRY-RUN: would push %s to %s:%s" % (local_ref, remote, remote_ref)) + return "" + + +# Set in main(): Git() normally, DryRunGit() for a dry run. Module-level default so clean_up and +# the top-level except handler have a usable client even if main() fails early. +git = Git() + + +class SkipCherryPick(Exception): + """Signals that the committer declined to resolve a conflicting cherry-pick. + + A backport conflict is routine, and by the time one is hit the merge into the target + branch (and any earlier cherry-picks) has already been pushed. Declining it must skip + only that one branch -- letting the caller offer another branch and, crucially, still + resolve the JIRA -- instead of aborting the whole merge the way a hard `fail()` would. + """ def continue_maybe(prompt, cherry=False): if get_input(f"{prompt} (y/N): ", ["y", "n", ""]) != "y": if cherry: try: - run_cmd("git cherry-pick --abort") + git.run("git cherry-pick --abort") except Exception: print_error("Unable to abort and get back to the state before cherry-pick") + print("Skipping this cherry-pick; the merge continues.") + clean_up() + raise SkipCherryPick() fail("Okay, exiting") def clean_up(): if "original_head" in globals(): print("Restoring head pointer to %s" % original_head) - run_cmd("git checkout %s" % original_head) + git.run("git checkout %s" % original_head) - branches = run_cmd("git branch").replace(" ", "").split("\n") + branches = git.run("git branch").replace(" ", "").split("\n") for branch in list(filter(lambda x: x.startswith(BRANCH_PREFIX), branches)): print("Deleting local branch %s" % branch) - run_cmd("git branch -D %s" % branch) + git.run("git branch -D %s" % branch) # merge the requested PR and return the merge hash def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, pr_author, co_authors): pr_branch_name = "%s_MERGE_PR_%s" % (BRANCH_PREFIX, pr_num) target_branch_name = "%s_MERGE_PR_%s_%s" % (BRANCH_PREFIX, pr_num, target_ref.upper()) - run_cmd("git fetch %s pull/%s/head:%s" % (PR_REMOTE_NAME, pr_num, pr_branch_name)) - run_cmd("git fetch %s %s:%s" % (PUSH_REMOTE_NAME, target_ref, target_branch_name)) - run_cmd("git checkout %s" % target_branch_name) + git.run("git fetch %s pull/%s/head:%s" % (PR_REMOTE_NAME, pr_num, pr_branch_name)) + git.run("git fetch %s %s:%s" % (PUSH_REMOTE_NAME, target_ref, target_branch_name)) + git.run("git checkout %s" % target_branch_name) had_conflicts = False try: - run_cmd(["git", "merge", pr_branch_name, "--squash"]) + git.run(["git", "merge", pr_branch_name, "--squash"]) except Exception as e: msg = "Error merging: %s\nWould you like to manually fix-up this merge?" % e continue_maybe(msg) @@ -531,8 +728,8 @@ def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, pr_author, co_author # to people every time someone creates a public fork of Spark. merge_message_flags += ["-m", body.replace("@", "")] - committer_name = run_cmd("git config --get user.name").strip() - committer_email = run_cmd("git config --get user.email").strip() + committer_name = git.run("git config --get user.name").strip() + committer_email = git.run("git config --get user.email").strip() if had_conflicts: message = "This patch had conflicts when merged, resolved by\nCommitter: %s <%s>" % ( @@ -552,19 +749,19 @@ def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, pr_author, co_author merge_message_flags += ["-m", authors] - run_cmd(["git", "commit", '--author="%s"' % primary_author] + merge_message_flags) + git.run(["git", "commit", '--author="%s"' % primary_author] + merge_message_flags) continue_maybe( "Merge complete (local ref %s). Push to %s?" % (target_branch_name, PUSH_REMOTE_NAME) ) try: - run_cmd("git push %s %s:%s" % (PUSH_REMOTE_NAME, target_branch_name, target_ref)) + git.push(PUSH_REMOTE_NAME, target_branch_name, target_ref) except Exception as e: clean_up() print_error("Exception while pushing: %s" % e) - merge_hash = run_cmd("git rev-parse %s" % target_branch_name).strip() + merge_hash = git.run("git rev-parse %s" % target_branch_name).strip() clean_up() print("Pull request #%s merged!" % pr_num) print("Merge hash: %s" % merge_hash) @@ -574,31 +771,54 @@ def merge_pr(pr_num, target_ref, title, body, pr_repo_desc, pr_author, co_author def _do_cherry_pick(pr_num, merge_hash, pick_ref): """Cherry-pick `merge_hash` onto `pick_ref` and push. - Returns the (pushed ref, pushed commit hash) pair. + Returns the (pushed ref, pushed commit hash) pair. Raises `SkipCherryPick` if the + cherry-pick conflicts and the committer declines to resolve it, after attempting to + abort the cherry-pick and restore the working tree. """ pick_branch_name = "%s_PICK_PR_%s_%s" % (BRANCH_PREFIX, pr_num, pick_ref.upper()) - run_cmd("git fetch %s %s:%s" % (PUSH_REMOTE_NAME, pick_ref, pick_branch_name)) - run_cmd("git checkout %s" % pick_branch_name) + git.run("git fetch %s %s:%s" % (PUSH_REMOTE_NAME, pick_ref, pick_branch_name)) + git.run("git checkout %s" % pick_branch_name) try: - run_cmd("git cherry-pick -sx %s" % merge_hash) + git.run( + [ + "git", + "-c", + "commit.cleanup=scissors", + "cherry-pick", + "-sx", + merge_hash, + ] + ) except Exception as e: msg = "Error cherry-picking: %s\nWould you like to manually fix-up this merge?" % e continue_maybe(msg, True) - msg = "Okay, please fix any conflicts and finish the cherry-pick. Finished?" + msg = "Okay, please fix any conflicts and 'git add' conflicting files... Finished?" continue_maybe(msg, True) + # Important to use `scissors` and `--edit` otherwise git will strip lines starting with `#` + # when calling `--continue`. See: https://github.com/apache/spark/pull/58214 + git.run( + [ + "git", + "-c", + "commit.cleanup=scissors", + "cherry-pick", + "--continue", + "--edit", + ] + ) continue_maybe( "Pick complete (local ref %s). Push to %s?" % (pick_branch_name, PUSH_REMOTE_NAME) ) try: - run_cmd("git push %s %s:%s" % (PUSH_REMOTE_NAME, pick_branch_name, pick_ref)) + git.push(PUSH_REMOTE_NAME, pick_branch_name, pick_ref) except Exception as e: fail("Exception while pushing: %s" % e) - pick_hash = run_cmd("git rev-parse %s" % pick_branch_name).strip() + pick_hash = git.run("git rev-parse %s" % pick_branch_name).strip() clean_up() print("Pull request #%s picked into %s!" % (pr_num, pick_ref)) @@ -606,6 +826,45 @@ def _do_cherry_pick(pr_num, merge_hash, pick_ref): return pick_ref, pick_hash +def branches_with_merge_footer(pr_num, branch_names): + """Release branches from `branch_names` that already carry `pr_num`'s merge footer. + + Thin wrapper over the shared reader in `spark_merge_footer`, adding this script's own + policy: a git failure here must not abort a merge that may already have pushed, so it + warns and reports nothing rather than exiting. Per that module's refresh policy no fetch + is issued, so a backport not yet fetched into PUSH_REMOTE_NAME's tracking refs is simply + not reported -- the committer is still prompted and can type any branch. + """ + try: + landed = _branches_with_merge_footer( + pr_num, PUSH_REMOTE_NAME, lambda args: git.run(["git"] + args) + ) + except Exception as e: + print_error("Could not scan for existing backports of #%s (%s)." % (pr_num, e)) + return [] + # Keep branch_names' newest-first order, and drop anything not a known release branch. + return [b for b in branch_names if b in landed] + + +def default_pick_branch(branch_names, already_picked): + """Highest-ranked release branch that has not already received the change, or None. + + `branch_names` is ordered newest-first (see `semver_branch_rank`) and `already_picked` + holds the branches the change is known to be on, so the prompt never defaults to a + branch where the cherry-pick would come up empty. Returns None when every known branch + already has it, so callers can say so instead of offering an empty pick. + + >>> default_pick_branch(["branch-4.x", "branch-4.3", "branch-4.2"], ("branch-4.x",)) + 'branch-4.3' + >>> default_pick_branch(["branch-4.x", "branch-4.3"], ()) + 'branch-4.x' + >>> default_pick_branch(["branch-4.x"], ("branch-4.x",)) is None + True + """ + remaining = [b for b in branch_names if b not in already_picked] + return remaining[0] if remaining else None + + def _upstream_first_sibling(target_ref, pick_ref, branch_names, already_picked): """Return the sibling branch-M.x if Upstream-First should prompt, else None. @@ -645,7 +904,9 @@ def cherry_pick(pr_num, merge_hash, default_branch, branch_names, target_ref, al BOTH (the policy-compliant default) or branch-M.N only (treated as a maintenance-only bugfix). Returns the list of (ref, commit_hash) pairs actually picked into, so the main loop can advance its remaining-branches list correctly - and record each backport commit for the merge comment. + and record each backport commit for the merge comment. The list is empty (or, for the + Upstream-First two-branch path, holds only what landed) when the committer declines to + resolve a conflict, so the caller simply offers the next branch and still resolves JIRA. """ while True: pick_ref = bold_input(f"Enter a branch name [{default_branch}]: ") @@ -677,17 +938,30 @@ def cherry_pick(pr_num, merge_hash, default_branch, branch_names, target_ref, al {"b": ["b", "both", ""], "o": ["o", "only"], "a": ["a", "abort"]}, ) if choice == "b": - picked_x = _do_cherry_pick(pr_num, merge_hash, sibling_x) - picked_n = _do_cherry_pick(pr_num, merge_hash, pick_ref) - return [picked_x, picked_n] + # Preserve any pick that was already pushed: if the branch-M.N pick is skipped after + # branch-M.x landed, still return branch-M.x so it's recorded and JIRA/comment + # reflect it. + picked = [] + try: + picked.append(_do_cherry_pick(pr_num, merge_hash, sibling_x)) + picked.append(_do_cherry_pick(pr_num, merge_hash, pick_ref)) + except SkipCherryPick: + pass + return picked elif choice == "o": - return [_do_cherry_pick(pr_num, merge_hash, pick_ref)] + try: + return [_do_cherry_pick(pr_num, merge_hash, pick_ref)] + except SkipCherryPick: + return [] elif choice == "a": fail("Aborted by user at Upstream-First policy prompt.") else: fail("Unrecognized choice %r; aborting." % choice) - return [_do_cherry_pick(pr_num, merge_hash, pick_ref)] + try: + return [_do_cherry_pick(pr_num, merge_hash, pick_ref)] + except SkipCherryPick: + return [] # Common words carry no signal when comparing a PR title to a JIRA summary, so they are @@ -771,7 +1045,15 @@ def tokens(text): def format_jira_verification( - pr_num, pr_title, jira_id, summary, status, issuetype, use_color=False, is_followup=False + pr_num, + pr_title, + jira_id, + summary, + status, + issuetype, + resolution=None, + use_color=False, + is_followup=False, ): """Render the JIRA-vs-PR match block shown before merging. @@ -838,6 +1120,7 @@ def format_jira_verification( Match: 0.00 (FOLLOWUP: title intentionally differs, not checked) """ status_warning = "" + resolution_suffix = " (%s)" % resolution if resolution else "" if status in ("Resolved", "Closed"): status_warning = " <-- WARNING: already Resolved/Closed" if use_color: @@ -858,7 +1141,7 @@ def format_jira_verification( "=== Verify JIRA matches PR #%s ===" % pr_num, "PR title: %s" % pr_title, "JIRA %s: %s" % (jira_id, summary), - " Status: %s%s" % (status, status_warning), + " Status: %s%s%s" % (status, resolution_suffix, status_warning), " Type: %s" % issuetype, " Match: %.2f%s" % (score, match_suffix), ] @@ -872,6 +1155,9 @@ def print_jira_issue_summary(issue): assignee = assignee.displayName assignee = "Assignee\t%s\n" % assignee status = "Status\t\t%s\n" % issue.fields.status.name + resolution = "" + if issue.fields.resolution is not None: + resolution = "Resolution\t%s\n" % issue.fields.resolution.name components = "Components\t%s\n" % [x.name for x in issue.fields.components] url = "Url\t\t%s/%s\n" % (JIRA_BASE, issue.key) target_versions = "Affected\t%s\n" % [x.name for x in issue.fields.versions] @@ -880,8 +1166,8 @@ def print_jira_issue_summary(issue): fix_versions = "Fixed\t\t%s\n" % [x.name for x in issue.fields.fixVersions] print("=== JIRA %s ===" % issue.key) print( - "%s%s%s%s%s%s%s" - % (summary, assignee, status, components, url, target_versions, fix_versions) + "%s%s%s%s%s%s%s%s" + % (summary, assignee, status, resolution, components, url, target_versions, fix_versions) ) @@ -901,6 +1187,8 @@ def jira_components_from_title_tags(tags): ['PySpark', 'Documentation'] >>> jira_components_from_title_tags(["SQL", "TEST"]) ['SQL', 'Tests'] + >>> jira_components_from_title_tags(["UDF"]) + ['UDF'] >>> jira_components_from_title_tags(["SQL", "FOLLOWUP", "4.X", "BOGUS"]) ['SQL'] >>> jira_components_from_title_tags(["SQL", "SQL"]) @@ -952,11 +1240,7 @@ def reconcile_jira_components(issue, title_components): # Append the PR title's components, keeping the existing ones first. new_names = list(dict.fromkeys(current + title_jira_components)) - try: - issue.update(fields={"components": [{"name": n} for n in new_names]}) - print("Updated JIRA %s components to: %s" % (issue.key, ", ".join(new_names))) - except Exception as e: - print_error("Failed to update components on JIRA %s: %s" % (issue.key, e)) + jira_ops.update_components(issue, new_names) def get_jira_issue(prompt, default_jira_id=""): @@ -971,26 +1255,45 @@ def get_jira_issue(prompt, default_jira_id=""): print_jira_issue_summary(issue) status = issue.fields.status.name if status == "Resolved" or status == "Closed": - print("JIRA issue %s already has status '%s'" % (jira_id, status)) - return None + resolution = issue.fields.resolution + resolution_name = resolution.name if resolution is not None else None + print("JIRA issue %s already has status '%s' (%s)" % (jira_id, status, resolution_name)) + # Only a ticket an earlier merge resolved as Fixed can legitimately gain + # another Fix Version. Duplicate / Won't Fix / Invalid tickets must not be + # touched. + if resolution_name != "Fixed": + return None if get_input("Check if the JIRA information is as expected (y/N): ", ["y", "n", ""]) == "y": return issue else: - return get_jira_issue("Enter the revised JIRA ID again or leave blank to skip") + return get_jira_issue( + "Enter the revised JIRA ID again or leave blank to skip", + ) except Exception as e: print_error("ASF JIRA could not find %s: %s" % (jira_id, e)) - return get_jira_issue("Enter the revised JIRA ID again or leave blank to skip") + return get_jira_issue( + "Enter the revised JIRA ID again or leave blank to skip", + ) -def resolve_jira_issue(merge_branches, comment, default_jira_id="", title_components=()): +def resolve_jira_issue( + merge_branches, + comment, + default_jira_id="", + title_components=(), +): issue = get_jira_issue("Enter a JIRA id", default_jira_id) if issue is None: return - if issue.fields.assignee is None: - choose_jira_assignee(issue) + status = issue.fields.status.name + is_resolved = status == "Resolved" or status == "Closed" + + if not is_resolved: + if issue.fields.assignee is None: + choose_jira_assignee(issue) - reconcile_jira_components(issue, title_components) + reconcile_jira_components(issue, title_components) versions = asf_jira.project_versions("SPARK") # Consider only x.y.z, unreleased, unarchived versions @@ -1007,17 +1310,48 @@ def resolve_jira_issue(merge_branches, comment, default_jira_id="", title_compon ) for w in infer_warnings: print_error(w) + + existing_fix_versions = list(issue.fields.fixVersions) if is_resolved else [] + existing_fix_version_names = [v.name for v in existing_fix_versions] + if is_resolved: + # A later backport run must preserve the versions recorded by the original merge and + # only add versions inferred from the newly discovered branches. + default_fix_list, all_inferred_present = fix_version_additions( + default_fix_list, existing_fix_version_names + ) + if all_inferred_present: + print( + "JIRA issue %s already contains all inferred fix versions; no update needed." + % issue.key + ) + return + if default_fix_list: + print( + "JIRA issue %s has fix version(s) %s; inferred addition(s): %s" + % (issue.key, existing_fix_version_names, default_fix_list) + ) + if get_input("Add these fix version(s)? (y/N): ", ["y", "n", ""]) != "y": + return + else: + # Nothing inferred, so there is nothing to confirm; fall through to the prompt. + print( + "JIRA issue %s has fix version(s) %s; no additional fix version could be " + "inferred." % (issue.key, existing_fix_version_names) + ) default_fix_versions = ",".join(default_fix_list) available_versions = set(list(map(lambda v: v.name, versions))) while True: try: - fix_versions = bold_input( - "Enter comma-separated fix version(s) [%s]: " % default_fix_versions + prompt = "Enter comma-separated fix version(s) [%s]: " + if is_resolved: + prompt = "Enter comma-separated additional fix version(s) [%s]: " + fix_versions = fix_versions_from_input( + bold_input(prompt % default_fix_versions), default_fix_versions ) - if fix_versions == "": - fix_versions = default_fix_versions - fix_versions = fix_versions.replace(" ", "").split(",") + if not fix_versions: + print("No fix version entered; update %s manually." % issue.key) + return if set(fix_versions).issubset(available_versions): break else: @@ -1036,24 +1370,23 @@ def get_version_json(version_str): jira_fix_versions = list(map(lambda v: get_version_json(v), fix_versions)) + if is_resolved: + existing_names = set(existing_fix_version_names) + jira_fix_versions = [v for v in jira_fix_versions if v["name"] not in existing_names] + if not jira_fix_versions: + print("No new fix versions selected for JIRA issue %s; no update needed." % issue.key) + return + jira_ops.add_fix_versions(issue, existing_fix_versions, jira_fix_versions) + return + resolve = list(filter(lambda a: a["name"] == "Resolve Issue", asf_jira.transitions(issue.key)))[ 0 ] resolution = list(filter(lambda r: r.raw["name"] == "Fixed", asf_jira.resolutions()))[0] - asf_jira.transition_issue( - issue.key, - resolve["id"], - fixVersions=jira_fix_versions, - comment=comment, - resolution={"id": resolution.raw["id"]}, + jira_ops.resolve_issue( + issue, resolve["id"], jira_fix_versions, comment, resolution.raw["id"], fix_versions ) - try: - print_jira_issue_summary(asf_jira.issue(issue.key)) - except Exception: - print("Unable to fetch JIRA issue %s after resolving" % issue.key) - print("Successfully resolved %s with fixVersions=%s!" % (issue.key, fix_versions)) - def choose_jira_assignee(issue): """ @@ -1088,7 +1421,7 @@ def choose_jira_assignee(issue): # assume it's a user id, and try to assign (might fail, we just prompt again) assignee = asf_jira.user(raw_assignee) try: - assign_issue(issue.key, assignee.name) + jira_ops.assign(issue.key, assignee.name) except Exception as e: if ( e.__class__.__name__ == "JIRAError" @@ -1099,8 +1432,8 @@ def choose_jira_assignee(issue): "User '%s' cannot be assigned, add to contributors role and try again?" % assignee.name ) - grant_contributor_role(assignee.name) - assign_issue(issue.key, assignee.name) + jira_ops.grant_contributor(assignee.name) + jira_ops.assign(issue.key, assignee.name) else: raise e return assignee @@ -1111,32 +1444,153 @@ def choose_jira_assignee(issue): print("Error assigning JIRA, try again (or leave blank and fix manually)") -def grant_contributor_role(user: str): - role = asf_jira.project_role("SPARK", 10010) - role.add_user(user) - print("Successfully added user '%s' to contributors role" % user) - +class Jira: + """ASF JIRA writes used by the merge script -- the single seam for JIRA mutations. -def assign_issue(issue: int, assignee: str) -> bool: - """ - Assign an issue to a user, which is a shorthand for jira.client.JIRA.assign_issue. - The original one has an issue that it will search users again and only choose the assignee - from 20 candidates. If it's unmatched, it picks the head blindly. In our case, the assignee - is already resolved. + The pre-write JIRA lookups stay on the module-level asf_jira client; the writes (components, + fix versions, the resolve transition, assignment, and the contributor-role grant) go through + here, so DryRunJira can log them. Two Production writes (add_fix_versions, resolve_issue) also + read the issue back afterward to print the updated summary; DryRunJira skips the write and that + read-back alike. main() builds Jira(asf_jira) normally and DryRunJira(asf_jira) for a dry run, + after initialize_jira() sets asf_jira. """ - url = getattr(asf_jira, "_get_latest_url")(f"issue/{issue}/assignee") - payload = {"name": assignee} - getattr(asf_jira, "_session").put(url, data=json.dumps(payload)) - return True + + def __init__(self, client): + self._client = client + + def update_components(self, issue, new_names): + try: + issue.update(fields={"components": [{"name": n} for n in new_names]}) + print("Updated JIRA %s components to: %s" % (issue.key, ", ".join(new_names))) + except Exception as e: + print_error("Failed to update components on JIRA %s: %s" % (issue.key, e)) + + def add_fix_versions(self, issue, existing_fix_versions, new_version_jsons): + issue.update( + fields={"fixVersions": [v.raw for v in existing_fix_versions] + new_version_jsons} + ) + try: + print_jira_issue_summary(self._client.issue(issue.key)) + except Exception: + print("Unable to fetch JIRA issue %s after updating fix versions" % issue.key) + print( + "Successfully updated %s with additional fixVersions=%s!" + % (issue.key, [v["name"] for v in new_version_jsons]) + ) + + def resolve_issue( + self, issue, resolve_id, fix_version_jsons, comment, resolution_id, fix_version_names + ): + self._client.transition_issue( + issue.key, + resolve_id, + fixVersions=fix_version_jsons, + comment=comment, + resolution={"id": resolution_id}, + ) + try: + print_jira_issue_summary(self._client.issue(issue.key)) + except Exception: + print("Unable to fetch JIRA issue %s after resolving" % issue.key) + print("Successfully resolved %s with fixVersions=%s!" % (issue.key, fix_version_names)) + + def assign(self, issue_key, assignee): + # Shorthand for jira.client.JIRA.assign_issue. The library's own assign_issue re-searches + # users and blindly picks the first of 20 candidates when unmatched; here the assignee is + # already resolved, so PUT it directly. + url = getattr(self._client, "_get_latest_url")(f"issue/{issue_key}/assignee") + getattr(self._client, "_session").put(url, data=json.dumps({"name": assignee})) + return True + + def grant_contributor(self, user): + role = self._client.project_role("SPARK", 10010) + role.add_user(user) + print("Successfully added user '%s' to contributors role" % user) + + +class DryRunJira(Jira): + """Logs the intended JIRA writes instead of calling the API; reads still go through.""" + + def update_components(self, issue, new_names): + print("DRY-RUN: would set JIRA %s components to: %s" % (issue.key, ", ".join(new_names))) + + def add_fix_versions(self, issue, existing_fix_versions, new_version_jsons): + print( + "DRY-RUN: would add fixVersions=%s to JIRA %s." + % ([v["name"] for v in new_version_jsons], issue.key) + ) + + def resolve_issue( + self, issue, resolve_id, fix_version_jsons, comment, resolution_id, fix_version_names + ): + print( + "DRY-RUN: would resolve JIRA %s as Fixed with fixVersions=%s and add comment:\n%s" + % (issue.key, fix_version_names, comment) + ) + + def assign(self, issue_key, assignee): + print("DRY-RUN: would assign JIRA %s to '%s'." % (issue_key, assignee)) + return True + + def grant_contributor(self, user): + print("DRY-RUN: would add user '%s' to the SPARK contributors role." % user) + + +# Set in main() after initialize_jira(): Jira(asf_jira) normally, DryRunJira(asf_jira) for a dry +# run. None until then; only the JIRA-write flow, well after construction, uses it. +jira_ops = None def resolve_jira_issues(title, merge_branches, comment, title_components=()): jira_ids = re.findall("SPARK-[0-9]{4,5}", title) if len(jira_ids) == 0: - resolve_jira_issue(merge_branches, comment, title_components=title_components) + resolve_jira_issue( + merge_branches, + comment, + title_components=title_components, + ) for jira_id in jira_ids: - resolve_jira_issue(merge_branches, comment, jira_id, title_components=title_components) + resolve_jira_issue( + merge_branches, + comment, + jira_id, + title_components=title_components, + ) + + +def update_jira_for_pr(pr_num, title, merge_branches, title_components): + skip_jira_title_tags = ("MINOR", "TRIVIAL", "FOLLOWUP") + tags = set(title_components) + try: + parsed = Title.parse(title) + tags.update(parsed.leading) + tags.update(parsed.components) + except ValueError: + pass + skipped = [tag for tag in skip_jira_title_tags if tag in tags] + if skipped: + print() + print_error( + "Skipping JIRA operations for PR #%s because title has %s." + % (pr_num, ", ".join("[%s]" % tag for tag in skipped)) + ) + return + + # asf_jira is guaranteed to be set here: initialize_jira() fails fast otherwise. + print() + continue_maybe("Would you like to update an associated JIRA?") + jira_comment = "Issue resolved by pull request %s\n[%s/%s]" % ( + pr_num, + GITHUB_BASE, + pr_num, + ) + resolve_jira_issues( + title, + merge_branches, + jira_comment, + title_components, + ) class Component: @@ -1197,7 +1651,7 @@ def find(cls, token): Component("DOC", ("DOCS", "DOCUMENTATION"), primary=True, jira_name="Documentation"), Component("DOCKER", primary=True, jira_name="Spark Docker"), Component("EC2", jira_name="EC2"), - Component("EXAMPLE", ("EXAMPLES",), jira_name="Examples"), + Component("EXAMPLES", ("EXAMPLE",), primary=True, jira_name="Examples"), Component("GRAPHX", primary=True, jira_name="GraphX"), Component("INFRA", ("PROJECT_INFRA",), primary=True, jira_name="Project Infra"), Component("IO", jira_name="Input/Output"), @@ -1221,6 +1675,7 @@ def find(cls, token): Component("STREAMING", ("DSTREAM", "DSTREAMS"), primary=True, jira_name="DStreams"), Component("SUBMIT", jira_name="Spark Submit"), Component("TEST", ("TESTS", "TEST-ONLY", "TESTS-ONLY"), jira_name="Tests"), + Component("UDF", primary=True, jira_name="UDF"), Component("UI", ("WEBUI", "WEB_UI"), primary=True, jira_name="Web UI"), Component("WINDOWS", primary=True, jira_name="Windows"), Component("YARN", primary=True, jira_name="YARN"), @@ -1409,10 +1864,10 @@ def prompt_for_components(): def get_current_ref(): - ref = run_cmd("git rev-parse --abbrev-ref HEAD").strip() + ref = git.run("git rev-parse --abbrev-ref HEAD").strip() if ref == "HEAD": # The current ref is a detached HEAD, so grab its SHA. - return run_cmd("git rev-parse HEAD").strip() + return git.run("git rev-parse HEAD").strip() else: return ref @@ -1516,10 +1971,49 @@ def check_script_up_to_date(): ) +def build_arg_parser(): + """CLI parser: an optional PR number plus the --dry-run/-n flag. + + --dry-run defaults to on when the DRY_RUN env var is set to any non-empty value (see the + DRY_RUN_ENV note near the top of this file), so either the flag or the env var enables a + dry run; passing the flag forces it on regardless. + """ + parser = argparse.ArgumentParser(description="Merge a Spark pull request.") + parser.add_argument( + "pr_num", + nargs="?", + help="Pull request number to merge; prompted for interactively if omitted.", + ) + parser.add_argument( + "--dry-run", + "-n", + action="store_true", + default=DRY_RUN_ENV, + help="Run every read-only step for real but suppress all outbound effects (git push, " + "GitHub PR close/comment, JIRA writes), logging each as a 'DRY-RUN: would ...' line. " + "Also enabled by setting the DRY_RUN environment variable.", + ) + return parser + + def main(): + global git, github, jira_ops, original_head + args = build_arg_parser().parse_args() + if args.dry_run: + # Swap the outbound clients for their logging-only variants; everything else is unchanged. + git = DryRunGit() + github = DryRunGitHub() + print( + bold( + "=== DRY-RUN: read-only steps run for real, but the git push to %s, the GitHub " + "PR close/comment, and all JIRA writes are suppressed and only logged. ===" + % PUSH_REMOTE_NAME + ) + ) + check_script_up_to_date() initialize_jira() - global original_head + jira_ops = DryRunJira(asf_jira) if args.dry_run else Jira(asf_jira) os.chdir(SPARK_HOME) original_head = get_current_ref() @@ -1528,10 +2022,10 @@ def main(): branch_names = list(filter(lambda x: x.startswith("branch-"), [x["name"] for x in branches])) branch_names = sorted(branch_names, key=semver_branch_rank, reverse=True) - if len(sys.argv) == 1: + if args.pr_num is None: pr_num = get_input("Which pull request would you like to merge? (e.g. 34): ", r"^\d+$") else: - pr_num = sys.argv[1] + pr_num = args.pr_num print("Start to merge pull request #%s" % (pr_num)) pr = get_json("%s/pulls/%s" % (GITHUB_API_BASE, pr_num)) pr_events = get_json("%s/issues/%s/events" % (GITHUB_API_BASE, pr_num)) @@ -1654,31 +2148,62 @@ def main(): # Merged pull requests don't appear as merged in the GitHub API; # Instead, they're closed by committers. - merge_commits = [e for e in pr_events if e["event"] == "closed" and e["commit_id"] is not None] - - if merge_commits and pr["state"] == "closed": - # A PR might have multiple merge commits, if it's reopened and merged again. We shall - # cherry-pick PRs in closed state with the latest merge hash. - # If the PR is still open(reopened), we shall not cherry-pick it but perform the normal - # merge as it could have been reverted earlier. - merge_commits = sorted(merge_commits, key=lambda x: x["created_at"]) - merge_hash = merge_commits[-1]["commit_id"] - message = get_json("%s/commits/%s" % (GITHUB_API_BASE, merge_hash))["commit"]["message"] - + # A PR might have multiple merge commits, if it's reopened and merged again. We shall + # cherry-pick PRs in closed state with the latest merge hash. + # If the PR is still open (reopened), we shall not cherry-pick it but perform the normal + # merge as it could have been reverted earlier. + merge_hash, message = (None, None) + if pr["state"] == "closed": + merge_hash, message = find_merge_commit(pr_num, pr_events) + + if merge_hash is not None: print("Pull request %s has already been merged, assuming you want to backport" % pr_num) commit_is_downloaded = ( - run_cmd(["git", "rev-parse", "--quiet", "--verify", "%s^{commit}" % merge_hash]).strip() + git.run(["git", "rev-parse", "--quiet", "--verify", "%s^{commit}" % merge_hash]).strip() != "" ) if not commit_is_downloaded: fail("Couldn't find any merge commit for #%s, you may need to update HEAD." % pr_num) print("Found commit %s:\n%s" % (merge_hash, message)) - default = branch_names[0] - picked = cherry_pick( - pr_num, merge_hash, default, branch_names, target_ref, already_picked=() - ) - post_merge_comment(pr_num, picked) + # The change is already on target_ref and on any branch a previous run backported it + # to, so exclude all of them: defaulting to one would cherry-pick an empty commit. + picked_refs = [target_ref] + [ + b for b in branches_with_merge_footer(pr_num, branch_names) if b != target_ref + ] + if len(picked_refs) > 1: + print("Already backported to: %s" % ", ".join(picked_refs[1:])) + # Loop so one invocation can reach several maintenance branches, as the merge path does. + picked_commits = [] + try: + while True: + default = default_pick_branch(branch_names, tuple(picked_refs)) + if default is None: + print( + "Every known release branch already contains #%s; nothing to pick." % pr_num + ) + break + picked = cherry_pick( + pr_num, + merge_hash, + default, + branch_names, + target_ref, + already_picked=tuple(picked_refs), + ) + picked_refs = picked_refs + [ref for ref, _ in picked] + picked_commits = picked_commits + picked + prompt = "Would you like to pick %s into another branch?" % merge_hash + if get_input(f"\n{prompt} (y/N): ", ["y", "n", ""]) != "y": + break + finally: + # Report whatever was pushed even if a later pick is aborted, since the earlier + # pushes have already landed. + if picked_commits: + post_merge_comment(pr_num, picked_commits) + # Backport mode may be the first chance to resolve a JIRA after an interrupted + # original merge. If it was already resolved, add any newly inferred fix versions. + update_jira_for_pr(pr_num, title, picked_refs, title_components) sys.exit(0) if not bool(pr["mergeable"]): @@ -1734,6 +2259,9 @@ def main(): issue.fields.summary, issue.fields.status.name, issue.fields.issuetype.name, + resolution=( + issue.fields.resolution.name if issue.fields.resolution is not None else None + ), use_color=True, is_followup=is_followup, ) @@ -1756,27 +2284,20 @@ def main(): # then each cherry-pick target as it is picked. merged_commits = [(target_ref, merge_hash)] - # The "Closes #N" keyword in the commit message only auto-closes the PR when the commit - # lands on the default branch. For merges into other branches (e.g. branch-X.Y backport - # PRs), GitHub leaves the PR open, so close it explicitly through the API. - pr_state = get_json("%s/pulls/%s" % (GITHUB_API_BASE, pr_num)).get("state") - if pr_state != "closed": - print("PR #%s is still open after push; closing it explicitly." % pr_num) - close_pr(pr_num) - - # Walk a mutable remaining-branches list so the next default correctly skips any - # branches already picked, including branches consumed by the Upstream-First two-branch - # path inside cherry_pick (e.g. picking branch-M.x + branch-M.N in a single prompt). - # merged_refs doubles as the already_picked set passed to cherry_pick: it starts with - # target_ref (the merge sink, never to be re-picked) and grows with every cherry-pick. - remaining_branches = [b for b in branch_names if b != target_ref] + # merged_refs drives both the next prompt default and the already_picked set passed to + # cherry_pick, so each grows with every cherry-pick -- including branches consumed by the + # Upstream-First two-branch path inside cherry_pick (e.g. picking branch-M.x + branch-M.N + # in a single prompt). It starts with target_ref, the merge sink, never to be re-picked. pick_prompt = "Would you like to pick %s into another branch?" % merge_hash # Always record the merge summary for what actually landed, even if a later # cherry-pick is aborted or cancelled: the merge into the target branch has # already been pushed, so cancelling a backport must not drop that line. try: while get_input(f"\n{pick_prompt} (y/N): ", ["y", "n", ""]) == "y": - default = remaining_branches[0] if remaining_branches else branch_names[0] + default = default_pick_branch(branch_names, tuple(merged_refs)) + if default is None: + print("Every known release branch already contains #%s; nothing to pick." % pr_num) + break picked = cherry_pick( pr_num, merge_hash, @@ -1785,23 +2306,23 @@ def main(): target_ref, already_picked=tuple(merged_refs), ) - picked_refs = [ref for ref, _ in picked] - merged_refs = merged_refs + picked_refs + merged_refs = merged_refs + [ref for ref, _ in picked] merged_commits = merged_commits + picked - for b in picked_refs: - if b in remaining_branches: - remaining_branches.remove(b) finally: - post_merge_comment(pr_num, merged_commits) - - # asf_jira is guaranteed to be set here: initialize_jira() fails fast otherwise. - continue_maybe("Would you like to update an associated JIRA?") - jira_comment = "Issue resolved by pull request %s\n[%s/%s]" % ( - pr_num, - GITHUB_BASE, - pr_num, - ) - resolve_jira_issues(title, merged_refs, jira_comment, title_components) + if merged_commits: + # The "Closes #N" keyword in the commit message only auto-closes the PR when the + # commit lands on the default branch. For merges into other branches (e.g. + # branch-X.Y backport PRs), GitHub leaves the PR open, so close it through the API. + pr_state = get_json("%s/pulls/%s" % (GITHUB_API_BASE, pr_num)).get("state") + if pr_state != "closed": + close_message = "PR #%s is still open after push; closing it explicitly." % pr_num + print("\n%s\n" % bold(close_message)) + github.close_pr(pr_num) + # Record every branch that successfully received the change on the PR. + post_merge_comment(pr_num, merged_commits) + # This is deliberately in the finally block: once the target branch has been pushed, + # cancelling a later cherry-pick must not bypass the JIRA update decision. + update_jira_for_pr(pr_num, title, merged_refs, title_components) if __name__ == "__main__": diff --git a/dev/mima b/dev/mima index 39be02a1dd557..eecf1ad1081db 100755 --- a/dev/mima +++ b/dev/mima @@ -24,7 +24,7 @@ set -e FWDIR="$(cd "`dirname "$0"`"/..; pwd)" cd "$FWDIR" -SPARK_PROFILES=${1:-"-Pkubernetes -Pyarn -Pspark-ganglia-lgpl -Pkinesis-asl -Phive-thriftserver -Phive"} +SPARK_PROFILES=${1:-"-Pkubernetes -Pyarn -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Phive-thriftserver -Phive"} # Capture sbt output to show errors if the command fails OLD_DEPS_OUTPUT="$(build/sbt -DcopyDependencies=false $SPARK_PROFILES "export oldDeps/fullClasspath" 2>&1)" || { diff --git a/dev/package-lock.json b/dev/package-lock.json index 1075fcdbd4731..f4b80a113c573 100644 --- a/dev/package-lock.json +++ b/dev/package-lock.json @@ -237,9 +237,9 @@ "dev": true }, "node_modules/brace-expansion": { - "version": "1.1.16", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.16.tgz", - "integrity": "sha512-IDw48K2/2kRkg9LdJxurvq3lV3aBgq0REY89duEqFRthjlPdXHKMj7EnQOXVckxzgisinf3nHfrcE2FufFLXMw==", + "version": "1.1.18", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.18.tgz", + "integrity": "sha512-Edep/X9fGqVNmzKBVsDYIOtD+z1tuezV70LBjdCst9Tqu76lsnvRiZ6oTic1n+/BIwX6QDGAO94PN4N2SADvtw==", "dev": true, "license": "MIT", "dependencies": { @@ -593,9 +593,9 @@ "dev": true }, "node_modules/fast-uri": { - "version": "3.1.4", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.4.tgz", - "integrity": "sha512-8JnbkQ4juDyvYs4mgFGQqg4yCYtFDtUtmp2QIQq11ZZe5CFQ5wcqm1rqDgAh/QdMySuBnPzMUiJUNZG5N/AiQw==", + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", + "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", "dev": true, "funding": [ { @@ -802,9 +802,9 @@ "dev": true }, "node_modules/js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "funding": [ { @@ -1416,9 +1416,9 @@ "dev": true }, "brace-expansion": { - "version": "1.1.16", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.16.tgz", - "integrity": "sha512-IDw48K2/2kRkg9LdJxurvq3lV3aBgq0REY89duEqFRthjlPdXHKMj7EnQOXVckxzgisinf3nHfrcE2FufFLXMw==", + "version": "1.1.18", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.18.tgz", + "integrity": "sha512-Edep/X9fGqVNmzKBVsDYIOtD+z1tuezV70LBjdCst9Tqu76lsnvRiZ6oTic1n+/BIwX6QDGAO94PN4N2SADvtw==", "dev": true, "requires": { "balanced-match": "^1.0.0", @@ -1692,9 +1692,9 @@ "dev": true }, "fast-uri": { - "version": "3.1.4", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.4.tgz", - "integrity": "sha512-8JnbkQ4juDyvYs4mgFGQqg4yCYtFDtUtmp2QIQq11ZZe5CFQ5wcqm1rqDgAh/QdMySuBnPzMUiJUNZG5N/AiQw==", + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", + "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", "dev": true }, "file-entry-cache": { @@ -1844,9 +1844,9 @@ "dev": true }, "js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "requires": { "argparse": "^2.0.1" diff --git a/dev/pip-sanity-check.py b/dev/pip-sanity-check.py index fdb1107f930f4..27443b4602ef4 100644 --- a/dev/pip-sanity-check.py +++ b/dev/pip-sanity-check.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import SparkSession import sys +from pyspark.sql import SparkSession + if __name__ == "__main__": spark = SparkSession.builder.appName("PipSanityCheck").getOrCreate() sc = spark.sparkContext diff --git a/dev/pr_merge_status.py b/dev/pr_merge_status.py index 510409f22500d..590659be093ef 100755 --- a/dev/pr_merge_status.py +++ b/dev/pr_merge_status.py @@ -60,6 +60,10 @@ import subprocess import sys +# Shared with dev/merge_spark_pr.py so the two committer tools cannot disagree about where a +# PR landed. Importable because Python puts this script's own directory first on sys.path. +from spark_merge_footer import branches_with_merge_footer, merge_footer_trailer + REPO = "apache/spark" @@ -177,38 +181,6 @@ def fetch_branches(remote): ) -def commits_with_trailer(trailer, remote): - """Returns the full SHAs of commits on `remote`'s branches whose message contains - `trailer`. Scoping to the one remote (rather than `--all`) keeps fork refs and tags - from adding noise or walk cost.""" - out = git("log", "--remotes=%s" % remote, "--fixed-strings", "--grep", trailer, "--format=%H") - return list(dict.fromkeys(out.split())) - - -def official_branches_containing(commit, remote): - """Returns the `remote` branch names (e.g. 'master', 'branch-4.x') that contain - `commit`, ignoring the remote's HEAD alias and any non-branch refs.""" - out = git( - "for-each-ref", - "--contains", - commit, - "--format=%(refname:short)", - "refs/remotes/%s/" % remote, - ) - prefix = remote + "/" - branches = set() - for ref in out.splitlines(): - # Real branches are "<remote>/<branch>"; the remote's HEAD symref shortens to the - # bare remote name (e.g. "upstream") -- skip anything without the "<remote>/" prefix, - # and the explicit "<remote>/HEAD" form for good measure. - if not ref.startswith(prefix): - continue - name = ref[len(prefix) :] - if name != "HEAD": - branches.add(name) - return branches - - def display_key(name): """Sorts `master` first, then branch-<major>.<minor> ascending, with branch-<N>.x (the active dev line for the next feature release) after its numeric siblings.""" @@ -257,19 +229,25 @@ def main(): # its merge there, since a merge always lands on the base branch. majors = {m for m in (latest_major(remote), branch_major(base)) if m is not None} - trailer = "Closes #%s from " % pr - landed = {} - for commit in commits_with_trailer(trailer, remote): - for branch in official_branches_containing(commit, remote): - if all_branches or is_relevant(branch, majors): - landed[branch] = commit[:11] + # The shared reader consumes local remote-tracking refs only, so fetch_branches above is + # what refreshes them -- best-effort, since it warns and continues on a failed fetch, + # leaving the refs possibly stale (a recently merged PR could look unmerged). + all_landed = branches_with_merge_footer(pr, remote, lambda args: git(*args)) + landed = { + branch: commit[:11] + for branch, commit in all_landed.items() + if all_branches or is_relevant(branch, majors) + } if landed: print("merged: yes") for branch in sorted(landed, key=display_key): print(" %-12s %s" % (branch, landed[branch])) else: - print('closed without merging -- no "%s" commit found (rejected or superseded).' % trailer) + print( + 'closed without merging -- no "%s" commit found (rejected or superseded).' + % merge_footer_trailer(pr) + ) if __name__ == "__main__": diff --git a/dev/reformat-python b/dev/reformat-python index e7686f0b85ee4..102c74a929923 100755 --- a/dev/reformat-python +++ b/dev/reformat-python @@ -30,3 +30,4 @@ if [ $? -ne 0 ]; then fi $RUFF_BUILD format python/pyspark dev python/packaging python/benchmarks +$RUFF_BUILD check --select I --fix diff --git a/dev/requirements.txt b/dev/requirements.txt index a5ac9f06a668f..edcaf053a449a 100644 --- a/dev/requirements.txt +++ b/dev/requirements.txt @@ -21,6 +21,9 @@ openpyxl asv coverage tabulate +# Used by the opt-in UDF transpile hypothesis test suite (gated on +# RUN_HYPOTHESIS in the environment). See SPARK-54783. +hypothesis>=6,<7 # Kafka streaming test dependencies (optional) # Required for running Kafka integration tests with Docker test containers diff --git a/dev/run-tests.py b/dev/run-tests.py index d50307513bb01..35b3a70589278 100755 --- a/dev/run-tests.py +++ b/dev/run-tests.py @@ -18,22 +18,23 @@ # import itertools -from argparse import ArgumentParser import os import re -import sys import subprocess +import sys +import tempfile +from argparse import ArgumentParser from contextlib import contextmanager -from sparktestsupport import SPARK_HOME, USER_HOME, ERROR_CODES -from sparktestsupport.shellutils import exit_from_command_with_retcode, run_cmd, rm_r, which +import sparktestsupport.modules as modules +from sparktestsupport import SPARK_HOME, USER_HOME +from sparktestsupport.shellutils import exit_from_command_with_retcode, rm_r, run_cmd, which from sparktestsupport.utils import ( determine_dangling_python_tests, determine_modules_for_files, determine_modules_to_test, identify_changed_files_from_git_commits, ) -import sparktestsupport.modules as modules def setup_test_environ(environ): @@ -69,43 +70,38 @@ def determine_java_executable(): # ------------------------------------------------------------------------------------------------- -def set_title_and_block(title, err_block): - os.environ["CURRENT_BLOCK"] = str(ERROR_CODES[err_block]) +@contextmanager +def titled_block(title): + if getattr(titled_block, "_entered", False): + raise RuntimeError(f"titled_block({title!r}) cannot be nested") + titled_block._entered = True line_str = "=" * 72 - + if "GITHUB_ACTIONS" in os.environ: + print(f"::group::{title}", flush=True) print("") print(line_str) print(title) print(line_str) - - -@contextmanager -def group_in_github_actions(title): - if "GITHUB_ACTIONS" in os.environ: - print(f"::group::{title}", flush=True) - try: - yield - finally: - print("::endgroup::", flush=True) - else: + try: yield + finally: + titled_block._entered = False + if "GITHUB_ACTIONS" in os.environ: + print("::endgroup::", flush=True) def run_apache_rat_checks(): - set_title_and_block("Running Apache RAT checks", "BLOCK_RAT") run_cmd([os.path.join(SPARK_HOME, "dev", "check-license")]) def run_scala_style_checks(extra_profiles): build_profiles = extra_profiles + modules.root.build_profile_flags - set_title_and_block("Running Scala style checks", "BLOCK_SCALA_STYLE") profiles = " ".join(build_profiles) print("[info] Checking Scala style using SBT with these profiles: ", profiles) run_cmd([os.path.join(SPARK_HOME, "dev", "lint-scala"), profiles]) def run_java_style_checks(build_profiles): - set_title_and_block("Running Java style checks", "BLOCK_JAVA_STYLE") # The same profiles used for building are used to run Checkstyle by SBT as well because # the previous build looks reused for Checkstyle and affecting Checkstyle. See SPARK-27130. profiles = " ".join(build_profiles) @@ -114,13 +110,10 @@ def run_java_style_checks(build_profiles): def run_python_style_checks(): - set_title_and_block("Running Python style checks", "BLOCK_PYTHON_STYLE") run_cmd([os.path.join(SPARK_HOME, "dev", "lint-python")]) def run_sparkr_style_checks(): - set_title_and_block("Running R style checks", "BLOCK_R_STYLE") - if which("R"): # R style check should be executed after `install-dev.sh`. # Since warnings about `no visible global function definition` appear @@ -130,27 +123,6 @@ def run_sparkr_style_checks(): print("Ignoring SparkR style check as R was not found in PATH") -def build_spark_documentation(): - set_title_and_block("Building Spark Documentation", "BLOCK_DOCUMENTATION") - os.environ["PRODUCTION"] = "1" - - os.chdir(os.path.join(SPARK_HOME, "docs")) - - bundle_bin = which("bundle") - - if not bundle_bin: - print( - "[error] Cannot find a version of `bundle` on the system; please", - " install one with `gem install bundler` and retry to build documentation.", - ) - sys.exit(int(os.environ.get("CURRENT_BLOCK", 255))) - else: - run_cmd([bundle_bin, "install"]) - run_cmd([bundle_bin, "exec", "jekyll", "build"]) - - os.chdir(SPARK_HOME) - - def exec_maven(mvn_args=()): """Will call Maven in the current directory with the list of mvn_args passed in and returns the subprocess for any further processing""" @@ -208,15 +180,13 @@ def get_scala_profiles(scala_version): " are", sbt_maven_scala_profiles.keys(), ) - sys.exit(int(os.environ.get("CURRENT_BLOCK", 255))) + sys.exit(1) def switch_scala_version(scala_version): """ Switch the code base to use the given Scala version. """ - set_title_and_block("Switch the Scala version to %s" % scala_version, "BLOCK_SCALA_VERSION") - assert scala_version is not None ver_num = scala_version[-4:] # Simply extract. e.g.) 2.13 from scala2.13 command = [os.path.join(SPARK_HOME, "dev", "change-scala-version.sh"), ver_num] @@ -243,7 +213,7 @@ def get_hadoop_profiles(hadoop_version): " are", sbt_maven_hadoop_profiles.keys(), ) - sys.exit(int(os.environ.get("CURRENT_BLOCK", 255))) + sys.exit(1) def build_spark_maven(extra_profiles): @@ -269,12 +239,10 @@ def build_spark_sbt(extra_profiles): print("[info] Building Spark using SBT with these arguments: ", " ".join(profiles_and_goals)) - with group_in_github_actions("sbt build spark"): - exec_sbt(profiles_and_goals) + exec_sbt(profiles_and_goals) def build_spark_unidoc_sbt(extra_profiles): - set_title_and_block("Building Unidoc API Documentation", "BLOCK_DOCUMENTATION") # Enable all of the profiles for the build: build_profiles = extra_profiles + modules.root.build_profile_flags sbt_goals = ["unidoc"] @@ -288,7 +256,7 @@ def build_spark_unidoc_sbt(extra_profiles): exec_sbt(profiles_and_goals) -def build_spark_assembly_sbt(extra_profiles, checkstyle=False): +def build_spark_assembly_sbt(extra_profiles): # Enable all of the profiles for the build: build_profiles = extra_profiles + modules.root.build_profile_flags sbt_goals = ["assembly/package"] @@ -298,22 +266,13 @@ def build_spark_assembly_sbt(extra_profiles, checkstyle=False): " ".join(profiles_and_goals), ) - with group_in_github_actions("sbt build spark assembly"): - exec_sbt(profiles_and_goals) - - if checkstyle: - run_java_style_checks(build_profiles) - - if not os.environ.get("SKIP_UNIDOC"): - build_spark_unidoc_sbt(extra_profiles) + exec_sbt(profiles_and_goals) def build_apache_spark(build_tool, extra_profiles): """Will build Spark with the extra profiles and the passed in build tool (either `sbt` or `maven`). Defaults to using `sbt`.""" - set_title_and_block("Building Spark", "BLOCK_BUILD") - rm_r("lib_managed") if build_tool == "maven": @@ -324,7 +283,6 @@ def build_apache_spark(build_tool, extra_profiles): def detect_binary_inop_with_mima(extra_profiles): build_profiles = extra_profiles + modules.root.build_profile_flags - set_title_and_block("Detecting binary incompatibilities with MiMa", "BLOCK_MIMA") profiles = " ".join(build_profiles) print( "[info] Detecting binary incompatibilities with MiMa using SBT with these profiles: ", @@ -364,7 +322,6 @@ def run_scala_tests_sbt(test_modules, test_profiles): def run_scala_tests(build_tool, extra_profiles, test_modules, excluded_tags, included_tags): """Function to properly execute all tests passed in as a set from the `determine_test_suites` function""" - set_title_and_block("Running Spark unit tests", "BLOCK_SPARK_UNIT_TESTS") # Remove duplicates while keeping the test module order test_modules = list(dict.fromkeys(test_modules)) @@ -392,9 +349,9 @@ def run_scala_tests(build_tool, extra_profiles, test_modules, excluded_tags, inc run_scala_tests_sbt(test_modules, test_profiles) -def run_python_tests(test_modules, test_pythons, parallelism, with_coverage=False): - set_title_and_block("Running PySpark tests", "BLOCK_PYSPARK_UNIT_TESTS") - +def run_python_tests( + test_modules, test_pythons, parallelism, changed_files=None, with_coverage=False +): if with_coverage: # Coverage makes the PySpark tests flaky due to heavy parallelism. # When we run PySpark tests with coverage, it uses 4 for now as @@ -408,24 +365,26 @@ def run_python_tests(test_modules, test_pythons, parallelism, with_coverage=Fals command.append("--modules=%s" % ",".join(m.name for m in test_modules)) command.append("--parallelism=%i" % parallelism) command.append("--python-executables=%s" % test_pythons) - run_cmd(command) + if changed_files: + with tempfile.NamedTemporaryFile("w") as f: + f.write("\n".join(changed_files)) + f.flush() + command.append("--changed-files=%s" % f.name) + run_cmd(command) + else: + run_cmd(command) def run_python_packaging_tests(): - if os.environ.get("SKIP_PACKAGING", "false") != "true": - set_title_and_block("Running PySpark packaging tests", "BLOCK_PYSPARK_PIP_TESTS") - command = [os.path.join(SPARK_HOME, "dev", "run-pip-tests")] - run_cmd(command) + command = [os.path.join(SPARK_HOME, "dev", "run-pip-tests")] + run_cmd(command) def run_build_tests(): - set_title_and_block("Running build tests", "BLOCK_BUILD_TESTS") run_cmd([os.path.join(SPARK_HOME, "dev", "test-dependencies.sh")]) def run_sparkr_tests(): - set_title_and_block("Running SparkR tests", "BLOCK_SPARKR_UNIT_TESTS") - if which("R"): run_cmd([os.path.join(SPARK_HOME, "R", "run-tests.sh")]) else: @@ -498,8 +457,6 @@ def main(): rm_r(os.path.join(USER_HOME, ".ivy2.5.2", "local", "org.apache.spark")) rm_r(os.path.join(USER_HOME, ".ivy2.5.2", "cache", "org.apache.spark")) - os.environ["CURRENT_BLOCK"] = str(ERROR_CODES["BLOCK_GENERAL"]) - java_exe = determine_java_executable() if not java_exe: @@ -509,7 +466,6 @@ def main(): ) sys.exit(2) - # Install SparkR should_only_test_modules = opts.modules is not None test_modules = [] if should_only_test_modules: @@ -520,7 +476,8 @@ def main(): # If tests modules are specified, we will not run R linter. # SparkR needs the manual SparkR installation. if which("R"): - run_cmd([os.path.join(SPARK_HOME, "R", "install-dev.sh")]) + with titled_block("Installing SparkR"): + run_cmd([os.path.join(SPARK_HOME, "R", "install-dev.sh")]) else: print("Cannot install SparkR as R was not found in PATH") @@ -610,18 +567,21 @@ def main(): if scala_version is not None: # If not set, assume this is default and doesn't need to change. - switch_scala_version(scala_version) + with titled_block(f"Switching to Scala version: {scala_version}"): + switch_scala_version(scala_version) should_run_java_style_checks = False if not should_only_test_modules: # license checks - run_apache_rat_checks() + with titled_block("Running Apache RAT checks"): + run_apache_rat_checks() # style checks if not changed_files or any( f.endswith(".scala") or f.endswith("scalastyle-config.xml") for f in changed_files ): - run_scala_style_checks(extra_profiles) + with titled_block("Running Scala style checks"): + run_scala_style_checks(extra_profiles) if not changed_files or any( f.endswith(".java") or f.endswith("checkstyle.xml") @@ -634,50 +594,87 @@ def main(): f.endswith("lint-python") or f.endswith("pyproject.toml") or f.endswith(".py") for f in changed_files ): - run_python_style_checks() + with titled_block("Running Python style checks"): + run_python_style_checks() if not changed_files or any( f.endswith(".R") or f.endswith("lint-r") or f.endswith(".lintr") for f in changed_files ): - run_sparkr_style_checks() + with titled_block("Running R style checks"): + run_sparkr_style_checks() if any(m.should_run_build_tests for m in test_modules): - run_build_tests() + with titled_block("Running build tests"): + run_build_tests() # spark build if os.environ.get("SKIP_SCALA_BUILD", "false") != "true": - build_apache_spark(build_tool, extra_profiles) + with titled_block("Building Spark"): + build_apache_spark(build_tool, extra_profiles) # backwards compatibility checks if build_tool == "sbt": # Note: compatibility tests only supported in sbt for now if not os.environ.get("SKIP_MIMA"): - detect_binary_inop_with_mima(extra_profiles) + with titled_block("Detecting binary incompatibilities with MiMa"): + detect_binary_inop_with_mima(extra_profiles) # Since we did not build assembly/package before running dev/mima, we need to # do it here because the tests still rely on it; see SPARK-13294 for details. if os.environ.get("SKIP_SCALA_BUILD", "false") != "true": - build_spark_assembly_sbt(extra_profiles, should_run_java_style_checks) + with titled_block("Building Spark assembly"): + build_spark_assembly_sbt(extra_profiles) + if should_run_java_style_checks: + with titled_block("Running Java style checks"): + run_java_style_checks(extra_profiles + modules.root.build_profile_flags) + if not os.environ.get("SKIP_UNIDOC"): + with titled_block("Building Unidoc API Documentation"): + build_spark_unidoc_sbt(extra_profiles) # run the test suites - run_scala_tests(build_tool, extra_profiles, test_modules, excluded_tags, included_tags) + with titled_block("Running Spark unit tests"): + run_scala_tests(build_tool, extra_profiles, test_modules, excluded_tags, included_tags) modules_with_python_tests = [m for m in test_modules if m.python_test_goals] if modules_with_python_tests and not os.environ.get("SKIP_PYTHON"): - run_python_tests( - modules_with_python_tests, - opts.python_executables, - opts.parallelism, - with_coverage=os.environ.get("PYSPARK_CODECOV", "false") == "true", - ) - run_python_packaging_tests() + relevant_changed_files = None + # We only do smart test selection on push action of apache/spark + # If APACHE_SPARK_REF is set, we are in a forked repository. + # Otherwise if we have a list of changed files, we must be in post-merge CI. + if not os.environ.get("APACHE_SPARK_REF", "") and changed_files: + relevant_changed_files = [f for f in changed_files if not modules.is_ignored_file(f)] + # If there are relevant changed files that are not pyspark, we don't do smart test + if any( + not (f.endswith(".py") and f.startswith("python/pyspark/")) + for f in relevant_changed_files + ): + relevant_changed_files = None + with titled_block("Running PySpark tests"): + run_python_tests( + modules_with_python_tests, + opts.python_executables, + opts.parallelism, + changed_files=relevant_changed_files, + with_coverage=os.environ.get("PYSPARK_CODECOV", "false") == "true", + ) + if os.environ.get("SKIP_PACKAGING", "false") != "true": + with titled_block("Running PySpark packaging tests"): + run_python_packaging_tests() if any(m.should_run_r_tests for m in test_modules) and not os.environ.get("SKIP_R"): - run_sparkr_tests() + with titled_block("Running SparkR tests"): + run_sparkr_tests() def _test(): import doctest + + import sparktestsupport.modules import sparktestsupport.utils - failure_count = doctest.testmod(sparktestsupport.utils)[0] + doctest.testmod()[0] + test_results = ( + doctest.testmod(sparktestsupport.modules), + doctest.testmod(sparktestsupport.utils), + doctest.testmod(), + ) + failure_count = sum([num_failures for (num_failures, num_tests) in test_results]) if failure_count: sys.exit(-1) diff --git a/dev/sbt-checkstyle b/dev/sbt-checkstyle index f2d5a0fa304ac..a148af6039181 100755 --- a/dev/sbt-checkstyle +++ b/dev/sbt-checkstyle @@ -17,7 +17,7 @@ # limitations under the License. # -SPARK_PROFILES=${1:-"-Pkinesis-asl -Pspark-ganglia-lgpl -Pkubernetes -Pyarn -Phive -Phive-thriftserver -Pjvm-profiler"} +SPARK_PROFILES=${1:-"-Pkinesis-asl -Pcredential-aws -Pspark-ganglia-lgpl -Pkubernetes -Pyarn -Phive -Phive-thriftserver -Pjvm-profiler"} # NOTE: echo "q" is needed because SBT prompts the user for input on encountering a build file # with failure (either resolution or compilation); the "q" makes SBT quit. diff --git a/dev/scalastyle b/dev/scalastyle index 09e6c2372614d..f77fb3e1ee224 100755 --- a/dev/scalastyle +++ b/dev/scalastyle @@ -17,7 +17,7 @@ # limitations under the License. # -SPARK_PROFILES=${1:-"-Pkubernetes -Pyarn -Pspark-ganglia-lgpl -Pkinesis-asl -Phive-thriftserver -Phive -Pvolcano -Pjvm-profiler -Phadoop-cloud -Pdocker-integration-tests -Pkubernetes-integration-tests"} +SPARK_PROFILES=${1:-"-Pkubernetes -Pyarn -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Phive-thriftserver -Phive -Pvolcano -Pjvm-profiler -Phadoop-cloud -Pdocker-integration-tests -Pkubernetes-integration-tests"} # NOTE: echo "q" is needed because SBT prompts the user for input on encountering a build file # with failure (either resolution or compilation); the "q" makes SBT quit. diff --git a/dev/spark-test-image-util/docs/build-docs b/dev/spark-test-image-util/docs/build-docs index ca59769f24231..c2ee826832ef6 100755 --- a/dev/spark-test-image-util/docs/build-docs +++ b/dev/spark-test-image-util/docs/build-docs @@ -35,7 +35,7 @@ FWDIR="$(cd "`dirname "${BASH_SOURCE[0]}"`"; pwd)" SPARK_HOME="$(cd "`dirname "${BASH_SOURCE[0]}"`"/../../..; pwd)" # 1.Compile spark outside the container to prepare for generating documents inside the container. -build/sbt -Phive -Pkinesis-asl clean unidoc package +build/sbt -Phive -Pkinesis-asl -Pcredential-aws clean unidoc package # 2.Build container image. docker buildx build \ diff --git a/dev/spark-test-image/python-312-classic-only/Dockerfile b/dev/spark-test-image/python-312-classic-only/Dockerfile index ceb4694b2dc9d..7dfa024f4aabd 100644 --- a/dev/spark-test-image/python-312-classic-only/Dockerfile +++ b/dev/spark-test-image/python-312-classic-only/Dockerfile @@ -59,11 +59,10 @@ ENV VIRTUAL_ENV=/opt/spark-venv RUN python3.12 -m venv $VIRTUAL_ENV ENV PATH="$VIRTUAL_ENV/bin:$PATH" -ARG BASIC_PIP_PKGS="numpy pyarrow>=23.0.0 pandas==2.3.3 plotly<6.0.0 matplotlib openpyxl memory-profiler>=0.61.0 mlflow>=2.8.1 scipy scikit-learn>=1.3.2 pystack>=1.6.0 psutil" -ARG TEST_PIP_PKGS="coverage unittest-xml-reporting" +COPY --from=root pyproject.toml ./pyproject.toml -RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 -RUN python3.12 -m pip install $BASIC_PIP_PKGS $TEST_PIP_PKGS && \ - python3.12 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \ - python3.12 -m pip install deepspeed torcheval && \ +RUN python3.12 -m pip install -U pip + +RUN python3.12 -m pip install --group ml_torch --index-url https://download.pytorch.org/whl/cpu && \ + python3.12 -m pip install --group ci_classic_standard && \ python3.12 -m pip cache purge diff --git a/dev/spark-test-image/python-minimum/Dockerfile b/dev/spark-test-image/python-minimum/Dockerfile index 89da6f618124f..292f4f328ced9 100644 --- a/dev/spark-test-image/python-minimum/Dockerfile +++ b/dev/spark-test-image/python-minimum/Dockerfile @@ -64,8 +64,9 @@ ENV VIRTUAL_ENV=/opt/spark-venv RUN python3.11 -m venv $VIRTUAL_ENV ENV PATH="$VIRTUAL_ENV/bin:$PATH" -ARG BASIC_PIP_PKGS="numpy==1.23.2 pyarrow==18.0.0 pandas==2.2.0 six==1.16.0 scipy scikit-learn coverage unittest-xml-reporting psutil" -ARG CONNECT_PIP_PKGS="grpcio==1.76.0 grpcio-status==1.76.0 googleapis-common-protos==1.71.0 zstandard==0.25.0 graphviz==0.20 protobuf==6.33.5" +COPY --from=root pyproject.toml ./pyproject.toml -RUN python3.11 -m pip install --force $BASIC_PIP_PKGS $CONNECT_PIP_PKGS && \ +RUN python3.11 -m pip install -U pip + +RUN python3.11 -m pip install --group ci_classic_minimum --group ci_connect_minimum && \ python3.11 -m pip cache purge diff --git a/dev/spark_merge_footer.py b/dev/spark_merge_footer.py new file mode 100644 index 0000000000000..567895b716ad5 --- /dev/null +++ b/dev/spark_merge_footer.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 + +# +# 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. +# + +""" +Shared reader for the merge footer that `dev/merge_spark_pr.py` writes into every commit +it creates, so the committer tools cannot disagree about where a pull request landed. + +`merge_pr` ends each message it generates with + + Closes #<pr> from <author>/<branch>. + + Authored-by: A <a@example.org> + Signed-off-by: C <c@example.org> + +and `git cherry-pick -x` copies that footer verbatim into every backport, appending its own +provenance lines after it. The footer is therefore the signal that identifies both a merge +and its backports -- `git ... --contains <merge_hash>` cannot, because a cherry-pick is a +new commit that no other branch contains. + +Two properties make reading it reliable, and both are easy to get wrong: + +- A PR body is passed through as its own `git commit -m` paragraph, so it may quote another + commit's footer in full, structure included. Only *position* distinguishes the generated + footer: `merge_pr` appends it last, so the generated one is the final "Closes" paragraph. +- `git log --grep` matches its pattern anywhere in a message, so it can only narrow the + walk; every candidate it returns must still be validated with `has_merge_footer`. + +When imported, nothing here exits, prints, or runs git: callers pass a `run_git` callable and +so keep their own error-handling policy -- `dev/pr_merge_status.py` exits on a git failure, +while `dev/merge_spark_pr.py` must not abort a merge in progress. (Running this file directly +executes its doctests and exits nonzero if any fail.) + +Refresh policy: `branches_with_merge_footer` reads local remote-tracking refs only and +never fetches. A caller that needs current data fetches first (as `pr_merge_status.py` +does, best-effort); a caller that must not touch the network mid-run simply accepts that a +branch not yet fetched goes unreported. +""" + +import re + +# The generated footer: a "Closes #<pr> from <ref>" line alone on its paragraph, followed by +# the authors paragraph. `\s*$` tolerates trailing whitespace. Requiring the blank line and +# the authors line rejects prose that merely mentions a PR; taking the *last* match (see +# `merge_footer_pr`) rejects a body that quotes a real footer. +_MERGE_FOOTER_RE = re.compile( + r"^Closes #(\d+) from \S+\s*$\n\n(?:Lead-authored-by|Authored-by):", + re.MULTILINE, +) + + +def merge_footer_trailer(pr_num): + """The literal fragment to pass to `git log --fixed-strings --grep`. + + Only a prefilter to narrow the walk: it matches anywhere in a message, so callers + validate each candidate with `has_merge_footer`. + + >>> merge_footer_trailer(1) + 'Closes #1 from ' + """ + return "Closes #%s from " % pr_num + + +def merge_footer_pr(message): + """The PR number in `message`'s generated merge footer, or None if it has none. + + Reads the *last* "Closes" paragraph, since a PR body copied into the message may quote + an earlier one. Cherry-pick provenance lines may follow the footer, but no later + "Closes" paragraph can. + + >>> footer = "Closes #1 from a/b.\\n\\nAuthored-by: A <a@e.org>\\nSigned-off-by: C <c@e.org>" + >>> merge_footer_pr("[SPARK-1][SQL] Title\\n\\nSome body.\\n\\n" + footer) + 1 + >>> merge_footer_pr("[SPARK-1][SQL] Title\\n\\n" + footer.replace("Authored", "Lead-authored")) + 1 + >>> merge_footer_pr("[SPARK-1][SQL] Title\\n\\nNo footer here.") is None + True + + A cherry-pick keeps the footer, with `-x` provenance appended after it: + + >>> pick = footer + "\\n(cherry picked from commit abc123)\\nSigned-off-by: C <c@e.org>" + >>> merge_footer_pr("[SPARK-1][SQL] Title\\n\\n" + pick) + 1 + + A body quoting another PR's complete footer does not shadow the real one: + + >>> quoted = "Reverting:\\n\\n" + footer + "\\n\\nSee above." + >>> own = footer.replace("#1", "#2") + >>> merge_footer_pr("[SPARK-2][SQL] Later\\n\\n%s\\n\\n%s" % (quoted, own)) + 2 + """ + matches = _MERGE_FOOTER_RE.findall(message) + return int(matches[-1]) if matches else None + + +def has_merge_footer(message, pr_num): + """Whether `message`'s generated merge footer closes `pr_num`. See `merge_footer_pr`. + + `pr_num` may be an int or a string of digits: callers get the PR number from argv or from + the GitHub API, and comparing those two forms directly would silently never match. + + >>> footer = "Closes #1 from a/b.\\n\\nAuthored-by: A <a@e.org>\\nSigned-off-by: C <c@e.org>" + >>> has_merge_footer("[SPARK-1][SQL] Title\\n\\n" + footer, 1) + True + >>> has_merge_footer("[SPARK-1][SQL] Title\\n\\n" + footer, "1") + True + >>> has_merge_footer("[SPARK-1][SQL] Title\\n\\n" + footer, 2) + False + + A commit whose body quotes another PR's full footer is not taken for that PR's merge: + + >>> quoted = "Reverting:\\n\\n" + footer + "\\n\\nSee above." + >>> later = "[SPARK-2][SQL] Later\\n\\n%s\\n\\n%s" % (quoted, footer.replace("#1", "#2")) + >>> has_merge_footer(later, 1) + False + >>> has_merge_footer(later, 2) + True + """ + return merge_footer_pr(message) == int(pr_num) + + +def parse_commit_records(out): + """Parse `git log --format='%H %B%x00'` output into (commit_hash, message) pairs. + + A commit message spans lines, so records are NUL-delimited rather than newline-delimited. + + >>> parse_commit_records("abc first\\nline two\\x00def second\\x00") + [('abc', 'first\\nline two'), ('def', 'second')] + >>> parse_commit_records("") + [] + """ + records = [] + for record in out.split("\0"): + record = record.strip("\n") + if not record: + continue + commit_hash, _, message = record.partition(" ") + records.append((commit_hash, message)) + return records + + +def branch_names_from_refs(out, remote): + """Branch names in `git for-each-ref --format='%(refname:short)'` output for `remote`. + + Real branches are "<remote>/<branch>"; the remote's HEAD symref shortens to the bare + remote name, so anything without the "<remote>/" prefix is skipped, as is the explicit + "<remote>/HEAD" form. + + >>> sorted(branch_names_from_refs("up/master\\nup/branch-4.x\\nup\\nup/HEAD\\n", "up")) + ['branch-4.x', 'master'] + """ + prefix = remote + "/" + names = set() + for ref in out.splitlines(): + if not ref.startswith(prefix): + continue + name = ref[len(prefix) :] + if name != "HEAD": + names.add(name) + return names + + +def branches_with_merge_footer(pr_num, remote, run_git): + """Map each `remote` branch carrying `pr_num`'s merge footer to the commit that has it. + + `run_git(args)` runs `git` with `args` and returns its stdout; the caller supplies it so + this module imposes no error-handling or exit policy of its own. Reads local + remote-tracking refs only -- see this module's refresh policy. + + Scoping the walk to `--remotes=<remote>` keeps fork refs and tags from adding noise or + cost. Every commit `--grep` returns is validated before its branches count, so a commit + that merely quotes the trailer cannot make a branch look like it has the change. + """ + out = run_git( + [ + "log", + "--remotes=%s" % remote, + "--fixed-strings", + "--grep", + merge_footer_trailer(pr_num), + "--format=%H %B%x00", + ] + ) + landed = {} + for commit_hash, message in parse_commit_records(out): + if not has_merge_footer(message, pr_num): + continue + refs = run_git( + [ + "for-each-ref", + "--contains", + commit_hash, + "--format=%(refname:short)", + "refs/remotes/%s/" % remote, + ] + ) + for branch in branch_names_from_refs(refs, remote): + landed[branch] = commit_hash + return landed + + +if __name__ == "__main__": + import doctest + import sys + + failure_count, test_count = doctest.testmod() + if failure_count: + sys.exit(-1) diff --git a/dev/sparktestsupport/__init__.py b/dev/sparktestsupport/__init__.py index b4edb6b62c20b..12696d98fb988 100644 --- a/dev/sparktestsupport/__init__.py +++ b/dev/sparktestsupport/__init__.py @@ -19,21 +19,3 @@ SPARK_HOME = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../../")) USER_HOME = os.environ.get("HOME") -ERROR_CODES = { - "BLOCK_GENERAL": 10, - "BLOCK_RAT": 11, - "BLOCK_SCALA_STYLE": 12, - "BLOCK_PYTHON_STYLE": 13, - "BLOCK_R_STYLE": 14, - "BLOCK_DOCUMENTATION": 15, - "BLOCK_BUILD": 16, - "BLOCK_MIMA": 17, - "BLOCK_SPARK_UNIT_TESTS": 18, - "BLOCK_PYSPARK_UNIT_TESTS": 19, - "BLOCK_SPARKR_UNIT_TESTS": 20, - "BLOCK_JAVA_STYLE": 21, - "BLOCK_BUILD_TESTS": 22, - "BLOCK_PYSPARK_PIP_TESTS": 23, - "BLOCK_SCALA_VERSION": 24, - "BLOCK_TIMEOUT": 124, -} diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 77aa43eb8d947..5d3efb657b9cd 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -15,14 +15,92 @@ # limitations under the License. # -from functools import total_ordering import itertools import os import re -from pathlib import Path +import sys +from functools import total_ordering +from pathlib import Path, PurePath all_modules = [] +# These are `pathlib.PurePath` glob-style patterns with some customization: +# - Bare patterns match a file name at any depth. +# - Leading slashes anchor at the repository root. +# - A trailing slash denotes an ignored directory subtree. +# See: https://docs.python.org/3/library/pathlib.html#pathlib.PurePath.match +# +# Rejected alternatives: +# - Regexes present a footgun in that `.` needs to be carefully escaped but +# often isn't. +# - `.gitignore`-style patterns would be ideal but don't have support in the +# standard library. +ignored_file_patterns = ( + ".asf.yaml", + ".gitignore", + "AGENTS.md", + "CONTRIBUTING.md", + "README.md", + "/LICENSE-binary", + "/NOTICE-binary", + "/scalastyle-config.xml", + "/SECURITY.md", + "/dev/checkstyle-suppressions.xml", + "/dev/checkstyle.xml", + "/dev/create_jira_and_branch.py", + "/dev/create_spark_jira.py", + "/dev/create-release/", + "/dev/lint-python", + "/dev/lint-scala", + "/dev/make-distribution.sh", + "/dev/merge_spark_pr.py", + "/dev/pr_merge_status.py", + "/dev/reformat-python", + "/dev/requirements.txt", + "/dev/spark_merge_footer.py", + "/dev/spark-test-image/lint/Dockerfile", + "/dev/structured_logging_style.py", + "/ui-test/package-lock.json", + "/ui-test/package.json", +) + + +def is_ignored_file(filename: str) -> bool: + """ + Return whether a repository-relative path should be ignored when selecting + test modules. + + Bare patterns match a file name at any depth: + >>> is_ignored_file("python/README.md") + True + + Leading slashes anchor at the repository root: + >>> is_ignored_file("SECURITY.md") + True + >>> is_ignored_file("docs/SECURITY.md") + False + + A trailing slash ignores a directory subtree: + >>> is_ignored_file("dev/create-release/spark-rm/Dockerfile") + True + + Non-matches fall through: + >>> is_ignored_file("xasfZyaml") + False + """ + path = PurePath("/") / filename + for pattern in ignored_file_patterns: + # TODO: When Python 3.13 becomes the minimum supported version, migrate + # to `PurePath.full_match` and use `**` patterns instead of this custom + # trailing slash behavior. + if pattern.endswith("/"): + ignored_directory = PurePath(pattern) + if path == ignored_directory or ignored_directory in path.parents: + return True + elif path.match(pattern): + return True + return False + @total_ordering class Module(object): @@ -410,6 +488,21 @@ def __hash__(self): ) +credential_aws = Module( + name="credential-aws", + dependencies=[tags, core], + source_file_regexes=[ + "connector/credential-aws/", + ], + build_profile_flags=[ + "-Pcredential-aws", + ], + sbt_test_goals=[ + "credential-aws/test", + ], +) + + streaming_kafka_0_10 = Module( name="streaming-kafka-0-10", dependencies=[streaming, core], @@ -522,12 +615,17 @@ def __hash__(self): "pyspark.tests.test_zero_copy_byte_stream", # unittests for upstream projects "pyspark.tests.upstream.pyarrow.test_pyarrow_array_cast", + "pyspark.tests.upstream.pyarrow.test_pyarrow_array_from_pandas_default", + "pyspark.tests.upstream.pyarrow.test_pyarrow_array_from_pandas_non_default", "pyspark.tests.upstream.pyarrow.test_pyarrow_array_type_inference", "pyspark.tests.upstream.pyarrow.test_pyarrow_arrow_to_pandas_default", "pyspark.tests.upstream.pyarrow.test_pyarrow_arrow_to_pandas_non_default", + "pyspark.tests.upstream.pyarrow.test_pyarrow_dataframe_from_pandas", "pyspark.tests.upstream.pyarrow.test_pyarrow_ignore_timezone", "pyspark.tests.upstream.pyarrow.test_pyarrow_scalar_type_coercion", "pyspark.tests.upstream.pyarrow.test_pyarrow_scalar_type_inference", + "pyspark.tests.upstream.pyarrow.test_pyarrow_table_cast", + "pyspark.tests.upstream.pyarrow.test_pyarrow_table_to_pandas", "pyspark.tests.upstream.pyarrow.test_pyarrow_type_coercion", ], ) @@ -593,6 +691,7 @@ def __hash__(self): "pyspark.sql.tests.arrow.test_arrow_cogrouped_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc", "pyspark.sql.tests.arrow.test_arrow_grouped_map", + "pyspark.sql.tests.arrow.test_arrow_python_aggregator", "pyspark.sql.tests.arrow.test_arrow_python_udf", "pyspark.sql.tests.arrow.test_arrow_python_udf_cached", "pyspark.sql.tests.arrow.test_arrow_udf", @@ -626,7 +725,11 @@ def __hash__(self): "pyspark.sql.tests.test_geometrytype", "pyspark.sql.tests.test_udf", "pyspark.sql.tests.test_udf_combinations", + "pyspark.sql.tests.test_udf_in_higher_order_function", "pyspark.sql.tests.test_udf_profiler", + "pyspark.sql.tests.test_udf_transpile_hypothesis", + "pyspark.sql.tests.test_udf_transpile_parity", + "pyspark.sql.tests.test_udf_transpile_unit", "pyspark.sql.tests.test_unified_udf", "pyspark.sql.tests.test_udtf", "pyspark.sql.tests.test_tvf", @@ -653,6 +756,7 @@ def __hash__(self): "pyspark.testing.utils", "pyspark.testing.pandasutils", # unittests + "pyspark.testing.tests.test_changed_files", "pyspark.testing.tests.test_fail", "pyspark.testing.tests.test_fail_in_set_up_class", "pyspark.testing.tests.test_no_tests", @@ -1178,6 +1282,8 @@ def __hash__(self): "pyspark.sql.tests.connect.test_connect_readwriter", "pyspark.sql.tests.connect.test_connect_retry", "pyspark.sql.tests.connect.test_connect_session", + "pyspark.sql.tests.connect.test_connect_local_server", + "pyspark.sql.tests.connect.test_connect_local_server_pool", "pyspark.sql.tests.connect.test_connect_stat", "pyspark.sql.tests.connect.test_parity_geographytype", "pyspark.sql.tests.connect.test_parity_geometrytype", @@ -1206,6 +1312,7 @@ def __hash__(self): "pyspark.sql.tests.connect.test_parity_column", "pyspark.sql.tests.connect.test_parity_readwriter", "pyspark.sql.tests.connect.test_parity_udf", + "pyspark.sql.tests.connect.test_parity_udf_in_higher_order_function", "pyspark.sql.tests.connect.test_parity_udf_combinations", "pyspark.sql.tests.connect.test_parity_udf_profiler", "pyspark.sql.tests.connect.test_parity_unified_udf", @@ -1230,6 +1337,7 @@ def __hash__(self): "pyspark.sql.tests.connect.arrow.test_parity_arrow_grouped_map", "pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map", "pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map_misc", + "pyspark.sql.tests.connect.arrow.test_parity_arrow_python_aggregator", "pyspark.sql.tests.connect.arrow.test_parity_arrow_python_udf", "pyspark.sql.tests.connect.arrow.test_parity_arrow_udf", "pyspark.sql.tests.connect.arrow.test_parity_arrow_udf_scalar", @@ -1707,45 +1815,6 @@ def __hash__(self): test_tags=["org.apache.spark.tags.DockerTest"], ) - -# dev_tools is a pseudo module that contains all the dev related files that -# won't impact the CI build and tests (except for CI which is forced to -# run anyway). -# This module is created so modifying files in this module won't trigger any -# tests to run. -dev_tools = Module( - name="dev-tools", - dependencies=[], - source_file_regexes=[ - ".*README.md", - ".*AGENTS.md", - r".*\.gitignore", - "CONTRIBUTING.md", - ".asf.yaml", - "SECURITY.md", - "NOTICE-binary", - "LICENSE-binary", - "ui-test/package.json", - "ui-test/package-lock.json", - "scalastyle-config.xml", - "dev/checkstyle.xml", - "dev/checkstyle-suppressions.xml", - "dev/create_jira_and_branch.py", - "dev/create_spark_jira.py", - "dev/spark-test-image/lint/Dockerfile", - "dev/lint-python", - "dev/lint-scala", - "dev/reformat-python", - "dev/structured_logging_style.py", - "dev/make-distribution.sh", - "dev/merge_spark_pr.py", - "dev/requirements.txt", - "dev/pr_merge_status.py", - "dev/create_spark_jira.py", - "dev/create-release/", - ], -) - # The root module is a dummy module which is used to run all of the tests. # No other modules should directly depend on this module. root = Module( @@ -1763,3 +1832,15 @@ def __hash__(self): should_run_r_tests=True, should_run_build_tests=True, ) + + +def _test(): + import doctest + + failure_count = doctest.testmod()[0] + if failure_count: + sys.exit(-1) + + +if __name__ == "__main__": + _test() diff --git a/dev/sparktestsupport/shellutils.py b/dev/sparktestsupport/shellutils.py index 1d40ae9f2718c..662b8c8cc0ec3 100644 --- a/dev/sparktestsupport/shellutils.py +++ b/dev/sparktestsupport/shellutils.py @@ -28,7 +28,7 @@ def exit_from_command_with_retcode(cmd, retcode): print("[error] running", " ".join(cmd), "; process was terminated by signal", -retcode) else: print("[error] running", " ".join(cmd), "; received return code", retcode) - sys.exit(int(os.environ.get("CURRENT_BLOCK", 255))) + sys.exit(1) def rm_r(path): diff --git a/dev/sparktestsupport/utils.py b/dev/sparktestsupport/utils.py index 029b1627bd0bd..b177952eb31dd 100755 --- a/dev/sparktestsupport/utils.py +++ b/dev/sparktestsupport/utils.py @@ -18,8 +18,9 @@ # import os -import sys import subprocess +import sys + from sparktestsupport import modules from sparktestsupport.shellutils import run_cmd from sparktestsupport.toposort import toposort_flatten @@ -39,9 +40,13 @@ def determine_modules_for_files(filenames): ['pyspark-core', 'pyspark-install', 'sql'] >>> [x.name for x in determine_modules_for_files(["file_not_matched_by_any_subproject"])] ['root'] + >>> [x.name for x in determine_modules_for_files(["python/README.md"])] + [] """ changed_modules = set() for filename in filenames: + if modules.is_ignored_file(filename): + continue if ("GITHUB_ACTIONS" not in os.environ) and filename.startswith(".github"): continue matched_at_least_one_module = False @@ -121,7 +126,8 @@ def determine_modules_to_test(changed_modules, deduplicated=True): >>> sorted([x.name for x in determine_modules_to_test( ... [modules.sql, modules.core], deduplicated=False)]) ... # doctest: +NORMALIZE_WHITESPACE - ['avro', 'catalyst', 'connect', 'core', 'docker-integration-tests', 'examples', 'graphx', + ['avro', 'catalyst', 'connect', 'core', 'credential-aws', 'docker-integration-tests', + 'examples', 'graphx', 'hive', 'hive-thriftserver', 'mllib', 'mllib-local', 'pipelines', 'protobuf', 'pyspark-connect', 'pyspark-core', 'pyspark-errors', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', diff --git a/dev/structured_logging_style.py b/dev/structured_logging_style.py index b0d1b2a76e4ae..9689a6455bcc8 100755 --- a/dev/structured_logging_style.py +++ b/dev/structured_logging_style.py @@ -17,10 +17,10 @@ # limitations under the License. # +import glob import os -import sys import re -import glob +import sys def main(): diff --git a/dev/test-dependencies.sh b/dev/test-dependencies.sh index 68c61232ea2af..21997b4ac6251 100755 --- a/dev/test-dependencies.sh +++ b/dev/test-dependencies.sh @@ -31,7 +31,7 @@ export LC_ALL=C # NOTE: These should match those in the release publishing script, and be kept in sync with # dev/create-release/release-build.sh HADOOP_MODULE_PROFILES="-Phive-thriftserver -Pkubernetes -Pyarn -Phive \ - -Pspark-ganglia-lgpl -Pkinesis-asl -Phadoop-cloud -Pjvm-profiler" + -Pspark-ganglia-lgpl -Pkinesis-asl -Pcredential-aws -Phadoop-cloud -Pjvm-profiler" MVN="build/mvn" HADOOP_HIVE_PROFILES=( hadoop-3-hive-2.3 @@ -115,7 +115,16 @@ for HADOOP_HIVE_PROFILE in "${HADOOP_HIVE_PROFILES[@]}"; do classifier_end_index=index(jar_name, ".jar") - 1; classifier=substr(jar_name, classifier_start_index, classifier_end_index - classifier_start_index + 1); print artifact_id"/"version"/"classifier"/"jar_name - }' | sort | grep -v spark > dev/pr-deps/spark-deps-$HADOOP_HIVE_PROFILE + }' | sort | grep -v spark > dev/pr-deps/spark-deps-$HADOOP_HIVE_PROFILE.tmp + { + printf '%s\n' \ + '#' \ + '# Generated by `dev/test-dependencies.sh --replace-manifest`.' \ + '# Do not edit manually.' \ + '#' + cat dev/pr-deps/spark-deps-$HADOOP_HIVE_PROFILE.tmp + } > dev/pr-deps/spark-deps-$HADOOP_HIVE_PROFILE + rm dev/pr-deps/spark-deps-$HADOOP_HIVE_PROFILE.tmp done if [[ $@ == **replace-manifest** ]]; then diff --git a/docs/Gemfile.lock b/docs/Gemfile.lock index 30851e513a8cb..5fcc37ed97cfe 100644 --- a/docs/Gemfile.lock +++ b/docs/Gemfile.lock @@ -45,7 +45,7 @@ GEM sass-embedded (~> 1.75) jekyll-watch (2.2.1) listen (~> 3.0) - json (2.21.1) + json (2.21.2) kramdown (2.5.1) rexml (>= 3.3.9) kramdown-parser-gfm (1.1.0) diff --git a/docs/_plugins/build-error-docs.py b/docs/_plugins/build-error-docs.py index df6b9e3c05270..068ec0da50c96 100644 --- a/docs/_plugins/build-error-docs.py +++ b/docs/_plugins/build-error-docs.py @@ -89,7 +89,7 @@ def generate_doc_rows(condition_name, condition_details): sub_condition_rows.append( """ <tr id="{anchor}"> - <td></td> + <td>{sql_state}</td> <td class="error-sub-condition"> <span class="error-condition-name"> <code> @@ -103,6 +103,9 @@ def generate_doc_rows(condition_name, condition_details): """ .format( anchor=anchor_name(condition_name, sub_condition_name), + sql_state=( + condition_details["subClass"][sub_condition_name].get("sqlState", "") + ), # See comment above for explanation of `<wbr />`. sub_condition_name=sub_condition_name.replace("_", "<wbr />_"), message=condition_details["subClass"][sub_condition_name]["message"], diff --git a/docs/_plugins/build_api_docs.rb b/docs/_plugins/build_api_docs.rb index 429cef5aa026c..4b6251b804237 100644 --- a/docs/_plugins/build_api_docs.rb +++ b/docs/_plugins/build_api_docs.rb @@ -45,7 +45,7 @@ def build_spark_if_necessary print_header "Building Spark." cd(SPARK_PROJECT_ROOT) - command = "NO_PROVIDED_SPARK_JARS=0 build/sbt -Phive -Pkinesis-asl clean package" + command = "NO_PROVIDED_SPARK_JARS=0 build/sbt -Phive -Pkinesis-asl -Pcredential-aws clean package" puts "Running '#{command}'; this may take a few minutes..." system(command) || raise("Failed to build Spark") # SPARK-53327: Use the modified ResourceImpl.class in spark-catalyst which is compatible with Java 25 @@ -129,7 +129,7 @@ def build_spark_scala_and_java_docs_if_necessary return end - command = "build/sbt -Pkinesis-asl unidoc" + command = "build/sbt -Pkinesis-asl -Pcredential-aws unidoc" puts "Running '#{command}'..." # Two filter passes on the unidoc output, plus an additive fatal-error summary: diff --git a/docs/building-spark.md b/docs/building-spark.md index ad95b452ca31e..4fc0724e83ffc 100644 --- a/docs/building-spark.md +++ b/docs/building-spark.md @@ -271,6 +271,15 @@ or ./build/sbt -Pdocker-integration-tests docker-integration-tests/test +## Local network binding + +On a machine with multiple network interfaces (for example a VPN), Spark may bind to a +non-loopback address, causing local tests to fail with errors such as +`RemoteClassLoaderError` or Netty `Connection reset by peer`. Forcing the loopback +interface usually resolves this: + + export SPARK_LOCAL_IP=localhost + <!--- ## Change Scala Version diff --git a/docs/configuration.md b/docs/configuration.md index 38460273a823b..616367592b626 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1605,6 +1605,21 @@ Apart from these, the following properties are also available, and may be useful </td> <td>1.0.0</td> </tr> +<tr> + <td><code>spark.ui.holdEnabled</code></td> + <td>true</td> + <td> + Allows the whole application to be held and resumed from the web UI. Holding gracefully + decommissions all executors and stops requesting new ones. Cached blocks are not preserved + and are recomputed after resuming. This takes effect only when + <code>spark.decommission.enabled</code> is true, the shuffle data is kept outside the + executors (through either <code>spark.shuffle.service.enabled</code> or a + <code>ShuffleDataIO</code> with reliable storage), and the cluster manager can hold + executors: Standalone, YARN, and Kubernetes with + <code>spark.kubernetes.allocation.pods.allocator=direct</code>. + </td> + <td>4.4.0</td> +</tr> <tr> <td><code>spark.ui.threadDumpsEnabled</code></td> <td>true</td> diff --git a/docs/declarative-pipelines-programming-guide.md b/docs/declarative-pipelines-programming-guide.md index e1c2c078212ae..cb6dd6bcd0965 100644 --- a/docs/declarative-pipelines-programming-guide.md +++ b/docs/declarative-pipelines-programming-guide.md @@ -517,6 +517,370 @@ AS INSERT INTO customers_us SELECT * FROM STREAM(customers_us_east); ``` +## Change Data Capture (CDC) with Auto CDC + +Many source systems emit a stream of *change events* rather than a snapshot of the current data: each record describes an insert, update, or delete to a row, identified by a key. Applying these events correctly to a target table by hand is tricky. Matching events to existing rows, applying them in the right order, and handling out-of-order and duplicate events without corrupting the table is the hard part. + +**Auto CDC** does this automatically. Given a source of change events and a rule for identifying and ordering them, SDP maintains a target streaming table that always reflects the latest state for each key. + +### Keys and Sequencing + +Auto CDC needs two things to make sense of a change feed: + +- The **keys** are the columns that identify a row across events. Events sharing a key describe the same logical row over time. In a customer feed, that is usually the customer id. +- The **sequencing expression** says what order the events for a key happened in. Its value for an event is that event's **sequence value**, and Auto CDC treats the highest sequence value it has seen for a key as the most recent state. Change feeds normally carry something suitable already: a monotonically increasing version or commit number, or a commit timestamp. + +Sequencing matters because a change feed does not have to arrive in order. It is what lets Auto CDC apply events by their intended order rather than the order they happen to land in. + +### What Auto CDC Does + +Given a stream of change events, Auto CDC keeps the target table in sync with the source: + +- **Inserts and updates** - For each key, the event with the highest sequence value wins. If no row exists for the key, it's inserted; if one exists, it's overwritten with the latest values. +- **Deletes** - Events that match a supplied delete condition remove the corresponding row from the target. +- **Out-of-order and duplicate events** - Events don't have to arrive in order, and the source does not need to be de-duplicated. An event whose sequence value is older than the state already applied for its key is discarded, and a re-delivered event converges to the same result. + +This section covers **Slowly Changing Dimensions (SCD) Type 1**, where the target keeps only the current version of each row, with no history of prior values. SCD Type 2 is also supported; [SPARK-58570](https://issues.apache.org/jira/browse/SPARK-58570) tracks documenting it. + +For example, take these change events. They are **not** in `version` order - the last event for `id 1` is a stale re-delivery of version 1, arriving after version 2: + +| id | name | version | op | +|----|--------|---------|--------| +| 1 | alice | 1 | UPSERT | +| 2 | bob | 1 | UPSERT | +| 1 | alicia | 2 | UPSERT | +| 2 | bob | 2 | DELETE | +| 3 | carol | 1 | UPSERT | +| 1 | alice | 1 | UPSERT | + +Auto CDC keyed on `id` and sequenced by `version` produces this target table. It drops the `op` column, and omits the internal metadata column Auto CDC appends (see [Considerations](#auto-cdc-considerations)): + +| id | name | version | +|----|--------|---------| +| 1 | alicia | 2 | +| 3 | carol | 1 | + +By key: `id 1` was inserted as `alice`, then updated to `alicia` at version 2, and the stale version 1 event arriving last is discarded because a higher sequence value has already been applied; `id 2` was inserted and then deleted at version 2, so it is gone; `id 3` was inserted and never changed. + +### Requirements + +- The **target must be a streaming table**, not a materialized view or an external table. Define it with `create_streaming_table` (Python) or `CREATE STREAMING TABLE` (SQL), or declare it together with the flow using the combined SQL form below. Definition order within the source file does not matter, since the graph is resolved after everything is registered. +- The **target's format must support row-level operations.** Auto CDC maintains the target with MERGE, so the table must be backed by a connector implementing the DSv2 `SupportsRowLevelOperations` interface. A target that does not fails at startup with `AUTOCDC_TARGET_DOES_NOT_SUPPORT_MERGE`. Spark's built-in file formats, including Parquet, do **not** qualify; see [Choosing a Target Format](#choosing-a-target-format). +- The **target must have exactly one input flow.** An Auto CDC flow cannot share a target with other flows; a target fed by more than one flow, one of which is Auto CDC, fails with `AUTOCDC_MULTIPLE_FLOWS_TO_TARGET`. This rules out the fan-in pattern shown in "Using Multiple Flows to Write to a Single Target" above. +- The **source must be a streaming source** (read with `spark.readStream` in Python or `STREAM(...)` in SQL). CDC is an incremental operation over newly arriving change events. +- The **catalog's metadata must survive between runs**, since each incremental run re-resolves the tables an earlier run created. This is a property of the connector, not of setting `spark.sql.catalog.*`. Spark's default session catalog keeps metadata only for the life of the session, so a second `spark-pipelines run` against it fails with `LOCATION_ALREADY_EXISTS`; use a catalog whose metastore persists. +- The flow must specify **keys** (one or more columns that identify a row) and a **sequencing expression** (used to order events per key). + +The sequencing expression may be any SQL expression over the source columns, not just a bare column reference; sequencing by a struct such as `(commit_ts, seq_no)`, or by a cast, is fine. It must satisfy three constraints: + +- **Its type must be orderable.** A non-orderable type fails with `AUTOCDC_MICROBATCH_VALIDATION.NON_ORDERABLE_SEQUENCE`. +- **It must never be null.** A microbatch containing a null sequence value fails with `AUTOCDC_MICROBATCH_VALIDATION.NULL_SEQUENCE` rather than guessing an order. +- **Its result type must stay the same across runs.** The expression may change between incremental runs, but not its type, or recorded values would stop being comparable; that fails with `AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT` and needs a full refresh. + +Within a microbatch, ties on the sequence value for a key are broken arbitrarily. For deterministic results, use an expression that is unique per event within a key. + +### Defining an Auto CDC Flow in Python + +Use `create_auto_cdc_flow` to write change events into a target streaming table, which is defined with `create_streaming_table`. + +```python +from pyspark import pipelines as dp +from pyspark.sql import DataFrame + +# The source of change events: a streaming read of the CDC feed. +@dp.table +def cdc_events() -> DataFrame: + return spark.readStream.table("cdc_source") + +# The target that Auto CDC keeps in sync. It must be a streaming table, in a +# catalog whose format supports row-level operations (see "Choosing a Target +# Format" below). +dp.create_streaming_table("customers") + +# The Auto CDC flow that applies the change events to the target. +dp.create_auto_cdc_flow( + target="customers", + source="cdc_events", + keys=["id"], + sequence_by="version", + apply_as_deletes="op = 'DELETE'", + except_column_list=["op"], + stored_as_scd_type=1, +) +``` + +`create_auto_cdc_flow` accepts the following arguments: + +| Parameter | Required | Description | +|-----------|----------|-------------| +| `target` | Yes | Name of the target streaming table that receives the changes. It must be defined in the pipeline. | +| `source` | Yes | Name of the CDC source dataset to stream change events from. | +| `keys` | Yes | The column or columns that uniquely identify a row. A list of column names (strings) or `Column` objects, given as unqualified identifiers: for example `"id"` or `col("id")`, but not `"source.id"`. | +| `sequence_by` | Yes | An expression used to order change events for each key. The highest value wins. A SQL expression string or a `Column`. | +| `apply_as_deletes` | No | A boolean expression identifying events that represent deletes. Matching rows are removed from the target. A SQL expression string or a `Column`. | +| `column_list` | No | The columns to include in the target. Mutually exclusive with `except_column_list`. | +| `except_column_list` | No | The columns to exclude from the target; all other columns are included. Mutually exclusive with `column_list`. Commonly used to drop operation/metadata columns such as `op`. | +| `stored_as_scd_type` | No | The SCD type of the target. Pass `1` (or `"1"`) for the Type 1 behavior described here. Type 2 is also accepted; see [SPARK-58570](https://issues.apache.org/jira/browse/SPARK-58570). | +| `name` | No | The name of the flow. Defaults to the target table name. | +| `spark_conf` | No | Spark configuration to apply while the flow runs. Overrides configuration set on the destination, the pipeline, or the cluster. | + +With neither `column_list` nor `except_column_list`, every source column is written to the target - including the operation column and any other change-feed bookkeeping. That is rarely right for a CDC feed, so the examples here drop `op` with `except_column_list=["op"]`. See [Selecting Which Columns Land in the Target](#selecting-which-columns-land-in-the-target). + +`keys`, `column_list`, and `except_column_list` must be given as unqualified column identifiers: `"id"` or `col("id")`, but not `"cdc_events.id"` or `col("cdc_events.id")`. (`sequence_by` is not restricted this way; it may be any expression over the source columns.) + +### Defining an Auto CDC Flow in SQL + +SQL provides two forms. The first attaches an Auto CDC flow to an already-declared streaming table: + +```sql +CREATE STREAMING TABLE customers; + +CREATE FLOW customers_cdc AS AUTO CDC INTO customers +FROM STREAM(cdc_events) +KEYS (id) +APPLY AS DELETE WHEN op = 'DELETE' +SEQUENCE BY version +COLUMNS * EXCEPT (op); +``` + +The second declares the streaming table and its Auto CDC flow together: + +```sql +CREATE STREAMING TABLE customers +FLOW AUTO CDC +FROM STREAM(cdc_events) +KEYS (id) +APPLY AS DELETE WHEN op = 'DELETE' +SEQUENCE BY version +COLUMNS * EXCEPT (op); +``` + +`FROM STREAM(source)` and `KEYS (col, ...)` come first, in that order. The remaining clauses may appear in any order after them: + +- `FROM STREAM(source)` - the streaming CDC source. **Required, first.** +- `KEYS (col, ...)` - the key columns that identify a row. **Required, second.** +- `SEQUENCE BY expr` - the expression that orders events per key. **Required.** +- `APPLY AS DELETE WHEN condition` - marks events that represent deletes. Optional. +- `COLUMNS (col, ...)` or `COLUMNS * EXCEPT (col, ...)` - selects or excludes columns. Optional; if omitted, all source columns are written. +- `STORED AS SCD TYPE 1` - selects the SCD type. Optional; defaults to Type 1 when omitted, which is the behavior described here. + +`CREATE FLOW ... AS AUTO CDC INTO` also accepts an optional `COMMENT`, and `CREATE STREAMING TABLE ... FLOW AUTO CDC` accepts `IF NOT EXISTS`. + +### Choosing a Target Format + +Auto CDC applies each microbatch to the target with a MERGE, so the target has to be a table that supports row-level updates and deletes. Concretely, its connector must implement the DSv2 `SupportsRowLevelOperations` interface. + +Spark's built-in file-based formats do not. Pointing an Auto CDC flow at a plain Parquet target - which is what a `create_streaming_table("customers")` with no `format` produces - fails when the flow starts: + +``` +[AUTOCDC_TARGET_DOES_NOT_SUPPORT_MERGE] Cannot start AutoCDC flow: the target table +`spark_catalog`.`default`.`customers` (format: parquet) does not support row-level +operations. AutoCDC requires a target backed by a connector that supports MERGE. +``` + +Auto CDC needs a table provider that implements row-level operations, configured as a catalog in the pipeline. Lakehouse connectors such as Apache Iceberg are the usual choice; consult the connector's documentation for whether it implements the DSv2 row-level operation interfaces and how to register its catalog. Supporting the `MERGE INTO` SQL statement is not on its own sufficient: a connector can implement MERGE through its own planner extension without implementing the DSv2 interface Auto CDC requires. The catalog must also persist metadata across runs, per [Requirements](#requirements). + +The examples below use a catalog named `lakehouse` to stand in for such a connector. Add the catalog to the generated `spark-pipeline.yml` - substituting the catalog name and connector class - without disturbing the `name`, `storage`, and `libraries` values `spark-pipelines init` already wrote: + +```yaml +catalog: lakehouse +database: cdc_demo +configuration: + spark.sql.catalog.lakehouse: <your connector's catalog class> +``` + +### End-to-End Example + +This example builds a small pipeline that ingests customer change events and maintains a `customers` table holding the latest state of each customer. Running it in three passes shows updates and deletes applied incrementally, out-of-order events reconciled within a batch, and a stale re-delivery in a later run discarded against committed state. + +It reads the change feed from a directory of JSON files, so appending a batch and re-running shows the incremental behavior. The target uses the `lakehouse` catalog from [Choosing a Target Format](#choosing-a-target-format); substitute a row-level-operation-capable catalog. + +Create a pipeline project: + +```bash +spark-pipelines init --name cdc_demo +cd cdc_demo +``` + +`spark-pipelines init` seeds `transformations/` with `example_python_materialized_view.py` and `example_sql_materialized_view.sql`. Delete them so only the CDC datasets run, then edit `spark-pipeline.yml` to add the catalog as shown in [Choosing a Target Format](#choosing-a-target-format), and put the following in `transformations/customers_cdc.py`: + +```python +from pyspark import pipelines as dp +from pyspark.sql import DataFrame +from pyspark.sql.types import ( + IntegerType, LongType, StringType, StructField, StructType) + +# An explicit schema keeps the streaming JSON read from having to infer one. +SCHEMA = StructType([ + StructField("id", IntegerType()), + StructField("name", StringType()), + StructField("version", LongType()), + StructField("op", StringType()), +]) + +# Ingest the raw change events. In a real pipeline this would read from Kafka, +# cloud storage, or a database CDC feed; here it tails a directory of JSON files. +@dp.table +def cdc_events() -> DataFrame: + return spark.readStream.schema(SCHEMA).json("file:///tmp/cdc_demo/events") + +# Declare the target streaming table that Auto CDC maintains. +dp.create_streaming_table("customers") + +# Apply the change events to the target. +dp.create_auto_cdc_flow( + target="customers", + source="cdc_events", + keys=["id"], + sequence_by="version", + apply_as_deletes="op = 'DELETE'", + except_column_list=["op"], + stored_as_scd_type=1, +) +``` + +Write the first batch of change events, inserting two customers: + +```bash +mkdir -p /tmp/cdc_demo/events +cat > /tmp/cdc_demo/events/batch1.json <<'EOF' +{"id": 1, "name": "alice", "version": 1, "op": "UPSERT"} +{"id": 2, "name": "bob", "version": 1, "op": "UPSERT"} +EOF +``` + +Run the pipeline: + +```bash +spark-pipelines run +``` + +`customers` now holds both rows, with the `op` column excluded: + +| id | name | version | +|----|-------|---------| +| 1 | alice | 1 | +| 2 | bob | 1 | + +Now add a second batch. It updates `id 1`, deletes `id 2`, and inserts `id 3`. The two events for `id 1` also arrive out of order within the batch - the newer value first, its stale predecessor after: + +```bash +cat > /tmp/cdc_demo/events/batch2.json <<'EOF' +{"id": 1, "name": "alicia", "version": 2, "op": "UPSERT"} +{"id": 1, "name": "alice", "version": 1, "op": "UPSERT"} +{"id": 2, "name": "bob", "version": 2, "op": "DELETE"} +{"id": 3, "name": "carol", "version": 1, "op": "UPSERT"} +EOF +``` + +Run the pipeline again. Because `customers` is a streaming table, this run processes only the new file: + +```bash +spark-pipelines run +``` + +| id | name | version | +|----|--------|---------| +| 1 | alicia | 2 | +| 3 | carol | 1 | + +`alice` became `alicia`, `bob` is gone, and `carol` was inserted. The two `id 1` events landed in the same batch out of order; Auto CDC collapsed them to the highest sequence value per key before applying, so the row is `alicia`, not the earlier `alice`. + +Ordering also holds *across* runs, against state already committed. Add a third batch that re-delivers the original `id 1` event at version 1, as a re-sending source might: + +```bash +cat > /tmp/cdc_demo/events/batch3.json <<'EOF' +{"id": 1, "name": "alice", "version": 1, "op": "UPSERT"} +EOF +``` + +Run once more: + +```bash +spark-pipelines run +``` + +| id | name | version | +|----|--------|---------| +| 1 | alicia | 2 | +| 3 | carol | 1 | + +The table is unchanged. This time the stale event is alone in its batch, so there is no newer event to supersede it within the batch; instead Auto CDC compares its version against the state already recorded for `id 1` and discards it, because version 1 is below the committed version 2. Applied in arrival order, this event would have wrongly reverted `id 1` to `alice`. + +### How-Tos + +#### Handling Deletes + +Change feeds usually mark deletes with an operation column or a tombstone flag rather than removing the row. Give Auto CDC a boolean expression that identifies delete events with `apply_as_deletes` (Python) or `APPLY AS DELETE WHEN` (SQL): + +```python +dp.create_auto_cdc_flow( + target="customers", + source="cdc_events", + keys=["id"], + sequence_by="version", + apply_as_deletes="op = 'DELETE'", +) +``` + +When an event matches the delete condition, the row for its key is removed from the target. If no delete condition is supplied, every event is treated as an insert or update. + +#### Selecting Which Columns Land in the Target + +CDC feeds often carry metadata columns (the operation type, a timestamp, source offsets) that do not belong in the target table. Use `except_column_list` / `COLUMNS * EXCEPT` to drop them, or `column_list` / `COLUMNS` to name exactly the columns to keep. The two options are mutually exclusive. + +```python +# Keep everything except the operation column. +dp.create_auto_cdc_flow( + target="customers", + source="cdc_events", + keys=["id"], + sequence_by="version", + except_column_list=["op"], +) + +# Or keep only an explicit set of columns. +dp.create_auto_cdc_flow( + target="customers", + source="cdc_events", + keys=["id"], + sequence_by="version", + column_list=["id", "name"], +) +``` + +#### Handling Out-of-Order and Duplicate Events + +The source need not be sorted or de-duplicated. Auto CDC orders events for a key by their sequence value - both within a microbatch and against state from earlier runs - so a late or repeated event lands correctly, as the end-to-end example above demonstrates. Just choose a `sequence_by` expression that strictly orders changes for a key, such as a monotonically increasing version number or a commit timestamp. + +#### Using a Composite Key + +Pass multiple columns to `keys` when a single column doesn't uniquely identify a row: + +```python +dp.create_auto_cdc_flow( + target="orders", + source="order_events", + keys=["region", "order_id"], + sequence_by="event_ts", +) +``` + +#### Changing the Key Set + +The set and types of `keys` are part of the flow's persisted state. Changing keys across incremental runs - renaming, swapping, adding, removing, or changing the type of a key column - is not supported: the run fails with `AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT`. To change the key set, [fully refresh](#spark-pipelines-run) the target table so it is recomputed from scratch: + +```bash +spark-pipelines run --full-refresh customers +``` + +### Auto CDC Considerations + +Beyond the [Requirements](#requirements) above, two engine details are worth knowing: + +- **Reserved column names** - Auto CDC owns the `__spark_autocdc_` column-name prefix; it projects an internal `__spark_autocdc_metadata` column onto the target, and SCD Type 2 adds more. The source cannot contain any column whose name begins with `__spark_autocdc_`; a collision fails when the flow is constructed. Treat these as engine details rather than columns to query. +- **Declaring a schema explicitly** - When the target is declared separately - `create_streaming_table(schema=...)` in Python, or a standalone `CREATE STREAMING TABLE` in SQL - the declared schema must currently include the reserved metadata column, because it has to match the flow's output schema exactly. The combined `CREATE STREAMING TABLE ... FLOW AUTO CDC` form takes no column list at all; it infers the schema from the flow. Either way, letting Auto CDC derive the schema avoids the issue. [SPARK-58118](https://issues.apache.org/jira/browse/SPARK-58118) tracks relaxing it. + ## Writing Data to External Targets with Sinks Sinks in SDP provide a way to write transformed data to external destinations beyond the default streaming tables and materialized views. Sinks are particularly useful for operational use cases that require low-latency data processing, reverse ETL operations, or writing to external systems. diff --git a/docs/graphx-programming-guide.md b/docs/graphx-programming-guide.md index 4791e215f458c..5d599bf39da0a 100644 --- a/docs/graphx-programming-guide.md +++ b/docs/graphx-programming-guide.md @@ -189,7 +189,7 @@ val users: RDD[(VertexId, (String, String))] = val relationships: RDD[Edge[String]] = sc.parallelize(Seq(Edge(3L, 7L, "collab"), Edge(5L, 3L, "advisor"), Edge(2L, 5L, "colleague"), Edge(5L, 7L, "pi"))) -// Define a default user in case there are relationship with missing user +// Define a default user in case there is a relationship with a missing user val defaultUser = ("John Doe", "Missing") // Build the initial Graph val graph = Graph(users, relationships, defaultUser) @@ -413,7 +413,7 @@ implemented efficiently without data movement or duplication. The [`subgraph`][Graph.subgraph] operator takes vertex and edge predicates and returns the graph containing only the vertices that satisfy the vertex predicate (evaluate to true) and edges that satisfy the edge predicate *and connect vertices that satisfy the vertex predicate*. The `subgraph` -operator can be used in number of situations to restrict the graph to the vertices and edges of +operator can be used in a number of situations to restrict the graph to the vertices and edges of interest or eliminate broken links. For example in the following code we remove broken links: @@ -428,7 +428,7 @@ val relationships: RDD[Edge[String]] = sc.parallelize(Seq(Edge(3L, 7L, "collab"), Edge(5L, 3L, "advisor"), Edge(2L, 5L, "colleague"), Edge(5L, 7L, "pi"), Edge(4L, 0L, "student"), Edge(5L, 0L, "colleague"))) -// Define a default user in case there are relationship with missing user +// Define a default user in case there is a relationship with a missing user val defaultUser = ("John Doe", "Missing") // Build the initial Graph val graph = Graph(users, relationships, defaultUser) diff --git a/docs/ml-guide.md b/docs/ml-guide.md index 64c469d906e30..34c7e7d72491d 100644 --- a/docs/ml-guide.md +++ b/docs/ml-guide.md @@ -40,7 +40,7 @@ The primary Machine Learning API for Spark is now the [DataFrame](sql-programmin * MLlib will still support the RDD-based API in `spark.mllib` with bug fixes. * MLlib will not add new features to the RDD-based API. -* In the Spark 2.x releases, MLlib will add features to the DataFrames-based API to reach feature parity with the RDD-based API. +* MLlib added features to the DataFrame-based API in the Spark 2.x releases, reaching feature parity with the RDD-based API. *Why is MLlib switching to the DataFrame-based API?* @@ -74,32 +74,6 @@ To use MLlib in Python, you will need [NumPy](http://www.numpy.org) version 1.23 [^1]: To learn more about the benefits and background of system optimised natives, you may wish to watch Sam Halliday's ScalaX talk on [High Performance Linear Algebra in Scala](http://fommil.github.io/scalax14/). -# Highlights in 3.0 - -The list below highlights some of the new features and enhancements added to MLlib in the `3.0` -release of Spark: - -* Multiple columns support was added to `Binarizer` ([SPARK-23578](https://issues.apache.org/jira/browse/SPARK-23578)), `StringIndexer` ([SPARK-11215](https://issues.apache.org/jira/browse/SPARK-11215)), `StopWordsRemover` ([SPARK-29808](https://issues.apache.org/jira/browse/SPARK-29808)) and PySpark `QuantileDiscretizer` ([SPARK-22796](https://issues.apache.org/jira/browse/SPARK-22796)). -* Tree-Based Feature Transformation was added -([SPARK-13677](https://issues.apache.org/jira/browse/SPARK-13677)). -* Two new evaluators `MultilabelClassificationEvaluator` ([SPARK-16692](https://issues.apache.org/jira/browse/SPARK-16692)) and `RankingEvaluator` ([SPARK-28045](https://issues.apache.org/jira/browse/SPARK-28045)) were added. -* Sample weights support was added in `DecisionTreeClassifier/Regressor` ([SPARK-19591](https://issues.apache.org/jira/browse/SPARK-19591)), `RandomForestClassifier/Regressor` ([SPARK-9478](https://issues.apache.org/jira/browse/SPARK-9478)), `GBTClassifier/Regressor` ([SPARK-9612](https://issues.apache.org/jira/browse/SPARK-9612)), `MulticlassClassificationEvaluator` ([SPARK-24101](https://issues.apache.org/jira/browse/SPARK-24101)), `RegressionEvaluator` ([SPARK-24102](https://issues.apache.org/jira/browse/SPARK-24102)), `BinaryClassificationEvaluator` ([SPARK-24103](https://issues.apache.org/jira/browse/SPARK-24103)), `BisectingKMeans` ([SPARK-30351](https://issues.apache.org/jira/browse/SPARK-30351)), `KMeans` ([SPARK-29967](https://issues.apache.org/jira/browse/SPARK-29967)) and `GaussianMixture` ([SPARK-30102](https://issues.apache.org/jira/browse/SPARK-30102)). -* R API for `PowerIterationClustering` was added -([SPARK-19827](https://issues.apache.org/jira/browse/SPARK-19827)). -* Added Spark ML listener for tracking ML pipeline status -([SPARK-23674](https://issues.apache.org/jira/browse/SPARK-23674)). -* Fit with validation set was added to Gradient Boosted Trees in Python -([SPARK-24333](https://issues.apache.org/jira/browse/SPARK-24333)). -* [`RobustScaler`](ml-features.html#robustscaler) transformer was added -([SPARK-28399](https://issues.apache.org/jira/browse/SPARK-28399)). -* [`Factorization Machines`](ml-classification-regression.html#factorization-machines) classifier and regressor were added -([SPARK-29224](https://issues.apache.org/jira/browse/SPARK-29224)). -* Gaussian Naive Bayes Classifier ([SPARK-16872](https://issues.apache.org/jira/browse/SPARK-16872)) and Complement Naive Bayes Classifier ([SPARK-29942](https://issues.apache.org/jira/browse/SPARK-29942)) were added. -* ML function parity between Scala and Python -([SPARK-28958](https://issues.apache.org/jira/browse/SPARK-28958)). -* `predictRaw` is made public in all the Classification models. `predictProbability` is made public in all the Classification models except `LinearSVCModel` -([SPARK-30358](https://issues.apache.org/jira/browse/SPARK-30358)). - # Migration Guide The migration guide is now archived [on this page](ml-migration-guide.html). diff --git a/docs/mllib-clustering.md b/docs/mllib-clustering.md index fde354457f23d..da7839c7b6d5a 100644 --- a/docs/mllib-clustering.md +++ b/docs/mllib-clustering.md @@ -384,7 +384,7 @@ Refer to the [`LDA` Java docs](api/java/org/apache/spark/mllib/clustering/LDA.ht Bisecting K-means can often be much faster than regular K-means, but it will generally produce a different clustering. Bisecting k-means is a kind of [hierarchical clustering](https://en.wikipedia.org/wiki/Hierarchical_clustering). -Hierarchical clustering is one of the most commonly used method of cluster analysis which seeks to build a hierarchy of clusters. +Hierarchical clustering is one of the most commonly used methods of cluster analysis which seeks to build a hierarchy of clusters. Strategies for hierarchical clustering generally fall into two types: - Agglomerative: This is a "bottom up" approach: each observation starts in its own cluster, and pairs of clusters are merged as one moves up the hierarchy. diff --git a/docs/monitoring.md b/docs/monitoring.md index 64552ba28512a..7b8d61c525f34 100644 --- a/docs/monitoring.md +++ b/docs/monitoring.md @@ -458,6 +458,14 @@ Security options for the Spark History Server are covered more detail in the </td> <td>4.1.0</td> </tr> + <tr> + <td>spark.history.fs.eventLog.onDemandLoadEnabled</td> + <td>true</td> + <td> + Whether to look up single event log locations on demand manner before listing files. + </td> + <td>4.3.0</td> + </tr> <tr> <td>spark.history.store.hybridStore.enabled</td> <td>false</td> @@ -648,6 +656,17 @@ can be identified by their `[attempt-id]`. In the API listed below, when running <td><code>/applications/[app-id]/allexecutors</code></td> <td>A list of all(active and dead) executors for the given application.</td> </tr> + <tr> + <td><code>/applications/[app-id]/holdstatus</code></td> + <td> + Whether the given application is held, as <code>supported</code> (whether the deployment + allows holding), <code>held</code>, and <code>draining</code> (the number of executors + that have not exited yet). An application is held and resumed through the + <code>/jobs/hold/</code> and <code>/jobs/resume/</code> POST endpoints of its web UI, + which require modify permissions, while reading this status only requires view + permissions. Not available via the history server. + </td> + </tr> <tr> <td><code>/applications/[app-id]/storage/rdd</code></td> <td>A list of stored RDDs for the given application.</td> @@ -1328,6 +1347,11 @@ This is the component with the largest amount of instrumented metrics - queue.eventLog.numDroppedEvents.count - queue.eventLog.size - queue.executorManagement.listenerProcessingTime (timer) + - queue.executorManagement.numDroppedEvents.count + - queue.executorManagement.size + - queue.shared.listenerProcessingTime (timer) + - queue.shared.numDroppedEvents.count + - queue.shared.size - namespace=appStatus (all metrics of type=counter) - **note:** Introduced in Spark 3.0. Conditional to a configuration parameter: diff --git a/docs/quick-start.md b/docs/quick-start.md index f9fd4921a9baa..4e24b94695b98 100644 --- a/docs/quick-start.md +++ b/docs/quick-start.md @@ -51,13 +51,13 @@ Or if PySpark is installed with pip in your current environment: pyspark -Spark's primary abstracted is called a **Dataset**. A Dataset is a structured set of information. You can create datasets from Hadoop InputFormats (such as HDFS +Spark's primary abstraction is called a **Dataset**. A Dataset is a structured set of information. You can create datasets from Hadoop InputFormats (such as HDFS files) or by transforming other Datasets. Datasets behave differently in some languages. Because Python allows for dynamic typing, Datasets in Python are all `Dataset[Row]` on an implementation level. This leads to another key Spark concept: a `DataFrame`, or a Dataset with named columns. If you're familiar with DataFrames from pandas or R, you'll be familiar with how DataFrames work in Spark. In other languages, like Java, the difference between a Dataset and DataFrame is larger, but for now let's proceed with Python. -Let's make a new DataFrame using the `README.md` file in the Spark souce directory via the command line: +Let's make a new DataFrame using the `README.md` file in the Spark source directory via the command line: {% highlight python %} >>> textFile = spark.read.text("README.md") @@ -92,10 +92,10 @@ You can also chain together transformations and actions: ./bin/spark-shell -Spark's primary abstracted is called a **Dataset**. A Dataset is a structured set of information. You can create datasets from Hadoop InputFormats (such as HDFS +Spark's primary abstraction is called a **Dataset**. A Dataset is a structured set of information. You can create datasets from Hadoop InputFormats (such as HDFS files) or by transforming other Datasets. -Let's make a new DataFrame using the `README.md` file in the Spark souce directory via the command line: +Let's make a new DataFrame using the `README.md` file in the Spark source directory via the command line: {% highlight scala %} scala> val textFile = spark.read.textFile("README.md") diff --git a/docs/rdd-programming-guide.md b/docs/rdd-programming-guide.md index deff45ddc852f..eb03d418d6b4c 100644 --- a/docs/rdd-programming-guide.md +++ b/docs/rdd-programming-guide.md @@ -70,7 +70,7 @@ PySpark requires the same minor version of Python in both driver and workers. It you can specify which version of Python you want to use by `PYSPARK_PYTHON`, for example: {% highlight bash %} -$ PYSPARK_PYTHON=python3.8 bin/pyspark +$ PYSPARK_PYTHON=python3.11 bin/pyspark {% endhighlight %} </div> diff --git a/docs/running-on-kubernetes.md b/docs/running-on-kubernetes.md index e9fc26fd00143..42e2601eab41d 100644 --- a/docs/running-on-kubernetes.md +++ b/docs/running-on-kubernetes.md @@ -905,8 +905,13 @@ See the [configuration page](configuration.html) for information on Spark config <td><code>default</code></td> <td> Service account that is used when running the driver pod. The driver pod uses this service account when requesting - executor pods from the API server. Note that this cannot be specified alongside a CA cert file, client key file, - client cert file, and/or OAuth token. In client mode, use <code>spark.kubernetes.authenticate.serviceAccountName</code> instead. + executor pods from the API server. Note that this cannot be specified alongside a submitted CA cert file, client key + file, client cert file, and/or OAuth token: Spark mounts those as a secret and they take precedence, so the driver + pod is left with the service account its spec already names, or the namespace's default. Spark logs a warning when + the account is dropped. To have Spark apply this configuration anyway, put the + credentials inside the driver pod and point the + <code>spark.kubernetes.authenticate.driver.mounted.*</code> configurations at them instead, which does not mount a + secret. In client mode, use <code>spark.kubernetes.authenticate.serviceAccountName</code> instead. </td> <td>2.3.0</td> </tr> @@ -1973,7 +1978,8 @@ See the below table for the full list of pod specifications that will be overwri <td>Value of <code>spark.kubernetes.authenticate.driver.serviceAccountName</code></td> <td> Spark will override <code>serviceAccount</code> with the value of the spark configuration for only - driver pods, and only if the spark configuration is specified. Executor pods will remain unaffected. + driver pods, and only if the spark configuration is specified and no driver credentials are + submitted for Spark to mount as a secret. Executor pods will remain unaffected. </td> </tr> <tr> @@ -1981,7 +1987,8 @@ See the below table for the full list of pod specifications that will be overwri <td>Value of <code>spark.kubernetes.authenticate.driver.serviceAccountName</code></td> <td> Spark will override <code>serviceAccountName</code> with the value of the spark configuration for only - driver pods, and only if the spark configuration is specified. Executor pods will remain unaffected. + driver pods, and only if the spark configuration is specified and no driver credentials are + submitted for Spark to mount as a secret. Executor pods will remain unaffected. </td> </tr> <tr> @@ -2188,10 +2195,10 @@ Install Apache YuniKorn: ```bash helm repo add yunikorn https://apache.github.io/yunikorn-release helm repo update -helm install yunikorn yunikorn/yunikorn --namespace yunikorn --version 1.8.0 --create-namespace --set embedAdmissionController=false +helm install yunikorn yunikorn/yunikorn --namespace yunikorn --version 1.9.0 --create-namespace --set embedAdmissionController=false ``` -The above steps will install YuniKorn v1.8.0 on an existing Kubernetes cluster. +The above steps will install YuniKorn v1.9.0 on an existing Kubernetes cluster. ##### Get started diff --git a/docs/running-on-yarn.md b/docs/running-on-yarn.md index f7f356822f63f..bd82a84453888 100644 --- a/docs/running-on-yarn.md +++ b/docs/running-on-yarn.md @@ -773,7 +773,7 @@ Stage level scheduling is supported on YARN: - When dynamic allocation is disabled: It allows users to specify different task resource requirements at the stage level and will use the same executors requested at startup. - When dynamic allocation is enabled: It allows users to specify task and executor resource requirements at the stage level and will request the extra executors. -One thing to note that is YARN specific is that each ResourceProfile requires a different container priority on YARN. The mapping is simply the ResourceProfile id becomes the priority, on YARN lower numbers are higher priority. This means that profiles created earlier will have a higher priority in YARN. Normally this won't matter as Spark finishes one stage before starting another one, the only case this might have an affect is in a job server type scenario, so its something to keep in mind. +One thing to note that is YARN specific is that each ResourceProfile requires a different container priority on YARN. The mapping is simply the ResourceProfile id becomes the priority, on YARN lower numbers are higher priority. This means that profiles created earlier will have a higher priority in YARN. Normally this won't matter as Spark finishes one stage before starting another one, the only case this might have an effect is in a job server type scenario, so it's something to keep in mind. Note there is a difference in the way custom resources are handled between the base default profile and custom ResourceProfiles. To allow for the user to request YARN containers with extra resources without Spark scheduling on them, the user can specify resources via the <code>spark.yarn.executor.resource.</code> config. Those configs are only used in the base default profile though and do not get propagated into any other custom ResourceProfiles. This is because there would be no way to remove them if you wanted a stage to not have them. This results in your default profile getting custom resources defined in <code>spark.yarn.executor.resource.</code> plus spark defined resources of GPU or FPGA. Spark converts GPU and FPGA resources into the YARN built in types <code>yarn.io/gpu</code>) and <code>yarn.io/fpga</code>, but does not know the mapping of any other resources. Any other Spark custom resources are not propagated to YARN for the default profile. So if you want Spark to schedule based off a custom resource and have it requested from YARN, you must specify it in both YARN (<code>spark.yarn.{driver/executor}.resource.</code>) and Spark (<code>spark.{driver/executor}.resource.</code>) configs. Leave the Spark config off if you only want YARN containers with the extra resources but Spark not to schedule using them. Now for custom ResourceProfiles, it doesn't currently have a way to only specify YARN resources without Spark scheduling off of them. This means for custom ResourceProfiles we propagate all the resources defined in the ResourceProfile to YARN. We still convert GPU and FPGA to the YARN build in types as well. This requires that the name of any custom resources you specify match what they are defined as in YARN. # Important notes diff --git a/docs/security.md b/docs/security.md index d802502560664..15b8e498617f9 100644 --- a/docs/security.md +++ b/docs/security.md @@ -98,7 +98,7 @@ Kubernetes admin to ensure that Spark authentication is secure. <td><code>spark.authenticate.secret</code></td> <td>None</td> <td> - The secret key used authentication. See above for when this configuration should be set. + The secret key used for authentication. See above for when this configuration should be set. </td> <td>1.0.0</td> </tr> diff --git a/docs/spark-connect-overview.md b/docs/spark-connect-overview.md index 4d455a8e4680b..de698a357600c 100644 --- a/docs/spark-connect-overview.md +++ b/docs/spark-connect-overview.md @@ -111,17 +111,19 @@ latest release in the release drop down at the top of the page. Then choose you Now extract the Spark package you just downloaded on your computer, for example: -{% highlight bash %} +```bash tar -xvf spark-{{site.SPARK_VERSION_SHORT}}-bin-hadoop3.tgz -{% endhighlight %} +``` In a terminal window, go to the `spark` folder in the location where you extracted Spark before and run the `start-connect-server.sh` script to start Spark server with Spark Connect, like in this example: -{% highlight bash %} +```bash ./sbin/start-connect-server.sh -{% endhighlight %} +``` + +Alternatively, `./bin/spark-connect-shell` starts an interactive Scala shell with the Connect server hosted inside the shell process itself. Make sure to use the same version of the package as the Spark version you downloaded previously. In this example, Spark {{site.SPARK_VERSION_SHORT}} with Scala 2.13. @@ -277,6 +279,101 @@ The connection may also be programmatically created using _SparkSession#builder_ </div> </div> +## Faster local iteration with a persistent Connect server + +When you develop or test locally with + +```python +from pyspark.sql import SparkSession +spark = SparkSession.builder.remote("local[*]").getOrCreate() +``` + +PySpark boots a fresh in-process Spark Connect server that lives only as long as that Python +process. Every `python script.py` invocation (or every new test worker process) therefore re-pays +the one-time startup cost -- JVM warmup, `SparkContext` construction, and Connect server boot -- +which can take a few seconds and makes a quick edit/run loop feel slow. + +To amortize that cost across runs, start one persistent local Spark Connect server and point +every run at it: + +```bash +# Start once; it stays up across runs. (--master is optional; it defaults to local[*].) +$SPARK_HOME/sbin/start-connect-server.sh --master "local[*]" + +# Every run reconnects instead of booting a new server. +python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("sc://localhost:15002").getOrCreate()' + +# Stop it when you are done. +$SPARK_HOME/sbin/stop-connect-server.sh +``` + +Alternatively, on POSIX systems PySpark can manage this persistent server for you. With +`SPARK_LOCAL_CONNECT_REUSE=1` set (or `spark.local.connect.reuse=true` on the builder), +`SparkSession.builder.remote("local[*]").getOrCreate()` starts a persistent server through +`sbin/start-connect-server.sh` on the first run and reconnects to it on later runs, so scripts +keep the plain `local[*]` URL: + +```bash +export SPARK_LOCAL_CONNECT_REUSE=1 + +# The first run starts the server; later runs reconnect to it. +python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("local[*]").getOrCreate()' + +# Stop the managed server when you are done. +python -m pyspark.sql.connect.local_server --stop +``` + +The managed server is an ordinary `spark-daemon.sh` daemon, but it runs with a per-user pid +directory and ident string so it cannot collide with a server you started by hand -- which also +means a plain `sbin/stop-connect-server.sh` does not find it. The `--stop` command signals the +recorded server and cleans up the discovery file; killing the server's pid directly also works, +and the next run notices the dead server and starts a fresh one. + +This managed-server workflow is experimental. The `--stop` command and the discovery file +location and format may change in a future release, for example if local server management is +folded into a unified `spark connect` CLI. + +The connection details (host, port, auth token, pid, Spark version) are recorded in a discovery +file in a private per-user directory; set `SPARK_LOCAL_CONNECT_DISCOVERY` to override its +location. A run reconnects only to a server whose Spark version matches. After upgrading Spark, +the server from the previous version cannot be reused and the next run fails with an error asking +you to stop it; run the `--stop` command above and rerun to start a fresh server. + +Each run connects as its own Connect session, so session-local state -- temp views, runtime SQL +configurations, and session artifacts -- is fresh on every run and never leaks between runs. State +backed by the shared `SparkContext` (the persistent catalog/warehouse, global temp views, and +cached datasets) *is* shared across runs, so namespace per-run databases or clear that state +yourself if your runs must be fully isolated. + +### Fully isolated runs with a pool of single-use servers + +If your runs must be fully isolated from each other but you still want to skip the per-run +startup cost, a second experimental opt-in keeps a small pool of booted servers that have +never been assigned to an application run instead of one shared one. With +`SPARK_LOCAL_CONNECT_POOL=1` set (or +`spark.local.connect.pool=true` on the builder), each run *claims* a fresh server from the pool, +a replacement is booted in the background, and the claimed server is torn down when the run's +session stops -- no server ever serves two runs, so runs are as isolated as with the default +in-process server: + +```bash +export SPARK_LOCAL_CONNECT_POOL=1 + +# Each run claims a booted server not previously assigned to another run. +python -c 'from pyspark.sql import SparkSession; SparkSession.builder.remote("local[*]").getOrCreate()' + +# Force-stop all pool servers and start over from a clean slate. +python -m pyspark.sql.connect.local_server_pool --purge +``` + +The trade-off relative to the reuse mode above is memory: `spark.local.connect.pool.size` +(default 2) idle servers stay resident while you iterate, several hundred MB each. They shut +down on their own after sitting unclaimed for `SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT` seconds +(default 1800), so an idle machine drains back to zero servers. Pool servers are only handed to +runs whose master, startup configurations, working directory, and Python environment match the +ones they were booted with; runs that differ in any of these boot their own pool members. If +both this and `spark.local.connect.reuse` are set, the pool takes precedence. + ## Use Spark Connect in standalone applications <div class="codetabs"> @@ -371,7 +468,7 @@ one may implement their own class extending `ClassFinder` for customized search </div> For more information on application development with Spark Connect as well as extending Spark Connect -with custom functionality, see [Application Development with Spark Connect](app-dev-spark-connect.html). +with custom functionality, see [Application Development with Spark Connect](app-dev-spark-connect.html). # Client application authentication While Spark Connect does not have built-in authentication, it is designed to @@ -413,3 +510,81 @@ APIs such as [SparkContext](api/scala/org/apache/spark/SparkContext.html) and [RDD](api/scala/org/apache/spark/rdd/RDD.html) are unsupported in Spark Connect. Support for more APIs is planned for upcoming Spark releases. + +# Routing through a shared ingress or reverse proxy + +When several services share a single hostname behind a Kubernetes Ingress (or another +reverse proxy), it is tempting to give each service a URL path prefix (for example +`sc://host/sparkConnect`) and route on that path. This does not work for gRPC: the +gRPC method name *is* the HTTP/2 `:path` (`/spark.connect.SparkConnectService/ExecutePlan`), +so a connection string cannot carry a separate routing path. Prepending a prefix to the +`:path` produces an unknown method and the server responds with `UNIMPLEMENTED`, unless the +proxy is configured to strip the prefix back off before forwarding, a two-sided contract +that is not part of gRPC's design. + +The routing dimension that *is* free is the HTTP/2 `:authority` (the virtual host). Set it +to a routing tag with the `grpc.default_authority` channel option, and route on it at the +proxy (for example, an Ingress `host:` rule). The client still dials the shared hostname +(`default_authority` only overrides the `:authority` header used for routing, not the address +it connects to), and the gRPC method `:path` is never touched, so no path rewrite is needed. +This is the approach the gRPC maintainers recommend for this scenario +([grpc/grpc#14900](https://github.com/grpc/grpc/issues/14900)). + +The example below uses the Python client, which exposes gRPC channel options directly: + +{% highlight python %} +from pyspark.sql.connect.session import SparkSession +from pyspark.sql.connect.client import DefaultChannelBuilder + +cb = DefaultChannelBuilder("sc://myhost.com:443") +cb.setChannelOption("grpc.default_authority", "sparkconnect") # routing tag +spark = SparkSession.builder.channelBuilder(cb).getOrCreate() +{% endhighlight %} + +A corresponding Kubernetes Ingress routes on that tag and keeps the service at path `/` +(no subpath, no rewrite). Note the routing tag must be a lowercase +[RFC 1123](https://datatracker.ietf.org/doc/html/rfc1123) name, since a Kubernetes Ingress +`host:` requires one: + +{% highlight yaml %} +spec: + rules: + - host: sparkconnect # matches grpc.default_authority + http: + paths: + - path: / + pathType: Prefix + backend: + service: + name: spark-connect-server + port: { number: 15002 } +{% endhighlight %} + +**Over TLS**, gRPC uses the `:authority` as the certificate-verification name by default, +so overriding it with a routing tag would fail verification (the tag is not in the server +certificate's SAN). Keep the two names separate: set `grpc.ssl_target_name_override` to the +real server hostname (used for certificate verification and SNI) and `grpc.default_authority` +to the routing tag. + +{% highlight python %} +cb = DefaultChannelBuilder("sc://myhost.com:443/;use_ssl=true") +cb.setChannelOption("grpc.ssl_target_name_override", "myhost.com") # certificate verification / SNI +cb.setChannelOption("grpc.default_authority", "sparkconnect") # routing tag +spark = SparkSession.builder.channelBuilder(cb).getOrCreate() +{% endhighlight %} + +`use_ssl=true` verifies the server certificate against the system's trusted CA store. If your +gateway's certificate is issued by a CA the client does not trust by default (a self-signed or +internal CA), make that CA trusted on the client side (for example, via +`GRPC_DEFAULT_SSL_ROOTS_FILE_PATH`) rather than through the connection string. Configuring the +proxy's own TLS (which certificate it presents for which hostname) is part of your ingress +setup and is out of scope here. + +Two notes on `grpc.ssl_target_name_override`: +gRPC documents it as testing-oriented because its typical misuse is to *mask* a certificate +name mismatch (verifying against a name the server does not actually present, which defeats +hostname verification). Here it is set to the real, verified hostname, so the certificate is +still checked correctly; it only prevents the routing tag from being used as the verification +name. If you prefer to avoid the option entirely, issue the server certificate with the routing +tag included in its SAN; then `grpc.default_authority` alone is sufficient and no override is +needed. diff --git a/docs/spark-standalone.md b/docs/spark-standalone.md index ec1656b0348c8..125ae01428cc9 100644 --- a/docs/spark-standalone.md +++ b/docs/spark-standalone.md @@ -21,7 +21,7 @@ license: | * This will become a table of contents (this text will be scraped). {:toc} -In addition to running on the YARN cluster manager, Spark also provides a simple standalone deploy mode. You can launch a standalone cluster either manually, by starting a master and workers by hand, or use our provided [launch scripts](#cluster-launch-scripts). It is also possible to run these daemons on a single machine for testing. +In addition to running on the [YARN cluster manager](running-on-yarn.html) and [Kubernetes](running-on-kubernetes.html), Spark also provides a simple standalone deploy mode. You can launch a standalone cluster either manually, by starting a master and workers by hand, or use our provided [launch scripts](#cluster-launch-scripts). It is also possible to run these daemons on a single machine for testing. # Security @@ -100,13 +100,16 @@ Once you've set up this file, you can launch or stop your cluster with the follo - `sbin/start-master.sh` - Starts a master instance on the machine the script is executed on. - `sbin/start-workers.sh` - Starts a worker instance on each machine specified in the `conf/workers` file. - `sbin/start-worker.sh` - Starts a worker instance on the machine the script is executed on. -- `sbin/start-connect-server.sh` - Starts a Spark Connect server on the machine the script is executed on. +- `sbin/start-connect-server.sh` - Starts a [Spark Connect](spark-connect-overview.html) server on the machine the script is executed on. +- `sbin/start-history-server.sh` - Starts the History Server, which lets you view logs for completed applications. Requires event logging to be enabled (see [Monitoring and Instrumentation](monitoring.html)). - `sbin/start-all.sh` - Starts both a master and a number of workers as described above. - `sbin/stop-master.sh` - Stops the master that was started via the `sbin/start-master.sh` script. - `sbin/stop-worker.sh` - Stops all worker instances on the machine the script is executed on. - `sbin/stop-workers.sh` - Stops all worker instances on the machines specified in the `conf/workers` file. - `sbin/stop-connect-server.sh` - Stops all Spark Connect server instances on the machine the script is executed on. +- `sbin/stop-history-server.sh` - Stops the History Server. - `sbin/stop-all.sh` - Stops both the master and the workers as described above. +- `sbin/decommission-worker.sh` - Gracefully decommissions a worker, allowing in-progress tasks to finish and shuffles to be migrated before the worker exits. Note that these scripts must be executed on the machine you want to run the Spark master on, not your local machine. @@ -687,7 +690,7 @@ configurations, <code>curl</code> CLI command can provide the required header li ```bash $ curl -XPOST http://IP:PORT/v1/submissions/create \ ---header "Authorization: Bearer USER-PROVIDED-WEB-TOEN-SIGNED-BY-THE-SAME-SHARED-KEY" +--header "Authorization: Bearer USER-PROVIDED-WEB-TOKEN-SIGNED-BY-THE-SAME-SHARED-KEY" ... ``` @@ -748,7 +751,7 @@ worker during one single schedule iteration. Stage level scheduling is supported on Standalone: - When dynamic allocation is disabled: It allows users to specify different task resource requirements at the stage level and will use the same executors requested at startup. -- When dynamic allocation is enabled: Currently, when the Master allocates executors for one application, it will schedule based on the order of the ResourceProfile ids for multiple ResourceProfiles. The ResourceProfile with smaller id will be scheduled firstly. Normally this won’t matter as Spark finishes one stage before starting another one, the only case this might have an affect is in a job server type scenario, so its something to keep in mind. For scheduling, we will only take executor memory and executor cores from built-in executor resources and all other custom resources from a ResourceProfile, other built-in executor resources such as offHeap and memoryOverhead won't take any effect. The base default profile will be created based on the spark configs when you submit an application. Executor memory and executor cores from the base default profile can be propagated to custom ResourceProfiles, but all other custom resources can not be propagated. +- When dynamic allocation is enabled: Currently, when the Master allocates executors for one application, it will schedule based on the order of the ResourceProfile ids for multiple ResourceProfiles. The ResourceProfile with smaller id will be scheduled first. Normally this won’t matter as Spark finishes one stage before starting another one; the only case this might have an effect is in a job server type scenario, so it’s something to keep in mind. For scheduling, we will only take executor memory and executor cores from built-in executor resources and all other custom resources from a ResourceProfile; other built-in executor resources such as `offHeap` and `memoryOverhead` won't take any effect. The base default profile will be created based on the Spark configs when you submit an application. Executor memory and executor cores from the base default profile can be propagated to custom ResourceProfiles, but all other custom resources can not be propagated. ## Caveats @@ -760,6 +763,20 @@ Spark's standalone mode offers a web-based user interface to monitor the cluster In addition, detailed log output for each job is also written to the work directory of each worker node (`SPARK_HOME/work` by default). You will see two files for each job, `stdout` and `stderr`, with all output it wrote to its console. +To track and review logs across completed applications, [enable event logging and start the History Server](monitoring.html#viewing-after-the-fact). + +## Held Applications + +An application that can be held reports to the Master whether it currently is, and the Master web +UI annotates the application state accordingly, for example `RUNNING (held, draining 2 executors)`. +An executor that has not exited yet is still finishing its running tasks, and the hold is complete +once no executor is left. The Master's `/json/` endpoint reports the same in the `holdsupported`, +`held`, and `draining` fields of each application. + +Only applications whose driver reports that it can be held are annotated, per the preconditions +described in [Web UI](web-ui.html#jobs-tab). Holding and resuming an application is done from its +own driver web UI. + # Running Alongside Hadoop @@ -835,7 +852,7 @@ In order to enable this recovery mode, you can set SPARK_DAEMON_JAVA_OPTS in spa <td><code>spark.deploy.recoveryDirectory</code></td> <td>""</td> <td>The directory in which Spark will store recovery state, accessible from the Master's perspective. - Note that the directory should be clearly manually if <code>spark.deploy.recoveryMode</code> + Note that the directory should be cleared manually if <code>spark.deploy.recoveryMode</code> or <code>spark.deploy.recoveryCompressionCodec</code> is changed. </td> <td>0.8.1</td> diff --git a/docs/sql-data-sources-csv.md b/docs/sql-data-sources-csv.md index 9dfe9739b7ac7..bf63fe2be5d3e 100644 --- a/docs/sql-data-sources-csv.md +++ b/docs/sql-data-sources-csv.md @@ -207,9 +207,9 @@ Data source options of CSV can be set via: <tr> <td><code>mode</code></td> <td>PERMISSIVE</td> - <td>Allows a mode for dealing with corrupt records during parsing. It supports the following case-insensitive modes. Note that Spark tries to parse only required columns in CSV under column pruning. Therefore, corrupt records can be different based on required set of fields. This behavior can be controlled by <code>spark.sql.csv.parser.columnPruning.enabled</code> (enabled by default).<br> + <td>Allows a mode for dealing with corrupt records during parsing. It supports the following case-insensitive modes. Note that Spark tries to parse only required columns in CSV under column pruning. Therefore, corrupt records can be different based on required set of fields. This behavior can be controlled by <code>spark.sql.csv.parser.columnPruning.enabled</code> (enabled by default). In particular, when <code>multiLine</code> is disabled, a quoted value containing a line separator (<code>lineSep</code>) splits one source record into two before <code>mode</code> is applied. Each of the two is then evaluated on its own, so a record whose token count happens to match the schema is retained as valid even though its value was truncated. This holds under <code>DROPMALFORMED</code> as well, since that mode drops only the records it finds malformed. An action requiring no columns (a bare <code>count()</code>, for example) may surface none of this because of column pruning.<br> <ul> - <li><code>PERMISSIVE</code>: when it meets a corrupted record, puts the malformed string into a field configured by <code>columnNameOfCorruptRecord</code>, and sets malformed fields to <code>null</code>. To keep corrupt records, an user can set a string type field named <code>columnNameOfCorruptRecord</code> in an user-defined schema. If a schema does not have the field, it drops corrupt records during parsing. A record with less/more tokens than schema is not a corrupted record to CSV. When it meets a record having fewer tokens than the length of the schema, sets <code>null</code> to extra fields. When the record has more tokens than the length of the schema, it drops extra tokens.</li> + <li><code>PERMISSIVE</code>: when it meets a corrupted record, puts the malformed string into a field configured by <code>columnNameOfCorruptRecord</code>, and sets malformed fields to <code>null</code>. To capture the malformed string, a user can set a string type field named <code>columnNameOfCorruptRecord</code> in a user-defined schema. If a schema does not have the field, the corrupt record is still retained with its malformed fields set to <code>null</code>, but the malformed string is not available. A record with a different number of tokens than the schema is a corrupted record, and is handled by this <code>mode</code> like any other; it is still parsed as far as it can be, with <code>null</code> set for tokens the record does not have and extra tokens dropped.</li> <li><code>DROPMALFORMED</code>: ignores the whole corrupted records. This mode is unsupported in the CSV built-in functions.</li> <li><code>FAILFAST</code>: throws an exception when it meets corrupted records.</li> </ul> @@ -231,7 +231,8 @@ Data source options of CSV can be set via: <tr> <td><code>multiLine</code></td> <td>false</td> - <td>Allows a row to span multiple lines, by parsing line breaks within quoted values as part of the value itself. CSV built-in functions ignore this option.</td> + <td>Allows a row to span multiple lines, by parsing line breaks within quoted values as part of the value itself. CSV built-in functions ignore this option.<br> + When this option is disabled (the default), a line break inside a quoted value terminates the record at that break: the value is truncated, and the rest of the value begins a new record. Either or both of the resulting records may then be malformed (for example when the token count no longer matches the schema); how they are handled is controlled by <code>mode</code>.</td> <td>read</td> </tr> <tr> diff --git a/docs/sql-data-sources-jdbc.md b/docs/sql-data-sources-jdbc.md index a205ed1a4e541..9ba3eacef63fe 100644 --- a/docs/sql-data-sources-jdbc.md +++ b/docs/sql-data-sources-jdbc.md @@ -374,6 +374,14 @@ logging into the data sources. </td> <td>read</td> </tr> + <tr> + <td><code>preferTimestampNanos</code></td> + <td>false</td> + <td> + When the option is set to <code>true</code>, a driver TIMESTAMP column that reports a sub-microsecond fractional-second scale (7-9) is inferred as one of the nanosecond-capable timestamp types (<code>TIMESTAMP_NTZ(p)</code> when <code>preferTimestampNTZ</code> is also <code>true</code>, otherwise <code>TIMESTAMP_LTZ(p)</code>). Otherwise such columns keep the historical microsecond mapping. This option only takes effect when the <code>spark.sql.timestampNanosTypes.enabled</code> preview flag is enabled; when that flag is off, setting this option alone leaves the inferred schema unchanged. + </td> + <td>read</td> + </tr> <tr> <td><code>hint</code></td> <td>(none)</td> diff --git a/docs/sql-data-sources-json.md b/docs/sql-data-sources-json.md index ba8691320f902..6ac5d690868c8 100644 --- a/docs/sql-data-sources-json.md +++ b/docs/sql-data-sources-json.md @@ -170,7 +170,7 @@ Data source options of JSON can be set via: <td><code>PERMISSIVE</code></td> <td>Allows a mode for dealing with corrupt records during parsing.<br> <ul> - <li><code>PERMISSIVE</code>: when it meets a corrupted record, puts the malformed string into a field configured by <code>columnNameOfCorruptRecord</code>, and sets malformed fields to <code>null</code>. To keep corrupt records, an user can set a string type field named <code>columnNameOfCorruptRecord</code> in an user-defined schema. If a schema does not have the field, it drops corrupt records during parsing. When inferring a schema, it implicitly adds a <code>columnNameOfCorruptRecord</code> field in an output schema.</li> + <li><code>PERMISSIVE</code>: when it meets a corrupted record, puts the malformed string into a field configured by <code>columnNameOfCorruptRecord</code>, and sets malformed fields to <code>null</code>. To capture the malformed string, a user can set a string type field named <code>columnNameOfCorruptRecord</code> in a user-defined schema. If a schema does not have the field, the corrupt record is still retained with its malformed fields set to <code>null</code>, but the malformed string is not available. When inferring a schema, it implicitly adds a <code>columnNameOfCorruptRecord</code> field in an output schema.</li> <li><code>DROPMALFORMED</code>: ignores the whole corrupted records. This mode is unsupported in the JSON built-in functions.</li> <li><code>FAILFAST</code>: throws an exception when it meets corrupted records.</li> </ul> @@ -249,6 +249,18 @@ Data source options of JSON can be set via: <td>For reading, allows to forcibly set one of standard basic or extended encoding for the JSON files. For example UTF-16BE, UTF-32LE. For writing, Specifies encoding (charset) of saved json files. JSON built-in functions ignore this option.</td> <td>read/write</td> </tr> + <tr> + <td><code>pretty</code></td> + <td><code>false</code></td> + <td>If true, writes the generated JSON with the default pretty printer, indenting nested structures over multiple lines instead of emitting each record on a single line.</td> + <td>write</td> + </tr> + <tr> + <td><code>writeNonAsciiCharacterAsCodePoint</code></td> + <td><code>false</code></td> + <td>If true, writes non-ASCII characters as \uXXXX escape sequences instead of emitting them literally.</td> + <td>write</td> + </tr> <tr> <td><code>lineSep</code></td> <td><code>\r</code>, <code>\r\n</code>, <code>\n</code> (for reading), <code>\n</code> (for writing)</td> diff --git a/docs/sql-data-sources-xml.md b/docs/sql-data-sources-xml.md index e714ded0ee285..8a07a486b70d7 100644 --- a/docs/sql-data-sources-xml.md +++ b/docs/sql-data-sources-xml.md @@ -248,6 +248,30 @@ Data source options of XML can be set via: <td>Compression codec to use when saving to file. This can be one of the known case-insensitive shortened names (none, bzip2, gzip, lz4, snappy and deflate). XML built-in functions ignore this option.</td> <td>write</td> </tr> + <tr> + <td><code>indent</code></td> + <td>four spaces</td> + <td>String used to indent each nested level of the generated XML. Setting it to an empty string disables indentation, writing each row on a new line.</td> + <td>write</td> + </tr> + <tr> + <td><code>multiLine</code></td> + <td><code>true</code></td> + <td>Whether to parse one record, which may span multiple lines, per file.</td> + <td>read</td> + </tr> + <tr> + <td><code>prefersDecimal</code></td> + <td><code>false</code></td> + <td>During schema inference, infers floating-point values as <code>DecimalType</code> rather than <code>DoubleType</code> when they fit.</td> + <td>read</td> + </tr> + <tr> + <td><code>preferDate</code></td> + <td><code>true</code></td> + <td>During schema inference, tries to infer string columns that contain dates as <code>DateType</code>. Disabled when <code>spark.sql.legacy.timeParserPolicy</code> is set to <code>LEGACY</code>.</td> + <td>read</td> + </tr> <tr> <td><code>validateName</code></td> diff --git a/docs/sql-distributed-sql-engine.md b/docs/sql-distributed-sql-engine.md index ae8fd9c7211bd..53fc8c1b5f77e 100644 --- a/docs/sql-distributed-sql-engine.md +++ b/docs/sql-distributed-sql-engine.md @@ -40,23 +40,23 @@ specify Hive properties. You may run `./sbin/start-thriftserver.sh --help` for a all available options. By default, the server listens on localhost:10000. You may override this behaviour via either environment variables, i.e.: -{% highlight bash %} +```bash export HIVE_SERVER2_THRIFT_PORT=<listening-port> export HIVE_SERVER2_THRIFT_BIND_HOST=<listening-host> ./sbin/start-thriftserver.sh \ --master <master-uri> \ ... -{% endhighlight %} +``` or system properties: -{% highlight bash %} +```bash ./sbin/start-thriftserver.sh \ --hiveconf hive.server2.thrift.port=<listening-port> \ --hiveconf hive.server2.thrift.bind.host=<listening-host> \ - --master <master-uri> + --master <master-uri> \ ... -{% endhighlight %} +``` Now you can use beeline to test the Thrift JDBC/ODBC server: @@ -85,13 +85,20 @@ To test, use beeline to connect to the JDBC/ODBC server in http mode with: beeline> !connect jdbc:hive2://<host>:<port>/<database>;transportMode=http;httpPath=<http_endpoint> -If you closed a session and do CTAS, you must set `fs.%s.impl.disable.cache` to true in `hive-site.xml`. -See more details in [[SPARK-21067]](https://issues.apache.org/jira/browse/SPARK-21067). +If you closed a session and do CTAS, you must set `fs.<scheme>.impl.disable.cache` to `true` in `hive-site.xml`. `<scheme>` is the Hadoop filesystem you are using, like `hdfs`, `s3a`, etc. For example: + + fs.hdfs.impl.disable.cache=true + +See more details in [SPARK-21067](https://issues.apache.org/jira/browse/SPARK-21067). + +To stop the Thrift JDBC/ODBC server, run: + + ./sbin/stop-thriftserver.sh ## Running the Spark SQL CLI To use the Spark SQL command line interface (CLI) from the shell: ./bin/spark-sql - -For details, please refer to [Spark SQL CLI](sql-distributed-sql-engine-spark-sql-cli.html) + +For details, please refer to [Spark SQL CLI](sql-distributed-sql-engine-spark-sql-cli.html). diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index 6efed7224c0d5..40c9784b986f6 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -22,22 +22,34 @@ license: | * Table of contents {:toc} +## Upgrading from Spark SQL 4.3 to 4.4 + +- Since Spark 4.4, for storage-partitioned joins, `spark.sql.requireAllClusterKeysForCoPartition` requires every join key to be covered by some partition key instead of matching the partition keys positionally. As a result, a join-key column partitioned by more than one transform no longer prevents shuffle elimination, and `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` no longer additionally requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false` when the join keys are a subset of the partition keys. As before, when the partition keys cover only part of the join keys, eliminating the shuffle still requires `spark.sql.requireAllClusterKeysForCoPartition` to be `false`. + ## Upgrading from Spark SQL 4.2 to 4.3 +- Since Spark 4.3, map-typed arguments to distinct aggregates are normalized by key order. As a result, `COUNT(DISTINCT m)` treats maps with the same entries in different orders as the same value, and aggregates that return distinct input values, such as `COLLECT_LIST(DISTINCT m)`, return maps sorted by key. - Since Spark 4.3, [ASOF JOIN](sql-ref-syntax-qry-select-asof-join.html) is available as an opt-in SQL feature gated by `spark.sql.join.asofJoin.enabled` (default `false`). When disabled, `ASOF JOIN` fails at parse time with `UNSUPPORTED_FEATURE.ASOF_JOIN`. - Since Spark 4.3, zero-length files are skipped during Parquet schema inference instead of failing with a `FAILED_READ_FILE.CANNOT_READ_FILE_FOOTER` error. - Since Spark 4.3, metrics produced by `Dataset.observe` include values from only the last successful task attempts instead of aggregating values from all attempts. To restore the previous behavior, set `spark.sql.legacy.observeMetricsAggregateAllAttempts` to `true`. - Since Spark 4.3, the configuration key `spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled` has been renamed to `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` to reflect that it now applies to storage-partitioned joins, aggregates, and windows. The old key continues to work as an alias. - Since Spark 4.3, the Spark Thrift Server rejects setting JVM system properties through the `set:system:` session configuration overlay (for example, in a JDBC connection string). To restore the previous behavior, set `spark.sql.legacy.hive.thriftServer.allowSettingSystemProperties` to `true`. - Since Spark 4.3, the adaptive execution rule `org.apache.spark.sql.execution.adaptive.DynamicJoinSelection` has been renamed to `DemoteBroadcastHashJoin`, which now only demotes broadcast hash joins (emitting `NO_BROADCAST_HASH`). Its selection of shuffled hash join over sort merge join has moved to a new physical rule gated by `spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled` (default `true`). If you previously disabled the shuffled-hash-join preference by listing `org.apache.spark.sql.execution.adaptive.DynamicJoinSelection` in `spark.sql.adaptive.optimizer.excludedRules`, that name no longer matches any rule (unknown names are silently ignored); set `spark.sql.adaptive.convertSortMergeJoinToShuffledHashJoin.enabled` to `false` instead. +- Since Spark 4.3, the optimizer rule `MergeSubplans` has moved from package `org.apache.spark.sql.catalyst.optimizer` to `org.apache.spark.sql.execution.planmerging`, so its fully qualified name is now `org.apache.spark.sql.execution.planmerging.MergeSubplans` (the rule was itself renamed from `MergeScalarSubqueries` to `MergeSubplans` in Spark 4.2). If you previously disabled it by listing `org.apache.spark.sql.catalyst.optimizer.MergeSubplans`, or the older `org.apache.spark.sql.catalyst.optimizer.MergeScalarSubqueries`, in `spark.sql.optimizer.excludedRules`, that name no longer matches any rule (unknown names are silently ignored); use `org.apache.spark.sql.execution.planmerging.MergeSubplans` instead. - Since Spark 4.3, `spark.sql.execution.replaceHashWithSortAgg` defaults to `true`. Spark now replaces a hash-based aggregate with a sort aggregate when the aggregate's child is already sorted on the grouping keys. To restore the previous behavior, set `spark.sql.execution.replaceHashWithSortAgg` to `false`. - Since Spark 4.3, `spark.sql.execution.combineAdjacentAggregation` defaults to `true`. Spark now merges an adjacent partial/final aggregate pair (with no shuffle between them) into a single complete-mode aggregate. This setting is independent of `spark.sql.execution.replaceHashWithSortAgg`, so disabling only `replaceHashWithSortAgg` still leaves adjacent aggregation combined; to fully restore the previous partial/final staging, set both `spark.sql.execution.replaceHashWithSortAgg` and `spark.sql.execution.combineAdjacentAggregation` to `false`. - Since Spark 4.3, the exact `percentile`, `percentile_cont`, and `median` aggregate functions (including their `WITHIN GROUP (ORDER BY ...)` forms) compute the linear interpolation between two neighboring values as `lower + fraction * (higher - lower)` instead of `(1 - fraction) * lower + fraction * higher`. The two are equal in exact arithmetic, but the new form is monotonically non-decreasing in the requested percentage and avoids a rounding error the old form could introduce. As a result these functions may return a value that differs from earlier releases in the last ULP. `percentile_disc` and `percentile_approx` are unaffected. -- Since Spark 4.3, non-deterministic filters (for example predicates involving `rand()`) are no longer pushed down to DataSource V2 sources that implement `SupportsPushDownV2Filters`; they are evaluated by Spark after the scan instead. This prevents a source from evaluating such a predicate a different number of times than Spark, or using it for pruning while also returning it for post-scan re-evaluation. +- Since Spark 4.3, non-deterministic filters (for example predicates involving `rand()`) are no longer pushed down to DataSource V2 sources: neither at query compilation to sources that implement `SupportsPushDownV2Filters`, nor as runtime filters to sources that implement `SupportsRuntimeV2Filtering`; they are evaluated by Spark after the scan instead. This prevents a source from evaluating such a predicate a different number of times than Spark, or using it for pruning while also returning it for post-scan re-evaluation. - Since Spark 4.3, the new `COMMENT ON COLUMN ... IS NULL` syntax removes a column comment by passing a `null` comment to `TableChange.updateColumnComment(String[], String)`, so a `UpdateColumnComment` table change may now carry a `null` `newComment()`. Previously `newComment()` was always non-null. DataSource V2 catalogs that handle `UpdateColumnComment` should null-check `newComment()` and treat `null` as "remove the column comment" (mirroring how `UpdateColumnDefaultValue` already carries a `null` value to drop a default). - Since Spark 4.3, `unix_seconds`, `unix_millis`, and `unix_micros` accept `TIMESTAMP_NTZ` and the nanosecond-precision timestamp types directly, reading them with no time-zone shift. Previously these functions accepted only `TIMESTAMP_LTZ`; a `TIMESTAMP_NTZ` or nanosecond-timestamp argument was rejected with a `DATATYPE_MISMATCH` error. This is a new capability and does not change the result of any query that previously succeeded. - Since Spark 4.3, `hash()` and `xxhash64()` include the `days` field of `CalendarInterval` when computing the hash, so their output for interval values differs from earlier releases. Previously the codegen path dropped `days`, disagreeing with interpreted evaluation. - Since Spark 4.3, when a `SELECT` or `INSERT` statement references the same table more than once with different `WITH (...)` options (for example a self-join, or `INSERT INTO t WITH (...) SELECT * FROM t WITH (...)`), each reference now uses its own options instead of the second reference silently inheriting the first reference's options via the analyzer's relation cache. +- Since Spark 4.3, `DataFrameWriter.save()` on a `SupportsCatalogOptions` source loads its write target with the required `TableWritePrivilege`s (`INSERT`, or `INSERT` and `DELETE` for `SaveMode.Overwrite`), matching `insertInto` and `saveAsTable`. A catalog that enforces these privileges may now reject an unauthorized `save()` that previously bypassed this authorization check. +- Since Spark 4.3, the Spark-recognized time-travel options (`versionAsOf` and `timestampAsOf`, or the keys configured by `spark.sql.timeTravelVersionKey` and `spark.sql.timeTravelTimestampKey`) are rejected with `UNSUPPORTED_FEATURE.TIME_TRAVEL` on catalog-backed Data Source V2 writes, including table creation and replacement through `DataFrameWriterV2`, because writes must target the current table state rather than a historical version. Previously, Spark passed them to the connector as ordinary write options. +- Since Spark 4.3, `HAVING` is evaluated before window functions when the `SELECT` list also contains generator functions such as `explode`. Previously, window functions could include groups removed by `HAVING` and produce incorrect results. +- Since Spark 4.3, the Spark Connect session errors `INVALID_HANDLE.SESSION_CHANGED`/`SESSION_CLOSED`/`SESSION_NOT_FOUND` carry SQLSTATE `08003` instead of `HY000`; the condition names are unchanged. Code matching these errors on SQLSTATE should match `08003` or class `08`. +- Since Spark 4.3, [Declarative Pipelines](declarative-pipelines-programming-guide.html) honors `spark.sql.caseSensitive` when inferring and evolving pipeline table schemas. Under case-insensitive resolution (the default), column names that differ only in case now identify the same column: flows writing to one table contribute a single column rather than one per spelling, and a column that differs only in case from one already persisted in the target is written to that column instead of being added alongside it. Previously such names were always treated as distinct, producing a table schema that Spark's own resolver could not disambiguate and that could fail later with errors such as `COLUMN_ALREADY_EXISTS` or `AMBIGUOUS_REFERENCE`. When two flows' columns fold together but their types are incompatible, the update now fails at validation with `UNABLE_TO_INFER_PIPELINE_TABLE_SCHEMA`. Where the spellings differ, the surviving one comes from the flow with the lowest identifier. An explicitly declared table schema keeps its spelling over the inferred one. Incremental streaming tables keep the persisted spelling of an existing column; materialized views re-infer the schema on every update. Set `spark.sql.caseSensitive` to `true` to keep names differing only in case distinct, as before. +- Since Spark 4.3, all flows writing to the same pipeline table must agree on the effective `spark.sql.caseSensitive`, which each flow takes from its own SQL configuration (a `SET` in pipeline source, which never reaches the session) and otherwise from the session. A disagreement fails the update with `CONFLICTING_PIPELINE_FLOW_CASE_SENSITIVITY`, because that value decides whether names differing only in case identify the same column and would otherwise make the table's schema depend on the order the flows are evaluated in. Set `spark.sql.caseSensitive` to the same value for every flow writing to the table, remembering that a flow which does not set it inherits the session's value. ## Upgrading from Spark SQL 4.1 to 4.2 diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md index 3ba3a6c749ca7..fee8801086c4c 100644 --- a/docs/sql-performance-tuning.md +++ b/docs/sql-performance-tuning.md @@ -30,7 +30,7 @@ Spark SQL can cache tables using an in-memory columnar format by calling `spark. Then Spark SQL will scan only required columns and will automatically tune compression to minimize memory usage and GC pressure. You can call `spark.catalog.uncacheTable("tableName")` or `dataFrame.unpersist()` to remove the table from memory. -To list relations cached with an explicit name, use `spark.catalog.listCachedTables()`. Entries cached only via `Dataset.cache()` without a name are not included. +To check whether a specific table or view is cached, use `spark.catalog.isCached("tableName")`. To inspect the storage level of an arbitrary `Dataset`, read its `storageLevel` property, which returns `StorageLevel.NONE` when the data is not currently cached. For an overview of everything persisted in the running application, including data cached directly via `Dataset.cache()`, use the [Storage tab](web-ui.html#storage-tab) of the web UI, which shows the storage levels, sizes and partitions of each persisted relation once an action has materialized it. Spark supports two cache formats: - **Default cache format**: The standard in-memory columnar cache (used by default). @@ -101,7 +101,7 @@ Configuration of in-memory caching can be done via `spark.conf.set` or by runnin <td>None</td> <td> The suggested (not guaranteed) maximum number of split file partitions. If it is set, - Spark will rescale each partition to make the number of partitions is close to this + Spark will rescale each partition to make the number of partitions close to this value if the initial number of partitions exceeds this value. This configuration is effective only when using file-based sources such as Parquet, JSON and ORC. </td> @@ -181,6 +181,74 @@ Missing or inaccurate statistics will hinder Spark's ability to select an optima - **Query plan estimates**: You can inspect Spark's cost estimates in the optimized query plan via [`EXPLAIN COST`](sql-ref-syntax-qry-explain.html) or `DataFrame.explain(mode="cost")`. - **Runtime statistics**: You can inspect these statistics in the [SQL UI](web-ui.html#sql-tab) under the "Details" section as a query is running. Look for `Statistics(..., isRuntime=true)` in the plan. +## Optimizing the Aggregate + +### Adaptive Partial Aggregation + +A grouping aggregation normally runs in two phases: a partial aggregation before the shuffle and a +final aggregation after it. The partial aggregation is only worthwhile when it actually reduces the +number of rows; when the grouping keys are close to unique it maintains, and possibly spills, an +aggregation map roughly as large as its input while emitting almost as many rows as it consumed. + +When adaptive partial aggregation is enabled, hash aggregation measures the compaction ratio (the +number of processed rows divided by the number of keys held in its aggregation maps) at runtime +and, if the partial aggregation is not collapsing enough rows to be worthwhile, stops populating +the aggregation map and passes the remaining rows through as single-row partial aggregation buffers +for the final aggregation to merge. Once pass-through is active the map is frozen and its output +always comes before the passed-through rows: rows that collide with the frozen map are held in a +queue behind it and flushed only after it drains, so a duplicate of a key already in the map still +merges after that key's accumulated rows, keeping an order-sensitive aggregate such as +`first`/`last` consistent with a run that never bypasses. The ratio is evaluated periodically, and +again right before the aggregation map would spill, in which case the spill is skipped entirely. + +Both evaluations use the cumulative rows and keys of the current map epoch and, once triggered, +pass-through is not reversed for the rest of the task, so a skewed prefix biases the outcome in +either direction. A favorable prefix can mask a later distinct-heavy tail, keeping the aggregation +on until a spill restarts the accounting; conversely, a distinct-heavy prefix can trip the +pass-through early and keep it tripped even where the rest of the input would aggregate well. To +catch the former without waiting for a spill, raise `minCompaction` so the cumulative ratio trips +the threshold on a weaker late turn (a higher threshold also bypasses more readily on other inputs). +To avoid committing the task to pass-through on the latter, raise `minRows` so the periodic +evaluations start later. + +<table class="spark-config"> + <thead><tr><th>Property Name</th><th>Default</th><th>Meaning</th><th>Since Version</th></tr></thead> + <tr> + <td><code>spark.sql.execution.aggregate.adaptivePartialAggregation.enabled</code></td> + <td>false</td> + <td> + When true, hash aggregation adaptively bypasses the pre-shuffle partial aggregation at runtime + when it observes that the partial aggregation is not reducing the number of rows enough to be + worthwhile. This applies only to hash aggregation with grouping keys. + </td> + <td>4.4.0</td> + </tr> + <tr> + <td><code>spark.sql.execution.aggregate.adaptivePartialAggregation.minRows</code></td> + <td>100000</td> + <td> + The number of rows between periodic compaction-ratio evaluations. Setting this to + <code>0</code> disables the periodic evaluation. The ratio may still be evaluated when the + aggregation map is about to spill. A larger value also delays the periodic evaluations, so + when one of them trips pass-through the frozen map tends to hold more rows; the frozen map + stays resident until its output is drained, so a larger value raises that transient memory + peak. + </td> + <td>4.4.0</td> + </tr> + <tr> + <td><code>spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction</code></td> + <td>1.05</td> + <td> + The minimum compaction ratio required to keep the partial aggregation. A ratio of 10 means the + partial aggregation collapses ten rows into one key; when an evaluation finds the ratio below + this value the partial aggregation is bypassed for the rest of the input. A larger value + bypasses more aggressively. + </td> + <td>4.4.0</td> + </tr> +</table> + ## Optimizing the Join Strategy ### Automatically Broadcasting Joins @@ -258,6 +326,78 @@ SELECT /*+ BROADCAST(r) */ * FROM src s JOIN records r ON s.key = r.key For more details please refer to the documentation of [Join Hints](sql-ref-syntax-qry-select-hints.html#join-hints). +## Merging Subplans + +Spark merges subplans that return a single row and read the same input, so that input is scanned once instead of once per subplan. The candidates are non-correlated deterministic scalar subqueries and non-grouping aggregates (aggregates without `GROUP BY`). The merged subplan is evaluated once and outputs a single struct, and each original site reads its own field out of that struct. This optimization is enabled by default. + +For example, in the following query both subqueries scan `store_sales`: + +```sql +SELECT + (SELECT min(ss_net_paid) FROM store_sales), + (SELECT max(ss_net_paid) FROM store_sales) +``` + +They are merged into one aggregate that computes `min` and `max` together, so `store_sales` is read once. In `EXPLAIN` output a merged subplan shows up as a subquery whose single output column is named `mergedValue`, and the sites that share it as `ReusedSubquery`. + +Two subplans are merged when their plans match node by node: `Project` lists are unioned, `Aggregate`s must have the same grouping and use the same aggregation implementation (so a `min` is not merged with a `collect_list`), `Filter`s must have the same condition, `Join`s must have the same type, condition and hints, and the leaves must read the same input. Subplans that differ only in their `WHERE` conditions can be merged as well, by turning each side's condition into a boolean column and giving each side's aggregate expressions a `FILTER (WHERE ...)` clause. That is controlled by the configurations below. Queries that still contain a `WITH` clause when this rule runs (one that was not inlined) are skipped. + +When only one of the two subplans has a filter, merging is always beneficial, because the unfiltered side reads all the data anyway. This case is on by default, unless the filter has to cross a `Join` to reach the aggregate, which needs the through-join configuration below. When both sides have a filter (the symmetric case), the merged scan filter becomes `OR(f1, f2)`, which is less selective than either original filter and can therefore read more data - for example when the filters prune partitions or Parquet row groups. That is why the symmetric case is disabled by default. + +Still, it is worth considering on queries that compute several differently filtered aggregates over the same table, which is a common analytical shape: + +```sql +SELECT + (SELECT avg(ss_net_paid) FROM store_sales WHERE ss_quantity BETWEEN 1 AND 20), + (SELECT avg(ss_net_paid) FROM store_sales WHERE ss_quantity BETWEEN 21 AND 40) +``` + +In TPC-DS benchmark runs, enabling symmetric filter propagation made `q9` and `q28` about 3.5x faster, and enabling it together with propagation through joins made `q88` about 7x and `q90` about 2x faster. See [SPARK-40193](https://issues.apache.org/jira/browse/SPARK-40193) and [SPARK-56677](https://issues.apache.org/jira/browse/SPARK-56677) for these measurements. The trade-off depends on the tables: the gain is largest when the differing filters are on columns the data source cannot prune on, and the risk is highest on heavily partitioned or file-pruned tables where the widened filter loses that pruning. Validate it on your own workload before enabling it in production. + +<table class="spark-config"> + <thead><tr><th>Property Name</th><th>Default</th><th>Meaning</th><th>Since Version</th></tr></thead> + <tr> + <td><code>spark.sql.optimizer.mergeSubplans.filterPropagation.enabled</code></td> + <td>true</td> + <td> + When true, subplans that differ only in their filter conditions can be merged by propagating the filters up to the enclosing non-grouping aggregates. This is the umbrella configuration of the three below: none of them has any effect when this is false. Subplans with the same filter condition are merged regardless of this configuration. + </td> + <td>4.2.0</td> + </tr> + <tr> + <td><code>spark.sql.optimizer.mergeSubplans.filterPropagation.symmetricFilterPropagation.enabled</code></td> + <td>false</td> + <td> + When true, two non-grouping aggregate subplans that both have a filter condition can also be merged. Disabled by default because the merged filter is widened to <code>OR(f1, f2)</code>, which may read more data than the two original filters, especially on heavily partitioned or file-pruned tables. + </td> + <td>4.2.0</td> + </tr> + <tr> + <td><code>spark.sql.optimizer.mergeSubplans.filterPropagation.throughJoin.enabled</code></td> + <td>false</td> + <td> + When true, filter conditions can also propagate through <code>Join</code> nodes, which lets subplans that differ only in their filter conditions and share a common join be merged. When false, no filter is propagated across a join, not even when only one of the two subplans has a filter. A filter only propagates from the preserved side of the join: the left side of <code>LEFT OUTER</code>/<code>LEFT SEMI</code>/<code>LEFT ANTI</code>, the right side of <code>RIGHT OUTER</code>, or either side of <code>INNER</code>/<code>CROSS</code>. <code>FULL OUTER</code> joins are never eligible. Subplans that differ in their filters typically have a filter on both sides, so this is usually set together with <code>spark.sql.optimizer.mergeSubplans.filterPropagation.symmetricFilterPropagation.enabled</code>. + </td> + <td>4.2.0</td> + </tr> + <tr> + <td><code>spark.sql.optimizer.mergeSubplans.filterPropagation.dsv2SymmetricFilterPropagation.enabled</code></td> + <td>false</td> + <td> + When true, two DataSource V2 scans that pushed the same strictly enforced filters but carry different best-effort (post-scan) filters can be merged even when <code>spark.sql.optimizer.mergeSubplans.filterPropagation.symmetricFilterPropagation.enabled</code> is false. In this case widening cannot change the set of rows the scan is required to return, as the strict filters are re-pushed unchanged and the enclosing <code>Filter</code> re-checks the rest above the scan. This applies only to V2 sources that opt in to scan merging with the <code>SCAN_MERGING</code> table capability; no built-in source does. + </td> + <td>4.3.0</td> + </tr> +</table> + +Subplan merging can be turned off entirely by adding the rule to `spark.sql.optimizer.excludedRules`: + +``` +spark.sql.optimizer.excludedRules=org.apache.spark.sql.execution.planmerging.MergeSubplans +``` + +Use exactly that name: the rule was called `MergeScalarSubqueries` before Spark 4.2 and sat in a different package before Spark 4.3, and unknown names in `spark.sql.optimizer.excludedRules` are silently ignored, so an older name carried over from a previous version does not turn the rule off. See the [SQL migration guide](sql-migration-guide.html) for the old names. + ## Adaptive Query Execution Adaptive Query Execution (AQE) is an optimization technique in Spark SQL that makes use of the runtime statistics to choose the most efficient query execution plan, which is enabled by default since Apache Spark 3.2.0. Spark SQL can turn on and off AQE by `spark.sql.adaptive.enabled` as an umbrella configuration. @@ -289,7 +429,7 @@ This feature coalesces the post shuffle partitions based on the map output stati <td><code>spark.sql.adaptive.coalescePartitions.parallelismFirst</code></td> <td>true</td> <td> - When true, Spark ignores the target size specified by <code>spark.sql.adaptive.advisoryPartitionSizeInBytes</code> (default 64MB) when coalescing contiguous shuffle partitions, and only respect the minimum partition size specified by <code>spark.sql.adaptive.coalescePartitions.minPartitionSize</code> (default 1MB), to maximize the parallelism. This is to avoid performance regressions when enabling adaptive query execution. It's recommended to set this config to false on a busy cluster to make resource utilization more efficient (not many small tasks). + When true, Spark ignores the target size specified by <code>spark.sql.adaptive.advisoryPartitionSizeInBytes</code> (default 64MB) when coalescing contiguous shuffle partitions, and only respects the minimum partition size specified by <code>spark.sql.adaptive.coalescePartitions.minPartitionSize</code> (default 1MB), to maximize the parallelism. This is to avoid performance regressions when enabling adaptive query execution. It's recommended to set this config to false on a busy cluster to make resource utilization more efficient (not many small tasks). </td> <td>3.2.0</td> </tr> @@ -342,7 +482,7 @@ This feature coalesces the post shuffle partitions based on the map output stati <td><code>spark.sql.adaptive.rebalancePartitionsSmallPartitionFactor</code></td> <td>0.2</td> <td> - A partition will be merged during splitting if its size is small than this factor multiply <code>spark.sql.adaptive.advisoryPartitionSizeInBytes</code>. + A partition will be merged during splitting if its size is smaller than this factor multiplying <code>spark.sql.adaptive.advisoryPartitionSizeInBytes</code>. </td> <td>3.3.0</td> </tr> @@ -482,7 +622,7 @@ You can control the details of how AQE works by providing your own cost evaluato ## Storage Partition Join -Storage Partition Join (SPJ) is an optimization technique in Spark SQL that makes use the existing storage layout to avoid the shuffle phase. +Storage Partition Join (SPJ) is an optimization technique in Spark SQL that makes use of the existing storage layout to avoid the shuffle phase. This is a generalization of the concept of Bucket Joins, which is only applicable for [bucketed](sql-data-sources-load-save-functions.html#bucketing-sorting-and-partitioning) tables, to tables partitioned by functions registered in FunctionCatalog. Storage Partition Joins are currently supported for compatible V2 DataSources. @@ -510,9 +650,9 @@ The following SQL properties enable Storage Partition Join in different join que <td><code>spark.sql.requireAllClusterKeysForCoPartition</code></td> <td>true</td> <td> - When true, require the join or MERGE keys to be same and in the same order as the partition keys to eliminate shuffle. Hence, set to <b>false</b> in this situation to eliminate shuffle. + When true, storage-partitioned join requires every join or MERGE key to be covered by some partition key (rather than matching the partition keys positionally) to eliminate shuffle. When the partition keys cover only part of the join or MERGE keys, set to <b>false</b> to eliminate shuffle, at the risk of data skew and reduced parallelism from the coarser storage partitioning. </td> - <td>3.4.0</td> + <td>3.3.0</td> </tr> <tr> <td><code>spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled</code></td> @@ -523,10 +663,10 @@ The following SQL properties enable Storage Partition Join in different join que <td>3.4.0</td> </tr> <tr> - <td><code>spark.sql.sources.v2.bucketing.allowJoinKeysSubsetOfPartitionKeys.enabled</code></td> + <td><code>spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled</code></td> <td>false</td> <td> - When enabled, try to avoid shuffle if join or MERGE condition does not include all partition columns. This config requires both <code>spark.sql.sources.v2.bucketing.enabled</code> and <code>spark.sql.sources.v2.bucketing.pushPartValues.enabled</code> to be true, and <code>spark.sql.requireAllClusterKeysForCoPartition</code> to be false. + When enabled, try to avoid shuffle if join or MERGE condition does not include all partition columns. This config requires both <code>spark.sql.sources.v2.bucketing.enabled</code> and <code>spark.sql.sources.v2.bucketing.pushPartValues.enabled</code> to be true. </td> <td>4.0.0</td> </tr> @@ -582,7 +722,6 @@ ON t.dep = s.dep AND t.id = s.id SET 'spark.sql.sources.v2.bucketing.enabled' 'true' SET 'spark.sql.iceberg.planning.preserve-data-grouping' 'true' SET 'spark.sql.sources.v2.bucketing.pushPartValues.enabled' 'true' -SET 'spark.sql.requireAllClusterKeysForCoPartition' 'false' SET 'spark.sql.sources.v2.bucketing.partiallyClusteredDistribution.enabled' 'true' -- Plan with Storage Partition Join diff --git a/docs/sql-ref-ansi-compliance.md b/docs/sql-ref-ansi-compliance.md index f4ac58a4e3f02..5ab05fd026f75 100644 --- a/docs/sql-ref-ansi-compliance.md +++ b/docs/sql-ref-ansi-compliance.md @@ -484,6 +484,7 @@ Below is a list of all the keywords in Spark SQL. |COMPUTE|non-reserved|non-reserved|non-reserved| |CONCATENATE|non-reserved|non-reserved|non-reserved| |CONDITION|non-reserved|non-reserved|non-reserved| +|CONDITIONAL|non-reserved|non-reserved|non-reserved| |CONSTRAINT|reserved|non-reserved|reserved| |CONTAINS|non-reserved|non-reserved|non-reserved| |CONTINUE|non-reserved|non-reserved|non-reserved| @@ -537,8 +538,10 @@ Below is a list of all the keywords in Spark SQL. |DROP|non-reserved|non-reserved|reserved| |ELSE|reserved|non-reserved|reserved| |ELSEIF|non-reserved|non-reserved|non-reserved| +|EMPTY|non-reserved|non-reserved|reserved| |END|reserved|non-reserved|reserved| |ENFORCED|non-reserved|non-reserved|non-reserved| +|ERROR|non-reserved|non-reserved|non-reserved| |ESCAPE|reserved|non-reserved|reserved| |ESCAPED|non-reserved|non-reserved|non-reserved| |EVOLUTION|non-reserved|non-reserved|non-reserved| @@ -616,6 +619,11 @@ Below is a list of all the keywords in Spark SQL. |ITERATE|non-reserved|non-reserved|non-reserved| |JOIN|reserved|strict-non-reserved|reserved| |JSON|non-reserved|non-reserved|non-reserved| +|JSON_EXISTS|non-reserved|non-reserved|reserved| +|JSON_QUERY|non-reserved|non-reserved|reserved| +|JSON_TABLE|non-reserved|non-reserved|reserved| +|JSON_VALUE|non-reserved|non-reserved|reserved| +|KEEP|non-reserved|non-reserved|non-reserved| |KEY|non-reserved|non-reserved|non-reserved| |KEYS|non-reserved|non-reserved|non-reserved| |LANGUAGE|non-reserved|non-reserved|reserved| @@ -675,8 +683,10 @@ Below is a list of all the keywords in Spark SQL. |NULL|reserved|non-reserved|reserved| |NULLS|non-reserved|non-reserved|non-reserved| |NUMERIC|non-reserved|non-reserved|non-reserved| +|OBJECT|non-reserved|non-reserved|non-reserved| |OF|non-reserved|non-reserved|reserved| |OFFSET|reserved|non-reserved|reserved| +|OMIT|non-reserved|non-reserved|reserved| |ON|reserved|strict-non-reserved|reserved| |ONLY|reserved|non-reserved|reserved| |OPEN|non-reserved|non-reserved|reserved| @@ -684,6 +694,7 @@ Below is a list of all the keywords in Spark SQL. |OPTIONS|non-reserved|non-reserved|non-reserved| |OR|reserved|non-reserved|reserved| |ORDER|reserved|non-reserved|reserved| +|ORDINALITY|non-reserved|non-reserved|non-reserved| |OUT|non-reserved|non-reserved|reserved| |OUTER|reserved|non-reserved|reserved| |OUTPUTFORMAT|non-reserved|non-reserved|non-reserved| @@ -709,6 +720,7 @@ Below is a list of all the keywords in Spark SQL. |QUALIFY|non-reserved|non-reserved|non-reserved| |QUARTER|non-reserved|non-reserved|non-reserved| |QUERY|non-reserved|non-reserved|non-reserved| +|QUOTES|non-reserved|non-reserved|non-reserved| |RANGE|non-reserved|non-reserved|reserved| |READ|non-reserved|non-reserved|non-reserved| |READS|non-reserved|non-reserved|non-reserved| @@ -732,6 +744,7 @@ Below is a list of all the keywords in Spark SQL. |RESPECT|non-reserved|non-reserved|non-reserved| |RESTRICT|non-reserved|non-reserved|non-reserved| |RETURN|non-reserved|non-reserved|reserved| +|RETURNING|non-reserved|non-reserved|non-reserved| |RETURNS|non-reserved|non-reserved|reserved| |REVOKE|non-reserved|non-reserved|reserved| |RIGHT|reserved|strict-non-reserved|reserved| @@ -818,11 +831,13 @@ Below is a list of all the keywords in Spark SQL. |UNARCHIVE|non-reserved|non-reserved|non-reserved| |UNBOUNDED|non-reserved|non-reserved|non-reserved| |UNCACHE|non-reserved|non-reserved|non-reserved| +|UNCONDITIONAL|non-reserved|non-reserved|non-reserved| |UNIFORM|non-reserved|non-reserved|non-reserved| |UNION|reserved|strict-non-reserved|reserved| |UNIQUE|reserved|non-reserved|reserved| |UNKNOWN|reserved|non-reserved|reserved| |UNLOCK|non-reserved|non-reserved|non-reserved| +|UNNEST|non-reserved|non-reserved|non-reserved| |UNPIVOT|non-reserved|non-reserved|non-reserved| |UNSET|non-reserved|non-reserved|non-reserved| |UNTIL|non-reserved|non-reserved|non-reserved| @@ -851,6 +866,7 @@ Below is a list of all the keywords in Spark SQL. |WITH|reserved|non-reserved|reserved| |WITHIN|reserved|non-reserved|reserved| |WITHOUT|non-reserved|non-reserved|non-reserved| +|WRAPPER|non-reserved|non-reserved|non-reserved| |X|non-reserved|non-reserved|non-reserved| |YEAR|non-reserved|non-reserved|non-reserved| |YEARS|non-reserved|non-reserved|non-reserved| diff --git a/docs/sql-ref-name-resolution.md b/docs/sql-ref-name-resolution.md index 3d574e58a9ad2..797e12acde639 100644 --- a/docs/sql-ref-name-resolution.md +++ b/docs/sql-ref-name-resolution.md @@ -36,7 +36,7 @@ Identifiers in expressions can be references to any one of the following: Name resolution applies the following principles: - The _closest_ matching reference wins, and -- Columns and parameter win over fields and keys. +- Columns and parameters win over fields and keys. In detail, resolution of identifiers to a specific reference follows these rules in order: @@ -62,7 +62,7 @@ In detail, resolution of identifiers to a specific reference follows these rules A. Remove the last identifier and treat it as a field or key. - B. Match the remainder to a column in table reference of the `FROM clause`. + B. Match the remainder to a column in a table reference of the `FROM clause`. - If there is more than one such match, raise an AMBIGUOUS_COLUMN_OR_FIELD error. @@ -72,7 +72,7 @@ In detail, resolution of identifiers to a specific reference follows these rules If the field cannot be matched, raise a FIELD_NOT_FOUND error. - If there is more than one field, raise a AMBIGUOUS_COLUMN_OR_FIELD error. + If there is more than one field, raise an AMBIGUOUS_COLUMN_OR_FIELD error. - **`MAP`**: Raise an error if the key is qualified. @@ -161,15 +161,15 @@ This restriction also applies to parameter references in SQL functions. > SELECT t.a FROM VALUES(named_struct('a', 1)) AS t(t); 1 --- A column takes precendece over a field +-- A column takes precedence over a field > SELECT t.a FROM VALUES(named_struct('a', 1), 2) AS t(t, a); 2 --- Implict lateral column alias +-- Implicit lateral column alias > SELECT c1 AS a, a + c1 FROM VALUES(2) AS T(c1); 2 4 --- A local column reference takes precedence, over a lateral column alias +-- A local column reference takes precedence over a lateral column alias > SELECT c1 AS a, a + c1 FROM VALUES(2, 3) AS T(c1, a); 2 5 @@ -204,7 +204,7 @@ This restriction also applies to parameter references in SQL functions. WHERE c4 = c2 * 2); [UNRESOLVED_COLUMN] `c2` --- Successsful usage of lateral correlation with keyword LATERAL +-- Successful usage of lateral correlation with keyword LATERAL > SELECT c1, c2, c3 FROM VALUES(1, 2) AS t(c1, c2), LATERAL(SELECT c3 FROM VALUES(3, 4) AS s(c3, c4) @@ -318,7 +318,7 @@ the effective search path, for example > CREATE TABLE rel(c1 int); > INSERT INTO rel VALUES(1); --- An fully qualified reference to rel: +-- A fully qualified reference to rel: > SELECT c1 FROM spark_catalog.default.rel; 1 @@ -428,7 +428,7 @@ effective search path, for example > CREATE FUNCTION func(a INT, b INT) RETURNS INT RETURN a / b; --- The temporary function takes precedent +-- The temporary function takes precedence > SELECT func(4, 2); 2 diff --git a/docs/sql-ref-syntax-aux-describe-table.md b/docs/sql-ref-syntax-aux-describe-table.md index 5c417689dde2b..e2db0ab3d88cd 100644 --- a/docs/sql-ref-syntax-aux-describe-table.md +++ b/docs/sql-ref-syntax-aux-describe-table.md @@ -164,6 +164,16 @@ to return the metadata pertaining to a partition or column respectively. | MapType | `{ "name" : "map", "key_type": <type_json>, "value_type": <type_json>, "value_nullable": <boolean> }` | | StructType | `{ "name" : "struct", "fields": [ {"name" : "field1", "type" : <type_json>, “nullable”: <boolean>, "comment": “<comment>”, "default": “<default_val>”}, ... ] }` | +**Note** +- If a V1 (session catalog) table's declared partition columns do not match the last columns + of its schema, the catalog metadata is considered inconsistent. `DESCRIBE TABLE` still + returns the rest of the metadata and reports both column lists under an + `# Invalid Partition Information` section in place of `# Partition Information`. Repair + the table metadata so the declared partition columns match the last columns of the + schema. Until then, commands that inspect a specific partition may still fail: + `DESCRIBE TABLE ... PARTITION`, `SHOW TABLE EXTENDED ... PARTITION`, and + `SHOW PARTITIONS ... PARTITION`. + ### Examples ```sql diff --git a/docs/sql-ref-syntax-ddl-create-table-datasource.md b/docs/sql-ref-syntax-ddl-create-table-datasource.md index f645732a15df9..07bdbe4f4173f 100644 --- a/docs/sql-ref-syntax-ddl-create-table-datasource.md +++ b/docs/sql-ref-syntax-ddl-create-table-datasource.md @@ -26,7 +26,7 @@ The `CREATE TABLE` statement defines a new table using a Data Source. ### Syntax ```sql -CREATE TABLE [ IF NOT EXISTS ] table_identifier +CREATE [ EXTERNAL ] TABLE [ IF NOT EXISTS ] table_identifier [ ( col_name1 col_type1 [ COMMENT col_comment1 ], ... ) ] USING data_source [ OPTIONS ( key1=val1, key2=val2, ... ) ] @@ -78,6 +78,14 @@ as any order. For example, you can write COMMENT table_comment after TBLPROPERTI Specifies buckets numbers, which is used in `CLUSTERED BY` clause. +* **EXTERNAL** + + The table is defined using the path provided as `LOCATION` and does not use the default location for this table. + Dropping an external table removes catalog metadata and leaves the data files in place. + `CREATE EXTERNAL TABLE` for a data source table must include `LOCATION`. + A data source table created with `LOCATION` is also external even without the + `EXTERNAL` keyword. + * **LOCATION** Path to the directory where table data is stored, which could be a path on distributed storage like HDFS, etc. @@ -116,6 +124,11 @@ input query, to make sure the table gets created contains exactly the same data --Use data source CREATE TABLE student (id INT, name STRING, age INT) USING CSV; +-- External data source table. DROP TABLE keeps the files at the location. +CREATE EXTERNAL TABLE student (id INT, name STRING, age INT) + USING parquet + LOCATION '/tmp/student'; + --Use data from another table CREATE TABLE student_copy USING CSV AS SELECT * FROM student; diff --git a/docs/sql-ref-syntax-qry-select-json-exists.md b/docs/sql-ref-syntax-qry-select-json-exists.md new file mode 100644 index 0000000000000..48d307d4fac3a --- /dev/null +++ b/docs/sql-ref-syntax-qry-select-json-exists.md @@ -0,0 +1,127 @@ +--- +layout: global +title: JSON_EXISTS +displayTitle: JSON_EXISTS +license: | + 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. +--- + +### Description + +The `JSON_EXISTS` predicate tests whether a SQL/JSON path matches at least one item in a JSON +document, returning a `BOOLEAN`. This is the SQL-standard (SQL:2016) way to test for the presence +of a JSON value, and is commonly used to migrate queries from other systems such as Oracle, DB2, +and PostgreSQL. + +Unlike `get_json_object(json_expr, path) IS NOT NULL`, `JSON_EXISTS` distinguishes a path that is +_present but whose value is JSON `null`_ (which is `true`) from a path that is _absent_ (which is +`false`). + +### Syntax + +```sql +JSON_EXISTS ( json_expr, path [ { TRUE | FALSE | UNKNOWN | ERROR } ON ERROR ] ) +``` + +### Parameters + +* **json_expr** + + An expression that evaluates to a `STRING` containing the JSON document. A `NULL` input yields + `NULL` (SQL Unknown), regardless of the `ON ERROR` clause. + +* **path** + + A constant SQL/JSON path literal (for example `'$.a.b'`, `'$.tags[0]'`, or `'$.a[*].b'`). Paths + are evaluated in **lax** mode, matching Oracle and PostgreSQL: array wildcards (`[*]`) and member + wildcards (`.*` / `['*']`) are supported, and arrays are auto-wrapped/unwrapped (a member, index, + or wildcard step applied to an array is applied to each element, and a non-array value is treated + as a single-element array). A structural mismatch is a non-match, not an error. A syntactically + invalid path is rejected during analysis. + +* **{ TRUE | FALSE | UNKNOWN | ERROR } ON ERROR** + + Controls the result when `json_expr` is not a single well-formed JSON value (malformed input, + or a valid value followed by trailing content). `TRUE`, `FALSE`, and `UNKNOWN` produce that + value (`UNKNOWN` is a `BOOLEAN` `NULL`); `ERROR` raises an error. The default is + `FALSE ON ERROR`. + +### Result + +* The path matches at least one item (including a match whose value is JSON `null`) → `true`. +* The path matches nothing → `false`. +* `json_expr` is SQL `NULL` → `NULL`. +* `json_expr` is not a single well-formed JSON value → the `ON ERROR` behavior. + +A structural mismatch is treated as "no match" (`false`), not an error -- for example reading an +absent key, reading a key from a scalar, an out-of-range array index, or `[*]` over an empty array. + +### Examples + +```sql +SELECT json_exists('{"a":{"b":1}}', '$.a.b') AS matched; ++-------+ +|matched| ++-------+ +| true| ++-------+ + +-- Present but JSON null -> true; absent -> false +SELECT json_exists('{"a":null}', '$.a') AS present_null, + json_exists('{"a":1}', '$.b') AS absent; ++------------+------+ +|present_null|absent| ++------------+------+ +| true| false| ++------------+------+ + +-- NULL input -> NULL (Unknown), regardless of the ON ERROR clause +SELECT json_exists(CAST(NULL AS STRING), '$.a' TRUE ON ERROR) AS r; ++----+ +| r| ++----+ +|NULL| ++----+ + +-- Malformed input follows the ON ERROR clause (default FALSE) +SELECT json_exists('not json', '$.a') AS default_false, + json_exists('not json', '$.a' TRUE ON ERROR) AS true_on_error, + json_exists('not json', '$.a' UNKNOWN ON ERROR) AS unknown_on_error; ++-------------+-------------+----------------+ +|default_false|true_on_error|unknown_on_error| ++-------------+-------------+----------------+ +| false| true| NULL| ++-------------+-------------+----------------+ + +-- Lax wildcards: [*] is true iff the array has elements; auto-unwrap applies a step to each element +SELECT json_exists('{"a":[1,2]}', '$.a[*]') AS has_elems, + json_exists('{"a":[]}', '$.a[*]') AS empty_array, + json_exists('{"a":[{"b":1},{"c":2}]}', '$.a[*].b') AS any_elem_has_b; ++---------+-----------+--------------+ +|has_elems|empty_array|any_elem_has_b| ++---------+-----------+--------------+ +| true| false| true| ++---------+-----------+--------------+ + +-- Use as a predicate in WHERE +SELECT id FROM docs WHERE json_exists(doc, '$.address.city'); +``` + +### Related Statements + +* [SELECT](sql-ref-syntax-qry-select.html) +* [WHERE Clause](sql-ref-syntax-qry-select-where.html) +* [JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) diff --git a/docs/sql-ref-syntax-qry-select-json-query.md b/docs/sql-ref-syntax-qry-select-json-query.md new file mode 100644 index 0000000000000..2cbc60674dfbe --- /dev/null +++ b/docs/sql-ref-syntax-qry-select-json-query.md @@ -0,0 +1,174 @@ +--- +layout: global +title: JSON_QUERY +displayTitle: JSON_QUERY +license: | + 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. +--- + +### Description + +The `JSON_QUERY` function extracts the JSON value located by a SQL/JSON path from a JSON document +and returns it as JSON text (a `STRING`). This is the SQL-standard way (SQL:2016) to pull an object, +array, or scalar fragment out of JSON, and is commonly used to migrate queries from other systems +such as Oracle, SQL Server, and Trino. Unlike +[JSON_TABLE](sql-ref-syntax-qry-select-json-table.html), which produces rows in a `FROM` clause, +`JSON_QUERY` is an expression that can appear anywhere a value is allowed. + +Where [JSON_VALUE](sql-ref-syntax-qry-select-json-value.html) returns a single scalar (and treats an +object or array match as an error), `JSON_QUERY` returns the matched value serialized as JSON text, +whether it is an object, an array, or a scalar. + +This implementation supports simple, wildcard-free SQL/JSON paths only. The `PASSING` clause, path +predicates and filters, and explicit `lax` / `strict` path modes defined by SQL:2016 are not +supported. + +### Syntax + +```sql +JSON_QUERY ( json_expr, path + [ RETURNING data_type ] + [ wrapper_behavior ] + [ quotes_behavior ] + [ empty_behavior ON EMPTY ] + [ error_behavior ON ERROR ] ) + +wrapper_behavior + { WITHOUT [ ARRAY ] WRAPPER + | WITH [ CONDITIONAL | UNCONDITIONAL ] [ ARRAY ] WRAPPER } + +quotes_behavior + { KEEP QUOTES | OMIT QUOTES } + +empty_behavior + { NULL | ERROR | EMPTY ARRAY | EMPTY OBJECT } + +error_behavior + { NULL | ERROR | EMPTY ARRAY | EMPTY OBJECT } +``` + +### Parameters + +* **json_expr** + + An expression that evaluates to a `STRING` containing the JSON document. A `NULL` input yields + `NULL` directly (it triggers neither the `ON EMPTY` nor the `ON ERROR` behavior). + +* **path** + + A SQL/JSON path literal that locates the value, for example `'$.a.b'` or `'$.items[0]'`. The + path must be wildcard-free; a path containing `[*]` is rejected at analysis time. + +* **RETURNING data_type** + + The type of the result. It must be a string type; the result is JSON text. If `RETURNING` is + omitted, the result type is `STRING`. + +* **wrapper_behavior** + + Whether to wrap the result in a JSON array: + * `WITHOUT ARRAY WRAPPER` (the default) returns the value unwrapped. + * `WITH UNCONDITIONAL ARRAY WRAPPER` (or simply `WITH ARRAY WRAPPER`) always wraps the value in a + one-element array. + * `WITH CONDITIONAL ARRAY WRAPPER` wraps the value only when it is a scalar; an object or array + is returned unwrapped. + +* **quotes_behavior** + + Whether to keep the surrounding quotes of a scalar string result: + * `KEEP QUOTES` (the default) leaves them, so a string is returned as a quoted JSON string. + * `OMIT QUOTES` strips them, returning the raw string content. It is a no-op for objects, + arrays, and non-string scalars, and cannot be combined with an array wrapper. + +* **empty_behavior ON EMPTY** + + What to produce when `path` matches nothing: + * `NULL` (the default) returns SQL `NULL`. + * `ERROR` raises an error. + * `EMPTY ARRAY` returns the JSON text `[]`. + * `EMPTY OBJECT` returns the JSON text `{}`. + +* **error_behavior ON ERROR** + + What to produce when the input is not well-formed JSON. The same four choices as `ON EMPTY` + apply, defaulting to `NULL`. + +A path that matches an explicit JSON `null` is a present scalar value and returns the JSON text +`null` (it is neither the `ON EMPTY` nor the `ON ERROR` case). + +Returning a scalar under the default `WITHOUT ARRAY WRAPPER` is an intentional convenience: the +matched scalar is emitted as JSON text (for example, `JSON_QUERY('{"id":7}', '$.id')` returns `7`), +whereas strict SQL:2016 treats a scalar without a wrapper as an error. The wrapper clauses behave the +standard way: `WITH CONDITIONAL ARRAY WRAPPER` wraps a scalar in a one-element array (`7` becomes +`[7]`) while leaving a single object or array unwrapped, and `WITH UNCONDITIONAL ARRAY WRAPPER` +always wraps. + +### Examples + +```sql +-- Extract an object as JSON text +SELECT json_query('{"id":7,"addr":{"city":"NYC"}}', '$.addr'); ++---------------------------------------------------+ +|json_query({"id":7,"addr":{"city":"NYC"}}, $.addr) | ++---------------------------------------------------+ +|{"city":"NYC"} | ++---------------------------------------------------+ + +-- Extract an array +SELECT json_query('{"tags":["x","y"]}', '$.tags'); ++-------------------------------------------+ +|json_query({"tags":["x","y"]}, $.tags) | ++-------------------------------------------+ +|["x","y"] | ++-------------------------------------------+ + +-- Wrap a scalar in an array with WITH ARRAY WRAPPER +-- (WITH ARRAY WRAPPER is a shorthand; the column name shows the canonical +-- WITH UNCONDITIONAL ARRAY WRAPPER form) +SELECT json_query('{"tags":["x","y"]}', '$.tags[0]' WITH ARRAY WRAPPER); ++----------------------------------------------------------------------------+ +|json_query({"tags":["x","y"]}, $.tags[0] WITH UNCONDITIONAL ARRAY WRAPPER) | ++----------------------------------------------------------------------------+ +|["x"] | ++----------------------------------------------------------------------------+ + +-- Strip the quotes from a scalar string with OMIT QUOTES +SELECT json_query('{"name":"Ada"}', '$.name' OMIT QUOTES); ++---------------------------------------------------+ +|json_query({"name":"Ada"}, $.name OMIT QUOTES) | ++---------------------------------------------------+ +|Ada | ++---------------------------------------------------+ + +-- A missing path defaults to NULL; supply a fallback with EMPTY ARRAY ON EMPTY +SELECT json_query('{"id":7}', '$.missing' EMPTY ARRAY ON EMPTY); ++---------------------------------------------------------+ +|json_query({"id":7}, $.missing EMPTY ARRAY ON EMPTY) | ++---------------------------------------------------------+ +|[] | ++---------------------------------------------------------+ + +-- ERROR ON ERROR raises instead of returning a value +SELECT json_query('not json', '$.a' ERROR ON ERROR); +[JSON_QUERY_ON_ERROR.ERROR] ... +``` + +### Related Statements + +* [SELECT](sql-ref-syntax-qry-select.html) +* [JSON_VALUE](sql-ref-syntax-qry-select-json-value.html) +* [JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) +* [Built-in Functions](sql-ref-functions-builtin.html) diff --git a/docs/sql-ref-syntax-qry-select-json-table.md b/docs/sql-ref-syntax-qry-select-json-table.md new file mode 100644 index 0000000000000..a8f40513778c2 --- /dev/null +++ b/docs/sql-ref-syntax-qry-select-json-table.md @@ -0,0 +1,130 @@ +--- +layout: global +title: JSON_TABLE +displayTitle: JSON_TABLE +license: | + 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. +--- + +### Description + +The `JSON_TABLE` table-valued function shreds a JSON document into a relational table. A +_row path_ selects a sequence of JSON items, and a `COLUMNS` clause projects a value out of each +item into a typed column. This is the SQL-standard way (SQL:2016) to turn JSON into rows and +columns, and is commonly used to migrate queries from other systems such as Oracle, DB2, and +MySQL. + +Only the flat (non-nested) form is currently supported. `NESTED PATH` columns are not yet +supported. + +### Syntax + +```sql +JSON_TABLE ( json_expr, row_path COLUMNS ( column_definition [ , ... ] ) [ error_clause ] ) [ table_alias ] + +column_definition + { column_name FOR ORDINALITY + | column_name data_type [ PATH json_path ] + | column_name data_type EXISTS [ PATH json_path ] } + +error_clause + { NULL | ERROR } ON ERROR +``` + +### Parameters + +* **json_expr** + + An expression that evaluates to a `STRING` containing the JSON document. + +* **row_path** + + A SQL/JSON path literal that selects the row source. A path ending in `[*]` (for example + `'$.items[*]'`) selects each element of the matched array as a separate row. Any other path + (for example `'$'`) selects a single value as one row. If the path matches nothing, no rows + are produced. + +* **column_name FOR ORDINALITY** + + Declares a `BIGINT` column that is a 1-based sequential counter of the generated rows. + +* **column_name data_type [ PATH json_path ]** + + A value column. The value at `json_path` (relative to a row item) is extracted and cast to + `data_type`. If `PATH` is omitted, the path defaults to the column name read as a single + object key: a simple identifier maps like `$.name`, while a name containing special characters + such as a dot is treated as one literal key (for example a column named `a.b` reads the key + `"a.b"`, equivalent to `$['a.b']`, not the nested path `a` -> `b`). If the path matches nothing, + the column is `NULL`. + +* **column_name data_type EXISTS [ PATH json_path ]** + + An existence column. Evaluates to a truthy value when `json_path` matches and a falsy value + otherwise, cast to `data_type` (for example `BOOLEAN`). + +* **{ NULL | ERROR } ON ERROR** + + Controls behavior when `json_expr` is `NULL` or not well-formed JSON. `NULL ON ERROR` (the + default) produces no rows. `ERROR ON ERROR` raises an error. + +* **table_alias** + + An optional alias for the output, optionally followed by a column alias list. + +### Examples + +```sql +-- Expand a JSON array into rows with typed columns and an ordinality counter +SELECT t.* FROM JSON_TABLE( + '{"items":[{"id":1,"n":"a"},{"id":2,"n":"b"}]}', + '$.items[*]' + COLUMNS ( + seq FOR ORDINALITY, + id INT PATH '$.id', + name STRING PATH '$.n' + ) +) AS t; ++---+---+----+ +|seq| id|name| ++---+---+----+ +| 1| 1| a| +| 2| 2| b| ++---+---+----+ + +-- Implicit column path derived from the column name, and an EXISTS column +SELECT * FROM JSON_TABLE( + '{"rows":[{"id":10,"opt":1},{"id":20}]}', + '$.rows[*]' + COLUMNS (id INT, hasOpt BOOLEAN EXISTS PATH '$.opt') +) AS t; ++---+------+ +| id|hasOpt| ++---+------+ +| 10| true| +| 20| false| ++---+------+ + +-- Join JSON_TABLE output against a base table using LATERAL +SELECT d.id, t.k +FROM docs d, +LATERAL JSON_TABLE(d.doc, '$.tags[*]' COLUMNS (k STRING PATH '$.k')) AS t; +``` + +### Related Statements + +* [SELECT](sql-ref-syntax-qry-select.html) +* [Table-valued Function](sql-ref-syntax-qry-select-tvf.html) +* [LATERAL VIEW Clause](sql-ref-syntax-qry-select-lateral-view.html) diff --git a/docs/sql-ref-syntax-qry-select-json-value.md b/docs/sql-ref-syntax-qry-select-json-value.md new file mode 100644 index 0000000000000..6f45a80417b51 --- /dev/null +++ b/docs/sql-ref-syntax-qry-select-json-value.md @@ -0,0 +1,138 @@ +--- +layout: global +title: JSON_VALUE +displayTitle: JSON_VALUE +license: | + 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. +--- + +### Description + +The `JSON_VALUE` scalar function extracts a single scalar value located by a SQL/JSON path from a +JSON document and returns it cast to the `RETURNING` type (`STRING` by default). This is the +SQL-standard way (SQL:2016) to pull an individual value out of JSON, and is commonly used to +migrate queries from other systems such as Oracle, DB2, and MySQL. Unlike +[JSON_TABLE](sql-ref-syntax-qry-select-json-table.html), which produces rows in a `FROM` clause, +`JSON_VALUE` is an expression that can appear anywhere a scalar is allowed. + +The function returns a scalar only. A path that matches an object or array is an *error* case (see +`ON ERROR`), not a value. To extract an object or array as a JSON fragment, use +[JSON_QUERY](sql-ref-syntax-qry-select-json-query.html); to produce rows from a JSON array, use +[JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) (the built-in `get_json_object` function +also extracts fragments). + +### Syntax + +```sql +JSON_VALUE ( json_expr, path + [ RETURNING data_type ] + [ empty_behavior ON EMPTY ] + [ error_behavior ON ERROR ] ) + +empty_behavior + { NULL | ERROR | DEFAULT default_expr } + +error_behavior + { NULL | ERROR | DEFAULT default_expr } +``` + +### Parameters + +* **json_expr** + + An expression that evaluates to a `STRING` containing the JSON document. A `NULL` input yields + `NULL` directly (it triggers neither the `ON EMPTY` nor the `ON ERROR` behavior). + +* **path** + + A SQL/JSON path literal that locates the value, for example `'$.a.b'` or `'$.items[0]'`. The + path must be wildcard-free, since `JSON_VALUE` returns a single value; a path containing `[*]` + is rejected at analysis time. + +* **RETURNING data_type** + + The type the extracted value is cast to. It must be a scalar (atomic) type: a string, numeric, + boolean, or datetime type. Non-atomic types (`STRUCT`, `ARRAY`, `MAP`) and `VARIANT` / `BINARY` + are not supported. If `RETURNING` is omitted, the result type is `STRING`. + +* **empty_behavior ON EMPTY** + + What to produce when `path` matches nothing: + * `NULL` (the default) returns SQL `NULL`. + * `ERROR` raises an error. + * `DEFAULT default_expr` returns `default_expr`, cast to the `RETURNING` type. + +* **error_behavior ON ERROR** + + What to produce when the extraction fails: the input is not well-formed JSON, the path matches a + non-scalar (object or array) value, or casting the matched scalar to the `RETURNING` type fails. + * `NULL` (the default) returns SQL `NULL`. + * `ERROR` raises an error. + * `DEFAULT default_expr` returns `default_expr`, cast to the `RETURNING` type. + + The cast of the matched scalar to the `RETURNING` type always follows ANSI semantics (a failed + conversion routes to `ON ERROR`), independently of the session's `spark.sql.ansi.enabled` + setting. + +A path that matches an explicit JSON `null` is a present scalar value and returns SQL `NULL` (it is +neither the `ON EMPTY` nor the `ON ERROR` case). + +### Examples + +```sql +-- Extract a scalar as STRING (the default) +SELECT json_value('{"id":7,"name":"Ada"}', '$.name'); ++-------------------------------------------+ +|json_value({"id":7,"name":"Ada"}, $.name) | ++-------------------------------------------+ +|Ada | ++-------------------------------------------+ + +-- Cast the extracted value with RETURNING +SELECT json_value('{"id":7}', '$.id' RETURNING INT) + 1; ++---------------------------------------------+ +|(json_value({"id":7}, $.id) + 1) | ++---------------------------------------------+ +|8 | ++---------------------------------------------+ + +-- A missing path defaults to NULL; supply a fallback with DEFAULT ... ON EMPTY +-- (RETURNING, when present, comes before the ON EMPTY / ON ERROR clauses) +SELECT json_value('{"id":7}', '$.missing' RETURNING INT DEFAULT -1 ON EMPTY); ++---------------------------------------------------------------+ +|json_value({"id":7}, $.missing RETURNING INT DEFAULT -1 ON EMPTY)| ++---------------------------------------------------------------+ +|-1 | ++---------------------------------------------------------------+ + +-- A non-scalar match or malformed input is an ON ERROR case +SELECT json_value('{"addr":{"city":"NYC"}}', '$.addr' DEFAULT 'n/a' ON ERROR); ++---------------------------------------------------------------+ +|json_value({"addr":{"city":"NYC"}}, $.addr DEFAULT n/a ON ERROR)| ++---------------------------------------------------------------+ +|n/a | ++---------------------------------------------------------------+ + +-- ERROR ON ERROR raises instead of returning a value +SELECT json_value('not json', '$.a' ERROR ON ERROR); +[JSON_VALUE_ON_ERROR.ERROR] ... +``` + +### Related Statements + +* [SELECT](sql-ref-syntax-qry-select.html) +* [JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) +* [Built-in Functions](sql-ref-functions-builtin.html) diff --git a/docs/sql-ref-syntax-qry-select-unnest.md b/docs/sql-ref-syntax-qry-select-unnest.md new file mode 100644 index 0000000000000..29430998874c5 --- /dev/null +++ b/docs/sql-ref-syntax-qry-select-unnest.md @@ -0,0 +1,113 @@ +--- +layout: global +title: UNNEST Clause +displayTitle: UNNEST Clause +license: | + 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. +--- + +### Description + +The `UNNEST` clause expands one or more arrays into a table that can be referenced in the +`FROM` clause, producing one row per array element. It is the ANSI SQL collection derived table +and is a standard alternative to `explode` / `LATERAL VIEW`. + +When several arrays are supplied, they are expanded in parallel: the number of output rows equals +the length of the longest array, and shorter arrays are padded with `NULL`s. A `NULL` array is +treated as an empty array and contributes no elements. + +To reference a column of another `FROM` item (a correlated array), use `UNNEST` on the right-hand +side of a `LATERAL` join. A `LEFT JOIN LATERAL ... ON true` preserves outer rows whose array is +empty or `NULL`. + +`UNNEST` is a non-reserved keyword, so it can still be used as a regular table or column name. +However, when it appears unquoted at the start of a `FROM` relation followed by `(`, it is parsed +as the `UNNEST` clause described here rather than as a call to a table-valued function named +`unnest`. To invoke such a function instead, quote the name: `` `unnest`(...) ``. + +### Syntax + +```sql +UNNEST ( expression [ , ... ] ) [ WITH ORDINALITY ] [ table_alias ] +``` + +### Parameters + +* **expression** + + One or more array-typed expressions to expand. Each array contributes one output column, + holding its element as-is (an array of structs is not expanded into one column per field). + +* **WITH ORDINALITY** + + Appends a trailing 1-based `BIGINT` column giving the position of each element. + +* **table_alias** + + Specifies a temporary name with an optional column name list. + + **Syntax:** `[ AS ] table_name [ ( column_name [ , ... ] ) ]` + +### Examples + +```sql +-- single array +SELECT * FROM UNNEST(array(10, 20, 30)); ++---+ +|col| ++---+ +| 10| +| 20| +| 30| ++---+ + +-- WITH ORDINALITY and a table alias +SELECT * FROM UNNEST(array('a', 'b')) WITH ORDINALITY AS t(value, pos); ++-----+---+ +|value|pos| ++-----+---+ +| a| 1| +| b| 2| ++-----+---+ + +-- multiple arrays are expanded in parallel and padded with NULLs +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) AS t(a, b); ++----+---+ +| a| b| ++----+---+ +| 1| 10| +| 2| 20| +|NULL| 30| ++----+---+ + +-- correlated UNNEST over a table column, via LATERAL +SELECT id, elem +FROM VALUES (1, array(10, 20)), (2, array(30)) AS data(id, arr), + LATERAL UNNEST(arr) AS t(elem); ++---+----+ +| id|elem| ++---+----+ +| 1| 10| +| 1| 20| +| 2| 30| ++---+----+ +``` + +### Related Statements + +* [SELECT](sql-ref-syntax-qry-select.html) +* [LATERAL VIEW Clause](sql-ref-syntax-qry-select-lateral-view.html) +* [Table-valued Function](sql-ref-syntax-qry-select-tvf.html) diff --git a/docs/sql-ref-syntax-qry-select.md b/docs/sql-ref-syntax-qry-select.md index 57bd4f6002606..2d01b6123ac10 100644 --- a/docs/sql-ref-syntax-qry-select.md +++ b/docs/sql-ref-syntax-qry-select.md @@ -93,7 +93,9 @@ SELECT [ hints , ... ] [ ALL | DISTINCT ] { [ [ named_expression | regex_column_ * [Pivot relation](sql-ref-syntax-qry-select-pivot.html) * [Unpivot relation](sql-ref-syntax-qry-select-unpivot.html) * [Table-value function](sql-ref-syntax-qry-select-tvf.html) + * [JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) * [Inline table](sql-ref-syntax-qry-select-inline-table.html) + * [UNNEST relation](sql-ref-syntax-qry-select-unnest.html) * [ [LATERAL](sql-ref-syntax-qry-select-lateral-subquery.html) ] ( Subquery ) * [File](sql-ref-syntax-qry-select-file.html) @@ -211,6 +213,9 @@ SELECT [ hints , ... ] [ ALL | DISTINCT ] { [ [ named_expression | regex_column_ * [Set Operators](sql-ref-syntax-qry-select-setops.html) * [TABLESAMPLE](sql-ref-syntax-qry-select-sampling.html) * [Table-valued Function](sql-ref-syntax-qry-select-tvf.html) +* [JSON_QUERY](sql-ref-syntax-qry-select-json-query.html) +* [JSON_VALUE](sql-ref-syntax-qry-select-json-value.html) +* [JSON_EXISTS](sql-ref-syntax-qry-select-json-exists.html) * [Window Function](sql-ref-syntax-qry-select-window.html) * [CASE Clause](sql-ref-syntax-qry-select-case.html) * [PIVOT Clause](sql-ref-syntax-qry-select-pivot.html) diff --git a/docs/sql-ref-syntax.md b/docs/sql-ref-syntax.md index 31ac500f594da..bc168fed60650 100644 --- a/docs/sql-ref-syntax.md +++ b/docs/sql-ref-syntax.md @@ -71,6 +71,7 @@ ability to generate logical and physical plan for a given query using * [HAVING Clause](sql-ref-syntax-qry-select-having.html) * [Hints](sql-ref-syntax-qry-select-hints.html) * [Inline Table](sql-ref-syntax-qry-select-inline-table.html) + * [UNNEST Clause](sql-ref-syntax-qry-select-unnest.html) * [File](sql-ref-syntax-qry-select-file.html) * [JOIN](sql-ref-syntax-qry-select-join.html) * [ASOF JOIN](sql-ref-syntax-qry-select-asof-join.html) @@ -82,6 +83,10 @@ ability to generate logical and physical plan for a given query using * [SORT BY Clause](sql-ref-syntax-qry-select-sortby.html) * [TABLESAMPLE](sql-ref-syntax-qry-select-sampling.html) * [Table-valued Function](sql-ref-syntax-qry-select-tvf.html) + * [JSON_QUERY](sql-ref-syntax-qry-select-json-query.html) + * [JSON_TABLE](sql-ref-syntax-qry-select-json-table.html) + * [JSON_VALUE](sql-ref-syntax-qry-select-json-value.html) + * [JSON_EXISTS](sql-ref-syntax-qry-select-json-exists.html) * [WHERE Clause](sql-ref-syntax-qry-select-where.html) * [Aggregate Function](sql-ref-syntax-qry-select-aggregate.html) * [Window Function](sql-ref-syntax-qry-select-window.html) diff --git a/docs/streaming-custom-receivers.md b/docs/streaming-custom-receivers.md index 11a52232510fd..fdde5f0c0f8e8 100644 --- a/docs/streaming-custom-receivers.md +++ b/docs/streaming-custom-receivers.md @@ -75,7 +75,7 @@ class CustomReceiver(host: String, port: Int) def onStop() { // There is nothing much to do as the thread calling receive() - // is designed to stop by itself if isStopped() returns false + // is designed to stop by itself when isStopped() returns true } /** Create a socket connection and receive data until receiver is stopped */ @@ -137,7 +137,7 @@ public class JavaCustomReceiver extends Receiver<String> { @Override public void onStop() { // There is nothing much to do as the thread calling receive() - // is designed to stop by itself if isStopped() returns false + // is designed to stop by itself when isStopped() returns true } /** Create a socket connection and receive data until receiver is stopped */ diff --git a/docs/streaming/apis-on-dataframes-and-datasets.md b/docs/streaming/apis-on-dataframes-and-datasets.md index 86585caead51f..e946a8e38b891 100644 --- a/docs/streaming/apis-on-dataframes-and-datasets.md +++ b/docs/streaming/apis-on-dataframes-and-datasets.md @@ -1851,7 +1851,7 @@ Here are the configs regarding to RocksDB instance of the state store provider: </tr> <tr> <td>spark.sql.streaming.stateStore.rocksdb.maxOpenFiles</td> - <td>The number of open files that can be used by the RocksDB instance. Value of -1 means that files opened are always kept open. If the open file limit is reached, RocksDB will evict entries from the open file cache and close those file descriptors and remove the entries from the cache.</td> + <td>The number of open files that can be used by the RocksDB instance. Value of -1 means that files opened are always kept open. If the open file limit is reached, RocksDB will evict entries from the open file cache and close those file descriptors and remove the entries from the cache. Set this configuration to a positive value if a query fails when opening RocksDB with <code>Too many open files</code>.</td> <td>-1</td> </tr> <tr> diff --git a/docs/streaming/real-time-mode.md b/docs/streaming/real-time-mode.md index 51a8ba04c0737..d7d5a67fcc9e0 100644 --- a/docs/streaming/real-time-mode.md +++ b/docs/streaming/real-time-mode.md @@ -25,20 +25,27 @@ license: | # Real-time Mode **Real-time Mode** is a new streaming execution mode introduced in Spark 4.1.0 that -targets ultra-low end-to-end latency with the exact same API and processing guarantees / semantics as the current structured streaming engine. +targets ultra-low end-to-end latency using the Structured Streaming APIs and processing +guarantees. Some scheduling and callback details differ from micro-batch execution and are +documented below. It is intended for operational workloads that must react to data the moment it arrives, such as fraud detection, real-time alerting, and live personalization. -In this release, Real-time Mode in Apache Spark supports **stateless queries** only -- projections, -filters and other map-like operations, unions, and stream-static joins. Stateful operations such as -streaming aggregations, deduplication, stream-stream joins, and `transformWithState` are not yet -supported, but support for them is planned starting in Spark 4.3. See +Real-time Mode in Apache Spark supports **stateless queries** -- projections, filters and other +map-like operations, unions, and stream-static joins -- and, starting in Spark 4.3.0, a first set of +**stateful queries**: streaming **deduplication** (`dropDuplicates`) and streaming **aggregations** +(`groupBy(...).agg(...)`), and the JVM (Scala/Java) **`transformWithState`** operator. These stateful +operations require a shuffle, which Real-time Mode runs as a *pipelined shuffle* so that records +still stream through without waiting for a batch boundary; see +[How Stateful Queries Work](#how-stateful-queries-work). Other stateful operations, including +stream-stream joins and `flatMapGroupsWithState`, are not yet supported. See [Supported Queries](#supported-queries) for the full list. -The most important thing to know: **the duration you pass to the trigger (default 5 minutes) is a -checkpoint interval, not a latency target.** Records are processed and emitted continuously rather -than at batch boundaries, so the trigger duration does not set latency the way a micro-batch interval does. See +The most important thing to know: **the duration you pass to the trigger (default 5 minutes) is +primarily a checkpoint interval, not the latency target for input-driven output.** Records can be +processed and emitted continuously rather than waiting for batch boundaries. Batch-scoped effects, +such as watermark advancement and timers on idle partitions, can still wait for the boundary. See [Batch Duration Is a Checkpoint Interval](#batch-duration-is-a-checkpoint-interval). You enable Real-time Mode by setting a Real-time trigger on the streaming write; the rest of your @@ -62,12 +69,101 @@ the time needed to process and ship one record (often a few milliseconds). Since records never wait for a batch boundary, the batch duration mainly controls how often the query checkpoints progress -- as the next section explains. +## How Stateful Queries Work + +Stateless operations are per-record: each long-running task reads a partition, transforms records, +and ships them without ever needing data from another partition. Stateful operations are different. +A streaming aggregation, `dropDuplicates`, or `transformWithState` groups records **by key**, so +every record for a given key must reach the same task, no matter which partition it arrived on. In +the micro-batch engine that regrouping is done by a **shuffle**: the batch's producer stage writes +shuffle files, and only once those files are fully materialized does the consumer stage read them +back, grouped by key. + +That "materialize, then read" boundary is exactly what Real-time Mode avoids for latency, and a +long-running Real-time task never finishes its batch, so an ordinary shuffle would deadlock -- the +consumer would wait forever for a producer that never completes. Real-time Mode therefore runs the +shuffle differently, using two cooperating pieces: + +- **Pipelined shuffle** changes *scheduling*. Normally the scheduler runs the consumer stage only + after the producer stage has completed. For a Real-time stateful query the producer (the source + scan) and the consumer (the stateful operator) are marked as a single **pipelined group** and + scheduled to run **at the same time**, for the whole batch. Neither stage waits for the other to + finish; they run concurrently as long-running tasks, just like the stateless case. + +- **Streaming shuffle** changes *data transport*. Because the two stages run concurrently, shuffle + data cannot be written to files and read back after the fact. Instead the producer's output is + streamed directly to the consumer tasks over the network, record by record, so a record is + regrouped by key and handed to the stateful operator as soon as it is produced -- there is no + intermediate file and no wait for the batch to end. + +Together these let a stateful Real-time query keep the same continuous, per-record flow as a +stateless one: a record is read, routed to the task that owns its key through the streaming shuffle, +merged into state, and emitted -- all without a batch boundary. The regrouped state itself is kept +in Spark's usual state store and checkpointed each batch, so the exactly-once and recovery +guarantees are the same as the micro-batch engine (see [Fault Tolerance](#fault-tolerance)). + +This mechanism is enabled automatically in Real-time Mode; there is nothing to configure. It applies +to every shuffle on the streaming path -- the shuffle a stateful operator needs, and also a bare +`repartition` (a hash or round-robin shuffle with no stateful operator), which runs as a pipelined +shuffle in the same way. The one exception is a shuffle that would require a separate preparatory +job: range partitioning (`repartitionByRange`, or an `ORDER BY` that plans to a range shuffle) needs +a sampling job to compute range bounds, and that job cannot complete while the source keeps +producing, so such a query fails to start with +`STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST`. + +## `transformWithState` in Real-time Mode + +Starting in Spark 4.3.0, the JVM `transformWithState` API can run in Real-time Mode. The Scala and +Java APIs are supported; the PySpark `transformWithState` and `transformWithStateInPandas` APIs are +not. `TimeMode.None`, `TimeMode.ProcessingTime`, and `TimeMode.EventTime` are supported. See the +[`transformWithState` guide](./structured-streaming-transform-with-state.html) for the stateful +processor API, state variables, timers, TTL, and initial state. `TimeMode.EventTime` requires an +input event-time watermark declared with `withWatermark`. + +The API is the same in micro-batch and Real-time Mode, but the input callback granularity differs. +In micro-batch mode, one `handleInputRows` invocation receives all input rows for a grouping key in +that batch. In Real-time Mode, Spark invokes `handleInputRows` once for each non-late input row, with +a single row in the iterator. A processor used in both modes must therefore work correctly whether +rows for the same key arrive in one invocation or in repeated invocations. + +Time-based operations also run incrementally while the long Real-time batch remains open: + +- **Processing-time timers** use the current executor clock during the long-running input batch. + Spark checks for expired timers after each input row reaches the state partition and once more at + batch completion. +- **Event-time timers** require an input watermark declared with `withWatermark` and use the + watermark established at the beginning of the batch. That watermark remains fixed for the whole + batch and advances only between batches. Timers already expired against it are checked after each + input row and again at batch completion. +- **TTL state** requires `TimeMode.ProcessingTime`. Expired values are not returned when the state + is accessed. Spark also removes expired values periodically while input rows are processed and + performs a final cleanup at batch completion. + +Timer checks and TTL cleanup are driven by input or batch completion; there is no independent +background polling while a state partition is idle. A processing-time timer on an idle partition +can therefore wait until another row arrives or the current batch completes. Similarly, event-time +watermark progress is bounded by the Real-time batch duration. + +When initial state is provided, Spark loads and commits it in a finite bootstrap batch before +starting the first long-running Real-time input batch. Input that is already available waits until +the initial state is durable. The bootstrap uses a regular shuffle; pipelined shuffle begins with +the following input batch. `TimerValues.getCurrentProcessingTimeInMs()` returns the finite bootstrap +batch timestamp while `handleInitialState` is running; it uses the live executor clock after the +long-running input batch starts. + +State variables, TTL information, and registered timers are checkpointed and restored after a +restart. A query can resume the same compatible checkpoint in micro-batch or Real-time Mode when it +uses RocksDB and state-store checkpoint format v2; see +[State store defaults](#state-store-defaults). + ## Batch Duration Is a Checkpoint Interval -In Real-time Mode, the batch duration is a **checkpoint interval, not a latency interval.** With the -default 5-minute duration, the query still emits records within milliseconds; the 5 minutes only -controls how often it commits progress and starts the next long-running batch. This is the opposite -of the micro-batch engine, where a longer batch interval directly increases latency. +In Real-time Mode, the batch duration is primarily a **checkpoint interval, not the latency +interval for input-driven output.** With the default 5-minute duration, a query can still emit +results produced from input records within milliseconds. The duration controls how often it commits +progress and starts the next long-running batch. It can also bound the delay for batch-scoped work, +including watermark advancement and processing-time timers on idle partitions. This differs from +the micro-batch engine, where all output waits for the batch interval. Do not confuse the 5-minute default trigger duration with the 5-second minimum allowed duration described under [Requirements](#requirements): the former is the checkpoint cadence used when you do @@ -85,11 +181,13 @@ Choosing the batch duration is a trade-off: The duration is set on the Real-time trigger, as shown under [Enabling Real-time Mode](#enabling-real-time-mode). -Progress is committed using **asynchronous progress tracking**: the offset and commit logs are -written off the record-processing path so that checkpointing does not stall processing. This is -enabled automatically for Real-time Mode queries and every batch is checkpointed (the async -progress tracking checkpoint interval is fixed at 0 in Real-time Mode). It can be turned off with -the `asyncProgressTrackingEnabled` writer option, in which case progress is committed synchronously. +For stateless Real-time queries, progress is committed using **asynchronous progress tracking**: +the offset and commit logs are written off the record-processing path so that checkpointing does +not stall processing. It is enabled automatically for stateless Real-time queries, and every batch +is checkpointed (the async progress tracking checkpoint interval is fixed at 0 in Real-time Mode). +It can be turned off with the `asyncProgressTrackingEnabled` writer option. Stateful queries, +including `transformWithState`, do not support asynchronous progress tracking and commit progress +synchronously. ## Comparison with Other Modes @@ -100,8 +198,8 @@ experimental [Continuous Processing](./performance-tips.html#continuous-processi | Mode | Latency | Processing Guarantees | Supported operations | When to use | |---|---|---|---|---| -| Micro-batch (default) | ~100 ms | Exactly-once | All streaming operations, including stateful | Stateful or higher-throughput workloads, or queries Real-time Mode does not yet support | -| Real-time Mode | millisecond-scale | Exactly-once | Stateless today (map-like operations, unions, and stream-static joins); designed to support all query shapes, including stateful | Low-latency workloads | +| Micro-batch (default) | ~100 ms | Exactly-once | All streaming operations, including all stateful ones | Stateful or higher-throughput workloads, or queries Real-time Mode does not yet support | +| Real-time Mode | millisecond-scale | Exactly-once | Stateless operations (map-like operations, unions, and stream-static joins) plus stateful deduplication, aggregation, and JVM `transformWithState`; more stateful operations planned | Low-latency workloads | | Continuous Processing (experimental) | ~1 ms | At-least-once | Map-like only (projections and selections); no stateful operations | Legacy; use Real-time Mode instead | The **Processing Guarantees** column refers to processing semantics, defined under @@ -118,8 +216,9 @@ substantially: apply to it. These constraints have limited its adoption. - **Real-time Mode** is designed to support all query shapes, including stateful operations, while reusing Spark's mature components such as state management, the Catalyst optimizer, and the - existing SQL operators. It provides exactly-once processing semantics. It currently supports - stateless queries, with stateful support planned starting in Spark 4.3. + existing SQL operators. It provides exactly-once processing semantics. It supports stateless + queries and, starting in Spark 4.3.0, stateful deduplication, aggregation, and JVM + `transformWithState`; support for the remaining stateful operations is ongoing. For new low-latency workloads, prefer Real-time Mode over Continuous Processing. @@ -225,14 +324,17 @@ the query starts: interval, and month-based intervals (for example, `"1 month"`) are not accepted. (This 5-second minimum is distinct from the 5-minute default; see [Batch Duration Is a Checkpoint Interval](#batch-duration-is-a-checkpoint-interval).) +- Stateful queries do not support asynchronous progress tracking. Do not set the + `asyncProgressTrackingEnabled` writer option to `true` for a query with a stateful operator. ## Supported Queries -Real-time Mode supports stateless, map-like queries only. +Real-time Mode supports stateless, map-like queries and a first set of stateful queries: +deduplication, aggregation, and JVM `transformWithState`. -The following operations, sources, and sinks are supported as of Spark 4.1.0: +The following operations, sources, and sinks are supported: -- *Operations*: stateless, map-like operations are supported: +- *Stateless operations* (supported since Spark 4.1.0): + Projections: `select`, `selectExpr`, `withColumn`, `drop`, and the typed `map` / `flatMap` Dataset operations. + Selections: `where` / `filter`. + Expressions that compile to a projection -- including functions such as `from_json` / `to_json` @@ -241,13 +343,30 @@ The following operations, sources, and sinks are supported as of Spark 4.1.0: + `union` of two or more *distinct* streaming sources. Referencing the same source DataFrame more than once is not supported and fails with `STREAMING_REAL_TIME_MODE.IDENTICAL_SOURCES_IN_UNION_NOT_SUPPORTED`; create a separate DataFrame - for each source instead. + for each source instead. A union may feed a stateful operator, but a stateful operator cannot + appear on an input branch before the union; that shape fails with + `STREAMING_REAL_TIME_MODE.STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED`. + Stream-static joins, where a streaming DataFrame is joined with a static DataFrame. The static - side must be broadcast (use the `broadcast(...)` hint), because Real-time Mode does not support - shuffles. - + `withWatermark` (event-time watermark declaration) is allowed, although it has no effect because - stateful operators are not supported. This lets queries that already declare a watermark run in - Real-time Mode without modification. + side must be broadcast (use the `broadcast(...)` hint), because a stream-static join must not + introduce a shuffle. + +- *Stateful operations* (supported since Spark 4.3.0): these regroup records by key through a + pipelined shuffle (see [How Stateful Queries Work](#how-stateful-queries-work)) and keep their + state in the state store, checkpointed each batch. + + **Deduplication**: `dropDuplicates`. (`dropDuplicatesWithinWatermark` is not yet supported in + Real-time Mode.) + + **Streaming aggregation**: `groupBy(...).agg(...)` (and the SQL `GROUP BY` equivalent), including + windowed aggregations with `window(...)`. Distinct aggregates such as `count(distinct ...)` are + not supported (see [Not supported](#not-supported)). + + **JVM `transformWithState`**: the Scala and Java APIs, including value, list, and map state; + processing-time TTL; processing-time and event-time timers; optional initial state; and output + event-time columns. See [`transformWithState` in Real-time Mode](#transformwithstate-in-real-time-mode) + for its incremental execution and timing semantics. + + `withWatermark` (event-time watermark declaration) is supported and now takes effect: it lets a + windowed aggregation drop late input and evict the state for windows that have closed, bounding + how much state a long-running query accumulates. (Real-time Mode always runs in `update` output + mode -- see [Requirements](#requirements) -- so a windowed aggregation emits each window's + running result as it changes, and the watermark governs when that window's state is evicted.) - *Sources*: the source must support Real-time Mode. In Apache Spark, the **Kafka** source supports Real-time Mode. An unsupported source fails with @@ -270,11 +389,24 @@ starts; anything outside the allowlist fails with ### Not supported -Stateful operations of any kind are not supported in this release. This includes streaming -aggregations, `dropDuplicates` / `dropDuplicatesWithinWatermark`, stream-stream joins, `repartition` -and other operations that introduce a shuffle, and stateful operators such as -`flatMapGroupsWithState` and `transformWithState`. Support for stateful operations is planned -starting in Spark 4.3. +The following are not yet supported in Real-time Mode. Unless noted otherwise, a query that uses one +fails to start with `STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST`: + +- Stateful operations other than those listed above: **stream-stream joins**, + `flatMapGroupsWithState`, session-window aggregation, and `dropDuplicatesWithinWatermark`. +- The PySpark `transformWithState` and `transformWithStateInPandas` APIs. Real-time Mode currently + supports only the JVM (Scala/Java) `transformWithState` implementation. +- **Range partitioning**: `repartitionByRange`, or an `ORDER BY` / sort that plans to a range + shuffle, because computing range bounds needs a separate sampling job that cannot complete while + the source keeps producing. (A plain `repartition` -- hash or round-robin -- is supported; it runs + as a pipelined shuffle. See [How Stateful Queries Work](#how-stateful-queries-work).) + +**Distinct aggregates** such as `count(distinct ...)` are not supported either, but this is a +general Structured Streaming restriction rather than a Real-time Mode one: any streaming distinct +aggregate is rejected during analysis (with a message suggesting `approx_count_distinct()`), +regardless of the trigger. + +Support for more stateful operations is ongoing. ## Fault Tolerance @@ -299,8 +431,9 @@ tolerate duplicates -- for example, with idempotent writes -- where exactly-once ## Examples -The following examples read from Kafka and assume a running Kafka cluster. Each example shows the -same query in Python, Scala, and Java. +The following examples read from Kafka and assume a running Kafka cluster. Most show the same query +in Python, Scala, and Java. The `transformWithState` example is shown in Scala and also applies to +Java; its PySpark APIs are not yet supported in Real-time Mode. ### Stream-static join @@ -391,6 +524,237 @@ spark </div> +### Deduplication + +Drop duplicate records by key. This is a stateful operation: Real-time Mode regroups records by the +deduplication key through a pipelined shuffle and keeps the set of seen keys in the state store (see +[How Stateful Queries Work](#how-stateful-queries-work)). No code changes are needed beyond running +under a Real-time trigger. + +<div class="codetabs"> + +<div data-lang="python" markdown="1"> +{% highlight python %} +spark \ + .readStream \ + .format("kafka") \ + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") \ + .option("subscribe", "input-topic") \ + .load() \ + .selectExpr("CAST(key AS STRING) AS id", "CAST(value AS STRING) AS value") \ + .dropDuplicates("id") \ + .writeStream \ + .format("kafka") \ + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") \ + .option("topic", "output-topic") \ + .option("checkpointLocation", "/path/to/checkpoint") \ + .outputMode("update") \ + .trigger(realTime="5 minutes") \ + .start() +{% endhighlight %} +</div> + +<div data-lang="scala" markdown="1"> +{% highlight scala %} +import org.apache.spark.sql.streaming.Trigger + +spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("subscribe", "input-topic") + .load() + .selectExpr("CAST(key AS STRING) AS id", "CAST(value AS STRING) AS value") + .dropDuplicates("id") + .writeStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start() +{% endhighlight %} +</div> + +<div data-lang="java" markdown="1"> +{% highlight java %} +import org.apache.spark.sql.streaming.Trigger; + +spark + .readStream() + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("subscribe", "input-topic") + .load() + .selectExpr("CAST(key AS STRING) AS id", "CAST(value AS STRING) AS value") + .dropDuplicates("id") + .writeStream() + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start(); +{% endhighlight %} +</div> + +</div> + +This keeps every distinct key it has seen in the state store. `dropDuplicatesWithinWatermark`, which +bounds how long keys are retained, is not yet supported in Real-time Mode. + +### Streaming aggregation + +Maintain a running aggregate per key. Real-time Mode regroups input by the grouping key through a +pipelined shuffle, merges each record into the running aggregate in the state store, and emits the +updated result. Because the output mode is `update`, each key is emitted as it changes rather than +only at the end of the batch. + +<div class="codetabs"> + +<div data-lang="python" markdown="1"> +{% highlight python %} +from pyspark.sql.functions import count + +spark \ + .readStream \ + .format("kafka") \ + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") \ + .option("subscribe", "input-topic") \ + .load() \ + .selectExpr("CAST(key AS STRING) AS id") \ + .groupBy("id") \ + .agg(count("*").alias("cnt")) \ + .selectExpr("CAST(id AS STRING) AS key", "CAST(cnt AS STRING) AS value") \ + .writeStream \ + .format("kafka") \ + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") \ + .option("topic", "output-topic") \ + .option("checkpointLocation", "/path/to/checkpoint") \ + .outputMode("update") \ + .trigger(realTime="5 minutes") \ + .start() +{% endhighlight %} +</div> + +<div data-lang="scala" markdown="1"> +{% highlight scala %} +import org.apache.spark.sql.functions.count +import org.apache.spark.sql.streaming.Trigger + +spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("subscribe", "input-topic") + .load() + .selectExpr("CAST(key AS STRING) AS id") + .groupBy("id") + .agg(count("*").as("cnt")) + .selectExpr("CAST(id AS STRING) AS key", "CAST(cnt AS STRING) AS value") + .writeStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start() +{% endhighlight %} +</div> + +<div data-lang="java" markdown="1"> +{% highlight java %} +import static org.apache.spark.sql.functions.count; +import org.apache.spark.sql.streaming.Trigger; + +spark + .readStream() + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("subscribe", "input-topic") + .load() + .selectExpr("CAST(key AS STRING) AS id") + .groupBy("id") + .agg(count("*").as("cnt")) + .selectExpr("CAST(id AS STRING) AS key", "CAST(cnt AS STRING) AS value") + .writeStream() + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start(); +{% endhighlight %} +</div> + +</div> + +### JVM `transformWithState` + +Run a JVM stateful processor continuously under a Real-time trigger. This processor maintains a +running total for each Kafka key. It sums the iterator so the same implementation works when a +micro-batch invocation contains several rows and when a Real-time invocation contains one row. See +the [`transformWithState` guide](./structured-streaming-transform-with-state.html) for the complete +API. Java applications use the same `TimeMode`, `OutputMode`, and `Trigger.RealTime` settings; the +Java `transformWithState` overload also takes an output encoder. + +{% highlight scala %} +import org.apache.spark.sql.Encoders +import org.apache.spark.sql.streaming.{ + OutputMode, StatefulProcessor, TTLConfig, TimeMode, TimerValues, Trigger, ValueState} + +import spark.implicits._ + +class RunningTotalProcessor + extends StatefulProcessor[String, (String, Int), String] { + @transient private var total: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + total = getHandle.getValueState("total", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[String] = { + val previous = if (total.exists()) total.get() else 0L + val updated = previous + inputRows.map(_._2.toLong).sum + total.update(updated) + Iterator.single(s"$key,$updated") + } +} + +val totals = spark + .readStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("subscribe", "input-topic") + .load() + .selectExpr( + "CAST(key AS STRING) AS id", + "CAST(CAST(value AS STRING) AS INT) AS delta") + .as[(String, Int)] + .groupByKey(_._1) + .transformWithState( + statefulProcessor = new RunningTotalProcessor, + timeMode = TimeMode.ProcessingTime(), + outputMode = OutputMode.Update()) + +totals + .writeStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start() +{% endhighlight %} + ### Writing to the console for development The console sink prints output to the driver's standard output and is handy while developing a @@ -471,6 +835,44 @@ spark |---|---|---| | `spark.sql.streaming.realTimeMode.minBatchDuration` | `5000` (ms, 5 seconds) | The minimum batch duration, in milliseconds, allowed for a Real-time trigger. See the batch-duration requirement under [Requirements](#requirements). | | `spark.sql.streaming.realTimeMode.allowlistCheck` | `true` | Whether to verify that all operators and sinks used by a Real-time query are in the supported allowlist. Disabling this check (not recommended) lets unsupported operators and sinks run at your own risk. | +| `spark.sql.streaming.realTimeMode.dangerouslyAllowCheckpointV1.enabled` | `false` | Whether to allow a Real-time query to use state-store checkpoint format version 1. This is unsafe for stateful queries: format v1 can reuse state-file names when a failed batch is rerun, so the rerun can load stale state and lose updates. Prefer format v2 and a fresh checkpoint location. See [State store defaults](#state-store-defaults). | + +### State store defaults + +A stateful Real-time query needs a low-latency, recovery-correct state store configuration. Because +that configuration is not the right default for the engine as a whole, Real-time Mode applies it +automatically at query start, only for Real-time queries. These are **soft defaults**: each is set +only when you have not set the config yourself, so an explicit value is preserved -- with the +exception of a few explicit values that are incompatible with Real-time Mode and are rejected at +query start rather than kept (see [Incompatible configurations](#incompatible-configurations)). + +| Configuration | Real-time default | Meaning | +|---|---|---| +| `spark.sql.streaming.stateStore.providerClass` | `RocksDBStateStoreProvider` | Real-time Mode defaults to the RocksDB state store, which checkpoint format v2 (below) requires. | +| `spark.sql.streaming.stateStore.checkpointFormatVersion` | `2` | Format v2 gives each batch its own state store checkpoint ids, which is what lets a failed batch be rerun correctly from committed offsets. Real-time Mode requires v2 (see below). | +| `spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled` | `true` | Writes a changelog instead of a full snapshot at each commit, shortening the state-commit step that sits on the critical path between Real-time batches. Applied only when the state store is RocksDB. | +| `spark.sql.execution.sortBeforeRepartition` | `false` | The local sort inserted before a round-robin repartition never drains an unbounded stream and would hang a Real-time query, so Real-time Mode defaults it off. Determinism from the sort is not needed because Real-time Mode does not retry tasks. Like the others this is a soft default -- but an explicit `true` is incompatible, so rather than being kept it is rejected at query start (see [Incompatible configurations](#incompatible-configurations)). | + +A Real-time query requires state-store checkpoint format v2. Starting a Real-time query with +format version 1 -- for example, when switching an existing micro-batch query to Real-time Mode, or +when +`spark.sql.streaming.stateStore.checkpointFormatVersion` is pinned to `1` -- fails to start with +`STREAMING_REAL_TIME_MODE.CHECKPOINT_FORMAT_V1_NOT_SUPPORTED`. Use a fresh checkpoint location, or, +accepting the risk of state loss on failure, set +`spark.sql.streaming.realTimeMode.dangerouslyAllowCheckpointV1.enabled=true`. + +### Incompatible configurations + +The defaults above are applied only when you have not set the config. If you instead set one of the +following to a value that Real-time Mode cannot run with, the query fails to start with +`STREAMING_REAL_TIME_MODE.SQL_CONFIGURATION_NOT_SUPPORTED` rather than having your value silently +overridden: + +- `spark.sql.streaming.stateStore.checkpointFormatVersion` set below `2` (unless + `spark.sql.streaming.realTimeMode.dangerouslyAllowCheckpointV1.enabled=true`). +- `spark.sql.streaming.stateStore.providerClass` set to a provider other than + `RocksDBStateStoreProvider`. +- `spark.sql.execution.sortBeforeRepartition` set to `true`. ## Best Practices @@ -481,6 +883,13 @@ spark topic with 10 partitions requires at least 10 cores for the query to make progress. Real-time Mode uses a fixed 1:1 mapping between Kafka topic partitions and reader tasks; the `minPartitions` option is not supported in Real-time Mode. +- Stateful queries need cores for **both** stages at once. A stateful Real-time query runs its + producer stage (the source scan) and its consumer stage (the stateful operator) concurrently as a + pipelined group (see [How Stateful Queries Work](#how-stateful-queries-work)), and both hold their + tasks for the whole batch. Size the cluster for the sum: the source's reader tasks plus the + stateful operator's tasks (`spark.sql.shuffle.partitions` post-shuffle tasks by default). For + example, a 10-partition Kafka source feeding an aggregation with 5 shuffle partitions needs at + least 15 cores to make progress. - Run a single Real-time query per cluster. Because Real-time Mode holds its task slots for the entire batch duration, any other queries sharing the cluster compete for the same slots, which can starve the Real-time query of resources and increase its latency. diff --git a/docs/streaming/structured-streaming-kafka-integration.md b/docs/streaming/structured-streaming-kafka-integration.md index 364fc77f1ef1b..5be6dfa741739 100644 --- a/docs/streaming/structured-streaming-kafka-integration.md +++ b/docs/streaming/structured-streaming-kafka-integration.md @@ -403,14 +403,24 @@ The following configurations are optional: <tr> <td>startingOffsets</td> <td>"earliest", "latest" (streaming only), or json string - """ {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} """ + """ {"topicA":{"0":23,"1":-1},"topicB":{"0":-2}} """ or + """ {"topicA":"earliest","topicB":"latest"} """ or a mix of both forms + """ {"topicA":"earliest","topicB":{"0":23,"1":-1},"topicC":"latest"} """ </td> <td>"latest" for streaming, "earliest" for batch</td> <td>streaming and batch</td> <td>The start point when a query is started, either "earliest" which is from the earliest offsets, "latest" which is just from the latest offsets, or a json string specifying a starting offset for each TopicPartition. In the json, -2 as an offset can be used to refer to earliest, -1 to latest. - Note: For batch queries, latest (either implicitly or by using -1 in json) is not allowed. + A topic may also map to the string "earliest" or "latest" instead of an object, which applies that + offset to every partition of the topic without having to enumerate them; the two forms can be mixed + in the same json. Topic-level values are expanded against the partitions discovered when the offsets + are resolved, so they keep working when a topic is repartitioned. As with the per-partition form, + the json must account for every topic subscribed at the time the offsets are resolved, so with + <code>subscribePattern</code> the topics matched by the pattern have to be known then; topics + matched later on are discovered as new partitions and start at earliest, as usual. + Note: For batch queries, latest (either implicitly, by using -1 in json, or by binding a topic to + "latest") is not allowed. For streaming queries, this only applies when a new query is started, and that resuming will always pick up from where the query left off. Newly discovered partitions during a query will start at earliest.</td> @@ -442,13 +452,16 @@ The following configurations are optional: <tr> <td>endingOffsets</td> <td>latest or json string - {"topicA":{"0":23,"1":-1},"topicB":{"0":-1}} + {"topicA":{"0":23,"1":-1},"topicB":{"0":-1}} or {"topicA":"latest","topicB":"latest"} + or a mix of both forms {"topicA":"latest","topicB":{"0":23,"1":-1}} </td> <td>latest</td> <td>batch query</td> <td>The end point when a batch query is ended, either "latest" which is just referred to the latest, or a json string specifying an ending offset for each TopicPartition. In the json, -1 - as an offset can be used to refer to latest, and -2 (earliest) as an offset is not allowed.</td> + as an offset can be used to refer to latest, and -2 (earliest) as an offset is not allowed. + As with <code>startingOffsets</code>, a topic may map to the string "latest" instead of an object to + apply that offset to every partition of the topic; binding a topic to "earliest" is not allowed.</td> </tr> <tr> <td>failOnDataLoss</td> @@ -1045,7 +1058,7 @@ For experimenting on `spark-shell`, you can also use `--packages` to add `spark- ./bin/spark-shell --packages org.apache.spark:spark-sql-kafka-0-10_{{site.SCALA_BINARY_VERSION}}:{{site.SPARK_VERSION_SHORT}} ... -See [Application Submission Guide](submitting-applications.html) for more details about submitting +See [Application Submission Guide](../submitting-applications.html) for more details about submitting applications with external dependencies. ## Security @@ -1226,7 +1239,18 @@ For possible Kafka parameters, see [Kafka adminclient config docs](http://kafka. #### Caveats -- Obtaining delegation token for proxy user is not yet supported ([KAFKA-6945](https://issues.apache.org/jira/browse/KAFKA-6945)). +- Obtaining delegation token for [proxy user](../security.html#proxy-user) requires Kafka broker 3.3.0 or higher + ([KAFKA-6945](https://issues.apache.org/jira/browse/KAFKA-6945)). The token is requested with the real user's + credentials but owned by the proxy user, therefore the real user must be granted the `CreateTokens` operation + on the `User:<proxy user>` resource, such as, + + ./bin/kafka-acls.sh --bootstrap-server <KAFKA_SERVERS> --add \ + --allow-principal User:<real user> --operation CreateTokens --user-principal "User:<proxy user>" + + Since `--proxy-user` cannot be combined with `--principal`/`--keytab`, the real user's Kerberos credentials + come from the ticket cache, and the token is obtained once at startup and not renewed: the application must + finish before the token's max lifetime, unless direct credential providers + (`spark.security.directCredentialProviders.enabled`) with non-Kerberos Kafka authentication are used. ### JAAS login configuration diff --git a/docs/streaming/structured-streaming-transform-with-state.md b/docs/streaming/structured-streaming-transform-with-state.md index c74c3a0909c31..184f61b18a0b3 100644 --- a/docs/streaming/structured-streaming-transform-with-state.md +++ b/docs/streaming/structured-streaming-transform-with-state.md @@ -29,6 +29,11 @@ This operator has support for an umbrella of features such as object-oriented st `TransformWithState` is available in Scala, Java and Python. +Starting in Spark 4.3.0, the Scala and Java APIs can also run with a Real-time trigger. The PySpark +`transformWithState` and `transformWithStateInPandas` APIs are not supported in Real-time Mode. See +[`transformWithState` in Real-time Mode](./real-time-mode.html#transformwithstate-in-real-time-mode) +for callback, timer, TTL, watermark, and initial-state behavior in that mode. + Note that in Python, there are two operators named `transformWithStateInPandas` which works with Pandas interface, and `transformWithState` which works with Row interface. Based on popularity of Pandas and its rich set of API with vectorization, `transformWithStateInPandas` may be the preferred API for most users. The `transformWithState` API is more suitable to handle high key cardinality use case, since the cost of conversion is considerably high for Pandas API. If users aren't familiar with Pandas, Row type API might be easier to learn. @@ -99,9 +104,15 @@ In Scala, implicit encoders can be provided for case classes and primitive types State variables can be configured with an optional TTL (Time-To-Live) value. The TTL value is used to automatically evict the state variable after the specified duration. The TTL value can be provided as a Duration. +TTL can be used only with `TimeMode.ProcessingTime`. + ### Handling input rows -The `handleInputRows` method is used to process input rows belonging to a grouping key and emit output if needed. The method is invoked by the Spark query engine for each grouping key value received by the operator. If multiple rows belong to the same grouping key, the provided iterator will include all those rows. +The `handleInputRows` method is used to process input rows belonging to a grouping key and emit +output if needed. In micro-batch mode, the method is invoked once for each grouping key received in +the batch, and the provided iterator includes all rows for that key. In Real-time Mode, the method +is invoked once for each non-late input row, and the iterator contains one row. A processor that can +run in either mode must handle repeated invocations for the same grouping key. ### Handling expired timers @@ -287,15 +298,15 @@ class DowntimeDetector(duration: Duration) extends key: String, timerValues: TimerValues, expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, Duration)] = { - val latestTimestamp = _lastSeen.get() - val downtimeDuration = new Duration( - timerValues.getCurrentProcessingTimeInMs() - latestTimestamp.getTime) + val latestTimestamp = _lastSeen.get() + val downtimeDuration = Duration.ofMillis( + timerValues.getCurrentProcessingTimeInMs() - latestTimestamp.getTime) - // Register another timer that will fire in 10 seconds. - // Timers can be registered anywhere but init() - getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + 10000) + // Register another timer that will fire in 10 seconds. + // Timers can be registered anywhere but init() + getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + 10000) - Iterator((key, downtimeDuration)) + Iterator((key, downtimeDuration)) } } @@ -319,7 +330,7 @@ q = (df.groupBy("key") statefulProcessor=DownTimeDetector(), outputStructType=output_schema, outputMode="Update", - timeMode="None", + timeMode="ProcessingTime", ) .writeStream... @@ -335,7 +346,7 @@ q = (df.groupBy("key") statefulProcessor=DownTimeDetector(), outputStructType=output_schema, outputMode="Update", - timeMode="None", + timeMode="ProcessingTime", ) .writeStream... @@ -345,16 +356,53 @@ q = (df.groupBy("key") <div data-lang="scala" markdown="1"> {% highlight scala %} -val query = df.groupBy("key") +val query = df.as[(String, Timestamp)] + .groupByKey(_._1) .transformWithState( - statefulProcessor = new DownTimeDetector(), - outputMode = OutputMode.Update, - timeMode = TimeMode.None) + statefulProcessor = new DowntimeDetector(Duration.ofSeconds(5)), + outputMode = OutputMode.Update(), + timeMode = TimeMode.ProcessingTime()) .writeStream... {% endhighlight %} </div> </div> +## Running in Real-time Mode + +For a Scala or Java `transformWithState` query, select `OutputMode.Update` and set +`Trigger.RealTime` on the streaming write. The query must use a +[supported source, sink, and plan](./real-time-mode.html#supported-queries), satisfy the +[Real-time Mode requirements](./real-time-mode.html#requirements), and use the required +[state-store configuration](./real-time-mode.html#state-store-defaults). For example, after +constructing a result `Dataset` of `(key, value)` tuples with `transformWithState`: + +{% highlight scala %} +import org.apache.spark.sql.streaming.Trigger + +val query = result + .selectExpr("CAST(_1 AS STRING) AS key", "CAST(_2 AS STRING) AS value") + .writeStream + .format("kafka") + .option("kafka.bootstrap.servers", "host1:port1,host2:port2") + .option("topic", "output-topic") + .option("checkpointLocation", "/path/to/checkpoint") + .outputMode("update") + .trigger(Trigger.RealTime("5 minutes")) + .start() +{% endhighlight %} + +The trigger duration is the checkpoint interval, not the per-record latency target. During a +long-running input batch, processing-time timers and TTL use the live clock. Timer scans and +physical TTL cleanup are driven by arriving rows and by batch completion, so an idle state +partition does not receive an independent timer wakeup. Event-time watermarks remain fixed during +a Real-time batch and advance between batches. + +If the query supplies initial state, Spark commits that state in a finite bootstrap batch before it +starts processing Real-time input. Processing-time values passed to `handleInitialState` use that +bootstrap batch's timestamp rather than the live clock. See +[`transformWithState` in Real-time Mode](./real-time-mode.html#transformwithstate-in-real-time-mode) +for the complete execution and recovery behavior. + ## State Schema Evolution TransformWithState also allows for performing schema evolution of the managed state. There are 2 parts here: diff --git a/docs/web-ui.md b/docs/web-ui.md index 6ae0a363d1873..085e8be755372 100644 --- a/docs/web-ui.md +++ b/docs/web-ui.md @@ -70,6 +70,24 @@ The information displayed at the top of the page includes: The current user, application start time, and total uptime are shown in the footer at the bottom of every page. +When the application can be held, the summary shows an **Application** line with a **(hold)** +link; clicking it stops requesting new executors and gracefully decommissions the running ones, +so each finishes its tasks and then exits (unless `spark.executor.decommission.forceKillTimeout` +is set, which kills a still-busy executor after that timeout). The line then reads `Held` with a +**(resume)** link that restores the executor requirement. Shuffle output written before the hold +stays available, but cached blocks are recomputed after resuming, and an RDD created while held +is sized against a default parallelism of 2 (no executors are alive) and keeps that partition +count afterwards. Pipelined-shuffle jobs and workloads with long-running tasks (streaming +receivers, continuous processing) are outside the hold's scope: a hold requested while a +pipelined job is running is rejected, a pipelined job submitted while held fails immediately +and should be resubmitted after the resume (unless the internal +`spark.scheduler.pipelinedGroup.slotCheck.enabled` is false, in which case there is no +admission check and it waits for the resume), and a long-running task never finishes, so the +drain cannot complete. The control only appears when `spark.ui.holdEnabled` is true, +`spark.decommission.enabled` is true, the shuffle data is kept outside the executors, and the +cluster manager can hold executors (Standalone, YARN, and Kubernetes with the `direct` pods +allocator); see [Configuration](configuration.html#spark-ui). + <p style="text-align: center;"> <img src="img/AllJobsPage.png" title="All Jobs page" alt="All Jobs page" width="100%"/> </p> diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java index f84a1978de1ad..71fd665b69063 100644 --- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java +++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaCustomReceiver.java @@ -38,12 +38,12 @@ import java.util.regex.Pattern; /** - * Custom Receiver that receives data over a socket. Received bytes is interpreted as + * Custom Receiver that receives data over a socket. Received bytes are interpreted as * text and \n delimited lines are considered as records. They are then counted and printed. * - * Usage: JavaCustomReceiver <master> <hostname> <port> - * <master> is the Spark master URL. In local mode, <master> should be 'local[n]' with n > 1. - * <hostname> and <port> of the TCP server that Spark Streaming would connect to receive data. + * Usage: JavaCustomReceiver <hostname> <port> + * <hostname> and <port> describe the TCP server that Spark Streaming would connect to receive + * data. * * To run this on your local machine, you need to first run a Netcat server * `$ nc -lk 9999` @@ -99,7 +99,7 @@ public void onStart() { @Override public void onStop() { // There is nothing much to do as the thread calling receive() - // is designed to stop by itself isStopped() returns false + // is designed to stop by itself when isStopped() returns true } /** Create a socket connection and receive data until receiver is stopped */ diff --git a/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java b/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java index 633b1ed5f5edf..e22daced96744 100644 --- a/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java +++ b/examples/src/main/java/org/apache/spark/examples/streaming/JavaRecoverableNetworkWordCount.java @@ -159,7 +159,7 @@ private static JavaStreamingContext createContext(String ip, public static void main(String[] args) throws Exception { if (args.length != 4) { - System.err.println("You arguments were " + Arrays.asList(args)); + System.err.println("Your arguments were " + Arrays.asList(args)); System.err.println( "Usage: JavaRecoverableNetworkWordCount <hostname> <port> <checkpoint-directory>\n" + " <output-file>. <hostname> and <port> describe the TCP server that Spark\n" + diff --git a/examples/src/main/python/als.py b/examples/src/main/python/als.py index 5bd1807cce1b0..37d3ed22c6bd6 100755 --- a/examples/src/main/python/als.py +++ b/examples/src/main/python/als.py @@ -24,8 +24,8 @@ import sys import numpy as np -from numpy.random import rand from numpy import matrix +from numpy.random import rand from pyspark.sql import SparkSession LAMBDA = 0.01 # regularization diff --git a/examples/src/main/python/avro_inputformat.py b/examples/src/main/python/avro_inputformat.py index 97475dde14ec6..19c02f8999232 100644 --- a/examples/src/main/python/avro_inputformat.py +++ b/examples/src/main/python/avro_inputformat.py @@ -44,9 +44,9 @@ {u'favorite_color': u'red', u'name': u'Ben'} """ import sys +from functools import reduce from typing import Any, Tuple -from functools import reduce from pyspark import RDD from pyspark.sql import SparkSession diff --git a/examples/src/main/python/logistic_regression.py b/examples/src/main/python/logistic_regression.py index 9645af619b1e3..788729bf8150c 100755 --- a/examples/src/main/python/logistic_regression.py +++ b/examples/src/main/python/logistic_regression.py @@ -25,11 +25,9 @@ import sys from typing import Iterable, List - import numpy as np from pyspark.sql import SparkSession - D = 10 # Number of dimensions diff --git a/examples/src/main/python/ml/aft_survival_regression.py b/examples/src/main/python/ml/aft_survival_regression.py index 2040a7876c7fa..7f40927848923 100644 --- a/examples/src/main/python/ml/aft_survival_regression.py +++ b/examples/src/main/python/ml/aft_survival_regression.py @@ -21,8 +21,9 @@ bin/spark-submit examples/src/main/python/ml/aft_survival_regression.py """ # $example on$ -from pyspark.ml.regression import AFTSurvivalRegression from pyspark.ml.linalg import Vectors +from pyspark.ml.regression import AFTSurvivalRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/als_example.py b/examples/src/main/python/ml/als_example.py index b39263978402b..7eebc225c5ed1 100644 --- a/examples/src/main/python/ml/als_example.py +++ b/examples/src/main/python/ml/als_example.py @@ -15,12 +15,11 @@ # limitations under the License. # -from pyspark.sql import SparkSession - # $example on$ from pyspark.ml.evaluation import RegressionEvaluator from pyspark.ml.recommendation import ALS -from pyspark.sql import Row +from pyspark.sql import Row, SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/binarizer_example.py b/examples/src/main/python/ml/binarizer_example.py index 5d5ae4122e1d4..265266b40de06 100644 --- a/examples/src/main/python/ml/binarizer_example.py +++ b/examples/src/main/python/ml/binarizer_example.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import Binarizer +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/bisecting_k_means_example.py b/examples/src/main/python/ml/bisecting_k_means_example.py index 513f80a09ef05..fc4fb84957a6a 100644 --- a/examples/src/main/python/ml/bisecting_k_means_example.py +++ b/examples/src/main/python/ml/bisecting_k_means_example.py @@ -23,6 +23,7 @@ # $example on$ from pyspark.ml.clustering import BisectingKMeans from pyspark.ml.evaluation import ClusteringEvaluator + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/bucketed_random_projection_lsh_example.py b/examples/src/main/python/ml/bucketed_random_projection_lsh_example.py index f5836091f35ba..9d156e8720cdd 100644 --- a/examples/src/main/python/ml/bucketed_random_projection_lsh_example.py +++ b/examples/src/main/python/ml/bucketed_random_projection_lsh_example.py @@ -23,9 +23,10 @@ # $example on$ from pyspark.ml.feature import BucketedRandomProjectionLSH from pyspark.ml.linalg import Vectors -from pyspark.sql.functions import col + # $example off$ from pyspark.sql import SparkSession +from pyspark.sql.functions import col if __name__ == "__main__": spark = SparkSession \ diff --git a/examples/src/main/python/ml/bucketizer_example.py b/examples/src/main/python/ml/bucketizer_example.py index bad5d787e1349..71d867046a371 100644 --- a/examples/src/main/python/ml/bucketizer_example.py +++ b/examples/src/main/python/ml/bucketizer_example.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import Bucketizer +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/chi_square_test_example.py b/examples/src/main/python/ml/chi_square_test_example.py index 0360742faf6f0..5939ef27c1200 100644 --- a/examples/src/main/python/ml/chi_square_test_example.py +++ b/examples/src/main/python/ml/chi_square_test_example.py @@ -20,10 +20,11 @@ Run with: bin/spark-submit examples/src/main/python/ml/chi_square_test_example.py """ -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.linalg import Vectors from pyspark.ml.stat import ChiSquareTest +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/chisq_selector_example.py b/examples/src/main/python/ml/chisq_selector_example.py index c83a8c1bc7b27..4933c1a403fae 100644 --- a/examples/src/main/python/ml/chisq_selector_example.py +++ b/examples/src/main/python/ml/chisq_selector_example.py @@ -15,10 +15,11 @@ # limitations under the License. # -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import ChiSqSelector from pyspark.ml.linalg import Vectors +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/correlation_example.py b/examples/src/main/python/ml/correlation_example.py index b15535a59882e..1d56ec0e9128d 100644 --- a/examples/src/main/python/ml/correlation_example.py +++ b/examples/src/main/python/ml/correlation_example.py @@ -23,6 +23,7 @@ # $example on$ from pyspark.ml.linalg import Vectors from pyspark.ml.stat import Correlation + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/count_vectorizer_example.py b/examples/src/main/python/ml/count_vectorizer_example.py index b3ddfb128c3d0..83f4233807ed9 100644 --- a/examples/src/main/python/ml/count_vectorizer_example.py +++ b/examples/src/main/python/ml/count_vectorizer_example.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import CountVectorizer +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/cross_validator.py b/examples/src/main/python/ml/cross_validator.py index 0ad0865486959..a3bc2b55be9b6 100644 --- a/examples/src/main/python/ml/cross_validator.py +++ b/examples/src/main/python/ml/cross_validator.py @@ -28,6 +28,7 @@ from pyspark.ml.evaluation import BinaryClassificationEvaluator from pyspark.ml.feature import HashingTF, Tokenizer from pyspark.ml.tuning import CrossValidator, ParamGridBuilder + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/dataframe_example.py b/examples/src/main/python/ml/dataframe_example.py index d2bf93744113b..16e4b3a5523e2 100644 --- a/examples/src/main/python/ml/dataframe_example.py +++ b/examples/src/main/python/ml/dataframe_example.py @@ -20,13 +20,13 @@ bin/spark-submit examples/src/main/python/ml/dataframe_example.py <input_path> """ import os +import shutil import sys import tempfile -import shutil -from pyspark.sql import SparkSession from pyspark.mllib.stat import Statistics from pyspark.mllib.util import MLUtils +from pyspark.sql import SparkSession if __name__ == "__main__": if len(sys.argv) > 2: diff --git a/examples/src/main/python/ml/dct_example.py b/examples/src/main/python/ml/dct_example.py index 37da4f5e8f1cb..e8dba41823740 100644 --- a/examples/src/main/python/ml/dct_example.py +++ b/examples/src/main/python/ml/dct_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import DCT from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/decision_tree_classification_example.py b/examples/src/main/python/ml/decision_tree_classification_example.py index eb7177b845357..82a1bf9ac0f65 100644 --- a/examples/src/main/python/ml/decision_tree_classification_example.py +++ b/examples/src/main/python/ml/decision_tree_classification_example.py @@ -21,8 +21,9 @@ # $example on$ from pyspark.ml import Pipeline from pyspark.ml.classification import DecisionTreeClassifier -from pyspark.ml.feature import StringIndexer, VectorIndexer from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import StringIndexer, VectorIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/decision_tree_regression_example.py b/examples/src/main/python/ml/decision_tree_regression_example.py index 1ed1636a3d962..9eafcc9ef0ece 100644 --- a/examples/src/main/python/ml/decision_tree_regression_example.py +++ b/examples/src/main/python/ml/decision_tree_regression_example.py @@ -20,9 +20,10 @@ """ # $example on$ from pyspark.ml import Pipeline -from pyspark.ml.regression import DecisionTreeRegressor -from pyspark.ml.feature import VectorIndexer from pyspark.ml.evaluation import RegressionEvaluator +from pyspark.ml.feature import VectorIndexer +from pyspark.ml.regression import DecisionTreeRegressor + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/elementwise_product_example.py b/examples/src/main/python/ml/elementwise_product_example.py index 71eec8d432998..e972db907a3aa 100644 --- a/examples/src/main/python/ml/elementwise_product_example.py +++ b/examples/src/main/python/ml/elementwise_product_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import ElementwiseProduct from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/estimator_transformer_param_example.py b/examples/src/main/python/ml/estimator_transformer_param_example.py index b34861a97fa7f..e981415f0dd40 100644 --- a/examples/src/main/python/ml/estimator_transformer_param_example.py +++ b/examples/src/main/python/ml/estimator_transformer_param_example.py @@ -19,8 +19,9 @@ Estimator Transformer Param Example. """ # $example on$ -from pyspark.ml.linalg import Vectors from pyspark.ml.classification import LogisticRegression +from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/feature_hasher_example.py b/examples/src/main/python/ml/feature_hasher_example.py index 4fe573d19dfbc..64e8a047b0173 100644 --- a/examples/src/main/python/ml/feature_hasher_example.py +++ b/examples/src/main/python/ml/feature_hasher_example.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import FeatureHasher +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/fm_classifier_example.py b/examples/src/main/python/ml/fm_classifier_example.py index da49e5fc2baa9..d3fba01b85f54 100644 --- a/examples/src/main/python/ml/fm_classifier_example.py +++ b/examples/src/main/python/ml/fm_classifier_example.py @@ -21,8 +21,9 @@ # $example on$ from pyspark.ml import Pipeline from pyspark.ml.classification import FMClassifier -from pyspark.ml.feature import MinMaxScaler, StringIndexer from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import MinMaxScaler, StringIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/fm_regressor_example.py b/examples/src/main/python/ml/fm_regressor_example.py index 47544b6324203..02527a9a34f69 100644 --- a/examples/src/main/python/ml/fm_regressor_example.py +++ b/examples/src/main/python/ml/fm_regressor_example.py @@ -20,9 +20,10 @@ """ # $example on$ from pyspark.ml import Pipeline -from pyspark.ml.regression import FMRegressor -from pyspark.ml.feature import MinMaxScaler from pyspark.ml.evaluation import RegressionEvaluator +from pyspark.ml.feature import MinMaxScaler +from pyspark.ml.regression import FMRegressor + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/fpgrowth_example.py b/examples/src/main/python/ml/fpgrowth_example.py index 39092e616d429..d2ab9bdd7c96f 100644 --- a/examples/src/main/python/ml/fpgrowth_example.py +++ b/examples/src/main/python/ml/fpgrowth_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.fpm import FPGrowth + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/gaussian_mixture_example.py b/examples/src/main/python/ml/gaussian_mixture_example.py index 1441faa792983..a8901f43f0a76 100644 --- a/examples/src/main/python/ml/gaussian_mixture_example.py +++ b/examples/src/main/python/ml/gaussian_mixture_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.clustering import GaussianMixture + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/generalized_linear_regression_example.py b/examples/src/main/python/ml/generalized_linear_regression_example.py index 06a8a5a2e9428..df3801607977e 100644 --- a/examples/src/main/python/ml/generalized_linear_regression_example.py +++ b/examples/src/main/python/ml/generalized_linear_regression_example.py @@ -20,9 +20,10 @@ Run with: bin/spark-submit examples/src/main/python/ml/generalized_linear_regression_example.py """ -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.regression import GeneralizedLinearRegression +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py b/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py index a7efa2170a069..803ae70e163f9 100644 --- a/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py +++ b/examples/src/main/python/ml/gradient_boosted_tree_classifier_example.py @@ -21,8 +21,9 @@ # $example on$ from pyspark.ml import Pipeline from pyspark.ml.classification import GBTClassifier -from pyspark.ml.feature import StringIndexer, VectorIndexer from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import StringIndexer, VectorIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py b/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py index 5e09b96c1ea3a..821ae6c6f8d8b 100644 --- a/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py +++ b/examples/src/main/python/ml/gradient_boosted_tree_regressor_example.py @@ -20,9 +20,10 @@ """ # $example on$ from pyspark.ml import Pipeline -from pyspark.ml.regression import GBTRegressor -from pyspark.ml.feature import VectorIndexer from pyspark.ml.evaluation import RegressionEvaluator +from pyspark.ml.feature import VectorIndexer +from pyspark.ml.regression import GBTRegressor + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/imputer_example.py b/examples/src/main/python/ml/imputer_example.py index 9ba0147763618..cec250a172776 100644 --- a/examples/src/main/python/ml/imputer_example.py +++ b/examples/src/main/python/ml/imputer_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.feature import Imputer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/index_to_string_example.py b/examples/src/main/python/ml/index_to_string_example.py index 98bdb89ce3039..16e5a6510568b 100644 --- a/examples/src/main/python/ml/index_to_string_example.py +++ b/examples/src/main/python/ml/index_to_string_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import IndexToString, StringIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/interaction_example.py b/examples/src/main/python/ml/interaction_example.py index ac365179b0c20..bc8f4441fcb76 100644 --- a/examples/src/main/python/ml/interaction_example.py +++ b/examples/src/main/python/ml/interaction_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import Interaction, VectorAssembler + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/isotonic_regression_example.py b/examples/src/main/python/ml/isotonic_regression_example.py index d7b893894fc71..76f462389b5d8 100644 --- a/examples/src/main/python/ml/isotonic_regression_example.py +++ b/examples/src/main/python/ml/isotonic_regression_example.py @@ -23,6 +23,7 @@ """ # $example on$ from pyspark.ml.regression import IsotonicRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/kmeans_example.py b/examples/src/main/python/ml/kmeans_example.py index 47223fd953d17..939e101179e82 100644 --- a/examples/src/main/python/ml/kmeans_example.py +++ b/examples/src/main/python/ml/kmeans_example.py @@ -25,8 +25,8 @@ # $example on$ from pyspark.ml.clustering import KMeans from pyspark.ml.evaluation import ClusteringEvaluator -# $example off$ +# $example off$ from pyspark.sql import SparkSession if __name__ == "__main__": diff --git a/examples/src/main/python/ml/lda_example.py b/examples/src/main/python/ml/lda_example.py index a47dfa383c895..4999314f032e6 100644 --- a/examples/src/main/python/ml/lda_example.py +++ b/examples/src/main/python/ml/lda_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.clustering import LDA + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/linear_regression_with_elastic_net.py b/examples/src/main/python/ml/linear_regression_with_elastic_net.py index 864fc76cff132..8282614ec4c4c 100644 --- a/examples/src/main/python/ml/linear_regression_with_elastic_net.py +++ b/examples/src/main/python/ml/linear_regression_with_elastic_net.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.regression import LinearRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/linearsvc.py b/examples/src/main/python/ml/linearsvc.py index 61d726cf3f1ae..f04106a3214b7 100644 --- a/examples/src/main/python/ml/linearsvc.py +++ b/examples/src/main/python/ml/linearsvc.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.classification import LinearSVC + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/logistic_regression_summary_example.py b/examples/src/main/python/ml/logistic_regression_summary_example.py index 6d045108da0aa..10b24bd848e7b 100644 --- a/examples/src/main/python/ml/logistic_regression_summary_example.py +++ b/examples/src/main/python/ml/logistic_regression_summary_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.classification import LogisticRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/logistic_regression_with_elastic_net.py b/examples/src/main/python/ml/logistic_regression_with_elastic_net.py index 916fdade27623..63c4a9760e93f 100644 --- a/examples/src/main/python/ml/logistic_regression_with_elastic_net.py +++ b/examples/src/main/python/ml/logistic_regression_with_elastic_net.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.classification import LogisticRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/max_abs_scaler_example.py b/examples/src/main/python/ml/max_abs_scaler_example.py index d7ff3561ce429..76e6afff1f9c1 100644 --- a/examples/src/main/python/ml/max_abs_scaler_example.py +++ b/examples/src/main/python/ml/max_abs_scaler_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import MaxAbsScaler from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/min_hash_lsh_example.py b/examples/src/main/python/ml/min_hash_lsh_example.py index 683f97a055ede..bbaa69bafaeae 100644 --- a/examples/src/main/python/ml/min_hash_lsh_example.py +++ b/examples/src/main/python/ml/min_hash_lsh_example.py @@ -23,9 +23,10 @@ # $example on$ from pyspark.ml.feature import MinHashLSH from pyspark.ml.linalg import Vectors -from pyspark.sql.functions import col + # $example off$ from pyspark.sql import SparkSession +from pyspark.sql.functions import col if __name__ == "__main__": spark = SparkSession \ diff --git a/examples/src/main/python/ml/min_max_scaler_example.py b/examples/src/main/python/ml/min_max_scaler_example.py index cd74243699894..5b3ae052b16b2 100644 --- a/examples/src/main/python/ml/min_max_scaler_example.py +++ b/examples/src/main/python/ml/min_max_scaler_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import MinMaxScaler from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/multiclass_logistic_regression_with_elastic_net.py b/examples/src/main/python/ml/multiclass_logistic_regression_with_elastic_net.py index 3bb4a72864101..d142b1946c110 100644 --- a/examples/src/main/python/ml/multiclass_logistic_regression_with_elastic_net.py +++ b/examples/src/main/python/ml/multiclass_logistic_regression_with_elastic_net.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.classification import LogisticRegression + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/multilayer_perceptron_classification.py b/examples/src/main/python/ml/multilayer_perceptron_classification.py index 74f532193573d..e5c337e14c204 100644 --- a/examples/src/main/python/ml/multilayer_perceptron_classification.py +++ b/examples/src/main/python/ml/multilayer_perceptron_classification.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.classification import MultilayerPerceptronClassifier from pyspark.ml.evaluation import MulticlassClassificationEvaluator + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/n_gram_example.py b/examples/src/main/python/ml/n_gram_example.py index 8c8031b939458..2872c29b9c2a1 100644 --- a/examples/src/main/python/ml/n_gram_example.py +++ b/examples/src/main/python/ml/n_gram_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import NGram + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/naive_bayes_example.py b/examples/src/main/python/ml/naive_bayes_example.py index 8d1777c6f9e39..a098962e1b4e3 100644 --- a/examples/src/main/python/ml/naive_bayes_example.py +++ b/examples/src/main/python/ml/naive_bayes_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.classification import NaiveBayes from pyspark.ml.evaluation import MulticlassClassificationEvaluator + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/normalizer_example.py b/examples/src/main/python/ml/normalizer_example.py index 2aa012961a2ee..fd17a3a7c5d3e 100644 --- a/examples/src/main/python/ml/normalizer_example.py +++ b/examples/src/main/python/ml/normalizer_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import Normalizer from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/one_vs_rest_example.py b/examples/src/main/python/ml/one_vs_rest_example.py index 4cae1a99808e8..e8d8fc5e59c40 100644 --- a/examples/src/main/python/ml/one_vs_rest_example.py +++ b/examples/src/main/python/ml/one_vs_rest_example.py @@ -24,6 +24,7 @@ # $example on$ from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.evaluation import MulticlassClassificationEvaluator + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/onehot_encoder_example.py b/examples/src/main/python/ml/onehot_encoder_example.py index 6deb84ed785ca..f52ecade792cc 100644 --- a/examples/src/main/python/ml/onehot_encoder_example.py +++ b/examples/src/main/python/ml/onehot_encoder_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import OneHotEncoder + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/pca_example.py b/examples/src/main/python/ml/pca_example.py index 03fb709c8e91d..32b2263efc457 100644 --- a/examples/src/main/python/ml/pca_example.py +++ b/examples/src/main/python/ml/pca_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import PCA from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/pipeline_example.py b/examples/src/main/python/ml/pipeline_example.py index 5fff3d61d5c48..b8cf59af285bf 100644 --- a/examples/src/main/python/ml/pipeline_example.py +++ b/examples/src/main/python/ml/pipeline_example.py @@ -23,6 +23,7 @@ from pyspark.ml import Pipeline from pyspark.ml.classification import LogisticRegression from pyspark.ml.feature import HashingTF, Tokenizer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/polynomial_expansion_example.py b/examples/src/main/python/ml/polynomial_expansion_example.py index 75f436e768dc5..09e89924d1c4c 100644 --- a/examples/src/main/python/ml/polynomial_expansion_example.py +++ b/examples/src/main/python/ml/polynomial_expansion_example.py @@ -18,6 +18,7 @@ # $example on$ from pyspark.ml.feature import PolynomialExpansion from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/power_iteration_clustering_example.py b/examples/src/main/python/ml/power_iteration_clustering_example.py index c983c4ad2b0d6..3a2bbc856d8ad 100644 --- a/examples/src/main/python/ml/power_iteration_clustering_example.py +++ b/examples/src/main/python/ml/power_iteration_clustering_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.clustering import PowerIterationClustering + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/prefixspan_example.py b/examples/src/main/python/ml/prefixspan_example.py index 88d1d4197341b..25cf0880a10e2 100644 --- a/examples/src/main/python/ml/prefixspan_example.py +++ b/examples/src/main/python/ml/prefixspan_example.py @@ -22,6 +22,7 @@ """ # $example on$ from pyspark.ml.fpm import PrefixSpan + # $example off$ from pyspark.sql import Row, SparkSession diff --git a/examples/src/main/python/ml/quantile_discretizer_example.py b/examples/src/main/python/ml/quantile_discretizer_example.py index 82be3936d2598..fa7cc74ae1e61 100644 --- a/examples/src/main/python/ml/quantile_discretizer_example.py +++ b/examples/src/main/python/ml/quantile_discretizer_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import QuantileDiscretizer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/random_forest_classifier_example.py b/examples/src/main/python/ml/random_forest_classifier_example.py index 8983d1f2e979b..a07ab40a2e7eb 100644 --- a/examples/src/main/python/ml/random_forest_classifier_example.py +++ b/examples/src/main/python/ml/random_forest_classifier_example.py @@ -21,8 +21,9 @@ # $example on$ from pyspark.ml import Pipeline from pyspark.ml.classification import RandomForestClassifier -from pyspark.ml.feature import IndexToString, StringIndexer, VectorIndexer from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import IndexToString, StringIndexer, VectorIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/random_forest_regressor_example.py b/examples/src/main/python/ml/random_forest_regressor_example.py index b9306ddf2f82c..98cf22e49b9eb 100644 --- a/examples/src/main/python/ml/random_forest_regressor_example.py +++ b/examples/src/main/python/ml/random_forest_regressor_example.py @@ -20,9 +20,10 @@ """ # $example on$ from pyspark.ml import Pipeline -from pyspark.ml.regression import RandomForestRegressor -from pyspark.ml.feature import VectorIndexer from pyspark.ml.evaluation import RegressionEvaluator +from pyspark.ml.feature import VectorIndexer +from pyspark.ml.regression import RandomForestRegressor + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/rformula_example.py b/examples/src/main/python/ml/rformula_example.py index 25bb6dac56e81..795e3202025b9 100644 --- a/examples/src/main/python/ml/rformula_example.py +++ b/examples/src/main/python/ml/rformula_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import RFormula + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/robust_scaler_example.py b/examples/src/main/python/ml/robust_scaler_example.py index 9f7c6d6507c78..6972e6c1aeac1 100644 --- a/examples/src/main/python/ml/robust_scaler_example.py +++ b/examples/src/main/python/ml/robust_scaler_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import RobustScaler + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/sql_transformer.py b/examples/src/main/python/ml/sql_transformer.py index c8ac5c46aa5e9..2de00320611dd 100644 --- a/examples/src/main/python/ml/sql_transformer.py +++ b/examples/src/main/python/ml/sql_transformer.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import SQLTransformer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/standard_scaler_example.py b/examples/src/main/python/ml/standard_scaler_example.py index 9021c10075d81..0194d768f6f88 100644 --- a/examples/src/main/python/ml/standard_scaler_example.py +++ b/examples/src/main/python/ml/standard_scaler_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import StandardScaler + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/stopwords_remover_example.py b/examples/src/main/python/ml/stopwords_remover_example.py index 832a7c7d0ad88..e30ddfc18eba5 100644 --- a/examples/src/main/python/ml/stopwords_remover_example.py +++ b/examples/src/main/python/ml/stopwords_remover_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import StopWordsRemover + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/string_indexer_example.py b/examples/src/main/python/ml/string_indexer_example.py index f2ac63eabd71c..a0a4f303152dd 100644 --- a/examples/src/main/python/ml/string_indexer_example.py +++ b/examples/src/main/python/ml/string_indexer_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import StringIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/summarizer_example.py b/examples/src/main/python/ml/summarizer_example.py index 4982746450132..5414a9bd744cf 100644 --- a/examples/src/main/python/ml/summarizer_example.py +++ b/examples/src/main/python/ml/summarizer_example.py @@ -20,11 +20,12 @@ Run with: bin/spark-submit examples/src/main/python/ml/summarizer_example.py """ -from pyspark.sql import SparkSession +from pyspark.ml.linalg import Vectors + # $example on$ from pyspark.ml.stat import Summarizer -from pyspark.sql import Row -from pyspark.ml.linalg import Vectors +from pyspark.sql import Row, SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/tf_idf_example.py b/examples/src/main/python/ml/tf_idf_example.py index b4bb0dfa3183c..43fd2e991d6d0 100644 --- a/examples/src/main/python/ml/tf_idf_example.py +++ b/examples/src/main/python/ml/tf_idf_example.py @@ -16,7 +16,8 @@ # # $example on$ -from pyspark.ml.feature import HashingTF, IDF, Tokenizer +from pyspark.ml.feature import IDF, HashingTF, Tokenizer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/tokenizer_example.py b/examples/src/main/python/ml/tokenizer_example.py index c6b5fac227315..b04b2ee5ed8ae 100644 --- a/examples/src/main/python/ml/tokenizer_example.py +++ b/examples/src/main/python/ml/tokenizer_example.py @@ -16,11 +16,12 @@ # # $example on$ -from pyspark.ml.feature import Tokenizer, RegexTokenizer -from pyspark.sql.functions import col, udf -from pyspark.sql.types import IntegerType +from pyspark.ml.feature import RegexTokenizer, Tokenizer + # $example off$ from pyspark.sql import SparkSession +from pyspark.sql.functions import col, udf +from pyspark.sql.types import IntegerType if __name__ == "__main__": spark = SparkSession\ diff --git a/examples/src/main/python/ml/train_validation_split.py b/examples/src/main/python/ml/train_validation_split.py index 5e3dc7b3ec2fa..21e110419cfa2 100644 --- a/examples/src/main/python/ml/train_validation_split.py +++ b/examples/src/main/python/ml/train_validation_split.py @@ -26,6 +26,7 @@ from pyspark.ml.evaluation import RegressionEvaluator from pyspark.ml.regression import LinearRegression from pyspark.ml.tuning import ParamGridBuilder, TrainValidationSplit + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/univariate_feature_selector_example.py b/examples/src/main/python/ml/univariate_feature_selector_example.py index 6dc293e49643b..3b5f2db7f9c09 100644 --- a/examples/src/main/python/ml/univariate_feature_selector_example.py +++ b/examples/src/main/python/ml/univariate_feature_selector_example.py @@ -20,10 +20,11 @@ Run with: bin/spark-submit examples/src/main/python/ml/univariate_feature_selector_example.py """ -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import UnivariateFeatureSelector from pyspark.ml.linalg import Vectors +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/variance_threshold_selector_example.py b/examples/src/main/python/ml/variance_threshold_selector_example.py index 0a996e0e28264..8bbcbffa22714 100644 --- a/examples/src/main/python/ml/variance_threshold_selector_example.py +++ b/examples/src/main/python/ml/variance_threshold_selector_example.py @@ -20,10 +20,11 @@ Run with: bin/spark-submit examples/src/main/python/ml/variance_threshold_selector_example.py """ -from pyspark.sql import SparkSession # $example on$ from pyspark.ml.feature import VarianceThresholdSelector from pyspark.ml.linalg import Vectors +from pyspark.sql import SparkSession + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/ml/vector_assembler_example.py b/examples/src/main/python/ml/vector_assembler_example.py index 0ce31cf0eabc9..8de9cad4114a9 100644 --- a/examples/src/main/python/ml/vector_assembler_example.py +++ b/examples/src/main/python/ml/vector_assembler_example.py @@ -16,8 +16,9 @@ # # $example on$ -from pyspark.ml.linalg import Vectors from pyspark.ml.feature import VectorAssembler +from pyspark.ml.linalg import Vectors + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/vector_indexer_example.py b/examples/src/main/python/ml/vector_indexer_example.py index 51a4191606fb8..bf99ec299d19a 100644 --- a/examples/src/main/python/ml/vector_indexer_example.py +++ b/examples/src/main/python/ml/vector_indexer_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import VectorIndexer + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/vector_size_hint_example.py b/examples/src/main/python/ml/vector_size_hint_example.py index 355d85aee8729..f376f06441929 100644 --- a/examples/src/main/python/ml/vector_size_hint_example.py +++ b/examples/src/main/python/ml/vector_size_hint_example.py @@ -16,8 +16,9 @@ # # $example on$ +from pyspark.ml.feature import VectorAssembler, VectorSizeHint from pyspark.ml.linalg import Vectors -from pyspark.ml.feature import (VectorSizeHint, VectorAssembler) + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/ml/vector_slicer_example.py b/examples/src/main/python/ml/vector_slicer_example.py index 86e089d152c5a..368ba54bdb5d6 100644 --- a/examples/src/main/python/ml/vector_slicer_example.py +++ b/examples/src/main/python/ml/vector_slicer_example.py @@ -18,9 +18,10 @@ # $example on$ from pyspark.ml.feature import VectorSlicer from pyspark.ml.linalg import Vectors -from pyspark.sql.types import Row + # $example off$ from pyspark.sql import SparkSession +from pyspark.sql.types import Row if __name__ == "__main__": spark = SparkSession\ diff --git a/examples/src/main/python/ml/word2vec_example.py b/examples/src/main/python/ml/word2vec_example.py index 0eabeda3dce4b..a2fc5a5f02db8 100644 --- a/examples/src/main/python/ml/word2vec_example.py +++ b/examples/src/main/python/ml/word2vec_example.py @@ -17,6 +17,7 @@ # $example on$ from pyspark.ml.feature import Word2Vec + # $example off$ from pyspark.sql import SparkSession diff --git a/examples/src/main/python/mllib/binary_classification_metrics_example.py b/examples/src/main/python/mllib/binary_classification_metrics_example.py index 741746e6e35ae..941f8f64bc582 100644 --- a/examples/src/main/python/mllib/binary_classification_metrics_example.py +++ b/examples/src/main/python/mllib/binary_classification_metrics_example.py @@ -18,10 +18,12 @@ Binary Classification Metrics Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.classification import LogisticRegressionWithLBFGS from pyspark.mllib.evaluation import BinaryClassificationMetrics from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/bisecting_k_means_example.py b/examples/src/main/python/mllib/bisecting_k_means_example.py index d7b6ad9d424a6..0cd500fbe95f7 100644 --- a/examples/src/main/python/mllib/bisecting_k_means_example.py +++ b/examples/src/main/python/mllib/bisecting_k_means_example.py @@ -17,11 +17,13 @@ # $example on$ from numpy import array -# $example off$ +# $example off$ from pyspark import SparkContext + # $example on$ from pyspark.mllib.clustering import BisectingKMeans + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/correlations.py b/examples/src/main/python/mllib/correlations.py index 27d07b22a5645..14c5449a8ae9b 100755 --- a/examples/src/main/python/mllib/correlations.py +++ b/examples/src/main/python/mllib/correlations.py @@ -25,7 +25,6 @@ from pyspark.mllib.stat import Statistics from pyspark.mllib.util import MLUtils - if __name__ == "__main__": if len(sys.argv) not in [1, 2]: print("Usage: correlations (<file>)", file=sys.stderr) diff --git a/examples/src/main/python/mllib/correlations_example.py b/examples/src/main/python/mllib/correlations_example.py index bb71b968687cb..52c108f4608f4 100644 --- a/examples/src/main/python/mllib/correlations_example.py +++ b/examples/src/main/python/mllib/correlations_example.py @@ -16,10 +16,11 @@ # import numpy as np - from pyspark import SparkContext + # $example on$ from pyspark.mllib.stat import Statistics + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/decision_tree_classification_example.py b/examples/src/main/python/mllib/decision_tree_classification_example.py index 009e393226c01..b29dea1b96e89 100644 --- a/examples/src/main/python/mllib/decision_tree_classification_example.py +++ b/examples/src/main/python/mllib/decision_tree_classification_example.py @@ -19,9 +19,11 @@ Decision Tree Classification Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import DecisionTree, DecisionTreeModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/decision_tree_regression_example.py b/examples/src/main/python/mllib/decision_tree_regression_example.py index 71dfbf0790175..8cf0388fe8220 100644 --- a/examples/src/main/python/mllib/decision_tree_regression_example.py +++ b/examples/src/main/python/mllib/decision_tree_regression_example.py @@ -19,9 +19,11 @@ Decision Tree Regression Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import DecisionTree, DecisionTreeModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/elementwise_product_example.py b/examples/src/main/python/mllib/elementwise_product_example.py index 15e6a43f736cf..48918a5431ae8 100644 --- a/examples/src/main/python/mllib/elementwise_product_example.py +++ b/examples/src/main/python/mllib/elementwise_product_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.feature import ElementwiseProduct from pyspark.mllib.linalg import Vectors + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/fpgrowth_example.py b/examples/src/main/python/mllib/fpgrowth_example.py index 715f5268206cb..d1c9016f81d13 100644 --- a/examples/src/main/python/mllib/fpgrowth_example.py +++ b/examples/src/main/python/mllib/fpgrowth_example.py @@ -16,9 +16,9 @@ # # $example on$ -from pyspark.mllib.fpm import FPGrowth # $example off$ from pyspark import SparkContext +from pyspark.mllib.fpm import FPGrowth if __name__ == "__main__": sc = SparkContext(appName="FPGrowth") diff --git a/examples/src/main/python/mllib/gaussian_mixture_example.py b/examples/src/main/python/mllib/gaussian_mixture_example.py index 3b19478f457ec..5c0f65f60d3e6 100644 --- a/examples/src/main/python/mllib/gaussian_mixture_example.py +++ b/examples/src/main/python/mllib/gaussian_mixture_example.py @@ -17,11 +17,13 @@ # $example on$ from numpy import array -# $example off$ +# $example off$ from pyspark import SparkContext + # $example on$ from pyspark.mllib.clustering import GaussianMixture, GaussianMixtureModel + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/gaussian_mixture_model.py b/examples/src/main/python/mllib/gaussian_mixture_model.py index 96ce6b6f6ab25..df910bb064f5b 100644 --- a/examples/src/main/python/mllib/gaussian_mixture_model.py +++ b/examples/src/main/python/mllib/gaussian_mixture_model.py @@ -19,10 +19,10 @@ A Gaussian Mixture Model clustering program using MLlib. """ -import random import argparse -import numpy as np +import random +import numpy as np from pyspark import SparkConf, SparkContext from pyspark.mllib.clustering import GaussianMixture diff --git a/examples/src/main/python/mllib/gradient_boosting_classification_example.py b/examples/src/main/python/mllib/gradient_boosting_classification_example.py index eb12f206196fe..1a7cb0f8a9322 100644 --- a/examples/src/main/python/mllib/gradient_boosting_classification_example.py +++ b/examples/src/main/python/mllib/gradient_boosting_classification_example.py @@ -19,9 +19,11 @@ Gradient Boosted Trees Classification Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/gradient_boosting_regression_example.py b/examples/src/main/python/mllib/gradient_boosting_regression_example.py index eb59a992df539..62b82113db953 100644 --- a/examples/src/main/python/mllib/gradient_boosting_regression_example.py +++ b/examples/src/main/python/mllib/gradient_boosting_regression_example.py @@ -19,9 +19,11 @@ Gradient Boosted Trees Regression Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import GradientBoostedTrees, GradientBoostedTreesModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/hypothesis_testing_example.py b/examples/src/main/python/mllib/hypothesis_testing_example.py index 321be8b76f1b9..6a92d3d95ebd7 100644 --- a/examples/src/main/python/mllib/hypothesis_testing_example.py +++ b/examples/src/main/python/mllib/hypothesis_testing_example.py @@ -16,10 +16,12 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.linalg import Matrices, Vectors from pyspark.mllib.regression import LabeledPoint from pyspark.mllib.stat import Statistics + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/hypothesis_testing_kolmogorov_smirnov_test_example.py b/examples/src/main/python/mllib/hypothesis_testing_kolmogorov_smirnov_test_example.py index 12a186900e358..0f686086d66c7 100644 --- a/examples/src/main/python/mllib/hypothesis_testing_kolmogorov_smirnov_test_example.py +++ b/examples/src/main/python/mllib/hypothesis_testing_kolmogorov_smirnov_test_example.py @@ -16,8 +16,10 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.stat import Statistics + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/isotonic_regression_example.py b/examples/src/main/python/mllib/isotonic_regression_example.py index a5a0cfeae9d75..10035b5b76197 100644 --- a/examples/src/main/python/mllib/isotonic_regression_example.py +++ b/examples/src/main/python/mllib/isotonic_regression_example.py @@ -18,11 +18,13 @@ """ Isotonic Regression Example. """ -from pyspark import SparkContext # $example on$ import math + +from pyspark import SparkContext from pyspark.mllib.regression import IsotonicRegression, IsotonicRegressionModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/k_means_example.py b/examples/src/main/python/mllib/k_means_example.py index ead1e56de55c6..c2901c154f683 100644 --- a/examples/src/main/python/mllib/k_means_example.py +++ b/examples/src/main/python/mllib/k_means_example.py @@ -16,13 +16,16 @@ # # $example on$ -from numpy import array from math import sqrt -# $example off$ +from numpy import array + +# $example off$ from pyspark import SparkContext + # $example on$ from pyspark.mllib.clustering import KMeans, KMeansModel + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/kernel_density_estimation_example.py b/examples/src/main/python/mllib/kernel_density_estimation_example.py index 22d191716057c..c3dc2d027eeac 100644 --- a/examples/src/main/python/mllib/kernel_density_estimation_example.py +++ b/examples/src/main/python/mllib/kernel_density_estimation_example.py @@ -16,8 +16,10 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.stat import KernelDensity + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/latent_dirichlet_allocation_example.py b/examples/src/main/python/mllib/latent_dirichlet_allocation_example.py index f82a28aadc5a3..81b039e56895c 100644 --- a/examples/src/main/python/mllib/latent_dirichlet_allocation_example.py +++ b/examples/src/main/python/mllib/latent_dirichlet_allocation_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.clustering import LDA, LDAModel from pyspark.mllib.linalg import Vectors + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/linear_regression_with_sgd_example.py b/examples/src/main/python/mllib/linear_regression_with_sgd_example.py index cb67396332312..1813a838d6355 100644 --- a/examples/src/main/python/mllib/linear_regression_with_sgd_example.py +++ b/examples/src/main/python/mllib/linear_regression_with_sgd_example.py @@ -19,8 +19,10 @@ Linear Regression With SGD Example. """ from pyspark import SparkContext + # $example on$ -from pyspark.mllib.regression import LabeledPoint, LinearRegressionWithSGD, LinearRegressionModel +from pyspark.mllib.regression import LabeledPoint, LinearRegressionModel, LinearRegressionWithSGD + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/logistic_regression.py b/examples/src/main/python/mllib/logistic_regression.py index 7b90615a53424..8ae5b0ab8b7d7 100755 --- a/examples/src/main/python/mllib/logistic_regression.py +++ b/examples/src/main/python/mllib/logistic_regression.py @@ -23,8 +23,8 @@ import sys from pyspark import SparkContext -from pyspark.mllib.regression import LabeledPoint from pyspark.mllib.classification import LogisticRegressionWithSGD +from pyspark.mllib.regression import LabeledPoint def parsePoint(line): diff --git a/examples/src/main/python/mllib/logistic_regression_with_lbfgs_example.py b/examples/src/main/python/mllib/logistic_regression_with_lbfgs_example.py index ac5ab1d1b5d91..014fbfca1a88b 100644 --- a/examples/src/main/python/mllib/logistic_regression_with_lbfgs_example.py +++ b/examples/src/main/python/mllib/logistic_regression_with_lbfgs_example.py @@ -19,9 +19,11 @@ Logistic Regression With LBFGS Example. """ from pyspark import SparkContext + # $example on$ -from pyspark.mllib.classification import LogisticRegressionWithLBFGS, LogisticRegressionModel +from pyspark.mllib.classification import LogisticRegressionModel, LogisticRegressionWithLBFGS from pyspark.mllib.regression import LabeledPoint + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/multi_class_metrics_example.py b/examples/src/main/python/mllib/multi_class_metrics_example.py index 03a564e75be90..e429e8795b36f 100644 --- a/examples/src/main/python/mllib/multi_class_metrics_example.py +++ b/examples/src/main/python/mllib/multi_class_metrics_example.py @@ -16,12 +16,11 @@ # # $example on$ -from pyspark.mllib.classification import LogisticRegressionWithLBFGS -from pyspark.mllib.util import MLUtils -from pyspark.mllib.evaluation import MulticlassMetrics # $example off$ - from pyspark import SparkContext +from pyspark.mllib.classification import LogisticRegressionWithLBFGS +from pyspark.mllib.evaluation import MulticlassMetrics +from pyspark.mllib.util import MLUtils if __name__ == "__main__": sc = SparkContext(appName="MultiClassMetricsExample") diff --git a/examples/src/main/python/mllib/multi_label_metrics_example.py b/examples/src/main/python/mllib/multi_label_metrics_example.py index 960ade6597379..be689b9565d26 100644 --- a/examples/src/main/python/mllib/multi_label_metrics_example.py +++ b/examples/src/main/python/mllib/multi_label_metrics_example.py @@ -16,9 +16,9 @@ # # $example on$ -from pyspark.mllib.evaluation import MultilabelMetrics # $example off$ from pyspark import SparkContext +from pyspark.mllib.evaluation import MultilabelMetrics if __name__ == "__main__": sc = SparkContext(appName="MultiLabelMetricsExample") diff --git a/examples/src/main/python/mllib/naive_bayes_example.py b/examples/src/main/python/mllib/naive_bayes_example.py index 74d18233d533a..f76ab891a79ad 100644 --- a/examples/src/main/python/mllib/naive_bayes_example.py +++ b/examples/src/main/python/mllib/naive_bayes_example.py @@ -25,11 +25,11 @@ import shutil from pyspark import SparkContext + # $example on$ from pyspark.mllib.classification import NaiveBayes, NaiveBayesModel from pyspark.mllib.util import MLUtils - # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/normalizer_example.py b/examples/src/main/python/mllib/normalizer_example.py index d46110d9a0300..df0d9faf9aa5b 100644 --- a/examples/src/main/python/mllib/normalizer_example.py +++ b/examples/src/main/python/mllib/normalizer_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.feature import Normalizer from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/pca_rowmatrix_example.py b/examples/src/main/python/mllib/pca_rowmatrix_example.py index 49b9b1bbe08e9..99650cc631c7c 100644 --- a/examples/src/main/python/mllib/pca_rowmatrix_example.py +++ b/examples/src/main/python/mllib/pca_rowmatrix_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.linalg import Vectors from pyspark.mllib.linalg.distributed import RowMatrix + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/power_iteration_clustering_example.py b/examples/src/main/python/mllib/power_iteration_clustering_example.py index 60eedef5fab30..3901cce1d2d67 100644 --- a/examples/src/main/python/mllib/power_iteration_clustering_example.py +++ b/examples/src/main/python/mllib/power_iteration_clustering_example.py @@ -16,8 +16,10 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.clustering import PowerIterationClustering, PowerIterationClusteringModel + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/random_forest_classification_example.py b/examples/src/main/python/mllib/random_forest_classification_example.py index a929c10d5a573..3880c77f65e23 100644 --- a/examples/src/main/python/mllib/random_forest_classification_example.py +++ b/examples/src/main/python/mllib/random_forest_classification_example.py @@ -19,9 +19,11 @@ Random Forest Classification Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import RandomForest, RandomForestModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/random_forest_regression_example.py b/examples/src/main/python/mllib/random_forest_regression_example.py index 4e05937768211..a45bd63c12944 100644 --- a/examples/src/main/python/mllib/random_forest_regression_example.py +++ b/examples/src/main/python/mllib/random_forest_regression_example.py @@ -19,9 +19,11 @@ Random Forest Regression Example. """ from pyspark import SparkContext + # $example on$ from pyspark.mllib.tree import RandomForest, RandomForestModel from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/random_rdd_generation.py b/examples/src/main/python/mllib/random_rdd_generation.py index 49afcfe9391ab..2141d7a2f4004 100755 --- a/examples/src/main/python/mllib/random_rdd_generation.py +++ b/examples/src/main/python/mllib/random_rdd_generation.py @@ -23,7 +23,6 @@ from pyspark import SparkContext from pyspark.mllib.random import RandomRDDs - if __name__ == "__main__": if len(sys.argv) not in [1, 2]: print("Usage: random_rdd_generation", file=sys.stderr) diff --git a/examples/src/main/python/mllib/ranking_metrics_example.py b/examples/src/main/python/mllib/ranking_metrics_example.py index 0913bb34cf9d7..2e5738d49d006 100644 --- a/examples/src/main/python/mllib/ranking_metrics_example.py +++ b/examples/src/main/python/mllib/ranking_metrics_example.py @@ -16,10 +16,10 @@ # # $example on$ -from pyspark.mllib.recommendation import ALS, Rating -from pyspark.mllib.evaluation import RegressionMetrics # $example off$ from pyspark import SparkContext +from pyspark.mllib.evaluation import RegressionMetrics +from pyspark.mllib.recommendation import ALS, Rating if __name__ == "__main__": sc = SparkContext(appName="Ranking Metrics Example") diff --git a/examples/src/main/python/mllib/recommendation_example.py b/examples/src/main/python/mllib/recommendation_example.py index 719f3f904b246..b8204caa7b9a8 100644 --- a/examples/src/main/python/mllib/recommendation_example.py +++ b/examples/src/main/python/mllib/recommendation_example.py @@ -22,6 +22,7 @@ # $example on$ from pyspark.mllib.recommendation import ALS, MatrixFactorizationModel, Rating + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/regression_metrics_example.py b/examples/src/main/python/mllib/regression_metrics_example.py index a3a83aafd7a1f..4ac0b4ed258d2 100644 --- a/examples/src/main/python/mllib/regression_metrics_example.py +++ b/examples/src/main/python/mllib/regression_metrics_example.py @@ -15,12 +15,11 @@ # limitations under the License. # # $example on$ -from pyspark.mllib.regression import LabeledPoint, LinearRegressionWithSGD -from pyspark.mllib.evaluation import RegressionMetrics -from pyspark.mllib.linalg import DenseVector # $example off$ - from pyspark import SparkContext +from pyspark.mllib.evaluation import RegressionMetrics +from pyspark.mllib.linalg import DenseVector +from pyspark.mllib.regression import LabeledPoint, LinearRegressionWithSGD if __name__ == "__main__": sc = SparkContext(appName="Regression Metrics Example") diff --git a/examples/src/main/python/mllib/sampled_rdds.py b/examples/src/main/python/mllib/sampled_rdds.py index 9095c2b2d70d6..abf804fbaec72 100755 --- a/examples/src/main/python/mllib/sampled_rdds.py +++ b/examples/src/main/python/mllib/sampled_rdds.py @@ -23,7 +23,6 @@ from pyspark import SparkContext from pyspark.mllib.util import MLUtils - if __name__ == "__main__": if len(sys.argv) not in [1, 2]: print("Usage: sampled_rdds <libsvm data file>", file=sys.stderr) diff --git a/examples/src/main/python/mllib/standard_scaler_example.py b/examples/src/main/python/mllib/standard_scaler_example.py index c8fd64dfbbf4a..9de3b67968413 100644 --- a/examples/src/main/python/mllib/standard_scaler_example.py +++ b/examples/src/main/python/mllib/standard_scaler_example.py @@ -16,10 +16,12 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.feature import StandardScaler from pyspark.mllib.linalg import Vectors from pyspark.mllib.util import MLUtils + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/streaming_k_means_example.py b/examples/src/main/python/mllib/streaming_k_means_example.py index 4904a9ebcf544..0b1df65277f04 100644 --- a/examples/src/main/python/mllib/streaming_k_means_example.py +++ b/examples/src/main/python/mllib/streaming_k_means_example.py @@ -16,11 +16,13 @@ # from pyspark import SparkContext -from pyspark.streaming import StreamingContext +from pyspark.mllib.clustering import StreamingKMeans + # $example on$ from pyspark.mllib.linalg import Vectors from pyspark.mllib.regression import LabeledPoint -from pyspark.mllib.clustering import StreamingKMeans +from pyspark.streaming import StreamingContext + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/streaming_linear_regression_example.py b/examples/src/main/python/mllib/streaming_linear_regression_example.py index 1d52e00fbfb5e..fe1b79414acb9 100644 --- a/examples/src/main/python/mllib/streaming_linear_regression_example.py +++ b/examples/src/main/python/mllib/streaming_linear_regression_example.py @@ -20,14 +20,15 @@ """ # $example on$ import sys -# $example off$ +# $example off$ from pyspark import SparkContext -from pyspark.streaming import StreamingContext + # $example on$ from pyspark.mllib.linalg import Vectors -from pyspark.mllib.regression import LabeledPoint -from pyspark.mllib.regression import StreamingLinearRegressionWithSGD +from pyspark.mllib.regression import LabeledPoint, StreamingLinearRegressionWithSGD +from pyspark.streaming import StreamingContext + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/summary_statistics_example.py b/examples/src/main/python/mllib/summary_statistics_example.py index d86e841145501..fad7cd4c5a6fe 100644 --- a/examples/src/main/python/mllib/summary_statistics_example.py +++ b/examples/src/main/python/mllib/summary_statistics_example.py @@ -15,11 +15,11 @@ # limitations under the License. # -from pyspark import SparkContext # $example on$ import numpy as np - +from pyspark import SparkContext from pyspark.mllib.stat import Statistics + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/svd_example.py b/examples/src/main/python/mllib/svd_example.py index 5b220fdb3fd67..5469bb6710ee0 100644 --- a/examples/src/main/python/mllib/svd_example.py +++ b/examples/src/main/python/mllib/svd_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.linalg import Vectors from pyspark.mllib.linalg.distributed import RowMatrix + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/svm_with_sgd_example.py b/examples/src/main/python/mllib/svm_with_sgd_example.py index 24b8f431e059e..456350bae7e9b 100644 --- a/examples/src/main/python/mllib/svm_with_sgd_example.py +++ b/examples/src/main/python/mllib/svm_with_sgd_example.py @@ -16,9 +16,11 @@ # from pyspark import SparkContext + # $example on$ -from pyspark.mllib.classification import SVMWithSGD, SVMModel +from pyspark.mllib.classification import SVMModel, SVMWithSGD from pyspark.mllib.regression import LabeledPoint + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/tf_idf_example.py b/examples/src/main/python/mllib/tf_idf_example.py index 4449066f5b0a6..ca46d62dacc90 100644 --- a/examples/src/main/python/mllib/tf_idf_example.py +++ b/examples/src/main/python/mllib/tf_idf_example.py @@ -16,8 +16,10 @@ # from pyspark import SparkContext + # $example on$ -from pyspark.mllib.feature import HashingTF, IDF +from pyspark.mllib.feature import IDF, HashingTF + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/mllib/word2vec_example.py b/examples/src/main/python/mllib/word2vec_example.py index d37a6e7137b8f..175eecb92a105 100644 --- a/examples/src/main/python/mllib/word2vec_example.py +++ b/examples/src/main/python/mllib/word2vec_example.py @@ -16,8 +16,10 @@ # from pyspark import SparkContext + # $example on$ from pyspark.mllib.feature import Word2Vec + # $example off$ if __name__ == "__main__": diff --git a/examples/src/main/python/pi.py b/examples/src/main/python/pi.py index e61740ad58832..5425e5d9142af 100755 --- a/examples/src/main/python/pi.py +++ b/examples/src/main/python/pi.py @@ -16,12 +16,11 @@ # import sys -from random import random from operator import add +from random import random from pyspark.sql import SparkSession - if __name__ == "__main__": """ Usage: pi [partitions] diff --git a/examples/src/main/python/sort.py b/examples/src/main/python/sort.py index 9ef2d5dbaff4c..139a9d5dc28aa 100755 --- a/examples/src/main/python/sort.py +++ b/examples/src/main/python/sort.py @@ -21,7 +21,6 @@ from pyspark import RDD from pyspark.sql import SparkSession - if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: sort <file>", file=sys.stderr) diff --git a/examples/src/main/python/sql/arrow.py b/examples/src/main/python/sql/arrow.py index 62d8385849ca1..0780e6c2bd872 100644 --- a/examples/src/main/python/sql/arrow.py +++ b/examples/src/main/python/sql/arrow.py @@ -22,10 +22,11 @@ """ # NOTE that this file is imported in tutorials in PySpark documentation. -# The codes are referred via line numbers. See also `literalinclude` directive in Sphinx. -import pandas as pd +# The code blocks are referred to via `$example on/off$` markers, included in the +# docs with the `literalinclude` directive's `:start-after:`/`:end-before:` options. from typing import Iterator +import pandas as pd from pyspark.sql import SparkSession from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version @@ -34,8 +35,9 @@ def dataframe_to_from_arrow_table_example(spark: SparkSession) -> None: - import pyarrow as pa + # $example on:dataframe_to_from_arrow_table$ import numpy as np + import pyarrow as pa # Create a PyArrow Table table = pa.table([pa.array(np.random.rand(100)) for i in range(3)], names=["a", "b", "c"]) @@ -50,9 +52,11 @@ def dataframe_to_from_arrow_table_example(spark: SparkSession) -> None: # a: double # b: double # c: double + # $example off:dataframe_to_from_arrow_table$ def dataframe_with_arrow_example(spark: SparkSession) -> None: + # $example on:dataframe_with_arrow$ import numpy as np import pandas as pd @@ -69,11 +73,12 @@ def dataframe_with_arrow_example(spark: SparkSession) -> None: result_pdf = df.select("*").toPandas() print("Pandas DataFrame result statistics:\n%s\n" % str(result_pdf.describe())) + # $example off:dataframe_with_arrow$ def ser_to_frame_pandas_udf_example(spark: SparkSession) -> None: + # $example on:ser_to_frame_pandas_udf$ import pandas as pd - from pyspark.sql.functions import pandas_udf @pandas_udf("col1 string, col2 long") # type: ignore[call-overload] @@ -97,11 +102,12 @@ def func(s1: pd.Series, s2: pd.Series, s3: pd.DataFrame) -> pd.DataFrame: # |-- func(long_col, string_col, struct_col): struct (nullable = true) # | |-- col1: string (nullable = true) # | |-- col2: long (nullable = true) + # $example off:ser_to_frame_pandas_udf$ def ser_to_ser_pandas_udf_example(spark: SparkSession) -> None: + # $example on:ser_to_ser_pandas_udf$ import pandas as pd - from pyspark.sql.functions import col, pandas_udf from pyspark.sql.types import LongType @@ -131,13 +137,14 @@ def multiply_func(a: pd.Series, b: pd.Series) -> pd.Series: # | 4| # | 9| # +-------------------+ + # $example off:ser_to_ser_pandas_udf$ def iter_ser_to_iter_ser_pandas_udf_example(spark: SparkSession) -> None: + # $example on:iter_ser_to_iter_ser_pandas_udf$ from typing import Iterator import pandas as pd - from pyspark.sql.functions import pandas_udf pdf = pd.DataFrame([1, 2, 3], columns=["x"]) @@ -157,13 +164,14 @@ def plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: # | 3| # | 4| # +-----------+ + # $example off:iter_ser_to_iter_ser_pandas_udf$ def iter_sers_to_iter_ser_pandas_udf_example(spark: SparkSession) -> None: + # $example on:iter_sers_to_iter_ser_pandas_udf$ from typing import Iterator, Tuple import pandas as pd - from pyspark.sql.functions import pandas_udf pdf = pd.DataFrame([1, 2, 3], columns=["x"]) @@ -184,13 +192,14 @@ def multiply_two_cols( # | 4| # | 9| # +-----------------------+ + # $example off:iter_sers_to_iter_ser_pandas_udf$ def ser_to_scalar_pandas_udf_example(spark: SparkSession) -> None: + # $example on:ser_to_scalar_pandas_udf$ import pandas as pd - - from pyspark.sql.functions import pandas_udf from pyspark.sql import Window + from pyspark.sql.functions import pandas_udf df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], @@ -229,9 +238,11 @@ def mean_udf(v: pd.Series) -> float: # | 2| 5.0| 6.0| # | 2|10.0| 6.0| # +---+----+------+ + # $example off:ser_to_scalar_pandas_udf$ def grouped_apply_in_pandas_example(spark: SparkSession) -> None: + # $example on:grouped_apply_in_pandas$ df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v")) @@ -251,9 +262,11 @@ def subtract_mean(pdf: pd.DataFrame) -> pd.DataFrame: # | 2|-1.0| # | 2| 4.0| # +---+----+ + # $example off:grouped_apply_in_pandas$ def map_in_pandas_example(spark: SparkSession) -> None: + # $example on:map_in_pandas$ df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age")) def filter_func(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: @@ -266,9 +279,11 @@ def filter_func(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: # +---+---+ # | 1| 21| # +---+---+ + # $example off:map_in_pandas$ def cogrouped_apply_in_pandas_example(spark: SparkSession) -> None: + # $example on:cogrouped_apply_in_pandas$ import pandas as pd df1 = spark.createDataFrame( @@ -292,9 +307,11 @@ def merge_ordered(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame: # |20000101| 2|2.0| y| # |20000102| 2|4.0|null| # +--------+---+---+----+ + # $example off:cogrouped_apply_in_pandas$ def arrow_python_udf_example(spark: SparkSession) -> None: + # $example on:arrow_python_udf$ from pyspark.sql.functions import udf @udf(returnType='int') # A default, pickled Python UDF @@ -314,6 +331,8 @@ def arrow_slen(s): # type: ignore[no-untyped-def] # | 8| 8| # +----------+----------------+ + # $example off:arrow_python_udf$ + if __name__ == "__main__": spark = SparkSession \ diff --git a/examples/src/main/python/sql/basic.py b/examples/src/main/python/sql/basic.py index 4f7ec7ba267df..8514c0b77a5cc 100644 --- a/examples/src/main/python/sql/basic.py +++ b/examples/src/main/python/sql/basic.py @@ -21,16 +21,15 @@ ./bin/spark-submit examples/src/main/python/sql/basic.py """ # $example on:init_session$ -from pyspark.sql import SparkSession # $example off:init_session$ - # $example on:schema_inferring$ -from pyspark.sql import Row -# $example off:schema_inferring$ +from pyspark.sql import Row, SparkSession +# $example off:schema_inferring$ # $example on:programmatic_schema$ # Import data types -from pyspark.sql.types import StringType, StructType, StructField +from pyspark.sql.types import StringType, StructField, StructType + # $example off:programmatic_schema$ diff --git a/examples/src/main/python/sql/datasource.py b/examples/src/main/python/sql/datasource.py index 0a76376e2c813..0820f09b28535 100644 --- a/examples/src/main/python/sql/datasource.py +++ b/examples/src/main/python/sql/datasource.py @@ -20,9 +20,9 @@ Run with: ./bin/spark-submit examples/src/main/python/sql/datasource.py """ -from pyspark.sql import SparkSession # $example on:schema_merging$ -from pyspark.sql import Row +from pyspark.sql import Row, SparkSession + # $example off:schema_merging$ diff --git a/examples/src/main/python/sql/hive.py b/examples/src/main/python/sql/hive.py index fa1b975e2bfdc..58c5d44e4d53f 100644 --- a/examples/src/main/python/sql/hive.py +++ b/examples/src/main/python/sql/hive.py @@ -23,8 +23,8 @@ # $example on:spark_hive$ from os.path import abspath -from pyspark.sql import SparkSession -from pyspark.sql import Row +from pyspark.sql import Row, SparkSession + # $example off:spark_hive$ diff --git a/examples/src/main/python/sql/jdbc.py b/examples/src/main/python/sql/jdbc.py index 7e67ab351f634..0956adb072f8f 100644 --- a/examples/src/main/python/sql/jdbc.py +++ b/examples/src/main/python/sql/jdbc.py @@ -21,8 +21,8 @@ ./bin/spark-submit examples/src/main/python/sql/jdbc.py [jdbc_url] """ import sys -from pyspark.sql import SparkSession +from pyspark.sql import SparkSession if __name__ == "__main__": if len(sys.argv) < 2: diff --git a/examples/src/main/python/sql/streaming/structured_kafka_wordcount.py b/examples/src/main/python/sql/streaming/structured_kafka_wordcount.py index 40a955a46c9b9..d241f556022aa 100644 --- a/examples/src/main/python/sql/streaming/structured_kafka_wordcount.py +++ b/examples/src/main/python/sql/streaming/structured_kafka_wordcount.py @@ -39,8 +39,7 @@ import sys from pyspark.sql import SparkSession -from pyspark.sql.functions import explode -from pyspark.sql.functions import split +from pyspark.sql.functions import explode, split if __name__ == "__main__": if len(sys.argv) != 4: diff --git a/examples/src/main/python/sql/streaming/structured_network_wordcount.py b/examples/src/main/python/sql/streaming/structured_network_wordcount.py index c8f43c9dcf2eb..6242b541b72e0 100644 --- a/examples/src/main/python/sql/streaming/structured_network_wordcount.py +++ b/examples/src/main/python/sql/streaming/structured_network_wordcount.py @@ -30,8 +30,7 @@ import sys from pyspark.sql import SparkSession -from pyspark.sql.functions import explode -from pyspark.sql.functions import split +from pyspark.sql.functions import explode, split if __name__ == "__main__": if len(sys.argv) != 3: diff --git a/examples/src/main/python/sql/streaming/structured_network_wordcount_session_window.py b/examples/src/main/python/sql/streaming/structured_network_wordcount_session_window.py index 722d409792bf7..f27ad44db55a3 100644 --- a/examples/src/main/python/sql/streaming/structured_network_wordcount_session_window.py +++ b/examples/src/main/python/sql/streaming/structured_network_wordcount_session_window.py @@ -31,21 +31,19 @@ localhost 9999` """ import sys -from typing import Iterator, Any +from typing import Any, Iterator import pandas as pd - from pyspark.sql import SparkSession -from pyspark.sql.functions import explode -from pyspark.sql.functions import split +from pyspark.sql.functions import explode, split +from pyspark.sql.streaming.state import GroupState, GroupStateTimeout from pyspark.sql.types import ( LongType, StringType, - TimestampType, - StructType, StructField, + StructType, + TimestampType, ) -from pyspark.sql.streaming.state import GroupStateTimeout, GroupState if __name__ == "__main__": if len(sys.argv) != 3: diff --git a/examples/src/main/python/sql/streaming/structured_network_wordcount_windowed.py b/examples/src/main/python/sql/streaming/structured_network_wordcount_windowed.py index cc39d8afa6be9..28fba29ad04ef 100644 --- a/examples/src/main/python/sql/streaming/structured_network_wordcount_windowed.py +++ b/examples/src/main/python/sql/streaming/structured_network_wordcount_windowed.py @@ -42,9 +42,7 @@ import sys from pyspark.sql import SparkSession -from pyspark.sql.functions import explode -from pyspark.sql.functions import split -from pyspark.sql.functions import window +from pyspark.sql.functions import explode, split, window if __name__ == "__main__": if len(sys.argv) != 5 and len(sys.argv) != 4: diff --git a/examples/src/main/python/sql/streaming/structured_sessionization.py b/examples/src/main/python/sql/streaming/structured_sessionization.py index 78cb406650e6a..97e441e77c690 100644 --- a/examples/src/main/python/sql/streaming/structured_sessionization.py +++ b/examples/src/main/python/sql/streaming/structured_sessionization.py @@ -33,9 +33,7 @@ import sys from pyspark.sql import SparkSession -from pyspark.sql.functions import explode -from pyspark.sql.functions import split -from pyspark.sql.functions import count, session_window +from pyspark.sql.functions import count, explode, session_window, split if __name__ == "__main__": if len(sys.argv) != 3 and len(sys.argv) != 2: diff --git a/examples/src/main/python/sql/udtf.py b/examples/src/main/python/sql/udtf.py index bff5182f8e16e..8a36c8d9c47c4 100644 --- a/examples/src/main/python/sql/udtf.py +++ b/examples/src/main/python/sql/udtf.py @@ -22,7 +22,8 @@ """ # NOTE that this file is imported in the tutorials in PySpark documentation. -# The codes are referred via line numbers. See also `literalinclude` directive in Sphinx. +# The code blocks are referred to via `$example on/off$` markers, included in the +# docs with the `literalinclude` directive's `:start-after:`/`:end-before:` options. from pyspark.sql import SparkSession from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version @@ -32,13 +33,15 @@ def python_udtf_simple_example(spark: SparkSession) -> None: - + # $example on:python_udtf_simple_class$ # Define the UDTF class and implement the required `eval` method. class SquareNumbers: def eval(self, start: int, end: int): for num in range(start, end + 1): yield (num, num * num) + # $example off:python_udtf_simple_class$ + # $example on:python_udtf_simple_udtf$ from pyspark.sql.functions import lit, udtf # Create a UDTF using the class definition and the `udtf` function. @@ -53,10 +56,11 @@ def eval(self, start: int, end: int): # | 2| 4| # | 3| 9| # +---+-------+ + # $example off:python_udtf_simple_udtf$ def python_udtf_decorator_example(spark: SparkSession) -> None: - + # $example on:python_udtf_decorator$ from pyspark.sql.functions import lit, udtf # Define a UDTF using the `udtf` decorator directly on the class. @@ -75,10 +79,11 @@ def eval(self, start: int, end: int): # | 2| 4| # | 3| 9| # +---+-------+ + # $example off:python_udtf_decorator$ def python_udtf_registration(spark: SparkSession) -> None: - + # $example on:python_udtf_registration$ from pyspark.sql.functions import udtf @udtf(returnType="word: string") @@ -114,21 +119,24 @@ def eval(self, text: str): # |Apache Spark|Apache| # |Apache Spark| Spark| # +------------+------+ + # $example off:python_udtf_registration$ def python_udtf_arrow_example(spark: SparkSession) -> None: - + # $example on:python_udtf_arrow$ from pyspark.sql.functions import udtf @udtf(returnType="c1: int, c2: int", useArrow=True) class PlusOne: def eval(self, x: int): yield x, x + 1 + # $example off:python_udtf_arrow$ def python_udtf_date_expander_example(spark: SparkSession) -> None: - + # $example on:python_udtf_date_expander$ from datetime import datetime, timedelta + from pyspark.sql.functions import lit, udtf @udtf(returnType="date: string") @@ -150,10 +158,11 @@ def eval(self, start_date: str, end_date: str): # |2023-02-28| # |2023-03-01| # +----------+ + # $example off:python_udtf_date_expander$ def python_udtf_terminate_example(spark: SparkSession) -> None: - + # $example on:python_udtf_terminate$ from pyspark.sql.functions import udtf @udtf(returnType="cnt: int") @@ -184,10 +193,11 @@ def terminate(self): # | 4| 5| # | 9| 5| # +---+---+ + # $example off:python_udtf_terminate$ def python_udtf_table_argument(spark: SparkSession) -> None: - + # $example on:python_udtf_table_argument$ from pyspark.sql.functions import udtf from pyspark.sql.types import Row @@ -208,10 +218,11 @@ def eval(self, row: Row): # | 8| # | 9| # +---+ + # $example off:python_udtf_table_argument$ def python_udtf_table_argument_with_partitioning(spark: SparkSession) -> None: - + # $example on:python_udtf_table_argument_with_partitioning$ from pyspark.sql.functions import udtf from pyspark.sql.types import Row @@ -285,6 +296,7 @@ def terminate(self): # Clean up. spark.sql("DROP TABLE values_table") + # $example off:python_udtf_table_argument_with_partitioning$ if __name__ == "__main__": diff --git a/examples/src/main/python/status_api_demo.py b/examples/src/main/python/status_api_demo.py index 3bf96ca4466fa..15d22ddc8d9c6 100644 --- a/examples/src/main/python/status_api_demo.py +++ b/examples/src/main/python/status_api_demo.py @@ -15,9 +15,9 @@ # limitations under the License. # -import time -import threading import queue as Queue +import threading +import time from typing import Any, Callable, List, Tuple from pyspark import SparkConf, SparkContext diff --git a/examples/src/main/python/streaming/network_wordjoinsentiments.py b/examples/src/main/python/streaming/network_wordjoinsentiments.py index fae40a77acaff..a21bf9285341a 100644 --- a/examples/src/main/python/streaming/network_wordjoinsentiments.py +++ b/examples/src/main/python/streaming/network_wordjoinsentiments.py @@ -33,8 +33,7 @@ import sys from typing import Tuple -from pyspark import SparkContext -from pyspark import RDD +from pyspark import RDD, SparkContext from pyspark.streaming import DStream, StreamingContext diff --git a/examples/src/main/python/streaming/recoverable_network_wordcount.py b/examples/src/main/python/streaming/recoverable_network_wordcount.py index 147d3c646799c..e9d8513c2eb4d 100644 --- a/examples/src/main/python/streaming/recoverable_network_wordcount.py +++ b/examples/src/main/python/streaming/recoverable_network_wordcount.py @@ -40,7 +40,7 @@ import sys from typing import List, Tuple -from pyspark import SparkContext, Accumulator, Broadcast, RDD +from pyspark import RDD, Accumulator, Broadcast, SparkContext from pyspark.streaming import StreamingContext diff --git a/examples/src/main/python/streaming/sql_network_wordcount.py b/examples/src/main/python/streaming/sql_network_wordcount.py index bba398c0d6109..a1b1530637369 100644 --- a/examples/src/main/python/streaming/sql_network_wordcount.py +++ b/examples/src/main/python/streaming/sql_network_wordcount.py @@ -27,12 +27,12 @@ and then run the example `$ bin/spark-submit examples/src/main/python/streaming/sql_network_wordcount.py localhost 9999` """ -import sys import datetime +import sys -from pyspark import SparkConf, SparkContext, RDD -from pyspark.streaming import StreamingContext +from pyspark import RDD, SparkConf, SparkContext from pyspark.sql import Row, SparkSession +from pyspark.streaming import StreamingContext def getSparkSessionInstance(sparkConf: SparkConf) -> SparkSession: diff --git a/examples/src/main/python/wordcount.py b/examples/src/main/python/wordcount.py index 037c1e8aa379d..33ba6982368e1 100755 --- a/examples/src/main/python/wordcount.py +++ b/examples/src/main/python/wordcount.py @@ -20,7 +20,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": if len(sys.argv) != 2: print("Usage: wordcount <file>", file=sys.stderr) diff --git a/examples/src/main/scala/org/apache/spark/examples/SkewedGroupByTest.scala b/examples/src/main/scala/org/apache/spark/examples/SkewedGroupByTest.scala index 9001ab0bec995..7b3969e2ca892 100644 --- a/examples/src/main/scala/org/apache/spark/examples/SkewedGroupByTest.scala +++ b/examples/src/main/scala/org/apache/spark/examples/SkewedGroupByTest.scala @@ -23,7 +23,7 @@ import java.util.Random import org.apache.spark.sql.SparkSession /** - * Usage: GroupByTest [numMappers] [numKVPairs] [KeySize] [numReducers] + * Usage: SkewedGroupByTest [numMappers] [numKVPairs] [KeySize] [numReducers] */ object SkewedGroupByTest { def main(args: Array[String]): Unit = { diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/CustomReceiver.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/CustomReceiver.scala index 59d0641b62cef..d666c78605ffc 100644 --- a/examples/src/main/scala/org/apache/spark/examples/streaming/CustomReceiver.scala +++ b/examples/src/main/scala/org/apache/spark/examples/streaming/CustomReceiver.scala @@ -73,7 +73,7 @@ class CustomReceiver(host: String, port: Int) def onStop(): Unit = { // There is nothing much to do as the thread calling receive() - // is designed to stop by itself isStopped() returns false + // is designed to stop by itself when isStopped() returns true } /** Create a socket connection and receive data until receiver is stopped */ diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKafkaWordCount.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKafkaWordCount.scala index 6fdb37194ea7d..1c0e577b66c86 100644 --- a/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKafkaWordCount.scala +++ b/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKafkaWordCount.scala @@ -27,7 +27,7 @@ import org.apache.spark.streaming.kafka010._ /** * Consumes messages from one or more topics in Kafka and does wordcount. - * Usage: DirectKafkaWordCount <brokers> <topics> + * Usage: DirectKafkaWordCount <brokers> <groupId> <topics> * <brokers> is a list of one or more Kafka brokers * <groupId> is a consumer group name to consume from topics * <topics> is a list of one or more kafka topics to consume from diff --git a/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKerberizedKafkaWordCount.scala b/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKerberizedKafkaWordCount.scala index 6a35ce9b2a293..3736b0c989a06 100644 --- a/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKerberizedKafkaWordCount.scala +++ b/examples/src/main/scala/org/apache/spark/examples/streaming/DirectKerberizedKafkaWordCount.scala @@ -29,7 +29,7 @@ import org.apache.spark.streaming.kafka010._ /** * Consumes messages from one or more topics in Kafka and does wordcount. - * Usage: DirectKerberizedKafkaWordCount <brokers> <topics> + * Usage: DirectKerberizedKafkaWordCount <brokers> <groupId> <topics> * <brokers> is a list of one or more Kafka brokers * <groupId> is a consumer group name to consume from topics * <topics> is a list of one or more kafka topics to consume from diff --git a/graphx/src/main/scala/org/apache/spark/graphx/lib/SVDPlusPlus.scala b/graphx/src/main/scala/org/apache/spark/graphx/lib/SVDPlusPlus.scala index 7b282ec6d24ee..beb15b0ccbc56 100644 --- a/graphx/src/main/scala/org/apache/spark/graphx/lib/SVDPlusPlus.scala +++ b/graphx/src/main/scala/org/apache/spark/graphx/lib/SVDPlusPlus.scala @@ -38,6 +38,17 @@ object SVDPlusPlus { var gamma7: Double) extends Serializable + /** Sums two Phase 2 training messages component-wise: both vectors and the scalar. */ + private[graphx] def combineTrainMessages( + g1: (Array[Double], Array[Double], Double), + g2: (Array[Double], Array[Double], Double)): (Array[Double], Array[Double], Double) = { + val out1 = g1._1.clone() + BLAS.nativeBLAS.daxpy(out1.length, 1.0, g2._1, 1, out1, 1) + val out2 = g1._2.clone() + BLAS.nativeBLAS.daxpy(out2.length, 1.0, g2._2, 1, out2, 1) + (out1, out2, g1._3 + g2._3) + } + // scalastyle:off line.size.limit /** * Implement SVD++ based on "Factorization Meets the Neighborhood: a Multifaceted @@ -152,14 +163,7 @@ object SVDPlusPlus { g.cache() val t2 = g.aggregateMessages( sendMsgTrainF(conf, u), - (g1: (Array[Double], Array[Double], Double), g2: (Array[Double], Array[Double], Double)) => - { - val out1 = g1._1.clone() - BLAS.nativeBLAS.daxpy(out1.length, 1.0, g2._1, 1, out1, 1) - val out2 = g1._2.clone() - BLAS.nativeBLAS.daxpy(out2.length, 1.0, g2._2, 1, out2, 1) - (out1, out2, g1._3 + g2._3) - }) + combineTrainMessages) val gJoinT2 = g.outerJoinVertices(t2) { (vid: VertexId, vd: (Array[Double], Array[Double], Double, Double), diff --git a/graphx/src/test/scala/org/apache/spark/graphx/lib/SVDPlusPlusSuite.scala b/graphx/src/test/scala/org/apache/spark/graphx/lib/SVDPlusPlusSuite.scala index da139cec5c3d5..62f257d2cce3c 100644 --- a/graphx/src/test/scala/org/apache/spark/graphx/lib/SVDPlusPlusSuite.scala +++ b/graphx/src/test/scala/org/apache/spark/graphx/lib/SVDPlusPlusSuite.scala @@ -49,4 +49,19 @@ class SVDPlusPlusSuite extends SparkFunSuite with LocalSparkContext { assert(graph.edges.count() == 0) } } + + test("SPARK-58177: training-phase message combiner sums both vectors element-wise") { + val g1 = (Array(1.0, 2.0), Array(10.0, 20.0), 100.0) + val g2 = (Array(3.0, 4.0), Array(30.0, 40.0), 200.0) + val (out1, out2, out3) = SVDPlusPlus.combineTrainMessages(g1, g2) + // Each component is the element-wise sum of the two messages. + assert(out1.toSeq === Seq(4.0, 6.0)) + assert(out2.toSeq === Seq(40.0, 60.0)) + assert(out3 === 300.0) + // Inputs are cloned, not mutated. + assert(g1._1.toSeq === Seq(1.0, 2.0)) + assert(g1._2.toSeq === Seq(10.0, 20.0)) + assert(g2._1.toSeq === Seq(3.0, 4.0)) + assert(g2._2.toSeq === Seq(30.0, 40.0)) + } } diff --git a/hadoop-cloud/pom.xml b/hadoop-cloud/pom.xml index e87fbfae61a6a..853ebc4a678e3 100644 --- a/hadoop-cloud/pom.xml +++ b/hadoop-cloud/pom.xml @@ -34,8 +34,6 @@ </description> <properties> <sbt.project.name>hadoop-cloud</sbt.project.name> - <okhttp.version>3.12.12</okhttp.version> - <okio.version>1.17.6</okio.version> <wildfly-openssl.version>2.3.0.Final</wildfly-openssl.version> </properties> @@ -158,12 +156,6 @@ <version>${hadoop.version}</version> <scope>${hadoop.deps.scope}</scope> </dependency> - <dependency> - <groupId>org.apache.hadoop</groupId> - <artifactId>hadoop-huaweicloud</artifactId> - <version>${hadoop.version}</version> - <scope>${huaweicloud.deps.scope}</scope> - </dependency> <!-- There's now a hadoop-cloud-storage which transitively pulls in the store JARs, but it still needs some selective exclusion across versions, especially 3.0.x. @@ -197,24 +189,16 @@ <groupId>org.apache.hadoop</groupId> <artifactId>hadoop-tos</artifactId> </exclusion> + <!-- + SPARK-58671: Exclude `hadoop-huaweicloud` to keep the vulnerable okhttp 3.x + dependency off the classpath. + --> <exclusion> <groupId>org.apache.hadoop</groupId> <artifactId>hadoop-huaweicloud</artifactId> </exclusion> </exclusions> </dependency> - <dependency> - <groupId>com.squareup.okhttp3</groupId> - <artifactId>okhttp</artifactId> - <version>${okhttp.version}</version> - <scope>${huaweicloud.deps.scope}</scope> - </dependency> - <dependency> - <groupId>com.squareup.okio</groupId> - <artifactId>okio</artifactId> - <version>${okio.version}</version> - <scope>${huaweicloud.deps.scope}</scope> - </dependency> </dependencies> <build> @@ -242,12 +226,6 @@ </build> <profiles> - <profile> - <id>huaweicloud-provided</id> - <properties> - <huaweicloud.deps.scope>provided</huaweicloud.deps.scope> - </properties> - </profile> <profile> <id>integration-test</id> <build> diff --git a/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala b/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala index 17d6be0ce7cd5..38d4c43bc154b 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala @@ -19,7 +19,7 @@ package org.apache.spark.ml import org.apache.spark.annotation.Since import org.apache.spark.internal.{LogKeys} -import org.apache.spark.ml.linalg.VectorUDT +import org.apache.spark.ml.linalg.SQLDataTypes import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util.SchemaUtils @@ -135,7 +135,7 @@ abstract class Predictor[ * * The default value is VectorUDT, but it may be overridden if FeaturesType is not Vector. */ - private[ml] def featuresDataType: DataType = new VectorUDT + private[ml] def featuresDataType: DataType = SQLDataTypes.VectorType override def transformSchema(schema: StructType): StructType = { validateAndTransformSchema(schema, fitting = true, featuresDataType) @@ -171,7 +171,7 @@ abstract class PredictionModel[FeaturesType, M <: PredictionModel[FeaturesType, * * The default value is VectorUDT, but it may be overridden if FeaturesType is not Vector. */ - protected def featuresDataType: DataType = new VectorUDT + protected def featuresDataType: DataType = SQLDataTypes.VectorType override def transformSchema(schema: StructType): StructType = { var outputSchema = validateAndTransformSchema(schema, fitting = false, featuresDataType) diff --git a/mllib/src/main/scala/org/apache/spark/ml/ann/Layer.scala b/mllib/src/main/scala/org/apache/spark/ml/ann/Layer.scala index 889f6febbe40d..7b7ec5f21e56a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/ann/Layer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/ann/Layer.scala @@ -479,10 +479,13 @@ private[ml] class FeedForwardModel private( val layers = topology.layers val layerModels = new Array[LayerModel](layers.length) + // Trained weights are normally dense, but Vector and persisted models do not guarantee it. + // Convert once so all layer views share one dense backing array. + private val denseWeights = weights.toArray private var offset = 0 for (i <- layers.indices) { layerModels(i) = layers(i).createModel( - new BDV[Double](weights.toArray, offset, 1, layers(i).weightSize)) + new BDV[Double](denseWeights, offset, 1, layers(i).weightSize)) offset += layers(i).weightSize } private var outputs: Array[BDM[Double]] = null diff --git a/mllib/src/main/scala/org/apache/spark/ml/attribute/AttributeGroup.scala b/mllib/src/main/scala/org/apache/spark/ml/attribute/AttributeGroup.scala index f2fe125db67f9..9bb94355e65f2 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/attribute/AttributeGroup.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/attribute/AttributeGroup.scala @@ -19,7 +19,7 @@ package org.apache.spark.ml.attribute import scala.collection.mutable.ArrayBuffer -import org.apache.spark.ml.linalg.VectorUDT +import org.apache.spark.ml.linalg.SQLDataTypes import org.apache.spark.sql.types.{Metadata, MetadataBuilder, StructField} import org.apache.spark.util.ArrayImplicits._ @@ -155,7 +155,7 @@ class AttributeGroup private ( /** Converts to a StructField with some existing metadata. */ def toStructField(existingMetadata: Metadata): StructField = { - StructField(name, new VectorUDT, nullable = false, toMetadata(existingMetadata)) + StructField(name, SQLDataTypes.VectorType, nullable = false, toMetadata(existingMetadata)) } /** Converts to a StructField. */ @@ -237,7 +237,7 @@ object AttributeGroup { * Creates an attribute group from a `StructField` instance. */ def fromStructField(field: StructField): AttributeGroup = { - require(field.dataType == new VectorUDT) + require(field.dataType == SQLDataTypes.VectorType) if (field.metadata.contains(ML_ATTR)) { fromMetadata(field.metadata.getMetadata(ML_ATTR), field.name) } else { diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala index d9238479e8031..2772bdb8c3f60 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala @@ -20,7 +20,7 @@ package org.apache.spark.ml.classification import org.apache.spark.annotation.Since import org.apache.spark.internal.{LogKeys} import org.apache.spark.ml.{PredictionModel, Predictor, PredictorParams} -import org.apache.spark.ml.linalg.{Vector, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector} import org.apache.spark.ml.param.ParamMap import org.apache.spark.ml.param.shared.HasRawPredictionCol import org.apache.spark.ml.util._ @@ -39,7 +39,7 @@ private[spark] trait ClassifierParams fitting: Boolean, featuresDataType: DataType): StructType = { val parentSchema = super.validateAndTransformSchema(schema, fitting, featuresDataType) - SchemaUtils.appendColumn(parentSchema, $(rawPredictionCol), new VectorUDT) + SchemaUtils.appendColumn(parentSchema, $(rawPredictionCol), SQLDataTypes.VectorType) } } @@ -199,13 +199,12 @@ abstract class ClassificationModel[FeaturesType, M <: ClassificationModel[Featur (ClassificationModel[FeaturesType, M], String, String) = { val model = if ($(rawPredictionCol).isEmpty && $(predictionCol).isEmpty) { copy(ParamMap.empty) - .setRawPredictionCol("rawPrediction_" + java.util.UUID.randomUUID.toString) - .setPredictionCol("prediction_" + java.util.UUID.randomUUID.toString) + .setRawPredictionCol(Identifiable.randomUID("rawPrediction")) + .setPredictionCol(Identifiable.randomUID("prediction")) } else if ($(rawPredictionCol).isEmpty) { - copy(ParamMap.empty).setRawPredictionCol("rawPrediction_" + - java.util.UUID.randomUUID.toString) + copy(ParamMap.empty).setRawPredictionCol(Identifiable.randomUID("rawPrediction")) } else if ($(predictionCol).isEmpty) { - copy(ParamMap.empty).setPredictionCol("prediction_" + java.util.UUID.randomUUID.toString) + copy(ParamMap.empty).setPredictionCol(Identifiable.randomUID("prediction")) } else { this } diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala index 3506c1a2502d0..11d43d1b31fb1 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala @@ -201,6 +201,8 @@ class DecisionTreeClassificationModel private[ml] ( // For ml connect only private[ml] def this() = this("", Node.dummyNode, -1, -1) + override private[ml] val treeStats: NodeStats = rootNode.computeStats + private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize() override def predict(features: Vector): Double = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala index 38c68fd7000ca..f819266238098 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala @@ -17,8 +17,6 @@ package org.apache.spark.ml.classification -import java.util.UUID - import scala.concurrent.Future import scala.concurrent.duration.Duration import scala.language.existentials @@ -150,8 +148,14 @@ final class OneVsRestModel private[ml] ( val numFeatures: Int = models.head.numFeatures private[spark] override def estimatedSize: Long = { - estimateMatadataSize + SizeEstimator.estimate(labelMetadata) + - models.iterator.map(_.estimatedSize).sum + var size = estimateMatadataSize(excluded = Seq( + // classifier: Param[ClassifierType] + classifier)) + // labelMetadata: Metadata + size += SizeEstimator.estimate(labelMetadata) + // models: Array[_ <: ClassificationModel[_, _]] + size += models.iterator.map(_.estimatedSize).sum + size } /** @group setParam */ @@ -195,10 +199,10 @@ final class OneVsRestModel private[ml] ( val isProbModel = models.head.isInstanceOf[ProbabilisticClassificationModel[_, _]] // use a temporary raw prediction column to avoid column conflict - val tmpRawPredName = "mbc$raw" + UUID.randomUUID().toString + val tmpRawPredName = Identifiable.randomUID("mbc$raw") // add an accumulator column to store predictions of all the models - val accColName = "mbc$acc" + UUID.randomUUID().toString + val accColName = Identifiable.randomUID("mbc$acc") val newDataset = dataset.withColumn(accColName, lit(Array.emptyDoubleArray)) val columns = newDataset.schema.fieldNames.map(col) diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala index ea2c79d8a2181..a6648f4fab7b0 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala @@ -19,10 +19,10 @@ package org.apache.spark.ml.classification import org.apache.spark.annotation.Since import org.apache.spark.internal.{LogKeys} -import org.apache.spark.ml.linalg.{DenseVector, Vector, VectorUDT} +import org.apache.spark.ml.linalg.{DenseVector, SQLDataTypes, Vector} import org.apache.spark.ml.param.ParamMap import org.apache.spark.ml.param.shared._ -import org.apache.spark.ml.util.SchemaUtils +import org.apache.spark.ml.util.{Identifiable, SchemaUtils} import org.apache.spark.sql.{DataFrame, Dataset} import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.{DataType, StructType} @@ -37,7 +37,7 @@ private[ml] trait ProbabilisticClassifierParams fitting: Boolean, featuresDataType: DataType): StructType = { val parentSchema = super.validateAndTransformSchema(schema, fitting, featuresDataType) - SchemaUtils.appendColumn(parentSchema, $(probabilityCol), new VectorUDT) + SchemaUtils.appendColumn(parentSchema, $(probabilityCol), SQLDataTypes.VectorType) } } @@ -241,12 +241,12 @@ abstract class ProbabilisticClassificationModel[ (ProbabilisticClassificationModel[FeaturesType, M], String, String) = { val model = if ($(probabilityCol).isEmpty && $(predictionCol).isEmpty) { copy(ParamMap.empty) - .setProbabilityCol("probability_" + java.util.UUID.randomUUID.toString) - .setPredictionCol("prediction_" + java.util.UUID.randomUUID.toString) + .setProbabilityCol(Identifiable.randomUID("probability")) + .setPredictionCol(Identifiable.randomUID("prediction")) } else if ($(probabilityCol).isEmpty) { - copy(ParamMap.empty).setProbabilityCol("probability_" + java.util.UUID.randomUUID.toString) + copy(ParamMap.empty).setProbabilityCol(Identifiable.randomUID("probability")) } else if ($(predictionCol).isEmpty) { - copy(ParamMap.empty).setPredictionCol("prediction_" + java.util.UUID.randomUUID.toString) + copy(ParamMap.empty).setPredictionCol(Identifiable.randomUID("prediction")) } else { this } diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala index e63a3eaa6455a..6b149102e253d 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/BisectingKMeans.scala @@ -182,11 +182,13 @@ class BisectingKMeansModel private[ml] ( private[spark] override def estimatedSize: Long = { var size = estimateMatadataSize if (parentModel != null) { - // parentModel contains: - // - root: ClusteringTreeNode containing centers, costs, and children. - // - distanceMeasure: String and trainingCost: Double. - // - distanceMeasureInstance, derived from distanceMeasure. - size += SizeEstimator.estimate(parentModel) + // root: ClusteringTreeNode + // distanceMeasure: String + // trainingCost: Double + size += SizeEstimator.estimate(( + parentModel.root, + parentModel.distanceMeasure, + parentModel.trainingCost)) } size } diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala index 1b85f7f998751..817802746bbf5 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/GaussianMixture.scala @@ -74,7 +74,7 @@ private[clustering] trait GaussianMixtureParams extends Params with HasMaxIter w protected def validateAndTransformSchema(schema: StructType): StructType = { SchemaUtils.validateVectorCompatibleColumn(schema, getFeaturesCol) val schemaWithPredictionCol = SchemaUtils.appendColumn(schema, $(predictionCol), IntegerType) - SchemaUtils.appendColumn(schemaWithPredictionCol, $(probabilityCol), new VectorUDT) + SchemaUtils.appendColumn(schemaWithPredictionCol, $(probabilityCol), SQLDataTypes.VectorType) } } @@ -225,8 +225,9 @@ class GaussianMixtureModel private[ml] ( private[spark] override def estimatedSize: Long = { var size = estimateMatadataSize // weights: Array[Double] - // gaussians: Array[MultivariateGaussian], each containing a mean Vector and covariance Matrix - size += SizeEstimator.estimate((weights, gaussians)) + size += SizeEstimator.estimate(weights) + // gaussians: Array[MultivariateGaussian], each with mean: Vector and cov: Matrix + gaussians.foreach(gaussian => size += SizeEstimator.estimate((gaussian.mean, gaussian.cov))) size } diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala index 2430773742967..5d410f77c7eba 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala @@ -216,9 +216,15 @@ class KMeansModel private[ml] ( private[spark] override def estimatedSize: Long = { var size = estimateMatadataSize if (parentModel != null) { - // parentModel contains clusterCenters, distanceMeasure, trainingCost, and numIter. - // It also has transient derived fields for distance calculation and statistics. - size += SizeEstimator.estimate(parentModel) + // clusterCenters: Array[Vector] + // distanceMeasure: String + // trainingCost: Double + // numIter: Int + size += SizeEstimator.estimate(( + parentModel.clusterCenters, + parentModel.distanceMeasure, + parentModel.trainingCost, + parentModel.numIter)) } size } diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala index e4736718430b8..1f072bd852d2a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/LDA.scala @@ -157,7 +157,7 @@ private[clustering] trait LDAParams extends Params with HasFeaturesCol with HasM /** Supported values for Param [[optimizer]]. */ @Since("1.6.0") - final val supportedOptimizers: Array[String] = Array("online", "em") + final val supportedOptimizers: Array[String] = LDA.supportedOptimizers /** * Optimizer or inference algorithm used to estimate the LDA model. @@ -180,7 +180,7 @@ private[clustering] trait LDAParams extends Params with HasFeaturesCol with HasM @Since("1.6.0") final val optimizer = new Param[String](this, "optimizer", "Optimizer or inference" + " algorithm used to estimate the LDA model. Supported: " + supportedOptimizers.mkString(", "), - (value: String) => supportedOptimizers.contains(value.toLowerCase(Locale.ROOT))) + (value: String) => LDA.supportedOptimizers.contains(value.toLowerCase(Locale.ROOT))) /** @group getParam */ @Since("1.6.0") @@ -355,7 +355,7 @@ private[clustering] trait LDAParams extends Params with HasFeaturesCol with HasM } } SchemaUtils.validateVectorCompatibleColumn(schema, getFeaturesCol) - SchemaUtils.appendColumn(schema, $(topicDistributionCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(topicDistributionCol), SQLDataTypes.VectorType) } private[clustering] def getOldOptimizer: OldLDAOptimizer = @@ -640,6 +640,21 @@ class LocalLDAModel private[ml] ( override def toString: String = { s"LocalLDAModel: uid=$uid, k=${$(k)}, numFeatures=$vocabSize" } + + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + if (oldLocalModel != null) { + // topicsMatrix: Matrix + if (oldLocalModel.topicsMatrix != null) { + size += oldLocalModel.topicsMatrix.asML.getSizeInBytes + } + // docConcentration: Vector + if (oldLocalModel.docConcentration != null) { + size += oldLocalModel.docConcentration.asML.getSizeInBytes + } + } + size + } } @Since("1.6.0") @@ -827,12 +842,15 @@ class DistributedLDAModel private[ml] ( } private[spark] override def estimatedSize: Long = { - this.oldDistributedModel.toInternals.map { + var size = estimateMatadataSize + // oldDistributedModel: metadata, global topic totals, graph vertices, and graph edges. + oldDistributedModel.toInternals.foreach { case df: org.apache.spark.sql.classic.DataFrame => - df.toArrowBatchRdd.map(_.length.toLong).reduce(_ + _) + size += df.toArrowBatchRdd.map(_.length.toLong).reduce(_ + _) case o => throw new UnsupportedOperationException( s"Unsupported dataframe type: ${o.getClass.getName}") - }.sum + } + size } } @@ -1056,6 +1074,8 @@ class LDA @Since("1.6.0") ( @Since("2.0.0") object LDA extends MLReadable[LDA] { + private[clustering] val supportedOptimizers: Array[String] = Array("online", "em") + /** Get dataset for spark.mllib LDA */ private[clustering] def getOldDataset( dataset: Dataset[_], diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala index 1a97eb2910056..8183cd34365d3 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluator.scala @@ -18,7 +18,7 @@ package org.apache.spark.ml.evaluation import org.apache.spark.annotation.Since -import org.apache.spark.ml.linalg.{Vector, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util._ @@ -115,7 +115,8 @@ class BinaryClassificationEvaluator @Since("1.4.0") (@Since("1.4.0") override va @Since("3.1.0") def getMetrics(dataset: Dataset[_]): BinaryClassificationMetrics = { val schema = dataset.schema - SchemaUtils.checkColumnTypes(schema, $(rawPredictionCol), Seq(DoubleType, new VectorUDT)) + SchemaUtils.checkColumnTypes(schema, $(rawPredictionCol), + Seq(DoubleType, SQLDataTypes.VectorType)) SchemaUtils.checkNumericType(schema, $(labelCol)) if (isDefined(weightCol)) { SchemaUtils.checkNumericType(schema, $(weightCol)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/evaluation/ClusteringMetrics.scala b/mllib/src/main/scala/org/apache/spark/ml/evaluation/ClusteringMetrics.scala index 98fbe471f2977..12735f835e68c 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/evaluation/ClusteringMetrics.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/evaluation/ClusteringMetrics.scala @@ -330,13 +330,10 @@ private[evaluation] object SquaredEuclideanSilhouette extends Silhouette { } ) - clustersStatsRDD - .collectAsMap() - .toMap - .transform { - case (_, (featureSum: DenseVector, squaredNormSum: Double, weightSum: Double)) => - SquaredEuclideanSilhouette.ClusterStats(featureSum, squaredNormSum, weightSum) - } + clustersStatsRDD.mapValues { + case (featureSum: DenseVector, squaredNormSum: Double, weightSum: Double) => + SquaredEuclideanSilhouette.ClusterStats(featureSum, squaredNormSum, weightSum) + }.collectAsMap().toMap } /** diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala index 52ed90415f1cd..50952010581fe 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Binarizer.scala @@ -218,7 +218,7 @@ final class Binarizer @Since("1.4.0") (@Since("1.4.0") override val uid: String) SchemaUtils.getSchemaField(schema, inputColName) ).size if (size < 0) { - StructField(outputColName, new VectorUDT) + StructField(outputColName, SQLDataTypes.VectorType) } else { new AttributeGroup(outputColName, numAttributes = size).toStructField() } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala index 3d1765417775e..76a01fd3e346f 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/BucketedRandomProjectionLSH.scala @@ -95,10 +95,14 @@ class BucketedRandomProjectionLSHModel private[ml]( @Since("2.1.0") override protected[ml] def hashFunction(elems: Vector): Array[Vector] = { - val hashVec = new DenseVector(Array.ofDim[Double](randMatrix.numRows)) - BLAS.gemv(1.0 / $(bucketLength), randMatrix, elems, 0.0, hashVec) - // TODO: Output vectors of dimension numHashFunctions in SPARK-18450 - hashVec.values.map(h => Vectors.dense(h.floor)) + BucketedRandomProjectionLSHModel.hashFunction(elems, randMatrix, $(bucketLength)) + } + + override protected[ml] def createTransformFunc: Vector => Array[Vector] = { + val localRandMatrix = randMatrix + val localBucketLength = $(bucketLength) + elems => BucketedRandomProjectionLSHModel.hashFunction( + elems, localRandMatrix, localBucketLength) } @Since("2.1.0") @@ -205,7 +209,7 @@ class BucketedRandomProjectionLSH(override val uid: String) @Since("2.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) validateAndTransformSchema(schema) } @@ -222,6 +226,17 @@ object BucketedRandomProjectionLSH extends DefaultParamsReadable[BucketedRandomP @Since("2.1.0") object BucketedRandomProjectionLSHModel extends MLReadable[BucketedRandomProjectionLSHModel] { + + private def hashFunction( + elems: Vector, + randMatrix: Matrix, + bucketLength: Double): Array[Vector] = { + val hashVec = new DenseVector(Array.ofDim[Double](randMatrix.numRows)) + BLAS.gemv(1.0 / bucketLength, randMatrix, elems, 0.0, hashVec) + // TODO: Output vectors of dimension numHashFunctions in SPARK-18450 + hashVec.values.map(h => Vectors.dense(h.floor)) + } + // TODO: Save using the existing format of Array[Vector] once SPARK-12878 is resolved. private[ml] case class Data(randUnitVectors: Matrix) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala index a85e9236517f7..a0da92fa7d7e4 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala @@ -17,6 +17,7 @@ package org.apache.spark.ml.feature import java.io.{DataInputStream, DataOutputStream} +import java.util.{HashMap => JHashMap} import org.apache.hadoop.fs.Path @@ -24,17 +25,16 @@ import org.apache.spark.annotation.Since import org.apache.spark.broadcast.Broadcast import org.apache.spark.ml.{Estimator, Model} import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute} -import org.apache.spark.ml.linalg.{Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vectors} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol} import org.apache.spark.ml.util._ import org.apache.spark.sql.{DataFrame, Dataset} import org.apache.spark.sql.functions._ import org.apache.spark.sql.types._ -import org.apache.spark.storage.StorageLevel import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.SizeEstimator -import org.apache.spark.util.collection.{OpenHashMap, Utils} +import org.apache.spark.util.collection.OpenHashMap /** * Params for [[CountVectorizer]] and [[CountVectorizerModel]]. @@ -99,7 +99,7 @@ private[feature] trait CountVectorizerParams extends Params with HasInputCol wit protected def validateAndTransformSchema(schema: StructType): StructType = { val typeCandidates = List(new ArrayType(StringType, true), new ArrayType(StringType, false)) SchemaUtils.checkColumnTypes(schema, $(inputCol), typeCandidates) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } /** @@ -192,61 +192,40 @@ class CountVectorizer @Since("1.5.0") (@Since("1.5.0") override val uid: String) } val vocSize = $(vocabSize) - val input = dataset.select($(inputCol)).rdd.map(_.getSeq[String](0)) - val countingRequired = $(minDF) < 1.0 || $(maxDF) < 1.0 - val maybeInputSize = if (countingRequired) { - if (dataset.storageLevel == StorageLevel.NONE) { - input.persist(StorageLevel.MEMORY_AND_DISK) - } - Some(input.count()) - } else { - None - } - val minDf = if ($(minDF) >= 1.0) { - $(minDF) - } else { - $(minDF) * maybeInputSize.get - } - val maxDf = if ($(maxDF) >= 1.0) { - $(maxDF) - } else { - $(maxDF) * maybeInputSize.get - } - require(maxDf >= minDf, "maxDF must be >= minDF.") - val allWordCounts = input.flatMap { tokens => - val wc = new OpenHashMap[String, Long] - tokens.foreach { w => - wc.changeValue(w, 1L, _ + 1L) - } - wc.map { case (word, count) => (word, (count, 1)) } - }.reduceByKey { (wcdf1, wcdf2) => - (wcdf1._1 + wcdf2._1, wcdf1._2 + wcdf2._2) - } - + val input = dataset.select($(inputCol)) + val inputSizeCol = input.select(count(lit(0))).scalar() + val minDfCol = if ($(minDF) >= 1.0) lit($(minDF)) else inputSizeCol * lit($(minDF)) + val maxDfCol = if ($(maxDF) >= 1.0) lit($(maxDF)) else inputSizeCol * lit($(maxDF)) val filteringRequired = isSet(minDF) || isSet(maxDF) - val maybeFilteredWordCounts = if (filteringRequired) { - allWordCounts.filter { case (_, (_, df)) => df >= minDf && df <= maxDf } + val wordCounts = if (filteringRequired) { + input + .select( + monotonically_increasing_id().as("docId"), + col($(inputCol)).as("doc")) + .select( + col("docId"), + explode(col("doc")).as("word")) + .groupBy("docId", "word") + .agg(count(lit(0)).as("wordCountInDoc")) + .groupBy("word") + .agg( + sum("wordCountInDoc").as("count"), + count(lit(0)).as("docCount")) + .filter(minDfCol <= col("docCount") && col("docCount") <= maxDfCol) + .select("word", "count") } else { - allWordCounts + input + .select(explode(col($(inputCol))).as("word")) + .groupBy("word") + .agg(count(lit(0)).as("count")) } - val wordCounts = maybeFilteredWordCounts - .map { case (word, (count, _)) => (word, count) } - .persist(StorageLevel.MEMORY_AND_DISK) - - val fullVocabSize = wordCounts.count() - - val ordering = Ordering.Tuple2(Ordering.Long, Ordering.String.reverse) - .on[(String, Long)] { case (word, count) => (count, word) } - val vocab = wordCounts - .top(math.min(fullVocabSize, vocSize).toInt)(ordering) - .map(_._1) - - if (input.getStorageLevel != StorageLevel.NONE) { - input.unpersist() - } - wordCounts.unpersist() + .orderBy(col("count").desc, col("word").asc) + .limit(vocSize) + .select("word") + .collect() + .map(_.getString(0)) if (vocab.isEmpty) { this.logWarning("The vocabulary size is empty. " + @@ -316,38 +295,43 @@ class CountVectorizerModel( def setBinary(value: Boolean): this.type = set(binary, value) /** Dictionary created from [[vocabulary]] and its indices, broadcast once for [[transform()]] */ - private var broadcastDict: Option[Broadcast[Map[String, Int]]] = None + private var broadcastDict: Option[Broadcast[JHashMap[String, Integer]]] = None @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { val outputSchema = transformSchema(dataset.schema, logging = true) if (broadcastDict.isEmpty) { - val dict = Utils.toMapWithIndex(vocabulary) + val dict = new JHashMap[String, Integer](math.ceil(vocabulary.length / 0.75).toInt) + var index = 0 + while (index < vocabulary.length) { + dict.put(vocabulary(index), index) + index += 1 + } broadcastDict = Some(dataset.sparkSession.sparkContext.broadcast(dict)) } val dictBr = broadcastDict.get - // SPARK-48837: capture parameter values here so that we only evaulate once-per-transform - // rather than once-per-row: - val minTf = $(minTF) - val isBinary = $(binary) + val localMinTF = $(minTF) + val localBinary = $(binary) val vectorizer = udf { document: Seq[String] => - val termCounts = new OpenHashMap[Int, Double] + val dict = dictBr.value + val dictSize = dict.size() + val termCounts = new OpenHashMap[Int, Int] var tokenCount = 0L document.foreach { term => - dictBr.value.get(term) match { - case Some(index) => termCounts.changeValue(index, 1.0, _ + 1.0) - case None => // ignore terms not in the vocabulary + val index = dict.get(term) + if (index != null) { + termCounts.changeValue(index.intValue(), 1, _ + 1) } tokenCount += 1 } - val effectiveMinTF = if (minTf >= 1.0) minTf else tokenCount * minTf - val effectiveCounts = if (isBinary) { + val effectiveMinTF = if (localMinTF >= 1.0) localMinTF else tokenCount * localMinTF + val effectiveCounts = if (localBinary) { termCounts.filter(_._2 >= effectiveMinTF).map(p => (p._1, 1.0)).toSeq } else { - termCounts.filter(_._2 >= effectiveMinTF).toSeq + termCounts.filter(_._2 >= effectiveMinTF).map(p => (p._1, p._2.toDouble)).toSeq } - Vectors.sparse(dictBr.value.size, effectiveCounts) + Vectors.sparse(dictSize, effectiveCounts) } dataset.withColumn($(outputCol), vectorizer(col($(inputCol))), outputSchema($(outputCol)).metadata) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala index 9a8bfb195666b..d6e66d83aed62 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/DCT.scala @@ -22,7 +22,7 @@ import org.jtransforms.dct._ import org.apache.spark.annotation.Since import org.apache.spark.ml.UnaryTransformer import org.apache.spark.ml.attribute.AttributeGroup -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors, VectorUDT} import org.apache.spark.ml.param.BooleanParam import org.apache.spark.ml.util._ import org.apache.spark.sql.types._ @@ -62,19 +62,32 @@ class DCT @Since("1.5.0") (@Since("1.5.0") override val uid: String) setDefault(inverse -> false) - override protected def createTransformFunc: Vector => Vector = { vec => - val result = vec.toArray - val jTransformer = new DoubleDCT_1D(result.length) - if ($(inverse)) jTransformer.inverse(result, true) else jTransformer.forward(result, true) - Vectors.dense(result) + override protected def createTransformFunc: Vector => Vector = { + $(inverse) match { + case true => + (vec: Vector) => { + val result = vec.toArray + val jTransformer = new DoubleDCT_1D(result.length) + jTransformer.inverse(result, true) + Vectors.dense(result) + } + case false => + (vec: Vector) => { + val result = vec.toArray + val jTransformer = new DoubleDCT_1D(result.length) + jTransformer.forward(result, true) + Vectors.dense(result) + } + } } override protected def validateInputType(inputType: DataType): Unit = { require(inputType.isInstanceOf[VectorUDT], - s"Input type must be ${(new VectorUDT).catalogString} but got ${inputType.catalogString}.") + s"Input type must be ${SQLDataTypes.VectorType.catalogString} " + + s"but got ${inputType.catalogString}.") } - override protected def outputDataType: DataType = new VectorUDT + override protected def outputDataType: DataType = SQLDataTypes.VectorType override def transformSchema(schema: StructType): StructType = { var outputSchema = super.transformSchema(schema) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/ElementwiseProduct.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/ElementwiseProduct.scala index 6dac09d13b99b..4e38a91cdb660 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/ElementwiseProduct.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/ElementwiseProduct.scala @@ -79,10 +79,11 @@ class ElementwiseProduct @Since("1.4.0") (@Since("1.4.0") override val uid: Stri override protected def validateInputType(inputType: DataType): Unit = { require(inputType.isInstanceOf[VectorUDT], - s"Input type must be ${(new VectorUDT).catalogString} but got ${inputType.catalogString}.") + s"Input type must be ${SQLDataTypes.VectorType.catalogString} " + + s"but got ${inputType.catalogString}.") } - override protected def outputDataType: DataType = new VectorUDT() + override protected def outputDataType: DataType = SQLDataTypes.VectorType override def transformSchema(schema: StructType): StructType = { var outputSchema = super.transformSchema(schema) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala index dab0a6494fdb9..0610465fe7fa1 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/HashingTF.scala @@ -91,12 +91,37 @@ class HashingTF @Since("3.0.0") private[ml] ( override def transform(dataset: Dataset[_]): DataFrame = { val outputSchema = transformSchema(dataset.schema) val n = $(numFeatures) - val updateFunc = if ($(binary)) (v: Double) => 1.0 else (v: Double) => v + 1.0 + def binaryHashUDF(hashFunc: Any => Int) = udf { terms: Seq[_] => + val map = new OpenHashMap[Int, Int] + terms.foreach { term => + map.update(Utils.nonNegativeMod(hashFunc(term), n), 1) + } + Vectors.sparse(n, map.iterator.map { case (index, count) => + (index, count.toDouble) + }.toSeq) + } + + def countHashUDF(hashFunc: Any => Int) = udf { terms: Seq[_] => + val map = new OpenHashMap[Int, Int] + terms.foreach { term => + val index = Utils.nonNegativeMod(hashFunc(term), n) + map.changeValue(index, 1, _ + 1) + } + Vectors.sparse(n, map.iterator.map { case (index, count) => + (index, count.toDouble) + }.toSeq) + } - val hashUDF = udf { terms: Seq[_] => - val map = new OpenHashMap[Int, Double]() - terms.foreach { term => map.changeValue(indexOf(term), 1.0, updateFunc) } - Vectors.sparse(n, map.toSeq) + val hashUDF = ($(binary), hashFuncVersion) match { + case (true, HashingTF.SPARK_2_MURMUR3_HASH) => + binaryHashUDF(OldHashingTF.murmur3Hash) + case (false, HashingTF.SPARK_2_MURMUR3_HASH) => + countHashUDF(OldHashingTF.murmur3Hash) + case (true, HashingTF.SPARK_3_MURMUR3_HASH) => + binaryHashUDF(FeatureHasher.murmur3Hash) + case (false, HashingTF.SPARK_3_MURMUR3_HASH) => + countHashUDF(FeatureHasher.murmur3Hash) + case _ => throw new IllegalArgumentException("Illegal hash function version setting.") } dataset.withColumn($(outputCol), hashUDF(col($(inputCol))), diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala index 3ef916bc200c0..6d1dae9392d0e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala @@ -60,8 +60,8 @@ private[feature] trait IDFBase extends Params with HasInputCol with HasOutputCol * Validate and transform the input schema. */ protected def validateAndTransformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } } @@ -152,20 +152,8 @@ class IDFModel private[ml] ( override def transform(dataset: Dataset[_]): DataFrame = { val outputSchema = transformSchema(dataset.schema, logging = true) - val func = { vector: Vector => - vector match { - case SparseVector(size, indices, values) => - val (newIndices, newValues) = feature.IDFModel.transformSparse(idfModel.idf, - indices, values) - Vectors.sparse(size, newIndices, newValues) - case DenseVector(values) => - val newValues = feature.IDFModel.transformDense(idfModel.idf, values) - Vectors.dense(newValues) - case other => - throw new UnsupportedOperationException( - s"Only sparse and dense vectors are supported but got ${other.getClass}.") - } - } + val localIdf = idfModel.idf.toArray + val func = (vector: Vector) => IDFModel.predict(localIdf, vector) val transformer = udf(func) dataset.withColumn($(outputCol), transformer(col($(inputCol))), @@ -213,6 +201,45 @@ class IDFModel private[ml] ( object IDFModel extends MLReadable[IDFModel] { private[ml] case class Data(idf: Vector, docFreq: Array[Long], numDocs: Long) + private def predict(idf: Array[Double], v: Vector): Vector = { + v match { + case SparseVector(size, indices, values) => + val (newIndices, newValues) = predictSparse(idf, indices, values) + Vectors.sparse(size, newIndices, newValues) + case DenseVector(values) => + val newValues = predictDense(idf, values) + Vectors.dense(newValues) + case other => + throw new UnsupportedOperationException( + s"Only sparse and dense vectors are supported but got ${other.getClass}.") + } + } + + private[spark] def predictDense(idf: Array[Double], values: Array[Double]): Array[Double] = { + val n = values.length + val newValues = new Array[Double](n) + var j = 0 + while (j < n) { + newValues(j) = values(j) * idf(j) + j += 1 + } + newValues + } + + private[spark] def predictSparse( + idf: Array[Double], + indices: Array[Int], + values: Array[Double]): (Array[Int], Array[Double]) = { + val nnz = indices.length + val newValues = new Array[Double](nnz) + var k = 0 + while (k < nnz) { + newValues(k) = values(k) * idf(indices(k)) + k += 1 + } + (indices, newValues) + } + private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = { import ReadWriteUtils._ serializeVector(data.idf, dos) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala index 3311231e6d830..7fa1f8e6a1de4 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Interaction.scala @@ -23,7 +23,7 @@ import org.apache.spark.SparkException import org.apache.spark.annotation.Since import org.apache.spark.ml.Transformer import org.apache.spark.ml.attribute._ -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors, VectorUDT} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util._ @@ -64,7 +64,7 @@ class Interaction @Since("1.6.0") (@Since("1.6.0") override val uid: String) ext require(get(outputCol).isDefined, "Output col must be defined first.") require($(inputCols).length > 0, "Input cols must have non-zero length.") require($(inputCols).distinct.length == $(inputCols).length, "Input cols must be distinct.") - StructType(schema.fields :+ StructField($(outputCol), new VectorUDT, false)) + StructType(schema.fields :+ StructField($(outputCol), SQLDataTypes.VectorType, false)) } @Since("2.0.0") diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala index 9c3b39b12bdc6..620b5b1ff8c16 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/LSH.scala @@ -18,7 +18,7 @@ package org.apache.spark.ml.feature import org.apache.spark.ml.{Estimator, Model} -import org.apache.spark.ml.linalg.{Vector, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector} import org.apache.spark.ml.param.{IntParam, ParamValidators} import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol} import org.apache.spark.ml.util._ @@ -53,7 +53,8 @@ private[ml] trait LSHParams extends HasInputCol with HasOutputCol { * @return A derived schema with [[outputCol]] added. */ protected[this] final def validateAndTransformSchema(schema: StructType): StructType = { - SchemaUtils.appendColumn(schema, $(outputCol), DataTypes.createArrayType(new VectorUDT)) + SchemaUtils.appendColumn(schema, $(outputCol), + DataTypes.createArrayType(SQLDataTypes.VectorType)) } } @@ -76,6 +77,12 @@ private[spark] abstract class LSHModel[T <: LSHModel[T]] */ protected[ml] def hashFunction(elems: Vector): Array[Vector] + /** + * Returns the hash function used by [[transform]]. The returned function must not capture this + * model. + */ + protected[ml] def createTransformFunc: Vector => Array[Vector] + /** * Calculate the distance between two different keys using the distance metric corresponding * to the hashFunction. @@ -96,7 +103,7 @@ private[spark] abstract class LSHModel[T <: LSHModel[T]] override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema, logging = true) - val transformUDF = udf(hashFunction(_: Vector)) + val transformUDF = udf(createTransformFunc) dataset.withColumn($(outputCol), transformUDF(dataset($(inputCol)))) } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MaxAbsScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MaxAbsScaler.scala index a962ce3784c9a..05e38f7e70489 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/MaxAbsScaler.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MaxAbsScaler.scala @@ -23,7 +23,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since import org.apache.spark.ml.{Estimator, Model} -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.param.{ParamMap, Params} import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol} import org.apache.spark.ml.stat.Summarizer @@ -39,10 +39,10 @@ private[feature] trait MaxAbsScalerParams extends Params with HasInputCol with H /** Validates and transforms the input schema. */ protected def validateAndTransformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) require(!schema.fieldNames.contains($(outputCol)), s"Output column ${$(outputCol)} already exists.") - val outputFields = schema.fields :+ StructField($(outputCol), new VectorUDT, false) + val outputFields = schema.fields :+ StructField($(outputCol), SQLDataTypes.VectorType, false) StructType(outputFields) } } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala index 4b0e5ca4fb31a..ea7b2d773e43f 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala @@ -24,7 +24,7 @@ import scala.util.Random import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.param.ParamMap import org.apache.spark.ml.param.shared.HasSeed import org.apache.spark.ml.util._ @@ -71,14 +71,12 @@ class MinHashLSHModel private[ml]( @Since("2.1.0") override protected[ml] def hashFunction(elems: Vector): Array[Vector] = { - require(elems.nonZeroIterator.nonEmpty, "Must have at least 1 non zero entry.") - val hashValues = randCoefficients.map { case (a, b) => - elems.nonZeroIterator.map { case (i, _) => - ((1L + i) * a + b) % MinHashLSH.HASH_PRIME - }.min.toDouble - } - // TODO: Output vectors of dimension numHashFunctions in SPARK-18450 - hashValues.map(Vectors.dense(_)) + MinHashLSHModel.hashFunction(elems, randCoefficients) + } + + override protected[ml] def createTransformFunc: Vector => Array[Vector] = { + val localRandCoefficients = randCoefficients + elems => MinHashLSHModel.hashFunction(elems, localRandCoefficients) } @Since("2.1.0") @@ -201,7 +199,7 @@ class MinHashLSH(override val uid: String) extends LSH[MinHashLSHModel] with Has @Since("2.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) validateAndTransformSchema(schema) } @@ -220,6 +218,20 @@ object MinHashLSH extends DefaultParamsReadable[MinHashLSH] { @Since("2.1.0") object MinHashLSHModel extends MLReadable[MinHashLSHModel] { + + private def hashFunction( + elems: Vector, + randCoefficients: Array[(Int, Int)]): Array[Vector] = { + require(elems.nonZeroIterator.nonEmpty, "Must have at least 1 non zero entry.") + val hashValues = randCoefficients.map { case (a, b) => + elems.nonZeroIterator.map { case (i, _) => + ((1L + i) * a + b) % MinHashLSH.HASH_PRIME + }.min.toDouble + } + // TODO: Output vectors of dimension numHashFunctions in SPARK-18450 + hashValues.map(Vectors.dense(_)) + } + private[ml] case class Data(randCoefficients: Array[Int]) private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala index 9bf13c48b22aa..339710a6f280d 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinMaxScaler.scala @@ -23,7 +23,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since import org.apache.spark.ml.{Estimator, Model} -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.param.{DoubleParam, ParamMap, Params} import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol} import org.apache.spark.ml.stat.Summarizer @@ -64,10 +64,10 @@ private[feature] trait MinMaxScalerParams extends Params with HasInputCol with H /** Validates and transforms the input schema. */ protected def validateAndTransformSchema(schema: StructType): StructType = { require($(min) < $(max), s"The specified min(${$(min)}) is larger or equal to max(${$(max)})") - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) require(!schema.fieldNames.contains($(outputCol)), s"Output column ${$(outputCol)} already exists.") - val outputFields = schema.fields :+ StructField($(outputCol), new VectorUDT, false) + val outputFields = schema.fields :+ StructField($(outputCol), SQLDataTypes.VectorType, false) StructType(outputFields) } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala index d72fb6ecc76d3..3ef7a5a887198 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/NGram.scala @@ -61,7 +61,8 @@ class NGram @Since("1.5.0") (@Since("1.5.0") override val uid: String) setDefault(n -> 2) override protected def createTransformFunc: Seq[String] => Seq[String] = { - _.iterator.sliding($(n)).withPartial(false).map(_.mkString(" ")).toSeq + val localN = $(n) + _.iterator.sliding(localN).withPartial(false).map(_.mkString(" ")).toSeq } override protected def validateInputType(inputType: DataType): Unit = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala index c7b7164e42f36..bdd749e7c485c 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Normalizer.scala @@ -20,11 +20,10 @@ package org.apache.spark.ml.feature import org.apache.spark.annotation.Since import org.apache.spark.ml.UnaryTransformer import org.apache.spark.ml.attribute.AttributeGroup -import org.apache.spark.ml.linalg.{Vector, VectorUDT} +import org.apache.spark.ml.linalg.{DenseVector, SparseVector, SQLDataTypes, Vector, Vectors, + VectorUDT} import org.apache.spark.ml.param.{DoubleParam, ParamValidators} import org.apache.spark.ml.util._ -import org.apache.spark.mllib.feature -import org.apache.spark.mllib.linalg.{Vectors => OldVectors} import org.apache.spark.sql.types._ /** @@ -56,16 +55,51 @@ class Normalizer @Since("1.4.0") (@Since("1.4.0") override val uid: String) def setP(value: Double): this.type = set(p, value) override protected def createTransformFunc: Vector => Vector = { - val normalizer = new feature.Normalizer($(p)) - vector => normalizer.transform(OldVectors.fromML(vector)).asML + val localP = $(p) + vector => { + val norm = Vectors.norm(vector, localP) + if (norm != 0.0) { + val scale = 1.0 / norm + // For dense vector, we've to allocate new memory for new output vector. + // However, for sparse vector, the `index` array will not be changed, + // so we can re-use it to save memory. + vector match { + case DenseVector(vs) => + val values = vs.clone() + val size = values.length + var i = 0 + while (i < size) { + values(i) *= scale + i += 1 + } + Vectors.dense(values) + case SparseVector(size, ids, vs) => + val values = vs.clone() + val nnz = values.length + var i = 0 + while (i < nnz) { + values(i) *= scale + i += 1 + } + Vectors.sparse(size, ids, values) + case v => throw new IllegalArgumentException("Do not support vector type " + v.getClass) + } + } else { + // Since the norm is zero, return the input vector object itself. + // Note that it's safe since we always assume that the data in RDD + // should be immutable. + vector + } + } } override protected def validateInputType(inputType: DataType): Unit = { require(inputType.isInstanceOf[VectorUDT], - s"Input type must be ${(new VectorUDT).catalogString} but got ${inputType.catalogString}.") + s"Input type must be ${SQLDataTypes.VectorType.catalogString} " + + s"but got ${inputType.catalogString}.") } - override protected def outputDataType: DataType = new VectorUDT() + override protected def outputDataType: DataType = SQLDataTypes.VectorType @Since("1.4.0") override def transformSchema(schema: StructType): StructType = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala index c5120d026d476..5599aa73a4e8b 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala @@ -31,8 +31,8 @@ import org.apache.spark.ml.param.shared.{HasHandleInvalid, HasInputCol, HasInput import org.apache.spark.ml.util._ import org.apache.spark.sql.{DataFrame, Dataset} import org.apache.spark.sql.expressions.UserDefinedFunction -import org.apache.spark.sql.functions.{col, lit, udf} -import org.apache.spark.sql.types.{DoubleType, StructField, StructType} +import org.apache.spark.sql.functions.{printf => fprintf, _} +import org.apache.spark.sql.types.{DoubleType, IntegerType, StructField, StructType} import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.SizeEstimator @@ -529,34 +529,24 @@ private[feature] object OneHotEncoderCommon { inputColNames: Seq[String], outputColNames: Seq[String], dropLast: Boolean): Seq[AttributeGroup] = { - // The RDD approach has advantage of early-stop if any values are invalid. It seems that - // DataFrame ops don't have equivalent functions. val columns = inputColNames.map { inputColName => - col(inputColName).cast(DoubleType) + val doubleCol = col(inputColName).cast(DoubleType) + val intCol = doubleCol.cast(IntegerType) + val invalidIndexError = raise_error(fprintf( + lit(s"Values from column $inputColName must be indices, but got %s."), doubleCol)) + val maxIndexError = raise_error(fprintf( + lit(s"OneHotEncoder only supports up to ${Int.MaxValue} indices, but got %s."), + doubleCol)) + when( + isnull(doubleCol) || isnan(doubleCol) || doubleCol < 0.0 || doubleCol =!= intCol, + invalidIndexError) + .when(doubleCol > Int.MaxValue, maxIndexError) + .otherwise(intCol) } - val numOfColumns = columns.length - - val numAttrsArray = dataset.select(columns: _*).rdd.map { row => - (0 until numOfColumns).map(idx => row.getDouble(idx)).toArray - }.treeAggregate(new Array[Double](numOfColumns))( - (maxValues, curValues) => { - (0 until numOfColumns).foreach { idx => - val x = curValues(idx) - assert(x <= Int.MaxValue, - s"OneHotEncoder only supports up to ${Int.MaxValue} indices, but got $x.") - assert(x >= 0.0 && x == x.toInt, - s"Values from column ${inputColNames(idx)} must be indices, but got $x.") - maxValues(idx) = math.max(maxValues(idx), x) - } - maxValues - }, - (m0, m1) => { - (0 until numOfColumns).foreach { idx => - m0(idx) = math.max(m0(idx), m1(idx)) - } - m0 - } - ).map(_.toInt + 1) + + val maxValues = columns.map(c => coalesce(max(c) + 1, lit(1))) + val maxValuesRow = dataset.select(maxValues: _*).head() + val numAttrsArray = Array.tabulate(maxValues.length)(maxValuesRow.getInt) outputColNames.zip(numAttrsArray).map { case (outputColName, numAttrs) => createAttrGroupForAttrNames(outputColName, numAttrs, dropLast, keepInvalid = false) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala index a731b6750573c..f6067f03e76fa 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala @@ -51,7 +51,7 @@ private[feature] trait PCAParams extends Params with HasInputCol with HasOutputC /** Validates and transforms the input schema. */ protected def validateAndTransformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) require(!schema.fieldNames.contains($(outputCol)), s"Output column ${$(outputCol)} already exists.") SchemaUtils.updateAttributeGroupSize(schema, $(outputCol), $(k)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala index 592ca001a2467..c89da4394ce85 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PolynomialExpansion.scala @@ -64,16 +64,18 @@ class PolynomialExpansion @Since("1.4.0") (@Since("1.4.0") override val uid: Str @Since("1.4.0") def setDegree(value: Int): this.type = set(degree, value) - override protected def createTransformFunc: Vector => Vector = { v => - PolynomialExpansion.expand(v, $(degree)) + override protected def createTransformFunc: Vector => Vector = { + val localDegree = $(degree) + v => PolynomialExpansion.expand(v, localDegree) } override protected def validateInputType(inputType: DataType): Unit = { require(inputType.isInstanceOf[VectorUDT], - s"Input type must be ${(new VectorUDT).catalogString} but got ${inputType.catalogString}.") + s"Input type must be ${SQLDataTypes.VectorType.catalogString} " + + s"but got ${inputType.catalogString}.") } - override protected def outputDataType: DataType = new VectorUDT() + override protected def outputDataType: DataType = SQLDataTypes.VectorType @Since("1.4.1") override def copy(extra: ParamMap): PolynomialExpansion = defaultCopy(extra) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala index 844cedb4a302d..3b760d1783a14 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala @@ -27,7 +27,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since import org.apache.spark.ml.{Estimator, Model, Pipeline, PipelineModel, PipelineStage, Transformer} import org.apache.spark.ml.attribute.AttributeGroup -import org.apache.spark.ml.linalg.{Vector, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, VectorUDT} import org.apache.spark.ml.param.{BooleanParam, Param, ParamMap, ParamValidators} import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasHandleInvalid, HasLabelCol} import org.apache.spark.ml.util._ @@ -314,9 +314,9 @@ class RFormula @Since("1.5.0") (@Since("1.5.0") override val uid: String) require(!hasLabelCol(schema) || !$(forceIndexLabel), "If label column already exists, forceIndexLabel can not be set with true.") if (hasLabelCol(schema)) { - StructType(schema.fields :+ StructField($(featuresCol), new VectorUDT, true)) + StructType(schema.fields :+ StructField($(featuresCol), SQLDataTypes.VectorType, true)) } else { - StructType(schema.fields :+ StructField($(featuresCol), new VectorUDT, true) :+ + StructType(schema.fields :+ StructField($(featuresCol), SQLDataTypes.VectorType, true) :+ StructField($(labelCol), DoubleType, true)) } } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/RobustScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/RobustScaler.scala index 0b520542566c8..7a44076b6037c 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/RobustScaler.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/RobustScaler.scala @@ -92,10 +92,10 @@ private[feature] trait RobustScalerParams extends Params with HasInputCol with H protected def validateAndTransformSchema(schema: StructType): StructType = { require($(lower) < $(upper), s"The specified lower quantile(${$(lower)}) is " + s"larger or equal to upper quantile(${$(upper)})") - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) require(!schema.fieldNames.contains($(outputCol)), s"Output column ${$(outputCol)} already exists.") - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala index 1914a98014daa..a654c74c3582e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala @@ -257,9 +257,9 @@ private[ml] abstract class Selector[T <: SelectorModel[T]] @Since("3.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) SchemaUtils.checkNumericType(schema, $(labelCol)) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } @Since("3.1.0") @@ -301,7 +301,7 @@ private[ml] abstract class SelectorModel[T <: SelectorModel[T]] ( @Since("3.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) val newField = SelectorModel.prepOutputField(schema, selectedFeatures, $(outputCol), $(featuresCol), isNumericAttribute) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala index fd61753c25dda..d21384787fe52 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StandardScaler.scala @@ -62,10 +62,10 @@ private[feature] trait StandardScalerParams extends Params with HasInputCol with /** Validates and transforms the input schema. */ protected def validateAndTransformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) require(!schema.fieldNames.contains($(outputCol)), s"Output column ${$(outputCol)} already exists.") - val outputFields = schema.fields :+ StructField($(outputCol), new VectorUDT, false) + val outputFields = schema.fields :+ StructField($(outputCol), SQLDataTypes.VectorType, false) StructType(outputFields) } diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala index 8bb07d3f02608..c363d3ca2f326 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala @@ -316,20 +316,8 @@ class StringIndexerModel ( labelsArray(0) } - // Prepares the maps for string values to corresponding index values. - private val labelsToIndexArray: Array[OpenHashMap[String, Double]] = { - for (labels <- labelsArray) yield { - val n = labels.length - val map = new OpenHashMap[String, Double](n) - labels.zipWithIndex.foreach { case (label, idx) => - map.update(label, idx) - } - map - } - } - private[spark] override def estimatedSize: Long = - estimateMatadataSize + SizeEstimator.estimate((labelsArray, labelsToIndexArray)) + estimateMatadataSize + SizeEstimator.estimate(labelsArray) /** @group setParam */ @Since("1.6.0") @@ -353,7 +341,10 @@ class StringIndexerModel ( // This filters out any null values and also the input labels which are not in // the dataset used for fitting. - private def filterInvalidData(dataset: Dataset[_], inputColNames: Seq[String]): Dataset[_] = { + private def filterInvalidData( + dataset: Dataset[_], + inputColNames: Seq[String], + labelsToIndexArray: Array[OpenHashMap[String, Int]]): Dataset[_] = { val conditions: Seq[Column] = inputColNames.indices.map { i => val inputColName = inputColNames(i) val labelToIndex = labelsToIndexArray(i) @@ -372,28 +363,32 @@ class StringIndexerModel ( .where(conditions.reduce(_ and _)) } - private def getIndexer(labels: Seq[String], labelToIndex: OpenHashMap[String, Double]) = { - val keepInvalid = (getHandleInvalid == StringIndexer.KEEP_INVALID) - - udf { label: String => - if (label == null) { - if (keepInvalid) { - labels.length + private def getIndexer( + labels: Seq[String], + labelToIndex: OpenHashMap[String, Int], + keepInvalid: Boolean) = { + val unknownIndex = labels.length + if (keepInvalid) { + udf { label: String => + if (label == null) { + unknownIndex } else { + labelToIndex.get(label).getOrElse(unknownIndex) + } + }.asNondeterministic() + } else { + udf { label: String => + if (label == null) { throw new SparkException("StringIndexer encountered NULL value. To handle or skip " + "NULLS, try setting StringIndexer.handleInvalid.") - } - } else { - if (labelToIndex.contains(label)) { - labelToIndex(label) - } else if (keepInvalid) { - labels.length } else { - throw new SparkException(s"Unseen label: $label. To handle unseen labels, " + - s"set Param handleInvalid to ${StringIndexer.KEEP_INVALID}.") + labelToIndex.get(label).getOrElse { + throw new SparkException(s"Unseen label: $label. To handle unseen labels, " + + s"set Param handleInvalid to ${StringIndexer.KEEP_INVALID}.") + } } - } - }.asNondeterministic() + }.asNondeterministic() + } } @Since("2.0.0") @@ -401,11 +396,19 @@ class StringIndexerModel ( transformSchema(dataset.schema, logging = true) val (inputColNames, outputColNames) = getInOutCols() + val labelsToIndexArray = labelsArray.map { labels => + val map = new OpenHashMap[String, Int](labels.length) + labels.zipWithIndex.foreach { case (label, idx) => + map.update(label, idx) + } + map + } val outputColumns = new Array[Column](outputColNames.length) + val keepInvalid = getHandleInvalid == StringIndexer.KEEP_INVALID // Skips invalid rows if `handleInvalid` is set to `StringIndexer.SKIP_INVALID`. val filteredDataset = if (getHandleInvalid == StringIndexer.SKIP_INVALID) { - filterInvalidData(dataset, inputColNames.toImmutableArraySeq) + filterInvalidData(dataset, inputColNames.toImmutableArraySeq, labelsToIndexArray) } else { dataset } @@ -418,18 +421,15 @@ class StringIndexerModel ( try { dataset.col(inputColName) - val filteredLabels = getHandleInvalid match { - case StringIndexer.KEEP_INVALID => labels :+ "__unknown" - case _ => labels - } + val filteredLabels = if (keepInvalid) labels :+ "__unknown" else labels val metadata = NominalAttribute.defaultAttr .withName(outputColName) .withValues(filteredLabels) .toMetadata() - val indexer = getIndexer(labels.toImmutableArraySeq, labelToIndex) + val indexer = getIndexer(labels.toImmutableArraySeq, labelToIndex, keepInvalid) - outputColumns(i) = indexer(dataset(inputColName).cast(StringType)) + outputColumns(i) = indexer(dataset(inputColName).cast(StringType)).cast(DoubleType) .as(outputColName, metadata) } catch { case _: AnalysisException => diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala index 6eb5696520cda..088b1086014bd 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala @@ -26,7 +26,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since import org.apache.spark.ml.{Estimator, Model} import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NominalAttribute, NumericAttribute} -import org.apache.spark.ml.linalg.{DenseVector, SparseVector, Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{DenseVector, SparseVector, SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared.{HasFeaturesCol, HasLabelCol, HasOutputCol} import org.apache.spark.ml.stat.{ANOVATest, ChiSquareTest, FValueTest} @@ -266,9 +266,9 @@ final class UnivariateFeatureSelector @Since("3.1.1")(@Since("3.1.1") override v } } require(isSet(featureType) && isSet(labelType), "featureType and labelType need to be set") - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) SchemaUtils.checkNumericType(schema, $(labelCol)) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } @Since("3.1.1") @@ -320,7 +320,7 @@ class UnivariateFeatureSelectorModel private[ml]( @Since("3.1.1") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) val newField = UnivariateFeatureSelectorModel .prepOutputField(schema, selectedFeatures, $(outputCol), $(featuresCol), isNumericAttribute) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala index cdbdf122dc05e..23b9a4bc67b75 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala @@ -104,8 +104,8 @@ with DefaultParamsWritable { @Since("3.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } @Since("3.1.0") @@ -159,7 +159,7 @@ class VarianceThresholdSelectorModel private[ml]( @Since("3.1.0") override def transformSchema(schema: StructType): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) val newField = SelectorModel.prepOutputField(schema, selectedFeatures, $(outputCol), $(featuresCol), true) SchemaUtils.appendColumn(schema, newField) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala index 831a8a33afecb..5d4156b69065e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorAssembler.scala @@ -25,7 +25,7 @@ import org.apache.spark.SparkException import org.apache.spark.annotation.Since import org.apache.spark.ml.Transformer import org.apache.spark.ml.attribute.{Attribute, AttributeGroup, NumericAttribute, UnresolvedAttribute} -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors, VectorUDT} import org.apache.spark.ml.param.{Param, ParamMap, ParamValidators} import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util._ @@ -173,7 +173,7 @@ class VectorAssembler @Since("1.4.0") (@Since("1.4.0") override val uid: String) if (schema.fieldNames.contains(outputColName)) { throw new IllegalArgumentException(s"Output column $outputColName already exists.") } - StructType(schema.fields :+ new StructField(outputColName, new VectorUDT, true)) + StructType(schema.fields :+ new StructField(outputColName, SQLDataTypes.VectorType, true)) } @Since("1.4.1") diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala index 57d28bba2eeb5..6827464219947 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala @@ -29,7 +29,7 @@ import org.apache.spark.SparkException import org.apache.spark.annotation.Since import org.apache.spark.ml.{Estimator, Model} import org.apache.spark.ml.attribute._ -import org.apache.spark.ml.linalg.{DenseVector, SparseVector, Vector, VectorUDT} +import org.apache.spark.ml.linalg.{DenseVector, SparseVector, SQLDataTypes, Vector} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util._ @@ -160,11 +160,10 @@ class VectorIndexer @Since("1.4.0") ( override def transformSchema(schema: StructType): StructType = { // We do not transfer feature metadata since we do not know what types of features we will // produce in transform(). - val dataType = new VectorUDT require(isDefined(inputCol), s"VectorIndexer requires input column parameter: $inputCol") require(isDefined(outputCol), s"VectorIndexer requires output column parameter: $outputCol") - SchemaUtils.checkColumnType(schema, $(inputCol), dataType) - SchemaUtils.appendColumn(schema, $(outputCol), dataType) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } @Since("1.4.1") @@ -366,17 +365,15 @@ class VectorIndexerModel private[ml] ( attrs } - // TODO: Check more carefully about whether this whole class will be included in a closure. - /** Per-vector transform function */ - private lazy val transformFunc: Vector => Vector = { + private def getTransformFunc: Vector => Vector = { val sortedCatFeatureIndices = categoryMaps.keys.toArray.sorted val localVectorMap = categoryMaps val localNumFeatures = numFeatures val localHandleInvalid = getHandleInvalid val f: Vector => Vector = { (v: Vector) => assert(v.size == localNumFeatures, "VectorIndexerModel expected vector of length" + - s" $numFeatures but found length ${v.size}") + s" $localNumFeatures but found length ${v.size}") v match { case dv: DenseVector => var hasInvalid = false @@ -449,7 +446,7 @@ class VectorIndexerModel private[ml] ( override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema, logging = true) val newField = prepOutputField(dataset.schema) - val transformUDF = udf { vector: Vector => transformFunc(vector) } + val transformUDF = udf(getTransformFunc) val newCol = transformUDF(dataset($(inputCol))) val ds = dataset.withColumn($(outputCol), newCol, newField.metadata) if (getHandleInvalid == VectorIndexer.SKIP_INVALID) { @@ -461,12 +458,11 @@ class VectorIndexerModel private[ml] ( @Since("1.4.0") override def transformSchema(schema: StructType): StructType = { - val dataType = new VectorUDT require(isDefined(inputCol), s"VectorIndexerModel requires input column parameter: $inputCol") require(isDefined(outputCol), s"VectorIndexerModel requires output column parameter: $outputCol") - SchemaUtils.checkColumnType(schema, $(inputCol), dataType) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) // If the input metadata specifies numFeatures, compare with expected numFeatures. val origAttrGroup = AttributeGroup.fromStructField( diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala index 58a44a41f0e84..c5a49cf24374f 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala @@ -150,7 +150,7 @@ final class VectorSlicer @Since("1.5.0") (@Since("1.5.0") override val uid: Stri override def transformSchema(schema: StructType): StructType = { require($(indices).length > 0 || $(names).length > 0, s"VectorSlicer requires that at least one feature be selected.") - SchemaUtils.checkColumnType(schema, $(inputCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(inputCol), SQLDataTypes.VectorType) if (schema.fieldNames.contains($(outputCol))) { throw new IllegalArgumentException(s"Output column ${$(outputCol)} already exists.") diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala index 2baf6260479f5..ae1039fc81c09 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala @@ -24,7 +24,7 @@ import org.apache.hadoop.fs.Path import org.apache.spark.annotation.Since import org.apache.spark.internal.config.Kryo.KRYO_SERIALIZER_MAX_BUFFER_SIZE import org.apache.spark.ml.{Estimator, Model} -import org.apache.spark.ml.linalg.{BLAS, Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{BLAS, SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util._ @@ -113,7 +113,7 @@ private[feature] trait Word2VecBase extends Params protected def validateAndTransformSchema(schema: StructType): StructType = { val typeCandidates = List(new ArrayType(StringType, true), new ArrayType(StringType, false)) SchemaUtils.checkColumnTypes(schema, $(inputCol), typeCandidates) - SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(outputCol), SQLDataTypes.VectorType) } } diff --git a/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala b/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala index 0a5213228a8b2..378bbcac95b13 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/fpm/FPGrowth.scala @@ -273,35 +273,30 @@ class FPGrowthModel private[ml] ( * from all the applicable rules as prediction. The prediction column has the same data type as * the input column(Array[T]) and will not contain existing items in the input column. The null * values in the itemsCol columns are treated as empty sets. - * WARNING: internally it collects association rules to the driver and uses broadcast for - * efficiency. This may bring pressure to driver memory for large set of association rules. + * Internally, transform aggregates association rules into a single-row DataFrame and joins it + * with the input dataset. */ @Since("2.2.0") override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema, logging = true) - genericTransform(dataset) - } - - private def genericTransform(dataset: Dataset[_]): DataFrame = { - val rules: Array[(Seq[Any], Seq[Any])] = associationRules.select("antecedent", "consequent") - .rdd.map(r => (r.getSeq(0), r.getSeq(1))) - .collect().asInstanceOf[Array[(Seq[Any], Seq[Any])]] - val brRules = dataset.sparkSession.sparkContext.broadcast(rules) - val dt = dataset.schema($(itemsCol)).dataType - // For each rule, examine the input items and summarize the consequents - val predictUDF = SparkUserDefinedFunction((items: Seq[Any]) => { + val rulesCol = Identifiable.randomUID("rules") + // For each rule, examine the input items and summarize the consequents. + val predictFunc = (items: Seq[Any], rules: Seq[Row]) => { if (items != null) { val itemset = items.toSet - brRules.value.filter(_._1.forall(itemset.contains)) - .flatMap(_._2.filter(!itemset.contains(_))).distinct + rules.filter(_.getSeq[Any](0).forall(itemset.contains)) + .flatMap(_.getSeq[Any](1).filter(!itemset.contains(_))).distinct } else { Seq.empty - }}, - dt, - Nil - ) - dataset.withColumn($(predictionCol), predictUDF(col($(itemsCol)))) + } + } + val predictUDF = SparkUserDefinedFunction(predictFunc, dt, Nil) + dataset.join( + associationRules.select("antecedent", "consequent") + .agg(collect_set(struct("antecedent", "consequent")).as(rulesCol))) + .withColumn($(predictionCol), predictUDF(col($(itemsCol)), col(rulesCol))) + .drop(rulesCol) } @Since("2.2.0") diff --git a/mllib/src/main/scala/org/apache/spark/ml/functions.scala b/mllib/src/main/scala/org/apache/spark/ml/functions.scala index 07db59e53ba7d..b98d0523ab510 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/functions.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/functions.scala @@ -44,6 +44,21 @@ object functions { */ def array_to_vector(v: Column): Column = Column.internalFn("array_to_vector", v) + /** + * Creates a new row for each index-value pair in the given vector column. This expression is + * dedicated only for Spark ML. It always emits a marker row with index `-1 - vector.size` and + * value `Double.NaN` before each non-null vector. + * @param v: the column of MLlib sparse/dense vectors + * @param mode: `dense` emits all elements, and `sparse` emits nonzero elements + * @return the index and value columns of the vector elements + * @since 4.4.0 + */ + private[ml] def vector_posexplode( + v: Column, + mode: String = "sparse"): Column = { + Column.internalFn("vector_posexplode", sf.unwrap_udt(v), sf.lit(mode)) + } + private[ml] def array_binary_search(a: Column, v: Column): Column = Column.internalFn("array_binary_search", a, v) diff --git a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala index 21ccea479c1d1..2e826e473e607 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/param/params.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/param/params.scala @@ -657,6 +657,37 @@ trait Params extends Identifiable with Serializable { SizeEstimator.estimate((this.paramMap, this.defaultParamMap, this.uid)) } + /** + * Estimates metadata size while omitting params that can retain shared runtime state. + * + * A `Param` may hold this state indirectly through a validation closure that captures its + * owning instance, which can retain a `SparkSession`. Its value may also be a complex object, + * such as an `Estimator` in a `Param[Estimator[_]]`, that retains a `SparkSession`. Exclude + * such params to avoid accounting for shared infrastructure. + * + * @param excluded params to omit from the estimate + */ + private[ml] def estimateMatadataSize(excluded: Seq[Param[_]]): Long = { + if (excluded.isEmpty) { + estimateMatadataSize + } else { + val filteredParamMap = mutable.Map.empty[Param[Any], Any] + paramMap.toSeq.foreach { pair => + if (!excluded.contains(pair.param)) { + filteredParamMap(pair.param.asInstanceOf[Param[Any]]) = pair.value + } + } + val filteredDefaultParamMap = mutable.Map.empty[Param[Any], Any] + defaultParamMap.toSeq.foreach { pair => + if (!excluded.contains(pair.param)) { + filteredDefaultParamMap(pair.param.asInstanceOf[Param[Any]]) = pair.value + } + } + SizeEstimator.estimate(( + new ParamMap(filteredParamMap), new ParamMap(filteredDefaultParamMap), uid)) + } + } + /** * Returns all params sorted by their names. The default implementation uses Java reflection to * list all public methods that have no arguments and return [[Param]]. diff --git a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala index 9a54b70f77772..b2c854b08a1c9 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/recommendation/ALS.scala @@ -293,20 +293,6 @@ class ALSModel private[ml] ( @Since("3.0.0") def setBlockSize(value: Int): this.type = set(blockSize, value) - private val predict = udf { (featuresA: Seq[Float], featuresB: Seq[Float]) => - if (featuresA != null && featuresB != null) { - var dotProduct = 0.0f - var i = 0 - while (i < rank) { - dotProduct += featuresA(i) * featuresB(i) - i += 1 - } - dotProduct - } else { - Float.NaN - } - } - @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema) @@ -326,7 +312,8 @@ class ALSModel private[ml] ( .join(itemFactors.alias(itemFactorsAlias), col(s"${validatedInputAlias}.${$(itemCol)}") === col(s"${itemFactorsAlias}.id"), "left") .select(col(s"${validatedInputAlias}.*"), - predict(col(s"${userFactorsAlias}.features"), col(s"${itemFactorsAlias}.features")) + ALSModel.getPredictUDF(rank)( + col(s"${userFactorsAlias}.features"), col(s"${itemFactorsAlias}.features")) .alias($(predictionCol))) getColdStartStrategy match { @@ -473,8 +460,10 @@ class ALSModel private[ml] ( var scores: Array[Float] = null var idxOrd: GuavaOrdering[Int] = null iter.flatMap { case (srcIds, srcMat, dstIds, dstMat) => - require(srcMat.length == srcIds.length * rank) - require(dstMat.length == dstIds.length * rank) + require(srcMat.length == srcIds.length * rank, + s"srcMat must have ${srcIds.length * rank} entries but has ${srcMat.length}.") + require(dstMat.length == dstIds.length * rank, + s"dstMat must have ${dstIds.length * rank} entries but has ${dstMat.length}.") val m = srcIds.length val n = dstIds.length if (scores == null || scores.length < n) { @@ -539,6 +528,22 @@ private[ml] case class FeatureData(id: Int, features: Array[Float]) @Since("1.6.0") object ALSModel extends MLReadable[ALSModel] { + private def getPredictUDF(rank: Int) = { + udf { (featuresA: Seq[Float], featuresB: Seq[Float]) => + if (featuresA != null && featuresB != null) { + var dotProduct = 0.0f + var i = 0 + while (i < rank) { + dotProduct += featuresA(i) * featuresB(i) + i += 1 + } + dotProduct + } else { + Float.NaN + } + } + } + private[ml] def serializeData(data: FeatureData, dos: DataOutputStream): Unit = { import ReadWriteUtils._ dos.writeInt(data.id) @@ -888,7 +893,8 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { ata = new Array[Double](rank * rank) initialized = true } else { - require(this.rank == rank) + require(this.rank == rank, + s"NNLSSolver was initialized with rank ${this.rank} but got $rank.") } } @@ -965,8 +971,8 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { /** Adds an observation. */ def add(a: Array[Float], b: Double, c: Double = 1.0): NormalEquation = { - require(c >= 0.0) - require(a.length == k) + require(c >= 0.0, s"Observation weight must be non-negative but found $c.") + require(a.length == k, s"Observation length ${a.length} must equal rank $k.") copyToDouble(a) BLAS.nativeBLAS.dspr(upper, k, c, da, 1, ata) if (b != 0.0) { @@ -977,7 +983,7 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { /** Merges another normal equation object. */ def merge(other: NormalEquation): NormalEquation = { - require(other.k == k) + require(other.k == k, s"Cannot merge normal equations of rank ${other.k} into rank $k.") BLAS.nativeBLAS.daxpy(ata.length, 1.0, other.ata, 1, ata, 1) BLAS.nativeBLAS.daxpy(atb.length, 1.0, other.atb, 1, atb, 1) this @@ -1327,8 +1333,10 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { ratings: Array[Float]) { /** Size of the block. */ def size: Int = ratings.length - require(dstEncodedIndices.length == size) - require(dstPtrs.length == srcIds.length + 1) + require(dstEncodedIndices.length == size, + s"dstEncodedIndices must have $size entries but has ${dstEncodedIndices.length}.") + require(dstPtrs.length == srcIds.length + 1, + s"dstPtrs must have ${srcIds.length + 1} entries but has ${dstPtrs.length}.") } /** @@ -1370,8 +1378,10 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { ratings: Array[Float]) { /** Size of the block. */ def size: Int = srcIds.length - require(dstIds.length == srcIds.length) - require(ratings.length == srcIds.length) + require(dstIds.length == srcIds.length, + s"dstIds must have ${srcIds.length} entries but has ${dstIds.length}.") + require(ratings.length == srcIds.length, + s"ratings must have ${srcIds.length} entries but has ${ratings.length}.") } /** @@ -1494,8 +1504,9 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { dstLocalIndices: Array[Int], ratings: Array[Float]): this.type = { val sz = srcIds.length - require(dstLocalIndices.length == sz) - require(ratings.length == sz) + require(dstLocalIndices.length == sz, + s"dstLocalIndices must have $sz entries but has ${dstLocalIndices.length}.") + require(ratings.length == sz, s"ratings must have $sz entries but has ${ratings.length}.") this.srcIds ++= srcIds this.ratings ++= ratings var j = 0 @@ -1875,8 +1886,9 @@ object ALS extends DefaultParamsReadable[ALS] with Logging { /** Encodes a (blockId, localIndex) into a single integer. */ def encode(blockId: Int, localIndex: Int): Int = { - require(blockId < numBlocks) - require((localIndex & ~localIndexMask) == 0) + require(blockId < numBlocks, s"blockId $blockId must be less than numBlocks $numBlocks.") + require((localIndex & ~localIndexMask) == 0, + s"localIndex $localIndex must be in [0, $localIndexMask].") (blockId << numLocalIndexBits) | localIndex } diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala index d96500ea84ab9..0e6846e6085d7 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala @@ -111,14 +111,14 @@ private[regression] trait AFTSurvivalRegressionParams extends PredictorParams protected def validateAndTransformSchema( schema: StructType, fitting: Boolean): StructType = { - SchemaUtils.checkColumnType(schema, $(featuresCol), new VectorUDT) + SchemaUtils.checkColumnType(schema, $(featuresCol), SQLDataTypes.VectorType) if (fitting) { SchemaUtils.checkNumericType(schema, $(censorCol)) SchemaUtils.checkNumericType(schema, $(labelCol)) } val schemaWithQuantilesCol = if (hasQuantilesCol) { - SchemaUtils.appendColumn(schema, $(quantilesCol), new VectorUDT) + SchemaUtils.appendColumn(schema, $(quantilesCol), SQLDataTypes.VectorType) } else schema SchemaUtils.appendColumn(schemaWithQuantilesCol, $(predictionCol), DoubleType) @@ -408,23 +408,13 @@ class AFTSurvivalRegressionModel private[ml] ( } } - private def lambda2Quantiles(lambda: Double): Vector = { - val quantiles = _quantiles.copy - BLAS.scal(lambda, quantiles) - quantiles - } - @Since("2.0.0") - def predictQuantiles(features: Vector): Vector = { - // scale parameter for the Weibull distribution of lifetime - val lambda = predict(features) - lambda2Quantiles(lambda) - } + def predictQuantiles(features: Vector): Vector = + AFTSurvivalRegressionModel.predictQuantiles(features, coefficients, intercept, _quantiles) @Since("2.0.0") - def predict(features: Vector): Double = { - math.exp(BLAS.dot(coefficients, features) + intercept) - } + def predict(features: Vector): Double = + AFTSurvivalRegressionModel.predict(features, coefficients, intercept) @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { @@ -432,19 +422,29 @@ class AFTSurvivalRegressionModel private[ml] ( var predictionColNames = Seq.empty[String] var predictionColumns = Seq.empty[Column] + val localCoefficients = coefficients + val localIntercept = intercept if ($(predictionCol).nonEmpty) { - val predCol = udf(predict _).apply(col($(featuresCol))) + val predCol = udf((features: Vector) => + AFTSurvivalRegressionModel.predict(features, localCoefficients, localIntercept)) + .apply(col($(featuresCol))) predictionColNames :+= $(predictionCol) predictionColumns :+= predCol .as($(predictionCol), outputSchema($(predictionCol)).metadata) } if (hasQuantilesCol) { + val localQuantiles = _quantiles val quanCol = if ($(predictionCol).nonEmpty) { - udf(lambda2Quantiles _).apply(predictionColumns.head) + udf((lambda: Double) => + AFTSurvivalRegressionModel.lambda2Quantiles(lambda, localQuantiles)) + .apply(predictionColumns.head) } else { - udf(predictQuantiles _).apply(col($(featuresCol))) + udf((features: Vector) => + AFTSurvivalRegressionModel.predictQuantiles( + features, localCoefficients, localIntercept, localQuantiles) + ).apply(col($(featuresCol))) } predictionColNames :+= $(quantilesCol) predictionColumns :+= quanCol @@ -501,6 +501,23 @@ class AFTSurvivalRegressionModel private[ml] ( object AFTSurvivalRegressionModel extends MLReadable[AFTSurvivalRegressionModel] { private[ml] case class Data(coefficients: Vector, intercept: Double, scale: Double) + private def lambda2Quantiles(lambda: Double, quantiles: Vector): Vector = { + val scaledQuantiles = quantiles.copy + BLAS.scal(lambda, scaledQuantiles) + scaledQuantiles + } + + private def predict(features: Vector, coefficients: Vector, intercept: Double): Double = + math.exp(BLAS.dot(coefficients, features) + intercept) + + private def predictQuantiles( + features: Vector, + coefficients: Vector, + intercept: Double, + quantiles: Vector): Vector = { + lambda2Quantiles(predict(features, coefficients, intercept), quantiles) + } + private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = { import ReadWriteUtils._ serializeVector(data.coefficients, dos) diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala index 135224d9a8f9f..4a5c3b5790354 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala @@ -190,6 +190,8 @@ class DecisionTreeRegressionModel private[ml] ( // For ml connect only private[ml] def this() = this("", Node.dummyNode, -1) + override private[ml] val treeStats: NodeStats = rootNode.computeStats + private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize() override def predict(features: Vector): Double = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala index b533953b91dcc..2119555d408e4 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala @@ -280,7 +280,7 @@ class GBTRegressionModel private[ml]( val predictUDF = udf { features: Vector => bcastModel.value.predict(features) } predictionColNames :+= $(predictionCol) predictionColumns :+= predictUDF(col($(featuresCol))) - .as($(featuresCol), outputSchema($(featuresCol)).metadata) + .as($(predictionCol), outputSchema($(predictionCol)).metadata) } if ($(leafCol).nonEmpty) { diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala index f4d6e4e3562af..ac9363511320e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala @@ -1032,22 +1032,22 @@ class GeneralizedLinearRegressionModel private[ml] ( private lazy val familyAndLink = FamilyAndLink(this) override def predict(features: Vector): Double = { - predict(features, 0.0) + GeneralizedLinearRegressionModel.predict(features, 0.0, coefficients, intercept, familyAndLink) } /** * Calculates the predicted value when offset is set. */ private def predict(features: Vector, offset: Double): Double = { - val eta = predictLink(features, offset) - familyAndLink.fitted(eta) + GeneralizedLinearRegressionModel.predict( + features, offset, coefficients, intercept, familyAndLink) } /** * Calculates the link prediction (linear predictor) of the given instance. */ private def predictLink(features: Vector, offset: Double): Double = { - BLAS.dot(features, coefficients) + intercept + offset + GeneralizedLinearRegressionModel.predictLink(features, offset, coefficients, intercept) } override def transform(dataset: Dataset[_]): DataFrame = { @@ -1061,9 +1061,14 @@ class GeneralizedLinearRegressionModel private[ml] ( val offset = if (!hasOffsetCol) lit(0.0) else col($(offsetCol)).cast(DoubleType) var outputData = dataset var numColsOutput = 0 + val localCoefficients = coefficients + val localIntercept = intercept + val localFamilyAndLink = familyAndLink if (hasLinkPredictionCol) { - val predLinkUDF = udf((features: Vector, offset: Double) => predictLink(features, offset)) + val predLinkUDF = udf((features: Vector, offset: Double) => + GeneralizedLinearRegressionModel.predictLink( + features, offset, localCoefficients, localIntercept)) outputData = outputData .withColumn($(linkPredictionCol), predLinkUDF(col($(featuresCol)), offset), outputSchema($(linkPredictionCol)).metadata) @@ -1076,7 +1081,9 @@ class GeneralizedLinearRegressionModel private[ml] ( outputData = outputData.withColumn($(predictionCol), predUDF(col($(linkPredictionCol))), outputSchema($(predictionCol)).metadata) } else { - val predUDF = udf((features: Vector, offset: Double) => predict(features, offset)) + val predUDF = udf((features: Vector, offset: Double) => + GeneralizedLinearRegressionModel.predict( + features, offset, localCoefficients, localIntercept, localFamilyAndLink)) outputData = outputData.withColumn($(predictionCol), predUDF(col($(featuresCol)), offset), outputSchema($(predictionCol)).metadata) } @@ -1177,6 +1184,24 @@ class GeneralizedLinearRegressionModel private[ml] ( object GeneralizedLinearRegressionModel extends MLReadable[GeneralizedLinearRegressionModel] { private[ml] case class Data(intercept: Double, coefficients: Vector) + private def predict( + features: Vector, + offset: Double, + coefficients: Vector, + intercept: Double, + familyAndLink: GeneralizedLinearRegression.FamilyAndLink): Double = { + val eta = predictLink(features, offset, coefficients, intercept) + familyAndLink.fitted(eta) + } + + private def predictLink( + features: Vector, + offset: Double, + coefficients: Vector, + intercept: Double): Double = { + BLAS.dot(features, coefficients) + intercept + offset + } + private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = { import ReadWriteUtils._ dos.writeDouble(data.intercept) @@ -1257,7 +1282,7 @@ class GeneralizedLinearRegressionSummary private[regression] ( if (origModel.isDefined(origModel.predictionCol) && origModel.getPredictionCol.nonEmpty) { origModel.getPredictionCol } else { - "prediction_" + java.util.UUID.randomUUID.toString + Identifiable.randomUID("prediction") } } @@ -1431,7 +1456,7 @@ class GeneralizedLinearRegressionSummary private[regression] ( link.link(glrSummary.getDouble(4)) } else { // Create empty feature column and fit intercept only model using param setting from model - val featureNull = "feature_" + java.util.UUID.randomUUID.toString + val featureNull = Identifiable.randomUID("feature") val paramMap = model.extractParamMap() paramMap.put(model.featuresCol, featureNull) if (family.name != "tweedie") { diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala index 6eddcb416d8ae..72fe7706d610e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala @@ -18,6 +18,7 @@ package org.apache.spark.ml.regression import java.io.{DataInputStream, DataOutputStream} +import java.util.Arrays.binarySearch import org.apache.hadoop.fs.Path @@ -38,6 +39,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, Row} import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.types.{DoubleType, StructType} import org.apache.spark.storage.StorageLevel +import org.apache.spark.util.SizeEstimator /** * Params for isotonic regression. @@ -246,22 +248,41 @@ class IsotonicRegressionModel private[ml] ( copyValues(new IsotonicRegressionModel(uid, oldModel), extra).setParent(parent) } + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + if (oldModel != null) { + // boundaries: Array[Double] + size += SizeEstimator.estimate(oldModel.boundaries) + // predictions: Array[Double] + size += SizeEstimator.estimate(oldModel.predictions) + } + size + } + @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { val outputSchema = transformSchema(dataset.schema, logging = true) + val localBoundaries = oldModel.boundaries + val localPredictions = oldModel.predictions val predict = dataset.schema($(featuresCol)).dataType match { case DoubleType => - udf { feature: Double => oldModel.predict(feature) } + udf { feature: Double => + IsotonicRegressionModel.predict(localBoundaries, localPredictions, feature) + } case _: VectorUDT => val idx = $(featureIndex) - udf { features: Vector => oldModel.predict(features(idx)) } + udf { features: Vector => + IsotonicRegressionModel.predict(localBoundaries, localPredictions, features(idx)) + } } dataset.withColumn($(predictionCol), predict(col($(featuresCol))), outputSchema($(predictionCol)).metadata) } @Since("3.0.0") - def predict(value: Double): Double = oldModel.predict(value) + def predict(value: Double): Double = { + IsotonicRegressionModel.predict(oldModel.boundaries, oldModel.predictions, value) + } @Since("1.5.0") override def transformSchema(schema: StructType): StructType = { @@ -292,6 +313,28 @@ object IsotonicRegressionModel extends MLReadable[IsotonicRegressionModel] { predictions: Array[Double], isotonic: Boolean) + private[spark] def predict( + boundaries: Array[Double], + predictions: Array[Double], + testData: Double): Double = { + val foundIndex = binarySearch(boundaries, testData) + val insertIndex = -foundIndex - 1 + + if (insertIndex == 0) { + predictions.head + } else if (insertIndex == boundaries.length) { + predictions.last + } else if (foundIndex < 0) { + val x1 = boundaries(insertIndex - 1) + val y1 = predictions(insertIndex - 1) + val x2 = boundaries(insertIndex) + val y2 = predictions(insertIndex) + y1 + (y2 - y1) * (testData - x1) / (x2 - x1) + } else { + predictions(foundIndex) + } + } + private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = { import ReadWriteUtils._ serializeDoubleArray(data.boundaries, dos) diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala index 822df270c0bf7..89e29a7fb26d3 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/LinearRegression.scala @@ -742,7 +742,7 @@ class LinearRegressionModel private[ml] ( private[regression] def findSummaryModelAndPredictionCol(): (LinearRegressionModel, String) = { $(predictionCol) match { case "" => - val predictionColName = "prediction_" + java.util.UUID.randomUUID.toString + val predictionColName = Identifiable.randomUID("prediction") (copy(ParamMap.empty).setPredictionCol(predictionColName), predictionColName) case p => (this, p) } @@ -1160,4 +1160,3 @@ class LinearRegressionSummary private[regression] ( } } } - diff --git a/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala b/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala index 7b0de1176fa80..bc11fb9a64b19 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/source/libsvm/LibSVMRelation.scala @@ -27,7 +27,7 @@ import org.apache.spark.TaskContext import org.apache.spark.internal.Logging import org.apache.spark.ml.attribute.AttributeGroup import org.apache.spark.ml.feature.LabeledPoint -import org.apache.spark.ml.linalg.{Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vectors, VectorUDT} import org.apache.spark.mllib.util.MLUtils import org.apache.spark.sql.{Row, SparkSession} import org.apache.spark.sql.catalyst.InternalRow @@ -82,7 +82,7 @@ private[libsvm] case class LibSVMFileFormat() if ( dataSchema.size != 2 || !DataTypeUtils.sameType(dataSchema(0).dataType, DataTypes.DoubleType) || - !DataTypeUtils.sameType(dataSchema(1).dataType, new VectorUDT()) || + !DataTypeUtils.sameType(dataSchema(1).dataType, SQLDataTypes.VectorType) || !(forWriting || dataSchema(1).metadata.getLong(LibSVMOptions.NUM_FEATURES).toInt > 0) ) { throw new IOException(s"Illegal schema for libsvm data, schema=$dataSchema") diff --git a/mllib/src/main/scala/org/apache/spark/ml/stat/ANOVATest.scala b/mllib/src/main/scala/org/apache/spark/ml/stat/ANOVATest.scala index 2a3470e38f6ef..c1ac69b6fcaed 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/stat/ANOVATest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/stat/ANOVATest.scala @@ -97,7 +97,7 @@ private[ml] object ANOVATest { val spark = dataset.sparkSession import spark.implicits._ - SchemaUtils.checkColumnType(dataset.schema, featuresCol, new VectorUDT) + SchemaUtils.checkColumnType(dataset.schema, featuresCol, SQLDataTypes.VectorType) SchemaUtils.checkNumericType(dataset.schema, labelCol) val points = dataset.select(col(labelCol).cast("double"), col(featuresCol)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/stat/ChiSquareTest.scala b/mllib/src/main/scala/org/apache/spark/ml/stat/ChiSquareTest.scala index cdbfb6090acf5..14b9e6a2be334 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/stat/ChiSquareTest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/stat/ChiSquareTest.scala @@ -18,7 +18,7 @@ package org.apache.spark.ml.stat import org.apache.spark.annotation.Since -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.util.SchemaUtils import org.apache.spark.mllib.linalg.{Vectors => OldVectors} import org.apache.spark.mllib.stat.test.{ChiSqTest => OldChiSqTest} @@ -71,7 +71,7 @@ object ChiSquareTest { featuresCol: String, labelCol: String, flatten: Boolean): DataFrame = { - SchemaUtils.checkColumnType(dataset.schema, featuresCol, new VectorUDT) + SchemaUtils.checkColumnType(dataset.schema, featuresCol, SQLDataTypes.VectorType) SchemaUtils.checkNumericType(dataset.schema, labelCol) val spark = dataset.sparkSession diff --git a/mllib/src/main/scala/org/apache/spark/ml/stat/FValueTest.scala b/mllib/src/main/scala/org/apache/spark/ml/stat/FValueTest.scala index 56b7c058a5379..16f25621654b1 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/stat/FValueTest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/stat/FValueTest.scala @@ -20,7 +20,7 @@ package org.apache.spark.ml.stat import org.apache.commons.math3.distribution.FDistribution import org.apache.spark.annotation.Since -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors} import org.apache.spark.ml.util.SchemaUtils import org.apache.spark.rdd.RDD import org.apache.spark.sql.{DataFrame, Dataset, Row} @@ -99,7 +99,7 @@ private[ml] object FValueTest { dataset: Dataset[_], featuresCol: String, labelCol: String): RDD[(Int, Double, Long, Double)] = { - SchemaUtils.checkColumnType(dataset.schema, featuresCol, new VectorUDT) + SchemaUtils.checkColumnType(dataset.schema, featuresCol, SQLDataTypes.VectorType) SchemaUtils.checkNumericType(dataset.schema, labelCol) val spark = dataset.sparkSession diff --git a/mllib/src/main/scala/org/apache/spark/ml/stat/Summarizer.scala b/mllib/src/main/scala/org/apache/spark/ml/stat/Summarizer.scala index 5a136b7578f36..19d20cafa5a60 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/stat/Summarizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/stat/Summarizer.scala @@ -22,7 +22,7 @@ import java.io._ import org.apache.spark.annotation.Since import org.apache.spark.internal.Logging import org.apache.spark.ml.feature.Instance -import org.apache.spark.ml.linalg.{Vector, Vectors, VectorUDT} +import org.apache.spark.ml.linalg.{SQLDataTypes, Vector, Vectors, VectorUDT} import org.apache.spark.rdd.RDD import org.apache.spark.sql.Column import org.apache.spark.sql.catalyst.InternalRow @@ -302,7 +302,7 @@ private[spark] object SummaryBuilderImpl extends Logging { } } - private val vectorUDT = new VectorUDT + private val vectorUDT = SQLDataTypes.VectorType.asInstanceOf[VectorUDT] /** * All the metrics that can be currently computed by Spark for vectors. @@ -648,7 +648,7 @@ private[spark] class SummarizerBuffer( * Sample mean of each dimension. */ def mean: Vector = { - require(requestedMetrics.contains(Mean)) + require(requestedMetrics.contains(Mean), "mean was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") val realMean = Array.ofDim[Double](n) @@ -664,7 +664,7 @@ private[spark] class SummarizerBuffer( * Sum of each dimension. */ def sum: Vector = { - require(requestedMetrics.contains(Sum)) + require(requestedMetrics.contains(Sum), "sum was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") val realSum = Array.ofDim[Double](n) @@ -680,7 +680,7 @@ private[spark] class SummarizerBuffer( * Unbiased estimate of sample variance of each dimension. */ def variance: Vector = { - require(requestedMetrics.contains(Variance)) + require(requestedMetrics.contains(Variance), "variance was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") val realVariance = computeVariance @@ -691,7 +691,7 @@ private[spark] class SummarizerBuffer( * Unbiased estimate of standard deviation of each dimension. */ def std: Vector = { - require(requestedMetrics.contains(Std)) + require(requestedMetrics.contains(Std), "std was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") val realVariance = computeVariance @@ -732,7 +732,7 @@ private[spark] class SummarizerBuffer( * */ def numNonzeros: Vector = { - require(requestedMetrics.contains(NumNonZeros)) + require(requestedMetrics.contains(NumNonZeros), "numNonZeros was not a requested metric.") require(totalCnt > 0, s"Nothing has been added to this summarizer.") Vectors.dense(nnz.map(_.toDouble)) @@ -742,7 +742,7 @@ private[spark] class SummarizerBuffer( * Maximum value of each dimension. */ def max: Vector = { - require(requestedMetrics.contains(Max)) + require(requestedMetrics.contains(Max), "max was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") var i = 0 @@ -757,7 +757,7 @@ private[spark] class SummarizerBuffer( * Minimum value of each dimension. */ def min: Vector = { - require(requestedMetrics.contains(Min)) + require(requestedMetrics.contains(Min), "min was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") var i = 0 @@ -772,7 +772,7 @@ private[spark] class SummarizerBuffer( * L2 (Euclidean) norm of each dimension. */ def normL2: Vector = { - require(requestedMetrics.contains(NormL2)) + require(requestedMetrics.contains(NormL2), "normL2 was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") val realMagnitude = Array.ofDim[Double](n) @@ -790,7 +790,7 @@ private[spark] class SummarizerBuffer( * L1 norm of each dimension. */ def normL1: Vector = { - require(requestedMetrics.contains(NormL1)) + require(requestedMetrics.contains(NormL1), "normL1 was not a requested metric.") require(totalWeightSum > 0, s"Nothing has been added to this summarizer.") Vectors.dense(currL1) diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala index b68a5e079dc44..0a7e37aa0c9ad 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala @@ -78,10 +78,35 @@ sealed abstract class Node extends Serializable { */ private[ml] def maxSplitFeatureIndex(): Int - /** Returns a deep copy of the subtree rooted at this node. */ - private[tree] def deepCopy(): Node + private[ml] def computeStats: NodeStats = { + var numDescendants = 0 + var subtreeDepth = 0 + var numLeaves = 0 + + def visit(node: Node, depth: Int): Unit = { + if (depth > 0) { + numDescendants += 1 + } + subtreeDepth = math.max(subtreeDepth, depth) + node match { + case _: LeafNode => + numLeaves += 1 + case internal: InternalNode => + visit(internal.leftChild, depth + 1) + visit(internal.rightChild, depth + 1) + } + } + + visit(this, 0) + NodeStats(subtreeDepth, numDescendants, numLeaves) + } } +private[ml] case class NodeStats( + subtreeDepth: Int, + numDescendants: Int, + numLeaves: Int) + private[ml] object Node { /** @@ -146,10 +171,6 @@ class LeafNode private[ml] ( } override private[ml] def maxSplitFeatureIndex(): Int = -1 - - override private[tree] def deepCopy(): Node = { - new LeafNode(prediction, impurity, impurityStats) - } } /** @@ -238,11 +259,6 @@ class InternalNode private[ml] ( math.max(split.featureIndex, math.max(leftChild.maxSplitFeatureIndex(), rightChild.maxSplitFeatureIndex())) } - - override private[tree] def deepCopy(): Node = { - new InternalNode(prediction, impurity, gain, leftChild.deepCopy(), rightChild.deepCopy(), - split, impurityStats) - } } private object InternalNode { @@ -393,11 +409,6 @@ private[tree] object LearningNode { */ def rightChildIndex(nodeIndex: Int): Int = (nodeIndex << 1) + 1 - /** - * Get the parent index of the given node, or 0 if it is the root. - */ - def parentIndex(nodeIndex: Int): Int = nodeIndex >> 1 - /** * Return the level of a tree which the given node is in. */ @@ -407,40 +418,4 @@ private[tree] object LearningNode { java.lang.Integer.numberOfTrailingZeros(java.lang.Integer.highestOneBit(nodeIndex)) } - /** - * Returns true if this is a left child. - * Note: Returns false for the root. - */ - def isLeftChild(nodeIndex: Int): Boolean = nodeIndex > 1 && nodeIndex % 2 == 0 - - /** - * Return the maximum number of nodes which can be in the given level of the tree. - * @param level Level of tree (0 = root). - */ - def maxNodesInLevel(level: Int): Int = 1 << level - - /** - * Return the index of the first node in the given level. - * @param level Level of tree (0 = root). - */ - def startIndexInLevel(level: Int): Int = 1 << level - - /** - * Traces down from a root node to get the node with the given node index. - * This assumes the node exists. - */ - def getNode(nodeIndex: Int, rootNode: LearningNode): LearningNode = { - var tmpNode: LearningNode = rootNode - var levelsToGo = indexToLevel(nodeIndex) - while (levelsToGo > 0) { - if ((nodeIndex & (1 << levelsToGo - 1)) == 0) { - tmpNode = tmpNode.leftChild.get - } else { - tmpNode = tmpNode.rightChild.get - } - levelsToGo -= 1 - } - tmpNode - } - } diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala index 4e9fa89cbde90..1cc64e58d9807 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala @@ -49,17 +49,13 @@ private[spark] trait DecisionTreeModel { def rootNode: Node /** Number of nodes in tree, including leaf nodes. */ - def numNodes: Int = { - 1 + rootNode.numDescendants - } + def numNodes: Int = treeStats.numDescendants + 1 /** * Depth of the tree. * E.g.: Depth 0 means 1 leaf node. Depth 1 means 1 internal node and 2 leaf nodes. */ - lazy val depth: Int = { - rootNode.subtreeDepth - } + def depth: Int = treeStats.subtreeDepth /** Summary of the model */ override def toString: String = { @@ -95,13 +91,11 @@ private[spark] trait DecisionTreeModel { } } - private[ml] lazy val numLeave: Int = - leafIterator(rootNode).size + private[ml] def treeStats: NodeStats - private[ml] lazy val leafAttr = { - NominalAttribute.defaultAttr - .withNumValues(numLeave) - } + private[ml] def numLeaves: Int = treeStats.numLeaves + + private[ml] def leafAttr = NominalAttribute.defaultAttr.withNumValues(numLeaves) private[ml] def getLeafField(leafCol: String) = { leafAttr.withName(leafCol).toStructField() diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala index 2244d49b2a35f..08dd0039a9913 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala @@ -24,7 +24,7 @@ import scala.util.Try import org.apache.spark.annotation.Since import org.apache.spark.ml.PredictorParams import org.apache.spark.ml.classification.ProbabilisticClassifierParams -import org.apache.spark.ml.linalg.VectorUDT +import org.apache.spark.ml.linalg.SQLDataTypes import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared._ import org.apache.spark.ml.util.SchemaUtils @@ -420,7 +420,7 @@ private[ml] trait TreeEnsembleClassifierParams featuresDataType: DataType): StructType = { var outputSchema = super.validateAndTransformSchema(schema, fitting, featuresDataType) if ($(leafCol).nonEmpty) { - outputSchema = SchemaUtils.appendColumn(outputSchema, $(leafCol), new VectorUDT) + outputSchema = SchemaUtils.appendColumn(outputSchema, $(leafCol), SQLDataTypes.VectorType) } outputSchema } @@ -438,7 +438,7 @@ private[ml] trait TreeEnsembleRegressorParams featuresDataType: DataType): StructType = { var outputSchema = super.validateAndTransformSchema(schema, fitting, featuresDataType) if ($(leafCol).nonEmpty) { - outputSchema = SchemaUtils.appendColumn(outputSchema, $(leafCol), new VectorUDT) + outputSchema = SchemaUtils.appendColumn(outputSchema, $(leafCol), SQLDataTypes.VectorType) } outputSchema } diff --git a/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala b/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala index 119cc47a56f4f..70bb2347cfc7e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala @@ -38,8 +38,8 @@ import org.apache.spark.ml.util.Instrumentation.instrumented import org.apache.spark.mllib.util.MLUtils import org.apache.spark.sql.{DataFrame, Dataset} import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.util.{SizeEstimator, ThreadUtils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.ThreadUtils /** * Params for [[CrossValidator]] and [[CrossValidatorModel]]. @@ -328,6 +328,31 @@ class CrossValidatorModel private[ml] ( @Since("2.3.0") def hasSubModels: Boolean = _subModels.isDefined + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize(excluded = Seq( + // estimator: Param[Estimator[_]] + estimator, + // estimatorParamMaps: Param[Array[ParamMap]] + estimatorParamMaps, + // evaluator: Param[Evaluator] + evaluator)) + // bestModel: Model[_] + size += bestModel.estimatedSize + // avgMetrics: Array[Double] + size += SizeEstimator.estimate(avgMetrics) + // _subModels: Option[Array[Array[Model[_]]]] + _subModels.foreach { models => + models.foreach { modelArray => + modelArray.foreach { model => + if (model != null) { + size += model.estimatedSize + } + } + } + } + size + } + @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema, logging = true) diff --git a/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala b/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala index 6ee64ef99a668..5b0b77abb2dd0 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala @@ -38,8 +38,8 @@ import org.apache.spark.ml.util._ import org.apache.spark.ml.util.Instrumentation.instrumented import org.apache.spark.sql.{DataFrame, Dataset} import org.apache.spark.sql.types.StructType +import org.apache.spark.util.{SizeEstimator, ThreadUtils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.ThreadUtils /** * Params for [[TrainValidationSplit]] and [[TrainValidationSplitModel]]. @@ -293,6 +293,29 @@ class TrainValidationSplitModel private[ml] ( @Since("2.3.0") def hasSubModels: Boolean = _subModels.isDefined + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize(excluded = Seq( + // estimator: Param[Estimator[_]] + estimator, + // estimatorParamMaps: Param[Array[ParamMap]] + estimatorParamMaps, + // evaluator: Param[Evaluator] + evaluator)) + // bestModel: Model[_] + size += bestModel.estimatedSize + // validationMetrics: Array[Double] + size += SizeEstimator.estimate(validationMetrics) + // _subModels: Option[Array[Model[_]]] + _subModels.foreach { modelArray => + modelArray.foreach { model => + if (model != null) { + size += model.estimatedSize + } + } + } + size + } + @Since("2.0.0") override def transform(dataset: Dataset[_]): DataFrame = { transformSchema(dataset.schema, logging = true) diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/MetadataUtils.scala b/mllib/src/main/scala/org/apache/spark/ml/util/MetadataUtils.scala index 631261af249f2..2e0dbe75dee7a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/util/MetadataUtils.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/util/MetadataUtils.scala @@ -20,7 +20,7 @@ package org.apache.spark.ml.util import scala.collection.immutable.HashMap import org.apache.spark.ml.attribute._ -import org.apache.spark.ml.linalg.VectorUDT +import org.apache.spark.ml.linalg.{SQLDataTypes, VectorUDT} import org.apache.spark.sql.types.StructField @@ -46,7 +46,7 @@ private[spark] object MetadataUtils { * Returns None if the number of features is not specified. */ def getNumFeatures(vectorSchema: StructField): Option[Int] = { - if (vectorSchema.dataType == new VectorUDT) { + if (vectorSchema.dataType == SQLDataTypes.VectorType) { val group = AttributeGroup.fromStructField(vectorSchema) val size = group.size if (size >= 0) { diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala b/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala index 5386641838726..3bb249f201b15 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala @@ -19,7 +19,7 @@ package org.apache.spark.ml.util import org.apache.spark.SparkIllegalArgumentException import org.apache.spark.ml.attribute._ -import org.apache.spark.ml.linalg.VectorUDT +import org.apache.spark.ml.linalg.SQLDataTypes import org.apache.spark.sql.catalyst.util.{AttributeNameParser, QuotingUtils} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -201,7 +201,8 @@ private[spark] object SchemaUtils { * @param colName column name */ def validateVectorCompatibleColumn(schema: StructType, colName: String): Unit = { - val typeCandidates = List( new VectorUDT, + val typeCandidates = List( + SQLDataTypes.VectorType, new ArrayType(DoubleType, false), new ArrayType(FloatType, false)) checkColumnTypes(schema, colName, typeCandidates) diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala index 083c3e3e77a9b..113445269ce76 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/BisectingKMeansModel.scala @@ -41,7 +41,7 @@ import org.apache.spark.util.ArrayImplicits._ */ @Since("1.6.0") class BisectingKMeansModel private[clustering] ( - private[clustering] val root: ClusteringTreeNode, + private[spark] val root: ClusteringTreeNode, @Since("2.4.0") val distanceMeasure: String, @Since("3.0.0") val trainingCost: Double ) extends Serializable with Saveable with Logging { diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/DistanceMeasure.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/DistanceMeasure.scala index 5b0fb5ef18c8b..f6bec405554cd 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/DistanceMeasure.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/DistanceMeasure.scala @@ -258,8 +258,8 @@ object DistanceMeasure { private[spark] def decodeFromString(distanceMeasure: String): DistanceMeasure = distanceMeasure match { - case EUCLIDEAN => new EuclideanDistanceMeasure - case COSINE => new CosineDistanceMeasure + case EUCLIDEAN => EuclideanDistanceMeasure + case COSINE => CosineDistanceMeasure case _ => throw new IllegalArgumentException(s"distanceMeasure must be one of: " + s"$EUCLIDEAN, $COSINE. $distanceMeasure provided.") } @@ -278,7 +278,7 @@ object DistanceMeasure { k.toLong * k * numFeatures < 1000000 } -private[spark] class EuclideanDistanceMeasure extends DistanceMeasure { +private[spark] object EuclideanDistanceMeasure extends DistanceMeasure { /** * Statistics used in triangle inequality to obtain useful bounds to find closest centers. @@ -400,10 +400,7 @@ private[spark] class EuclideanDistanceMeasure extends DistanceMeasure { centroid: VectorWithNorm): Double = { EuclideanDistanceMeasure.fastSquaredDistance(point, centroid) } -} - -private[spark] object EuclideanDistanceMeasure { /** * @return the squared Euclidean distance between two vectors computed by * [[org.apache.spark.mllib.util.MLUtils#fastSquaredDistance]]. @@ -415,7 +412,7 @@ private[spark] object EuclideanDistanceMeasure { } } -private[spark] class CosineDistanceMeasure extends DistanceMeasure { +private[spark] object CosineDistanceMeasure extends DistanceMeasure { /** * Statistics used in triangle inequality to obtain useful bounds to find closest centers. diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala index 562e5b3995cfb..9103d83db41dc 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala @@ -201,10 +201,22 @@ class GaussianMixture private ( val compute = sc.broadcast(ExpectationSum.add(weights, gaussians)_) // aggregate the cluster contribution for all sample points + // Avoid allocating and serializing a large zero value for empty partitions. val sums = breezeData.treeAggregate[ExpectationSum]( - zeroValue = ExpectationSum.zero(k, d), - seqOp = (agg: ExpectationSum, v: BV[Double]) => compute.value(agg, v), - combOp = (agg1: ExpectationSum, agg2: ExpectationSum) => agg1 += agg2, + zeroValue = null.asInstanceOf[ExpectationSum], + seqOp = (maybeAgg: ExpectationSum, v: BV[Double]) => { + val agg = if (maybeAgg == null) ExpectationSum.zero(k, d) else maybeAgg + compute.value(agg, v) + }, + combOp = (agg1: ExpectationSum, agg2: ExpectationSum) => { + if (agg1 == null) { + agg2 + } else if (agg2 == null) { + agg1 + } else { + agg1 += agg2 + } + }, depth = 2, finalAggregateOnExecutor = true) diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeansModel.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeansModel.scala index e5c0b27072d02..ad2b812684a0e 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeansModel.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeansModel.scala @@ -43,7 +43,7 @@ class KMeansModel (@Since("1.0.0") val clusterCenters: Array[Vector], private[spark] val numIter: Int) extends Saveable with Serializable with PMMLExportable { - @transient private lazy val distanceMeasureInstance: DistanceMeasure = + private val distanceMeasureInstance: DistanceMeasure = DistanceMeasure.decodeFromString(distanceMeasure) @transient private lazy val clusterCentersWithNorm = diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala index ac31c0d3be479..2fe15a59f4c7c 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/LocalKMeans.scala @@ -75,7 +75,7 @@ private[mllib] object LocalKMeans extends Logging { } - val distanceMeasureInstance = new EuclideanDistanceMeasure + val distanceMeasureInstance = EuclideanDistanceMeasure // Run up to maxIterations iterations of Lloyd's algorithm val oldClosest = Array.fill(points.length)(-1) diff --git a/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala b/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala index 5634f0e1d9b79..c7af3ca4a8b63 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/feature/IDF.scala @@ -21,6 +21,7 @@ import breeze.linalg.{DenseVector => BDV} import org.apache.spark.annotation.Since import org.apache.spark.api.java.JavaRDD +import org.apache.spark.ml.feature.{IDFModel => NewIDFModel} import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors} import org.apache.spark.rdd.RDD @@ -184,7 +185,10 @@ class IDFModel private[spark](@Since("1.1.0") val idf: Vector, @Since("1.1.0") def transform(dataset: RDD[Vector]): RDD[Vector] = { val bcIdf = dataset.context.broadcast(idf) - dataset.mapPartitions(iter => iter.map(v => IDFModel.transform(bcIdf.value, v))) + dataset.mapPartitions { iter => + val localIdf = bcIdf.value.toArray + iter.map(v => IDFModel.transform(localIdf, v)) + } } /** @@ -194,7 +198,7 @@ class IDFModel private[spark](@Since("1.1.0") val idf: Vector, * @return a TF-IDF vector */ @Since("1.3.0") - def transform(v: Vector): Vector = IDFModel.transform(idf, v) + def transform(v: Vector): Vector = IDFModel.transform(idf.toArray, v) /** * Transforms term frequency (TF) vectors to TF-IDF vectors (Java version). @@ -210,50 +214,23 @@ class IDFModel private[spark](@Since("1.1.0") val idf: Vector, private[spark] object IDFModel { /** - * Transforms a term frequency (TF) vector to a TF-IDF vector with a IDF vector + * Transforms a term frequency (TF) vector to a TF-IDF vector with IDF values * - * @param idf an IDF vector + * @param idf IDF values * @param v a term frequency vector * @return a TF-IDF vector */ - def transform(idf: Vector, v: Vector): Vector = { + private def transform(idf: Array[Double], v: Vector): Vector = { v match { case SparseVector(size, indices, values) => - val (newIndices, newValues) = transformSparse(idf, indices, values) + val (newIndices, newValues) = NewIDFModel.predictSparse(idf, indices, values) Vectors.sparse(size, newIndices, newValues) case DenseVector(values) => - val newValues = transformDense(idf, values) + val newValues = NewIDFModel.predictDense(idf, values) Vectors.dense(newValues) case other => throw new UnsupportedOperationException( s"Only sparse and dense vectors are supported but got ${other.getClass}.") } } - - private[spark] def transformDense( - idf: Vector, - values: Array[Double]): Array[Double] = { - val n = values.length - val newValues = new Array[Double](n) - var j = 0 - while (j < n) { - newValues(j) = values(j) * idf(j) - j += 1 - } - newValues - } - - private[spark] def transformSparse( - idf: Vector, - indices: Array[Int], - values: Array[Double]): (Array[Int], Array[Double]) = { - val nnz = indices.length - val newValues = new Array[Double](nnz) - var k = 0 - while (k < nnz) { - newValues(k) = values(k) * idf(indices(k)) - k += 1 - } - (indices, newValues) - } } diff --git a/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala b/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala index 456580ffa5315..e4aec6d577106 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala @@ -18,7 +18,6 @@ package org.apache.spark.mllib.regression import java.io.Serializable import java.lang.{Double => JDouble} -import java.util.Arrays.binarySearch import scala.collection.mutable.ArrayBuffer import scala.jdk.CollectionConverters._ @@ -30,6 +29,7 @@ import org.json4s.jackson.JsonMethods._ import org.apache.spark.{RangePartitioner, SparkContext} import org.apache.spark.annotation.Since import org.apache.spark.api.java.{JavaDoubleRDD, JavaRDD} +import org.apache.spark.ml.regression.{IsotonicRegressionModel => NewIsotonicRegressionModel} import org.apache.spark.mllib.linalg.{Vector, Vectors} import org.apache.spark.mllib.util.{Loader, Saveable} import org.apache.spark.rdd.RDD @@ -89,7 +89,11 @@ class IsotonicRegressionModel @Since("1.3.0") ( */ @Since("1.3.0") def predict(testData: RDD[Double]): RDD[Double] = { - testData.map(predict) + val localBoundaries = boundaries + val localPredictions = predictions + testData.map { value => + NewIsotonicRegressionModel.predict(localBoundaries, localPredictions, value) + } } /** @@ -124,30 +128,7 @@ class IsotonicRegressionModel @Since("1.3.0") ( */ @Since("1.3.0") def predict(testData: Double): Double = { - - def linearInterpolation(x1: Double, y1: Double, x2: Double, y2: Double, x: Double): Double = { - y1 + (y2 - y1) * (x - x1) / (x2 - x1) - } - - val foundIndex = binarySearch(boundaries, testData) - val insertIndex = -foundIndex - 1 - - // Find if the index was lower than all values, - // higher than all values, in between two values or exact match. - if (insertIndex == 0) { - predictions.head - } else if (insertIndex == boundaries.length) { - predictions.last - } else if (foundIndex < 0) { - linearInterpolation( - boundaries(insertIndex - 1), - predictions(insertIndex - 1), - boundaries(insertIndex), - predictions(insertIndex), - testData) - } else { - predictions(foundIndex) - } + NewIsotonicRegressionModel.predict(boundaries, predictions, testData) } /** A convenient method for boundaries called by the Python API. */ @@ -477,7 +458,8 @@ class IsotonicRegression private (private var isotonic: Boolean) extends Seriali private def parallelPoolAdjacentViolators( input: RDD[(Double, Double, Double)]): Array[(Double, Double, Double)] = { val keyedInput = input.keyBy(_._2) - val parallelStepResult = keyedInput + + keyedInput // Points with same or adjacent features must collocate within the same partition. .partitionBy(new RangePartitioner(keyedInput.getNumPartitions, keyedInput)) .values @@ -486,10 +468,10 @@ class IsotonicRegression private (private var isotonic: Boolean) extends Seriali // Aggregate points with equal features into a single point. .map(makeUnique) .flatMap(poolAdjacentViolators) + // Sort partial results with a spill-capable shuffle before the final PAV pass. + .sortBy(_._2, ascending = true, numPartitions = 1) + .mapPartitions(p => poolAdjacentViolators(p.toArray).iterator) .collect() - // Sort again because collect() doesn't promise ordering. - .sortBy(_._2) - poolAdjacentViolators(parallelStepResult) } /** diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Node.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Node.scala index a0eec8b2afb9c..6ac0bfdc7b634 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Node.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Node.scala @@ -80,23 +80,6 @@ class Node @Since("1.2.0") ( } } - /** - * Returns a deep copy of the subtree rooted at this node. - */ - private[tree] def deepCopy(): Node = { - val leftNodeCopy = if (leftNode.isEmpty) { - None - } else { - Some(leftNode.get.deepCopy()) - } - val rightNodeCopy = if (rightNode.isEmpty) { - None - } else { - Some(rightNode.get.deepCopy()) - } - new Node(id, predict, impurity, isLeaf, split, leftNodeCopy, rightNodeCopy, stats) - } - /** * Get the number of nodes in tree below this node, including leaf nodes. * E.g., if this is a leaf, returns 0. If both children are leaves, returns 2. @@ -157,12 +140,6 @@ class Node @Since("1.2.0") ( private[spark] object Node { - /** - * Return a node with the given node id (but nothing else set). - */ - def emptyNode(nodeIndex: Int): Node = new Node(nodeIndex, new Predict(Double.MinValue), -1.0, - false, None, None, None, None) - /** * Construct a node with nodeIndex, predict, impurity and isLeaf parameters. * This is used in `DecisionTree.findBestSplits` to construct child nodes @@ -192,54 +169,4 @@ private[spark] object Node { */ def rightChildIndex(nodeIndex: Int): Int = (nodeIndex << 1) + 1 - /** - * Get the parent index of the given node, or 0 if it is the root. - */ - def parentIndex(nodeIndex: Int): Int = nodeIndex >> 1 - - /** - * Return the level of a tree which the given node is in. - */ - def indexToLevel(nodeIndex: Int): Int = if (nodeIndex == 0) { - throw new IllegalArgumentException(s"0 is not a valid node index.") - } else { - java.lang.Integer.numberOfTrailingZeros(java.lang.Integer.highestOneBit(nodeIndex)) - } - - /** - * Returns true if this is a left child. - * Note: Returns false for the root. - */ - def isLeftChild(nodeIndex: Int): Boolean = nodeIndex > 1 && nodeIndex % 2 == 0 - - /** - * Return the maximum number of nodes which can be in the given level of the tree. - * @param level Level of tree (0 = root). - */ - def maxNodesInLevel(level: Int): Int = 1 << level - - /** - * Return the index of the first node in the given level. - * @param level Level of tree (0 = root). - */ - def startIndexInLevel(level: Int): Int = 1 << level - - /** - * Traces down from a root node to get the node with the given node index. - * This assumes the node exists. - */ - def getNode(nodeIndex: Int, rootNode: Node): Node = { - var tmpNode: Node = rootNode - var levelsToGo = indexToLevel(nodeIndex) - while (levelsToGo > 0) { - if ((nodeIndex & (1 << levelsToGo - 1)) == 0) { - tmpNode = tmpNode.leftNode.get - } else { - tmpNode = tmpNode.rightNode.get - } - levelsToGo -= 1 - } - tmpNode - } - } diff --git a/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala b/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala index 5a213e6803486..31ddf5b5a36b9 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/util/MLUtils.scala @@ -33,8 +33,10 @@ import org.apache.spark.rdd.{PartitionwiseSampledRDD, RDD} import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession} import org.apache.spark.sql.execution.datasources.DataSource import org.apache.spark.sql.execution.datasources.text.TextFileFormat -import org.apache.spark.sql.functions._ +import org.apache.spark.sql.functions.{col, length, lit, not, printf, raise_error, trim, + unwrap_udt, when, wrap_udt} import org.apache.spark.storage.StorageLevel +import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.random.BernoulliCellSampler /** @@ -319,15 +321,15 @@ object MLUtils extends Logging { * Converts vector columns in an input Dataset from the [[org.apache.spark.mllib.linalg.Vector]] * type to the new [[org.apache.spark.ml.linalg.Vector]] type under the `spark.ml` package. * @param dataset input dataset - * @param cols a list of vector columns to be converted. New vector columns will be ignored. If - * unspecified, all old vector columns will be converted except nested ones. + * @param cols a list of vector columns to be converted. If unspecified, all old vector columns + * will be converted except nested ones. * @return the input `DataFrame` with old vector columns converted to the new vector type */ @Since("2.0.0") @varargs def convertVectorColumnsToML(dataset: Dataset[_], cols: String*): DataFrame = { val schema = dataset.schema - val colSet = if (cols.nonEmpty) { + val colNames = if (cols.nonEmpty) { cols.flatMap { c => val dataType = schema(c).dataType if (dataType.getClass == classOf[VectorUDT]) { @@ -338,49 +340,38 @@ object MLUtils extends Logging { s"Column $c must be old Vector type to be converted to new type but got $dataType.") None } - }.toSet + } } else { schema.fields .filter(_.dataType.getClass == classOf[VectorUDT]) .map(_.name) - .toSet + .toImmutableArraySeq } - if (colSet.isEmpty) { + if (colNames.isEmpty) { return dataset.toDF() } - logWarning("Vector column conversion has serialization overhead. " + - "Please migrate your datasets and workflows to use the spark.ml package.") - - // TODO: This implementation has performance issues due to unnecessary serialization. - // TODO: It is better (but trickier) if we can cast the old vector type to new type directly. - val convertToML = udf { v: Vector => v.asML } - val exprs = schema.fields.map { field => - val c = field.name - if (colSet.contains(c)) { - convertToML(col(c)).as(c, field.metadata) - } else { - col(c) - } - } - import org.apache.spark.util.ArrayImplicits._ - dataset.select(exprs.toImmutableArraySeq: _*) + val fields = colNames.map(schema(_)) + dataset.withColumns( + fields.map(_.name), + fields.map(field => wrap_udt(unwrap_udt(col(field.name)), new MLVectorUDT)), + fields.map(_.metadata)) } /** * Converts vector columns in an input Dataset to the [[org.apache.spark.mllib.linalg.Vector]] * type from the new [[org.apache.spark.ml.linalg.Vector]] type under the `spark.ml` package. * @param dataset input dataset - * @param cols a list of vector columns to be converted. Old vector columns will be ignored. If - * unspecified, all new vector columns will be converted except nested ones. + * @param cols a list of vector columns to be converted. If unspecified, all new vector columns + * will be converted except nested ones. * @return the input `DataFrame` with new vector columns converted to the old vector type */ @Since("2.0.0") @varargs def convertVectorColumnsFromML(dataset: Dataset[_], cols: String*): DataFrame = { val schema = dataset.schema - val colSet = if (cols.nonEmpty) { + val colNames = if (cols.nonEmpty) { cols.flatMap { c => val dataType = schema(c).dataType if (dataType.getClass == classOf[MLVectorUDT]) { @@ -391,49 +382,38 @@ object MLUtils extends Logging { s"Column $c must be new Vector type to be converted to old type but got $dataType.") None } - }.toSet + } } else { schema.fields .filter(_.dataType.getClass == classOf[MLVectorUDT]) .map(_.name) - .toSet + .toImmutableArraySeq } - if (colSet.isEmpty) { + if (colNames.isEmpty) { return dataset.toDF() } - logWarning("Vector column conversion has serialization overhead. " + - "Please migrate your datasets and workflows to use the spark.ml package.") - - // TODO: This implementation has performance issues due to unnecessary serialization. - // TODO: It is better (but trickier) if we can cast the new vector type to old type directly. - val convertFromML = udf { Vectors.fromML _ } - val exprs = schema.fields.map { field => - val c = field.name - if (colSet.contains(c)) { - convertFromML(col(c)).as(c, field.metadata) - } else { - col(c) - } - } - import org.apache.spark.util.ArrayImplicits._ - dataset.select(exprs.toImmutableArraySeq: _*) + val fields = colNames.map(schema(_)) + dataset.withColumns( + fields.map(_.name), + fields.map(field => wrap_udt(unwrap_udt(col(field.name)), new VectorUDT)), + fields.map(_.metadata)) } /** * Converts Matrix columns in an input Dataset from the [[org.apache.spark.mllib.linalg.Matrix]] * type to the new [[org.apache.spark.ml.linalg.Matrix]] type under the `spark.ml` package. * @param dataset input dataset - * @param cols a list of matrix columns to be converted. New matrix columns will be ignored. If - * unspecified, all old matrix columns will be converted except nested ones. + * @param cols a list of matrix columns to be converted. If unspecified, all old matrix columns + * will be converted except nested ones. * @return the input `DataFrame` with old matrix columns converted to the new matrix type */ @Since("2.0.0") @varargs def convertMatrixColumnsToML(dataset: Dataset[_], cols: String*): DataFrame = { val schema = dataset.schema - val colSet = if (cols.nonEmpty) { + val colNames = if (cols.nonEmpty) { cols.flatMap { c => val dataType = schema(c).dataType if (dataType.getClass == classOf[MatrixUDT]) { @@ -444,47 +424,38 @@ object MLUtils extends Logging { s"Column $c must be old Matrix type to be converted to new type but got $dataType.") None } - }.toSet + } } else { schema.fields .filter(_.dataType.getClass == classOf[MatrixUDT]) .map(_.name) - .toSet + .toImmutableArraySeq } - if (colSet.isEmpty) { + if (colNames.isEmpty) { return dataset.toDF() } - logWarning("Matrix column conversion has serialization overhead. " + - "Please migrate your datasets and workflows to use the spark.ml package.") - - val convertToML = udf { v: Matrix => v.asML } - val exprs = schema.fields.map { field => - val c = field.name - if (colSet.contains(c)) { - convertToML(col(c)).as(c, field.metadata) - } else { - col(c) - } - } - import org.apache.spark.util.ArrayImplicits._ - dataset.select(exprs.toImmutableArraySeq: _*) + val fields = colNames.map(schema(_)) + dataset.withColumns( + fields.map(_.name), + fields.map(field => wrap_udt(unwrap_udt(col(field.name)), new MLMatrixUDT)), + fields.map(_.metadata)) } /** * Converts matrix columns in an input Dataset to the [[org.apache.spark.mllib.linalg.Matrix]] * type from the new [[org.apache.spark.ml.linalg.Matrix]] type under the `spark.ml` package. * @param dataset input dataset - * @param cols a list of matrix columns to be converted. Old matrix columns will be ignored. If - * unspecified, all new matrix columns will be converted except nested ones. + * @param cols a list of matrix columns to be converted. If unspecified, all new matrix columns + * will be converted except nested ones. * @return the input `DataFrame` with new matrix columns converted to the old matrix type */ @Since("2.0.0") @varargs def convertMatrixColumnsFromML(dataset: Dataset[_], cols: String*): DataFrame = { val schema = dataset.schema - val colSet = if (cols.nonEmpty) { + val colNames = if (cols.nonEmpty) { cols.flatMap { c => val dataType = schema(c).dataType if (dataType.getClass == classOf[MLMatrixUDT]) { @@ -495,35 +466,25 @@ object MLUtils extends Logging { s"Column $c must be new Matrix type to be converted to old type but got $dataType.") None } - }.toSet + } } else { schema.fields .filter(_.dataType.getClass == classOf[MLMatrixUDT]) .map(_.name) - .toSet + .toImmutableArraySeq } - if (colSet.isEmpty) { + if (colNames.isEmpty) { return dataset.toDF() } - logWarning("Matrix column conversion has serialization overhead. " + - "Please migrate your datasets and workflows to use the spark.ml package.") - - val convertFromML = udf { Matrices.fromML _ } - val exprs = schema.fields.map { field => - val c = field.name - if (colSet.contains(c)) { - convertFromML(col(c)).as(c, field.metadata) - } else { - col(c) - } - } - import org.apache.spark.util.ArrayImplicits._ - dataset.select(exprs.toImmutableArraySeq: _*) + val fields = colNames.map(schema(_)) + dataset.withColumns( + fields.map(_.name), + fields.map(field => wrap_udt(unwrap_udt(col(field.name)), new MatrixUDT)), + fields.map(_.metadata)) } - /** * Returns the squared Euclidean distance between two vectors. The following formula will be used * if it does not introduce too much numerical error: diff --git a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala index 7fcb1d2fbfbb9..c942a9a8c4f4a 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala @@ -19,16 +19,57 @@ package org.apache.spark.ml import org.apache.spark.SparkException import org.apache.spark.ml.functions._ -import org.apache.spark.ml.linalg.{Vector, Vectors} +import org.apache.spark.ml.linalg.{Matrices, MatrixUDT, Vector, Vectors, VectorUDT} import org.apache.spark.ml.util.MLTest -import org.apache.spark.mllib.linalg.{Vectors => OldVectors} -import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.functions.col +import org.apache.spark.mllib.linalg.{Matrices => OldMatrices, MatrixUDT => OldMatrixUDT, + Vector => OldVector, Vectors => OldVectors, VectorUDT => OldVectorUDT} +import org.apache.spark.sql.{AnalysisException, DataFrame, Row} +import org.apache.spark.sql.functions.{col, unwrap_udt, wrap_udt} +import org.apache.spark.sql.types.{StructField, StructType, UserDefinedType} class FunctionsSuite extends MLTest { import testImplicits._ + private def checkWrapUDTConversion( + df: DataFrame, + targetUDT: UserDefinedType[_], + expected: Any): Unit = { + val converted = df.select(wrap_udt(unwrap_udt(col("value")), targetUDT).as("value")) + assert(converted.schema("value").dataType === targetUDT) + assert(converted.first().get(0) === expected) + } + + private def checkWrapUDTTypeMismatch(df: DataFrame, targetUDT: UserDefinedType[_]): Unit = { + val e = intercept[AnalysisException] { + df.select(wrap_udt(unwrap_udt(col("value")), targetUDT).as("value")).collect() + } + assert(e.getCondition === "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE") + } + + private def checkNullableWrapUDTConversion( + value: Any, + sourceUDT: UserDefinedType[_], + targetUDT: UserDefinedType[_], + expected: Any): Unit = { + val schema = StructType(Seq(StructField("value", sourceUDT, nullable = true))) + val df = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(value), Row(null))), + schema) + val converted = df.select(wrap_udt(unwrap_udt(col("value")), targetUDT).as("value")) + + assert(converted.schema("value").dataType === targetUDT) + assert(converted.schema("value").nullable) + assert(converted.collect().map(_.get(0)).toSeq === Seq(expected, null)) + } + + private def normalizeNaN(rows: Seq[(Int, Int, Double)]): Seq[(Int, Int, String)] = { + rows.map { + case (id, index, value) if value.isNaN => (id, index, "NaN") + case (id, index, value) => (id, index, value.toString) + } + } + test("test vector_to_array") { val df = Seq( (Vectors.dense(1.0, 2.0, 3.0), OldVectors.dense(10.0, 20.0, 30.0)), @@ -103,6 +144,77 @@ class FunctionsSuite extends MLTest { assert(resultVec3 === Vectors.dense(Array(1.0, 2.0))) } + test("test vector_posexplode with vector UDT") { + val df = Seq( + (0, Vectors.dense(1.0, 0.0, 3.0), OldVectors.dense(10.0, 0.0, 30.0)), + (1, Vectors.sparse(4, Seq((1, 2.0), (2, 0.0), (3, 4.0))), + OldVectors.sparse(4, Seq((0, 20.0), (1, 0.0), (2, 30.0)))), + (2, null.asInstanceOf[Vector], null.asInstanceOf[OldVector]), + (3, Vectors.sparse(10, Array.emptyIntArray, Array.emptyDoubleArray), + OldVectors.sparse(10, Array.emptyIntArray, Array.emptyDoubleArray)), + (4, Vectors.dense(Array.emptyDoubleArray), + OldVectors.dense(Array.emptyDoubleArray)) + ).toDF("id", "vec", "oldVec") + + val result = df.select($"id", vector_posexplode($"vec")) + .as[(Int, Int, Double)] + .collect() + .toSeq + assert(normalizeNaN(result) === Seq( + (0, -4, "NaN"), + (0, 0, "1.0"), + (0, 2, "3.0"), + (1, -5, "NaN"), + (1, 1, "2.0"), + (1, 3, "4.0"), + (3, -11, "NaN"), + (4, -1, "NaN"))) + + val oldResult = df.select($"id", vector_posexplode($"oldVec")) + .as[(Int, Int, Double)] + .collect() + .toSeq + assert(normalizeNaN(oldResult) === Seq( + (0, -4, "NaN"), + (0, 0, "10.0"), + (0, 2, "30.0"), + (1, -5, "NaN"), + (1, 0, "20.0"), + (1, 2, "30.0"), + (3, -11, "NaN"), + (4, -1, "NaN"))) + + val denseResult = df + .where($"id" === 1) + .select($"id", vector_posexplode($"vec", mode = "dense")) + .as[(Int, Int, Double)] + .collect() + .toSeq + assert(normalizeNaN(denseResult) === Seq( + (1, -5, "NaN"), + (1, 0, "0.0"), + (1, 1, "2.0"), + (1, 2, "0.0"), + (1, 3, "4.0"))) + + val sparseResult = df.select($"id", vector_posexplode($"vec", mode = "sparse")) + .as[(Int, Int, Double)] + .collect() + .toSeq + assert(normalizeNaN(sparseResult) === Seq( + (0, -4, "NaN"), + (0, 0, "1.0"), + (0, 2, "3.0"), + (1, -5, "NaN"), + (1, 1, "2.0"), + (1, 3, "4.0"), + (3, -11, "NaN"), + (4, -1, "NaN"))) + + val schema = df.select(vector_posexplode($"vec")).schema + assert(schema.simpleString === "struct<index:int,value:double>") + } + test("test get_vector") { val df = Seq( (Vectors.dense(1.0, 2.0, 3.0), 0), @@ -130,4 +242,90 @@ class FunctionsSuite extends MLTest { val result = df.select(array_argmax(col("arr"))).as[Int].collect() assert(result === Array(2, 1, 0, 1, 0, -1)) } + + test("wrap and unwrap vector and matrix UDT columns") { + val oldVector = OldVectors.sparse(3, Array(1), Array(2.0)) + val oldVectorDF = Seq(Tuple1(oldVector)).toDF("value") + checkWrapUDTConversion( + oldVectorDF, + new OldVectorUDT, + oldVector) + checkWrapUDTConversion( + oldVectorDF, + new VectorUDT, + oldVector.asML) + checkWrapUDTTypeMismatch( + oldVectorDF, + new OldMatrixUDT) + checkWrapUDTTypeMismatch( + oldVectorDF, + new MatrixUDT) + + val mlVector = Vectors.dense(1.0, 2.0) + val mlVectorDF = Seq(Tuple1(mlVector)).toDF("value") + checkWrapUDTConversion( + mlVectorDF, + new VectorUDT, + mlVector) + checkWrapUDTConversion( + mlVectorDF, + new OldVectorUDT, + OldVectors.fromML(mlVector)) + checkWrapUDTTypeMismatch( + mlVectorDF, + new OldMatrixUDT) + checkWrapUDTTypeMismatch( + mlVectorDF, + new MatrixUDT) + + val oldMatrix = OldMatrices.dense(2, 2, Array(1.0, 2.0, 3.0, 4.0)) + val oldMatrixDF = Seq(Tuple1(oldMatrix)).toDF("value") + checkWrapUDTConversion( + oldMatrixDF, + new OldMatrixUDT, + oldMatrix) + checkWrapUDTConversion( + oldMatrixDF, + new MatrixUDT, + oldMatrix.asML) + checkWrapUDTTypeMismatch( + oldMatrixDF, + new OldVectorUDT) + checkWrapUDTTypeMismatch( + oldMatrixDF, + new VectorUDT) + + val mlMatrix = Matrices.dense(2, 2, Array(1.0, 2.0, 3.0, 4.0)) + val mlMatrixDF = Seq(Tuple1(mlMatrix)).toDF("value") + checkWrapUDTConversion( + mlMatrixDF, + new MatrixUDT, + mlMatrix) + checkWrapUDTConversion( + mlMatrixDF, + new OldMatrixUDT, + OldMatrices.fromML(mlMatrix)) + checkWrapUDTTypeMismatch( + mlMatrixDF, + new OldVectorUDT) + checkWrapUDTTypeMismatch( + mlMatrixDF, + new VectorUDT) + } + + test("wrap and unwrap nullable vector UDT columns") { + val oldVector = OldVectors.sparse(3, Array(1), Array(2.0)) + checkNullableWrapUDTConversion( + oldVector, + new OldVectorUDT, + new VectorUDT, + oldVector.asML) + + val mlVector = Vectors.dense(1.0, 2.0) + checkNullableWrapUDTConversion( + mlVector, + new VectorUDT, + new OldVectorUDT, + OldVectors.fromML(mlVector)) + } } diff --git a/mllib/src/test/scala/org/apache/spark/ml/ann/ANNSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/ann/ANNSuite.scala index e2d00c98f1ca8..c81d3763169b5 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/ann/ANNSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/ann/ANNSuite.scala @@ -25,6 +25,19 @@ import org.apache.spark.util.ArrayImplicits._ class ANNSuite extends SparkFunSuite with MLlibTestSparkContext { + test("layer models share one dense weights array") { + val topology = FeedForwardTopology.multiLayerPerceptron(Array(2, 3, 2)) + val numWeights = topology.layers.map(_.weightSize).sum + val weights = Vectors.sparse(numWeights, Array(0, numWeights - 1), Array(1.0, 2.0)) + + val model = topology.model(weights) + val affineLayers = model.layerModels.collect { case layer: AffineLayerModel => layer } + + assert(affineLayers.length === 2) + assert(affineLayers.forall(_.weights.data eq affineLayers.head.weights.data)) + assert(affineLayers.head.weights.data.length === numWeights) + } + // TODO: test for weights comparison with Weka MLP test("ANN with Sigmoid learns XOR function with LBFGS optimizer") { val inputs = Array( diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala index 765fccf6c6207..d298e625880f1 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala @@ -270,7 +270,7 @@ class DecisionTreeClassifierSuite extends MLTest with DefaultReadWriteTest { val transformed = newTree.transform(newData) checkNominalOnDF(transformed, "prediction", newTree.numClasses) - checkNominalOnDF(transformed, "predictedLeafId", newTree.numLeave) + checkNominalOnDF(transformed, "predictedLeafId", newTree.numLeaves) checkVectorSizeOnDF(transformed, "rawPrediction", newTree.numClasses) checkVectorSizeOnDF(transformed, "probability", newTree.numClasses) diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala index 70408731a20f6..75839a81c1f88 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala @@ -68,7 +68,7 @@ class OneVsRestSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(model) } - test("SPARK-58250: OneVsRestModel estimated size") { + test("SPARK-58250 and SPARK-58509: OneVsRestModel estimated size") { val trainingData = Seq( (0.0, Vectors.dense(0.0, 0.0)), (0.0, Vectors.dense(0.0, 1.0)), @@ -77,13 +77,16 @@ class OneVsRestSuite extends MLTest with DefaultReadWriteTest { (2.0, Vectors.dense(2.0, 0.0)), (2.0, Vectors.dense(2.0, 1.0))).toDF("label", "features") + val classifier = new LogisticRegression().setMaxIter(1) + // Initialize the classifier's logger before retaining it in the OneVsRestModel. + classifier.fit(trainingData) val model = new OneVsRest() - .setClassifier(new LogisticRegression().setMaxIter(1)) + .setClassifier(classifier) .fit(trainingData) val maxSize = 32 * 1024 assert(model.estimatedSize < maxSize, - s"Estimation (${model.estimatedSize}) should be less than $maxSize") + s"Estimation (${model.estimatedSize}) should not include shared runtime state") } test("one-vs-rest: default params") { diff --git a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala index a0223396da317..0b89495340039 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/LDASuite.scala @@ -236,6 +236,24 @@ class LDASuite extends MLTest with DefaultReadWriteTest { assert(lp <= 0.0 && lp != Double.NegativeInfinity) } + test("LocalLDAModel estimated size") { + val lda = new LDA().setK(2).setSeed(1).setOptimizer("online").setMaxIter(1) + val model = lda.fit(LDASuite.generateLDAData(spark, 3, 2, 3)) + val maxSize = 1024 * 16 + assert( + model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + + test("DistributedLDAModel estimated size") { + val lda = new LDA().setK(2).setSeed(1).setOptimizer("em").setMaxIter(1) + val model = lda.fit(LDASuite.generateLDAData(spark, 3, 2, 3)) + val maxSize = 1024 * 16 + assert( + model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + test("read/write LocalLDAModel") { def checkModelData(model: LDAModel, model2: LDAModel): Unit = { assert(model.vocabSize === model2.vocabSize) diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala index 295e96bcfe6a0..c007339980507 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala @@ -160,6 +160,23 @@ class CountVectorizerSuite extends MLTest with DefaultReadWriteTest { } } + test("CountVectorizer document frequency ignores duplicate tokens") { + val df = Seq( + Array("a", "a", "a", "b"), + Array("a", "b", "b", "b"), + Array("b") + ).toDF("words") + + val cvModel = new CountVectorizer() + .setInputCol("words") + .setOutputCol("features") + .setMinDF(2) + .setMaxDF(2) + .fit(df) + + assert(cvModel.vocabulary === Array("a")) + } + test("CountVectorizer using both minDF and maxDF") { // Ignore terms with count more than 3 AND less than 2 val df = Seq( diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala index 861bf1e0b1292..4eac3557cb931 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/HashingTFSuite.scala @@ -97,6 +97,18 @@ class HashingTFSuite extends MLTest with DefaultReadWriteTest { assert(loadedHashingTF.indexOf("c") === mLlibHashingTF.indexOf("c")) assert(loadedHashingTF.indexOf("d") === mLlibHashingTF.indexOf("d")) + mLlibHashingTF.setBinary(loadedHashingTF.getBinary) + val terms = "a a b b c d".split(" ").toSeq + val df = Seq(Tuple1(terms)).toDF("words") + val features = loadedHashingTF + .setInputCol("words") + .setOutputCol("features") + .transform(df) + .select("features") + .first() + .getAs[Vector](0) + assert(features ~== mLlibHashingTF.transform(terms).asML absTol 1e-14) + val metadata = spark.read.json(s"$hashingTFPath/metadata") val sparkVersionStr = metadata.select("sparkVersion").first().getString(0) assert(sparkVersionStr === "2.4.4") diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala index 1529b4c661cec..664a1937b381a 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala @@ -271,6 +271,13 @@ class VectorIndexerSuite extends MLTest with DefaultReadWriteTest with Logging { testTransformerByGlobalCheckFunc[FeatureData](points, model1, "indexed") { rows => assert(rows.map(_(0)) == expected) } + model1.set(model1.handleInvalid, "keep") + testTransformerByGlobalCheckFunc[FeatureData](pointsTestInvalid, model1, "indexed") { rows => + assert(rows.map(_(0)) == expected ++ Array( + Vectors.dense(2.0, 2.0, 0.0), + Vectors dense(0.0, 4.0, 2.0), + Vectors.dense(1.0, 3.0, 3.0))) + } val vectorIndexer2 = getIndexer.setMaxCategories(4).setHandleInvalid("keep") val model2 = vectorIndexer2.fit(points) testTransformerByGlobalCheckFunc[FeatureData](pointsTestInvalid, model2, "indexed") { rows => diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala index d7f15dc2cfe9e..07a11fa92caf1 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala @@ -119,6 +119,17 @@ class GBTRegressorSuite extends MLTest with DefaultReadWriteTest { testPredictionModelSinglePrediction(model, validationData.toDF()) } + test("prediction column has correct metadata") { + val tree = new DecisionTreeRegressionModel("dtr", TreeTests.root0, 3) + val model = new GBTRegressionModel("gbtr", Array(tree), Array(1.0), 3) + val df = sc.parallelize(TreeTests.getTwoTreesLeafData.toImmutableArraySeq, 1) + .toDF("leafId", "features") + + val expectedMetadata = model.transformSchema(df.schema)(model.getPredictionCol).metadata + val actualMetadata = model.transform(df).schema(model.getPredictionCol).metadata + assert(actualMetadata === expectedMetadata) + } + test("Checkpointing") { val tempDir = Utils.createTempDir() val path = tempDir.toURI.toString diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala index 3077a60b56b76..ae779a51e5999 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/regression/IsotonicRegressionSuite.scala @@ -22,6 +22,7 @@ import org.apache.spark.ml.param.ParamsSuite import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTest, MLTestingUtils} import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.functions.col +import org.apache.spark.util.SizeEstimator class IsotonicRegressionSuite extends MLTest with DefaultReadWriteTest { @@ -67,6 +68,16 @@ class IsotonicRegressionSuite extends MLTest with DefaultReadWriteTest { } } + test("model size estimation") { + val dataset = generateIsotonicInput(Seq(1, 2, 3, 1, 6, 17, 16, 17, 18)) + val model = new IsotonicRegression().fit(dataset) + + val expectedSize = model.estimateMatadataSize + + SizeEstimator.estimate(model.boundaries.toArray) + + SizeEstimator.estimate(model.predictions.toArray) + assert(model.estimatedSize === expectedSize) + } + test("params validation") { val dataset = generateIsotonicInput(Seq(1, 2, 3)) val ir = new IsotonicRegression diff --git a/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala index f97fefa245145..7673ec47c9f56 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala @@ -80,6 +80,22 @@ class CrossValidatorSuite } } + test("CrossValidatorModel estimated size") { + val estimator = new LogisticRegression().setMaxIter(1) + // Initialize the estimator's logger before retaining it in the CrossValidatorModel. + estimator.fit(dataset) + val model = new CrossValidator() + .setEstimator(estimator) + .setEstimatorParamMaps(Array(ParamMap.empty)) + .setEvaluator(new BinaryClassificationEvaluator()) + .setNumFolds(2) + .fit(dataset) + + val maxSize = 16 * 1024 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should not include shared runtime state") + } + test("cross validation with logistic regression with fold col") { val lr = new LogisticRegression val lrParamMaps = new ParamGridBuilder() diff --git a/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala index 20ba69a5adbb9..0b4b022f7e6c4 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala @@ -73,6 +73,21 @@ class TrainValidationSplitSuite } } + test("TrainValidationSplitModel estimated size") { + val estimator = new LogisticRegression().setMaxIter(1) + // Initialize the estimator's logger before retaining it in the TrainValidationSplitModel. + estimator.fit(dataset) + val model = new TrainValidationSplit() + .setEstimator(estimator) + .setEstimatorParamMaps(Array(ParamMap.empty)) + .setEvaluator(new BinaryClassificationEvaluator()) + .fit(dataset) + + val maxSize = 16 * 1024 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should not include shared runtime state") + } + test("train validation with linear regression") { val dataset = sc.parallelize( LinearDataGenerator.generateLinearInput( diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala index 2ba987b96ef79..10c86472ac0dd 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala @@ -40,12 +40,12 @@ class GaussianMixtureSuite extends SparkFunSuite with MLlibTestSparkContext { } } - test("single cluster") { + test("single cluster with empty partitions") { val data = sc.parallelize(Seq( Vectors.dense(6.0, 9.0), Vectors.dense(5.0, 10.0), Vectors.dense(4.0, 11.0) - )) + ), 4) // expectations val Ew = 1.0 diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala index 94afcb8c8e5c5..80ed707aec632 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/clustering/KMeansSuite.scala @@ -92,7 +92,7 @@ class KMeansSuite extends SparkFunSuite with MLlibTestSparkContext { .setInitializationSteps(10) .setSeed(seed) - val distanceMeasureInstance = new EuclideanDistanceMeasure + val distanceMeasureInstance = EuclideanDistanceMeasure val initialCenters = km.initKMeansParallel(normedData, distanceMeasureInstance).map(_.vector) assert(initialCenters.length === initialCenters.distinct.length) assert(initialCenters.length <= numDistinctPoints) diff --git a/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala index 0a4f81bd7aa9c..2500a52ac6320 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala @@ -176,6 +176,13 @@ class IsotonicRegressionSuite extends SparkFunSuite with MLlibTestSparkContext w assert(model.predictions === Array(1, 2, 3, 4, 5)) } + test("isotonic regression merges partial results across partitions") { + val model = runIsotonicRegressionOnInput(generateIsotonicInput(Seq(1, 3, 2, 4)), true, 2) + + assert(model.boundaries === Array(0, 1, 2, 3)) + assert(model.predictions === Array(1, 2.5, 2.5, 4)) + } + test("weighted isotonic regression") { val model = runIsotonicRegression(Seq(1, 2, 3, 4, 2), Seq(1, 1, 1, 1, 2), true) diff --git a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala index 5f6ac1b716cc9..376a51b872725 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/util/MLUtilsSuite.scala @@ -23,13 +23,15 @@ import java.nio.file.Files import scala.io.Source import org.apache.spark.{SparkException, SparkFunSuite, SparkRuntimeException} -import org.apache.spark.mllib.linalg.{DenseVector, Matrices, SparseVector, Vector, Vectors} +import org.apache.spark.ml.linalg.{MatrixUDT => MLMatrixUDT, VectorUDT => MLVectorUDT} +import org.apache.spark.mllib.linalg.{DenseVector, Matrices, MatrixUDT => OldMatrixUDT, + SparseVector, Vector, Vectors, VectorUDT => OldVectorUDT} import org.apache.spark.mllib.regression.LabeledPoint import org.apache.spark.mllib.util.MLUtils._ import org.apache.spark.mllib.util.TestingUtils._ import org.apache.spark.sql.Row import org.apache.spark.sql.functions.col -import org.apache.spark.sql.types.MetadataBuilder +import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StructField, StructType} import org.apache.spark.util.Utils class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { @@ -307,6 +309,30 @@ class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { } } + test("convert nullable vector columns") { + val oldVector = Vectors.sparse(2, Array(1), Array(1.0)) + val oldSchema = StructType(Seq( + StructField("id", IntegerType, nullable = false), + StructField("x", new OldVectorUDT, nullable = true))) + val oldDF = spark.createDataFrame( + sc.parallelize(Seq(Row(0, oldVector), Row(1, null))), + oldSchema) + val toML = convertVectorColumnsToML(oldDF) + assert(toML.schema("x").dataType === new MLVectorUDT) + assert(toML.collect().toSeq === Seq(Row(0, oldVector.asML), Row(1, null))) + + val mlVector = oldVector.asML + val mlSchema = StructType(Seq( + StructField("id", IntegerType, nullable = false), + StructField("x", new MLVectorUDT, nullable = true))) + val mlDF = spark.createDataFrame( + sc.parallelize(Seq(Row(0, mlVector), Row(1, null))), + mlSchema) + val fromML = convertVectorColumnsFromML(mlDF) + assert(fromML.schema("x").dataType === new OldVectorUDT) + assert(fromML.collect().toSeq === Seq(Row(0, oldVector), Row(1, null))) + } + test("convertMatrixColumnsToML") { val x = Matrices.sparse(3, 2, Array(0, 2, 3), Array(0, 2, 1), Array(0.0, -1.2, 0.0)) val metadata = new MetadataBuilder().putLong("numFeatures", 2L).build() @@ -357,6 +383,31 @@ class MLUtilsSuite extends SparkFunSuite with MLlibTestSparkContext { } } + test("convert nullable matrix columns") { + val oldMatrix = Matrices.sparse( + 3, 2, Array(0, 2, 3), Array(0, 2, 1), Array(0.0, -1.2, 0.0)) + val oldSchema = StructType(Seq( + StructField("id", IntegerType, nullable = false), + StructField("x", new OldMatrixUDT, nullable = true))) + val oldDF = spark.createDataFrame( + sc.parallelize(Seq(Row(0, oldMatrix), Row(1, null))), + oldSchema) + val toML = convertMatrixColumnsToML(oldDF) + assert(toML.schema("x").dataType === new MLMatrixUDT) + assert(toML.collect().toSeq === Seq(Row(0, oldMatrix.asML), Row(1, null))) + + val mlMatrix = oldMatrix.asML + val mlSchema = StructType(Seq( + StructField("id", IntegerType, nullable = false), + StructField("x", new MLMatrixUDT, nullable = true))) + val mlDF = spark.createDataFrame( + sc.parallelize(Seq(Row(0, mlMatrix), Row(1, null))), + mlSchema) + val fromML = convertMatrixColumnsFromML(mlDF) + assert(fromML.schema("x").dataType === new OldMatrixUDT) + assert(fromML.collect().toSeq === Seq(Row(0, oldMatrix), Row(1, null))) + } + test("kFold with fold column") { val data = sc.parallelize(1 to 100, 2).map(x => (x, if (x <= 50) 0 else 1)).toDF("i", "fold") val collectedData = data.collect().map(_.getInt(0)).sorted diff --git a/pom.xml b/pom.xml index 5e1a785456ae4..e40270844945e 100644 --- a/pom.xml +++ b/pom.xml @@ -129,7 +129,7 @@ <sbt.project.name>spark</sbt.project.name> <asm.version>9.10.1</asm.version> <slf4j.version>2.0.17</slf4j.version> - <log4j.version>2.26.0</log4j.version> + <log4j.version>2.26.1</log4j.version> <!-- make sure to update IsolatedClientLoader whenever this version is changed --> <hadoop.version>3.5.0</hadoop.version> <!-- SPARK-41247: When updating `protobuf.version`, also need to update `protoVersion` in `SparkBuild.scala` --> @@ -173,7 +173,7 @@ <commons.httpclient.version>4.5.14</commons.httpclient.version> <commons.httpcore.version>4.4.16</commons.httpcore.version> <commons.math3.version>3.6.1</commons.math3.version> - <commons.collections4.version>4.5.0</commons.collections4.version> + <commons.collections4.version>4.6.0</commons.collections4.version> <scala.version>2.13.18</scala.version> <scala.binary.version>2.13</scala.binary.version> <scalatest-maven-plugin.version>2.2.0</scalatest-maven-plugin.version> @@ -185,11 +185,11 @@ <scalafmt.validateOnly>true</scalafmt.validateOnly> <scalafmt.changedOnly>true</scalafmt.changedOnly> <!-- Should be consistent with SparkBuild.scala and docs --> - <fasterxml.jackson.version>2.22.0</fasterxml.jackson.version> + <fasterxml.jackson.version>2.22.1</fasterxml.jackson.version> <ws.xmlschema.version>2.3.1</ws.xmlschema.version> <snappy.version>1.1.10.8</snappy.version> <netlib.ludovic.dev.version>3.2.0</netlib.ludovic.dev.version> - <commons-codec.version>1.22.0</commons-codec.version> + <commons-codec.version>1.22.1</commons-codec.version> <commons-compress.version>1.28.0</commons-compress.version> <commons-io.version>2.22.0</commons-io.version> <!-- To support Hive UDF jars built by Hive 2.0.0 ~ 2.3.9 and 3.0.0 ~ 3.1.3. --> @@ -202,7 +202,7 @@ <guava.version>33.6.0-jre</guava.version> <guava.failureaccess.version>1.0.3</guava.failureaccess.version> <gson.version>2.14.0</gson.version> - <janino.version>3.1.9</janino.version> + <janino.version>3.1.12</janino.version> <jersey.version>3.1.11</jersey.version> <joda.version>2.14.3</joda.version> <jsr305.version>3.0.0</jsr305.version> @@ -222,7 +222,7 @@ SPARK-53327 workaround should be reverted. --> <datasketches.version>6.2.0</datasketches.version> - <netty.version>4.2.16.Final</netty.version> + <netty.version>4.2.17.Final</netty.version> <netty-tcnative.version>2.0.81.Final</netty-tcnative.version> <icu4j.version>78.3</icu4j.version> <junit.version>6.0.3</junit.version> @@ -286,7 +286,6 @@ --> <derby.deps.scope>compile</derby.deps.scope> <hadoop.deps.scope>compile</hadoop.deps.scope> - <huaweicloud.deps.scope>compile</huaweicloud.deps.scope> <hive.deps.scope>compile</hive.deps.scope> <hive.storage.version>2.8.1</hive.storage.version> <hive.storage.scope>compile</hive.storage.scope> @@ -879,12 +878,12 @@ <dependency> <groupId>at.yawk.lz4</groupId> <artifactId>lz4-java</artifactId> - <version>1.11.1</version> + <version>1.11.2</version> </dependency> <dependency> <groupId>com.github.luben</groupId> <artifactId>zstd-jni</artifactId> - <version>1.5.7-9</version> + <version>1.5.7-13</version> </dependency> <dependency> <groupId>com.clearspring.analytics</groupId> @@ -1008,7 +1007,7 @@ </exclusions> </dependency> <!-- SPARK-38885: After Netty 4.1.76, add the following `Netty` dependencies explicitly - to ensure `./dev/test-dependencies.sh` produce the same results on Linux and MacOS. + to ensure `./dev/test-dependencies.sh` produces the same results on Linux and macOS. --> <dependency> <groupId>io.netty</groupId> @@ -2315,7 +2314,7 @@ </exclusion> </exclusions> </dependency> - <!-- hive-llap-client is needed when run MapReduce test in Hive 2.3. --> + <!-- hive-llap-client is needed when running MapReduce tests in Hive 2.3. --> <dependency> <groupId>${hive.group}</groupId> <artifactId>hive-llap-client</artifactId> @@ -2642,7 +2641,7 @@ <artifactId>arpack</artifactId> <version>${netlib.ludovic.dev.version}</version> </dependency> - <!-- SPARK-16484 add `datasketches-java` for support Datasketches HllSketch --> + <!-- SPARK-16484 add `datasketches-java` to support Datasketches HllSketch --> <dependency> <groupId>org.apache.datasketches</groupId> <artifactId>datasketches-java</artifactId> @@ -3365,10 +3364,10 @@ </executions> </plugin> <!-- - Couple of dependencies are coming in bundle format (bundle is just a normal jar which - contains OSGi metadata in the manifest). If one don't use OSGi, then a bundle will work as + A couple of dependencies are coming in bundle format (bundle is just a normal jar which + contains OSGi metadata in the manifest). If one doesn't use OSGi, then a bundle will work as any other jar. Since maven doesn't have native bundle support it needs an external plugin - handle it. If the plugin is not added then the build can't resolve bundle dependencies. + to handle it. If the plugin is not added then the build can't resolve bundle dependencies. --> <plugin> <groupId>org.apache.felix</groupId> @@ -3519,6 +3518,13 @@ </modules> </profile> + <profile> + <id>credential-aws</id> + <modules> + <module>connector/credential-aws</module> + </modules> + </profile> + <profile> <id>test-java-home</id> <activation> @@ -3540,6 +3546,12 @@ <url>${env.MAVEN_MIRROR_URL}</url> </repository> </repositories> + <pluginRepositories> + <pluginRepository> + <id>maven-mirror</id> + <url>${env.MAVEN_MIRROR_URL}</url> + </pluginRepository> + </pluginRepositories> </profile> <!-- @@ -3564,7 +3576,7 @@ <!-- This is a profile to enable the use of the ASF snapshot and staging repositories during a build. It is useful when testing against nightly or RC releases of dependencies. - It MUST NOT be used when building copies of Spark to use in production of for distribution, + It MUST NOT be used when building copies of Spark to use in production or for distribution, --> <profile> <id>snapshots-and-staging</id> @@ -3618,7 +3630,6 @@ <id>hadoop-provided</id> <properties> <spark.yarn.isHadoopProvided>true</spark.yarn.isHadoopProvided> - <huaweicloud.deps.scope>provided</huaweicloud.deps.scope> </properties> </profile> <profile> @@ -3639,7 +3650,7 @@ <profile> <id>jjwt-provided</id> </profile> - <!-- use org.openlabtesting.leveldbjni on aarch64 platform except MacOS --> + <!-- use org.openlabtesting.leveldbjni on aarch64 platform except macOS --> <profile> <id>aarch64</id> <properties> diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala index d8e74fe32a108..c27ef26959edf 100644 --- a/project/MimaExcludes.scala +++ b/project/MimaExcludes.scala @@ -33,11 +33,24 @@ import com.typesafe.tools.mima.core.* */ object MimaExcludes { - // Exclude rules for 5.0.x from 4.3.0 (add 5.0-specific filters below as needed). - lazy val v50excludes: Seq[Problem => Boolean] = v43excludes + // Exclude rules for 5.0.x from 4.4.0 (add 5.0-specific filters below as needed). + lazy val v50excludes: Seq[Problem => Boolean] = v44excludes ++ Seq( + // [SPARK-58896] Decision tree leaf counts moved behind NodeStats. The old package-private + // numLeave accessors are removed from concrete models. + ProblemFilters.exclude[DirectMissingMethodProblem]( + "org.apache.spark.ml.classification.DecisionTreeClassificationModel.numLeave"), + ProblemFilters.exclude[DirectMissingMethodProblem]( + "org.apache.spark.ml.regression.DecisionTreeRegressionModel.numLeave") + ) + + // Exclude rules for 4.4.x from 4.3.0 (add 4.4-specific filters below as needed). + lazy val v44excludes: Seq[Problem => Boolean] = v43excludes // Exclude rules for 4.3.x from 4.2.0 (add 4.3-specific filters below as needed). lazy val v43excludes: Seq[Problem => Boolean] = v42excludes ++ Seq( + // [SPARK-54879] Add exitCode field to ApplicationAttemptInfo + ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.tupled"), + ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.curried"), // [SPARK-58192][CORE] Support fractional spark.task.cpus. The existing TaskContext.cpus(): Int // and TaskResourceRequests.cpus(Int) are retained (cpus() deprecated in favor of the new // TaskContext.cpuAmount(): BigDecimal), but the internal StatusUpdate message's taskCpus field @@ -62,9 +75,6 @@ object MimaExcludes { // Exclude rules for 4.2.x from 4.1.0 lazy val v42excludes = v41excludes ++ Seq( - // [SPARK-54879] Add exitCode field to ApplicationAttemptInfo - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.tupled"), - ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.curried"), // [SQL] SafeJsonSerializer.safeMapToJValue: second parameter widened from Function1 to // Function2 so the key is passed to the value serializer (progress.scala). Binary-incompatible // vs spark-sql-api 4.0.0; not part of the public supported API (private[streaming] package). @@ -198,6 +208,7 @@ object MimaExcludes { def excludes(version: String): Seq[Problem => Boolean] = version match { case v if v.startsWith("5.0") => v50excludes + case v if v.startsWith("4.4") => v44excludes case v if v.startsWith("4.3") => v43excludes case v if v.startsWith("4.2") => v42excludes case v if v.startsWith("4.1") => v41excludes diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index ef027de5699f6..6fffcaf47cb5b 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -72,10 +72,10 @@ object BuildCommons { udfWorkerProjects val optionallyEnabledProjects@Seq(kubernetes, yarn, - sparkGangliaLgpl, streamingKinesisAsl, profiler, + sparkGangliaLgpl, streamingKinesisAsl, profiler, credentialAws, dockerIntegrationTests, hadoopCloud, kubernetesIntegrationTests) = Seq("kubernetes", "yarn", - "ganglia-lgpl", "streaming-kinesis-asl", "profiler", + "ganglia-lgpl", "streaming-kinesis-asl", "profiler", "credential-aws", "docker-integration-tests", "hadoop-cloud", "kubernetes-integration-tests").map(ProjectRef(buildLocation, _)) val assemblyProjects@Seq(networkYarn, streamingKafka010Assembly, streamingKinesisAslAssembly) = @@ -416,7 +416,8 @@ object SparkBuild extends PomBuild { Seq( spark, hive, hiveThriftServer, repl, networkCommon, networkShuffle, networkYarn, unsafe, tags, tokenProviderKafka010, sqlKafka010, pipelines, connectCommon, connect, - connectJdbc, connectClient, variant, connectShims, profiler, commonUtilsJava, sparkConfig, + connectJdbc, connectClient, variant, connectShims, profiler, credentialAws, + commonUtilsJava, sparkConfig, udfWorkerProto, udfWorkerCore, udfWorkerGrpc ).contains(x) } @@ -553,7 +554,7 @@ object SparkParallelTestGrouping { // SBT project. Here, we take an opt-in approach where the default behavior is to run all // tests sequentially in a single JVM, requiring us to manually opt-in to the extra parallelism. // - // There are a reasons why such an opt-in approach is good: + // There are reasons why such an opt-in approach is good: // // 1. Launching one JVM per suite adds significant overhead for short-running suites. In // addition to JVM startup time and JIT warmup, it appears that initialization of Derby @@ -1307,7 +1308,7 @@ object KubernetesIntegrationTests { * Overrides to work around sbt's dependency resolution being different from Maven's. */ object DependencyOverrides { - lazy val jacksonVersion = sys.props.get("fasterxml.jackson.version").getOrElse("2.22.0") + lazy val jacksonVersion = sys.props.get("fasterxml.jackson.version").getOrElse("2.22.1") lazy val jacksonDeps = Bom.dependencies("com.fasterxml.jackson" % "jackson-bom" % jacksonVersion) lazy val grpcVersion = sys.props.get("io.grpc.version").getOrElse("1.76.0") lazy val grpcDeps = Bom.dependencies("io.grpc" % "grpc-bom" % grpcVersion) diff --git a/pyproject.toml b/pyproject.toml index 35de692cf5b0d..9fcd00e69bf9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ line-length = 100 [tool.ruff.lint] extend-select = [ + "I", # isort "G010", # logging-warn # ambiguous unicode character "RUF001", # string @@ -64,14 +65,14 @@ ignore = [ "python/pyspark/sql/tests/df_golden/scripts/**/*.py" = ["F821"] [dependency-groups] -_py4j = ["py4j>=0.10.9.9"] +internal_py4j = ["py4j>=0.10.9.9"] -_numpy = ["numpy>=1.23.2"] -_pandas = ["pandas>=2.2.0,<3.0.0"] -_protobuf = ["protobuf==6.33.5"] -_pyarrow = ["pyarrow>=18.0.0"] +internal_numpy = ["numpy>=1.23.2"] +internal_pandas = ["pandas>=2.2.0,<3.0.0"] +internal_protobuf = ["protobuf==6.33.5"] +internal_pyarrow = ["pyarrow>=18.0.0"] -_pyspark_devtool = [ +internal_pyspark_devtool = [ # Optional tools that are used by pyspark directly. "memory-profiler>=0.61.0", "flameprof==0.4", @@ -80,10 +81,10 @@ _pyspark_devtool = [ ] connect = [ - {include-group = "_py4j"}, - {include-group = "_pandas"}, - {include-group = "_protobuf"}, - {include-group = "_pyarrow"}, + {include-group = "internal_py4j"}, + {include-group = "internal_pandas"}, + {include-group = "internal_protobuf"}, + {include-group = "internal_pyarrow"}, "grpcio>=1.76.0", "grpcio-status>=1.76.0", "googleapis-common-protos>=1.71.0", @@ -92,15 +93,15 @@ connect = [ ] sql = [ - {include-group = "_py4j"}, - {include-group = "_pandas"}, - {include-group = "_pyarrow"}, + {include-group = "internal_py4j"}, + {include-group = "internal_pandas"}, + {include-group = "internal_pyarrow"}, ] pandas_on_spark = [ - {include-group = "_py4j"}, - {include-group = "_pandas"}, - {include-group = "_pyarrow"}, + {include-group = "internal_py4j"}, + {include-group = "internal_pandas"}, + {include-group = "internal_pyarrow"}, ] pandas_on_spark_extra = [ @@ -111,13 +112,13 @@ pandas_on_spark_extra = [ ] ml = [ - {include-group = "_py4j"}, - {include-group = "_numpy"}, + {include-group = "internal_py4j"}, + {include-group = "internal_numpy"}, ] mllib = [ - {include-group = "_py4j"}, - {include-group = "_numpy"}, + {include-group = "internal_py4j"}, + {include-group = "internal_numpy"}, ] ml_extra_base = [ @@ -145,7 +146,7 @@ pipelines = [ "pyyaml>=3.11", ] -_lint = [ +internal_lint = [ "ruff==0.14.8", "mypy==1.19.1", "pytest", @@ -159,9 +160,11 @@ _lint = [ "types-PyYAML" ] -_test = [ +internal_test = [ # Dependency for running tests "unittest-xml-reporting", + # Import graph support for smart test selection + "grimp", # Coverage tool "coverage", # 3rd party libraries that unittests use. @@ -170,9 +173,12 @@ _test = [ "lxml", "testcontainers[kafka]>=3.7.0", "kafka-python-ng>=2.0.2", + # Used by the opt-in UDF transpile hypothesis test suite (gated on + # the RUN_HYPOTHESIS env var); the suite skips cleanly if absent. + "hypothesis>=6,<7", ] -_docs = [ +internal_docs = [ "sphinx==8.2.3", "sphinx-plotly-directive", "sphinx-copybutton", @@ -185,7 +191,7 @@ _docs = [ "mkdocs", ] -_devtools = [ +internal_devtools = [ # Devtools that are used by pyspark developers. "debugpy", "viztracer", @@ -195,11 +201,11 @@ _devtools = [ ] dev = [ - {include-group = "_pyspark_devtool"}, - {include-group = "_lint"}, - {include-group = "_test"}, - {include-group = "_docs"}, - {include-group = "_devtools"}, + {include-group = "internal_pyspark_devtool"}, + {include-group = "internal_lint"}, + {include-group = "internal_test"}, + {include-group = "internal_docs"}, + {include-group = "internal_devtools"}, {include-group = "sql"}, {include-group = "connect"}, {include-group = "ml"}, @@ -210,7 +216,7 @@ dev = [ {include-group = "pipelines"}, ] -_ci_connect = [ +internal_ci_connect = [ {include-group = "connect"}, "grpcio==1.76.0", "grpcio-status==1.76.0", @@ -220,24 +226,24 @@ _ci_connect = [ ] ci_connect_standard = [ - {include-group = "_ci_connect"}, + {include-group = "internal_ci_connect"}, "graphviz==0.20.3", ] ci_connect_minimum = [ - {include-group = "_ci_connect"}, + {include-group = "internal_ci_connect"}, "graphviz==0.20", ] ci_classic_standard = [ - {include-group = "_pyspark_devtool"}, + {include-group = "internal_pyspark_devtool"}, {include-group = "sql"}, {include-group = "pandas_on_spark"}, {include-group = "pandas_on_spark_extra"}, {include-group = "ml"}, # torch should be installed with ml_torch group {include-group = "ml_extra_base"}, - {include-group = "_test"}, + {include-group = "internal_test"}, "pyarrow>=23.0.0", "pandas==2.3.3", ] @@ -257,10 +263,10 @@ ci_classic_minimum = [ ] ci_lint = [ - {include-group = "_lint"}, + {include-group = "internal_lint"}, {include-group = "ml_extra_base"}, {include-group = "pandas_on_spark_extra"}, - {include-group = "_ci_connect"}, + {include-group = "internal_ci_connect"}, "pyarrow>=23.0.0", "pandas==2.3.3", "pandas-stubs==2.3.3.260113", @@ -268,9 +274,9 @@ ci_lint = [ ] ci_docs = [ - {include-group = "_docs"}, - {include-group = "_ci_connect"}, - {include-group = "_numpy"}, + {include-group = "internal_docs"}, + {include-group = "internal_ci_connect"}, + {include-group = "internal_numpy"}, # These are required for sphinx to introspect the code "pyarrow>=23.0.0", "pandas==2.3.3", diff --git a/python/MANIFEST.in b/python/MANIFEST.in index 82979a0344c3c..e57d60b8db982 100644 --- a/python/MANIFEST.in +++ b/python/MANIFEST.in @@ -14,20 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Reference: https://setuptools.pypa.io/en/latest/userguide/miscellaneous.html - -recursive-include pyspark *.pyi py.typed *.json -recursive-include deps/jars *.jar -graft deps/bin -recursive-include deps/sbin spark-config.sh spark-daemon.sh start-history-server.sh stop-history-server.sh -recursive-include deps/data *.data *.txt -graft deps/licenses -recursive-include deps/examples *.py -recursive-include lib *.zip -include README.md -include LICENSE -include NOTICE +# NOTE +# +# This file exists only for the exclude below, which applies across every +# package in every setup.py. (setuptools's `exclude_package_data` is scoped +# to the individual package and setup.py it's declared in). To include non- +# `.py` files, use `package_data` in setup.py instead of adding them here. +# +# See: +# - https://setuptools.pypa.io/en/latest/userguide/datafiles.html +# - https://setuptools.pypa.io/en/latest/userguide/miscellaneous.html -# Note that these commands are processed in the order they appear, so keep -# this exclude at the end. +# Keep this exclude at the end. MANIFEST.in commands are processed in the order +# they appear. global-exclude *.py[cod] __pycache__ .DS_Store diff --git a/python/README.md b/python/README.md index ca4e493fdd39d..9fc7cf54f5106 100644 --- a/python/README.md +++ b/python/README.md @@ -1,33 +1,36 @@ # Apache Spark Spark is a unified analytics engine for large-scale data processing. It provides -high-level APIs in Scala, Java, Python, and R, and an optimized engine that +high-level APIs in Scala, Java, Python, and R (deprecated), and an optimized engine that supports general computation graphs for data analysis. It also supports a rich set of higher-level tools including Spark SQL for SQL and DataFrames, pandas API on Spark for pandas workloads, MLlib for machine learning, GraphX for graph processing, and Structured Streaming for stream processing. -<https://spark.apache.org/> +PySpark is the Python distribution of Spark. -## Online Documentation +Project home page: https://spark.apache.org/ -You can find the latest Spark documentation, including a programming -guide, on the [project web page](https://spark.apache.org/documentation.html) +Main documentation: https://spark.apache.org/docs/latest/ +PySpark documentation: https://spark.apache.org/docs/latest/api/python/index.html -## Python Packaging +## PySpark on PyPI -This README file only contains basic information related to pip installed PySpark. -This packaging is currently experimental and may change in future versions (although we will do our best to keep compatibility). -Using PySpark requires the Spark JARs, and if you are building this from source please see the builder instructions at -["Building Spark"](https://spark.apache.org/docs/latest/building-spark.html). +There are a few PySpark packages published by the Apache Spark project to PyPI: -The Python packaging for Spark is not intended to replace all of the other use cases. This Python packaged version of Spark is suitable for interacting with an existing cluster (be it Spark standalone, YARN) - but does not contain the tools required to set up your own standalone Spark cluster. You can download the full version of Spark from the [Apache Spark downloads page](https://spark.apache.org/downloads.html). +- `pyspark`: Classic PySpark, includes Spark assembly JARs +- `pyspark-connect`: Classic PySpark with Spark Connect configured as the default, includes `pyspark` +- `pyspark-client`: Pure Python Spark Connect client, no JARs or JRE needed +For more information, see the [installation guide][install]. If you're building PySpark from source, see [Building Spark][build] and [pyproject.toml][py]. -**NOTE:** If you are using this with a Spark standalone cluster you must ensure that the version (including minor version) matches or you may experience odd errors. +[install]: https://spark.apache.org/docs/latest/api/python/getting_started/install.html +[build]: https://spark.apache.org/docs/latest/building-spark.html +[py]: https://github.com/apache/spark/blob/master/pyproject.toml -## Python Requirements +## Python vs. "Full" Distribution of Spark -At its core PySpark depends on Py4J, but some additional sub-packages have their own extra requirements for some features (including numpy, pandas, and pyarrow). -See also [Dependencies](https://spark.apache.org/docs/latest/api/python/getting_started/install.html#dependencies) for production, and [pyproject.toml](https://github.com/apache/spark/blob/master/pyproject.toml) for development. +PySpark is not intended to be a complete distribution of Spark. It's meant for local development or for interacting with an existing cluster (be it Spark standalone, YARN, or Kubernetes). Using PySpark to set up a new standalone Spark cluster is not supported. To set up a standalone cluster please [use the full distribution of Spark](https://spark.apache.org/downloads.html). + +When using PySpark with an existing Spark cluster you must ensure that the major and minor version (e.g. `4.3.*`) match or you may experience odd errors. diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py index 3228d915f44e0..1c0a4e47dccf0 100644 --- a/python/benchmarks/bench_eval_type.py +++ b/python/benchmarks/bench_eval_type.py @@ -24,8 +24,8 @@ """ import io -import os import json +import os import socket import struct import sys @@ -35,9 +35,8 @@ import numpy as np import pyarrow as pa - from pyspark.cloudpickle import dumps as cloudpickle_dumps -from pyspark.serializers import CPickleSerializer, write_int, write_long, SpecialLengths +from pyspark.serializers import CPickleSerializer, SpecialLengths, write_int, write_long from pyspark.sql.types import ( BinaryType, BooleanType, @@ -52,7 +51,6 @@ from pyspark.util import PythonEvalType from pyspark.worker import main as worker_main - # --------------------------------------------------------------------------- # Mock helpers: protocol writer, data factory, UDF factory # --------------------------------------------------------------------------- @@ -119,6 +117,9 @@ def write_preamble(cls, buf: io.BytesIO) -> None: "attemptNumber": 0, "taskAttemptId": 0, "cpus": 1, + # Plain decimal string, matching CpuAmount.toDisplayString on + # the JVM side; TaskContextInfo.from_stream requires the key. + "cpuAmount": "1", "resources": {}, "localProperties": {}, } diff --git a/python/benchmarks/bench_local_data_to_arrow.py b/python/benchmarks/bench_local_data_to_arrow.py new file mode 100644 index 0000000000000..0c41498c8ec10 --- /dev/null +++ b/python/benchmarks/bench_local_data_to_arrow.py @@ -0,0 +1,162 @@ +# +# 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. +# + +""" +Microbenchmarks for ``LocalDataToArrowConversion.convert``, the hot path of +Spark Connect ``createDataFrame`` and Arrow-optimized Python UDF output. + +``string`` / ``binary`` columns are converted element by element: the per-element +converter is fast when the element is already the target type (``str`` / +``bytes``) and slower when it needs coercion. Each benchmark sweeps the fraction +of non-target ("other") values from 0% (all fast path) to 100% (all slow path), +for a scalar column, a one-level ``array`` column, and a two-level +``array<array<...>>`` column (which drives the inlined array fast path hardest). +""" + + +def _build(convert, is_string, leaf, other_val, other_frac, n_rows): + import random + + from pyspark.sql.types import ArrayType, BinaryType, StringType, StructField, StructType + + leaf_type = StringType() if is_string else BinaryType() + + def target(i): + return f"s{i}" if is_string else b"s%d" % i + + rnd = random.Random(0) + + def elem(i): + return other_val(i) if rnd.random() * 100 < other_frac else target(i) + + if leaf == "array2": + schema = StructType([StructField("c", ArrayType(ArrayType(leaf_type)), True)]) + data = [([[elem(i), elem(i + 1)], [elem(i + 2), elem(i + 3)]],) for i in range(n_rows)] + elif leaf == "array": + schema = StructType([StructField("c", ArrayType(leaf_type), True)]) + data = [([elem(i), elem(i + 1), elem(i + 2)],) for i in range(n_rows)] + else: # scalar + schema = StructType([StructField("c", leaf_type, True)]) + data = [(elem(i),) for i in range(n_rows)] + return data, schema + + +class LocalDataToArrowStringBenchmark: + """ + Benchmark ``convert`` on a ``string`` column, sweeping the kind and fraction + of non-target values. + + - ``leaf``: ``scalar``, one-level ``array``, or two-level ``array2`` + (``array<array<string>>``). + - ``other``: the non-target values -- ``none`` (null), ``int`` or ``bool`` + (a single non-str type, isolating one coercion branch), or ``mix`` (a + rotation of int / float / bool / Decimal / date / datetime). + - ``other_frac``: fraction of elements that are non-target, ``0`` to ``100``. + """ + + params = [ + [1000000], + ["scalar", "array", "array2"], + ["none", "int", "bool", "mix"], + [0, 30, 70, 100], + ] + param_names = ["n_rows", "leaf", "other", "other_frac"] + + def setup(self, n_rows, leaf, other, other_frac): + import datetime + import decimal + + from pyspark.sql.conversion import LocalDataToArrowConversion + + self.convert = LocalDataToArrowConversion.convert + + if other == "none": + + def other_val(i): + return None + elif other == "int": + + def other_val(i): + return i # int coerced to string via str() + elif other == "bool": + + def other_val(i): + return i % 2 == 0 # bool coerced to "true" / "false" + else: # mix: a rotation of non-str types, each coerced to string + _pool = [ + 123, + 4.5, + True, + False, + decimal.Decimal("1.50"), + datetime.date(2020, 1, 1), + datetime.datetime(2020, 1, 1, 3, 4, 5), + ] + + def other_val(i): + return _pool[i % len(_pool)] + + self.data, self.schema = _build(self.convert, True, leaf, other_val, other_frac, n_rows) + + def time_convert(self, n_rows, leaf, other, other_frac): + self.convert(self.data, self.schema, False) + + def peakmem_convert(self, n_rows, leaf, other, other_frac): + self.convert(self.data, self.schema, False) + + +class LocalDataToArrowBinaryBenchmark: + """ + Benchmark ``convert`` on a ``binary`` column, sweeping the fraction of + non-target values. + + - ``leaf``: ``scalar``, one-level ``array``, or two-level ``array2`` + (``array<array<binary>>``). + - ``other``: ``none`` (null) or ``bytearray`` (the only non-bytes input + binary accepts, copied to immutable bytes). + - ``other_frac``: fraction of elements that are non-target, ``0`` to ``100``. + """ + + params = [ + [1000000], + ["scalar", "array", "array2"], + ["none", "bytearray"], + [0, 30, 70, 100], + ] + param_names = ["n_rows", "leaf", "other", "other_frac"] + + def setup(self, n_rows, leaf, other, other_frac): + from pyspark.sql.conversion import LocalDataToArrowConversion + + self.convert = LocalDataToArrowConversion.convert + + if other == "none": + + def other_val(i): + return None + else: # bytearray coerced to immutable bytes + + def other_val(i): + return bytearray(b"s%d" % i) + + self.data, self.schema = _build(self.convert, False, leaf, other_val, other_frac, n_rows) + + def time_convert(self, n_rows, leaf, other, other_frac): + self.convert(self.data, self.schema, False) + + def peakmem_convert(self, n_rows, leaf, other, other_frac): + self.convert(self.data, self.schema, False) diff --git a/python/benchmarks/bench_pipelined_udf.py b/python/benchmarks/bench_pipelined_udf.py index 31eb9d0b16e39..17b5be6985236 100644 --- a/python/benchmarks/bench_pipelined_udf.py +++ b/python/benchmarks/bench_pipelined_udf.py @@ -24,7 +24,6 @@ """ import pandas as pd - from pyspark import SparkConf from pyspark.sql import SparkSession from pyspark.sql.functions import col, pandas_udf diff --git a/python/conf_viztracer/daemon_viztracer.py b/python/conf_viztracer/daemon_viztracer.py index 216f6368080d9..ea6d736f68aa5 100644 --- a/python/conf_viztracer/daemon_viztracer.py +++ b/python/conf_viztracer/daemon_viztracer.py @@ -18,11 +18,10 @@ import os import sys +import pyspark.worker import viztracer from viztracer.main import main -import pyspark.worker - def viztracer_wrapper(func): def wrapper(*args, **kwargs): diff --git a/python/conf_viztracer/worker_viztracer.py b/python/conf_viztracer/worker_viztracer.py index 362532c39dc40..7ffd832e53d9f 100644 --- a/python/conf_viztracer/worker_viztracer.py +++ b/python/conf_viztracer/worker_viztracer.py @@ -20,7 +20,6 @@ from viztracer.main import main - if __name__ == "__main__": if os.getenv("SPARK_VIZTRACER_OUTPUT_DIR") is not None: diff --git a/python/conf_vscode/sitecustomize.py b/python/conf_vscode/sitecustomize.py index 486d580efdc1e..113252c4c5c05 100644 --- a/python/conf_vscode/sitecustomize.py +++ b/python/conf_vscode/sitecustomize.py @@ -24,9 +24,10 @@ ): def install_debugpy(): - import debugpy import fcntl + import debugpy + lock_file = os.getenv("DEBUGPY_ADAPTER_ENDPOINTS") + ".lock" try: fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o600) diff --git a/python/docs/source/conf.py b/python/docs/source/conf.py index 53e940add09ec..36da7608c05f3 100644 --- a/python/docs/source/conf.py +++ b/python/docs/source/conf.py @@ -73,6 +73,8 @@ 'IPython.sphinxext.ipython_console_highlighting', 'numpydoc', # handle NumPy documentation formatted docstrings. 'sphinx_plotly_directive', # For visualize plot result + # Local: forbid ':lines:' on 'literalinclude' (fragile line-number pinning). + 'forbid_literalinclude_lines', ] # sphinx copy button diff --git a/python/docs/source/development/contributing.rst b/python/docs/source/development/contributing.rst index bbf7c49a17d7c..d422b69528df6 100644 --- a/python/docs/source/development/contributing.rst +++ b/python/docs/source/development/contributing.rst @@ -253,8 +253,8 @@ and third block is for another argument. As an example, please refer `DataFrame. These blocks should be consistently separated in PySpark doctests, and more doctests should be added if the coverage of the doctests or the number of examples to show is not enough. -Contributing Error and Exception --------------------------------- +Contributing Errors and Exceptions +---------------------------------- .. currentmodule:: pyspark.errors diff --git a/python/docs/source/development/debugging.rst b/python/docs/source/development/debugging.rst index 4b7e2e288301c..e69055b410a21 100644 --- a/python/docs/source/development/debugging.rst +++ b/python/docs/source/development/debugging.rst @@ -575,7 +575,7 @@ Solution: 1 6 dtype: int64 -**RuntimeError: Result vector from pandas_udf was not the required length** +**PySparkRuntimeError: [RESULT_ROWS_MISMATCH] The number of output rows must match the number of input rows** Exception: @@ -588,7 +588,7 @@ Exception: 22/04/12 13:46:39 ERROR Executor: Exception in task 2.0 in stage 16.0 (TID 88) org.apache.spark.api.python.PythonException: Traceback (most recent call last): ... - RuntimeError: Result vector from pandas_udf was not the required length: expected 1, got 0 + pyspark.errors.exceptions.base.PySparkRuntimeError: [RESULT_ROWS_MISMATCH] The number of output rows (0) must match the number of input rows (1). Solution: diff --git a/python/docs/source/forbid_literalinclude_lines.py b/python/docs/source/forbid_literalinclude_lines.py new file mode 100644 index 0000000000000..a8fa76672cb3d --- /dev/null +++ b/python/docs/source/forbid_literalinclude_lines.py @@ -0,0 +1,62 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +A tiny Sphinx extension that forbids the ``:lines:`` option on ``literalinclude``. + +Pinning a snippet to absolute line numbers rots silently: any edit to the included +source file (adding an import, re-sorting imports, inserting a blank line) shifts the +range, so the docs either render the wrong code or fail the build with a +"non-whitespace stripped by dedent" warning. Prefer selecting the region by name with +``:pyobject:`` or by sentinel comments with ``:start-after:`` / ``:end-before:``, both +of which are robust to line movement. + +This overrides the built-in ``literalinclude`` directive to raise a build error the +moment ``:lines:`` is used. As the docs build treats warnings as errors, the error is +surfaced immediately in CI. +""" + +from __future__ import annotations + +from typing import Any, Dict + +from docutils.nodes import Node +from sphinx.application import Sphinx +from sphinx.directives.code import LiteralInclude + + +class LiteralIncludeNoLines(LiteralInclude): + """``literalinclude`` that rejects the fragile ``:lines:`` option.""" + + def run(self) -> list[Node]: + if "lines" in self.options: + raise self.error( + "literalinclude with ':lines:' is not allowed: it pins absolute line " + "numbers, which silently break when the included file changes. Select " + "the region by name with ':pyobject:', or bracket it with sentinel " + "comments and use ':start-after:'/':end-before:' instead." + ) + return super().run() + + +def setup(app: Sphinx) -> Dict[str, Any]: + # override=True replaces the built-in ``literalinclude`` registration. + app.add_directive("literalinclude", LiteralIncludeNoLines, override=True) + return { + "parallel_read_safe": True, + "parallel_write_safe": True, + } diff --git a/python/docs/source/migration_guide/pyspark_upgrade.rst b/python/docs/source/migration_guide/pyspark_upgrade.rst index b257ee49b7ddc..35bfe5f23c961 100644 --- a/python/docs/source/migration_guide/pyspark_upgrade.rst +++ b/python/docs/source/migration_guide/pyspark_upgrade.rst @@ -19,6 +19,10 @@ Upgrading PySpark ================== +Upgrading from PySpark 4.2 to 4.3 +--------------------------------- +* In Spark 4.3, a ``mapInPandas`` UDF must return an iterator of ``pandas.DataFrame``\s; returning any other iterable such as a ``list`` now raises ``UDF_RETURN_TYPE``, matching the existing ``mapInArrow`` behavior and the declared ``Iterator[...]`` signature. To restore the previous behavior of accepting any iterable for both ``mapInPandas`` and ``mapInArrow``, set ``spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled`` to ``true``. + Upgrading from PySpark 4.1 to 4.2 --------------------------------- * In Spark 4.2, the minimum supported version for PyArrow has been raised from 15.0.0 to 18.0.0 in PySpark. diff --git a/python/docs/source/reference/pyspark.sql/datasource.rst b/python/docs/source/reference/pyspark.sql/datasource.rst index bb52ef26d94f7..65a6544d587c5 100644 --- a/python/docs/source/reference/pyspark.sql/datasource.rst +++ b/python/docs/source/reference/pyspark.sql/datasource.rst @@ -32,6 +32,7 @@ Python Data Source DataSource.writer DataSourceReader.partitions DataSourceReader.pushFilters + DataSourceReader.pushLimit DataSourceReader.read DataSourceRegistration.register DataSourceStreamReader.commit diff --git a/python/docs/source/reference/pyspark.sql/functions.rst b/python/docs/source/reference/pyspark.sql/functions.rst index 91b722d50efd6..94c5a9a87e13a 100644 --- a/python/docs/source/reference/pyspark.sql/functions.rst +++ b/python/docs/source/reference/pyspark.sql/functions.rst @@ -143,6 +143,7 @@ Mathematical Functions sqrt tan tanh + truncate try_add try_divide try_mod @@ -177,6 +178,7 @@ String Functions find_in_set format_number format_string + from_base32 initcap instr is_valid_utf8 @@ -191,6 +193,7 @@ String Functions ltrim make_valid_utf8 mask + normalize octet_length overlay position @@ -216,6 +219,7 @@ String Functions substr substring substring_index + to_base32 to_binary to_char to_number @@ -349,6 +353,8 @@ Hash Functions sha sha1 sha2 + xxh3_128 + xxh3_64 xxhash64 @@ -407,6 +413,7 @@ Array Functions shuffle slice sort_array + trim_array Struct Functions @@ -449,6 +456,7 @@ Aggregate Functions bit_xor bitmap_construct_agg bitmap_or_agg + bitmap_xor_agg bool_and bool_or collect_list @@ -587,6 +595,7 @@ JSON Functions json_array_length json_object_keys json_tuple + json_typeof schema_of_json to_json @@ -605,6 +614,9 @@ VARIANT Functions variant_array_append try_variant_array_append variant_delete + variant_from_arrays + variant_from_entries + variant_strip_nulls variant_get variant_insert try_variant_insert @@ -653,9 +665,13 @@ Misc Functions aes_decrypt aes_encrypt assert_true + bitmap_and + bitmap_andnot bitmap_bit_position bitmap_bucket_number bitmap_count + bitmap_or + bitmap_xor current_catalog current_database current_path @@ -759,9 +775,11 @@ UDF, UDTF and UDT arrow_udtf call_udf pandas_udf + udaf udf udtf unwrap_udt + wrap_udt Table-Valued Functions diff --git a/python/docs/source/tutorial/sql/arrow_pandas.rst b/python/docs/source/tutorial/sql/arrow_pandas.rst index 5e60090ec3db9..2ef306b0460a0 100644 --- a/python/docs/source/tutorial/sql/arrow_pandas.rst +++ b/python/docs/source/tutorial/sql/arrow_pandas.rst @@ -48,7 +48,8 @@ with :meth:`DataFrame.toArrow`. .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 37-52 + :start-after: $example on:dataframe_to_from_arrow_table$ + :end-before: $example off:dataframe_to_from_arrow_table$ :dedent: 4 Note that :meth:`DataFrame.toArrow` results in the collection of all records in the DataFrame to @@ -69,7 +70,8 @@ This can be controlled by ``spark.sql.execution.arrow.pyspark.fallback.enabled`` .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 56-71 + :start-after: $example on:dataframe_with_arrow$ + :end-before: $example off:dataframe_with_arrow$ :dedent: 4 Using the above optimizations with Arrow will produce the same results as when Arrow is not @@ -106,7 +108,8 @@ specify the type hints of ``pandas.Series`` and ``pandas.DataFrame`` as below: .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 75-99 + :start-after: $example on:ser_to_frame_pandas_udf$ + :end-before: $example off:ser_to_frame_pandas_udf$ :dedent: 4 In the following sections, it describes the combinations of the supported type hints. For simplicity, @@ -129,7 +132,8 @@ The following example shows how to create this Pandas UDF that computes the prod .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 103-133 + :start-after: $example on:ser_to_ser_pandas_udf$ + :end-before: $example off:ser_to_ser_pandas_udf$ :dedent: 4 For detailed usage, please see :func:`pandas_udf`. @@ -168,7 +172,8 @@ The following example shows how to create this Pandas UDF: .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 137-159 + :start-after: $example on:iter_ser_to_iter_ser_pandas_udf$ + :end-before: $example off:iter_ser_to_iter_ser_pandas_udf$ :dedent: 4 For detailed usage, please see :func:`pandas_udf`. @@ -190,7 +195,8 @@ The following example shows how to create this Pandas UDF: .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 163-186 + :start-after: $example on:iter_sers_to_iter_ser_pandas_udf$ + :end-before: $example off:iter_sers_to_iter_ser_pandas_udf$ :dedent: 4 For detailed usage, please see :func:`pandas_udf`. @@ -221,7 +227,8 @@ and window operations: .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 190-231 + :start-after: $example on:ser_to_scalar_pandas_udf$ + :end-before: $example off:ser_to_scalar_pandas_udf$ :dedent: 4 .. currentmodule:: pyspark.sql.functions @@ -286,7 +293,8 @@ in the group. .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 235-253 + :start-after: $example on:grouped_apply_in_pandas$ + :end-before: $example off:grouped_apply_in_pandas$ :dedent: 4 For detailed usage, please see please see :meth:`GroupedData.applyInPandas` @@ -304,7 +312,8 @@ The following example shows how to use :meth:`DataFrame.mapInPandas`: .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 257-268 + :start-after: $example on:map_in_pandas$ + :end-before: $example off:map_in_pandas$ :dedent: 4 For detailed usage, please see :meth:`DataFrame.mapInPandas`. @@ -343,7 +352,8 @@ The following example shows how to use ``DataFrame.groupby().cogroup().applyInPa .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 272-294 + :start-after: $example on:cogrouped_apply_in_pandas$ + :end-before: $example off:cogrouped_apply_in_pandas$ :dedent: 4 @@ -365,7 +375,8 @@ Here's an example that demonstrates the usage of both a default, pickled Python .. literalinclude:: ../../../../../examples/src/main/python/sql/arrow.py :language: python - :lines: 298-316 + :start-after: $example on:arrow_python_udf$ + :end-before: $example off:arrow_python_udf$ :dedent: 4 Type coercion: diff --git a/python/docs/source/tutorial/sql/arrow_python_udf.rst b/python/docs/source/tutorial/sql/arrow_python_udf.rst index d1430844a7053..0ed9212edf110 100644 --- a/python/docs/source/tutorial/sql/arrow_python_udf.rst +++ b/python/docs/source/tutorial/sql/arrow_python_udf.rst @@ -479,6 +479,19 @@ SQL boolean expressions do not short-circuit: in ``WHERE cond AND udf(x)``, the called on all rows regardless of ``cond``. If the function can fail on certain input values (e.g., division by zero), handle those cases inside the function itself. +The same holds for a Python UDF inside a higher-order function's lambda -- a plain, scalar pandas +(``pandas_udf``), scalar Arrow (``arrow_udf``) or iterator UDF. It is not evaluated element by +element inside the lambda; it is applied once, outside the lambda, and the lambda reads the +precomputed result. So the UDF runs over *every* element, even ones a lambda would otherwise skip +(the elements after ``exists`` matches, or the untaken branch of a ``when``) -- if it can fail on +some input, handle that inside the function. This covers ``transform``, ``filter``, ``exists``, +``forall``, ``zip_with``, ``array_sort`` and the map functions, as well as *nested* lambdas +(``transform(matrix, lambda row: transform(row, lambda x: udf(x)))``, including ones that capture +the enclosing variable). A UDF inside ``aggregate`` / ``reduce`` is not supported, because the fold +is sequential and cannot be applied once to the whole array. Toggle with +``spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled`` (default ``true``; set ``false`` to +reject at analysis). + The Arrow data type of the returned ``pyarrow.Array`` should match the declared ``returnType``. When there is a mismatch, Spark will attempt to convert the returned data to the expected type using Arrow's safe casting, which raises an error on overflow or precision loss. diff --git a/python/docs/source/tutorial/sql/python_data_source.rst b/python/docs/source/tutorial/sql/python_data_source.rst index 6e9b3d9bd63cb..63a41c65b1eb4 100644 --- a/python/docs/source/tutorial/sql/python_data_source.rst +++ b/python/docs/source/tutorial/sql/python_data_source.rst @@ -179,6 +179,50 @@ Define the reader logic to generate synthetic data. Use the `faker` library to p row.append(value) yield tuple(row) +Push Down a Limit to a Batch Reader +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When a query only needs the first few rows, a reader can implement ``pushLimit`` to fetch less +data, for example by adding a page size parameter to a REST request or a ``LIMIT`` clause to a +SQL query. ``pushLimit`` is called during planning, before ``partitions`` and ``read``, and +returns whether the reader will make use of the limit. It runs after ``pushFilters`` when the +query has filters to push down, so state it depends on belongs in ``__init__``. + +Pushing down a limit is only a hint: Spark always applies the limit again after the scan, so a +reader is free to return more rows than requested. Set +``spark.sql.python.limitPushdown.enabled`` to ``true`` to enable limit pushdown. + +.. code-block:: python + + from typing import Dict + + from pyspark.sql.datasource import DataSourceReader, InputPartition + from pyspark.sql.types import StructType + + class FakeDataSourceReader(DataSourceReader): + + def __init__(self, schema: StructType, options: Dict[str, str]): + self.schema: StructType = schema + self.options = options + self.limit = None + + def pushLimit(self, limit: int) -> bool: + self.limit = limit + return True + + def partitions(self): + # A limit makes a single request cheaper than a fan-out, since every + # partition opens its own connection to the data source. + if self.limit is not None: + return [InputPartition(None)] + return [InputPartition(i) for i in range(16)] + + def read(self, partition): + num_rows = int(self.options.get("numRows", 3)) + if self.limit is not None: + num_rows = min(num_rows, self.limit) + ... + Implement a Batch Writer ~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/python/docs/source/tutorial/sql/python_udtf.rst b/python/docs/source/tutorial/sql/python_udtf.rst index e2d240a6f9de1..8afc8dc007cdc 100644 --- a/python/docs/source/tutorial/sql/python_udtf.rst +++ b/python/docs/source/tutorial/sql/python_udtf.rst @@ -421,7 +421,8 @@ Python UDTFs can be registered and used in SQL queries. .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 82-116 + :start-after: $example on:python_udtf_registration$ + :end-before: $example off:python_udtf_registration$ :dedent: 4 @@ -440,7 +441,8 @@ when declaring the UDTF. .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 121-126 + :start-after: $example on:python_udtf_arrow$ + :end-before: $example off:python_udtf_arrow$ :dedent: 4 @@ -454,7 +456,8 @@ Here is a simple example of a UDTF class implementation: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 36-40 + :start-after: $example on:python_udtf_simple_class$ + :end-before: $example off:python_udtf_simple_class$ :dedent: 4 @@ -462,7 +465,8 @@ To make use of the UDTF, you'll first need to instantiate it using the ``@udtf`` .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 42-55 + :start-after: $example on:python_udtf_simple_udtf$ + :end-before: $example off:python_udtf_simple_udtf$ :dedent: 4 @@ -470,21 +474,24 @@ An alternative way to create a UDTF is to use the :func:`udtf` function: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 60-77 + :start-after: $example on:python_udtf_decorator$ + :end-before: $example off:python_udtf_decorator$ :dedent: 4 Here is a Python UDTF that expands date ranges into individual dates: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 131-152 + :start-after: $example on:python_udtf_date_expander$ + :end-before: $example off:python_udtf_date_expander$ :dedent: 4 Here is a Python UDTF with ``__init__`` and ``terminate``: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 157-186 + :start-after: $example on:python_udtf_terminate$ + :end-before: $example off:python_udtf_terminate$ :dedent: 4 @@ -510,7 +517,8 @@ For example: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 191-210 + :start-after: $example on:python_udtf_table_argument$ + :end-before: $example off:python_udtf_table_argument$ :dedent: 4 When calling a UDTF with a table argument, any SQL query can request that the input table be @@ -535,7 +543,8 @@ For example: .. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py :language: python - :lines: 215-287 + :start-after: $example on:python_udtf_table_argument_with_partitioning$ + :end-before: $example off:python_udtf_table_argument_with_partitioning$ :dedent: 4 Note that in for each of these ways of partitioning the input table when calling UDTFs in SQL diff --git a/python/packaging/classic/setup.py b/python/packaging/classic/setup.py index 97ec9dd90ddb6..add11863f000a 100755 --- a/python/packaging/classic/setup.py +++ b/python/packaging/classic/setup.py @@ -16,15 +16,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -import importlib.util +import ctypes import glob +import importlib.util import os import sys -import ctypes +from dataclasses import dataclass +from pathlib import Path +from shutil import copyfile, copytree, rmtree + from setuptools import setup from setuptools.command.install import install -from shutil import copyfile, copytree, rmtree -from pathlib import Path if ( # When we package, the parent directory 'classic' dir @@ -100,10 +102,6 @@ print(incorrect_invocation_message, file=sys.stderr) sys.exit(-1) -EXAMPLES_PATH = os.path.join(SPARK_HOME, "examples/src/main/python") -SCRIPTS_PATH = os.path.join(SPARK_HOME, "bin") -USER_SCRIPTS_PATH = os.path.join(SPARK_HOME, "sbin") -DATA_PATH = os.path.join(SPARK_HOME, "data") # The classic PySpark package bundles the assembly jars, so it ships the binary # license texts (licenses-binary), which enumerate those jars' licenses, mirroring # the binary distribution. The connect/client packages bundle no jars. @@ -113,12 +111,26 @@ # were already copied to licenses/ (see dev/make-distribution.sh). LICENSES_PATH = os.path.join(SPARK_HOME, "licenses") -SCRIPTS_TARGET = os.path.join(TEMP_PATH, "bin") -USER_SCRIPTS_TARGET = os.path.join(TEMP_PATH, "sbin") + +@dataclass(frozen=True) +class InstallPath: + source: str + target: str + + JARS_TARGET = os.path.join(TEMP_PATH, "jars") -EXAMPLES_TARGET = os.path.join(TEMP_PATH, "examples") -DATA_TARGET = os.path.join(TEMP_PATH, "data") -LICENSES_TARGET = os.path.join(TEMP_PATH, "licenses") +SCRIPTS_TARGET = os.path.join(TEMP_PATH, "bin") +INSTALL_PATHS = [ + InstallPath(JARS_PATH, JARS_TARGET), + InstallPath(os.path.join(SPARK_HOME, "bin"), SCRIPTS_TARGET), + InstallPath(os.path.join(SPARK_HOME, "sbin"), os.path.join(TEMP_PATH, "sbin")), + InstallPath( + os.path.join(SPARK_HOME, "examples/src/main/python"), + os.path.join(TEMP_PATH, "examples"), + ), + InstallPath(os.path.join(SPARK_HOME, "data"), os.path.join(TEMP_PATH, "data")), + InstallPath(LICENSES_PATH, os.path.join(TEMP_PATH, "licenses")), +] # Check and see if we are under the spark path in which case we need to build the symlink farm. # This is important because we only want to build the symlink farm while under Spark otherwise we @@ -206,39 +218,27 @@ def run(self): # We copy the shell script to be under pyspark/python/pyspark so that the launcher scripts # find it where expected. The rest of the files aren't copied because they are accessed # using Python imports instead which will be resolved correctly. - try: - os.makedirs("pyspark/python/pyspark") - except OSError: - # Don't worry if the directory already exists. - pass + os.makedirs("pyspark/python/pyspark", exist_ok=True) copyfile("pyspark/shell.py", "pyspark/python/pyspark/shell.py") if in_spark: - # !!HACK ALTERT!! + # !!HACK ALERT!! # `setup.py` has to be located with the same directory with the package. # Therefore, we copy the current file, and place it at `spark/python` directory. # After that, we remove it in the end. copyfile("packaging/classic/setup.py", "setup.py") copyfile("packaging/classic/setup.cfg", "setup.cfg") - # Construct the symlink farm - this is nein_sparkcessary since we can't refer to + # Construct the symlink farm - this is necessary since we can't refer to # the path above the package root and we need to copy the jars and scripts which # are up above the python root. if _supports_symlinks(): - os.symlink(JARS_PATH, JARS_TARGET) - os.symlink(SCRIPTS_PATH, SCRIPTS_TARGET) - os.symlink(USER_SCRIPTS_PATH, USER_SCRIPTS_TARGET) - os.symlink(EXAMPLES_PATH, EXAMPLES_TARGET) - os.symlink(DATA_PATH, DATA_TARGET) - os.symlink(LICENSES_PATH, LICENSES_TARGET) + for path in INSTALL_PATHS: + os.symlink(path.source, path.target) else: # For windows fall back to the slower copytree - copytree(JARS_PATH, JARS_TARGET) - copytree(SCRIPTS_PATH, SCRIPTS_TARGET) - copytree(USER_SCRIPTS_PATH, USER_SCRIPTS_TARGET) - copytree(EXAMPLES_PATH, EXAMPLES_TARGET) - copytree(DATA_PATH, DATA_TARGET) - copytree(LICENSES_PATH, LICENSES_TARGET) + for path in INSTALL_PATHS: + copytree(path.source, path.target) else: # If we are not inside of SPARK_HOME verify we have the required symlink farm if not os.path.exists(JARS_TARGET): @@ -340,21 +340,25 @@ def run(self): "pyspark.examples.src.main.python": "deps/examples", }, package_data={ - "pyspark.jars": ["*.jar"], - "pyspark.bin": ["*"], + "pyspark": ["**/*.pyi", "**/py.typed", "**/*.json"], + "pyspark.jars": ["**/*.jar"], + "pyspark.bin": ["**/*"], "pyspark.sbin": [ "spark-config.sh", "spark-daemon.sh", + "start-connect-server.sh", "start-history-server.sh", + "stop-connect-server.sh", "stop-history-server.sh", ], - "pyspark.python.lib": ["*.zip"], - "pyspark.data": ["*.txt", "*.data"], - "pyspark.licenses": ["*"], - "pyspark.examples.src.main.python": ["*.py", "*/*.py"], + "pyspark.python.lib": ["**/*.zip"], + "pyspark.data": ["**/*.txt", "**/*.data"], + "pyspark.licenses": ["**/*"], + "pyspark.examples.src.main.python": ["**/*.py"], }, scripts=scripts, license="Apache-2.0", + license_files=["LICENSE", "NOTICE"], # Don't forget to update python/docs/source/getting_started/install.rst # if you're updating the versions or dependencies. install_requires=["py4j>=0.10.9.7,<0.10.9.10"], @@ -411,19 +415,10 @@ def run(self): if in_spark: os.remove("setup.py") os.remove("setup.cfg") - # Depending on cleaning up the symlink farm or copied version - if _supports_symlinks(): - os.remove(os.path.join(TEMP_PATH, "jars")) - os.remove(os.path.join(TEMP_PATH, "bin")) - os.remove(os.path.join(TEMP_PATH, "sbin")) - os.remove(os.path.join(TEMP_PATH, "examples")) - os.remove(os.path.join(TEMP_PATH, "data")) - os.remove(os.path.join(TEMP_PATH, "licenses")) - else: - rmtree(os.path.join(TEMP_PATH, "jars")) - rmtree(os.path.join(TEMP_PATH, "bin")) - rmtree(os.path.join(TEMP_PATH, "sbin")) - rmtree(os.path.join(TEMP_PATH, "examples")) - rmtree(os.path.join(TEMP_PATH, "data")) - rmtree(os.path.join(TEMP_PATH, "licenses")) + for path in INSTALL_PATHS: + if os.path.islink(path.target): + # Remove the link and not the real source trees under SPARK_HOME. + os.remove(path.target) + else: + rmtree(path.target) os.rmdir(TEMP_PATH) diff --git a/python/packaging/client/setup.py b/python/packaging/client/setup.py index 6c97164937eea..af9e9475018c8 100755 --- a/python/packaging/client/setup.py +++ b/python/packaging/client/setup.py @@ -22,12 +22,13 @@ # cd python/packaging/classic # python setup.py sdist -import sys -from setuptools import setup -import os -from shutil import copyfile, move import glob +import os +import sys from pathlib import Path +from shutil import copyfile + +from setuptools import setup if ( # When we package, the parent directory 'client' dir @@ -118,13 +119,10 @@ try: if in_spark: - # !!HACK ALTERT!! - # 1. `setup.py` has to be located with the same directory with the package. - # Therefore, we copy the current file, and place it at `spark/python` directory. - # After that, we remove it in the end. - # 2. Here it renames `lib` to `lib.back` so MANIFEST.in does not pick `py4j` up. - # We rename it back in the end. - move("lib", "lib.back") + # !!HACK ALERT!! + # `setup.py` has to be located with the same directory with the package. + # Therefore, we copy the current file, and place it at `spark/python` directory. + # After that, we remove it in the end. copyfile("packaging/client/setup.py", "setup.py") copyfile("packaging/client/setup.cfg", "setup.cfg") @@ -206,7 +204,11 @@ url="https://github.com/apache/spark/tree/master/python", packages=connect_packages + test_packages, include_package_data=True, + package_data={ + "pyspark": ["**/*.pyi", "**/py.typed", "**/*.json"], + }, license="Apache-2.0", + license_files=["LICENSE", "NOTICE"], # Don't forget to update python/docs/source/getting_started/install.rst # if you're updating the versions or dependencies. install_requires=[ @@ -232,6 +234,5 @@ ) finally: if in_spark: - move("lib.back", "lib") os.remove("setup.py") os.remove("setup.cfg") diff --git a/python/packaging/connect/setup.py b/python/packaging/connect/setup.py index a4e83e41d1518..d24d0436891b4 100755 --- a/python/packaging/connect/setup.py +++ b/python/packaging/connect/setup.py @@ -22,12 +22,13 @@ # cd python/packaging/connect # python setup.py sdist -import sys -from setuptools import setup -import os -from shutil import copyfile, copytree, move, rmtree import glob +import os +import sys from pathlib import Path +from shutil import copyfile, copytree, rmtree + +from setuptools import setup if ( # When we package, the parent directory 'connect' dir @@ -57,18 +58,14 @@ try: if in_spark: - # !!HACK ALTERT!! - # 1. `setup.py` has to be located with the same directory with the package. - # Therefore, we copy the current file, and place it at `spark/python` directory. - # After that, we remove it in the end. - # 2. Here it renames `pyspark` and `lib` to `pyspark.back` and `lib.back` so MANIFEST.in - # does not pick `pyspark` and `py4j` up. We rename it back in the end. - move("pyspark", "pyspark.back") - move("lib", "lib.back") + # !!HACK ALERT!! + # `setup.py` has to be located with the same directory with the package. + # Therefore, we copy the current file, and place it at `spark/python` directory. + # After that, we remove it in the end. copyfile("packaging/connect/setup.py", "setup.py") copyfile("packaging/connect/setup.cfg", "setup.cfg") copytree("packaging/connect/pyspark_connect", "pyspark_connect") - copyfile("pyspark.back/version.py", "pyspark_connect/version.py") + copyfile("pyspark/version.py", "pyspark_connect/version.py") try: exec(open("pyspark_connect/version.py").read()) @@ -114,6 +111,7 @@ packages=connect_packages, include_package_data=True, license="Apache-2.0", + license_files=["LICENSE", "NOTICE"], # Don't forget to update python/docs/source/getting_started/install.rst # if you're updating the versions or dependencies. install_requires=[ @@ -140,8 +138,6 @@ ) finally: if in_spark: - move("pyspark.back", "pyspark") - move("lib.back", "lib") os.remove("setup.py") os.remove("setup.cfg") rmtree("pyspark_connect") diff --git a/python/pyspark/__init__.py b/python/pyspark/__init__.py index f2f8b08b9da7f..2dd58e06ec650 100644 --- a/python/pyspark/__init__.py +++ b/python/pyspark/__init__.py @@ -48,16 +48,16 @@ import sys from functools import wraps -from typing import cast, Any, Callable, TypeVar, Union +from typing import Any, Callable, TypeVar, Union, cast from pyspark.util import is_remote_only if not is_remote_only(): - from pyspark.core.rdd import RDD, RDDBarrier - from pyspark.core.files import SparkFiles - from pyspark.core.status import StatusTracker, SparkJobInfo, SparkStageInfo, SparkExecutorInfo + from pyspark.core import broadcast, files, rdd, status from pyspark.core.broadcast import Broadcast - from pyspark.core import rdd, files, status, broadcast + from pyspark.core.files import SparkFiles + from pyspark.core.rdd import RDD, RDDBarrier + from pyspark.core.status import SparkExecutorInfo, SparkJobInfo, SparkStageInfo, StatusTracker # for backward compatibility references. sys.modules["pyspark.rdd"] = rdd @@ -65,15 +65,15 @@ sys.modules["pyspark.status"] = status sys.modules["pyspark.broadcast"] = broadcast +from pyspark._globals import _NoValue # noqa: F401 +from pyspark.accumulators import Accumulator, AccumulatorParam from pyspark.conf import SparkConf -from pyspark.util import InheritableThread, inheritable_thread_target +from pyspark.profiler import BasicProfiler, Profiler +from pyspark.serializers import CPickleSerializer, MarshalSerializer from pyspark.storagelevel import StorageLevel -from pyspark.accumulators import Accumulator, AccumulatorParam -from pyspark.serializers import MarshalSerializer, CPickleSerializer -from pyspark.taskcontext import TaskContext, BarrierTaskContext, BarrierTaskInfo -from pyspark.profiler import Profiler, BasicProfiler +from pyspark.taskcontext import BarrierTaskContext, BarrierTaskInfo, TaskContext +from pyspark.util import InheritableThread, inheritable_thread_target from pyspark.version import __version__ -from pyspark._globals import _NoValue # noqa: F401 _F = TypeVar("_F", bound=Callable) @@ -119,14 +119,14 @@ def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: # To avoid circular dependencies if not is_remote_only(): - from pyspark.core.context import SparkContext from pyspark.core import context + from pyspark.core.context import SparkContext # for backward compatibility references. sys.modules["pyspark.context"] = context # for back compatibility - from pyspark.sql import SQLContext, HiveContext # noqa: F401 + from pyspark.sql import HiveContext, SQLContext # noqa: F401 from pyspark.sql import Row # noqa: F401 diff --git a/python/pyspark/_typing.pyi b/python/pyspark/_typing.pyi index fd4e1e2ad0195..a4002788f51bf 100644 --- a/python/pyspark/_typing.pyi +++ b/python/pyspark/_typing.pyi @@ -17,9 +17,9 @@ # under the License. from typing import Any, Callable, Iterable, Sized, TypeVar, Union -from typing_extensions import Literal, Protocol -from numpy import int32, int64, float32, float64, ndarray +from numpy import float32, float64, int32, int64, ndarray +from typing_extensions import Literal, Protocol F = TypeVar("F", bound=Callable) T_co = TypeVar("T_co", covariant=True) diff --git a/python/pyspark/accumulators.py b/python/pyspark/accumulators.py index fcfa347092ee2..e09e987ff4175 100644 --- a/python/pyspark/accumulators.py +++ b/python/pyspark/accumulators.py @@ -15,22 +15,23 @@ # limitations under the License. # -import os -import sys import hmac +import os import select -import struct import socketserver +import struct +import sys import threading -from typing import Callable, Dict, Generic, Tuple, Type, TYPE_CHECKING, TypeVar, Union, Optional +from typing import TYPE_CHECKING, Callable, Dict, Generic, Optional, Tuple, Type, TypeVar, Union -from pyspark.serializers import read_int, CPickleSerializer from pyspark.errors import PySparkRuntimeError +from pyspark.serializers import CPickleSerializer, read_int if TYPE_CHECKING: - from pyspark._typing import SupportsIAdd from socketserver import BaseRequestHandler + from pyspark._typing import SupportsIAdd + __all__ = ["Accumulator", "AccumulatorParam"] diff --git a/python/pyspark/conf.py b/python/pyspark/conf.py index 2f109b007b9c1..fef0c6fd4c456 100644 --- a/python/pyspark/conf.py +++ b/python/pyspark/conf.py @@ -18,13 +18,13 @@ __all__ = ["SparkConf"] import sys -from typing import Dict, List, Optional, Tuple, cast, overload, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, cast, overload -from pyspark.util import is_remote_only from pyspark.errors import PySparkRuntimeError +from pyspark.util import is_remote_only if TYPE_CHECKING: - from py4j.java_gateway import JVMView, JavaObject + from py4j.java_gateway import JavaObject, JVMView class SparkConf: diff --git a/python/pyspark/core/broadcast.py b/python/pyspark/core/broadcast.py index 9b0002bdccfe6..64cf0ac937002 100644 --- a/python/pyspark/core/broadcast.py +++ b/python/pyspark/core/broadcast.py @@ -17,29 +17,29 @@ import gc import os +import pickle import sys -from tempfile import NamedTemporaryFile import threading -import pickle +from tempfile import NamedTemporaryFile from typing import ( - overload, + IO, + TYPE_CHECKING, Any, BinaryIO, Callable, Dict, Generic, - IO, Iterator, Optional, Tuple, TypeVar, - TYPE_CHECKING, Union, + overload, ) -from pyspark.serializers import ChunkedStream, pickle_protocol -from pyspark.util import print_exec, local_connect_and_auth from pyspark.errors import PySparkRuntimeError +from pyspark.serializers import ChunkedStream, pickle_protocol +from pyspark.util import local_connect_and_auth, print_exec if TYPE_CHECKING: from pyspark import SparkContext @@ -361,8 +361,9 @@ def clear(self) -> None: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.core.broadcast + from pyspark.sql import SparkSession globs = pyspark.core.broadcast.__dict__.copy() spark = SparkSession.builder.master("local[4]").appName("broadcast tests").getOrCreate() diff --git a/python/pyspark/core/context.py b/python/pyspark/core/context.py index 2a6d8ab68e5f1..10e4fe14ab9a1 100644 --- a/python/pyspark/core/context.py +++ b/python/pyspark/core/context.py @@ -15,21 +15,21 @@ # limitations under the License. # -import uuid +import importlib import os import shutil import signal import sys import threading +import uuid import warnings -import importlib -from threading import RLock from tempfile import NamedTemporaryFile +from threading import RLock from types import TracebackType from typing import ( + TYPE_CHECKING, Any, Callable, - cast, ClassVar, Dict, Iterable, @@ -37,42 +37,42 @@ NoReturn, Optional, Sequence, + Set, Tuple, Type, - TYPE_CHECKING, TypeVar, - Set, + cast, ) from py4j.java_collections import JavaMap +from py4j.java_gateway import JavaGateway, JavaObject, JVMView, is_instance_of from py4j.protocol import Py4JError from pyspark import accumulators -from pyspark.conf import SparkConf from pyspark.accumulators import Accumulator +from pyspark.conf import SparkConf from pyspark.core.broadcast import Broadcast, BroadcastPickleRegistry from pyspark.core.files import SparkFiles +from pyspark.core.rdd import RDD +from pyspark.core.status import StatusTracker +from pyspark.errors import PySparkRuntimeError from pyspark.java_gateway import launch_gateway +from pyspark.profiler import BasicProfiler, MemoryProfiler, ProfilerCollector, UDFBasicProfiler +from pyspark.resource.information import ResourceInformation from pyspark.serializers import ( - CPickleSerializer, + AutoBatchedSerializer, BatchedSerializer, + ChunkedStream, + CPickleSerializer, + NoOpSerializer, + PairDeserializer, Serializer, UTF8Deserializer, - PairDeserializer, - AutoBatchedSerializer, - NoOpSerializer, - ChunkedStream, ) from pyspark.storagelevel import StorageLevel -from pyspark.resource.information import ResourceInformation -from pyspark.core.rdd import RDD -from pyspark.util import _load_from_socket, local_connect_and_auth from pyspark.taskcontext import TaskContext from pyspark.traceback_utils import CallSite, first_spark_call -from pyspark.core.status import StatusTracker -from pyspark.profiler import ProfilerCollector, BasicProfiler, UDFBasicProfiler, MemoryProfiler -from pyspark.errors import PySparkRuntimeError -from py4j.java_gateway import is_instance_of, JavaGateway, JavaObject, JVMView +from pyspark.util import _load_from_socket, local_connect_and_auth if TYPE_CHECKING: from pyspark.accumulators import AccumulatorParam @@ -2491,19 +2491,31 @@ def cancelJobsWithTag(self, tag: str) -> None: """ return self._jsc.cancelJobsWithTag(tag) - def cancelAllJobs(self) -> None: + def cancelAllJobs(self, reason: Optional[str] = None) -> None: """ Cancel all jobs that have been scheduled or are running. .. versionadded:: 1.1.0 + Parameters + ---------- + reason : str, optional + Reason for cancellation. It is surfaced in the error of every cancelled job, so that + a job aborted as collateral of a context-wide cancellation can be told apart from one + that failed on its own. + + .. versionadded:: 4.4.0 + See Also -------- :meth:`SparkContext.cancelJobGroup` :meth:`SparkContext.cancelJobsWithTag` :meth:`SparkContext.runJob` """ - self._jsc.sc().cancelAllJobs() + if reason is None: + self._jsc.sc().cancelAllJobs() + else: + self._jsc.sc().cancelAllJobs(reason) def statusTracker(self) -> StatusTracker: """ diff --git a/python/pyspark/core/files.py b/python/pyspark/core/files.py index aec67a0fe31c2..a3c59e7681500 100644 --- a/python/pyspark/core/files.py +++ b/python/pyspark/core/files.py @@ -19,7 +19,7 @@ __all__ = ["SparkFiles"] -from typing import cast, ClassVar, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, ClassVar, Optional, cast if TYPE_CHECKING: from pyspark import SparkContext @@ -149,6 +149,7 @@ def getRootDirectory(cls) -> str: def _test() -> None: import doctest import sys + from pyspark import SparkContext globs = globals().copy() diff --git a/python/pyspark/core/rdd.py b/python/pyspark/core/rdd.py index 190857bf3cecc..629a1bc93e6b1 100644 --- a/python/pyspark/core/rdd.py +++ b/python/pyspark/core/rdd.py @@ -15,22 +15,24 @@ # limitations under the License. # +import bisect import copy -import sys -import os +import heapq import operator +import os +import random import shlex +import sys import warnings -import heapq -import bisect -import random -from subprocess import Popen, PIPE -from threading import Thread from collections import defaultdict -from itertools import chain from functools import reduce -from math import sqrt, log, isinf, isnan, pow, ceil +from itertools import chain +from math import ceil, isinf, isnan, log, pow, sqrt +from subprocess import PIPE, Popen +from threading import Thread from typing import ( + IO, + TYPE_CHECKING, Any, Callable, Dict, @@ -38,73 +40,71 @@ Hashable, Iterable, Iterator, - IO, List, NoReturn, Optional, Sequence, Tuple, - Union, TypeVar, + Union, cast, overload, - TYPE_CHECKING, ) +from pyspark.errors import PySparkRuntimeError +from pyspark.join import ( + python_cogroup, + python_full_outer_join, + python_join, + python_left_outer_join, + python_right_outer_join, +) +from pyspark.rddsampler import RDDRangeSampler, RDDSampler, RDDStratifiedSampler +from pyspark.resource.profile import ResourceProfile +from pyspark.resource.requests import ExecutorResourceRequests, TaskResourceRequests +from pyspark.resultiterable import ResultIterable from pyspark.serializers import ( AutoBatchedSerializer, BatchedSerializer, - NoOpSerializer, CartesianDeserializer, CloudPickleSerializer, - PairDeserializer, CPickleSerializer, + NoOpSerializer, + PairDeserializer, Serializer, pack_long, ) -from pyspark.join import ( - python_join, - python_left_outer_join, - python_right_outer_join, - python_full_outer_join, - python_cogroup, -) -from pyspark.statcounter import StatCounter -from pyspark.rddsampler import RDDSampler, RDDRangeSampler, RDDStratifiedSampler -from pyspark.storagelevel import StorageLevel -from pyspark.resource.requests import ExecutorResourceRequests, TaskResourceRequests -from pyspark.resource.profile import ResourceProfile -from pyspark.resultiterable import ResultIterable from pyspark.shuffle import ( Aggregator, + ExternalGroupBy, ExternalMerger, - get_used_memory, ExternalSorter, - ExternalGroupBy, + get_used_memory, ) +from pyspark.statcounter import StatCounter +from pyspark.storagelevel import StorageLevel from pyspark.traceback_utils import SCCallSiteSync + +# for backward compatibility references. from pyspark.util import ( - fail_on_stopiteration, - _parse_memory, + PythonEvalType, # noqa: F401 _load_from_socket, _local_iterator_from_socket, + _parse_memory, + fail_on_stopiteration, ) -from pyspark.errors import PySparkRuntimeError - -# for backward compatibility references. -from pyspark.util import PythonEvalType # noqa: F401 if TYPE_CHECKING: from py4j.java_gateway import JavaObject - from pyspark._typing import S, SizedIterable, NumberOrArray + from pyspark._typing import NumberOrArray, S, SizedIterable from pyspark.core.context import SparkContext - from pyspark.sql.dataframe import DataFrame - from pyspark.sql.types import AtomicType, StructType from pyspark.sql._typing import ( AtomicValue, RowLike, ) + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.types import AtomicType, StructType T = TypeVar("T") T_co = TypeVar("T_co", covariant=True) @@ -5356,6 +5356,7 @@ def _is_barrier(self) -> bool: def _test() -> None: import doctest import tempfile + from pyspark.core.context import SparkContext try: diff --git a/python/pyspark/daemon.py b/python/pyspark/daemon.py index 0e321d442e816..ab761d911d144 100644 --- a/python/pyspark/daemon.py +++ b/python/pyspark/daemon.py @@ -14,25 +14,25 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import uuid +import faulthandler +import gc import os -import signal import select +import signal import socket import sys -import traceback import time -import gc -import faulthandler -from errno import EINTR, EAGAIN +import traceback +import uuid +from errno import EAGAIN, EINTR +from signal import SIG_DFL, SIG_IGN, SIGCHLD, SIGHUP, SIGINT, SIGTERM from socket import AF_INET, AF_INET6, SOCK_STREAM, SOMAXCONN -from signal import SIGHUP, SIGTERM, SIGCHLD, SIG_DFL, SIG_IGN, SIGINT from types import FrameType from typing import Any, Optional -from pyspark.serializers import read_int, write_int, write_with_length, UTF8Deserializer -from pyspark.util import enable_faulthandler from pyspark.errors import PySparkRuntimeError +from pyspark.serializers import UTF8Deserializer, read_int, write_int, write_with_length +from pyspark.util import enable_faulthandler def compute_real_exit_code(exit_code: Any) -> int: diff --git a/python/pyspark/errors/__init__.py b/python/pyspark/errors/__init__.py index cde12949c9ca6..22f75a6d1b27c 100644 --- a/python/pyspark/errors/__init__.py +++ b/python/pyspark/errors/__init__.py @@ -20,38 +20,38 @@ """ from pyspark.errors.exceptions.base import ( - PySparkException, AnalysisException, - SessionNotSameException, - TempTableAlreadyExistsException, - ParseException, - IllegalArgumentException, ArithmeticException, - UnsupportedOperationException, ArrayIndexOutOfBoundsException, DateTimeException, + IllegalArgumentException, NumberFormatException, - StreamingQueryException, - QueryExecutionException, - PythonException, - UnknownException, - SparkRuntimeException, - SparkUpgradeException, - SparkNoSuchElementException, - PySparkTypeError, - PySparkValueError, + ParseException, + PickleException, + PySparkAssertionError, + PySparkAttributeError, + PySparkException, PySparkImportError, PySparkIndexError, - PySparkAttributeError, - PySparkRuntimeError, - PySparkAssertionError, + PySparkKeyError, PySparkNotImplementedError, PySparkPicklingError, - PySparkKeyError, + PySparkRuntimeError, + PySparkTypeError, + PySparkValueError, + PythonException, QueryContext, QueryContextType, + QueryExecutionException, + SessionNotSameException, + SparkNoSuchElementException, + SparkRuntimeException, + SparkUpgradeException, StreamingPythonRunnerInitializationException, - PickleException, + StreamingQueryException, + TempTableAlreadyExistsException, + UnknownException, + UnsupportedOperationException, ) __all__ = [ diff --git a/python/pyspark/errors/error-conditions.json b/python/pyspark/errors/error-conditions.json index b8bbaf097279d..ae35e98a5e5c0 100644 --- a/python/pyspark/errors/error-conditions.json +++ b/python/pyspark/errors/error-conditions.json @@ -72,7 +72,7 @@ }, "CANNOT_CONVERT_COLUMN_INTO_BOOL": { "message": [ - "Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions." + "Cannot convert column into bool (offending column: <column>): please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions." ] }, "CANNOT_CONVERT_TYPE": { @@ -204,7 +204,7 @@ }, "DATA_SOURCE_PUSHDOWN_DISABLED": { "message": [ - "<type> implements pushFilters() but filter pushdown is disabled because configuration '<conf>' is false. Set it to true to enable filter pushdown." + "<type> implements <method>() but the corresponding pushdown is disabled because configuration '<conf>' is false. Set it to true to enable it." ] }, "DATA_SOURCE_RETURN_SCHEMA_MISMATCH": { @@ -546,6 +546,16 @@ "<arg1> and <arg2> should be of the same length, got <arg1_length> and <arg2_length>." ] }, + "LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE": { + "message": [ + "Cannot claim the per-user runtime directory <path> (was it created by another user?); remove it or point SPARK_LOCAL_CONNECT_DISCOVERY at a path you own." + ] + }, + "LOCAL_CONNECT_SERVER_START_FAILED": { + "message": [ + "Failed to start a persistent local Spark Connect server: <reason>." + ] + }, "LOCAL_RELATION_SIZE_LIMIT_EXCEEDED": { "message": [ "Local relation size (<actualSize> bytes) exceeds the limit (<sizeLimit> bytes)." @@ -851,11 +861,7 @@ "An Observation can be used with a DataFrame only once." ] }, - "SCHEMA_MISMATCH_FOR_PANDAS_UDF": { - "message": [ - "Result vector from <udf_type> was not the required length: expected <expected>, got <actual>." - ] - }, + "SESSION_ALREADY_EXIST": { "message": [ "Cannot start a remote Spark session because there is a regular Spark session already running." diff --git a/python/pyspark/errors/error_classes.py b/python/pyspark/errors/error_classes.py index 4a840bafb8ae1..2f0417b0c8db3 100644 --- a/python/pyspark/errors/error_classes.py +++ b/python/pyspark/errors/error_classes.py @@ -15,8 +15,8 @@ # limitations under the License. # -import json import importlib.resources +import json # Note: Though we call them "error classes" here, the proper name is "error conditions", # hence why the name of the JSON file is different. diff --git a/python/pyspark/errors/exceptions/__init__.py b/python/pyspark/errors/exceptions/__init__.py index 3de00c17500c1..b858a3b996729 100644 --- a/python/pyspark/errors/exceptions/__init__.py +++ b/python/pyspark/errors/exceptions/__init__.py @@ -19,6 +19,7 @@ def _write_self() -> None: import json from pathlib import Path + from pyspark.errors import error_classes ERRORS_DIR = Path(__file__).parents[1] diff --git a/python/pyspark/errors/exceptions/base.py b/python/pyspark/errors/exceptions/base.py index 1e9f0fef27a76..aeabb46c2579e 100644 --- a/python/pyspark/errors/exceptions/base.py +++ b/python/pyspark/errors/exceptions/base.py @@ -17,12 +17,12 @@ import warnings from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Dict, Optional, TypeVar, cast, Iterable, TYPE_CHECKING, List +from pickle import PicklingError +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, TypeVar, cast from pyspark.errors.exceptions.tblib import Traceback from pyspark.errors.utils import ErrorClassesReader from pyspark.logger import PySparkLogger -from pickle import PicklingError if TYPE_CHECKING: from pyspark.sql.types import Row diff --git a/python/pyspark/errors/exceptions/captured.py b/python/pyspark/errors/exceptions/captured.py index da755de0d6ec1..8aa9d9b624e1c 100644 --- a/python/pyspark/errors/exceptions/captured.py +++ b/python/pyspark/errors/exceptions/captured.py @@ -16,33 +16,65 @@ # import warnings from contextlib import contextmanager -from typing import Any, Callable, Dict, Iterator, Optional, cast, List, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, cast from pyspark.errors.exceptions.base import ( AnalysisException as BaseAnalysisException, - IllegalArgumentException as BaseIllegalArgumentException, +) +from pyspark.errors.exceptions.base import ( ArithmeticException as BaseArithmeticException, - UnsupportedOperationException as BaseUnsupportedOperationException, +) +from pyspark.errors.exceptions.base import ( ArrayIndexOutOfBoundsException as BaseArrayIndexOutOfBoundsException, +) +from pyspark.errors.exceptions.base import ( DateTimeException as BaseDateTimeException, +) +from pyspark.errors.exceptions.base import ( + IllegalArgumentException as BaseIllegalArgumentException, +) +from pyspark.errors.exceptions.base import ( NumberFormatException as BaseNumberFormatException, +) +from pyspark.errors.exceptions.base import ( ParseException as BaseParseException, +) +from pyspark.errors.exceptions.base import ( PySparkException, + QueryContextType, + recover_python_exception, +) +from pyspark.errors.exceptions.base import ( PythonException as BasePythonException, +) +from pyspark.errors.exceptions.base import ( + QueryContext as BaseQueryContext, +) +from pyspark.errors.exceptions.base import ( QueryExecutionException as BaseQueryExecutionException, +) +from pyspark.errors.exceptions.base import ( + SparkNoSuchElementException as BaseNoSuchElementException, +) +from pyspark.errors.exceptions.base import ( SparkRuntimeException as BaseSparkRuntimeException, +) +from pyspark.errors.exceptions.base import ( SparkUpgradeException as BaseSparkUpgradeException, - SparkNoSuchElementException as BaseNoSuchElementException, +) +from pyspark.errors.exceptions.base import ( StreamingQueryException as BaseStreamingQueryException, +) +from pyspark.errors.exceptions.base import ( UnknownException as BaseUnknownException, - QueryContext as BaseQueryContext, - QueryContextType, - recover_python_exception, +) +from pyspark.errors.exceptions.base import ( + UnsupportedOperationException as BaseUnsupportedOperationException, ) if TYPE_CHECKING: - from py4j.protocol import Py4JJavaError from py4j.java_gateway import JavaObject + from py4j.protocol import Py4JJavaError class CapturedException(PySparkException): @@ -53,9 +85,10 @@ def __init__( cause: Optional["Py4JJavaError"] = None, origin: Optional["Py4JJavaError"] = None, ): - from pyspark import SparkContext from py4j.protocol import Py4JJavaError + from pyspark import SparkContext + # desc & stackTrace vs origin are mutually exclusive. # cause is optional. assert (origin is not None and desc is None and stackTrace is None) or ( @@ -98,9 +131,10 @@ def __str__(self) -> str: return str(desc) def getCondition(self) -> Optional[str]: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert SparkContext._gateway is not None assert SparkContext._jvm is not None @@ -118,9 +152,10 @@ def getErrorClass(self) -> Optional[str]: return self.getCondition() def getMessageParameters(self) -> Optional[Dict[str, str]]: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert SparkContext._gateway is not None assert SparkContext._jvm is not None @@ -134,9 +169,10 @@ def getMessageParameters(self) -> Optional[Dict[str, str]]: return None def getSqlState(self) -> Optional[str]: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert SparkContext._gateway is not None assert SparkContext._jvm is not None gw = SparkContext._gateway @@ -149,9 +185,10 @@ def getSqlState(self) -> Optional[str]: return None def getMessage(self) -> str: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert SparkContext._gateway is not None assert SparkContext._jvm is not None gw = SparkContext._gateway @@ -172,9 +209,10 @@ def getMessage(self) -> str: return "" def getQueryContext(self) -> List[BaseQueryContext]: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert SparkContext._gateway is not None assert SparkContext._jvm is not None @@ -201,9 +239,10 @@ def convert_exception(e: "Py4JJavaError") -> CapturedException: def _convert_exception(e: "Py4JJavaError") -> CapturedException: - from pyspark import SparkContext from py4j.java_gateway import is_instance_of + from pyspark import SparkContext + assert e is not None assert SparkContext._jvm is not None assert SparkContext._gateway is not None @@ -268,9 +307,10 @@ def deco(*a: Any, **kw: Any) -> Any: @contextmanager def unwrap_spark_exception() -> Iterator[Any]: - from pyspark import SparkContext - from py4j.protocol import Py4JJavaError from py4j.java_gateway import is_instance_of + from py4j.protocol import Py4JJavaError + + from pyspark import SparkContext assert SparkContext._gateway is not None diff --git a/python/pyspark/errors/exceptions/connect.py b/python/pyspark/errors/exceptions/connect.py index 90537c2cc364e..794a9d7e73a70 100644 --- a/python/pyspark/errors/exceptions/connect.py +++ b/python/pyspark/errors/exceptions/connect.py @@ -14,45 +14,91 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import grpc import json +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import grpc from grpc import StatusCode -from typing import Any, Dict, List, Optional, TYPE_CHECKING from pyspark.errors.exceptions.base import ( AnalysisException as BaseAnalysisException, - IllegalArgumentException as BaseIllegalArgumentException, +) +from pyspark.errors.exceptions.base import ( ArithmeticException as BaseArithmeticException, - UnsupportedOperationException as BaseUnsupportedOperationException, +) +from pyspark.errors.exceptions.base import ( ArrayIndexOutOfBoundsException as BaseArrayIndexOutOfBoundsException, +) +from pyspark.errors.exceptions.base import ( DateTimeException as BaseDateTimeException, +) +from pyspark.errors.exceptions.base import ( + IllegalArgumentException as BaseIllegalArgumentException, +) +from pyspark.errors.exceptions.base import ( NumberFormatException as BaseNumberFormatException, +) +from pyspark.errors.exceptions.base import ( ParseException as BaseParseException, +) +from pyspark.errors.exceptions.base import ( + PickleException as BasePickleException, +) +from pyspark.errors.exceptions.base import ( PySparkException, + QueryContextType, + recover_python_exception, +) +from pyspark.errors.exceptions.base import ( PythonException as BasePythonException, - StreamingQueryException as BaseStreamingQueryException, +) +from pyspark.errors.exceptions.base import ( + QueryContext as BaseQueryContext, +) +from pyspark.errors.exceptions.base import ( QueryExecutionException as BaseQueryExecutionException, - SparkRuntimeException as BaseSparkRuntimeException, +) +from pyspark.errors.exceptions.base import ( SparkNoSuchElementException as BaseNoSuchElementException, +) +from pyspark.errors.exceptions.base import ( + SparkRuntimeException as BaseSparkRuntimeException, +) +from pyspark.errors.exceptions.base import ( SparkUpgradeException as BaseSparkUpgradeException, - QueryContext as BaseQueryContext, - QueryContextType, +) +from pyspark.errors.exceptions.base import ( StreamingPythonRunnerInitializationException as BaseStreamingPythonRunnerInitException, - PickleException as BasePickleException, +) +from pyspark.errors.exceptions.base import ( + StreamingQueryException as BaseStreamingQueryException, +) +from pyspark.errors.exceptions.base import ( UnknownException as BaseUnknownException, - recover_python_exception, +) +from pyspark.errors.exceptions.base import ( + UnsupportedOperationException as BaseUnsupportedOperationException, ) if TYPE_CHECKING: - import pyspark.sql.connect.proto as pb2 from google.rpc.error_details_pb2 import ErrorInfo + import pyspark.sql.connect.proto as pb2 + class SparkConnectException(PySparkException): """ Exception thrown from Spark Connect. """ + @property + def operation_id(self) -> Optional[str]: + """The Spark Connect ExecutePlan operation ID, when available. + + .. versionadded:: 4.3.0 + """ + return getattr(self, "_operation_id", None) + def convert_exception( info: "ErrorInfo", diff --git a/python/pyspark/errors/tests/test_connect_errors_conversion.py b/python/pyspark/errors/tests/test_connect_errors_conversion.py index d3bf96c4acb40..1bda1bc19be98 100644 --- a/python/pyspark/errors/tests/test_connect_errors_conversion.py +++ b/python/pyspark/errors/tests/test_connect_errors_conversion.py @@ -17,15 +17,15 @@ import unittest -from pyspark.testing.utils import should_test_connect, connect_requirement_message +from pyspark.testing.utils import connect_requirement_message, should_test_connect if should_test_connect: from pyspark.errors.exceptions.connect import ( - convert_exception, EXCEPTION_CLASS_MAPPING, - SparkConnectGrpcException, - PythonException, AnalysisException, + PythonException, + SparkConnectGrpcException, + convert_exception, ) @@ -116,6 +116,7 @@ def test_exception_class_mapping(self): def test_convert_exception_with_stacktrace(self): # Mock FetchErrorDetailsResponse with stacktrace from google.rpc.error_details_pb2 import ErrorInfo + from pyspark.sql.connect.proto import FetchErrorDetailsResponse as pb2 resp = pb2( @@ -190,10 +191,11 @@ def test_convert_exception_fallback(self): def test_convert_exception_with_breaking_change_info(self): """Test that breaking change info is correctly extracted from protobuf response.""" - import pyspark.sql.connect.proto as pb2 from google.rpc.error_details_pb2 import ErrorInfo from grpc import StatusCode + import pyspark.sql.connect.proto as pb2 + # Create mock FetchErrorDetailsResponse with breaking change info resp = pb2.FetchErrorDetailsResponse() resp.root_error_idx = 0 @@ -245,10 +247,11 @@ def test_convert_exception_with_breaking_change_info(self): def test_convert_exception_without_breaking_change_info(self): """Test that getBreakingChangeInfo returns None when no breaking change info.""" - import pyspark.sql.connect.proto as pb2 from google.rpc.error_details_pb2 import ErrorInfo from grpc import StatusCode + import pyspark.sql.connect.proto as pb2 + # Create mock FetchErrorDetailsResponse without breaking change info resp = pb2.FetchErrorDetailsResponse() resp.root_error_idx = 0 @@ -325,10 +328,11 @@ def test_breaking_change_info_inheritance(self): def test_breaking_change_info_without_mitigation_config(self): """Test breaking change info that only has migration messages.""" - import pyspark.sql.connect.proto as pb2 from google.rpc.error_details_pb2 import ErrorInfo from grpc import StatusCode + import pyspark.sql.connect.proto as pb2 + # Create mock FetchErrorDetailsResponse with breaking change info (no mitigation config) resp = pb2.FetchErrorDetailsResponse() resp.root_error_idx = 0 @@ -367,10 +371,11 @@ def test_breaking_change_info_without_mitigation_config(self): def test_convert_exception_error_class_from_fetch_error_details(self): """Test that errorClass is extracted from FetchErrorDetailsResponse when not present in ErrorInfo metadata (e.g., when messageParameters exceed limit).""" - import pyspark.sql.connect.proto as pb2 from google.rpc.error_details_pb2 import ErrorInfo from grpc import StatusCode + import pyspark.sql.connect.proto as pb2 + # Create mock FetchErrorDetailsResponse with errorClass resp = pb2.FetchErrorDetailsResponse() resp.root_error_idx = 0 diff --git a/python/pyspark/errors/tests/test_errors.py b/python/pyspark/errors/tests/test_errors.py index 17ca6a34c6dd5..e477f4a99b0e0 100644 --- a/python/pyspark/errors/tests/test_errors.py +++ b/python/pyspark/errors/tests/test_errors.py @@ -118,12 +118,35 @@ def test_sqlstate(self): error = PySparkRuntimeError(errorClass="APPLICATION_NAME_NOT_SET", messageParameters={}) self.assertIsNone(error.getSqlState()) + # Neither the sub-class nor the main class declares a sqlState. error = PySparkRuntimeError( errorClass="SESSION_MUTATION_IN_DECLARATIVE_PIPELINE.SET_RUNTIME_CONF", messageParameters={"method": "set"}, ) self.assertIsNone(error.getSqlState()) + # A sub-class inherits the main class's sqlState. + error = PySparkRuntimeError( + errorClass="NEAREST_BY_JOIN.UNSUPPORTED_MODE", + messageParameters={"mode": "invalid", "supported": "nearest"}, + ) + self.assertEqual(error.getSqlState(), "42604") + + def test_sqlstate_is_taken_from_the_main_class(self): + # The sqlState is looked up on the main class only, so a sub-class entry never + # takes precedence and an unknown sub-class name still resolves. + error_reader = ErrorClassesReader() + error_reader.error_info_map = { + "TEST_ERROR": { + "message": ["Error message."], + "sqlState": "42604", + "sub_class": {"SUBCLASS": {"message": ["Subclass message."], "sqlState": "08003"}}, + }, + } + self.assertEqual(error_reader.get_sqlstate("TEST_ERROR.SUBCLASS"), "42604") + self.assertEqual(error_reader.get_sqlstate("TEST_ERROR.NON_EXISTENT_SUB"), "42604") + self.assertIsNone(error_reader.get_sqlstate("NON_EXISTENT_ERROR.SUBCLASS")) + if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/errors/utils.py b/python/pyspark/errors/utils.py index c796e37a28c56..115d9f31cb3b1 100644 --- a/python/pyspark/errors/utils.py +++ b/python/pyspark/errors/utils.py @@ -14,12 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import re import functools import inspect import itertools import os +import re import threading +from types import FrameType from typing import ( Any, Callable, @@ -27,14 +28,13 @@ Iterator, List, Match, - TypeVar, - Type, Optional, + Type, + TypeVar, Union, - overload, cast, + overload, ) -from types import FrameType import pyspark from pyspark.errors.error_classes import ERROR_CLASSES_MAP @@ -98,20 +98,16 @@ def __init__(self) -> None: def get_sqlstate(self, errorClass: Optional[str]) -> Optional[str]: """ Returns the SQL state for the given error class. + + The SQL state is declared on the main class, so a sub-class inherits it. """ if errorClass is None: return None - error_classes = errorClass.split(".") - try: - if len(error_classes) == 1: - return self.error_info_map[errorClass]["sqlState"] - else: - return self.error_info_map[error_classes[0]]["sub_class"][error_classes[1]][ - "sqlState" - ] - except KeyError: + main_class_info = self.error_info_map.get(errorClass.split(".")[0]) + if main_class_info is None: return None + return main_class_info.get("sqlState") def get_error_message(self, errorClass: str, messageParameters: Dict[str, str]) -> str: """ @@ -268,10 +264,9 @@ def inspect_stack() -> Iterator[FrameType]: # We try import here since IPython is not a required dependency try: - import IPython - # ipykernel is required for IPython import ipykernel + import IPython ipython = IPython.get_ipython() # Filtering out IPython related frames diff --git a/python/pyspark/errors_doc_gen.py b/python/pyspark/errors_doc_gen.py index c541cd8fb7e2b..30bd1dc1be35d 100644 --- a/python/pyspark/errors_doc_gen.py +++ b/python/pyspark/errors_doc_gen.py @@ -2,25 +2,8 @@ from pyspark.errors.error_classes import ERROR_CLASSES_MAP - -def generate_errors_doc(output_rst_file_path: str) -> None: - """ - Generates a reStructuredText (RST) documentation file for PySpark error classes. - - This function fetches error classes defined in `pyspark.errors.error_classes` - and writes them into an RST file. The generated RST file provides an overview - of common, named error classes returned by PySpark. - - Parameters - ---------- - output_rst_file_path : str - The file path where the RST documentation will be written. - - Notes - ----- - The generated RST file can be rendered using Sphinx to visualize the documentation. - """ - header = """.. Licensed to the Apache Software Foundation (ASF) under one +_ERROR_DOC_HEADER = """ +.. 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 @@ -41,12 +24,34 @@ def generate_errors_doc(output_rst_file_path: str) -> None: Error classes in PySpark ======================== -This is a list of common, named error classes returned by PySpark which are defined at `error-conditions.json <https://github.com/apache/spark/blob/master/python/pyspark/errors/error-conditions.json>`_. +This is a list of common, named error classes returned by PySpark which are defined at +`error-conditions.json <https://github.com/apache/spark/blob/master/python/pyspark/errors/error-conditions.json>`_. -When writing PySpark errors, developers must use an error class from the list. If an appropriate error class is not available, add a new one into the list. For more information, please refer to `Contributing Error and Exception <contributing.rst#contributing-error-and-exception>`_. +When writing PySpark errors, developers must use an error class from the list. If an appropriate +error class is not available, add a new one into the list. For more information, please refer to +`Contributing Errors and Exceptions <contributing.rst#contributing-errors-and-exceptions>`_. """ + + +def generate_errors_doc(output_rst_file_path: str) -> None: + """ + Generates a reStructuredText (RST) documentation file for PySpark error classes. + + This function fetches error classes defined in `pyspark.errors.error_classes` + and writes them into an RST file. The generated RST file provides an overview + of common, named error classes returned by PySpark. + + Parameters + ---------- + output_rst_file_path : str + The file path where the RST documentation will be written. + + Notes + ----- + The generated RST file can be rendered using Sphinx to visualize the documentation. + """ with open(output_rst_file_path, "w", encoding="utf-8") as f: - f.write(header + "\n\n") + f.write(_ERROR_DOC_HEADER.strip() + "\n\n") for error_key, error_details in ERROR_CLASSES_MAP.items(): f.write(error_key + "\n") # The length of the error class name and underline must be the same diff --git a/python/pyspark/install.py b/python/pyspark/install.py index 4426a5899b324..fd157f3a08c2e 100644 --- a/python/pyspark/install.py +++ b/python/pyspark/install.py @@ -23,7 +23,6 @@ from shutil import rmtree from typing import TYPE_CHECKING - if TYPE_CHECKING: from http.client import HTTPResponse diff --git a/python/pyspark/instrumentation_utils.py b/python/pyspark/instrumentation_utils.py index 1f822207c4afc..0ee75d43378c3 100644 --- a/python/pyspark/instrumentation_utils.py +++ b/python/pyspark/instrumentation_utils.py @@ -16,12 +16,12 @@ # import functools +import importlib import inspect import threading -import importlib import time from types import ModuleType -from typing import Tuple, Union, List, Callable, Any, Type +from typing import Any, Callable, List, Tuple, Type, Union __all__: List[str] = [] diff --git a/python/pyspark/java_gateway.py b/python/pyspark/java_gateway.py index 6303a43618578..eccca1f412a76 100644 --- a/python/pyspark/java_gateway.py +++ b/python/pyspark/java_gateway.py @@ -17,20 +17,20 @@ import atexit import os -import signal +import platform import shlex import shutil -import platform +import signal import tempfile import time -from subprocess import Popen, PIPE +from subprocess import PIPE, Popen -from py4j.java_gateway import java_import, JavaGateway, JavaObject, GatewayParameters from py4j.clientserver import ClientServer, JavaParameters, PythonParameters -from pyspark.serializers import read_int, UTF8Deserializer +from py4j.java_gateway import GatewayParameters, JavaGateway, JavaObject, java_import -from pyspark.find_spark_home import _find_spark_home from pyspark.errors import PySparkRuntimeError +from pyspark.find_spark_home import _find_spark_home +from pyspark.serializers import UTF8Deserializer, read_int # for backward compatibility references. from pyspark.util import local_connect_and_auth # noqa: F401 diff --git a/python/pyspark/logger/__init__.py b/python/pyspark/logger/__init__.py index 9e9548e919833..fa01187748a1f 100644 --- a/python/pyspark/logger/__init__.py +++ b/python/pyspark/logger/__init__.py @@ -19,6 +19,6 @@ PySpark logging """ -from pyspark.logger.logger import PySparkLogger, SPARK_LOG_SCHEMA +from pyspark.logger.logger import SPARK_LOG_SCHEMA, PySparkLogger __all__ = ["PySparkLogger", "SPARK_LOG_SCHEMA"] diff --git a/python/pyspark/logger/logger.py b/python/pyspark/logger/logger.py index 8d08275fe03cd..bb390fe8b0cfe 100644 --- a/python/pyspark/logger/logger.py +++ b/python/pyspark/logger/logger.py @@ -15,11 +15,11 @@ # limitations under the License. # -import logging import json -import traceback +import logging import sys -from typing import cast, Mapping, Optional, TYPE_CHECKING +import traceback +from typing import TYPE_CHECKING, Mapping, Optional, cast if TYPE_CHECKING: from logging import _ArgsType, _ExcInfoType @@ -314,6 +314,7 @@ def _log( def _test() -> None: import doctest + import pyspark.logger.logger globs = pyspark.logger.logger.__dict__.copy() diff --git a/python/pyspark/logger/tests/test_logger.py b/python/pyspark/logger/tests/test_logger.py index 7c3fd4b1e61e3..c3227b01045dd 100644 --- a/python/pyspark/logger/tests/test_logger.py +++ b/python/pyspark/logger/tests/test_logger.py @@ -14,13 +14,15 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import logging import json +import logging import tempfile from io import StringIO + from pyspark.errors import ArithmeticException -from pyspark.logger.logger import PySparkLogger, SPARK_LOG_SCHEMA -from pyspark.sql import Row, functions as sf +from pyspark.logger.logger import SPARK_LOG_SCHEMA, PySparkLogger +from pyspark.sql import Row +from pyspark.sql import functions as sf from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/logger/worker_io.py b/python/pyspark/logger/worker_io.py index 4b8a7cf15912d..b573b38e3ecc9 100644 --- a/python/pyspark/logger/worker_io.py +++ b/python/pyspark/logger/worker_io.py @@ -15,15 +15,15 @@ # limitations under the License. # -from contextlib import contextmanager import inspect import io import logging import os import sys import time -from typing import BinaryIO, Callable, Generator, Iterable, Iterator, Optional, TextIO, Union +from contextlib import contextmanager from types import FrameType, TracebackType +from typing import BinaryIO, Callable, Generator, Iterable, Iterator, Optional, TextIO, Union from pyspark.logger.logger import JSONFormatter diff --git a/python/pyspark/memory_profiler_ext.py b/python/pyspark/memory_profiler_ext.py index df09377c9d544..a3e75c8e9f40e 100644 --- a/python/pyspark/memory_profiler_ext.py +++ b/python/pyspark/memory_profiler_ext.py @@ -15,10 +15,10 @@ # limitations under the License. # -from types import CodeType -from typing import Any, Optional, List, Iterator, Tuple, Type, TYPE_CHECKING, Callable import inspect import warnings +from types import CodeType +from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Optional, Tuple, Type if TYPE_CHECKING: has_memory_profiler: bool diff --git a/python/pyspark/messages/__init__.py b/python/pyspark/messages/__init__.py index 69cfbf6bd53a2..d99409c606ea8 100644 --- a/python/pyspark/messages/__init__.py +++ b/python/pyspark/messages/__init__.py @@ -15,9 +15,9 @@ # limitations under the License. # +from pyspark.messages.socket.spark_socket_message_receiver import SparkSocketMessageReceiver from pyspark.messages.spark_message_receiver import SparkMessageReceiver from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream -from pyspark.messages.socket.spark_socket_message_receiver import SparkSocketMessageReceiver __all__ = [ "SparkMessageReceiver", diff --git a/python/pyspark/messages/socket/spark_socket_message_receiver.py b/python/pyspark/messages/socket/spark_socket_message_receiver.py index 7099f0207c1f1..3a17a86e80698 100644 --- a/python/pyspark/messages/socket/spark_socket_message_receiver.py +++ b/python/pyspark/messages/socket/spark_socket_message_receiver.py @@ -17,11 +17,11 @@ from typing import BinaryIO -from pyspark.serializers import read_int, SpecialLengths -from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream from pyspark.messages.spark_message_receiver import ( SparkMessageReceiver, ) +from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream +from pyspark.serializers import SpecialLengths, read_int def _assert_message_id(message_id: int, expected: int) -> None: diff --git a/python/pyspark/messages/spark_message_receiver.py b/python/pyspark/messages/spark_message_receiver.py index ec6b6fc306243..2266558d4aa38 100644 --- a/python/pyspark/messages/spark_message_receiver.py +++ b/python/pyspark/messages/spark_message_receiver.py @@ -15,14 +15,13 @@ # limitations under the License. # +from abc import ABC, abstractmethod from enum import Enum from functools import wraps from typing import BinaryIO, Callable, TypeVar -from abc import ABC, abstractmethod from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream - T = TypeVar("T", bound="SparkMessageReceiver") R = TypeVar("R") diff --git a/python/pyspark/messages/zero_copy_byte_stream.py b/python/pyspark/messages/zero_copy_byte_stream.py index 611b4f928be08..69e58905c020d 100644 --- a/python/pyspark/messages/zero_copy_byte_stream.py +++ b/python/pyspark/messages/zero_copy_byte_stream.py @@ -16,8 +16,8 @@ # import threading -from typing import Optional from collections import deque +from typing import Optional class ZeroCopyByteStream: diff --git a/python/pyspark/ml/__init__.py b/python/pyspark/ml/__init__.py index 2f692ee2e569a..4197a5db26356 100644 --- a/python/pyspark/ml/__init__.py +++ b/python/pyspark/ml/__init__.py @@ -20,15 +20,6 @@ machine learning pipelines. """ -from pyspark.ml.base import ( - Estimator, - Model, - Predictor, - PredictionModel, - Transformer, - UnaryTransformer, -) -from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml import ( classification, clustering, @@ -36,14 +27,23 @@ feature, fpm, image, + linalg, + param, recommendation, regression, stat, tuning, util, - linalg, - param, ) +from pyspark.ml.base import ( + Estimator, + Model, + PredictionModel, + Predictor, + Transformer, + UnaryTransformer, +) +from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml.torch.distributor import TorchDistributor __all__ = [ diff --git a/python/pyspark/ml/_typing.pyi b/python/pyspark/ml/_typing.pyi index c24dfe577350e..f7224d5f038f4 100644 --- a/python/pyspark/ml/_typing.pyi +++ b/python/pyspark/ml/_typing.pyi @@ -16,19 +16,19 @@ # specific language governing permissions and limitations # under the License. -from typing import Any, Dict, List, TYPE_CHECKING, TypeVar, Tuple, Union -from typing_extensions import Literal +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, TypeVar, Union from numpy import ndarray from py4j.java_gateway import JavaObject +from typing_extensions import Literal import pyspark.ml.base import pyspark.ml.param -from pyspark.ml.linalg import Vector import pyspark.ml.wrapper +from pyspark.ml.linalg import Vector if TYPE_CHECKING: - from scipy.sparse import spmatrix, sparray + from scipy.sparse import sparray, spmatrix ParamMap = Dict[pyspark.ml.param.Param, Any] PipelineStage = Union[pyspark.ml.base.Estimator, pyspark.ml.base.Transformer] diff --git a/python/pyspark/ml/base.py b/python/pyspark/ml/base.py index 882c9b1f65a0c..928976eee257b 100644 --- a/python/pyspark/ml/base.py +++ b/python/pyspark/ml/base.py @@ -15,10 +15,11 @@ # limitations under the License. # -from abc import ABCMeta, abstractmethod import copy import threading +from abc import ABCMeta, abstractmethod from typing import ( + TYPE_CHECKING, Any, Callable, Generic, @@ -31,17 +32,16 @@ Union, cast, overload, - TYPE_CHECKING, ) from pyspark import since -from pyspark.ml.param import P from pyspark.ml.common import inherit_doc +from pyspark.ml.param import P from pyspark.ml.param.shared import ( + HasFeaturesCol, HasInputCol, - HasOutputCol, HasLabelCol, - HasFeaturesCol, + HasOutputCol, HasPredictionCol, Params, ) diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py index 4b7f2e4da2090..0006cbc619ee8 100644 --- a/python/pyspark/ml/classification.py +++ b/python/pyspark/ml/classification.py @@ -20,88 +20,90 @@ import uuid import warnings from abc import ABCMeta, abstractmethod -from multiprocessing.pool import ThreadPool from functools import cached_property +from multiprocessing.pool import ThreadPool from typing import ( + TYPE_CHECKING, Any, + Callable, Dict, Generic, List, Optional, + Tuple, Type, TypeVar, Union, cast, overload, - TYPE_CHECKING, - Tuple, - Callable, ) -from pyspark import keyword_only, since, inheritable_thread_target -from pyspark.ml import Estimator, Predictor, PredictionModel, Model, functions as MF +from pyspark import inheritable_thread_target, keyword_only, since +from pyspark.ml import functions as MF +from pyspark.ml.base import Estimator, Model, PredictionModel, Predictor, _PredictorParams +from pyspark.ml.common import inherit_doc +from pyspark.ml.linalg import Matrix, Vector from pyspark.ml.param.shared import ( - HasRawPredictionCol, + HasAggregationDepth, + HasBlockSize, + HasElasticNetParam, + HasFitIntercept, + HasMaxBlockSizeInMB, + HasMaxIter, + HasParallelism, HasProbabilityCol, - HasThresholds, + HasRawPredictionCol, HasRegParam, - HasMaxIter, - HasFitIntercept, - HasTol, + HasSeed, + HasSolver, HasStandardization, - HasWeightCol, - HasAggregationDepth, + HasStepSize, HasThreshold, - HasBlockSize, - HasMaxBlockSizeInMB, + HasThresholds, + HasTol, + HasWeightCol, Param, Params, TypeConverters, - HasElasticNetParam, - HasSeed, - HasStepSize, - HasSolver, - HasParallelism, ) +from pyspark.ml.regression import DecisionTreeRegressionModel, _FactorizationMachinesParams from pyspark.ml.tree import ( _DecisionTreeModel, _DecisionTreeParams, - _TreeEnsembleModel, - _RandomForestParams, _GBTParams, _HasVarianceImpurity, + _RandomForestParams, _TreeClassifierParams, + _TreeEnsembleModel, ) -from pyspark.ml.regression import _FactorizationMachinesParams, DecisionTreeRegressionModel -from pyspark.ml.base import _PredictorParams from pyspark.ml.util import ( DefaultParamsReader, DefaultParamsWriter, + HasTrainingSummary, JavaMLReadable, JavaMLWritable, JavaMLWriter, - MLReader, MLReadable, - MLWriter, + MLReader, MLWritable, - HasTrainingSummary, + MLWriter, + _cache_spark_dataset, + try_remote_attribute_relation, try_remote_read, try_remote_write, - try_remote_attribute_relation, - _cache_spark_dataset, ) -from pyspark.ml.wrapper import JavaParams, JavaPredictor, JavaPredictionModel, JavaWrapper -from pyspark.ml.common import inherit_doc -from pyspark.ml.linalg import Matrix, Vector -from pyspark.sql import DataFrame, Row, SparkSession, functions as F +from pyspark.ml.wrapper import JavaParams, JavaPredictionModel, JavaPredictor, JavaWrapper +from pyspark.sql import DataFrame, Row, SparkSession +from pyspark.sql import functions as F from pyspark.sql.internal import InternalFunction as SF -from pyspark.storagelevel import StorageLevel from pyspark.sql.utils import is_remote +from pyspark.storagelevel import StorageLevel if TYPE_CHECKING: - from pyspark.ml._typing import P, ParamMap from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext + from pyspark.ml._typing import P, ParamMap T = TypeVar("T") @@ -4352,6 +4354,7 @@ class FMClassificationTrainingSummary(FMClassificationSummary, _TrainingSummary) if __name__ == "__main__": import doctest + import pyspark.ml.classification from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/clustering.py b/python/pyspark/ml/clustering.py index ba3e486cb4ef6..525b3e7b3baaa 100644 --- a/python/pyspark/ml/clustering.py +++ b/python/pyspark/ml/clustering.py @@ -15,51 +15,52 @@ # limitations under the License. # +import functools import sys import warnings -from typing import Any, Dict, List, Optional, TYPE_CHECKING -import functools +from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np -from pyspark import since, keyword_only +from pyspark import keyword_only, since +from pyspark.ml.common import inherit_doc +from pyspark.ml.linalg import Matrix, Vector from pyspark.ml.param.shared import ( - HasMaxIter, + HasAggregationDepth, + HasCheckpointInterval, + HasDistanceMeasure, HasFeaturesCol, - HasSeed, + HasIntermediateStorageLevel, + HasMaxBlockSizeInMB, + HasMaxIter, HasPredictionCol, - HasAggregationDepth, - HasWeightCol, - HasTol, HasProbabilityCol, - HasDistanceMeasure, - HasCheckpointInterval, + HasSeed, HasSolver, - HasMaxBlockSizeInMB, - HasIntermediateStorageLevel, + HasTol, + HasWeightCol, Param, Params, TypeConverters, ) +from pyspark.ml.stat import MultivariateGaussian from pyspark.ml.util import ( - JavaMLWritable, - JavaMLReadable, GeneralJavaMLWritable, HasTrainingSummary, - try_remote_attribute_relation, + JavaMLReadable, + JavaMLWritable, invoke_helper_relation, + try_remote_attribute_relation, ) from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams, JavaWrapper -from pyspark.ml.common import inherit_doc -from pyspark.ml.stat import MultivariateGaussian from pyspark.sql import DataFrame -from pyspark.ml.linalg import Vector, Matrix from pyspark.sql.utils import is_remote if TYPE_CHECKING: - from pyspark.ml._typing import M from py4j.java_gateway import JavaObject + from pyspark.ml._typing import M + __all__ = [ "BisectingKMeans", @@ -2180,7 +2181,9 @@ def assignClusters(self, dataset: DataFrame) -> DataFrame: if __name__ == "__main__": import doctest + import numpy + import pyspark.ml.clustering from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/common.py b/python/pyspark/ml/common.py index 2417df6ab9eb3..a0a4b02ee13ae 100644 --- a/python/pyspark/ml/common.py +++ b/python/pyspark/ml/common.py @@ -15,19 +15,19 @@ # limitations under the License. # -from typing import Any, Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable -from pyspark.util import is_remote_only -from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer +from pyspark.serializers import AutoBatchedSerializer, CPickleSerializer from pyspark.sql import DataFrame, SparkSession +from pyspark.util import is_remote_only if TYPE_CHECKING: import py4j.protocol from py4j.java_gateway import JavaObject import pyspark.core.context - from pyspark.core.rdd import RDD from pyspark.core.context import SparkContext + from pyspark.core.rdd import RDD from pyspark.ml._typing import C, JavaObjectOrPickleDump @@ -80,8 +80,9 @@ def _to_java_object_rdd(rdd: "RDD") -> "JavaObject": def _py2java(sc: "SparkContext", obj: Any) -> "JavaObject": """Convert Python object into Java""" from py4j.java_gateway import JavaObject - from pyspark.core.rdd import RDD + from pyspark.core.context import SparkContext + from pyspark.core.rdd import RDD if isinstance(obj, RDD): obj = _to_java_object_rdd(obj) @@ -103,9 +104,9 @@ def _py2java(sc: "SparkContext", obj: Any) -> "JavaObject": def _java2py(sc: "SparkContext", r: "JavaObjectOrPickleDump", encoding: str = "bytes") -> Any: - from py4j.protocol import Py4JJavaError - from py4j.java_gateway import JavaObject from py4j.java_collections import JavaArray, JavaList + from py4j.java_gateway import JavaObject + from py4j.protocol import Py4JJavaError if isinstance(r, JavaObject): clsName = r.getClass().getSimpleName() diff --git a/python/pyspark/ml/connect/__init__.py b/python/pyspark/ml/connect/__init__.py index c4bc8c9d84d2a..0ce207e5d865d 100644 --- a/python/pyspark/ml/connect/__init__.py +++ b/python/pyspark/ml/connect/__init__.py @@ -21,16 +21,16 @@ check_dependencies() -from pyspark.ml.connect.base import ( - Estimator, - Transformer, - Model, -) from pyspark.ml.connect import ( - feature, evaluation, + feature, tuning, ) +from pyspark.ml.connect.base import ( + Estimator, + Model, + Transformer, +) from pyspark.ml.connect.evaluation import Evaluator from pyspark.ml.connect.pipeline import Pipeline, PipelineModel diff --git a/python/pyspark/ml/connect/base.py b/python/pyspark/ml/connect/base.py index 92fe1ae4e32c0..3497aa9c87faf 100644 --- a/python/pyspark/ml/connect/base.py +++ b/python/pyspark/ml/connect/base.py @@ -17,28 +17,28 @@ from abc import ABCMeta, abstractmethod from typing import ( + TYPE_CHECKING, Any, + Callable, Generic, List, Optional, + Tuple, TypeVar, Union, - TYPE_CHECKING, - Tuple, - Callable, ) import pandas as pd from pyspark import since from pyspark.ml.common import inherit_doc -from pyspark.sql.dataframe import DataFrame from pyspark.ml.param import Params from pyspark.ml.param.shared import ( - HasLabelCol, HasFeaturesCol, + HasLabelCol, HasPredictionCol, ) +from pyspark.sql.dataframe import DataFrame if TYPE_CHECKING: from pyspark.ml._typing import ParamMap diff --git a/python/pyspark/ml/connect/classification.py b/python/pyspark/ml/connect/classification.py index 3263f47e6135f..71efe29cdd305 100644 --- a/python/pyspark/ml/connect/classification.py +++ b/python/pyspark/ml/connect/classification.py @@ -14,31 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Dict, Union, List, Tuple, Callable, Optional import math +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd from pyspark import keyword_only -from pyspark.ml.connect.base import _PredictorParams -from pyspark.ml.param.shared import HasProbabilityCol -from pyspark.sql import DataFrame from pyspark.ml.common import inherit_doc -from pyspark.ml.torch.distributor import TorchDistributor +from pyspark.ml.connect.base import PredictionModel, Predictor, _PredictorParams +from pyspark.ml.connect.io_utils import CoreModelReadWrite, ParamsReadWrite from pyspark.ml.param.shared import ( - HasMaxIter, - HasFitIntercept, - HasTol, - HasWeightCol, - HasSeed, - HasNumTrainWorkers, HasBatchSize, + HasFitIntercept, HasLearningRate, + HasMaxIter, HasMomentum, + HasNumTrainWorkers, + HasProbabilityCol, + HasSeed, + HasTol, + HasWeightCol, ) -from pyspark.ml.connect.base import Predictor, PredictionModel -from pyspark.ml.connect.io_utils import ParamsReadWrite, CoreModelReadWrite +from pyspark.ml.torch.distributor import TorchDistributor +from pyspark.sql import DataFrame from pyspark.sql import functions as sf @@ -86,12 +85,13 @@ def _train_logistic_regression_model_worker_fn( fit_intercept: bool, seed: int, ) -> Any: - from pyspark.ml.torch.distributor import _get_spark_partition_data_loader import torch - import torch.nn as torch_nn - from torch.nn.parallel import DistributedDataParallel as DDP import torch.distributed + import torch.nn as torch_nn import torch.optim as optim + from torch.nn.parallel import DistributedDataParallel as DDP + + from pyspark.ml.torch.distributor import _get_spark_partition_data_loader # TODO: add a setting seed param. torch.manual_seed(seed) diff --git a/python/pyspark/ml/connect/evaluation.py b/python/pyspark/ml/connect/evaluation.py index f324bb193c0ce..e4d4acb9cf603 100644 --- a/python/pyspark/ml/connect/evaluation.py +++ b/python/pyspark/ml/connect/evaluation.py @@ -14,16 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Union, List, Tuple +from typing import Any, List, Tuple, Union import numpy as np import pandas as pd from pyspark import keyword_only -from pyspark.ml.param import Param, Params, TypeConverters -from pyspark.ml.param.shared import HasLabelCol, HasPredictionCol, HasProbabilityCol from pyspark.ml.connect.base import Evaluator from pyspark.ml.connect.io_utils import ParamsReadWrite +from pyspark.ml.param import Param, Params, TypeConverters +from pyspark.ml.param.shared import HasLabelCol, HasPredictionCol, HasProbabilityCol from pyspark.sql import DataFrame diff --git a/python/pyspark/ml/connect/feature.py b/python/pyspark/ml/connect/feature.py index 2184b3c7f332f..e018522098b61 100644 --- a/python/pyspark/ml/connect/feature.py +++ b/python/pyspark/ml/connect/feature.py @@ -15,26 +15,26 @@ # limitations under the License. # -from typing import Any, Union, List, Tuple, Callable, Dict, Optional +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import numpy as np import pandas as pd import pyarrow as pa from pyspark import keyword_only -from pyspark.sql import DataFrame +from pyspark.ml.connect.base import Estimator, Model, Transformer +from pyspark.ml.connect.io_utils import CoreModelReadWrite, ParamsReadWrite from pyspark.ml.param.shared import ( + HasFeatureSizes, + HasHandleInvalid, HasInputCol, HasInputCols, HasOutputCol, - HasFeatureSizes, - HasHandleInvalid, Param, Params, TypeConverters, ) -from pyspark.ml.connect.base import Estimator, Model, Transformer -from pyspark.ml.connect.io_utils import ParamsReadWrite, CoreModelReadWrite +from pyspark.sql import DataFrame class MaxAbsScaler(Estimator, HasInputCol, HasOutputCol, ParamsReadWrite): diff --git a/python/pyspark/ml/connect/functions.py b/python/pyspark/ml/connect/functions.py index 5f2738cb2b3e5..3a66253fd3571 100644 --- a/python/pyspark/ml/connect/functions.py +++ b/python/pyspark/ml/connect/functions.py @@ -15,7 +15,7 @@ # limitations under the License. # -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING, Any from pyspark.ml import functions as PyMLFunctions from pyspark.sql.column import Column @@ -64,8 +64,9 @@ def _test() -> None: sys.exit(0) import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.ml.connect.functions + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.ml.connect.functions.__dict__.copy() diff --git a/python/pyspark/ml/connect/io_utils.py b/python/pyspark/ml/connect/io_utils.py index f6108934f3994..c6104d12d0944 100644 --- a/python/pyspark/ml/connect/io_utils.py +++ b/python/pyspark/ml/connect/io_utils.py @@ -16,17 +16,17 @@ # import json -import shutil import os +import shutil import tempfile import time -from urllib.parse import urlparse from typing import Any, Dict, List +from urllib.parse import urlparse +from pyspark import __version__ as pyspark_version from pyspark.ml.base import Params from pyspark.sql import SparkSession from pyspark.sql.utils import is_remote -from pyspark import __version__ as pyspark_version _META_DATA_FILE_NAME = "metadata.json" diff --git a/python/pyspark/ml/connect/pipeline.py b/python/pyspark/ml/connect/pipeline.py index 55e57b88da047..029dc40c9b4b0 100644 --- a/python/pyspark/ml/connect/pipeline.py +++ b/python/pyspark/ml/connect/pipeline.py @@ -14,18 +14,18 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Dict, List, Optional, Union, cast, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast import pandas as pd from pyspark import keyword_only, since +from pyspark.ml.common import inherit_doc from pyspark.ml.connect.base import Estimator, Model, Transformer from pyspark.ml.connect.io_utils import ( - ParamsReadWrite, MetaAlgorithmReadWrite, + ParamsReadWrite, ) from pyspark.ml.param import Param, Params -from pyspark.ml.common import inherit_doc from pyspark.sql.dataframe import DataFrame if TYPE_CHECKING: diff --git a/python/pyspark/ml/connect/proto.py b/python/pyspark/ml/connect/proto.py index eecf971440fb0..6d274b30390ba 100644 --- a/python/pyspark/ml/connect/proto.py +++ b/python/pyspark/ml/connect/proto.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Optional, TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Optional import pyspark.sql.connect.proto as pb2 from pyspark.sql.connect.plan import LogicalPlan diff --git a/python/pyspark/ml/connect/readwrite.py b/python/pyspark/ml/connect/readwrite.py index 08802a0fceaee..8f95f55352fdf 100644 --- a/python/pyspark/ml/connect/readwrite.py +++ b/python/pyspark/ml/connect/readwrite.py @@ -14,17 +14,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import cast, Type, TYPE_CHECKING, Union, Dict, Any + +# TODO: The type design of RemoteMLWriter is not very good. In saveInstance, +# we try to call session() with a connect session, but the session() method +# accepts a classic session. What's make it worse is that mypy considers some +# of the branches as unreachable. We need a full re-design of this module. +# mypy: disable-error-code="arg-type" + +from typing import TYPE_CHECKING, Any, Dict, Type, Union, cast import pyspark.sql.connect.proto as pb2 -from pyspark.ml.connect.serialize import serialize_ml_params, deserialize, deserialize_param -from pyspark.ml.util import MLWriter, MLReader, RL +from pyspark.ml.connect.serialize import deserialize, deserialize_param, serialize_ml_params +from pyspark.ml.util import RL, MLReader, MLWriter from pyspark.ml.wrapper import JavaWrapper if TYPE_CHECKING: from pyspark.core.context import SparkContext - from pyspark.sql.connect.session import SparkSession from pyspark.ml.util import JavaMLReadable, JavaMLWritable + from pyspark.sql.connect.session import SparkSession class RemoteMLWriter(MLWriter): @@ -58,17 +65,17 @@ def saveInstance( shouldOverwrite: bool = False, optionMap: Dict[str, Any] = {}, ) -> None: - from pyspark.ml.wrapper import JavaModel, JavaEstimator, JavaTransformer - from pyspark.ml.evaluation import JavaEvaluator - from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml.classification import OneVsRest, OneVsRestModel from pyspark.ml.clustering import PowerIterationClustering + from pyspark.ml.evaluation import JavaEvaluator + from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml.tuning import ( CrossValidator, CrossValidatorModel, TrainValidationSplit, TrainValidationSplitModel, ) + from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaTransformer # Spark Connect ML is built on scala Spark.ML, that means we're only # supporting JavaModel or JavaEstimator or JavaEvaluator @@ -119,21 +126,21 @@ def saveInstance( RemoteMLWriter.handleOverwrite(path, shouldOverwrite) pl_writer = PipelineWriter(instance) - pl_writer.session(session) # type: ignore[arg-type] + pl_writer.session(session) pl_writer.save(path) elif isinstance(instance, PipelineModel): from pyspark.ml.pipeline import PipelineModelWriter RemoteMLWriter.handleOverwrite(path, shouldOverwrite) plm_writer = PipelineModelWriter(instance) - plm_writer.session(session) # type: ignore[arg-type] + plm_writer.session(session) plm_writer.save(path) elif isinstance(instance, CrossValidator): from pyspark.ml.tuning import CrossValidatorWriter RemoteMLWriter.handleOverwrite(path, shouldOverwrite) cv_writer = CrossValidatorWriter(instance) - cv_writer.session(session) # type: ignore[arg-type] + cv_writer.session(session) cv_writer.save(path) elif isinstance(instance, CrossValidatorModel): from pyspark.ml.tuning import CrossValidatorModelWriter @@ -141,7 +148,7 @@ def saveInstance( RemoteMLWriter.handleOverwrite(path, shouldOverwrite) cvm_writer = CrossValidatorModelWriter(instance) cvm_writer.optionMap = optionMap - cvm_writer.session(session) # type: ignore[arg-type] + cvm_writer.session(session) cvm_writer.save(path) elif isinstance(instance, TrainValidationSplit): from pyspark.ml.tuning import TrainValidationSplitWriter @@ -155,21 +162,21 @@ def saveInstance( RemoteMLWriter.handleOverwrite(path, shouldOverwrite) tvsm_writer = TrainValidationSplitModelWriter(instance) tvsm_writer.optionMap = optionMap - tvsm_writer.session(session) # type: ignore[arg-type] + tvsm_writer.session(session) tvsm_writer.save(path) elif isinstance(instance, OneVsRest): from pyspark.ml.classification import OneVsRestWriter RemoteMLWriter.handleOverwrite(path, shouldOverwrite) ovr_writer = OneVsRestWriter(instance) - ovr_writer.session(session) # type: ignore[arg-type] + ovr_writer.session(session) ovr_writer.save(path) elif isinstance(instance, OneVsRestModel): from pyspark.ml.classification import OneVsRestModelWriter RemoteMLWriter.handleOverwrite(path, shouldOverwrite) ovrm_writer = OneVsRestModelWriter(instance) - ovrm_writer.session(session) # type: ignore[arg-type] + ovrm_writer.session(session) ovrm_writer.save(path) elif isinstance(instance, PowerIterationClustering): @@ -179,7 +186,7 @@ def saveInstance( transformer._resetUid(instance.uid) transformer._paramMap = instance._paramMap RemoteMLWriter.saveInstance( - transformer, # type: ignore[arg-type] + transformer, path, session, shouldOverwrite, @@ -217,17 +224,17 @@ def loadInstance( path: str, session: "SparkSession", ) -> RL: - from pyspark.ml.wrapper import JavaModel, JavaEstimator, JavaTransformer - from pyspark.ml.evaluation import JavaEvaluator - from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml.classification import OneVsRest, OneVsRestModel from pyspark.ml.clustering import PowerIterationClustering + from pyspark.ml.evaluation import JavaEvaluator + from pyspark.ml.pipeline import Pipeline, PipelineModel from pyspark.ml.tuning import ( CrossValidator, CrossValidatorModel, TrainValidationSplit, TrainValidationSplitModel, ) + from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaTransformer if ( issubclass(clazz, JavaModel) diff --git a/python/pyspark/ml/connect/serialize.py b/python/pyspark/ml/connect/serialize.py index 42bedfb330b1b..edd95d5898982 100644 --- a/python/pyspark/ml/connect/serialize.py +++ b/python/pyspark/ml/connect/serialize.py @@ -14,20 +14,20 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, List, TYPE_CHECKING, Mapping, Dict +from typing import TYPE_CHECKING, Any, Dict, List, Mapping import pyspark.sql.connect.proto as pb2 -from pyspark.sql.types import DataType from pyspark.ml.linalg import ( - DenseVector, - SparseVector, DenseMatrix, + DenseVector, SparseMatrix, + SparseVector, ) +from pyspark.sql.types import DataType if TYPE_CHECKING: - from pyspark.sql.connect.client import SparkConnectClient from pyspark.ml.param import Params + from pyspark.sql.connect.client import SparkConnectClient def literal_null() -> pb2.Expression.Literal: diff --git a/python/pyspark/ml/connect/summarizer.py b/python/pyspark/ml/connect/summarizer.py index 2f3e16ef4d188..92f3aaed438ee 100644 --- a/python/pyspark/ml/connect/summarizer.py +++ b/python/pyspark/ml/connect/summarizer.py @@ -15,13 +15,13 @@ # limitations under the License. # -from typing import Any, Union, List, Dict +from typing import Any, Dict, List, Union import numpy as np import pandas as pd -from pyspark.sql import DataFrame from pyspark.ml.connect.util import aggregate_dataframe +from pyspark.sql import DataFrame class SummarizerAggState: diff --git a/python/pyspark/ml/connect/tuning.py b/python/pyspark/ml/connect/tuning.py index 9d6a008c99e29..d31e0fdad1dfd 100644 --- a/python/pyspark/ml/connect/tuning.py +++ b/python/pyspark/ml/connect/tuning.py @@ -17,6 +17,7 @@ from multiprocessing.pool import ThreadPool from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -26,24 +27,22 @@ Tuple, Union, cast, - TYPE_CHECKING, ) import numpy as np import pandas as pd -from pyspark import keyword_only, since, inheritable_thread_target -from pyspark.ml.connect import Estimator, Model -from pyspark.ml.connect.base import Evaluator +from pyspark import inheritable_thread_target, keyword_only, since +from pyspark.ml.connect.base import Estimator, Evaluator, Model from pyspark.ml.connect.io_utils import ( MetaAlgorithmReadWrite, ParamsReadWrite, ) -from pyspark.ml.param import Params, Param, TypeConverters +from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import HasParallelism, HasSeed -from pyspark.sql.functions import col, lit, rand -from pyspark.sql.dataframe import DataFrame from pyspark.sql import SparkSession +from pyspark.sql.dataframe import DataFrame +from pyspark.sql.functions import col, lit, rand from pyspark.sql.utils import is_remote if TYPE_CHECKING: diff --git a/python/pyspark/ml/connect/util.py b/python/pyspark/ml/connect/util.py index 17b072ff57eeb..4f22ba4948206 100644 --- a/python/pyspark/ml/connect/util.py +++ b/python/pyspark/ml/connect/util.py @@ -15,7 +15,7 @@ # limitations under the License. # -from typing import Any, TypeVar, Callable, List, Tuple, Union, Iterator, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Tuple, TypeVar, Union import pandas as pd diff --git a/python/pyspark/ml/deepspeed/deepspeed_distributor.py b/python/pyspark/ml/deepspeed/deepspeed_distributor.py index 162cd2adda523..cbfbf3702027b 100644 --- a/python/pyspark/ml/deepspeed/deepspeed_distributor.py +++ b/python/pyspark/ml/deepspeed/deepspeed_distributor.py @@ -18,12 +18,12 @@ import sys import tempfile from typing import ( - Union, + Any, Callable, - List, Dict, + List, Optional, - Any, + Union, ) from pyspark.ml.torch.distributor import TorchDistributor diff --git a/python/pyspark/ml/evaluation.py b/python/pyspark/ml/evaluation.py index 801d23db4eefd..4b1ed21c3f72b 100644 --- a/python/pyspark/ml/evaluation.py +++ b/python/pyspark/ml/evaluation.py @@ -16,32 +16,32 @@ # import sys -from abc import abstractmethod, ABCMeta -from typing import Any, Dict, Optional, TYPE_CHECKING +from abc import ABCMeta, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, Optional -from pyspark import since, keyword_only -from pyspark.ml.wrapper import JavaParams +from pyspark import keyword_only, since +from pyspark.ml.common import inherit_doc from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import ( + HasFeaturesCol, HasLabelCol, HasPredictionCol, HasProbabilityCol, HasRawPredictionCol, - HasFeaturesCol, HasWeightCol, ) -from pyspark.ml.common import inherit_doc from pyspark.ml.util import JavaMLReadable, JavaMLWritable, try_remote_evaluate +from pyspark.ml.wrapper import JavaParams from pyspark.sql.dataframe import DataFrame if TYPE_CHECKING: from pyspark.ml._typing import ( - ParamMap, BinaryClassificationEvaluatorMetricType, ClusteringEvaluatorDistanceMeasureType, ClusteringEvaluatorMetricType, MulticlassClassificationEvaluatorMetricType, MultilabelClassificationEvaluatorMetricType, + ParamMap, RankingEvaluatorMetricType, RegressionEvaluatorMetricType, ) @@ -1171,6 +1171,7 @@ def isLargerBetter(self) -> bool: if __name__ == "__main__": import doctest import tempfile + import pyspark.ml.evaluation from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py index 4024c0c07f4f1..4167249afaa87 100755 --- a/python/pyspark/ml/feature.py +++ b/python/pyspark/ml/feature.py @@ -15,8 +15,7 @@ # limitations under the License. # from typing import ( - cast, - overload, + TYPE_CHECKING, Any, Dict, Generic, @@ -25,36 +24,38 @@ Tuple, TypeVar, Union, - TYPE_CHECKING, + cast, + overload, ) from pyspark import keyword_only, since -from pyspark.ml.linalg import _convert_to_vector, DenseMatrix, DenseVector, Vector -from pyspark.sql.dataframe import DataFrame +from pyspark.ml.common import inherit_doc +from pyspark.ml.linalg import DenseMatrix, DenseVector, Vector, _convert_to_vector from pyspark.ml.param.shared import ( - HasThreshold, - HasThresholds, + HasFeaturesCol, + HasHandleInvalid, HasInputCol, - HasOutputCol, HasInputCols, + HasLabelCol, + HasMaxIter, + HasNumFeatures, + HasOutputCol, HasOutputCols, - HasHandleInvalid, HasRelativeError, - HasFeaturesCol, - HasLabelCol, HasSeed, - HasNumFeatures, HasStepSize, - HasMaxIter, - TypeConverters, + HasThreshold, + HasThresholds, Param, Params, + TypeConverters, ) from pyspark.ml.util import ( JavaMLReadable, JavaMLWritable, - try_remote_attribute_relation, + RemoteModelRef, invoke_helper_attr, + try_remote_attribute_relation, ) from pyspark.ml.wrapper import ( JavaEstimator, @@ -63,8 +64,7 @@ JavaTransformer, _jvm, ) -from pyspark.ml.common import inherit_doc -from pyspark.ml.util import RemoteModelRef +from pyspark.sql.dataframe import DataFrame from pyspark.sql.types import ArrayType, StringType from pyspark.sql.utils import is_remote diff --git a/python/pyspark/ml/fpm.py b/python/pyspark/ml/fpm.py index 7ad28f69b7c70..fc250db37fe04 100644 --- a/python/pyspark/ml/fpm.py +++ b/python/pyspark/ml/fpm.py @@ -16,18 +16,18 @@ # import sys -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, Optional from pyspark import keyword_only, since -from pyspark.sql import DataFrame +from pyspark.ml.param.shared import HasPredictionCol, Param, Params, TypeConverters from pyspark.ml.util import ( - JavaMLWritable, JavaMLReadable, - try_remote_attribute_relation, + JavaMLWritable, invoke_helper_relation, + try_remote_attribute_relation, ) from pyspark.ml.wrapper import JavaEstimator, JavaModel, JavaParams -from pyspark.ml.param.shared import HasPredictionCol, Param, TypeConverters, Params +from pyspark.sql import DataFrame if TYPE_CHECKING: from py4j.java_gateway import JavaObject @@ -533,6 +533,7 @@ def findFrequentSequentialPatterns(self, dataset: DataFrame) -> DataFrame: if __name__ == "__main__": import doctest + import pyspark.ml.fpm from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/functions.py b/python/pyspark/ml/functions.py index b3cd3978df445..51a735f0862f9 100644 --- a/python/pyspark/ml/functions.py +++ b/python/pyspark/ml/functions.py @@ -18,7 +18,7 @@ import inspect import uuid -from typing import Any, Callable, Iterator, List, Mapping, TYPE_CHECKING, Tuple, Union, Optional +from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Mapping, Optional, Tuple, Union import numpy as np @@ -27,8 +27,9 @@ except ImportError: pass # Let it throw a better error message later when the API is invoked. -from pyspark.sql.functions import pandas_udf +from pyspark.ml.util import try_remote_functions from pyspark.sql.column import Column +from pyspark.sql.functions import pandas_udf from pyspark.sql.types import ( ArrayType, ByteType, @@ -41,7 +42,6 @@ StringType, StructType, ) -from pyspark.ml.util import try_remote_functions if TYPE_CHECKING: from pyspark.sql._typing import UserDefinedFunctionLike @@ -833,10 +833,10 @@ def predict(data: Iterator[Union[pd.Series, pd.DataFrame]]) -> Iterator[pd.DataF def _test() -> None: import doctest - from pyspark.sql import SparkSession - import pyspark.ml.functions import sys + import pyspark.ml.functions + from pyspark.sql import SparkSession from pyspark.sql.pandas.utils import ( require_minimum_pandas_version, require_minimum_pyarrow_version, diff --git a/python/pyspark/ml/image.py b/python/pyspark/ml/image.py index d4f59d2d02738..713af8b58b3af 100644 --- a/python/pyspark/ml/image.py +++ b/python/pyspark/ml/image.py @@ -25,13 +25,13 @@ """ import sys -from typing import Any, Dict, List, NoReturn, cast from functools import cached_property +from typing import Any, Dict, List, NoReturn, cast import numpy as np -from pyspark.sql.types import Row, StructType, _create_row, _parse_datatype_json_string from pyspark.sql import SparkSession +from pyspark.sql.types import Row, StructType, _create_row, _parse_datatype_json_string __all__ = ["ImageSchema"] @@ -231,6 +231,7 @@ def _disallow_instance(_: Any) -> NoReturn: def _test() -> None: import doctest + import pyspark.ml.image globs = pyspark.ml.image.__dict__.copy() diff --git a/python/pyspark/ml/linalg/__init__.py b/python/pyspark/ml/linalg/__init__.py index a5bf8318c43b4..a1d7278bd2e62 100644 --- a/python/pyspark/ml/linalg/__init__.py +++ b/python/pyspark/ml/linalg/__init__.py @@ -23,36 +23,37 @@ SciPy is available in their environment. """ -import sys import array import struct +import sys from typing import ( + TYPE_CHECKING, Any, Callable, - cast, Dict, Iterable, List, Optional, - overload, Sequence, Tuple, Type, - TYPE_CHECKING, Union, + cast, + overload, ) import numpy as np from pyspark.sql.types import ( - UserDefinedType, - StructField, - StructType, ArrayType, + BooleanType, + ByteType, + DataTypeSingleton, DoubleType, IntegerType, - ByteType, - BooleanType, + StructField, + StructType, + UserDefinedType, ) __all__ = [ @@ -67,8 +68,8 @@ ] if TYPE_CHECKING: - from pyspark.mllib._typing import NormType from pyspark.ml._typing import VectorLike + from pyspark.mllib._typing import NormType # Check whether we have SciPy. MLlib works without it too, but if we have it, some methods, @@ -157,7 +158,7 @@ def _double_to_long_bits(value: float) -> int: return struct.unpack("Q", struct.pack("d", value))[0] -class VectorUDT(UserDefinedType): +class VectorUDT(UserDefinedType, metaclass=DataTypeSingleton): """ SQL user-defined type (UDT) for Vector. """ @@ -212,7 +213,7 @@ def simpleString(self) -> str: return "vector" -class MatrixUDT(UserDefinedType): +class MatrixUDT(UserDefinedType, metaclass=DataTypeSingleton): """ SQL user-defined type (UDT) for Matrix. """ diff --git a/python/pyspark/ml/param/__init__.py b/python/pyspark/ml/param/__init__.py index cf0da7dce8279..db4f3a5bbb17d 100644 --- a/python/pyspark/ml/param/__init__.py +++ b/python/pyspark/ml/param/__init__.py @@ -15,25 +15,25 @@ # limitations under the License. # import array -from abc import ABCMeta import copy +from abc import ABCMeta from typing import ( + TYPE_CHECKING, Any, Callable, Generic, List, Optional, - overload, TypeVar, Union, - TYPE_CHECKING, + overload, ) import numpy as np -from pyspark.util import is_remote_only -from pyspark.ml.linalg import DenseVector, Vector, Matrix +from pyspark.ml.linalg import DenseVector, Matrix, Vector from pyspark.ml.util import Identifiable +from pyspark.util import is_remote_only if TYPE_CHECKING: from pyspark.ml._typing import ParamMap diff --git a/python/pyspark/ml/pipeline.py b/python/pyspark/ml/pipeline.py index f5f0d12bc8363..e4facbcc7d5d4 100644 --- a/python/pyspark/ml/pipeline.py +++ b/python/pyspark/ml/pipeline.py @@ -15,33 +15,33 @@ # limitations under the License. # import os - -from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast from pyspark import keyword_only, since from pyspark.ml.base import Estimator, Model, Transformer +from pyspark.ml.common import inherit_doc from pyspark.ml.param import Param, Params from pyspark.ml.util import ( - MLReadable, - MLWritable, - JavaMLWriter, DefaultParamsReader, DefaultParamsWriter, - MLWriter, - MLReader, JavaMLWritable, + JavaMLWriter, + MLReadable, + MLReader, + MLWritable, + MLWriter, try_remote_read, try_remote_write, ) from pyspark.ml.wrapper import JavaParams -from pyspark.ml.common import inherit_doc from pyspark.sql import SparkSession from pyspark.sql.dataframe import DataFrame if TYPE_CHECKING: - from pyspark.ml._typing import ParamMap, PipelineStage from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext + from pyspark.ml._typing import ParamMap, PipelineStage @inherit_doc diff --git a/python/pyspark/ml/recommendation.py b/python/pyspark/ml/recommendation.py index ee4a0c8d2e968..273b37f83ab1c 100644 --- a/python/pyspark/ml/recommendation.py +++ b/python/pyspark/ml/recommendation.py @@ -16,22 +16,22 @@ # import sys -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, Optional -from pyspark import since, keyword_only +from pyspark import keyword_only, since +from pyspark.ml.common import inherit_doc +from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import ( - HasPredictionCol, HasBlockSize, + HasCheckpointInterval, + HasIntermediateStorageLevel, HasMaxIter, + HasPredictionCol, HasRegParam, - HasCheckpointInterval, HasSeed, - HasIntermediateStorageLevel, ) +from pyspark.ml.util import JavaMLReadable, JavaMLWritable, try_remote_attribute_relation from pyspark.ml.wrapper import JavaEstimator, JavaModel -from pyspark.ml.common import inherit_doc -from pyspark.ml.param import Params, TypeConverters, Param -from pyspark.ml.util import JavaMLWritable, JavaMLReadable, try_remote_attribute_relation from pyspark.sql import DataFrame if TYPE_CHECKING: @@ -720,6 +720,7 @@ def recommendForItemSubset(self, dataset: DataFrame, numUsers: int) -> DataFrame if __name__ == "__main__": import doctest + import pyspark.ml.recommendation from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py index d910db6c6d30b..f3f780aa360c0 100644 --- a/python/pyspark/ml/regression.py +++ b/python/pyspark/ml/regression.py @@ -16,61 +16,59 @@ # import sys -from typing import Any, Dict, Generic, List, Optional, TypeVar, TYPE_CHECKING from abc import ABCMeta from functools import cached_property +from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, TypeVar from pyspark import keyword_only, since -from pyspark.ml import Predictor, PredictionModel -from pyspark.ml.base import _PredictorParams +from pyspark.ml.base import PredictionModel, Predictor, Transformer, _PredictorParams +from pyspark.ml.common import inherit_doc +from pyspark.ml.linalg import Matrix, Vector from pyspark.ml.param.shared import ( + HasAggregationDepth, + HasElasticNetParam, HasFeaturesCol, - HasLabelCol, - HasPredictionCol, - HasWeightCol, - Param, - Params, - TypeConverters, - HasMaxIter, - HasTol, HasFitIntercept, - HasAggregationDepth, + HasLabelCol, + HasLoss, HasMaxBlockSizeInMB, + HasMaxIter, + HasPredictionCol, HasRegParam, - HasSolver, - HasStepSize, HasSeed, - HasElasticNetParam, + HasSolver, HasStandardization, - HasLoss, + HasStepSize, + HasTol, HasVarianceCol, + HasWeightCol, + Param, + Params, + TypeConverters, ) -from pyspark.ml.util import try_remote_attribute_relation from pyspark.ml.tree import ( _DecisionTreeModel, _DecisionTreeParams, - _TreeEnsembleModel, - _RandomForestParams, _GBTParams, + _RandomForestParams, + _TreeEnsembleModel, _TreeRegressorParams, ) -from pyspark.ml.base import Transformer -from pyspark.ml.linalg import Vector, Matrix from pyspark.ml.util import ( - JavaMLWritable, - JavaMLReadable, - HasTrainingSummary, GeneralJavaMLWritable, + HasTrainingSummary, + JavaMLReadable, + JavaMLWritable, + try_remote_attribute_relation, ) from pyspark.ml.wrapper import ( JavaEstimator, JavaModel, - JavaPredictor, JavaPredictionModel, + JavaPredictor, JavaTransformer, JavaWrapper, ) -from pyspark.ml.common import inherit_doc from pyspark.sql import DataFrame from pyspark.sql.utils import is_remote @@ -3314,6 +3312,7 @@ def factors(self) -> Matrix: if __name__ == "__main__": import doctest + import pyspark.ml.regression from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/stat.py b/python/pyspark/ml/stat.py index 99c4b80018cfb..50c0fcb456824 100644 --- a/python/pyspark/ml/stat.py +++ b/python/pyspark/ml/stat.py @@ -16,13 +16,13 @@ # import sys -from typing import Optional, Tuple, TYPE_CHECKING +from typing import TYPE_CHECKING, Optional, Tuple from pyspark import since from pyspark.ml.common import _java2py, _py2java from pyspark.ml.linalg import Matrix, Vector -from pyspark.ml.wrapper import JavaWrapper, _jvm from pyspark.ml.util import invoke_helper_relation +from pyspark.ml.wrapper import JavaWrapper, _jvm from pyspark.sql.column import Column from pyspark.sql.dataframe import DataFrame from pyspark.sql.functions import lit @@ -559,7 +559,9 @@ def __init__(self, mean: Vector, cov: Matrix): if __name__ == "__main__": import doctest + import numpy + import pyspark.ml.stat from pyspark.sql import SparkSession diff --git a/python/pyspark/ml/tests/connect/test_connect_cache.py b/python/pyspark/ml/tests/connect/test_connect_cache.py index ad3ee04e812e2..23c17bd2ecbda 100644 --- a/python/pyspark/ml/tests/connect/test_connect_cache.py +++ b/python/pyspark/ml/tests/connect/test_connect_cache.py @@ -17,8 +17,8 @@ import json -from pyspark.ml.linalg import Vectors from pyspark.ml.classification import LinearSVC +from pyspark.ml.linalg import Vectors from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/ml/tests/connect/test_connect_classification.py b/python/pyspark/ml/tests/connect/test_connect_classification.py index 7540c0e276a18..a173b49a3cff3 100644 --- a/python/pyspark/ml/tests/connect/test_connect_classification.py +++ b/python/pyspark/ml/tests/connect/test_connect_classification.py @@ -15,14 +15,17 @@ # limitations under the License. # -import unittest import os +import unittest -from pyspark.util import is_remote_only from pyspark.ml.tests.connect.test_legacy_mode_classification import ClassificationTestsMixin -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import ( + ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, +) from pyspark.testing.utils import have_torch, torch_requirement_message -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.util import is_remote_only @unittest.skipIf( diff --git a/python/pyspark/ml/tests/connect/test_connect_evaluation.py b/python/pyspark/ml/tests/connect/test_connect_evaluation.py index adf57ec842c21..693995cd8dbfe 100644 --- a/python/pyspark/ml/tests/connect/test_connect_evaluation.py +++ b/python/pyspark/ml/tests/connect/test_connect_evaluation.py @@ -18,8 +18,7 @@ import os import unittest -from pyspark.testing.connectutils import should_test_connect -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect if should_test_connect: from pyspark.ml.tests.connect.test_legacy_mode_evaluation import EvaluationTestsMixin diff --git a/python/pyspark/ml/tests/connect/test_connect_feature.py b/python/pyspark/ml/tests/connect/test_connect_feature.py index 0d0109d02c8d2..9ec9583a8caf5 100644 --- a/python/pyspark/ml/tests/connect/test_connect_feature.py +++ b/python/pyspark/ml/tests/connect/test_connect_feature.py @@ -18,9 +18,12 @@ import os import unittest -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import ( + ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, +) from pyspark.testing.utils import have_sklearn, sklearn_requirement_message -from pyspark.testing.connectutils import ReusedConnectTestCase if should_test_connect: from pyspark.ml.tests.connect.test_legacy_mode_feature import FeatureTestsMixin diff --git a/python/pyspark/ml/tests/connect/test_connect_function.py b/python/pyspark/ml/tests/connect/test_connect_function.py index a18aa619249c8..59df2bbfd13bf 100644 --- a/python/pyspark/ml/tests/connect/test_connect_function.py +++ b/python/pyspark/ml/tests/connect/test_connect_function.py @@ -16,13 +16,13 @@ # import unittest -from pyspark.util import is_remote_only from pyspark.ml import functions as SF from pyspark.testing.connectutils import ( - should_test_connect, ReusedMixedTestCase, + should_test_connect, ) from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.util import is_remote_only if should_test_connect: from pyspark.ml.connect import functions as CF diff --git a/python/pyspark/ml/tests/connect/test_connect_model_offloading.py b/python/pyspark/ml/tests/connect/test_connect_model_offloading.py index 9d78ab8e45da7..ddc0c0ef869c5 100644 --- a/python/pyspark/ml/tests/connect/test_connect_model_offloading.py +++ b/python/pyspark/ml/tests/connect/test_connect_model_offloading.py @@ -19,28 +19,28 @@ import numpy as np -from pyspark.sql import functions as sf -from pyspark.ml.linalg import Vectors from pyspark.ml.classification import ( LinearSVC, LinearSVCSummary, LinearSVCTrainingSummary, ) -from pyspark.ml.regression import ( - LinearRegression, - LinearRegressionSummary, - LinearRegressionTrainingSummary, -) from pyspark.ml.clustering import ( LDA, + DistributedLDAModel, LDAModel, LocalLDAModel, - DistributedLDAModel, ) from pyspark.ml.fpm import ( FPGrowth, FPGrowthModel, ) +from pyspark.ml.linalg import Vectors +from pyspark.ml.regression import ( + LinearRegression, + LinearRegressionSummary, + LinearRegressionTrainingSummary, +) +from pyspark.sql import functions as sf from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/ml/tests/connect/test_connect_pipeline.py b/python/pyspark/ml/tests/connect/test_connect_pipeline.py index dcdc1de62d29a..a98b6b3b6ff5e 100644 --- a/python/pyspark/ml/tests/connect/test_connect_pipeline.py +++ b/python/pyspark/ml/tests/connect/test_connect_pipeline.py @@ -18,10 +18,13 @@ import os import unittest -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import ( + ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, +) from pyspark.testing.utils import have_torch, torch_requirement_message -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.util import is_remote_only if should_test_connect: from pyspark.ml.tests.connect.test_legacy_mode_pipeline import PipelineTestsMixin diff --git a/python/pyspark/ml/tests/connect/test_connect_summarizer.py b/python/pyspark/ml/tests/connect/test_connect_summarizer.py index f68c18943bc48..2766a25bb0ed0 100644 --- a/python/pyspark/ml/tests/connect/test_connect_summarizer.py +++ b/python/pyspark/ml/tests/connect/test_connect_summarizer.py @@ -18,8 +18,11 @@ import os import unittest -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.connectutils import ( + ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, +) if should_test_connect: from pyspark.ml.tests.connect.test_legacy_mode_summarizer import SummarizerTestsMixin diff --git a/python/pyspark/ml/tests/connect/test_connect_tuning.py b/python/pyspark/ml/tests/connect/test_connect_tuning.py index 21ada54431127..3a25c4e7af9c6 100644 --- a/python/pyspark/ml/tests/connect/test_connect_tuning.py +++ b/python/pyspark/ml/tests/connect/test_connect_tuning.py @@ -18,10 +18,13 @@ import os import unittest -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import ( + ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, +) from pyspark.testing.utils import have_torch, torch_requirement_message -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.util import is_remote_only if should_test_connect: from pyspark.ml.tests.connect.test_legacy_mode_tuning import CrossValidatorTestsMixin diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py b/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py index 0028feb6cce72..854326101f8fd 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_classification.py @@ -20,17 +20,20 @@ import numpy as np -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.util import is_remote_only if should_test_connect: + import pandas as pd + from pyspark.ml.connect.classification import ( LogisticRegression as LORV2, + ) + from pyspark.ml.connect.classification import ( LogisticRegressionModel as LORV2Model, ) - import pandas as pd class ClassificationTestsMixin: diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py b/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py index cc7822c439300..7220adca15a7a 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py @@ -14,21 +14,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest import tempfile +import unittest import numpy as np -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.utils import have_torcheval, torcheval_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import have_torcheval, torcheval_requirement_message +from pyspark.util import is_remote_only if should_test_connect: from pyspark.ml.connect.evaluation import ( - RegressionEvaluator, BinaryClassificationEvaluator, MulticlassClassificationEvaluator, + RegressionEvaluator, ) diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py b/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py index de65858e2f513..5bf8b479ddef2 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_feature.py @@ -21,20 +21,21 @@ import numpy as np -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.util import is_remote_only if should_test_connect: + import pandas as pd + from pyspark.ml.connect.feature import ( + ArrayAssembler, MaxAbsScaler, MaxAbsScalerModel, StandardScaler, StandardScalerModel, - ArrayAssembler, ) - import pandas as pd class FeatureTestsMixin: diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py b/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py index bc85f98619a66..0cec1b6565d3f 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py @@ -20,16 +20,17 @@ import numpy as np -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import have_torch, torch_requirement_message +from pyspark.util import is_remote_only if should_test_connect: - from pyspark.ml.connect.feature import StandardScaler + import pandas as pd + from pyspark.ml.connect.classification import LogisticRegression as LORV2 + from pyspark.ml.connect.feature import StandardScaler from pyspark.ml.connect.pipeline import Pipeline - import pandas as pd class PipelineTestsMixin: diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py b/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py index 567f6d854d3bc..9eb03fa70ef2a 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py @@ -19,9 +19,9 @@ import numpy as np -from pyspark.util import is_remote_only -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.util import is_remote_only if should_test_connect: from pyspark.ml.connect.summarizer import summarize_dataframe diff --git a/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py b/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py index 4a819520963b6..2912b0d96ac4f 100644 --- a/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py +++ b/python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py @@ -15,35 +15,36 @@ # limitations under the License. # +import sys import tempfile import unittest -import sys import numpy as np -from pyspark.util import is_remote_only from pyspark.ml.param import Param, Params from pyspark.ml.tuning import ParamGridBuilder from pyspark.sql.functions import rand -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_sklearn, - sklearn_requirement_message, have_torch, - torch_requirement_message, have_torcheval, + sklearn_requirement_message, + torch_requirement_message, torcheval_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.util import is_remote_only if should_test_connect: import pandas as pd - from pyspark.ml.connect import Model, Estimator - from pyspark.ml.connect.feature import StandardScaler + + from pyspark.ml.connect import Estimator, Model from pyspark.ml.connect.classification import LogisticRegression as LORV2 + from pyspark.ml.connect.evaluation import BinaryClassificationEvaluator, RegressionEvaluator + from pyspark.ml.connect.feature import StandardScaler from pyspark.ml.connect.pipeline import Pipeline from pyspark.ml.connect.tuning import CrossValidator, CrossValidatorModel - from pyspark.ml.connect.evaluation import BinaryClassificationEvaluator, RegressionEvaluator class HasInducedError(Params): def __init__(self): diff --git a/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py b/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py index e7b96aab585d0..256c4319dc810 100644 --- a/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py +++ b/python/pyspark/ml/tests/connect/test_parity_torch_data_loader.py @@ -17,8 +17,8 @@ import unittest -from pyspark.util import is_remote_only from pyspark.sql import SparkSession +from pyspark.util import is_remote_only if not is_remote_only(): from pyspark.ml.torch.tests.test_data_loader import TorchDistributorDataLoaderUnitTests diff --git a/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py b/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py index ce0d64c66f886..3a6d271834539 100644 --- a/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py +++ b/python/pyspark/ml/tests/connect/test_parity_torch_distributor.py @@ -19,24 +19,24 @@ import shutil import unittest -from pyspark.util import is_remote_only from pyspark.sql import SparkSession from pyspark.testing.utils import ( + connect_requirement_message, have_torch, - torch_requirement_message, should_test_connect, - connect_requirement_message, + torch_requirement_message, ) +from pyspark.util import is_remote_only if not is_remote_only() and should_test_connect: from pyspark.ml.torch.tests.test_distributor import ( TorchDistributorBaselineUnitTestsMixin, - TorchDistributorLocalUnitTestsMixin, TorchDistributorDistributedUnitTestsMixin, + TorchDistributorLocalUnitTestsMixin, TorchWrapperUnitTestsMixin, - set_up_test_dirs, - get_local_mode_conf, get_distributed_mode_conf, + get_local_mode_conf, + set_up_test_dirs, ) @unittest.skipIf( diff --git a/python/pyspark/ml/tests/test_algorithms.py b/python/pyspark/ml/tests/test_algorithms.py index 8c9f18d627840..7d630f71e470e 100644 --- a/python/pyspark/ml/tests/test_algorithms.py +++ b/python/pyspark/ml/tests/test_algorithms.py @@ -15,8 +15,8 @@ # limitations under the License. # import os -from shutil import rmtree import tempfile +from shutil import rmtree import numpy as np @@ -26,9 +26,9 @@ MultilayerPerceptronClassifier, OneVsRest, ) -from pyspark.ml.clustering import DistributedLDAModel, KMeans, LocalLDAModel, LDA, LDAModel +from pyspark.ml.clustering import LDA, DistributedLDAModel, KMeans, LDAModel, LocalLDAModel from pyspark.ml.fpm import FPGrowth -from pyspark.ml.linalg import Vectors, DenseVector +from pyspark.ml.linalg import DenseVector, Vectors from pyspark.ml.recommendation import ALS from pyspark.ml.regression import GeneralizedLinearRegression, LinearRegression from pyspark.sql import Row diff --git a/python/pyspark/ml/tests/test_base.py b/python/pyspark/ml/tests/test_base.py index cf4cec224e246..7bf662dc3b7fb 100644 --- a/python/pyspark/ml/tests/test_base.py +++ b/python/pyspark/ml/tests/test_base.py @@ -21,8 +21,8 @@ from pyspark.testing.mlutils import ( MockDataset, MockEstimator, - MockUnaryTransformer, MockTransformer, + MockUnaryTransformer, SparkSessionTestCase, ) diff --git a/python/pyspark/ml/tests/test_classification.py b/python/pyspark/ml/tests/test_classification.py index c8785d0e9ce12..761ec650fac1d 100644 --- a/python/pyspark/ml/tests/test_classification.py +++ b/python/pyspark/ml/tests/test_classification.py @@ -21,11 +21,18 @@ import numpy as np from pyspark.errors import PySparkException -from pyspark.ml.linalg import Vectors, Matrices -from pyspark.sql import DataFrame, Row from pyspark.ml.classification import ( - NaiveBayes, - NaiveBayesModel, + BinaryLogisticRegressionSummary, + BinaryRandomForestClassificationSummary, + BinaryRandomForestClassificationTrainingSummary, + DecisionTreeClassificationModel, + DecisionTreeClassifier, + FMClassificationModel, + FMClassificationSummary, + FMClassificationTrainingSummary, + FMClassifier, + GBTClassificationModel, + GBTClassifier, LinearSVC, LinearSVCModel, LinearSVCSummary, @@ -33,27 +40,20 @@ LogisticRegression, LogisticRegressionModel, LogisticRegressionSummary, - BinaryLogisticRegressionSummary, - FMClassifier, - FMClassificationModel, - FMClassificationSummary, - FMClassificationTrainingSummary, - DecisionTreeClassifier, - DecisionTreeClassificationModel, - RandomForestClassifier, - RandomForestClassificationModel, - RandomForestClassificationSummary, - RandomForestClassificationTrainingSummary, - BinaryRandomForestClassificationSummary, - BinaryRandomForestClassificationTrainingSummary, - GBTClassifier, - GBTClassificationModel, - MultilayerPerceptronClassifier, MultilayerPerceptronClassificationModel, MultilayerPerceptronClassificationSummary, MultilayerPerceptronClassificationTrainingSummary, + MultilayerPerceptronClassifier, + NaiveBayes, + NaiveBayesModel, + RandomForestClassificationModel, + RandomForestClassificationSummary, + RandomForestClassificationTrainingSummary, + RandomForestClassifier, ) +from pyspark.ml.linalg import Matrices, Vectors from pyspark.ml.regression import DecisionTreeRegressionModel +from pyspark.sql import DataFrame, Row from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_clustering.py b/python/pyspark/ml/tests/test_clustering.py index e22e97a5e7f16..6b753d5910274 100644 --- a/python/pyspark/ml/tests/test_clustering.py +++ b/python/pyspark/ml/tests/test_clustering.py @@ -19,23 +19,23 @@ import numpy as np -from pyspark.ml.linalg import Vectors, SparseVector from pyspark.ml.clustering import ( - KMeans, - KMeansModel, - KMeansSummary, + LDA, BisectingKMeans, BisectingKMeansModel, BisectingKMeansSummary, + DistributedLDAModel, GaussianMixture, GaussianMixtureModel, GaussianMixtureSummary, - LDA, + KMeans, + KMeansModel, + KMeansSummary, LDAModel, LocalLDAModel, - DistributedLDAModel, PowerIterationClustering, ) +from pyspark.ml.linalg import SparseVector, Vectors from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_dl_util.py b/python/pyspark/ml/tests/test_dl_util.py index ba9dc31b204ec..7bff80dc606df 100644 --- a/python/pyspark/ml/tests/test_dl_util.py +++ b/python/pyspark/ml/tests/test_dl_util.py @@ -14,11 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from contextlib import contextmanager import os import textwrap -from typing import Any, BinaryIO, Callable, Iterator import unittest +from contextlib import contextmanager +from typing import Any, BinaryIO, Callable, Iterator from pyspark import cloudpickle from pyspark.ml.dl_util import FunctionPickler diff --git a/python/pyspark/ml/tests/test_evaluation.py b/python/pyspark/ml/tests/test_evaluation.py index 5c227ccb423e4..3c2947dbf30a9 100644 --- a/python/pyspark/ml/tests/test_evaluation.py +++ b/python/pyspark/ml/tests/test_evaluation.py @@ -19,12 +19,12 @@ import numpy as np from pyspark.ml.evaluation import ( - ClusteringEvaluator, - RegressionEvaluator, BinaryClassificationEvaluator, + ClusteringEvaluator, MulticlassClassificationEvaluator, MultilabelClassificationEvaluator, RankingEvaluator, + RegressionEvaluator, ) from pyspark.ml.linalg import Vectors from pyspark.sql import Row diff --git a/python/pyspark/ml/tests/test_feature.py b/python/pyspark/ml/tests/test_feature.py index 0047d08106993..13e36ce76c0ec 100644 --- a/python/pyspark/ml/tests/test_feature.py +++ b/python/pyspark/ml/tests/test_feature.py @@ -16,68 +16,68 @@ # import tempfile -from typing import List, Tuple, Any +from typing import Any, List, Tuple import numpy as np from pyspark.ml.feature import ( DCT, + IDF, + PCA, Binarizer, + BucketedRandomProjectionLSH, + BucketedRandomProjectionLSHModel, Bucketizer, - QuantileDiscretizer, + ChiSqSelector, + ChiSqSelectorModel, CountVectorizer, CountVectorizerModel, - OneHotEncoder, - OneHotEncoderModel, - FeatureHasher, ElementwiseProduct, + FeatureHasher, HashingTF, - IDF, IDFModel, Imputer, ImputerModel, - NGram, - Normalizer, + IndexToString, Interaction, - RFormula, - RFormulaModel, - Tokenizer, - SQLTransformer, - RegexTokenizer, - StandardScaler, - StandardScalerModel, MaxAbsScaler, MaxAbsScalerModel, + MinHashLSH, + MinHashLSHModel, MinMaxScaler, MinMaxScalerModel, + NGram, + Normalizer, + OneHotEncoder, + OneHotEncoderModel, + PCAModel, + PolynomialExpansion, + QuantileDiscretizer, + RegexTokenizer, + RFormula, + RFormulaModel, RobustScaler, RobustScalerModel, - ChiSqSelector, - ChiSqSelectorModel, + SQLTransformer, + StandardScaler, + StandardScalerModel, + StopWordsRemover, + StringIndexer, + StringIndexerModel, + TargetEncoder, + TargetEncoderModel, + Tokenizer, UnivariateFeatureSelector, UnivariateFeatureSelectorModel, VarianceThresholdSelector, VarianceThresholdSelectorModel, - StopWordsRemover, - StringIndexer, - StringIndexerModel, + VectorAssembler, VectorIndexer, VectorIndexerModel, - TargetEncoder, - TargetEncoderModel, VectorSizeHint, VectorSlicer, - VectorAssembler, - PCA, - PCAModel, Word2Vec, Word2VecModel, - BucketedRandomProjectionLSH, - BucketedRandomProjectionLSHModel, - MinHashLSH, - MinHashLSHModel, - IndexToString, - PolynomialExpansion, ) from pyspark.ml.linalg import DenseVector, SparseVector, Vectors from pyspark.sql import Row diff --git a/python/pyspark/ml/tests/test_fpm.py b/python/pyspark/ml/tests/test_fpm.py index b54f55738111a..5d486042a0062 100644 --- a/python/pyspark/ml/tests/test_fpm.py +++ b/python/pyspark/ml/tests/test_fpm.py @@ -17,13 +17,13 @@ import tempfile -from pyspark.sql import Row import pyspark.sql.functions as sf from pyspark.ml.fpm import ( FPGrowth, FPGrowthModel, PrefixSpan, ) +from pyspark.sql import Row from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_functions.py b/python/pyspark/ml/tests/test_functions.py index 6e1d25a0cfda6..82eb5bc02b03f 100644 --- a/python/pyspark/ml/tests/test_functions.py +++ b/python/pyspark/ml/tests/test_functions.py @@ -19,10 +19,10 @@ import numpy as np +from pyspark.ml.functions import array_to_vector, predict_batch_udf, vector_to_array from pyspark.ml.linalg import DenseVector -from pyspark.ml.functions import array_to_vector, vector_to_array, predict_batch_udf -from pyspark.sql.functions import array, struct, col -from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StructType, StructField, FloatType +from pyspark.sql.functions import array, col, struct +from pyspark.sql.types import ArrayType, DoubleType, FloatType, IntegerType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, diff --git a/python/pyspark/ml/tests/test_image.py b/python/pyspark/ml/tests/test_image.py index 7cacde935400e..8b98ad6f5bc89 100644 --- a/python/pyspark/ml/tests/test_image.py +++ b/python/pyspark/ml/tests/test_image.py @@ -16,8 +16,8 @@ # from pyspark.ml.image import ImageSchema -from pyspark.testing.mlutils import SparkSessionTestCase from pyspark.sql import Row +from pyspark.testing.mlutils import SparkSessionTestCase from pyspark.testing.utils import QuietTest, eventually diff --git a/python/pyspark/ml/tests/test_linalg.py b/python/pyspark/ml/tests/test_linalg.py index 4d5752cecd673..5b6677088ab77 100644 --- a/python/pyspark/ml/tests/test_linalg.py +++ b/python/pyspark/ml/tests/test_linalg.py @@ -19,7 +19,7 @@ from numpy import arange, array, array_equal, inf, ones, tile, zeros -from pyspark.serializers import CPickleSerializer +from pyspark.errors import AnalysisException from pyspark.ml.linalg import ( DenseMatrix, DenseVector, @@ -27,12 +27,14 @@ SparseMatrix, SparseVector, Vector, - VectorUDT, Vectors, + VectorUDT, ) -from pyspark.testing.mllibutils import MLlibTestCase +from pyspark.serializers import CPickleSerializer from pyspark.sql import Row -from pyspark.sql.functions import unwrap_udt +from pyspark.sql.functions import lit, unwrap_udt, wrap_udt +from pyspark.sql.types import StructField, StructType +from pyspark.testing.mllibutils import MLlibTestCase class VectorTests(MLlibTestCase): @@ -328,6 +330,10 @@ class VectorUDTTests(MLlibTestCase): def test_json_schema(self): self.assertEqual(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_singleton(self): + self.assertIs(VectorUDT(), VectorUDT()) + self.assertIs(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_serialization(self): for v in [self.dv0, self.dv1, self.sv0, self.sv1]: self.assertEqual(v, self.udt.deserialize(self.udt.serialize(v))) @@ -363,6 +369,49 @@ def test_unwrap_udt(self): ] self.assertEqual(results, expected) + def test_wrap_udt(self): + schema = StructType([StructField("vec", VectorUDT.sqlType(), True)]) + vector_struct = Row("type", "size", "indices", "values") + df = self.spark.createDataFrame( + [ + (vector_struct(1, None, None, [1.0, 2.0]),), + (vector_struct(0, 2, [1], [2.0]),), + ], + schema, + ) + wrapped = df.select(wrap_udt("vec", VectorUDT()).alias("vec")) + + self.assertEqual(wrapped.schema["vec"].dataType, VectorUDT()) + self.assertEqual(wrapped.collect(), [Row(vec=self.dv1), Row(vec=self.sv1)]) + + def test_wrap_udt_type_mismatch(self): + schema = StructType([StructField("vec", VectorUDT.sqlType(), True)]) + vector_struct = Row("type", "size", "indices", "values") + df = self.spark.createDataFrame( + [(vector_struct(1, None, None, [1.0, 2.0]),)], + schema, + ) + + with self.assertRaises(AnalysisException) as context: + df.select(wrap_udt("vec", MatrixUDT())).collect() + self.assertEqual( + context.exception.getCondition(), "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + ) + + def test_wrap_unwrap_udt_round_trip(self): + df = self.spark.createDataFrame([(self.dv1,), (self.sv1,)], ["vec"]) + round_trip = df.select(wrap_udt(unwrap_udt("vec"), VectorUDT()).alias("vec")) + + self.assertEqual(round_trip.schema["vec"].dataType, VectorUDT()) + self.assertEqual(round_trip.collect(), [Row(vec=self.dv1), Row(vec=self.sv1)]) + + def test_wrap_unwrap_udt_round_trip_with_udt_column(self): + df = self.spark.createDataFrame([(self.dv1,), (self.sv1,)], ["vec"]) + round_trip = df.select(wrap_udt(unwrap_udt("vec"), lit(VectorUDT().json())).alias("vec")) + + self.assertEqual(round_trip.schema["vec"].dataType, VectorUDT()) + self.assertEqual(round_trip.collect(), [Row(vec=self.dv1), Row(vec=self.sv1)]) + def test_hashable(self): _ = hash(VectorUDT()) @@ -377,6 +426,10 @@ class MatrixUDTTests(MLlibTestCase): def test_json_schema(self): self.assertEqual(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_singleton(self): + self.assertIs(MatrixUDT(), MatrixUDT()) + self.assertIs(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_serialization(self): for m in [self.dm1, self.dm2, self.sm1, self.sm2]: self.assertEqual(m, self.udt.deserialize(self.udt.serialize(m))) @@ -396,6 +449,65 @@ def test_infer_schema(self): else: raise ValueError("Expected a matrix but got type %r" % type(m)) + def test_wrap_udt(self): + schema = StructType([StructField("mat", MatrixUDT.sqlType(), True)]) + matrix_struct = Row( + "type", + "numRows", + "numCols", + "colPtrs", + "rowIndices", + "values", + "isTransposed", + ) + df = self.spark.createDataFrame( + [ + (matrix_struct(1, 3, 2, None, None, [0.0, 1.0, 4.0, 5.0, 9.0, 10.0], False),), + (matrix_struct(0, 1, 1, [0, 1], [0], [2.0], False),), + ], + schema, + ) + wrapped = df.select(wrap_udt("mat", MatrixUDT()).alias("mat")) + + self.assertEqual(wrapped.schema["mat"].dataType, MatrixUDT()) + self.assertEqual(wrapped.collect(), [Row(mat=self.dm1), Row(mat=self.sm1)]) + + def test_wrap_udt_type_mismatch(self): + schema = StructType([StructField("mat", MatrixUDT.sqlType(), True)]) + matrix_struct = Row( + "type", + "numRows", + "numCols", + "colPtrs", + "rowIndices", + "values", + "isTransposed", + ) + df = self.spark.createDataFrame( + [(matrix_struct(1, 3, 2, None, None, [0.0, 1.0, 4.0, 5.0, 9.0, 10.0], False),)], + schema, + ) + + with self.assertRaises(AnalysisException) as context: + df.select(wrap_udt("mat", VectorUDT())).collect() + self.assertEqual( + context.exception.getCondition(), "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE" + ) + + def test_wrap_unwrap_udt_round_trip(self): + df = self.spark.createDataFrame([(self.dm1,), (self.sm1,)], ["mat"]) + round_trip = df.select(wrap_udt(unwrap_udt("mat"), MatrixUDT()).alias("mat")) + + self.assertEqual(round_trip.schema["mat"].dataType, MatrixUDT()) + self.assertEqual(round_trip.collect(), [Row(mat=self.dm1), Row(mat=self.sm1)]) + + def test_wrap_unwrap_udt_round_trip_with_udt_column(self): + df = self.spark.createDataFrame([(self.dm1,), (self.sm1,)], ["mat"]) + round_trip = df.select(wrap_udt(unwrap_udt("mat"), lit(MatrixUDT().json())).alias("mat")) + + self.assertEqual(round_trip.schema["mat"].dataType, MatrixUDT()) + self.assertEqual(round_trip.collect(), [Row(mat=self.dm1), Row(mat=self.sm1)]) + def test_hashable(self): _ = hash(MatrixUDT()) diff --git a/python/pyspark/ml/tests/test_ovr.py b/python/pyspark/ml/tests/test_ovr.py index d68ea1d9737ba..3b10bfac42c16 100644 --- a/python/pyspark/ml/tests/test_ovr.py +++ b/python/pyspark/ml/tests/test_ovr.py @@ -19,13 +19,13 @@ import numpy as np -from pyspark.ml.linalg import Vectors from pyspark.ml.classification import ( LinearSVC, LinearSVCModel, OneVsRest, OneVsRestModel, ) +from pyspark.ml.linalg import Vectors from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_param.py b/python/pyspark/ml/tests/test_param.py index fc4fa07863e26..f71048edd6593 100644 --- a/python/pyspark/ml/tests/test_param.py +++ b/python/pyspark/ml/tests/test_param.py @@ -15,8 +15,8 @@ # limitations under the License. # -import inspect import array as pyarray +import inspect import numpy as np @@ -35,9 +35,9 @@ from pyspark.ml.linalg import DenseVector, SparseVector, Vectors from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import HasInputCol, HasMaxIter, HasSeed -from pyspark.ml.regression import LinearRegressionModel, GeneralizedLinearRegressionModel +from pyspark.ml.regression import GeneralizedLinearRegressionModel, LinearRegressionModel from pyspark.ml.wrapper import JavaParams -from pyspark.testing.mlutils import check_params, PySparkTestCase, SparkSessionTestCase +from pyspark.testing.mlutils import PySparkTestCase, SparkSessionTestCase, check_params class ParamTypeConversionTests(PySparkTestCase): @@ -386,10 +386,10 @@ class DefaultValuesTests(PySparkTestCase): def test_java_params(self): import re - import pyspark.ml.feature import pyspark.ml.classification import pyspark.ml.clustering import pyspark.ml.evaluation + import pyspark.ml.feature import pyspark.ml.pipeline import pyspark.ml.recommendation import pyspark.ml.regression diff --git a/python/pyspark/ml/tests/test_persistence.py b/python/pyspark/ml/tests/test_persistence.py index ee97ae164a97c..76eea039a3f83 100644 --- a/python/pyspark/ml/tests/test_persistence.py +++ b/python/pyspark/ml/tests/test_persistence.py @@ -16,22 +16,22 @@ # import json -from shutil import rmtree import tempfile +from shutil import rmtree from pyspark.ml import Transformer from pyspark.ml.classification import ( DecisionTreeClassifier, - FMClassifier, FMClassificationModel, + FMClassifier, LogisticRegression, - MultilayerPerceptronClassifier, MultilayerPerceptronClassificationModel, + MultilayerPerceptronClassifier, OneVsRest, OneVsRestModel, ) from pyspark.ml.clustering import KMeans -from pyspark.ml.feature import Binarizer, Bucketizer, HashingTF, PCA +from pyspark.ml.feature import PCA, Binarizer, Bucketizer, HashingTF from pyspark.ml.linalg import Vectors from pyspark.ml.param import Params from pyspark.ml.pipeline import Pipeline, PipelineModel diff --git a/python/pyspark/ml/tests/test_pipeline.py b/python/pyspark/ml/tests/test_pipeline.py index ce104b563772d..6e72624115b06 100644 --- a/python/pyspark/ml/tests/test_pipeline.py +++ b/python/pyspark/ml/tests/test_pipeline.py @@ -18,18 +18,18 @@ import tempfile import unittest -from pyspark.sql import Row -from pyspark.ml.pipeline import Pipeline, PipelineModel +from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel +from pyspark.ml.clustering import GaussianMixture, KMeans, KMeansModel from pyspark.ml.feature import ( - VectorAssembler, MaxAbsScaler, MaxAbsScalerModel, MinMaxScaler, MinMaxScalerModel, + VectorAssembler, ) from pyspark.ml.linalg import Vectors -from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel -from pyspark.ml.clustering import KMeans, KMeansModel, GaussianMixture +from pyspark.ml.pipeline import Pipeline, PipelineModel +from pyspark.sql import Row from pyspark.testing.mlutils import MockDataset, MockEstimator, MockTransformer from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_regression.py b/python/pyspark/ml/tests/test_regression.py index 1427f2a450050..929ded5d27619 100644 --- a/python/pyspark/ml/tests/test_regression.py +++ b/python/pyspark/ml/tests/test_regression.py @@ -23,24 +23,24 @@ from pyspark.ml.regression import ( AFTSurvivalRegression, AFTSurvivalRegressionModel, - IsotonicRegression, - IsotonicRegressionModel, - LinearRegression, - LinearRegressionModel, + DecisionTreeRegressionModel, + DecisionTreeRegressor, + FMRegressionModel, + FMRegressor, + GBTRegressionModel, + GBTRegressor, GeneralizedLinearRegression, GeneralizedLinearRegressionModel, GeneralizedLinearRegressionSummary, GeneralizedLinearRegressionTrainingSummary, + IsotonicRegression, + IsotonicRegressionModel, + LinearRegression, + LinearRegressionModel, LinearRegressionSummary, LinearRegressionTrainingSummary, - FMRegressor, - FMRegressionModel, - DecisionTreeRegressor, - DecisionTreeRegressionModel, - RandomForestRegressor, RandomForestRegressionModel, - GBTRegressor, - GBTRegressionModel, + RandomForestRegressor, ) from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_stat.py b/python/pyspark/ml/tests/test_stat.py index 6f3b82e1e12c4..60e7f152563b8 100644 --- a/python/pyspark/ml/tests/test_stat.py +++ b/python/pyspark/ml/tests/test_stat.py @@ -18,7 +18,7 @@ import numpy as np from pyspark.errors import PySparkException -from pyspark.ml.linalg import Vectors, DenseVector +from pyspark.ml.linalg import DenseVector, Vectors from pyspark.ml.stat import ( ChiSquareTest, Correlation, @@ -26,8 +26,8 @@ Summarizer, SummaryBuilder, ) -from pyspark.sql import functions as F from pyspark.sql import DataFrame, Row +from pyspark.sql import functions as F from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/ml/tests/test_training_summary.py b/python/pyspark/ml/tests/test_training_summary.py index ed43d60986b0e..33fdb7f735b4e 100644 --- a/python/pyspark/ml/tests/test_training_summary.py +++ b/python/pyspark/ml/tests/test_training_summary.py @@ -18,12 +18,12 @@ from pyspark.ml.classification import ( BinaryRandomForestClassificationSummary, - FMClassifier, FMClassificationSummary, + FMClassifier, LinearSVC, LinearSVCSummary, - MultilayerPerceptronClassifier, MultilayerPerceptronClassificationSummary, + MultilayerPerceptronClassifier, RandomForestClassificationSummary, RandomForestClassifier, ) diff --git a/python/pyspark/ml/tests/test_tuning.py b/python/pyspark/ml/tests/test_tuning.py index 72380f4d68d21..6e2a03a3049fa 100644 --- a/python/pyspark/ml/tests/test_tuning.py +++ b/python/pyspark/ml/tests/test_tuning.py @@ -19,13 +19,13 @@ import numpy as np +from pyspark.ml.classification import LogisticRegression, RandomForestClassifier from pyspark.ml.evaluation import BinaryClassificationEvaluator from pyspark.ml.linalg import Vectors -from pyspark.ml.classification import LogisticRegression, RandomForestClassifier from pyspark.ml.tuning import ( - ParamGridBuilder, CrossValidator, CrossValidatorModel, + ParamGridBuilder, TrainValidationSplit, TrainValidationSplitModel, ) diff --git a/python/pyspark/ml/tests/test_wrapper.py b/python/pyspark/ml/tests/test_wrapper.py index 4a58336bce81c..1e5d04769f000 100644 --- a/python/pyspark/ml/tests/test_wrapper.py +++ b/python/pyspark/ml/tests/test_wrapper.py @@ -21,10 +21,10 @@ from pyspark.ml.linalg import DenseVector, Vectors from pyspark.ml.regression import LinearRegression from pyspark.ml.wrapper import ( - _java2py, - _py2java, JavaParams, JavaWrapper, + _java2py, + _py2java, ) from pyspark.testing.mllibutils import MLlibTestCase from pyspark.testing.mlutils import SparkSessionTestCase diff --git a/python/pyspark/ml/tests/tuning/test_cv_io_basic.py b/python/pyspark/ml/tests/tuning/test_cv_io_basic.py index cd65b3fbb4f7b..891f1fc1765b0 100644 --- a/python/pyspark/ml/tests/tuning/test_cv_io_basic.py +++ b/python/pyspark/ml/tests/tuning/test_cv_io_basic.py @@ -20,6 +20,7 @@ from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel from pyspark.ml.evaluation import BinaryClassificationEvaluator from pyspark.ml.linalg import Vectors +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( CrossValidator, CrossValidatorModel, @@ -31,7 +32,6 @@ DummyLogisticRegressionModel, SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin class CrossValidatorIOBasicTests(SparkSessionTestCase, ValidatorTestUtilsMixin): diff --git a/python/pyspark/ml/tests/tuning/test_cv_io_nested.py b/python/pyspark/ml/tests/tuning/test_cv_io_nested.py index 5f25e6e23ccc4..efcd124bb526d 100644 --- a/python/pyspark/ml/tests/tuning/test_cv_io_nested.py +++ b/python/pyspark/ml/tests/tuning/test_cv_io_nested.py @@ -20,6 +20,7 @@ from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.evaluation import MulticlassClassificationEvaluator from pyspark.ml.linalg import Vectors +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( CrossValidator, CrossValidatorModel, @@ -29,7 +30,6 @@ DummyLogisticRegression, SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin class CrossValidatorIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): diff --git a/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py b/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py index a4eefc98f6b4c..96cb0c4823658 100644 --- a/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py +++ b/python/pyspark/ml/tests/tuning/test_cv_io_pipeline.py @@ -16,12 +16,12 @@ # import tempfile -from concurrent.futures import ThreadPoolExecutor -from pyspark.ml.feature import HashingTF, Tokenizer from pyspark.ml import Pipeline from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import HashingTF, Tokenizer +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( CrossValidator, CrossValidatorModel, @@ -31,7 +31,6 @@ DummyLogisticRegression, SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin class CrossValidatorIOPipelineTests(SparkSessionTestCase, ValidatorTestUtilsMixin): @@ -55,7 +54,7 @@ def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): tokenizer = Tokenizer(inputCol="text", outputCol="words") hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features") - ova = OneVsRest(classifier=LogisticRegressionCls(), parallelism=2) + ova = OneVsRest(classifier=LogisticRegressionCls()) lr1 = LogisticRegressionCls().setMaxIter(5) lr2 = LogisticRegressionCls().setMaxIter(10) @@ -73,7 +72,6 @@ def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): estimatorParamMaps=paramGrid, evaluator=MulticlassClassificationEvaluator(), numFolds=2, - parallelism=4, ) # use 3+ folds in practice cvPath = temp_path + "/cv" crossval.save(cvPath) @@ -102,7 +100,6 @@ def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): estimatorParamMaps=paramGrid, evaluator=MulticlassClassificationEvaluator(), numFolds=2, - parallelism=4, ) # use 3+ folds in practice cv2Path = temp_path + "/cv2" crossval2.save(cv2Path) @@ -129,13 +126,8 @@ def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): self.assertEqual(loadedStage.uid, originalStage.uid) def test_save_load_pipeline_estimator(self): - with ThreadPoolExecutor(max_workers=2) as executor: - list( - executor.map( - self._run_test_save_load_pipeline_estimator, - [LogisticRegression, DummyLogisticRegression], - ) - ) + self._run_test_save_load_pipeline_estimator(LogisticRegression) + self._run_test_save_load_pipeline_estimator(DummyLogisticRegression) if __name__ == "__main__": diff --git a/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py b/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py index b8ff98c12f330..044e6d38ee1b1 100644 --- a/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py +++ b/python/pyspark/ml/tests/tuning/test_tvs_io_basic.py @@ -20,6 +20,7 @@ from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel from pyspark.ml.evaluation import BinaryClassificationEvaluator from pyspark.ml.linalg import Vectors +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( ParamGridBuilder, TrainValidationSplit, @@ -31,7 +32,6 @@ DummyLogisticRegressionModel, SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin class TrainValidationSplitIOBasicTests(SparkSessionTestCase, ValidatorTestUtilsMixin): diff --git a/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py b/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py index 252afc344b4e7..7176784916a90 100644 --- a/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py +++ b/python/pyspark/ml/tests/tuning/test_tvs_io_nested.py @@ -20,6 +20,7 @@ from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.evaluation import MulticlassClassificationEvaluator from pyspark.ml.linalg import Vectors +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( ParamGridBuilder, TrainValidationSplit, @@ -29,7 +30,6 @@ DummyLogisticRegression, SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin class TrainValidationSplitIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): diff --git a/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py b/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py index 029b6990a3108..c13de85c5f891 100644 --- a/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py +++ b/python/pyspark/ml/tests/tuning/test_tvs_io_pipeline.py @@ -17,10 +17,11 @@ import tempfile -from pyspark.ml.feature import HashingTF, Tokenizer from pyspark.ml import Pipeline from pyspark.ml.classification import LogisticRegression, OneVsRest from pyspark.ml.evaluation import MulticlassClassificationEvaluator +from pyspark.ml.feature import HashingTF, Tokenizer +from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin from pyspark.ml.tuning import ( ParamGridBuilder, TrainValidationSplit, @@ -31,8 +32,6 @@ SparkSessionTestCase, ) -from pyspark.ml.tests.tuning.test_tuning import ValidatorTestUtilsMixin - class TrainValidationSplitIONestedTests(SparkSessionTestCase, ValidatorTestUtilsMixin): def _run_test_save_load_pipeline_estimator(self, LogisticRegressionCls): diff --git a/python/pyspark/ml/torch/data.py b/python/pyspark/ml/torch/data.py index 826526d56c8d3..cc49dc1909e62 100644 --- a/python/pyspark/ml/torch/data.py +++ b/python/pyspark/ml/torch/data.py @@ -17,8 +17,8 @@ from typing import Any, Callable, Iterator -import torch import numpy as np +import torch from pyspark.sql.types import StructType diff --git a/python/pyspark/ml/torch/distributor.py b/python/pyspark/ml/torch/distributor.py index 67ba6adb03c48..8d5b1a74b9ae5 100644 --- a/python/pyspark/ml/torch/distributor.py +++ b/python/pyspark/ml/torch/distributor.py @@ -14,9 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import json -from contextlib import contextmanager import collections +import json import logging import os import random @@ -27,27 +26,28 @@ import tempfile import textwrap import time +from contextlib import contextmanager from typing import ( - Union, + Any, Callable, - List, Dict, - Optional, - Any, - Tuple, Generator, Iterator, + List, + Optional, + Tuple, + Union, ) from pyspark import cloudpickle -from pyspark.resource.information import ResourceInformation -from pyspark.sql import DataFrame, SparkSession -from pyspark.taskcontext import BarrierTaskContext +from pyspark.ml.dl_util import FunctionPickler from pyspark.ml.torch.log_communication import ( # type: ignore LogStreamingClient, LogStreamingServer, ) -from pyspark.ml.dl_util import FunctionPickler +from pyspark.resource.information import ResourceInformation +from pyspark.sql import DataFrame, SparkSession +from pyspark.taskcontext import BarrierTaskContext def _get_resources(session: SparkSession) -> Dict[str, ResourceInformation]: @@ -660,8 +660,10 @@ def _get_spark_task_function( # Spark task program def wrapped_train_fn(iterator): # type: ignore[no-untyped-def] import os + import pandas as pd import pyarrow + from pyspark import BarrierTaskContext CUDA_VISIBLE_DEVICES = "CUDA_VISIBLE_DEVICES" @@ -855,10 +857,11 @@ def _setup_files( def _setup_spark_partition_data( partition_data_iterator: Iterator[Any], input_schema_json: Dict[str, Any] ) -> Iterator[Any]: - from pyspark.sql.pandas.serializers import ArrowStreamSerializer - from pyspark.core.files import SparkFiles import json + from pyspark.core.files import SparkFiles + from pyspark.sql.pandas.serializers import ArrowStreamSerializer + if input_schema_json is None: yield return @@ -1081,10 +1084,11 @@ def _get_spark_partition_data_loader( prefetch_factor: Number of batches loaded in advance by each worker """ - from pyspark.sql.types import StructType - from pyspark.ml.torch.data import _SparkPartitionTorchDataset from torch.utils.data import DataLoader + from pyspark.ml.torch.data import _SparkPartitionTorchDataset + from pyspark.sql.types import StructType + arrow_file = os.environ[SPARK_PARTITION_ARROW_DATA_FILE] schema_file = os.environ[SPARK_DATAFRAME_SCHEMA_FILE] diff --git a/python/pyspark/ml/torch/log_communication.py b/python/pyspark/ml/torch/log_communication.py index ad1fc810f3ef4..d6c19af1553e2 100644 --- a/python/pyspark/ml/torch/log_communication.py +++ b/python/pyspark/ml/torch/log_communication.py @@ -16,16 +16,16 @@ # # type: ignore -from contextlib import closing -import time import socket import socketserver -from struct import pack, unpack import sys import threading +import time import traceback -from typing import Generator import warnings +from contextlib import closing +from struct import pack, unpack +from typing import Generator # Use b'\x00' as separator instead of b'\n', because the bytes are encoded in utf-8 _SERVER_POLL_INTERVAL = 0.1 diff --git a/python/pyspark/ml/torch/tests/test_data_loader.py b/python/pyspark/ml/torch/tests/test_data_loader.py index c911188f58077..8b3b970483995 100644 --- a/python/pyspark/ml/torch/tests/test_data_loader.py +++ b/python/pyspark/ml/torch/tests/test_data_loader.py @@ -17,12 +17,12 @@ import unittest +from pyspark.ml.linalg import Vectors from pyspark.ml.torch.distributor import ( TorchDistributor, _get_spark_partition_data_loader, ) from pyspark.sql import SparkSession -from pyspark.ml.linalg import Vectors # @unittest.skipIf(not have_torch, torch_requirement_message) diff --git a/python/pyspark/ml/torch/tests/test_distributor.py b/python/pyspark/ml/torch/tests/test_distributor.py index 62399e2db68d5..0149b869fc538 100644 --- a/python/pyspark/ml/torch/tests/test_distributor.py +++ b/python/pyspark/ml/torch/tests/test_distributor.py @@ -18,20 +18,20 @@ import contextlib import os import shutil -from io import StringIO import stat import subprocess import sys -import time import tempfile import threading -from typing import Callable, Dict, Any +import time import unittest +from io import StringIO +from typing import Any, Callable, Dict from unittest.mock import patch from pyspark import SparkConf, SparkContext from pyspark.ml.torch.distributor import TorchDistributor, _get_gpus_owned -from pyspark.ml.torch.torch_run_process_wrapper import clean_and_terminate, check_parent_alive +from pyspark.ml.torch.torch_run_process_wrapper import check_parent_alive, clean_and_terminate from pyspark.sql import SparkSession from pyspark.testing.sqlutils import SPARK_HOME from pyspark.testing.utils import have_torch, torch_requirement_message @@ -52,7 +52,7 @@ def patch_stdout() -> StringIO: def create_training_function(mnist_dir_path: str) -> Callable: import torch.nn as nn import torch.nn.functional as F - from torchvision import transforms, datasets + from torchvision import datasets, transforms batch_size = 100 num_epochs = 1 @@ -87,8 +87,8 @@ def forward(self, x: Any) -> Any: def train_fn(learning_rate: float) -> Any: import torch - import torch.optim as optim import torch.distributed as dist + import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler diff --git a/python/pyspark/ml/torch/tests/test_log_communication.py b/python/pyspark/ml/torch/tests/test_log_communication.py index ef8d0c7398485..6485bc9c64012 100644 --- a/python/pyspark/ml/torch/tests/test_log_communication.py +++ b/python/pyspark/ml/torch/tests/test_log_communication.py @@ -18,18 +18,18 @@ from __future__ import absolute_import, division, print_function import contextlib -from io import StringIO import sys import time -from typing import Any, Callable import unittest +from io import StringIO +from typing import Any, Callable import pyspark.ml.torch.log_communication from pyspark.ml.torch.log_communication import ( - LogStreamingServer, + _SERVER_POLL_INTERVAL, LogStreamingClient, LogStreamingClientBase, - _SERVER_POLL_INTERVAL, + LogStreamingServer, ) diff --git a/python/pyspark/ml/tree.py b/python/pyspark/ml/tree.py index 92692ec225a76..070a37944f853 100644 --- a/python/pyspark/ml/tree.py +++ b/python/pyspark/ml/tree.py @@ -14,23 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import List, Sequence, TypeVar, TYPE_CHECKING +from typing import TYPE_CHECKING, List, Sequence, TypeVar from pyspark import since +from pyspark.ml.common import inherit_doc from pyspark.ml.linalg import Vector from pyspark.ml.param import Params from pyspark.ml.param.shared import ( HasCheckpointInterval, + HasMaxIter, HasSeed, + HasStepSize, + HasValidationIndicatorCol, HasWeightCol, Param, TypeConverters, - HasMaxIter, - HasStepSize, - HasValidationIndicatorCol, ) from pyspark.ml.wrapper import JavaPredictionModel -from pyspark.ml.common import inherit_doc if TYPE_CHECKING: from pyspark.ml._typing import P diff --git a/python/pyspark/ml/tuning.py b/python/pyspark/ml/tuning.py index 00e317375c401..e4bac8d302c59 100644 --- a/python/pyspark/ml/tuning.py +++ b/python/pyspark/ml/tuning.py @@ -15,12 +15,13 @@ # limitations under the License. # +import itertools import json import os import sys -import itertools from multiprocessing.pool import ThreadPool from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -33,39 +34,39 @@ Union, cast, overload, - TYPE_CHECKING, ) import numpy as np -from pyspark import keyword_only, since, inheritable_thread_target -from pyspark.ml import Estimator, Transformer, Model -from pyspark.ml.common import inherit_doc, _py2java, _java2py +from pyspark import inheritable_thread_target, keyword_only, since +from pyspark.ml.base import Estimator, Model, Transformer +from pyspark.ml.common import _java2py, _py2java, inherit_doc from pyspark.ml.evaluation import Evaluator, JavaEvaluator -from pyspark.ml.param import Params, Param, TypeConverters +from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import HasCollectSubModels, HasParallelism, HasSeed from pyspark.ml.util import ( DefaultParamsReader, DefaultParamsWriter, + JavaMLWriter, MetaAlgorithmReadWrite, MLReadable, MLReader, MLWritable, MLWriter, - JavaMLWriter, - try_remote_write, - try_remote_read, _cache_spark_dataset, + try_remote_read, + try_remote_write, ) -from pyspark.ml.wrapper import JavaParams, JavaEstimator, JavaWrapper +from pyspark.ml.wrapper import JavaEstimator, JavaParams, JavaWrapper from pyspark.sql import functions as F from pyspark.sql.dataframe import DataFrame if TYPE_CHECKING: - from pyspark.ml._typing import ParamMap - from py4j.java_gateway import JavaObject from py4j.java_collections import JavaArray + from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext + from pyspark.ml._typing import ParamMap __all__ = [ "ParamGridBuilder", diff --git a/python/pyspark/ml/util.py b/python/pyspark/ml/util.py index c289078cc43e1..447cf44c3c4c3 100644 --- a/python/pyspark/ml/util.py +++ b/python/pyspark/ml/util.py @@ -15,14 +15,16 @@ # limitations under the License. # +import functools import json import logging import os import threading import time import uuid -import functools +from contextlib import contextmanager from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -33,11 +35,9 @@ Sequence, Type, TypeVar, - cast, - TYPE_CHECKING, Union, + cast, ) -from contextlib import contextmanager from pyspark import since from pyspark.ml.common import inherit_doc @@ -48,13 +48,14 @@ if TYPE_CHECKING: from py4j.java_gateway import JavaGateway, JavaObject + + from pyspark.core.context import SparkContext from pyspark.ml._typing import PipelineStage from pyspark.ml.base import Params - from pyspark.core.context import SparkContext + from pyspark.ml.evaluation import JavaEvaluator + from pyspark.ml.wrapper import JavaEstimator, JavaWrapper from pyspark.sql import DataFrame from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame - from pyspark.ml.wrapper import JavaWrapper, JavaEstimator - from pyspark.ml.evaluation import JavaEvaluator T = TypeVar("T") RW = TypeVar("RW", bound="BaseReadWrite") @@ -89,15 +90,15 @@ def invoke_remote_attribute_relation( instance: "JavaWrapper", method: str, *args: Any ) -> "ConnectDataFrame": import pyspark.sql.connect.proto as pb2 - from pyspark.ml.connect.util import _extract_id_methods - from pyspark.ml.connect.serialize import serialize # The attribute returns a dataframe, we need to wrap it # in the AttributeRelation from pyspark.ml.connect.proto import AttributeRelation - from pyspark.sql.connect.session import SparkSession - from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + from pyspark.ml.connect.serialize import serialize + from pyspark.ml.connect.util import _extract_id_methods from pyspark.ml.wrapper import JavaModel + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + from pyspark.sql.connect.session import SparkSession session = SparkSession.getActiveSession() assert session is not None @@ -172,7 +173,7 @@ def try_remote_fit(f: FuncT) -> FuncT: def wrapped(self: "JavaEstimator", dataset: "ConnectDataFrame") -> Any: if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: import pyspark.sql.connect.proto as pb2 - from pyspark.ml.connect.serialize import serialize_ml_params, deserialize + from pyspark.ml.connect.serialize import deserialize, serialize_ml_params client = dataset.sparkSession.client input = dataset._plan.plan(client) @@ -219,8 +220,8 @@ def try_remote_transform_relation(f: FuncT) -> FuncT: def wrapped(self: "JavaWrapper", dataset: "ConnectDataFrame") -> Any: if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: from pyspark.ml import Model, Transformer - from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame from pyspark.ml.connect.serialize import serialize_ml_params + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame session = dataset.sparkSession assert session is not None @@ -278,15 +279,15 @@ def try_remote_call(f: FuncT) -> FuncT: @functools.wraps(f) def wrapped(self: "JavaWrapper", name: str, *args: Any) -> Any: if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: - from pyspark.errors.exceptions.connect import SparkException import pyspark.sql.connect.proto as pb2 + from pyspark.errors.exceptions.connect import SparkException from pyspark.sql.connect.session import SparkSession session = SparkSession.getActiveSession() def remote_call() -> Any: + from pyspark.ml.connect.serialize import deserialize, serialize from pyspark.ml.connect.util import _extract_id_methods - from pyspark.ml.connect.serialize import serialize, deserialize from pyspark.ml.wrapper import JavaModel assert session is not None @@ -460,7 +461,7 @@ def try_remote_evaluate(f: FuncT) -> FuncT: def wrapped(self: "JavaEvaluator", dataset: "ConnectDataFrame") -> Any: if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ: import pyspark.sql.connect.proto as pb2 - from pyspark.ml.connect.serialize import serialize_ml_params, deserialize + from pyspark.ml.connect.serialize import deserialize, serialize_ml_params client = dataset.sparkSession.client input = dataset._plan.plan(client) @@ -1137,8 +1138,8 @@ class MetaAlgorithmReadWrite: @staticmethod def isMetaEstimator(pyInstance: Any) -> bool: from pyspark.ml import Estimator, Pipeline - from pyspark.ml.tuning import _ValidatorParams from pyspark.ml.classification import OneVsRest + from pyspark.ml.tuning import _ValidatorParams return ( isinstance(pyInstance, Pipeline) @@ -1149,8 +1150,8 @@ def isMetaEstimator(pyInstance: Any) -> bool: @staticmethod def getAllNestedStages(pyInstance: Any) -> List["Params"]: from pyspark.ml import Pipeline, PipelineModel - from pyspark.ml.tuning import _ValidatorParams from pyspark.ml.classification import OneVsRest, OneVsRestModel + from pyspark.ml.tuning import _ValidatorParams # TODO: We need to handle `RFormulaModel.pipelineModel` here after Pyspark RFormulaModel # support pipelineModel property. diff --git a/python/pyspark/ml/wrapper.py b/python/pyspark/ml/wrapper.py index cd8df055daa85..09563c3c32da4 100644 --- a/python/pyspark/ml/wrapper.py +++ b/python/pyspark/ml/wrapper.py @@ -16,27 +16,34 @@ # from abc import ABCMeta, abstractmethod -from typing import Any, Generic, Optional, List, Type, TypeVar, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Generic, List, Optional, Type, TypeVar from pyspark import since +from pyspark.ml.base import ( + Estimator, + Model, + PredictionModel, + Predictor, + Transformer, + _PredictorParams, +) +from pyspark.ml.common import _java2py, _py2java, inherit_doc +from pyspark.ml.param import Param, Params from pyspark.ml.util import ( - try_remote_transform_relation, + _jvm, try_remote_call, - try_remote_fit, try_remote_del, - try_remote_return_java_class, + try_remote_fit, try_remote_intercept, + try_remote_return_java_class, + try_remote_transform_relation, ) from pyspark.sql import DataFrame, is_remote -from pyspark.ml import Estimator, Predictor, PredictionModel, Transformer, Model -from pyspark.ml.base import _PredictorParams -from pyspark.ml.param import Param, Params -from pyspark.ml.util import _jvm -from pyspark.ml.common import inherit_doc, _java2py, _py2java if TYPE_CHECKING: + from py4j.java_gateway import JavaClass, JavaObject + from pyspark.ml._typing import ParamMap - from py4j.java_gateway import JavaObject, JavaClass T = TypeVar("T") diff --git a/python/pyspark/mllib/_typing.pyi b/python/pyspark/mllib/_typing.pyi index c5af46eb1d601..2f158aeacadcb 100644 --- a/python/pyspark/mllib/_typing.pyi +++ b/python/pyspark/mllib/_typing.pyi @@ -16,16 +16,16 @@ # specific language governing permissions and limitations # under the License. -from typing import List, Tuple, TYPE_CHECKING, TypeVar, Union +from typing import TYPE_CHECKING, List, Tuple, TypeVar, Union -from typing_extensions import Literal from numpy import ndarray from py4j.java_gateway import JavaObject +from typing_extensions import Literal from pyspark.mllib.linalg import Vector if TYPE_CHECKING: - from scipy.sparse import spmatrix, sparray + from scipy.sparse import sparray, spmatrix C = TypeVar("C", bound=type) JavaObjectOrPickleDump = Union[JavaObject, bytearray, bytes] diff --git a/python/pyspark/mllib/classification.py b/python/pyspark/mllib/classification.py index c8b66947e07cb..0d0efa9fba9b0 100644 --- a/python/pyspark/mllib/classification.py +++ b/python/pyspark/mllib/classification.py @@ -15,25 +15,24 @@ # limitations under the License. # -from math import exp import sys import warnings -from typing import Any, Iterable, Optional, Union, overload, TYPE_CHECKING +from math import exp +from typing import TYPE_CHECKING, Any, Iterable, Optional, Union, overload import numpy from pyspark import RDD, SparkContext, since -from pyspark.streaming.dstream import DStream -from pyspark.mllib.common import callMLlibFunc, _py2java, _java2py -from pyspark.mllib.linalg import _convert_to_vector +from pyspark.mllib.common import _java2py, _py2java, callMLlibFunc +from pyspark.mllib.linalg import Vector, _convert_to_vector from pyspark.mllib.regression import ( LabeledPoint, LinearModel, - _regression_train_wrapper, StreamingLinearAlgorithm, + _regression_train_wrapper, ) -from pyspark.mllib.util import Saveable, Loader, inherit_doc -from pyspark.mllib.linalg import Vector +from pyspark.mllib.util import Loader, Saveable, inherit_doc +from pyspark.streaming.dstream import DStream if TYPE_CHECKING: from pyspark.mllib._typing import VectorLike @@ -950,12 +949,13 @@ def trainOn(self, dstream: "DStream[LabeledPoint]") -> None: def update(rdd: RDD[LabeledPoint]) -> None: # LogisticRegressionWithSGD.train raises an error for an empty RDD. if not rdd.isEmpty(): + assert self._model is not None self._model = LogisticRegressionWithSGD.train( rdd, self.numIterations, self.stepSize, self.miniBatchFraction, - self._model.weights, # type: ignore[union-attr] + self._model.weights, regParam=self.regParam, convergenceTol=self.convergenceTol, ) @@ -965,8 +965,9 @@ def update(rdd: RDD[LabeledPoint]) -> None: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.mllib.classification + from pyspark.sql import SparkSession globs = pyspark.mllib.classification.__dict__.copy() spark = ( diff --git a/python/pyspark/mllib/clustering.py b/python/pyspark/mllib/clustering.py index ca7335999a83f..4a0991485f347 100644 --- a/python/pyspark/mllib/clustering.py +++ b/python/pyspark/mllib/clustering.py @@ -15,25 +15,26 @@ # limitations under the License. # -import sys import array as pyarray -from math import exp, log +import sys from collections import namedtuple -from typing import Any, List, Optional, Tuple, TypeVar, Union, overload, TYPE_CHECKING +from math import exp, log +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, overload import numpy as np from numpy import array, random, tile from pyspark import SparkContext, since from pyspark.core.rdd import RDD -from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc, callJavaFunc, _py2java, _java2py -from pyspark.mllib.linalg import SparseVector, _convert_to_vector, DenseVector # noqa: F401 +from pyspark.mllib.common import JavaModelWrapper, _java2py, _py2java, callJavaFunc, callMLlibFunc +from pyspark.mllib.linalg import DenseVector, SparseVector, _convert_to_vector # noqa: F401 from pyspark.mllib.stat.distribution import MultivariateGaussian -from pyspark.mllib.util import Saveable, Loader, inherit_doc, JavaLoader, JavaSaveable +from pyspark.mllib.util import JavaLoader, JavaSaveable, Loader, Saveable, inherit_doc from pyspark.streaming import DStream if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.mllib._typing import VectorLike T = TypeVar("T") @@ -1075,7 +1076,8 @@ def trainOn(self, dstream: "DStream[VectorLike]") -> None: self._validate(dstream) def update(rdd: RDD["VectorLike"]) -> None: - self._model.update(rdd, self._decayFactor, self._timeUnit) # type: ignore[union-attr] + assert self._model is not None + self._model.update(rdd, self._decayFactor, self._timeUnit) dstream.foreachRDD(update) @@ -1284,7 +1286,9 @@ def train( def _test() -> None: import doctest + import numpy + import pyspark.mllib.clustering try: diff --git a/python/pyspark/mllib/common.py b/python/pyspark/mllib/common.py index bfab55b8552ca..fc0f9121fd4e2 100644 --- a/python/pyspark/mllib/common.py +++ b/python/pyspark/mllib/common.py @@ -15,19 +15,19 @@ # limitations under the License. # -from typing import Any, Callable, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable if TYPE_CHECKING: from pyspark.mllib._typing import C, JavaObjectOrPickleDump import py4j.protocol -from py4j.protocol import Py4JJavaError -from py4j.java_gateway import JavaObject from py4j.java_collections import JavaArray, JavaList +from py4j.java_gateway import JavaObject +from py4j.protocol import Py4JJavaError import pyspark.core.context from pyspark import RDD, SparkContext -from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer +from pyspark.serializers import AutoBatchedSerializer, CPickleSerializer from pyspark.sql import DataFrame, SparkSession # Hack for support float('inf') in Py4j diff --git a/python/pyspark/mllib/evaluation.py b/python/pyspark/mllib/evaluation.py index 1990668f852a6..e873fecc71b61 100644 --- a/python/pyspark/mllib/evaluation.py +++ b/python/pyspark/mllib/evaluation.py @@ -15,8 +15,8 @@ # limitations under the License. # -from typing import Generic, List, Optional, Tuple, TypeVar, Union import sys +from typing import Generic, List, Optional, Tuple, TypeVar, Union from pyspark import since from pyspark.core.rdd import RDD @@ -668,9 +668,11 @@ def accuracy(self) -> float: def _test() -> None: import doctest + import numpy - from pyspark.sql import SparkSession + import pyspark.mllib.evaluation + from pyspark.sql import SparkSession try: # Numpy 1.14+ changed it's string format. diff --git a/python/pyspark/mllib/feature.py b/python/pyspark/mllib/feature.py index 685a82cf66933..e92c664fd59df 100644 --- a/python/pyspark/mllib/feature.py +++ b/python/pyspark/mllib/feature.py @@ -21,19 +21,18 @@ import sys import warnings -from typing import Dict, Hashable, Iterable, List, Optional, Tuple, Union, overload, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Hashable, Iterable, List, Optional, Tuple, Union, overload +from py4j.java_collections import JavaMap from py4j.protocol import Py4JJavaError from pyspark import since -from pyspark.core.rdd import RDD -from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper -from pyspark.mllib.linalg import Vectors, _convert_to_vector -from pyspark.mllib.util import JavaLoader, JavaSaveable from pyspark.core.context import SparkContext -from pyspark.mllib.linalg import Vector +from pyspark.core.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc +from pyspark.mllib.linalg import Vector, Vectors, _convert_to_vector from pyspark.mllib.regression import LabeledPoint -from py4j.java_collections import JavaMap +from pyspark.mllib.util import JavaLoader, JavaSaveable if TYPE_CHECKING: from pyspark.mllib._typing import VectorLike @@ -1052,6 +1051,7 @@ def transform( def _test() -> None: import doctest + from pyspark.sql import SparkSession globs = globals().copy() diff --git a/python/pyspark/mllib/fpm.py b/python/pyspark/mllib/fpm.py index 3f4d36884d94e..cdecff471e8da 100644 --- a/python/pyspark/mllib/fpm.py +++ b/python/pyspark/mllib/fpm.py @@ -16,13 +16,12 @@ # import sys - from typing import Any, Generic, List, NamedTuple, TypeVar -from pyspark import since, SparkContext -from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc -from pyspark.mllib.util import JavaSaveable, JavaLoader, inherit_doc +from pyspark import SparkContext, since from pyspark.core.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc +from pyspark.mllib.util import JavaLoader, JavaSaveable, inherit_doc __all__ = ["FPGrowth", "FPGrowthModel", "PrefixSpan", "PrefixSpanModel"] @@ -205,8 +204,9 @@ class FreqSequence(NamedTuple): def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.mllib.fpm + from pyspark.sql import SparkSession globs = pyspark.mllib.fpm.__dict__.copy() spark = SparkSession.builder.master("local[4]").appName("mllib.fpm tests").getOrCreate() diff --git a/python/pyspark/mllib/linalg/__init__.py b/python/pyspark/mllib/linalg/__init__.py index b8b1cbba0fff8..9f19c8fbe2f1c 100644 --- a/python/pyspark/mllib/linalg/__init__.py +++ b/python/pyspark/mllib/linalg/__init__.py @@ -23,25 +23,25 @@ SciPy is available in their environment. """ -import sys import array import struct +import sys from typing import ( + TYPE_CHECKING, Any, Callable, - cast, Dict, Generic, Iterable, List, Optional, - overload, Sequence, Tuple, Type, TypeVar, - TYPE_CHECKING, Union, + cast, + overload, ) import numpy as np @@ -49,20 +49,22 @@ from pyspark import since from pyspark.ml import linalg as newlinalg from pyspark.sql.types import ( - UserDefinedType, - StructField, - StructType, ArrayType, + BooleanType, + ByteType, + DataTypeSingleton, DoubleType, IntegerType, - ByteType, - BooleanType, + StructField, + StructType, + UserDefinedType, ) if TYPE_CHECKING: - from pyspark.mllib._typing import VectorLike, NormType from numpy.typing import ArrayLike + from pyspark.mllib._typing import NormType, VectorLike + QT = TypeVar("QT") RT = TypeVar("RT") @@ -167,7 +169,7 @@ def _double_to_long_bits(value: float) -> int: return struct.unpack("Q", struct.pack("d", value))[0] -class VectorUDT(UserDefinedType): +class VectorUDT(UserDefinedType, metaclass=DataTypeSingleton): """ SQL user-defined type (UDT) for Vector. """ @@ -222,7 +224,7 @@ def simpleString(self) -> str: return "vector" -class MatrixUDT(UserDefinedType): +class MatrixUDT(UserDefinedType, metaclass=DataTypeSingleton): """ SQL user-defined type (UDT) for Matrix. """ @@ -1622,6 +1624,7 @@ def R(self) -> RT: def _test() -> None: import doctest + import numpy try: diff --git a/python/pyspark/mllib/linalg/distributed.py b/python/pyspark/mllib/linalg/distributed.py index e5d27814602af..299c52d073386 100644 --- a/python/pyspark/mllib/linalg/distributed.py +++ b/python/pyspark/mllib/linalg/distributed.py @@ -20,13 +20,13 @@ """ import sys -from typing import Any, Generic, Optional, Tuple, TypeVar, Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, TypeVar, Union from py4j.java_gateway import JavaObject from pyspark import RDD, since -from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper -from pyspark.mllib.linalg import _convert_to_vector, DenseMatrix, Matrix, QRDecomposition, Vector +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc +from pyspark.mllib.linalg import DenseMatrix, Matrix, QRDecomposition, Vector, _convert_to_vector from pyspark.mllib.stat import MultivariateStatisticalSummary from pyspark.sql import DataFrame from pyspark.storagelevel import StorageLevel @@ -1633,10 +1633,12 @@ def toCoordinateMatrix(self) -> CoordinateMatrix: def _test() -> None: import doctest + import numpy - from pyspark.sql import SparkSession - from pyspark.mllib.linalg import Matrices + import pyspark.mllib.linalg.distributed + from pyspark.mllib.linalg import Matrices + from pyspark.sql import SparkSession try: # Numpy 1.14+ changed it's string format. diff --git a/python/pyspark/mllib/random.py b/python/pyspark/mllib/random.py index faafb024abb4b..815f7c1f57692 100644 --- a/python/pyspark/mllib/random.py +++ b/python/pyspark/mllib/random.py @@ -25,9 +25,9 @@ import numpy as np -from pyspark.mllib.common import callMLlibFunc from pyspark.core.context import SparkContext from pyspark.core.rdd import RDD +from pyspark.mllib.common import callMLlibFunc from pyspark.mllib.linalg import Vector __all__ = [ @@ -680,6 +680,7 @@ def gammaVectorRDD( def _test() -> None: import doctest + from pyspark.sql import SparkSession globs = globals().copy() diff --git a/python/pyspark/mllib/recommendation.py b/python/pyspark/mllib/recommendation.py index f6a0d39a74c60..f305f7c36adc9 100644 --- a/python/pyspark/mllib/recommendation.py +++ b/python/pyspark/mllib/recommendation.py @@ -371,6 +371,7 @@ def trainImplicit( def _test() -> None: import doctest + import pyspark.mllib.recommendation from pyspark.sql import SQLContext diff --git a/python/pyspark/mllib/regression.py b/python/pyspark/mllib/regression.py index a5dde25d32fca..7fb5ff055da68 100644 --- a/python/pyspark/mllib/regression.py +++ b/python/pyspark/mllib/regression.py @@ -18,6 +18,7 @@ import sys import warnings from typing import ( + TYPE_CHECKING, Any, Callable, Iterable, @@ -27,19 +28,17 @@ TypeVar, Union, overload, - TYPE_CHECKING, ) import numpy as np from pyspark import since -from pyspark.streaming.dstream import DStream -from pyspark.mllib.common import callMLlibFunc, _py2java, _java2py, inherit_doc -from pyspark.mllib.linalg import _convert_to_vector -from pyspark.mllib.util import Saveable, Loader -from pyspark.core.rdd import RDD from pyspark.core.context import SparkContext -from pyspark.mllib.linalg import Vector +from pyspark.core.rdd import RDD +from pyspark.mllib.common import _java2py, _py2java, callMLlibFunc, inherit_doc +from pyspark.mllib.linalg import Vector, _convert_to_vector +from pyspark.mllib.util import Loader, Saveable +from pyspark.streaming.dstream import DStream if TYPE_CHECKING: from pyspark.mllib._typing import VectorLike @@ -1038,8 +1037,9 @@ def update(rdd: RDD[LabeledPoint]) -> None: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.mllib.regression + from pyspark.sql import SparkSession globs = pyspark.mllib.regression.__dict__.copy() spark = SparkSession.builder.master("local[2]").appName("mllib.regression tests").getOrCreate() diff --git a/python/pyspark/mllib/stat/KernelDensity.py b/python/pyspark/mllib/stat/KernelDensity.py index bb03d0ef13575..ce14a99ddf953 100644 --- a/python/pyspark/mllib/stat/KernelDensity.py +++ b/python/pyspark/mllib/stat/KernelDensity.py @@ -20,8 +20,8 @@ import numpy as np from numpy import ndarray -from pyspark.mllib.common import callMLlibFunc from pyspark.core.rdd import RDD +from pyspark.mllib.common import callMLlibFunc class KernelDensity: diff --git a/python/pyspark/mllib/stat/__init__.py b/python/pyspark/mllib/stat/__init__.py index dc7a6da7545f3..b09d56cc054b1 100644 --- a/python/pyspark/mllib/stat/__init__.py +++ b/python/pyspark/mllib/stat/__init__.py @@ -19,10 +19,10 @@ Python package for statistical functions in MLlib. """ -from pyspark.mllib.stat._statistics import Statistics, MultivariateStatisticalSummary +from pyspark.mllib.stat._statistics import MultivariateStatisticalSummary, Statistics from pyspark.mllib.stat.distribution import MultivariateGaussian -from pyspark.mllib.stat.test import ChiSqTestResult, KolmogorovSmirnovTestResult from pyspark.mllib.stat.KernelDensity import KernelDensity +from pyspark.mllib.stat.test import ChiSqTestResult, KolmogorovSmirnovTestResult __all__ = [ "Statistics", diff --git a/python/pyspark/mllib/stat/_statistics.py b/python/pyspark/mllib/stat/_statistics.py index bab36039a6101..633c1482adec1 100644 --- a/python/pyspark/mllib/stat/_statistics.py +++ b/python/pyspark/mllib/stat/_statistics.py @@ -16,13 +16,13 @@ # import sys -from typing import cast, overload, List, Optional, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, List, Optional, Union, cast, overload from numpy import ndarray from py4j.java_gateway import JavaObject from pyspark.core.rdd import RDD -from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc from pyspark.mllib.linalg import Matrix, Vector, _convert_to_vector from pyspark.mllib.regression import LabeledPoint from pyspark.mllib.stat.test import ChiSqTestResult, KolmogorovSmirnovTestResult @@ -381,7 +381,9 @@ def kolmogorovSmirnovTest( def _test() -> None: import doctest + import numpy + from pyspark.sql import SparkSession try: diff --git a/python/pyspark/mllib/stat/test.py b/python/pyspark/mllib/stat/test.py index 4afd23f8df060..20553f2053520 100644 --- a/python/pyspark/mllib/stat/test.py +++ b/python/pyspark/mllib/stat/test.py @@ -17,7 +17,7 @@ from typing import Generic, Tuple, TypeVar -from pyspark.mllib.common import inherit_doc, JavaModelWrapper +from pyspark.mllib.common import JavaModelWrapper, inherit_doc __all__ = ["ChiSqTestResult", "KolmogorovSmirnovTestResult"] diff --git a/python/pyspark/mllib/tests/test_algorithms.py b/python/pyspark/mllib/tests/test_algorithms.py index 2cd15a5bf8ab9..872ff3ce770b6 100644 --- a/python/pyspark/mllib/tests/test_algorithms.py +++ b/python/pyspark/mllib/tests/test_algorithms.py @@ -141,14 +141,14 @@ def test_gmm_with_initial_model(self): self.assertAlmostEqual((gmm1.weights - gmm2.weights).sum(), 0.0) def test_classification(self): - from pyspark.mllib.classification import LogisticRegressionWithSGD, SVMWithSGD, NaiveBayes + from pyspark.mllib.classification import LogisticRegressionWithSGD, NaiveBayes, SVMWithSGD from pyspark.mllib.tree import ( DecisionTree, DecisionTreeModel, - RandomForest, - RandomForestModel, GradientBoostedTrees, GradientBoostedTreesModel, + RandomForest, + RandomForestModel, ) data = [ @@ -232,11 +232,11 @@ def test_classification(self): def test_regression(self): from pyspark.mllib.regression import ( - LinearRegressionWithSGD, LassoWithSGD, + LinearRegressionWithSGD, RidgeRegressionWithSGD, ) - from pyspark.mllib.tree import DecisionTree, RandomForest, GradientBoostedTrees + from pyspark.mllib.tree import DecisionTree, GradientBoostedTrees, RandomForest data = [ LabeledPoint(-1.0, [0, -1]), diff --git a/python/pyspark/mllib/tests/test_feature.py b/python/pyspark/mllib/tests/test_feature.py index 1ece1d861d451..f5382cb6d977b 100644 --- a/python/pyspark/mllib/tests/test_feature.py +++ b/python/pyspark/mllib/tests/test_feature.py @@ -17,11 +17,11 @@ from math import sqrt -from numpy import array, abs, tile +from numpy import abs, array, tile -from pyspark.mllib.linalg import SparseVector, DenseVector, Vectors +from pyspark.mllib.feature import IDF, ElementwiseProduct, HashingTF, StandardScaler, Word2Vec +from pyspark.mllib.linalg import DenseVector, SparseVector, Vectors from pyspark.mllib.linalg.distributed import RowMatrix -from pyspark.mllib.feature import HashingTF, IDF, StandardScaler, ElementwiseProduct, Word2Vec from pyspark.testing.mllibutils import MLlibTestCase diff --git a/python/pyspark/mllib/tests/test_linalg.py b/python/pyspark/mllib/tests/test_linalg.py index 36a1e30c09100..6dba3f63d82d6 100644 --- a/python/pyspark/mllib/tests/test_linalg.py +++ b/python/pyspark/mllib/tests/test_linalg.py @@ -18,24 +18,24 @@ import array as pyarray import unittest -from numpy import array, array_equal, zeros, arange, tile, ones, inf +from numpy import arange, array, array_equal, inf, ones, tile, zeros import pyspark.ml.linalg as newlinalg -from pyspark.serializers import CPickleSerializer from pyspark.mllib.linalg import ( - Vector, - SparseVector, - DenseVector, - VectorUDT, - _convert_to_vector, DenseMatrix, - SparseMatrix, - Vectors, + DenseVector, Matrices, MatrixUDT, + SparseMatrix, + SparseVector, + Vector, + Vectors, + VectorUDT, + _convert_to_vector, ) -from pyspark.mllib.linalg.distributed import RowMatrix, IndexedRowMatrix, IndexedRow +from pyspark.mllib.linalg.distributed import IndexedRow, IndexedRowMatrix, RowMatrix from pyspark.mllib.regression import LabeledPoint +from pyspark.serializers import CPickleSerializer from pyspark.sql import Row from pyspark.testing.mllibutils import MLlibTestCase from pyspark.testing.utils import have_scipy @@ -418,6 +418,10 @@ class VectorUDTTests(MLlibTestCase): def test_json_schema(self): self.assertEqual(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_singleton(self): + self.assertIs(VectorUDT(), VectorUDT()) + self.assertIs(VectorUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_serialization(self): for v in [self.dv0, self.dv1, self.sv0, self.sv1]: self.assertEqual(v, self.udt.deserialize(self.udt.serialize(v))) @@ -479,6 +483,10 @@ class MatrixUDTTests(MLlibTestCase): def test_json_schema(self): self.assertEqual(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_singleton(self): + self.assertIs(MatrixUDT(), MatrixUDT()) + self.assertIs(MatrixUDT.fromJson(self.udt.jsonValue()), self.udt) + def test_serialization(self): for m in [self.dm1, self.dm2, self.sm1, self.sm2]: self.assertEqual(m, self.udt.deserialize(self.udt.serialize(m))) @@ -583,7 +591,7 @@ def test_clustering(self): self.assertEqual(clusters.predict(data[2]), clusters.predict(data[3])) def test_classification(self): - from pyspark.mllib.classification import LogisticRegressionWithSGD, SVMWithSGD, NaiveBayes + from pyspark.mllib.classification import LogisticRegressionWithSGD, NaiveBayes, SVMWithSGD from pyspark.mllib.tree import DecisionTree data = [ @@ -624,8 +632,8 @@ def test_classification(self): def test_regression(self): from pyspark.mllib.regression import ( - LinearRegressionWithSGD, LassoWithSGD, + LinearRegressionWithSGD, RidgeRegressionWithSGD, ) from pyspark.mllib.tree import DecisionTree diff --git a/python/pyspark/mllib/tests/test_stat.py b/python/pyspark/mllib/tests/test_stat.py index 6d1f4851e8973..54cda5342ae28 100644 --- a/python/pyspark/mllib/tests/test_stat.py +++ b/python/pyspark/mllib/tests/test_stat.py @@ -19,11 +19,11 @@ from numpy import array -from pyspark.mllib.linalg import Vectors, Matrices +from pyspark.errors import IllegalArgumentException +from pyspark.mllib.linalg import Matrices, Vectors from pyspark.mllib.random import RandomRDDs from pyspark.mllib.regression import LabeledPoint from pyspark.mllib.stat import Statistics -from pyspark.errors import IllegalArgumentException from pyspark.testing.mllibutils import MLlibTestCase diff --git a/python/pyspark/mllib/tests/test_streaming_algorithms.py b/python/pyspark/mllib/tests/test_streaming_algorithms.py index 35e27056a2ffe..ad61dcd872cc3 100644 --- a/python/pyspark/mllib/tests/test_streaming_algorithms.py +++ b/python/pyspark/mllib/tests/test_streaming_algorithms.py @@ -18,12 +18,12 @@ import os import unittest -from numpy import array, random, exp, dot, all, mean, abs +from numpy import abs, all, array, dot, exp, mean, random from numpy import sum as array_sum from pyspark import SparkContext -from pyspark.mllib.clustering import StreamingKMeans, StreamingKMeansModel from pyspark.mllib.classification import StreamingLogisticRegressionWithSGD +from pyspark.mllib.clustering import StreamingKMeans, StreamingKMeansModel from pyspark.mllib.linalg import Vectors from pyspark.mllib.regression import LabeledPoint, StreamingLinearRegressionWithSGD from pyspark.mllib.util import LinearDataGenerator diff --git a/python/pyspark/mllib/tests/test_util.py b/python/pyspark/mllib/tests/test_util.py index 6ff165e8dcbd8..8e68a589684ca 100644 --- a/python/pyspark/mllib/tests/test_util.py +++ b/python/pyspark/mllib/tests/test_util.py @@ -19,10 +19,9 @@ import tempfile from pyspark.mllib.common import _to_java_object_rdd -from pyspark.mllib.util import LinearDataGenerator -from pyspark.mllib.util import MLUtils -from pyspark.mllib.linalg import SparseVector, DenseVector, Vectors +from pyspark.mllib.linalg import DenseVector, SparseVector, Vectors from pyspark.mllib.random import RandomRDDs +from pyspark.mllib.util import LinearDataGenerator, MLUtils from pyspark.testing.mllibutils import MLlibTestCase diff --git a/python/pyspark/mllib/tree.py b/python/pyspark/mllib/tree.py index 2f3a85d83ad81..086fa410b18f0 100644 --- a/python/pyspark/mllib/tree.py +++ b/python/pyspark/mllib/tree.py @@ -15,16 +15,16 @@ # limitations under the License. # -import sys import random +import sys +from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union, overload from pyspark import since -from pyspark.mllib.common import callMLlibFunc, inherit_doc, JavaModelWrapper +from pyspark.core.rdd import RDD +from pyspark.mllib.common import JavaModelWrapper, callMLlibFunc, inherit_doc from pyspark.mllib.linalg import _convert_to_vector from pyspark.mllib.regression import LabeledPoint from pyspark.mllib.util import JavaLoader, JavaSaveable -from typing import Dict, Optional, Tuple, Union, overload, TYPE_CHECKING -from pyspark.core.rdd import RDD if TYPE_CHECKING: from pyspark.mllib._typing import VectorLike diff --git a/python/pyspark/mllib/util.py b/python/pyspark/mllib/util.py index 699909d430ea5..5c1bea3e2bf03 100644 --- a/python/pyspark/mllib/util.py +++ b/python/pyspark/mllib/util.py @@ -17,25 +17,25 @@ import sys from functools import reduce +from typing import TYPE_CHECKING, Generic, Iterable, List, Optional, Tuple, Type, TypeVar, cast import numpy as np from pyspark import since -from pyspark.mllib.common import callMLlibFunc, inherit_doc -from pyspark.mllib.linalg import Vectors, SparseVector, _convert_to_vector -from pyspark.sql import DataFrame -from typing import Generic, Iterable, List, Optional, Tuple, Type, TypeVar, cast, TYPE_CHECKING from pyspark.core.context import SparkContext -from pyspark.mllib.linalg import Vector from pyspark.core.rdd import RDD +from pyspark.mllib.common import callMLlibFunc, inherit_doc +from pyspark.mllib.linalg import SparseVector, Vector, Vectors, _convert_to_vector +from pyspark.sql import DataFrame T = TypeVar("T") L = TypeVar("L", bound="Loader") JL = TypeVar("JL", bound="JavaLoader") if TYPE_CHECKING: - from pyspark.mllib._typing import VectorLike from py4j.java_gateway import JavaObject + + from pyspark.mllib._typing import VectorLike from pyspark.mllib.regression import LabeledPoint @@ -636,6 +636,7 @@ def generateLinearRDD( def _test() -> None: import doctest + from pyspark.sql import SparkSession globs = globals().copy() diff --git a/python/pyspark/pandas/__init__.py b/python/pyspark/pandas/__init__.py index bb05ec3bf857b..aa2b2701acff4 100644 --- a/python/pyspark/pandas/__init__.py +++ b/python/pyspark/pandas/__init__.py @@ -40,13 +40,13 @@ raise from pyspark.pandas.frame import DataFrame +from pyspark.pandas.groupby import NamedAgg from pyspark.pandas.indexes.base import Index from pyspark.pandas.indexes.category import CategoricalIndex from pyspark.pandas.indexes.datetimes import DatetimeIndex from pyspark.pandas.indexes.multi import MultiIndex from pyspark.pandas.indexes.timedelta import TimedeltaIndex from pyspark.pandas.series import Series -from pyspark.pandas.groupby import NamedAgg __all__ = [ # noqa: F405 "read_csv", @@ -80,8 +80,8 @@ def _auto_patch_spark() -> None: - import os import logging + import os # Attach a usage logger. 'KOALAS_USAGE_LOGGER' is legacy, and it's for compatibility. logger_module = os.getenv("PYSPARK_PANDAS_USAGE_LOGGER", os.getenv("KOALAS_USAGE_LOGGER", "")) @@ -129,7 +129,7 @@ def _auto_patch_pandas() -> None: _auto_patch_pandas() # Import after the usage logger is attached. -from pyspark.pandas.config import get_option, options, option_context, reset_option, set_option +from pyspark.pandas.config import get_option, option_context, options, reset_option, set_option from pyspark.pandas.namespace import * # noqa: F403 from pyspark.pandas.sql_formatter import sql diff --git a/python/pyspark/pandas/_typing.py b/python/pyspark/pandas/_typing.py index 51d1233fae25d..6a18b73acb713 100644 --- a/python/pyspark/pandas/_typing.py +++ b/python/pyspark/pandas/_typing.py @@ -16,7 +16,7 @@ # import datetime import decimal -from typing import Any, Tuple, TypeVar, Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Tuple, TypeVar, Union import numpy as np from pandas.api.extensions import ExtensionDtype diff --git a/python/pyspark/pandas/accessors.py b/python/pyspark/pandas/accessors.py index a0edf34da5858..86c64a7277e55 100644 --- a/python/pyspark/pandas/accessors.py +++ b/python/pyspark/pandas/accessors.py @@ -19,32 +19,32 @@ """ import inspect -from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING, cast, List from types import FunctionType +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple, Union, cast import numpy as np # noqa: F401 import pandas as pd -from pyspark.sql import functions as F -from pyspark.sql.functions import pandas_udf -from pyspark.sql.types import DataType, LongType, StructField, StructType from pyspark.pandas._typing import DataFrameOrSeries, Name from pyspark.pandas.internal import ( - InternalField, - InternalFrame, - SPARK_INDEX_NAME_FORMAT, SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_FORMAT, SPARK_INDEX_NAME_PATTERN, + InternalField, + InternalFrame, ) -from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType +from pyspark.pandas.typedef import DataFrameType, ScalarType, SeriesType, infer_return_type from pyspark.pandas.utils import ( - is_name_like_value, is_name_like_tuple, + is_name_like_value, + log_advice, name_like_string, scol_for, verify_temp_column_name, - log_advice, ) +from pyspark.sql import functions as F +from pyspark.sql.functions import pandas_udf +from pyspark.sql.types import DataType, LongType, StructField, StructType if TYPE_CHECKING: from pyspark.pandas.frame import DataFrame @@ -328,9 +328,9 @@ def apply_batch( """ # TODO: codes here partially duplicate `DataFrame.apply`. Can we deduplicate? - from pyspark.pandas.groupby import GroupBy - from pyspark.pandas.frame import DataFrame from pyspark import pandas as ps + from pyspark.pandas.frame import DataFrame + from pyspark.pandas.groupby import GroupBy if not isinstance(func, FunctionType): assert callable(func), "the first argument should be a callable function." @@ -564,10 +564,10 @@ def transform_batch( 2 12 Name: B, dtype: int64 """ - from pyspark.pandas.groupby import GroupBy + from pyspark import pandas as ps from pyspark.pandas.frame import DataFrame + from pyspark.pandas.groupby import GroupBy from pyspark.pandas.series import first_series - from pyspark import pandas as ps assert callable(func), "the first argument should be a callable function." spec = inspect.getfullargspec(func) @@ -899,9 +899,9 @@ def transform_batch( def _transform_batch( self, func: Callable[..., pd.Series], return_type: Optional[Union[SeriesType, ScalarType]] ) -> "Series": + from pyspark import pandas as ps from pyspark.pandas.groupby import GroupBy from pyspark.pandas.series import Series, first_series - from pyspark import pandas as ps if not isinstance(func, FunctionType): f = func @@ -960,11 +960,12 @@ def pudf(*series: pd.Series) -> pd.Series: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.accessors + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/base.py b/python/pyspark/pandas/base.py index fe14f2becc92d..eca0cf074464c 100644 --- a/python/pyspark/pandas/base.py +++ b/python/pyspark/pandas/base.py @@ -21,56 +21,56 @@ import warnings from abc import ABCMeta, abstractmethod -from functools import wraps, partial +from functools import partial, wraps from itertools import chain -from typing import Any, Callable, ClassVar, Optional, Sequence, Tuple, Union, cast, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Sequence, Tuple, Union, cast import numpy as np import pandas as pd -from pandas.api.types import is_list_like, CategoricalDtype +from pandas.api.types import CategoricalDtype, is_list_like -from pyspark.sql import functions as F, Column, Window -from pyspark.sql.types import ( - BinaryType, - BooleanType, - CharType, - DataType, - DateType, - DayTimeIntervalType, - LongType, - NumericType, - StringType, - TimestampNTZType, - TimestampType, - TimeType, - VarcharType, -) from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark.pandas._typing import Axis, Dtype, IndexOpsLike, Label, SeriesOrIndex from pyspark.pandas.config import get_option, option_context +from pyspark.pandas.frame import DataFrame from pyspark.pandas.internal import ( - InternalField, - InternalFrame, NATURAL_ORDER_COLUMN_NAME, SPARK_DEFAULT_INDEX_NAME, + InternalField, + InternalFrame, ) from pyspark.pandas.spark.accessors import SparkIndexOpsMethods from pyspark.pandas.typedef.typehints import handle_dtype_as_extension_dtype from pyspark.pandas.utils import ( + ERROR_MESSAGE_CANNOT_COMBINE, ansi_mode_context, combine_frames, same_anchor, scol_for, validate_axis, - ERROR_MESSAGE_CANNOT_COMBINE, ) -from pyspark.pandas.frame import DataFrame +from pyspark.sql import Column, Window +from pyspark.sql import functions as F +from pyspark.sql.types import ( + BinaryType, + BooleanType, + CharType, + DataType, + DateType, + DayTimeIntervalType, + LongType, + NumericType, + StringType, + TimestampNTZType, + TimestampType, + TimeType, + VarcharType, +) if TYPE_CHECKING: - from pyspark.sql._typing import ColumnOrName - from pyspark.pandas.data_type_ops.base import DataTypeOps from pyspark.pandas.series import Series + from pyspark.sql._typing import ColumnOrName def should_alignment_for_column_op(self: SeriesOrIndex, other: SeriesOrIndex) -> bool: @@ -542,12 +542,14 @@ def __len__(self) -> int: # NDArray Compat def __array_ufunc__( self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any - ) -> SeriesOrIndex: + ) -> Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]: from pyspark.pandas import numpy_compat - # Try dunder methods first. - result = numpy_compat.maybe_dispatch_ufunc_to_dunder_op( - self, ufunc, method, *inputs, **kwargs + # Try dunder methods first. A multi-output ufunc (for example np.modf) yields a + # 2-tuple of results rather than a single Series or Index, so both `result` and this + # method's return type must admit that tuple. + result: Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]] = ( + numpy_compat.maybe_dispatch_ufunc_to_dunder_op(self, ufunc, method, *inputs, **kwargs) ) # After that, we try with PySpark APIs. @@ -557,7 +559,7 @@ def __array_ufunc__( ) if result is not NotImplemented: - return cast(SeriesOrIndex, result) + return result else: # TODO: support more APIs? raise NotImplementedError( @@ -1453,8 +1455,8 @@ def value_counts( 3 1 Name: count, dtype: int64 """ - from pyspark.pandas.series import first_series from pyspark.pandas.indexes.multi import MultiIndex + from pyspark.pandas.series import first_series if bins is not None: raise NotImplementedError("value_counts currently does not support bins") @@ -1791,11 +1793,12 @@ def factorize( def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.base + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/categorical.py b/python/pyspark/pandas/categorical.py index 36f10f3426a14..90ef691bca0d1 100644 --- a/python/pyspark/pandas/categorical.py +++ b/python/pyspark/pandas/categorical.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Callable, List, Optional, Union, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union, cast import pandas as pd from pandas.api.types import ( @@ -23,8 +23,8 @@ is_list_like, ) -from pyspark.pandas.internal import InternalField from pyspark.pandas.data_type_ops.categorical_ops import _to_cat +from pyspark.pandas.internal import InternalField from pyspark.sql import functions as F from pyspark.sql.types import StructField @@ -769,11 +769,12 @@ def set_categories( def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.categorical + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/config.py b/python/pyspark/pandas/config.py index 8db649997db60..9161810adde1c 100644 --- a/python/pyspark/pandas/config.py +++ b/python/pyspark/pandas/config.py @@ -19,13 +19,13 @@ Infrastructure of options for pandas-on-Spark. """ -from contextlib import contextmanager import json -from typing import Any, Callable, Dict, Iterator, List, Tuple, Union, Optional +from contextlib import contextmanager +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union from pyspark._globals import _NoValue, _NoValueType -from pyspark.sql.session import SparkSession from pyspark.pandas.utils import default_session +from pyspark.sql.session import SparkSession __all__ = ["get_option", "set_option", "reset_option", "options", "option_context"] @@ -554,11 +554,12 @@ def __dir__(self) -> List[str]: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.config + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/correlation.py b/python/pyspark/pandas/correlation.py index cf9d389612166..eb5a172758a48 100644 --- a/python/pyspark/pandas/correlation.py +++ b/python/pyspark/pandas/correlation.py @@ -17,9 +17,10 @@ from typing import List -from pyspark.sql import DataFrame as SparkDataFrame, functions as F +from pyspark.pandas.utils import is_ansi_mode_enabled, verify_temp_column_name +from pyspark.sql import DataFrame as SparkDataFrame +from pyspark.sql import functions as F from pyspark.sql.window import Window -from pyspark.pandas.utils import verify_temp_column_name, is_ansi_mode_enabled CORRELATION_VALUE_1_COLUMN = "__correlation_value_1_input__" CORRELATION_VALUE_2_COLUMN = "__correlation_value_2_input__" diff --git a/python/pyspark/pandas/data_type_ops/base.py b/python/pyspark/pandas/data_type_ops/base.py index 72ce6cf7d9301..b7b8d187fc2ec 100644 --- a/python/pyspark/pandas/data_type_ops/base.py +++ b/python/pyspark/pandas/data_type_ops/base.py @@ -17,15 +17,25 @@ import numbers from abc import ABCMeta -from typing import Any, Optional, Union, cast from itertools import chain +from typing import Any, Optional, Union, cast import numpy as np import pandas as pd from pandas.api.types import CategoricalDtype from pandas.core.dtypes.common import is_numeric_dtype -from pyspark.sql import functions as F, Column as PySparkColumn +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.typedef.typehints import ( + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, + handle_dtype_as_extension_dtype, + is_str_dtype, + spark_type_to_pandas_dtype, +) +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F from pyspark.sql.types import ( ArrayType, BinaryType, @@ -41,19 +51,10 @@ NumericType, StringType, StructType, - TimestampType, TimestampNTZType, + TimestampType, UserDefinedType, ) -from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.pandas.typedef.typehints import ( - extension_dtypes_available, - extension_float_dtypes_available, - extension_object_dtypes_available, - handle_dtype_as_extension_dtype, - is_str_dtype, - spark_type_to_pandas_dtype, -) if extension_dtypes_available: from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype @@ -115,9 +116,10 @@ def _should_return_all_false(left: IndexOpsLike, right: Any) -> bool: Determine if binary comparison should short-circuit to all False, based on incompatible dtypes: non-numeric vs. numeric (including bools). """ - from pyspark.pandas.base import IndexOpsMixin from pandas.api.types import is_list_like + from pyspark.pandas.base import IndexOpsMixin + def are_both_numeric(left_dtype: Dtype, right_dtype: Dtype) -> bool: return is_numeric_dtype(left_dtype) and is_numeric_dtype(right_dtype) @@ -262,11 +264,11 @@ class DataTypeOps(object, metaclass=ABCMeta): def __new__(cls, dtype: Dtype, spark_type: DataType) -> "DataTypeOps": from pyspark.pandas.data_type_ops.binary_ops import BinaryOps - from pyspark.pandas.data_type_ops.boolean_ops import BooleanOps, BooleanExtensionOps + from pyspark.pandas.data_type_ops.boolean_ops import BooleanExtensionOps, BooleanOps from pyspark.pandas.data_type_ops.categorical_ops import CategoricalOps from pyspark.pandas.data_type_ops.complex_ops import ArrayOps, MapOps, StructOps from pyspark.pandas.data_type_ops.date_ops import DateOps - from pyspark.pandas.data_type_ops.datetime_ops import DatetimeOps, DatetimeNTZOps + from pyspark.pandas.data_type_ops.datetime_ops import DatetimeNTZOps, DatetimeOps from pyspark.pandas.data_type_ops.null_ops import NullOps from pyspark.pandas.data_type_ops.num_ops import ( DecimalOps, @@ -275,7 +277,7 @@ def __new__(cls, dtype: Dtype, spark_type: DataType) -> "DataTypeOps": IntegralExtensionOps, IntegralOps, ) - from pyspark.pandas.data_type_ops.string_ops import StringOps, StringExtensionOps + from pyspark.pandas.data_type_ops.string_ops import StringExtensionOps, StringOps from pyspark.pandas.data_type_ops.timedelta_ops import TimedeltaOps from pyspark.pandas.data_type_ops.udt_ops import UDTOps @@ -434,9 +436,9 @@ def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: return cast(SeriesOrIndex, left_scol) if isinstance(right, (list, tuple)): - from pyspark.pandas.series import first_series, scol_for from pyspark.pandas.frame import DataFrame from pyspark.pandas.internal import NATURAL_ORDER_COLUMN_NAME, InternalField + from pyspark.pandas.series import first_series, scol_for if len(left) != len(right): raise ValueError("Lengths must be equal") @@ -526,7 +528,7 @@ def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: return column_op(PySparkColumn.__eq__)(left, right) def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex: - from pyspark.pandas.base import column_op, IndexOpsMixin + from pyspark.pandas.base import IndexOpsMixin, column_op _sanitize_list_like(right) diff --git a/python/pyspark/pandas/data_type_ops/binary_ops.py b/python/pyspark/pandas/data_type_ops/binary_ops.py index f528d3e9ae2a4..26eea9bfe622c 100644 --- a/python/pyspark/pandas/data_type_ops/binary_ops.py +++ b/python/pyspark/pandas/data_type_ops/binary_ops.py @@ -19,8 +19,8 @@ from pandas.api.types import CategoricalDtype -from pyspark.pandas.base import column_op, IndexOpsMixin from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin, column_op from pyspark.pandas.data_type_ops.base import ( DataTypeOps, _as_categorical_type, diff --git a/python/pyspark/pandas/data_type_ops/boolean_ops.py b/python/pyspark/pandas/data_type_ops/boolean_ops.py index c52bc9b5051c9..b34fafa273d50 100644 --- a/python/pyspark/pandas/data_type_ops/boolean_ops.py +++ b/python/pyspark/pandas/data_type_ops/boolean_ops.py @@ -22,19 +22,20 @@ from pandas.api.types import CategoricalDtype, is_integer_dtype from pandas.core.dtypes.common import is_numeric_dtype -from pyspark.pandas.base import column_op, IndexOpsMixin -from pyspark.pandas.config import get_option +from pyspark.errors import PySparkValueError from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin, column_op +from pyspark.pandas.config import get_option from pyspark.pandas.data_type_ops.base import ( DataTypeOps, - is_valid_operand_for_numeric_arithmetic, - transform_boolean_operand_to_numeric, _as_bool_type, _as_categorical_type, _as_other_type, - _sanitize_list_like, - _is_valid_for_logical_operator, _is_boolean_type, + _is_valid_for_logical_operator, + _sanitize_list_like, + is_valid_operand_for_numeric_arithmetic, + transform_boolean_operand_to_numeric, ) from pyspark.pandas.typedef.typehints import ( as_spark_type, @@ -43,9 +44,9 @@ pandas_on_spark_type, ) from pyspark.pandas.utils import is_ansi_mode_enabled -from pyspark.sql import functions as F, Column as PySparkColumn +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F from pyspark.sql.types import BooleanType, StringType -from pyspark.errors import PySparkValueError class BooleanOps(DataTypeOps): diff --git a/python/pyspark/pandas/data_type_ops/categorical_ops.py b/python/pyspark/pandas/data_type_ops/categorical_ops.py index 3a977f418641b..4253cf94c06e2 100644 --- a/python/pyspark/pandas/data_type_ops/categorical_ops.py +++ b/python/pyspark/pandas/data_type_ops/categorical_ops.py @@ -16,15 +16,15 @@ # from itertools import chain -from typing import cast, Any, Sequence, Union +from typing import Any, Sequence, Union, cast -import pandas as pd import numpy as np -from pandas.api.types import is_list_like, CategoricalDtype +import pandas as pd +from pandas.api.types import CategoricalDtype, is_list_like from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex from pyspark.pandas.base import IndexOpsMixin -from pyspark.pandas.data_type_ops.base import _sanitize_list_like, DataTypeOps +from pyspark.pandas.data_type_ops.base import DataTypeOps, _sanitize_list_like from pyspark.pandas.typedef import pandas_on_spark_type from pyspark.sql import functions as F from pyspark.sql.utils import pyspark_column_op diff --git a/python/pyspark/pandas/data_type_ops/complex_ops.py b/python/pyspark/pandas/data_type_ops/complex_ops.py index 415301e400e99..7971516324a61 100644 --- a/python/pyspark/pandas/data_type_ops/complex_ops.py +++ b/python/pyspark/pandas/data_type_ops/complex_ops.py @@ -20,7 +20,7 @@ from pandas.api.types import CategoricalDtype from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.base import IndexOpsMixin, column_op from pyspark.pandas.data_type_ops.base import ( DataTypeOps, _as_bool_type, @@ -30,7 +30,8 @@ _sanitize_list_like, ) from pyspark.pandas.typedef import pandas_on_spark_type -from pyspark.sql import functions as F, Column +from pyspark.sql import Column +from pyspark.sql import functions as F from pyspark.sql.types import ArrayType, BooleanType, NumericType, StringType diff --git a/python/pyspark/pandas/data_type_ops/date_ops.py b/python/pyspark/pandas/data_type_ops/date_ops.py index 9a0b82de6ce8b..d4e1dce693cdb 100644 --- a/python/pyspark/pandas/data_type_ops/date_ops.py +++ b/python/pyspark/pandas/data_type_ops/date_ops.py @@ -23,10 +23,8 @@ import pandas as pd from pandas.api.types import CategoricalDtype -from pyspark.sql import functions as F, Column as PySparkColumn -from pyspark.sql.types import BooleanType, DateType, StringType from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.base import IndexOpsMixin, column_op from pyspark.pandas.data_type_ops.base import ( DataTypeOps, _as_categorical_type, @@ -35,6 +33,9 @@ _sanitize_list_like, ) from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F +from pyspark.sql.types import BooleanType, DateType, StringType class DateOps(DataTypeOps): diff --git a/python/pyspark/pandas/data_type_ops/datetime_ops.py b/python/pyspark/pandas/data_type_ops/datetime_ops.py index d407b8512b4ed..d90ce08b0dfa4 100644 --- a/python/pyspark/pandas/data_type_ops/datetime_ops.py +++ b/python/pyspark/pandas/data_type_ops/datetime_ops.py @@ -24,18 +24,7 @@ from pandas.api.types import CategoricalDtype from pyspark.loose_version import LooseVersion -from pyspark.sql import Column, functions as F -from pyspark.sql.types import ( - BooleanType, - LongType, - StringType, - TimestampType, - TimestampNTZType, - NumericType, -) -from pyspark.sql.utils import pyspark_column_op from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.sql.internal import InternalFunction as SF from pyspark.pandas.base import IndexOpsMixin from pyspark.pandas.data_type_ops.base import ( DataTypeOps, @@ -45,6 +34,18 @@ _sanitize_list_like, ) from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql import Column +from pyspark.sql import functions as F +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + BooleanType, + LongType, + NumericType, + StringType, + TimestampNTZType, + TimestampType, +) +from pyspark.sql.utils import pyspark_column_op class DatetimeOps(DataTypeOps): diff --git a/python/pyspark/pandas/data_type_ops/null_ops.py b/python/pyspark/pandas/data_type_ops/null_ops.py index 1c3296011b616..955a2422f5d1e 100644 --- a/python/pyspark/pandas/data_type_ops/null_ops.py +++ b/python/pyspark/pandas/data_type_ops/null_ops.py @@ -19,7 +19,8 @@ from pandas.api.types import CategoricalDtype, is_list_like -from pyspark.pandas._typing import Dtype, IndexOpsLike +from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin from pyspark.pandas.data_type_ops.base import ( DataTypeOps, _as_bool_type, @@ -28,11 +29,9 @@ _as_string_type, _sanitize_list_like, ) -from pyspark.pandas._typing import SeriesOrIndex from pyspark.pandas.typedef import pandas_on_spark_type from pyspark.sql.types import BooleanType, StringType from pyspark.sql.utils import pyspark_column_op -from pyspark.pandas.base import IndexOpsMixin class NullOps(DataTypeOps): diff --git a/python/pyspark/pandas/data_type_ops/num_ops.py b/python/pyspark/pandas/data_type_ops/num_ops.py index e2a3180331eb9..8baeacf3b0187 100644 --- a/python/pyspark/pandas/data_type_ops/num_ops.py +++ b/python/pyspark/pandas/data_type_ops/num_ops.py @@ -17,34 +17,35 @@ import decimal import numbers -from typing import Any, Union, Callable, cast +from typing import Any, Callable, Union, cast import numpy as np import pandas as pd from pandas.api.types import ( + CategoricalDtype, is_bool_dtype, - is_integer_dtype, is_float_dtype, - is_numeric_dtype, - CategoricalDtype, + is_integer_dtype, is_list_like, + is_numeric_dtype, ) +from pyspark.errors import PySparkValueError from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.pandas.base import column_op, IndexOpsMixin, numpy_column_op +from pyspark.pandas.base import IndexOpsMixin, column_op, numpy_column_op from pyspark.pandas.config import get_option from pyspark.pandas.data_type_ops.base import ( DataTypeOps, - is_valid_operand_for_numeric_arithmetic, - transform_boolean_operand_to_numeric, _as_bool_type, _as_categorical_type, _as_other_type, _as_string_type, - _sanitize_list_like, - _is_valid_for_logical_operator, _is_boolean_type, + _is_valid_for_logical_operator, + _sanitize_list_like, _should_return_all_false, + is_valid_operand_for_numeric_arithmetic, + transform_boolean_operand_to_numeric, ) from pyspark.pandas.typedef.typehints import ( as_spark_type, @@ -52,14 +53,14 @@ pandas_on_spark_type, ) from pyspark.pandas.utils import is_ansi_mode_enabled -from pyspark.sql import functions as F, Column as PySparkColumn +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F from pyspark.sql.types import ( BooleanType, DataType, DecimalType, StringType, ) -from pyspark.errors import PySparkValueError # For Supporting Spark Connect from pyspark.sql.utils import pyspark_column_op diff --git a/python/pyspark/pandas/data_type_ops/string_ops.py b/python/pyspark/pandas/data_type_ops/string_ops.py index c416d03a9c8f6..78d2778038a82 100644 --- a/python/pyspark/pandas/data_type_ops/string_ops.py +++ b/python/pyspark/pandas/data_type_ops/string_ops.py @@ -22,11 +22,8 @@ from pandas.api.types import CategoricalDtype from pyspark.loose_version import LooseVersion -from pyspark.sql import functions as F -from pyspark.sql.types import IntegralType, StringType -from pyspark.sql.utils import pyspark_column_op from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex -from pyspark.pandas.base import column_op, IndexOpsMixin +from pyspark.pandas.base import IndexOpsMixin, column_op from pyspark.pandas.data_type_ops.base import ( DataTypeOps, _as_categorical_type, @@ -39,7 +36,9 @@ is_str_dtype, pandas_on_spark_type, ) -from pyspark.sql.types import BooleanType +from pyspark.sql import functions as F +from pyspark.sql.types import BooleanType, IntegralType, StringType +from pyspark.sql.utils import pyspark_column_op class StringOps(DataTypeOps): diff --git a/python/pyspark/pandas/data_type_ops/timedelta_ops.py b/python/pyspark/pandas/data_type_ops/timedelta_ops.py index 6f17474f61d80..bc38d1502ebf2 100644 --- a/python/pyspark/pandas/data_type_ops/timedelta_ops.py +++ b/python/pyspark/pandas/data_type_ops/timedelta_ops.py @@ -23,11 +23,6 @@ from pandas.api.types import CategoricalDtype from pyspark.loose_version import LooseVersion -from pyspark.sql.types import ( - BooleanType, - DayTimeIntervalType, - StringType, -) from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex from pyspark.pandas.base import IndexOpsMixin from pyspark.pandas.data_type_ops.base import ( @@ -38,6 +33,11 @@ _sanitize_list_like, ) from pyspark.pandas.typedef import pandas_on_spark_type +from pyspark.sql.types import ( + BooleanType, + DayTimeIntervalType, + StringType, +) from pyspark.sql.utils import pyspark_column_op diff --git a/python/pyspark/pandas/datetimes.py b/python/pyspark/pandas/datetimes.py index aee37e11ab963..4a1afca1fae66 100644 --- a/python/pyspark/pandas/datetimes.py +++ b/python/pyspark/pandas/datetimes.py @@ -26,12 +26,12 @@ from pandas.tseries.offsets import DateOffset import pyspark.pandas as ps -from pyspark.loose_version import LooseVersion import pyspark.sql.functions as F -from pyspark.sql.types import DateType, TimestampType, TimestampNTZType, IntegerType +from pyspark.loose_version import LooseVersion from pyspark.pandas import DataFrame -from pyspark.pandas.config import option_context from pyspark.pandas._typing import Dtype +from pyspark.pandas.config import option_context +from pyspark.sql.types import DateType, IntegerType, TimestampNTZType, TimestampType class DatetimeMethods: @@ -688,8 +688,6 @@ def round(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps. - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series @@ -748,8 +746,6 @@ def floor(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps. - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series @@ -808,8 +804,6 @@ def ceil(self, freq: Union[str, DateOffset], *args: Any, **kwargs: Any) -> "ps.S - 'raise' will raise an NonExistentTimeError if there are nonexistent times - .. note:: this option only works with pandas 0.24.0+ - Returns ------- Series @@ -919,11 +913,12 @@ def pandas_day_name(s) -> ps.Series[str]: # type: ignore[no-untyped-def] def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.datetimes + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/exceptions.py b/python/pyspark/pandas/exceptions.py index fce91deec2b2d..3289cd32f380d 100644 --- a/python/pyspark/pandas/exceptions.py +++ b/python/pyspark/pandas/exceptions.py @@ -113,11 +113,12 @@ def __init__( def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.exceptions + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/extensions.py b/python/pyspark/pandas/extensions.py index 1764f24770656..c36f61f6c0bec 100644 --- a/python/pyspark/pandas/extensions.py +++ b/python/pyspark/pandas/extensions.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Callable, Generic, Optional, Type, Union, TYPE_CHECKING import warnings +from typing import TYPE_CHECKING, Callable, Generic, Optional, Type, Union from pyspark.pandas._typing import T @@ -357,12 +357,14 @@ def bar(self): def _test() -> None: - import os import doctest + import os import sys + import numpy - from pyspark.sql import SparkSession + import pyspark.pandas.extensions + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py index d196a85070bd8..73919101f996e 100644 --- a/python/pyspark/pandas/frame.py +++ b/python/pyspark/pandas/frame.py @@ -19,24 +19,26 @@ A wrapper class for Spark DataFrame to behave like pandas DataFrame. """ -from collections import defaultdict, namedtuple -from collections.abc import Mapping -import re -import warnings +import datetime import inspect import json +import re +import sys import types +import warnings +from collections import defaultdict, namedtuple +from collections.abc import Mapping from functools import partial, reduce, wraps -import sys -from itertools import zip_longest, chain +from itertools import chain, zip_longest from types import TracebackType from typing import ( + IO, + TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, - IO, Iterable, Iterator, List, @@ -49,16 +51,14 @@ Union, cast, no_type_check, - TYPE_CHECKING, ) -import datetime import numpy as np import pandas as pd from pandas.api.types import ( is_bool_dtype, - is_list_like, is_dict_like, + is_list_like, is_scalar, ) from pandas.tseries.frequencies import DateOffset, to_offset # type: ignore[attr-defined] @@ -66,34 +66,15 @@ if TYPE_CHECKING: from pandas.io.formats.style import Styler -from pandas.core.dtypes.common import infer_dtype_from_object # type: ignore[attr-defined] from pandas.core.accessor import CachedAccessor # type: ignore[attr-defined] +from pandas.core.dtypes.common import infer_dtype_from_object # type: ignore[attr-defined] from pandas.core.dtypes.inference import is_sequence # type: ignore[attr-defined] -from pyspark._globals import _NoValue, _NoValueType -from pyspark.loose_version import LooseVersion -from pyspark.errors import PySparkValueError from pyspark import StorageLevel -from pyspark.sql import Column as PySparkColumn, DataFrame as PySparkDataFrame, functions as F -from pyspark.sql.functions import pandas_udf -from pyspark.sql.internal import InternalFunction as SF -from pyspark.sql.types import ( - ArrayType, - BooleanType, - DataType, - DoubleType, - NumericType, - Row, - StringType, - StructField, - StructType, - DecimalType, - TimestampType, - TimestampNTZType, - NullType, -) -from pyspark.sql.window import Window from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark._globals import _NoValue, _NoValueType +from pyspark.errors import PySparkValueError +from pyspark.loose_version import LooseVersion from pyspark.pandas._typing import ( Axis, DataFrameOrSeries, @@ -104,15 +85,38 @@ T, ) from pyspark.pandas.accessors import PandasOnSparkFrameMethods -from pyspark.pandas.config import option_context, get_option +from pyspark.pandas.config import get_option, option_context from pyspark.pandas.correlation import ( - compute, - CORRELATION_VALUE_1_COLUMN, - CORRELATION_VALUE_2_COLUMN, CORRELATION_CORR_OUTPUT_COLUMN, CORRELATION_COUNT_OUTPUT_COLUMN, + CORRELATION_VALUE_1_COLUMN, + CORRELATION_VALUE_2_COLUMN, + compute, +) +from pyspark.pandas.generic import Frame +from pyspark.pandas.internal import ( + HIDDEN_COLUMNS, + NATURAL_ORDER_COLUMN_NAME, + SPARK_DEFAULT_INDEX_NAME, + SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_FORMAT, + SPARK_INDEX_NAME_PATTERN, + InternalField, + InternalFrame, +) +from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame +from pyspark.pandas.plot import PandasOnSparkPlotAccessor +from pyspark.pandas.spark.accessors import CachedSparkFrameMethods, SparkFrameMethods +from pyspark.pandas.typedef.typehints import ( + DataFrameType, + ScalarType, + SeriesType, + as_spark_type, + create_tuple_for_frame_type, + infer_return_type, + pandas_on_spark_type, + spark_type_to_pandas_dtype, ) -from pyspark.pandas.spark.accessors import SparkFrameMethods, CachedSparkFrameMethods from pyspark.pandas.utils import ( align_diff_frames, ansi_mode_context, @@ -123,6 +127,7 @@ is_name_like_tuple, is_name_like_value, is_testing, + log_advice, name_like_string, same_anchor, scol_for, @@ -132,39 +137,35 @@ validate_how, validate_mode, verify_temp_column_name, - log_advice, ) -from pyspark.pandas.generic import Frame -from pyspark.pandas.internal import ( - InternalField, - InternalFrame, - HIDDEN_COLUMNS, - NATURAL_ORDER_COLUMN_NAME, - SPARK_INDEX_NAME_FORMAT, - SPARK_DEFAULT_INDEX_NAME, - SPARK_DEFAULT_SERIES_NAME, - SPARK_INDEX_NAME_PATTERN, -) -from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame -from pyspark.pandas.typedef.typehints import ( - as_spark_type, - infer_return_type, - pandas_on_spark_type, - spark_type_to_pandas_dtype, - DataFrameType, - SeriesType, - ScalarType, - create_tuple_for_frame_type, +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import DataFrame as PySparkDataFrame +from pyspark.sql import functions as F +from pyspark.sql.functions import pandas_udf +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DecimalType, + DoubleType, + NullType, + NumericType, + Row, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, ) -from pyspark.pandas.plot import PandasOnSparkPlotAccessor +from pyspark.sql.window import Window if TYPE_CHECKING: - from pyspark.sql._typing import OptionalPrimitiveType - from pyspark.pandas.groupby import DataFrameGroupBy - from pyspark.pandas.resample import DataFrameResampler from pyspark.pandas.indexes import Index + from pyspark.pandas.resample import DataFrameResampler from pyspark.pandas.series import Series + from pyspark.sql._typing import OptionalPrimitiveType # These regular expression patterns are compiled and defined here to avoid compiling the same @@ -2718,6 +2719,7 @@ def to_feather( # but PlanMetrics/PlanObservedMetrics objects from Spark Connect are not # JSON serializable. We filter these internal attrs only for affected versions. import pyarrow as pa + from pyspark.loose_version import LooseVersion if LooseVersion(pa.__version__) >= LooseVersion("22.0.0"): @@ -4946,7 +4948,6 @@ def shift(self, periods: int = 1, fill_value: Optional[Any] = None) -> "DataFram lambda psser: psser._shift(periods, fill_value), should_resolve=True ) - # TODO(SPARK-46161): axis should support 1 or 'columns' either at this moment def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame": """ First discrete difference of element. @@ -4954,7 +4955,7 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame": Calculates the difference of a DataFrame element compared with another element in the DataFrame (default is the element in the same column of the previous row). - .. note:: the current implementation of diff uses Spark's Window without + .. note:: When ``axis=0``, the current implementation of diff uses Spark's Window without specifying partition specification. This leads to moving all data into a single partition in a single machine and could cause serious performance degradation. Avoid this method with very large datasets. @@ -4963,8 +4964,8 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame": ---------- periods : int, default 1 Periods to shift for calculating difference, accepts negative values. - axis : int, default 0 or 'index' - Can only be set to 0 now. + axis : {0 or 'index', 1 or 'columns'}, default 0 + Take difference over rows (0) or columns (1). Returns ------- @@ -5014,12 +5015,41 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame": 3 -1.0 -2.0 -9.0 4 -1.0 -3.0 -11.0 5 NaN NaN NaN + + Difference with previous column + + >>> df.diff(axis=1) + a b c + 0 NaN 0 0 + 1 NaN -1 3 + 2 NaN -1 7 + 3 NaN -1 13 + 4 NaN 0 20 + 5 NaN 2 28 """ axis = validate_axis(axis) - if axis != 0: - raise NotImplementedError('axis should be either 0 or "index" currently.') - - return self._apply_series_op(lambda psser: psser._diff(periods), should_resolve=True) + if axis == 0: + return self._apply_series_op(lambda psser: psser._diff(periods), should_resolve=True) + else: + if not isinstance(periods, int): + raise TypeError( + "periods should be an int; however, got [%s]" % type(periods).__name__ + ) + column_labels = self._internal.column_labels + data_col_names = self._internal.data_spark_column_names + new_columns: list[PySparkColumn] = [] + for i, label in enumerate(column_labels): + prev_idx = i - periods + if 0 <= prev_idx < len(column_labels): + prev_label = column_labels[prev_idx] + cur_col = self._internal.spark_column_for(label) + prev_col = self._internal.spark_column_for(prev_label) + new_columns.append(cur_col - prev_col) + else: + col_type = self._internal.spark_type_for(label) + new_columns.append(F.lit(None).cast(col_type).alias(data_col_names[i])) + internal = self._internal.with_new_columns(new_columns) + return DataFrame(internal) def nunique( self, @@ -14370,7 +14400,7 @@ def __iter__(self) -> Iterator[Name]: # NDArray Compat def __array_ufunc__( self, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any - ) -> "DataFrame": + ) -> Union["DataFrame", Tuple["DataFrame", "DataFrame"]]: # TODO: is it possible to deduplicate it with '_map_series_op'? if all(isinstance(inp, DataFrame) for inp in inputs) and any( not same_anchor(inp, inputs[0]) for inp in inputs @@ -14399,14 +14429,25 @@ def apply_op( # DataFrame and Series this = inputs[0] assert all(inp is this for inp in inputs if isinstance(inp, DataFrame)) - applied = [ + column_labels = this._internal.column_labels + outputs = [ ufunc( *[inp[label] if isinstance(inp, DataFrame) else inp for inp in inputs], **kwargs - ).rename(label) - for label in this._internal.column_labels + ) + for label in column_labels ] - internal = this._internal.with_new_columns(applied) - return DataFrame(internal) + if outputs and isinstance(outputs[0], tuple): + # A multi-output ufunc (for example np.modf) yields a two-element tuple per + # column; regroup into one DataFrame per output, matching numpy/pandas. + first_cols = [out[0].rename(label) for out, label in zip(outputs, column_labels)] + second_cols = [out[1].rename(label) for out, label in zip(outputs, column_labels)] + first_df: DataFrame = DataFrame(this._internal.with_new_columns(first_cols)) + second_df: DataFrame = DataFrame(this._internal.with_new_columns(second_cols)) + return first_df, second_df + else: + applied = [out.rename(label) for out, label in zip(outputs, column_labels)] + internal = this._internal.with_new_columns(applied) + return DataFrame(internal) def __class_getitem__(cls, params: Any) -> object: # See https://github.com/python/typing/issues/193 @@ -14481,14 +14522,15 @@ def __exit__( def _test() -> None: - import os import doctest + import os import shutil import sys import tempfile import uuid - from pyspark.sql import SparkSession + import pyspark.pandas.frame + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/generic.py b/python/pyspark/pandas/generic.py index ab0dd63b0f4e4..004cd580645e9 100644 --- a/python/pyspark/pandas/generic.py +++ b/python/pyspark/pandas/generic.py @@ -19,38 +19,30 @@ A base class of DataFrame/Column to behave like pandas DataFrame/Series. """ +import warnings from abc import ABCMeta, abstractmethod from functools import reduce from typing import ( + IO, + TYPE_CHECKING, Any, Callable, Dict, - IO, List, - Optional, NoReturn, + Optional, Tuple, Union, - TYPE_CHECKING, cast, ) -import warnings import numpy as np import pandas as pd from pandas.api.types import is_list_like +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark._globals import _NoValue, _NoValueType from pyspark.loose_version import LooseVersion -from pyspark.sql import Column, functions as F -from pyspark.sql.internal import InternalFunction as SF -from pyspark.sql.types import ( - BooleanType, - DoubleType, - LongType, - NumericType, -) -from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark.pandas._typing import ( Axis, DataFrameOrSeries, @@ -60,28 +52,37 @@ Name, Scalar, ) -from pyspark.pandas.indexing import AtIndexer, iAtIndexer, iLocIndexer, LocIndexer +from pyspark.pandas.indexing import AtIndexer, LocIndexer, iAtIndexer, iLocIndexer from pyspark.pandas.internal import InternalFrame from pyspark.pandas.typedef import spark_type_to_pandas_dtype from pyspark.pandas.utils import ( + SPARK_CONF_ARROW_ENABLED, is_name_like_tuple, is_name_like_value, + log_advice, name_like_string, scol_for, sql_conf, validate_arguments_and_invoke_function, validate_axis, validate_mode, - SPARK_CONF_ARROW_ENABLED, - log_advice, +) +from pyspark.sql import Column +from pyspark.sql import functions as F +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + BooleanType, + DoubleType, + LongType, + NumericType, ) if TYPE_CHECKING: from pyspark.pandas.frame import DataFrame - from pyspark.pandas.indexes.base import Index from pyspark.pandas.groupby import GroupBy + from pyspark.pandas.indexes.base import Index from pyspark.pandas.series import Series - from pyspark.pandas.window import Rolling, Expanding, ExponentialMoving + from pyspark.pandas.window import Expanding, ExponentialMoving, Rolling bool_type = bool @@ -3659,13 +3660,14 @@ def _count_expr(psser: "Series") -> Column: def _test() -> None: - import os import doctest + import os import shutil import sys import tempfile - from pyspark.sql import SparkSession + import pyspark.pandas.generic + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/groupby.py b/python/pyspark/pandas/groupby.py index 9907d95532458..667e272a0baff 100644 --- a/python/pyspark/pandas/groupby.py +++ b/python/pyspark/pandas/groupby.py @@ -19,19 +19,21 @@ A wrapper for GroupedData to behave like pandas GroupBy. """ -from abc import ABCMeta, abstractmethod import inspect +import warnings +from abc import ABCMeta, abstractmethod from collections import defaultdict, namedtuple from functools import partial, wraps from itertools import product from typing import ( + TYPE_CHECKING, Any, Callable, Dict, Generic, Iterator, - Mapping, List, + Mapping, Optional, Sequence, Set, @@ -40,68 +42,68 @@ TypeVar, Union, cast, - TYPE_CHECKING, ) -import warnings import pandas as pd -from pandas.api.types import is_number, is_hashable, is_list_like +from pandas.api.types import is_hashable, is_list_like, is_number -from pyspark.sql import Column, DataFrame as SparkDataFrame, Window, functions as F -from pyspark.sql.internal import InternalFunction as SF -from pyspark.sql.types import ( - BooleanType, - DataType, - DoubleType, - NumericType, - StructField, - StructType, - StringType, -) from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark._globals import _NoValue, _NoValueType from pyspark.loose_version import LooseVersion from pyspark.pandas._typing import Axis, FrameLike, Label, Name -from pyspark.pandas.typedef import infer_return_type, DataFrameType, ScalarType, SeriesType +from pyspark.pandas.config import get_option +from pyspark.pandas.correlation import ( + CORRELATION_CORR_OUTPUT_COLUMN, + CORRELATION_COUNT_OUTPUT_COLUMN, + CORRELATION_VALUE_1_COLUMN, + CORRELATION_VALUE_2_COLUMN, + compute, +) +from pyspark.pandas.exceptions import DataError from pyspark.pandas.frame import DataFrame from pyspark.pandas.internal import ( - InternalField, - InternalFrame, HIDDEN_COLUMNS, NATURAL_ORDER_COLUMN_NAME, - SPARK_INDEX_NAME_FORMAT, SPARK_DEFAULT_SERIES_NAME, + SPARK_INDEX_NAME_FORMAT, SPARK_INDEX_NAME_PATTERN, + InternalField, + InternalFrame, ) from pyspark.pandas.missing.groupby import ( MissingPandasLikeDataFrameGroupBy, MissingPandasLikeSeriesGroupBy, ) from pyspark.pandas.series import Series, first_series -from pyspark.pandas.config import get_option -from pyspark.pandas.correlation import ( - compute, - CORRELATION_VALUE_1_COLUMN, - CORRELATION_VALUE_2_COLUMN, - CORRELATION_CORR_OUTPUT_COLUMN, - CORRELATION_COUNT_OUTPUT_COLUMN, -) +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale +from pyspark.pandas.typedef import DataFrameType, ScalarType, SeriesType, infer_return_type from pyspark.pandas.utils import ( align_diff_frames, ansi_mode_context, is_name_like_tuple, is_name_like_value, + log_advice, name_like_string, same_anchor, scol_for, verify_temp_column_name, - log_advice, ) -from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale -from pyspark.pandas.exceptions import DataError +from pyspark.sql import Column, Window +from pyspark.sql import DataFrame as SparkDataFrame +from pyspark.sql import functions as F +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + BooleanType, + DataType, + DoubleType, + NumericType, + StringType, + StructField, + StructType, +) if TYPE_CHECKING: - from pyspark.pandas.window import RollingGroupby, ExpandingGroupby, ExponentialMovingGroupby + from pyspark.pandas.window import ExpandingGroupby, ExponentialMovingGroupby, RollingGroupby FuncT = TypeVar("FuncT", bound=Callable[..., Any]) @@ -5029,12 +5031,14 @@ def normalize_keyword_aggregation( def _test() -> None: - import os import doctest + import os import sys + import numpy - from pyspark.sql import SparkSession + import pyspark.pandas.groupby + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/indexes/base.py b/python/pyspark/pandas/indexes/base.py index 97bed387864de..07fb897e3a698 100644 --- a/python/pyspark/pandas/indexes/base.py +++ b/python/pyspark/pandas/indexes/base.py @@ -15,8 +15,10 @@ # limitations under the License. # +import warnings from functools import partial from typing import ( + TYPE_CHECKING, Any, Callable, Iterator, @@ -26,62 +28,61 @@ Union, cast, no_type_check, - TYPE_CHECKING, ) -import warnings -import pandas as pd import numpy as np +import pandas as pd +from pandas._libs import lib from pandas.api.types import ( - is_list_like, + CategoricalDtype, is_bool_dtype, - is_integer_dtype, is_float_dtype, + is_hashable, + is_integer_dtype, + is_list_like, is_numeric_dtype, is_object_dtype, ) from pandas.core.accessor import CachedAccessor # type: ignore[attr-defined] from pandas.io.formats.printing import pprint_thing # type: ignore[import-not-found] -from pandas.api.types import CategoricalDtype, is_hashable -from pandas._libs import lib -from pyspark.loose_version import LooseVersion -from pyspark.sql.column import Column -from pyspark.sql import functions as F -from pyspark.sql.types import ( - DayTimeIntervalType, - IntegralType, - TimestampType, - TimestampNTZType, -) from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. +from pyspark.loose_version import LooseVersion from pyspark.pandas._typing import Dtype, Label, Name, Scalar -from pyspark.pandas.config import get_option, option_context from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.config import get_option, option_context from pyspark.pandas.frame import DataFrame +from pyspark.pandas.internal import ( + DEFAULT_SERIES_NAME, + SPARK_DEFAULT_INDEX_NAME, + SPARK_INDEX_NAME_FORMAT, + InternalField, + InternalFrame, +) from pyspark.pandas.missing.indexes import MissingPandasLikeIndex from pyspark.pandas.series import Series, first_series from pyspark.pandas.spark.accessors import SparkIndexMethods from pyspark.pandas.utils import ( + ERROR_MESSAGE_CANNOT_COMBINE, is_ansi_mode_enabled, is_name_like_tuple, is_name_like_value, + log_advice, name_like_string, same_anchor, scol_for, - verify_temp_column_name, validate_bool_kwarg, validate_index_loc, - ERROR_MESSAGE_CANNOT_COMBINE, - log_advice, + verify_temp_column_name, xor, ) -from pyspark.pandas.internal import ( - InternalField, - InternalFrame, - DEFAULT_SERIES_NAME, - SPARK_DEFAULT_INDEX_NAME, - SPARK_INDEX_NAME_FORMAT, +from pyspark.sql import functions as F +from pyspark.sql.column import Column +from pyspark.sql.types import ( + DayTimeIntervalType, + IntegralType, + TimestampNTZType, + TimestampType, ) if TYPE_CHECKING: @@ -1847,8 +1848,8 @@ def append(self, other: "Index") -> "Index": ('b', 'y')], ) """ - from pyspark.pandas.indexes.multi import MultiIndex from pyspark.pandas.indexes.category import CategoricalIndex + from pyspark.pandas.indexes.multi import MultiIndex if isinstance(self, MultiIndex) != isinstance(other, MultiIndex): raise NotImplementedError( @@ -2652,13 +2653,15 @@ def __bool__(self) -> bool: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession - import pyspark.pandas.indexes.base + from pandas.util.version import Version + import pyspark.pandas.indexes.base + from pyspark.sql import SparkSession + os.chdir(os.environ["SPARK_HOME"]) if Version(np.__version__) >= Version("2"): diff --git a/python/pyspark/pandas/indexes/category.py b/python/pyspark/pandas/indexes/category.py index d4c709fe6284c..96830c6fe8a74 100644 --- a/python/pyspark/pandas/indexes/category.py +++ b/python/pyspark/pandas/indexes/category.py @@ -17,7 +17,7 @@ from typing import Any, Callable, List, Optional, Union, cast, no_type_check import pandas as pd -from pandas.api.types import is_hashable, CategoricalDtype +from pandas.api.types import CategoricalDtype, is_hashable from pyspark import pandas as ps from pyspark.pandas.indexes.base import Index @@ -639,11 +639,12 @@ def all(self, *args, **kwargs) -> None: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.indexes.category + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/indexes/datetimes.py b/python/pyspark/pandas/indexes/datetimes.py index eae709c9a2615..d866d052e0d3c 100644 --- a/python/pyspark/pandas/indexes/datetimes.py +++ b/python/pyspark/pandas/indexes/datetimes.py @@ -22,10 +22,10 @@ import pandas as pd from pandas.api.types import is_hashable from pandas.tseries.offsets import DateOffset -from pyspark._globals import _NoValue -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark._globals import _NoValue +from pyspark.loose_version import LooseVersion from pyspark.pandas import DataFrame from pyspark.pandas.indexes.base import Index from pyspark.pandas.missing.indexes import MissingPandasLikeDatetimeIndex @@ -866,11 +866,12 @@ def disallow_nanoseconds(freq: Union[str, DateOffset]) -> None: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.indexes.datetimes + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/indexes/multi.py b/python/pyspark/pandas/indexes/multi.py index e903244df1616..08f4fa20dbbde 100644 --- a/python/pyspark/pandas/indexes/multi.py +++ b/python/pyspark/pandas/indexes/multi.py @@ -21,13 +21,17 @@ import pandas as pd from pandas.api.types import is_hashable, is_list_like -from pyspark.sql import functions as F, Column as PySparkColumn, Window -from pyspark.sql.types import DataType from pyspark import pandas as ps from pyspark.pandas._typing import Label, Name, Scalar from pyspark.pandas.exceptions import PandasNotImplementedError from pyspark.pandas.frame import DataFrame from pyspark.pandas.indexes.base import Index +from pyspark.pandas.internal import ( + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, + InternalField, + InternalFrame, +) from pyspark.pandas.missing.indexes import MissingPandasLikeMultiIndex from pyspark.pandas.series import Series, first_series from pyspark.pandas.utils import ( @@ -35,16 +39,14 @@ is_name_like_tuple, name_like_string, scol_for, - verify_temp_column_name, validate_index_loc, + verify_temp_column_name, xor, ) -from pyspark.pandas.internal import ( - InternalField, - InternalFrame, - NATURAL_ORDER_COLUMN_NAME, - SPARK_INDEX_NAME_FORMAT, -) +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import Window +from pyspark.sql import functions as F +from pyspark.sql.types import DataType class MultiIndex(Index): @@ -1261,12 +1263,14 @@ def map( def _test() -> None: - import os import doctest + import os import sys + import numpy - from pyspark.sql import SparkSession + import pyspark.pandas.indexes.multi + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/indexes/timedelta.py b/python/pyspark/pandas/indexes/timedelta.py index 87b57d84b435d..5eeab70f561e9 100644 --- a/python/pyspark/pandas/indexes/timedelta.py +++ b/python/pyspark/pandas/indexes/timedelta.py @@ -15,12 +15,12 @@ # limitations under the License. # import warnings -from typing import cast, no_type_check, Any from functools import partial +from typing import Any, cast, no_type_check +import numpy as np import pandas as pd from pandas.api.types import is_hashable -import numpy as np from pyspark import pandas as ps from pyspark._globals import _NoValue diff --git a/python/pyspark/pandas/indexing.py b/python/pyspark/pandas/indexing.py index 236118a6d6136..de487c4fc08f6 100644 --- a/python/pyspark/pandas/indexing.py +++ b/python/pyspark/pandas/indexing.py @@ -22,27 +22,24 @@ from abc import ABCMeta, abstractmethod from collections.abc import Iterable from functools import reduce -from typing import Any, Optional, List, Tuple, TYPE_CHECKING, Union, cast, Sized +from typing import TYPE_CHECKING, Any, List, Optional, Sized, Tuple, Union, cast +import numpy as np import pandas as pd from pandas.api.types import is_list_like -import numpy as np -from pyspark.loose_version import LooseVersion -from pyspark.sql import functions as F, Column as PySparkColumn -from pyspark.sql.types import BooleanType, LongType, DataType -from pyspark.sql.utils import is_remote -from pyspark.errors import AnalysisException from pyspark import pandas as ps # noqa: F401 +from pyspark.errors import AnalysisException +from pyspark.loose_version import LooseVersion from pyspark.pandas._typing import Label, Name, Scalar +from pyspark.pandas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError from pyspark.pandas.internal import ( DEFAULT_SERIES_NAME, - InternalField, - InternalFrame, NATURAL_ORDER_COLUMN_NAME, SPARK_DEFAULT_SERIES_NAME, + InternalField, + InternalFrame, ) -from pyspark.pandas.exceptions import SparkPandasIndexingError, SparkPandasNotImplementedError from pyspark.pandas.utils import ( is_name_like_tuple, is_name_like_value, @@ -53,6 +50,10 @@ spark_column_equals, verify_temp_column_name, ) +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F +from pyspark.sql.types import BooleanType, DataType, LongType +from pyspark.sql.utils import is_remote if TYPE_CHECKING: from pyspark.pandas.frame import DataFrame @@ -1887,13 +1888,15 @@ def __setitem__(self, key: Any, value: Any) -> None: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession - import pyspark.pandas.indexing + from pandas.util.version import Version + import pyspark.pandas.indexing + from pyspark.sql import SparkSession + os.chdir(os.environ["SPARK_HOME"]) if Version(np.__version__) >= Version("2"): diff --git a/python/pyspark/pandas/internal.py b/python/pyspark/pandas/internal.py index d24402c46b68b..5c324ac582093 100644 --- a/python/pyspark/pandas/internal.py +++ b/python/pyspark/pandas/internal.py @@ -20,34 +20,17 @@ """ import re -from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union, cast import numpy as np import pandas as pd from pandas.api.types import CategoricalDtype, is_integer_dtype # noqa: F401 -from pyspark._globals import _NoValue, _NoValueType -from pyspark.sql import ( - functions as F, - Column as PySparkColumn, - DataFrame as PySparkDataFrame, - Window, -) -from pyspark.sql.types import ( # noqa: F401 - _drop_metadata, - BooleanType, - DataType, - LongType, - StructField, - StructType, - StringType, -) -from pyspark.sql.utils import is_timestamp_ntz_preferred, is_remote from pyspark import pandas as ps -from pyspark.sql.internal import InternalFunction as SF +from pyspark._globals import _NoValue, _NoValueType from pyspark.pandas._typing import Label -from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale from pyspark.pandas.data_type_ops.base import DataTypeOps +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale from pyspark.pandas.typedef import ( Dtype, as_spark_type, @@ -65,6 +48,29 @@ scol_for, spark_column_equals, ) +from pyspark.sql import ( + Column as PySparkColumn, +) +from pyspark.sql import ( + DataFrame as PySparkDataFrame, +) +from pyspark.sql import ( + Window, +) +from pyspark.sql import ( + functions as F, +) +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( # noqa: F401 + BooleanType, + DataType, + LongType, + StringType, + StructField, + StructType, + _drop_metadata, +) +from pyspark.sql.utils import is_remote, is_timestamp_ntz_preferred if TYPE_CHECKING: from pyspark.pandas.series import Series @@ -1630,11 +1636,12 @@ def prepare_pandas_frame( def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.internal + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/missing/frame.py b/python/pyspark/pandas/missing/frame.py index bdfa7574dc3d3..6f04aab5e8219 100644 --- a/python/pyspark/pandas/missing/frame.py +++ b/python/pyspark/pandas/missing/frame.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.pandas.missing import unsupported_function, unsupported_property, common +from pyspark.pandas.missing import common, unsupported_function, unsupported_property def _unsupported_function(method_name, deprecated=False, reason=""): diff --git a/python/pyspark/pandas/missing/indexes.py b/python/pyspark/pandas/missing/indexes.py index 2419908b3129a..4d6daa7d33b65 100644 --- a/python/pyspark/pandas/missing/indexes.py +++ b/python/pyspark/pandas/missing/indexes.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.pandas.missing import unsupported_function, unsupported_property, common +from pyspark.pandas.missing import common, unsupported_function, unsupported_property def _unsupported_function(method_name, deprecated=False, reason="", cls="Index"): diff --git a/python/pyspark/pandas/missing/series.py b/python/pyspark/pandas/missing/series.py index 08f21f46b2cc1..20afd98923f0e 100644 --- a/python/pyspark/pandas/missing/series.py +++ b/python/pyspark/pandas/missing/series.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.pandas.missing import unsupported_function, unsupported_property, common +from pyspark.pandas.missing import common, unsupported_function, unsupported_property def _unsupported_function(method_name, deprecated=False, reason=""): diff --git a/python/pyspark/pandas/mlflow.py b/python/pyspark/pandas/mlflow.py index 06988a2871c6b..34a868a0c2d52 100644 --- a/python/pyspark/pandas/mlflow.py +++ b/python/pyspark/pandas/mlflow.py @@ -19,19 +19,18 @@ MLflow-related functions to load models and apply them to pandas-on-Spark dataframes. """ -from typing import List, Union -from typing import Any +from typing import Any, List, Union -import pandas as pd import numpy as np +import pandas as pd -from pyspark.sql.types import DataType -from pyspark.sql.functions import struct -from pyspark.pandas._typing import Label, Dtype -from pyspark.pandas.utils import lazy_property, default_session +from pyspark.pandas._typing import Dtype, Label from pyspark.pandas.frame import DataFrame from pyspark.pandas.series import Series, first_series from pyspark.pandas.typedef import as_spark_type +from pyspark.pandas.utils import default_session, lazy_property +from pyspark.sql.functions import struct +from pyspark.sql.types import DataType __all__ = ["PythonModelWrapper", "load_model"] @@ -204,11 +203,12 @@ def load_model( def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.mlflow + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/namespace.py b/python/pyspark/pandas/namespace.py index 05875347a3759..99e07ee90ae80 100644 --- a/python/pyspark/pandas/namespace.py +++ b/python/pyspark/pandas/namespace.py @@ -19,7 +19,15 @@ Wrappers around spark that correspond to common pandas functions. """ +import json +import pickle +import warnings +from collections.abc import Iterable +from datetime import tzinfo +from functools import reduce +from io import BytesIO from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -30,81 +38,74 @@ Sized, Tuple, Type, - TYPE_CHECKING, Union, cast, no_type_check, ) -from collections.abc import Iterable -from datetime import tzinfo -from functools import reduce -from io import BytesIO -import pickle -import json -import warnings import numpy as np import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq from pandas.api.types import ( is_datetime64_dtype, is_list_like, ) from pandas.tseries.offsets import DateOffset -import pyarrow as pa -import pyarrow.parquet as pq +from pyspark import pandas as ps from pyspark._globals import _NoValue, _NoValueType from pyspark.loose_version import LooseVersion -from pyspark.sql import functions as F, Column as PySparkColumn -from pyspark.sql.functions import pandas_udf -from pyspark.sql.types import ( - ByteType, - ShortType, - IntegerType, - LongType, - FloatType, - DoubleType, - BooleanType, - NumericType, - TimestampType, - TimestampNTZType, - DecimalType, - StringType, - DateType, - StructType, - StructField, - DataType, -) -from pyspark.sql.dataframe import DataFrame as PySparkDataFrame -from pyspark import pandas as ps from pyspark.pandas._typing import Axis, Dtype, Label, Name from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.config import get_option +from pyspark.pandas.frame import DataFrame, _reduce_spark_multi +from pyspark.pandas.indexes import DatetimeIndex, Index, TimedeltaIndex +from pyspark.pandas.indexes.multi import MultiIndex +from pyspark.pandas.internal import ( + DEFAULT_SERIES_NAME, + HIDDEN_COLUMNS, + NATURAL_ORDER_COLUMN_NAME, + SPARK_INDEX_NAME_FORMAT, + InternalField, + InternalFrame, +) +from pyspark.pandas.series import Series, first_series +from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale from pyspark.pandas.utils import ( align_diff_frames, default_session, is_ansi_mode_enabled, is_name_like_tuple, is_name_like_value, + log_advice, name_like_string, same_anchor, scol_for, validate_axis, - log_advice, ) -from pyspark.pandas.config import get_option -from pyspark.pandas.frame import DataFrame, _reduce_spark_multi -from pyspark.pandas.internal import ( - InternalFrame, - InternalField, - DEFAULT_SERIES_NAME, - HIDDEN_COLUMNS, - SPARK_INDEX_NAME_FORMAT, - NATURAL_ORDER_COLUMN_NAME, +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import functions as F +from pyspark.sql.dataframe import DataFrame as PySparkDataFrame +from pyspark.sql.functions import pandas_udf +from pyspark.sql.types import ( + BooleanType, + ByteType, + DataType, + DateType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + NumericType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, ) -from pyspark.pandas.series import Series, first_series -from pyspark.pandas.spark.utils import as_nullable_spark_type, force_decimal_precision_scale -from pyspark.pandas.indexes import Index, DatetimeIndex, TimedeltaIndex -from pyspark.pandas.indexes.multi import MultiIndex if TYPE_CHECKING: from pandas._typing import HTMLFlavors @@ -4046,16 +4047,18 @@ def _get_index_map( def _test() -> None: - import os import doctest + import os import shutil import sys import tempfile import uuid - from pyspark.sql import SparkSession - import pyspark.pandas.namespace + from pandas.util.version import Version + import pyspark.pandas.namespace + from pyspark.sql import SparkSession + os.chdir(os.environ["SPARK_HOME"]) if Version(np.__version__) >= Version("2"): diff --git a/python/pyspark/pandas/numpy_compat.py b/python/pyspark/pandas/numpy_compat.py index 09194cbee5e05..05dbc7c1c7b44 100644 --- a/python/pyspark/pandas/numpy_compat.py +++ b/python/pyspark/pandas/numpy_compat.py @@ -14,14 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Callable, no_type_check +from typing import Any, Callable, Tuple, Union, no_type_check import numpy as np +from pyspark.loose_version import LooseVersion +from pyspark.pandas._typing import SeriesOrIndex +from pyspark.pandas.base import IndexOpsMixin +from pyspark.sql import Column from pyspark.sql import functions as F from pyspark.sql.pandas.functions import pandas_udf -from pyspark.sql.types import DoubleType, LongType, BooleanType -from pyspark.pandas.base import IndexOpsMixin +from pyspark.sql.types import BooleanType, DoubleType unary_np_spark_mappings = { "abs": F.abs, @@ -50,25 +53,66 @@ "frexp": lambda _: NotImplemented, # 'frexp' output lengths become different # and it cannot be supported via pandas UDF. "invert": F.bitwise_not, - "isfinite": lambda c: c != float("inf"), - "isinf": lambda c: c == float("inf"), + "isfinite": lambda c: F.coalesce( + ~(F.isnan(c) | (c == float("inf")) | (c == float("-inf"))), F.lit(False) + ), + "isinf": lambda c: F.coalesce((c == float("inf")) | (c == float("-inf")), F.lit(False)), "isnan": F.isnan, "isnat": lambda c: NotImplemented, # pandas-on-Spark and PySpark does not have Nat concept. "log": F.log, "log10": F.log10, "log1p": F.log1p, + "log2": lambda c: F.when(c == 0, F.lit(float("-inf"))).otherwise(F.log2(c)), "logical_not": lambda c: ~(c.cast(BooleanType())), "matmul": lambda _: NotImplemented, # Can return a NumPy array in pandas. "negative": F.negative, "positive": F.positive, "rad2deg": F.degrees, "radians": F.radians, - "reciprocal": pandas_udf( # type: ignore[call-overload] - lambda s: np.reciprocal(s), DoubleType() + "reciprocal": lambda c: F.when( + # Floating-point and decimal inputs take a true reciprocal; numpy + # applies it element-wise to Decimal objects as well. + F.typeof(c).isin("float", "double") | F.typeof(c).startswith("decimal"), + F.when(c.isNull(), c.cast("double")) + .when( + # Cast to double so the zero check also analyzes for the integer and + # boolean columns that fall through to the otherwise branch (Spark + # type-checks every branch of the CASE, not just the taken one). + c.cast("double") == 0, + F.when(c.cast("string") == "-0.0", F.lit(float("-inf"))).otherwise(F.lit(float("inf"))), + ) + .otherwise(F.lit(1.0) / c.cast("double")), + ).otherwise( + # Integer and boolean inputs: numpy does integer division (truncated + # toward zero), so only +/-1 survive and every other magnitude -> 0. + # Dividing by 0 overflows to the width-specific integer minimum for int + # (int32) and bigint (int64), while narrower widths (tinyint, smallint, + # and boolean promoted to int8) return 0. Cast through long so boolean + # and narrower integers can take part in the division. + F.when( + c.cast("long") == 0, + F.when(F.typeof(c) == "int", F.lit(float(np.iinfo(np.int32).min))) + .when(F.typeof(c) == "bigint", F.lit(float(np.iinfo(np.int64).min))) + .otherwise(F.lit(0.0)), + ).otherwise((F.lit(1) / c.cast("long")).cast("long").cast("double")) ), - "rint": pandas_udf(lambda s: np.rint(s), DoubleType()), # type: ignore[call-overload] + "rint": lambda c: F.rint(c.cast("double")), "sign": F.signum, - "signbit": lambda c: F.when(c < 0, True).otherwise(False), + "signbit": lambda c: F.when( + # A genuine <NA> from a nullable dtype (e.g. Int64) arrives as a non-floating null + # and must propagate. A NaN from a default (numpy-backed) dtype arrives as a floating + # null and must map to False (np.signbit(nan) is False); it falls through to + # otherwise(False) below. Two cases this expression cannot match, seeing only the Spark + # value and not the pandas dtype: a nullable float dtype's <NA> (Float32 or Float64) is + # also a floating null, indistinguishable from a NaN after from_pandas, so it reads False + # instead of propagating; and the sign of a NaN never reaches here (from_pandas nulls a + # NaN, and a NaN computed in Spark arrives as +NaN), so a negative NaN reads False where + # np.signbit reports True. + c.isNull() & ~F.typeof(c).isin("float", "double"), + F.lit(None).cast("boolean"), + ) + .when((c < 0) | (c.cast("string") == "-0.0"), True) + .otherwise(False), "sin": F.sin, "sinh": F.sinh, "spacing": pandas_udf(lambda s: np.spacing(s), DoubleType()), # type: ignore[call-overload] @@ -76,44 +120,179 @@ "square": lambda c: c.cast("double") * c, "tan": F.tan, "tanh": F.tanh, - "trunc": pandas_udf(lambda s: np.trunc(s), DoubleType()), # type: ignore[call-overload] + "trunc": lambda c: F.when( + c.cast("double").isNull() + | F.isnan(c.cast("double")) + | c.cast("double").isin(float("-inf"), float("inf")), + c.cast("double"), + ).otherwise( + F.signum(c.cast("double")) + * (F.abs(c.cast("double")) - (F.abs(c.cast("double")) % F.lit(1.0))) + ), } + +def _copysign_func(c1: Column, c2: Column) -> Column: + # Sign of y is taken from its IEEE-754 sign bit, so -0.0 counts as negative. + # c2 < 0 misses -0.0, so detect it via the string cast, the same way the + # 'reciprocal' mapping distinguishes -0.0 from 0.0. NaN's sign bit is positive + # and c2 < 0 is already false for NaN, so it correctly falls through to +1.0. + sign = F.when((c2 < 0) | (c2.cast("string") == "-0.0"), F.lit(-1.0)).otherwise(F.lit(1.0)) + # An integer y column's NULL is a genuine missing value and propagates. A + # float/double column instead stores its missing value as NaN (surfaced as a + # Spark NULL by pandas-on-Spark), for which copysign(x, NaN) returns |x|. A + # nullable Float64 <NA> collapses to that same Spark NULL, so it is likewise + # treated as NaN and returns |x| rather than propagating; this cannot be + # distinguished here and matches the prior pandas_udf behavior. + return F.when( + c2.isNull() & ~F.typeof(c2).isin("float", "double"), F.lit(None).cast("double") + ).otherwise(F.abs(c1.cast("double")) * sign) + + +def _fmod_func(c1: Column, c2: Column) -> Column: + c1_double = c1.cast("double") + c2_double = c2.cast("double") + + return F.when( + F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), + F.when(c1.isNull() | F.isnan(c1), c1_double) + .when(c2.isNull() | F.isnan(c2), c2_double) + .when(c2_double == 0, F.lit(float("nan"))) + .otherwise(F.try_mod(c1_double, c2_double)), + ).otherwise( + F.when(c1.isNull() | F.isnan(c1), c1_double) + .when(c2.isNull() | F.isnan(c2), c2_double) + .when(c2_double == 0, F.lit(0.0)) + .otherwise(F.try_mod(c1_double, c2_double)) + ) + + +def _logaddexp_func(c1: Column, c2: Column, base2: bool = False) -> Column: + c1_double = c1.cast("double") + c2_double = c2.cast("double") + difference = F.abs(c1_double - c2_double) + maximum = F.greatest(c1_double, c2_double) + if base2: + log_term = F.log1p(F.pow(F.lit(2.0), -difference)) / F.log(F.lit(2.0)) + else: + log_term = F.log1p(F.exp(-difference)) + + return ( + F.when(c1_double.isNull() | F.isnan(c1_double), c1_double) + .when(c2_double.isNull() | F.isnan(c2_double), c2_double) + .when((c1_double == float("inf")) | (c2_double == float("inf")), F.lit(float("inf"))) + .when(c1_double == float("-inf"), c2_double + F.lit(0.0)) + .when(c2_double == float("-inf"), c1_double + F.lit(0.0)) + .otherwise(maximum + log_term) + ) + + +def _floor_divide_func(c1: Column, c2: Column) -> Column: + c1_double = c1.cast("double") + c2_double = c2.cast("double") + + return F.when( + F.typeof(c1).isin("float", "double") | F.typeof(c2).isin("float", "double"), + F.when(c1.isNull() | F.isnan(c1), c1_double) + .when(c2.isNull() | F.isnan(c2), c2_double) + .when( + c1_double.isin(float("-inf"), float("inf")), + F.when( + c2_double == 0, + F.when( + (c1_double < 0) != (c2_double.cast("string") == "-0.0"), + F.lit(float("-inf")), + ).otherwise(F.lit(float("inf"))), + ).otherwise(F.lit(float("nan"))), + ) + .when( + c2_double.isin(float("-inf"), float("inf")), + F.when(c1_double == 0, c1_double / c2_double) + .when((c1_double < 0) != (c2_double < 0), F.lit(-1.0)) + .otherwise(F.lit(0.0)), + ) + .when( + c2_double == 0, + F.when(c1_double == 0, F.lit(float("nan"))) + .when( + (c1_double < 0) != (c2_double.cast("string") == "-0.0"), + F.lit(float("-inf")), + ) + .otherwise(F.lit(float("inf"))), + ) + .when(c1_double == 0, c1_double / c2_double) + .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double, F.lit(1.0))), + ).otherwise( + # np.floor_divide on pandas Series returns IEEE values for an integral zero divisor. + F.when(c1.isNull() | F.isnan(c1), c1_double) + .when(c2.isNull() | F.isnan(c2), c2_double) + .when( + c2_double == 0, + F.when(c1_double == 0, F.lit(float("nan"))) + .when(c1_double < 0, F.lit(float("-inf"))) + .otherwise(F.lit(float("inf"))), + ) + .otherwise((c1_double / c2_double) - F.pmod(c1_double / c2_double, F.lit(1.0))) + ) + + +# NumPy 2.3.0 changed how fmax/fmin break a signed-zero tie: for equal operands +# (for example +0.0 and -0.0) it returns the first operand, while older versions +# returned the second. Track the installed NumPy so the result keeps the matching +# sign of zero. +_tie_returns_first_operand = LooseVersion(np.__version__) >= LooseVersion("2.3.0") + + +def _fmax_func(c1: Column, c2: Column) -> Column: + tie = c1 if _tie_returns_first_operand else c2 + return ( + F.when(F.isnan(c1.cast("double")), c2) + .when(F.isnan(c2.cast("double")), c1) + .when(c1 == c2, tie) + .otherwise(F.greatest(c1, c2)) + .cast("double") + ) + + +def _fmin_func(c1: Column, c2: Column) -> Column: + tie = c1 if _tie_returns_first_operand else c2 + return F.when(c1 == c2, tie).otherwise(F.least(c1, c2)).cast("double") + + binary_np_spark_mappings = { "arctan2": F.atan2, "bitwise_and": lambda c1, c2: c1.bitwiseAND(c2), "bitwise_or": lambda c1, c2: c1.bitwiseOR(c2), "bitwise_xor": lambda c1, c2: c1.bitwiseXOR(c2), - "copysign": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.copysign(s1, s2), DoubleType() - ), - "float_power": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.float_power(s1, s2), DoubleType() - ), - "floor_divide": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.floor_divide(s1, s2), DoubleType() - ), - "fmax": pandas_udf(lambda s1, s2: np.fmax(s1, s2), DoubleType()), # type: ignore[call-overload] - "fmin": pandas_udf(lambda s1, s2: np.fmin(s1, s2), DoubleType()), # type: ignore[call-overload] - "fmod": pandas_udf(lambda s1, s2: np.fmod(s1, s2), DoubleType()), # type: ignore[call-overload] + "copysign": _copysign_func, + "float_power": lambda c1, c2: F.pow(c1.cast("double"), c2.cast("double")), + # np.floor_divide dispatches to the pandas-on-Spark floordiv dunder operation + # before this registry is consulted, so this mapping is not used for that case. + "floor_divide": _floor_divide_func, + "fmax": _fmax_func, + "fmin": _fmin_func, + "fmod": _fmod_func, "gcd": pandas_udf(lambda s1, s2: np.gcd(s1, s2), DoubleType()), # type: ignore[call-overload] - "heaviside": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.heaviside(s1, s2), DoubleType() - ), + "heaviside": lambda c1, c2: F.when( + c1.isNull() | F.isnan(c1.cast("double")), + c1.cast("double"), + ) + .when(c1 < 0, F.lit(0.0)) + .when(c1 == 0, c2.cast("double")) + .otherwise(F.lit(1.0)), "hypot": F.hypot, "lcm": pandas_udf(lambda s1, s2: np.lcm(s1, s2), DoubleType()), # type: ignore[call-overload] - "ldexp": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.ldexp(s1, s2), DoubleType() - ), - "left_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.left_shift(s1, s2), LongType() - ), - "logaddexp": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.logaddexp(s1, s2), DoubleType() - ), - "logaddexp2": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.logaddexp2(s1, s2), DoubleType() + "ldexp": lambda c1, c2: F.when( + c1.cast("double").isin(0.0, float("-inf"), float("inf")), + c1.cast("double"), + ).otherwise(c1.cast("double") * F.pow(F.lit(2.0), c2)), + # F.shiftleft accepts literal counts only; call_function also accepts a column. + # NumPy returns zero for counts outside an int64's bit width, unlike JVM shifts. + "left_shift": lambda c1, c2: F.when((c2 < 0) | (c2 >= 64), F.lit(0)).otherwise( + F.call_function("shiftleft", c1, c2) ), + "logaddexp": _logaddexp_func, + "logaddexp2": lambda c1, c2: _logaddexp_func(c1, c2, base2=True), "logical_and": lambda c1, c2: c1.cast(BooleanType()) & c2.cast(BooleanType()), "logical_or": lambda c1, c2: c1.cast(BooleanType()) | c2.cast(BooleanType()), "logical_xor": lambda c1, c2: ( @@ -123,13 +302,38 @@ ), "maximum": F.greatest, "minimum": F.least, - "modf": pandas_udf(lambda s1, s2: np.modf(s1, s2), DoubleType()), # type: ignore[call-overload] "nextafter": pandas_udf( # type: ignore[call-overload] lambda s1, s2: np.nextafter(s1, s2), DoubleType() ), - "right_shift": pandas_udf( # type: ignore[call-overload] - lambda s1, s2: np.right_shift(s1, s2), LongType() - ), + # F.shiftright accepts literal counts only; call_function also accepts a column. + # NumPy sign-extends counts outside an int64's bit width, unlike JVM shifts. + "right_shift": lambda c1, c2: F.when( + (c2 < 0) | (c2 >= 64), F.call_function("shiftright", c1, F.lit(63)) + ).otherwise(F.call_function("shiftright", c1, c2)), +} + + +def _modf_fractional_func(c: Column) -> Column: + c_double = c.cast("double") + # signum * (abs % 1) keeps the fractional magnitude with the sign of the input, + # including the signed zero of a whole number (for example -2.0 -> -0.0), the same + # way the "trunc" mapping (reused below for the integral part) relies on signum. + fractional = F.signum(c_double) * (F.abs(c_double) % F.lit(1.0)) + return ( + F.when(c.isNull() | F.isnan(c_double), c_double) + # +-inf has no fractional part; numpy returns a zero with the input's sign. + .when(c_double == float("inf"), F.lit(0.0)) + .when(c_double == float("-inf"), F.lit(-0.0)) + .otherwise(fractional) + ) + + +# Every multi-output ufunc numpy ships (modf, frexp) has exactly two outputs, so each entry +# maps to a pair of Column->Column functions applied independently and returned as a 2-tuple +# that numpy's __array_ufunc__ unpacks (for example `fractional, integral = np.modf(series)`). +multi_output_np_spark_mappings = { + # np.modf(x) -> (fractional part, integral part); the integral part is exactly trunc. + "modf": (_modf_fractional_func, unary_np_spark_mappings["trunc"]), } @@ -137,7 +341,7 @@ # See also https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses def maybe_dispatch_ufunc_to_dunder_op( ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any -) -> IndexOpsMixin: +) -> SeriesOrIndex: special = { "add", "sub", @@ -203,11 +407,22 @@ def not_implemented(*args, **kwargs): # See also https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#standard-array-subclasses def maybe_dispatch_ufunc_to_spark_func( ser_or_index: IndexOpsMixin, ufunc: Callable, method: str, *inputs: Any, **kwargs: Any -) -> IndexOpsMixin: +) -> Union[SeriesOrIndex, Tuple[SeriesOrIndex, SeriesOrIndex]]: from pyspark.pandas.base import column_op op_name = ufunc.__name__ + if ( + method == "__call__" + and op_name in multi_output_np_spark_mappings + and kwargs.get("out") is None + ): + # These ufuncs are unary in their input, so the single input is always a Series + # that column_op unwraps to a Column -- no literal wrapping needed. Build one + # Series per output and return them as a 2-tuple (see the mapping's docstring). + first_func, second_func = multi_output_np_spark_mappings[op_name] + return column_op(first_func)(*inputs), column_op(second_func)(*inputs) + if ( method == "__call__" and (op_name in unary_np_spark_mappings or op_name in binary_np_spark_mappings) @@ -228,11 +443,12 @@ def convert_arguments(*args): def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.numpy_compat + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/plot/core.py b/python/pyspark/pandas/plot/core.py index c9998ffddfd8b..2fd6528b6ced1 100644 --- a/python/pyspark/pandas/plot/core.py +++ b/python/pyspark/pandas/plot/core.py @@ -18,16 +18,17 @@ import importlib import math -import pandas as pd import numpy as np +import pandas as pd from pandas.core.base import PandasObject # type: ignore[attr-defined] from pandas.core.dtypes.inference import is_integer -from pyspark.sql import functions as F, Column -from pyspark.sql.internal import InternalFunction as SF -from pyspark.pandas.missing import unsupported_function from pyspark.pandas.config import get_option +from pyspark.pandas.missing import unsupported_function from pyspark.pandas.utils import name_like_string +from pyspark.sql import Column +from pyspark.sql import functions as F +from pyspark.sql.internal import InternalFunction as SF class TopNPlotBase: @@ -527,6 +528,7 @@ def _get_plot_backend(backend=None): try: # test if matplotlib can be imported import matplotlib # noqa: F401 + from pyspark.pandas.plot import matplotlib as module except ImportError: raise ImportError( @@ -539,6 +541,7 @@ def _get_plot_backend(backend=None): try: # test if plotly can be imported import plotly # noqa: F401 + from pyspark.pandas.plot import plotly as module except ImportError: raise ImportError( diff --git a/python/pyspark/pandas/plot/matplotlib.py b/python/pyspark/pandas/plot/matplotlib.py index d0fcf65117cc9..5b2fb7eb7978e 100644 --- a/python/pyspark/pandas/plot/matplotlib.py +++ b/python/pyspark/pandas/plot/matplotlib.py @@ -17,36 +17,35 @@ from typing import final -from pyspark.loose_version import LooseVersion - import matplotlib as mat import numpy as np +import pandas as pd from matplotlib.axes._base import _process_plot_format # type: ignore[attr-defined] from matplotlib.figure import Figure -import pandas as pd from pandas.core.dtypes.inference import is_list_like from pandas.io.formats.printing import pprint_thing # type: ignore[import-not-found] -from pandas.plotting._matplotlib import ( # type: ignore[import-not-found] - BarPlot as PandasBarPlot, - BoxPlot as PandasBoxPlot, - HistPlot as PandasHistPlot, - PiePlot as PandasPiePlot, - AreaPlot as PandasAreaPlot, - LinePlot as PandasLinePlot, - BarhPlot as PandasBarhPlot, - ScatterPlot as PandasScatterPlot, - KdePlot as PandasKdePlot, -) from pandas.plotting._core import PlotAccessor -from pandas.plotting._matplotlib.core import MPLPlot as PandasMPLPlot # type: ignore[import-not-found] +from pandas.plotting._matplotlib import AreaPlot as PandasAreaPlot # type: ignore[import-not-found] +from pandas.plotting._matplotlib import BarhPlot as PandasBarhPlot +from pandas.plotting._matplotlib import BarPlot as PandasBarPlot +from pandas.plotting._matplotlib import BoxPlot as PandasBoxPlot +from pandas.plotting._matplotlib import HistPlot as PandasHistPlot +from pandas.plotting._matplotlib import KdePlot as PandasKdePlot +from pandas.plotting._matplotlib import LinePlot as PandasLinePlot +from pandas.plotting._matplotlib import PiePlot as PandasPiePlot +from pandas.plotting._matplotlib import ScatterPlot as PandasScatterPlot +from pandas.plotting._matplotlib.core import ( # type: ignore[import-not-found] + MPLPlot as PandasMPLPlot, +) +from pyspark.loose_version import LooseVersion from pyspark.pandas.plot import ( - TopNPlotBase, - SampledPlotBase, - HistogramPlotBase, BoxPlotBase, - unsupported_function, + HistogramPlotBase, KdePlotBase, + SampledPlotBase, + TopNPlotBase, + unsupported_function, ) from pyspark.pandas.series import Series, first_series diff --git a/python/pyspark/pandas/plot/plotly.py b/python/pyspark/pandas/plot/plotly.py index d8b09c97744b1..ffa343ee08bf8 100644 --- a/python/pyspark/pandas/plot/plotly.py +++ b/python/pyspark/pandas/plot/plotly.py @@ -21,11 +21,11 @@ import pandas as pd from pyspark.pandas.plot import ( - HistogramPlotBase, - name_like_string, - PandasOnSparkPlotAccessor, BoxPlotBase, + HistogramPlotBase, KdePlotBase, + PandasOnSparkPlotAccessor, + name_like_string, ) if TYPE_CHECKING: @@ -50,9 +50,9 @@ def plot_pandas_on_spark(data: Union["ps.DataFrame", "ps.Series"], kind: str, ** def plot_pie(data: Union["ps.DataFrame", "ps.Series"], **kwargs): + import plotly.graph_objs as go from plotly import express from plotly.subplots import make_subplots - import plotly.graph_objs as go data = PandasOnSparkPlotAccessor.pandas_plot_data_map["pie"](data) subplots = kwargs.pop("subplots", False) @@ -99,6 +99,7 @@ def plot_pie(data: Union["ps.DataFrame", "ps.Series"], **kwargs): def plot_histogram(data: Union["ps.DataFrame", "ps.Series"], **kwargs): import plotly.graph_objs as go + import pyspark.pandas as ps bins = kwargs.get("bins", 10) @@ -147,6 +148,7 @@ def plot_histogram(data: Union["ps.DataFrame", "ps.Series"], **kwargs): def plot_box(data: Union["ps.DataFrame", "ps.Series"], **kwargs): import plotly.graph_objs as go + import pyspark.pandas as ps from pyspark.sql.types import NumericType @@ -241,6 +243,7 @@ def plot_box(data: Union["ps.DataFrame", "ps.Series"], **kwargs): def plot_kde(data: Union["ps.DataFrame", "ps.Series"], **kwargs): from plotly import express + import pyspark.pandas as ps if isinstance(data, ps.DataFrame) and "color" not in kwargs: diff --git a/python/pyspark/pandas/resample.py b/python/pyspark/pandas/resample.py index bdb939cffd0ad..a4ce2ac6f2d7b 100644 --- a/python/pyspark/pandas/resample.py +++ b/python/pyspark/pandas/resample.py @@ -33,21 +33,13 @@ import pandas as pd from pandas.tseries.frequencies import to_offset -from pyspark.sql import Column, functions as F -from pyspark.sql.internal import InternalFunction as SF -from pyspark.sql.types import ( - NumericType, - StructField, - TimestampNTZType, - DataType, -) from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark.pandas._typing import FrameLike from pyspark.pandas.frame import DataFrame from pyspark.pandas.internal import ( + SPARK_DEFAULT_INDEX_NAME, InternalField, InternalFrame, - SPARK_DEFAULT_INDEX_NAME, ) from pyspark.pandas.missing.resample import ( MissingPandasLikeDataFrameResampler, @@ -58,6 +50,15 @@ scol_for, verify_temp_column_name, ) +from pyspark.sql import Column +from pyspark.sql import functions as F +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + DataType, + NumericType, + StructField, + TimestampNTZType, +) class Resampler(Generic[FrameLike], metaclass=ABCMeta): @@ -766,11 +767,12 @@ def _handle_output(self, psdf: DataFrame) -> Series: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.resample + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/series.py b/python/pyspark/pandas/series.py index 159a7d01a135f..0c9d73e313076 100644 --- a/python/pyspark/pandas/series.py +++ b/python/pyspark/pandas/series.py @@ -20,17 +20,18 @@ """ import datetime -import re import inspect +import re import warnings from collections.abc import Mapping from functools import partial, reduce, wraps from typing import ( + IO, + TYPE_CHECKING, Any, Callable, Dict, Generic, - IO, Iterable, List, Literal, @@ -43,78 +44,68 @@ cast, no_type_check, overload, - TYPE_CHECKING, ) import numpy as np import pandas as pd -from pandas.core.accessor import CachedAccessor # type: ignore[attr-defined] -from pandas.io.formats.printing import pprint_thing # type: ignore[import-not-found] from pandas.api.extensions import no_default from pandas.api.types import ( - is_list_like, + CategoricalDtype, is_hashable, + is_list_like, is_numeric_dtype, - CategoricalDtype, ) +from pandas.core.accessor import CachedAccessor # type: ignore[attr-defined] +from pandas.io.formats.printing import pprint_thing # type: ignore[import-not-found] from pandas.tseries.frequencies import DateOffset # type: ignore[attr-defined] +from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark._globals import _NoValue, _NoValueType from pyspark.loose_version import LooseVersion -from pyspark.sql import ( - functions as F, - Column as PySparkColumn, - DataFrame as SparkDataFrame, - Window as PySparkWindow, -) -from pyspark.sql.internal import InternalFunction as SF -from pyspark.sql.types import ( - ArrayType, - BooleanType, - DecimalType, - DoubleType, - FloatType, - IntegerType, - LongType, - NumericType, - Row, - StructType, - TimestampType, - NullType, -) -from pyspark.sql.window import Window -from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. from pyspark.pandas._typing import Axis, Dtype, Label, Name, Scalar, T from pyspark.pandas.accessors import PandasOnSparkSeriesMethods +from pyspark.pandas.base import IndexOpsMixin from pyspark.pandas.categorical import CategoricalAccessor from pyspark.pandas.config import get_option from pyspark.pandas.correlation import ( - compute, - CORRELATION_VALUE_1_COLUMN, - CORRELATION_VALUE_2_COLUMN, CORRELATION_CORR_OUTPUT_COLUMN, CORRELATION_COUNT_OUTPUT_COLUMN, + CORRELATION_VALUE_1_COLUMN, + CORRELATION_VALUE_2_COLUMN, + compute, ) -from pyspark.pandas.base import IndexOpsMixin +from pyspark.pandas.datetimes import DatetimeMethods from pyspark.pandas.exceptions import SparkPandasIndexingError from pyspark.pandas.frame import DataFrame from pyspark.pandas.generic import Frame from pyspark.pandas.internal import ( - InternalField, - InternalFrame, DEFAULT_SERIES_NAME, NATURAL_ORDER_COLUMN_NAME, SPARK_DEFAULT_INDEX_NAME, SPARK_DEFAULT_SERIES_NAME, + InternalField, + InternalFrame, ) from pyspark.pandas.missing.series import MissingPandasLikeSeries from pyspark.pandas.plot import PandasOnSparkPlotAccessor +from pyspark.pandas.spark.accessors import SparkSeriesMethods +from pyspark.pandas.strings import StringMethods +from pyspark.pandas.typedef import ( + ScalarType, + SeriesType, + create_type_for_series_type, + infer_return_type, + spark_type_to_pandas_dtype, +) +from pyspark.pandas.typedef.typehints import as_spark_type from pyspark.pandas.utils import ( + SPARK_CONF_ARROW_ENABLED, ansi_mode_context, combine_frames, is_ansi_mode_enabled, is_name_like_tuple, is_name_like_value, + log_advice, name_like_string, same_anchor, scol_for, @@ -123,28 +114,42 @@ validate_axis, validate_bool_kwarg, verify_temp_column_name, - SPARK_CONF_ARROW_ENABLED, - log_advice, ) -from pyspark.pandas.datetimes import DatetimeMethods -from pyspark.pandas.spark.accessors import SparkSeriesMethods -from pyspark.pandas.strings import StringMethods -from pyspark.pandas.typedef import ( - infer_return_type, - spark_type_to_pandas_dtype, - ScalarType, - SeriesType, - create_type_for_series_type, +from pyspark.sql import ( + Column as PySparkColumn, ) -from pyspark.pandas.typedef.typehints import as_spark_type +from pyspark.sql import ( + DataFrame as SparkDataFrame, +) +from pyspark.sql import ( + Window as PySparkWindow, +) +from pyspark.sql import ( + functions as F, +) +from pyspark.sql.internal import InternalFunction as SF +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + NullType, + NumericType, + Row, + StructType, + TimestampType, +) +from pyspark.sql.window import Window if TYPE_CHECKING: - from pyspark.sql._typing import ColumnOrName - from pyspark.pandas.groupby import SeriesGroupBy - from pyspark.pandas.resample import SeriesResampler from pyspark.pandas.indexes import Index + from pyspark.pandas.resample import SeriesResampler from pyspark.pandas.spark.accessors import SparkIndexOpsMethods + from pyspark.sql._typing import ColumnOrName # This regular expression pattern is compiled and defined here to avoid to compile the same # pattern every time it is used in _repr_ in Series. @@ -7549,11 +7554,12 @@ def first_series(df: Union[DataFrame, pd.DataFrame]) -> Union[Series, pd.Series] def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.series + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/spark/accessors.py b/python/pyspark/pandas/spark/accessors.py index efebdee001663..e9cb03f9d2ec5 100644 --- a/python/pyspark/pandas/spark/accessors.py +++ b/python/pyspark/pandas/spark/accessors.py @@ -24,17 +24,17 @@ from typing import TYPE_CHECKING, Callable, Generic, List, Optional, Union from pyspark import StorageLevel -from pyspark.sql import Column as PySparkColumn, DataFrame as PySparkDataFrame -from pyspark.sql.types import DataType, StructType from pyspark.pandas._typing import IndexOpsLike from pyspark.pandas.internal import InternalField +from pyspark.sql import Column as PySparkColumn +from pyspark.sql import DataFrame as PySparkDataFrame +from pyspark.sql.types import DataType, StructType if TYPE_CHECKING: - from pyspark.sql._typing import OptionalPrimitiveType - from pyspark._typing import PrimitiveType - import pyspark.pandas as ps + from pyspark._typing import PrimitiveType from pyspark.pandas.frame import CachedDataFrame + from pyspark.sql._typing import OptionalPrimitiveType class SparkIndexOpsMethods(Generic[IndexOpsLike], metaclass=ABCMeta): @@ -187,8 +187,8 @@ def apply(self, func: Callable[[PySparkColumn], PySparkColumn]) -> "ps.Series": Name: a, dtype: int64 """ from pyspark.pandas.frame import DataFrame - from pyspark.pandas.series import Series, first_series from pyspark.pandas.internal import HIDDEN_COLUMNS + from pyspark.pandas.series import Series, first_series output = func(self._data.spark.column) if not isinstance(output, PySparkColumn): @@ -1228,16 +1228,18 @@ def unpersist(self) -> None: def _test() -> None: - import os import doctest + import os import shutil import sys import tempfile import uuid + import numpy import pandas - from pyspark.sql import SparkSession + import pyspark.pandas.spark.accessors + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/spark/utils.py b/python/pyspark/pandas/spark/utils.py index d8c5aa41debcf..0424b7860a6c9 100644 --- a/python/pyspark/pandas/spark/utils.py +++ b/python/pyspark/pandas/spark/utils.py @@ -20,7 +20,7 @@ from typing import overload -from pyspark.sql.types import DecimalType, StructType, MapType, ArrayType, StructField, DataType +from pyspark.sql.types import ArrayType, DataType, DecimalType, MapType, StructField, StructType @overload @@ -179,6 +179,7 @@ def force_decimal_precision_scale( def _test() -> None: import doctest import sys + import pyspark.pandas.spark.utils globs = pyspark.pandas.spark.utils.__dict__.copy() diff --git a/python/pyspark/pandas/sql_formatter.py b/python/pyspark/pandas/sql_formatter.py index b6e383fc4cea8..b948090c65e4d 100644 --- a/python/pyspark/pandas/sql_formatter.py +++ b/python/pyspark/pandas/sql_formatter.py @@ -17,21 +17,20 @@ import os import string -from typing import Any, Dict, Optional, Union, List, Sequence, Mapping, Tuple import uuid import warnings +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union import pandas as pd -from pyspark.pandas.internal import InternalFrame -from pyspark.pandas.namespace import _get_index_map from pyspark import pandas as ps -from pyspark.sql import SparkSession -from pyspark.sql.utils import get_lit_sql_str -from pyspark.pandas.utils import default_session from pyspark.pandas.frame import DataFrame +from pyspark.pandas.internal import InternalFrame +from pyspark.pandas.namespace import _get_index_map from pyspark.pandas.series import Series -from pyspark.sql.utils import is_remote +from pyspark.pandas.utils import default_session +from pyspark.sql import SparkSession +from pyspark.sql.utils import get_lit_sql_str, is_remote __all__ = ["sql"] @@ -303,11 +302,12 @@ def clear(self) -> None: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.sql_formatter + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/sql_processor.py b/python/pyspark/pandas/sql_processor.py index afac0ca9cb338..3e109d816c3aa 100644 --- a/python/pyspark/pandas/sql_processor.py +++ b/python/pyspark/pandas/sql_processor.py @@ -16,18 +16,19 @@ # import _string # type: ignore[import-not-found] -from typing import Any, Dict, Optional, Union, List import inspect +from typing import Any, Dict, List, Optional, Union import pandas as pd -from pyspark.sql import SparkSession, DataFrame as SDataFrame from pyspark import pandas as ps # For running doctests and reference resolution in PyCharm. -from pyspark.pandas.utils import default_session from pyspark.pandas.frame import DataFrame -from pyspark.pandas.series import Series from pyspark.pandas.internal import InternalFrame from pyspark.pandas.namespace import _get_index_map +from pyspark.pandas.series import Series +from pyspark.pandas.utils import default_session +from pyspark.sql import DataFrame as SDataFrame +from pyspark.sql import SparkSession __all__ = ["sql"] @@ -359,11 +360,12 @@ def _convert_var(self, var: Any) -> Any: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.sql_processor + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/strings.py b/python/pyspark/pandas/strings.py index b29482d32d784..c1b1b6f68a8aa 100644 --- a/python/pyspark/pandas/strings.py +++ b/python/pyspark/pandas/strings.py @@ -36,14 +36,14 @@ import pandas as pd from pandas.api.extensions import no_default +import pyspark.pandas as ps from pyspark._globals import _NoValue, _NoValueType from pyspark.loose_version import LooseVersion +from pyspark.pandas.typedef.typehints import SeriesType, is_str_dtype from pyspark.pandas.utils import ansi_mode_context, is_ansi_mode_enabled -from pyspark.pandas.typedef.typehints import is_str_dtype, SeriesType -from pyspark.sql.types import StringType, BinaryType, ArrayType, LongType, MapType from pyspark.sql import functions as F from pyspark.sql.functions import pandas_udf -import pyspark.pandas as ps +from pyspark.sql.types import ArrayType, BinaryType, LongType, MapType, StringType FuncT = TypeVar("FuncT", bound=Callable[..., Any]) @@ -2408,11 +2408,12 @@ def get_dummies(self, sep: str = "|") -> "ps.DataFrame": def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.strings + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/supported_api_gen.py b/python/pyspark/pandas/supported_api_gen.py index 3a9c4f1d24687..81450f17818b9 100644 --- a/python/pyspark/pandas/supported_api_gen.py +++ b/python/pyspark/pandas/supported_api_gen.py @@ -22,16 +22,16 @@ import warnings from enum import Enum, unique from inspect import getmembers, isclass, isfunction, signature -from typing import Any, Dict, List, NamedTuple, Set, TextIO, Tuple from types import FunctionType +from typing import Any, Dict, List, NamedTuple, Set, TextIO, Tuple -import pyspark.pandas as ps -import pyspark.pandas.groupby as psg -import pyspark.pandas.window as psw import pandas as pd import pandas.core.groupby as pdg import pandas.core.window as pdw +import pyspark.pandas as ps +import pyspark.pandas.groupby as psg +import pyspark.pandas.window as psw from pyspark.loose_version import LooseVersion from pyspark.pandas.exceptions import PandasNotImplementedError diff --git a/python/pyspark/pandas/testing.py b/python/pyspark/pandas/testing.py index c2c10c09d7242..cf8fa426739ea 100644 --- a/python/pyspark/pandas/testing.py +++ b/python/pyspark/pandas/testing.py @@ -20,6 +20,7 @@ """ from typing import Literal, Union + import pyspark.pandas as ps try: diff --git a/python/pyspark/pandas/tests/computation/test_compute.py b/python/pyspark/pandas/tests/computation/test_compute.py index 4538509cd4d36..394d461731420 100644 --- a/python/pyspark/pandas/tests/computation/test_compute.py +++ b/python/pyspark/pandas/tests/computation/test_compute.py @@ -18,8 +18,8 @@ import numpy as np import pandas as pd -from pyspark.sql import functions as sf from pyspark import pandas as ps +from pyspark.sql import functions as sf from pyspark.testing.pandasutils import PandasOnSparkTestCase @@ -198,9 +198,14 @@ def test_diff(self): msg = "should be an int" with self.assertRaisesRegex(TypeError, msg): psdf.diff(1.5) - msg = 'axis should be either 0 or "index" currently.' - with self.assertRaisesRegex(NotImplementedError, msg): - psdf.diff(axis=1) + + # axis=1: difference across columns + self.assert_eq(pdf.diff(axis=1), psdf.diff(axis=1)) + self.assert_eq(pdf.diff(periods=2, axis=1), psdf.diff(periods=2, axis=1)) + self.assert_eq(pdf.diff(periods=-1, axis=1), psdf.diff(periods=-1, axis=1)) + + with self.assertRaisesRegex(TypeError, msg): + psdf.diff(1.5, axis=1) # multi-index columns columns = pd.MultiIndex.from_tuples([("x", "Col1"), ("x", "Col2"), ("y", "Col3")]) diff --git a/python/pyspark/pandas/tests/computation/test_corr.py b/python/pyspark/pandas/tests/computation/test_corr.py index 6b3ab85f3fd1c..1654c62cf7780 100644 --- a/python/pyspark/pandas/tests/computation/test_corr.py +++ b/python/pyspark/pandas/tests/computation/test_corr.py @@ -19,7 +19,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase, SPARK_CONF_ARROW_ENABLED +from pyspark.testing.pandasutils import SPARK_CONF_ARROW_ENABLED, PandasOnSparkTestCase class FrameCorrMixin: diff --git a/python/pyspark/pandas/tests/computation/test_describe.py b/python/pyspark/pandas/tests/computation/test_describe.py index ad3377b9df605..0c1edd87d30dc 100644 --- a/python/pyspark/pandas/tests/computation/test_describe.py +++ b/python/pyspark/pandas/tests/computation/test_describe.py @@ -19,8 +19,8 @@ import numpy as np import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/computation/test_idxmax_idxmin.py b/python/pyspark/pandas/tests/computation/test_idxmax_idxmin.py index 9a3ca66f4ed68..87f3f62947832 100644 --- a/python/pyspark/pandas/tests/computation/test_idxmax_idxmin.py +++ b/python/pyspark/pandas/tests/computation/test_idxmax_idxmin.py @@ -17,8 +17,8 @@ import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/computation/test_melt.py b/python/pyspark/pandas/tests/computation/test_melt.py index 73b0f12836f17..4cb517a71a017 100644 --- a/python/pyspark/pandas/tests/computation/test_melt.py +++ b/python/pyspark/pandas/tests/computation/test_melt.py @@ -19,8 +19,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.utils import name_like_string +from pyspark.testing.pandasutils import PandasOnSparkTestCase class FrameMeltMixin: diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_as_type.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_as_type.py index 4c3682b1e6bed..98143eeaa9ff4 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_as_type.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_as_type.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_as_type import AsTypeTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class AsTypeParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py index b51a0585efbd1..6005ab4b3dd40 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_binary_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_binary_ops import BinaryOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class BinaryOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py index 4cb7133346462..20c07d4a06949 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_boolean_ops.py @@ -16,12 +16,12 @@ # from pyspark.pandas.tests.data_type_ops.test_boolean_ops import ( - BooleanOpsTestsMixin, BooleanExtensionOpsTestsMixin, + BooleanOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class BooleanOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py index 12a8fe81fd622..03bf61eb6b151 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_categorical_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_categorical_ops import CategoricalOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class CategoricalOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py index 7eca857bc860f..a3c0ca3a418d4 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_complex_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_complex_ops import ComplexOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class ComplexOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py index 6fc3d66e9effe..a15fb697f21b9 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_date_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_date_ops import DateOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class DateOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py index 807904058ddfb..95b9de5a55140 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_datetime_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_datetime_ops import DatetimeOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class DatetimeOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py index ca323107c0d9f..9d55e27709d00 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_null_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_null_ops import NullOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class NullOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py index f5b9a56a3a532..69b286ee9de91 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_arithmetic.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_arithmetic import ArithmeticTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class ArithmeticParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mod.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mod.py index b666728e45918..638d5874a5e97 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mod.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mod.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_mod import NumModTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class NumModParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mul_div.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mul_div.py index cce619df93e13..45c7851e7d8aa 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mul_div.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_mul_div.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_mul_div import NumMulDivTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class NumMulDivParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py index faeef9b0dd4a1..5401547395407 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_ops import NumOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class NumOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py index da14d21cc3eda..a7d91cea9c592 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext.py @@ -19,8 +19,8 @@ FractionalExtensionOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class FractionalExtensionOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py index f140db0a13f2f..1ebab8894c38d 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_fractional_ext_astype_cmp.py @@ -19,8 +19,8 @@ FractionalExtensionAstypeCmpOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class FractionalExtensionAstypeCmpOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py index 863b09d1a737f..83b78faccf49e 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext.py @@ -19,8 +19,8 @@ IntegralExtensionOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class IntegralExtensionOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py index af85dc5c8fca0..d2d65d335f77f 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_ops_integral_ext_astype_cmp.py @@ -19,8 +19,8 @@ IntegralExtensionAstypeCmpOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class IntegralExtensionAstypeCmpOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_pow.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_pow.py index e70b4c57a5b56..5e7203e77589b 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_pow.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_pow.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_pow import NumPowTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class NumPowParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py index 127a87bf4a1c4..762f5f2965cc7 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_num_reverse.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_num_reverse import ReverseTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class ReverseParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py index 8a324f9f91f31..a41b8c23b5d3e 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_string_ops.py @@ -16,12 +16,12 @@ # from pyspark.pandas.tests.data_type_ops.test_string_ops import ( - StringOpsTestsMixin, StringExtensionOpsTestsMixin, + StringOpsTestsMixin, ) from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class StringOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py index 7e5e9ae28db7a..dc5205d79b01e 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_timedelta_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_timedelta_ops import TimedeltaOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class TimedeltaOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py index 5c0dda87618ed..645534112b17b 100644 --- a/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py +++ b/python/pyspark/pandas/tests/connect/data_type_ops/test_parity_udt_ops.py @@ -17,8 +17,8 @@ from pyspark.pandas.tests.data_type_ops.test_udt_ops import UDTOpsTestsMixin from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.pandasutils import PandasOnSparkTestUtils class UDTOpsParityTests( diff --git a/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic.py b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic.py index 870308686ed73..48e99fa32a0d5 100644 --- a/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic.py +++ b/python/pyspark/pandas/tests/connect/diff_frames_ops/test_parity_basic.py @@ -15,9 +15,9 @@ # limitations under the License. # +from pyspark.pandas.tests.diff_frames_ops.test_basic import BasicMixin from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.testing.pandasutils import PandasOnSparkTestUtils -from pyspark.pandas.tests.diff_frames_ops.test_basic import BasicMixin class BasicParityTests( diff --git a/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing_adv.py b/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing_adv.py index 0af2d091d87ed..4f8d78a0fc28e 100644 --- a/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing_adv.py +++ b/python/pyspark/pandas/tests/connect/indexes/test_parity_indexing_adv.py @@ -30,7 +30,6 @@ class IndexingAdvParityTests( if __name__ == "__main__": from pyspark.pandas.tests.connect.indexes.test_parity_indexing import * # noqa: F403 - from pyspark.testing import main main() diff --git a/python/pyspark/pandas/tests/data_type_ops/test_as_type.py b/python/pyspark/pandas/tests/data_type_ops/test_as_type.py index 3808da3a18f4e..7e63e529599f6 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_as_type.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_as_type.py @@ -16,17 +16,17 @@ # -import pandas as pd import numpy as np +import pandas as pd from pandas.api.types import CategoricalDtype from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import ( extension_float_dtypes_available, extension_object_dtypes_available, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase class AsTypeTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_base.py b/python/pyspark/pandas/tests/data_type_ops/test_base.py index 1cfce32405837..a594ee1f5d380 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_base.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_base.py @@ -17,18 +17,18 @@ import unittest -from pandas.api.types import CategoricalDtype from pandas.api.extensions import ExtensionDtype +from pandas.api.types import CategoricalDtype from pyspark.pandas.data_type_ops.base import DataTypeOps from pyspark.pandas.data_type_ops.binary_ops import BinaryOps -from pyspark.pandas.data_type_ops.boolean_ops import BooleanOps, BooleanExtensionOps +from pyspark.pandas.data_type_ops.boolean_ops import BooleanExtensionOps, BooleanOps from pyspark.pandas.data_type_ops.categorical_ops import CategoricalOps from pyspark.pandas.data_type_ops.complex_ops import ArrayOps, MapOps, StructOps from pyspark.pandas.data_type_ops.date_ops import DateOps -from pyspark.pandas.data_type_ops.datetime_ops import DatetimeOps, DatetimeNTZOps +from pyspark.pandas.data_type_ops.datetime_ops import DatetimeNTZOps, DatetimeOps from pyspark.pandas.data_type_ops.null_ops import NullOps -from pyspark.pandas.data_type_ops.num_ops import IntegralOps, FractionalOps, DecimalOps +from pyspark.pandas.data_type_ops.num_ops import DecimalOps, FractionalOps, IntegralOps from pyspark.pandas.data_type_ops.string_ops import StringOps from pyspark.pandas.data_type_ops.timedelta_ops import TimedeltaOps from pyspark.pandas.data_type_ops.udt_ops import UDTOps @@ -46,8 +46,8 @@ NullType, StringType, StructType, - TimestampType, TimestampNTZType, + TimestampType, UserDefinedType, ) diff --git a/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py index c3ea2c71a5401..c96c486e1de14 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_binary_ops.py @@ -19,8 +19,8 @@ from pandas.api.types import CategoricalDtype from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class BinaryOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py index 9911b5dc49767..ad0cf550ad567 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_boolean_ops.py @@ -18,19 +18,19 @@ import datetime import unittest -import pandas as pd import numpy as np +import pandas as pd from pandas.api.types import CategoricalDtype from pyspark import pandas as ps from pyspark.pandas import option_context -from pyspark.testing.utils import is_ansi_mode_test -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import ( extension_float_dtypes_available, extension_object_dtypes_available, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.utils import is_ansi_mode_test class BooleanOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py index f441ea2af9195..068ff74058627 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_categorical_ops.py @@ -15,14 +15,14 @@ # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pandas.api.types import CategoricalDtype from pyspark import pandas as ps from pyspark.pandas.config import option_context -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class CategoricalOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py index 1b5d18f0e8b89..84369afe6f17a 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_complex_ops.py @@ -15,14 +15,14 @@ # limitations under the License. # -import decimal import datetime +import decimal import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ComplexOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py index 9c47d931236e8..783cd0599f5db 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_date_ops.py @@ -21,8 +21,8 @@ from pandas.api.types import CategoricalDtype from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DateOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py index 8800387f657e2..8b15411bbdc52 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_datetime_ops.py @@ -21,8 +21,8 @@ from pandas.api.types import CategoricalDtype from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py index b54ae907dc957..fddfc8d3be0c6 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_null_ops.py @@ -19,8 +19,8 @@ from pandas.api.types import CategoricalDtype import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class NullOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py b/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py index 5d56b8275c4cf..6836f24fc56e0 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_arithmetic.py @@ -19,9 +19,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.utils import is_ansi_mode_test -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.utils import is_ansi_mode_test class ArithmeticTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_mod.py b/python/pyspark/pandas/tests/data_type_ops/test_num_mod.py index d40b83ffea8de..0f2d799bfb7c3 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_mod.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_mod.py @@ -16,12 +16,12 @@ # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class NumModTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py b/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py index 61721410bab4a..d272b22cc475c 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_mul_div.py @@ -16,13 +16,13 @@ # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.testing.utils import is_ansi_mode_test -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.utils import is_ansi_mode_test class NumMulDivTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py index 4e95ee1cb0aaa..b7b7df99633d6 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops.py @@ -15,15 +15,15 @@ # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps from pyspark.pandas.config import option_context -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.testing.utils import is_ansi_mode_test from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.sql.types import DecimalType, IntegralType +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.utils import is_ansi_mode_test class NumOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py index e507029ad8960..1c2b175f9106b 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext.py @@ -20,9 +20,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase @unittest.skipIf( diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py index aa429e6b4185e..f549a611a91c6 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_fractional_ext_astype_cmp.py @@ -17,14 +17,14 @@ import unittest -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps from pyspark.pandas.config import option_context -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase @unittest.skipIf( diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py index 6745b5a840f7f..c9b1e3de59f8b 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext.py @@ -20,9 +20,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py index 9f8224eb07985..9db13945a7c32 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_ops_integral_ext_astype_cmp.py @@ -17,14 +17,14 @@ import unittest -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps from pyspark.pandas.config import option_context -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase @unittest.skipIf(not extension_dtypes_available, "pandas extension dtypes are not available") diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_pow.py b/python/pyspark/pandas/tests/data_type_ops/test_num_pow.py index bc4db248c3628..cc6d30d3434fb 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_pow.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_pow.py @@ -19,8 +19,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class NumPowTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py b/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py index d3e1f74ff5a2c..f2b85085e7d1b 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_num_reverse.py @@ -20,8 +20,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ReverseTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py index f173d8ca95ad5..a2d8589e4bfdd 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_string_ops.py @@ -23,9 +23,9 @@ from pyspark import pandas as ps from pyspark.pandas.config import option_context -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase from pyspark.pandas.typedef.typehints import extension_object_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase if extension_object_dtypes_available: from pandas import StringDtype diff --git a/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py index 34190920017a2..051a5a4be40e5 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_timedelta_ops.py @@ -23,8 +23,8 @@ import pyspark.pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class TimedeltaOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py b/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py index 42b2f6544ca27..fd7088a744b71 100644 --- a/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py +++ b/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py @@ -19,8 +19,8 @@ import pyspark.pandas as ps from pyspark.ml.linalg import SparseVector -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase +from pyspark.testing.pandasutils import PandasOnSparkTestCase class UDTOpsTestsMixin: diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_align.py b/python/pyspark/pandas/tests/diff_frames_ops/test_align.py index f09f54d8579b4..6eae20e93134f 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_align.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_align.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic.py index 2b483c35ff8bc..abc0298b58c7c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic.py @@ -18,9 +18,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticTestingFuncMixin: diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain.py index 338bbe2dac823..c18ebe55d6aa4 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain.py @@ -18,9 +18,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticChainTestingFuncMixin: diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext.py index a64458f3e6622..3f739a96742ce 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext.py @@ -18,12 +18,12 @@ import pandas as pd -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_arithmetic_chain import ( ArithmeticChainTestingFuncMixin, ) +from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticChainExtMixin(ArithmeticChainTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext_float.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext_float.py index e3feb5d81ebd5..7e6525c4a3e23 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext_float.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_chain_ext_float.py @@ -18,12 +18,12 @@ import pandas as pd -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_arithmetic_chain import ( ArithmeticChainTestingFuncMixin, ) +from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticChainExtFloatMixin(ArithmeticChainTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext.py index 091e0280976ab..a71680002ad14 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext.py @@ -18,10 +18,10 @@ import pandas as pd -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_arithmetic import ArithmeticMixin +from pyspark.pandas.typedef.typehints import extension_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticExtMixin(ArithmeticMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext_float.py b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext_float.py index a761b47bb01fe..74af741878f80 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext_float.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_arithmetic_ext_float.py @@ -18,10 +18,10 @@ import pandas as pd -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_arithmetic import ArithmeticMixin +from pyspark.pandas.typedef.typehints import extension_float_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ArithmeticExtFloatMixin(ArithmeticMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_assign_frame.py b/python/pyspark/pandas/tests/diff_frames_ops/test_assign_frame.py index 794069980024d..3243b0ab4c90c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_assign_frame.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_assign_frame.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_assign_series.py b/python/pyspark/pandas/tests/diff_frames_ops/test_assign_series.py index 2eb3f7274d629..5118a4de72cbc 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_assign_series.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_assign_series.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_basic.py b/python/pyspark/pandas/tests/diff_frames_ops/test_basic.py index 79d471620f2cc..8399ff9e0791f 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_basic.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_basic.py @@ -19,7 +19,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py b/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py index da3e8754e1d84..5da76ca0a7485 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_basic_slow.py @@ -15,11 +15,11 @@ # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_bitwise.py b/python/pyspark/pandas/tests/diff_frames_ops/test_bitwise.py index 564bf9e195449..063dd42390249 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_bitwise.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_bitwise.py @@ -15,13 +15,14 @@ # limitations under the License. # import unittest + import numpy as np import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.typedef.typehints import extension_object_dtypes_available +from pyspark.testing.pandasutils import PandasOnSparkTestCase class BitwiseMixin: diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_combine_first.py b/python/pyspark/pandas/tests/diff_frames_ops/test_combine_first.py index da16d8fa6b067..19447b494cd17 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_combine_first.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_combine_first.py @@ -17,7 +17,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_compare_series.py b/python/pyspark/pandas/tests/diff_frames_ops/test_compare_series.py index dc8cb8398c481..e677fe8915d91 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_compare_series.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_compare_series.py @@ -19,7 +19,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_concat_inner.py b/python/pyspark/pandas/tests/diff_frames_ops/test_concat_inner.py index 1fe2f37b544eb..34e319650db47 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_concat_inner.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_concat_inner.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_concat_outer.py b/python/pyspark/pandas/tests/diff_frames_ops/test_concat_outer.py index 42e0ded722c82..57961f3eca75a 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_concat_outer.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_concat_outer.py @@ -18,9 +18,9 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_concat_inner import ConcatTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ConcatOuterMixin(ConcatTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_corrwith.py b/python/pyspark/pandas/tests/diff_frames_ops/test_corrwith.py index 5e0f0f8be4a91..c5f6f706c6d69 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_corrwith.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_corrwith.py @@ -15,11 +15,11 @@ # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_cov.py b/python/pyspark/pandas/tests/diff_frames_ops/test_cov.py index 0e22ed07c41e1..977f78d0e7907 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_cov.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_cov.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py b/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py index a3fea0ae8e664..239ca72467b3c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_dot_frame.py @@ -17,7 +17,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py b/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py index 80504da2eccae..a7e4a6e9c8236 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_dot_series.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_error.py b/python/pyspark/pandas/tests/diff_frames_ops/test_error.py index 1476ac48538d9..86b70bace8b9e 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_error.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_error.py @@ -16,11 +16,11 @@ # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py index 1d89adf4667c5..63c50aa67e81f 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_aggregate.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_aggregate.py index 3b60a4aa16955..795e3d19435d7 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_aggregate.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_aggregate.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_apply.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_apply.py index 179748c4872ef..7719b26d10ea3 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_apply.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_apply.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_cumulative.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_cumulative.py index bedc4594a4d55..49c9d218eb40c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_cumulative.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_cumulative.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff.py index 48423f81d2b2a..88a481302b473 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff_len.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff_len.py index 74efe92ceef99..d9f400d89d6d0 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff_len.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_diff_len.py @@ -19,7 +19,7 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding.py index 76c95f3f1e480..aa8a9f5bb2e4a 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_adv.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_adv.py index 1bd73738bd38c..3c9e00ec8b83b 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_adv.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_adv.py @@ -15,11 +15,11 @@ # limitations under the License. # -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_groupby_expanding import ( GroupByExpandingTestingFuncMixin, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupByExpandingAdvMixin(GroupByExpandingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_count.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_count.py index 393acd6993dda..edb066399a511 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_count.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_expanding_count.py @@ -15,11 +15,11 @@ # limitations under the License. # -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_groupby_expanding import ( GroupByExpandingTestingFuncMixin, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupByExpandingCountMixin(GroupByExpandingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_fillna.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_fillna.py index c61d213e6f84f..045bd85090661 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_fillna.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_fillna.py @@ -19,7 +19,7 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_filter.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_filter.py index 149d6170eb759..92f56e50da318 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_filter.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_filter.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling.py index 4e57a53be00d1..c69839be0ed00 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_adv.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_adv.py index e7d6a4988a1f3..b7b9835fcd49c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_adv.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_adv.py @@ -15,9 +15,9 @@ # limitations under the License. # -from pyspark.pandas.config import set_option, reset_option -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.tests.diff_frames_ops.test_groupby_rolling import GroupByRollingTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupByRollingAdvMixin(GroupByRollingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_count.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_count.py index 8cb04ff8b944e..c0f9838c74e54 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_count.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_rolling_count.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py index 45006258ce819..3190fb47f4080 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_shift.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py index 5f36bd1ac88b4..2c03c86fc933f 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_split_apply_combine.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py index 04bf25338f8d0..91d420d8a5f55 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_groupby_transform.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_index.py b/python/pyspark/pandas/tests/diff_frames_ops/test_index.py index 8056bc39504c3..225c2ac6ed19f 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_index.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_index.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_series.py b/python/pyspark/pandas/tests/diff_frames_ops/test_series.py index 57a9f967c2029..70f5b565ee56a 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_series.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_series.py @@ -14,11 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py b/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py index eeb8000d00e53..a9f62399e025c 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_frame.py @@ -18,8 +18,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option from pyspark.loose_version import LooseVersion +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py b/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py index ef8facd04634e..41d822c0309ef 100644 --- a/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py +++ b/python/pyspark/pandas/tests/diff_frames_ops/test_setitem_series.py @@ -18,7 +18,7 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/frame/test_constructor.py b/python/pyspark/pandas/tests/frame/test_constructor.py index 0c80d2bbb7955..96a1837d08b47 100644 --- a/python/pyspark/pandas/tests/frame/test_constructor.py +++ b/python/pyspark/pandas/tests/frame/test_constructor.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from datetime import datetime, timedelta import unittest +from datetime import datetime, timedelta import numpy as np import pandas as pd @@ -27,7 +27,6 @@ extension_object_dtypes_available, ) from pyspark.pandas.utils import is_testing - from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/frame/test_spark.py b/python/pyspark/pandas/tests/frame/test_spark.py index 58a1c2d478af3..019c8d129f60f 100644 --- a/python/pyspark/pandas/tests/frame/test_spark.py +++ b/python/pyspark/pandas/tests/frame/test_spark.py @@ -22,13 +22,13 @@ import pandas as pd from pyspark import StorageLevel -from pyspark.ml.linalg import SparseVector -from pyspark.sql.types import StructType from pyspark import pandas as ps -from pyspark.pandas.frame import CachedDataFrame +from pyspark.ml.linalg import SparseVector from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.frame import CachedDataFrame from pyspark.pandas.missing.frame import MissingPandasLikeDataFrame -from pyspark.testing.pandasutils import PandasOnSparkTestCase, SPARK_CONF_ARROW_ENABLED +from pyspark.sql.types import StructType +from pyspark.testing.pandasutils import SPARK_CONF_ARROW_ENABLED, PandasOnSparkTestCase # This file contains test cases for 'Spark-related' diff --git a/python/pyspark/pandas/tests/groupby/test_describe.py b/python/pyspark/pandas/tests/groupby/test_describe.py index e255c62389c85..ad152516fa5a3 100644 --- a/python/pyspark/pandas/tests/groupby/test_describe.py +++ b/python/pyspark/pandas/tests/groupby/test_describe.py @@ -16,6 +16,7 @@ # from itertools import product + import pandas as pd from pyspark import pandas as ps diff --git a/python/pyspark/pandas/tests/groupby/test_grouping.py b/python/pyspark/pandas/tests/groupby/test_grouping.py index a5b2d81d6bb80..21d48448a695d 100644 --- a/python/pyspark/pandas/tests/groupby/test_grouping.py +++ b/python/pyspark/pandas/tests/groupby/test_grouping.py @@ -16,11 +16,12 @@ # -import pandas as pd import numpy as np +import pandas as pd + import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.groupby import SeriesGroupBy +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupingTestsMixin: diff --git a/python/pyspark/pandas/tests/groupby/test_index.py b/python/pyspark/pandas/tests/groupby/test_index.py index e4f9902952a4a..11b852aed6a45 100644 --- a/python/pyspark/pandas/tests/groupby/test_index.py +++ b/python/pyspark/pandas/tests/groupby/test_index.py @@ -17,8 +17,8 @@ import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/groupby/test_missing.py b/python/pyspark/pandas/tests/groupby/test_missing.py index 55302f19e6c60..1f5a16d66e61e 100644 --- a/python/pyspark/pandas/tests/groupby/test_missing.py +++ b/python/pyspark/pandas/tests/groupby/test_missing.py @@ -19,11 +19,11 @@ import pyspark.pandas as ps from pyspark.pandas.exceptions import PandasNotImplementedError -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.missing.groupby import ( MissingPandasLikeDataFrameGroupBy, MissingPandasLikeSeriesGroupBy, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase class MissingTestsMixin: diff --git a/python/pyspark/pandas/tests/groupby/test_nlargest_nsmallest.py b/python/pyspark/pandas/tests/groupby/test_nlargest_nsmallest.py index 18b305d6bc46a..d2bcbba711c58 100644 --- a/python/pyspark/pandas/tests/groupby/test_nlargest_nsmallest.py +++ b/python/pyspark/pandas/tests/groupby/test_nlargest_nsmallest.py @@ -16,8 +16,9 @@ # -import pandas as pd import numpy as np +import pandas as pd + import pyspark.pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/groupby/test_rank.py b/python/pyspark/pandas/tests/groupby/test_rank.py index 8b548dac7edb6..34a78f1d004bb 100644 --- a/python/pyspark/pandas/tests/groupby/test_rank.py +++ b/python/pyspark/pandas/tests/groupby/test_rank.py @@ -16,8 +16,9 @@ # -import pandas as pd import numpy as np +import pandas as pd + import pyspark.pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/groupby/test_size.py b/python/pyspark/pandas/tests/groupby/test_size.py index 47f596a658698..f09af209095cf 100644 --- a/python/pyspark/pandas/tests/groupby/test_size.py +++ b/python/pyspark/pandas/tests/groupby/test_size.py @@ -17,6 +17,7 @@ import pandas as pd + import pyspark.pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_count.py b/python/pyspark/pandas/tests/groupby/test_split_apply_count.py index a8e026450c461..2d4ae6534f3fa 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_count.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_count.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyCountMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_first.py b/python/pyspark/pandas/tests/groupby/test_split_apply_first.py index 8e8637b01deb7..28b2d02ef06c3 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_first.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_first.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyFirstMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_last.py b/python/pyspark/pandas/tests/groupby/test_split_apply_last.py index 4657b332e1db3..2345677132309 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_last.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_last.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyLastMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py b/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py index dd260c5dd1871..cc43e0c4a262a 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_mean.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyMeanMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_min_max.py b/python/pyspark/pandas/tests/groupby/test_split_apply_min_max.py index 513993e355ab7..6b1b39f42faa5 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_min_max.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_min_max.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyMMMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_skew.py b/python/pyspark/pandas/tests/groupby/test_split_apply_skew.py index 4bc839c669544..10e00a6d1f352 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_skew.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_skew.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplySkewMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_std.py b/python/pyspark/pandas/tests/groupby/test_split_apply_std.py index 1d51c00e2ff98..4706e555026ba 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_std.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_std.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyStdMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_split_apply_var.py b/python/pyspark/pandas/tests/groupby/test_split_apply_var.py index f6c5c0b9cc704..463b3af934e66 100644 --- a/python/pyspark/pandas/tests/groupby/test_split_apply_var.py +++ b/python/pyspark/pandas/tests/groupby/test_split_apply_var.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_split_apply import GroupbySplitApplyTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbySplitApplyVarMixin(GroupbySplitApplyTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_stat_adv.py b/python/pyspark/pandas/tests/groupby/test_stat_adv.py index f06645c2a8652..b0c9ebc6eeb4c 100644 --- a/python/pyspark/pandas/tests/groupby/test_stat_adv.py +++ b/python/pyspark/pandas/tests/groupby/test_stat_adv.py @@ -18,10 +18,10 @@ import numpy as np import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.loose_version import LooseVersion from pyspark.pandas.tests.groupby.test_stat import GroupbyStatTestingFuncMixin, using_pandas3 +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbyStatAdvMixin(GroupbyStatTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_stat_func.py b/python/pyspark/pandas/tests/groupby/test_stat_func.py index 7eca1e53918ac..56ff3aa75b5a4 100644 --- a/python/pyspark/pandas/tests/groupby/test_stat_func.py +++ b/python/pyspark/pandas/tests/groupby/test_stat_func.py @@ -20,8 +20,8 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_stat import GroupbyStatTestingFuncMixin, using_pandas3 +from pyspark.testing.pandasutils import PandasOnSparkTestCase class FuncTestsMixin(GroupbyStatTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_stat_median.py b/python/pyspark/pandas/tests/groupby/test_stat_median.py index a555e4b06e44b..40de71669a2cf 100644 --- a/python/pyspark/pandas/tests/groupby/test_stat_median.py +++ b/python/pyspark/pandas/tests/groupby/test_stat_median.py @@ -19,8 +19,8 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_stat import GroupbyStatTestingFuncMixin, using_pandas3 +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupbyStatMedianMixin(GroupbyStatTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_stat_prod.py b/python/pyspark/pandas/tests/groupby/test_stat_prod.py index c4564cc0949a1..51e5354921f12 100644 --- a/python/pyspark/pandas/tests/groupby/test_stat_prod.py +++ b/python/pyspark/pandas/tests/groupby/test_stat_prod.py @@ -21,8 +21,8 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.groupby.test_stat import GroupbyStatTestingFuncMixin, using_pandas3 +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ProdTestsMixin(GroupbyStatTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/groupby/test_value_counts.py b/python/pyspark/pandas/tests/groupby/test_value_counts.py index 28d6004ffaec1..254deeceb5872 100644 --- a/python/pyspark/pandas/tests/groupby/test_value_counts.py +++ b/python/pyspark/pandas/tests/groupby/test_value_counts.py @@ -16,8 +16,9 @@ # -import pandas as pd import numpy as np +import pandas as pd + import pyspark.pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/indexes/test_conversion.py b/python/pyspark/pandas/tests/indexes/test_conversion.py index 8146915ae359a..88ca75ac995b7 100644 --- a/python/pyspark/pandas/tests/indexes/test_conversion.py +++ b/python/pyspark/pandas/tests/indexes/test_conversion.py @@ -19,11 +19,11 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.utils import ( SPARK_CONF_ARROW_ENABLED, SPARK_CONF_PANDAS_STRUCT_MODE, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.testing.utils import is_ansi_mode_test diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_at.py b/python/pyspark/pandas/tests/indexes/test_datetime_at.py index a71a20fe85f90..6858ffb39a013 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_at.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_at.py @@ -20,8 +20,8 @@ import pandas as pd import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexAtMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_between.py b/python/pyspark/pandas/tests/indexes/test_datetime_between.py index 2268f7245184a..6c725a8a2ebfa 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_between.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_between.py @@ -19,8 +19,8 @@ import pandas as pd -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexBetweenMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_ceil.py b/python/pyspark/pandas/tests/indexes/test_datetime_ceil.py index d2d5beee8d047..0c5ee8b8140fa 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_ceil.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_ceil.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexCeilMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_floor.py b/python/pyspark/pandas/tests/indexes/test_datetime_floor.py index 1c41f7b32103b..477ba9a6d5e28 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_floor.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_floor.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexFloorMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_iso.py b/python/pyspark/pandas/tests/indexes/test_datetime_iso.py index f3073ee0b6d33..63141b7439d24 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_iso.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_iso.py @@ -16,8 +16,8 @@ # import numpy as np -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexISOMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_map.py b/python/pyspark/pandas/tests/indexes/test_datetime_map.py index af06f50a5a494..1200b1b67d23f 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_map.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_map.py @@ -20,8 +20,8 @@ import pandas as pd import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexMapMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_property.py b/python/pyspark/pandas/tests/indexes/test_datetime_property.py index df8df3b04b887..1edca646215d9 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_property.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_property.py @@ -17,8 +17,8 @@ import numpy as np import pandas as pd -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexPropertyTestsMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_datetime_round.py b/python/pyspark/pandas/tests/indexes/test_datetime_round.py index e4eb84ccf0e25..3c290048b49ff 100644 --- a/python/pyspark/pandas/tests/indexes/test_datetime_round.py +++ b/python/pyspark/pandas/tests/indexes/test_datetime_round.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.indexes.test_datetime import DatetimeIndexTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class DatetimeIndexRoundMixin(DatetimeIndexTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/indexes/test_default.py b/python/pyspark/pandas/tests/indexes/test_default.py index e45e9904e1ff2..9f7e1e273ab9b 100644 --- a/python/pyspark/pandas/tests/indexes/test_default.py +++ b/python/pyspark/pandas/tests/indexes/test_default.py @@ -17,8 +17,8 @@ import pandas as pd -from pyspark.sql import functions as F from pyspark import pandas as ps +from pyspark.sql import functions as F from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/indexes/test_missing.py b/python/pyspark/pandas/tests/indexes/test_missing.py index 2218cdc422f6b..17f79e8e54d5b 100644 --- a/python/pyspark/pandas/tests/indexes/test_missing.py +++ b/python/pyspark/pandas/tests/indexes/test_missing.py @@ -20,7 +20,6 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.exceptions import PandasNotImplementedError from pyspark.pandas.missing.indexes import ( MissingPandasLikeDatetimeIndex, @@ -28,6 +27,7 @@ MissingPandasLikeMultiIndex, MissingPandasLikeTimedeltaIndex, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase class MissingMixin: diff --git a/python/pyspark/pandas/tests/indexes/test_reset_index.py b/python/pyspark/pandas/tests/indexes/test_reset_index.py index ed7dc49ac17b3..5ef9ac4b32c10 100644 --- a/python/pyspark/pandas/tests/indexes/test_reset_index.py +++ b/python/pyspark/pandas/tests/indexes/test_reset_index.py @@ -19,8 +19,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.testing import assert_frame_equal, assert_index_equal, assert_series_equal +from pyspark.testing.pandasutils import PandasOnSparkTestCase class FrameResetIndexMixin: diff --git a/python/pyspark/pandas/tests/io/test_csv.py b/python/pyspark/pandas/tests/io/test_csv.py index dd89b5eb09258..69d5b7f2edcee 100644 --- a/python/pyspark/pandas/tests/io/test_csv.py +++ b/python/pyspark/pandas/tests/io/test_csv.py @@ -18,8 +18,8 @@ import os from contextlib import contextmanager -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils diff --git a/python/pyspark/pandas/tests/io/test_dataframe_conversion.py b/python/pyspark/pandas/tests/io/test_dataframe_conversion.py index 363e8aa6960d5..a6aff768fef63 100644 --- a/python/pyspark/pandas/tests/io/test_dataframe_conversion.py +++ b/python/pyspark/pandas/tests/io/test_dataframe_conversion.py @@ -17,8 +17,8 @@ import os import string -import unittest import sys +import unittest import numpy as np import pandas as pd @@ -26,10 +26,10 @@ from pyspark import pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.testing.utils import ( - have_openpyxl, - openpyxl_requirement_message, have_jinja2, + have_openpyxl, jinja2_requirement_message, + openpyxl_requirement_message, ) diff --git a/python/pyspark/pandas/tests/io/test_io.py b/python/pyspark/pandas/tests/io/test_io.py index 8931591b8df84..01670e5eaf50a 100644 --- a/python/pyspark/pandas/tests/io/test_io.py +++ b/python/pyspark/pandas/tests/io/test_io.py @@ -25,8 +25,8 @@ from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.testing.utils import ( have_jinja2, - jinja2_requirement_message, have_tabulate, + jinja2_requirement_message, tabulate_requirement_message, ) diff --git a/python/pyspark/pandas/tests/io/test_series_conversion.py b/python/pyspark/pandas/tests/io/test_series_conversion.py index 1f15b6f12c15b..31acd70085810 100644 --- a/python/pyspark/pandas/tests/io/test_series_conversion.py +++ b/python/pyspark/pandas/tests/io/test_series_conversion.py @@ -15,8 +15,8 @@ # limitations under the License. # -import unittest import sys +import unittest import pandas as pd diff --git a/python/pyspark/pandas/tests/plot/test_frame_plot.py b/python/pyspark/pandas/tests/plot/test_frame_plot.py index b1917c6e2b13d..09cfaa12cf84f 100644 --- a/python/pyspark/pandas/tests/plot/test_frame_plot.py +++ b/python/pyspark/pandas/tests/plot/test_frame_plot.py @@ -15,13 +15,13 @@ # limitations under the License. # -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option, option_context -from pyspark.pandas.plot import TopNPlotBase, SampledPlotBase, HistogramPlotBase, BoxPlotBase +from pyspark.pandas.config import option_context, reset_option, set_option from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.plot import BoxPlotBase, HistogramPlotBase, SampledPlotBase, TopNPlotBase from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py b/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py index f8ac0a326b504..e7af87985c94d 100644 --- a/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py +++ b/python/pyspark/pandas/tests/plot/test_frame_plot_matplotlib.py @@ -16,14 +16,14 @@ # import base64 -from io import BytesIO import unittest +from io import BytesIO -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.testing.utils import have_matplotlib, matplotlib_requirement_message diff --git a/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py b/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py index 727d8549bf677..97a6f25fe0623 100644 --- a/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/pandas/tests/plot/test_frame_plot_plotly.py @@ -15,21 +15,21 @@ # limitations under the License. # -import unittest import pprint +import unittest -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option +from pyspark.pandas.utils import name_like_string from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.testing.utils import have_plotly, plotly_requirement_message -from pyspark.pandas.utils import name_like_string if have_plotly: - from plotly import express import plotly.graph_objs as go + from plotly import express @unittest.skipIf(not have_plotly, plotly_requirement_message) diff --git a/python/pyspark/pandas/tests/plot/test_series_plot.py b/python/pyspark/pandas/tests/plot/test_series_plot.py index f1743b6627736..1548968843c5d 100644 --- a/python/pyspark/pandas/tests/plot/test_series_plot.py +++ b/python/pyspark/pandas/tests/plot/test_series_plot.py @@ -17,11 +17,11 @@ import unittest -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.plot import PandasOnSparkPlotAccessor, BoxPlotBase +from pyspark.pandas.plot import BoxPlotBase, PandasOnSparkPlotAccessor from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.testing.utils import have_plotly, plotly_requirement_message diff --git a/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py b/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py index 1975cc1b899b6..847d70d2cbe78 100644 --- a/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py +++ b/python/pyspark/pandas/tests/plot/test_series_plot_matplotlib.py @@ -16,14 +16,14 @@ # import base64 -from io import BytesIO import unittest +from io import BytesIO import numpy as np import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.testing.utils import have_matplotlib, matplotlib_requirement_message diff --git a/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py b/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py index 757e3344feef6..4abdd8abbb8a6 100644 --- a/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py +++ b/python/pyspark/pandas/tests/plot/test_series_plot_plotly.py @@ -15,21 +15,21 @@ # limitations under the License. # -import unittest import pprint +import unittest -import pandas as pd import numpy as np +import pandas as pd from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option +from pyspark.pandas.config import reset_option, set_option from pyspark.pandas.utils import name_like_string from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.testing.utils import have_plotly, plotly_requirement_message if have_plotly: - from plotly import express import plotly.graph_objs as go + from plotly import express @unittest.skipIf(not have_plotly, plotly_requirement_message) diff --git a/python/pyspark/pandas/tests/resample/test_missing.py b/python/pyspark/pandas/tests/resample/test_missing.py index 530c5de23ac3a..a7645617983c2 100644 --- a/python/pyspark/pandas/tests/resample/test_missing.py +++ b/python/pyspark/pandas/tests/resample/test_missing.py @@ -16,8 +16,8 @@ # -import inspect import datetime +import inspect import numpy as np import pandas as pd diff --git a/python/pyspark/pandas/tests/series/test_arg_ops.py b/python/pyspark/pandas/tests/series/test_arg_ops.py index 9780296d5f9d3..a8ed19c989c1b 100644 --- a/python/pyspark/pandas/tests/series/test_arg_ops.py +++ b/python/pyspark/pandas/tests/series/test_arg_ops.py @@ -18,8 +18,8 @@ import numpy as np import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/series/test_as_type.py b/python/pyspark/pandas/tests/series/test_as_type.py index 2b17fbb4a1052..d76e4a2baa8b4 100644 --- a/python/pyspark/pandas/tests/series/test_as_type.py +++ b/python/pyspark/pandas/tests/series/test_as_type.py @@ -19,13 +19,13 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.testing.utils import is_ansi_mode_test from pyspark.pandas.typedef.typehints import ( extension_dtypes_available, extension_float_dtypes_available, extension_object_dtypes_available, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase +from pyspark.testing.utils import is_ansi_mode_test class SeriesAsTypeMixin: diff --git a/python/pyspark/pandas/tests/series/test_conversion.py b/python/pyspark/pandas/tests/series/test_conversion.py index 1405cac690d8b..51353a405568f 100644 --- a/python/pyspark/pandas/tests/series/test_conversion.py +++ b/python/pyspark/pandas/tests/series/test_conversion.py @@ -18,8 +18,8 @@ import pandas as pd -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.testing.utils import have_tabulate, tabulate_requirement_message diff --git a/python/pyspark/pandas/tests/series/test_series.py b/python/pyspark/pandas/tests/series/test_series.py index c4ee13a240f08..b5823e60c0361 100644 --- a/python/pyspark/pandas/tests/series/test_series.py +++ b/python/pyspark/pandas/tests/series/test_series.py @@ -15,24 +15,24 @@ # limitations under the License. # +import inspect import unittest from collections import defaultdict -import inspect from datetime import datetime, timedelta import numpy as np import pandas as pd -from pyspark.ml.linalg import SparseVector from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import ( - PandasOnSparkTestCase, - SPARK_CONF_ARROW_ENABLED, -) +from pyspark.ml.linalg import SparseVector from pyspark.pandas.exceptions import PandasNotImplementedError from pyspark.pandas.missing.series import MissingPandasLikeSeries from pyspark.pandas.typedef.typehints import extension_object_dtypes_available +from pyspark.testing.pandasutils import ( + SPARK_CONF_ARROW_ENABLED, + PandasOnSparkTestCase, +) class SeriesTestsMixin: diff --git a/python/pyspark/pandas/tests/series/test_string_ops_adv.py b/python/pyspark/pandas/tests/series/test_string_ops_adv.py index 9835ca1a6e4ea..079db39fe4316 100644 --- a/python/pyspark/pandas/tests/series/test_string_ops_adv.py +++ b/python/pyspark/pandas/tests/series/test_string_ops_adv.py @@ -14,10 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import pandas as pd -import numpy as np import re +import numpy as np +import pandas as pd + from pyspark import pandas as ps from pyspark.loose_version import LooseVersion from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/series/test_string_ops_basic.py b/python/pyspark/pandas/tests/series/test_string_ops_basic.py index 746ce0f3c48e6..0b9afbe3ff563 100644 --- a/python/pyspark/pandas/tests/series/test_string_ops_basic.py +++ b/python/pyspark/pandas/tests/series/test_string_ops_basic.py @@ -15,10 +15,11 @@ # limitations under the License. # -import pandas as pd -import numpy as np import re +import numpy as np +import pandas as pd + from pyspark import pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_arrow_interface.py b/python/pyspark/pandas/tests/test_arrow_interface.py index 304c38c756c41..249cc9fc82218 100644 --- a/python/pyspark/pandas/tests/test_arrow_interface.py +++ b/python/pyspark/pandas/tests/test_arrow_interface.py @@ -16,14 +16,15 @@ # import ctypes import unittest + +import pandas as pd + +from pyspark import pandas as ps +from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.testing.utils import ( have_pyarrow, pyarrow_requirement_message, ) -from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase - -import pandas as pd @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) diff --git a/python/pyspark/pandas/tests/test_config.py b/python/pyspark/pandas/tests/test_config.py index f95fde7979216..8db721b8c03e6 100644 --- a/python/pyspark/pandas/tests/test_config.py +++ b/python/pyspark/pandas/tests/test_config.py @@ -17,7 +17,7 @@ from pyspark import pandas as ps from pyspark.pandas import config -from pyspark.pandas.config import Option, DictWrapper +from pyspark.pandas.config import DictWrapper, Option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_extension.py b/python/pyspark/pandas/tests/test_extension.py index af34540815516..cd9478b37ab24 100644 --- a/python/pyspark/pandas/tests/test_extension.py +++ b/python/pyspark/pandas/tests/test_extension.py @@ -21,12 +21,12 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import assert_produces_warning, PandasOnSparkTestCase from pyspark.pandas.extensions import ( register_dataframe_accessor, - register_series_accessor, register_index_accessor, + register_series_accessor, ) +from pyspark.testing.pandasutils import PandasOnSparkTestCase, assert_produces_warning @contextlib.contextmanager diff --git a/python/pyspark/pandas/tests/test_indexops_spark.py b/python/pyspark/pandas/tests/test_indexops_spark.py index 70ebbed476831..5a0a345d6f777 100644 --- a/python/pyspark/pandas/tests/test_indexops_spark.py +++ b/python/pyspark/pandas/tests/test_indexops_spark.py @@ -17,9 +17,9 @@ import pandas as pd +from pyspark import pandas as ps from pyspark.errors import AnalysisException from pyspark.sql import functions as F -from pyspark import pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_internal.py b/python/pyspark/pandas/tests/test_internal.py index b82b89735e6bc..fdcd1800f9411 100644 --- a/python/pyspark/pandas/tests/test_internal.py +++ b/python/pyspark/pandas/tests/test_internal.py @@ -17,13 +17,13 @@ import pandas as pd -from pyspark.sql.types import LongType, StructType, StructField from pyspark.pandas.internal import ( - InternalFrame, SPARK_DEFAULT_INDEX_NAME, SPARK_INDEX_NAME_FORMAT, + InternalFrame, ) from pyspark.pandas.utils import spark_column_equals +from pyspark.sql.types import LongType, StructField, StructType from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_namespace.py b/python/pyspark/pandas/tests/test_namespace.py index 152b2804193ff..3f7f2a3bfd117 100644 --- a/python/pyspark/pandas/tests/test_namespace.py +++ b/python/pyspark/pandas/tests/test_namespace.py @@ -15,26 +15,26 @@ # limitations under the License. # -import itertools import inspect import io +import itertools import json import os import tempfile from contextlib import redirect_stdout -import pandas as pd import numpy as np +import pandas as pd import pyspark -from pyspark.loose_version import LooseVersion from pyspark import pandas as ps +from pyspark.loose_version import LooseVersion from pyspark.pandas.exceptions import PandasNotImplementedError +from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions from pyspark.pandas.namespace import _get_index_map, read_delta +from pyspark.pandas.testing import assert_frame_equal from pyspark.pandas.utils import spark_column_equals -from pyspark.pandas.missing.general_functions import MissingPandasLikeGeneralFunctions from pyspark.testing.pandasutils import PandasOnSparkTestCase -from pyspark.pandas.testing import assert_frame_equal class NamespaceTestsMixin: diff --git a/python/pyspark/pandas/tests/test_numpy_compat.py b/python/pyspark/pandas/tests/test_numpy_compat.py index a3234ed4169da..b211d5d91aaa9 100644 --- a/python/pyspark/pandas/tests/test_numpy_compat.py +++ b/python/pyspark/pandas/tests/test_numpy_compat.py @@ -15,13 +15,35 @@ # limitations under the License. # +import platform +import unittest +from decimal import Decimal + import numpy as np import pandas as pd from pyspark import pandas as ps -from pyspark.pandas import set_option, reset_option +from pyspark.loose_version import LooseVersion +from pyspark.pandas import reset_option, set_option +from pyspark.sql import functions as F from pyspark.testing.pandasutils import PandasOnSparkTestCase +# np.reciprocal(int 0) and the fmax/fmin signed-zero tie are unspecified by C/IEEE, so NumPy's +# own answer varies by CPU architecture and NumPy version. pandas-on-Spark returns one fixed +# value everywhere, which matches NumPy only on the environment it was verified against, so the +# tests comparing the two run only there. +_numpy_matches_spark = ( + platform.system() == "Linux" + and platform.machine() == "x86_64" + and LooseVersion(np.__version__) >= LooseVersion("2.3.0") +) +_skip_if_numpy_differs = unittest.skipIf( + not _numpy_matches_spark, + "NumPy's reciprocal(int 0) and fmax/fmin signed-zero tie vary by architecture and NumPy " + "version, while pandas-on-Spark returns one fixed value that matches NumPy only on " + "Linux x86_64 with NumPy >= 2.3.0", +) + class NumPyCompatTestsMixin: @classmethod @@ -41,8 +63,10 @@ def setUpClass(cls): "log", # flaky "log10", # flaky "log1p", # flaky - "modf", ] + # The sweeps below draw random integers including 0, where reciprocal diverges. + if not _numpy_matches_spark: + blacklist = blacklist + ["reciprocal"] @property def pdf(self): @@ -99,13 +123,39 @@ def test_np_math_functions(self): (np.fabs, [np.iinfo(np.int64).min, -2, 0, 2]), (np.fabs, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.invert, [np.iinfo(np.int64).min, -2, -1, 0, 1, 2, np.iinfo(np.int64).max]), + (np.isfinite, [-np.inf, -64.0, -0.0, 0.0, 64.0, np.inf, np.nan]), + (np.isinf, [-np.inf, -64.0, -0.0, 0.0, 64.0, np.inf, np.nan]), + (np.log2, [-np.inf, -64.0, -1.0, -0.0, 0.0, 1.0, 2.0, 64.0, np.inf, np.nan]), (np.negative, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.positive, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.rad2deg, [-np.inf, -64.0, -np.pi, 0.0, np.pi, 64.0, np.inf, np.nan]), + (np.rint, [-np.inf, -2.5, -1.5, -0.5, -0.0, 0.0, 0.5, 1.5, 2.5, np.inf, np.nan]), + ( + np.reciprocal, + [-np.inf, -64.0, -2.0, -1.0, -0.0, 0.0, 1.0, 2.0, 64.0, np.inf, np.nan], + ), (np.sign, [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.sinh, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.square, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), (np.tanh, [-np.inf, -64.0, -2.0, 0.0, 2.0, 64.0, np.inf, np.nan]), + ( + np.trunc, + [ + -np.inf, + -64.0, + -2.0, + -1.5, + -0.5, + -0.0, + 0.0, + 0.5, + 1.5, + 2.0, + 64.0, + np.inf, + np.nan, + ], + ), ): with self.subTest(name=np_func.__name__, values=values): pdf = pd.DataFrame({"a": values}) @@ -113,8 +163,421 @@ def test_np_math_functions(self): self.assert_eq(np_func(psdf.a), np_func(pdf.a), almost=True) + @_skip_if_numpy_differs + def test_np_reciprocal_integer(self): + # np.reciprocal on an integer column does integer division (truncated + # toward zero): 1 -> 1, -1 -> -1, and every other magnitude -> 0. The + # value 0 overflows to the int64 minimum. Cover positive, negative, and + # zero inputs to lock in parity with pandas. + for values in ( + [1, 2, 3, 64, 100], + [-1, -2, -3, -64, -100], + [-2, -1, 0, 1, 2], + [np.iinfo(np.int64).min, -1, 1, np.iinfo(np.int64).max], + ): + with self.subTest(values=values): + pdf = pd.DataFrame({"a": values}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(np.reciprocal(psdf.a), np.reciprocal(pdf.a), almost=True) + + @_skip_if_numpy_differs + def test_np_reciprocal_non_default_dtypes(self): + # The non-floating reciprocal branch also serves narrower integers, + # booleans, and decimals. numpy divides integers (and booleans, as + # int8) toward zero, so 0 overflows to the width-specific minimum + # (0 for int8/int16, int32 min for int32), while decimals take a true + # floating reciprocal. Lock in parity with pandas for each. + for dtype in ("int8", "int16", "int32"): + with self.subTest(dtype=dtype): + pdf = pd.DataFrame({"a": np.array([-2, -1, 0, 1, 2], dtype=dtype)}) + psdf = ps.from_pandas(pdf) + + self.assert_eq(np.reciprocal(psdf.a), np.reciprocal(pdf.a), almost=True) + + # Boolean: numpy promotes to int8 (True -> 1, False -> 0). + pdf = pd.DataFrame({"a": [True, False, True]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(np.reciprocal(psdf.a), np.reciprocal(pdf.a), almost=True) + + # Decimal: numpy takes a floating reciprocal. 0 is excluded because + # numpy raises DivisionByZero on Decimal('0'). + pdf = pd.DataFrame({"a": [Decimal("2.5"), Decimal("-4"), Decimal("0.5")]}) + psdf = ps.from_pandas(pdf) + self.assert_eq(np.reciprocal(psdf.a), np.reciprocal(pdf.a), almost=True) + + def test_np_bitwise_shift_functions(self): + pdf = pd.DataFrame( + { + "value": [np.iinfo(np.int64).min, -2, -1, 0, 1, 2, np.iinfo(np.int64).max], + "bits": [-1, 0, 1, 63, 64, 65, 2], + } + ) + psdf = ps.from_pandas(pdf) + + for np_func in (np.left_shift, np.right_shift): + with self.subTest(name=np_func.__name__): + self.assert_eq( + np_func(psdf.value, psdf.bits), np_func(pdf.value, pdf.bits), almost=True + ) + + def test_np_float_power(self): + for pdf in ( + pd.DataFrame({"base": [-64, -2, -1, 0, 1, 2, 64], "exponent": [-2, -1, 0, 1, 2, 3, 2]}), + pd.DataFrame( + { + "base": [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan], + "exponent": [2.0, 3.0, -2.0, -3.0, -3.0, 0.5, -2.0, 2.0, 2.0], + } + ), + ): + psdf = ps.from_pandas(pdf) + self.assert_eq( + np.float_power(psdf.base, psdf.exponent), + np.float_power(pdf.base, pdf.exponent), + almost=True, + ) + + def test_np_ldexp(self): + pdf = pd.DataFrame( + { + "x": [ + -np.inf, + -64.0, + -2.0, + -0.0, + 0.0, + 1.0, + 2.0, + 64.0, + np.inf, + np.nan, + 1.0, + 1.0, + 1.0, + 0.0, + -0.0, + np.inf, + -np.inf, + ], + "exp": [ + 2, + 3, + -2, + -3, + -3, + 0, + -2, + 2, + 2, + 2, + -1074, + -1075, + 1024, + 1024, + 1024, + -1075, + -1075, + ], + } + ) + psdf = ps.from_pandas(pdf) + + result = np.ldexp(psdf.x, psdf.exp) + expected = np.ldexp(pdf.x, pdf.exp) + self.assert_eq(result, expected, almost=True) + self.assert_eq(np.signbit(result.to_pandas()), np.signbit(expected)) + + def test_np_fmod(self): + for pdf in ( + pd.DataFrame( + { + "x1": [-64, -2, -1, 0, 1, 2, 64], + "x2": [2, 3, -2, -3, -3, 0, 2], + } + ), + pd.DataFrame( + { + "x1": [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan, 1.0], + "x2": [2.0, 3.0, -2.0, -3.0, -3.0, 0.0, -np.inf, np.inf, 2.0, 0.0], + } + ), + pd.DataFrame( + { + "x1": pd.array([1, 2, None, None], dtype="Int64"), + "x2": pd.array([2, None, 2, 0], dtype="Int64"), + } + ), + ): + psdf = ps.from_pandas(pdf) + + self.assert_eq(np.fmod(psdf.x1, psdf.x2), np.fmod(pdf.x1, pdf.x2), almost=True) + + def test_np_modf(self): + # np.modf(x) returns a tuple (fractional part, integral part). + for pdf in ( + pd.DataFrame({"a": [-64, -2, -1, 0, 1, 2, 64]}), + pd.DataFrame( + {"a": [-np.inf, -64.0, -2.0, -0.5, -0.0, 0.0, 0.5, 2.0, 64.0, np.inf, np.nan]} + ), + pd.DataFrame({"a": pd.array([1, -2, None], dtype="Int64")}), + ): + psdf = ps.from_pandas(pdf) + ps_fractional, ps_integral = np.modf(psdf.a) + pd_fractional, pd_integral = np.modf(pdf.a) + self.assert_eq(ps_fractional, pd_fractional, almost=True) + self.assert_eq(ps_integral, pd_integral, almost=True) + + # almost=True treats -0.0 and 0.0 as equal, so verify the sign of zero explicitly: + # the fractional part of a negative whole number (-2.0 -> -0.0) and the fractional + # part of -inf (-> -0.0) must keep the input's sign, as must the integral part of a + # value in (-1, 0) (-0.5 -> -0.0). + pdf = pd.DataFrame({"a": [-2.0, -0.5, -0.0, 0.0, 0.5, 2.0, -np.inf, np.inf]}) + psdf = ps.from_pandas(pdf) + ps_fractional, ps_integral = np.modf(psdf.a) + pd_fractional, pd_integral = np.modf(pdf.a) + self.assert_eq(np.signbit(ps_fractional.to_pandas()), np.signbit(pd_fractional)) + self.assert_eq(np.signbit(ps_integral.to_pandas()), np.signbit(pd_integral)) + + # DataFrame input: np.modf returns a tuple of DataFrames, one per output. + pdf = pd.DataFrame( + { + "a": [-3.5, -2.0, -0.5, 0.0, 2.7], + "b": [1.5, -0.0, np.inf, -np.inf, np.nan], + } + ) + psdf = ps.from_pandas(pdf) + ps_fractional, ps_integral = np.modf(psdf) + pd_fractional, pd_integral = np.modf(pdf) + self.assert_eq(ps_fractional, pd_fractional, almost=True) + self.assert_eq(ps_integral, pd_integral, almost=True) + self.assert_eq(np.signbit(ps_fractional.to_pandas()), np.signbit(pd_fractional)) + self.assert_eq(np.signbit(ps_integral.to_pandas()), np.signbit(pd_integral)) + + # Index input: np.modf returns a tuple of Index objects. + pidx = pd.Index([-3.5, -2.0, -0.5, 0.0, 2.7]) + psidx = ps.from_pandas(pidx) + ps_fractional, ps_integral = np.modf(psidx) + pd_fractional, pd_integral = np.modf(pidx) + self.assert_eq(ps_fractional, pd_fractional, almost=True) + self.assert_eq(ps_integral, pd_integral, almost=True) + + def test_floor_divide_func(self): + from pyspark.pandas.numpy_compat import _floor_divide_func + + for pdf in ( + pd.DataFrame( + { + "x1": [-64, -2, -1, 0, 1, 2, 64, -1, 0], + "x2": [2, 3, -2, -3, -3, 0, 2, 0, 0], + } + ), + pd.DataFrame( + { + "x1": [ + -np.inf, + -64.0, + -2.0, + -0.0, + 0.0, + 2.0, + 64.0, + np.inf, + np.nan, + 1.0, + -1.0, + np.inf, + -np.inf, + np.inf, + -np.inf, + 1.0, + ], + "x2": [ + 2.0, + 3.0, + -2.0, + -3.0, + -3.0, + 0.0, + -np.inf, + np.inf, + 2.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0, + -0.0, + np.nan, + ], + } + ), + pd.DataFrame( + { + "x1": pd.array([1, None, None], dtype="Int64"), + "x2": pd.array([None, 2, 0], dtype="Int64"), + } + ), + ): + psdf = ps.from_pandas(pdf) + result = ( + psdf.spark.frame() + .select(_floor_divide_func(F.col("x1"), F.col("x2")).alias("result")) + .toPandas()["result"] + .rename(None) + ) + self.assert_eq(result, np.floor_divide(pdf.x1, pdf.x2), almost=True) + + def test_np_logaddexp(self): + for pdf in ( + pd.DataFrame( + { + "x1": [-64, -2, -1, 0, 1, 2, 64], + "x2": [2, 3, -2, -3, -3, 0, 2], + } + ), + pd.DataFrame( + { + "x1": [ + -np.inf, + -np.inf, + -2.0, + -2.0, + -0.0, + 0.0, + 2.0, + np.inf, + np.inf, + np.nan, + -1000.0, + -np.inf, + -0.0, + ], + "x2": [ + -np.inf, + 3.0, + -np.inf, + 2.0, + 0.0, + -0.0, + np.inf, + 2.0, + np.inf, + 2.0, + 1000.0, + -0.0, + -np.inf, + ], + } + ), + ): + psdf = ps.from_pandas(pdf) + for np_func in (np.logaddexp, np.logaddexp2): + result = np_func(psdf.x1, psdf.x2) + expected = np_func(pdf.x1, pdf.x2) + self.assert_eq(result, expected, almost=True) + self.assert_eq(np.signbit(result.to_pandas()), np.signbit(expected)) + + @_skip_if_numpy_differs + def test_np_fmax_fmin(self): + for pdf in ( + pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [2, 1, 0, -1, -2]}), + pd.DataFrame( + { + "x1": [np.nan, 2.0, np.nan, -np.inf, -2.0, -0.0, 0.0, 2.0, np.inf], + "x2": [2.0, np.nan, np.nan, np.inf, -np.inf, 0.0, -0.0, np.inf, -np.inf], + } + ), + pd.DataFrame({"x1": [-0.0, 0.0], "x2": [0.0, -0.0]}), + ): + psdf = ps.from_pandas(pdf) + for np_func in (np.fmax, np.fmin): + result = np_func(psdf.x1, psdf.x2) + expected = np_func(pdf.x1, pdf.x2) + self.assert_eq(result, expected, almost=True) + # NumPy's vectorized implementation may select either zero operand, whereas + # its scalar implementation consistently selects the first one. + expected_signbit = pd.Series( + [np.signbit(np_func(x1, x2)) for x1, x2 in zip(pdf.x1, pdf.x2)] + ) + self.assert_eq(np.signbit(result.to_pandas()), expected_signbit) + + def test_np_copysign(self): + for pdf in ( + pd.DataFrame( + { + "x1": [-64, -2, -1, 0, 1, 2, 64], + "x2": [2, -3, -2, -3, 3, -1, 2], + } + ), + pd.DataFrame( + { + "x1": [-np.inf, -64.0, -2.0, -0.0, 0.0, 2.0, 64.0, np.inf, np.nan, 1.0], + "x2": [2.0, -3.0, -2.0, 0.0, -0.0, -1.0, np.inf, -np.inf, 2.0, np.nan], + } + ), + pd.DataFrame( + { + "x1": pd.array([1, -2, 3, None, None], dtype="Int64"), + "x2": pd.array([-2, 3, None, 2, None], dtype="Int64"), + } + ), + ): + psdf = ps.from_pandas(pdf) + result = np.copysign(psdf.x1, psdf.x2) + expected = np.copysign(pdf.x1, pdf.x2) + self.assert_eq(result, expected, almost=True) + # copysign only differs from |x| in the sign bit, so assert on signbit + # explicitly -- 0.0 == -0.0 numerically and would hide a wrong sign. + self.assert_eq(np.signbit(result.to_pandas()), np.signbit(expected)) + + def test_np_copysign_signed_zero(self): + # np.copysign takes the sign from y's IEEE-754 sign bit, not from y < 0: + # copysign(1.0, -0.0) == -1.0 and copysign(1.0, 0.0) == 1.0. + pdf = pd.DataFrame( + { + "x1": [1.0, 1.0, -0.0, -0.0, 3.0], + "x2": [0.0, -0.0, 0.0, -0.0, -0.0], + } + ) + psdf = ps.from_pandas(pdf) + result = np.copysign(psdf.x1, psdf.x2).to_pandas() + expected = np.copysign(pdf.x1, pdf.x2) + self.assert_eq(result, expected) + self.assert_eq(np.signbit(result), np.signbit(expected)) + + def test_np_heaviside(self): + for pdf in ( + pd.DataFrame({"x1": [-2, -1, 0, 1, 2], "x2": [-2, -1, 0, 1, 2]}), + pd.DataFrame( + { + "x1": [-np.inf, -2.0, -0.0, 0.0, 0.0, 2.0, np.inf, np.nan], + "x2": [2.0, -2.0, -0.0, 0.5, np.nan, np.nan, -0.0, 2.0], + } + ), + ): + psdf = ps.from_pandas(pdf) + self.assert_eq( + np.heaviside(psdf.x1, psdf.x2), np.heaviside(pdf.x1, pdf.x2), almost=True + ) + + def test_np_signbit(self): + # np.signbit returns the IEEE-754 sign bit, which differs from (x < 0) only + # at -0.0: the sign bit is set even though -0.0 is not less than zero. A + # missing value in a default (numpy-backed) dtype arrives as a NaN and maps + # to False (np.signbit(nan) is False), whereas a genuine <NA> in a nullable + # dtype (e.g. Int64) propagates. A nullable Float64 <NA> is indistinguishable + # from a NaN after from_pandas, so it is deliberately not covered here. + for pdf in ( + pd.DataFrame({"a": [-0.0, 0.0, -1.0, 1.0, -np.inf, np.inf, np.nan]}), + pd.DataFrame({"a": [1, -2, None]}), + pd.DataFrame({"a": pd.array([1, -2, None], dtype="Int64")}), + ): + psdf = ps.from_pandas(pdf) + self.assert_eq(np.signbit(psdf.a), np.signbit(pdf.a)) + def test_np_spark_compat_series(self): - from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings + from pyspark.pandas.numpy_compat import binary_np_spark_mappings, unary_np_spark_mappings # Use randomly generated dataFrame pdf = pd.DataFrame( @@ -164,7 +627,7 @@ def test_np_spark_compat_series(self): reset_option("compute.ops_on_diff_frames") def test_np_spark_compat_frame(self): - from pyspark.pandas.numpy_compat import unary_np_spark_mappings, binary_np_spark_mappings + from pyspark.pandas.numpy_compat import binary_np_spark_mappings, unary_np_spark_mappings # Use randomly generated dataFrame pdf = pd.DataFrame( diff --git a/python/pyspark/pandas/tests/test_repr.py b/python/pyspark/pandas/tests/test_repr.py index 9ed2b34bc9c78..5cc6a6064fd2d 100644 --- a/python/pyspark/pandas/tests/test_repr.py +++ b/python/pyspark/pandas/tests/test_repr.py @@ -18,7 +18,7 @@ import numpy as np from pyspark import pandas as ps -from pyspark.pandas.config import set_option, reset_option, option_context +from pyspark.pandas.config import option_context, reset_option, set_option from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_typedef.py b/python/pyspark/pandas/tests/test_typedef.py index 1b9db0a6a29ac..70c14ac66cedc 100644 --- a/python/pyspark/pandas/tests/test_typedef.py +++ b/python/pyspark/pandas/tests/test_typedef.py @@ -15,44 +15,43 @@ # limitations under the License. # -import unittest import datetime import decimal +import unittest from typing import List +import numpy as np import pandas import pandas as pd from pandas.api.types import CategoricalDtype -import numpy as np +from pyspark import pandas as ps from pyspark.loose_version import LooseVersion +from pyspark.pandas.typedef import ( + as_spark_type, + extension_dtypes_available, + extension_float_dtypes_available, + extension_object_dtypes_available, + infer_return_type, + pandas_on_spark_type, +) from pyspark.sql.types import ( ArrayType, BinaryType, BooleanType, + ByteType, + DateType, + DecimalType, + DoubleType, FloatType, IntegerType, LongType, + ShortType, StringType, StructField, StructType, - ByteType, - ShortType, - DateType, - DecimalType, - DoubleType, TimestampType, ) - -from pyspark.pandas.typedef import ( - as_spark_type, - extension_dtypes_available, - extension_float_dtypes_available, - extension_object_dtypes_available, - infer_return_type, - pandas_on_spark_type, -) -from pyspark import pandas as ps from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/test_utils.py b/python/pyspark/pandas/tests/test_utils.py index 85e0f52afe97c..50954ceea5435 100644 --- a/python/pyspark/pandas/tests/test_utils.py +++ b/python/pyspark/pandas/tests/test_utils.py @@ -17,6 +17,7 @@ import pandas as pd +from pyspark.errors import PySparkAssertionError from pyspark.pandas.indexes.base import Index from pyspark.pandas.utils import ( lazy_property, @@ -27,10 +28,9 @@ ) from pyspark.testing.pandasutils import ( PandasOnSparkTestCase, - _assert_pandas_equal, _assert_pandas_almost_equal, + _assert_pandas_equal, ) -from pyspark.errors import PySparkAssertionError some_global_variable = 0 diff --git a/python/pyspark/pandas/tests/window/test_ewm_error.py b/python/pyspark/pandas/tests/window/test_ewm_error.py index 32987532db6f3..17a4380ff6829 100644 --- a/python/pyspark/pandas/tests/window/test_ewm_error.py +++ b/python/pyspark/pandas/tests/window/test_ewm_error.py @@ -16,8 +16,8 @@ # import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.pandas.window import ExponentialMoving +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils class EWMErrorMixin: diff --git a/python/pyspark/pandas/tests/window/test_expanding_adv.py b/python/pyspark/pandas/tests/window/test_expanding_adv.py index 086e53a003e5f..decfa35f9640b 100644 --- a/python/pyspark/pandas/tests/window/test_expanding_adv.py +++ b/python/pyspark/pandas/tests/window/test_expanding_adv.py @@ -19,8 +19,8 @@ from pyspark import pandas as ps from pyspark.loose_version import LooseVersion -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.window.test_expanding import ExpandingTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class ExpandingAdvMixin(ExpandingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/window/test_groupby_expanding_adv.py b/python/pyspark/pandas/tests/window/test_groupby_expanding_adv.py index a59a97eb2667b..9f6555d91e576 100644 --- a/python/pyspark/pandas/tests/window/test_groupby_expanding_adv.py +++ b/python/pyspark/pandas/tests/window/test_groupby_expanding_adv.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.window.test_groupby_expanding import GroupByExpandingTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class GroupByExpandingAdvMixin(GroupByExpandingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/window/test_groupby_rolling_adv.py b/python/pyspark/pandas/tests/window/test_groupby_rolling_adv.py index 7515c7be9cadb..7ebbecb332446 100644 --- a/python/pyspark/pandas/tests/window/test_groupby_rolling_adv.py +++ b/python/pyspark/pandas/tests/window/test_groupby_rolling_adv.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.pandas.tests.window.test_groupby_rolling import GroupByRollingTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils class GroupByRollingAdvMixin(GroupByRollingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/window/test_missing.py b/python/pyspark/pandas/tests/window/test_missing.py index 99fd6a8230fd9..7269b56a99978 100644 --- a/python/pyspark/pandas/tests/window/test_missing.py +++ b/python/pyspark/pandas/tests/window/test_missing.py @@ -21,11 +21,11 @@ from pyspark.pandas.exceptions import PandasNotImplementedError from pyspark.pandas.missing.window import ( MissingPandasLikeExpanding, - MissingPandasLikeRolling, MissingPandasLikeExpandingGroupby, - MissingPandasLikeRollingGroupby, MissingPandasLikeExponentialMoving, MissingPandasLikeExponentialMovingGroupby, + MissingPandasLikeRolling, + MissingPandasLikeRollingGroupby, ) from pyspark.testing.pandasutils import PandasOnSparkTestCase diff --git a/python/pyspark/pandas/tests/window/test_rolling_adv.py b/python/pyspark/pandas/tests/window/test_rolling_adv.py index 5f5907ed20cca..bcef6ad595af9 100644 --- a/python/pyspark/pandas/tests/window/test_rolling_adv.py +++ b/python/pyspark/pandas/tests/window/test_rolling_adv.py @@ -18,8 +18,8 @@ import pandas as pd from pyspark import pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase from pyspark.pandas.tests.window.test_rolling import RollingTestingFuncMixin +from pyspark.testing.pandasutils import PandasOnSparkTestCase class RollingAdvMixin(RollingTestingFuncMixin): diff --git a/python/pyspark/pandas/tests/window/test_rolling_error.py b/python/pyspark/pandas/tests/window/test_rolling_error.py index cf1e2fe4d6e09..18f449db7825c 100644 --- a/python/pyspark/pandas/tests/window/test_rolling_error.py +++ b/python/pyspark/pandas/tests/window/test_rolling_error.py @@ -16,8 +16,8 @@ # import pyspark.pandas as ps -from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils from pyspark.pandas.window import Rolling +from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils class RollingErrorMixin: diff --git a/python/pyspark/pandas/typedef/typehints.py b/python/pyspark/pandas/typedef/typehints.py index 8d99584f81a99..ad2fb4c2b35f1 100644 --- a/python/pyspark/pandas/typedef/typehints.py +++ b/python/pyspark/pandas/typedef/typehints.py @@ -24,12 +24,12 @@ import typing from collections.abc import Iterable from inspect import isclass -from typing import Any, Callable, Generic, List, Optional, Tuple, Union, Type, get_type_hints +from typing import Any, Callable, Generic, List, Optional, Tuple, Type, Union, get_type_hints import numpy as np import pandas as pd -from pandas.api.types import CategoricalDtype, pandas_dtype from pandas.api.extensions import ExtensionDtype +from pandas.api.types import CategoricalDtype, pandas_dtype from pyspark.loose_version import LooseVersion @@ -63,12 +63,13 @@ extension_dtypes = () import pyarrow as pa + import pyspark.sql.types as types -from pyspark.sql.pandas.types import to_arrow_type, from_arrow_type # For running doctests and reference resolution in PyCharm. from pyspark import pandas as ps # noqa: F401 from pyspark.pandas._typing import Dtype, T +from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type if typing.TYPE_CHECKING: from pyspark.pandas.internal import InternalField @@ -626,8 +627,8 @@ def infer_return_type(f: Callable) -> Union[SeriesType, DataFrameType, ScalarTyp # We should re-import to make sure the class 'SeriesType' is not treated as a class # within this module locally. See Series.__class_getitem__ which imports this class # canonically. - from pyspark.pandas.internal import InternalField, SPARK_INDEX_NAME_FORMAT - from pyspark.pandas.typedef import SeriesType, NameTypeHolder, IndexNameTypeHolder + from pyspark.pandas.internal import SPARK_INDEX_NAME_FORMAT, InternalField + from pyspark.pandas.typedef import IndexNameTypeHolder, NameTypeHolder, SeriesType from pyspark.pandas.utils import name_like_string tpe = get_type_hints(f).get("return", None) @@ -798,7 +799,7 @@ def create_tuple_for_frame_type(params: Any) -> object: def _to_type_holders(params: Any) -> Tuple: - from pyspark.pandas.typedef import NameTypeHolder, IndexNameTypeHolder + from pyspark.pandas.typedef import IndexNameTypeHolder, NameTypeHolder is_with_index = ( isinstance(params, tuple) @@ -941,6 +942,7 @@ def _new_type_holders( def _test() -> None: import doctest import sys + import pyspark.pandas.typedef.typehints globs = pyspark.pandas.typedef.typehints.__dict__.copy() diff --git a/python/pyspark/pandas/usage_logging/__init__.py b/python/pyspark/pandas/usage_logging/__init__.py index f02ee7efebb26..7a028a6329b08 100644 --- a/python/pyspark/pandas/usage_logging/__init__.py +++ b/python/pyspark/pandas/usage_logging/__init__.py @@ -20,10 +20,11 @@ import pandas as pd +from pyspark.instrumentation_utils import _attach from pyspark.pandas import config, namespace, sql_formatter from pyspark.pandas.accessors import PandasOnSparkFrameMethods -from pyspark.pandas.frame import DataFrame from pyspark.pandas.datetimes import DatetimeMethods +from pyspark.pandas.frame import DataFrame from pyspark.pandas.groupby import DataFrameGroupBy, SeriesGroupBy from pyspark.pandas.indexes.base import Index from pyspark.pandas.indexes.category import CategoricalIndex @@ -43,11 +44,11 @@ from pyspark.pandas.missing.series import MissingPandasLikeSeries from pyspark.pandas.missing.window import ( MissingPandasLikeExpanding, - MissingPandasLikeRolling, MissingPandasLikeExpandingGroupby, - MissingPandasLikeRollingGroupby, MissingPandasLikeExponentialMoving, MissingPandasLikeExponentialMovingGroupby, + MissingPandasLikeRolling, + MissingPandasLikeRollingGroupby, ) from pyspark.pandas.series import Series from pyspark.pandas.spark.accessors import ( @@ -59,12 +60,11 @@ from pyspark.pandas.window import ( Expanding, ExpandingGroupby, - Rolling, - RollingGroupby, ExponentialMoving, ExponentialMovingGroupby, + Rolling, + RollingGroupby, ) -from pyspark.instrumentation_utils import _attach def attach(logger_module: Union[str, ModuleType]) -> None: diff --git a/python/pyspark/pandas/usage_logging/usage_logger.py b/python/pyspark/pandas/usage_logging/usage_logger.py index a17c52a157a6c..052cd6722a2a1 100644 --- a/python/pyspark/pandas/usage_logging/usage_logger.py +++ b/python/pyspark/pandas/usage_logging/usage_logger.py @@ -19,8 +19,8 @@ The reference implementation of usage logger using the Python standard logging library. """ -from inspect import Signature import logging +from inspect import Signature from typing import Any, Optional diff --git a/python/pyspark/pandas/utils.py b/python/pyspark/pandas/utils.py index 489cfaa7dc3d3..7d36f8431b813 100644 --- a/python/pyspark/pandas/utils.py +++ b/python/pyspark/pandas/utils.py @@ -19,11 +19,13 @@ """ import functools -from contextlib import contextmanager import json import os import threading +import warnings +from contextlib import contextmanager from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -33,33 +35,33 @@ Optional, Tuple, Union, - TYPE_CHECKING, cast, no_type_check, overload, ) -import warnings import pandas as pd from pandas.api.types import is_list_like -from pyspark.sql import functions as F, Column, DataFrame as PySparkDataFrame, SparkSession -from pyspark.sql.types import DoubleType -from pyspark.sql.utils import is_remote -from pyspark.errors import PySparkTypeError, UnsupportedOperationException from pyspark import pandas as ps +from pyspark.errors import PySparkTypeError, UnsupportedOperationException from pyspark.pandas._typing import ( Axis, + DataFrameOrSeries, Label, Name, - DataFrameOrSeries, ) from pyspark.pandas.typedef.typehints import as_spark_type +from pyspark.sql import Column, SparkSession +from pyspark.sql import DataFrame as PySparkDataFrame +from pyspark.sql import functions as F +from pyspark.sql.types import DoubleType +from pyspark.sql.utils import is_remote if TYPE_CHECKING: - from pyspark.pandas.indexes.base import Index from pyspark.pandas.base import IndexOpsMixin from pyspark.pandas.frame import DataFrame + from pyspark.pandas.indexes.base import Index from pyspark.pandas.internal import InternalFrame from pyspark.pandas.series import Series @@ -133,11 +135,11 @@ def combine_frames( from pyspark.pandas.config import get_option from pyspark.pandas.frame import DataFrame from pyspark.pandas.internal import ( - InternalField, - InternalFrame, HIDDEN_COLUMNS, NATURAL_ORDER_COLUMN_NAME, SPARK_INDEX_NAME_FORMAT, + InternalField, + InternalFrame, ) from pyspark.pandas.series import Series @@ -1153,8 +1155,8 @@ def ansi_mode_context(spark: SparkSession) -> Iterator[None]: def is_ansi_mode_enabled(spark: SparkSession) -> bool: def _is_ansi_mode_enabled() -> bool: if is_remote(): - from pyspark.sql.connect.session import SparkSession as ConnectSession from pyspark.pandas.config import _key_format, _options_dict + from pyspark.sql.connect.session import SparkSession as ConnectSession client = cast(ConnectSession, spark).client ansi_mode_support, ansi_enabled = client.get_config_with_defaults( @@ -1186,11 +1188,12 @@ def _is_ansi_mode_enabled() -> bool: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.utils + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pandas/window.py b/python/pyspark/pandas/window.py index 20ebbddd163b9..2a99144abadd2 100644 --- a/python/pyspark/pandas/window.py +++ b/python/pyspark/pandas/window.py @@ -21,24 +21,24 @@ import numpy as np import pandas as pd +from pyspark import pandas as ps # noqa: F401 from pyspark.loose_version import LooseVersion -from pyspark.sql import Window -from pyspark.sql import functions as F -from pyspark.sql.internal import InternalFunction as SF +from pyspark.pandas._typing import FrameLike +from pyspark.pandas.groupby import DataFrameGroupBy, GroupBy +from pyspark.pandas.internal import NATURAL_ORDER_COLUMN_NAME, SPARK_INDEX_NAME_FORMAT from pyspark.pandas.missing.window import ( - MissingPandasLikeRolling, - MissingPandasLikeRollingGroupby, MissingPandasLikeExpanding, MissingPandasLikeExpandingGroupby, MissingPandasLikeExponentialMoving, MissingPandasLikeExponentialMovingGroupby, + MissingPandasLikeRolling, + MissingPandasLikeRollingGroupby, ) -from pyspark import pandas as ps # noqa: F401 -from pyspark.pandas._typing import FrameLike -from pyspark.pandas.groupby import GroupBy, DataFrameGroupBy -from pyspark.pandas.internal import NATURAL_ORDER_COLUMN_NAME, SPARK_INDEX_NAME_FORMAT from pyspark.pandas.utils import scol_for +from pyspark.sql import Window +from pyspark.sql import functions as F from pyspark.sql.column import Column +from pyspark.sql.internal import InternalFunction as SF from pyspark.sql.types import ( DoubleType, ) @@ -3036,11 +3036,12 @@ def __repr__(self) -> str: def _test() -> None: - import os import doctest + import os import sys - from pyspark.sql import SparkSession + import pyspark.pandas.window + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/pipelines/__init__.py b/python/pyspark/pipelines/__init__.py index bd41c9ecd6b2e..ebaa6e6a65ce5 100644 --- a/python/pyspark/pipelines/__init__.py +++ b/python/pyspark/pipelines/__init__.py @@ -17,11 +17,11 @@ from pyspark.pipelines.api import ( append_flow, create_auto_cdc_flow, + create_sink, create_streaming_table, materialized_view, table, temporary_view, - create_sink, ) __all__ = [ diff --git a/python/pyspark/pipelines/add_pipeline_analysis_context.py b/python/pyspark/pipelines/add_pipeline_analysis_context.py index 6d0bd4dd7308c..d9c5a1a26107e 100644 --- a/python/pyspark/pipelines/add_pipeline_analysis_context.py +++ b/python/pyspark/pipelines/add_pipeline_analysis_context.py @@ -15,10 +15,9 @@ # limitations under the License. # from contextlib import contextmanager -from typing import Generator, Optional -from pyspark.sql import SparkSession +from typing import Any, Generator, Optional, cast -from typing import Any, cast +from pyspark.sql import SparkSession @contextmanager @@ -34,9 +33,10 @@ def add_pipeline_analysis_context( # Likely related to SPARK-47544. client = cast(Any, spark).client try: - import pyspark.sql.connect.proto as pb2 from google.protobuf import any_pb2 + import pyspark.sql.connect.proto as pb2 + analysis_context = pb2.PipelineAnalysisContext( dataflow_graph_id=dataflow_graph_id, flow_name=flow_name ) diff --git a/python/pyspark/pipelines/api.py b/python/pyspark/pipelines/api.py index 19bd47e4af765..da3a6d07ea932 100644 --- a/python/pyspark/pipelines/api.py +++ b/python/pyspark/pipelines/api.py @@ -17,18 +17,18 @@ from typing import Callable, Dict, List, Literal, Optional, Union, overload from pyspark.errors import PySparkTypeError, PySparkValueError -from pyspark.pipelines.graph_element_registry import get_active_graph_element_registry -from pyspark.pipelines.type_error_utils import validate_optional_list_of_str_arg from pyspark.pipelines.flow import AutoCdcFlow, Flow, QueryFunction -from pyspark.pipelines.source_code_location import ( - get_caller_source_code_location, -) +from pyspark.pipelines.graph_element_registry import get_active_graph_element_registry from pyspark.pipelines.output import ( MaterializedView, + Sink, StreamingTable, TemporaryView, - Sink, ) +from pyspark.pipelines.source_code_location import ( + get_caller_source_code_location, +) +from pyspark.pipelines.type_error_utils import validate_optional_list_of_str_arg from pyspark.sql import Column from pyspark.sql.types import StructType @@ -541,6 +541,7 @@ def create_auto_cdc_flow( *, track_history_column_list: Optional[Union[List[str], List[Column]]] = None, track_history_except_column_list: Optional[Union[List[str], List[Column]]] = None, + spark_conf: Optional[Dict[str, str]] = None, ) -> None: """ Create an Auto CDC flow into the target table from the Change Data Capture (CDC) source. @@ -600,6 +601,9 @@ def create_auto_cdc_flow( to be 2. :param name: The name of the flow for this create_auto_cdc_flow command. When unspecified, \ this will build a "default flow" with name equal to the target name. + :param spark_conf: A dict whose keys are the conf names and values are the conf values. \ + These confs will be set when the flow is executed; they can override confs set for the \ + destination, for the pipeline, or on the cluster. """ # Lazy import: pyspark.sql.connect.functions.builtin transitively imports grpc, which is # not available in the docs-build environment. pyspark.pipelines.api is loaded eagerly @@ -736,6 +740,7 @@ def create_auto_cdc_flow( stored_as_scd_type=stored_as_scd_type, track_history_column_list=track_history_column_list, track_history_except_column_list=track_history_except_column_list, + spark_conf=spark_conf or {}, source_code_location=source_code_location, ) diff --git a/python/pyspark/pipelines/block_session_mutations.py b/python/pyspark/pipelines/block_session_mutations.py index df63d2023a4ba..b8bf9ff1ac0c7 100644 --- a/python/pyspark/pipelines/block_session_mutations.py +++ b/python/pyspark/pipelines/block_session_mutations.py @@ -15,7 +15,7 @@ # limitations under the License. # from contextlib import contextmanager -from typing import Generator, NoReturn, List, Callable +from typing import Callable, Generator, List, NoReturn from pyspark.errors import PySparkException from pyspark.sql.connect.catalog import Catalog diff --git a/python/pyspark/pipelines/cli.py b/python/pyspark/pipelines/cli.py index 9598987fc3f35..4c6dfa9adc7de 100644 --- a/python/pyspark/pipelines/cli.py +++ b/python/pyspark/pipelines/cli.py @@ -22,22 +22,23 @@ $ bin/spark-pipelines run --spec /path/to/pipeline.yaml """ -from contextlib import contextmanager import argparse import glob import importlib.util import os -import yaml +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Any, Generator, List, Mapping, Optional, Sequence +import yaml + from pyspark.errors import PySparkException, PySparkTypeError -from pyspark.sql import SparkSession +from pyspark.pipelines.add_pipeline_analysis_context import add_pipeline_analysis_context from pyspark.pipelines.block_session_mutations import block_session_mutations from pyspark.pipelines.graph_element_registry import ( - graph_element_registration_context, GraphElementRegistry, + graph_element_registration_context, ) from pyspark.pipelines.init_cli import init from pyspark.pipelines.logging_utils import log_with_curr_timestamp @@ -46,11 +47,10 @@ ) from pyspark.pipelines.spark_connect_pipeline import ( create_dataflow_graph, - start_run, handle_pipeline_events, + start_run, ) - -from pyspark.pipelines.add_pipeline_analysis_context import add_pipeline_analysis_context +from pyspark.sql import SparkSession PIPELINE_SPEC_FILE_NAMES = ["spark-pipeline.yaml", "spark-pipeline.yml"] diff --git a/python/pyspark/pipelines/flow.py b/python/pyspark/pipelines/flow.py index 8a0f44e2abe96..85aa50dc2e6e1 100644 --- a/python/pyspark/pipelines/flow.py +++ b/python/pyspark/pipelines/flow.py @@ -17,9 +17,8 @@ from dataclasses import dataclass from typing import Callable, Dict, List, Literal, Optional -from pyspark.sql import DataFrame -from pyspark.sql import Column from pyspark.pipelines.source_code_location import SourceCodeLocation +from pyspark.sql import Column, DataFrame QueryFunction = Callable[[], DataFrame] @@ -65,6 +64,8 @@ class AutoCdcFlow: history record. :param track_history_except_column_list: Optional SCD2-only columns excluded from history \ tracking. + :param spark_conf: A dict where the keys are the Spark configuration property names and the + values are the property values. These properties will be set on the flow. :param source_code_location: The location of the source code that created this flow. """ @@ -79,4 +80,5 @@ class AutoCdcFlow: stored_as_scd_type: Optional[Literal[1, 2, "1", "2"]] track_history_column_list: Optional[List[Column]] track_history_except_column_list: Optional[List[Column]] + spark_conf: Dict[str, str] source_code_location: SourceCodeLocation diff --git a/python/pyspark/pipelines/graph_element_registry.py b/python/pyspark/pipelines/graph_element_registry.py index 4eddabaabda0e..2cd6caa67bd2b 100644 --- a/python/pyspark/pipelines/graph_element_registry.py +++ b/python/pyspark/pipelines/graph_element_registry.py @@ -16,15 +16,14 @@ # from abc import ABC, abstractmethod -from pathlib import Path - -from pyspark.pipelines.output import Output -from pyspark.pipelines.flow import AutoCdcFlow, Flow from contextlib import contextmanager from contextvars import ContextVar +from pathlib import Path from typing import Generator, Optional from pyspark.errors import PySparkRuntimeError +from pyspark.pipelines.flow import AutoCdcFlow, Flow +from pyspark.pipelines.output import Output class GraphElementRegistry(ABC): diff --git a/python/pyspark/pipelines/spark_connect_graph_element_registry.py b/python/pyspark/pipelines/spark_connect_graph_element_registry.py index a46eaca5f4cb3..c2ff4b6c8c0cc 100644 --- a/python/pyspark/pipelines/spark_connect_graph_element_registry.py +++ b/python/pyspark/pipelines/spark_connect_graph_element_registry.py @@ -15,26 +15,26 @@ # limitations under the License. # from pathlib import Path +from typing import Any, List, Optional, cast +import pyspark.sql.connect.proto as pb2 from pyspark.errors import PySparkTypeError -from pyspark.sql import SparkSession, Column -from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame +from pyspark.pipelines.add_pipeline_analysis_context import add_pipeline_analysis_context +from pyspark.pipelines.flow import AutoCdcFlow, Flow +from pyspark.pipelines.graph_element_registry import GraphElementRegistry from pyspark.pipelines.output import ( - Output, MaterializedView, - Table, + Output, Sink, StreamingTable, + Table, TemporaryView, ) -from pyspark.pipelines.flow import AutoCdcFlow, Flow -from pyspark.pipelines.graph_element_registry import GraphElementRegistry from pyspark.pipelines.source_code_location import SourceCodeLocation +from pyspark.sql import Column, SparkSession +from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame from pyspark.sql.connect.types import pyspark_types_to_proto_types from pyspark.sql.types import StructType -from typing import Any, List, Optional, cast -import pyspark.sql.connect.proto as pb2 -from pyspark.pipelines.add_pipeline_analysis_context import add_pipeline_analysis_context class SparkConnectGraphElementRegistry(GraphElementRegistry): @@ -165,7 +165,7 @@ def to_plans(cols: Optional[List[Column]]) -> list: flow_name=flow.name, target_dataset_name=flow.target, auto_cdc_flow_details=auto_cdc_details, - sql_conf={}, + sql_conf=flow.spark_conf, source_code_location=source_code_location_to_proto(flow.source_code_location), ) diff --git a/python/pyspark/pipelines/spark_connect_pipeline.py b/python/pyspark/pipelines/spark_connect_pipeline.py index 5a9e5802753e8..6b5858c087292 100644 --- a/python/pyspark/pipelines/spark_connect_pipeline.py +++ b/python/pyspark/pipelines/spark_connect_pipeline.py @@ -15,12 +15,12 @@ # limitations under the License. # from datetime import timezone -from typing import Any, Dict, Mapping, Iterator, Optional, cast, Sequence +from typing import Any, Dict, Iterator, Mapping, Optional, Sequence, cast import pyspark.sql.connect.proto as pb2 -from pyspark.sql import SparkSession from pyspark.errors.exceptions.base import PySparkValueError from pyspark.pipelines.logging_utils import log_with_provided_timestamp +from pyspark.sql import SparkSession def create_dataflow_graph( diff --git a/python/pyspark/pipelines/tests/local_graph_element_registry.py b/python/pyspark/pipelines/tests/local_graph_element_registry.py index 3b9ea15a1ed6b..b4eb46bb30a8b 100644 --- a/python/pyspark/pipelines/tests/local_graph_element_registry.py +++ b/python/pyspark/pipelines/tests/local_graph_element_registry.py @@ -19,9 +19,9 @@ from pathlib import Path from typing import List, Sequence -from pyspark.pipelines.output import Output from pyspark.pipelines.flow import AutoCdcFlow, Flow from pyspark.pipelines.graph_element_registry import GraphElementRegistry +from pyspark.pipelines.output import Output @dataclass(frozen=True) diff --git a/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py b/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py index f99e7d56c3d51..bc9b5674d06b7 100644 --- a/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py +++ b/python/pyspark/pipelines/tests/test_add_pipeline_analysis_context.py @@ -19,8 +19,8 @@ from pyspark.testing.connectutils import ( ReusedConnectTestCase, - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: @@ -94,9 +94,10 @@ def test_setup_failure_does_not_mask_original_error(self): # If any setup step fails before the extension is registered, extension_id stays None and # the finally block must skip remove_user_context_extension(None) - which would raise # AttributeError and mask the original error. Cover each step that can fail. - import pyspark.sql.connect.proto as pb2 from google.protobuf import any_pb2 + import pyspark.sql.connect.proto as pb2 + failing_any = mock.MagicMock() failing_any.Pack.side_effect = ValueError("boom") diff --git a/python/pyspark/pipelines/tests/test_auto_cdc_flow.py b/python/pyspark/pipelines/tests/test_auto_cdc_flow.py index c142e1b1562c2..4af5c9d1fbabe 100644 --- a/python/pyspark/pipelines/tests/test_auto_cdc_flow.py +++ b/python/pyspark/pipelines/tests/test_auto_cdc_flow.py @@ -21,14 +21,14 @@ from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.sql import Column from pyspark.testing.connectutils import ( - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: from pyspark import pipelines as dp - from pyspark.pipelines.graph_element_registry import graph_element_registration_context from pyspark.pipelines.flow import AutoCdcFlow + from pyspark.pipelines.graph_element_registry import graph_element_registration_context from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry from pyspark.sql.connect.functions.builtin import col, expr @@ -72,11 +72,27 @@ def test_create_auto_cdc_flow_with_all_args(self): column_list=[col("id"), col("val")], stored_as_scd_type=1, name="my_flow", + spark_conf={"spark.sql.shuffle.partitions": "8"}, ) flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) self.assertEqual(flow.name, "my_flow") self.assertEqual(flow.stored_as_scd_type, 1) + self.assertEqual(flow.spark_conf, {"spark.sql.shuffle.partitions": "8"}) + + def test_create_auto_cdc_flow_spark_conf_defaults_to_empty(self): + registry = LocalGraphElementRegistry() + with graph_element_registration_context(registry): + dp.create_streaming_table("target") + dp.create_auto_cdc_flow( + target="target", + source="source", + keys=[col("key")], + sequence_by=expr("seq"), + ) + + flow = cast(AutoCdcFlow, registry.auto_cdc_flows[0]) + self.assertEqual(flow.spark_conf, {}) def test_create_auto_cdc_flow_with_string_args(self): # Verify that string forms of column / expression arguments are normalized to diff --git a/python/pyspark/pipelines/tests/test_block_session_mutations.py b/python/pyspark/pipelines/tests/test_block_session_mutations.py index 922308810d55f..81ddbe13eba81 100644 --- a/python/pyspark/pipelines/tests/test_block_session_mutations.py +++ b/python/pyspark/pipelines/tests/test_block_session_mutations.py @@ -21,15 +21,15 @@ from pyspark.sql.types import StringType from pyspark.testing.connectutils import ( ReusedConnectTestCase, - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: from pyspark.pipelines.block_session_mutations import ( - block_session_mutations, BLOCKED_METHODS, ERROR_CLASS, + block_session_mutations, ) diff --git a/python/pyspark/pipelines/tests/test_cli.py b/python/pyspark/pipelines/tests/test_cli.py index e183df6c347de..afe109ac52532 100644 --- a/python/pyspark/pipelines/tests/test_cli.py +++ b/python/pyspark/pipelines/tests/test_cli.py @@ -15,29 +15,29 @@ # limitations under the License. # -import unittest import tempfile import textwrap +import unittest from pathlib import Path from pyspark.errors import PySparkException from pyspark.testing.connectutils import ( ReusedConnectTestCase, - should_test_connect, connect_requirement_message, + should_test_connect, ) from pyspark.testing.utils import have_yaml, yaml_requirement_message if should_test_connect and have_yaml: from pyspark.pipelines.cli import ( + LibrariesGlob, + PipelineSpec, change_dir, find_pipeline_spec, load_pipeline_spec, register_definitions, - unpack_pipeline_spec, - LibrariesGlob, - PipelineSpec, run, + unpack_pipeline_spec, ) from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry diff --git a/python/pyspark/pipelines/tests/test_decorators.py b/python/pyspark/pipelines/tests/test_decorators.py index 6544c430424ff..f7282c9b4e6bc 100644 --- a/python/pyspark/pipelines/tests/test_decorators.py +++ b/python/pyspark/pipelines/tests/test_decorators.py @@ -17,8 +17,8 @@ import unittest -from pyspark.errors import PySparkTypeError from pyspark import pipelines as dp +from pyspark.errors import PySparkTypeError class DecoratorsTest(unittest.TestCase): diff --git a/python/pyspark/pipelines/tests/test_graph_element_registry.py b/python/pyspark/pipelines/tests/test_graph_element_registry.py index 1e6fcf224a0ac..4d5b0b5be0a1c 100644 --- a/python/pyspark/pipelines/tests/test_graph_element_registry.py +++ b/python/pyspark/pipelines/tests/test_graph_element_registry.py @@ -16,13 +16,13 @@ # import unittest +from typing import cast +from pyspark import pipelines as dp from pyspark.errors import PySparkException from pyspark.pipelines.graph_element_registry import graph_element_registration_context -from pyspark import pipelines as dp from pyspark.pipelines.output import Sink from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry -from typing import cast class GraphElementRegistryTest(unittest.TestCase): diff --git a/python/pyspark/pipelines/tests/test_init_cli.py b/python/pyspark/pipelines/tests/test_init_cli.py index f4f48ade35aa1..1fbba0684a771 100644 --- a/python/pyspark/pipelines/tests/test_init_cli.py +++ b/python/pyspark/pipelines/tests/test_init_cli.py @@ -14,14 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest import tempfile +import unittest from pathlib import Path from pyspark.testing.connectutils import ( ReusedConnectTestCase, - should_test_connect, connect_requirement_message, + should_test_connect, ) from pyspark.testing.utils import have_yaml, yaml_requirement_message @@ -29,8 +29,8 @@ from pyspark.pipelines.cli import ( change_dir, find_pipeline_spec, - load_pipeline_spec, init, + load_pipeline_spec, register_definitions, ) from pyspark.pipelines.tests.local_graph_element_registry import LocalGraphElementRegistry diff --git a/python/pyspark/pipelines/tests/test_spark_connect.py b/python/pyspark/pipelines/tests/test_spark_connect.py index b6b3935107d2d..9cd0376b05490 100644 --- a/python/pyspark/pipelines/tests/test_spark_connect.py +++ b/python/pyspark/pipelines/tests/test_spark_connect.py @@ -24,8 +24,8 @@ from pyspark import pipelines as dp from pyspark.testing.connectutils import ( ReusedConnectTestCase, - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: @@ -36,8 +36,8 @@ ) from pyspark.pipelines.spark_connect_pipeline import ( create_dataflow_graph, - start_run, handle_pipeline_events, + start_run, ) diff --git a/python/pyspark/pipelines/type_error_utils.py b/python/pyspark/pipelines/type_error_utils.py index 6ec0f4e71fd0b..fa8027e3e88ff 100644 --- a/python/pyspark/pipelines/type_error_utils.py +++ b/python/pyspark/pipelines/type_error_utils.py @@ -15,6 +15,7 @@ # limitations under the License. # from typing import List, Optional + from pyspark.errors import PySparkTypeError diff --git a/python/pyspark/profiler.py b/python/pyspark/profiler.py index 45ab1e5afdbfe..b112dee0f4b16 100644 --- a/python/pyspark/profiler.py +++ b/python/pyspark/profiler.py @@ -15,7 +15,14 @@ # limitations under the License. # +import atexit +import cProfile +import linecache +import os +import pstats +import sys from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -23,16 +30,9 @@ Optional, Tuple, Type, - TYPE_CHECKING, Union, cast, ) -import cProfile -import pstats -import linecache -import os -import atexit -import sys import pyspark from pyspark.accumulators import AccumulatorParam diff --git a/python/pyspark/rddsampler.py b/python/pyspark/rddsampler.py index 14db5af9e6537..1c6b50df992aa 100644 --- a/python/pyspark/rddsampler.py +++ b/python/pyspark/rddsampler.py @@ -15,12 +15,11 @@ # limitations under the License. # -import sys -import random import math +import random +import sys from typing import Generic, Hashable, Iterable, Iterator, Optional, TypeVar - T = TypeVar("T") K = TypeVar("K", bound=Hashable) diff --git a/python/pyspark/resource/__init__.py b/python/pyspark/resource/__init__.py index b3b33942827d9..5c5b4fd8cbbfb 100644 --- a/python/pyspark/resource/__init__.py +++ b/python/pyspark/resource/__init__.py @@ -20,13 +20,13 @@ """ from pyspark.resource.information import ResourceInformation +from pyspark.resource.profile import ResourceProfile, ResourceProfileBuilder from pyspark.resource.requests import ( - TaskResourceRequest, - TaskResourceRequests, ExecutorResourceRequest, ExecutorResourceRequests, + TaskResourceRequest, + TaskResourceRequests, ) -from pyspark.resource.profile import ResourceProfile, ResourceProfileBuilder __all__ = [ "TaskResourceRequest", diff --git a/python/pyspark/resource/profile.py b/python/pyspark/resource/profile.py index 9fb7b42a8a733..f66d51e1c360b 100644 --- a/python/pyspark/resource/profile.py +++ b/python/pyspark/resource/profile.py @@ -15,13 +15,13 @@ # limitations under the License. # from threading import RLock -from typing import overload, Dict, Union, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Optional, Union, overload from pyspark.resource.requests import ( + ExecutorResourceRequest, + ExecutorResourceRequests, TaskResourceRequest, TaskResourceRequests, - ExecutorResourceRequests, - ExecutorResourceRequest, ) if TYPE_CHECKING: @@ -322,6 +322,7 @@ def build(self) -> ResourceProfile: def _test() -> None: import doctest import sys + from pyspark import SparkContext globs = globals().copy() diff --git a/python/pyspark/resource/requests.py b/python/pyspark/resource/requests.py index 999faf6b83140..7d502805f0a99 100644 --- a/python/pyspark/resource/requests.py +++ b/python/pyspark/resource/requests.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import overload, Optional, Dict, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Optional, overload from pyspark.util import _parse_memory diff --git a/python/pyspark/resource/tests/test_connect_resources.py b/python/pyspark/resource/tests/test_connect_resources.py index a80e00814504b..6f0ce62652f26 100644 --- a/python/pyspark/resource/tests/test_connect_resources.py +++ b/python/pyspark/resource/tests/test_connect_resources.py @@ -14,15 +14,15 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest import os +import unittest from pyspark.errors import PySparkException -from pyspark.resource import ResourceProfileBuilder, TaskResourceRequests, ExecutorResourceRequests +from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests from pyspark.sql import SparkSession from pyspark.testing.connectutils import ( - should_test_connect, connect_requirement_message, + should_test_connect, ) from pyspark.testing.utils import eventually diff --git a/python/pyspark/resource/tests/test_resources.py b/python/pyspark/resource/tests/test_resources.py index 8e9b304774acf..1815b586946eb 100644 --- a/python/pyspark/resource/tests/test_resources.py +++ b/python/pyspark/resource/tests/test_resources.py @@ -15,6 +15,7 @@ # limitations under the License. # import unittest + from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests from pyspark.sql import SparkSession from pyspark.testing.utils import ( @@ -56,7 +57,7 @@ def assert_request_contents(exec_reqs, task_reqs): assert_request_contents(ereqs.requests, treqs.requests) rp = rpb.require(ereqs).require(treqs).build assert_request_contents(rp.executorResources, rp.taskResources) - from pyspark import SparkContext, SparkConf + from pyspark import SparkConf, SparkContext sc = SparkContext(conf=SparkConf()) rdd = sc.parallelize(range(10)).withResources(rp) diff --git a/python/pyspark/resultiterable.py b/python/pyspark/resultiterable.py index d866ce0f092c5..9dae2a98d63b7 100644 --- a/python/pyspark/resultiterable.py +++ b/python/pyspark/resultiterable.py @@ -15,7 +15,7 @@ # limitations under the License. # -from typing import TypeVar, TYPE_CHECKING, Iterator, Iterable +from typing import TYPE_CHECKING, Iterable, Iterator, TypeVar if TYPE_CHECKING: from pyspark._typing import SizedIterable diff --git a/python/pyspark/serializers.py b/python/pyspark/serializers.py index 48166c948b5b1..8b9cea2979066 100644 --- a/python/pyspark/serializers.py +++ b/python/pyspark/serializers.py @@ -53,17 +53,16 @@ >>> sc.stop() """ -import sys -import os -from itertools import chain, product +import codecs +import collections +import itertools import marshal +import os +import pickle import struct +import sys import types -import collections import zlib -import itertools -import pickle -import codecs pickle_protocol = pickle.HIGHEST_PROTOCOL @@ -125,9 +124,6 @@ def _load_stream_without_unbatching(self, stream): def __eq__(self, other): return isinstance(other, self.__class__) and self.__dict__ == other.__dict__ - def __ne__(self, other): - return not self.__eq__(other) - def __repr__(self): return "%s()" % self.__class__.__name__ @@ -172,13 +168,6 @@ def _read_with_length(self, stream): raise EOFError return self.loads(obj) - def dumps(self, obj): - """ - Serialize an object into a byte array. - When batching is used, this will be called with an array of objects. - """ - raise NotImplementedError - def loads(self, obj): """ Deserialize an object from a byte array. @@ -211,7 +200,7 @@ def dump_stream(self, iterator, stream): self.serializer.dump_stream(self._batched(iterator), stream) def load_stream(self, stream): - return chain.from_iterable(self._load_stream_without_unbatching(stream)) + return itertools.chain.from_iterable(self._load_stream_without_unbatching(stream)) def _load_stream_without_unbatching(self, stream): return self.serializer.load_stream(stream) @@ -290,10 +279,10 @@ def _load_stream_without_unbatching(self, stream): val_batch_stream = self.val_ser._load_stream_without_unbatching(stream) for key_batch, val_batch in zip(key_batch_stream, val_batch_stream): # for correctness with repeated cartesian/zip this must be returned as one batch - yield product(key_batch, val_batch) + yield itertools.product(key_batch, val_batch) def load_stream(self, stream): - return chain.from_iterable(self._load_stream_without_unbatching(stream)) + return itertools.chain.from_iterable(self._load_stream_without_unbatching(stream)) def __repr__(self): return "CartesianDeserializer(%s, %s)" % (str(self.key_ser), str(self.val_ser)) @@ -327,7 +316,7 @@ def _load_stream_without_unbatching(self, stream): yield zip(key_batch, val_batch) def load_stream(self, stream): - return chain.from_iterable(self._load_stream_without_unbatching(stream)) + return itertools.chain.from_iterable(self._load_stream_without_unbatching(stream)) def __repr__(self): return "PairDeserializer(%s, %s)" % (str(self.key_ser), str(self.val_ser)) @@ -478,41 +467,12 @@ def loads(self, obj): return marshal.loads(obj) -class AutoSerializer(FramedSerializer): - """ - Choose marshal or pickle as serialization protocol automatically - """ - - def __init__(self): - FramedSerializer.__init__(self) - self._type = None - - def dumps(self, obj): - if self._type is not None: - return b"P" + pickle.dumps(obj, -1) - try: - return b"M" + marshal.dumps(obj) - except Exception: - self._type = b"P" - return b"P" + pickle.dumps(obj, -1) - - def loads(self, obj): - _type = obj[0] - if _type == b"M": - return marshal.loads(obj[1:]) - elif _type == b"P": - return pickle.loads(obj[1:]) - else: - raise ValueError("invalid serialization type: %s" % _type) - - class CompressedSerializer(FramedSerializer): """ Compress the serialized data """ def __init__(self, serializer): - FramedSerializer.__init__(self) assert isinstance(serializer, FramedSerializer), "serializer must be a FramedSerializer" self.serializer = serializer diff --git a/python/pyspark/shell.py b/python/pyspark/shell.py index 435cc9de3b2e1..3e503f13ff211 100644 --- a/python/pyspark/shell.py +++ b/python/pyspark/shell.py @@ -25,8 +25,9 @@ import builtins import os import platform -import warnings import sys +import warnings +from urllib.parse import urlparse import pyspark from pyspark.core.context import SparkContext @@ -34,7 +35,6 @@ from pyspark.sql import SparkSession from pyspark.sql.context import SQLContext from pyspark.sql.utils import is_remote -from urllib.parse import urlparse if getattr(builtins, "__IPYTHON__", False): # (Only) during PYTHONSTARTUP execution, IPython temporarily adds the parent diff --git a/python/pyspark/shuffle.py b/python/pyspark/shuffle.py index c37bb7fbce304..dc4288b3ad5a7 100644 --- a/python/pyspark/shuffle.py +++ b/python/pyspark/shuffle.py @@ -15,36 +15,36 @@ # limitations under the License. # -import os -import platform -import shutil -import warnings import gc +import heapq import itertools import operator +import os +import platform import random +import shutil import sys -import heapq +import warnings from typing import ( + IO, + TYPE_CHECKING, Any, Callable, Generic, Hashable, - IO, Iterable, Iterator, Optional, - TYPE_CHECKING, TypeVar, Union, ) from pyspark.serializers import ( + AutoBatchedSerializer, BatchedSerializer, + CompressedSerializer, CPickleSerializer, FlattenedValuesSerializer, - CompressedSerializer, - AutoBatchedSerializer, Serializer, ) from pyspark.util import fail_on_stopiteration diff --git a/python/pyspark/sql/__init__.py b/python/pyspark/sql/__init__.py index 117f57f29d3fc..f21e9e3823341 100644 --- a/python/pyspark/sql/__init__.py +++ b/python/pyspark/sql/__init__.py @@ -40,19 +40,19 @@ For working with window functions. """ -from pyspark.sql.types import Geography, Geometry, Row, VariantVal -from pyspark.sql.context import SQLContext, HiveContext, UDFRegistration, UDTFRegistration -from pyspark.sql.session import SparkSession -from pyspark.sql.column import Column from pyspark.sql.catalog import Catalog +from pyspark.sql.column import Column +from pyspark.sql.context import HiveContext, SQLContext, UDFRegistration, UDTFRegistration from pyspark.sql.dataframe import DataFrame, DataFrameNaFunctions, DataFrameStatFunctions from pyspark.sql.group import GroupedData -from pyspark.sql.observation import Observation -from pyspark.sql.readwriter import DataFrameReader, DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter -from pyspark.sql.window import Window, WindowSpec +from pyspark.sql.observation import Observation from pyspark.sql.pandas.group_ops import PandasCogroupedOps +from pyspark.sql.readwriter import DataFrameReader, DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.session import SparkSession +from pyspark.sql.types import Geography, Geometry, Row, VariantVal from pyspark.sql.utils import is_remote +from pyspark.sql.window import Window, WindowSpec __all__ = [ "SparkSession", diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi index 94e3ccf770939..604e7394aa0ee 100644 --- a/python/pyspark/sql/_typing.pyi +++ b/python/pyspark/sql/_typing.pyi @@ -16,6 +16,9 @@ # specific language governing permissions and limitations # under the License. +import datetime +import decimal +import pstats from typing import ( Any, Callable, @@ -27,15 +30,12 @@ from typing import ( TypeVar, Union, ) -from typing_extensions import Literal, Protocol -import datetime -import decimal -import pstats +from typing_extensions import Literal, Protocol +import pyspark.sql.types from pyspark._typing import PrimitiveType from pyspark.profiler import CodeMapDict -import pyspark.sql.types from pyspark.sql.column import Column from pyspark.sql.tvf_argument import TableValuedFunctionArgument @@ -64,6 +64,11 @@ RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...], pyspark.sql.types.Row) SQLBatchedUDFType = Literal[100] SQLArrowBatchedUDFType = Literal[101] +SQLArrowElementwiseUDFType = Literal[102] +SQLScalarPandasElementwiseUDFType = Literal[103] +SQLScalarPandasIterElementwiseUDFType = Literal[104] +SQLScalarArrowElementwiseUDFType = Literal[105] +SQLScalarArrowIterElementwiseUDFType = Literal[106] SQLTableUDFType = Literal[300] SQLArrowTableUDFType = Literal[301] SQLArrowUDTFType = Literal[302] diff --git a/python/pyspark/sql/aggregator.py b/python/pyspark/sql/aggregator.py new file mode 100644 index 0000000000000..8ce383b9b4f2c --- /dev/null +++ b/python/pyspark/sql/aggregator.py @@ -0,0 +1,139 @@ +# +# 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. +# +""" +Incremental user-defined aggregators for PySpark, the Python analog of Scala's +``org.apache.spark.sql.expressions.Aggregator``. +""" + +from abc import ABC, abstractmethod +from typing import Any, Tuple + +from pyspark.errors import PySparkNotImplementedError +from pyspark.sql.types import DataType, StructType + +__all__ = ["Aggregator"] + + +class Aggregator(ABC): + """ + Base class for a user-defined *incremental* aggregator, the Python analog of Scala's + :class:`org.apache.spark.sql.expressions.Aggregator`. + + Unlike a grouped-aggregate ``pandas_udf`` (which materializes the whole group and is invoked + once), an :class:`Aggregator` is executed as a genuine two-stage aggregation with map-side + combine: :meth:`reduce` folds input rows into a per-group *buffer* on the map side, the buffers + are shuffled by the grouping key, :meth:`merge` combines the partial buffers of each group, and + :meth:`finish` produces the final output value. + + The buffer is represented as a Python :class:`tuple` whose elements correspond, in order, to the + fields of :attr:`bufferSchema`. An input row is likewise a tuple of the argument values passed + to the aggregator call. :meth:`merge` must be associative and commutative (the framework may + combine partial buffers in any order), and :meth:`zero` must be its identity element -- see + :meth:`zero` for the identity law that makes the result independent of the partition count. + + .. versionadded:: 4.4.0 + + Examples + -------- + A mean aggregator:: + + from pyspark.sql.aggregator import Aggregator + from pyspark.sql.functions import udaf + from pyspark.sql.types import StructType, StructField, DoubleType, LongType + + class Mean(Aggregator): + @property + def bufferSchema(self): + return StructType([ + StructField("sum", DoubleType()), + StructField("count", LongType()), + ]) + + @property + def outputType(self): + return DoubleType() + + def zero(self): + return (0.0, 0) + + def reduce(self, buffer, value): + (v,) = value + if v is None: # ignore null inputs, like SQL aggregates do + return buffer + return (buffer[0] + v, buffer[1] + 1) + + def merge(self, b1, b2): + return (b1[0] + b2[0], b1[1] + b2[1]) + + def finish(self, buffer): + return buffer[0] / buffer[1] if buffer[1] else None + + mean = udaf(Mean()) + df.groupBy("k").agg(mean(df.v)).show() + """ + + @property + @abstractmethod + def bufferSchema(self) -> StructType: + """The schema of the intermediate buffer that crosses the shuffle.""" + ... + + @property + @abstractmethod + def outputType(self) -> DataType: + """The data type of the aggregator's output value.""" + ... + + @abstractmethod + def zero(self) -> Tuple[Any, ...]: + """The initial (identity) buffer value, as a tuple matching :attr:`bufferSchema`. + + This must be the identity element for :meth:`merge`:: + + merge(buffer, zero()) == buffer + merge(zero(), buffer) == buffer + + A fresh ``zero()`` seeds every partition -- and every early-flushed chunk of the map-side + combine -- so associativity and commutativity of :meth:`merge` alone do not guarantee a + partition-independent result; the identity law above is what makes the aggregate value + independent of how the input is split across partitions and batches. + """ + ... + + @abstractmethod + def reduce(self, buffer: Tuple[Any, ...], value: Tuple[Any, ...]) -> Tuple[Any, ...]: + """Fold a single input row ``value`` into ``buffer`` and return the updated buffer.""" + ... + + @abstractmethod + def merge(self, buffer1: Tuple[Any, ...], buffer2: Tuple[Any, ...]) -> Tuple[Any, ...]: + """Merge two partial buffers into one. Must be associative and commutative.""" + ... + + @abstractmethod + def finish(self, buffer: Tuple[Any, ...]) -> Any: + """Produce the output value from the final merged buffer.""" + ... + + # The aggregator instance is shipped to the worker as the UDF "function"; making it callable + # lets it satisfy ``UserDefinedFunction``'s ``callable`` check. It is never actually invoked as + # a function -- the worker calls :meth:`zero`/:meth:`reduce`/:meth:`merge`/:meth:`finish`. + def __call__(self, *args: Any, **kwargs: Any) -> Any: + raise PySparkNotImplementedError( + errorClass="NOT_IMPLEMENTED", + messageParameters={"feature": "calling an Aggregator directly; wrap it with udaf(...)"}, + ) diff --git a/python/pyspark/sql/avro/functions.py b/python/pyspark/sql/avro/functions.py index 79e301d08d2a0..f547f17a676c3 100644 --- a/python/pyspark/sql/avro/functions.py +++ b/python/pyspark/sql/avro/functions.py @@ -19,7 +19,7 @@ A collections of builtin avro functions """ -from typing import Dict, Optional, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Dict, Optional, cast from pyspark.errors import PySparkTypeError from pyspark.sql.column import Column @@ -78,6 +78,7 @@ def from_avro( [Row(value=Row(avro=Row(age=2, name='Alice')))] """ from py4j.java_gateway import JVMView + from pyspark.sql.classic.column import _to_java_column if not isinstance(data, (Column, str)): @@ -148,6 +149,7 @@ def to_avro(data: "ColumnOrName", jsonFormatSchema: str = "") -> Column: [Row(suite=b'\\x02\\x00')] """ from py4j.java_gateway import JVMView + from pyspark.sql.classic.column import _to_java_column if not isinstance(data, (Column, str)): @@ -184,6 +186,7 @@ def to_avro(data: "ColumnOrName", jsonFormatSchema: str = "") -> Column: def _test() -> None: import os import sys + from pyspark.testing.sqlutils import search_jar avro_jar = search_jar("connector/avro", "spark-avro", "spark-avro") @@ -201,8 +204,9 @@ def _test() -> None: os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) import doctest - from pyspark.sql import SparkSession + import pyspark.sql.avro.functions + from pyspark.sql import SparkSession globs = pyspark.sql.avro.functions.__dict__.copy() spark = ( diff --git a/python/pyspark/sql/catalog.py b/python/pyspark/sql/catalog.py index d698c458b0155..ed7d6ac7482b5 100644 --- a/python/pyspark/sql/catalog.py +++ b/python/pyspark/sql/catalog.py @@ -17,17 +17,16 @@ import sys import warnings -from typing import Any, Callable, Dict, NamedTuple, List, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional from pyspark.errors import PySparkTypeError -from pyspark.storagelevel import StorageLevel from pyspark.sql.dataframe import DataFrame from pyspark.sql.session import SparkSession from pyspark.sql.types import StructType +from pyspark.storagelevel import StorageLevel if TYPE_CHECKING: - from pyspark.sql._typing import UserDefinedFunctionLike - from pyspark.sql._typing import DataTypeOrString + from pyspark.sql._typing import DataTypeOrString, UserDefinedFunctionLike class CatalogMetadata(NamedTuple): @@ -1570,10 +1569,11 @@ def refreshByPath(self, path: str) -> None: def _test() -> None: - import os import doctest - from pyspark.sql import SparkSession + import os + import pyspark.sql.catalog + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/classic/column.py b/python/pyspark/sql/classic/column.py index 5f54443b69e46..4a3ac86968ffb 100644 --- a/python/pyspark/sql/classic/column.py +++ b/python/pyspark/sql/classic/column.py @@ -15,32 +15,33 @@ # limitations under the License. # -import sys import json +import sys import warnings from typing import ( - cast, - overload, + TYPE_CHECKING, Any, Callable, Iterable, List, Optional, Tuple, - TYPE_CHECKING, Union, + cast, + overload, ) -from pyspark.sql.column import Column as ParentColumn from pyspark.errors import PySparkAttributeError, PySparkTypeError, PySparkValueError from pyspark.errors.utils import with_origin_to_class +from pyspark.sql.column import Column as ParentColumn from pyspark.sql.types import DataType -from pyspark.sql.utils import get_active_spark_context, enum_to_value +from pyspark.sql.utils import enum_to_value, get_active_spark_context if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext - from pyspark.sql._typing import ColumnOrName, LiteralType, DecimalLiteral, DateTimeLiteral + from pyspark.sql._typing import ColumnOrName, DateTimeLiteral, DecimalLiteral, LiteralType from pyspark.sql.window import WindowSpec __all__ = ["Column"] @@ -62,6 +63,13 @@ def _create_column_from_name(name: str) -> "JavaObject": return cast(JVMView, sc._jvm).functions.col(name) +def _to_java_column_opt(col: Optional["ColumnOrName"]) -> Optional["JavaObject"]: + if col is None: + return None + else: + return _to_java_column(col) + + def _to_java_column(col: "ColumnOrName") -> "JavaObject": if isinstance(col, Column): jcol = col._jc @@ -650,9 +658,13 @@ def outer(self) -> ParentColumn: return Column(jc) def __nonzero__(self) -> None: + try: + column_repr = self._jc.toString() + except Exception: + column_repr = "<unknown>" raise PySparkValueError( errorClass="CANNOT_CONVERT_COLUMN_INTO_BOOL", - messageParameters={}, + messageParameters={"column": column_repr}, ) __bool__ = __nonzero__ @@ -663,8 +675,9 @@ def __repr__(self) -> str: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.column + from pyspark.sql import SparkSession # It inherits docstrings but doctests cannot detect them so we run # the parent classe's doctests here directly. diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py index eeed6fef44aa4..aea6869f455f8 100644 --- a/python/pyspark/sql/classic/dataframe.py +++ b/python/pyspark/sql/classic/dataframe.py @@ -15,14 +15,15 @@ # limitations under the License. # -import os import json -import sys +import os import random +import sys import warnings from collections.abc import Iterable -from functools import reduce, cached_property +from functools import cached_property, reduce from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -35,53 +36,57 @@ Union, cast, overload, - TYPE_CHECKING, ) from pyspark import _NoValue -from pyspark.resource import ResourceProfile from pyspark._globals import _NoValueType from pyspark.errors import ( AnalysisException, + PySparkAttributeError, + PySparkIndexError, PySparkTypeError, PySparkValueError, - PySparkIndexError, - PySparkAttributeError, -) -from pyspark.util import ( - _load_from_socket, - _local_iterator_from_socket, ) +from pyspark.resource import ResourceProfile from pyspark.serializers import BatchedSerializer, CPickleSerializer, UTF8Deserializer -from pyspark.storagelevel import StorageLevel -from pyspark.traceback_utils import SCCallSiteSync +from pyspark.sql.classic.column import _to_java_column, _to_list, _to_seq from pyspark.sql.column import Column -from pyspark.sql.functions import builtin as F -from pyspark.sql.classic.column import _to_seq, _to_list, _to_java_column -from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 -from pyspark.sql.merge import MergeIntoWriter -from pyspark.sql.streaming import DataStreamWriter -from pyspark.sql.types import ( - StructType, - Row, - _parse_datatype_json_string, -) from pyspark.sql.dataframe import ( DataFrame as ParentDataFrame, +) +from pyspark.sql.dataframe import ( DataFrameNaFunctions as ParentDataFrameNaFunctions, +) +from pyspark.sql.dataframe import ( DataFrameStatFunctions as ParentDataFrameStatFunctions, ) -from pyspark.sql.utils import get_active_spark_context, to_java_array, to_scala_map +from pyspark.sql.functions import builtin as F +from pyspark.sql.merge import MergeIntoWriter from pyspark.sql.pandas.conversion import PandasConversionMixin from pyspark.sql.pandas.map_ops import PandasMapOpsMixin +from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.table_arg import TableArg +from pyspark.sql.types import ( + Row, + StructType, + _parse_datatype_json_string, +) +from pyspark.sql.utils import get_active_spark_context, to_java_array, to_scala_map +from pyspark.storagelevel import StorageLevel +from pyspark.traceback_utils import SCCallSiteSync +from pyspark.util import ( + _load_from_socket, + _local_iterator_from_socket, +) if TYPE_CHECKING: - from py4j.java_gateway import JavaObject import pyarrow as pa - from pyspark.core.rdd import RDD - from pyspark.core.context import SparkContext + from py4j.java_gateway import JavaObject + from pyspark._typing import PrimitiveType + from pyspark.core.context import SparkContext + from pyspark.core.rdd import RDD from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.sql._typing import ( ColumnOrName, @@ -89,17 +94,19 @@ LiteralType, OptionalPrimitiveType, ) + from pyspark.sql.context import SQLContext + from pyspark.sql.group import GroupedData + from pyspark.sql.metrics import ExecutionInfo + from pyspark.sql.observation import Observation from pyspark.sql.pandas._typing import ( - PandasMapIterFunction, ArrowMapIterFunction, + PandasMapIterFunction, + ) + from pyspark.sql.pandas._typing import ( DataFrameLike as PandasDataFrameLike, ) - from pyspark.sql.context import SQLContext - from pyspark.sql.session import SparkSession - from pyspark.sql.group import GroupedData - from pyspark.sql.observation import Observation - from pyspark.sql.metrics import ExecutionInfo from pyspark.sql.plot import PySparkPlotAccessor + from pyspark.sql.session import SparkSession class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin): @@ -1946,9 +1953,9 @@ def mergeInto(self, table: str, condition: Column) -> "MergeIntoWriter": def pandas_api( self, index_col: Optional[Union[str, List[str]]] = None ) -> "PandasOnSparkDataFrame": - from pyspark.pandas.namespace import _get_index_map from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.pandas.internal import InternalFrame + from pyspark.pandas.namespace import _get_index_map index_spark_columns, index_names = _get_index_map(self, index_col) internal = InternalFrame( @@ -2123,8 +2130,9 @@ def sampleBy( def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.dataframe + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow # It inherits docstrings but doctests cannot detect them so we run @@ -2141,6 +2149,7 @@ def _test() -> None: del pyspark.sql.dataframe.DataFrame.mapInArrow.__doc__ else: import pyarrow as pa + from pyspark.loose_version import LooseVersion if LooseVersion(pa.__version__) < LooseVersion("21.0.0"): diff --git a/python/pyspark/sql/classic/table_arg.py b/python/pyspark/sql/classic/table_arg.py index 97f6fa79e1894..8f45a90cc851e 100644 --- a/python/pyspark/sql/classic/table_arg.py +++ b/python/pyspark/sql/classic/table_arg.py @@ -23,6 +23,7 @@ if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql._typing import ColumnOrName diff --git a/python/pyspark/sql/classic/window.py b/python/pyspark/sql/classic/window.py index 0a9e0b3c96c64..eda4a7de64090 100644 --- a/python/pyspark/sql/classic/window.py +++ b/python/pyspark/sql/classic/window.py @@ -15,16 +15,19 @@ # limitations under the License. # import sys -from typing import cast, Iterable, Sequence, Tuple, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Iterable, Sequence, Tuple, Union, cast +from pyspark.sql.utils import get_active_spark_context from pyspark.sql.window import ( Window as ParentWindow, +) +from pyspark.sql.window import ( WindowSpec as ParentWindowSpec, ) -from pyspark.sql.utils import get_active_spark_context if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql._typing import ColumnOrName @@ -34,7 +37,7 @@ def _to_java_cols( cols: Tuple[Union["ColumnOrName", Sequence["ColumnOrName"]], ...], ) -> "JavaObject": - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq if len(cols) == 1 and isinstance(cols[0], list): cols = cols[0] # type: ignore[assignment] @@ -125,8 +128,9 @@ def rangeBetween(self, start: int, end: int) -> ParentWindowSpec: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.window + from pyspark.sql import SparkSession # It inherits docstrings but doctests cannot detect them so we run # the parent classe's doctests here directly. diff --git a/python/pyspark/sql/column.py b/python/pyspark/sql/column.py index 7c39a0ffbc333..f5e1cc9510cf4 100644 --- a/python/pyspark/sql/column.py +++ b/python/pyspark/sql/column.py @@ -19,20 +19,20 @@ import sys from typing import ( - overload, + TYPE_CHECKING, Any, Callable, - TYPE_CHECKING, Union, + overload, ) +from pyspark.errors import PySparkValueError from pyspark.sql.tvf_argument import TableValuedFunctionArgument -from pyspark.sql.utils import dispatch_col_method from pyspark.sql.types import DataType -from pyspark.errors import PySparkValueError +from pyspark.sql.utils import dispatch_col_method if TYPE_CHECKING: - from pyspark.sql._typing import LiteralType, DecimalLiteral, DateTimeLiteral + from pyspark.sql._typing import DateTimeLiteral, DecimalLiteral, LiteralType from pyspark.sql.window import WindowSpec __all__ = ["Column"] @@ -1615,8 +1615,9 @@ def __repr__(self) -> str: ... def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.column + from pyspark.sql import SparkSession globs = pyspark.sql.column.__dict__.copy() spark = SparkSession.builder.master("local[4]").appName("sql.column tests").getOrCreate() diff --git a/python/pyspark/sql/conf.py b/python/pyspark/sql/conf.py index d4991cd2a4125..dabdbb4b3e022 100644 --- a/python/pyspark/sql/conf.py +++ b/python/pyspark/sql/conf.py @@ -16,7 +16,7 @@ # import sys -from typing import Any, Dict, Optional, Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, Optional, Union from pyspark import _NoValue from pyspark._globals import _NoValueType @@ -182,10 +182,11 @@ def isModifiable(self, key: str) -> bool: def _test() -> None: - import os import doctest - from pyspark.sql.session import SparkSession + import os + import pyspark.sql.conf + from pyspark.sql.session import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/connect/_typing.py b/python/pyspark/sql/connect/_typing.py index 50bb02be09484..707059b225016 100644 --- a/python/pyspark/sql/connect/_typing.py +++ b/python/pyspark/sql/connect/_typing.py @@ -14,10 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from types import FunctionType -from typing import Any, Callable, Iterable, Union, Optional, NewType, Protocol, Tuple import datetime import decimal +from types import FunctionType +from typing import Any, Callable, Iterable, Iterator, NewType, Optional, Protocol, Tuple, Union import pyarrow from pandas.core.frame import DataFrame as PandasDataFrame @@ -44,9 +44,9 @@ DataFrameLike = PandasDataFrame -PandasMapIterFunction = Callable[[Iterable[DataFrameLike]], Iterable[DataFrameLike]] +PandasMapIterFunction = Callable[[Iterator[DataFrameLike]], Iterator[DataFrameLike]] -ArrowMapIterFunction = Callable[[Iterable[pyarrow.RecordBatch]], Iterable[pyarrow.RecordBatch]] +ArrowMapIterFunction = Callable[[Iterator[pyarrow.RecordBatch]], Iterator[pyarrow.RecordBatch]] PandasGroupedMapFunction = Union[ Callable[[DataFrameLike], DataFrameLike], diff --git a/python/pyspark/sql/connect/avro/functions.py b/python/pyspark/sql/connect/avro/functions.py index d1edfc813a80f..69b91d261926b 100644 --- a/python/pyspark/sql/connect/avro/functions.py +++ b/python/pyspark/sql/connect/avro/functions.py @@ -19,13 +19,12 @@ A collections of builtin avro functions """ -from pyspark.errors import PySparkTypeError - -from typing import Dict, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Optional +from pyspark.errors import PySparkTypeError from pyspark.sql.avro import functions as PyAvroFunctions from pyspark.sql.column import Column -from pyspark.sql.connect.functions.builtin import _invoke_function, _to_col, _options_to_col, lit +from pyspark.sql.connect.functions.builtin import _invoke_function, _options_to_col, _to_col, lit if TYPE_CHECKING: from pyspark.sql.connect._typing import ColumnOrName @@ -91,6 +90,7 @@ def to_avro(data: "ColumnOrName", jsonFormatSchema: str = "") -> Column: def _test() -> None: import os import sys + from pyspark.testing.sqlutils import search_jar avro_jar = search_jar("connector/avro", "spark-avro", "spark-avro") @@ -108,8 +108,9 @@ def _test() -> None: os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.avro.functions + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.avro.functions.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/catalog.py b/python/pyspark/sql/connect/catalog.py index 51448c813e8d5..9859fef90caaf 100644 --- a/python/pyspark/sql/connect/catalog.py +++ b/python/pyspark/sql/connect/catalog.py @@ -14,30 +14,31 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.errors import PySparkTypeError - -from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING - import warnings +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + import pyarrow as pa -from pyspark.storagelevel import StorageLevel -from pyspark.sql.types import StructType -from pyspark.sql.connect.dataframe import DataFrame +from pyspark.errors import PySparkTypeError from pyspark.sql.catalog import ( Catalog as PySparkCatalog, +) +from pyspark.sql.catalog import ( CatalogMetadata, + Column, Database, + Function, Table, TablePartition, - Function, - Column, ) from pyspark.sql.connect import plan +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.types import StructType +from pyspark.storagelevel import StorageLevel if TYPE_CHECKING: - from pyspark.sql.connect.session import SparkSession from pyspark.sql.connect._typing import DataTypeOrString, UserDefinedFunctionLike + from pyspark.sql.connect.session import SparkSession class Catalog: @@ -410,11 +411,12 @@ def registerFunction( def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.catalog + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.catalog.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/client/artifact.py b/python/pyspark/sql/connect/client/artifact.py index 94879171b41f0..3ab3df4ff5938 100644 --- a/python/pyspark/sql/connect/client/artifact.py +++ b/python/pyspark/sql/connect/client/artifact.py @@ -14,27 +14,26 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.errors import PySparkRuntimeError, PySparkValueError -from pyspark.sql.connect.logging import logger - +import abc import hashlib import importlib import io -import sys import os +import sys import zlib +from functools import cached_property from itertools import chain -from typing import List, Iterable, BinaryIO, Iterator, Optional, Tuple -import abc from pathlib import Path, PureWindowsPath +from typing import BinaryIO, Iterable, Iterator, List, Optional, Tuple from urllib.parse import urlparse from urllib.request import url2pathname -from functools import cached_property import grpc import pyspark.sql.connect.proto as proto import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib +from pyspark.errors import PySparkRuntimeError, PySparkValueError +from pyspark.sql.connect.logging import logger JAR_PREFIX: str = "jars" PYFILE_PREFIX: str = "pyfiles" diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py index 43a22c4998f2b..7685571c7253a 100644 --- a/python/pyspark/sql/connect/client/core.py +++ b/python/pyspark/sql/connect/client/core.py @@ -22,114 +22,112 @@ ] import atexit -from dataclasses import dataclass, fields - -import pyspark -from pyspark.sql.connect.proto.base_pb2 import FetchErrorDetailsResponse - import concurrent.futures +import copy import logging -import threading import os -import copy import platform -import urllib.parse -import uuid import sys +import threading import time import traceback +import urllib.parse +import uuid import weakref +from dataclasses import dataclass, fields from typing import ( + TYPE_CHECKING, + Any, + Dict, Iterable, Iterator, - Optional, - Any, - Union, List, - Tuple, - Dict, - Set, - NoReturn, Mapping, - cast, - TYPE_CHECKING, + NoReturn, + Optional, + Set, + Tuple, Type, + Union, + cast, ) -import pandas as pd -import pyarrow as pa - import google.protobuf.message -from grpc_status import rpc_status import grpc -from google.protobuf import text_format, any_pb2 +import pandas as pd +import pyarrow as pa +from google.protobuf import any_pb2, text_format from google.rpc import error_details_pb2 +from grpc_status import rpc_status -from pyspark.util import is_remote_only, disable_gc +import pyspark +import pyspark.sql.connect.proto as pb2 +import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib +import pyspark.sql.connect.types as types from pyspark.accumulators import SpecialAccumulatorIds, pickleSer -from pyspark.version import __version__ -from pyspark.traceback_utils import CallSite +from pyspark.errors import ( + PySparkAssertionError, + PySparkNotImplementedError, + PySparkValueError, +) +from pyspark.errors.exceptions.connect import ( + SparkConnectException, + SparkConnectGrpcException, + convert_exception, + convert_observation_errors, +) from pyspark.resource.information import ResourceInformation -from pyspark.sql.metrics import MetricValue, PlanMetrics, ExecutionInfo, ObservedMetrics from pyspark.sql.connect.client.artifact import ArtifactManager -from pyspark.sql.connect.logging import logger -from pyspark.sql.connect.profiler import ConnectProfilerCollector from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator from pyspark.sql.connect.client.retries import ( - RetryPolicy, - Retrying, - DefaultPolicy, DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME, + DefaultPolicy, + Retrying, + RetryPolicy, ) from pyspark.sql.connect.conversion import ( - storage_level_to_proto, - proto_to_storage_level, proto_to_remote_cached_dataframe, -) -import pyspark.sql.connect.proto as pb2 -import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib -import pyspark.sql.connect.types as types -from pyspark.errors.exceptions.connect import ( - convert_exception, - convert_observation_errors, - SparkConnectException, - SparkConnectGrpcException, + proto_to_storage_level, + storage_level_to_proto, ) from pyspark.sql.connect.expressions import ( - LiteralExpression, - PythonUDF, CommonInlineUserDefinedFunction, JavaUDF, + LiteralExpression, + PythonUDF, ) +from pyspark.sql.connect.logging import logger +from pyspark.sql.connect.observation import Observation from pyspark.sql.connect.plan import ( - CommonInlineUserDefinedTableFunction, CommonInlineUserDefinedDataSource, - PythonUDTF, + CommonInlineUserDefinedTableFunction, PythonDataSource, + PythonUDTF, ) -from pyspark.sql.connect.observation import Observation +from pyspark.sql.connect.profiler import ConnectProfilerCollector +from pyspark.sql.connect.proto.base_pb2 import FetchErrorDetailsResponse +from pyspark.sql.connect.shell.progress import Progress, ProgressHandler, from_proto from pyspark.sql.connect.utils import get_python_ver -from pyspark.sql.pandas.types import from_arrow_schema +from pyspark.sql.metrics import ExecutionInfo, MetricValue, ObservedMetrics, PlanMetrics from pyspark.sql.pandas.conversion import _convert_arrow_table_to_pandas +from pyspark.sql.pandas.types import from_arrow_schema from pyspark.sql.types import DataType, StructType -from pyspark.util import PythonEvalType from pyspark.storagelevel import StorageLevel -from pyspark.errors import ( - PySparkAssertionError, - PySparkNotImplementedError, - PySparkValueError, -) -from pyspark.sql.connect.shell.progress import Progress, ProgressHandler, from_proto +from pyspark.traceback_utils import CallSite +from pyspark.util import PythonEvalType, disable_gc, is_remote_only +from pyspark.version import __version__ if TYPE_CHECKING: from google.rpc.error_details_pb2 import ErrorInfo from google.rpc.status_pb2 import Status + from pyspark.sql.connect._typing import DataTypeOrString from pyspark.sql.connect.session import SparkSession from pyspark.sql.datasource import DataSource PYSPARK_ROOT = os.path.dirname(pyspark.__file__) +_OPERATION_ID_METADATA_KEY = "spark-connect-operation-id" @dataclass(frozen=True) @@ -145,6 +143,14 @@ class RpcDeadlines: fires, the server-side operation continues running; the client opens a new ReattachExecute stream to resume receiving results. Non-reattachable ExecutePlan has no deadline because a timeout there would kill the execution with no recovery path. + + Note on ``release_relation``: the RemoveRemoteCachedRelation cleanup command is sent over a + blocking, non-reattachable ExecutePlan call issued from + :meth:`CachedRemoteRelation.__del__`. Unlike a query ExecutePlan, a timeout here does not kill + any recoverable execution -- it only abandons a best-effort cache eviction that the server also + performs independently -- so this call is given a bounded deadline. Without it the finalizer + can block forever if the release response is never delivered, which (on the foreachBatch + Connect path) stalls the streaming query indefinitely. """ reattachable_execute_plan: Optional[float] = 10 * 60 # 10 min @@ -154,6 +160,7 @@ class RpcDeadlines: config: Optional[float] = 10 * 60 # 10 min interrupt: Optional[float] = 10 * 60 # 10 min release_session: Optional[float] = 10 * 60 # 10 min + release_relation: Optional[float] = 60 # 1 min; short: per-batch finalizer artifact_status: Optional[float] = 10 * 60 # 10 min clone_session: Optional[float] = 10 * 60 # 10 min get_status: Optional[float] = 10 * 60 # 10 min @@ -184,6 +191,7 @@ def disabled(cls) -> "RpcDeadlines": config=None, interrupt=None, release_session=None, + release_relation=None, artifact_status=None, clone_session=None, get_status=None, @@ -887,6 +895,16 @@ def __init__( if isinstance(connection, ChannelBuilder) else DefaultChannelBuilder(connection, channel_options) ) + metadata = list(self._builder.metadata()) + if any(key.lower() == _OPERATION_ID_METADATA_KEY for key, _ in metadata): + logger.warning( + "Connection option %s is ignored because Spark Connect sets it for each " + "ExecutePlan request.", + _OPERATION_ID_METADATA_KEY, + ) + artifact_manager_metadata = [ + (key, value) for key, value in metadata if key.lower() != _OPERATION_ID_METADATA_KEY + ] self._user_id = None self._retry_policies: List[RetryPolicy] = [] @@ -927,7 +945,7 @@ def __init__( self._user_id, self._session_id, self._channel, - self._builder.metadata(), + artifact_manager_metadata, add_artifacts_timeout=self._rpc_deadlines.add_artifacts, artifact_status_timeout=self._rpc_deadlines.artifact_status, ) @@ -1081,6 +1099,7 @@ def register_udf( name: Optional[str] = None, eval_type: int = PythonEvalType.SQL_BATCHED_UDF, deterministic: bool = True, + buffer_type: Optional["DataType"] = None, ) -> str: """ Create a temporary UDF in the session catalog on the other side. We generate a @@ -1096,6 +1115,8 @@ def register_udf( eval_type=eval_type, func=function, python_ver="%d.%d" % sys.version_info[:2], + # Set for the incremental aggregator (see pyspark.sql.aggregator). + buffer_type=buffer_type, ) # construct a CommonInlineUserDefinedFunction @@ -1239,7 +1260,7 @@ def to_table( table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(req, observations) # Create a query execution object. - ei = ExecutionInfo(metrics, observed_metrics) + ei = ExecutionInfo(metrics, observed_metrics, req.operation_id) assert table is not None return table, schema, ei @@ -1275,7 +1296,7 @@ def to_pandas( req, observations, selfDestruct == "true" ) assert table is not None - ei = ExecutionInfo(metrics, observed_metrics) + ei = ExecutionInfo(metrics, observed_metrics, req.operation_id) schema = schema or from_arrow_schema(table.schema, prefer_timestamp_ntz=True) assert schema is not None and isinstance(schema, StructType) @@ -1418,7 +1439,7 @@ def execute_command( req, observations or {} ) # Create a query execution object. - ei = ExecutionInfo(metrics, observed_metrics) + ei = ExecutionInfo(metrics, observed_metrics, req.operation_id) if data is not None: return (data.to_pandas(), properties, ei) else: @@ -1535,7 +1556,9 @@ def _execute_plan_request_with_metadata( ) ) ) - if operation_id is not None: + if operation_id is None: + operation_id = str(uuid.uuid4()) + else: try: uuid.UUID(operation_id, version=4) except ValueError as ve: @@ -1543,7 +1566,7 @@ def _execute_plan_request_with_metadata( errorClass="INVALID_OPERATION_UUID_ID", messageParameters={"arg_name": "operation_id", "origin": str(ve)}, ) - req.operation_id = operation_id + req.operation_id = operation_id self._update_request_with_user_context_extensions(req) if call_stack_trace := self.__class__._build_call_stack_trace(): @@ -1652,7 +1675,7 @@ def _analyze(self, method: str, **kwargs: Any) -> AnalyzeResult: with attempt: resp = self._stub.AnalyzePlan( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.analyze_plan, ) self._verify_response_integrity(resp) @@ -1673,8 +1696,10 @@ def _execute(self, req: pb2.ExecutePlanRequest) -> None: """ logger.debug("Execute") + operation_id = req.operation_id for hook in self._session_hooks: req = hook.on_execute_plan(req) + req.operation_id = operation_id def handle_response(b: pb2.ExecutePlanResponse) -> None: self._verify_response_integrity(b) @@ -1686,7 +1711,7 @@ def handle_response(b: pb2.ExecutePlanResponse) -> None: req, self._stub, self._retrying, - self._builder.metadata(), + self._execute_plan_metadata(req.operation_id), reattachable_execute_plan_timeout=self._rpc_deadlines.reattachable_execute_plan, reattach_execute_timeout=self._rpc_deadlines.reattach_execute, ) @@ -1700,10 +1725,24 @@ def handle_response(b: pb2.ExecutePlanResponse) -> None: for attempt in self._retrying(): with attempt: with disable_gc(): - for b in self._stub.ExecutePlan(req, metadata=self._builder.metadata()): + for b in self._stub.ExecutePlan( + req, metadata=self._execute_plan_metadata(req.operation_id) + ): handle_response(b) except Exception as error: - self._handle_error(error) + self._handle_error(error, req.operation_id) + + def _builder_metadata(self) -> List[Tuple[str, str]]: + return [ + (key, value) + for key, value in self._builder.metadata() + if key.lower() != _OPERATION_ID_METADATA_KEY + ] + + def _execute_plan_metadata(self, operation_id: str) -> List[Tuple[str, str]]: + metadata = self._builder_metadata() + metadata.append((_OPERATION_ID_METADATA_KEY, operation_id)) + return metadata def _execute_and_fetch_as_iterator( self, @@ -1724,9 +1763,10 @@ def _execute_and_fetch_as_iterator( # when not at debug log level. logger.debug(f"ExecuteAndFetchAsIterator. Request: {self._proto_to_string(req)}") + operation_id = req.operation_id for hook in self._session_hooks: req = hook.on_execute_plan(req) - + req.operation_id = operation_id num_records = 0 arrow_batch_chunks_to_assemble: List[bytes] = [] @@ -1901,7 +1941,7 @@ def handle_response( req, self._stub, self._retrying, - self._builder.metadata(), + self._execute_plan_metadata(req.operation_id), reattachable_execute_plan_timeout=self._rpc_deadlines.reattachable_execute_plan, reattach_execute_timeout=self._rpc_deadlines.reattach_execute, ) @@ -1916,7 +1956,9 @@ def handle_response( with attempt: with disable_gc(): it = iter( - self._stub.ExecutePlan(req, metadata=self._builder.metadata()) + self._stub.ExecutePlan( + req, metadata=self._execute_plan_metadata(req.operation_id) + ) ) while True: try: @@ -1932,7 +1974,7 @@ def handle_response( self.interrupt_operation(req.operation_id) raise kb except Exception as error: - self._handle_error(error) + self._handle_error(error, req.operation_id) def _execute_and_fetch( self, @@ -2058,7 +2100,7 @@ def config(self, operation: pb2.ConfigRequest.Operation) -> ConfigResult: with disable_gc(): resp = self._stub.Config( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.config, ) self._verify_response_integrity(resp) @@ -2104,7 +2146,7 @@ def interrupt_all(self) -> Optional[List[str]]: with attempt: resp = self._stub.Interrupt( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.interrupt, ) self._verify_response_integrity(resp) @@ -2120,7 +2162,7 @@ def interrupt_tag(self, tag: str) -> Optional[List[str]]: with attempt: resp = self._stub.Interrupt( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.interrupt, ) self._verify_response_integrity(resp) @@ -2136,7 +2178,7 @@ def interrupt_operation(self, op_id: str) -> Optional[List[str]]: with attempt: resp = self._stub.Interrupt( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.interrupt, ) self._verify_response_integrity(resp) @@ -2156,7 +2198,7 @@ def release_session(self) -> None: with attempt: resp = self._stub.ReleaseSession( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.release_session, ) self._verify_response_integrity(resp) @@ -2212,7 +2254,7 @@ def _get_operation_statuses( with attempt: resp = self._stub.GetStatus( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.get_status, ) self._verify_response_integrity(resp) @@ -2231,7 +2273,9 @@ def remove_tag(self, tag: str) -> None: self._throw_if_invalid_tag(tag) if not hasattr(self.thread_local, "tags"): self.thread_local.tags = set() - self.thread_local.tags.remove(tag) + # Use discard, not remove: removing an absent tag is a documented no-op + # (see SparkSession.removeTag), matching the Classic behavior. + self.thread_local.tags.discard(tag) def get_tags(self) -> Set[str]: if not hasattr(self.thread_local, "tags"): @@ -2249,7 +2293,7 @@ def _throw_if_invalid_tag(self, tag: str) -> None: spark_job_tags_sep = "," if tag is None: raise PySparkValueError( - errorClass="CANNOT_BE_NONE", message_paramters={"arg_name": "Spark Connect tag"} + errorClass="CANNOT_BE_NONE", messageParameters={"arg_name": "Spark Connect tag"} ) if spark_job_tags_sep in tag: raise PySparkValueError( @@ -2297,7 +2341,7 @@ def clear_user_context_extensions(self) -> None: with self.global_user_context_extensions_lock: self.global_user_context_extensions = list() - def _handle_error(self, error: Exception) -> NoReturn: + def _handle_error(self, error: Exception, operation_id: Optional[str] = None) -> NoReturn: """ Handle errors that occur during RPC calls. @@ -2318,9 +2362,14 @@ def _handle_error(self, error: Exception) -> NoReturn: try: self.thread_local.inside_error_handling = True - if isinstance(error, grpc.RpcError): - self._handle_rpc_error(error) - raise error + try: + if isinstance(error, grpc.RpcError): + self._handle_rpc_error(error) + raise error + except BaseException as handled_error: + if operation_id: + handled_error._operation_id = operation_id # type: ignore[attr-defined] + raise finally: self.thread_local.inside_error_handling = False @@ -2341,7 +2390,7 @@ def _fetch_enriched_error(self, info: "ErrorInfo") -> Optional[pb2.FetchErrorDet try: return self._stub.FetchErrorDetails( req, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.fetch_error_details, ) except grpc.RpcError: @@ -2760,7 +2809,7 @@ def clone(self, new_session_id: Optional[str] = None) -> "SparkConnectClient": with attempt: response: pb2.CloneSessionResponse = self._stub.CloneSession( request, - metadata=self._builder.metadata(), + metadata=self._builder_metadata(), timeout=self._rpc_deadlines.clone_session, ) diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py index f1d06320866e0..3127278883527 100644 --- a/python/pyspark/sql/connect/client/reattach.py +++ b/python/pyspark/sql/connect/client/reattach.py @@ -14,23 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.sql.connect.client.retries import Retrying, RetryException - -from threading import RLock +import os import uuid +import weakref from collections.abc import Generator -from typing import Optional, Any, Iterator, Iterable, Tuple, Callable, cast, ClassVar from concurrent.futures import Future, ThreadPoolExecutor -import os -import weakref +from threading import RLock +from typing import Any, Callable, ClassVar, Iterable, Iterator, Optional, Tuple, cast import grpc from grpc_status import rpc_status -from pyspark.sql.connect.logging import logger import pyspark.sql.connect.proto as pb2 import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib from pyspark.errors import PySparkRuntimeError +from pyspark.sql.connect.client.retries import RetryException, Retrying +from pyspark.sql.connect.logging import logger from pyspark.util import disable_gc diff --git a/python/pyspark/sql/connect/client/retries.py b/python/pyspark/sql/connect/client/retries.py index f16968f6391fa..362774e10c855 100644 --- a/python/pyspark/sql/connect/client/retries.py +++ b/python/pyspark/sql/connect/client/retries.py @@ -15,17 +15,19 @@ # limitations under the License. # -import grpc import random import time import typing import warnings +from types import TracebackType +from typing import Callable, Generator, List, Optional, Type, cast + +import grpc from google.rpc import error_details_pb2 from grpc_status import rpc_status -from typing import Optional, Callable, Generator, List, Type, cast -from types import TracebackType -from pyspark.sql.connect.logging import logger + from pyspark.errors import PySparkRuntimeError +from pyspark.sql.connect.logging import logger """ This module contains retry system. The system is designed to be diff --git a/python/pyspark/sql/connect/column.py b/python/pyspark/sql/connect/column.py index 0eece36c950d0..dda3158c57b46 100644 --- a/python/pyspark/sql/connect/column.py +++ b/python/pyspark/sql/connect/column.py @@ -17,46 +17,44 @@ import datetime import decimal import warnings - from typing import ( TYPE_CHECKING, Any, Callable, - Union, Optional, Tuple, + Union, ) -from pyspark.sql.column import Column as ParentColumn +import pyspark.sql.connect.proto as proto from pyspark.errors import ( - PySparkTypeError, PySparkAttributeError, + PySparkTypeError, PySparkValueError, ) -from pyspark.sql.types import DataType -from pyspark.sql.utils import enum_to_value - -import pyspark.sql.connect.proto as proto +from pyspark.errors.utils import with_origin_to_class +from pyspark.sql.column import Column as ParentColumn from pyspark.sql.connect.expressions import ( + CaseWhen, + CastExpression, + DropField, Expression, - UnresolvedFunction, - UnresolvedExtractValue, LiteralExpression, - CaseWhen, SortOrder, SubqueryExpression, - CastExpression, + UnresolvedExtractValue, + UnresolvedFunction, WindowExpression, WithField, - DropField, ) -from pyspark.errors.utils import with_origin_to_class +from pyspark.sql.types import DataType +from pyspark.sql.utils import enum_to_value if TYPE_CHECKING: from pyspark.sql.connect._typing import ( - LiteralType, DateTimeLiteral, DecimalLiteral, + LiteralType, ) from pyspark.sql.connect.client import SparkConnectClient from pyspark.sql.connect.window import WindowSpec @@ -622,20 +620,25 @@ def __iter__(self) -> None: ) def __nonzero__(self) -> None: + try: + column_repr = repr(self._expr) + except Exception: + column_repr = "<unknown>" raise PySparkValueError( errorClass="CANNOT_CONVERT_COLUMN_INTO_BOOL", - messageParameters={}, + messageParameters={"column": column_repr}, ) __bool__ = __nonzero__ def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.column + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.column.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/conf.py b/python/pyspark/sql/connect/conf.py index 433e3e5c10b7c..8507bdf7a7252 100644 --- a/python/pyspark/sql/connect/conf.py +++ b/python/pyspark/sql/connect/conf.py @@ -14,13 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.errors import PySparkValueError, PySparkTypeError - -from typing import Any, Dict, Optional, Union, cast import warnings +from typing import Any, Dict, Optional, Union, cast from pyspark import _NoValue from pyspark._globals import _NoValueType +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.conf import RuntimeConfig as PySparkRuntimeConfig from pyspark.sql.connect import proto from pyspark.sql.connect.client import SparkConnectClient @@ -134,11 +133,12 @@ def _checkType(self, obj: Any, identifier: str) -> None: def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.conf + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.conf.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/context.py b/python/pyspark/sql/connect/context.py index eacb22a25cbd3..7f11ae1cc05ef 100644 --- a/python/pyspark/sql/connect/context.py +++ b/python/pyspark/sql/connect/context.py @@ -17,34 +17,35 @@ import warnings from typing import ( - Optional, - Union, - Callable, + TYPE_CHECKING, Any, + Callable, + ClassVar, Iterable, List, + Optional, Tuple, - ClassVar, - TYPE_CHECKING, + Union, ) from pyspark import _NoValue from pyspark._globals import _NoValueType from pyspark.errors import PySparkNotImplementedError -from pyspark.sql.dataframe import DataFrame from pyspark.sql.connect.readwriter import DataFrameReader -from pyspark.sql.connect.streaming.readwriter import DataStreamReader from pyspark.sql.connect.streaming.query import StreamingQueryManager +from pyspark.sql.connect.streaming.readwriter import DataStreamReader +from pyspark.sql.dataframe import DataFrame from pyspark.sql.types import AtomicType, BooleanType, DataType, StringType, StructField, StructType if TYPE_CHECKING: import numpy as np import pandas as pd import pyarrow as pa + + from pyspark.sql._typing import UserDefinedFunctionLike from pyspark.sql.connect.session import SparkSession from pyspark.sql.connect.udf import UDFRegistration from pyspark.sql.connect.udtf import UDTFRegistration - from pyspark.sql._typing import UserDefinedFunctionLike # Internal module - not part of the public PySpark API surface. # The public SQLContext/HiveContext are in pyspark.sql.context; this module diff --git a/python/pyspark/sql/connect/conversion.py b/python/pyspark/sql/connect/conversion.py index 137b04510f0f1..220ac36a8a001 100644 --- a/python/pyspark/sql/connect/conversion.py +++ b/python/pyspark/sql/connect/conversion.py @@ -48,9 +48,9 @@ def proto_to_storage_level(storage_level: pb2.StorageLevel) -> StorageLevel: def proto_to_remote_cached_dataframe(relation: pb2.CachedRemoteRelation) -> "DataFrame": assert relation is not None and isinstance(relation, pb2.CachedRemoteRelation) + import pyspark.sql.connect.plan as plan from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect.session import SparkSession - import pyspark.sql.connect.plan as plan session = SparkSession.active() return DataFrame( diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 093489757115b..95d6cddd50ba6 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -16,68 +16,51 @@ # # mypy: disable-error-code="override" -from pyspark.errors.exceptions.base import ( - SessionNotSameException, - PySparkIndexError, -) -from pyspark.resource import ResourceProfile -from pyspark.sql.connect.logging import logger - +import copy +import functools +import json +import os +import random +import sys +import warnings +from collections.abc import Iterable from typing import ( + TYPE_CHECKING, Any, + Callable, Dict, Iterator, List, NoReturn, Optional, + Sequence, Tuple, + Type, Union, - Sequence, - TYPE_CHECKING, - overload, - Callable, cast, - Type, + overload, ) -import copy -import os -import sys -import random import pyarrow as pa -import json -import warnings -from collections.abc import Iterable -import functools +import pyspark.sql.connect.plan as plan from pyspark import _NoValue from pyspark._globals import _NoValueType -from pyspark.util import is_remote_only -from pyspark.sql.types import Row, StructType, _create_row -from pyspark.sql.dataframe import ( - DataFrame as ParentDataFrame, - DataFrameNaFunctions as ParentDataFrameNaFunctions, - DataFrameStatFunctions as ParentDataFrameStatFunctions, -) - from pyspark.errors import ( - PySparkTypeError, PySparkAttributeError, - PySparkValueError, PySparkNotImplementedError, PySparkRuntimeError, + PySparkTypeError, + PySparkValueError, +) +from pyspark.errors.exceptions.base import ( + PySparkIndexError, + SessionNotSameException, ) -from pyspark.util import PythonEvalType +from pyspark.resource import ResourceProfile from pyspark.serializers import CPickleSerializer -from pyspark.storagelevel import StorageLevel -import pyspark.sql.connect.plan as plan -from pyspark.sql.conversion import ArrowTableToRowsConversion -from pyspark.sql.connect.group import GroupedData -from pyspark.sql.connect.merge import MergeIntoWriter -from pyspark.sql.connect.readwriter import DataFrameWriter, DataFrameWriterV2 -from pyspark.sql.connect.streaming.readwriter import DataStreamWriter -from pyspark.sql.connect.column import Column as ConnectColumn from pyspark.sql.column import Column +from pyspark.sql.connect.column import Column as ConnectColumn from pyspark.sql.connect.expressions import ( ColumnReference, DirectShufflePartitionID, @@ -87,26 +70,44 @@ UnresolvedStar, ) from pyspark.sql.connect.functions import builtin as F -from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema +from pyspark.sql.connect.group import GroupedData +from pyspark.sql.connect.logging import logger +from pyspark.sql.connect.merge import MergeIntoWriter +from pyspark.sql.connect.readwriter import DataFrameWriter, DataFrameWriterV2 +from pyspark.sql.connect.streaming.readwriter import DataStreamWriter +from pyspark.sql.conversion import ArrowTableToRowsConversion +from pyspark.sql.dataframe import ( + DataFrame as ParentDataFrame, +) +from pyspark.sql.dataframe import ( + DataFrameNaFunctions as ParentDataFrameNaFunctions, +) +from pyspark.sql.dataframe import ( + DataFrameStatFunctions as ParentDataFrameStatFunctions, +) from pyspark.sql.pandas.functions import _validate_vectorized_udf # type: ignore[attr-defined] +from pyspark.sql.pandas.types import from_arrow_schema, to_arrow_schema from pyspark.sql.table_arg import TableArg +from pyspark.sql.types import Row, StructType, _create_row +from pyspark.storagelevel import StorageLevel +from pyspark.util import PythonEvalType, is_remote_only if TYPE_CHECKING: + from pyspark.core.rdd import RDD + from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.sql.connect._typing import ( + ArrowMapIterFunction, ColumnOrName, ColumnOrNameOrOrdinal, LiteralType, - PrimitiveType, OptionalPrimitiveType, PandasMapIterFunction, - ArrowMapIterFunction, + PrimitiveType, ) - from pyspark.core.rdd import RDD - from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike from pyspark.sql.connect.observation import Observation from pyspark.sql.connect.session import SparkSession - from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.sql.metrics import ExecutionInfo + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike from pyspark.sql.plot import PySparkPlotAccessor @@ -2243,9 +2244,9 @@ def toLocalIterator(self, prefetchPartitions: bool = False) -> Iterator[Row]: def pandas_api( self, index_col: Optional[Union[str, List[str]]] = None ) -> "PandasOnSparkDataFrame": - from pyspark.pandas.namespace import _get_index_map from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.pandas.internal import InternalFrame + from pyspark.pandas.namespace import _get_index_map index_spark_columns, index_names = _get_index_map(self, index_col) internal = InternalFrame( @@ -2332,7 +2333,7 @@ def foreachPartition(self, f: Callable[[Iterator[Row]], None]) -> None: for f in schema.fields ] - def foreach_partition_func(itr: Iterable[pa.RecordBatch]) -> Iterable[pa.RecordBatch]: + def foreach_partition_func(itr: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: def flatten() -> Iterator[Row]: for table in itr: columnar_data = [ @@ -2505,13 +2506,14 @@ def sampleBy( def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.util import is_remote_only - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.dataframe + from pyspark.sql import SparkSession as PySparkSession from pyspark.testing.utils import have_pandas, have_pyarrow + from pyspark.util import is_remote_only # It inherits docstrings but doctests cannot detect them so we run # the parent classe's doctests here directly. @@ -2533,6 +2535,7 @@ def _test() -> None: del pyspark.sql.dataframe.DataFrame.mapInArrow.__doc__ else: import pyarrow as pa + from pyspark.loose_version import LooseVersion if LooseVersion(pa.__version__) < LooseVersion("21.0.0"): diff --git a/python/pyspark/sql/connect/datasource.py b/python/pyspark/sql/connect/datasource.py index c9c3ffcc85fcf..36bfedcc4f012 100644 --- a/python/pyspark/sql/connect/datasource.py +++ b/python/pyspark/sql/connect/datasource.py @@ -14,13 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Type, TYPE_CHECKING +from typing import TYPE_CHECKING, Type from pyspark.sql.datasource import DataSourceRegistration as PySparkDataSourceRegistration if TYPE_CHECKING: - from pyspark.sql.datasource import DataSource from pyspark.sql.connect.session import SparkSession + from pyspark.sql.datasource import DataSource class DataSourceRegistration: diff --git a/python/pyspark/sql/connect/expressions.py b/python/pyspark/sql/connect/expressions.py index 57270398118f7..cc960454f0082 100644 --- a/python/pyspark/sql/connect/expressions.py +++ b/python/pyspark/sql/connect/expressions.py @@ -14,75 +14,73 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import datetime +import decimal +import json +import warnings +from threading import Lock from typing import ( - cast, TYPE_CHECKING, Any, Callable, - Union, + Optional, Sequence, Tuple, - Optional, + Union, + cast, ) -import json -import decimal -import datetime -import warnings -from threading import Lock - import numpy as np +import pyspark.sql.connect.proto as proto +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.errors.utils import current_origin from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.connect.types import ( + UnparsedDataType, + proto_schema_to_pyspark_data_type, + pyspark_types_to_proto_types, +) from pyspark.sql.types import ( - _create_row, - _from_numpy_type, - DateType, ArrayType, - NullType, - BooleanType, BinaryType, + BooleanType, ByteType, - ShortType, - IntegerType, - LongType, - FloatType, - DoubleType, - DecimalType, - StringType, DataType, - TimeType, - TimestampType, - TimestampNTZType, + DateType, DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, MapType, + NullType, + ShortType, + StringType, StructType, + TimestampNTZType, + TimestampType, + TimeType, + _create_row, + _from_numpy_type, ) - -import pyspark.sql.connect.proto as proto +from pyspark.sql.utils import enum_to_value, is_timestamp_ntz_preferred from pyspark.util import ( - JVM_BYTE_MIN, JVM_BYTE_MAX, - JVM_SHORT_MIN, - JVM_SHORT_MAX, - JVM_INT_MIN, + JVM_BYTE_MIN, JVM_INT_MAX, - JVM_LONG_MIN, + JVM_INT_MIN, JVM_LONG_MAX, + JVM_LONG_MIN, + JVM_SHORT_MAX, + JVM_SHORT_MIN, ) -from pyspark.sql.connect.types import ( - UnparsedDataType, - pyspark_types_to_proto_types, - proto_schema_to_pyspark_data_type, -) -from pyspark.errors import PySparkTypeError, PySparkValueError -from pyspark.errors.utils import current_origin -from pyspark.sql.utils import is_timestamp_ntz_preferred, enum_to_value if TYPE_CHECKING: from pyspark.sql.connect.client import SparkConnectClient - from pyspark.sql.connect.window import WindowSpec from pyspark.sql.connect.plan import LogicalPlan + from pyspark.sql.connect.window import WindowSpec class Expression: @@ -741,6 +739,7 @@ def __init__( eval_type: int, func: Callable[..., Any], python_ver: str, + buffer_type: Optional[DataType] = None, ) -> None: self._output_type: DataType = ( UnparsedDataType(output_type) if isinstance(output_type, str) else output_type @@ -748,6 +747,8 @@ def __init__( self._eval_type = eval_type self._func = func self._python_ver = python_ver + # Intermediate buffer schema for an incremental Python aggregator; None otherwise. + self._buffer_type = buffer_type def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF: if isinstance(self._output_type, UnparsedDataType): @@ -763,6 +764,8 @@ def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF: expr.eval_type = self._eval_type expr.command = CloudPickleSerializer().dumps((self._func, output_type)) expr.python_ver = self._python_ver + if self._buffer_type is not None: + expr.buffer_type.CopyFrom(pyspark_types_to_proto_types(self._buffer_type)) return expr def __repr__(self) -> str: diff --git a/python/pyspark/sql/connect/functions/__init__.py b/python/pyspark/sql/connect/functions/__init__.py index 7110860521dfd..eb42de94c5c31 100644 --- a/python/pyspark/sql/connect/functions/__init__.py +++ b/python/pyspark/sql/connect/functions/__init__.py @@ -20,5 +20,5 @@ from pyspark.testing.utils import should_test_connect if should_test_connect: - from pyspark.sql.connect.functions.builtin import * # noqa: F403 from pyspark.sql.connect.functions import partitioning # noqa: F401 + from pyspark.sql.connect.functions.builtin import * # noqa: F403 diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py index f3e1a0c058f28..922c8f26e0aec 100644 --- a/python/pyspark/sql/connect/functions/builtin.py +++ b/python/pyspark/sql/connect/functions/builtin.py @@ -15,62 +15,71 @@ # limitations under the License. # import decimal +import functools import inspect +import random as py_random +import sys import warnings -import functools from typing import ( - Any, - Mapping, TYPE_CHECKING, - Union, - Sequence, + Any, + Callable, List, - overload, + Mapping, Optional, + Sequence, Tuple, Type, - Callable, + Union, ValuesView, cast, + overload, ) -import random as py_random -import sys import numpy as np from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.errors.utils import _with_origin from pyspark.sql import Column +from pyspark.sql import functions as pysparkfuncs from pyspark.sql.connect.expressions import ( + CallFunction, CaseWhen, - SortOrder, + ColumnReference, Expression, + LambdaFunction, LiteralExpression, - ColumnReference, - UnresolvedFunction, - UnresolvedStar, + SortOrder, SQLExpression, - LambdaFunction, + UnresolvedFunction, UnresolvedNamedLambdaVariable, - CallFunction, + UnresolvedStar, ) from pyspark.sql.connect.udf import _create_py_udf -from pyspark.sql.connect.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401 -from pyspark.sql.connect.udtf import _create_py_udtf, _create_pyarrow_udtf -from pyspark.sql import functions as pysparkfuncs +from pyspark.sql.connect.udtf import ( # noqa: F401 + AnalyzeArgument, + AnalyzeResult, + _create_py_udtf, + _create_pyarrow_udtf, +) from pyspark.sql.types import ( - _from_numpy_type, - DataType, - StructType, ArrayType, + DataType, MapType, StringType, + StructType, + _from_numpy_type, +) +from pyspark.sql.types import ( + UserDefinedType as _UserDefinedType, ) -from pyspark.sql.utils import enum_to_value as _enum_to_value -# The implementation of pandas_udf is embedded in pyspark.sql.function.pandas_udf -# for code reuse. -from pyspark.sql.functions import arrow_udf, pandas_udf # noqa: F401 +if TYPE_CHECKING: + from pyspark.sql.types import UserDefinedType +# The implementations of pandas_udf, arrow_udf and udaf are embedded in pyspark.sql.functions +# (they select the classic vs Connect UserDefinedFunction internally), so reuse them here. +from pyspark.sql.functions import arrow_udf, pandas_udf, udaf # noqa: F401 +from pyspark.sql.utils import enum_to_value as _enum_to_value if TYPE_CHECKING: from pyspark.sql.connect._typing import ( @@ -78,8 +87,8 @@ DataTypeOrString, UserDefinedFunctionLike, ) - from pyspark.sql.dataframe import DataFrame from pyspark.sql.connect.udtf import UserDefinedTableFunction + from pyspark.sql.dataframe import DataFrame def _to_col(col: "ColumnOrName") -> Column: @@ -847,6 +856,18 @@ def round(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Co round.__doc__ = pysparkfuncs.round.__doc__ +def truncate(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Column: + if scale is None: + return _invoke_function_over_columns("truncate", col) + else: + scale = _enum_to_value(scale) + scale = lit(scale) if isinstance(scale, int) else scale + return _invoke_function_over_columns("truncate", col, scale) + + +truncate.__doc__ = pysparkfuncs.truncate.__doc__ + + def sec(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("sec", col) @@ -1082,6 +1103,13 @@ def collect_set(col: "ColumnOrName") -> Column: collect_set.__doc__ = pysparkfuncs.collect_set.__doc__ +def collect_union(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("collect_union", col) + + +collect_union.__doc__ = pysparkfuncs.collect_union.__doc__ + + def listagg(col: "ColumnOrName", delimiter: Optional[Union[Column, str, bytes]] = None) -> Column: if delimiter is None: return _invoke_function_over_columns("listagg", col) @@ -2060,6 +2088,13 @@ def json_object_keys(col: "ColumnOrName") -> Column: json_object_keys.__doc__ = pysparkfuncs.json_object_keys.__doc__ +def json_typeof(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("json_typeof", col) + + +json_typeof.__doc__ = pysparkfuncs.json_typeof.__doc__ + + def inline(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("inline", col) @@ -2185,6 +2220,20 @@ def to_variant_object(col: "ColumnOrName") -> Column: to_variant_object.__doc__ = pysparkfuncs.to_variant_object.__doc__ +def variant_from_arrays(keys: "ColumnOrName", values: "ColumnOrName") -> Column: + return _invoke_function_over_columns("variant_from_arrays", keys, values) + + +variant_from_arrays.__doc__ = pysparkfuncs.variant_from_arrays.__doc__ + + +def variant_from_entries(entries: "ColumnOrName") -> Column: + return _invoke_function_over_columns("variant_from_entries", entries) + + +variant_from_entries.__doc__ = pysparkfuncs.variant_from_entries.__doc__ + + def parse_json(col: "ColumnOrName") -> Column: return _invoke_function("parse_json", _to_col(col)) @@ -2287,6 +2336,13 @@ def try_variant_array_append( try_variant_array_append.__doc__ = pysparkfuncs.try_variant_array_append.__doc__ +def variant_strip_nulls(v: "ColumnOrName", include_arrays: bool = True) -> Column: + return _invoke_function("variant_strip_nulls", _to_col(v), lit(include_arrays)) + + +variant_strip_nulls.__doc__ = pysparkfuncs.variant_strip_nulls.__doc__ + + def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str) -> Column: assert isinstance(path, (Column, str)) if isinstance(path, str): @@ -2473,6 +2529,28 @@ def slice( slice.__doc__ = pysparkfuncs.slice.__doc__ +def trim_array(x: "ColumnOrName", n: Union["ColumnOrName", int]) -> Column: + n = _enum_to_value(n) + if isinstance(n, (Column, str)): + _n = n + elif isinstance(n, int): + _n = lit(n) + else: + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "Column, int or str", + "arg_name": "n", + "arg_type": type(n).__name__, + }, + ) + + return _invoke_function_over_columns("trim_array", x, _n) + + +trim_array.__doc__ = pysparkfuncs.trim_array.__doc__ + + def sort_array(col: "ColumnOrName", asc: bool = True) -> Column: return _invoke_function("sort_array", _to_col(col), lit(asc)) @@ -2594,6 +2672,13 @@ def base64(col: "ColumnOrName") -> Column: base64.__doc__ = pysparkfuncs.base64.__doc__ +def to_base32(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("to_base32", col) + + +to_base32.__doc__ = pysparkfuncs.to_base32.__doc__ + + def unbase64(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("unbase64", col) @@ -2601,6 +2686,13 @@ def unbase64(col: "ColumnOrName") -> Column: unbase64.__doc__ = pysparkfuncs.unbase64.__doc__ +def from_base32(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("from_base32", col) + + +from_base32.__doc__ = pysparkfuncs.from_base32.__doc__ + + def ltrim(col: "ColumnOrName", trim: Optional["ColumnOrName"] = None) -> Column: if trim is not None: return _invoke_function_over_columns("ltrim", trim, col) @@ -2680,6 +2772,16 @@ def try_validate_utf8(str: "ColumnOrName") -> Column: try_validate_utf8.__doc__ = pysparkfuncs.try_validate_utf8.__doc__ +def normalize(str: "ColumnOrName", form: Optional["ColumnOrName"] = None) -> Column: + if form is None: + return _invoke_function_over_columns("normalize", str) + else: + return _invoke_function_over_columns("normalize", str, form) + + +normalize.__doc__ = pysparkfuncs.normalize.__doc__ + + def format_number(col: "ColumnOrName", d: int) -> Column: return _invoke_function("format_number", _to_col(col), lit(d)) @@ -4689,6 +4791,20 @@ def md5(col: "ColumnOrName") -> Column: md5.__doc__ = pysparkfuncs.md5.__doc__ +def xxh3_64(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xxh3_64", col) + + +xxh3_64.__doc__ = pysparkfuncs.xxh3_64.__doc__ + + +def xxh3_128(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("xxh3_128", col) + + +xxh3_128.__doc__ = pysparkfuncs.xxh3_128.__doc__ + + def sha1(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("sha1", col) @@ -5556,6 +5672,34 @@ def bitmap_count(col: "ColumnOrName") -> Column: bitmap_count.__doc__ = pysparkfuncs.bitmap_count.__doc__ +def bitmap_and(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_and", left, right) + + +bitmap_and.__doc__ = pysparkfuncs.bitmap_and.__doc__ + + +def bitmap_or(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_or", left, right) + + +bitmap_or.__doc__ = pysparkfuncs.bitmap_or.__doc__ + + +def bitmap_andnot(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_andnot", left, right) + + +bitmap_andnot.__doc__ = pysparkfuncs.bitmap_andnot.__doc__ + + +def bitmap_xor(left: "ColumnOrName", right: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_xor", left, right) + + +bitmap_xor.__doc__ = pysparkfuncs.bitmap_xor.__doc__ + + def bitmap_or_agg(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("bitmap_or_agg", col) @@ -5570,6 +5714,13 @@ def bitmap_and_agg(col: "ColumnOrName") -> Column: bitmap_and_agg.__doc__ = pysparkfuncs.bitmap_and_agg.__doc__ +def bitmap_xor_agg(col: "ColumnOrName") -> Column: + return _invoke_function_over_columns("bitmap_xor_agg", col) + + +bitmap_xor_agg.__doc__ = pysparkfuncs.bitmap_xor_agg.__doc__ + + # Geospatial ST Functions @@ -5638,6 +5789,26 @@ def unwrap_udt(col: "ColumnOrName") -> Column: unwrap_udt.__doc__ = pysparkfuncs.unwrap_udt.__doc__ +def wrap_udt(col: "ColumnOrName", udt: "Union[UserDefinedType, Column]") -> Column: + if isinstance(udt, _UserDefinedType): + udt_col = lit(udt.json()) + elif isinstance(udt, Column): + udt_col = udt + else: + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "UserDefinedType or Column", + "arg_name": "udt", + "arg_type": type(udt).__name__, + }, + ) + return _invoke_function("wrap_udt", _to_col(col), _to_col(udt_col)) + + +wrap_udt.__doc__ = pysparkfuncs.wrap_udt.__doc__ + + def udf( f: Optional[Union[Callable[..., Any], "DataTypeOrString"]] = None, returnType: "DataTypeOrString" = StringType(), @@ -5758,11 +5929,12 @@ def vector_sum(col: "ColumnOrName") -> Column: def _test() -> None: - import sys - import os import doctest - from pyspark.sql import SparkSession as PySparkSession + import os + import sys + import pyspark.sql.connect.functions.builtin + from pyspark.sql import SparkSession as PySparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.connect.functions.builtin.__dict__.copy() diff --git a/python/pyspark/sql/connect/functions/partitioning.py b/python/pyspark/sql/connect/functions/partitioning.py index b5fd0442318d0..22c9977e445e2 100644 --- a/python/pyspark/sql/connect/functions/partitioning.py +++ b/python/pyspark/sql/connect/functions/partitioning.py @@ -14,13 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Union from pyspark.errors import PySparkTypeError from pyspark.sql import functions as pysparkfuncs from pyspark.sql.column import Column -from pyspark.sql.connect.functions.builtin import _to_col, _invoke_function_over_columns -from pyspark.sql.connect.functions.builtin import lit, _invoke_function +from pyspark.sql.connect.functions.builtin import ( + _invoke_function, + _invoke_function_over_columns, + _to_col, + lit, +) if TYPE_CHECKING: from pyspark.sql.connect._typing import ColumnOrName @@ -76,11 +80,12 @@ def hours(col: "ColumnOrName") -> Column: def _test() -> None: - import sys - import os import doctest - from pyspark.sql import SparkSession as PySparkSession + import os + import sys + import pyspark.sql.connect.functions.partitioning + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.functions.partitioning.__dict__.copy() diff --git a/python/pyspark/sql/connect/group.py b/python/pyspark/sql/connect/group.py index 7fee9400fabd5..8c21b8154e39a 100644 --- a/python/pyspark/sql/connect/group.py +++ b/python/pyspark/sql/connect/group.py @@ -17,37 +17,36 @@ import warnings from typing import ( + TYPE_CHECKING, Dict, List, + Optional, Sequence, Union, - TYPE_CHECKING, - Optional, - overload, cast, + overload, ) -from pyspark.util import PythonEvalType -from pyspark.sql.group import GroupedData as PySparkGroupedData -from pyspark.sql.pandas.group_ops import PandasCogroupedOps as PySparkPandasCogroupedOps -from pyspark.sql.pandas.functions import _validate_vectorized_udf # type: ignore[attr-defined] -from pyspark.sql.pandas.typehints import infer_group_arrow_eval_type_from_func -from pyspark.sql.types import NumericType, StructType - import pyspark.sql.connect.plan as plan +from pyspark.errors import PySparkNotImplementedError, PySparkTypeError from pyspark.sql.column import Column from pyspark.sql.connect.functions import builtin as F -from pyspark.errors import PySparkNotImplementedError, PySparkTypeError +from pyspark.sql.group import GroupedData as PySparkGroupedData +from pyspark.sql.pandas.functions import _validate_vectorized_udf # type: ignore[attr-defined] +from pyspark.sql.pandas.group_ops import PandasCogroupedOps as PySparkPandasCogroupedOps +from pyspark.sql.pandas.typehints import infer_group_arrow_eval_type_from_func from pyspark.sql.streaming.stateful_processor import StatefulProcessor +from pyspark.sql.types import NumericType, StructType +from pyspark.util import PythonEvalType if TYPE_CHECKING: from pyspark.sql.connect._typing import ( - LiteralType, - PandasGroupedMapFunction, - GroupedMapPandasUserDefinedFunction, - PandasCogroupedMapFunction, ArrowCogroupedMapFunction, ArrowGroupedMapFunction, + GroupedMapPandasUserDefinedFunction, + LiteralType, + PandasCogroupedMapFunction, + PandasGroupedMapFunction, PandasGroupedMapFunctionWithState, ) from pyspark.sql.connect.dataframe import DataFrame @@ -296,8 +295,8 @@ def apply(self, udf: "GroupedMapPandasUserDefinedFunction") -> "DataFrame": def applyInPandas( self, func: "PandasGroupedMapFunction", schema: Union["StructType", str] ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.pandas.typehints import infer_group_pandas_eval_type_from_func # Try to infer the eval type from type hints @@ -342,8 +341,8 @@ def applyInPandasWithState( outputMode: str, timeoutConf: str, ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction _validate_vectorized_udf(func, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE) udf_obj = UserDefinedFunction( @@ -387,8 +386,8 @@ def transformWithStateInPandas( initialState: Optional["GroupedData"] = None, eventTimeColumnName: str = "", ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.streaming.stateful_processor_util import ( TransformWithStateInPandasUdfUtils, ) @@ -439,8 +438,8 @@ def transformWithState( initialState: Optional["GroupedData"] = None, eventTimeColumnName: str = "", ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.streaming.stateful_processor_util import ( TransformWithStateInPandasUdfUtils, ) @@ -485,8 +484,8 @@ def transformWithState( def applyInArrow( self, func: "ArrowGroupedMapFunction", schema: Union[StructType, str] ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction try: # Try to infer the eval type from type hints @@ -539,8 +538,8 @@ def __init__(self, gd1: "GroupedData", gd2: "GroupedData"): def applyInPandas( self, func: "PandasCogroupedMapFunction", schema: Union["StructType", str] ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction _validate_vectorized_udf(func, PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF) if isinstance(schema, str): @@ -570,8 +569,8 @@ def applyInPandas( def applyInArrow( self, func: "ArrowCogroupedMapFunction", schema: Union[StructType, str] ) -> "DataFrame": - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.udf import UserDefinedFunction _validate_vectorized_udf(func, PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF) if isinstance(schema, str): @@ -603,11 +602,12 @@ def applyInArrow( def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.group + from pyspark.sql import SparkSession as PySparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.connect.group.__dict__.copy() diff --git a/python/pyspark/sql/connect/local_server.py b/python/pyspark/sql/connect/local_server.py new file mode 100644 index 0000000000000..a765f4cc117da --- /dev/null +++ b/python/pyspark/sql/connect/local_server.py @@ -0,0 +1,628 @@ +# +# 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. +# + +""" +Opt-in reuse of a persistent local Spark Connect server (``spark.local.connect.reuse`` / +``SPARK_LOCAL_CONNECT_REUSE``). + +By default ``SparkSession.builder.remote("local[*]").getOrCreate()`` boots a fresh in-process +Connect server in every Python process. With reuse enabled, the first run starts one +long-lived server through ``sbin/start-connect-server.sh`` and records how to reach it (host, +port, auth token, pid, Spark version) in a discovery file; later runs reconnect to it if the +version matches, the pid is alive, and the port accepts connections. Each run still gets its +own server-side session, so session-local state does not leak between runs. + +The discovery file, the daemon's pid file, and the logs live in a per-user ``0700`` directory +under the system temp dir; ``SPARK_LOCAL_CONNECT_DISCOVERY`` overrides the discovery file +location. The auth token is stored with ``0600`` and the server always binds IPv4 loopback, +overriding any configured binding address, so other users on the machine can neither read the +token nor authenticate to the server. Processes of the same user share the server by design. + +The server runs until stopped with ``python -m pyspark.sql.connect.local_server --stop``. +(A plain ``sbin/stop-connect-server.sh`` cannot find it: the daemon runs with a custom pid +dir and ident string.) Windows is not supported, as this relies on the POSIX scripts under +``sbin/``. + +This module is experimental. The discovery file location and format and the ``--stop`` +entry point are internal details that may change or move server-side (e.g. into a unified +``spark connect`` CLI); only the reuse opt-in itself is meant to be a stable surface. +""" + +import argparse +import contextlib +import getpass +import json +import os +import signal +import socket +import subprocess +import sys +import tempfile +import time +import uuid +from typing import Any, Dict, Iterator, Optional, TextIO + +from pyspark.errors import PySparkRuntimeError + +_SERVER_CLASS = "org.apache.spark.sql.connect.service.SparkConnectServer" +# A fixed SPARK_IDENT_STRING keeps the spark-daemon.sh pid and log file names stable +# regardless of $USER. +_SPARK_IDENT = "local-connect" +_LINUX_ZOMBIE_STATE = "Z" + + +def _pid_alive(pid: int) -> bool: + """Whether ``pid`` is running. A process we cannot signal counts as alive. Linux zombies + count as terminated: they remain signalable until their parent reaps them, but cannot own + or serve a managed server. + + Off POSIX this returns ``True`` without probing: ``os.kill`` there terminates the target for + any signal other than ``CTRL_C_EVENT`` / ``CTRL_BREAK_EVENT``, so signal 0 is not a safe + liveness probe, and callers fall through to the port check instead. Guarding here rather than + at each call site keeps every caller (reuse and pool) safe. (The pool path needs ``fcntl`` and + so never runs off POSIX regardless.) + """ + if os.name != "posix": + return True + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except OverflowError: + return False + except OSError: + pass + if sys.platform.startswith("linux"): + try: + with open(f"/proc/{pid}/status", encoding="utf-8") as status_file: + for line in status_file: + key, separator, value = line.partition(":") + if separator and key == "State": + state, _, _ = value.strip().partition(" ") + if state == _LINUX_ZOMBIE_STATE: + return False + break + except FileNotFoundError: + return False + except OSError: + pass + return True + + +def _port_open(host: str, port: int, timeout: float = 0.5) -> bool: + """Whether a TCP connection to ``host``:``port`` succeeds within ``timeout`` seconds. A + socket error or a host that fails to resolve counts as closed. + """ + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(timeout) + return sock.connect_ex((host, port)) == 0 + except (OSError, UnicodeError): + return False + + +def _process_command(pid: int) -> Optional[str]: + """The command of ``pid``, an empty string if it is gone, or ``None`` if inspection fails.""" + try: + result = subprocess.run( + ["ps", "-ww", "-p", str(pid), "-o", "command="], + capture_output=True, + text=True, + timeout=5, + ) + except (OSError, subprocess.SubprocessError): + return None + return result.stdout if result.returncode == 0 else "" + + +def _is_local_connect_server(pid: int) -> Optional[bool]: + """Whether ``pid`` is still the managed Connect server recorded in discovery. + + Returns ``None`` when the process cannot be inspected, so callers do not discard the + discovery information needed to retry later. + """ + command = _process_command(pid) + return None if command is None else _SERVER_CLASS in command + + +def runtime_dir() -> str: + """Return the private per-user directory holding local-server state.""" + path = os.path.join(tempfile.gettempdir(), "spark-connect-{}".format(getpass.getuser())) + try: + # exist_ok also covers two first runs racing to create the directory; chmod + # re-asserts 0700 and fails if another user owns the path. + os.makedirs(path, mode=0o700, exist_ok=True) + os.chmod(path, 0o700) + except OSError as e: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_RUNTIME_DIR_UNAVAILABLE", + messageParameters={"path": path}, + ) from e + return path + + +class Discovery: + """Reads and writes the discovery file recording the persistent local server. + + The file lives in a per-user directory under the system temp dir, or wherever + ``SPARK_LOCAL_CONNECT_DISCOVERY`` points; the daemon's pid file and logs sit next to it. + """ + + def __init__(self, path: Optional[str] = None): + self.path = os.path.abspath( + path + or os.environ.get("SPARK_LOCAL_CONNECT_DISCOVERY") + or os.path.join(runtime_dir(), "connect-local.json") + ) + self._lock_file: Optional[TextIO] = None + + @property + def directory(self) -> str: + return os.path.dirname(self.path) + + @property + def daemon_pid_path(self) -> str: + return os.path.join(self.directory, "spark-{}-{}-1.pid".format(_SPARK_IDENT, _SERVER_CLASS)) + + def daemon_pid(self) -> Optional[int]: + """The pid recorded by ``spark-daemon.sh``, or ``None`` if absent or unreadable. + The daemon writes this file outside our lock, so no lock is required to read it. + """ + try: + with open(self.daemon_pid_path, "r") as f: + return int(f.read().strip()) + except (OSError, ValueError): + return None + + def __enter__(self) -> "Discovery": + os.makedirs(self.directory, exist_ok=True) + self._lock_file = open(self.path + ".lock", "a+") + import fcntl + + fcntl.flock(self._lock_file.fileno(), fcntl.LOCK_EX) + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + assert self._lock_file is not None + self._lock_file.close() + self._lock_file = None + + def _assert_locked(self) -> None: + assert self._lock_file is not None, "Discovery must be used as a context manager" + + def load(self) -> Optional[Dict[str, Any]]: + """Read the discovery file, returning ``None`` if it is absent or malformed.""" + self._assert_locked() + try: + with open(self.path, "r") as f: + data = json.load(f) + except (OSError, ValueError): + return None + if not isinstance(data, dict): + return None + try: + data["port"] = int(data["port"]) + data["pid"] = int(data["pid"]) + except (KeyError, TypeError, ValueError): + return None + if not all(isinstance(data[k], str) for k in ("host", "token", "spark_version")): + return None + return data + + def save(self, data: Dict[str, Any]) -> None: + """Write the discovery file with ``0600`` perms; it holds the auth token. Readers + and writers all hold the exclusive lock, so the write does not need to be atomic. + """ + self._assert_locked() + fd = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + # O_CREAT applies the 0600 mode only when it creates the file; re-assert it in + # case a pre-existing file had wider permissions. + os.fchmod(fd, 0o600) + f.write(json.dumps(data)) + + def clear(self) -> None: + self._assert_locked() + with contextlib.suppress(OSError): + os.remove(self.path) + with contextlib.suppress(OSError): + os.remove(self.daemon_pid_path) + + +class LocalConnectServer: + """The persistent server described by a locked ``Discovery``.""" + + def __init__(self, discovery: Discovery): + self._discovery = discovery + self._reload() + + def _reload(self) -> None: + # ``data`` is None when no server is recorded yet (first run, or the file was + # cleared): the None fields make ``is_reusable()`` False, so ``reuse_or_start()`` + # launches a fresh server. + data = self._discovery.load() + self.host = data["host"] if data else None + self.port = data["port"] if data else None + self.token = data["token"] if data else None + self.pid = data["pid"] if data else None + self.spark_version = data["spark_version"] if data else None + + @property + def url(self) -> str: + assert self.host is not None and self.port is not None + return "sc://{}:{}".format(self.host, self.port) + + def is_listening(self) -> bool: + if self.host is None or self.port is None: + return False + return _port_open(self.host, self.port) + + def is_reusable(self) -> bool: + from pyspark.version import __version__ + + if self.spark_version != __version__ or self.pid is None: + return False + if not _pid_alive(self.pid): + return False + return self.is_listening() + + def reuse_or_start(self, master: str, opts: Dict[str, Any]) -> str: + if not self.is_reusable(): + self.start(master, opts) + assert self.token is not None + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = self.token + return self.url + + def start( + self, + master: str, + opts: Dict[str, Any], + *, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ) -> None: + """Start this server and reload its discovery record. + + Callers starting isolated daemons can request an ephemeral port and provide a + precomputed startup configuration while sharing the standard launch path. + """ + ServerLauncher( + master, + opts, + self._discovery, + use_ephemeral_port=use_ephemeral_port, + seed_conf=seed_conf, + ).launch() + self._reload() + + def stop(self) -> Optional[bool]: + stopped = False + if self.pid is not None: + is_server = _is_local_connect_server(self.pid) + if is_server is None: + return None + if is_server: + try: + os.kill(self.pid, signal.SIGTERM) + stopped = True + except OSError: + pass + self._discovery.clear() + return stopped + + +def _strip_launcher_conf(conf: Dict[str, Any]) -> Dict[str, Any]: + """Drop keys the launcher sets itself (master, binding port, auth token) and the + ``spark.local.connect.*`` opt-in keys, returning a new dict. Idempotent, so it is safe to + apply to a conf that is already sanitized. + """ + stripped = dict(conf) + for k in list(stripped): + if k in ( + "spark.remote", + "spark.api.mode", + "spark.master", + "spark.connect.authenticate.token", + "spark.connect.grpc.binding.address", + "spark.connect.grpc.binding.port", + ) or k.startswith("spark.local.connect."): + stripped.pop(k) + return stripped + + +def startup_seed_conf(opts: Dict[str, Any]) -> Dict[str, Any]: + """Compute startup confs using the same merge as the in-process server path, then strip + the keys the launcher manages (see ``_strip_launcher_conf``). + """ + conf: Dict[str, Any] = {} + for i in range(int(os.environ.get("PYSPARK_REMOTE_INIT_CONF_LEN", "0"))): + conf = json.loads(os.environ["PYSPARK_REMOTE_INIT_CONF_{}".format(i)]) + conf.update(opts) + return _strip_launcher_conf(conf) + + +class ServerLauncher: + """Starts a persistent local server via ``sbin/start-connect-server.sh`` and waits until + it accepts connections. Callers must hold a ``Discovery`` context. + + ``use_ephemeral_port`` and ``seed_conf`` allow other managed local servers to share this + launch path without duplicating its process and readiness handling. + """ + + _READY_TIMEOUT = 120 + + def __init__( + self, + master: str, + opts: Dict[str, Any], + discovery: Discovery, + use_ephemeral_port: bool = False, + seed_conf: Optional[Dict[str, Any]] = None, + ): + self._master = master + self._opts = opts + self._discovery = discovery + self._use_ephemeral_port = use_ephemeral_port + self._seed_override = seed_conf + self._log_dir = os.path.join(discovery.directory, "logs") + + def launch(self) -> None: + token = self._token() + port = self._pick_port() + # The conf file must outlive _await_ready: spark-daemon.sh backgrounds the JVM, + # which reads --properties-file while starting up. + with self._seed_properties_file() as conf_file: + self._run_script(port, token, conf_file) + self._await_ready(port, token) + + def _token(self) -> str: + # Same precedence as the in-process _start_connect_server: explicit env token, then + # conf, then a fresh one. Passed via the environment so it never shows up in `ps`. + return ( + os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN") + or self._opts.get("spark.connect.authenticate.token") + or str(uuid.uuid4()) + ) + + def _pick_port(self) -> int: + """Use an OS-assigned free port when requested or under SPARK_TESTING so suites can + run in parallel. Otherwise honor the configured/default port, falling back to a free + one if another process holds it. (A live stale server of ours also holds the port, but + that start fails later at spark-daemon.sh's pid-file check regardless of port.) The + sbin script cannot report an ephemeral port back, so the free port is picked and + released here, with a small race until the server binds it. + """ + + def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + return sock.getsockname()[1] + + if self._use_ephemeral_port or "SPARK_TESTING" in os.environ: + return free_port() + from pyspark.sql.connect.client import DefaultChannelBuilder + + port = int( + self._opts.get("spark.local.connect.server.port", DefaultChannelBuilder.default_port()) + ) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + try: + sock.bind(("localhost", port)) + return port + except OSError: + return free_port() + + def _seed_conf(self) -> Dict[str, Any]: + """Startup confs for the new server, minus the keys the launcher sets itself and the + ``spark.local.connect.*`` opt-in keys. Only the run that starts the server can seed + static confs; later runs find the JVM already warm. + + With no ``seed_conf`` override, this merges ``PYSPARK_REMOTE_INIT_CONF_*`` with the + builder opts like the in-process ``_start_connect_server`` does. When an override is + given, it is used verbatim instead of the merge. Either way the result is run through + ``_strip_launcher_conf``, so the launcher-managed keys never reach + ``--properties-file`` even if a caller passes raw opts as the override. + """ + if self._seed_override is not None: + return _strip_launcher_conf(self._seed_override) + return startup_seed_conf(self._opts) + + @contextlib.contextmanager + def _seed_properties_file(self) -> Iterator[Optional[str]]: + # NamedTemporaryFile creates the file with 0600 perms since confs may hold sensitive + # values; a --properties-file keeps them off the server's argv where they would show + # up in `ps`. Yields None when there is nothing to seed. + seed = self._seed_conf() + if not seed: + yield None + return + with tempfile.NamedTemporaryFile( + mode="w", + prefix="connect-local-conf-", + suffix=".properties", + dir=self._discovery.directory, + ) as f: + for key, value in seed.items(): + escaped = str(value).replace("\\", "\\\\").replace("\n", "\\n") + f.write("{}={}\n".format(key, escaped)) + f.flush() + yield f.name + + def _run_script(self, port: int, token: str, conf_file: Optional[str]) -> None: + from pyspark.find_spark_home import _find_spark_home + + spark_home = os.environ.get("SPARK_HOME") or _find_spark_home() + script = os.path.join(spark_home, "sbin", "start-connect-server.sh") + if not os.path.isfile(script): + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={"reason": "cannot find {}".format(script)}, + ) + + env = dict(os.environ) + for var in ("SPARK_REMOTE", "SPARK_LOCAL_REMOTE", "SPARK_CONNECT_MODE_ENABLED"): + env.pop(var, None) + env["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = token + env["SPARK_PID_DIR"] = self._discovery.directory + env["SPARK_LOG_DIR"] = self._log_dir + env["SPARK_IDENT_STRING"] = _SPARK_IDENT + + cmd = [ + script, + "--master", + self._master, + "--conf", + "spark.connect.grpc.binding.address=127.0.0.1", + "--conf", + "spark.connect.grpc.binding.port={}".format(port), + ] + if conf_file is not None: + cmd += ["--properties-file", conf_file] + + result = subprocess.run( + cmd, + env=env, + stdin=subprocess.DEVNULL, + capture_output=True, + text=True, + timeout=120, + ) + if result.returncode != 0: + stale_pid = self._discovery.daemon_pid() + if stale_pid is not None and _pid_alive(stale_pid): + # spark-daemon.sh refuses to start while its pid file points at a live + # process -- here a server this client just rejected as not reusable + # (e.g. after a Spark upgrade). + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "a local Connect server that is not reusable by this client is " + "already running (pid {}); stop it with " + "`python -m pyspark.sql.connect.local_server --stop`".format(stale_pid) + }, + ) + output = (result.stderr or "") + (result.stdout or "") + last_line = output.strip().splitlines()[-1] if output.strip() else "" + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "start-connect-server.sh exited with code {}: {}".format( + result.returncode, last_line + ) + }, + ) + + def _await_ready(self, port: int, token: str) -> None: + from pyspark.version import __version__ + + deadline = time.time() + self._READY_TIMEOUT + while time.time() < deadline: + pid = self._discovery.daemon_pid() + if pid is not None: + if _port_open("localhost", port): + self._discovery.save( + { + "host": "localhost", + "port": port, + "token": token, + "pid": pid, + "spark_version": __version__, + } + ) + return + if not _pid_alive(pid): + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server exited during start-up; see logs under {}".format( + self._log_dir + ) + }, + ) + time.sleep(0.25) + + pid = self._discovery.daemon_pid() + if pid is not None: + with contextlib.suppress(OSError): + os.kill(pid, signal.SIGTERM) + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "the server did not become ready within {} seconds; " + "see logs under {}".format(self._READY_TIMEOUT, self._log_dir) + }, + ) + + +def reuse_or_start_local_connect_server(master: str, opts: Dict[str, Any]) -> str: + """Reuse a running persistent local Connect server, or start one if none is reusable. + + Returns the ``sc://host:port`` endpoint and sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so + the client authenticates against that server. Only reached for a ``local`` master when + the reuse opt-in is set; see ``SparkSession.getOrCreate`` in ``pyspark.sql.session``. + """ + if os.name != "posix": + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "spark.local.connect.reuse relies on the POSIX scripts under sbin/; " + "on this platform start a server manually (sbin/start-connect-server.sh) and " + 'connect with .remote("sc://...")' + }, + ) + with Discovery() as discovery: + return LocalConnectServer(discovery).reuse_or_start(master, opts) + + +def stop_local_connect_server() -> Optional[bool]: + """Stop the recorded persistent local Connect server, if any; safe to call when none is + running. Returns ``True`` when the server was signalled, ``False`` when no matching server + was found, and ``None`` when the process could not be inspected. Also available as + ``python -m pyspark.sql.connect.local_server --stop``. + """ + with Discovery() as discovery: + return LocalConnectServer(discovery).stop() + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Manage the persistent local Spark Connect server used by the opt-in " + "spark.local.connect.reuse path. The server itself is started on demand through " + "sbin/start-connect-server.sh." + ) + parser.add_argument( + "--stop", action="store_true", help="stop the recorded running server, if any" + ) + args = parser.parse_args() + + if not args.stop: + parser.print_help(sys.stderr) + sys.exit(2) + stopped = stop_local_connect_server() + if stopped: + print("Stopped the persistent local Spark Connect server.") + elif stopped is None: + print("Could not verify the persistent local Spark Connect server; try again later.") + sys.exit(1) + else: + print("No running persistent local Spark Connect server found.") + + +if __name__ == "__main__": + main() diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py new file mode 100644 index 0000000000000..df24644d5f373 --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,1443 @@ +# +# 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. +# + +""" +Opt-in pool of single-use local Spark Connect servers +(``spark.local.connect.pool`` / ``SPARK_LOCAL_CONNECT_POOL``). + +The reuse mode (``spark.local.connect.reuse``, see ``pyspark.sql.connect.local_server``) makes +local runs fast by sharing one long-lived server, at the price of state backed by the shared +``SparkContext`` (persistent catalog, global temp views, cached data) carrying across runs. +The pool keeps the speed without the sharing: it maintains a small set of booted servers that +have never been assigned to an application run, and +``SparkSession.builder.remote("local[*]").getOrCreate()`` *claims* one exclusively, spawns a +replacement in the background, and tears the claimed server down when the session stops or the +client exits. No server ever serves two application runs, so runs are as isolated from each +other as with the default in-process server -- at the cost of the idle servers' memory while +you iterate. If both opt-ins are set, the pool takes precedence. + +The pool lives in a ``pool`` subdirectory of the per-user runtime directory (override with +``SPARK_LOCAL_CONNECT_POOL_DIR``). A member is a set of files named by a random ``<uid>``; +every access happens under the directory's ``.lock`` file lock, so readers always observe +complete states: + + pending-<uid>.json an in-flight launch (the attendant process booting the server) + conf-<uid>.json startup confs for that launch, read once by its attendant + server-<uid>.json a ready, unclaimed server (host/port/token/pid/version) + claimed-<pid>-<uid>.json a server owned by the live client process <pid> + retired-<uid>.json a server being torn down; hard-killed if it hangs + member-<uid>/ the server's pid file and logs (spark-daemon.sh directories) + +Each launch runs an *attendant* (``python -m pyspark.sql.connect.local_server_pool --attend``), +a small detached process that boots the server through ``sbin/start-connect-server.sh``, +publishes its ``server-<uid>.json``, and supervises it. It retires the server once it has sat +unclaimed past the idle timeout, or once the client that claimed it has died without releasing +it. A janitor pass on every acquire is the backstop for members whose attendant itself died. + +Servers are only handed to runs they were built for: each member carries a fingerprint of its +master, seeded confs, working directory, and Python executable, and a run only claims members +whose fingerprint matches its own. ``python -m pyspark.sql.connect.local_server_pool --purge`` +force-stops every member and empties the pool directory. + +This mode is experimental. Everything but the opt-in itself -- the pool directory layout, the +attendant, and the ``--purge`` entry point -- is an internal detail that may change or move, +for example into a unified ``spark connect`` CLI. POSIX only, like the reuse mode. +""" + +import argparse +import atexit +import contextlib +import hashlib +import json +import math +import os +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import uuid +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple + +from pyspark.errors import PySparkRuntimeError, PySparkValueError +from pyspark.sql.connect.local_server import ( + Discovery, + _is_local_connect_server, + _pid_alive, + _port_open, + _process_command, + runtime_dir, +) + +# Environment variables that shape the JVM the launcher boots through +# sbin/start-connect-server.sh -> spark-daemon.sh -> load-spark-env.sh / spark-submit. +# SPARK_CONF_DIR selects the spark-env.sh and spark-defaults.conf that seed the server; the +# rest feed the classpath, heap, and JVM options. This is a curated set, not an exhaustive +# one: the launcher inherits the whole environment, so it names the inputs that most commonly +# differ between runs rather than every variable a server could read. +# +# PATH is a deliberate omission: bin/spark-class prefers ${JAVA_HOME}/bin/java and only falls +# back to the first java on PATH, so two runs with different JDKs first on PATH and no JAVA_HOME +# would share a member. The identity already tracks PATH indirectly through shutil.which for the +# interpreters, and PATH is too volatile to fingerprint whole; a run needing a specific JDK +# should set JAVA_HOME. +_JVM_ENV_VARS = ( + "SPARK_CONF_DIR", + "JAVA_HOME", + "SPARK_DIST_CLASSPATH", + "SPARK_DAEMON_MEMORY", + "SPARK_DRIVER_MEMORY", + "SPARK_SUBMIT_OPTS", + "SPARK_DAEMON_JAVA_OPTS", +) + + +def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str: + """The identity of a pool member: a curated set of inputs that shape the server a run would + have booted for itself. A run only claims members whose fingerprint equals its own, so a + pre-booted JVM is never handed to a run it would not have produced. The set is curated + rather than complete because the launcher inherits the full environment (see + ``_JVM_ENV_VARS``) -- it covers the inputs that most commonly differ between runs. + + Besides the master and the seeded confs, this covers the working directory (unset warehouse + and Derby metastore locations resolve relative to it), the PySpark installation, the Python + interpreters the server would run UDFs and Python data sources with, and the environment + variables that shape the launched JVM. + """ + + def resolved(command: str) -> str: + # Relative commands resolve through PATH server-side; fold that in so equal command + # strings cannot stand for different interpreters on different PATHs. + return shutil.which(command) or command + + # Two server code paths resolve the Python interpreter with opposite precedence, so a run + # changing only one of these variables would still have booted a different server. Include + # both resolutions: SparkConnectPlanner.pythonExec (Connect Python UDFs) prefers + # PYSPARK_PYTHON, while PythonUtils.defaultPythonExec (Python data sources) prefers + # PYSPARK_DRIVER_PYTHON. Both fall back to python3 and treat an empty value as set, matching + # the Scala sys.env.getOrElse chains. + udf_python = resolved( + os.environ.get("PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3")) + ) + data_source_python = resolved( + os.environ.get("PYSPARK_DRIVER_PYTHON", os.environ.get("PYSPARK_PYTHON", "python3")) + ) + spark_home = os.environ.get("SPARK_HOME") + identity = [ + master, + sorted((str(k), str(v)) for k, v in seed_conf.items()), + os.getcwd(), + sys.executable, + udf_python, + data_source_python, + os.path.realpath(__file__), + os.path.realpath(spark_home) if spark_home else "", + os.environ.get("PYTHONPATH", ""), + [os.environ.get(var, "") for var in _JVM_ENV_VARS], + ] + return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest() + + +# How long one acquire may wait for a member to become ready. A cold launch takes at most +# 120s; this leaves room for one relaunch of a failed one. +_ACQUIRE_TIMEOUT = 180 + +_DEFAULT_POOL_SIZE = 2 + + +def _pool_size(opts: Dict[str, Any]) -> int: + """The number of ready or in-flight members to keep per fingerprint; at least one. + Malformed values fall back to the default rather than failing session creation over a + tuning knob. + """ + value = opts.get( + "spark.local.connect.pool.size", os.environ.get("SPARK_LOCAL_CONNECT_POOL_SIZE") + ) + try: + return max(1, int(value)) if value is not None else _DEFAULT_POOL_SIZE + except (TypeError, ValueError, OverflowError): + return _DEFAULT_POOL_SIZE + + +class _PoolStateRecord: + """Validation shared by JSON-backed pool state records.""" + + # The end of year 9999 UTC, as a Unix timestamp. Pool timestamps are wall-clock + # ``time.time()`` readings, so rejecting larger values keeps corrupt far-future records + # from looking perpetually fresh to age-based reaping. + _MAX_TIMESTAMP = 253402300799 + + @staticmethod + def _positive_pid(value: Any) -> Optional[int]: + """A positive integer process id, or ``None`` for malformed persisted data.""" + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return None + return value + + @classmethod + def _timestamp(cls, value: Any) -> Optional[float]: + """A finite persisted wall-clock timestamp, or ``None`` when malformed.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + timestamp = float(value) + except OverflowError: + return None + if not math.isfinite(timestamp) or not 0 <= timestamp <= cls._MAX_TIMESTAMP: + return None + return timestamp + + +@dataclass(frozen=True) +class PendingState(_PoolStateRecord): + """Validated fields of a ``pending-<uid>.json`` launch record.""" + + attendant_pid: int + created: float + fingerprint: str + + @classmethod + def attendant_pid_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[int]: + """Recover a valid attendant pid even when another record field is malformed.""" + return cls._positive_pid(data.get("attendant_pid")) if data is not None else None + + @classmethod + def from_data(cls, data: Optional[Dict[str, Any]]) -> Optional["PendingState"]: + if data is None: + return None + attendant_pid = cls.attendant_pid_from_data(data) + created = cls._timestamp(data.get("created")) + fingerprint = data.get("fingerprint") + if ( + attendant_pid is None + or created is None + or not isinstance(fingerprint, str) + or not fingerprint + ): + return None + return cls(attendant_pid, created, fingerprint) + + +@dataclass(frozen=True) +class RetiredState(_PoolStateRecord): + """Validated fields of a ``retired-<uid>.json`` shutdown record.""" + + pid: int + process_start_id: str + retired: float + signalled: bool = False + + @classmethod + def pid_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[int]: + """Recover a valid server pid even when the retirement time is malformed.""" + return cls._positive_pid(data.get("pid")) if data is not None else None + + @staticmethod + def process_start_id_from_data(data: Optional[Dict[str, Any]]) -> Optional[str]: + """Recover a process generation identifier from a malformed state record.""" + value = data.get("process_start_id") if data is not None else None + return value if isinstance(value, str) and value else None + + @classmethod + def retired_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[float]: + """Recover a valid retirement timestamp from a malformed state record.""" + return cls._timestamp(data.get("retired")) if data is not None else None + + @staticmethod + def signalled_from_data(data: Optional[Dict[str, Any]]) -> Optional[bool]: + """Recover SIGTERM delivery state, defaulting legacy records to unsignalled.""" + if data is None: + return None + value = data.get("signalled", False) + return value if isinstance(value, bool) else None + + @classmethod + def from_data(cls, data: Optional[Dict[str, Any]]) -> Optional["RetiredState"]: + if data is None: + return None + pid = cls.pid_from_data(data) + process_start_id = cls.process_start_id_from_data(data) + retired = cls.retired_from_data(data) + signalled = cls.signalled_from_data(data) + if pid is None or process_start_id is None or retired is None or signalled is None: + return None + return cls(pid, process_start_id, retired, signalled) + + def as_data(self) -> Dict[str, Any]: + return { + "pid": self.pid, + "process_start_id": self.process_start_id, + "retired": self.retired, + "signalled": self.signalled, + } + + +class PoolMember(_PoolStateRecord): + """One published pool server, wrapping its ``server-<uid>.json`` record.""" + + def __init__(self, data: Dict[str, Any]): + record = dict(data) + for key in ("host", "token", "spark_version", "fingerprint", "process_start_id"): + if not isinstance(record[key], str) or not record[key]: + raise PySparkValueError(f"{key} must be a nonempty string") + for key in ("port", "pid"): + value = record[key] + if isinstance(value, bool) or not isinstance(value, int): + raise PySparkValueError(f"{key} must be an integer") + created = record["created"] + if isinstance(created, bool) or not isinstance(created, (int, float)): + raise PySparkValueError("created must be a number") + created = float(created) + if not 1 <= record["port"] <= 65535: + raise PySparkValueError("port is out of range") + if record["pid"] <= 0: + raise PySparkValueError("pid must be positive") + if not math.isfinite(created) or not 0 <= created <= self._MAX_TIMESTAMP: + raise PySparkValueError( + f"created must be a finite timestamp in [0, {self._MAX_TIMESTAMP}]" + ) + self.host: str = record["host"] + self.port: int = record["port"] + self.token: str = record["token"] + self.pid: int = record["pid"] + self.spark_version: str = record["spark_version"] + self.fingerprint: str = record["fingerprint"] + self.process_start_id: str = record["process_start_id"] + self.created: float = created + # Set when this process claims the member; the path of its claimed-<pid>-<uid>.json. + self.claim_path: Optional[str] = None + + @classmethod + def from_data(cls, data: Dict[str, Any]) -> Optional["PoolMember"]: + """Parse a published member record, returning ``None`` when it is malformed.""" + try: + return cls(data) + except (KeyError, TypeError, ValueError, OverflowError): + return None + + @classmethod + def pid_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[int]: + """Recover a valid server pid even when another member field is malformed.""" + return cls._positive_pid(data.get("pid")) if data is not None else None + + @staticmethod + def process_start_id_from_data(data: Optional[Dict[str, Any]]) -> Optional[str]: + """Recover a process generation identifier from a malformed member record.""" + value = data.get("process_start_id") if data is not None else None + return value if isinstance(value, str) and value else None + + @staticmethod + def client_process_start_id_from_data(data: Optional[Dict[str, Any]]) -> Optional[str]: + """Recover the claiming client's process generation from a claimed record.""" + value = data.get("client_process_start_id") if data is not None else None + return value if isinstance(value, str) and value else None + + @property + def url(self) -> str: + return f"sc://{self.host}:{self.port}" + + def is_usable(self) -> bool: + """Whether this member has a matching Spark version, live process, and open port. Uses + the same liveness and reachability probes as the reuse path (see ``local_server``), so + the pool and reuse discovery agree on when a recorded server is still good.""" + from pyspark.version import __version__ + + if ( + self.spark_version != __version__ + or not _pid_alive(self.pid) + or ServerPool._process_start_id(self.pid) != self.process_start_id + ): + return False + return _port_open(self.host, self.port) + + +class PoolDirectory: + """Path layout, file access, and the cross-process lock of one pool directory. + + Used as a context manager that holds the directory's exclusive lock: + + directory = PoolDirectory() + with directory: + path = directory.pending_path(uid) + directory.write_json(path, data) + stored = directory.read_json(path) + directory.rename(path, directory.server_path(uid)) + + Entering the context creates the directory and acquires its lock. Callers then use path + builders and the locked accessors to enumerate, read, write, rename, or remove state. Exiting + the context releases the lock. + + Pool operations are infrequent, so one exclusive lock for every state transition is simpler + than a finer-grained scheme. A context can be entered again after exiting, allowing callers + to release the lock between polling attempts so other processes can update the directory. + """ + + _STATE_TEMP_PREFIX = ".pool-state-" + + def __init__(self, path: Optional[str] = None): + if path is None: + path = os.environ.get("SPARK_LOCAL_CONNECT_POOL_DIR") + if path is None: + path = os.path.join(runtime_dir(), "pool") + self.path = os.path.abspath(path) + self._lock_fd: Optional[int] = None + + def __enter__(self) -> "PoolDirectory": + import fcntl + + # Not reentrant: a nested enter would os.open a second fd and flock(LOCK_EX) would block + # forever against the fd this process already holds. Fail loudly instead of deadlocking. + assert self._lock_fd is None, "PoolDirectory is not reentrant" + os.makedirs(self.path, mode=0o700, exist_ok=True) + # Re-assert privacy for an existing override directory: state files contain auth tokens, + # and directory write access would allow replacing them or bypassing the shared lock. + os.chmod(self.path, 0o700) + lock_fd = os.open(os.path.join(self.path, ".lock"), os.O_RDWR | os.O_CREAT, 0o600) + try: + os.fchmod(lock_fd, 0o600) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + except BaseException: + os.close(lock_fd) + raise + self._lock_fd = lock_fd + try: + # A process killed between writing and replacing an atomic state-file update can + # leave its private temporary file behind. No writer can still be active once this + # lock is acquired, so these leftovers are always safe to discard. + for name in self._entries(): + if name.startswith(self._STATE_TEMP_PREFIX): + with contextlib.suppress(OSError): + os.remove(os.path.join(self.path, name)) + except BaseException: + os.close(lock_fd) + self._lock_fd = None + raise + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + assert self._lock_fd is not None + os.close(self._lock_fd) # closing releases the lock + self._lock_fd = None + + def _assert_locked(self) -> None: + assert self._lock_fd is not None, "PoolDirectory must be used as a context manager" + + # Path builders; these do not touch the filesystem and need no lock. + + def pending_path(self, uid: str) -> str: + return os.path.join(self.path, f"pending-{uid}.json") + + def conf_path(self, uid: str) -> str: + return os.path.join(self.path, f"conf-{uid}.json") + + def server_path(self, uid: str) -> str: + return os.path.join(self.path, f"server-{uid}.json") + + def claimed_path(self, client_pid: int, uid: str) -> str: + return os.path.join(self.path, f"claimed-{client_pid}-{uid}.json") + + def retired_path(self, uid: str) -> str: + return os.path.join(self.path, f"retired-{uid}.json") + + def member_dir(self, uid: str) -> str: + return os.path.join(self.path, f"member-{uid}") + + # uids are generated as ``uuid.uuid4().hex[:12]`` (see the acquisition layer), so a valid + # uid is a nonempty run of lowercase hex. Validating the shape keeps editor droppings such + # as ``member-abc.json.swp`` and empty stems like ``server-.json`` from becoming phantom uids. + _UID_CHARS = frozenset("0123456789abcdef") + + @classmethod + def _is_uid(cls, uid: str) -> bool: + return bool(uid) and all(c in cls._UID_CHARS for c in uid) + + @classmethod + def _split_claimed(cls, stem: str) -> Optional[Tuple[str, str]]: + """Split a well-formed ``claimed-<pid>-<uid>`` stem (without the ``.json`` suffix) into + ``(client_pid, uid)`` as strings, or ``None`` otherwise. The pid is returned unparsed: + ``parse_entry`` classifies over every directory entry and must never raise, and + ``str.isdigit()`` accepts characters ``int()`` rejects (e.g. superscripts), so the + ``isascii()`` guard keeps the eventual ``int()`` in ``claiming_pid`` total.""" + if not stem.startswith("claimed-"): + return None + client_pid, sep, uid = stem[len("claimed-") :].partition("-") + if not sep or not (client_pid.isascii() and client_pid.isdigit()) or not cls._is_uid(uid): + return None + return client_pid, uid + + @classmethod + def parse_entry(cls, name: str) -> Tuple[Optional[str], Optional[str]]: + """The ``(kind, uid)`` of a pool directory entry, ``(None, None)`` for anything + else (the lock file, editor droppings, entries with a malformed uid, ...).""" + if name.startswith("member-"): + uid = name[len("member-") :] + return ("member", uid) if cls._is_uid(uid) else (None, None) + if not name.endswith(".json"): + return None, None + stem = name[: -len(".json")] + for kind in ("pending", "conf", "server", "retired"): + if stem.startswith(kind + "-"): + uid = stem[len(kind) + 1 :] + return (kind, uid) if cls._is_uid(uid) else (None, None) + claimed = cls._split_claimed(stem) + return ("claimed", claimed[1]) if claimed is not None else (None, None) + + @classmethod + def claiming_pid(cls, claimed_path: str) -> int: + """The client pid recorded in a ``claimed-<pid>-<uid>.json`` file name.""" + name = os.path.basename(claimed_path) + stem = name[: -len(".json")] if name.endswith(".json") else name + claimed = cls._split_claimed(stem) + assert claimed is not None, f"not a claimed entry: {claimed_path!r}" + return int(claimed[0]) + + # Locked accessors. + + def uids(self) -> List[str]: + self._assert_locked() + seen = [] + for name in self._entries(): + _, uid = self.parse_entry(name) + if uid is not None and uid not in seen: + seen.append(uid) + return seen + + def states(self, uid: str) -> Dict[str, str]: + """The state entries currently existing for ``uid``, as ``{kind: path}`` with kinds + ``pending``, ``conf``, ``server``, ``claimed``, ``retired``, and ``member`` (the + member's directory).""" + self._assert_locked() + found: Dict[str, str] = {} + for name in self._entries(): + kind, entry_uid = self.parse_entry(name) + if kind is not None and entry_uid == uid: + # At most one entry per kind. Claiming renames a single file into place (see the + # claiming layer), so two claimed entries for one uid means the pid a reaper would + # read via claiming_pid is ambiguous; surface that rather than pick one silently. + assert kind not in found, f"duplicate {kind} entries for uid {uid}" + found[kind] = os.path.join(self.path, name) + return found + + def paths_of_kind(self, kind: str) -> List[Tuple[str, str]]: + """All ``(uid, path)`` of one state kind.""" + self._assert_locked() + return [ + (uid, os.path.join(self.path, name)) + for name in self._entries() + for entry_kind, uid in (self.parse_entry(name),) + if entry_kind == kind and uid is not None + ] + + def _entries(self) -> List[str]: + try: + return sorted(os.listdir(self.path)) + except FileNotFoundError: + return [] + + def read_json(self, path: str) -> Optional[Dict[str, Any]]: + """``None`` for files that are missing or malformed. + + Other I/O failures propagate so lifecycle callers retry instead of mistaking temporarily + unavailable state for a completed transition. + """ + self._assert_locked() + try: + with open(path, "r") as f: + data = json.load(f) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError, ValueError): + return None + return data if isinstance(data, dict) else None + + def write_json(self, path: str, data: Dict[str, Any]) -> None: + """Atomically replace ``path`` with private JSON state. + + ``path`` must be on the same filesystem as this directory. The replacement is atomic + against process failure; durability across power loss is not promised. + """ + self._assert_locked() + # Write through a temporary file so a process dying during a state transition leaves + # either the old record or the complete new one. Server entries hold an auth token, so + # the temporary and final files both remain private. + fd, temp_path = tempfile.mkstemp(prefix=self._STATE_TEMP_PREFIX, dir=self.path) + try: + try: + state_file = os.fdopen(fd, "w") + except BaseException: + # fdopen only takes ownership after it returns successfully. + os.close(fd) + raise + with state_file as f: + os.fchmod(fd, 0o600) + f.write(json.dumps(data)) + os.replace(temp_path, path) + except BaseException: + with contextlib.suppress(FileNotFoundError): + os.remove(temp_path) + raise + + def rename(self, src: str, dst: str) -> None: + self._assert_locked() + os.rename(src, dst) + + def remove(self, path: str) -> None: + self._assert_locked() + with contextlib.suppress(FileNotFoundError): + os.remove(path) + + def remove_member_dir(self, uid: str) -> None: + self._assert_locked() + shutil.rmtree(self.member_dir(uid), ignore_errors=True) + + +class ServerPool: + """Acquires, reaps, and retires members of one pool directory.""" + + # A pending marker older than this belongs to a launch that hung. Keep this above the local + # server startup timeout so a slow but healthy launch is never stopped by the janitor. + _LAUNCH_TIMEOUT_SECONDS = 180 + # A retired server still alive after the grace period is hard-killed. With a process handle, + # tracking is removed only once it is gone, replaced, or successfully signalled. PID-less + # malformed state uses the give-up age as its bounded recovery window. + _RETIRE_KILL_AFTER_SECONDS = 30 + _RETIRE_GIVE_UP_AFTER_SECONDS = 600 + # Preserve a failed launch's logs for diagnosis before collecting its unreferenced directory. + _MEMBER_DIR_GC_AGE_SECONDS = 24 * 3600 + _DEFAULT_IDLE_TIMEOUT_SECONDS = 1800 + _PROCESS_INSPECTION_TIMEOUT_SECONDS = 5 + _PROC_STAT_START_TIME_INDEX = 19 + _ATTENDANT_MODULE = "pyspark.sql.connect.local_server_pool" + + def __init__(self, directory: Optional[PoolDirectory] = None): + self._directory = directory or PoolDirectory() + + def acquire(self, master: str, opts: Dict[str, Any]) -> PoolMember: + """Claim one ready, fingerprint-matching member and top the pool back up. + + When a member is ready this returns after one janitor-claim-refill pass; on a cold pool + (first run, conf change, or all members consumed) the refill starts a full complement + and the loop waits for the first member to become ready, which costs one ordinary + cold start. The lock is released between polls so attendants can publish; launches + that die are relaunched by later passes, and only the overall deadline fails. + """ + from pyspark.sql.connect.local_server import startup_seed_conf + + seed_conf = startup_seed_conf(opts) + fingerprint = pool_fingerprint(master, seed_conf) + target = _pool_size(opts) + deadline = time.monotonic() + _ACQUIRE_TIMEOUT + while True: + with self._directory: + self.janitor() + member = self.claim(fingerprint) + self.refill(master, seed_conf, fingerprint, target) + if member is not None: + return member + if time.monotonic() >= deadline: + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": f"no pooled local server became ready within " + f"{_ACQUIRE_TIMEOUT}s; see the attendant and server logs under " + f"{self._directory.path}" + }, + ) + time.sleep(0.25) + + @classmethod + def _process_start_id(cls, pid: int) -> Optional[str]: + """An identifier for this generation of ``pid``, or ``None`` if it cannot be read. + + Linux exposes a boot id and a process start tick, which together survive PID reuse and + distinguish records left across a reboot. Other POSIX systems use ``ps``'s absolute + start time. The fallback has one-second precision but still closes the long-lived + stale-record window; signalling performs this check immediately before acting. + """ + if pid <= 0: + return None + if sys.platform.startswith("linux"): + try: + with open("/proc/sys/kernel/random/boot_id", encoding="ascii") as boot_id_file: + boot_id = boot_id_file.read().strip() + with open(f"/proc/{pid}/stat", encoding="utf-8") as stat_file: + stat = stat_file.read() + except (OSError, UnicodeError): + return None + _, separator, fields_text = stat.rpartition(") ") + fields = fields_text.split() + if not boot_id or not separator or len(fields) <= cls._PROC_STAT_START_TIME_INDEX: + return None + start_tick = fields[cls._PROC_STAT_START_TIME_INDEX] + if not (start_tick.isascii() and start_tick.isdigit()): + return None + return f"linux:{boot_id}:{start_tick}" + + env = dict(os.environ) + env["LC_ALL"] = "C" + try: + result = subprocess.run( + ["ps", "-ww", "-p", str(pid), "-o", "lstart="], + capture_output=True, + text=True, + timeout=cls._PROCESS_INSPECTION_TIMEOUT_SECONDS, + env=env, + ) + except (OSError, subprocess.SubprocessError): + return None + started = " ".join(result.stdout.split()) + return f"ps:{started}" if result.returncode == 0 and started else None + + @classmethod + def _same_process_generation( + cls, pid: Optional[int], process_start_id: Optional[str] + ) -> Optional[bool]: + """Whether ``pid`` is alive and still has its recorded process generation.""" + if pid is None or not _pid_alive(pid): + return False + if process_start_id is None: + return None + current_start_id = cls._process_start_id(pid) + if current_start_id is None: + return None + return current_start_id == process_start_id + + @classmethod + def _same_server_instance(cls, pid: int, process_start_id: str) -> Optional[bool]: + """Whether ``pid`` is still the recorded Connect server process generation.""" + same_generation = cls._same_process_generation(pid, process_start_id) + if same_generation is not True: + return same_generation + return _is_local_connect_server(pid) + + @staticmethod + def _signal(pid: int, sig: int) -> bool: + """Best-effort signal; ``False`` when the process is already gone or not ours.""" + if pid <= 0: + return False + try: + os.kill(pid, sig) + return True + except (OSError, OverflowError): + return False + + @classmethod + def _signal_server(cls, pid: int, process_start_id: str, sig: int) -> bool: + """Signal only the recorded generation of the managed Connect server.""" + return cls._same_server_instance(pid, process_start_id) is True and cls._signal(pid, sig) + + @classmethod + def _idle_timeout(cls) -> int: + """Seconds an unclaimed member may sit before it is retired. + + Zero or a negative value disables idle retirement. Read the environment wherever + reaping runs so clients and attendants use the same source of truth. + """ + try: + return int(os.environ["SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT"]) + except (KeyError, ValueError): + return cls._DEFAULT_IDLE_TIMEOUT_SECONDS + + @classmethod + def _is_pool_attendant(cls, pid: int, uid: str) -> Optional[bool]: + """Whether ``pid`` is still the pool attendant recorded for ``uid``. + + Returns ``None`` when the process cannot be inspected. A stale pending record can + outlive its attendant long enough for the pid to be reused, so liveness alone is not + sufficient before a janitor signals it. + """ + command = _process_command(pid) + if command is None: + return None + args = command.split() + try: + module_index = args.index(cls._ATTENDANT_MODULE) + uid_index = args.index("--uid") + except ValueError: + return False + return ( + module_index > 0 + and args[module_index - 1] == "-m" + and "--attend" in args + and uid_index + 1 < len(args) + and args[uid_index + 1] == uid + ) + + @staticmethod + def _signal_attendant_group(pid: int, sig: int) -> bool: + """Signal a detached attendant and the launch subprocesses in its process group.""" + if pid <= 0 or pid == os.getpgrp(): + return False + try: + if os.getpgid(pid) != pid: + return False + os.killpg(pid, sig) + return True + except (OSError, OverflowError): + return False + + def claim(self, fingerprint: str) -> Optional[PoolMember]: + """Claim the oldest usable member with this fingerprint, or ``None``. The rename to + ``claimed-<pid>-<uid>.json`` marks the member as owned by this process; the reaping + rules use that pid to retire members whose client died without releasing them. The + caller must hold the directory lock so selection and rename form one transition. + + Ordering is by ``created``, a wall-clock ``time.time()`` reading. It is comparable + across the independent processes that publish members, which ``time.monotonic()`` is + not, at the cost that a backward clock step (NTP, suspend/resume) can perturb the order. + Ties break by the stable ``sorted()`` over the sorted directory listing, so the order is + well defined but only approximately FIFO, not guaranteed. + + ``is_usable`` runs under the held lock and does blocking process inspection and network + I/O -- potentially one ``ps`` and up to a 0.5s connect for each candidate. The candidate + count is bounded by ``spark.local.connect.pool.size``, which is user-tunable, so a large + pool widens the window the lock is held; the reaping rules keep stale members from + accumulating without bound.""" + candidates = [] + for uid, path in self._directory.paths_of_kind("server"): + data = self._directory.read_json(path) + member = PoolMember.from_data(data) if data is not None else None + if member is not None and member.fingerprint == fingerprint: + assert data is not None + candidates.append((member, uid, path, data)) + candidates.sort(key=lambda c: c[0].created) + for member, uid, path, data in candidates: + if not member.is_usable(): + continue # left for the reaping rules to retire + client_process_start_id = self._process_start_id(os.getpid()) + if client_process_start_id is not None: + claimed_data = dict(data) + claimed_data["client_process_start_id"] = client_process_start_id + # Persist ownership before the atomic rename. If this process dies between the + # write and rename, the extra field is harmless on the still-unclaimed record. + self._directory.write_json(path, claimed_data) + claim_path = self._directory.claimed_path(os.getpid(), uid) + self._directory.rename(path, claim_path) + member.claim_path = claim_path + return member + return None + + def refill(self, master: str, seed_conf: Dict[str, Any], fingerprint: str, target: int) -> None: + """Launch members until ready or in-flight ones with this fingerprint reach + ``target``. Running under the directory lock is what makes concurrent cold starters + share one complement of launches instead of each spawning their own.""" + available = 0 + for kind in ("server", "pending"): + for _, path in self._directory.paths_of_kind(kind): + data = self._directory.read_json(path) + if data is not None and data.get("fingerprint") == fingerprint: + available += 1 + for _ in range(target - available): + MemberAttendant.spawn(self._directory, master, seed_conf, fingerprint) + + def janitor(self) -> None: + """Reap leftovers of launches, clients, and attendants that died uncleanly. Every + rule is idempotent, so successive passes from any process are safe.""" + for uid in self._directory.uids(): + self.reap(uid) + + def reap(self, uid: str) -> bool: + """Apply the reaping rules to one member; ``True`` when nothing of it remains. + Shared by the janitor (all members) and by each attendant supervising its own member. + """ + states = self._directory.states(uid) + if "conf" in states and "pending" not in states: + # A later state proves the attendant consumed the seed. A conf-only record can be + # left if its spawning client dies before starting or recording the attendant; use + # the launch deadline to avoid accumulating those records forever. + later_state = any(kind in states for kind in ("server", "claimed", "retired")) + try: + conf_expired = ( + time.time() - os.path.getmtime(states["conf"]) > self._LAUNCH_TIMEOUT_SECONDS + ) + except FileNotFoundError: + conf_expired = True + if later_state or conf_expired: + self._directory.remove(states["conf"]) + states = self._directory.states(uid) + had_retired = "retired" in states + if "pending" in states: + self._reap_pending(uid, states["pending"]) + states = self._directory.states(uid) + if "server" in states: + self._reap_server(uid, states["server"]) + states = self._directory.states(uid) + if "claimed" in states: + self._reap_claimed(uid, states["claimed"]) + states = self._directory.states(uid) + if had_retired and "retired" in states: + self._reap_retired(uid, states["retired"]) + + remaining = self._directory.states(uid) + if set(remaining) == {"member"}: + # Nothing references the member directory anymore. The age gate keeps the logs + # of a freshly failed launch around long enough to be looked at. + try: + expired = ( + time.time() - os.path.getmtime(remaining["member"]) + > self._MEMBER_DIR_GC_AGE_SECONDS + ) + except FileNotFoundError: + expired = True + if expired: + self._directory.remove_member_dir(uid) + remaining = self._directory.states(uid) + return not remaining + + def _reap_pending(self, uid: str, path: str) -> None: + """A launch whose attendant died or hung: kill the attendant and whatever server + spark-daemon.sh may have recorded for it, and withdraw the launch's bookkeeping so + refills stop counting it.""" + data = self._directory.read_json(path) + pending = PendingState.from_data(data) + parsed_pid = pending.attendant_pid if pending is not None else None + created = pending.created if pending is not None else None + if pending is None and data is not None: + # Preserve an independently valid pid when another field is corrupt. + parsed_pid = PendingState.attendant_pid_from_data(data) + age = time.time() - created if created is not None else self._LAUNCH_TIMEOUT_SECONDS + 1 + attendant_pid = parsed_pid if parsed_pid is not None else -1 + attendant_alive = _pid_alive(attendant_pid) + if not attendant_alive: + self.abort_launch(uid) + elif age > self._LAUNCH_TIMEOUT_SECONDS: + is_attendant = self._is_pool_attendant(attendant_pid, uid) + if is_attendant is None: + return + if is_attendant and not self._signal_attendant_group(attendant_pid, signal.SIGKILL): + # Keep the record when an attendant that still appears live could not be + # stopped; a later pass can retry without losing its only process handle. + if _pid_alive(attendant_pid): + return + self.abort_launch(uid) + + def abort_launch(self, uid: str) -> None: + """Withdraw a failed launch and retire any server it started before failing.""" + states = self._directory.states(uid) + pending_path = states.get("pending") + server_path = states.get("server") + if server_path is not None: + data = self._directory.read_json(server_path) + server_pid, process_start_id = self._recover_server_handle(uid, data) + else: + server_pid = self._recorded_daemon_pid(uid) + process_start_id = None + retirement_source = server_path or pending_path + retired_source = False + if server_pid is not None and retirement_source is not None: + # Keep shutdown state so a half-started JVM that ignores SIGTERM is escalated. + self._retire(retirement_source, server_pid, process_start_id) + retired_source = True + if pending_path is not None and (not retired_source or pending_path != retirement_source): + self._directory.remove(pending_path) + self._directory.remove(self._directory.conf_path(uid)) + + def _reap_server(self, uid: str, path: str) -> None: + """A ready member that is unusable (dead, unreachable, version-mismatched after an + upgrade, or an unreadable record) or has sat unclaimed past the idle timeout: retire + it.""" + data = self._directory.read_json(path) + member = PoolMember.from_data(data) if data is not None else None + server_pid, process_start_id = self._recover_server_handle(uid, data) + idle = self._idle_timeout() + expired = member is not None and idle > 0 and time.time() - member.created > idle + if member is None or expired or not member.is_usable(): + self._retire(path, server_pid, process_start_id) + + def _reap_claimed(self, uid: str, path: str) -> None: + """A claimed member whose client died without releasing it (e.g. SIGKILL), or whose + server died under its client: retire it. Claims of this live process are its own.""" + data = self._directory.read_json(path) + server_pid, process_start_id = self._recover_server_handle(uid, data) + client_pid = self._directory.claiming_pid(path) + client_process_start_id = PoolMember.client_process_start_id_from_data(data) + if client_process_start_id is None: + # Records written by older clients have no process generation. Preserve their + # liveness-only behavior rather than risking retirement of a live claim. + client_alive = client_pid == os.getpid() or _pid_alive(client_pid) + else: + client_alive = ( + self._same_process_generation(client_pid, client_process_start_id) is not False + ) + if not client_alive or self._same_process_generation(server_pid, process_start_id) is False: + self._retire(path, server_pid, process_start_id) + + def _reap_retired(self, uid: str, path: str) -> None: + """A retiring member: drop it once its server is gone, hard-kill the server if it + hangs in shutdown, and stop tracking a record whose process generation was reused.""" + data = self._directory.read_json(path) + record_pid = RetiredState.pid_from_data(data) + record_process_start_id = RetiredState.process_start_id_from_data(data) + retired = RetiredState.retired_from_data(data) + signalled = RetiredState.signalled_from_data(data) + server_pid = self._recover_server_pid(uid, record_pid) + # A generation id is meaningful only with the pid from the same record. The daemon pid + # file has no companion generation id, so never pair a recovered daemon pid with + # unrelated record data. + process_start_id = record_process_start_id if record_pid is not None else None + now = time.time() + if server_pid is None: + # A daemon pid may appear after a partial publication, so retain the state for one + # recovery window. After that there is no process handle left to act on or observe. + if retired is None or retired > now: + self._directory.write_json( + path, + { + "retired": now, + "signalled": False, + }, + ) + elif now - retired > self._RETIRE_GIVE_UP_AFTER_SECONDS: + self._remove_retired(uid, path) + return + if not _pid_alive(server_pid): + self._remove_retired(uid, path) + return + if process_start_id is None: + # A pid without its persisted generation cannot safely be adopted: the pid may have + # been recycled to an unrelated local Connect server. Retain the handle until it dies + # rather than synthesizing authority to signal its current owner. + return + current_start_id = self._process_start_id(server_pid) + if current_start_id is None: + return + if current_start_id != process_start_id: + # The original server is gone and its PID now belongs to another process. Forget the + # stale record without signalling the new owner. + self._remove_retired(uid, path) + return + if retired is None or retired > now: + # A crash while _retire rewrites the atomically renamed state can leave its old + # payload. Restore a shutdown clock while preserving its process identity. + self._directory.write_json( + path, + RetiredState( + server_pid, + process_start_id, + now, + signalled=signalled is True, + ).as_data(), + ) + return + age = now - retired + if age > self._RETIRE_GIVE_UP_AFTER_SECONDS: + # Drop tracking only after proving the process was replaced or successfully issuing + # the hard kill. Transient inspection or signalling failures remain retryable. + is_server = self._same_server_instance(server_pid, process_start_id) + if is_server is None or (is_server and not self._signal(server_pid, signal.SIGKILL)): + return + self._remove_retired(uid, path) + elif age > self._RETIRE_KILL_AFTER_SECONDS: + is_server = self._same_server_instance(server_pid, process_start_id) + if is_server is False: + self._remove_retired(uid, path) + elif is_server is True: + self._signal(server_pid, signal.SIGKILL) + elif signalled is not True: + # Retry a SIGTERM that was not confirmed during retirement. Persist success without + # refreshing the retirement clock so repeated passes remain free and escalation is + # still measured from the original transition. + if self._signal_server(server_pid, process_start_id, signal.SIGTERM): + self._directory.write_json( + path, + RetiredState( + server_pid, + process_start_id, + retired, + signalled=True, + ).as_data(), + ) + + def _remove_retired(self, uid: str, path: str) -> None: + self._directory.remove(path) + self._directory.remove_member_dir(uid) + + def _retire( + self, + state_path: str, + server_pid: Optional[int], + process_start_id: Optional[str], + ) -> None: + """Move a member into the retired state: signal its server and track the shutdown so + :meth:`_reap_retired` can escalate if the JVM hangs.""" + _, uid = self._directory.parse_entry(os.path.basename(state_path)) + assert uid is not None + signalled = False + if server_pid is not None and process_start_id is not None: + signalled = self._signal_server(server_pid, process_start_id, signal.SIGTERM) + retired_path = self._directory.retired_path(uid) + # Rename instead of removing the old state so a crash cannot leave a live server with + # no state. If rewriting is interrupted, _reap_retired preserves the recoverable pid. + self._directory.rename(state_path, retired_path) + retired_data: Dict[str, Any] = { + "retired": time.time(), + "signalled": signalled, + } + if server_pid is not None: + retired_data["pid"] = server_pid + if process_start_id is not None: + retired_data["process_start_id"] = process_start_id + self._directory.write_json(retired_path, retired_data) + + def _recover_server_handle( + self, uid: str, data: Optional[Dict[str, Any]] + ) -> Tuple[Optional[int], Optional[str]]: + """Recover a pid and only the process identity paired with that pid's record.""" + record_pid = PoolMember.pid_from_data(data) + if record_pid is None: + return self._recorded_daemon_pid(uid), None + return record_pid, PoolMember.process_start_id_from_data(data) + + def _recover_server_pid(self, uid: str, record_pid: Optional[int]) -> Optional[int]: + """Use a record pid when present, otherwise fall back to the daemon pid file. + + Full member validation intentionally rejects corrupt records, including out-of-range + timestamps. Its independently valid pid remains paired with the record's generation id; + a daemon pid has no generation id and is used only to retain state while it remains live. + """ + return record_pid if record_pid is not None else self._recorded_daemon_pid(uid) + + def _recorded_daemon_pid(self, uid: str) -> Optional[int]: + """The positive server pid recorded by spark-daemon.sh, if readable.""" + discovery = Discovery(os.path.join(self._directory.member_dir(uid), "connect-local.json")) + return _PoolStateRecord._positive_pid(discovery.daemon_pid()) + + def release(self, member: PoolMember) -> None: + """Retire this process's claimed member; the shutdown completes in the background, + watched by the member's attendant with the janitor as backstop. + + This method acquires the pool-directory lock and must not be called while the same pool + directory is already locked, including through a different ``PoolDirectory`` instance. + """ + assert member.claim_path is not None + kind, uid = self._directory.parse_entry(os.path.basename(member.claim_path)) + assert kind == "claimed" and uid is not None + if self._directory.claiming_pid(member.claim_path) != os.getpid(): + # A forked child inherits module globals and atexit handlers, but it must not retire + # the server still claimed by its parent process. + return + with self._directory: + # A janitor or concurrent purge may already have moved or removed the claim. + # Release is idempotent with respect to that completed lifecycle transition. + if self._directory.states(uid).get("claimed") == member.claim_path: + self._retire(member.claim_path, member.pid, member.process_start_id) + + def purge(self) -> int: + """Force-stop every member -- ready, in-flight, or claimed -- and empty the pool + directory; the escape hatch back to a clean slate. Returns the number of processes + signalled. SIGKILL rather than SIGTERM because nothing tracks a member once its + state files are gone, so a shutdown that hangs would leak. Supervising attendants + hold no state file; they notice the emptied directory and exit on their own.""" + signalled_pids: List[int] = [] + + def hard_kill_server(pid: Optional[int], process_start_id: Optional[str] = None) -> None: + if pid is None or pid in signalled_pids: + return + if process_start_id is not None: + signalled = self._signal_server(pid, process_start_id, signal.SIGKILL) + else: + # Malformed and half-published states can have only spark-daemon.sh's pid. + # Purge is explicitly destructive, but still verify the current command before + # signalling rather than trusting a potentially reused numeric pid alone. + signalled = _is_local_connect_server(pid) is True and self._signal( + pid, signal.SIGKILL + ) + if signalled: + signalled_pids.append(pid) + + def hard_kill_attendant(pid: Optional[int], uid: str) -> None: + if ( + pid is not None + and pid not in signalled_pids + and self._is_pool_attendant(pid, uid) is True + and self._signal_attendant_group(pid, signal.SIGKILL) + ): + signalled_pids.append(pid) + + with self._directory: + uids = self._directory.uids() + # Iterate entries directly by kind rather than through states(uid): duplicate + # claimed records violate the normal state invariant, but purge is the corruption + # escape hatch and must still clear and stop every recoverable member. + for kind in ("pending", "conf", "server", "claimed", "retired"): + for uid, path in self._directory.paths_of_kind(kind): + data = self._directory.read_json(path) + if kind == "pending": + pending = PendingState.from_data(data) + pid = ( + pending.attendant_pid + if pending is not None + else PendingState.attendant_pid_from_data(data) + ) + hard_kill_attendant(pid, uid) + elif kind in ("server", "claimed"): + pid, process_start_id = self._recover_server_handle(uid, data) + hard_kill_server(pid, process_start_id) + elif kind == "retired": + record_pid = RetiredState.pid_from_data(data) + pid = self._recover_server_pid(uid, record_pid) + process_start_id = ( + RetiredState.process_start_id_from_data(data) + if record_pid is not None + else None + ) + hard_kill_server(pid, process_start_id) + self._directory.remove(path) + for uid in uids: + hard_kill_server(self._recorded_daemon_pid(uid)) + self._directory.remove_member_dir(uid) + return len(signalled_pids) + + +# The member this client process has claimed, if any. Module-level so the session's stop +# callback and atexit share one idempotent release path. +_claimed_member: Optional[PoolMember] = None +_release_registered = False + + +def acquire_pooled_local_connect_server(master: str, opts: Dict[str, Any]) -> str: + """Claim a local Connect server not previously assigned to an application run. + + Returns the ``sc://host:port`` endpoint and sets ``SPARK_CONNECT_AUTHENTICATE_TOKEN`` so + the client authenticates against that server. Only reached for a ``local`` master when the + pool opt-in is set; see ``SparkSession.getOrCreate`` in ``pyspark.sql.session``. + """ + global _claimed_member, _release_registered + if os.name != "posix": + raise PySparkRuntimeError( + errorClass="LOCAL_CONNECT_SERVER_START_FAILED", + messageParameters={ + "reason": "spark.local.connect.pool relies on the POSIX scripts under sbin/; " + "on this platform start a server manually (sbin/start-connect-server.sh) and " + 'connect with .remote("sc://...")' + }, + ) + # getOrCreate() may be re-entered while this process already holds a live claimed member + # (the connect layer then returns the existing session); claiming again would strand a + # second server. + if _claimed_member is not None and _pid_alive(_claimed_member.pid): + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = _claimed_member.token + return _claimed_member.url + + member = ServerPool().acquire(master, opts) + _claimed_member = member + if not _release_registered: + # A client that exits without stopping its session still releases its member; the + # member's attendant and the janitor cover clients that die uncleanly. + atexit.register(release_pooled_local_connect_server) + _release_registered = True + os.environ["SPARK_CONNECT_AUTHENTICATE_TOKEN"] = member.token + return member.url + + +def release_pooled_local_connect_server() -> None: + """Retire this process's claimed pooled server; safe to call when there is none. The + server winds down in the background while this client moves on.""" + global _claimed_member + member = _claimed_member + if member is not None: + assert member.claim_path is not None + directory = PoolDirectory(os.path.dirname(member.claim_path)) + ServerPool(directory).release(member) + if _claimed_member is member: + _claimed_member = None + + +def purge_local_connect_pool() -> int: + """See :meth:`ServerPool.purge`. Also available as + ``python -m pyspark.sql.connect.local_server_pool --purge``.""" + return ServerPool().purge() + + +class MemberAttendant: + """Boots one pool member, publishes it, and supervises it until it is gone. + + Runs detached from the spawning client (``--attend``). Every phase is appended to + ``member-<uid>/attendant.log`` so a member that misbehaved can be debugged after the + fact. A failed boot reaps whatever half-started server spark-daemon.sh recorded and + withdraws the launch's bookkeeping, so refills stop counting it. + """ + + @classmethod + def spawn( + cls, + directory: PoolDirectory, + master: str, + seed_conf: Dict[str, Any], + fingerprint: str, + ) -> None: + """Start one detached attendant and publish its in-flight launch. + + Callers hold ``directory``'s lock, so another refill sees the pending marker before it + can start a duplicate launch for the same pool complement. + """ + uid = uuid.uuid4().hex[:12] + directory.write_json(directory.conf_path(uid), seed_conf) + cmd = [ + sys.executable, + "-m", + "pyspark.sql.connect.local_server_pool", + "--attend", + "--pool-dir", + directory.path, + "--uid", + uid, + "--master", + master, + "--fingerprint", + fingerprint, + ] + env = dict(os.environ) + # The attendant must neither see this client's Connect mode nor inherit its auth + # token: each member gets its own token from LocalConnectServer. + for var in ( + "SPARK_REMOTE", + "SPARK_LOCAL_REMOTE", + "SPARK_CONNECT_MODE_ENABLED", + "SPARK_CONNECT_AUTHENTICATE_TOKEN", + ): + env.pop(var, None) + try: + proc = subprocess.Popen( + cmd, + env=env, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + except OSError: + directory.remove(directory.conf_path(uid)) + raise + directory.write_json( + directory.pending_path(uid), + {"attendant_pid": proc.pid, "created": time.time(), "fingerprint": fingerprint}, + ) + + def __init__(self, directory: PoolDirectory, uid: str, master: str, fingerprint: str): + self._directory = directory + self._pool = ServerPool(directory) + self._uid = uid + self._master = master + self._fingerprint = fingerprint + self._member_dir = directory.member_dir(uid) + + def run(self) -> int: + os.makedirs(self._member_dir, mode=0o700, exist_ok=True) + member = self._boot() + if member is None: + return 1 + self._log("supervising") + self._supervise(member.pid) + self._log("done") + return 0 + + def _log(self, message: str) -> None: + with open(os.path.join(self._member_dir, "attendant.log"), "a") as f: + f.write(f"{time.time():.3f} {message}\n") + + def _boot(self) -> Optional[PoolMember]: + """Start the server on a fresh ephemeral port and publish it as ready-to-claim. + Publishing consumes the launch's pending marker and conf seed: from that moment the + member counts as a server, and this process's job shifts to supervising it.""" + from pyspark.sql.connect.local_server import Discovery, LocalConnectServer + + with self._directory: + seed_conf = self._directory.read_json(self._directory.conf_path(self._uid)) or {} + discovery = Discovery(os.path.join(self._member_dir, "connect-local.json")) + self._log(f"booting master={self._master}") + try: + with discovery: + server = LocalConnectServer(discovery) + server.start( + self._master, + {}, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + data = discovery.load() + assert data is not None # launch() saved it + process_start_id = self._pool._process_start_id(data["pid"]) + if process_start_id is None: + raise RuntimeError(f"could not identify launched server pid {data['pid']}") + data["process_start_id"] = process_start_id + except Exception as e: + self._log(f"boot failed: {e!r}") + with self._directory: + self._pool.abort_launch(self._uid) + return None + data["fingerprint"] = self._fingerprint + data["created"] = time.time() + with self._directory: + self._directory.write_json(self._directory.server_path(self._uid), data) + self._directory.remove(self._directory.pending_path(self._uid)) + self._directory.remove(self._directory.conf_path(self._uid)) + self._log(f"published pid={data['pid']} port={data['port']}") + return PoolMember(data) + + def _supervise(self, server_pid: int) -> None: + """Watch this member until nothing of it remains, applying the pool's reaping rules. + + This is what lets an idle machine drain to zero servers with no further Spark run: the + idle timeout, dead-client cleanup, and hard-kill escalation all fire from here even + when no client ever runs again. + """ + while True: + with self._directory: + if self._pool.reap(self._uid): + return + states = set(self._directory.states(self._uid)) + if states <= {"member"} and not _pid_alive(server_pid): + # A purge or another process's janitor already tore the member down; only + # the young member directory remains. + self._directory.remove_member_dir(self._uid) + return + time.sleep(2) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Manage the opt-in pool of single-use local Spark Connect servers " + "(spark.local.connect.pool)." + ) + parser.add_argument( + "--purge", + action="store_true", + help="force-stop every pool member and empty the pool directory", + ) + # Internal entry point spawned by ServerPool. + parser.add_argument("--attend", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--pool-dir", help=argparse.SUPPRESS) + parser.add_argument("--uid", help=argparse.SUPPRESS) + parser.add_argument("--master", help=argparse.SUPPRESS) + parser.add_argument("--fingerprint", help=argparse.SUPPRESS) + args = parser.parse_args() + + if args.purge: + print(f"Signalled {purge_local_connect_pool()} pool process(es).") + elif args.attend: + attendant = MemberAttendant( + PoolDirectory(args.pool_dir), args.uid, args.master, args.fingerprint + ) + sys.exit(attendant.run()) + else: + parser.print_help(sys.stderr) + sys.exit(2) + + +if __name__ == "__main__": + main() diff --git a/python/pyspark/sql/connect/logging.py b/python/pyspark/sql/connect/logging.py index 099193fd7ce45..6b8de0a400ec1 100644 --- a/python/pyspark/sql/connect/logging.py +++ b/python/pyspark/sql/connect/logging.py @@ -17,10 +17,11 @@ import logging -from pyspark.logger import PySparkLogger import os from typing import Optional +from pyspark.logger import PySparkLogger + __all__ = ["configureLogging", "getLogLevel"] diff --git a/python/pyspark/sql/connect/merge.py b/python/pyspark/sql/connect/merge.py index 9464e33c014c4..6f2a864852cb9 100644 --- a/python/pyspark/sql/connect/merge.py +++ b/python/pyspark/sql/connect/merge.py @@ -15,7 +15,7 @@ # limitations under the License. # import sys -from typing import Dict, Optional, TYPE_CHECKING, Callable +from typing import TYPE_CHECKING, Callable, Dict, Optional from pyspark.sql.connect import proto from pyspark.sql.connect.column import Column @@ -228,8 +228,9 @@ def delete(self) -> "MergeIntoWriter": def _test() -> None: import doctest import os - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.merge + from pyspark.sql import SparkSession as PySparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/connect/observation.py b/python/pyspark/sql/connect/observation.py index 8e16e2e94a663..3c765f1774212 100644 --- a/python/pyspark/sql/connect/observation.py +++ b/python/pyspark/sql/connect/observation.py @@ -14,19 +14,19 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Dict, Optional import uuid +from typing import Any, Dict, Optional +import pyspark.sql.connect.plan as plan from pyspark.errors import ( - PySparkTypeError, - PySparkValueError, IllegalArgumentException, PySparkAssertionError, + PySparkTypeError, + PySparkValueError, ) from pyspark.sql.column import Column from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.observation import Observation as PySparkObservation -import pyspark.sql.connect.plan as plan __all__ = ["Observation"] @@ -92,11 +92,12 @@ def get(self) -> Dict[str, Any]: def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.observation + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.observation.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/plan.py b/python/pyspark/sql/connect/plan.py index 51f1175166631..297f20e2de054 100644 --- a/python/pyspark/sql/connect/plan.py +++ b/python/pyspark/sql/connect/plan.py @@ -17,52 +17,50 @@ # mypy: disable-error-code="operator" -from pyspark.resource import ResourceProfile - +import functools +import json +import pickle +from inspect import isclass, signature +from threading import Lock from typing import ( + TYPE_CHECKING, Any, + Dict, Iterator, List, + Mapping, Optional, - Type, Sequence, + Tuple, + Type, Union, cast, - TYPE_CHECKING, - Mapping, - Dict, - Tuple, ) -import functools -import json -import pickle -from threading import Lock -from inspect import signature, isclass import pyarrow as pa -from pyspark.serializers import CloudPickleSerializer -from pyspark.storagelevel import StorageLevel -from pyspark.sql.types import DataType, StructType - import pyspark.sql.connect.proto as proto -from pyspark.sql.column import Column -from pyspark.sql.connect.logging import logger -from pyspark.sql.connect.proto import base_pb2 as spark_dot_connect_dot_base__pb2 -from pyspark.sql.connect.conversion import storage_level_to_proto -from pyspark.sql.connect.expressions import Expression, SubqueryExpression -from pyspark.sql.connect.types import pyspark_types_to_proto_types, UnparsedDataType from pyspark.errors import ( AnalysisException, - PySparkValueError, PySparkPicklingError, + PySparkValueError, ) +from pyspark.resource import ResourceProfile +from pyspark.serializers import CloudPickleSerializer +from pyspark.sql.column import Column +from pyspark.sql.connect.conversion import storage_level_to_proto +from pyspark.sql.connect.expressions import Expression, SubqueryExpression +from pyspark.sql.connect.logging import logger +from pyspark.sql.connect.proto import base_pb2 as spark_dot_connect_dot_base__pb2 +from pyspark.sql.connect.types import UnparsedDataType, pyspark_types_to_proto_types +from pyspark.sql.types import DataType, StructType +from pyspark.storagelevel import StorageLevel if TYPE_CHECKING: from pyspark.sql.connect.client import SparkConnectClient - from pyspark.sql.connect.udf import UserDefinedFunction from pyspark.sql.connect.observation import Observation from pyspark.sql.connect.session import SparkSession + from pyspark.sql.connect.udf import UserDefinedFunction class LogicalPlan: @@ -784,8 +782,16 @@ def __del__(self) -> None: request_serializer=request_serializer, response_deserializer=response_deserializer, ) - metadata = session.client._builder.metadata() - channel(req, metadata=metadata) # type: ignore[arg-type] + metadata = session.client._execute_plan_metadata(req.operation_id) + # Bound this blocking call with a client-side deadline. It is issued from a + # finalizer with no other timeout at any layer; without a deadline it can + # block forever if the response is never delivered, which stalls the + # foreachBatch Connect handshake (the Python worker never sends its + # completion signal and the driver JVM blocks on the per-batch read). A + # timeout here is non-fatal: the eviction is best effort and the server + # performs it independently, so on timeout we log and move on. + timeout = session.client._rpc_deadlines.release_relation + channel(req, metadata=metadata, timeout=timeout) # type: ignore[arg-type] except Exception as e: logger.warning(f"RemoveRemoteCachedRelation failed with exception: {e}.") diff --git a/python/pyspark/sql/connect/proto/base_pb2.py b/python/pyspark/sql/connect/proto/base_pb2.py index a77c61ca6d2b4..ee3ddb48ba5c6 100644 --- a/python/pyspark/sql/connect/proto/base_pb2.py +++ b/python/pyspark/sql/connect/proto/base_pb2.py @@ -46,7 +46,7 @@ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"\xe3\x03\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.CommandH\x00R\x07\x63ommand\x12\\\n\x14\x63ompressed_operation\x18\x03 \x01(\x0b\x32\'.spark.connect.Plan.CompressedOperationH\x00R\x13\x63ompressedOperation\x1a\x8e\x02\n\x13\x43ompressedOperation\x12\x12\n\x04\x64\x61ta\x18\x01 \x01(\x0cR\x04\x64\x61ta\x12G\n\x07op_type\x18\x02 \x01(\x0e\x32..spark.connect.Plan.CompressedOperation.OpTypeR\x06opType\x12L\n\x11\x63ompression_codec\x18\x03 \x01(\x0e\x32\x1f.spark.connect.CompressionCodecR\x10\x63ompressionCodec"L\n\x06OpType\x12\x17\n\x13OP_TYPE_UNSPECIFIED\x10\x00\x12\x14\n\x10OP_TYPE_RELATION\x10\x01\x12\x13\n\x0fOP_TYPE_COMMAND\x10\x02\x42\t\n\x07op_type"z\n\x0bUserContext\x12\x17\n\x07user_id\x18\x01 \x01(\tR\x06userId\x12\x1b\n\tuser_name\x18\x02 \x01(\tR\x08userName\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions"\xf5\x14\n\x12\x41nalyzePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x11 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x02R\nclientType\x88\x01\x01\x12\x42\n\x06schema\x18\x04 \x01(\x0b\x32(.spark.connect.AnalyzePlanRequest.SchemaH\x00R\x06schema\x12\x45\n\x07\x65xplain\x18\x05 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.ExplainH\x00R\x07\x65xplain\x12O\n\x0btree_string\x18\x06 \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.TreeStringH\x00R\ntreeString\x12\x46\n\x08is_local\x18\x07 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.IsLocalH\x00R\x07isLocal\x12R\n\x0cis_streaming\x18\x08 \x01(\x0b\x32-.spark.connect.AnalyzePlanRequest.IsStreamingH\x00R\x0bisStreaming\x12O\n\x0binput_files\x18\t \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.InputFilesH\x00R\ninputFiles\x12U\n\rspark_version\x18\n \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SparkVersionH\x00R\x0csparkVersion\x12I\n\tddl_parse\x18\x0b \x01(\x0b\x32*.spark.connect.AnalyzePlanRequest.DDLParseH\x00R\x08\x64\x64lParse\x12X\n\x0esame_semantics\x18\x0c \x01(\x0b\x32/.spark.connect.AnalyzePlanRequest.SameSemanticsH\x00R\rsameSemantics\x12U\n\rsemantic_hash\x18\r \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SemanticHashH\x00R\x0csemanticHash\x12\x45\n\x07persist\x18\x0e \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.PersistH\x00R\x07persist\x12K\n\tunpersist\x18\x0f \x01(\x0b\x32+.spark.connect.AnalyzePlanRequest.UnpersistH\x00R\tunpersist\x12_\n\x11get_storage_level\x18\x10 \x01(\x0b\x32\x31.spark.connect.AnalyzePlanRequest.GetStorageLevelH\x00R\x0fgetStorageLevel\x12M\n\x0bjson_to_ddl\x18\x12 \x01(\x0b\x32+.spark.connect.AnalyzePlanRequest.JsonToDDLH\x00R\tjsonToDdl\x1a\x31\n\x06Schema\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\xbb\x02\n\x07\x45xplain\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12X\n\x0c\x65xplain_mode\x18\x02 \x01(\x0e\x32\x35.spark.connect.AnalyzePlanRequest.Explain.ExplainModeR\x0b\x65xplainMode"\xac\x01\n\x0b\x45xplainMode\x12\x1c\n\x18\x45XPLAIN_MODE_UNSPECIFIED\x10\x00\x12\x17\n\x13\x45XPLAIN_MODE_SIMPLE\x10\x01\x12\x19\n\x15\x45XPLAIN_MODE_EXTENDED\x10\x02\x12\x18\n\x14\x45XPLAIN_MODE_CODEGEN\x10\x03\x12\x15\n\x11\x45XPLAIN_MODE_COST\x10\x04\x12\x1a\n\x16\x45XPLAIN_MODE_FORMATTED\x10\x05\x1aZ\n\nTreeString\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12\x19\n\x05level\x18\x02 \x01(\x05H\x00R\x05level\x88\x01\x01\x42\x08\n\x06_level\x1a\x32\n\x07IsLocal\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x36\n\x0bIsStreaming\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x35\n\nInputFiles\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x0e\n\x0cSparkVersion\x1a)\n\x08\x44\x44LParse\x12\x1d\n\nddl_string\x18\x01 \x01(\tR\tddlString\x1ay\n\rSameSemantics\x12\x34\n\x0btarget_plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\ntargetPlan\x12\x32\n\nother_plan\x18\x02 \x01(\x0b\x32\x13.spark.connect.PlanR\totherPlan\x1a\x37\n\x0cSemanticHash\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x97\x01\n\x07Persist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x45\n\rstorage_level\x18\x02 \x01(\x0b\x32\x1b.spark.connect.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level\x1an\n\tUnpersist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x1f\n\x08\x62locking\x18\x02 \x01(\x08H\x00R\x08\x62locking\x88\x01\x01\x42\x0b\n\t_blocking\x1a\x46\n\x0fGetStorageLevel\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x1a,\n\tJsonToDDL\x12\x1f\n\x0bjson_string\x18\x01 \x01(\tR\njsonStringB\t\n\x07\x61nalyzeB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xca\x0e\n\x13\x41nalyzePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x0f \x01(\tR\x13serverSideSessionId\x12\x43\n\x06schema\x18\x02 \x01(\x0b\x32).spark.connect.AnalyzePlanResponse.SchemaH\x00R\x06schema\x12\x46\n\x07\x65xplain\x18\x03 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.ExplainH\x00R\x07\x65xplain\x12P\n\x0btree_string\x18\x04 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.TreeStringH\x00R\ntreeString\x12G\n\x08is_local\x18\x05 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.IsLocalH\x00R\x07isLocal\x12S\n\x0cis_streaming\x18\x06 \x01(\x0b\x32..spark.connect.AnalyzePlanResponse.IsStreamingH\x00R\x0bisStreaming\x12P\n\x0binput_files\x18\x07 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.InputFilesH\x00R\ninputFiles\x12V\n\rspark_version\x18\x08 \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SparkVersionH\x00R\x0csparkVersion\x12J\n\tddl_parse\x18\t \x01(\x0b\x32+.spark.connect.AnalyzePlanResponse.DDLParseH\x00R\x08\x64\x64lParse\x12Y\n\x0esame_semantics\x18\n \x01(\x0b\x32\x30.spark.connect.AnalyzePlanResponse.SameSemanticsH\x00R\rsameSemantics\x12V\n\rsemantic_hash\x18\x0b \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SemanticHashH\x00R\x0csemanticHash\x12\x46\n\x07persist\x18\x0c \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.PersistH\x00R\x07persist\x12L\n\tunpersist\x18\r \x01(\x0b\x32,.spark.connect.AnalyzePlanResponse.UnpersistH\x00R\tunpersist\x12`\n\x11get_storage_level\x18\x0e \x01(\x0b\x32\x32.spark.connect.AnalyzePlanResponse.GetStorageLevelH\x00R\x0fgetStorageLevel\x12N\n\x0bjson_to_ddl\x18\x10 \x01(\x0b\x32,.spark.connect.AnalyzePlanResponse.JsonToDDLH\x00R\tjsonToDdl\x1a\x39\n\x06Schema\x12/\n\x06schema\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1a\x30\n\x07\x45xplain\x12%\n\x0e\x65xplain_string\x18\x01 \x01(\tR\rexplainString\x1a-\n\nTreeString\x12\x1f\n\x0btree_string\x18\x01 \x01(\tR\ntreeString\x1a$\n\x07IsLocal\x12\x19\n\x08is_local\x18\x01 \x01(\x08R\x07isLocal\x1a\x30\n\x0bIsStreaming\x12!\n\x0cis_streaming\x18\x01 \x01(\x08R\x0bisStreaming\x1a"\n\nInputFiles\x12\x14\n\x05\x66iles\x18\x01 \x03(\tR\x05\x66iles\x1a(\n\x0cSparkVersion\x12\x18\n\x07version\x18\x01 \x01(\tR\x07version\x1a;\n\x08\x44\x44LParse\x12/\n\x06parsed\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06parsed\x1a\'\n\rSameSemantics\x12\x16\n\x06result\x18\x01 \x01(\x08R\x06result\x1a&\n\x0cSemanticHash\x12\x16\n\x06result\x18\x01 \x01(\x05R\x06result\x1a\t\n\x07Persist\x1a\x0b\n\tUnpersist\x1aS\n\x0fGetStorageLevel\x12@\n\rstorage_level\x18\x01 \x01(\x0b\x32\x1b.spark.connect.StorageLevelR\x0cstorageLevel\x1a*\n\tJsonToDDL\x12\x1d\n\nddl_string\x18\x01 \x01(\tR\tddlStringB\x08\n\x06result"\x83\x06\n\x12\x45xecutePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x08 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12&\n\x0coperation_id\x18\x06 \x01(\tH\x01R\x0boperationId\x88\x01\x01\x12\'\n\x04plan\x18\x03 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x02R\nclientType\x88\x01\x01\x12X\n\x0frequest_options\x18\x05 \x03(\x0b\x32/.spark.connect.ExecutePlanRequest.RequestOptionR\x0erequestOptions\x12\x12\n\x04tags\x18\x07 \x03(\tR\x04tags\x1a\x85\x02\n\rRequestOption\x12K\n\x10reattach_options\x18\x01 \x01(\x0b\x32\x1e.spark.connect.ReattachOptionsH\x00R\x0freattachOptions\x12^\n\x17result_chunking_options\x18\x02 \x01(\x0b\x32$.spark.connect.ResultChunkingOptionsH\x00R\x15resultChunkingOptions\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textensionB\x10\n\x0erequest_optionB)\n\'_client_observed_server_side_session_idB\x0f\n\r_operation_idB\x0e\n\x0c_client_type"\x87\x1c\n\x13\x45xecutePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x0f \x01(\tR\x13serverSideSessionId\x12!\n\x0coperation_id\x18\x0c \x01(\tR\x0boperationId\x12\x1f\n\x0bresponse_id\x18\r \x01(\tR\nresponseId\x12P\n\x0b\x61rrow_batch\x18\x02 \x01(\x0b\x32-.spark.connect.ExecutePlanResponse.ArrowBatchH\x00R\narrowBatch\x12\x63\n\x12sql_command_result\x18\x05 \x01(\x0b\x32\x33.spark.connect.ExecutePlanResponse.SqlCommandResultH\x00R\x10sqlCommandResult\x12~\n#write_stream_operation_start_result\x18\x08 \x01(\x0b\x32..spark.connect.WriteStreamOperationStartResultH\x00R\x1fwriteStreamOperationStartResult\x12q\n\x1estreaming_query_command_result\x18\t \x01(\x0b\x32*.spark.connect.StreamingQueryCommandResultH\x00R\x1bstreamingQueryCommandResult\x12k\n\x1cget_resources_command_result\x18\n \x01(\x0b\x32(.spark.connect.GetResourcesCommandResultH\x00R\x19getResourcesCommandResult\x12\x87\x01\n&streaming_query_manager_command_result\x18\x0b \x01(\x0b\x32\x31.spark.connect.StreamingQueryManagerCommandResultH\x00R"streamingQueryManagerCommandResult\x12\x87\x01\n&streaming_query_listener_events_result\x18\x10 \x01(\x0b\x32\x31.spark.connect.StreamingQueryListenerEventsResultH\x00R"streamingQueryListenerEventsResult\x12\\\n\x0fresult_complete\x18\x0e \x01(\x0b\x32\x31.spark.connect.ExecutePlanResponse.ResultCompleteH\x00R\x0eresultComplete\x12\x87\x01\n&create_resource_profile_command_result\x18\x11 \x01(\x0b\x32\x31.spark.connect.CreateResourceProfileCommandResultH\x00R"createResourceProfileCommandResult\x12\x65\n\x12\x65xecution_progress\x18\x12 \x01(\x0b\x32\x34.spark.connect.ExecutePlanResponse.ExecutionProgressH\x00R\x11\x65xecutionProgress\x12\x64\n\x19\x63heckpoint_command_result\x18\x13 \x01(\x0b\x32&.spark.connect.CheckpointCommandResultH\x00R\x17\x63heckpointCommandResult\x12L\n\x11ml_command_result\x18\x14 \x01(\x0b\x32\x1e.spark.connect.MlCommandResultH\x00R\x0fmlCommandResult\x12X\n\x15pipeline_event_result\x18\x15 \x01(\x0b\x32".spark.connect.PipelineEventResultH\x00R\x13pipelineEventResult\x12^\n\x17pipeline_command_result\x18\x16 \x01(\x0b\x32$.spark.connect.PipelineCommandResultH\x00R\x15pipelineCommandResult\x12\x8d\x01\n(pipeline_query_function_execution_signal\x18\x17 \x01(\x0b\x32\x33.spark.connect.PipelineQueryFunctionExecutionSignalH\x00R$pipelineQueryFunctionExecutionSignal\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x44\n\x07metrics\x18\x04 \x01(\x0b\x32*.spark.connect.ExecutePlanResponse.MetricsR\x07metrics\x12]\n\x10observed_metrics\x18\x06 \x03(\x0b\x32\x32.spark.connect.ExecutePlanResponse.ObservedMetricsR\x0fobservedMetrics\x12/\n\x06schema\x18\x07 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1aG\n\x10SqlCommandResult\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x1a\xf8\x01\n\nArrowBatch\x12\x1b\n\trow_count\x18\x01 \x01(\x03R\x08rowCount\x12\x12\n\x04\x64\x61ta\x18\x02 \x01(\x0cR\x04\x64\x61ta\x12&\n\x0cstart_offset\x18\x03 \x01(\x03H\x00R\x0bstartOffset\x88\x01\x01\x12$\n\x0b\x63hunk_index\x18\x04 \x01(\x03H\x01R\nchunkIndex\x88\x01\x01\x12\x32\n\x13num_chunks_in_batch\x18\x05 \x01(\x03H\x02R\x10numChunksInBatch\x88\x01\x01\x42\x0f\n\r_start_offsetB\x0e\n\x0c_chunk_indexB\x16\n\x14_num_chunks_in_batch\x1a\x85\x04\n\x07Metrics\x12Q\n\x07metrics\x18\x01 \x03(\x0b\x32\x37.spark.connect.ExecutePlanResponse.Metrics.MetricObjectR\x07metrics\x1a\xcc\x02\n\x0cMetricObject\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x17\n\x07plan_id\x18\x02 \x01(\x03R\x06planId\x12\x16\n\x06parent\x18\x03 \x01(\x03R\x06parent\x12z\n\x11\x65xecution_metrics\x18\x04 \x03(\x0b\x32M.spark.connect.ExecutePlanResponse.Metrics.MetricObject.ExecutionMetricsEntryR\x10\x65xecutionMetrics\x1a{\n\x15\x45xecutionMetricsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ExecutePlanResponse.Metrics.MetricValueR\x05value:\x02\x38\x01\x1aX\n\x0bMetricValue\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x14\n\x05value\x18\x02 \x01(\x03R\x05value\x12\x1f\n\x0bmetric_type\x18\x03 \x01(\tR\nmetricType\x1a\x93\x02\n\x0fObservedMetrics\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x12\x12\n\x04keys\x18\x03 \x03(\tR\x04keys\x12\x17\n\x07plan_id\x18\x04 \x01(\x03R\x06planId\x12)\n\x0eroot_error_idx\x18\x05 \x01(\x05H\x00R\x0crootErrorIdx\x88\x01\x01\x12\x46\n\x06\x65rrors\x18\x06 \x03(\x0b\x32..spark.connect.FetchErrorDetailsResponse.ErrorR\x06\x65rrorsB\x11\n\x0f_root_error_idx\x1a\x10\n\x0eResultComplete\x1a\xcd\x02\n\x11\x45xecutionProgress\x12V\n\x06stages\x18\x01 \x03(\x0b\x32>.spark.connect.ExecutePlanResponse.ExecutionProgress.StageInfoR\x06stages\x12,\n\x12num_inflight_tasks\x18\x02 \x01(\x03R\x10numInflightTasks\x1a\xb1\x01\n\tStageInfo\x12\x19\n\x08stage_id\x18\x01 \x01(\x03R\x07stageId\x12\x1b\n\tnum_tasks\x18\x02 \x01(\x03R\x08numTasks\x12.\n\x13num_completed_tasks\x18\x03 \x01(\x03R\x11numCompletedTasks\x12(\n\x10input_bytes_read\x18\x04 \x01(\x03R\x0einputBytesRead\x12\x12\n\x04\x64one\x18\x05 \x01(\x08R\x04\x64oneB\x0f\n\rresponse_type"A\n\x08KeyValue\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x19\n\x05value\x18\x02 \x01(\tH\x00R\x05value\x88\x01\x01\x42\x08\n\x06_value"\xaf\t\n\rConfigRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x08 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12\x44\n\toperation\x18\x03 \x01(\x0b\x32&.spark.connect.ConfigRequest.OperationR\toperation\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x1a\xf2\x03\n\tOperation\x12\x34\n\x03set\x18\x01 \x01(\x0b\x32 .spark.connect.ConfigRequest.SetH\x00R\x03set\x12\x34\n\x03get\x18\x02 \x01(\x0b\x32 .spark.connect.ConfigRequest.GetH\x00R\x03get\x12W\n\x10get_with_default\x18\x03 \x01(\x0b\x32+.spark.connect.ConfigRequest.GetWithDefaultH\x00R\x0egetWithDefault\x12G\n\nget_option\x18\x04 \x01(\x0b\x32&.spark.connect.ConfigRequest.GetOptionH\x00R\tgetOption\x12>\n\x07get_all\x18\x05 \x01(\x0b\x32#.spark.connect.ConfigRequest.GetAllH\x00R\x06getAll\x12:\n\x05unset\x18\x06 \x01(\x0b\x32".spark.connect.ConfigRequest.UnsetH\x00R\x05unset\x12P\n\ris_modifiable\x18\x07 \x01(\x0b\x32).spark.connect.ConfigRequest.IsModifiableH\x00R\x0cisModifiableB\t\n\x07op_type\x1a\\\n\x03Set\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x12\x1b\n\x06silent\x18\x02 \x01(\x08H\x00R\x06silent\x88\x01\x01\x42\t\n\x07_silent\x1a\x19\n\x03Get\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a?\n\x0eGetWithDefault\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x1a\x1f\n\tGetOption\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a\x30\n\x06GetAll\x12\x1b\n\x06prefix\x18\x01 \x01(\tH\x00R\x06prefix\x88\x01\x01\x42\t\n\x07_prefix\x1a\x1b\n\x05Unset\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a"\n\x0cIsModifiable\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keysB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xaf\x01\n\x0e\x43onfigResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x04 \x01(\tR\x13serverSideSessionId\x12-\n\x05pairs\x18\x02 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x12\x1a\n\x08warnings\x18\x03 \x03(\tR\x08warnings"\xea\x07\n\x13\x41\x64\x64\x41rtifactsRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12$\n\x0b\x63lient_type\x18\x06 \x01(\tH\x02R\nclientType\x88\x01\x01\x12@\n\x05\x62\x61tch\x18\x03 \x01(\x0b\x32(.spark.connect.AddArtifactsRequest.BatchH\x00R\x05\x62\x61tch\x12Z\n\x0b\x62\x65gin_chunk\x18\x04 \x01(\x0b\x32\x37.spark.connect.AddArtifactsRequest.BeginChunkedArtifactH\x00R\nbeginChunk\x12H\n\x05\x63hunk\x18\x05 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkH\x00R\x05\x63hunk\x1a\x35\n\rArtifactChunk\x12\x12\n\x04\x64\x61ta\x18\x01 \x01(\x0cR\x04\x64\x61ta\x12\x10\n\x03\x63rc\x18\x02 \x01(\x03R\x03\x63rc\x1ao\n\x13SingleChunkArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x44\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x04\x64\x61ta\x1a]\n\x05\x42\x61tch\x12T\n\tartifacts\x18\x01 \x03(\x0b\x32\x36.spark.connect.AddArtifactsRequest.SingleChunkArtifactR\tartifacts\x1a\xc1\x01\n\x14\x42\x65ginChunkedArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x1f\n\x0btotal_bytes\x18\x02 \x01(\x03R\ntotalBytes\x12\x1d\n\nnum_chunks\x18\x03 \x01(\x03R\tnumChunks\x12U\n\rinitial_chunk\x18\x04 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x0cinitialChunkB\t\n\x07payloadB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\x90\x02\n\x14\x41\x64\x64\x41rtifactsResponse\x12\x1d\n\nsession_id\x18\x02 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12Q\n\tartifacts\x18\x01 \x03(\x0b\x32\x33.spark.connect.AddArtifactsResponse.ArtifactSummaryR\tartifacts\x1aQ\n\x0f\x41rtifactSummary\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12*\n\x11is_crc_successful\x18\x02 \x01(\x08R\x0fisCrcSuccessful"\xc6\x02\n\x17\x41rtifactStatusesRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12\x14\n\x05names\x18\x04 \x03(\tR\x05namesB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xe0\x02\n\x18\x41rtifactStatusesResponse\x12\x1d\n\nsession_id\x18\x02 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12Q\n\x08statuses\x18\x01 \x03(\x0b\x32\x35.spark.connect.ArtifactStatusesResponse.StatusesEntryR\x08statuses\x1as\n\rStatusesEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ArtifactStatusesResponse.ArtifactStatusR\x05value:\x02\x38\x01\x1a(\n\x0e\x41rtifactStatus\x12\x16\n\x06\x65xists\x18\x01 \x01(\x08R\x06\x65xists"\xdb\x04\n\x10InterruptRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x02R\nclientType\x88\x01\x01\x12T\n\x0einterrupt_type\x18\x04 \x01(\x0e\x32-.spark.connect.InterruptRequest.InterruptTypeR\rinterruptType\x12%\n\roperation_tag\x18\x05 \x01(\tH\x00R\x0coperationTag\x12#\n\x0coperation_id\x18\x06 \x01(\tH\x00R\x0boperationId"\x80\x01\n\rInterruptType\x12\x1e\n\x1aINTERRUPT_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12INTERRUPT_TYPE_ALL\x10\x01\x12\x16\n\x12INTERRUPT_TYPE_TAG\x10\x02\x12\x1f\n\x1bINTERRUPT_TYPE_OPERATION_ID\x10\x03\x42\x0b\n\tinterruptB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\x90\x01\n\x11InterruptResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12\'\n\x0finterrupted_ids\x18\x02 \x03(\tR\x0einterruptedIds"5\n\x0fReattachOptions\x12"\n\x0creattachable\x18\x01 \x01(\x08R\x0creattachable"\xb5\x01\n\x15ResultChunkingOptions\x12;\n\x1a\x61llow_arrow_batch_chunking\x18\x01 \x01(\x08R\x17\x61llowArrowBatchChunking\x12@\n\x1apreferred_arrow_chunk_size\x18\x02 \x01(\x03H\x00R\x17preferredArrowChunkSize\x88\x01\x01\x42\x1d\n\x1b_preferred_arrow_chunk_size"\x96\x03\n\x16ReattachExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x06 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x12-\n\x10last_response_id\x18\x05 \x01(\tH\x02R\x0elastResponseId\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_typeB\x13\n\x11_last_response_id"\xc9\x04\n\x15ReleaseExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x02R\nclientType\x88\x01\x01\x12R\n\x0brelease_all\x18\x05 \x01(\x0b\x32/.spark.connect.ReleaseExecuteRequest.ReleaseAllH\x00R\nreleaseAll\x12X\n\rrelease_until\x18\x06 \x01(\x0b\x32\x31.spark.connect.ReleaseExecuteRequest.ReleaseUntilH\x00R\x0creleaseUntil\x1a\x0c\n\nReleaseAll\x1a/\n\x0cReleaseUntil\x12\x1f\n\x0bresponse_id\x18\x01 \x01(\tR\nresponseIdB\t\n\x07releaseB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xa5\x01\n\x16ReleaseExecuteResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12&\n\x0coperation_id\x18\x02 \x01(\tH\x00R\x0boperationId\x88\x01\x01\x42\x0f\n\r_operation_id"\xd4\x01\n\x15ReleaseSessionRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x00R\nclientType\x88\x01\x01\x12\'\n\x0f\x61llow_reconnect\x18\x04 \x01(\x08R\x0e\x61llowReconnectB\x0e\n\x0c_client_type"l\n\x16ReleaseSessionResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId"\xcc\x02\n\x18\x46\x65tchErrorDetailsRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12\x19\n\x08\x65rror_id\x18\x03 \x01(\tR\x07\x65rrorId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xd9\x0f\n\x19\x46\x65tchErrorDetailsResponse\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12\x1d\n\nsession_id\x18\x04 \x01(\tR\tsessionId\x12)\n\x0eroot_error_idx\x18\x01 \x01(\x05H\x00R\x0crootErrorIdx\x88\x01\x01\x12\x46\n\x06\x65rrors\x18\x02 \x03(\x0b\x32..spark.connect.FetchErrorDetailsResponse.ErrorR\x06\x65rrors\x1a\xae\x01\n\x11StackTraceElement\x12\'\n\x0f\x64\x65\x63laring_class\x18\x01 \x01(\tR\x0e\x64\x65\x63laringClass\x12\x1f\n\x0bmethod_name\x18\x02 \x01(\tR\nmethodName\x12 \n\tfile_name\x18\x03 \x01(\tH\x00R\x08\x66ileName\x88\x01\x01\x12\x1f\n\x0bline_number\x18\x04 \x01(\x05R\nlineNumberB\x0c\n\n_file_name\x1a\xf0\x02\n\x0cQueryContext\x12\x64\n\x0c\x63ontext_type\x18\n \x01(\x0e\x32\x41.spark.connect.FetchErrorDetailsResponse.QueryContext.ContextTypeR\x0b\x63ontextType\x12\x1f\n\x0bobject_type\x18\x01 \x01(\tR\nobjectType\x12\x1f\n\x0bobject_name\x18\x02 \x01(\tR\nobjectName\x12\x1f\n\x0bstart_index\x18\x03 \x01(\x05R\nstartIndex\x12\x1d\n\nstop_index\x18\x04 \x01(\x05R\tstopIndex\x12\x1a\n\x08\x66ragment\x18\x05 \x01(\tR\x08\x66ragment\x12\x1b\n\tcall_site\x18\x06 \x01(\tR\x08\x63\x61llSite\x12\x18\n\x07summary\x18\x07 \x01(\tR\x07summary"%\n\x0b\x43ontextType\x12\x07\n\x03SQL\x10\x00\x12\r\n\tDATAFRAME\x10\x01\x1a\xa6\x04\n\x0eSparkThrowable\x12$\n\x0b\x65rror_class\x18\x01 \x01(\tH\x00R\nerrorClass\x88\x01\x01\x12}\n\x12message_parameters\x18\x02 \x03(\x0b\x32N.spark.connect.FetchErrorDetailsResponse.SparkThrowable.MessageParametersEntryR\x11messageParameters\x12\\\n\x0equery_contexts\x18\x03 \x03(\x0b\x32\x35.spark.connect.FetchErrorDetailsResponse.QueryContextR\rqueryContexts\x12 \n\tsql_state\x18\x04 \x01(\tH\x01R\x08sqlState\x88\x01\x01\x12r\n\x14\x62reaking_change_info\x18\x05 \x01(\x0b\x32;.spark.connect.FetchErrorDetailsResponse.BreakingChangeInfoH\x02R\x12\x62reakingChangeInfo\x88\x01\x01\x1a\x44\n\x16MessageParametersEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\x0e\n\x0c_error_classB\x0c\n\n_sql_stateB\x17\n\x15_breaking_change_info\x1a\xfa\x01\n\x12\x42reakingChangeInfo\x12+\n\x11migration_message\x18\x01 \x03(\tR\x10migrationMessage\x12k\n\x11mitigation_config\x18\x02 \x01(\x0b\x32\x39.spark.connect.FetchErrorDetailsResponse.MitigationConfigH\x00R\x10mitigationConfig\x88\x01\x01\x12$\n\x0bneeds_audit\x18\x03 \x01(\x08H\x01R\nneedsAudit\x88\x01\x01\x42\x14\n\x12_mitigation_configB\x0e\n\x0c_needs_audit\x1a:\n\x10MitigationConfig\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value\x1a\xdb\x02\n\x05\x45rror\x12\x30\n\x14\x65rror_type_hierarchy\x18\x01 \x03(\tR\x12\x65rrorTypeHierarchy\x12\x18\n\x07message\x18\x02 \x01(\tR\x07message\x12[\n\x0bstack_trace\x18\x03 \x03(\x0b\x32:.spark.connect.FetchErrorDetailsResponse.StackTraceElementR\nstackTrace\x12 \n\tcause_idx\x18\x04 \x01(\x05H\x00R\x08\x63\x61useIdx\x88\x01\x01\x12\x65\n\x0fspark_throwable\x18\x05 \x01(\x0b\x32\x37.spark.connect.FetchErrorDetailsResponse.SparkThrowableH\x01R\x0esparkThrowable\x88\x01\x01\x42\x0c\n\n_cause_idxB\x12\n\x10_spark_throwableB\x11\n\x0f_root_error_idx"Z\n\x17\x43heckpointCommandResult\x12?\n\x08relation\x18\x01 \x01(\x0b\x32#.spark.connect.CachedRemoteRelationR\x08relation"\xea\x02\n\x13\x43loneSessionRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12)\n\x0enew_session_id\x18\x04 \x01(\tH\x02R\x0cnewSessionId\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_typeB\x11\n\x0f_new_session_id"\xcc\x01\n\x14\x43loneSessionResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId\x12$\n\x0enew_session_id\x18\x03 \x01(\tR\x0cnewSessionId\x12:\n\x1anew_server_side_session_id\x18\x04 \x01(\tR\x16newServerSideSessionId"\xd3\x04\n\x10GetStatusRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x00R\nclientType\x88\x01\x01\x12V\n&client_observed_server_side_session_id\x18\x04 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12\x66\n\x10operation_status\x18\x05 \x01(\x0b\x32\x36.spark.connect.GetStatusRequest.OperationStatusRequestH\x02R\x0foperationStatus\x88\x01\x01\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions\x1at\n\x16OperationStatusRequest\x12#\n\roperation_ids\x18\x01 \x03(\tR\x0coperationIds\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensionsB\x0e\n\x0c_client_typeB)\n\'_client_observed_server_side_session_idB\x13\n\x11_operation_status"\xad\x05\n\x11GetStatusResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId\x12_\n\x12operation_statuses\x18\x03 \x03(\x0b\x32\x30.spark.connect.GetStatusResponse.OperationStatusR\x11operationStatuses\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions\x1a\xab\x03\n\x0fOperationStatus\x12!\n\x0coperation_id\x18\x01 \x01(\tR\x0boperationId\x12U\n\x05state\x18\x02 \x01(\x0e\x32?.spark.connect.GetStatusResponse.OperationStatus.OperationStateR\x05state\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions"\xe6\x01\n\x0eOperationState\x12\x1f\n\x1bOPERATION_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17OPERATION_STATE_UNKNOWN\x10\x01\x12\x1b\n\x17OPERATION_STATE_RUNNING\x10\x02\x12\x1f\n\x1bOPERATION_STATE_TERMINATING\x10\x03\x12\x1d\n\x19OPERATION_STATE_SUCCEEDED\x10\x04\x12\x1a\n\x16OPERATION_STATE_FAILED\x10\x05\x12\x1d\n\x19OPERATION_STATE_CANCELLED\x10\x06*Q\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZSTD\x10\x01\x32\xdf\x08\n\x13SparkConnectService\x12X\n\x0b\x45xecutePlan\x12!.spark.connect.ExecutePlanRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12V\n\x0b\x41nalyzePlan\x12!.spark.connect.AnalyzePlanRequest\x1a".spark.connect.AnalyzePlanResponse"\x00\x12G\n\x06\x43onfig\x12\x1c.spark.connect.ConfigRequest\x1a\x1d.spark.connect.ConfigResponse"\x00\x12[\n\x0c\x41\x64\x64\x41rtifacts\x12".spark.connect.AddArtifactsRequest\x1a#.spark.connect.AddArtifactsResponse"\x00(\x01\x12\x63\n\x0e\x41rtifactStatus\x12&.spark.connect.ArtifactStatusesRequest\x1a\'.spark.connect.ArtifactStatusesResponse"\x00\x12P\n\tInterrupt\x12\x1f.spark.connect.InterruptRequest\x1a .spark.connect.InterruptResponse"\x00\x12`\n\x0fReattachExecute\x12%.spark.connect.ReattachExecuteRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12_\n\x0eReleaseExecute\x12$.spark.connect.ReleaseExecuteRequest\x1a%.spark.connect.ReleaseExecuteResponse"\x00\x12_\n\x0eReleaseSession\x12$.spark.connect.ReleaseSessionRequest\x1a%.spark.connect.ReleaseSessionResponse"\x00\x12h\n\x11\x46\x65tchErrorDetails\x12\'.spark.connect.FetchErrorDetailsRequest\x1a(.spark.connect.FetchErrorDetailsResponse"\x00\x12Y\n\x0c\x43loneSession\x12".spark.connect.CloneSessionRequest\x1a#.spark.connect.CloneSessionResponse"\x00\x12P\n\tGetStatus\x12\x1f.spark.connect.GetStatusRequest\x1a .spark.connect.GetStatusResponse"\x00\x42\x36\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' + b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"\xe3\x03\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.CommandH\x00R\x07\x63ommand\x12\\\n\x14\x63ompressed_operation\x18\x03 \x01(\x0b\x32\'.spark.connect.Plan.CompressedOperationH\x00R\x13\x63ompressedOperation\x1a\x8e\x02\n\x13\x43ompressedOperation\x12\x12\n\x04\x64\x61ta\x18\x01 \x01(\x0cR\x04\x64\x61ta\x12G\n\x07op_type\x18\x02 \x01(\x0e\x32..spark.connect.Plan.CompressedOperation.OpTypeR\x06opType\x12L\n\x11\x63ompression_codec\x18\x03 \x01(\x0e\x32\x1f.spark.connect.CompressionCodecR\x10\x63ompressionCodec"L\n\x06OpType\x12\x17\n\x13OP_TYPE_UNSPECIFIED\x10\x00\x12\x14\n\x10OP_TYPE_RELATION\x10\x01\x12\x13\n\x0fOP_TYPE_COMMAND\x10\x02\x42\t\n\x07op_type"z\n\x0bUserContext\x12\x17\n\x07user_id\x18\x01 \x01(\tR\x06userId\x12\x1b\n\tuser_name\x18\x02 \x01(\tR\x08userName\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions"\xf5\x14\n\x12\x41nalyzePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x11 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x02R\nclientType\x88\x01\x01\x12\x42\n\x06schema\x18\x04 \x01(\x0b\x32(.spark.connect.AnalyzePlanRequest.SchemaH\x00R\x06schema\x12\x45\n\x07\x65xplain\x18\x05 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.ExplainH\x00R\x07\x65xplain\x12O\n\x0btree_string\x18\x06 \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.TreeStringH\x00R\ntreeString\x12\x46\n\x08is_local\x18\x07 \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.IsLocalH\x00R\x07isLocal\x12R\n\x0cis_streaming\x18\x08 \x01(\x0b\x32-.spark.connect.AnalyzePlanRequest.IsStreamingH\x00R\x0bisStreaming\x12O\n\x0binput_files\x18\t \x01(\x0b\x32,.spark.connect.AnalyzePlanRequest.InputFilesH\x00R\ninputFiles\x12U\n\rspark_version\x18\n \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SparkVersionH\x00R\x0csparkVersion\x12I\n\tddl_parse\x18\x0b \x01(\x0b\x32*.spark.connect.AnalyzePlanRequest.DDLParseH\x00R\x08\x64\x64lParse\x12X\n\x0esame_semantics\x18\x0c \x01(\x0b\x32/.spark.connect.AnalyzePlanRequest.SameSemanticsH\x00R\rsameSemantics\x12U\n\rsemantic_hash\x18\r \x01(\x0b\x32..spark.connect.AnalyzePlanRequest.SemanticHashH\x00R\x0csemanticHash\x12\x45\n\x07persist\x18\x0e \x01(\x0b\x32).spark.connect.AnalyzePlanRequest.PersistH\x00R\x07persist\x12K\n\tunpersist\x18\x0f \x01(\x0b\x32+.spark.connect.AnalyzePlanRequest.UnpersistH\x00R\tunpersist\x12_\n\x11get_storage_level\x18\x10 \x01(\x0b\x32\x31.spark.connect.AnalyzePlanRequest.GetStorageLevelH\x00R\x0fgetStorageLevel\x12M\n\x0bjson_to_ddl\x18\x12 \x01(\x0b\x32+.spark.connect.AnalyzePlanRequest.JsonToDDLH\x00R\tjsonToDdl\x1a\x31\n\x06Schema\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\xbb\x02\n\x07\x45xplain\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12X\n\x0c\x65xplain_mode\x18\x02 \x01(\x0e\x32\x35.spark.connect.AnalyzePlanRequest.Explain.ExplainModeR\x0b\x65xplainMode"\xac\x01\n\x0b\x45xplainMode\x12\x1c\n\x18\x45XPLAIN_MODE_UNSPECIFIED\x10\x00\x12\x17\n\x13\x45XPLAIN_MODE_SIMPLE\x10\x01\x12\x19\n\x15\x45XPLAIN_MODE_EXTENDED\x10\x02\x12\x18\n\x14\x45XPLAIN_MODE_CODEGEN\x10\x03\x12\x15\n\x11\x45XPLAIN_MODE_COST\x10\x04\x12\x1a\n\x16\x45XPLAIN_MODE_FORMATTED\x10\x05\x1aZ\n\nTreeString\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12\x19\n\x05level\x18\x02 \x01(\x05H\x00R\x05level\x88\x01\x01\x42\x08\n\x06_level\x1a\x32\n\x07IsLocal\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x36\n\x0bIsStreaming\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x35\n\nInputFiles\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x0e\n\x0cSparkVersion\x1a)\n\x08\x44\x44LParse\x12\x1d\n\nddl_string\x18\x01 \x01(\tR\tddlString\x1ay\n\rSameSemantics\x12\x34\n\x0btarget_plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\ntargetPlan\x12\x32\n\nother_plan\x18\x02 \x01(\x0b\x32\x13.spark.connect.PlanR\totherPlan\x1a\x37\n\x0cSemanticHash\x12\'\n\x04plan\x18\x01 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x1a\x97\x01\n\x07Persist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x45\n\rstorage_level\x18\x02 \x01(\x0b\x32\x1b.spark.connect.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level\x1an\n\tUnpersist\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x12\x1f\n\x08\x62locking\x18\x02 \x01(\x08H\x00R\x08\x62locking\x88\x01\x01\x42\x0b\n\t_blocking\x1a\x46\n\x0fGetStorageLevel\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x1a,\n\tJsonToDDL\x12\x1f\n\x0bjson_string\x18\x01 \x01(\tR\njsonStringB\t\n\x07\x61nalyzeB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\x81\x0f\n\x13\x41nalyzePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x0f \x01(\tR\x13serverSideSessionId\x12\x43\n\x06schema\x18\x02 \x01(\x0b\x32).spark.connect.AnalyzePlanResponse.SchemaH\x00R\x06schema\x12\x46\n\x07\x65xplain\x18\x03 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.ExplainH\x00R\x07\x65xplain\x12P\n\x0btree_string\x18\x04 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.TreeStringH\x00R\ntreeString\x12G\n\x08is_local\x18\x05 \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.IsLocalH\x00R\x07isLocal\x12S\n\x0cis_streaming\x18\x06 \x01(\x0b\x32..spark.connect.AnalyzePlanResponse.IsStreamingH\x00R\x0bisStreaming\x12P\n\x0binput_files\x18\x07 \x01(\x0b\x32-.spark.connect.AnalyzePlanResponse.InputFilesH\x00R\ninputFiles\x12V\n\rspark_version\x18\x08 \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SparkVersionH\x00R\x0csparkVersion\x12J\n\tddl_parse\x18\t \x01(\x0b\x32+.spark.connect.AnalyzePlanResponse.DDLParseH\x00R\x08\x64\x64lParse\x12Y\n\x0esame_semantics\x18\n \x01(\x0b\x32\x30.spark.connect.AnalyzePlanResponse.SameSemanticsH\x00R\rsameSemantics\x12V\n\rsemantic_hash\x18\x0b \x01(\x0b\x32/.spark.connect.AnalyzePlanResponse.SemanticHashH\x00R\x0csemanticHash\x12\x46\n\x07persist\x18\x0c \x01(\x0b\x32*.spark.connect.AnalyzePlanResponse.PersistH\x00R\x07persist\x12L\n\tunpersist\x18\r \x01(\x0b\x32,.spark.connect.AnalyzePlanResponse.UnpersistH\x00R\tunpersist\x12`\n\x11get_storage_level\x18\x0e \x01(\x0b\x32\x32.spark.connect.AnalyzePlanResponse.GetStorageLevelH\x00R\x0fgetStorageLevel\x12N\n\x0bjson_to_ddl\x18\x10 \x01(\x0b\x32,.spark.connect.AnalyzePlanResponse.JsonToDDLH\x00R\tjsonToDdl\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions\x1a\x39\n\x06Schema\x12/\n\x06schema\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1a\x30\n\x07\x45xplain\x12%\n\x0e\x65xplain_string\x18\x01 \x01(\tR\rexplainString\x1a-\n\nTreeString\x12\x1f\n\x0btree_string\x18\x01 \x01(\tR\ntreeString\x1a$\n\x07IsLocal\x12\x19\n\x08is_local\x18\x01 \x01(\x08R\x07isLocal\x1a\x30\n\x0bIsStreaming\x12!\n\x0cis_streaming\x18\x01 \x01(\x08R\x0bisStreaming\x1a"\n\nInputFiles\x12\x14\n\x05\x66iles\x18\x01 \x03(\tR\x05\x66iles\x1a(\n\x0cSparkVersion\x12\x18\n\x07version\x18\x01 \x01(\tR\x07version\x1a;\n\x08\x44\x44LParse\x12/\n\x06parsed\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06parsed\x1a\'\n\rSameSemantics\x12\x16\n\x06result\x18\x01 \x01(\x08R\x06result\x1a&\n\x0cSemanticHash\x12\x16\n\x06result\x18\x01 \x01(\x05R\x06result\x1a\t\n\x07Persist\x1a\x0b\n\tUnpersist\x1aS\n\x0fGetStorageLevel\x12@\n\rstorage_level\x18\x01 \x01(\x0b\x32\x1b.spark.connect.StorageLevelR\x0cstorageLevel\x1a*\n\tJsonToDDL\x12\x1d\n\nddl_string\x18\x01 \x01(\tR\tddlStringB\x08\n\x06result"\x83\x06\n\x12\x45xecutePlanRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x08 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12&\n\x0coperation_id\x18\x06 \x01(\tH\x01R\x0boperationId\x88\x01\x01\x12\'\n\x04plan\x18\x03 \x01(\x0b\x32\x13.spark.connect.PlanR\x04plan\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x02R\nclientType\x88\x01\x01\x12X\n\x0frequest_options\x18\x05 \x03(\x0b\x32/.spark.connect.ExecutePlanRequest.RequestOptionR\x0erequestOptions\x12\x12\n\x04tags\x18\x07 \x03(\tR\x04tags\x1a\x85\x02\n\rRequestOption\x12K\n\x10reattach_options\x18\x01 \x01(\x0b\x32\x1e.spark.connect.ReattachOptionsH\x00R\x0freattachOptions\x12^\n\x17result_chunking_options\x18\x02 \x01(\x0b\x32$.spark.connect.ResultChunkingOptionsH\x00R\x15resultChunkingOptions\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textensionB\x10\n\x0erequest_optionB)\n\'_client_observed_server_side_session_idB\x0f\n\r_operation_idB\x0e\n\x0c_client_type"\x87\x1c\n\x13\x45xecutePlanResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x0f \x01(\tR\x13serverSideSessionId\x12!\n\x0coperation_id\x18\x0c \x01(\tR\x0boperationId\x12\x1f\n\x0bresponse_id\x18\r \x01(\tR\nresponseId\x12P\n\x0b\x61rrow_batch\x18\x02 \x01(\x0b\x32-.spark.connect.ExecutePlanResponse.ArrowBatchH\x00R\narrowBatch\x12\x63\n\x12sql_command_result\x18\x05 \x01(\x0b\x32\x33.spark.connect.ExecutePlanResponse.SqlCommandResultH\x00R\x10sqlCommandResult\x12~\n#write_stream_operation_start_result\x18\x08 \x01(\x0b\x32..spark.connect.WriteStreamOperationStartResultH\x00R\x1fwriteStreamOperationStartResult\x12q\n\x1estreaming_query_command_result\x18\t \x01(\x0b\x32*.spark.connect.StreamingQueryCommandResultH\x00R\x1bstreamingQueryCommandResult\x12k\n\x1cget_resources_command_result\x18\n \x01(\x0b\x32(.spark.connect.GetResourcesCommandResultH\x00R\x19getResourcesCommandResult\x12\x87\x01\n&streaming_query_manager_command_result\x18\x0b \x01(\x0b\x32\x31.spark.connect.StreamingQueryManagerCommandResultH\x00R"streamingQueryManagerCommandResult\x12\x87\x01\n&streaming_query_listener_events_result\x18\x10 \x01(\x0b\x32\x31.spark.connect.StreamingQueryListenerEventsResultH\x00R"streamingQueryListenerEventsResult\x12\\\n\x0fresult_complete\x18\x0e \x01(\x0b\x32\x31.spark.connect.ExecutePlanResponse.ResultCompleteH\x00R\x0eresultComplete\x12\x87\x01\n&create_resource_profile_command_result\x18\x11 \x01(\x0b\x32\x31.spark.connect.CreateResourceProfileCommandResultH\x00R"createResourceProfileCommandResult\x12\x65\n\x12\x65xecution_progress\x18\x12 \x01(\x0b\x32\x34.spark.connect.ExecutePlanResponse.ExecutionProgressH\x00R\x11\x65xecutionProgress\x12\x64\n\x19\x63heckpoint_command_result\x18\x13 \x01(\x0b\x32&.spark.connect.CheckpointCommandResultH\x00R\x17\x63heckpointCommandResult\x12L\n\x11ml_command_result\x18\x14 \x01(\x0b\x32\x1e.spark.connect.MlCommandResultH\x00R\x0fmlCommandResult\x12X\n\x15pipeline_event_result\x18\x15 \x01(\x0b\x32".spark.connect.PipelineEventResultH\x00R\x13pipelineEventResult\x12^\n\x17pipeline_command_result\x18\x16 \x01(\x0b\x32$.spark.connect.PipelineCommandResultH\x00R\x15pipelineCommandResult\x12\x8d\x01\n(pipeline_query_function_execution_signal\x18\x17 \x01(\x0b\x32\x33.spark.connect.PipelineQueryFunctionExecutionSignalH\x00R$pipelineQueryFunctionExecutionSignal\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x44\n\x07metrics\x18\x04 \x01(\x0b\x32*.spark.connect.ExecutePlanResponse.MetricsR\x07metrics\x12]\n\x10observed_metrics\x18\x06 \x03(\x0b\x32\x32.spark.connect.ExecutePlanResponse.ObservedMetricsR\x0fobservedMetrics\x12/\n\x06schema\x18\x07 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema\x1aG\n\x10SqlCommandResult\x12\x33\n\x08relation\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x08relation\x1a\xf8\x01\n\nArrowBatch\x12\x1b\n\trow_count\x18\x01 \x01(\x03R\x08rowCount\x12\x12\n\x04\x64\x61ta\x18\x02 \x01(\x0cR\x04\x64\x61ta\x12&\n\x0cstart_offset\x18\x03 \x01(\x03H\x00R\x0bstartOffset\x88\x01\x01\x12$\n\x0b\x63hunk_index\x18\x04 \x01(\x03H\x01R\nchunkIndex\x88\x01\x01\x12\x32\n\x13num_chunks_in_batch\x18\x05 \x01(\x03H\x02R\x10numChunksInBatch\x88\x01\x01\x42\x0f\n\r_start_offsetB\x0e\n\x0c_chunk_indexB\x16\n\x14_num_chunks_in_batch\x1a\x85\x04\n\x07Metrics\x12Q\n\x07metrics\x18\x01 \x03(\x0b\x32\x37.spark.connect.ExecutePlanResponse.Metrics.MetricObjectR\x07metrics\x1a\xcc\x02\n\x0cMetricObject\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x17\n\x07plan_id\x18\x02 \x01(\x03R\x06planId\x12\x16\n\x06parent\x18\x03 \x01(\x03R\x06parent\x12z\n\x11\x65xecution_metrics\x18\x04 \x03(\x0b\x32M.spark.connect.ExecutePlanResponse.Metrics.MetricObject.ExecutionMetricsEntryR\x10\x65xecutionMetrics\x1a{\n\x15\x45xecutionMetricsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ExecutePlanResponse.Metrics.MetricValueR\x05value:\x02\x38\x01\x1aX\n\x0bMetricValue\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x14\n\x05value\x18\x02 \x01(\x03R\x05value\x12\x1f\n\x0bmetric_type\x18\x03 \x01(\tR\nmetricType\x1a\x93\x02\n\x0fObservedMetrics\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x12\x12\n\x04keys\x18\x03 \x03(\tR\x04keys\x12\x17\n\x07plan_id\x18\x04 \x01(\x03R\x06planId\x12)\n\x0eroot_error_idx\x18\x05 \x01(\x05H\x00R\x0crootErrorIdx\x88\x01\x01\x12\x46\n\x06\x65rrors\x18\x06 \x03(\x0b\x32..spark.connect.FetchErrorDetailsResponse.ErrorR\x06\x65rrorsB\x11\n\x0f_root_error_idx\x1a\x10\n\x0eResultComplete\x1a\xcd\x02\n\x11\x45xecutionProgress\x12V\n\x06stages\x18\x01 \x03(\x0b\x32>.spark.connect.ExecutePlanResponse.ExecutionProgress.StageInfoR\x06stages\x12,\n\x12num_inflight_tasks\x18\x02 \x01(\x03R\x10numInflightTasks\x1a\xb1\x01\n\tStageInfo\x12\x19\n\x08stage_id\x18\x01 \x01(\x03R\x07stageId\x12\x1b\n\tnum_tasks\x18\x02 \x01(\x03R\x08numTasks\x12.\n\x13num_completed_tasks\x18\x03 \x01(\x03R\x11numCompletedTasks\x12(\n\x10input_bytes_read\x18\x04 \x01(\x03R\x0einputBytesRead\x12\x12\n\x04\x64one\x18\x05 \x01(\x08R\x04\x64oneB\x0f\n\rresponse_type"A\n\x08KeyValue\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x19\n\x05value\x18\x02 \x01(\tH\x00R\x05value\x88\x01\x01\x42\x08\n\x06_value"\xaf\t\n\rConfigRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x08 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12\x44\n\toperation\x18\x03 \x01(\x0b\x32&.spark.connect.ConfigRequest.OperationR\toperation\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x1a\xf2\x03\n\tOperation\x12\x34\n\x03set\x18\x01 \x01(\x0b\x32 .spark.connect.ConfigRequest.SetH\x00R\x03set\x12\x34\n\x03get\x18\x02 \x01(\x0b\x32 .spark.connect.ConfigRequest.GetH\x00R\x03get\x12W\n\x10get_with_default\x18\x03 \x01(\x0b\x32+.spark.connect.ConfigRequest.GetWithDefaultH\x00R\x0egetWithDefault\x12G\n\nget_option\x18\x04 \x01(\x0b\x32&.spark.connect.ConfigRequest.GetOptionH\x00R\tgetOption\x12>\n\x07get_all\x18\x05 \x01(\x0b\x32#.spark.connect.ConfigRequest.GetAllH\x00R\x06getAll\x12:\n\x05unset\x18\x06 \x01(\x0b\x32".spark.connect.ConfigRequest.UnsetH\x00R\x05unset\x12P\n\ris_modifiable\x18\x07 \x01(\x0b\x32).spark.connect.ConfigRequest.IsModifiableH\x00R\x0cisModifiableB\t\n\x07op_type\x1a\\\n\x03Set\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x12\x1b\n\x06silent\x18\x02 \x01(\x08H\x00R\x06silent\x88\x01\x01\x42\t\n\x07_silent\x1a\x19\n\x03Get\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a?\n\x0eGetWithDefault\x12-\n\x05pairs\x18\x01 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x1a\x1f\n\tGetOption\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a\x30\n\x06GetAll\x12\x1b\n\x06prefix\x18\x01 \x01(\tH\x00R\x06prefix\x88\x01\x01\x42\t\n\x07_prefix\x1a\x1b\n\x05Unset\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keys\x1a"\n\x0cIsModifiable\x12\x12\n\x04keys\x18\x01 \x03(\tR\x04keysB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xaf\x01\n\x0e\x43onfigResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x04 \x01(\tR\x13serverSideSessionId\x12-\n\x05pairs\x18\x02 \x03(\x0b\x32\x17.spark.connect.KeyValueR\x05pairs\x12\x1a\n\x08warnings\x18\x03 \x03(\tR\x08warnings"\xea\x07\n\x13\x41\x64\x64\x41rtifactsRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12$\n\x0b\x63lient_type\x18\x06 \x01(\tH\x02R\nclientType\x88\x01\x01\x12@\n\x05\x62\x61tch\x18\x03 \x01(\x0b\x32(.spark.connect.AddArtifactsRequest.BatchH\x00R\x05\x62\x61tch\x12Z\n\x0b\x62\x65gin_chunk\x18\x04 \x01(\x0b\x32\x37.spark.connect.AddArtifactsRequest.BeginChunkedArtifactH\x00R\nbeginChunk\x12H\n\x05\x63hunk\x18\x05 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkH\x00R\x05\x63hunk\x1a\x35\n\rArtifactChunk\x12\x12\n\x04\x64\x61ta\x18\x01 \x01(\x0cR\x04\x64\x61ta\x12\x10\n\x03\x63rc\x18\x02 \x01(\x03R\x03\x63rc\x1ao\n\x13SingleChunkArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x44\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x04\x64\x61ta\x1a]\n\x05\x42\x61tch\x12T\n\tartifacts\x18\x01 \x03(\x0b\x32\x36.spark.connect.AddArtifactsRequest.SingleChunkArtifactR\tartifacts\x1a\xc1\x01\n\x14\x42\x65ginChunkedArtifact\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12\x1f\n\x0btotal_bytes\x18\x02 \x01(\x03R\ntotalBytes\x12\x1d\n\nnum_chunks\x18\x03 \x01(\x03R\tnumChunks\x12U\n\rinitial_chunk\x18\x04 \x01(\x0b\x32\x30.spark.connect.AddArtifactsRequest.ArtifactChunkR\x0cinitialChunkB\t\n\x07payloadB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\x90\x02\n\x14\x41\x64\x64\x41rtifactsResponse\x12\x1d\n\nsession_id\x18\x02 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12Q\n\tartifacts\x18\x01 \x03(\x0b\x32\x33.spark.connect.AddArtifactsResponse.ArtifactSummaryR\tartifacts\x1aQ\n\x0f\x41rtifactSummary\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12*\n\x11is_crc_successful\x18\x02 \x01(\x08R\x0fisCrcSuccessful"\xc6\x02\n\x17\x41rtifactStatusesRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12\x14\n\x05names\x18\x04 \x03(\tR\x05namesB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xe0\x02\n\x18\x41rtifactStatusesResponse\x12\x1d\n\nsession_id\x18\x02 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12Q\n\x08statuses\x18\x01 \x03(\x0b\x32\x35.spark.connect.ArtifactStatusesResponse.StatusesEntryR\x08statuses\x1as\n\rStatusesEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12L\n\x05value\x18\x02 \x01(\x0b\x32\x36.spark.connect.ArtifactStatusesResponse.ArtifactStatusR\x05value:\x02\x38\x01\x1a(\n\x0e\x41rtifactStatus\x12\x16\n\x06\x65xists\x18\x01 \x01(\x08R\x06\x65xists"\xdb\x04\n\x10InterruptRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x02R\nclientType\x88\x01\x01\x12T\n\x0einterrupt_type\x18\x04 \x01(\x0e\x32-.spark.connect.InterruptRequest.InterruptTypeR\rinterruptType\x12%\n\roperation_tag\x18\x05 \x01(\tH\x00R\x0coperationTag\x12#\n\x0coperation_id\x18\x06 \x01(\tH\x00R\x0boperationId"\x80\x01\n\rInterruptType\x12\x1e\n\x1aINTERRUPT_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12INTERRUPT_TYPE_ALL\x10\x01\x12\x16\n\x12INTERRUPT_TYPE_TAG\x10\x02\x12\x1f\n\x1bINTERRUPT_TYPE_OPERATION_ID\x10\x03\x42\x0b\n\tinterruptB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\x90\x01\n\x11InterruptResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12\'\n\x0finterrupted_ids\x18\x02 \x03(\tR\x0einterruptedIds"5\n\x0fReattachOptions\x12"\n\x0creattachable\x18\x01 \x01(\x08R\x0creattachable"\xb5\x01\n\x15ResultChunkingOptions\x12;\n\x1a\x61llow_arrow_batch_chunking\x18\x01 \x01(\x08R\x17\x61llowArrowBatchChunking\x12@\n\x1apreferred_arrow_chunk_size\x18\x02 \x01(\x03H\x00R\x17preferredArrowChunkSize\x88\x01\x01\x42\x1d\n\x1b_preferred_arrow_chunk_size"\x96\x03\n\x16ReattachExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x06 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x12-\n\x10last_response_id\x18\x05 \x01(\tH\x02R\x0elastResponseId\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_typeB\x13\n\x11_last_response_id"\xc9\x04\n\x15ReleaseExecuteRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x07 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12!\n\x0coperation_id\x18\x03 \x01(\tR\x0boperationId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x02R\nclientType\x88\x01\x01\x12R\n\x0brelease_all\x18\x05 \x01(\x0b\x32/.spark.connect.ReleaseExecuteRequest.ReleaseAllH\x00R\nreleaseAll\x12X\n\rrelease_until\x18\x06 \x01(\x0b\x32\x31.spark.connect.ReleaseExecuteRequest.ReleaseUntilH\x00R\x0creleaseUntil\x1a\x0c\n\nReleaseAll\x1a/\n\x0cReleaseUntil\x12\x1f\n\x0bresponse_id\x18\x01 \x01(\tR\nresponseIdB\t\n\x07releaseB)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xa5\x01\n\x16ReleaseExecuteResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12&\n\x0coperation_id\x18\x02 \x01(\tH\x00R\x0boperationId\x88\x01\x01\x42\x0f\n\r_operation_id"\xd4\x01\n\x15ReleaseSessionRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x00R\nclientType\x88\x01\x01\x12\'\n\x0f\x61llow_reconnect\x18\x04 \x01(\x08R\x0e\x61llowReconnectB\x0e\n\x0c_client_type"l\n\x16ReleaseSessionResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId"\xcc\x02\n\x18\x46\x65tchErrorDetailsRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12\x19\n\x08\x65rror_id\x18\x03 \x01(\tR\x07\x65rrorId\x12$\n\x0b\x63lient_type\x18\x04 \x01(\tH\x01R\nclientType\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_type"\xd9\x0f\n\x19\x46\x65tchErrorDetailsResponse\x12\x33\n\x16server_side_session_id\x18\x03 \x01(\tR\x13serverSideSessionId\x12\x1d\n\nsession_id\x18\x04 \x01(\tR\tsessionId\x12)\n\x0eroot_error_idx\x18\x01 \x01(\x05H\x00R\x0crootErrorIdx\x88\x01\x01\x12\x46\n\x06\x65rrors\x18\x02 \x03(\x0b\x32..spark.connect.FetchErrorDetailsResponse.ErrorR\x06\x65rrors\x1a\xae\x01\n\x11StackTraceElement\x12\'\n\x0f\x64\x65\x63laring_class\x18\x01 \x01(\tR\x0e\x64\x65\x63laringClass\x12\x1f\n\x0bmethod_name\x18\x02 \x01(\tR\nmethodName\x12 \n\tfile_name\x18\x03 \x01(\tH\x00R\x08\x66ileName\x88\x01\x01\x12\x1f\n\x0bline_number\x18\x04 \x01(\x05R\nlineNumberB\x0c\n\n_file_name\x1a\xf0\x02\n\x0cQueryContext\x12\x64\n\x0c\x63ontext_type\x18\n \x01(\x0e\x32\x41.spark.connect.FetchErrorDetailsResponse.QueryContext.ContextTypeR\x0b\x63ontextType\x12\x1f\n\x0bobject_type\x18\x01 \x01(\tR\nobjectType\x12\x1f\n\x0bobject_name\x18\x02 \x01(\tR\nobjectName\x12\x1f\n\x0bstart_index\x18\x03 \x01(\x05R\nstartIndex\x12\x1d\n\nstop_index\x18\x04 \x01(\x05R\tstopIndex\x12\x1a\n\x08\x66ragment\x18\x05 \x01(\tR\x08\x66ragment\x12\x1b\n\tcall_site\x18\x06 \x01(\tR\x08\x63\x61llSite\x12\x18\n\x07summary\x18\x07 \x01(\tR\x07summary"%\n\x0b\x43ontextType\x12\x07\n\x03SQL\x10\x00\x12\r\n\tDATAFRAME\x10\x01\x1a\xa6\x04\n\x0eSparkThrowable\x12$\n\x0b\x65rror_class\x18\x01 \x01(\tH\x00R\nerrorClass\x88\x01\x01\x12}\n\x12message_parameters\x18\x02 \x03(\x0b\x32N.spark.connect.FetchErrorDetailsResponse.SparkThrowable.MessageParametersEntryR\x11messageParameters\x12\\\n\x0equery_contexts\x18\x03 \x03(\x0b\x32\x35.spark.connect.FetchErrorDetailsResponse.QueryContextR\rqueryContexts\x12 \n\tsql_state\x18\x04 \x01(\tH\x01R\x08sqlState\x88\x01\x01\x12r\n\x14\x62reaking_change_info\x18\x05 \x01(\x0b\x32;.spark.connect.FetchErrorDetailsResponse.BreakingChangeInfoH\x02R\x12\x62reakingChangeInfo\x88\x01\x01\x1a\x44\n\x16MessageParametersEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\x0e\n\x0c_error_classB\x0c\n\n_sql_stateB\x17\n\x15_breaking_change_info\x1a\xfa\x01\n\x12\x42reakingChangeInfo\x12+\n\x11migration_message\x18\x01 \x03(\tR\x10migrationMessage\x12k\n\x11mitigation_config\x18\x02 \x01(\x0b\x32\x39.spark.connect.FetchErrorDetailsResponse.MitigationConfigH\x00R\x10mitigationConfig\x88\x01\x01\x12$\n\x0bneeds_audit\x18\x03 \x01(\x08H\x01R\nneedsAudit\x88\x01\x01\x42\x14\n\x12_mitigation_configB\x0e\n\x0c_needs_audit\x1a:\n\x10MitigationConfig\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value\x1a\xdb\x02\n\x05\x45rror\x12\x30\n\x14\x65rror_type_hierarchy\x18\x01 \x03(\tR\x12\x65rrorTypeHierarchy\x12\x18\n\x07message\x18\x02 \x01(\tR\x07message\x12[\n\x0bstack_trace\x18\x03 \x03(\x0b\x32:.spark.connect.FetchErrorDetailsResponse.StackTraceElementR\nstackTrace\x12 \n\tcause_idx\x18\x04 \x01(\x05H\x00R\x08\x63\x61useIdx\x88\x01\x01\x12\x65\n\x0fspark_throwable\x18\x05 \x01(\x0b\x32\x37.spark.connect.FetchErrorDetailsResponse.SparkThrowableH\x01R\x0esparkThrowable\x88\x01\x01\x42\x0c\n\n_cause_idxB\x12\n\x10_spark_throwableB\x11\n\x0f_root_error_idx"Z\n\x17\x43heckpointCommandResult\x12?\n\x08relation\x18\x01 \x01(\x0b\x32#.spark.connect.CachedRemoteRelationR\x08relation"\xea\x02\n\x13\x43loneSessionRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12V\n&client_observed_server_side_session_id\x18\x05 \x01(\tH\x00R!clientObservedServerSideSessionId\x88\x01\x01\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x01R\nclientType\x88\x01\x01\x12)\n\x0enew_session_id\x18\x04 \x01(\tH\x02R\x0cnewSessionId\x88\x01\x01\x42)\n\'_client_observed_server_side_session_idB\x0e\n\x0c_client_typeB\x11\n\x0f_new_session_id"\xcc\x01\n\x14\x43loneSessionResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId\x12$\n\x0enew_session_id\x18\x03 \x01(\tR\x0cnewSessionId\x12:\n\x1anew_server_side_session_id\x18\x04 \x01(\tR\x16newServerSideSessionId"\xd3\x04\n\x10GetStatusRequest\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12=\n\x0cuser_context\x18\x02 \x01(\x0b\x32\x1a.spark.connect.UserContextR\x0buserContext\x12$\n\x0b\x63lient_type\x18\x03 \x01(\tH\x00R\nclientType\x88\x01\x01\x12V\n&client_observed_server_side_session_id\x18\x04 \x01(\tH\x01R!clientObservedServerSideSessionId\x88\x01\x01\x12\x66\n\x10operation_status\x18\x05 \x01(\x0b\x32\x36.spark.connect.GetStatusRequest.OperationStatusRequestH\x02R\x0foperationStatus\x88\x01\x01\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions\x1at\n\x16OperationStatusRequest\x12#\n\roperation_ids\x18\x01 \x03(\tR\x0coperationIds\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensionsB\x0e\n\x0c_client_typeB)\n\'_client_observed_server_side_session_idB\x13\n\x11_operation_status"\xad\x05\n\x11GetStatusResponse\x12\x1d\n\nsession_id\x18\x01 \x01(\tR\tsessionId\x12\x33\n\x16server_side_session_id\x18\x02 \x01(\tR\x13serverSideSessionId\x12_\n\x12operation_statuses\x18\x03 \x03(\x0b\x32\x30.spark.connect.GetStatusResponse.OperationStatusR\x11operationStatuses\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions\x1a\xab\x03\n\x0fOperationStatus\x12!\n\x0coperation_id\x18\x01 \x01(\tR\x0boperationId\x12U\n\x05state\x18\x02 \x01(\x0e\x32?.spark.connect.GetStatusResponse.OperationStatus.OperationStateR\x05state\x12\x35\n\nextensions\x18\xe7\x07 \x03(\x0b\x32\x14.google.protobuf.AnyR\nextensions"\xe6\x01\n\x0eOperationState\x12\x1f\n\x1bOPERATION_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17OPERATION_STATE_UNKNOWN\x10\x01\x12\x1b\n\x17OPERATION_STATE_RUNNING\x10\x02\x12\x1f\n\x1bOPERATION_STATE_TERMINATING\x10\x03\x12\x1d\n\x19OPERATION_STATE_SUCCEEDED\x10\x04\x12\x1a\n\x16OPERATION_STATE_FAILED\x10\x05\x12\x1d\n\x19OPERATION_STATE_CANCELLED\x10\x06*Q\n\x10\x43ompressionCodec\x12!\n\x1d\x43OMPRESSION_CODEC_UNSPECIFIED\x10\x00\x12\x1a\n\x16\x43OMPRESSION_CODEC_ZSTD\x10\x01\x32\xdf\x08\n\x13SparkConnectService\x12X\n\x0b\x45xecutePlan\x12!.spark.connect.ExecutePlanRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12V\n\x0b\x41nalyzePlan\x12!.spark.connect.AnalyzePlanRequest\x1a".spark.connect.AnalyzePlanResponse"\x00\x12G\n\x06\x43onfig\x12\x1c.spark.connect.ConfigRequest\x1a\x1d.spark.connect.ConfigResponse"\x00\x12[\n\x0c\x41\x64\x64\x41rtifacts\x12".spark.connect.AddArtifactsRequest\x1a#.spark.connect.AddArtifactsResponse"\x00(\x01\x12\x63\n\x0e\x41rtifactStatus\x12&.spark.connect.ArtifactStatusesRequest\x1a\'.spark.connect.ArtifactStatusesResponse"\x00\x12P\n\tInterrupt\x12\x1f.spark.connect.InterruptRequest\x1a .spark.connect.InterruptResponse"\x00\x12`\n\x0fReattachExecute\x12%.spark.connect.ReattachExecuteRequest\x1a".spark.connect.ExecutePlanResponse"\x00\x30\x01\x12_\n\x0eReleaseExecute\x12$.spark.connect.ReleaseExecuteRequest\x1a%.spark.connect.ReleaseExecuteResponse"\x00\x12_\n\x0eReleaseSession\x12$.spark.connect.ReleaseSessionRequest\x1a%.spark.connect.ReleaseSessionResponse"\x00\x12h\n\x11\x46\x65tchErrorDetails\x12\'.spark.connect.FetchErrorDetailsRequest\x1a(.spark.connect.FetchErrorDetailsResponse"\x00\x12Y\n\x0c\x43loneSession\x12".spark.connect.CloneSessionRequest\x1a#.spark.connect.CloneSessionResponse"\x00\x12P\n\tGetStatus\x12\x1f.spark.connect.GetStatusRequest\x1a .spark.connect.GetStatusResponse"\x00\x42\x36\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' ) _globals = globals() @@ -71,8 +71,8 @@ _globals[ "_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE_MESSAGEPARAMETERSENTRY" ]._serialized_options = b"8\001" - _globals["_COMPRESSIONCODEC"]._serialized_start = 19991 - _globals["_COMPRESSIONCODEC"]._serialized_end = 20072 + _globals["_COMPRESSIONCODEC"]._serialized_start = 20046 + _globals["_COMPRESSIONCODEC"]._serialized_end = 20127 _globals["_PLAN"]._serialized_start = 275 _globals["_PLAN"]._serialized_end = 758 _globals["_PLAN_COMPRESSEDOPERATION"]._serialized_start = 477 @@ -114,173 +114,173 @@ _globals["_ANALYZEPLANREQUEST_JSONTODDL"]._serialized_start = 3448 _globals["_ANALYZEPLANREQUEST_JSONTODDL"]._serialized_end = 3492 _globals["_ANALYZEPLANRESPONSE"]._serialized_start = 3565 - _globals["_ANALYZEPLANRESPONSE"]._serialized_end = 5431 - _globals["_ANALYZEPLANRESPONSE_SCHEMA"]._serialized_start = 4806 - _globals["_ANALYZEPLANRESPONSE_SCHEMA"]._serialized_end = 4863 - _globals["_ANALYZEPLANRESPONSE_EXPLAIN"]._serialized_start = 4865 - _globals["_ANALYZEPLANRESPONSE_EXPLAIN"]._serialized_end = 4913 - _globals["_ANALYZEPLANRESPONSE_TREESTRING"]._serialized_start = 4915 - _globals["_ANALYZEPLANRESPONSE_TREESTRING"]._serialized_end = 4960 - _globals["_ANALYZEPLANRESPONSE_ISLOCAL"]._serialized_start = 4962 - _globals["_ANALYZEPLANRESPONSE_ISLOCAL"]._serialized_end = 4998 - _globals["_ANALYZEPLANRESPONSE_ISSTREAMING"]._serialized_start = 5000 - _globals["_ANALYZEPLANRESPONSE_ISSTREAMING"]._serialized_end = 5048 - _globals["_ANALYZEPLANRESPONSE_INPUTFILES"]._serialized_start = 5050 - _globals["_ANALYZEPLANRESPONSE_INPUTFILES"]._serialized_end = 5084 - _globals["_ANALYZEPLANRESPONSE_SPARKVERSION"]._serialized_start = 5086 - _globals["_ANALYZEPLANRESPONSE_SPARKVERSION"]._serialized_end = 5126 - _globals["_ANALYZEPLANRESPONSE_DDLPARSE"]._serialized_start = 5128 - _globals["_ANALYZEPLANRESPONSE_DDLPARSE"]._serialized_end = 5187 - _globals["_ANALYZEPLANRESPONSE_SAMESEMANTICS"]._serialized_start = 5189 - _globals["_ANALYZEPLANRESPONSE_SAMESEMANTICS"]._serialized_end = 5228 - _globals["_ANALYZEPLANRESPONSE_SEMANTICHASH"]._serialized_start = 5230 - _globals["_ANALYZEPLANRESPONSE_SEMANTICHASH"]._serialized_end = 5268 + _globals["_ANALYZEPLANRESPONSE"]._serialized_end = 5486 + _globals["_ANALYZEPLANRESPONSE_SCHEMA"]._serialized_start = 4861 + _globals["_ANALYZEPLANRESPONSE_SCHEMA"]._serialized_end = 4918 + _globals["_ANALYZEPLANRESPONSE_EXPLAIN"]._serialized_start = 4920 + _globals["_ANALYZEPLANRESPONSE_EXPLAIN"]._serialized_end = 4968 + _globals["_ANALYZEPLANRESPONSE_TREESTRING"]._serialized_start = 4970 + _globals["_ANALYZEPLANRESPONSE_TREESTRING"]._serialized_end = 5015 + _globals["_ANALYZEPLANRESPONSE_ISLOCAL"]._serialized_start = 5017 + _globals["_ANALYZEPLANRESPONSE_ISLOCAL"]._serialized_end = 5053 + _globals["_ANALYZEPLANRESPONSE_ISSTREAMING"]._serialized_start = 5055 + _globals["_ANALYZEPLANRESPONSE_ISSTREAMING"]._serialized_end = 5103 + _globals["_ANALYZEPLANRESPONSE_INPUTFILES"]._serialized_start = 5105 + _globals["_ANALYZEPLANRESPONSE_INPUTFILES"]._serialized_end = 5139 + _globals["_ANALYZEPLANRESPONSE_SPARKVERSION"]._serialized_start = 5141 + _globals["_ANALYZEPLANRESPONSE_SPARKVERSION"]._serialized_end = 5181 + _globals["_ANALYZEPLANRESPONSE_DDLPARSE"]._serialized_start = 5183 + _globals["_ANALYZEPLANRESPONSE_DDLPARSE"]._serialized_end = 5242 + _globals["_ANALYZEPLANRESPONSE_SAMESEMANTICS"]._serialized_start = 5244 + _globals["_ANALYZEPLANRESPONSE_SAMESEMANTICS"]._serialized_end = 5283 + _globals["_ANALYZEPLANRESPONSE_SEMANTICHASH"]._serialized_start = 5285 + _globals["_ANALYZEPLANRESPONSE_SEMANTICHASH"]._serialized_end = 5323 _globals["_ANALYZEPLANRESPONSE_PERSIST"]._serialized_start = 3111 _globals["_ANALYZEPLANRESPONSE_PERSIST"]._serialized_end = 3120 _globals["_ANALYZEPLANRESPONSE_UNPERSIST"]._serialized_start = 3264 _globals["_ANALYZEPLANRESPONSE_UNPERSIST"]._serialized_end = 3275 - _globals["_ANALYZEPLANRESPONSE_GETSTORAGELEVEL"]._serialized_start = 5294 - _globals["_ANALYZEPLANRESPONSE_GETSTORAGELEVEL"]._serialized_end = 5377 - _globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_start = 5379 - _globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_end = 5421 - _globals["_EXECUTEPLANREQUEST"]._serialized_start = 5434 - _globals["_EXECUTEPLANREQUEST"]._serialized_end = 6205 - _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_start = 5868 - _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_end = 6129 - _globals["_EXECUTEPLANRESPONSE"]._serialized_start = 6208 - _globals["_EXECUTEPLANRESPONSE"]._serialized_end = 9799 - _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_start = 8308 - _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_end = 8379 - _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_start = 8382 - _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_end = 8630 - _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_start = 8633 - _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_end = 9150 - _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_start = 8728 - _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_end = 9060 + _globals["_ANALYZEPLANRESPONSE_GETSTORAGELEVEL"]._serialized_start = 5349 + _globals["_ANALYZEPLANRESPONSE_GETSTORAGELEVEL"]._serialized_end = 5432 + _globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_start = 5434 + _globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_end = 5476 + _globals["_EXECUTEPLANREQUEST"]._serialized_start = 5489 + _globals["_EXECUTEPLANREQUEST"]._serialized_end = 6260 + _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_start = 5923 + _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_end = 6184 + _globals["_EXECUTEPLANRESPONSE"]._serialized_start = 6263 + _globals["_EXECUTEPLANRESPONSE"]._serialized_end = 9854 + _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_start = 8363 + _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_end = 8434 + _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_start = 8437 + _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_end = 8685 + _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_start = 8688 + _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_end = 9205 + _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_start = 8783 + _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_end = 9115 _globals[ "_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY" - ]._serialized_start = 8937 + ]._serialized_start = 8992 _globals[ "_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY" - ]._serialized_end = 9060 - _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_start = 9062 - _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_end = 9150 - _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_start = 9153 - _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_end = 9428 - _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_start = 9430 - _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_end = 9446 - _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_start = 9449 - _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_end = 9782 - _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_start = 9605 - _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_end = 9782 - _globals["_KEYVALUE"]._serialized_start = 9801 - _globals["_KEYVALUE"]._serialized_end = 9866 - _globals["_CONFIGREQUEST"]._serialized_start = 9869 - _globals["_CONFIGREQUEST"]._serialized_end = 11068 - _globals["_CONFIGREQUEST_OPERATION"]._serialized_start = 10177 - _globals["_CONFIGREQUEST_OPERATION"]._serialized_end = 10675 - _globals["_CONFIGREQUEST_SET"]._serialized_start = 10677 - _globals["_CONFIGREQUEST_SET"]._serialized_end = 10769 - _globals["_CONFIGREQUEST_GET"]._serialized_start = 10771 - _globals["_CONFIGREQUEST_GET"]._serialized_end = 10796 - _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_start = 10798 - _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_end = 10861 - _globals["_CONFIGREQUEST_GETOPTION"]._serialized_start = 10863 - _globals["_CONFIGREQUEST_GETOPTION"]._serialized_end = 10894 - _globals["_CONFIGREQUEST_GETALL"]._serialized_start = 10896 - _globals["_CONFIGREQUEST_GETALL"]._serialized_end = 10944 - _globals["_CONFIGREQUEST_UNSET"]._serialized_start = 10946 - _globals["_CONFIGREQUEST_UNSET"]._serialized_end = 10973 - _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_start = 10975 - _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_end = 11009 - _globals["_CONFIGRESPONSE"]._serialized_start = 11071 - _globals["_CONFIGRESPONSE"]._serialized_end = 11246 - _globals["_ADDARTIFACTSREQUEST"]._serialized_start = 11249 - _globals["_ADDARTIFACTSREQUEST"]._serialized_end = 12251 - _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_start = 11724 - _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_end = 11777 - _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_start = 11779 - _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_end = 11890 - _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_start = 11892 - _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_end = 11985 - _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_start = 11988 - _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_end = 12181 - _globals["_ADDARTIFACTSRESPONSE"]._serialized_start = 12254 - _globals["_ADDARTIFACTSRESPONSE"]._serialized_end = 12526 - _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_start = 12445 - _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_end = 12526 - _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_start = 12529 - _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_end = 12855 - _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_start = 12858 - _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_end = 13210 - _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_start = 13053 - _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_end = 13168 - _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_start = 13170 - _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_end = 13210 - _globals["_INTERRUPTREQUEST"]._serialized_start = 13213 - _globals["_INTERRUPTREQUEST"]._serialized_end = 13816 - _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_start = 13616 - _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_end = 13744 - _globals["_INTERRUPTRESPONSE"]._serialized_start = 13819 - _globals["_INTERRUPTRESPONSE"]._serialized_end = 13963 - _globals["_REATTACHOPTIONS"]._serialized_start = 13965 - _globals["_REATTACHOPTIONS"]._serialized_end = 14018 - _globals["_RESULTCHUNKINGOPTIONS"]._serialized_start = 14021 - _globals["_RESULTCHUNKINGOPTIONS"]._serialized_end = 14202 - _globals["_REATTACHEXECUTEREQUEST"]._serialized_start = 14205 - _globals["_REATTACHEXECUTEREQUEST"]._serialized_end = 14611 - _globals["_RELEASEEXECUTEREQUEST"]._serialized_start = 14614 - _globals["_RELEASEEXECUTEREQUEST"]._serialized_end = 15199 - _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_start = 15068 - _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_end = 15080 - _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_start = 15082 - _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_end = 15129 - _globals["_RELEASEEXECUTERESPONSE"]._serialized_start = 15202 - _globals["_RELEASEEXECUTERESPONSE"]._serialized_end = 15367 - _globals["_RELEASESESSIONREQUEST"]._serialized_start = 15370 - _globals["_RELEASESESSIONREQUEST"]._serialized_end = 15582 - _globals["_RELEASESESSIONRESPONSE"]._serialized_start = 15584 - _globals["_RELEASESESSIONRESPONSE"]._serialized_end = 15692 - _globals["_FETCHERRORDETAILSREQUEST"]._serialized_start = 15695 - _globals["_FETCHERRORDETAILSREQUEST"]._serialized_end = 16027 - _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_start = 16030 - _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_end = 18039 - _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_start = 16259 - _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_end = 16433 - _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_start = 16436 - _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_end = 16804 - _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_start = 16767 - _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_end = 16804 - _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_start = 16807 - _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_end = 17357 + ]._serialized_end = 9115 + _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_start = 9117 + _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_end = 9205 + _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_start = 9208 + _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_end = 9483 + _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_start = 9485 + _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_end = 9501 + _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_start = 9504 + _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_end = 9837 + _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_start = 9660 + _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_end = 9837 + _globals["_KEYVALUE"]._serialized_start = 9856 + _globals["_KEYVALUE"]._serialized_end = 9921 + _globals["_CONFIGREQUEST"]._serialized_start = 9924 + _globals["_CONFIGREQUEST"]._serialized_end = 11123 + _globals["_CONFIGREQUEST_OPERATION"]._serialized_start = 10232 + _globals["_CONFIGREQUEST_OPERATION"]._serialized_end = 10730 + _globals["_CONFIGREQUEST_SET"]._serialized_start = 10732 + _globals["_CONFIGREQUEST_SET"]._serialized_end = 10824 + _globals["_CONFIGREQUEST_GET"]._serialized_start = 10826 + _globals["_CONFIGREQUEST_GET"]._serialized_end = 10851 + _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_start = 10853 + _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_end = 10916 + _globals["_CONFIGREQUEST_GETOPTION"]._serialized_start = 10918 + _globals["_CONFIGREQUEST_GETOPTION"]._serialized_end = 10949 + _globals["_CONFIGREQUEST_GETALL"]._serialized_start = 10951 + _globals["_CONFIGREQUEST_GETALL"]._serialized_end = 10999 + _globals["_CONFIGREQUEST_UNSET"]._serialized_start = 11001 + _globals["_CONFIGREQUEST_UNSET"]._serialized_end = 11028 + _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_start = 11030 + _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_end = 11064 + _globals["_CONFIGRESPONSE"]._serialized_start = 11126 + _globals["_CONFIGRESPONSE"]._serialized_end = 11301 + _globals["_ADDARTIFACTSREQUEST"]._serialized_start = 11304 + _globals["_ADDARTIFACTSREQUEST"]._serialized_end = 12306 + _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_start = 11779 + _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_end = 11832 + _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_start = 11834 + _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_end = 11945 + _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_start = 11947 + _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_end = 12040 + _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_start = 12043 + _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_end = 12236 + _globals["_ADDARTIFACTSRESPONSE"]._serialized_start = 12309 + _globals["_ADDARTIFACTSRESPONSE"]._serialized_end = 12581 + _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_start = 12500 + _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_end = 12581 + _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_start = 12584 + _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_end = 12910 + _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_start = 12913 + _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_end = 13265 + _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_start = 13108 + _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_end = 13223 + _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_start = 13225 + _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_end = 13265 + _globals["_INTERRUPTREQUEST"]._serialized_start = 13268 + _globals["_INTERRUPTREQUEST"]._serialized_end = 13871 + _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_start = 13671 + _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_end = 13799 + _globals["_INTERRUPTRESPONSE"]._serialized_start = 13874 + _globals["_INTERRUPTRESPONSE"]._serialized_end = 14018 + _globals["_REATTACHOPTIONS"]._serialized_start = 14020 + _globals["_REATTACHOPTIONS"]._serialized_end = 14073 + _globals["_RESULTCHUNKINGOPTIONS"]._serialized_start = 14076 + _globals["_RESULTCHUNKINGOPTIONS"]._serialized_end = 14257 + _globals["_REATTACHEXECUTEREQUEST"]._serialized_start = 14260 + _globals["_REATTACHEXECUTEREQUEST"]._serialized_end = 14666 + _globals["_RELEASEEXECUTEREQUEST"]._serialized_start = 14669 + _globals["_RELEASEEXECUTEREQUEST"]._serialized_end = 15254 + _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_start = 15123 + _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_end = 15135 + _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_start = 15137 + _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_end = 15184 + _globals["_RELEASEEXECUTERESPONSE"]._serialized_start = 15257 + _globals["_RELEASEEXECUTERESPONSE"]._serialized_end = 15422 + _globals["_RELEASESESSIONREQUEST"]._serialized_start = 15425 + _globals["_RELEASESESSIONREQUEST"]._serialized_end = 15637 + _globals["_RELEASESESSIONRESPONSE"]._serialized_start = 15639 + _globals["_RELEASESESSIONRESPONSE"]._serialized_end = 15747 + _globals["_FETCHERRORDETAILSREQUEST"]._serialized_start = 15750 + _globals["_FETCHERRORDETAILSREQUEST"]._serialized_end = 16082 + _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_start = 16085 + _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_end = 18094 + _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_start = 16314 + _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_end = 16488 + _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_start = 16491 + _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_end = 16859 + _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_start = 16822 + _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_end = 16859 + _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_start = 16862 + _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_end = 17412 _globals[ "_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE_MESSAGEPARAMETERSENTRY" - ]._serialized_start = 17234 + ]._serialized_start = 17289 _globals[ "_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE_MESSAGEPARAMETERSENTRY" - ]._serialized_end = 17302 - _globals["_FETCHERRORDETAILSRESPONSE_BREAKINGCHANGEINFO"]._serialized_start = 17360 - _globals["_FETCHERRORDETAILSRESPONSE_BREAKINGCHANGEINFO"]._serialized_end = 17610 - _globals["_FETCHERRORDETAILSRESPONSE_MITIGATIONCONFIG"]._serialized_start = 17612 - _globals["_FETCHERRORDETAILSRESPONSE_MITIGATIONCONFIG"]._serialized_end = 17670 - _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_start = 17673 - _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_end = 18020 - _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_start = 18041 - _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_end = 18131 - _globals["_CLONESESSIONREQUEST"]._serialized_start = 18134 - _globals["_CLONESESSIONREQUEST"]._serialized_end = 18496 - _globals["_CLONESESSIONRESPONSE"]._serialized_start = 18499 - _globals["_CLONESESSIONRESPONSE"]._serialized_end = 18703 - _globals["_GETSTATUSREQUEST"]._serialized_start = 18706 - _globals["_GETSTATUSREQUEST"]._serialized_end = 19301 - _globals["_GETSTATUSREQUEST_OPERATIONSTATUSREQUEST"]._serialized_start = 19105 - _globals["_GETSTATUSREQUEST_OPERATIONSTATUSREQUEST"]._serialized_end = 19221 - _globals["_GETSTATUSRESPONSE"]._serialized_start = 19304 - _globals["_GETSTATUSRESPONSE"]._serialized_end = 19989 - _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS"]._serialized_start = 19562 - _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS"]._serialized_end = 19989 - _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS_OPERATIONSTATE"]._serialized_start = 19759 - _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS_OPERATIONSTATE"]._serialized_end = 19989 - _globals["_SPARKCONNECTSERVICE"]._serialized_start = 20075 - _globals["_SPARKCONNECTSERVICE"]._serialized_end = 21194 + ]._serialized_end = 17357 + _globals["_FETCHERRORDETAILSRESPONSE_BREAKINGCHANGEINFO"]._serialized_start = 17415 + _globals["_FETCHERRORDETAILSRESPONSE_BREAKINGCHANGEINFO"]._serialized_end = 17665 + _globals["_FETCHERRORDETAILSRESPONSE_MITIGATIONCONFIG"]._serialized_start = 17667 + _globals["_FETCHERRORDETAILSRESPONSE_MITIGATIONCONFIG"]._serialized_end = 17725 + _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_start = 17728 + _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_end = 18075 + _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_start = 18096 + _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_end = 18186 + _globals["_CLONESESSIONREQUEST"]._serialized_start = 18189 + _globals["_CLONESESSIONREQUEST"]._serialized_end = 18551 + _globals["_CLONESESSIONRESPONSE"]._serialized_start = 18554 + _globals["_CLONESESSIONRESPONSE"]._serialized_end = 18758 + _globals["_GETSTATUSREQUEST"]._serialized_start = 18761 + _globals["_GETSTATUSREQUEST"]._serialized_end = 19356 + _globals["_GETSTATUSREQUEST_OPERATIONSTATUSREQUEST"]._serialized_start = 19160 + _globals["_GETSTATUSREQUEST_OPERATIONSTATUSREQUEST"]._serialized_end = 19276 + _globals["_GETSTATUSRESPONSE"]._serialized_start = 19359 + _globals["_GETSTATUSRESPONSE"]._serialized_end = 20044 + _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS"]._serialized_start = 19617 + _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS"]._serialized_end = 20044 + _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS_OPERATIONSTATE"]._serialized_start = 19814 + _globals["_GETSTATUSRESPONSE_OPERATIONSTATUS_OPERATIONSTATE"]._serialized_end = 20044 + _globals["_SPARKCONNECTSERVICE"]._serialized_start = 20130 + _globals["_SPARKCONNECTSERVICE"]._serialized_end = 21249 # @@protoc_insertion_point(module_scope) diff --git a/python/pyspark/sql/connect/proto/base_pb2.pyi b/python/pyspark/sql/connect/proto/base_pb2.pyi index 2db3132cd0c01..7750dc6726546 100644 --- a/python/pyspark/sql/connect/proto/base_pb2.pyi +++ b/python/pyspark/sql/connect/proto/base_pb2.pyi @@ -1015,6 +1015,7 @@ class AnalyzePlanResponse(google.protobuf.message.Message): UNPERSIST_FIELD_NUMBER: builtins.int GET_STORAGE_LEVEL_FIELD_NUMBER: builtins.int JSON_TO_DDL_FIELD_NUMBER: builtins.int + EXTENSIONS_FIELD_NUMBER: builtins.int session_id: builtins.str server_side_session_id: builtins.str """Server-side generated idempotency key that the client can use to assert that the server side @@ -1048,6 +1049,13 @@ class AnalyzePlanResponse(google.protobuf.message.Message): def get_storage_level(self) -> global___AnalyzePlanResponse.GetStorageLevel: ... @property def json_to_ddl(self) -> global___AnalyzePlanResponse.JsonToDDL: ... + @property + def extensions( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + google.protobuf.any_pb2.Any + ]: + """Support arbitrary result objects.""" def __init__( self, *, @@ -1067,6 +1075,7 @@ class AnalyzePlanResponse(google.protobuf.message.Message): unpersist: global___AnalyzePlanResponse.Unpersist | None = ..., get_storage_level: global___AnalyzePlanResponse.GetStorageLevel | None = ..., json_to_ddl: global___AnalyzePlanResponse.JsonToDDL | None = ..., + extensions: collections.abc.Iterable[google.protobuf.any_pb2.Any] | None = ..., ) -> None: ... def HasField( self, @@ -1110,6 +1119,8 @@ class AnalyzePlanResponse(google.protobuf.message.Message): b"ddl_parse", "explain", b"explain", + "extensions", + b"extensions", "get_storage_level", b"get_storage_level", "input_files", diff --git a/python/pyspark/sql/connect/proto/expressions_pb2.py b/python/pyspark/sql/connect/proto/expressions_pb2.py index aa51c393c043f..5bb6335cfe3b9 100644 --- a/python/pyspark/sql/connect/proto/expressions_pb2.py +++ b/python/pyspark/sql/connect/proto/expressions_pb2.py @@ -41,7 +41,7 @@ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x1fspark/connect/expressions.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x19spark/connect/types.proto\x1a\x1aspark/connect/common.proto"\x90<\n\nExpression\x12\x37\n\x06\x63ommon\x18\x12 \x01(\x0b\x32\x1f.spark.connect.ExpressionCommonR\x06\x63ommon\x12=\n\x07literal\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x07literal\x12\x62\n\x14unresolved_attribute\x18\x02 \x01(\x0b\x32-.spark.connect.Expression.UnresolvedAttributeH\x00R\x13unresolvedAttribute\x12_\n\x13unresolved_function\x18\x03 \x01(\x0b\x32,.spark.connect.Expression.UnresolvedFunctionH\x00R\x12unresolvedFunction\x12Y\n\x11\x65xpression_string\x18\x04 \x01(\x0b\x32*.spark.connect.Expression.ExpressionStringH\x00R\x10\x65xpressionString\x12S\n\x0funresolved_star\x18\x05 \x01(\x0b\x32(.spark.connect.Expression.UnresolvedStarH\x00R\x0eunresolvedStar\x12\x37\n\x05\x61lias\x18\x06 \x01(\x0b\x32\x1f.spark.connect.Expression.AliasH\x00R\x05\x61lias\x12\x34\n\x04\x63\x61st\x18\x07 \x01(\x0b\x32\x1e.spark.connect.Expression.CastH\x00R\x04\x63\x61st\x12V\n\x10unresolved_regex\x18\x08 \x01(\x0b\x32).spark.connect.Expression.UnresolvedRegexH\x00R\x0funresolvedRegex\x12\x44\n\nsort_order\x18\t \x01(\x0b\x32#.spark.connect.Expression.SortOrderH\x00R\tsortOrder\x12S\n\x0flambda_function\x18\n \x01(\x0b\x32(.spark.connect.Expression.LambdaFunctionH\x00R\x0elambdaFunction\x12:\n\x06window\x18\x0b \x01(\x0b\x32 .spark.connect.Expression.WindowH\x00R\x06window\x12l\n\x18unresolved_extract_value\x18\x0c \x01(\x0b\x32\x30.spark.connect.Expression.UnresolvedExtractValueH\x00R\x16unresolvedExtractValue\x12M\n\rupdate_fields\x18\r \x01(\x0b\x32&.spark.connect.Expression.UpdateFieldsH\x00R\x0cupdateFields\x12\x82\x01\n unresolved_named_lambda_variable\x18\x0e \x01(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableH\x00R\x1dunresolvedNamedLambdaVariable\x12~\n#common_inline_user_defined_function\x18\x0f \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionH\x00R\x1f\x63ommonInlineUserDefinedFunction\x12\x42\n\rcall_function\x18\x10 \x01(\x0b\x32\x1b.spark.connect.CallFunctionH\x00R\x0c\x63\x61llFunction\x12\x64\n\x19named_argument_expression\x18\x11 \x01(\x0b\x32&.spark.connect.NamedArgumentExpressionH\x00R\x17namedArgumentExpression\x12?\n\x0cmerge_action\x18\x13 \x01(\x0b\x32\x1a.spark.connect.MergeActionH\x00R\x0bmergeAction\x12g\n\x1atyped_aggregate_expression\x18\x14 \x01(\x0b\x32\'.spark.connect.TypedAggregateExpressionH\x00R\x18typedAggregateExpression\x12T\n\x13subquery_expression\x18\x15 \x01(\x0b\x32!.spark.connect.SubqueryExpressionH\x00R\x12subqueryExpression\x12s\n\x1b\x64irect_shuffle_partition_id\x18\x16 \x01(\x0b\x32\x32.spark.connect.Expression.DirectShufflePartitionIDH\x00R\x18\x64irectShufflePartitionId\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x1a\x8f\x06\n\x06Window\x12\x42\n\x0fwindow_function\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0ewindowFunction\x12@\n\x0epartition_spec\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\rpartitionSpec\x12\x42\n\norder_spec\x18\x03 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\torderSpec\x12K\n\nframe_spec\x18\x04 \x01(\x0b\x32,.spark.connect.Expression.Window.WindowFrameR\tframeSpec\x1a\xed\x03\n\x0bWindowFrame\x12U\n\nframe_type\x18\x01 \x01(\x0e\x32\x36.spark.connect.Expression.Window.WindowFrame.FrameTypeR\tframeType\x12P\n\x05lower\x18\x02 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05lower\x12P\n\x05upper\x18\x03 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05upper\x1a\x91\x01\n\rFrameBoundary\x12!\n\x0b\x63urrent_row\x18\x01 \x01(\x08H\x00R\ncurrentRow\x12\x1e\n\tunbounded\x18\x02 \x01(\x08H\x00R\tunbounded\x12\x31\n\x05value\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionH\x00R\x05valueB\n\n\x08\x62oundary"O\n\tFrameType\x12\x18\n\x14\x46RAME_TYPE_UNDEFINED\x10\x00\x12\x12\n\x0e\x46RAME_TYPE_ROW\x10\x01\x12\x14\n\x10\x46RAME_TYPE_RANGE\x10\x02\x1a\xa9\x03\n\tSortOrder\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12O\n\tdirection\x18\x02 \x01(\x0e\x32\x31.spark.connect.Expression.SortOrder.SortDirectionR\tdirection\x12U\n\rnull_ordering\x18\x03 \x01(\x0e\x32\x30.spark.connect.Expression.SortOrder.NullOrderingR\x0cnullOrdering"l\n\rSortDirection\x12\x1e\n\x1aSORT_DIRECTION_UNSPECIFIED\x10\x00\x12\x1c\n\x18SORT_DIRECTION_ASCENDING\x10\x01\x12\x1d\n\x19SORT_DIRECTION_DESCENDING\x10\x02"U\n\x0cNullOrdering\x12\x1a\n\x16SORT_NULLS_UNSPECIFIED\x10\x00\x12\x14\n\x10SORT_NULLS_FIRST\x10\x01\x12\x13\n\x0fSORT_NULLS_LAST\x10\x02\x1aK\n\x18\x44irectShufflePartitionID\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x1a\xbb\x02\n\x04\x43\x61st\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12-\n\x04type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04type\x12\x1b\n\x08type_str\x18\x03 \x01(\tH\x00R\x07typeStr\x12\x44\n\teval_mode\x18\x04 \x01(\x0e\x32\'.spark.connect.Expression.Cast.EvalModeR\x08\x65valMode"b\n\x08\x45valMode\x12\x19\n\x15\x45VAL_MODE_UNSPECIFIED\x10\x00\x12\x14\n\x10\x45VAL_MODE_LEGACY\x10\x01\x12\x12\n\x0e\x45VAL_MODE_ANSI\x10\x02\x12\x11\n\rEVAL_MODE_TRY\x10\x03\x42\x0e\n\x0c\x63\x61st_to_type\x1a\x9c\x15\n\x07Literal\x12-\n\x04null\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04null\x12\x18\n\x06\x62inary\x18\x02 \x01(\x0cH\x00R\x06\x62inary\x12\x1a\n\x07\x62oolean\x18\x03 \x01(\x08H\x00R\x07\x62oolean\x12\x14\n\x04\x62yte\x18\x04 \x01(\x05H\x00R\x04\x62yte\x12\x16\n\x05short\x18\x05 \x01(\x05H\x00R\x05short\x12\x1a\n\x07integer\x18\x06 \x01(\x05H\x00R\x07integer\x12\x14\n\x04long\x18\x07 \x01(\x03H\x00R\x04long\x12\x16\n\x05\x66loat\x18\n \x01(\x02H\x00R\x05\x66loat\x12\x18\n\x06\x64ouble\x18\x0b \x01(\x01H\x00R\x06\x64ouble\x12\x45\n\x07\x64\x65\x63imal\x18\x0c \x01(\x0b\x32).spark.connect.Expression.Literal.DecimalH\x00R\x07\x64\x65\x63imal\x12\x18\n\x06string\x18\r \x01(\tH\x00R\x06string\x12\x14\n\x04\x64\x61te\x18\x10 \x01(\x05H\x00R\x04\x64\x61te\x12\x1e\n\ttimestamp\x18\x11 \x01(\x03H\x00R\ttimestamp\x12%\n\rtimestamp_ntz\x18\x12 \x01(\x03H\x00R\x0ctimestampNtz\x12\x61\n\x11\x63\x61lendar_interval\x18\x13 \x01(\x0b\x32\x32.spark.connect.Expression.Literal.CalendarIntervalH\x00R\x10\x63\x61lendarInterval\x12\x30\n\x13year_month_interval\x18\x14 \x01(\x05H\x00R\x11yearMonthInterval\x12,\n\x11\x64\x61y_time_interval\x18\x15 \x01(\x03H\x00R\x0f\x64\x61yTimeInterval\x12?\n\x05\x61rray\x18\x16 \x01(\x0b\x32\'.spark.connect.Expression.Literal.ArrayH\x00R\x05\x61rray\x12\x39\n\x03map\x18\x17 \x01(\x0b\x32%.spark.connect.Expression.Literal.MapH\x00R\x03map\x12\x42\n\x06struct\x18\x18 \x01(\x0b\x32(.spark.connect.Expression.Literal.StructH\x00R\x06struct\x12\x61\n\x11specialized_array\x18\x19 \x01(\x0b\x32\x32.spark.connect.Expression.Literal.SpecializedArrayH\x00R\x10specializedArray\x12<\n\x04time\x18\x1a \x01(\x0b\x32&.spark.connect.Expression.Literal.TimeH\x00R\x04time\x12\x65\n\x13timestamp_ntz_nanos\x18\x1d \x01(\x0b\x32\x33.spark.connect.Expression.Literal.TimestampNTZNanosH\x00R\x11timestampNtzNanos\x12\x65\n\x13timestamp_ltz_nanos\x18\x1e \x01(\x0b\x32\x33.spark.connect.Expression.Literal.TimestampLTZNanosH\x00R\x11timestampLtzNanos\x12\x34\n\tdata_type\x18\x64 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x08\x64\x61taType\x1au\n\x07\x44\x65\x63imal\x12\x14\n\x05value\x18\x01 \x01(\tR\x05value\x12!\n\tprecision\x18\x02 \x01(\x05H\x00R\tprecision\x88\x01\x01\x12\x19\n\x05scale\x18\x03 \x01(\x05H\x01R\x05scale\x88\x01\x01\x42\x0c\n\n_precisionB\x08\n\x06_scale\x1a\x62\n\x10\x43\x61lendarInterval\x12\x16\n\x06months\x18\x01 \x01(\x05R\x06months\x12\x12\n\x04\x64\x61ys\x18\x02 \x01(\x05R\x04\x64\x61ys\x12"\n\x0cmicroseconds\x18\x03 \x01(\x03R\x0cmicroseconds\x1a\x86\x01\n\x05\x41rray\x12>\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\x0b\x65lementType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lements\x1a\xeb\x01\n\x03Map\x12\x36\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\x07keyType\x12:\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\tvalueType\x12\x35\n\x04keys\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x04keys\x12\x39\n\x06values\x18\x04 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1a\x85\x01\n\x06Struct\x12<\n\x0bstruct_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\nstructType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lements\x1a\xc0\x02\n\x10SpecializedArray\x12,\n\x05\x62ools\x18\x01 \x01(\x0b\x32\x14.spark.connect.BoolsH\x00R\x05\x62ools\x12)\n\x04ints\x18\x02 \x01(\x0b\x32\x13.spark.connect.IntsH\x00R\x04ints\x12,\n\x05longs\x18\x03 \x01(\x0b\x32\x14.spark.connect.LongsH\x00R\x05longs\x12/\n\x06\x66loats\x18\x04 \x01(\x0b\x32\x15.spark.connect.FloatsH\x00R\x06\x66loats\x12\x32\n\x07\x64oubles\x18\x05 \x01(\x0b\x32\x16.spark.connect.DoublesH\x00R\x07\x64oubles\x12\x32\n\x07strings\x18\x06 \x01(\x0b\x32\x16.spark.connect.StringsH\x00R\x07stringsB\x0c\n\nvalue_type\x1aK\n\x04Time\x12\x12\n\x04nano\x18\x01 \x01(\x03R\x04nano\x12!\n\tprecision\x18\x02 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precision\x1a\x95\x01\n\x11TimestampNTZNanos\x12!\n\x0c\x65poch_micros\x18\x01 \x01(\x03R\x0b\x65pochMicros\x12,\n\x12nanos_within_micro\x18\x02 \x01(\x05R\x10nanosWithinMicro\x12!\n\tprecision\x18\x03 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precision\x1a\x95\x01\n\x11TimestampLTZNanos\x12!\n\x0c\x65poch_micros\x18\x01 \x01(\x03R\x0b\x65pochMicros\x12,\n\x12nanos_within_micro\x18\x02 \x01(\x05R\x10nanosWithinMicro\x12!\n\tprecision\x18\x03 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precisionB\x0e\n\x0cliteral_typeJ\x04\x08\x1b\x10\x1cJ\x04\x08\x1c\x10\x1d\x1a\xba\x01\n\x13UnresolvedAttribute\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x12\x31\n\x12is_metadata_column\x18\x03 \x01(\x08H\x01R\x10isMetadataColumn\x88\x01\x01\x42\n\n\x08_plan_idB\x15\n\x13_is_metadata_column\x1a\x82\x02\n\x12UnresolvedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x1f\n\x0bis_distinct\x18\x03 \x01(\x08R\nisDistinct\x12\x37\n\x18is_user_defined_function\x18\x04 \x01(\x08R\x15isUserDefinedFunction\x12$\n\x0bis_internal\x18\x05 \x01(\x08H\x00R\nisInternal\x88\x01\x01\x42\x0e\n\x0c_is_internal\x1a\x32\n\x10\x45xpressionString\x12\x1e\n\nexpression\x18\x01 \x01(\tR\nexpression\x1a|\n\x0eUnresolvedStar\x12,\n\x0funparsed_target\x18\x01 \x01(\tH\x00R\x0eunparsedTarget\x88\x01\x01\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x01R\x06planId\x88\x01\x01\x42\x12\n\x10_unparsed_targetB\n\n\x08_plan_id\x1aV\n\x0fUnresolvedRegex\x12\x19\n\x08\x63ol_name\x18\x01 \x01(\tR\x07\x63olName\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x42\n\n\x08_plan_id\x1a\x84\x01\n\x16UnresolvedExtractValue\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12\x39\n\nextraction\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\nextraction\x1a\xbb\x01\n\x0cUpdateFields\x12\x46\n\x11struct_expression\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x10structExpression\x12\x1d\n\nfield_name\x18\x02 \x01(\tR\tfieldName\x12\x44\n\x10value_expression\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0fvalueExpression\x1ax\n\x05\x41lias\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12\x12\n\x04name\x18\x02 \x03(\tR\x04name\x12\x1f\n\x08metadata\x18\x03 \x01(\tH\x00R\x08metadata\x88\x01\x01\x42\x0b\n\t_metadata\x1a\x9e\x01\n\x0eLambdaFunction\x12\x35\n\x08\x66unction\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08\x66unction\x12U\n\targuments\x18\x02 \x03(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableR\targuments\x1a>\n\x1dUnresolvedNamedLambdaVariable\x12\x1d\n\nname_parts\x18\x01 \x03(\tR\tnamePartsB\x0b\n\texpr_type"A\n\x10\x45xpressionCommon\x12-\n\x06origin\x18\x01 \x01(\x0b\x32\x15.spark.connect.OriginR\x06origin"\x8d\x03\n\x1f\x43ommonInlineUserDefinedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x39\n\npython_udf\x18\x04 \x01(\x0b\x32\x18.spark.connect.PythonUDFH\x00R\tpythonUdf\x12I\n\x10scalar_scala_udf\x18\x05 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFH\x00R\x0escalarScalaUdf\x12\x33\n\x08java_udf\x18\x06 \x01(\x0b\x32\x16.spark.connect.JavaUDFH\x00R\x07javaUdf\x12\x1f\n\x0bis_distinct\x18\x07 \x01(\x08R\nisDistinctB\n\n\x08\x66unction"\xcc\x01\n\tPythonUDF\x12\x38\n\x0boutput_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVer\x12/\n\x13\x61\x64\x64itional_includes\x18\x05 \x03(\tR\x12\x61\x64\x64itionalIncludes"\xd6\x01\n\x0eScalarScalaUDF\x12\x18\n\x07payload\x18\x01 \x01(\x0cR\x07payload\x12\x37\n\ninputTypes\x18\x02 \x03(\x0b\x32\x17.spark.connect.DataTypeR\ninputTypes\x12\x37\n\noutputType\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1a\n\x08nullable\x18\x04 \x01(\x08R\x08nullable\x12\x1c\n\taggregate\x18\x05 \x01(\x08R\taggregate"\x95\x01\n\x07JavaUDF\x12\x1d\n\nclass_name\x18\x01 \x01(\tR\tclassName\x12=\n\x0boutput_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\noutputType\x88\x01\x01\x12\x1c\n\taggregate\x18\x03 \x01(\x08R\taggregateB\x0e\n\x0c_output_type"c\n\x18TypedAggregateExpression\x12G\n\x10scalar_scala_udf\x18\x01 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFR\x0escalarScalaUdf"l\n\x0c\x43\x61llFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments"\\\n\x17NamedArgumentExpression\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value"\x80\x04\n\x0bMergeAction\x12\x46\n\x0b\x61\x63tion_type\x18\x01 \x01(\x0e\x32%.spark.connect.MergeAction.ActionTypeR\nactionType\x12<\n\tcondition\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionH\x00R\tcondition\x88\x01\x01\x12G\n\x0b\x61ssignments\x18\x03 \x03(\x0b\x32%.spark.connect.MergeAction.AssignmentR\x0b\x61ssignments\x1aj\n\nAssignment\x12+\n\x03key\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value"\xa7\x01\n\nActionType\x12\x17\n\x13\x41\x43TION_TYPE_INVALID\x10\x00\x12\x16\n\x12\x41\x43TION_TYPE_DELETE\x10\x01\x12\x16\n\x12\x41\x43TION_TYPE_INSERT\x10\x02\x12\x1b\n\x17\x41\x43TION_TYPE_INSERT_STAR\x10\x03\x12\x16\n\x12\x41\x43TION_TYPE_UPDATE\x10\x04\x12\x1b\n\x17\x41\x43TION_TYPE_UPDATE_STAR\x10\x05\x42\x0c\n\n_condition"\xc5\x05\n\x12SubqueryExpression\x12\x17\n\x07plan_id\x18\x01 \x01(\x03R\x06planId\x12S\n\rsubquery_type\x18\x02 \x01(\x0e\x32..spark.connect.SubqueryExpression.SubqueryTypeR\x0csubqueryType\x12\x62\n\x11table_arg_options\x18\x03 \x01(\x0b\x32\x31.spark.connect.SubqueryExpression.TableArgOptionsH\x00R\x0ftableArgOptions\x88\x01\x01\x12G\n\x12in_subquery_values\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x10inSubqueryValues\x1a\xea\x01\n\x0fTableArgOptions\x12@\n\x0epartition_spec\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\rpartitionSpec\x12\x42\n\norder_spec\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\torderSpec\x12\x37\n\x15with_single_partition\x18\x03 \x01(\x08H\x00R\x13withSinglePartition\x88\x01\x01\x42\x18\n\x16_with_single_partition"\x90\x01\n\x0cSubqueryType\x12\x19\n\x15SUBQUERY_TYPE_UNKNOWN\x10\x00\x12\x18\n\x14SUBQUERY_TYPE_SCALAR\x10\x01\x12\x18\n\x14SUBQUERY_TYPE_EXISTS\x10\x02\x12\x1b\n\x17SUBQUERY_TYPE_TABLE_ARG\x10\x03\x12\x14\n\x10SUBQUERY_TYPE_IN\x10\x04\x42\x14\n\x12_table_arg_optionsB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' + b'\n\x1fspark/connect/expressions.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x19spark/connect/types.proto\x1a\x1aspark/connect/common.proto"\x90<\n\nExpression\x12\x37\n\x06\x63ommon\x18\x12 \x01(\x0b\x32\x1f.spark.connect.ExpressionCommonR\x06\x63ommon\x12=\n\x07literal\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x07literal\x12\x62\n\x14unresolved_attribute\x18\x02 \x01(\x0b\x32-.spark.connect.Expression.UnresolvedAttributeH\x00R\x13unresolvedAttribute\x12_\n\x13unresolved_function\x18\x03 \x01(\x0b\x32,.spark.connect.Expression.UnresolvedFunctionH\x00R\x12unresolvedFunction\x12Y\n\x11\x65xpression_string\x18\x04 \x01(\x0b\x32*.spark.connect.Expression.ExpressionStringH\x00R\x10\x65xpressionString\x12S\n\x0funresolved_star\x18\x05 \x01(\x0b\x32(.spark.connect.Expression.UnresolvedStarH\x00R\x0eunresolvedStar\x12\x37\n\x05\x61lias\x18\x06 \x01(\x0b\x32\x1f.spark.connect.Expression.AliasH\x00R\x05\x61lias\x12\x34\n\x04\x63\x61st\x18\x07 \x01(\x0b\x32\x1e.spark.connect.Expression.CastH\x00R\x04\x63\x61st\x12V\n\x10unresolved_regex\x18\x08 \x01(\x0b\x32).spark.connect.Expression.UnresolvedRegexH\x00R\x0funresolvedRegex\x12\x44\n\nsort_order\x18\t \x01(\x0b\x32#.spark.connect.Expression.SortOrderH\x00R\tsortOrder\x12S\n\x0flambda_function\x18\n \x01(\x0b\x32(.spark.connect.Expression.LambdaFunctionH\x00R\x0elambdaFunction\x12:\n\x06window\x18\x0b \x01(\x0b\x32 .spark.connect.Expression.WindowH\x00R\x06window\x12l\n\x18unresolved_extract_value\x18\x0c \x01(\x0b\x32\x30.spark.connect.Expression.UnresolvedExtractValueH\x00R\x16unresolvedExtractValue\x12M\n\rupdate_fields\x18\r \x01(\x0b\x32&.spark.connect.Expression.UpdateFieldsH\x00R\x0cupdateFields\x12\x82\x01\n unresolved_named_lambda_variable\x18\x0e \x01(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableH\x00R\x1dunresolvedNamedLambdaVariable\x12~\n#common_inline_user_defined_function\x18\x0f \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionH\x00R\x1f\x63ommonInlineUserDefinedFunction\x12\x42\n\rcall_function\x18\x10 \x01(\x0b\x32\x1b.spark.connect.CallFunctionH\x00R\x0c\x63\x61llFunction\x12\x64\n\x19named_argument_expression\x18\x11 \x01(\x0b\x32&.spark.connect.NamedArgumentExpressionH\x00R\x17namedArgumentExpression\x12?\n\x0cmerge_action\x18\x13 \x01(\x0b\x32\x1a.spark.connect.MergeActionH\x00R\x0bmergeAction\x12g\n\x1atyped_aggregate_expression\x18\x14 \x01(\x0b\x32\'.spark.connect.TypedAggregateExpressionH\x00R\x18typedAggregateExpression\x12T\n\x13subquery_expression\x18\x15 \x01(\x0b\x32!.spark.connect.SubqueryExpressionH\x00R\x12subqueryExpression\x12s\n\x1b\x64irect_shuffle_partition_id\x18\x16 \x01(\x0b\x32\x32.spark.connect.Expression.DirectShufflePartitionIDH\x00R\x18\x64irectShufflePartitionId\x12\x35\n\textension\x18\xe7\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x1a\x8f\x06\n\x06Window\x12\x42\n\x0fwindow_function\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0ewindowFunction\x12@\n\x0epartition_spec\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\rpartitionSpec\x12\x42\n\norder_spec\x18\x03 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\torderSpec\x12K\n\nframe_spec\x18\x04 \x01(\x0b\x32,.spark.connect.Expression.Window.WindowFrameR\tframeSpec\x1a\xed\x03\n\x0bWindowFrame\x12U\n\nframe_type\x18\x01 \x01(\x0e\x32\x36.spark.connect.Expression.Window.WindowFrame.FrameTypeR\tframeType\x12P\n\x05lower\x18\x02 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05lower\x12P\n\x05upper\x18\x03 \x01(\x0b\x32:.spark.connect.Expression.Window.WindowFrame.FrameBoundaryR\x05upper\x1a\x91\x01\n\rFrameBoundary\x12!\n\x0b\x63urrent_row\x18\x01 \x01(\x08H\x00R\ncurrentRow\x12\x1e\n\tunbounded\x18\x02 \x01(\x08H\x00R\tunbounded\x12\x31\n\x05value\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionH\x00R\x05valueB\n\n\x08\x62oundary"O\n\tFrameType\x12\x18\n\x14\x46RAME_TYPE_UNDEFINED\x10\x00\x12\x12\n\x0e\x46RAME_TYPE_ROW\x10\x01\x12\x14\n\x10\x46RAME_TYPE_RANGE\x10\x02\x1a\xa9\x03\n\tSortOrder\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12O\n\tdirection\x18\x02 \x01(\x0e\x32\x31.spark.connect.Expression.SortOrder.SortDirectionR\tdirection\x12U\n\rnull_ordering\x18\x03 \x01(\x0e\x32\x30.spark.connect.Expression.SortOrder.NullOrderingR\x0cnullOrdering"l\n\rSortDirection\x12\x1e\n\x1aSORT_DIRECTION_UNSPECIFIED\x10\x00\x12\x1c\n\x18SORT_DIRECTION_ASCENDING\x10\x01\x12\x1d\n\x19SORT_DIRECTION_DESCENDING\x10\x02"U\n\x0cNullOrdering\x12\x1a\n\x16SORT_NULLS_UNSPECIFIED\x10\x00\x12\x14\n\x10SORT_NULLS_FIRST\x10\x01\x12\x13\n\x0fSORT_NULLS_LAST\x10\x02\x1aK\n\x18\x44irectShufflePartitionID\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x1a\xbb\x02\n\x04\x43\x61st\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12-\n\x04type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04type\x12\x1b\n\x08type_str\x18\x03 \x01(\tH\x00R\x07typeStr\x12\x44\n\teval_mode\x18\x04 \x01(\x0e\x32\'.spark.connect.Expression.Cast.EvalModeR\x08\x65valMode"b\n\x08\x45valMode\x12\x19\n\x15\x45VAL_MODE_UNSPECIFIED\x10\x00\x12\x14\n\x10\x45VAL_MODE_LEGACY\x10\x01\x12\x12\n\x0e\x45VAL_MODE_ANSI\x10\x02\x12\x11\n\rEVAL_MODE_TRY\x10\x03\x42\x0e\n\x0c\x63\x61st_to_type\x1a\x9c\x15\n\x07Literal\x12-\n\x04null\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x04null\x12\x18\n\x06\x62inary\x18\x02 \x01(\x0cH\x00R\x06\x62inary\x12\x1a\n\x07\x62oolean\x18\x03 \x01(\x08H\x00R\x07\x62oolean\x12\x14\n\x04\x62yte\x18\x04 \x01(\x05H\x00R\x04\x62yte\x12\x16\n\x05short\x18\x05 \x01(\x05H\x00R\x05short\x12\x1a\n\x07integer\x18\x06 \x01(\x05H\x00R\x07integer\x12\x14\n\x04long\x18\x07 \x01(\x03H\x00R\x04long\x12\x16\n\x05\x66loat\x18\n \x01(\x02H\x00R\x05\x66loat\x12\x18\n\x06\x64ouble\x18\x0b \x01(\x01H\x00R\x06\x64ouble\x12\x45\n\x07\x64\x65\x63imal\x18\x0c \x01(\x0b\x32).spark.connect.Expression.Literal.DecimalH\x00R\x07\x64\x65\x63imal\x12\x18\n\x06string\x18\r \x01(\tH\x00R\x06string\x12\x14\n\x04\x64\x61te\x18\x10 \x01(\x05H\x00R\x04\x64\x61te\x12\x1e\n\ttimestamp\x18\x11 \x01(\x03H\x00R\ttimestamp\x12%\n\rtimestamp_ntz\x18\x12 \x01(\x03H\x00R\x0ctimestampNtz\x12\x61\n\x11\x63\x61lendar_interval\x18\x13 \x01(\x0b\x32\x32.spark.connect.Expression.Literal.CalendarIntervalH\x00R\x10\x63\x61lendarInterval\x12\x30\n\x13year_month_interval\x18\x14 \x01(\x05H\x00R\x11yearMonthInterval\x12,\n\x11\x64\x61y_time_interval\x18\x15 \x01(\x03H\x00R\x0f\x64\x61yTimeInterval\x12?\n\x05\x61rray\x18\x16 \x01(\x0b\x32\'.spark.connect.Expression.Literal.ArrayH\x00R\x05\x61rray\x12\x39\n\x03map\x18\x17 \x01(\x0b\x32%.spark.connect.Expression.Literal.MapH\x00R\x03map\x12\x42\n\x06struct\x18\x18 \x01(\x0b\x32(.spark.connect.Expression.Literal.StructH\x00R\x06struct\x12\x61\n\x11specialized_array\x18\x19 \x01(\x0b\x32\x32.spark.connect.Expression.Literal.SpecializedArrayH\x00R\x10specializedArray\x12<\n\x04time\x18\x1a \x01(\x0b\x32&.spark.connect.Expression.Literal.TimeH\x00R\x04time\x12\x65\n\x13timestamp_ntz_nanos\x18\x1d \x01(\x0b\x32\x33.spark.connect.Expression.Literal.TimestampNTZNanosH\x00R\x11timestampNtzNanos\x12\x65\n\x13timestamp_ltz_nanos\x18\x1e \x01(\x0b\x32\x33.spark.connect.Expression.Literal.TimestampLTZNanosH\x00R\x11timestampLtzNanos\x12\x34\n\tdata_type\x18\x64 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x08\x64\x61taType\x1au\n\x07\x44\x65\x63imal\x12\x14\n\x05value\x18\x01 \x01(\tR\x05value\x12!\n\tprecision\x18\x02 \x01(\x05H\x00R\tprecision\x88\x01\x01\x12\x19\n\x05scale\x18\x03 \x01(\x05H\x01R\x05scale\x88\x01\x01\x42\x0c\n\n_precisionB\x08\n\x06_scale\x1a\x62\n\x10\x43\x61lendarInterval\x12\x16\n\x06months\x18\x01 \x01(\x05R\x06months\x12\x12\n\x04\x64\x61ys\x18\x02 \x01(\x05R\x04\x64\x61ys\x12"\n\x0cmicroseconds\x18\x03 \x01(\x03R\x0cmicroseconds\x1a\x86\x01\n\x05\x41rray\x12>\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\x0b\x65lementType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lements\x1a\xeb\x01\n\x03Map\x12\x36\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\x07keyType\x12:\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\tvalueType\x12\x35\n\x04keys\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x04keys\x12\x39\n\x06values\x18\x04 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1a\x85\x01\n\x06Struct\x12<\n\x0bstruct_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeB\x02\x18\x01R\nstructType\x12=\n\x08\x65lements\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x08\x65lements\x1a\xc0\x02\n\x10SpecializedArray\x12,\n\x05\x62ools\x18\x01 \x01(\x0b\x32\x14.spark.connect.BoolsH\x00R\x05\x62ools\x12)\n\x04ints\x18\x02 \x01(\x0b\x32\x13.spark.connect.IntsH\x00R\x04ints\x12,\n\x05longs\x18\x03 \x01(\x0b\x32\x14.spark.connect.LongsH\x00R\x05longs\x12/\n\x06\x66loats\x18\x04 \x01(\x0b\x32\x15.spark.connect.FloatsH\x00R\x06\x66loats\x12\x32\n\x07\x64oubles\x18\x05 \x01(\x0b\x32\x16.spark.connect.DoublesH\x00R\x07\x64oubles\x12\x32\n\x07strings\x18\x06 \x01(\x0b\x32\x16.spark.connect.StringsH\x00R\x07stringsB\x0c\n\nvalue_type\x1aK\n\x04Time\x12\x12\n\x04nano\x18\x01 \x01(\x03R\x04nano\x12!\n\tprecision\x18\x02 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precision\x1a\x95\x01\n\x11TimestampNTZNanos\x12!\n\x0c\x65poch_micros\x18\x01 \x01(\x03R\x0b\x65pochMicros\x12,\n\x12nanos_within_micro\x18\x02 \x01(\x05R\x10nanosWithinMicro\x12!\n\tprecision\x18\x03 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precision\x1a\x95\x01\n\x11TimestampLTZNanos\x12!\n\x0c\x65poch_micros\x18\x01 \x01(\x03R\x0b\x65pochMicros\x12,\n\x12nanos_within_micro\x18\x02 \x01(\x05R\x10nanosWithinMicro\x12!\n\tprecision\x18\x03 \x01(\x05H\x00R\tprecision\x88\x01\x01\x42\x0c\n\n_precisionB\x0e\n\x0cliteral_typeJ\x04\x08\x1b\x10\x1cJ\x04\x08\x1c\x10\x1d\x1a\xba\x01\n\x13UnresolvedAttribute\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x12\x31\n\x12is_metadata_column\x18\x03 \x01(\x08H\x01R\x10isMetadataColumn\x88\x01\x01\x42\n\n\x08_plan_idB\x15\n\x13_is_metadata_column\x1a\x82\x02\n\x12UnresolvedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x1f\n\x0bis_distinct\x18\x03 \x01(\x08R\nisDistinct\x12\x37\n\x18is_user_defined_function\x18\x04 \x01(\x08R\x15isUserDefinedFunction\x12$\n\x0bis_internal\x18\x05 \x01(\x08H\x00R\nisInternal\x88\x01\x01\x42\x0e\n\x0c_is_internal\x1a\x32\n\x10\x45xpressionString\x12\x1e\n\nexpression\x18\x01 \x01(\tR\nexpression\x1a|\n\x0eUnresolvedStar\x12,\n\x0funparsed_target\x18\x01 \x01(\tH\x00R\x0eunparsedTarget\x88\x01\x01\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x01R\x06planId\x88\x01\x01\x42\x12\n\x10_unparsed_targetB\n\n\x08_plan_id\x1aV\n\x0fUnresolvedRegex\x12\x19\n\x08\x63ol_name\x18\x01 \x01(\tR\x07\x63olName\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x42\n\n\x08_plan_id\x1a\x84\x01\n\x16UnresolvedExtractValue\x12/\n\x05\x63hild\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05\x63hild\x12\x39\n\nextraction\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\nextraction\x1a\xbb\x01\n\x0cUpdateFields\x12\x46\n\x11struct_expression\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x10structExpression\x12\x1d\n\nfield_name\x18\x02 \x01(\tR\tfieldName\x12\x44\n\x10value_expression\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x0fvalueExpression\x1ax\n\x05\x41lias\x12-\n\x04\x65xpr\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x04\x65xpr\x12\x12\n\x04name\x18\x02 \x03(\tR\x04name\x12\x1f\n\x08metadata\x18\x03 \x01(\tH\x00R\x08metadata\x88\x01\x01\x42\x0b\n\t_metadata\x1a\x9e\x01\n\x0eLambdaFunction\x12\x35\n\x08\x66unction\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08\x66unction\x12U\n\targuments\x18\x02 \x03(\x0b\x32\x37.spark.connect.Expression.UnresolvedNamedLambdaVariableR\targuments\x1a>\n\x1dUnresolvedNamedLambdaVariable\x12\x1d\n\nname_parts\x18\x01 \x03(\tR\tnamePartsB\x0b\n\texpr_type"A\n\x10\x45xpressionCommon\x12-\n\x06origin\x18\x01 \x01(\x0b\x32\x15.spark.connect.OriginR\x06origin"\x8d\x03\n\x1f\x43ommonInlineUserDefinedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12\x39\n\npython_udf\x18\x04 \x01(\x0b\x32\x18.spark.connect.PythonUDFH\x00R\tpythonUdf\x12I\n\x10scalar_scala_udf\x18\x05 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFH\x00R\x0escalarScalaUdf\x12\x33\n\x08java_udf\x18\x06 \x01(\x0b\x32\x16.spark.connect.JavaUDFH\x00R\x07javaUdf\x12\x1f\n\x0bis_distinct\x18\x07 \x01(\x08R\nisDistinctB\n\n\x08\x66unction"\x9b\x02\n\tPythonUDF\x12\x38\n\x0boutput_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVer\x12/\n\x13\x61\x64\x64itional_includes\x18\x05 \x03(\tR\x12\x61\x64\x64itionalIncludes\x12=\n\x0b\x62uffer_type\x18\x06 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\nbufferType\x88\x01\x01\x42\x0e\n\x0c_buffer_type"\xd6\x01\n\x0eScalarScalaUDF\x12\x18\n\x07payload\x18\x01 \x01(\x0cR\x07payload\x12\x37\n\ninputTypes\x18\x02 \x03(\x0b\x32\x17.spark.connect.DataTypeR\ninputTypes\x12\x37\n\noutputType\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeR\noutputType\x12\x1a\n\x08nullable\x18\x04 \x01(\x08R\x08nullable\x12\x1c\n\taggregate\x18\x05 \x01(\x08R\taggregate"\x95\x01\n\x07JavaUDF\x12\x1d\n\nclass_name\x18\x01 \x01(\tR\tclassName\x12=\n\x0boutput_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\noutputType\x88\x01\x01\x12\x1c\n\taggregate\x18\x03 \x01(\x08R\taggregateB\x0e\n\x0c_output_type"c\n\x18TypedAggregateExpression\x12G\n\x10scalar_scala_udf\x18\x01 \x01(\x0b\x32\x1d.spark.connect.ScalarScalaUDFR\x0escalarScalaUdf"l\n\x0c\x43\x61llFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments"\\\n\x17NamedArgumentExpression\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value"\x80\x04\n\x0bMergeAction\x12\x46\n\x0b\x61\x63tion_type\x18\x01 \x01(\x0e\x32%.spark.connect.MergeAction.ActionTypeR\nactionType\x12<\n\tcondition\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionH\x00R\tcondition\x88\x01\x01\x12G\n\x0b\x61ssignments\x18\x03 \x03(\x0b\x32%.spark.connect.MergeAction.AssignmentR\x0b\x61ssignments\x1aj\n\nAssignment\x12+\n\x03key\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value"\xa7\x01\n\nActionType\x12\x17\n\x13\x41\x43TION_TYPE_INVALID\x10\x00\x12\x16\n\x12\x41\x43TION_TYPE_DELETE\x10\x01\x12\x16\n\x12\x41\x43TION_TYPE_INSERT\x10\x02\x12\x1b\n\x17\x41\x43TION_TYPE_INSERT_STAR\x10\x03\x12\x16\n\x12\x41\x43TION_TYPE_UPDATE\x10\x04\x12\x1b\n\x17\x41\x43TION_TYPE_UPDATE_STAR\x10\x05\x42\x0c\n\n_condition"\xc5\x05\n\x12SubqueryExpression\x12\x17\n\x07plan_id\x18\x01 \x01(\x03R\x06planId\x12S\n\rsubquery_type\x18\x02 \x01(\x0e\x32..spark.connect.SubqueryExpression.SubqueryTypeR\x0csubqueryType\x12\x62\n\x11table_arg_options\x18\x03 \x01(\x0b\x32\x31.spark.connect.SubqueryExpression.TableArgOptionsH\x00R\x0ftableArgOptions\x88\x01\x01\x12G\n\x12in_subquery_values\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x10inSubqueryValues\x1a\xea\x01\n\x0fTableArgOptions\x12@\n\x0epartition_spec\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\rpartitionSpec\x12\x42\n\norder_spec\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\torderSpec\x12\x37\n\x15with_single_partition\x18\x03 \x01(\x08H\x00R\x13withSinglePartition\x88\x01\x01\x42\x18\n\x16_with_single_partition"\x90\x01\n\x0cSubqueryType\x12\x19\n\x15SUBQUERY_TYPE_UNKNOWN\x10\x00\x12\x18\n\x14SUBQUERY_TYPE_SCALAR\x10\x01\x12\x18\n\x14SUBQUERY_TYPE_EXISTS\x10\x02\x12\x1b\n\x17SUBQUERY_TYPE_TABLE_ARG\x10\x03\x12\x14\n\x10SUBQUERY_TYPE_IN\x10\x04\x42\x14\n\x12_table_arg_optionsB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' ) _globals = globals() @@ -135,27 +135,27 @@ _globals["_COMMONINLINEUSERDEFINEDFUNCTION"]._serialized_start = 7899 _globals["_COMMONINLINEUSERDEFINEDFUNCTION"]._serialized_end = 8296 _globals["_PYTHONUDF"]._serialized_start = 8299 - _globals["_PYTHONUDF"]._serialized_end = 8503 - _globals["_SCALARSCALAUDF"]._serialized_start = 8506 - _globals["_SCALARSCALAUDF"]._serialized_end = 8720 - _globals["_JAVAUDF"]._serialized_start = 8723 - _globals["_JAVAUDF"]._serialized_end = 8872 - _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_start = 8874 - _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_end = 8973 - _globals["_CALLFUNCTION"]._serialized_start = 8975 - _globals["_CALLFUNCTION"]._serialized_end = 9083 - _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_start = 9085 - _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_end = 9177 - _globals["_MERGEACTION"]._serialized_start = 9180 - _globals["_MERGEACTION"]._serialized_end = 9692 - _globals["_MERGEACTION_ASSIGNMENT"]._serialized_start = 9402 - _globals["_MERGEACTION_ASSIGNMENT"]._serialized_end = 9508 - _globals["_MERGEACTION_ACTIONTYPE"]._serialized_start = 9511 - _globals["_MERGEACTION_ACTIONTYPE"]._serialized_end = 9678 - _globals["_SUBQUERYEXPRESSION"]._serialized_start = 9695 - _globals["_SUBQUERYEXPRESSION"]._serialized_end = 10404 - _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_start = 10001 - _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_end = 10235 - _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_start = 10238 - _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_end = 10382 + _globals["_PYTHONUDF"]._serialized_end = 8582 + _globals["_SCALARSCALAUDF"]._serialized_start = 8585 + _globals["_SCALARSCALAUDF"]._serialized_end = 8799 + _globals["_JAVAUDF"]._serialized_start = 8802 + _globals["_JAVAUDF"]._serialized_end = 8951 + _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_start = 8953 + _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_end = 9052 + _globals["_CALLFUNCTION"]._serialized_start = 9054 + _globals["_CALLFUNCTION"]._serialized_end = 9162 + _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_start = 9164 + _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_end = 9256 + _globals["_MERGEACTION"]._serialized_start = 9259 + _globals["_MERGEACTION"]._serialized_end = 9771 + _globals["_MERGEACTION_ASSIGNMENT"]._serialized_start = 9481 + _globals["_MERGEACTION_ASSIGNMENT"]._serialized_end = 9587 + _globals["_MERGEACTION_ACTIONTYPE"]._serialized_start = 9590 + _globals["_MERGEACTION_ACTIONTYPE"]._serialized_end = 9757 + _globals["_SUBQUERYEXPRESSION"]._serialized_start = 9774 + _globals["_SUBQUERYEXPRESSION"]._serialized_end = 10483 + _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_start = 10080 + _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_end = 10314 + _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_start = 10317 + _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_end = 10461 # @@protoc_insertion_point(module_scope) diff --git a/python/pyspark/sql/connect/proto/expressions_pb2.pyi b/python/pyspark/sql/connect/proto/expressions_pb2.pyi index c613ade2f43f0..22b65357c2345 100644 --- a/python/pyspark/sql/connect/proto/expressions_pb2.pyi +++ b/python/pyspark/sql/connect/proto/expressions_pb2.pyi @@ -1826,6 +1826,7 @@ class PythonUDF(google.protobuf.message.Message): COMMAND_FIELD_NUMBER: builtins.int PYTHON_VER_FIELD_NUMBER: builtins.int ADDITIONAL_INCLUDES_FIELD_NUMBER: builtins.int + BUFFER_TYPE_FIELD_NUMBER: builtins.int @property def output_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: """(Required) Output type of the Python UDF""" @@ -1840,6 +1841,11 @@ class PythonUDF(google.protobuf.message.Message): self, ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: """(Optional) Additional includes for the Python UDF.""" + @property + def buffer_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType: + """(Optional) Intermediate buffer schema for an incremental Python aggregator + (see PythonAggregate). Set only for the incremental aggregator eval types. + """ def __init__( self, *, @@ -1848,15 +1854,28 @@ class PythonUDF(google.protobuf.message.Message): command: builtins.bytes = ..., python_ver: builtins.str = ..., additional_includes: collections.abc.Iterable[builtins.str] | None = ..., + buffer_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ..., ) -> None: ... def HasField( - self, field_name: typing_extensions.Literal["output_type", b"output_type"] + self, + field_name: typing_extensions.Literal[ + "_buffer_type", + b"_buffer_type", + "buffer_type", + b"buffer_type", + "output_type", + b"output_type", + ], ) -> builtins.bool: ... def ClearField( self, field_name: typing_extensions.Literal[ + "_buffer_type", + b"_buffer_type", "additional_includes", b"additional_includes", + "buffer_type", + b"buffer_type", "command", b"command", "eval_type", @@ -1867,6 +1886,9 @@ class PythonUDF(google.protobuf.message.Message): b"python_ver", ], ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_buffer_type", b"_buffer_type"] + ) -> typing_extensions.Literal["buffer_type"] | None: ... global___PythonUDF = PythonUDF diff --git a/python/pyspark/sql/connect/protobuf/functions.py b/python/pyspark/sql/connect/protobuf/functions.py index df6837e2bc809..120d896eaa551 100644 --- a/python/pyspark/sql/connect/protobuf/functions.py +++ b/python/pyspark/sql/connect/protobuf/functions.py @@ -19,12 +19,11 @@ A collections of builtin protobuf functions """ -from typing import Dict, Optional, TYPE_CHECKING - -from pyspark.sql.protobuf import functions as PyProtobufFunctions +from typing import TYPE_CHECKING, Dict, Optional from pyspark.sql.column import Column -from pyspark.sql.connect.functions.builtin import _invoke_function, _to_col, _options_to_col, lit +from pyspark.sql.connect.functions.builtin import _invoke_function, _options_to_col, _to_col, lit +from pyspark.sql.protobuf import functions as PyProtobufFunctions if TYPE_CHECKING: from pyspark.sql.connect._typing import ColumnOrName @@ -116,6 +115,7 @@ def _read_descriptor_set_file(filePath: str) -> bytes: def _test() -> None: import os import sys + from pyspark.testing.sqlutils import search_jar protobuf_jar = search_jar("connector/protobuf", "spark-protobuf-assembly-", "spark-protobuf") @@ -133,8 +133,9 @@ def _test() -> None: os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.protobuf.functions + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.protobuf.functions.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/readwriter.py b/python/pyspark/sql/connect/readwriter.py index 5c2c0c80ccdfb..561447a437e6d 100644 --- a/python/pyspark/sql/connect/readwriter.py +++ b/python/pyspark/sql/connect/readwriter.py @@ -14,38 +14,41 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Dict -from typing import Optional, Union, List, overload, Tuple, cast, Callable -from typing import TYPE_CHECKING +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast, overload +import pyspark.sql.connect.proto as proto +from pyspark.errors import ( + AnalysisException, + PySparkAttributeError, + PySparkTypeError, + PySparkValueError, +) +from pyspark.sql.connect.functions import builtin as F from pyspark.sql.connect.plan import ( - Read, - RelationChanges, DataSource, LogicalPlan, + Parse, + Read, + RelationChanges, WriteOperation, WriteOperationV2, - Parse, ) -import pyspark.sql.connect.proto as proto -from pyspark.sql.types import StructType -from pyspark.sql.utils import to_str from pyspark.sql.readwriter import ( - DataFrameWriter as PySparkDataFrameWriter, DataFrameReader as PySparkDataFrameReader, - DataFrameWriterV2 as PySparkDataFrameWriterV2, ) -from pyspark.errors import ( - AnalysisException, - PySparkAttributeError, - PySparkTypeError, - PySparkValueError, +from pyspark.sql.readwriter import ( + DataFrameWriter as PySparkDataFrameWriter, ) -from pyspark.sql.connect.functions import builtin as F +from pyspark.sql.readwriter import ( + DataFrameWriterV2 as PySparkDataFrameWriterV2, +) +from pyspark.sql.types import StructType +from pyspark.sql.utils import to_str if TYPE_CHECKING: - from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType + from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect.session import SparkSession from pyspark.sql.metrics import ExecutionInfo @@ -55,7 +58,10 @@ TupleOrListOfString = Union[List[str], Tuple[str, ...]] -class OptionUtils: +class OptionUtils(ABC): + @abstractmethod + def option(self, key: str, value: "OptionalPrimitiveType") -> Any: ... + def _set_opts( self, schema: Optional[Union[StructType, str]] = None, @@ -68,7 +74,7 @@ def _set_opts( self.schema(schema) # type: ignore[attr-defined] for k, v in options.items(): if v is not None: - self.option(k, v) # type: ignore[attr-defined] + self.option(k, v) class DataFrameReader(OptionUtils): @@ -130,15 +136,17 @@ def load( self.schema(schema) self.options(**options) - paths = path + paths: Optional[List[str]] if isinstance(path, str): paths = [path] + else: + paths = path plan = DataSource( format=self._format, schema=self._schema, options=self._options, - paths=paths, # type: ignore[arg-type] + paths=paths, ) return self._df(plan) @@ -1070,11 +1078,12 @@ def overwritePartitions(self) -> None: def _test() -> None: - import sys - import os import doctest - from pyspark.sql import SparkSession as PySparkSession + import os + import sys + import pyspark.sql.connect.readwriter + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.readwriter.__dict__.copy() diff --git a/python/pyspark/sql/connect/resource/profile.py b/python/pyspark/sql/connect/resource/profile.py index c97b75476a58d..300d5a054b799 100644 --- a/python/pyspark/sql/connect/resource/profile.py +++ b/python/pyspark/sql/connect/resource/profile.py @@ -15,11 +15,10 @@ # limitations under the License. # -from typing import Optional, Dict - -from pyspark.resource import ExecutorResourceRequest, TaskResourceRequest +from typing import Dict, Optional import pyspark.sql.connect.proto as pb2 +from pyspark.resource import ExecutorResourceRequest, TaskResourceRequest class ResourceProfile: diff --git a/python/pyspark/sql/connect/session.py b/python/pyspark/sql/connect/session.py index 1a3e1c1aebb39..180a36569f704 100644 --- a/python/pyspark/sql/connect/session.py +++ b/python/pyspark/sql/connect/session.py @@ -14,104 +14,104 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import uuid - +import functools import json -import threading import os import sys +import threading +import urllib +import uuid import warnings from collections.abc import Callable, Sized -import functools from threading import RLock from types import TracebackType from typing import ( - Optional, + TYPE_CHECKING, Any, - Iterator, - Union, + ClassVar, Dict, + Iterable, + Iterator, List, + Mapping, + Optional, + Set, Tuple, Type, - Set, + Union, cast, overload, - Iterable, - Mapping, - TYPE_CHECKING, - ClassVar, ) import numpy as np import pandas as pd import pyarrow as pa from pandas.api.types import is_datetime64_dtype, is_timedelta64_dtype -import urllib +from pyspark.errors import ( + AnalysisException, + PySparkAssertionError, + PySparkAttributeError, + PySparkNotImplementedError, + PySparkRuntimeError, + PySparkTypeError, + PySparkValueError, +) +from pyspark.sql.connect.client import DefaultChannelBuilder, SparkConnectClient +from pyspark.sql.connect.conf import RuntimeConf from pyspark.sql.connect.dataframe import DataFrame -from pyspark.sql.dataframe import DataFrame as ParentDataFrame +from pyspark.sql.connect.functions import builtin as F from pyspark.sql.connect.logging import logger -from pyspark.sql.connect.client import SparkConnectClient, DefaultChannelBuilder -from pyspark.sql.connect.conf import RuntimeConf from pyspark.sql.connect.plan import ( SQL, - Range, - LocalRelation, - LogicalPlan, - ChunkedCachedLocalRelation, CachedRelation, CachedRemoteRelation, + ChunkedCachedLocalRelation, + LocalRelation, + LogicalPlan, + Range, SubqueryAlias, ) -from pyspark.sql.connect.functions import builtin as F from pyspark.sql.connect.profiler import ProfilerCollector from pyspark.sql.connect.readwriter import DataFrameReader -from pyspark.sql.connect.streaming.readwriter import DataStreamReader from pyspark.sql.connect.streaming.query import StreamingQueryManager +from pyspark.sql.connect.streaming.readwriter import DataStreamReader +from pyspark.sql.dataframe import DataFrame as ParentDataFrame from pyspark.sql.pandas.conversion import create_arrow_table_from_pandas from pyspark.sql.pandas.types import ( - to_arrow_schema, + _check_arrow_table_timestamps_localize, _deduplicate_field_names, from_arrow_schema, from_arrow_type, - _check_arrow_table_timestamps_localize, + to_arrow_schema, ) from pyspark.sql.profiler import Profile -from pyspark.sql.session import classproperty, SparkSession as PySparkSession +from pyspark.sql.session import SparkSession as PySparkSession +from pyspark.sql.session import classproperty from pyspark.sql.types import ( - _infer_schema, - _has_nulltype, - _merge_type, - Row, + AtomicType, DataType, DayTimeIntervalType, - StructType, - AtomicType, - TimestampType, MapType, + Row, StringType, + StructType, + TimestampType, + _has_nulltype, + _infer_schema, + _merge_type, ) from pyspark.sql.utils import to_str -from pyspark.errors import ( - AnalysisException, - PySparkAttributeError, - PySparkNotImplementedError, - PySparkRuntimeError, - PySparkValueError, - PySparkTypeError, - PySparkAssertionError, -) if TYPE_CHECKING: import pyspark.sql.connect.proto as pb2 from pyspark.sql.connect._typing import OptionalPrimitiveType from pyspark.sql.connect.catalog import Catalog + from pyspark.sql.connect.datasource import DataSourceRegistration + from pyspark.sql.connect.shell.progress import ProgressHandler + from pyspark.sql.connect.tvf import TableValuedFunction from pyspark.sql.connect.udf import UDFRegistration from pyspark.sql.connect.udtf import UDTFRegistration - from pyspark.sql.connect.tvf import TableValuedFunction - from pyspark.sql.connect.shell.progress import ProgressHandler - from pyspark.sql.connect.datasource import DataSourceRegistration class SparkSession: @@ -310,8 +310,19 @@ def __init__( ) self._session_id = self._client._session_id + self._initialize_lifecycle_state() + + def _initialize_lifecycle_state(self) -> None: + """Initialize state shared by normally constructed and derived sessions.""" + # Set to false to prevent client.release_session on close() (testing only) self.release_session_on_close = True + self._on_stop_callbacks: List[Callable[[], None]] = [] + + def _register_on_stop_callback(self, callback: Callable[[], None]) -> None: + """Register internal lifecycle cleanup owned by the code that created this session.""" + if callback not in self._on_stop_callbacks: + self._on_stop_callbacks.append(callback) @classmethod def _set_default_and_active_session(cls, session: "SparkSession") -> None: @@ -991,6 +1002,13 @@ def stop(self) -> None: if "SPARK_REMOTE" in os.environ: del os.environ["SPARK_REMOTE"] + callbacks, self._on_stop_callbacks = self._on_stop_callbacks, [] + for callback in callbacks: + try: + callback() + except Exception as e: + logger.warning(f"session.stop(): Cleanup callback failed. Error: {e}") + def __enter__(self) -> "SparkSession": """ Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. @@ -1232,7 +1250,7 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None: Returns the authentication token that should be used to connect to this session. """ - from pyspark import SparkContext, SparkConf + from pyspark import SparkConf, SparkContext session = PySparkSession._instantiatedSession if session is None or session._sc._jsc is None: @@ -1370,7 +1388,7 @@ def cloneSession(self, new_session_id: Optional[str] = None) -> "SparkSession": new_session = object.__new__(SparkSession) new_session._client = cloned_client new_session._session_id = cloned_client._session_id - new_session.release_session_on_close = True + new_session._initialize_lifecycle_state() return new_session def newSession(self) -> "SparkSession": @@ -1398,7 +1416,7 @@ def newSession(self) -> "SparkSession": new_session = object.__new__(SparkSession) new_session._client = new_client new_session._session_id = new_client._session_id - new_session.release_session_on_close = True + new_session._initialize_lifecycle_state() return new_session @@ -1406,11 +1424,12 @@ def newSession(self) -> "SparkSession": def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.session + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.session.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/connect/shell/progress.py b/python/pyspark/sql/connect/shell/progress.py index ada3b4d177024..6391474850158 100644 --- a/python/pyspark/sql/connect/shell/progress.py +++ b/python/pyspark/sql/connect/shell/progress.py @@ -18,12 +18,12 @@ """Implementation of a progress bar that is displayed while a query is running.""" import abc -from dataclasses import dataclass -import time import sys +import time import typing +from dataclasses import dataclass from types import TracebackType -from typing import Iterable, Any +from typing import Any, Iterable from pyspark.sql.connect.proto import ExecutePlanResponse diff --git a/python/pyspark/sql/connect/sql_formatter.py b/python/pyspark/sql/connect/sql_formatter.py index 8fced80081ad1..6ab85c6145c73 100644 --- a/python/pyspark/sql/connect/sql_formatter.py +++ b/python/pyspark/sql/connect/sql_formatter.py @@ -17,14 +17,14 @@ import string import typing -from typing import Any, Optional, List, Tuple, Sequence, Mapping import uuid +from typing import Any, List, Mapping, Optional, Sequence, Tuple from pyspark.errors import PySparkValueError if typing.TYPE_CHECKING: - from pyspark.sql.connect.session import SparkSession from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.session import SparkSession class SQLStringFormatter(string.Formatter): @@ -46,8 +46,8 @@ def _convert_value(self, val: Any, field_name: str) -> Optional[str]: """ Converts the given value into a SQL string. """ - from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect.column import Column + from pyspark.sql.connect.dataframe import DataFrame from pyspark.sql.connect.expressions import ColumnReference from pyspark.sql.utils import get_lit_sql_str diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index 8418e080fa7aa..e5d16073f4d5e 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -17,26 +17,28 @@ import json import sys import warnings -from typing import TYPE_CHECKING, Any, cast, Dict, List, Optional, Union, Iterator -from threading import Thread, Lock +from threading import Lock, Thread +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union, cast -from pyspark.errors import StreamingQueryException, PySparkValueError import pyspark.sql.connect.proto as pb2 +from pyspark.errors import PySparkValueError, StreamingQueryException +from pyspark.errors.exceptions.connect import ( + StreamingQueryException as CapturedStreamingQueryException, +) from pyspark.sql.connect import proto from pyspark.sql.streaming import StreamingQueryListener from pyspark.sql.streaming.listener import ( - QueryStartedEvent, - QueryProgressEvent, QueryIdleEvent, + QueryProgressEvent, + QueryStartedEvent, QueryTerminatedEvent, StreamingQueryProgress, ) from pyspark.sql.streaming.query import ( StreamingQuery as PySparkStreamingQuery, - StreamingQueryManager as PySparkStreamingQueryManager, ) -from pyspark.errors.exceptions.connect import ( - StreamingQueryException as CapturedStreamingQueryException, +from pyspark.sql.streaming.query import ( + StreamingQueryManager as PySparkStreamingQueryManager, ) if TYPE_CHECKING: @@ -432,8 +434,9 @@ def post_to_all( def _test() -> None: import doctest import os - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.streaming.query + from pyspark.sql import SparkSession as PySparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/connect/streaming/readwriter.py b/python/pyspark/sql/connect/streaming/readwriter.py index 130844309ae4c..6cb19c6b37bc9 100644 --- a/python/pyspark/sql/connect/streaming/readwriter.py +++ b/python/pyspark/sql/connect/streaming/readwriter.py @@ -15,11 +15,18 @@ # limitations under the License. # import json +import pickle import re import sys -import pickle -from typing import cast, overload, Callable, Dict, List, Optional, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union, cast, overload +import pyspark.sql.connect.proto as pb2 +from pyspark.errors import ( + AnalysisException, + PySparkPicklingError, + PySparkTypeError, + PySparkValueError, +) from pyspark.serializers import CloudPickleSerializer from pyspark.sql.connect.plan import ( DataSource, @@ -28,28 +35,23 @@ RelationChanges, WriteStreamOperation, ) -import pyspark.sql.connect.proto as pb2 from pyspark.sql.connect.readwriter import OptionUtils, to_str from pyspark.sql.connect.streaming.query import StreamingQuery +from pyspark.sql.connect.utils import get_python_ver +from pyspark.sql.streaming.listener import QueryStartedEvent from pyspark.sql.streaming.readwriter import ( DataStreamReader as PySparkDataStreamReader, +) +from pyspark.sql.streaming.readwriter import ( DataStreamWriter as PySparkDataStreamWriter, ) -from pyspark.sql.streaming.listener import QueryStartedEvent -from pyspark.sql.connect.utils import get_python_ver from pyspark.sql.types import Row, StructType -from pyspark.errors import ( - AnalysisException, - PySparkTypeError, - PySparkValueError, - PySparkPicklingError, -) if TYPE_CHECKING: - from pyspark.sql.connect.session import SparkSession + from pyspark.sql._typing import SupportsProcess from pyspark.sql.connect._typing import OptionalPrimitiveType from pyspark.sql.connect.dataframe import DataFrame - from pyspark.sql._typing import SupportsProcess + from pyspark.sql.connect.session import SparkSession class DataStreamReader(OptionUtils): @@ -640,7 +642,7 @@ def foreach(self, f: Callable[[Row], None]) -> "DataStreamWriter": ... def foreach(self, f: "SupportsProcess") -> "DataStreamWriter": ... def foreach(self, f: Union[Callable[[Row], None], "SupportsProcess"]) -> "DataStreamWriter": - from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer + from pyspark.serializers import AutoBatchedSerializer, CPickleSerializer func = PySparkDataStreamWriter._construct_foreach_function(f) serializer = AutoBatchedSerializer(CPickleSerializer()) @@ -766,11 +768,12 @@ def toTable( def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.streaming.readwriter + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.readwriter.__dict__.copy() diff --git a/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py b/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py index 5aa7ba29ad6cf..998805aa8cfc1 100644 --- a/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py +++ b/python/pyspark/sql/connect/streaming/worker/foreach_batch_worker.py @@ -21,19 +21,18 @@ """ import os +from typing import IO -from pyspark.worker_util import get_sock_file_to_executor +from pyspark import worker from pyspark.serializers import ( - write_int, - read_long, - UTF8Deserializer, CPickleSerializer, + UTF8Deserializer, + read_long, + write_int, ) -from pyspark import worker from pyspark.sql.connect.session import SparkSession from pyspark.util import handle_worker_exception -from typing import IO -from pyspark.worker_util import check_python_version +from pyspark.worker_util import check_python_version, get_sock_file_to_executor pickle_ser = CPickleSerializer() utf8_deserializer = UTF8Deserializer() diff --git a/python/pyspark/sql/connect/streaming/worker/listener_worker.py b/python/pyspark/sql/connect/streaming/worker/listener_worker.py index 1d2776d2cf4f3..166f41d4f02db 100644 --- a/python/pyspark/sql/connect/streaming/worker/listener_worker.py +++ b/python/pyspark/sql/connect/streaming/worker/listener_worker.py @@ -20,28 +20,26 @@ Usually this is ran on the driver side of the Spark Connect Server. """ -import os import json +import os +from typing import IO -from pyspark.worker_util import get_sock_file_to_executor +from pyspark import worker from pyspark.serializers import ( + CPickleSerializer, + UTF8Deserializer, read_int, write_int, - UTF8Deserializer, - CPickleSerializer, ) -from pyspark import worker from pyspark.sql.connect.session import SparkSession -from pyspark.util import handle_worker_exception -from typing import IO - from pyspark.sql.streaming.listener import ( - QueryStartedEvent, + QueryIdleEvent, QueryProgressEvent, + QueryStartedEvent, QueryTerminatedEvent, - QueryIdleEvent, ) -from pyspark.worker_util import check_python_version +from pyspark.util import handle_worker_exception +from pyspark.worker_util import check_python_version, get_sock_file_to_executor pickle_ser = CPickleSerializer() utf8_deserializer = UTF8Deserializer() diff --git a/python/pyspark/sql/connect/table_arg.py b/python/pyspark/sql/connect/table_arg.py index 789f241de509b..b0fc7012036a2 100644 --- a/python/pyspark/sql/connect/table_arg.py +++ b/python/pyspark/sql/connect/table_arg.py @@ -16,22 +16,21 @@ # from typing import ( - Iterable, TYPE_CHECKING, - Union, - Sequence, + Iterable, List, + Sequence, Tuple, + Union, cast, ) import pyspark.sql.connect.proto as proto +from pyspark.errors import IllegalArgumentException from pyspark.sql.column import Column -from pyspark.sql.table_arg import TableArg as ParentTableArg -from pyspark.sql.connect.expressions import Expression, SubqueryExpression, SortOrder +from pyspark.sql.connect.expressions import Expression, SortOrder, SubqueryExpression from pyspark.sql.connect.functions import builtin as F - -from pyspark.errors import IllegalArgumentException +from pyspark.sql.table_arg import TableArg as ParentTableArg if TYPE_CHECKING: from pyspark.sql._typing import ColumnOrName diff --git a/python/pyspark/sql/connect/tvf.py b/python/pyspark/sql/connect/tvf.py index b313ce15d9f00..395ade38bdc5c 100644 --- a/python/pyspark/sql/connect/tvf.py +++ b/python/pyspark/sql/connect/tvf.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Optional from pyspark.errors import PySparkValueError from pyspark.sql.tvf import TableValuedFunction as PySparkTableValuedFunction @@ -116,8 +116,8 @@ def python_worker_logs(self) -> "DataFrame": def _fn(self, name: str, *args: "Column") -> "DataFrame": from pyspark.sql.connect.dataframe import DataFrame - from pyspark.sql.connect.plan import UnresolvedTableValuedFunction from pyspark.sql.connect.functions.builtin import _to_col + from pyspark.sql.connect.plan import UnresolvedTableValuedFunction return DataFrame( UnresolvedTableValuedFunction(name, [_to_col(arg) for arg in args]), self._sparkSession @@ -139,8 +139,9 @@ def _test() -> None: sys.exit(0) import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.connect.tvf + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.connect.tvf.__dict__.copy() diff --git a/python/pyspark/sql/connect/types.py b/python/pyspark/sql/connect/types.py index e8b292eb090bc..78270ca4c2280 100644 --- a/python/pyspark/sql/connect/types.py +++ b/python/pyspark/sql/connect/types.py @@ -15,44 +15,42 @@ # limitations under the License. # import json +from typing import Any, Dict, List, Optional -from typing import Any, Dict, Optional, List - +import pyspark.sql.connect.proto as pb2 +from pyspark.errors import PySparkAssertionError, PySparkValueError from pyspark.sql.types import ( - DataType, + ArrayType, + BinaryType, + BooleanType, ByteType, - ShortType, - IntegerType, - FloatType, - DateType, - TimeType, - TimestampType, - TimestampNTZType, - DayTimeIntervalType, - YearMonthIntervalType, CalendarIntervalType, - MapType, - StringType, CharType, - VarcharType, - StructType, - StructField, - ArrayType, + DataType, + DateType, + DayTimeIntervalType, + DecimalType, DoubleType, + FloatType, + GeographyType, + GeometryType, + IntegerType, LongType, - DecimalType, - BinaryType, - BooleanType, + MapType, NullType, NumericType, - VariantType, - GeographyType, - GeometryType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, + TimeType, UserDefinedType, + VarcharType, + VariantType, + YearMonthIntervalType, ) -from pyspark.errors import PySparkAssertionError, PySparkValueError - -import pyspark.sql.connect.proto as pb2 class UnparsedDataType(DataType): diff --git a/python/pyspark/sql/connect/udf.py b/python/pyspark/sql/connect/udf.py index c9848634eb303..9eb2bf5e05107 100644 --- a/python/pyspark/sql/connect/udf.py +++ b/python/pyspark/sql/connect/udf.py @@ -18,12 +18,13 @@ User-defined function related classes and functions """ -import warnings -import sys import functools -from typing import cast, Callable, Any, List, TYPE_CHECKING, Optional, Union +import sys +import warnings +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union, cast -from pyspark.util import PythonEvalType +from pyspark.errors import PySparkRuntimeError, PySparkTypeError +from pyspark.sql.connect.column import Column from pyspark.sql.connect.expressions import ( ColumnReference, CommonInlineUserDefinedFunction, @@ -31,14 +32,15 @@ NamedArgumentExpression, PythonUDF, ) -from pyspark.sql.connect.column import Column -from pyspark.sql.types import DataType, StringType, _parse_datatype_string +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version +from pyspark.sql.types import DataType, StringType, StructType, _parse_datatype_string from pyspark.sql.udf import ( UDFRegistration as PySparkUDFRegistration, +) +from pyspark.sql.udf import ( UserDefinedFunction as PySparkUserDefinedFunction, ) -from pyspark.sql.pandas.utils import require_minimum_pyarrow_version, require_minimum_pandas_version -from pyspark.errors import PySparkTypeError, PySparkRuntimeError +from pyspark.util import PythonEvalType if TYPE_CHECKING: from pyspark.sql.connect._typing import ( @@ -113,10 +115,16 @@ def _create_udf( evalType: int, name: Optional[str] = None, deterministic: bool = True, + bufferSchema: Optional[StructType] = None, ) -> "UserDefinedFunctionLike": # Set the name of the UserDefinedFunction object to be the name of function f udf_obj = UserDefinedFunction( - f, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + f, + returnType=returnType, + name=name, + evalType=evalType, + deterministic=deterministic, + bufferSchema=bufferSchema, ) return udf_obj._wrapped() @@ -139,6 +147,7 @@ def __init__( name: Optional[str] = None, evalType: int = PythonEvalType.SQL_BATCHED_UDF, deterministic: bool = True, + bufferSchema: Optional[StructType] = None, ): if not callable(func): raise PySparkTypeError( @@ -178,6 +187,10 @@ def __init__( ) self.evalType = evalType self.deterministic = deterministic + # Intermediate aggregation buffer schema, set only for an incremental Python aggregator + # (see :class:`pyspark.sql.aggregator.Aggregator`); ``None`` otherwise. A first-class field + # so it survives ``_wrapped()``, ``asNondeterministic()`` and ``spark.udf.register``. + self.bufferSchema = bufferSchema @property def returnType(self) -> DataType: @@ -210,6 +223,8 @@ def to_expr(col: "ColumnOrName") -> Expression: eval_type=self.evalType, func=self.func, python_ver="%d.%d" % sys.version_info[:2], + # Set for incremental Python aggregators (see pyspark.sql.aggregator). + buffer_type=self.bufferSchema, ) return CommonInlineUserDefinedFunction( function_name=self._name, @@ -253,6 +268,7 @@ def wrapper(*args: "ColumnOrName", **kwargs: "ColumnOrName") -> Column: wrapper.returnType = self.returnType # type: ignore[attr-defined] wrapper.evalType = self.evalType # type: ignore[attr-defined] wrapper.deterministic = self.deterministic # type: ignore[attr-defined] + wrapper.bufferSchema = self.bufferSchema # type: ignore[attr-defined] wrapper.asNondeterministic = functools.wraps( # type: ignore[attr-defined] self.asNondeterministic )(lambda: self.asNondeterministic()._wrapped()) @@ -303,6 +319,7 @@ def register( PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, ]: raise PySparkTypeError( errorClass="INVALID_UDF_EVAL_TYPE", @@ -311,11 +328,18 @@ def register( "SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, " "SQL_SCALAR_PANDAS_ITER_UDF, SQL_SCALAR_ARROW_ITER_UDF, " "SQL_GROUPED_AGG_PANDAS_UDF, SQL_GROUPED_AGG_ARROW_UDF, " - "SQL_GROUPED_AGG_PANDAS_ITER_UDF or SQL_GROUPED_AGG_ARROW_ITER_UDF" + "SQL_GROUPED_AGG_PANDAS_ITER_UDF, SQL_GROUPED_AGG_ARROW_ITER_UDF " + "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF" }, ) self.sparkSession._client.register_udf( - f.func, f.returnType, name, f.evalType, f.deterministic + f.func, + f.returnType, + name, + f.evalType, + f.deterministic, + # Set for the incremental aggregator (see pyspark.sql.aggregator). + buffer_type=getattr(f, "bufferSchema", None), ) return f else: diff --git a/python/pyspark/sql/connect/udtf.py b/python/pyspark/sql/connect/udtf.py index fa99c331084a4..aa7d86c0f72d7 100644 --- a/python/pyspark/sql/connect/udtf.py +++ b/python/pyspark/sql/connect/udtf.py @@ -19,9 +19,9 @@ """ import warnings -from typing import List, Type, TYPE_CHECKING, Optional, Union, Any +from typing import TYPE_CHECKING, Any, List, Optional, Type, Union -from pyspark.util import PythonEvalType +from pyspark.errors import PySparkAttributeError, PySparkRuntimeError, PySparkTypeError from pyspark.sql.connect.column import Column from pyspark.sql.connect.expressions import ColumnReference, Expression, NamedArgumentExpression from pyspark.sql.connect.plan import ( @@ -31,11 +31,11 @@ from pyspark.sql.connect.table_arg import TableArg from pyspark.sql.connect.types import UnparsedDataType from pyspark.sql.connect.utils import get_python_ver -from pyspark.sql.pandas.utils import require_minimum_pyarrow_version, require_minimum_pandas_version -from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401 -from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration, _validate_udtf_handler +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version from pyspark.sql.types import DataType, StructType -from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkAttributeError +from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult, _validate_udtf_handler # noqa: F401 +from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration +from pyspark.util import PythonEvalType if TYPE_CHECKING: from pyspark.sql.connect._typing import ColumnOrName @@ -196,8 +196,8 @@ def to_expr(col: "ColumnOrName") -> Expression: ) def __call__(self, *args: "ColumnOrName", **kwargs: "ColumnOrName") -> "DataFrame": - from pyspark.sql.connect.session import SparkSession from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.session import SparkSession session = SparkSession.active() diff --git a/python/pyspark/sql/connect/utils.py b/python/pyspark/sql/connect/utils.py index 2c3130f2c3eee..1861770bb4cae 100644 --- a/python/pyspark/sql/connect/utils.py +++ b/python/pyspark/sql/connect/utils.py @@ -16,9 +16,9 @@ # import sys +from pyspark.errors import PySparkImportError from pyspark.loose_version import LooseVersion from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version -from pyspark.errors import PySparkImportError def check_dependencies() -> None: @@ -29,7 +29,7 @@ def check_dependencies() -> None: and not hasattr(sys, "ps1") ): # The main module is not initialized at all at this point. We must be running doctests. - from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + from pyspark.testing.connectutils import connect_requirement_message, should_test_connect if not should_test_connect: print( diff --git a/python/pyspark/sql/connect/window.py b/python/pyspark/sql/connect/window.py index 55a43c2158a3b..3dd9b5cbdcfae 100644 --- a/python/pyspark/sql/connect/window.py +++ b/python/pyspark/sql/connect/window.py @@ -14,15 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import TYPE_CHECKING, Any, Union, Sequence, List, Optional, Tuple, cast, Iterable +from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Sequence, Tuple, Union, cast from pyspark.sql.column import Column +from pyspark.sql.connect.expressions import Expression, SortOrder +from pyspark.sql.connect.functions import builtin as F from pyspark.sql.window import ( Window as ParentWindow, +) +from pyspark.sql.window import ( WindowSpec as ParentWindowSpec, ) -from pyspark.sql.connect.expressions import Expression, SortOrder -from pyspark.sql.connect.functions import builtin as F if TYPE_CHECKING: from pyspark.sql.connect._typing import ColumnOrName @@ -154,11 +156,12 @@ def rangeBetween(start: int, end: int) -> "WindowSpec": def _test() -> None: + import doctest import os import sys - import doctest - from pyspark.sql import SparkSession as PySparkSession + import pyspark.sql.window + from pyspark.sql import SparkSession as PySparkSession globs = pyspark.sql.window.__dict__.copy() globs["spark"] = ( diff --git a/python/pyspark/sql/context.py b/python/pyspark/sql/context.py index 01d8eb5583c93..07af10e75627a 100644 --- a/python/pyspark/sql/context.py +++ b/python/pyspark/sql/context.py @@ -18,39 +18,39 @@ import sys import warnings from typing import ( - Optional, - Union, - Callable, + TYPE_CHECKING, Any, + Callable, + ClassVar, Iterable, List, + Optional, Tuple, - overload, Type, - ClassVar, - TYPE_CHECKING, + Union, cast, + overload, ) from pyspark import _NoValue from pyspark._globals import _NoValueType from pyspark.errors import PySparkNotImplementedError, PySparkValueError -from pyspark.sql.session import _monkey_patch_RDD, SparkSession +from pyspark.errors.exceptions.captured import install_exception_handler from pyspark.sql.dataframe import DataFrame from pyspark.sql.readwriter import DataFrameReader -from pyspark.sql.streaming import DataStreamReader +from pyspark.sql.session import SparkSession, _monkey_patch_RDD +from pyspark.sql.streaming import DataStreamReader, StreamingQueryManager +from pyspark.sql.streaming.query import StreamingCheckpointManager +from pyspark.sql.types import AtomicType, DataType, StructType from pyspark.sql.udf import UDFRegistration from pyspark.sql.udtf import UDTFRegistration -from pyspark.errors.exceptions.captured import install_exception_handler -from pyspark.sql.types import AtomicType, DataType, StructType -from pyspark.sql.streaming import StreamingQueryManager -from pyspark.sql.streaming.query import StreamingCheckpointManager if TYPE_CHECKING: - from py4j.java_gateway import JavaObject import pyarrow as pa - from pyspark.core.rdd import RDD + from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext + from pyspark.core.rdd import RDD from pyspark.sql._typing import ( AtomicValue, RowLike, @@ -816,12 +816,13 @@ def refreshTable(self, tableName: str) -> None: def _test() -> None: - import os import doctest + import os import tempfile + + import pyspark.sql.context from pyspark.core.context import SparkContext from pyspark.sql import Row, SQLContext - import pyspark.sql.context os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index e472c7304e2e6..0585a730fb087 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -24,11 +24,11 @@ import pyspark from pyspark.errors import PySparkNotImplementedError, PySparkRuntimeError, PySparkValueError from pyspark.sql.pandas.types import ( + _create_converter_to_pandas, _dedup_names, _deduplicate_field_names, - _create_converter_to_pandas, - to_arrow_schema, from_arrow_schema, + to_arrow_schema, ) from pyspark.sql.pandas.utils import require_minimum_pyarrow_version from pyspark.sql.types import ( @@ -36,39 +36,39 @@ BinaryType, BooleanType, ByteType, - ShortType, - IntegerType, - LongType, DataType, - FloatType, - DoubleType, + DateType, + DayTimeIntervalType, DecimalType, - GeographyType, + DoubleType, + FloatType, Geography, - GeometryType, + GeographyType, Geometry, + GeometryType, + IntegerType, + LongType, MapType, NullType, Row, + ShortType, StringType, StructField, StructType, - DateType, - TimeType, TimestampNTZType, TimestampType, - DayTimeIntervalType, - YearMonthIntervalType, + TimeType, UserDefinedType, VariantType, VariantVal, + YearMonthIntervalType, _create_row, _has_type, ) if TYPE_CHECKING: - import pyarrow as pa import pandas as pd + import pyarrow as pa class ArrowBatchTransformer: @@ -359,11 +359,11 @@ def convert( ------- pa.RecordBatch """ - import pyarrow as pa import pandas as pd + import pyarrow as pa from pyspark.errors import PySparkTypeError, PySparkValueError - from pyspark.sql.pandas.types import to_arrow_type, _create_converter_from_pandas + from pyspark.sql.pandas.types import _create_converter_from_pandas, to_arrow_type # Handle empty schema (0 columns) # Use dummy column + select([]) to preserve row count (PyArrow limitation workaround) @@ -503,7 +503,10 @@ def convert_column( ) raise PySparkValueError(error_msg) from e - arrays = [convert_column(col, field) for col, field in zip(columns, schema.fields)] + converted = [convert_column(col, field) for col, field in zip(columns, schema.fields)] + # pa.Array.from_pandas returns a pa.ChunkedArray for a chunked arrow-backed Series + # (e.g. a pyarrow-backed extension dtype), which pa.RecordBatch.from_arrays rejects. + arrays = [a.combine_chunks() if isinstance(a, pa.ChunkedArray) else a for a in converted] return pa.RecordBatch.from_arrays(arrays, schema.names) @@ -682,6 +685,27 @@ def convert_array(value: Any) -> Any: assert isinstance(value, (list, array.array)) return list(value) + elif isinstance(dataType.elementType, (StringType, BinaryType)): + # Inline the scalar identity fast path so elements that are + # already the target Python type skip the per-element converter + # call entirely: `convert_string`/`convert_binary` return such + # elements unchanged. `str` and immutable `bytes` are the two + # element types whose converter is a no-op on a matching value. + # Any other element -- including `None` (whose nullability is + # enforced by `element_conv`) and values that need coercion (e.g. + # a bool to string) -- falls back to `element_conv`, reused + # unchanged. + fast_type = str if isinstance(dataType.elementType, StringType) else bytes + + def convert_array(value: Any) -> Any: + if value is None: + if not nullable: + raise PySparkValueError(f"input for {dataType} must not be None") + return None + else: + assert isinstance(value, (list, array.array)) + return [v if type(v) is fast_type else element_conv(v) for v in value] + else: def convert_array(value: Any) -> Any: @@ -739,6 +763,12 @@ def convert_binary(value: Any) -> Any: if not nullable: raise PySparkValueError(f"input for {dataType} must not be None") return None + elif type(value) is bytes: + # Fast path: `bytes(value)` returns `value` itself for a `bytes` + # input (no copy, as `bytes` is immutable), but still pays the + # constructor dispatch per element. Returning it directly skips + # that. `bytearray` falls through and is copied into `bytes`. + return value else: assert isinstance(value, (bytes, bytearray)) return bytes(value) @@ -801,13 +831,19 @@ def convert_string(value: Any) -> Any: if not nullable: raise PySparkValueError(f"input for {dataType} must not be None") return None + elif type(value) is str: + # Fast path: `str(value)` returns `value` itself for a `str` + # input (no copy), but still pays the constructor dispatch per + # element. Returning it directly skips that and the bool check. + return value + elif value is True: + # To match the PySpark Classic which convert bool to string in + # the JVM side (python.EvaluatePython.makeFromJava) + return "true" + elif value is False: + return "false" else: - if isinstance(value, bool): - # To match the PySpark Classic which convert bool to string in - # the JVM side (python.EvaluatePython.makeFromJava) - return str(value).lower() - else: - return str(value) + return str(value) return convert_string @@ -997,6 +1033,7 @@ def _should_manual_bulk() -> bool: the minimum supported PyArrow version contains the fix. """ import pyarrow as pa + from pyspark.loose_version import LooseVersion if LooseVersion(pa.__version__) >= LooseVersion("25.0.1"): @@ -1616,8 +1653,8 @@ def localize_tz( doesn't need this conversion. """ import pyarrow as pa - import pyarrow.types as types import pyarrow.compute as pc + import pyarrow.types as types def check_type_func(pa_type: pa.DataType) -> bool: # match timezone-aware TimestampType @@ -1653,8 +1690,8 @@ def preprocess_time( 2, coerce_temporal_nanoseconds: coerce timestamp time units to nanoseconds """ import pyarrow as pa - import pyarrow.types as types import pyarrow.compute as pc + import pyarrow.types as types def check_type_func(pa_type: pa.DataType) -> bool: return types.is_timestamp(pa_type) and (pa_type.unit != "ns" or pa_type.tz is not None) @@ -1798,8 +1835,8 @@ def convert_legacy( This method handles date type columns specially to avoid overflow issues with datetime64[ns] intermediate representations. """ - import pyarrow as pa import pandas as pd + import pyarrow as pa assert isinstance(arr, (pa.Array, pa.ChunkedArray)) @@ -1894,8 +1931,8 @@ def convert_numpy( prefer_int_ext_dtype: bool = False, df_for_struct: bool = False, ) -> Union["pd.Series", "pd.DataFrame"]: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert isinstance(arr, (pa.Array, pa.ChunkedArray)) diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py index 6c4d32ea1797f..d92290b8443e2 100644 --- a/python/pyspark/sql/dataframe.py +++ b/python/pyspark/sql/dataframe.py @@ -17,9 +17,10 @@ # mypy: disable-error-code="empty-body" -import sys import random +import sys from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -30,28 +31,28 @@ Tuple, Union, overload, - TYPE_CHECKING, ) from pyspark import _NoValue from pyspark._globals import _NoValueType -from pyspark.util import is_remote_only -from pyspark.storagelevel import StorageLevel from pyspark.resource import ResourceProfile from pyspark.sql.column import Column -from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.merge import MergeIntoWriter +from pyspark.sql.readwriter import DataFrameWriter, DataFrameWriterV2 from pyspark.sql.streaming import DataStreamWriter from pyspark.sql.table_arg import TableArg -from pyspark.sql.types import StructType, Row +from pyspark.sql.types import Row, StructType from pyspark.sql.utils import dispatch_df_method +from pyspark.storagelevel import StorageLevel +from pyspark.util import is_remote_only if TYPE_CHECKING: - from py4j.java_gateway import JavaObject import pyarrow as pa + from py4j.java_gateway import JavaObject + + from pyspark._typing import PrimitiveType from pyspark.core.context import SparkContext from pyspark.core.rdd import RDD - from pyspark._typing import PrimitiveType from pyspark.pandas.frame import DataFrame as PandasOnSparkDataFrame from pyspark.sql._typing import ( ColumnOrName, @@ -60,16 +61,18 @@ OptionalPrimitiveType, ) from pyspark.sql.context import SQLContext - from pyspark.sql.session import SparkSession from pyspark.sql.group import GroupedData + from pyspark.sql.metrics import ExecutionInfo from pyspark.sql.observation import Observation from pyspark.sql.pandas._typing import ( - PandasMapIterFunction, ArrowMapIterFunction, + PandasMapIterFunction, + ) + from pyspark.sql.pandas._typing import ( DataFrameLike as PandasDataFrameLike, ) from pyspark.sql.plot import PySparkPlotAccessor - from pyspark.sql.metrics import ExecutionInfo + from pyspark.sql.session import SparkSession __all__ = ["DataFrame", "DataFrameNaFunctions", "DataFrameStatFunctions"] @@ -3297,7 +3300,7 @@ def _preapare_cols_for_sort( cols: Sequence[Union[Sequence["ColumnOrNameOrOrdinal"], "ColumnOrNameOrOrdinal"]], kwargs: Dict[str, Any], ) -> Sequence[Column]: - from pyspark.errors import PySparkTypeError, PySparkValueError, PySparkIndexError + from pyspark.errors import PySparkIndexError, PySparkTypeError, PySparkValueError if not cols: raise PySparkValueError( @@ -3351,7 +3354,7 @@ def _get_col( raise PySparkTypeError( errorClass="NOT_EXPECTED_TYPE", messageParameters={ - "expected_type": "Column, int or str", + "expected_type": "bool, int or list", "arg_name": "ascending", "arg_type": type(ascending).__name__, }, diff --git a/python/pyspark/sql/datasource.py b/python/pyspark/sql/datasource.py index ad28136e8c361..2afdd2ebc260a 100644 --- a/python/pyspark/sql/datasource.py +++ b/python/pyspark/sql/datasource.py @@ -19,6 +19,7 @@ from collections.abc import MutableMapping from dataclasses import dataclass from typing import ( + TYPE_CHECKING, Any, Iterable, Iterator, @@ -28,16 +29,16 @@ Tuple, Type, Union, - TYPE_CHECKING, ) +from pyspark.errors import PySparkNotImplementedError from pyspark.sql import Row from pyspark.sql.streaming.datasource import ReadAllAvailable, ReadLimit from pyspark.sql.types import StructType -from pyspark.errors import PySparkNotImplementedError if TYPE_CHECKING: from pyarrow import RecordBatch + from pyspark.sql.session import SparkSession __all__ = [ @@ -527,9 +528,18 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]: can improve performance by reducing the amount of data that needs to be processed by Spark. - This method is called once during query planning. By default, it returns - all filters, indicating that no filters can be pushed down. Subclasses can - override this method to implement filter pushdown. + This method may be called more than once during query planning, on separate reader + instances (see the note below). By default, it returns all filters, indicating that no + filters can be pushed down. Subclasses can override this method to implement filter + pushdown. + + .. note:: + When limit pushdown is enabled (see :meth:`pushLimit`), planning may create + additional reader instances and call this method on each with the same filters, so + that a reader reaches the same state before a limit is pushed and read planning runs + only after the limit is known. Implementations must therefore be deterministic -- + returning a different set of supported filters across calls fails the query -- and + must not rely on side effects outside of `self`. It's recommended to implement this method only for data sources that natively support filtering, such as databases and GraphQL APIs. @@ -578,6 +588,83 @@ def pushFilters(self, filters: List["Filter"]) -> Iterable["Filter"]: """ return filters + def pushLimit(self, limit: int) -> bool: + """ + Called with the maximum number of rows that the query needs from this data source. + + Limit pushdown allows the data source to fetch less data, for example by adding a + `LIMIT` clause to a SQL query or a page size parameter to a REST request. + + This method is called once during query planning, before :meth:`partitions` and + :meth:`read`. By default, it returns False, indicating that the limit cannot be + pushed down. Subclasses can override this method to implement limit pushdown. + + :meth:`pushFilters` is called before this method only when the query has filters that + Spark can push down; for a query without them, :meth:`pushFilters` is not called at + all. This method may use state that :meth:`pushFilters` set when filters were pushed, + but must not assume :meth:`pushFilters` ran: initialize defaults in `__init__` so this + method works whether or not it did. + + A limit is only pushed down when every filter was pushed down, because Spark cannot + apply a limit before a filter it still has to evaluate itself. To benefit from limit + pushdown alongside filters, :meth:`pushFilters` should return an empty iterable. + + Pushing down a limit is only a hint: Spark always applies the limit again after the + scan, so it is safe to return True even if `read()` yields more than `limit` rows. + Returning True never causes the query to see fewer rows than it requires. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + limit : int + The maximum number of rows the query needs. Always positive: `LIMIT 0` is + optimized into an empty relation and never reaches the data source. + + Returns + ------- + bool + True if the data source will use the limit to reduce the amount of data it + reads, False otherwise. + + Side effects + ------------ + This method is allowed to modify `self`. The object must remain picklable. + Modifications to `self` are visible to the `partitions()` and `read()` methods + only when this method returns True. When it returns False, planning proceeds as if + `pushLimit` was never called -- `partitions()` and `read()` run on a reader that did + not observe these modifications -- so any state they rely on must be initialized in + `__init__` instead. + + Notes + ----- + This method is only called when the configuration + `spark.sql.python.limitPushdown.enabled` is set to true. + + Examples + -------- + Implement pushLimit to fetch fewer rows from the data source. Initialize the limit in + `__init__`, because :meth:`partitions` and :meth:`read` may run even when `pushLimit` + was not called -- for a query without a limit, or when this method returned False: + + >>> class MyReader(DataSourceReader): + ... def __init__(self): + ... self.limit = None + ... + ... def pushLimit(self, limit): + ... # Save the limit for handling in partitions() and read(). + ... self.limit = limit + ... return True + ... + ... def partitions(self): + ... # A limit can reduce the number of partitions, since every partition opens + ... # its own connection to the data source. + ... if self.limit is not None: + ... return [InputPartition(None)] + ... return [InputPartition(i) for i in range(16)] + """ + return False + def partitions(self) -> Sequence[InputPartition]: """ Returns an iterator of partitions for this data source. diff --git a/python/pyspark/sql/datasource_internal.py b/python/pyspark/sql/datasource_internal.py index e71dfb25ef9c5..932ea8c101941 100644 --- a/python/pyspark/sql/datasource_internal.py +++ b/python/pyspark/sql/datasource_internal.py @@ -16,11 +16,13 @@ # -import json import copy +import json from itertools import chain -from typing import Iterator, List, Optional, Sequence, Tuple, Type, Dict +from typing import Dict, Iterator, List, Optional, Sequence, Tuple, Type +from pyspark.errors import PySparkNotImplementedError +from pyspark.errors.exceptions.base import PySparkException from pyspark.sql.datasource import ( DataSource, DataSourceStreamReader, @@ -31,13 +33,11 @@ ReadAllAvailable, ReadLimit, ReadMaxBytes, + ReadMaxFiles, ReadMaxRows, ReadMinRows, - ReadMaxFiles, ) from pyspark.sql.types import StructType -from pyspark.errors import PySparkNotImplementedError -from pyspark.errors.exceptions.base import PySparkException def _streamReader(datasource: DataSource, schema: StructType) -> "DataSourceStreamReader": diff --git a/python/pyspark/sql/functions/__init__.py b/python/pyspark/sql/functions/__init__.py index 50f0fd440dd94..861f5902f977a 100644 --- a/python/pyspark/sql/functions/__init__.py +++ b/python/pyspark/sql/functions/__init__.py @@ -17,8 +17,8 @@ """PySpark Functions""" -from pyspark.sql.functions.builtin import * # noqa: F403 from pyspark.sql.functions import partitioning # noqa: F401 +from pyspark.sql.functions.builtin import * # noqa: F403 __all__ = [ # noqa: F405 # Normal functions @@ -110,6 +110,7 @@ "sqrt", "tan", "tanh", + "truncate", "try_add", "try_divide", "try_mod", @@ -138,6 +139,7 @@ "find_in_set", "format_number", "format_string", + "from_base32", "initcap", "instr", "is_valid_utf8", @@ -152,6 +154,7 @@ "ltrim", "make_valid_utf8", "mask", + "normalize", "octet_length", "overlay", "position", @@ -177,6 +180,7 @@ "substr", "substring", "substring_index", + "to_base32", "to_binary", "to_char", "to_number", @@ -292,6 +296,8 @@ "sha", "sha1", "sha2", + "xxh3_128", + "xxh3_64", "xxhash64", # Collection Functions "aggregate", @@ -338,6 +344,7 @@ "shuffle", "slice", "sort_array", + "trim_array", # Struct Functions "named_struct", "struct", @@ -363,10 +370,12 @@ "bitmap_and_agg", "bitmap_construct_agg", "bitmap_or_agg", + "bitmap_xor_agg", "bool_and", "bool_or", "collect_list", "collect_set", + "collect_union", "corr", "count", "count_distinct", @@ -471,6 +480,7 @@ "json_array_length", "json_object_keys", "json_tuple", + "json_typeof", "schema_of_json", "to_json", # VARIANT Functions @@ -483,6 +493,9 @@ "variant_array_append", "try_variant_array_append", "variant_delete", + "variant_from_arrays", + "variant_from_entries", + "variant_strip_nulls", "variant_get", "variant_insert", "try_variant_insert", @@ -513,9 +526,13 @@ "aes_decrypt", "aes_encrypt", "assert_true", + "bitmap_and", + "bitmap_andnot", "bitmap_bit_position", "bitmap_bucket_number", "bitmap_count", + "bitmap_or", + "bitmap_xor", "current_catalog", "current_database", "current_path", @@ -606,8 +623,10 @@ # Call Functions "call_udf", "pandas_udf", + "udaf", "udf", "udtf", "arrow_udtf", "unwrap_udt", + "wrap_udt", ] diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py index 319f1a2f2b022..e54927c36e651 100644 --- a/python/pyspark/sql/functions/builtin.py +++ b/python/pyspark/sql/functions/builtin.py @@ -19,25 +19,25 @@ A collections of builtin functions """ -import inspect import decimal -import sys import functools +import inspect +import sys import warnings from typing import ( + TYPE_CHECKING, Any, - cast, Callable, - Mapping, - Sequence, Iterable, - overload, + Mapping, Optional, + Sequence, Tuple, Type, - TYPE_CHECKING, Union, ValuesView, + cast, + overload, ) from pyspark.errors import PySparkTypeError, PySparkValueError @@ -47,43 +47,61 @@ ArrayType, ByteType, DataType, - StringType, - StructType, MapType, NumericType, + StringType, + StructType, _from_numpy_type, ) +from pyspark.sql.types import ( + UserDefinedType as _UserDefinedType, +) -# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409 -from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401 -from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401 -from pyspark.sql.udtf import OrderingColumn, PartitioningColumn, SelectedColumn # noqa: F401 -from pyspark.sql.udtf import SkipRestOfInputTableException # noqa: F401 -from pyspark.sql.udtf import UserDefinedTableFunction, _create_py_udtf, _create_pyarrow_udtf +if TYPE_CHECKING: + from pyspark.sql.types import UserDefinedType +# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409 # Keep pandas_udf and PandasUDFType import for backwards compatible import; moved in SPARK-28264 from pyspark.sql.pandas.functions import ( # noqa: F401 - arrow_udf, - pandas_udf, ArrowUDFType, PandasUDFType, + arrow_udf, + pandas_udf, +) +from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401 +from pyspark.sql.udtf import ( # noqa: F401 + AnalyzeArgument, + AnalyzeResult, + OrderingColumn, + PartitioningColumn, + SelectedColumn, + SkipRestOfInputTableException, + UserDefinedTableFunction, + _create_py_udtf, + _create_pyarrow_udtf, +) +from pyspark.sql.utils import ( + enum_to_value as _enum_to_value, +) +from pyspark.sql.utils import ( + get_active_spark_context as _get_active_spark_context, ) - from pyspark.sql.utils import ( to_str as _to_str, +) +from pyspark.sql.utils import ( try_remote_functions as _try_remote_functions, - get_active_spark_context as _get_active_spark_context, - enum_to_value as _enum_to_value, ) if TYPE_CHECKING: from pyspark import SparkContext - from pyspark.sql.dataframe import DataFrame from pyspark.sql._typing import ( ColumnOrName, DataTypeOrString, UserDefinedFunctionLike, ) + from pyspark.sql.aggregator import Aggregator + from pyspark.sql.dataframe import DataFrame # Note to developers: all of PySpark functions here take string as column names whenever possible. @@ -140,7 +158,7 @@ def _invoke_binary_math_function(name: str, col1: Any, col2: Any) -> Column: Invokes binary JVM math function identified by name and wraps the result with :class:`~pyspark.sql.Column`. """ - from pyspark.sql.classic.column import _to_java_column, _create_column_from_literal + from pyspark.sql.classic.column import _create_column_from_literal, _to_java_column # For legacy reasons, the arguments here can be implicitly converted into column cols = [ @@ -6284,6 +6302,74 @@ def collect_set(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("collect_set", col) +@_try_remote_functions +def collect_union(col: "ColumnOrName") -> Column: + """ + Aggregate function: given an array-typed column, collects the distinct union of the + elements of the arrays across rows and returns it as an array. + + The aggregation buffer holds only the distinct elements, so its size is bounded by the + element universe rather than by the number of input rows. Null elements are dropped by + default (``IGNORE NULLS``), matching :func:`collect_set`. With ``RESPECT NULLS`` a single + null element is kept, in which case this is equivalent to + ``array_distinct(flatten(collect_list(col)))``. The ``RESPECT NULLS`` clause is only + available through SQL, e.g. ``expr("collect_union(col) RESPECT NULLS")``. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The target array column on which the function is computed. + + Returns + ------- + :class:`~pyspark.sql.Column` + A new Column object representing the distinct union of the array elements. + + See Also + -------- + :meth:`pyspark.sql.functions.collect_set` + :meth:`pyspark.sql.functions.collect_list` + :meth:`pyspark.sql.functions.array_distinct` + :meth:`pyspark.sql.functions.flatten` + + Notes + ----- + This function is non-deterministic as the order of collected results depends + on the order of the rows, which may be non-deterministic after any shuffle operations. + + Examples + -------- + Example 1: Union the elements of array columns across rows + + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame( + ... [([1, 2],), ([2, 3],), ([1],)], ('value',)) + >>> df.select(sf.sort_array(sf.collect_union('value')).alias('u')).show() + +---------+ + | u| + +---------+ + |[1, 2, 3]| + +---------+ + + Example 2: Union per group + + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame( + ... [("a", [1, 2]), ("a", [2, 3]), ("b", [4])], ("k", "value")) + >>> df = df.groupBy("k").agg(sf.sort_array(sf.collect_union('value')).alias('u')) + >>> df.orderBy("k").show() + +---+---------+ + | k| u| + +---+---------+ + | a|[1, 2, 3]| + | b| [4]| + +---+---------+ + """ + return _invoke_function_over_columns("collect_union", col) + + @_try_remote_functions def degrees(col: "ColumnOrName") -> Column: """ @@ -7061,6 +7147,7 @@ def broadcast(df: "DataFrame") -> "DataFrame": +-----+---+ """ from py4j.java_gateway import JVMView + from pyspark.sql.dataframe import DataFrame sc = _get_active_spark_context() @@ -7343,7 +7430,7 @@ def count_distinct(col: "ColumnOrName", *cols: "ColumnOrName") -> Column: | 2| +------------------------------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function( @@ -8243,6 +8330,67 @@ def round(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Co return _invoke_function_over_columns("round", col, scale) +@_try_remote_functions +def truncate(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Column: + """ + Truncate the given value toward zero to `scale` decimal places when `scale` >= 0, + or to the left of the decimal point when `scale` < 0. `scale` defaults to 0. + + Unlike :func:`round`, the result is always rounded toward zero, and unlike :func:`floor` + negative values are not rounded toward negative infinity. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The target column or column name to truncate. + A column that evaluates to a numeric. + scale : :class:`~pyspark.sql.Column` or int, optional + An optional parameter to control the number of decimal places to keep. + A column that evaluates to an integer. Must be a constant. Defaults to 0. + + Returns + ------- + :class:`~pyspark.sql.Column` + A column for the truncated value, of the same type as the input, except that a decimal + input may return a decimal of different precision and scale. + + See Also + -------- + :meth:`pyspark.sql.functions.round` + :meth:`pyspark.sql.functions.trunc` + :meth:`pyspark.sql.functions.floor` + :meth:`pyspark.sql.functions.ceil` + + Examples + -------- + Example 1: Truncate toward zero to a given number of decimal places + + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.truncate(sf.lit(15.79), sf.lit(1)).alias("r")).collect() + [Row(r=15.7)] + + Example 2: Truncation rounds toward zero for negative values + + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.truncate(sf.lit(-2.99), sf.lit(0)).alias("r")).collect() + [Row(r=-2.0)] + + Example 3: The scale argument defaults to 0 when omitted + + >>> import pyspark.sql.functions as sf + >>> spark.range(1).select(sf.truncate(sf.lit(1234.5678)).alias("r")).collect() + [Row(r=1234.0)] + """ + if scale is None: + return _invoke_function_over_columns("truncate", col) + else: + scale = _enum_to_value(scale) + scale = lit(scale) if isinstance(scale, int) else scale + return _invoke_function_over_columns("truncate", col, scale) + + @_try_remote_functions def bround(col: "ColumnOrName", scale: Optional[Union[Column, int]] = None) -> Column: """ @@ -14697,6 +14845,71 @@ def md5(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("md5", col) +@_try_remote_functions +def xxh3_64(col: "ColumnOrName") -> Column: + """Returns a 64-bit hash value of the argument using the XXH3 algorithm. + + Unlike :func:`xxhash64`, which hashes one or more columns structurally, this hashes the raw + bytes of a single value with seed 0, so its result is byte compatible with the reference XXH3. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The target column to hash, which must have string or binary type. + + Returns + ------- + :class:`~pyspark.sql.Column` + Returns a column that evaluates to a long. + + See Also + -------- + :meth:`pyspark.sql.functions.xxh3_128` + :meth:`pyspark.sql.functions.xxhash64` + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([('Spark',)], ['a']) + >>> df.select(sf.xxh3_64('a').alias('h')).collect() + [Row(h=80997306238743657)] + """ + return _invoke_function_over_columns("xxh3_64", col) + + +@_try_remote_functions +def xxh3_128(col: "ColumnOrName") -> Column: + """Returns a 128-bit XXH3 hash of the argument as a 32-character hex string. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The target column to hash, which must have string or binary type. + + Returns + ------- + :class:`~pyspark.sql.Column` + Returns a column that evaluates to a string. + + See Also + -------- + :meth:`pyspark.sql.functions.xxh3_64` + :meth:`pyspark.sql.functions.md5` + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([('Spark',)], ['a']) + >>> df.select(sf.xxh3_128('a').alias('h')).collect() + [Row(h='7d57dd84c60c86ca1f4e82ab91a12b5e')] + """ + return _invoke_function_over_columns("xxh3_128", col) + + @_try_remote_functions def sha1(col: "ColumnOrName") -> Column: """Returns the hex string result of SHA-1. @@ -14872,6 +15085,7 @@ def xxhash64(*cols: "ColumnOrName") -> Column: See Also -------- :meth:`pyspark.sql.functions.hash` + :meth:`pyspark.sql.functions.xxh3_64` Examples -------- @@ -15175,6 +15389,7 @@ def base64(col: "ColumnOrName") -> Column: See Also -------- :meth:`pyspark.sql.functions.unbase64` + :meth:`pyspark.sql.functions.to_base32` Examples -------- @@ -15192,6 +15407,41 @@ def base64(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("base64", col) +@_try_remote_functions +def to_base32(col: "ColumnOrName") -> Column: + """ + Computes the BASE32 (RFC 4648) encoding of a binary column and returns it as a + string column. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + target column to work on. + A column that evaluates to a binary. + + Returns + ------- + :class:`~pyspark.sql.Column` + BASE32 encoding of the binary value. + Returns a column that evaluates to a string. + + See Also + -------- + :meth:`pyspark.sql.functions.from_base32` + :meth:`pyspark.sql.functions.base64` + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([(b"foobar",)], ["value"]) + >>> df.select(sf.to_base32("value").alias("r")).collect() + [Row(r='MZXW6YTBOI======')] + """ + return _invoke_function_over_columns("to_base32", col) + + @_try_remote_functions def unbase64(col: "ColumnOrName") -> Column: """ @@ -15217,6 +15467,7 @@ def unbase64(col: "ColumnOrName") -> Column: See Also -------- :meth:`pyspark.sql.functions.base64` + :meth:`pyspark.sql.functions.from_base32` Examples -------- @@ -15234,6 +15485,41 @@ def unbase64(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("unbase64", col) +@_try_remote_functions +def from_base32(col: "ColumnOrName") -> Column: + """ + Decodes a BASE32 (RFC 4648) encoded string column and returns it as a binary + column. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + target column to work on. + A column that evaluates to a string. + + Returns + ------- + :class:`~pyspark.sql.Column` + decoded binary value. + Returns a column that evaluates to a binary. + + See Also + -------- + :meth:`pyspark.sql.functions.to_base32` + :meth:`pyspark.sql.functions.unbase64` + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([("MZXW6YTBOI======",)], ["value"]) + >>> df.select(sf.from_base32("value").alias("r")).collect() + [Row(r=b'foobar')] + """ + return _invoke_function_over_columns("from_base32", col) + + @_try_remote_functions def ltrim(col: "ColumnOrName", trim: Optional["ColumnOrName"] = None) -> Column: """ @@ -15509,7 +15795,7 @@ def concat_ws(sep: str, *cols: "ColumnOrName") -> Column: |abcd|123| abcd-123-xyz| +----+---+-----------------------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function("concat_ws", _enum_to_value(sep), _to_seq(sc, cols, _to_java_column)) @@ -15757,6 +16043,46 @@ def try_validate_utf8(str: "ColumnOrName") -> Column: return _invoke_function_over_columns("try_validate_utf8", str) +@_try_remote_functions +def normalize(str: "ColumnOrName", form: Optional["ColumnOrName"] = None) -> Column: + """ + Returns the Unicode normalization of ``str`` using the given normalization ``form``, as + defined by Unicode Standard Annex #15. Normalization is backed by Spark's bundled ICU4J + library rather than the JVM's own Unicode data, so results are stable across JVM vendors + and versions. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + str : :class:`~pyspark.sql.Column` or column name + the input string to normalize. + form : :class:`~pyspark.sql.Column` or column name, optional + the normalization form, one of 'NFC', 'NFD', 'NFKC', 'NFKD' (case-insensitive). + If omitted, 'NFC' is used. + + Returns + ------- + :class:`~pyspark.sql.Column` + the normalized string. + + Examples + -------- + >>> import pyspark.sql.functions as sf + >>> df = spark.createDataFrame([("\ufb01",)], ["s"]) + >>> df.select(sf.normalize(df.s, sf.lit("NFKC"))).show() + +------------------+ + |normalize(s, NFKC)| + +------------------+ + | fi| + +------------------+ + """ + if form is None: + return _invoke_function_over_columns("normalize", str) + else: + return _invoke_function_over_columns("normalize", str, form) + + @_try_remote_functions def format_number(col: "ColumnOrName", d: int) -> Column: """ @@ -15839,7 +16165,7 @@ def format_string(format: str, *cols: "ColumnOrName") -> Column: | 5|hello| 5 hello| +---+-----+--------------------------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function( @@ -18517,7 +18843,7 @@ def printf(format: "ColumnOrName", *cols: "ColumnOrName") -> Column: | aa123cc| +---------------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function("printf", _to_java_column(format), _to_seq(sc, cols, _to_java_column)) @@ -19194,7 +19520,7 @@ def elt(*inputs: "ColumnOrName") -> Column: >>> df.select(elt(df.a, df.b, df.c).alias('r')).collect() [Row(r='scala')] """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function("elt", _to_seq(sc, inputs, _to_java_column)) @@ -20146,6 +20472,62 @@ def slice( return _invoke_function_over_columns("slice", x, start, length) +@_try_remote_functions +def trim_array(x: "ColumnOrName", n: Union["ColumnOrName", int]) -> Column: + """ + Array function: Returns the given array column with the last ``n`` elements removed. + Raises an error if ``n`` is negative or greater than the number of elements in the array. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + x : :class:`~pyspark.sql.Column` or str + Input array column or column name to be trimmed. + A column that evaluates to an array. + n : :class:`~pyspark.sql.Column`, str, or int + The number of elements to remove from the end of the array. Must be between 0 and + the number of elements in the array (inclusive). + A column that evaluates to an integer. + + Returns + ------- + :class:`~pyspark.sql.Column` + A new Column object of Array type, where each value is the corresponding input array + with its last ``n`` elements removed. + Returns a column that evaluates to an array. + + Examples + -------- + Example 1: Basic usage of the trim_array function. + + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([([1, 2, 3, 4, 5],), ([4, 5],)], ['x']) + >>> df.select(sf.trim_array(df.x, 2)).show() + +----------------+ + |trim_array(x, 2)| + +----------------+ + | [1, 2, 3]| + | []| + +----------------+ + + Example 2: trim_array function with a column input for n. + + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([([1, 2, 3, 4, 5], 1), ([4, 5], 0)], ['x', 'n']) + >>> df.select(sf.trim_array(df.x, df.n)).show() + +----------------+ + |trim_array(x, n)| + +----------------+ + | [1, 2, 3, 4]| + | [4, 5]| + +----------------+ + """ + n = _enum_to_value(n) + n = lit(n) if isinstance(n, int) else n + return _invoke_function_over_columns("trim_array", x, n) + + @_try_remote_functions def array_join( col: "ColumnOrName", delimiter: str, null_replacement: Optional[str] = None @@ -22290,7 +22672,7 @@ def json_tuple(col: "ColumnOrName", *fields: str) -> Column: >>> df.select(df.key, json_tuple(df.jstring, 'f1', 'f2')).collect() [Row(key='1', c0='value1', c1='value2'), Row(key='2', c0='value12', c1=None)] """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq if len(fields) == 0: raise PySparkValueError( @@ -22510,6 +22892,75 @@ def to_variant_object( return _invoke_function("to_variant_object", _to_java_column(col)) +@_try_remote_functions +def variant_from_arrays(keys: "ColumnOrName", values: "ColumnOrName") -> Column: + """ + Creates a variant object from the given arrays of keys and values. The keys must be non-null + strings and the two arrays must have the same length. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + keys : :class:`~pyspark.sql.Column` or column name + an array of string keys. + values : :class:`~pyspark.sql.Column` or column name + an array of values. + + Returns + ------- + :class:`~pyspark.sql.Column` + a new column of VariantType. + + See Also + -------- + :meth:`pyspark.sql.functions.variant_from_entries` + :meth:`pyspark.sql.functions.to_variant_object` + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.sql("SELECT array('a', 'b') AS keys, array(1, 2) AS values") + >>> df.select(sf.variant_from_arrays("keys", "values").cast("string").alias("r")).collect() + [Row(r='{"a":1,"b":2}')] + """ + return _invoke_function_over_columns("variant_from_arrays", keys, values) + + +@_try_remote_functions +def variant_from_entries(entries: "ColumnOrName") -> Column: + """ + Creates a variant object from an array of key/value struct entries. The keys must be non-null + strings. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + entries : :class:`~pyspark.sql.Column` or column name + an array of key/value structs, where the first field is a string key and the second field + is the value. + + Returns + ------- + :class:`~pyspark.sql.Column` + a new column of VariantType. + + See Also + -------- + :meth:`pyspark.sql.functions.variant_from_arrays` + :meth:`pyspark.sql.functions.to_variant_object` + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.sql("SELECT array(struct('a', 1), struct('b', 2)) AS entries") + >>> df.select(sf.variant_from_entries("entries").cast("string").alias("r")).collect() + [Row(r='{"a":1,"b":2}')] + """ + return _invoke_function_over_columns("variant_from_entries", entries) + + @_try_remote_functions def parse_json( col: "ColumnOrName", @@ -23047,6 +23498,52 @@ def try_variant_array_append( ) +@_try_remote_functions +def variant_strip_nulls(v: "ColumnOrName", include_arrays: bool = True) -> Column: + """ + Recursively removes object fields and array elements whose value is a variant null, unless + `include_arrays` is False, in which case null array elements are kept. Returns NULL if any + argument is NULL. + + .. versionadded:: 4.3.0 + + Parameters + ---------- + v : :class:`~pyspark.sql.Column` or str + a variant column or column name + include_arrays : bool, optional + whether null elements are also removed from arrays (default True). + + Returns + ------- + :class:`~pyspark.sql.Column` + a variant column with variant null fields/elements removed + + Examples + -------- + >>> from pyspark.sql.functions import lit, parse_json, to_json, variant_strip_nulls + >>> df = spark.createDataFrame([{ + ... 'json': '''{ "a" : 1, "b" : null, "c" : [1, null], "d" : { "e" : null, "f" : 4 } }''' + ... }]) + >>> v = parse_json(df.json) + >>> df.select(to_json(variant_strip_nulls(v)).alias("r")).collect() + [Row(r='{"a":1,"c":[1],"d":{"f":4}}')] + >>> df.select(to_json(variant_strip_nulls(v, False)).alias("r")).collect() + [Row(r='{"a":1,"c":[1,null],"d":{"f":4}}')] + >>> df.select(variant_strip_nulls(lit(None)).alias("r")).collect() + [Row(r=None)] + >>> df2 = spark.createDataFrame([{'json': '{"a": null}'}, {'json': 'null'}]) + >>> v2 = parse_json(df2.json) + >>> df2.select(to_json(variant_strip_nulls(v2)).alias("r")).collect() + [Row(r='{}'), Row(r='null')] + """ + from pyspark.sql.classic.column import _to_java_column + + return _invoke_function( + "variant_strip_nulls", _to_java_column(v), _enum_to_value(include_arrays) + ) + + @_try_remote_functions def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str) -> Column: """ @@ -23451,6 +23948,42 @@ def json_object_keys(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("json_object_keys", col) +@_try_remote_functions +def json_typeof(col: "ColumnOrName") -> Column: + """ + Returns the type of the outermost JSON value as a string: one of 'object', 'array', + 'string', 'number', 'boolean', or 'null'. Returns null if the input is not a valid JSON + string or is an empty string. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + col: :class:`~pyspark.sql.Column` or str + target column to compute on. + A column that evaluates to a string. + + Returns + ------- + :class:`~pyspark.sql.Column` + the type of the outermost JSON value. + Returns a column that evaluates to a string. + + See Also + -------- + :meth:`pyspark.sql.functions.json_object_keys` + :meth:`pyspark.sql.functions.get_json_object` + :meth:`pyspark.sql.functions.json_array_length` + + Examples + -------- + >>> df = spark.createDataFrame([('{"a": 1}',), ('[1, 2, 3]',), ('123',), ('',)], ['data']) + >>> df.select(json_typeof(df.data).alias('r')).collect() + [Row(r='object'), Row(r='array'), Row(r='number'), Row(r=None)] + """ + return _invoke_function_over_columns("json_typeof", col) + + # TODO: Fix and add an example for StructType with Spark Connect # e.g., StructType([StructField("a", IntegerType())]) @_try_remote_functions @@ -25504,6 +26037,7 @@ def _create_lambda(f: Callable) -> Callable: - (Column, Column, Column) -> Column: ... """ from py4j.java_gateway import JVMView + from pyspark.sql.classic.column import _to_seq parameters = _get_lambda_parameters(f) @@ -25543,7 +26077,8 @@ def _invoke_higher_order_function( :return: a Column """ from py4j.java_gateway import JVMView - from pyspark.sql.classic.column import _to_seq, _to_java_column + + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() jfuns = [_create_lambda(f) for f in funs] @@ -28591,7 +29126,7 @@ def call_udf(udfName: str, *cols: "ColumnOrName") -> Column: | cc| +-----------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function("call_udf", udfName, _to_seq(sc, cols, _to_java_column)) @@ -28662,7 +29197,7 @@ def call_function(funcName: str, *cols: "ColumnOrName") -> Column: | 102.0| +------------------------------------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = _get_active_spark_context() return _invoke_function("call_function", funcName, _to_seq(sc, cols, _to_java_column)) @@ -28684,6 +29219,10 @@ def unwrap_udt(col: "ColumnOrName") -> Column: :class:`~pyspark.sql.Column` The underlying representation. + See Also + -------- + :meth:`pyspark.sql.functions.wrap_udt` + Examples -------- Example 1: Unwrap ML-specific UDT - VectorUDT @@ -28729,6 +29268,107 @@ def unwrap_udt(col: "ColumnOrName") -> Column: return _invoke_function("unwrap_udt", _to_java_column(col)) +@_try_remote_functions +def wrap_udt(col: "ColumnOrName", udt: "Union[UserDefinedType, Column]") -> Column: + """ + Wrap a column as a user-defined type. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The column to wrap. The column data type must match the UDT's underlying SQL type. + udt : :class:`~pyspark.sql.types.UserDefinedType` or :class:`~pyspark.sql.Column` + The target user-defined type, or a constant string column containing its JSON + representation. + + Returns + ------- + :class:`~pyspark.sql.Column` + A column of the target user-defined type. + + See Also + -------- + :meth:`pyspark.sql.functions.unwrap_udt` + + Examples + -------- + Example 1: Wrapping a vector struct as VectorUDT + + >>> from pyspark.sql import functions as sf + >>> from pyspark.sql import Row + >>> from pyspark.sql.types import StructField, StructType + >>> from pyspark.ml.linalg import VectorUDT + >>> vector_schema = StructType([StructField("vec", VectorUDT.sqlType(), True)]) + >>> df = spark.createDataFrame( + ... [(Row(type=1, size=None, indices=None, values=[1.0, 2.0, 3.0]),)], + ... vector_schema) + >>> df.select("*", sf.wrap_udt("vec", VectorUDT())).show() + +--------------------+...+ + | vec|wrap_udt(vec...| + +--------------------+...+ + |{1, NULL, NULL, [...|...[1.0,2.0,3.0]| + +--------------------+...+ + >>> row = df.select(sf.wrap_udt("vec", VectorUDT())).first() + >>> type(row[0]) + <class 'pyspark.ml.linalg.DenseVector'> + + Example 2: Wrapping a matrix struct as MatrixUDT + + >>> from pyspark.sql import functions as sf + >>> from pyspark.sql import Row + >>> from pyspark.sql.types import StructField, StructType + >>> from pyspark.mllib.linalg import MatrixUDT + >>> matrix_schema = StructType([StructField("mat", MatrixUDT.sqlType(), True)]) + >>> df = spark.createDataFrame( + ... [( + ... Row( + ... type=1, + ... numRows=2, + ... numCols=2, + ... colPtrs=None, + ... rowIndices=None, + ... values=[1.0, 2.0, 3.0, 4.0], + ... isTransposed=False), + ... )], + ... matrix_schema) + >>> df.select("*", sf.wrap_udt("mat", MatrixUDT())).printSchema() + root + |-- mat: struct (nullable = true) + | |-- type: byte (nullable = false) + | |-- numRows: integer (nullable = false) + | |-- numCols: integer (nullable = false) + | |-- colPtrs: array (nullable = true) + | | |-- element: integer (containsNull = false) + | |-- rowIndices: array (nullable = true) + | | |-- element: integer (containsNull = false) + | |-- values: array (nullable = true) + | | |-- element: double (containsNull = false) + | |-- isTransposed: boolean (nullable = false) + |-- wrap_udt(mat...: matrix... (nullable = true) + >>> row = df.select(sf.wrap_udt("mat", MatrixUDT())).first() + >>> type(row[0]) + <class 'pyspark.mllib.linalg.DenseMatrix'> + """ + from pyspark.sql.classic.column import _to_java_column + + if isinstance(udt, _UserDefinedType): + udt_col = lit(udt.json()) + elif isinstance(udt, Column): + udt_col = udt + else: + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "UserDefinedType or Column", + "arg_name": "udt", + "arg_type": type(udt).__name__, + }, + ) + return _invoke_function("wrap_udt", _to_java_column(col), _to_java_column(udt_col)) + + # ---------------------- Datasketch functions ------------------------------ @@ -32528,6 +33168,194 @@ def bitmap_count(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("bitmap_count", col) +@_try_remote_functions +def bitmap_and(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise AND of two input bitmaps. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or column name + The left input bitmap. + right : :class:`~pyspark.sql.Column` or column name + The right input bitmap. + + See Also + -------- + :meth:`pyspark.sql.functions.bitmap_andnot` + :meth:`pyspark.sql.functions.bitmap_or` + :meth:`pyspark.sql.functions.bitmap_xor` + + Notes + ----- + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + ``BITMAP_INPUT_TOO_LARGE``. Both inputs must use the same bit-position mapping. If they were + constructed by grouping ``bitmap_bit_position`` values by ``bitmap_bucket_number``, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use ``bitmap_*_agg`` to combine bitmaps across + rows. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([("F0", "70")], ["left", "right"]) + >>> df.select(sf.bitmap_and( + ... sf.to_binary("left", sf.lit("hex")), + ... sf.to_binary("right", sf.lit("hex"))).alias("bitmap")).show() + +--------------------+ + | bitmap| + +--------------------+ + |[70 00 00 00 00 0...| + +--------------------+ + """ + return _invoke_function_over_columns("bitmap_and", left, right) + + +@_try_remote_functions +def bitmap_or(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise OR of two input bitmaps. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or column name + The left input bitmap. + right : :class:`~pyspark.sql.Column` or column name + The right input bitmap. + + See Also + -------- + :meth:`pyspark.sql.functions.bitmap_and` + :meth:`pyspark.sql.functions.bitmap_andnot` + :meth:`pyspark.sql.functions.bitmap_xor` + + Notes + ----- + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + ``BITMAP_INPUT_TOO_LARGE``. Both inputs must use the same bit-position mapping. If they were + constructed by grouping ``bitmap_bit_position`` values by ``bitmap_bucket_number``, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use ``bitmap_*_agg`` to combine bitmaps across + rows. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([("10", "20")], ["left", "right"]) + >>> df.select(sf.bitmap_or( + ... sf.to_binary("left", sf.lit("hex")), + ... sf.to_binary("right", sf.lit("hex"))).alias("bitmap")).show() + +--------------------+ + | bitmap| + +--------------------+ + |[30 00 00 00 00 0...| + +--------------------+ + """ + return _invoke_function_over_columns("bitmap_or", left, right) + + +@_try_remote_functions +def bitmap_andnot(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise AND NOT of two input bitmaps. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or column name + The left input bitmap. + right : :class:`~pyspark.sql.Column` or column name + The right input bitmap. + + See Also + -------- + :meth:`pyspark.sql.functions.bitmap_and` + :meth:`pyspark.sql.functions.bitmap_or` + :meth:`pyspark.sql.functions.bitmap_xor` + + Notes + ----- + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + ``BITMAP_INPUT_TOO_LARGE``. Both inputs must use the same bit-position mapping. If they were + constructed by grouping ``bitmap_bit_position`` values by ``bitmap_bucket_number``, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use ``bitmap_*_agg`` to combine bitmaps across + rows. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([("F0", "70")], ["left", "right"]) + >>> df.select(sf.bitmap_andnot( + ... sf.to_binary("left", sf.lit("hex")), + ... sf.to_binary("right", sf.lit("hex"))).alias("bitmap")).show() + +--------------------+ + | bitmap| + +--------------------+ + |[80 00 00 00 00 0...| + +--------------------+ + """ + return _invoke_function_over_columns("bitmap_andnot", left, right) + + +@_try_remote_functions +def bitmap_xor(left: "ColumnOrName", right: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise XOR of two input bitmaps. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + left : :class:`~pyspark.sql.Column` or column name + The left input bitmap. + right : :class:`~pyspark.sql.Column` or column name + The right input bitmap. + + See Also + -------- + :meth:`pyspark.sql.functions.bitmap_and` + :meth:`pyspark.sql.functions.bitmap_andnot` + :meth:`pyspark.sql.functions.bitmap_or` + + Notes + ----- + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + ``BITMAP_INPUT_TOO_LARGE``. Both inputs must use the same bit-position mapping. If they were + constructed by grouping ``bitmap_bit_position`` values by ``bitmap_bucket_number``, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use ``bitmap_*_agg`` to combine bitmaps across + rows. + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([("F0", "70")], ["left", "right"]) + >>> df.select(sf.bitmap_xor( + ... sf.to_binary("left", sf.lit("hex")), + ... sf.to_binary("right", sf.lit("hex"))).alias("bitmap")).show() + +--------------------+ + | bitmap| + +--------------------+ + |[80 00 00 00 00 0...| + +--------------------+ + """ + return _invoke_function_over_columns("bitmap_xor", left, right) + + @_try_remote_functions def bitmap_or_agg(col: "ColumnOrName") -> Column: """ @@ -32598,9 +33426,169 @@ def bitmap_and_agg(col: "ColumnOrName") -> Column: return _invoke_function_over_columns("bitmap_and_agg", col) +@_try_remote_functions +def bitmap_xor_agg(col: "ColumnOrName") -> Column: + """ + Returns a bitmap that is the bitwise XOR of all of the bitmaps from the input column. + The input column should be bitmaps created from bitmap_construct_agg(). + + .. versionadded:: 4.4.0 + + See Also + -------- + :meth:`pyspark.sql.functions.bitmap_bit_position` + :meth:`pyspark.sql.functions.bitmap_bucket_number` + :meth:`pyspark.sql.functions.bitmap_construct_agg` + :meth:`pyspark.sql.functions.bitmap_count` + :meth:`pyspark.sql.functions.bitmap_or_agg` + :meth:`pyspark.sql.functions.bitmap_and_agg` + + Parameters + ---------- + col : :class:`~pyspark.sql.Column` or column name + The input column should be bitmaps created from bitmap_construct_agg(). + + Examples + -------- + >>> from pyspark.sql import functions as sf + >>> df = spark.createDataFrame([("10",), ("30",), ("40",)], ["a"]) + >>> df.select(sf.bitmap_xor_agg(sf.to_binary(df.a, sf.lit("hex")))).show() + +---------------------------------+ + |bitmap_xor_agg(to_binary(a, hex))| + +---------------------------------+ + | [60 00 00 00 00 0...| + +---------------------------------+ + """ + return _invoke_function_over_columns("bitmap_xor_agg", col) + + # ---------------------------- User Defined Function ---------------------------------- +def udaf(agg: "Aggregator") -> "UserDefinedFunctionLike": + """Turn an :class:`~pyspark.sql.aggregator.Aggregator` instance into a callable usable in + ``groupBy().agg(...)`` (and as a window function), the Python counterpart of Scala's + ``functions.udaf``. + + The aggregator is executed with true incremental (partial) aggregation and transfers its + intermediate buffer as Arrow; PyArrow is therefore required. + + .. versionadded:: 4.4.0 + + Parameters + ---------- + agg : :class:`~pyspark.sql.aggregator.Aggregator` + The aggregator instance. + + Returns + ------- + function + A callable that, applied to input columns, produces an aggregate + :class:`~pyspark.sql.Column`. + + Raises + ------ + :class:`PySparkImportError` + If a supported version of PyArrow is not installed. + :class:`PySparkTypeError` + If ``agg`` is not an :class:`~pyspark.sql.aggregator.Aggregator`, or its ``bufferSchema`` is + not a :class:`StructType`. + + Examples + -------- + >>> from pyspark.sql.aggregator import Aggregator + >>> from pyspark.sql.functions import udaf + >>> from pyspark.sql.types import StructType, StructField, DoubleType, LongType + >>> class Mean(Aggregator): + ... @property + ... def bufferSchema(self): + ... return StructType( + ... [StructField("sum", DoubleType()), StructField("count", LongType())] + ... ) + ... + ... @property + ... def outputType(self): + ... return DoubleType() + ... + ... def zero(self): + ... return (0.0, 0) + ... + ... def reduce(self, buffer, value): + ... (v,) = value + ... return buffer if v is None else (buffer[0] + v, buffer[1] + 1) + ... + ... def merge(self, b1, b2): + ... return (b1[0] + b2[0], b1[1] + b2[1]) + ... + ... def finish(self, buffer): + ... return buffer[0] / buffer[1] if buffer[1] else None + >>> df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0)], ("k", "v")) + >>> df.groupBy("k").agg(udaf(Mean())(df.v).alias("m")).orderBy("k").show() + +---+---+ + | k| m| + +---+---+ + | 1|1.5| + | 2|3.0| + +---+---+ + """ + from pyspark.sql.aggregator import Aggregator + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + from pyspark.sql.utils import is_remote + from pyspark.util import PythonEvalType + + require_minimum_pyarrow_version() + + if is_remote(): + from pyspark.sql.connect.udf import UserDefinedFunction + else: + # The classic UserDefinedFunction is a distinct class from the Connect one above; + # both provide the same interface used below, so silence mypy's reassignment check. + from pyspark.sql.udf import UserDefinedFunction # type: ignore[assignment] + + if not isinstance(agg, Aggregator): + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "arg_name": "agg", + "expected_type": "Aggregator", + "arg_type": type(agg).__name__, + }, + ) + if not isinstance(agg.bufferSchema, StructType): + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "arg_name": "bufferSchema", + "expected_type": "StructType", + "arg_type": type(agg.bufferSchema).__name__, + }, + ) + # The buffer crosses the shuffle as an Arrow struct whose children are matched by name, and the + # worker keys the buffer tuple back by field name. Duplicate names would silently collapse + # fields on the map side and then fail with an opaque Arrow error post-shuffle, so reject them + # up front where the aggregator is created. + field_names = [field.name for field in agg.bufferSchema.fields] + if len(field_names) != len(set(field_names)): + duplicates = sorted({name for name in field_names if field_names.count(name) > 1}) + raise PySparkValueError( + errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT", + messageParameters={"field_names": ", ".join(duplicates)}, + ) + + # ``bufferSchema`` is a first-class ``UserDefinedFunction`` field (threaded to the JVM in + # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage aggregation), so it survives + # ``_wrapped()`` and ``spark.udf.register`` without being re-attached. + udf_obj = UserDefinedFunction( + agg, + returnType=agg.outputType, + name=agg.__class__.__name__, + evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, + deterministic=True, + bufferSchema=agg.bufferSchema, + ) + return udf_obj._wrapped() + + @overload def udf( f: Callable[..., Any], @@ -33276,8 +34264,9 @@ def vector_sum(col: "ColumnOrName") -> Column: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.functions.builtin + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.functions.builtin.__dict__.copy() @@ -33285,6 +34274,7 @@ def _test() -> None: if not have_pandas or not have_pyarrow: del pyspark.sql.functions.builtin.udf.__doc__ del pyspark.sql.functions.builtin.arrow_udtf.__doc__ + del pyspark.sql.functions.builtin.udaf.__doc__ spark = ( SparkSession.builder.master("local[4]").appName("sql.functions.builtin tests").getOrCreate() diff --git a/python/pyspark/sql/functions/partitioning.py b/python/pyspark/sql/functions/partitioning.py index 74f1ea8126be3..84e6cf4ea5533 100644 --- a/python/pyspark/sql/functions/partitioning.py +++ b/python/pyspark/sql/functions/partitioning.py @@ -27,11 +27,13 @@ from pyspark.errors import PySparkTypeError from pyspark.sql.column import Column -from pyspark.sql.functions.builtin import _invoke_function_over_columns, _invoke_function +from pyspark.sql.functions.builtin import _invoke_function, _invoke_function_over_columns from pyspark.sql.utils import ( - try_partitioning_remote_functions as _try_partitioning_remote_functions, get_active_spark_context as _get_active_spark_context, ) +from pyspark.sql.utils import ( + try_partitioning_remote_functions as _try_partitioning_remote_functions, +) if TYPE_CHECKING: from pyspark.sql._typing import ColumnOrName @@ -206,7 +208,7 @@ def bucket(numBuckets: Union[Column, int], col: "ColumnOrName") -> Column: method of the `DataFrameWriterV2`. """ - from pyspark.sql.classic.column import _to_java_column, _create_column_from_literal + from pyspark.sql.classic.column import _create_column_from_literal, _to_java_column if not isinstance(numBuckets, (int, Column)): raise PySparkTypeError( @@ -229,8 +231,9 @@ def bucket(numBuckets: Union[Column, int], col: "ColumnOrName") -> Column: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.functions.partitioning + from pyspark.sql import SparkSession globs = pyspark.sql.functions.partitioning.__dict__.copy() spark = ( diff --git a/python/pyspark/sql/group.py b/python/pyspark/sql/group.py index 221105d96783a..8086e043db4c8 100644 --- a/python/pyspark/sql/group.py +++ b/python/pyspark/sql/group.py @@ -18,16 +18,16 @@ # mypy: disable-error-code="empty-body" import sys - -from typing import Callable, List, Optional, TYPE_CHECKING, overload, Dict, Union, cast, Tuple +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union, cast, overload from pyspark.sql.column import Column -from pyspark.sql.session import SparkSession from pyspark.sql.dataframe import DataFrame from pyspark.sql.pandas.group_ops import PandasGroupedOpsMixin +from pyspark.sql.session import SparkSession if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql._typing import LiteralType __all__ = ["GroupedData"] @@ -534,8 +534,9 @@ def pivot(self, pivot_col: str, values: Optional[List["LiteralType"]] = None) -> def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.group + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.group.__dict__.copy() diff --git a/python/pyspark/sql/interchange.py b/python/pyspark/sql/interchange.py index 141d9f37148e1..f01daa405ccae 100644 --- a/python/pyspark/sql/interchange.py +++ b/python/pyspark/sql/interchange.py @@ -19,8 +19,8 @@ import pyarrow as pa import pyspark.sql -from pyspark.sql.types import StructType, StructField, BinaryType from pyspark.sql.pandas.types import to_arrow_schema +from pyspark.sql.types import BinaryType, StructField, StructType def _get_arrow_array_partition_stream(df: pyspark.sql.DataFrame) -> Iterator[pa.RecordBatch]: diff --git a/python/pyspark/sql/internal.py b/python/pyspark/sql/internal.py index b18ee66168041..f9605875c34b9 100644 --- a/python/pyspark/sql/internal.py +++ b/python/pyspark/sql/internal.py @@ -15,9 +15,10 @@ # limitations under the License. # -from pyspark.sql import Column, functions as F, is_remote +from typing import TYPE_CHECKING, Union -from typing import Union, TYPE_CHECKING +from pyspark.sql import Column, is_remote +from pyspark.sql import functions as F if TYPE_CHECKING: from pyspark.sql._typing import ColumnOrName @@ -36,8 +37,8 @@ def _invoke_internal_function_over_columns(name: str, *cols: "ColumnOrName") -> return _invoke_function_over_columns(name, *cols) else: - from pyspark.sql.classic.column import Column, _to_seq, _to_java_column from pyspark import SparkContext + from pyspark.sql.classic.column import Column, _to_java_column, _to_seq sc = SparkContext._active_spark_context return Column( diff --git a/python/pyspark/sql/merge.py b/python/pyspark/sql/merge.py index dbbc8692ba50b..c0931a6602170 100644 --- a/python/pyspark/sql/merge.py +++ b/python/pyspark/sql/merge.py @@ -15,7 +15,7 @@ # limitations under the License. # import sys -from typing import Dict, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Dict, Optional from pyspark.sql.column import Column from pyspark.sql.utils import to_scala_map @@ -224,10 +224,12 @@ def delete(self) -> "MergeIntoWriter": def _test() -> None: import doctest import os + import py4j + + import pyspark.sql.merge from pyspark.core.context import SparkContext from pyspark.sql import SparkSession - import pyspark.sql.merge os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/metrics.py b/python/pyspark/sql/metrics.py index a258a4f1db70e..49db19adf64f4 100644 --- a/python/pyspark/sql/metrics.py +++ b/python/pyspark/sql/metrics.py @@ -16,7 +16,7 @@ # import abc import dataclasses -from typing import Optional, List, Tuple, Dict, Any, Union, TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union from pyspark.errors import PySparkValueError @@ -297,10 +297,14 @@ class ExecutionInfo: data frame. This value is only set in the data frame if it was executed.""" def __init__( - self, metrics: Optional[list[PlanMetrics]], obs: Optional[Sequence[ObservedMetrics]] + self, + metrics: Optional[list[PlanMetrics]], + obs: Optional[Sequence[ObservedMetrics]], + operation_id: Optional[str] = None, ): self._metrics = CollectedMetrics(metrics) if metrics else None self._observations = obs if obs else [] + self._operation_id = operation_id @property def metrics(self) -> Optional[CollectedMetrics]: @@ -309,3 +313,11 @@ def metrics(self) -> Optional[CollectedMetrics]: @property def flows(self) -> List[Tuple[str, Dict[str, Any]]]: return [(f.name, f.pairs) for f in self._observations] + + @property + def operation_id(self) -> Optional[str]: + """The Spark Connect ExecutePlan operation ID, when available. + + .. versionadded:: 4.3.0 + """ + return self._operation_id diff --git a/python/pyspark/sql/observation.py b/python/pyspark/sql/observation.py index 721bfa5de8a8c..79b578378f73b 100644 --- a/python/pyspark/sql/observation.py +++ b/python/pyspark/sql/observation.py @@ -15,13 +15,13 @@ # limitations under the License. # import os -from typing import Any, Dict, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, Optional -from pyspark.errors import PySparkTypeError, PySparkValueError, PySparkAssertionError +from pyspark.errors import PySparkAssertionError, PySparkTypeError, PySparkValueError from pyspark.serializers import CPickleSerializer -from pyspark.sql import Row from pyspark.sql.column import Column from pyspark.sql.dataframe import DataFrame +from pyspark.sql.types import Row from pyspark.sql.utils import is_remote if TYPE_CHECKING: @@ -160,9 +160,10 @@ def get(self) -> Dict[str, Any]: def _test() -> None: import doctest import sys + + import pyspark.sql.observation from pyspark.core.context import SparkContext from pyspark.sql import SparkSession - import pyspark.sql.observation globs = pyspark.sql.observation.__dict__.copy() sc = SparkContext("local[4]", "PythonTest") diff --git a/python/pyspark/sql/pandas/_typing/__init__.pyi b/python/pyspark/sql/pandas/_typing/__init__.pyi index f989443a4dd90..da1b36205f2d9 100644 --- a/python/pyspark/sql/pandas/_typing/__init__.pyi +++ b/python/pyspark/sql/pandas/_typing/__init__.pyi @@ -16,6 +16,7 @@ # specific language governing permissions and limitations # under the License. +from types import FunctionType from typing import ( Any, Callable, @@ -25,16 +26,14 @@ from typing import ( TypeVar, Union, ) -from typing_extensions import Protocol, Literal -from types import FunctionType -from pyspark.sql._typing import LiteralType -from pyspark.sql.streaming.state import GroupState +import pyarrow +from numpy import ndarray as NDArray from pandas.core.frame import DataFrame as PandasDataFrame from pandas.core.series import Series as PandasSeries -from numpy import ndarray as NDArray - -import pyarrow +from pyspark.sql._typing import LiteralType +from pyspark.sql.streaming.state import GroupState +from typing_extensions import Literal, Protocol ArrayLike = NDArray DataFrameLike = PandasDataFrame @@ -68,6 +67,9 @@ ArrowScalarIterUDFType = Literal[251] ArrowGroupedAggUDFType = Literal[252] ArrowWindowAggUDFType = Literal[253] ArrowGroupedAggIterUDFType = Literal[254] +ArrowGroupedAggIncrementalPartialUDFType = Literal[255] +ArrowGroupedAggIncrementalFinalUDFType = Literal[256] +ArrowWindowAggIncrementalUDFType = Literal[257] # Arrow stream types # A single group of Arrow batches (e.g., one key group in groupBy). diff --git a/python/pyspark/sql/pandas/conversion.py b/python/pyspark/sql/pandas/conversion.py index d0b89354e1005..3bdaa67b91e64 100644 --- a/python/pyspark/sql/pandas/conversion.py +++ b/python/pyspark/sql/pandas/conversion.py @@ -16,6 +16,7 @@ # import sys from typing import ( + TYPE_CHECKING, Any, Callable, Iterable, @@ -27,26 +28,25 @@ cast, no_type_check, overload, - TYPE_CHECKING, ) from warnings import warn +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.errors.exceptions.captured import unwrap_spark_exception -from pyspark.util import _load_from_socket from pyspark.sql.pandas.serializers import ArrowCollectSerializer from pyspark.sql.pandas.types import _dedup_names from pyspark.sql.types import ( ArrayType, + DataType, MapType, - TimestampType, + StringType, StructType, - _has_type, - DataType, + TimestampType, _create_row, - StringType, + _has_type, ) from pyspark.traceback_utils import SCCallSiteSync -from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.util import _load_from_socket if TYPE_CHECKING: import numpy as np @@ -54,8 +54,8 @@ import pyarrow as pa from py4j.java_gateway import JavaObject - from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike from pyspark.sql import DataFrame + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike def create_arrow_array_from_pandas( @@ -86,10 +86,11 @@ def create_arrow_array_from_pandas( ------- pyarrow.Array """ - import pyarrow as pa import pandas as pd + import pyarrow as pa + from pyspark.loose_version import LooseVersion - from pyspark.sql.pandas.types import to_arrow_type, _create_converter_from_pandas + from pyspark.sql.pandas.types import _create_converter_from_pandas, to_arrow_type if isinstance(series.dtype, pd.CategoricalDtype): series = series.astype(series.dtype.categories.dtype) @@ -240,6 +241,7 @@ def _convert_arrow_table_to_pandas( The converted pandas DataFrame """ import pandas as pd + from pyspark.sql.pandas.types import _create_converter_to_pandas # Build pandas options @@ -392,7 +394,6 @@ def _to_pandas(self, **kwargs: Any) -> "PandasDataFrameLike": batches = self._collect_as_arrow( split_batches=arrowPySparkSelfDestructEnabled == "true", - prefers_large_var_types=prefers_large_var_types, ) if len(batches) > 0: @@ -491,12 +492,12 @@ def toArrow(self) -> "pa.Table": import pyarrow as pa self_destruct = arrowPySparkSelfDestructEnabled == "true" - batches = self._collect_as_arrow( - split_batches=self_destruct, - empty_list_if_zero_records=False, - prefers_large_var_types=prefers_large_var_types, - ) - table = pa.Table.from_batches(batches).cast(schema) + batches = self._collect_as_arrow(split_batches=self_destruct) + if batches: + table = pa.Table.from_batches(batches).cast(schema) + else: + # empty dataset + table = schema.empty_table() # Ensure only the table has a reference to the batches, so that # self_destruct (if enabled) is effective del batches @@ -505,21 +506,15 @@ def toArrow(self) -> "pa.Table": def _collect_as_arrow( self, split_batches: bool = False, - empty_list_if_zero_records: bool = True, - prefers_large_var_types: bool = False, ) -> List["pa.RecordBatch"]: """ - Returns all records as a list of Arrow RecordBatches. PyArrow must be installed - and available on driver and worker Python environments. - This is an experimental feature. + Returns all records as a list of Arrow RecordBatches, which is empty if the result has + 0 records. PyArrow must be installed and available on driver and worker Python + environments. This is an experimental feature. :param split_batches: split batches such that each column is in its own allocation, so that the selfDestruct optimization is effective; default False. - :param empty_list_if_zero_records: If True (the default), returns an empty list if the - result has 0 records. Otherwise, returns a list of length 1 containing an empty - Arrow RecordBatch which includes the schema. - .. note:: Experimental. """ from pyspark.sql.dataframe import DataFrame @@ -533,7 +528,7 @@ def _collect_as_arrow( jsocket_auth_server, ) = self._jdf.collectAsArrowToPython() - # Collect list of un-ordered batches where last element is a list of correct order indices + # Collect the batches already reordered by ArrowCollectSerializer. try: with _load_from_socket((port, auth_secret), ArrowCollectSerializer()) as batch_stream: if split_batches: @@ -545,41 +540,22 @@ def _collect_as_arrow( # converted. import pyarrow as pa - results = [] - for batch_or_indices in batch_stream: - if isinstance(batch_or_indices, pa.RecordBatch): - batch_or_indices = pa.RecordBatch.from_arrays( - [ - # This call actually reallocates the array - pa.concat_arrays([array]) - for array in batch_or_indices - ], - schema=batch_or_indices.schema, - ) - results.append(batch_or_indices) + batches = [ + pa.RecordBatch.from_arrays( + # This call actually reallocates the array + [pa.concat_arrays([array]) for array in batch], + schema=batch.schema, + ) + for batch in batch_stream + ] else: - results = list(batch_stream) + batches = list(batch_stream) finally: with unwrap_spark_exception(): # Join serving thread and raise any exceptions from collectAsArrowToPython jsocket_auth_server.getResult() - # Separate RecordBatches from batch order indices in results - batches = results[:-1] - batch_order = results[-1] - - if len(batches) or empty_list_if_zero_records: - # Re-order the batch list using the correct order - return [batches[i] for i in batch_order] - else: - from pyspark.sql.pandas.types import to_arrow_schema - import pyarrow as pa - - schema = to_arrow_schema( - self.schema, timezone="UTC", prefers_large_types=prefers_large_var_types - ) - empty_arrays = [pa.array([], type=field.type) for field in schema] - return [pa.RecordBatch.from_arrays(empty_arrays, schema=schema)] + return batches class SparkConversionMixin: @@ -744,12 +720,13 @@ def _convert_from_pandas( assert isinstance(self, SparkSession) if timezone is not None: + import pandas as pd + from pandas.core.dtypes.common import is_timedelta64_dtype + from pyspark.sql.pandas.types import ( _check_series_convert_timestamps_tz_local, _get_local_timezone, ) - import pandas as pd - from pandas.core.dtypes.common import is_timedelta64_dtype copied = False if isinstance(schema, StructType): @@ -949,22 +926,22 @@ def _create_from_pandas_with_arrow( assert isinstance(self, SparkSession) from pyspark.sql.pandas.serializers import ArrowStreamSerializer - from pyspark.sql.types import TimestampType from pyspark.sql.pandas.types import ( - from_arrow_type, _deduplicate_field_names, + from_arrow_type, ) from pyspark.sql.pandas.utils import ( require_minimum_pandas_version, require_minimum_pyarrow_version, ) + from pyspark.sql.types import TimestampType require_minimum_pandas_version() require_minimum_pyarrow_version() import pandas as pd - from pandas.api.types import is_datetime64_dtype import pyarrow as pa + from pandas.api.types import is_datetime64_dtype # Create the Spark schema from list of names passed in with Arrow types if isinstance(schema, (list, tuple)): @@ -1078,10 +1055,10 @@ def _create_from_arrow_table( from pyspark.sql.pandas.serializers import ArrowStreamSerializer from pyspark.sql.pandas.types import ( - from_arrow_type, + _check_arrow_table_timestamps_localize, from_arrow_schema, + from_arrow_type, to_arrow_schema, - _check_arrow_table_timestamps_localize, ) from pyspark.sql.pandas.utils import require_minimum_pyarrow_version @@ -1141,8 +1118,9 @@ def create_iter_server(): def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.pandas.conversion + from pyspark.sql import SparkSession globs = pyspark.sql.pandas.conversion.__dict__.copy() spark = ( diff --git a/python/pyspark/sql/pandas/functions.py b/python/pyspark/sql/pandas/functions.py index b14b10b44859e..ef48d8dc13d68 100644 --- a/python/pyspark/sql/pandas/functions.py +++ b/python/pyspark/sql/pandas/functions.py @@ -20,13 +20,13 @@ from inspect import getfullargspec, signature from typing import get_type_hints -from pyspark.util import PythonEvalType +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.pandas.typehints import infer_eval_type from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version from pyspark.sql.types import DataType from pyspark.sql.udf import _create_udf from pyspark.sql.utils import is_remote -from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.util import PythonEvalType class PandasUDFType: @@ -1031,10 +1031,11 @@ def _create_vectorized_udf(f, returnType, evalType, kind): def _test() -> None: - import sys import doctest - from pyspark.sql import SparkSession + import sys + import pyspark.sql.pandas.functions + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.column.__dict__.copy() diff --git a/python/pyspark/sql/pandas/group_ops.py b/python/pyspark/sql/pandas/group_ops.py index bbbe55ef5037b..2920500b27e1d 100644 --- a/python/pyspark/sql/pandas/group_ops.py +++ b/python/pyspark/sql/pandas/group_ops.py @@ -15,27 +15,27 @@ # limitations under the License. # import sys -from typing import List, Optional, Union, TYPE_CHECKING, cast, Any import warnings +from typing import TYPE_CHECKING, Any, List, Optional, Union, cast from pyspark.errors import PySparkTypeError -from pyspark.util import PythonEvalType from pyspark.sql.column import Column from pyspark.sql.dataframe import DataFrame from pyspark.sql.streaming.state import GroupStateTimeout from pyspark.sql.streaming.stateful_processor import StatefulProcessor from pyspark.sql.types import StructType +from pyspark.util import PythonEvalType if TYPE_CHECKING: + from pyspark.sql.group import GroupedData from pyspark.sql.pandas._typing import ( + ArrowCogroupedMapFunction, + ArrowGroupedMapFunction, GroupedMapPandasUserDefinedFunction, + PandasCogroupedMapFunction, PandasGroupedMapFunction, PandasGroupedMapFunctionWithState, - PandasCogroupedMapFunction, - ArrowGroupedMapFunction, - ArrowCogroupedMapFunction, ) - from pyspark.sql.group import GroupedData class PandasGroupedOpsMixin: @@ -1196,8 +1196,9 @@ def _extract_cols(gd: "GroupedData") -> List[Column]: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.pandas.group_ops + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.pandas.group_ops.__dict__.copy() diff --git a/python/pyspark/sql/pandas/map_ops.py b/python/pyspark/sql/pandas/map_ops.py index 65e73abc79562..923c8a9ee200a 100644 --- a/python/pyspark/sql/pandas/map_ops.py +++ b/python/pyspark/sql/pandas/map_ops.py @@ -15,17 +15,18 @@ # limitations under the License. # import sys -from typing import Union, TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, Union -from pyspark.resource.requests import ExecutorResourceRequests, TaskResourceRequests from pyspark.resource import ResourceProfile -from pyspark.util import PythonEvalType +from pyspark.resource.requests import ExecutorResourceRequests, TaskResourceRequests from pyspark.sql.types import StructType +from pyspark.util import PythonEvalType if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql.dataframe import DataFrame - from pyspark.sql.pandas._typing import PandasMapIterFunction, ArrowMapIterFunction + from pyspark.sql.pandas._typing import ArrowMapIterFunction, PandasMapIterFunction class PandasMapOpsMixin: @@ -105,8 +106,9 @@ def _build_java_profile( def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.pandas.map_ops + from pyspark.sql import SparkSession globs = pyspark.sql.pandas.map_ops.__dict__.copy() spark = ( diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py index 3b2bb187ee4dc..9eed884dc7b7f 100644 --- a/python/pyspark/sql/pandas/serializers.py +++ b/python/pyspark/sql/pandas/serializers.py @@ -24,9 +24,9 @@ from pyspark.errors import PySparkRuntimeError, PySparkValueError from pyspark.serializers import ( Serializer, + UTF8Deserializer, read_int, write_int, - UTF8Deserializer, ) if TYPE_CHECKING: @@ -42,49 +42,6 @@ class SpecialLengths: START_ARROW_STREAM = -6 -class ArrowCollectSerializer(Serializer): - """ - Deserialize a stream of batches followed by batch order information. Used in - PandasConversionMixin._collect_as_arrow() after invoking Dataset.collectAsArrowToPython() - in the JVM. - """ - - def __init__(self): - self.serializer = ArrowStreamSerializer() - - def dump_stream(self, iterator, stream): - return self.serializer.dump_stream(iterator, stream) - - def load_stream(self, stream): - """ - Load a stream of un-ordered Arrow RecordBatches, where the last iteration yields - a list of indices that can be used to put the RecordBatches in the correct order. - """ - # load the batches - for batch in self.serializer.load_stream(stream): - yield batch - - # load the batch order indices or propagate any error that occurred in the JVM - num = read_int(stream) - if num == -1: - error_msg = UTF8Deserializer().loads(stream) - raise PySparkRuntimeError( - errorClass="ERROR_OCCURRED_WHILE_CALLING", - messageParameters={ - "func_name": "ArrowCollectSerializer.load_stream", - "error_msg": error_msg, - }, - ) - batch_order = [] - for i in range(num): - index = read_int(stream) - batch_order.append(index) - yield batch_order - - def __repr__(self): - return "ArrowCollectSerializer(%s)" % self.serializer - - class ArrowStreamSerializer(Serializer): """ Serializes Arrow record batches as a plain stream. @@ -149,6 +106,44 @@ def __repr__(self) -> str: return "ArrowStreamSerializer(write_start_stream=%s)" % self._write_start_stream +class ArrowCollectSerializer(ArrowStreamSerializer): + """ + Extends :class:`ArrowStreamSerializer` to load Arrow RecordBatches that the JVM + sends out of order, followed by the indices giving their correct order, and yields + the batches already reordered. Used in PandasConversionMixin._collect_as_arrow() + after invoking Dataset.collectAsArrowToPython() in the JVM. + """ + + def load_stream(self, stream: IO[bytes]) -> Iterator["pa.RecordBatch"]: + """Load the out-of-order batches, then yield them in the correct order.""" + batches = list(super().load_stream(stream)) + + # Load the batch order indices, or propagate any error that occurred in the JVM. + num = read_int(stream) + if num == -1: + error_msg = UTF8Deserializer().loads(stream) + raise PySparkRuntimeError( + errorClass="ERROR_OCCURRED_WHILE_CALLING", + messageParameters={ + "func_name": "ArrowCollectSerializer.load_stream", + "error_msg": error_msg, + }, + ) + # Yield the batches in order, dropping our reference to each as it goes so + # that, when selfDestruct is enabled, the caller's reallocated copy is the + # only remaining reference and the original batch can be freed immediately + # rather than staying pinned here until the stream is fully consumed. The + # indices are a permutation, so each batch is yielded exactly once. + for _ in range(num): + i = read_int(stream) + batch = batches[i] + batches[i] = None + yield batch + + def __repr__(self) -> str: + return "ArrowCollectSerializer()" + + class ArrowStreamGroupSerializer(ArrowStreamSerializer): """ Extends :class:`ArrowStreamSerializer` with group-count protocol for loading diff --git a/python/pyspark/sql/pandas/typehints.py b/python/pyspark/sql/pandas/typehints.py index 4490748568e7a..18578b94f4dd8 100644 --- a/python/pyspark/sql/pandas/typehints.py +++ b/python/pyspark/sql/pandas/typehints.py @@ -14,29 +14,28 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from inspect import Signature -from typing import Any, Callable, Dict, Optional, Union, TYPE_CHECKING, get_type_hints -from inspect import getfullargspec, signature +from inspect import Signature, getfullargspec, signature +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union, get_type_hints -from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version from pyspark.errors import PySparkNotImplementedError, PySparkValueError +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version if TYPE_CHECKING: from pyspark.sql.pandas._typing import ( - PandasScalarUDFType, - PandasScalarIterUDFType, - PandasGroupedAggUDFType, - PandasGroupedAggIterUDFType, - ArrowScalarUDFType, - ArrowScalarIterUDFType, - ArrowGroupedAggUDFType, ArrowGroupedAggIterUDFType, + ArrowGroupedAggUDFType, + ArrowGroupedMapFunction, ArrowGroupedMapIterUDFType, ArrowGroupedMapUDFType, - ArrowGroupedMapFunction, + ArrowScalarIterUDFType, + ArrowScalarUDFType, + PandasGroupedAggIterUDFType, + PandasGroupedAggUDFType, PandasGroupedMapFunction, - PandasGroupedMapUDFType, PandasGroupedMapIterUDFType, + PandasGroupedMapUDFType, + PandasScalarIterUDFType, + PandasScalarUDFType, ) diff --git a/python/pyspark/sql/pandas/types.py b/python/pyspark/sql/pandas/types.py index 1a5c7a33c2406..c4facb3e3a8b4 100644 --- a/python/pyspark/sql/pandas/types.py +++ b/python/pyspark/sql/pandas/types.py @@ -21,56 +21,56 @@ """ import datetime -import itertools import functools +import itertools import json from decimal import Decimal -from typing import Any, Callable, Dict, Iterable, List, Optional, Union, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Union -from pyspark.errors import PySparkTypeError, UnsupportedOperationException, PySparkValueError +from pyspark.errors import PySparkTypeError, PySparkValueError, UnsupportedOperationException from pyspark.loose_version import LooseVersion from pyspark.sql.types import ( - cast, + ArrayType, + BinaryType, BooleanType, ByteType, - ShortType, + DataType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, + Geography, + GeographyType, + Geometry, + GeometryType, IntegerType, IntegralType, LongType, - FloatType, - DoubleType, - DecimalType, - StringType, - BinaryType, - DateType, - TimeType, - TimestampType, - TimestampNTZType, - DayTimeIntervalType, - YearMonthIntervalType, - ArrayType, MapType, - StructType, - StructField, NullType, - DataType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, + TimeType, UserDefinedType, VariantType, VariantVal, - GeometryType, - Geometry, - GeographyType, - Geography, + YearMonthIntervalType, _create_row, + cast, ) if TYPE_CHECKING: + import numpy as np import pandas as pd import pyarrow as pa - import numpy as np - from pyspark.sql.pandas._typing import SeriesLike as PandasSeriesLike from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike + from pyspark.sql.pandas._typing import SeriesLike as PandasSeriesLike # Should keep in line with org.apache.spark.sql.util.ArrowUtils.metadataKey @@ -538,9 +538,9 @@ def _check_arrow_array_timestamps_localize( ------- :class:`pyarrow.Array` or :class:`pyarrow.ChunkedArray` """ - import pyarrow.types as types import pyarrow as pa import pyarrow.compute as pc + import pyarrow.types as types if isinstance(a, pa.ChunkedArray) and (types.is_nested(a.type) or types.is_dictionary(a.type)): return pa.chunked_array( @@ -647,8 +647,8 @@ def _check_arrow_table_timestamps_localize( ------- :class:`pyarrow.Table` """ - import pyarrow.types as types import pyarrow as pa + import pyarrow.types as types # Return the table as-is if it contains no nested fields or timestamps if all([not types.is_nested(at) and not types.is_timestamp(at) for at in table.schema.types]): diff --git a/python/pyspark/sql/pandas/utils.py b/python/pyspark/sql/pandas/utils.py index 875c77dedada8..9cfad556f135e 100644 --- a/python/pyspark/sql/pandas/utils.py +++ b/python/pyspark/sql/pandas/utils.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.loose_version import LooseVersion from pyspark.errors import PySparkImportError, PySparkRuntimeError +from pyspark.loose_version import LooseVersion def require_minimum_pandas_version() -> None: diff --git a/python/pyspark/sql/plot/core.py b/python/pyspark/sql/plot/core.py index 99ebcec29fe00..dfa4998b95f1d 100644 --- a/python/pyspark/sql/plot/core.py +++ b/python/pyspark/sql/plot/core.py @@ -16,20 +16,22 @@ # import math - -from typing import Any, TYPE_CHECKING, List, Optional, Union, Sequence from types import ModuleType +from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union + from pyspark.errors import PySparkValueError -from pyspark.sql import Column, functions as F +from pyspark.sql import Column +from pyspark.sql import functions as F from pyspark.sql.internal import InternalFunction as SF from pyspark.sql.pandas.utils import require_minimum_pandas_version from pyspark.sql.utils import NumpyHelper, require_minimum_plotly_version if TYPE_CHECKING: - from pyspark.sql import DataFrame, Row import pandas as pd from plotly.graph_objs import Figure + from pyspark.sql import DataFrame, Row + class PySparkTopNPlotBase: def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": @@ -48,7 +50,8 @@ def get_top_n(self, sdf: "DataFrame") -> "pd.DataFrame": class PySparkSampledPlotBase: def get_sampled(self, sdf: "DataFrame") -> "pd.DataFrame": - from pyspark.sql import Observation, functions as F + from pyspark.sql import Observation + from pyspark.sql import functions as F max_rows = int( sdf._session.conf.get("spark.sql.pyspark.plotting.max_rows") # type: ignore[arg-type] diff --git a/python/pyspark/sql/plot/plotly.py b/python/pyspark/sql/plot/plotly.py index 584d3869fe361..408dfcdab7c53 100644 --- a/python/pyspark/sql/plot/plotly.py +++ b/python/pyspark/sql/plot/plotly.py @@ -20,17 +20,18 @@ from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.plot import ( - PySparkPlotAccessor, PySparkBoxPlotBase, - PySparkKdePlotBase, PySparkHistogramPlotBase, + PySparkKdePlotBase, + PySparkPlotAccessor, ) from pyspark.sql.types import NumericType if TYPE_CHECKING: - from pyspark.sql import DataFrame from plotly.graph_objs import Figure + from pyspark.sql import DataFrame + def plot_pyspark(data: "DataFrame", kind: str, **kwargs: Any) -> "Figure": import plotly @@ -164,8 +165,8 @@ def plot_box(data: "DataFrame", **kwargs: Any) -> "Figure": def plot_kde(data: "DataFrame", **kwargs: Any) -> "Figure": - from pyspark.testing.utils import have_numpy from pyspark.sql.pandas.utils import require_minimum_pandas_version + from pyspark.testing.utils import have_numpy require_minimum_pandas_version() diff --git a/python/pyspark/sql/profiler.py b/python/pyspark/sql/profiler.py index eae9b44f449b2..8917e7173e0ab 100644 --- a/python/pyspark/sql/profiler.py +++ b/python/pyspark/sql/profiler.py @@ -14,14 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from abc import ABC, abstractmethod -from io import StringIO import cProfile import os import pstats +import warnings +from abc import ABC, abstractmethod +from io import StringIO from threading import RLock from types import CodeType, TracebackType from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -30,10 +32,8 @@ Optional, Tuple, Union, - TYPE_CHECKING, overload, ) -import warnings import pyspark.memory_profiler_ext from pyspark.accumulators import ( diff --git a/python/pyspark/sql/protobuf/functions.py b/python/pyspark/sql/protobuf/functions.py index 42592bd563d30..23177bea6d062 100644 --- a/python/pyspark/sql/protobuf/functions.py +++ b/python/pyspark/sql/protobuf/functions.py @@ -19,7 +19,7 @@ A collections of builtin protobuf functions """ -from typing import Dict, Optional, TYPE_CHECKING, cast +from typing import TYPE_CHECKING, Dict, Optional, cast from pyspark.sql.column import Column from pyspark.sql.utils import get_active_spark_context, try_remote_protobuf_functions @@ -138,6 +138,7 @@ def from_protobuf( +------------------+ """ from py4j.java_gateway import JVMView + from pyspark.sql.classic.column import _to_java_column sc = get_active_spark_context() @@ -260,6 +261,7 @@ def to_protobuf( +----------------------------+ """ from py4j.java_gateway import JVMView + from pyspark.sql.classic.column import _to_java_column sc = get_active_spark_context() @@ -294,6 +296,7 @@ def _read_descriptor_set_file(filePath: str) -> bytes: def _test() -> None: import os import sys + from pyspark.testing.sqlutils import search_jar protobuf_jar = search_jar("connector/protobuf", "spark-protobuf-assembly-", "spark-protobuf") @@ -311,8 +314,9 @@ def _test() -> None: os.environ["PYSPARK_SUBMIT_ARGS"] = " ".join([jars_args, existing_args]) import doctest - from pyspark.sql import SparkSession + import pyspark.sql.protobuf.functions + from pyspark.sql import SparkSession globs = pyspark.sql.protobuf.functions.__dict__.copy() spark = ( diff --git a/python/pyspark/sql/readwriter.py b/python/pyspark/sql/readwriter.py index afe8000b5c456..73ac96701b823 100644 --- a/python/pyspark/sql/readwriter.py +++ b/python/pyspark/sql/readwriter.py @@ -15,20 +15,22 @@ # limitations under the License. # import sys -from typing import cast, overload, Dict, Iterable, List, Optional, Tuple, TYPE_CHECKING, Union +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union, cast, overload -from pyspark.util import is_remote_only -from pyspark.sql.types import StructType +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql import utils +from pyspark.sql.types import StructType from pyspark.sql.utils import to_str -from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.util import is_remote_only if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.core.rdd import RDD - from pyspark.sql._typing import OptionalPrimitiveType, ColumnOrName - from pyspark.sql.session import SparkSession + from pyspark.sql._typing import ColumnOrName, OptionalPrimitiveType from pyspark.sql.dataframe import DataFrame + from pyspark.sql.session import SparkSession from pyspark.sql.streaming import StreamingQuery __all__ = ["DataFrameReader", "DataFrameWriter", "DataFrameWriterV2"] @@ -37,7 +39,10 @@ TupleOrListOfString = Union[List[str], Tuple[str, ...]] -class OptionUtils: +class OptionUtils(ABC): + @abstractmethod + def option(self, key: str, value: "OptionalPrimitiveType") -> Any: ... + def _set_opts( self, schema: Optional[Union[StructType, str]] = None, @@ -50,7 +55,7 @@ def _set_opts( self.schema(schema) # type: ignore[attr-defined] for k, v in options.items(): if v is not None: - self.option(k, v) # type: ignore[attr-defined] + self.option(k, v) class DataFrameReader(OptionUtils): @@ -2527,7 +2532,7 @@ def partitionedBy(self, col: "ColumnOrName", *cols: "ColumnOrName") -> "DataFram .. versionadded:: 3.1.0 """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq col = _to_java_column(col) cols = _to_seq(self._spark._sc, [_to_java_column(c) for c in cols]) @@ -2614,10 +2619,12 @@ def overwritePartitions(self) -> None: def _test() -> None: import doctest import os + import py4j + + import pyspark.sql.readwriter from pyspark.core.context import SparkContext from pyspark.sql import SparkSession - import pyspark.sql.readwriter os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index c36260b9d13ea..be9d80e80a272 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -14,36 +14,37 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import filecmp import os import sys import warnings -import filecmp from collections.abc import Sized -from functools import reduce, cached_property +from functools import cached_property, reduce from threading import RLock from types import TracebackType from typing import ( + TYPE_CHECKING, Any, Callable, ClassVar, Dict, - Iterable, Generic, + Iterable, List, Optional, + Set, Tuple, Type, TypeVar, Union, - Set, cast, no_type_check, overload, - TYPE_CHECKING, ) from pyspark.conf import SparkConf -from pyspark.util import default_api_mode, is_remote_only +from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError +from pyspark.errors.exceptions.captured import install_exception_handler from pyspark.sql.conf import RuntimeConfig from pyspark.sql.dataframe import DataFrame from pyspark.sql.functions import lit @@ -57,40 +58,41 @@ DataType, StructType, VariantVal, - _make_type_verifier, - _infer_schema, + _create_converter, _has_nulltype, + _infer_schema, + _make_type_verifier, _merge_type, - _create_converter, ) -from pyspark.errors.exceptions.captured import install_exception_handler from pyspark.sql.utils import ( + remote_only, to_str, try_remote_session_classmethod, - remote_only, ) -from pyspark.errors import PySparkValueError, PySparkTypeError, PySparkRuntimeError +from pyspark.util import default_api_mode, is_remote_only if TYPE_CHECKING: - from py4j.java_gateway import JavaClass, JavaObject, JVMView import pyarrow as pa + from py4j.java_gateway import JavaClass, JavaObject, JVMView + from pyspark.core.context import SparkContext from pyspark.core.rdd import RDD - from pyspark.sql._typing import AtomicValue, RowLike, OptionalPrimitiveType + from pyspark.sql._typing import AtomicValue, OptionalPrimitiveType, RowLike from pyspark.sql.catalog import Catalog - from pyspark.sql.pandas._typing import ArrayLike, DataFrameLike as PandasDataFrameLike - from pyspark.sql.streaming import StreamingQueryManager - from pyspark.sql.streaming.query import StreamingCheckpointManager - from pyspark.sql.tvf import TableValuedFunction - from pyspark.sql.udf import UDFRegistration - from pyspark.sql.udtf import UDTFRegistration - from pyspark.sql.datasource import DataSourceRegistration - from pyspark.sql.dataframe import DataFrame as ParentDataFrame # Running MyPy type checks will always require pandas and # other dependencies so importing here is fine. from pyspark.sql.connect.client import SparkConnectClient from pyspark.sql.connect.shell.progress import ProgressHandler + from pyspark.sql.dataframe import DataFrame as ParentDataFrame + from pyspark.sql.datasource import DataSourceRegistration + from pyspark.sql.pandas._typing import ArrayLike + from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike + from pyspark.sql.streaming import StreamingQueryManager + from pyspark.sql.streaming.query import StreamingCheckpointManager + from pyspark.sql.tvf import TableValuedFunction + from pyspark.sql.udf import UDFRegistration + from pyspark.sql.udtf import UDTFRegistration __all__ = ["SparkSession"] @@ -523,7 +525,50 @@ def getOrCreate(self) -> "SparkSession": messageParameters={}, ) - if url.startswith("local") or ( + pool_cleanup: Optional[Callable[[], None]] = None + pool_local = str( + opts.get( + "spark.local.connect.pool", + os.environ.get("SPARK_LOCAL_CONNECT_POOL", ""), + ) + ).lower() in ("1", "true") + + reuse_local = str( + opts.get( + "spark.local.connect.reuse", + os.environ.get("SPARK_LOCAL_CONNECT_REUSE", ""), + ) + ).lower() in ("1", "true") + + if url.startswith("local") and pool_local: + from pyspark.sql.connect.local_server_pool import ( + acquire_pooled_local_connect_server, + release_pooled_local_connect_server, + ) + + # Opt-in: claim a server not previously assigned to an + # application run from a pool of booted local Connect servers. + # It is torn down when this run's session stops, so no state + # carries across runs. Takes precedence over reuse. See + # `pyspark.sql.connect.local_server_pool`. + url = acquire_pooled_local_connect_server(url, opts) + pool_cleanup = release_pooled_local_connect_server + for k in list(opts): + if k.startswith("spark.local.connect."): + opts.pop(k) + elif url.startswith("local") and reuse_local: + from pyspark.sql.connect.local_server import ( + reuse_or_start_local_connect_server, + ) + + # Opt-in: reconnect to a persistent local Connect server (starting + # one on the first run) instead of booting a fresh in-process server + # every process. See `pyspark.sql.connect.local_server`. + url = reuse_or_start_local_connect_server(url, opts) + for k in list(opts): + if k.startswith("spark.local.connect."): + opts.pop(k) + elif url.startswith("local") or ( is_api_mode_connect and not url.startswith("sc://") ): os.environ["SPARK_LOCAL_REMOTE"] = "1" @@ -532,9 +577,27 @@ def getOrCreate(self) -> "SparkSession": os.environ["SPARK_CONNECT_MODE_ENABLED"] = "1" opts["spark.remote"] = url + try: + remote_session = RemoteSparkSession.builder.config( + map=opts + ).getOrCreate() + except Exception: + if pool_cleanup is not None: + try: + pool_cleanup() + except Exception as cleanup_error: + warnings.warn( + "Failed to clean up a pooled local Connect server " + f"after session creation failed: {cleanup_error}", + RuntimeWarning, + stacklevel=2, + ) + raise + if pool_cleanup is not None: + remote_session._register_on_stop_callback(pool_cleanup) return cast( SparkSession, - RemoteSparkSession.builder.config(map=opts).getOrCreate(), + remote_session, ) elif "SPARK_LOCAL_REMOTE" in os.environ: url = "sc://localhost" @@ -1289,6 +1352,7 @@ def _create_shell_session() -> "SparkSession": that script, which would expose those to users. """ import py4j + from pyspark.core.context import SparkContext try: @@ -2555,8 +2619,9 @@ def _parse_ddl(self, ddl: str) -> DataType: def _test() -> None: - import os import doctest + import os + import pyspark.sql.session os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/sql_formatter.py b/python/pyspark/sql/sql_formatter.py index 1366ef277c47c..0914b3716f48c 100644 --- a/python/pyspark/sql/sql_formatter.py +++ b/python/pyspark/sql/sql_formatter.py @@ -17,13 +17,13 @@ import string import typing -from typing import Any, Optional, List, Tuple, Sequence, Mapping import uuid +from typing import Any, List, Mapping, Optional, Sequence, Tuple if typing.TYPE_CHECKING: - from pyspark.sql import SparkSession, DataFrame -from pyspark.sql.utils import get_lit_sql_str + from pyspark.sql import DataFrame, SparkSession from pyspark.errors import PySparkValueError +from pyspark.sql.utils import get_lit_sql_str class SQLStringFormatter(string.Formatter): diff --git a/python/pyspark/sql/streaming/__init__.py b/python/pyspark/sql/streaming/__init__.py index fe7f0159806ff..2f49fb0196f12 100644 --- a/python/pyspark/sql/streaming/__init__.py +++ b/python/pyspark/sql/streaming/__init__.py @@ -16,12 +16,12 @@ # # TODO: Add StreamingCheckpointManager to this when we want to make it public +from pyspark.errors import StreamingQueryException # noqa: F401 +from pyspark.sql.streaming.listener import StreamingQueryListener # noqa: F401 from pyspark.sql.streaming.query import StreamingQuery, StreamingQueryManager # noqa: F401 from pyspark.sql.streaming.readwriter import DataStreamReader, DataStreamWriter # noqa: F401 -from pyspark.sql.streaming.listener import StreamingQueryListener # noqa: F401 from pyspark.sql.streaming.stateful_processor import ( # noqa: F401 StatefulProcessor, StatefulProcessorHandle, ) from pyspark.sql.streaming.tws_tester import TwsTester # noqa: F401 -from pyspark.errors import StreamingQueryException # noqa: F401 diff --git a/python/pyspark/sql/streaming/benchmark/benchmark_tws_state_server.py b/python/pyspark/sql/streaming/benchmark/benchmark_tws_state_server.py index 7f7e0db24aa35..3cdc41a67bd49 100644 --- a/python/pyspark/sql/streaming/benchmark/benchmark_tws_state_server.py +++ b/python/pyspark/sql/streaming/benchmark/benchmark_tws_state_server.py @@ -15,29 +15,28 @@ # limitations under the License. # -import sys import os +import sys # Required to run the script easily on PySpark's root directory on the Spark repo. sys.path.append(os.getcwd()) -import uuid -import time import random +import time +import uuid from typing import List -from pyspark.sql.types import ( - StringType, - StructType, - StructField, -) +from pyspark.sql.streaming.benchmark.tws_utils import get_list_state, get_map_state, get_value_state +from pyspark.sql.streaming.benchmark.utils import print_percentiles from pyspark.sql.streaming.stateful_processor_api_client import ( ListTimerIterator, StatefulProcessorApiClient, ) - -from pyspark.sql.streaming.benchmark.utils import print_percentiles -from pyspark.sql.streaming.benchmark.tws_utils import get_list_state, get_map_state, get_value_state +from pyspark.sql.types import ( + StringType, + StructField, + StructType, +) def benchmark_value_state(api_client: StatefulProcessorApiClient, params: List[str]) -> None: diff --git a/python/pyspark/sql/streaming/benchmark/tws_utils.py b/python/pyspark/sql/streaming/benchmark/tws_utils.py index 6e7e5cf681107..18d43befa0f5d 100644 --- a/python/pyspark/sql/streaming/benchmark/tws_utils.py +++ b/python/pyspark/sql/streaming/benchmark/tws_utils.py @@ -17,9 +17,9 @@ from pyspark.sql.streaming.list_state_client import ListStateClient from pyspark.sql.streaming.map_state_client import MapStateClient -from pyspark.sql.streaming.value_state_client import ValueStateClient from pyspark.sql.streaming.stateful_processor import ListState, MapState, ValueState from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient +from pyspark.sql.streaming.value_state_client import ValueStateClient from pyspark.sql.types import StructType diff --git a/python/pyspark/sql/streaming/benchmark/utils.py b/python/pyspark/sql/streaming/benchmark/utils.py index b47ddf66caad1..647d4899d6105 100644 --- a/python/pyspark/sql/streaming/benchmark/utils.py +++ b/python/pyspark/sql/streaming/benchmark/utils.py @@ -15,10 +15,10 @@ # limitations under the License. # -import numpy as np - from typing import List +import numpy as np + def print_percentiles(values: List[float], percentiles: List[float]) -> None: percentile_values = np.percentile(values, q=percentiles) diff --git a/python/pyspark/sql/streaming/list_state_client.py b/python/pyspark/sql/streaming/list_state_client.py index 509e1cfb60b80..97e80cf46791c 100644 --- a/python/pyspark/sql/streaming/list_state_client.py +++ b/python/pyspark/sql/streaming/list_state_client.py @@ -14,12 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Dict, Iterator, List, Union, Tuple +import uuid +from typing import Any, Dict, Iterator, List, Tuple, Union +from pyspark.errors import PySparkRuntimeError from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient from pyspark.sql.types import StructType -from pyspark.errors import PySparkRuntimeError -import uuid __all__ = ["ListStateClient"] diff --git a/python/pyspark/sql/streaming/listener.py b/python/pyspark/sql/streaming/listener.py index 8799c4a11373a..bb247e750a48c 100644 --- a/python/pyspark/sql/streaming/listener.py +++ b/python/pyspark/sql/streaming/listener.py @@ -14,13 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import uuid import json -from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING +import uuid from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set -from pyspark.sql import Row from pyspark import cloudpickle +from pyspark.sql.types import Row __all__ = ["StreamingQueryListener"] @@ -1112,13 +1112,15 @@ def __repr__(self) -> str: def _test() -> None: - import sys import doctest import os + import sys + + from py4j.protocol import Py4JError + + import pyspark.sql.streaming.listener from pyspark.core.context import SparkContext from pyspark.sql import SparkSession - import pyspark.sql.streaming.listener - from py4j.protocol import Py4JError os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/streaming/map_state_client.py b/python/pyspark/sql/streaming/map_state_client.py index 0484ef792cf30..a972f9bf4fc55 100644 --- a/python/pyspark/sql/streaming/map_state_client.py +++ b/python/pyspark/sql/streaming/map_state_client.py @@ -14,12 +14,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, Dict, Iterator, Union, Tuple, Optional +import uuid +from typing import Any, Dict, Iterator, Optional, Tuple, Union +from pyspark.errors import PySparkRuntimeError from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient from pyspark.sql.types import StructType -from pyspark.errors import PySparkRuntimeError -import uuid __all__ = ["MapStateClient"] diff --git a/python/pyspark/sql/streaming/python_streaming_source_runner.py b/python/pyspark/sql/streaming/python_streaming_source_runner.py index d0b6cfbe32340..3f4bb9b4a8ab8 100644 --- a/python/pyspark/sql/streaming/python_streaming_source_runner.py +++ b/python/pyspark/sql/streaming/python_streaming_source_runner.py @@ -15,45 +15,45 @@ # limitations under the License. # +import dataclasses +import json import os import sys -import json from typing import IO, Iterator, Tuple -import dataclasses from pyspark.accumulators import _accumulatorRegistry from pyspark.errors import IllegalArgumentException, PySparkAssertionError from pyspark.errors.exceptions.base import PySparkException from pyspark.serializers import ( + SpecialLengths, read_int, write_int, write_with_length, - SpecialLengths, ) from pyspark.sql.datasource import ( DataSource, DataSourceStreamReader, ) -from pyspark.sql.streaming.datasource import ( - SupportsTriggerAvailableNow, -) from pyspark.sql.datasource_internal import ( + ReadLimitRegistry, _SimpleStreamReaderWrapper, _streamReader, - ReadLimitRegistry, ) from pyspark.sql.pandas.serializers import ArrowStreamSerializer +from pyspark.sql.streaming.datasource import ( + SupportsTriggerAvailableNow, +) from pyspark.sql.types import ( - _parse_datatype_json_string, StructType, + _parse_datatype_json_string, ) from pyspark.sql.worker.plan_data_source_read import records_to_arrow_batches from pyspark.util import handle_worker_exception from pyspark.worker_util import ( - get_sock_file_to_executor, check_python_version, - read_command, + get_sock_file_to_executor, pickleSer, + read_command, send_accumulator_updates, setup_memory_limits, setup_spark_files, diff --git a/python/pyspark/sql/streaming/query.py b/python/pyspark/sql/streaming/query.py index 03d04cbbc4a5b..065f117534d47 100644 --- a/python/pyspark/sql/streaming/query.py +++ b/python/pyspark/sql/streaming/query.py @@ -16,9 +16,9 @@ # import json -from typing import Any, Dict, List, Optional, TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Dict, List, Optional -from pyspark.errors import StreamingQueryException, PySparkValueError +from pyspark.errors import PySparkValueError, StreamingQueryException from pyspark.errors.exceptions.captured import ( StreamingQueryException as CapturedStreamingQueryException, ) @@ -672,6 +672,7 @@ def addListener(self, listener: StreamingQueryListener) -> None: >>> spark.streams.removeListener(test_listener) """ from py4j.java_gateway import java_import + from pyspark import SparkContext from pyspark.java_gateway import ensure_callback_server_started @@ -781,10 +782,12 @@ def _test() -> None: import doctest import os import sys + + from py4j.protocol import Py4JError + + import pyspark.sql.streaming.query from pyspark.core.context import SparkContext from pyspark.sql import SparkSession - import pyspark.sql.streaming.query - from py4j.protocol import Py4JError os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/streaming/readwriter.py b/python/pyspark/sql/streaming/readwriter.py index 6b7faa6222076..b8cf4ab05d196 100644 --- a/python/pyspark/sql/streaming/readwriter.py +++ b/python/pyspark/sql/streaming/readwriter.py @@ -18,24 +18,25 @@ import re import sys from collections.abc import Iterator -from typing import cast, overload, Any, Callable, List, Optional, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union, cast, overload -from pyspark.sql.readwriter import OptionUtils, to_str -from pyspark.sql.streaming.query import StreamingQuery -from pyspark.sql.types import Row, StructType -from pyspark.sql.utils import ForeachBatchFunction from pyspark.errors import ( - PySparkTypeError, - PySparkValueError, PySparkAttributeError, PySparkRuntimeError, + PySparkTypeError, + PySparkValueError, ) +from pyspark.sql.readwriter import OptionUtils, to_str +from pyspark.sql.streaming.query import StreamingQuery +from pyspark.sql.types import Row, StructType +from pyspark.sql.utils import ForeachBatchFunction if TYPE_CHECKING: from py4j.java_gateway import JavaObject - from pyspark.sql.session import SparkSession - from pyspark.sql._typing import SupportsProcess, OptionalPrimitiveType + + from pyspark.sql._typing import OptionalPrimitiveType, SupportsProcess from pyspark.sql.dataframe import DataFrame + from pyspark.sql.session import SparkSession __all__ = ["DataStreamReader", "DataStreamWriter"] @@ -1687,7 +1688,7 @@ def foreach(self, f: Union[Callable[[Row], None], "SupportsProcess"]) -> "DataSt """ from pyspark.core.rdd import _wrap_function - from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer + from pyspark.serializers import AutoBatchedSerializer, CPickleSerializer func = self._construct_foreach_function(f) serializer = AutoBatchedSerializer(CPickleSerializer()) @@ -1738,6 +1739,7 @@ def foreachBatch(self, func: Callable[["DataFrame", int], None]) -> "DataStreamW >>> # if in Spark Connect, my_value = -1, else my_value = 100 """ from py4j.java_gateway import java_import + from pyspark.java_gateway import ensure_callback_server_started gw = self._spark._sc._gateway @@ -1934,8 +1936,9 @@ def toTable( def _test() -> None: import doctest import os - from pyspark.sql import SparkSession + import pyspark.sql.streaming.readwriter + from pyspark.sql import SparkSession os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/streaming/state.py b/python/pyspark/sql/streaming/state.py index 7db667c4314a9..10895118819b3 100644 --- a/python/pyspark/sql/streaming/state.py +++ b/python/pyspark/sql/streaming/state.py @@ -16,10 +16,10 @@ # import datetime import json -from typing import Tuple, Optional +from typing import Optional, Tuple +from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.sql.types import Row, StructType, TimestampType -from pyspark.errors import PySparkTypeError, PySparkValueError, PySparkRuntimeError __all__ = ["GroupState", "GroupStateTimeout"] diff --git a/python/pyspark/sql/streaming/stateful_processor.py b/python/pyspark/sql/streaming/stateful_processor.py index dbece47a93ceb..a2eeca089d7ff 100644 --- a/python/pyspark/sql/streaming/stateful_processor.py +++ b/python/pyspark/sql/streaming/stateful_processor.py @@ -16,20 +16,20 @@ # from abc import ABC, abstractmethod -from typing import Any, List, TYPE_CHECKING, Iterator, Optional, Union, Tuple +from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union -from pyspark.sql.streaming.stateful_processor_api_client import ( - StatefulProcessorApiClient, - ListTimerIterator, -) from pyspark.sql.streaming.list_state_client import ListStateClient, ListStateIterator from pyspark.sql.streaming.map_state_client import ( MapStateClient, MapStateIterator, MapStateKeyValuePairIterator, ) +from pyspark.sql.streaming.stateful_processor_api_client import ( + ListTimerIterator, + StatefulProcessorApiClient, +) from pyspark.sql.streaming.value_state_client import ValueStateClient -from pyspark.sql.types import StructType, Row +from pyspark.sql.types import Row, StructType if TYPE_CHECKING: from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike diff --git a/python/pyspark/sql/streaming/stateful_processor_api_client.py b/python/pyspark/sql/streaming/stateful_processor_api_client.py index 4be92ff38dadd..0be61e0f20706 100644 --- a/python/pyspark/sql/streaming/stateful_processor_api_client.py +++ b/python/pyspark/sql/streaming/stateful_processor_api_client.py @@ -14,23 +14,22 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from datetime import datetime -from enum import Enum import json import os import socket -from typing import IO, Any, Dict, List, Union, Optional, Tuple, Iterator, cast +import uuid +from datetime import datetime +from enum import Enum +from typing import IO, Any, Dict, Iterator, List, Optional, Tuple, Union, cast -from pyspark.serializers import write_int, read_int, UTF8Deserializer +from pyspark.errors import PySparkRuntimeError +from pyspark.serializers import PickleSerializer, UTF8Deserializer, read_int, write_int from pyspark.sql.pandas.serializers import ArrowStreamSerializer +from pyspark.sql.pandas.types import convert_pandas_using_numpy_type from pyspark.sql.types import ( - StructType, Row, + StructType, ) -from pyspark.sql.pandas.types import convert_pandas_using_numpy_type -from pyspark.serializers import PickleSerializer -from pyspark.errors import PySparkRuntimeError -import uuid __all__ = ["StatefulProcessorApiClient", "StatefulProcessorHandleState"] @@ -547,8 +546,8 @@ def _deserialize_from_bytes(self, value: bytes) -> Any: return self.pickleSer.loads(value) def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None: - import pyarrow as pa import pandas as pd + import pyarrow as pa column_names = [field.name for field in schema.fields] pandas_df = convert_pandas_using_numpy_type( diff --git a/python/pyspark/sql/streaming/stateful_processor_util.py b/python/pyspark/sql/streaming/stateful_processor_util.py index c926f45c2c15d..f33737f064530 100644 --- a/python/pyspark/sql/streaming/stateful_processor_util.py +++ b/python/pyspark/sql/streaming/stateful_processor_util.py @@ -15,20 +15,21 @@ # limitations under the License. # -from enum import Enum import itertools -from typing import Any, Iterator, Optional, TYPE_CHECKING, Union -from pyspark.sql.streaming.stateful_processor_api_client import ( - StatefulProcessorApiClient, - StatefulProcessorHandleState, -) +from enum import Enum +from typing import TYPE_CHECKING, Any, Iterator, Optional, Union + from pyspark.sql.streaming.stateful_processor import ( ExpiredTimerInfo, StatefulProcessor, StatefulProcessorHandle, TimerValues, ) -from pyspark.sql.streaming.stateful_processor_api_client import ExpiredTimerIterator +from pyspark.sql.streaming.stateful_processor_api_client import ( + ExpiredTimerIterator, + StatefulProcessorApiClient, + StatefulProcessorHandleState, +) from pyspark.sql.types import Row if TYPE_CHECKING: diff --git a/python/pyspark/sql/streaming/transform_with_state_driver_worker.py b/python/pyspark/sql/streaming/transform_with_state_driver_worker.py index a05d616eda2cd..19583035631d4 100644 --- a/python/pyspark/sql/streaming/transform_with_state_driver_worker.py +++ b/python/pyspark/sql/streaming/transform_with_state_driver_worker.py @@ -16,22 +16,20 @@ # import json -from typing import Any, Iterator, TYPE_CHECKING +from typing import IO, TYPE_CHECKING, Any, Iterator -from pyspark.worker_util import get_sock_file_to_executor +from pyspark import worker from pyspark.serializers import ( - write_int, - read_int, - UTF8Deserializer, CPickleSerializer, + UTF8Deserializer, + read_int, + write_int, ) -from pyspark import worker -from pyspark.util import handle_worker_exception -from typing import IO -from pyspark.worker_util import check_python_version from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode from pyspark.sql.types import StructType +from pyspark.util import handle_worker_exception +from pyspark.worker_util import check_python_version, get_sock_file_to_executor if TYPE_CHECKING: from pyspark.sql.pandas._typing import ( diff --git a/python/pyspark/sql/streaming/tws_tester.py b/python/pyspark/sql/streaming/tws_tester.py index 5851b9955c9f5..9af4c24773255 100644 --- a/python/pyspark/sql/streaming/tws_tester.py +++ b/python/pyspark/sql/streaming/tws_tester.py @@ -16,8 +16,10 @@ # from __future__ import annotations -from typing import Any, Callable, cast, Iterator, Optional, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, Union, cast +from pyspark.errors import PySparkAssertionError, PySparkValueError +from pyspark.errors.exceptions.base import IllegalArgumentException from pyspark.sql.streaming.stateful_processor import ( ExpiredTimerInfo, ListState, @@ -28,8 +30,6 @@ ValueState, ) from pyspark.sql.types import Row, StructType -from pyspark.errors import PySparkValueError, PySparkAssertionError -from pyspark.errors.exceptions.base import IllegalArgumentException if TYPE_CHECKING: from pyspark.sql.pandas._typing import DataFrameLike as PandasDataFrameLike diff --git a/python/pyspark/sql/streaming/value_state_client.py b/python/pyspark/sql/streaming/value_state_client.py index 6c8e0faffdfde..f5df23303a39e 100644 --- a/python/pyspark/sql/streaming/value_state_client.py +++ b/python/pyspark/sql/streaming/value_state_client.py @@ -14,11 +14,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Union, Tuple, Optional +from typing import Optional, Tuple, Union +from pyspark.errors import PySparkRuntimeError from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient from pyspark.sql.types import StructType -from pyspark.errors import PySparkRuntimeError __all__ = ["ValueStateClient"] diff --git a/python/pyspark/sql/tests/arrow/test_arrow.py b/python/pyspark/sql/tests/arrow/test_arrow.py index 63a21aed1bddf..df4239651d344 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow.py +++ b/python/pyspark/sql/tests/arrow/test_arrow.py @@ -15,46 +15,47 @@ # limitations under the License. # +import calendar import datetime import os import threading -import calendar import time import unittest from collections import namedtuple from pyspark import SparkConf +from pyspark.errors import ArithmeticException, PySparkTypeError, UnsupportedOperationException from pyspark.sql import Row, SparkSession -from pyspark.sql.functions import rand, udf, assert_true, lit +from pyspark.sql.functions import assert_true, lit, rand, udf +from pyspark.sql.pandas.types import ( + from_arrow_schema, + from_arrow_type, + to_arrow_schema, + to_arrow_type, +) from pyspark.sql.types import ( + ArrayType, + BinaryType, BooleanType, ByteType, - StructType, - StringType, - ShortType, + DateType, + DayTimeIntervalType, + DecimalType, + DoubleType, + FloatType, IntegerType, LongType, - FloatType, - DoubleType, - DecimalType, - DateType, - TimeType, - TimestampType, - TimestampNTZType, - BinaryType, - StructField, - ArrayType, MapType, NullType, - DayTimeIntervalType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, + TimeType, VariantType, ) -from pyspark.sql.pandas.types import ( - from_arrow_type, - to_arrow_type, - from_arrow_schema, - to_arrow_schema, -) from pyspark.testing.objects import ExamplePoint, ExamplePointUDT from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( @@ -63,7 +64,6 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.errors import ArithmeticException, PySparkTypeError, UnsupportedOperationException from pyspark.util import is_remote_only if have_pandas: @@ -1427,6 +1427,12 @@ def test_toArrow_duplicate_field_names(self): ): df.toArrow() + # An empty result must reject duplicated struct field names just like a non-empty one. + with self.assertRaisesRegex( + UnsupportedOperationException, "DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT" + ): + df.limit(0).toArrow() + def test_createDataFrame_pandas_duplicate_field_names(self): for arrow_enabled in [True, False]: with self.subTest(arrow_enabled=arrow_enabled): @@ -1856,7 +1862,7 @@ def test_toArrow_with_compression_codec(self): def test_toPandas_with_compression_codec_large_dataset(self): # Test compression with a larger dataset to verify memory savings # Create a dataset with repetitive data that compresses well - from pyspark.sql.functions import lit, col + from pyspark.sql.functions import col, lit df = self.spark.range(10000).select( col("id"), @@ -1873,7 +1879,7 @@ def test_toPandas_with_compression_codec_large_dataset(self): def test_toArrow_with_compression_codec_large_dataset(self): # Test compression with a larger dataset for toArrow - from pyspark.sql.functions import lit, col + from pyspark.sql.functions import col, lit df = self.spark.range(10000).select( col("id"), diff --git a/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py index 15eb0b2320cbc..e52bba8f35451 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_cogrouped_map_misc.py @@ -14,10 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging import os import time import unittest -import logging from pyspark.sql import Row from pyspark.sql import functions as sf diff --git a/python/pyspark/sql/tests/arrow/test_arrow_grouped_map.py b/python/pyspark/sql/tests/arrow/test_arrow_grouped_map.py index e0d40cfebe59e..8af706d8a9c50 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_grouped_map.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_grouped_map.py @@ -15,18 +15,24 @@ # limitations under the License. # import inspect +import logging import os import time -import logging -from typing import Iterator, Tuple import unittest +from typing import Iterator, Tuple from pyspark.errors import PythonException -from pyspark.sql import Row, functions as sf +from pyspark.sql import Row +from pyspark.sql import functions as sf from pyspark.sql.functions import array, col, explode, lit, mean, stddev from pyspark.sql.window import Window from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.utils import assertDataFrameEqual, have_pyarrow, pyarrow_requirement_message +from pyspark.testing.utils import ( + assertDataFrameEqual, + eventually, + have_pyarrow, + pyarrow_requirement_message, +) from pyspark.util import is_remote_only if have_pyarrow: @@ -471,20 +477,27 @@ def func_with_logging(group): df, ) - logs = self.spark.tvf.python_worker_logs() + # Worker logs are captured asynchronously from the worker's stdout and only + # become visible once the trailing block is flushed, so they may not all be + # present immediately after the query completes. Poll until they show up. + @eventually(timeout=5, catch_assertions=True) + def check_logs(): + logs = self.spark.tvf.python_worker_logs() + + assertDataFrameEqual( + logs.select("level", "msg", "context", "logger"), + [ + Row( + level="WARNING", + msg=f"arrow grouped map: {dict(id=lst, value=[v * 10 for v in lst])}", + context={"func_name": func_with_logging.__name__}, + logger="test_arrow_grouped_map", + ) + for lst in [[0, 2, 4, 6, 8], [1, 3, 5, 7]] + ], + ) - assertDataFrameEqual( - logs.select("level", "msg", "context", "logger"), - [ - Row( - level="WARNING", - msg=f"arrow grouped map: {dict(id=lst, value=[v * 10 for v in lst])}", - context={"func_name": func_with_logging.__name__}, - logger="test_arrow_grouped_map", - ) - for lst in [[0, 2, 4, 6, 8], [1, 3, 5, 7]] - ], - ) + check_logs() @unittest.skipIf(is_remote_only(), "Requires JVM access") def test_apply_in_arrow_iter_with_logging(self): @@ -511,20 +524,27 @@ def func_with_logging(group: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatc df, ) - logs = self.spark.tvf.python_worker_logs() + # Worker logs are captured asynchronously from the worker's stdout and only + # become visible once the trailing block is flushed, so they may not all be + # present immediately after the query completes. Poll until they show up. + @eventually(timeout=5, catch_assertions=True) + def check_logs(): + logs = self.spark.tvf.python_worker_logs() + + assertDataFrameEqual( + logs.select("level", "msg", "context", "logger"), + [ + Row( + level="WARNING", + msg=f"arrow grouped map: {dict(id=lst, value=[v * 10 for v in lst])}", + context={"func_name": func_with_logging.__name__}, + logger="test_arrow_grouped_map", + ) + for lst in [[0, 2, 4], [6, 8], [1, 3, 5], [7]] + ], + ) - assertDataFrameEqual( - logs.select("level", "msg", "context", "logger"), - [ - Row( - level="WARNING", - msg=f"arrow grouped map: {dict(id=lst, value=[v * 10 for v in lst])}", - context={"func_name": func_with_logging.__name__}, - logger="test_arrow_grouped_map", - ) - for lst in [[0, 2, 4], [6, 8], [1, 3, 5], [7]] - ], - ) + check_logs() class ApplyInArrowTests(ApplyInArrowTestsMixin, ReusedSQLTestCase): diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py index 5119e0e827f6d..f2c5644059b64 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_map.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py @@ -14,14 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging import os import time import unittest -import logging +from pyspark.sql import Row from pyspark.sql.utils import PythonException from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.sql import Row from pyspark.testing.utils import ( assertDataFrameEqual, have_pandas, @@ -55,6 +55,43 @@ def func(iterator): expected = df.collect() self.assertEqual(actual, expected) + def test_map_in_arrow_legacy_accept_any_iterable(self): + # With the legacy flag enabled, returning a non-Iterator iterable (e.g. list) is accepted. + def list_not_iter(iterator): + return [batch for batch in iterator] + + with self.sql_conf( + {"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled": True} + ): + df = self.spark.range(10) + actual = df.mapInArrow(list_not_iter, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_map_in_arrow_legacy_accept_sequence_protocol(self): + # A sequence-protocol object (implements __getitem__ but not __iter__) is iterable via + # iter(...) even though it is not a collections.abc.Iterable, so the legacy flag must + # accept it too. + class SequenceOnly: + def __init__(self, items): + self._items = items + + def __getitem__(self, index): + return self._items[index] + + self.assertFalse(hasattr(SequenceOnly([]), "__iter__")) + + def returns_sequence(iterator): + return SequenceOnly([batch for batch in iterator]) + + with self.sql_conf( + {"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled": True} + ): + df = self.spark.range(10) + actual = df.mapInArrow(returns_sequence, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + def test_map_in_arrow_with_limit(self): def get_size(iterator): for batch in iterator: @@ -185,8 +222,18 @@ def test_self_join(self): def test_map_in_arrow_with_barrier_mode(self): df = self.spark.range(10) + def func0(iterator): + from pyspark import BarrierTaskContext + + BarrierTaskContext.get() + for batch in iterator: + yield batch + + with self.assertRaisesRegex(PythonException, "\\[NOT_IN_BARRIER_STAGE\\]"): + df.mapInArrow(func0, "id long", False).collect() + def func1(iterator): - from pyspark import TaskContext, BarrierTaskContext + from pyspark import BarrierTaskContext, TaskContext tc = TaskContext.get() assert tc is not None @@ -197,7 +244,7 @@ def func1(iterator): df.mapInArrow(func1, "id long", False).collect() def func2(iterator): - from pyspark import TaskContext, BarrierTaskContext + from pyspark import BarrierTaskContext, TaskContext tc = TaskContext.get() assert tc is not None diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py new file mode 100644 index 0000000000000..96b16aea9a999 --- /dev/null +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py @@ -0,0 +1,536 @@ +# +# 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. +# +import unittest +from decimal import Decimal + +from pyspark.errors import AnalysisException, PySparkValueError +from pyspark.sql import functions as sf +from pyspark.sql.types import ( + DecimalType, + DoubleType, + LongType, + StructField, + StructType, +) +from pyspark.sql.window import Window +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import ( + have_pyarrow, + pyarrow_requirement_message, +) +from pyspark.util import PythonEvalType + +if have_pyarrow: + from pyspark.sql.aggregator import Aggregator + from pyspark.sql.functions import udaf + + class Mean(Aggregator): + @property + def bufferSchema(self): + return StructType([StructField("sum", DoubleType()), StructField("count", LongType())]) + + @property + def outputType(self): + return DoubleType() + + def zero(self): + return (0.0, 0) + + def reduce(self, buffer, value): + (v,) = value + if v is None: # ignore nulls, like SQL avg + return buffer + return (buffer[0] + v, buffer[1] + 1) + + def merge(self, b1, b2): + return (b1[0] + b2[0], b1[1] + b2[1]) + + def finish(self, buffer): + return buffer[0] / buffer[1] if buffer[1] else None + + class DecimalSum(Aggregator): + # Non-trivial output/buffer type: the result column and the intermediate buffer are both + # DecimalType, exercising explicit Arrow typing of the emitted arrays (a bare + # ``pa.array([Decimal(...)])`` would infer a decimal type whose precision/scale need not + # match the declared one). + @property + def bufferSchema(self): + return StructType([StructField("total", DecimalType(20, 4))]) + + @property + def outputType(self): + return DecimalType(20, 4) + + def zero(self): + return (Decimal(0),) + + def reduce(self, buffer, value): + (v,) = value + return buffer if v is None else (buffer[0] + Decimal(str(v)),) + + def merge(self, b1, b2): + return (b1[0] + b2[0],) + + def finish(self, buffer): + return buffer[0] + + class SumSquares(Aggregator): + @property + def bufferSchema(self): + return StructType([StructField("sumsq", DoubleType())]) + + @property + def outputType(self): + return DoubleType() + + def zero(self): + return (0.0,) + + def reduce(self, buffer, value): + (v,) = value + return (buffer[0] + float(v) * float(v),) + + def merge(self, b1, b2): + return (b1[0] + b2[0],) + + def finish(self, buffer): + return buffer[0] + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowPythonAggregatorTestsMixin: + def _data(self): + # 100 rows across 5 keys; repartition so each key is split across partitions, + # exercising map-side PARTIAL combine + post-shuffle FINAL merge. + return ( + self.spark.range(0, 100) + .select((sf.col("id") % 5).alias("k"), sf.col("id").cast("double").alias("v")) + .repartition(4, sf.col("v") % 3) + ) + + def test_incremental_aggregator_matches_builtin_mean(self): + df = self._data() + result = df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect() + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + got = {r["k"]: r["m"] for r in result} + exp = {r["k"]: r["m"] for r in expected} + self.assertEqual(got, exp) + + def test_incremental_aggregator_no_group(self): + df = self._data() + result = df.agg(udaf(Mean())(sf.col("v")).alias("m")).collect() + expected = df.agg(sf.avg("v").alias("m")).collect() + self.assertAlmostEqual(result[0]["m"], expected[0]["m"], places=6) + + def test_incremental_aggregator_empty_global_input(self): + # A global aggregation over empty input must still return one identity row: finish(zero). + empty = self._data().limit(0) + result = empty.agg(udaf(Mean())(sf.col("v")).alias("m")).collect() + expected = empty.agg(sf.avg("v").alias("m")).collect() + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["m"], expected[0]["m"]) + + def test_incremental_aggregator_custom_buffer(self): + df = self._data() + result = ( + df.groupBy("k").agg(udaf(SumSquares())(sf.col("v")).alias("s")).orderBy("k").collect() + ) + expected = ( + df.groupBy("k").agg(sf.sum(sf.col("v") * sf.col("v")).alias("s")).orderBy("k").collect() + ) + got = {r["k"]: r["s"] for r in result} + exp = {r["k"]: r["s"] for r in expected} + for k in exp: + self.assertAlmostEqual(got[k], exp[k], places=6) + + def test_incremental_aggregator_decimal_output(self): + # Non-trivial output/buffer type (DecimalType), crossing the shuffle as a decimal buffer + # and emitted as a decimal result -- guards the explicit Arrow typing of both stages. + df = self._data() + result = ( + df.groupBy("k").agg(udaf(DecimalSum())(sf.col("v")).alias("s")).orderBy("k").collect() + ) + expected = df.groupBy("k").agg(sf.sum("v").alias("s")).orderBy("k").collect() + got = {r["k"]: r["s"] for r in result} + exp = {r["k"]: r["s"] for r in expected} + for k in exp: + self.assertIsInstance(got[k], Decimal) + self.assertAlmostEqual(float(got[k]), exp[k], places=4) + + def test_incremental_aggregator_null_inputs(self): + # reduce must tolerate null input values; the null-skipping Mean should match SQL avg, + # including a group whose values are all null (identity buffer -> finish returns None). + df = self.spark.createDataFrame( + [("a", 1.0), ("a", None), ("a", 3.0), ("b", None), ("b", None)], + "k string, v double", + ) + result = df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect() + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + got = {r["k"]: r["m"] for r in result} + exp = {r["k"]: r["m"] for r in expected} + self.assertEqual(got, exp) + + def test_multiple_incremental_aggregators(self): + # Two aggregators with different buffer schemas over the same input in one agg call. + df = self._data() + result = ( + df.groupBy("k") + .agg( + udaf(Mean())(sf.col("v")).alias("m"), + udaf(SumSquares())(sf.col("v")).alias("s"), + ) + .orderBy("k") + .collect() + ) + expected = ( + df.groupBy("k") + .agg( + sf.avg("v").alias("m"), + sf.sum(sf.col("v") * sf.col("v")).alias("s"), + ) + .orderBy("k") + .collect() + ) + got_m = {r["k"]: r["m"] for r in result} + got_s = {r["k"]: r["s"] for r in result} + for r in expected: + self.assertAlmostEqual(got_m[r["k"]], r["m"], places=6) + self.assertAlmostEqual(got_s[r["k"]], r["s"], places=6) + + def test_result_independent_of_partition_count(self): + # Partial buffers must merge to the same result regardless of how keys are split. + base = self.spark.range(0, 60).select( + (sf.col("id") % 3).alias("k"), sf.col("id").cast("double").alias("v") + ) + results = [] + for n in (1, 2, 7): + rows = ( + base.repartition(n, sf.col("v")) + .groupBy("k") + .agg(udaf(Mean())(sf.col("v")).alias("m")) + .orderBy("k") + .collect() + ) + results.append({r["k"]: r["m"] for r in rows}) + self.assertEqual(results[0], results[1]) + self.assertEqual(results[1], results[2]) + + def test_sql_registration(self): + # Register the aggregator and invoke it from SQL text. + df = self._data() + df.createOrReplaceTempView("agg_input") + self.spark.udf.register("my_mean", udaf(Mean())) + result = self.spark.sql( + "SELECT k, my_mean(v) AS m FROM agg_input GROUP BY k ORDER BY k" + ).collect() + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + got = {r["k"]: r["m"] for r in result} + exp = {r["k"]: r["m"] for r in expected} + self.assertEqual(got, exp) + + def test_named_arguments(self): + # A named argument (both DataFrame and SQL forms) must feed the aggregator's value tuple, + # not be silently dropped. + df = self._data() + result = df.groupBy("k").agg(udaf(Mean())(v=sf.col("v")).alias("m")).orderBy("k").collect() + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"] for r in expected}) + + df.createOrReplaceTempView("agg_input") + self.spark.udf.register("my_mean", udaf(Mean())) + sql_result = self.spark.sql( + "SELECT k, my_mean(v => v) AS m FROM agg_input GROUP BY k ORDER BY k" + ).collect() + self.assertEqual({r["k"]: r["m"] for r in sql_result}, {r["k"]: r["m"] for r in expected}) + + def test_distinct_and_filter_rejected(self): + # Neither DISTINCT nor FILTER is honored by the two-stage operator, so both must be + # rejected at analysis rather than silently returning the non-distinct/unfiltered result. + df = self._data() + df.createOrReplaceTempView("agg_input") + self.spark.udf.register("my_mean", udaf(Mean())) + with self.assertRaises(AnalysisException): + self.spark.sql("SELECT my_mean(DISTINCT v) FROM agg_input GROUP BY k").collect() + with self.assertRaises(AnalysisException): + self.spark.sql( + "SELECT my_mean(v) FILTER (WHERE v > 0) FROM agg_input GROUP BY k" + ).collect() + + def test_pivot_rejected(self): + # ResolvePivot must reject the incremental aggregator (its null-ignoring fallback rewrite + # would produce wrong results), like it already rejects pandas UDAFs. + df = self.spark.createDataFrame( + [("a", "x", 1.0), ("a", "y", 2.0), ("b", "x", 3.0)], + "k string, p string, v double", + ) + with self.assertRaises(AnalysisException): + df.groupBy("k").pivot("p").agg(udaf(Mean())(sf.col("v"))).collect() + + def test_window_unbounded(self): + # Unbounded partition frame: every row gets its whole group's aggregate. Cross-checked + # against the equivalent SQL window aggregate. + df = self._data() + w = Window.partitionBy("k") + result = df.withColumn("m", udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect() + expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k", "v").collect() + self.assertEqual(len(result), len(expected)) + for r, e in zip(result, expected): + self.assertAlmostEqual(r["m"], e["m"], places=6) + + def test_window_running_frame(self): + # Ordered, growing frame (unbounded preceding .. current row): a running aggregate that + # exercises the per-row bounded-frame path in the worker. + df = self._data() + w = ( + Window.partitionBy("k") + .orderBy("v") + .rowsBetween(Window.unboundedPreceding, Window.currentRow) + ) + result = df.withColumn("m", udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect() + expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k", "v").collect() + self.assertEqual(len(result), len(expected)) + for r, e in zip(result, expected): + self.assertAlmostEqual(r["m"], e["m"], places=6) + + def test_window_sliding_frame(self): + # Sliding frame (1 preceding .. 1 following) with a custom single-field buffer aggregator. + df = self._data() + w = Window.partitionBy("k").orderBy("v").rowsBetween(-1, 1) + result = ( + df.withColumn("s", udaf(SumSquares())(sf.col("v")).over(w)).orderBy("k", "v").collect() + ) + expected = ( + df.withColumn("s", sf.sum(sf.col("v") * sf.col("v")).over(w)) + .orderBy("k", "v") + .collect() + ) + self.assertEqual(len(result), len(expected)) + for r, e in zip(result, expected): + self.assertAlmostEqual(r["s"], e["s"], places=4) + + def test_window_bounded_preceding_frame(self): + # A fixed number of preceding rows exercises both branches of the running-buffer + # optimization: the lower bound is clamped to 0 for the first rows (running buffer is + # extended in place) and then advances (each frame is refolded from zero). + df = self._data() + w = Window.partitionBy("k").orderBy("v").rowsBetween(-3, Window.currentRow) + result = df.withColumn("m", udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect() + expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k", "v").collect() + self.assertEqual(len(result), len(expected)) + for r, e in zip(result, expected): + self.assertAlmostEqual(r["m"], e["m"], places=6) + + def test_window_decimal_output(self): + # A non-trivial (Decimal) output type over a window, exercising the explicit + # ``pa.array(..., type=result_type)`` typing on the window path. + df = self._data() + w = Window.partitionBy("k") + result = ( + df.withColumn("s", udaf(DecimalSum())(sf.col("v")).over(w)).orderBy("k", "v").collect() + ) + expected = df.withColumn("s", sf.sum("v").over(w)).orderBy("k", "v").collect() + self.assertEqual(len(result), len(expected)) + for r, e in zip(result, expected): + self.assertIsInstance(r["s"], Decimal) + self.assertAlmostEqual(float(r["s"]), e["s"], places=4) + + def test_window_mixed_python_udf_rejected(self): + # An incremental aggregator and a grouped-agg pandas UDF over the same window are both + # Python window functions but use different eval types, so they cannot share one operator. + # This must raise a clear analysis error rather than an internal assertion. + from pyspark.sql.functions import PandasUDFType, pandas_udf + + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def pandas_mean(v): + return v.mean() + + df = self._data() + w = Window.partitionBy("k") + with self.assertRaises(AnalysisException) as ctx: + df.select(udaf(Mean())(sf.col("v")).over(w), pandas_mean(sf.col("v")).over(w)).collect() + self.assertEqual( + ctx.exception.getCondition(), + "UNSUPPORTED_FEATURE.MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW", + ) + + def test_mixed_with_other_aggregate_rejected(self): + # An incremental aggregator mixed with another aggregate in one Aggregate is unsupported; + # the error must be the dedicated (non-pandas) placement error. + df = self._data() + with self.assertRaises(AnalysisException) as ctx: + df.groupBy("k").agg( + udaf(Mean())(sf.col("v")).alias("m"), sf.count("*").alias("c") + ).collect() + self.assertEqual(ctx.exception.getCondition(), "INVALID_PYTHON_UDF_PLACEMENT") + + def test_duplicate_buffer_field_names_rejected(self): + # Duplicate buffer field names silently collapse on the map side and then fail with an + # opaque Arrow error post-shuffle; reject them where the aggregator is created. + class DupBuffer(Mean): + @property + def bufferSchema(self): + return StructType([StructField("x", DoubleType()), StructField("x", LongType())]) + + with self.assertRaises(PySparkValueError) as ctx: + udaf(DupBuffer()) + self.check_error( + exception=ctx.exception, + errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT", + messageParameters={"field_names": "x"}, + ) + + def test_float_grouping_keys_normalized(self): + # 0.0 / -0.0 must fall into one group, and NaN keys (unequal to themselves) must group + # together, matching SQL aggregate semantics -- guards grouping-key normalization and the + # NaN handling of the map-side hash combine. + zeros = self.spark.createDataFrame( + [(0.0, 1.0), (-0.0, 2.0), (0.0, 3.0)], "k double, v double" + ) + zero_rows = zeros.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect() + self.assertEqual(len(zero_rows), 1) + self.assertAlmostEqual(zero_rows[0]["m"], 2.0, places=6) + + nans = self.spark.createDataFrame( + [(float("nan"), 1.0), (float("nan"), 3.0)], "k double, v double" + ) + nan_rows = nans.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect() + self.assertEqual(len(nan_rows), 1) + self.assertAlmostEqual(nan_rows[0]["m"], 2.0, places=6) + + def test_complex_grouping_key(self): + # A struct grouping key exercises the map-side hash combine's canonicalization of complex + # keys (dict -> hashable) as well as the authoritative FINAL re-grouping. + df = self.spark.createDataFrame( + [(1, "a", 1.0), (1, "a", 3.0), (2, "b", 10.0)], + "i int, s string, v double", + ).select(sf.struct("i", "s").alias("k"), sf.col("v")) + result = df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect() + got = {(r["k"]["i"], r["k"]["s"]): r["m"] for r in result} + self.assertEqual(got, {(1, "a"): 2.0, (2, "b"): 10.0}) + + def test_bounded_map_side_combine(self): + # A small maxRecordsPerBatch forces the map-side PARTIAL stage to flush its per-key buffer + # in bounded chunks (emitting duplicate keys that the FINAL stage re-merges authoritatively) + # instead of holding every key for the whole partition. The result must be unchanged -- + # this guards the cap/flush + chunked-emission path against OOM on high-cardinality keys. + df = self._data() + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 2}): + result = ( + df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect() + ) + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"] for r in expected}) + + def test_missing_buffer_schema_rejected(self): + # udaf() enforces a struct buffer schema up front, but a low-level construction (or a + # malformed Connect proto) can build an incremental aggregator UDF without one. That must + # surface a classed planner error, not a bare IllegalArgumentException / ClassCastException. + from pyspark.sql.utils import is_remote + + if is_remote(): + from pyspark.sql.connect.udf import UserDefinedFunction + else: + from pyspark.sql.udf import UserDefinedFunction # type: ignore[assignment] + + bad = UserDefinedFunction( + Mean(), + returnType=DoubleType(), + name="bad_mean", + evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, + deterministic=True, + )._wrapped() + df = self._data() + with self.assertRaises(AnalysisException) as ctx: + df.groupBy("k").agg(bad(sf.col("v"))).collect() + self.assertEqual(ctx.exception.getCondition(), "INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA") + + def test_mixed_pandas_udaf_and_incremental_rejected(self): + # Mixing a grouped-agg pandas UDAF with an incremental aggregator in one Aggregate is + # unsupported. It falls through to the dedicated placement error, which must name BOTH + # offending functions rather than dropping the co-offending pandas UDAF. + from pyspark.sql.functions import PandasUDFType, pandas_udf + + @pandas_udf("double", PandasUDFType.GROUPED_AGG) + def pandas_mean(v): + return v.mean() + + df = self._data() + with self.assertRaises(AnalysisException) as ctx: + df.groupBy("k").agg( + udaf(Mean())(sf.col("v")).alias("m"), + pandas_mean(sf.col("v")).alias("pm"), + ).collect() + self.assertEqual(ctx.exception.getCondition(), "INVALID_PYTHON_UDF_PLACEMENT") + message = ctx.exception.getMessage() + self.assertIn("Mean", message) + self.assertIn("pandas_mean", message) + + +class ArrowPythonAggregatorTests(ArrowPythonAggregatorTestsMixin, ReusedSQLTestCase): + pass + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowPythonAggregatorProfilerTests(unittest.TestCase): + # Profiling is not supported for incremental aggregators: their ``func`` is an ``Aggregator`` + # object, so the profiler wrappers would either break the worker (a plain function has no + # ``bufferSchema`` / ``zero`` / ``reduce``) or fail on the driver in + # ``inspect.getsourcelines(f.__code__)``. Enabling a profiler must therefore fall back to the + # non-profiled path (with a warning) and still compute the correct result, not crash. These + # confs are set at session creation, so this needs its own session (classic only). + def _run(self, conf_key): + import warnings + + from pyspark import SparkConf + from pyspark.sql import SparkSession + + conf = SparkConf().set(conf_key, "true") + spark = ( + SparkSession.builder.master("local[4]") + .config(conf=conf) + .appName(self.__class__.__name__) + .getOrCreate() + ) + try: + df = spark.range(0, 20).select( + (sf.col("id") % 3).alias("k"), sf.col("id").cast("double").alias("v") + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = ( + df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect() + ) + expected = df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect() + self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"] for r in expected}) + self.assertTrue( + any("incremental Python aggregators" in str(w.message) for w in caught), + "expected an unsupported-profiling warning", + ) + finally: + spark.stop() + + def test_cpu_profiler_falls_back(self): + self._run("spark.python.profile") + + def test_memory_profiler_falls_back(self): + self._run("spark.python.profile.memory") + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py index f4f1219ee7cd1..e56fc33ae1630 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf.py @@ -15,13 +15,13 @@ # limitations under the License. # -from decimal import Decimal import unittest +from decimal import Decimal -from pyspark.errors import AnalysisException, PythonException, PySparkNotImplementedError +from pyspark.errors import AnalysisException, PySparkNotImplementedError, PythonException from pyspark.loose_version import LooseVersion from pyspark.sql import Row -from pyspark.sql.functions import udf, col +from pyspark.sql.functions import col, udf from pyspark.sql.tests.test_udf import BaseUDFTestsMixin from pyspark.sql.types import ( ArrayType, diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py index 4367f6b959d97..04b5afec944cd 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py @@ -18,8 +18,8 @@ import datetime import unittest -from pyspark.sql.functions import col, udf from pyspark.sql import Row +from pyspark.sql.functions import col, udf from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( assertDataFrameEqual, diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf.py b/python/pyspark/sql/tests/arrow/test_arrow_udf.py index a7f2ef197fa70..132a2e3d9c92d 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udf.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udf.py @@ -15,29 +15,30 @@ # limitations under the License. # +import datetime import os import time import unittest -import datetime from typing import Iterator -from pyspark.sql.functions import arrow_udf, ArrowUDFType, PandasUDFType -from pyspark.sql import functions as F, Row +from pyspark.errors import ParseException, PySparkTypeError +from pyspark.sql import Row +from pyspark.sql import functions as F +from pyspark.sql.functions import ArrowUDFType, PandasUDFType, arrow_udf from pyspark.sql.types import ( + DayTimeIntervalType, DoubleType, - StructType, - StructField, LongType, - DayTimeIntervalType, + StructField, + StructType, VariantType, ) -from pyspark.errors import ParseException, PySparkTypeError -from pyspark.util import PythonEvalType from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pyarrow, pyarrow_requirement_message, ) +from pyspark.util import PythonEvalType @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py b/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py index 0227524d43414..7a992295d6120 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py @@ -15,31 +15,31 @@ # limitations under the License. # -import unittest import logging +import unittest +from typing import Iterator, Tuple -from pyspark.sql.functions import arrow_udf, ArrowUDFType -from pyspark.util import PythonEvalType, is_remote_only +from pyspark.errors import AnalysisException, PythonException from pyspark.sql import Row +from pyspark.sql import functions as sf +from pyspark.sql.functions import ArrowUDFType, arrow_udf from pyspark.sql.types import ( ArrayType, - YearMonthIntervalType, - StructType, StructField, + StructType, VariantType, VariantVal, + YearMonthIntervalType, ) -from pyspark.sql import functions as sf -from pyspark.errors import AnalysisException, PythonException +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( + assertDataFrameEqual, have_numpy, - numpy_requirement_message, have_pyarrow, + numpy_requirement_message, pyarrow_requirement_message, - assertDataFrameEqual, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from typing import Iterator, Tuple +from pyspark.util import PythonEvalType, is_remote_only @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) @@ -111,8 +111,8 @@ def sum(v): @property def arrow_agg_weighted_mean_udf(self): - import pyarrow as pa import numpy as np + import pyarrow as pa @arrow_udf("double", ArrowUDFType.GROUPED_AGG) def weighted_mean(v, w): @@ -1062,9 +1062,10 @@ def test_iterator_grouped_agg_single_column(self): """ Test iterator API for grouped aggregation with single column. """ - import pyarrow as pa from typing import Iterator + import pyarrow as pa + @arrow_udf("double") def arrow_mean_iter(it: Iterator[pa.Array]) -> float: sum_val = 0.0 @@ -1089,8 +1090,8 @@ def test_iterator_grouped_agg_multiple_columns(self): """ Test iterator API for grouped aggregation with multiple columns. """ - import pyarrow as pa import numpy as np + import pyarrow as pa @arrow_udf("double") def arrow_weighted_mean_iter(it: Iterator[Tuple[pa.Array, pa.Array]]) -> float: @@ -1129,9 +1130,10 @@ def test_iterator_grouped_agg_eval_type(self): """ Test that the eval type is correctly inferred for iterator grouped agg UDFs. """ - import pyarrow as pa from typing import Iterator + import pyarrow as pa + @arrow_udf("double") def arrow_sum_iter(it: Iterator[pa.Array]) -> float: total = 0.0 @@ -1146,9 +1148,10 @@ def test_iterator_grouped_agg_partial_consumption(self): Test that iterator grouped agg UDF can partially consume batches. This ensures that batches are processed one by one without loading all data into memory. """ - import pyarrow as pa from typing import Iterator + import pyarrow as pa + # Create a dataset with multiple batches per group # Use small batch size to ensure multiple batches per group # Use same value for all data points to avoid ordering issues diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py b/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py index d574c7266537a..761654455870e 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py @@ -15,55 +15,53 @@ # limitations under the License. # +import datetime +import logging import os import random import time import unittest -import datetime -import logging from decimal import Decimal from typing import Iterator, Tuple -from pyspark.util import PythonEvalType - -from pyspark.sql.functions import arrow_udf, ArrowUDFType +from pyspark.errors import AnalysisException, PythonException from pyspark.sql import functions as F +from pyspark.sql.functions import ArrowUDFType, arrow_udf from pyspark.sql.types import ( - IntegerType, - ByteType, - StructType, - ShortType, + ArrayType, + BinaryType, BooleanType, - LongType, - FloatType, - DoubleType, + ByteType, DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + MapType, + Row, + ShortType, StringType, - ArrayType, StructField, - Row, - MapType, - BinaryType, + StructType, YearMonthIntervalType, ) -from pyspark.errors import AnalysisException, PythonException +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( + assertDataFrameEqual, have_numpy, - numpy_requirement_message, have_pyarrow, + numpy_requirement_message, pyarrow_requirement_message, - assertDataFrameEqual, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.util import is_remote_only +from pyspark.util import PythonEvalType, is_remote_only @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ScalarArrowUDFTestsMixin: @property def nondeterministic_arrow_udf(self): - import pyarrow as pa import numpy as np + import pyarrow as pa @arrow_udf("double") def random_udf(v): @@ -73,8 +71,8 @@ def random_udf(v): @property def nondeterministic_arrow_iter_udf(self): - import pyarrow as pa import numpy as np + import pyarrow as pa @arrow_udf("double", ArrowUDFType.SCALAR_ITER) def random_udf(it): @@ -343,6 +341,7 @@ def extract_second(d): def test_arrow_udf_output_timestamps_ltz(self): from zoneinfo import ZoneInfo + import pyarrow as pa tz = self.spark.conf.get("spark.sql.session.timeZone") @@ -1267,6 +1266,33 @@ def return_one(iterator): result = df.select(return_one("id").alias("one")).collect() self.assertEqual(expected, result) + def test_arrow_udf_task_context(self): + import pyarrow as pa + + from pyspark import TaskContext + + def partition_ids(size): + task_context = TaskContext.get() + assert task_context is not None + return pa.array([task_context.partitionId()] * size, type=pa.int32()) + + @arrow_udf("int") + def scalar_task_context(values): + return partition_ids(len(values)) + + @arrow_udf("int", ArrowUDFType.SCALAR_ITER) + def scalar_iter_task_context(iterator): + for values in iterator: + yield partition_ids(len(values)) + + for task_context_udf in [scalar_task_context, scalar_iter_task_context]: + rows = ( + self.spark.range(10, numPartitions=2) + .select(task_context_udf("id").alias("partition_id")) + .collect() + ) + self.assertEqual({0, 1}, {row.partition_id for row in rows}) + class ScalarArrowUDFTests(ScalarArrowUDFTestsMixin, ReusedSQLTestCase): @classmethod diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_typehints.py b/python/pyspark/sql/tests/arrow/test_arrow_udf_typehints.py index b39b4e5993e61..3192679c17b89 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udf_typehints.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_typehints.py @@ -16,21 +16,21 @@ # import unittest from inspect import signature -from typing import Union, Iterator, Tuple, get_type_hints +from typing import Iterator, Tuple, Union, get_type_hints +from pyspark.sql import Row from pyspark.sql import functions as sf +from pyspark.sql.pandas.functions import ArrowUDFType, arrow_udf +from pyspark.sql.pandas.typehints import infer_eval_type, infer_group_arrow_eval_type +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( + have_numpy, have_pandas, - pandas_requirement_message, have_pyarrow, - pyarrow_requirement_message, - have_numpy, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.sql.pandas.typehints import infer_eval_type, infer_group_arrow_eval_type -from pyspark.sql.pandas.functions import arrow_udf, ArrowUDFType -from pyspark.sql import Row from pyspark.util import PythonEvalType if have_pyarrow: diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_window.py b/python/pyspark/sql/tests/arrow/test_arrow_udf_window.py index 50793a50cd879..dac9a3bf011c5 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udf_window.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_window.py @@ -15,22 +15,23 @@ # limitations under the License. # -import unittest import logging +import unittest -from pyspark.sql.functions import arrow_udf, ArrowUDFType -from pyspark.util import PythonEvalType, is_remote_only -from pyspark.sql import Row, functions as sf +from pyspark.errors import AnalysisException, PySparkTypeError, PythonException +from pyspark.sql import Row +from pyspark.sql import functions as sf +from pyspark.sql.functions import ArrowUDFType, arrow_udf from pyspark.sql.window import Window -from pyspark.errors import AnalysisException, PythonException, PySparkTypeError +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( + assertDataFrameEqual, have_numpy, - numpy_requirement_message, have_pyarrow, + numpy_requirement_message, pyarrow_requirement_message, - assertDataFrameEqual, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.util import PythonEvalType, is_remote_only @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udtf.py b/python/pyspark/sql/tests/arrow/test_arrow_udtf.py index b82523005ac72..ac695a9f8001e 100644 --- a/python/pyspark/sql/tests/arrow/test_arrow_udtf.py +++ b/python/pyspark/sql/tests/arrow/test_arrow_udtf.py @@ -14,17 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest import logging +import unittest from typing import Iterator, Optional -from pyspark.errors import PySparkAttributeError -from pyspark.errors import PythonException +from pyspark.errors import PySparkAttributeError, PythonException from pyspark.sql.functions import arrow_udtf, lit -from pyspark.sql.types import Row, StructType, StructField, IntegerType +from pyspark.sql.types import IntegerType, Row, StructField, StructType +from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message -from pyspark.testing import assertDataFrameEqual from pyspark.util import is_remote_only if have_pyarrow: @@ -689,7 +688,7 @@ def eval(self, input_table: "pa.Table") -> Iterator["pa.Table"]: yield result_table from pyspark.sql.functions import udtf - from pyspark.sql.types import StructType, StructField, IntegerType + from pyspark.sql.types import IntegerType, StructField, StructType @udtf(returnType=StructType([StructField("multiplied", IntegerType())])) class MultiplyUDTF: diff --git a/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.csv b/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.csv index 4f5d13daaf10a..ae47d24a24dba 100644 --- a/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.csv +++ b/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.csv @@ -1,4 +1,4 @@ -SQL Type \ Value@Type [None, None]@list [True, False]@list ['a', 'b']@list ['12', '34']@list [Decimal('1'), Decimal('2')]@list [{'a': 1}, {'b': 2}]@list [1 2]@ndarray[int8] [1 2]@ndarray[int16] [1 2]@ndarray[int32] [1 2]@ndarray[int64] [1 2]@ndarray[uint8] [1 2]@ndarray[uint16] [1 2]@ndarray[uint32] [1 2]@ndarray[uint64] [1. 2.]@ndarray[float16] [1. 2.]@ndarray[float32] [1. 2.]@ndarray[float64] [1.+0.j 2.+0.j]@ndarray[complex64] [1.+0.j 2.+0.j]@ndarray[complex128] [array([1, 2, 3], dtype=int32), @list ['1970-01-01T00:00:00.000000000'@ndarray[datetime64[ns]] ['1970-01-01T05:00:00.000000000'@ndarray[datetime64[ns]] [Timedelta('1 days 00:00:00'), T@list ['A', 'B'] Categories (2, object@Categorical "{""_1"":{""0"":1,""1"":2}}@Dataframe[_1 int64]" +SQL Type \ Value@Type [None, None]@list [True, False]@list ['a', 'b']@list ['12', '34']@list [Decimal('1'), Decimal('2')]@list [{'a': 1}, {'b': 2}]@list [1 2]@ndarray[int8] [1 2]@ndarray[int16] [1 2]@ndarray[int32] [1 2]@ndarray[int64] [1 2]@ndarray[uint8] [1 2]@ndarray[uint16] [1 2]@ndarray[uint32] [1 2]@ndarray[uint64] [1. 2.]@ndarray[float16] [1. 2.]@ndarray[float32] [1. 2.]@ndarray[float64] [1.+0.j 2.+0.j]@ndarray[complex64] [1.+0.j 2.+0.j]@ndarray[complex128] [array([1, 2, 3], dtype=int32), @list ['1970-01-01T00:00:00.000000000'@ndarray[datetime64[ns]] ['1970-01-01T05:00:00.000000000'@ndarray[datetime64[ns]] [Timedelta('1 days 00:00:00'), T@list ['A', 'B'] Categories (2, object@Categorical {'_1': [1, 2]}@Dataframe[_1 int64] boolean [None, None] [True, False] X X X X [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] X [True, True] [True, True] X X X X X X X X tinyint [None, None] [1, 0] X [12, 34] [1, 2] X [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] X X X X X X X X smallint [None, None] [1, 0] X [12, 34] [1, 2] X [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] X X X X X X X X diff --git a/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.md b/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.md index f8d7f11310fa1..54a0d16389d22 100644 --- a/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.md +++ b/python/pyspark/sql/tests/coercion/pandas_2/golden_pandas_udf_return_type_coercion_base.md @@ -1,17 +1,17 @@ -| SQL Type \ Value@Type | [None, None]@list | [True, False]@list | ['a', 'b']@list | ['12', '34']@list | [Decimal('1'), Decimal('2')]@list | [{'a': 1}, {'b': 2}]@list | [1 2]@ndarray[int8] | [1 2]@ndarray[int16] | [1 2]@ndarray[int32] | [1 2]@ndarray[int64] | [1 2]@ndarray[uint8] | [1 2]@ndarray[uint16] | [1 2]@ndarray[uint32] | [1 2]@ndarray[uint64] | [1. 2.]@ndarray[float16] | [1. 2.]@ndarray[float32] | [1. 2.]@ndarray[float64] | [1.+0.j 2.+0.j]@ndarray[complex64] | [1.+0.j 2.+0.j]@ndarray[complex128] | [array([1, 2, 3], dtype=int32), @list | ['1970-01-01T00:00:00.000000000'@ndarray[datetime64[ns]] | ['1970-01-01T05:00:00.000000000'@ndarray[datetime64[ns]] | [Timedelta('1 days 00:00:00'), T@list | ['A', 'B'] Categories (2, object@Categorical | {"_1":{"0":1,"1":2}}@Dataframe[_1 int64] | -|--------------------------|---------------------|----------------------|-------------------|---------------------|------------------------------------------|-----------------------------|------------------------------|------------------------------|------------------------------------------|-----------------------------------------|------------------------------|------------------------------|------------------------------|-------------------------|----------------------------|------------------------------|------------------------------|--------------------------------------|---------------------------------------|-----------------------------------------|------------------------------------------------------------|------------------------------------------------------------|-----------------------------------------|------------------------------------------------|--------------------------------------------| -| boolean | [None, None] | [True, False] | X | X | X | X | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | X | [True, True] | [True, True] | X | X | X | X | X | X | X | X | -| tinyint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| smallint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| int | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| bigint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | [0, 86400000000000] | [18000000000000, 104400000000000] | [86400000000000, 172800000000000] | X | X | -| string | [None, None] | X | ['a', 'b'] | ['12', '34'] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | ['A', 'B'] | X | -| date | [None, None] | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | X | X | X | X | X | X | X | X | [datetime.date(1970, 1, 1), datetime.dat | [datetime.date(1970, 1, 1), datetime.dat | X | X | X | -| timestamp | [None, None] | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | X | X | X | X | X | X | [datetime.datetime(1970, 1, 1, 0, 0), da | [datetime.datetime(1970, 1, 1, 5, 0), da | X | X | X | -| float | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | -| double | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | -| array<int> | [None, None] | X | X | [[1, 2], [3, 4]] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [[1, 2, 3], [1, 2, 3]] | X | X | X | X | X | -| binary | [None, None] | [b'\x01', b''] | [b'a', b'b'] | [b'12', b'34'] | X | X | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | X | [b'', b''] | [b'', b''] | [b'', b''] | [b'A', b'B'] | X | -| decimal(10,0) | [None, None] | X | X | X | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | X | X | X | X | X | X | -| map<string,int> | [None, None] | X | X | X | X | [{'a': 1}, {'b': 2}] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | -| struct<_1:int> | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [Row(_1=1), Row(_1=2)] | \ No newline at end of file +| SQL Type \ Value@Type | [None, None]@list | [True, False]@list | ['a', 'b']@list | ['12', '34']@list | [Decimal('1'), Decimal('2')]@list | [{'a': 1}, {'b': 2}]@list | [1 2]@ndarray[int8] | [1 2]@ndarray[int16] | [1 2]@ndarray[int32] | [1 2]@ndarray[int64] | [1 2]@ndarray[uint8] | [1 2]@ndarray[uint16] | [1 2]@ndarray[uint32] | [1 2]@ndarray[uint64] | [1. 2.]@ndarray[float16] | [1. 2.]@ndarray[float32] | [1. 2.]@ndarray[float64] | [1.+0.j 2.+0.j]@ndarray[complex64] | [1.+0.j 2.+0.j]@ndarray[complex128] | [array([1, 2, 3], dtype=int32), @list | ['1970-01-01T00:00:00.000000000'@ndarray[datetime64[ns]] | ['1970-01-01T05:00:00.000000000'@ndarray[datetime64[ns]] | [Timedelta('1 days 00:00:00'), T@list | ['A', 'B'] Categories (2, object@Categorical | {'_1': [1, 2]}@Dataframe[_1 int64] | +|--------------------------|---------------------|----------------------|-------------------|---------------------|------------------------------------------|-----------------------------|------------------------------|------------------------------|------------------------------------------|-----------------------------------------|------------------------------|------------------------------|------------------------------|-------------------------|----------------------------|------------------------------|------------------------------|--------------------------------------|---------------------------------------|-----------------------------------------|------------------------------------------------------------|------------------------------------------------------------|-----------------------------------------|------------------------------------------------|--------------------------------------| +| boolean | [None, None] | [True, False] | X | X | X | X | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | X | [True, True] | [True, True] | X | X | X | X | X | X | X | X | +| tinyint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| smallint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| int | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| bigint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | [0, 86400000000000] | [18000000000000, 104400000000000] | [86400000000000, 172800000000000] | X | X | +| string | [None, None] | X | ['a', 'b'] | ['12', '34'] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | ['A', 'B'] | X | +| date | [None, None] | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | X | X | X | X | X | X | X | X | [datetime.date(1970, 1, 1), datetime.dat | [datetime.date(1970, 1, 1), datetime.dat | X | X | X | +| timestamp | [None, None] | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | X | X | X | X | X | X | [datetime.datetime(1970, 1, 1, 0, 0), da | [datetime.datetime(1970, 1, 1, 5, 0), da | X | X | X | +| float | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | +| double | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | +| array<int> | [None, None] | X | X | [[1, 2], [3, 4]] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [[1, 2, 3], [1, 2, 3]] | X | X | X | X | X | +| binary | [None, None] | [b'\x01', b''] | [b'a', b'b'] | [b'12', b'34'] | X | X | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | X | [b'', b''] | [b'', b''] | [b'', b''] | [b'A', b'B'] | X | +| decimal(10,0) | [None, None] | X | X | X | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | X | X | X | X | X | X | +| map<string,int> | [None, None] | X | X | X | X | [{'a': 1}, {'b': 2}] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | +| struct<_1:int> | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [Row(_1=1), Row(_1=2)] | \ No newline at end of file diff --git a/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.csv b/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.csv index 5e12cc57c84af..0b955d06a5522 100644 --- a/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.csv +++ b/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.csv @@ -1,4 +1,4 @@ -SQL Type \ Value@Type [None, None]@list [True, False]@list ['a', 'b']@list ['12', '34']@list [Decimal('1'), Decimal('2')]@list [{'a': 1}, {'b': 2}]@list [1 2]@ndarray[int8] [1 2]@ndarray[int16] [1 2]@ndarray[int32] [1 2]@ndarray[int64] [1 2]@ndarray[uint8] [1 2]@ndarray[uint16] [1 2]@ndarray[uint32] [1 2]@ndarray[uint64] [1. 2.]@ndarray[float16] [1. 2.]@ndarray[float32] [1. 2.]@ndarray[float64] [1.+0.j 2.+0.j]@ndarray[complex64] [1.+0.j 2.+0.j]@ndarray[complex128] [array([1, 2, 3], dtype=int32), @list ['1970-01-01T00:00:00.000000' '1@ndarray[datetime64[us]] ['1970-01-01T05:00:00.000000' '1@ndarray[datetime64[us]] [Timedelta('1 days 00:00:00'), T@list ['A', 'B'] Categories (2, str): @Categorical "{""_1"":{""0"":1,""1"":2}}@Dataframe[_1 int64]" +SQL Type \ Value@Type [None, None]@list [True, False]@list ['a', 'b']@list ['12', '34']@list [Decimal('1'), Decimal('2')]@list [{'a': 1}, {'b': 2}]@list [1 2]@ndarray[int8] [1 2]@ndarray[int16] [1 2]@ndarray[int32] [1 2]@ndarray[int64] [1 2]@ndarray[uint8] [1 2]@ndarray[uint16] [1 2]@ndarray[uint32] [1 2]@ndarray[uint64] [1. 2.]@ndarray[float16] [1. 2.]@ndarray[float32] [1. 2.]@ndarray[float64] [1.+0.j 2.+0.j]@ndarray[complex64] [1.+0.j 2.+0.j]@ndarray[complex128] [array([1, 2, 3], dtype=int32), @list ['1970-01-01T00:00:00.000000' '1@ndarray[datetime64[us]] ['1970-01-01T05:00:00.000000' '1@ndarray[datetime64[us]] [Timedelta('1 days 00:00:00'), T@list ['A', 'B'] Categories (2, str): @Categorical {'_1': [1, 2]}@Dataframe[_1 int64] boolean [None, None] [True, False] X X X X [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] [True, True] X [True, True] [True, True] X X X X X X X X tinyint [None, None] [1, 0] X [12, 34] [1, 2] X [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] X X X X X X X X smallint [None, None] [1, 0] X [12, 34] [1, 2] X [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] [1, 2] X X X X X X X X diff --git a/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.md b/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.md index aa57ca6eec96a..257ae0e7548ca 100644 --- a/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.md +++ b/python/pyspark/sql/tests/coercion/pandas_3/golden_pandas_udf_return_type_coercion_base.md @@ -1,17 +1,17 @@ -| SQL Type \ Value@Type | [None, None]@list | [True, False]@list | ['a', 'b']@list | ['12', '34']@list | [Decimal('1'), Decimal('2')]@list | [{'a': 1}, {'b': 2}]@list | [1 2]@ndarray[int8] | [1 2]@ndarray[int16] | [1 2]@ndarray[int32] | [1 2]@ndarray[int64] | [1 2]@ndarray[uint8] | [1 2]@ndarray[uint16] | [1 2]@ndarray[uint32] | [1 2]@ndarray[uint64] | [1. 2.]@ndarray[float16] | [1. 2.]@ndarray[float32] | [1. 2.]@ndarray[float64] | [1.+0.j 2.+0.j]@ndarray[complex64] | [1.+0.j 2.+0.j]@ndarray[complex128] | [array([1, 2, 3], dtype=int32), @list | ['1970-01-01T00:00:00.000000' '1@ndarray[datetime64[us]] | ['1970-01-01T05:00:00.000000' '1@ndarray[datetime64[us]] | [Timedelta('1 days 00:00:00'), T@list | ['A', 'B'] Categories (2, str): @Categorical | {"_1":{"0":1,"1":2}}@Dataframe[_1 int64] | -|--------------------------|---------------------|----------------------|-------------------|--------------------------------|------------------------------------------|-----------------------------|------------------------------|------------------------------|------------------------------------------|-----------------------------------------|------------------------------|------------------------------|------------------------------|-------------------------|----------------------------|------------------------------|------------------------------|--------------------------------------|---------------------------------------|-----------------------------------------|------------------------------------------------------------|------------------------------------------------------------|-----------------------------------------|------------------------------------------------|--------------------------------------------| -| boolean | [None, None] | [True, False] | X | X | X | X | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | X | [True, True] | [True, True] | X | X | X | X | X | X | X | X | -| tinyint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| smallint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| int | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | -| bigint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | [0, 86400000000] | [18000000000, 104400000000] | [86400000000, 172800000000] | X | X | -| string | [None, None] | X | ['a', 'b'] | ['12', '34'] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | ['A', 'B'] | X | -| date | [None, None] | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | X | X | X | X | X | X | X | X | [datetime.date(1970, 1, 1), datetime.dat | [datetime.date(1970, 1, 1), datetime.dat | X | X | X | -| timestamp | [None, None] | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | X | X | X | X | X | X | [datetime.datetime(1970, 1, 1, 0, 0), da | [datetime.datetime(1970, 1, 1, 5, 0), da | X | X | X | -| float | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | -| double | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | -| array<int> | [None, None] | X | X | [[1, 2], [3, 4]] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [[1, 2, 3], [1, 2, 3]] | X | X | X | X | X | -| binary | [None, None] | [b'\x01', b''] | [b'a', b'b'] | [b'12', b'34'] | X | X | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | X | [b'', b''] | [b'', b''] | [b'', b''] | [b'A', b'B'] | X | -| decimal(10,0) | [None, None] | X | X | [Decimal('12'), Decimal('34')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | X | X | X | X | X | X | -| map<string,int> | [None, None] | X | X | X | X | [{'a': 1}, {'b': 2}] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | -| struct<_1:int> | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [Row(_1=1), Row(_1=2)] | \ No newline at end of file +| SQL Type \ Value@Type | [None, None]@list | [True, False]@list | ['a', 'b']@list | ['12', '34']@list | [Decimal('1'), Decimal('2')]@list | [{'a': 1}, {'b': 2}]@list | [1 2]@ndarray[int8] | [1 2]@ndarray[int16] | [1 2]@ndarray[int32] | [1 2]@ndarray[int64] | [1 2]@ndarray[uint8] | [1 2]@ndarray[uint16] | [1 2]@ndarray[uint32] | [1 2]@ndarray[uint64] | [1. 2.]@ndarray[float16] | [1. 2.]@ndarray[float32] | [1. 2.]@ndarray[float64] | [1.+0.j 2.+0.j]@ndarray[complex64] | [1.+0.j 2.+0.j]@ndarray[complex128] | [array([1, 2, 3], dtype=int32), @list | ['1970-01-01T00:00:00.000000' '1@ndarray[datetime64[us]] | ['1970-01-01T05:00:00.000000' '1@ndarray[datetime64[us]] | [Timedelta('1 days 00:00:00'), T@list | ['A', 'B'] Categories (2, str): @Categorical | {'_1': [1, 2]}@Dataframe[_1 int64] | +|--------------------------|---------------------|----------------------|-------------------|--------------------------------|------------------------------------------|-----------------------------|------------------------------|------------------------------|------------------------------------------|-----------------------------------------|------------------------------|------------------------------|------------------------------|-------------------------|----------------------------|------------------------------|------------------------------|--------------------------------------|---------------------------------------|-----------------------------------------|------------------------------------------------------------|------------------------------------------------------------|-----------------------------------------|------------------------------------------------|--------------------------------------| +| boolean | [None, None] | [True, False] | X | X | X | X | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | [True, True] | X | [True, True] | [True, True] | X | X | X | X | X | X | X | X | +| tinyint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| smallint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| int | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | X | X | X | X | X | +| bigint | [None, None] | [1, 0] | X | [12, 34] | [1, 2] | X | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | [1, 2] | X | X | X | [0, 86400000000] | [18000000000, 104400000000] | [86400000000, 172800000000] | X | X | +| string | [None, None] | X | ['a', 'b'] | ['12', '34'] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | ['A', 'B'] | X | +| date | [None, None] | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | [datetime.date(1970, 1, 2), datetime.dat | X | X | X | X | X | X | X | X | X | X | X | [datetime.date(1970, 1, 1), datetime.dat | [datetime.date(1970, 1, 1), datetime.dat | X | X | X | +| timestamp | [None, None] | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | [datetime.datetime(1969, 12, 31, 16, 0, | X | X | X | X | X | X | X | X | X | X | [datetime.datetime(1970, 1, 1, 0, 0), da | [datetime.datetime(1970, 1, 1, 5, 0), da | X | X | X | +| float | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | +| double | [None, None] | [1.0, 0.0] | X | [12.0, 34.0] | [1.0, 2.0] | X | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | [1.0, 2.0] | X | X | X | X | X | X | X | X | +| array<int> | [None, None] | X | X | [[1, 2], [3, 4]] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [[1, 2, 3], [1, 2, 3]] | X | X | X | X | X | +| binary | [None, None] | [b'\x01', b''] | [b'a', b'b'] | [b'12', b'34'] | X | X | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'\x01', b'\x02'] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | [b'', b''] | X | [b'', b''] | [b'', b''] | [b'', b''] | [b'A', b'B'] | X | +| decimal(10,0) | [None, None] | X | X | [Decimal('12'), Decimal('34')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | [Decimal('1'), Decimal('2')] | [Decimal('1'), Decimal('2')] | X | X | X | X | X | X | X | X | +| map<string,int> | [None, None] | X | X | X | X | [{'a': 1}, {'b': 2}] | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | +| struct<_1:int> | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | X | [Row(_1=1), Row(_1=2)] | \ No newline at end of file diff --git a/python/pyspark/sql/tests/coercion/test_pandas_udf_input_type.py b/python/pyspark/sql/tests/coercion/test_pandas_udf_input_type.py index 0f36e142f3591..79bded3d636d9 100644 --- a/python/pyspark/sql/tests/coercion/test_pandas_udf_input_type.py +++ b/python/pyspark/sql/tests/coercion/test_pandas_udf_input_type.py @@ -15,14 +15,14 @@ # limitations under the License. # -from decimal import Decimal import datetime import os import unittest +from decimal import Decimal +from pyspark.loose_version import LooseVersion from pyspark.sql.functions import pandas_udf from pyspark.sql.types import ( - Row, ArrayType, BinaryType, BooleanType, @@ -34,23 +34,23 @@ IntegerType, LongType, MapType, + Row, ShortType, StringType, StructField, StructType, TimestampType, ) -from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_numpy: import numpy as np diff --git a/python/pyspark/sql/tests/coercion/test_pandas_udf_return_type.py b/python/pyspark/sql/tests/coercion/test_pandas_udf_return_type.py index e26efd2b999ca..a752b4941e071 100644 --- a/python/pyspark/sql/tests/coercion/test_pandas_udf_return_type.py +++ b/python/pyspark/sql/tests/coercion/test_pandas_udf_return_type.py @@ -16,11 +16,12 @@ # import concurrent.futures -from decimal import Decimal import itertools import os import unittest +from decimal import Decimal +from pyspark.loose_version import LooseVersion from pyspark.sql.functions import pandas_udf from pyspark.sql.types import ( ArrayType, @@ -40,17 +41,16 @@ StructType, TimestampType, ) -from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_numpy: import numpy as np diff --git a/python/pyspark/sql/tests/coercion/test_python_udf_input_type.py b/python/pyspark/sql/tests/coercion/test_python_udf_input_type.py index 5f00d97b7d12b..1eac3e45ba69e 100644 --- a/python/pyspark/sql/tests/coercion/test_python_udf_input_type.py +++ b/python/pyspark/sql/tests/coercion/test_python_udf_input_type.py @@ -15,14 +15,14 @@ # limitations under the License. # -from decimal import Decimal import datetime import os import unittest +from decimal import Decimal +from pyspark.loose_version import LooseVersion from pyspark.sql.functions import udf from pyspark.sql.types import ( - Row, ArrayType, BinaryType, BooleanType, @@ -34,23 +34,23 @@ IntegerType, LongType, MapType, + Row, ShortType, StringType, StructField, StructType, TimestampType, ) -from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_numpy: import numpy as np diff --git a/python/pyspark/sql/tests/coercion/test_python_udf_return_type.py b/python/pyspark/sql/tests/coercion/test_python_udf_return_type.py index e5281e9ad3d6a..d88651a8a6fe1 100644 --- a/python/pyspark/sql/tests/coercion/test_python_udf_return_type.py +++ b/python/pyspark/sql/tests/coercion/test_python_udf_return_type.py @@ -18,12 +18,13 @@ import array import concurrent.futures import datetime -from decimal import Decimal import itertools import os import re import unittest +from decimal import Decimal +from pyspark.loose_version import LooseVersion from pyspark.sql import Row from pyspark.sql.functions import udf from pyspark.sql.types import ( @@ -44,17 +45,16 @@ StructType, TimestampType, ) -from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_numpy: import numpy as np diff --git a/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py new file mode 100644 index 0000000000000..5b49582e8391a --- /dev/null +++ b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py @@ -0,0 +1,29 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.sql.tests.arrow.test_arrow_python_aggregator import ArrowPythonAggregatorTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class ArrowPythonAggregatorParityTests(ArrowPythonAggregatorTestsMixin, ReusedConnectTestCase): + pass + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_udf.py b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_udf.py index 0a7a545216869..51955fc94e997 100644 --- a/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_udf.py +++ b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_udf.py @@ -17,8 +17,8 @@ import unittest -from pyspark.sql.tests.connect.test_parity_udf import UDFParityTests from pyspark.sql.tests.arrow.test_arrow_python_udf import ArrowPythonUDFTestsMixin +from pyspark.sql.tests.connect.test_parity_udf import UDFParityTests class ArrowPythonUDFParityTests(UDFParityTests, ArrowPythonUDFTestsMixin): diff --git a/python/pyspark/sql/tests/connect/client/test_artifact.py b/python/pyspark/sql/tests/connect/client/test_artifact.py index 87ea570204e71..6e5ed71745950 100644 --- a/python/pyspark/sql/tests/connect/client/test_artifact.py +++ b/python/pyspark/sql/tests/connect/client/test_artifact.py @@ -15,22 +15,22 @@ # limitations under the License. # import hashlib +import os import shutil import tempfile import unittest -import os import zipfile import zlib -from pyspark.util import is_remote_only from pyspark.sql import SparkSession +from pyspark.sql.functions import assert_true, lit, udf from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect -from pyspark.sql.functions import udf, assert_true, lit +from pyspark.util import is_remote_only if should_test_connect: - from pyspark.sql.connect.client.artifact import ArtifactManager - from pyspark.sql.connect.client import DefaultChannelBuilder from pyspark.errors import SparkRuntimeException + from pyspark.sql.connect.client import DefaultChannelBuilder + from pyspark.sql.connect.client.artifact import ArtifactManager class ArtifactTestsMixin: diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py index b5bf76d86df48..b1463642dfdc7 100644 --- a/python/pyspark/sql/tests/connect/client/test_client.py +++ b/python/pyspark/sql/tests/connect/client/test_client.py @@ -18,30 +18,31 @@ import unittest import uuid from collections.abc import Generator -from typing import Optional, Any, Union +from typing import Any, Optional, Union -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect from pyspark.testing.utils import eventually if should_test_connect: - import grpc import google.protobuf.any_pb2 as any_pb2 import google.protobuf.wrappers_pb2 as wrappers_pb2 - from google.rpc import status_pb2 - from google.rpc.error_details_pb2 import ErrorInfo + import grpc import pandas as pd import pyarrow as pa - from pyspark.sql.connect.client import SparkConnectClient, DefaultChannelBuilder + from google.rpc import status_pb2 + from google.rpc.error_details_pb2 import ErrorInfo + + import pyspark.sql.connect.proto as proto + from pyspark.errors import PySparkRuntimeError + from pyspark.errors.exceptions.connect import SparkConnectGrpcException + from pyspark.sql.connect.client import DefaultChannelBuilder, SparkConnectClient from pyspark.sql.connect.client.core import RpcDeadlines + from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator from pyspark.sql.connect.client.retries import ( - Retrying, DefaultPolicy, + Retrying, ) - from pyspark.sql.connect.client.reattach import ExecutePlanResponseReattachableIterator from pyspark.sql.connect.session import SparkSession as RemoteSparkSession - from pyspark.errors import PySparkRuntimeError - from pyspark.errors.exceptions.connect import SparkConnectGrpcException - import pyspark.sql.connect.proto as proto class TestPolicy(DefaultPolicy): def __init__(self): @@ -113,16 +114,22 @@ def __init__(self, execute_ops=None, attach_ops=None): self.release_calls = 0 self.release_until_calls = 0 self.attach_calls = 0 + self.execute_metadata = [] + self.attach_metadata = [] + self.release_metadata = [] def ExecutePlan(self, *args, **kwargs): self.execute_calls += 1 + self.execute_metadata.append(kwargs.get("metadata")) return self._execute_ops def ReattachExecute(self, *args, **kwargs): self.attach_calls += 1 + self.attach_metadata.append(kwargs.get("metadata")) return self._attach_ops def ReleaseExecute(self, req: proto.ReleaseExecuteRequest, *args, **kwargs): + self.release_metadata.append(kwargs.get("metadata")) if req.HasField("release_all"): self.release_calls += 1 elif req.HasField("release_until"): @@ -149,6 +156,7 @@ class MockService: def __init__(self, session_id: str, operation_statuses=None): self._session_id = session_id self.req = None + self.metadata = None self.client_user_context_extensions = [] if operation_statuses is None: operation_statuses = self.DEFAULT_OPERATION_STATUSES @@ -156,6 +164,7 @@ def __init__(self, session_id: str, operation_statuses=None): def ExecutePlan(self, req: proto.ExecutePlanRequest, metadata, timeout=None): self.req = req + self.metadata = metadata self.client_user_context_extensions = list(req.user_context.extensions) resp = proto.ExecutePlanResponse() resp.session_id = self._session_id @@ -297,6 +306,30 @@ def userId(self) -> Optional[str]: self.assertEqual(client._user_id, "abc") + def test_channel_builder_metadata_is_filtered_per_call(self): + class CustomChannelBuilder(DefaultChannelBuilder): + def __init__(self): + super().__init__("sc://foo/") + self.metadata_calls = 0 + + def metadata(self): + self.metadata_calls += 1 + return iter( + [ + ("authorization", f"token-{self.metadata_calls}"), + ("spark-connect-operation-id", "ignored"), + ] + ) + + builder = CustomChannelBuilder() + client = SparkConnectClient(builder, use_reattachable_execute=False) + try: + self.assertEqual(client._artifact_manager._metadata, [("authorization", "token-1")]) + self.assertEqual(client._builder_metadata(), [("authorization", "token-2")]) + self.assertEqual(client._builder_metadata(), [("authorization", "token-3")]) + finally: + client.close() + def test_user_context_extension(self): client = SparkConnectClient("sc://foo/", use_reattachable_execute=False) mock = MockService(client._session_id) @@ -479,6 +512,43 @@ def on_execute_plan(self, req): session.client.close() session.stop() + def test_session_hook_preserves_operation_id(self): + execute_plan_req = None + + class TestHook(RemoteSparkSession.Hook): + def __init__(self, _session): + pass + + def on_execute_plan(self, req): + replacement = proto.ExecutePlanRequest() + replacement.CopyFrom(req) + replacement.ClearField("operation_id") + return replacement + + class TestService(MockService): + def ExecutePlan(self, req, metadata, timeout=None): + nonlocal execute_plan_req + execute_plan_req = req + return super().ExecutePlan(req, metadata, timeout) + + session = ( + RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate() + ) + try: + mock = TestService(session.client._session_id) + session.client._stub = mock + session.client.disable_reattachable_execute() + + df = session.range(1) + df.collect() + self.assertIsNotNone(df.executionInfo) + self.assertIsNotNone(execute_plan_req) + assert execute_plan_req is not None + self.assertEqual(execute_plan_req.operation_id, df.executionInfo.operation_id) + uuid.UUID(execute_plan_req.operation_id) + finally: + session.stop() + def test_new_session_preserves_custom_channel_builder(self): class CustomChannelBuilder(DefaultChannelBuilder): pass @@ -509,6 +579,29 @@ def test_custom_operation_id(self): for resp in client._stub.ExecutePlan(req, metadata=None): assert resp.operation_id == "10a4c38e-7e87-40ee-9d6f-60ff0751e63b" + def test_execute_plan_request_generates_operation_id(self): + client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False) + try: + req = client._execute_plan_request_with_metadata() + uuid.UUID(req.operation_id) + finally: + client.close() + + def test_execute_plan_sends_operation_id_metadata(self): + client = SparkConnectClient( + "sc://foo/;spark-connect-operation-id=ignored", use_reattachable_execute=False + ) + mock = MockService(client._session_id) + client._stub = mock + try: + req = client._execute_plan_request_with_metadata() + client._execute(req) + self.assertIsNotNone(mock.metadata) + values = [v for k, v in mock.metadata if k == "spark-connect-operation-id"] + self.assertEqual(values, [req.operation_id]) + finally: + client.close() + def test_on_exit_calls_release_and_close_when_enabled(self): client = SparkConnectClient("sc://foo/", use_reattachable_execute=False) client._release_session_on_exit = True @@ -875,6 +968,49 @@ def GetStatus(self, req, metadata, timeout=None): client._get_operation_statuses() self.assertEqual(mock.captured_timeouts["GetStatus"], 55.0) + def test_remove_cached_relation_uses_release_relation_deadline(self): + """CachedRemoteRelation.__del__ must bound its release RPC with the release_relation + deadline. + + Regression test: the RemoveRemoteCachedRelation cleanup is a blocking, non-reattachable + ExecutePlan call issued from a finalizer with no other timeout at any layer. Without a + client-side deadline it can hang forever if the response is never delivered, which on the + foreachBatch Connect path stalls the streaming query indefinitely (the Python worker never + sends its completion signal and the driver JVM blocks on the per-batch read). + """ + from types import SimpleNamespace + from unittest.mock import MagicMock + + from pyspark.sql.connect.plan import CachedRemoteRelation + + def make_relation(deadlines): + client = SparkConnectClient( + "sc://foo/", + use_reattachable_execute=False, + rpc_deadlines=deadlines, + retry_policy=dict(max_retries=0), + ) + captured = {} + + def fake_call(req, metadata=None, timeout="unset"): + captured["timeout"] = timeout + return proto.ExecutePlanResponse() + + client._channel = MagicMock() + client._channel.unary_unary.return_value = fake_call + rel = CachedRemoteRelation("df-id-123", spark_session=SimpleNamespace(client=client)) + return rel, captured + + # A configured deadline is forwarded as the gRPC call timeout. + rel, captured = make_relation(RpcDeadlines(release_relation=44.0)) + rel.__del__() + self.assertEqual(captured["timeout"], 44.0) + + # Disabled deadlines forward None (opt out; rely on transport/server-side timeouts). + rel, captured = make_relation(RpcDeadlines.disabled()) + rel.__del__() + self.assertIsNone(captured["timeout"]) + @unittest.skipIf(not should_test_connect, connect_requirement_message) class SparkConnectClientReattachTestCase(unittest.TestCase): @@ -909,6 +1045,21 @@ def check_all(): eventually(timeout=1, catch_assertions=True)(check_all)() + def test_operation_id_metadata_is_sent_on_all_rpcs(self): + metadata = [("spark-connect-operation-id", "operation-id")] + stub = self._stub_with([self.response], [self.finished]) + ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.retrying, metadata) + for _ in ite: + pass + + def check_all(): + self.assertEqual(stub.execute_metadata, [metadata]) + self.assertEqual(stub.attach_metadata, [metadata]) + self.assertTrue(stub.release_metadata) + self.assertTrue(all(value == metadata for value in stub.release_metadata)) + + eventually(timeout=1, catch_assertions=True)(check_all)() + def test_fail_during_execute(self): def fatal(): raise TestException("Fatal") diff --git a/python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py b/python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py index 100ea58a88fe6..f17b993c2d8e2 100644 --- a/python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py +++ b/python/pyspark/sql/tests/connect/client/test_client_call_stack_trace.py @@ -20,7 +20,7 @@ from unittest.mock import patch import pyspark -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect if should_test_connect: import pyspark.sql.connect.proto as pb2 diff --git a/python/pyspark/sql/tests/connect/client/test_client_retries.py b/python/pyspark/sql/tests/connect/client/test_client_retries.py index ebfac74114068..f6d352ad27280 100644 --- a/python/pyspark/sql/tests/connect/client/test_client_retries.py +++ b/python/pyspark/sql/tests/connect/client/test_client_retries.py @@ -18,25 +18,25 @@ import unittest import warnings -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect if should_test_connect: - import grpc import google.protobuf.any_pb2 as any_pb2 import google.protobuf.duration_pb2 as duration_pb2 - from google.rpc import status_pb2 - from google.rpc import error_details_pb2 + import grpc + from google.rpc import error_details_pb2, status_pb2 + from pyspark.sql.connect.client import SparkConnectClient from pyspark.sql.connect.client.core import RpcDeadlines from pyspark.sql.connect.client.retries import ( - Retrying, + DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME, DefaultPolicy, RetryException, - DEFAULT_MAX_RETRY_EXCEPTION_ELAPSED_TIME, + Retrying, ) from pyspark.sql.tests.connect.client.test_client import ( - TestPolicy, TestException, + TestPolicy, ) class SleepTimeTracker: diff --git a/python/pyspark/sql/tests/connect/client/test_reattach.py b/python/pyspark/sql/tests/connect/client/test_reattach.py index 7e08c4697d903..48a5fa3e3a8e0 100644 --- a/python/pyspark/sql/tests/connect/client/test_reattach.py +++ b/python/pyspark/sql/tests/connect/client/test_reattach.py @@ -17,11 +17,11 @@ import unittest -from pyspark.util import is_remote_only from pyspark.sql import SparkSession as PySparkSession from pyspark.testing.connectutils import ReusedMixedTestCase from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.utils import eventually +from pyspark.util import is_remote_only @unittest.skipIf(is_remote_only(), "Requires JVM access") diff --git a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state.py b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state.py index 878268c5006bd..7276ccc60fa05 100644 --- a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state.py +++ b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state.py @@ -16,10 +16,10 @@ # import unittest +from pyspark import SparkConf from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state import ( TransformWithStateInPandasTestsMixin, ) -from pyspark import SparkConf from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state_state_variable.py b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state_state_variable.py index 4ac1ed9ac37ff..c9ff06a78179e 100644 --- a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state_state_variable.py +++ b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_pandas_transform_with_state_state_variable.py @@ -15,10 +15,10 @@ # limitations under the License. # +from pyspark import SparkConf from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state_state_variable import ( TransformWithStateInPandasStateVariableTestsMixin, ) -from pyspark import SparkConf from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state.py b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state.py index 063376ad41282..4fd0c9a800e8f 100644 --- a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state.py +++ b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state.py @@ -16,10 +16,10 @@ # import unittest +from pyspark import SparkConf from pyspark.sql.tests.pandas.streaming.test_transform_with_state import ( TransformWithStateInPySparkTestsMixin, ) -from pyspark import SparkConf from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state_state_variable.py b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state_state_variable.py index d45064614af44..a50fb12fa94b5 100644 --- a/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state_state_variable.py +++ b/python/pyspark/sql/tests/connect/pandas/streaming/test_parity_transform_with_state_state_variable.py @@ -15,10 +15,10 @@ # limitations under the License. # +from pyspark import SparkConf from pyspark.sql.tests.pandas.streaming.test_transform_with_state_state_variable import ( TransformWithStateInPySparkStateVariableTestsMixin, ) -from pyspark import SparkConf from pyspark.testing.connectutils import ReusedConnectTestCase diff --git a/python/pyspark/sql/tests/connect/shell/test_progress.py b/python/pyspark/sql/tests/connect/shell/test_progress.py index f567320bda087..a810c64cc3508 100644 --- a/python/pyspark/sql/tests/connect/shell/test_progress.py +++ b/python/pyspark/sql/tests/connect/shell/test_progress.py @@ -15,14 +15,14 @@ # limitations under the License. # -from io import StringIO import unittest +from io import StringIO from typing import Iterable from pyspark.testing.connectutils import ( - should_test_connect, - connect_requirement_message, ReusedConnectTestCase, + connect_requirement_message, + should_test_connect, ) from pyspark.testing.utils import PySparkErrorTestUtils diff --git a/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py b/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py index aa8c504966088..b6dd23353eb91 100644 --- a/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py +++ b/python/pyspark/sql/tests/connect/streaming/test_parity_foreach_batch.py @@ -16,10 +16,11 @@ # import time + +from pyspark.errors import PySparkPicklingError from pyspark.sql.tests.streaming.test_streaming_foreach_batch import StreamingTestsForeachBatchMixin from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect from pyspark.testing.utils import eventually, timeout -from pyspark.errors import PySparkPicklingError if should_test_connect: from pyspark.errors.exceptions.connect import StreamingPythonRunnerInitializationException diff --git a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py index a8b46c38db4fe..81a116eff2285 100644 --- a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py +++ b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py @@ -19,9 +19,9 @@ import pyspark.cloudpickle from pyspark.errors import AnalysisException -from pyspark.sql.tests.streaming.test_streaming_listener import StreamingListenerTestsMixin -from pyspark.sql.streaming.listener import StreamingQueryListener from pyspark.sql.functions import count, lit +from pyspark.sql.streaming.listener import StreamingQueryListener +from pyspark.sql.tests.streaming.test_streaming_listener import StreamingListenerTestsMixin from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.testing.utils import eventually diff --git a/python/pyspark/sql/tests/connect/test_connect_basic.py b/python/pyspark/sql/tests/connect/test_connect_basic.py index b8d258267d267..6a28f101b9b3d 100755 --- a/python/pyspark/sql/tests/connect/test_connect_basic.py +++ b/python/pyspark/sql/tests/connect/test_connect_basic.py @@ -15,44 +15,47 @@ # limitations under the License. # -import os +import datetime import gc -import unittest +import io +import os import shutil import tempfile -import io +import unittest from contextlib import redirect_stdout -import datetime -from pyspark.util import is_remote_only from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.types import ( - StructType, - StructField, - LongType, - StringType, + ArrayType, + CharType, IntegerType, + LongType, MapType, - ArrayType, Row, + StringType, + StructField, + StructType, + VarcharType, ) from pyspark.testing import assertDataFrameEqual -from pyspark.testing.utils import eventually from pyspark.testing.connectutils import ( - should_test_connect, - connect_requirement_message, ReusedMixedTestCase, + connect_requirement_message, + should_test_connect, ) from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.testing.utils import eventually +from pyspark.util import is_remote_only if should_test_connect: - from pyspark.sql.connect.proto import ExecutePlanResponse, Expression as ProtoExpression - from pyspark.sql.connect.column import Column - from pyspark.sql.dataframe import DataFrame - from pyspark.sql.connect.dataframe import DataFrame as CDataFrame + from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException from pyspark.sql import functions as SF from pyspark.sql.connect import functions as CF - from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException + from pyspark.sql.connect.column import Column + from pyspark.sql.connect.dataframe import DataFrame as CDataFrame + from pyspark.sql.connect.proto import ExecutePlanResponse + from pyspark.sql.connect.proto import Expression as ProtoExpression + from pyspark.sql.dataframe import DataFrame @unittest.skipIf( @@ -154,8 +157,8 @@ def test_serialization_II(self): self.assertEqual(cdf.collect(), cdf2.collect()) def test_window_spec_serialization(self): - from pyspark.sql.connect.window import Window from pyspark.serializers import CPickleSerializer + from pyspark.sql.connect.window import Window pickle_ser = CPickleSerializer() w = Window.partitionBy("some_string").orderBy("value") @@ -487,6 +490,19 @@ def test_schema(self): ) self._check_print_schema(query) + def test_char_varchar_result_schema(self): + # SPARK-58794: Python Connect maps first-class CHAR/VARCHAR the same as classic. + query = "SELECT CAST('ab' AS CHAR(4)) AS c, CAST('cd' AS VARCHAR(6)) AS v" + conf = {"spark.sql.charVarchar.standardSemantics.enabled": "true"} + with self.both_conf(conf): + classic_df = self.spark.sql(query) + connect_df = self.connect.sql(query) + self.assertEqual(classic_df.schema, connect_df.schema) + self.assertEqual(classic_df.schema["c"].dataType, CharType(4)) + self.assertEqual(classic_df.schema["v"].dataType, VarcharType(6)) + self.assertEqual(classic_df.collect(), connect_df.collect()) + self.assertEqual(connect_df.collect(), [Row(c="ab ", v="cd")]) + def test_to(self): # SPARK-41464: test DataFrame.to() diff --git a/python/pyspark/sql/tests/connect/test_connect_channel.py b/python/pyspark/sql/tests/connect/test_connect_channel.py index 66a8f0d8abc87..19c4866654f4f 100644 --- a/python/pyspark/sql/tests/connect/test_connect_channel.py +++ b/python/pyspark/sql/tests/connect/test_connect_channel.py @@ -20,15 +20,16 @@ from pyspark.errors import PySparkValueError from pyspark.testing.connectutils import ( - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: import grpc - from pyspark.sql.connect.client import DefaultChannelBuilder, ChannelBuilder - from pyspark.sql.connect.client.core import SparkConnectClient + from pyspark.errors.exceptions.connect import SparkConnectException + from pyspark.sql.connect.client import ChannelBuilder, DefaultChannelBuilder + from pyspark.sql.connect.client.core import SparkConnectClient @unittest.skipIf(not should_test_connect, connect_requirement_message) diff --git a/python/pyspark/sql/tests/connect/test_connect_collection.py b/python/pyspark/sql/tests/connect/test_connect_collection.py index 555058513efaf..805300f881ae5 100644 --- a/python/pyspark/sql/tests/connect/test_connect_collection.py +++ b/python/pyspark/sql/tests/connect/test_connect_collection.py @@ -16,7 +16,7 @@ # from pyspark.testing import assertDataFrameEqual -from pyspark.testing.connectutils import should_test_connect, ReusedMixedTestCase +from pyspark.testing.connectutils import ReusedMixedTestCase, should_test_connect from pyspark.testing.pandasutils import PandasOnSparkTestUtils if should_test_connect: diff --git a/python/pyspark/sql/tests/connect/test_connect_column.py b/python/pyspark/sql/tests/connect/test_connect_column.py index 6fa6c4686c527..a842f6c57fd54 100644 --- a/python/pyspark/sql/tests/connect/test_connect_column.py +++ b/python/pyspark/sql/tests/connect/test_connect_column.py @@ -15,53 +15,54 @@ # limitations under the License. # -import decimal import datetime +import decimal +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.types import ( - Row, - StructField, - StructType, - MapType, - NullType, - DateType, - TimeType, - TimestampType, - TimestampNTZType, - ByteType, BinaryType, - ShortType, - IntegerType, - FloatType, + BooleanType, + ByteType, + DateType, DayTimeIntervalType, - StringType, + DecimalType, DoubleType, + FloatType, + IntegerType, LongType, - DecimalType, - BooleanType, + MapType, + NullType, + Row, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, + TimeType, ) -from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.testing import assertDataFrameEqual -from pyspark.testing.connectutils import should_test_connect, ReusedMixedTestCase +from pyspark.testing.connectutils import ReusedMixedTestCase, should_test_connect from pyspark.testing.pandasutils import PandasOnSparkTestUtils if should_test_connect: import pandas as pd + + from pyspark.errors.exceptions.connect import SparkConnectException from pyspark.sql import functions as SF from pyspark.sql.connect import functions as CF from pyspark.sql.connect.column import Column from pyspark.sql.connect.expressions import DistributedSequenceID, LiteralExpression from pyspark.util import ( - JVM_BYTE_MIN, JVM_BYTE_MAX, - JVM_SHORT_MIN, - JVM_SHORT_MAX, - JVM_INT_MIN, + JVM_BYTE_MIN, JVM_INT_MAX, - JVM_LONG_MIN, + JVM_INT_MIN, JVM_LONG_MAX, + JVM_LONG_MIN, + JVM_SHORT_MAX, + JVM_SHORT_MIN, ) - from pyspark.errors.exceptions.connect import SparkConnectException class SparkConnectColumnTests(ReusedMixedTestCase, PandasOnSparkTestUtils): diff --git a/python/pyspark/sql/tests/connect/test_connect_creation.py b/python/pyspark/sql/tests/connect/test_connect_creation.py index 83848319e9b91..e6d0172a3dad7 100644 --- a/python/pyspark/sql/tests/connect/test_connect_creation.py +++ b/python/pyspark/sql/tests/connect/test_connect_creation.py @@ -22,25 +22,26 @@ from pyspark.errors import PySparkValueError from pyspark.sql.types import ( - StructType, - StructField, - StringType, + ArrayType, IntegerType, LongType, MapType, - ArrayType, Row, + StringType, + StructField, + StructType, ) +from pyspark.testing.connectutils import ReusedMixedTestCase, should_test_connect from pyspark.testing.objects import MyObject, PythonOnlyUDT -from pyspark.testing.connectutils import should_test_connect, ReusedMixedTestCase from pyspark.testing.pandasutils import PandasOnSparkTestUtils if should_test_connect: - import pandas as pd import numpy as np + import pandas as pd + + from pyspark.errors.exceptions.connect import ParseException from pyspark.sql import functions as SF from pyspark.sql.connect import functions as CF - from pyspark.errors.exceptions.connect import ParseException class SparkConnectCreationTests(ReusedMixedTestCase, PandasOnSparkTestUtils): @@ -549,9 +550,10 @@ def test_create_df_nullability(self): def test_create_dataframe_from_pandas_with_ns_timestamp(self): """Truncate the timestamps for nanoseconds.""" - from datetime import datetime, timezone, timedelta - from pandas import Timestamp + from datetime import datetime, timedelta, timezone + import pandas as pd + from pandas import Timestamp # Nanoseconds are truncated to microseconds in the serializer # Arrow will throw an error if precision is lost diff --git a/python/pyspark/sql/tests/connect/test_connect_dataframe_property.py b/python/pyspark/sql/tests/connect/test_connect_dataframe_property.py index 44f56828685c0..0655b7ae18cda 100644 --- a/python/pyspark/sql/tests/connect/test_connect_dataframe_property.py +++ b/python/pyspark/sql/tests/connect/test_connect_dataframe_property.py @@ -18,19 +18,19 @@ import unittest from pyspark.loose_version import LooseVersion +from pyspark.sql import functions as SF from pyspark.sql.types import ( - StructType, - StructField, - StringType, + DoubleType, IntegerType, LongType, - DoubleType, Row, + StringType, + StructField, + StructType, ) from pyspark.sql.utils import is_remote -from pyspark.sql import functions as SF from pyspark.testing import assertDataFrameEqual -from pyspark.testing.connectutils import should_test_connect, ReusedMixedTestCase +from pyspark.testing.connectutils import ReusedMixedTestCase, should_test_connect from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.utils import ( have_pandas, diff --git a/python/pyspark/sql/tests/connect/test_connect_error.py b/python/pyspark/sql/tests/connect/test_connect_error.py index 8f142dd556325..4c1a3bc3f24c1 100644 --- a/python/pyspark/sql/tests/connect/test_connect_error.py +++ b/python/pyspark/sql/tests/connect/test_connect_error.py @@ -17,11 +17,10 @@ import unittest -from pyspark.errors import PySparkAttributeError +from pyspark.errors import PySparkAttributeError, PySparkNotImplementedError, PySparkTypeError from pyspark.errors.exceptions.base import SessionNotSameException -from pyspark.sql.types import Row from pyspark.sql import functions as F -from pyspark.errors import PySparkNotImplementedError, PySparkTypeError +from pyspark.sql.types import Row from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.util import is_remote_only diff --git a/python/pyspark/sql/tests/connect/test_connect_function.py b/python/pyspark/sql/tests/connect/test_connect_function.py index 5c19cb895c7c0..ad9da9272fcfc 100644 --- a/python/pyspark/sql/tests/connect/test_connect_function.py +++ b/python/pyspark/sql/tests/connect/test_connect_function.py @@ -17,28 +17,28 @@ import unittest from inspect import getmembers, isfunction -from pyspark.util import is_remote_only from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.types import ( - _drop_metadata, - StringType, - StructType, - StructField, ArrayType, IntegerType, MapType, + StringType, + StructField, + StructType, + _drop_metadata, ) from pyspark.testing import assertDataFrameEqual -from pyspark.testing.pandasutils import PandasOnSparkTestUtils from pyspark.testing.connectutils import ReusedMixedTestCase, should_test_connect +from pyspark.testing.pandasutils import PandasOnSparkTestUtils +from pyspark.util import is_remote_only if should_test_connect: - from pyspark.sql.connect.column import Column + from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException from pyspark.sql import functions as SF - from pyspark.sql.window import Window as SW from pyspark.sql.connect import functions as CF + from pyspark.sql.connect.column import Column from pyspark.sql.connect.window import Window as CW - from pyspark.errors.exceptions.connect import AnalysisException, SparkConnectException + from pyspark.sql.window import Window as SW @unittest.skipIf(is_remote_only(), "Requires JVM access") @@ -615,6 +615,13 @@ def test_aggregation_functions(self): check_exact=False, ) + # collect_union takes an array-typed column; build one via array(b, c). + self.assert_eq( + cdf.groupBy("a").agg(CF.sort_array(CF.collect_union(CF.array("b", "c")))).toPandas(), + sdf.groupBy("a").agg(SF.sort_array(SF.collect_union(SF.array("b", "c")))).toPandas(), + check_exact=False, + ) + for cfunc, sfunc in [ (CF.corr, SF.corr), (CF.covar_pop, SF.covar_pop), diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server.py b/python/pyspark/sql/tests/connect/test_connect_local_server.py new file mode 100644 index 0000000000000..9133d19b114cc --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server.py @@ -0,0 +1,642 @@ +# +# 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. +# + +import contextlib +import getpass +import json +import os +import shutil +import signal +import socket +import subprocess +import sys +import tempfile +import textwrap +import time +import unittest + +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect +from pyspark.util import is_remote_only + +if should_test_connect: + from pyspark.sql import SparkSession as PySparkSession + from pyspark.sql.connect import local_server + from pyspark.sql.connect.local_server import Discovery, LocalConnectServer + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession + from pyspark.version import __version__ + + +@contextlib.contextmanager +def _listening_socket(): + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + yield listener.getsockname()[1] + finally: + listener.close() + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires JVM access to start a local Connect server", +) +class LocalConnectServerReuseTests(unittest.TestCase): + """Tests for the opt-in persistent local Spark Connect server (SPARK_LOCAL_CONNECT_REUSE).""" + + def setUp(self) -> None: + # Point discovery at a throwaway path so the real per-user file is never touched. + self._tmpdir = tempfile.mkdtemp() + self._discovery_path = os.path.join(self._tmpdir, "connect-local.json") + self._saved_env = { + k: os.environ.get(k) + for k in ("SPARK_LOCAL_CONNECT_DISCOVERY", "SPARK_CONNECT_AUTHENTICATE_TOKEN") + } + os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery_path + + def tearDown(self) -> None: + try: + # Only stop a real, separately-spawned server. Several tests fabricate discovery + # files pointing at this very process, which must never be signalled. + server = self._discovered_server() + if server.pid is not None and server.pid != os.getpid(): + local_server.stop_local_connect_server() + # Wait for the JVM to release the port so the next test starts clean. + self._wait_port_closed(server.host, server.port) + finally: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _server(self, **overrides) -> "LocalConnectServer": + from unittest import mock + + fields = { + "host": "localhost", + "port": 0, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + } + fields.update(overrides) + discovery = mock.Mock() + discovery.load.return_value = fields + return LocalConnectServer(discovery) + + def _discovered_server(self) -> "LocalConnectServer": + with Discovery() as discovery: + return LocalConnectServer(discovery) + + def _launcher_discovery(self): + # A stand-in Discovery for ServerLauncher unit tests: only its directory is read + # (for the log dir and the seed properties file), so point it at the temp dir. + from unittest import mock + + discovery = mock.Mock() + discovery.directory = self._tmpdir + return discovery + + @contextlib.contextmanager + def _without_spark_testing(self): + # _pick_port's ephemeral branch is a no-op when SPARK_TESTING is set (as it is under + # the test runner), so drop it to exercise the production behavior. + saved = os.environ.pop("SPARK_TESTING", None) + try: + yield + finally: + if saved is not None: + os.environ["SPARK_TESTING"] = saved + + def test_discovery_location(self) -> None: + self.assertEqual(Discovery().path, self._discovery_path) + # Without the override the file lives in a per-user 0700 dir under the temp dir. + os.environ.pop("SPARK_LOCAL_CONNECT_DISCOVERY") + default = Discovery() + self.assertTrue(default.directory.startswith(tempfile.gettempdir())) + if os.name == "posix": + self.assertIn("spark-connect-{}".format(getpass.getuser()), default.directory) + self.assertEqual(os.stat(default.directory).st_mode & 0o777, 0o700) + + def test_startup_seed_conf(self) -> None: + from unittest import mock + + initial = { + "spark.sql.shuffle.partitions": "8", + "spark.master": "local[1]", + } + opts = { + "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "warehouse"), + "spark.local.connect.reuse": "true", + "spark.connect.grpc.binding.port": "0", + } + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps(initial), + } + with mock.patch.dict(os.environ, env): + self.assertEqual( + local_server.startup_seed_conf(opts), + { + "spark.sql.shuffle.partitions": "8", + "spark.sql.warehouse.dir": opts["spark.sql.warehouse.dir"], + }, + ) + + def test_start_delegates_launch_options(self) -> None: + from unittest import mock + + discovery = mock.Mock() + discovery.load.side_effect = [ + None, + { + "host": "localhost", + "port": 15002, + "token": "t", + "pid": os.getpid(), + "spark_version": __version__, + }, + ] + server = LocalConnectServer(discovery) + seed_conf = {"spark.sql.shuffle.partitions": "4"} + with mock.patch.object(local_server, "ServerLauncher") as launcher: + server.start( + "local[2]", + {"spark.local.connect.reuse": "true"}, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + + launcher.assert_called_once_with( + "local[2]", + {"spark.local.connect.reuse": "true"}, + discovery, + use_ephemeral_port=True, + seed_conf=seed_conf, + ) + launcher.return_value.launch.assert_called_once_with() + self.assertEqual(server.port, 15002) + + def test_pick_port_uses_ephemeral_port_when_requested(self) -> None: + # This is the production path for pool attendants, which run without SPARK_TESTING. + # A non-integer configured port would raise int() in the configured/default branch; + # the ephemeral branch never reads it, so returning a clean OS-assigned port proves + # the free-port path was taken even with SPARK_TESTING unset. + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.local.connect.server.port": "not-a-port"}, + self._launcher_discovery(), + use_ephemeral_port=True, + ) + with self._without_spark_testing(): + port = launcher._pick_port() + self.assertGreater(port, 0) + + def test_pick_port_honors_configured_port_without_testing(self) -> None: + # With neither the ephemeral flag nor SPARK_TESTING, a free configured port is used + # as-is rather than replaced by an OS-assigned one. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + free = sock.getsockname()[1] + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.local.connect.server.port": str(free)}, + self._launcher_discovery(), + use_ephemeral_port=False, + ) + with self._without_spark_testing(): + self.assertEqual(launcher._pick_port(), free) + + def test_seed_conf_override_is_used_and_sanitized(self) -> None: + # The override is taken verbatim except for launcher-managed keys, which are stripped + # so raw builder opts passed as a seed cannot land in --properties-file. + launcher = local_server.ServerLauncher( + "local[2]", + {}, + self._launcher_discovery(), + seed_conf={ + "spark.sql.shuffle.partitions": "4", + "spark.master": "local[9]", + "spark.local.connect.reuse": "true", + }, + ) + self.assertEqual(launcher._seed_conf(), {"spark.sql.shuffle.partitions": "4"}) + + def test_seed_conf_empty_override_does_not_fall_through_to_env(self) -> None: + from unittest import mock + + # Load-bearing for the pool attendant: an empty override means "seed nothing", and + # must not silently pick up PYSPARK_REMOTE_INIT_CONF_* the way opts=None would. + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps({"spark.sql.shuffle.partitions": "8"}), + } + launcher = local_server.ServerLauncher( + "local[2]", {}, self._launcher_discovery(), seed_conf={} + ) + with mock.patch.dict(os.environ, env): + self.assertEqual(launcher._seed_conf(), {}) + + def test_seed_conf_none_override_merges_env_and_opts(self) -> None: + from unittest import mock + + # No override: the env-plus-opts merge (minus launcher-managed keys) is used. + env = { + "PYSPARK_REMOTE_INIT_CONF_LEN": "1", + "PYSPARK_REMOTE_INIT_CONF_0": json.dumps({"spark.sql.shuffle.partitions": "8"}), + } + launcher = local_server.ServerLauncher( + "local[2]", + {"spark.sql.warehouse.dir": os.path.join(self._tmpdir, "wh")}, + self._launcher_discovery(), + seed_conf=None, + ) + with mock.patch.dict(os.environ, env): + self.assertEqual( + launcher._seed_conf(), + { + "spark.sql.shuffle.partitions": "8", + "spark.sql.warehouse.dir": os.path.join(self._tmpdir, "wh"), + }, + ) + + def test_seed_properties_file_reflects_seed_conf(self) -> None: + # An empty seed yields no properties file, so start-connect-server.sh gets no + # --properties-file; a non-empty seed writes a 0600 file with the seeded confs. + launcher = local_server.ServerLauncher( + "local[2]", {}, self._launcher_discovery(), seed_conf={} + ) + with launcher._seed_properties_file() as path: + self.assertIsNone(path) + + launcher = local_server.ServerLauncher( + "local[2]", + {}, + self._launcher_discovery(), + seed_conf={"spark.sql.shuffle.partitions": "4"}, + ) + with launcher._seed_properties_file() as path: + self.assertIsNotNone(path) + self.assertEqual(os.stat(path).st_mode & 0o777, 0o600) + with open(path) as f: + contents = f.read() + self.assertIn("spark.sql.shuffle.partitions=4", contents) + + def test_discovery_roundtrip(self) -> None: + with Discovery() as discovery: + saved = self._server(port=15002) + discovery.save( + {k: getattr(saved, k) for k in ("host", "port", "token", "pid", "spark_version")} + ) + # The file holds the auth token and must not be readable by other users. + self.assertEqual(os.stat(discovery.path).st_mode & 0o777, 0o600) + loaded = LocalConnectServer(discovery) + for attr in ("host", "port", "token", "pid", "spark_version", "url"): + self.assertEqual(getattr(loaded, attr), getattr(saved, attr), attr) + discovery.clear() + self.assertIsNone(discovery.load()) + discovery.clear() # clearing again is a no-op + + def test_discovery_load_rejects_malformed_files(self) -> None: + malformed = [ + "not json", + json.dumps(["a", "list"]), + json.dumps({"host": "localhost"}), # missing required keys + json.dumps( + { + "host": "localhost", + "port": 1, + "token": "t", + "pid": "not-a-pid", + "spark_version": __version__, + } + ), + json.dumps( + {"host": None, "port": 1, "token": "t", "pid": 1, "spark_version": __version__} + ), + ] + with Discovery() as discovery: + for content in malformed: + with self.subTest(content=content): + with open(discovery.path, "w") as f: + f.write(content) + self.assertIsNone(discovery.load()) + + def test_server_is_reusable(self) -> None: + with _listening_socket() as port: + with self.subTest("alive process listening on the port with a matching version"): + self.assertTrue(self._server(port=port).is_reusable()) + with self.subTest("version mismatch"): + self.assertFalse( + self._server(port=port, spark_version="0.0.0-not-this-build").is_reusable() + ) + if os.name == "posix": # the pid probe only runs on POSIX (see the test below) + with self.subTest("dead pid"): + # PID 2**31 - 1 is effectively guaranteed not to exist. + self.assertFalse(self._server(port=port, pid=2**31 - 1).is_reusable()) + server = self._server(port=port) + with self.subTest("port no longer listening"): + self.assertFalse(server.is_reusable()) + + def test_pid_probe_is_skipped_on_windows(self) -> None: + # On Windows os.kill(pid, 0) terminates the target instead of probing it, so the + # reuse check would kill the very server it is examining. + from unittest import mock + + with _listening_socket() as port: + server = self._server(port=port) + with mock.patch.object(os, "name", "nt"), mock.patch.object(os, "kill") as kill: + self.assertTrue(server.is_reusable()) + kill.assert_not_called() + + def test_stop_when_no_server_is_safe(self) -> None: + self.assertFalse(local_server.stop_local_connect_server()) + + def test_stop_signals_recorded_server_and_clears_discovery(self) -> None: + from unittest import mock + + with Discovery() as discovery: + server = self._server(pid=12345) + discovery.save( + {k: getattr(server, k) for k in ("host", "port", "token", "pid", "spark_version")} + ) + # Avoid inspecting or signaling a real process while exercising the stop path. + ps_result = subprocess.CompletedProcess([], 0, stdout=local_server._SERVER_CLASS) + with ( + mock.patch.object(subprocess, "run", return_value=ps_result) as run, + mock.patch.object(os, "kill") as kill, + ): + self.assertTrue(local_server.stop_local_connect_server()) + run.assert_called_once_with( + ["ps", "-ww", "-p", "12345", "-o", "command="], + capture_output=True, + text=True, + timeout=5, + ) + kill.assert_called_once_with(12345, signal.SIGTERM) + self.assertIsNone(self._discovered_server().pid) + + def test_stop_does_not_signal_reused_pid(self) -> None: + from unittest import mock + + with Discovery() as discovery: + server = self._server(pid=12345) + discovery.save( + {k: getattr(server, k) for k in ("host", "port", "token", "pid", "spark_version")} + ) + # Model a recycled pid without depending on host process state. + ps_result = subprocess.CompletedProcess([], 0, stdout="unrelated process") + with ( + mock.patch.object(subprocess, "run", return_value=ps_result), + mock.patch.object(os, "kill") as kill, + ): + self.assertFalse(local_server.stop_local_connect_server()) + kill.assert_not_called() + self.assertIsNone(self._discovered_server().pid) + + def test_stop_preserves_discovery_when_process_cannot_be_inspected(self) -> None: + from unittest import mock + + with Discovery() as discovery: + server = self._server(pid=12345) + discovery.save( + {k: getattr(server, k) for k in ("host", "port", "token", "pid", "spark_version")} + ) + with ( + mock.patch.object(subprocess, "run", side_effect=subprocess.TimeoutExpired("ps", 5)), + mock.patch.object(os, "kill") as kill, + ): + self.assertIsNone(local_server.stop_local_connect_server()) + kill.assert_not_called() + self.assertEqual(self._discovered_server().pid, 12345) + with Discovery() as discovery: + discovery.clear() + + def test_server_launcher_binds_to_loopback(self) -> None: + from unittest import mock + + with Discovery() as discovery: + launcher = local_server.ServerLauncher("local[2]", {}, discovery) + # Capture the launcher argv without starting an external daemon. + with ( + mock.patch.dict(os.environ, {"SPARK_HOME": self._tmpdir}), + mock.patch.object(os.path, "isfile", return_value=True) as isfile, + mock.patch.object( + subprocess, "run", return_value=subprocess.CompletedProcess([], 0) + ) as run, + ): + launcher._run_script(15002, "token", None) + isfile.assert_called_once_with( + os.path.join(self._tmpdir, "sbin", "start-connect-server.sh") + ) + self.assertIn("spark.connect.grpc.binding.address=127.0.0.1", run.call_args.args[0]) + + def test_stop_cli_reports_when_no_server(self) -> None: + result = subprocess.run( + [sys.executable, "-m", "pyspark.sql.connect.local_server", "--stop"], + env=dict(os.environ), + capture_output=True, + text=True, + timeout=120, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("No running persistent local Spark Connect server", result.stdout) + + def test_stop_cli_fails_when_process_cannot_be_inspected(self) -> None: + from unittest import mock + + with ( + mock.patch.object(sys, "argv", ["local_server", "--stop"]), + mock.patch.object(local_server, "stop_local_connect_server", return_value=None), + self.assertRaises(SystemExit) as raised, + ): + local_server.main() + self.assertEqual(raised.exception.code, 1) + + def test_reuse_or_start_requires_posix(self) -> None: + from unittest import mock + + from pyspark.errors import PySparkRuntimeError + + with mock.patch.object(os, "name", "nt"): + with self.assertRaises(PySparkRuntimeError) as ctx: + local_server.reuse_or_start_local_connect_server("local[2]", {}) + self.assertIn("POSIX", str(ctx.exception)) + + def _release(self, session) -> None: + """Close one client session without stopping the shared server.""" + try: + session.client.release_session() + except Exception: + pass + try: + session.client.close() + except Exception: + pass + + def _wait_port_closed(self, host, port, timeout=30) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.settimeout(0.5) + if sock.connect_ex((host, int(port))) != 0: + return True + time.sleep(0.5) + return False + + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") + def test_builder_remote_local_uses_reuse_flag(self) -> None: + spark = None + try: + spark = ( + PySparkSession.builder.remote("local[2]") + .config("spark.local.connect.reuse", "true") + .getOrCreate() + ) + self.assertEqual(spark.range(2).count(), 2) + + server = self._discovered_server() + self.assertIsNotNone(server.pid) + self.assertEqual(server.spark_version, __version__) + self.assertNotEqual(server.pid, os.getpid()) + finally: + if spark is not None: + spark.stop() + + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") + def test_concurrent_startup_reuses_one_server(self) -> None: + script = textwrap.dedent(""" + import json + import os + + from pyspark.sql import SparkSession + + spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.local.connect.reuse", "true") + .getOrCreate() + ) + try: + count = spark.range(1).count() + with open(os.environ["SPARK_LOCAL_CONNECT_DISCOVERY"], "r") as f: + disc = json.load(f) + print(json.dumps({"count": count, "pid": disc["pid"], "port": disc["port"]})) + finally: + spark.stop() + """) + env = dict(os.environ) + env["SPARK_LOCAL_CONNECT_DISCOVERY"] = self._discovery_path + env["SPARK_LOCAL_CONNECT_REUSE"] = "1" + + procs = [ + subprocess.Popen( + [sys.executable, "-c", script], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for _ in range(3) + ] + outputs = [] + try: + for proc in procs: + stdout, stderr = proc.communicate(timeout=180) + self.assertEqual(proc.returncode, 0, stderr) + lines = stdout.strip().splitlines() + self.assertTrue(lines, stderr) + outputs.append(json.loads(lines[-1])) + finally: + for proc in procs: + if proc.poll() is None: + proc.kill() + proc.communicate() + + self.assertEqual({o["count"] for o in outputs}, {1}) + self.assertEqual(len({o["pid"] for o in outputs}), 1) + self.assertEqual(len({o["port"] for o in outputs}), 1) + + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") + def test_start_reuse_and_session_isolation(self) -> None: + endpoint = local_server.reuse_or_start_local_connect_server("local[2]", {}) + self.assertTrue(endpoint.startswith("sc://localhost:")) + + server = self._discovered_server() + self.assertIsNotNone(server.pid) + self.assertEqual(server.url, endpoint) + self.assertEqual(server.spark_version, __version__) + self.assertEqual(os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN"), server.token) + first_pid = server.pid + + s1 = s2 = None + try: + # A second call reuses the running server instead of spawning a new one. + endpoint2 = local_server.reuse_or_start_local_connect_server("local[2]", {}) + self.assertEqual(endpoint2, endpoint) + self.assertEqual(self._discovered_server().pid, first_pid) + + s1 = RemoteSparkSession.builder.remote(endpoint).create() + s2 = RemoteSparkSession.builder.remote(endpoint).create() + self.assertEqual(s1.range(5).count(), 5) + self.assertEqual(s2.range(3).count(), 3) + + # Session-local state must not leak across connections. + s1.range(1).createOrReplaceTempView("only_in_s1") + self.assertIn("only_in_s1", [t.name for t in s1.catalog.listTables()]) + self.assertNotIn("only_in_s1", [t.name for t in s2.catalog.listTables()]) + finally: + if s1 is not None: + self._release(s1) + if s2 is not None: + self._release(s2) + + self.assertTrue(local_server.stop_local_connect_server()) + self.assertIsNone(self._discovered_server().pid) + # Check the port rather than the pid, which can linger while the JVM shuts down. + self.assertTrue( + self._wait_port_closed(server.host, server.port), + "server port {} still open after stop".format(server.port), + ) + + @unittest.skipUnless(os.name == "posix", "the reuse path relies on the POSIX sbin scripts") + def test_start_seeds_static_conf_on_the_server(self) -> None: + # spark.local.connect.* and spark.master must be stripped from the seed, not + # forwarded; startup succeeding with them present covers that. + warehouse = os.path.join(self._tmpdir, "seeded-wh") + opts = { + "spark.sql.warehouse.dir": warehouse, + "spark.local.connect.reuse": "true", + "spark.master": "local[2]", + } + endpoint = local_server.reuse_or_start_local_connect_server("local[2]", opts) + spark = None + try: + spark = RemoteSparkSession.builder.remote(endpoint).create() + # A static conf cannot be set per-session after the JVM is up, so seeing it here + # proves the seed reached the server's SparkConf. + self.assertTrue(spark.conf.get("spark.sql.warehouse.dir").endswith(warehouse)) + finally: + if spark is not None: + self._release(spark) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py new file mode 100644 index 0000000000000..3c0f4f158aaa5 --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,1794 @@ +# +# 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. +# + +import contextlib +import json +import os +import shutil +import subprocess +import sys +import tempfile +import textwrap +import time +import unittest +from typing import Tuple +from unittest import mock + +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect +from pyspark.util import is_remote_only + +if should_test_connect: + from pyspark.sql.connect import local_server_pool + from pyspark.sql.connect.local_server import _SERVER_CLASS, _pid_alive + from pyspark.sql.connect.local_server_pool import ( + _JVM_ENV_VARS, + MemberAttendant, + PendingState, + PoolDirectory, + PoolMember, + RetiredState, + ServerPool, + pool_fingerprint, + ) + from pyspark.version import __version__ + + +@contextlib.contextmanager +def _listening_socket(): + import socket + + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + listener.bind(("localhost", 0)) + listener.listen(1) + yield listener.getsockname()[1] + finally: + listener.close() + + +@contextlib.contextmanager +def _non_listening_socket(): + """Reserve a port without listening on it, so connection attempts are rejected.""" + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("localhost", 0)) + yield sock.getsockname()[1] + + +def _spawn_live_process() -> "subprocess.Popen": + """A child blocked on its parent pipe, standing in for a live pool server.""" + return subprocess.Popen( + [sys.executable, "-c", "import sys; sys.stdin.buffer.read()", _SERVER_CLASS], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + +def _spawn_stubborn_process() -> "subprocess.Popen": + """A pipe-blocked child that ignores SIGTERM, standing in for a hung server. It reports + when its handler is installed so tests do not signal it too early.""" + proc = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, sys\n" + "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + "print('ready', flush=True)\n" + "sys.stdin.buffer.read()", + _SERVER_CLASS, + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + assert proc.stdout is not None + assert proc.stdout.readline() == "ready\n" + return proc + + +def _wait_proc_dead(proc: "subprocess.Popen", timeout: float = 30.0) -> bool: + try: + proc.wait(timeout=timeout) + return True + except subprocess.TimeoutExpired: + return False + + +def _wait_pid_dead(pid: int, timeout: float = 30.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _pid_alive(pid): + return True + time.sleep(0.05) + return False + + +def _wait_pid_gone(pid: int, timeout: float = 60.0) -> bool: + """Wait for a non-child pid to stop running. Linux zombies count as terminated.""" + deadline = time.time() + timeout + while time.time() < deadline: + if not local_server_pool._pid_alive(pid): + return True + time.sleep(0.2) + return False + + +_SAVED_ENV_KEYS = ( + "SPARK_LOCAL_CONNECT_POOL", + "SPARK_LOCAL_CONNECT_POOL_DIR", + "SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT", + "SPARK_LOCAL_CONNECT_POOL_SIZE", + "SPARK_CONNECT_AUTHENTICATE_TOKEN", + "PYSPARK_DRIVER_PYTHON", + "PYSPARK_PYTHON", +) + + +# These tests start no server and exercise only stdlib filesystem code, so they do not need JVM +# access (no is_remote_only gate). should_test_connect is still required because importing +# PoolDirectory pulls in the pyspark.sql.connect package, which checks Connect dependencies. +@unittest.skipIf( + not should_test_connect, + connect_requirement_message or "Requires Spark Connect dependencies to import the pool module", +) +@unittest.skipUnless(os.name == "posix", "the pool relies on POSIX file locks") +class LocalConnectServerPoolUnitTests(unittest.TestCase): + """Tests for the pool filesystem model; no real servers are started.""" + + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp() + self._saved_env = {k: os.environ.get(k) for k in _SAVED_ENV_KEYS} + for k in _SAVED_ENV_KEYS: + os.environ.pop(k, None) + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = os.path.join(self._tmpdir, "pool") + self._directory = PoolDirectory() + self._pool = ServerPool(self._directory) + self._procs = [] + local_server_pool._claimed_member = None + + def tearDown(self) -> None: + local_server_pool._claimed_member = None + for proc in self._procs: + try: + proc.kill() + proc.communicate(timeout=10) + except Exception: + pass + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _live_process(self) -> "subprocess.Popen": + proc = _spawn_live_process() + self._procs.append(proc) + return proc + + def _stubborn_process(self) -> "subprocess.Popen": + proc = _spawn_stubborn_process() + self._procs.append(proc) + return proc + + def _attendant(self, uid: str) -> "subprocess.Popen": + proc = subprocess.Popen( + [ + sys.executable, + "-c", + "import sys; sys.stdin.buffer.read()", + "-m", + "pyspark.sql.connect.local_server_pool", + "--attend", + "--pool-dir", + self._directory.path, + "--uid", + uid, + ], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + self._procs.append(proc) + return proc + + def _attendant_with_launch_child(self, uid: str) -> Tuple["subprocess.Popen", int]: + proc = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys\n" + "child = subprocess.Popen([sys.executable, '-c', " + "'import time; time.sleep(300)'])\n" + "print(child.pid, flush=True)\n" + "sys.stdin.buffer.read()", + "-m", + "pyspark.sql.connect.local_server_pool", + "--attend", + "--pool-dir", + self._directory.path, + "--uid", + uid, + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + start_new_session=True, + ) + self._procs.append(proc) + assert proc.stdout is not None + return proc, int(proc.stdout.readline()) + + def _server_data(self, port: int, pid: int, fingerprint: str = "fp", **overrides) -> dict: + process_start_id = None + if isinstance(pid, int) and not isinstance(pid, bool): + process_start_id = ServerPool._process_start_id(pid) + data = { + "host": "localhost", + "port": port, + "token": "t", + "pid": pid, + "spark_version": __version__, + "fingerprint": fingerprint, + "process_start_id": process_start_id or f"unobserved:{pid}", + "created": time.time(), + } + data.update(overrides) + return data + + def _retired_data(self, pid: int, retired: object) -> dict: + process_start_id = ServerPool._process_start_id(pid) + assert process_start_id is not None + return { + "pid": pid, + "process_start_id": process_start_id, + "retired": retired, + } + + def _write_state(self, path: str, data: dict) -> str: + with self._directory as directory: + directory.write_json(path, data) + return path + + def _states(self, uid: str) -> dict: + with self._directory as directory: + return directory.states(uid) + + def _write_daemon_pid(self, uid: str, pid: int) -> None: + from pyspark.sql.connect.local_server import Discovery + + member_dir = self._directory.member_dir(uid) + os.makedirs(member_dir, exist_ok=True) + discovery = Discovery(os.path.join(member_dir, "connect-local.json")) + with open(discovery.daemon_pid_path, "w") as pid_file: + pid_file.write(str(pid)) + + def test_pool_directory_location(self) -> None: + self.assertEqual(self._directory.path, os.path.join(self._tmpdir, "pool")) + os.environ.pop("SPARK_LOCAL_CONNECT_POOL_DIR") + default = PoolDirectory() + self.assertEqual(os.path.basename(default.path), "pool") + self.assertTrue(default.path.startswith(tempfile.gettempdir())) + + def test_pool_directory_lock_and_state_file_permissions(self) -> None: + os.makedirs(self._directory.path, mode=0o755) + os.chmod(self._directory.path, 0o755) + state_path = self._directory.server_path("abcdef") + with self._directory as directory: + self.assertEqual(os.stat(directory.path).st_mode & 0o777, 0o700) + lock_path = os.path.join(directory.path, ".lock") + self.assertEqual(os.stat(lock_path).st_mode & 0o777, 0o600) + directory.write_json(state_path, {"token": "secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + self.assertEqual(directory.read_json(state_path), {"token": "secret"}) + + # Replacing an existing file with wider permissions must restore the private mode. + os.chmod(state_path, 0o644) + directory.write_json(state_path, {"token": "new-secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + + with open(state_path, "w") as state_file: + state_file.write("not json") + with self._directory as directory: + self.assertIsNone(directory.read_json(state_path)) + + directory.write_json(state_path, {"token": "old-secret"}) + with mock.patch("builtins.open", side_effect=OSError("temporary read failure")): + with self.assertRaisesRegex(OSError, "temporary read failure"): + directory.read_json(state_path) + with mock.patch.object( + local_server_pool.os, "replace", side_effect=OSError("interrupted replace") + ): + with self.assertRaisesRegex(OSError, "interrupted replace"): + directory.write_json(state_path, {"token": "new-secret"}) + self.assertEqual(directory.read_json(state_path), {"token": "old-secret"}) + self.assertFalse( + [ + name + for name in os.listdir(directory.path) + if name.startswith(directory._STATE_TEMP_PREFIX) + ] + ) + + stale_temp = os.path.join(self._directory.path, ".pool-state-orphan") + with open(stale_temp, "w") as temp_file: + temp_file.write("partial") + with self._directory: + self.assertFalse(os.path.exists(stale_temp)) + + def test_failed_state_write_does_not_close_a_reused_descriptor(self) -> None: + real_fdopen = os.fdopen + victim_fd = None + test_case = self + + class FailingStateFile: + def __init__(self, fd: int): + self.fd = fd + self.file = real_fdopen(fd, "w") + + def __enter__(self): + return self + + def write(self, data: str) -> None: + raise OSError("interrupted write") + + def __exit__(self, exc_type, exc_value, traceback) -> None: + nonlocal victim_fd + self.file.close() + # Reuse the just-released descriptor before write_json handles the failure. Once + # fdopen succeeds, it owns the original descriptor, so cleanup must not close the + # new file that happens to receive the same number. + victim_fd = os.open(os.devnull, os.O_RDONLY) + test_case.assertEqual(victim_fd, self.fd) + + try: + with self._directory as directory: + with mock.patch.object( + local_server_pool.os, + "fdopen", + side_effect=lambda fd, mode: FailingStateFile(fd), + ): + with self.assertRaisesRegex(OSError, "interrupted write"): + directory.write_json(directory.server_path("f00d"), {"a": 1}) + assert victim_fd is not None + os.fstat(victim_fd) + finally: + if victim_fd is not None: + with contextlib.suppress(OSError): + os.close(victim_fd) + + def test_pool_directory_does_not_hide_listing_failures(self) -> None: + state_path = self._write_state( + self._directory.retired_path("abc123"), + {"pid": os.getpid(), "process_start_id": "unused", "retired": time.time()}, + ) + with self._directory: + with mock.patch.object( + local_server_pool.os, "listdir", side_effect=OSError("temporary listing failure") + ): + with self.assertRaisesRegex(OSError, "temporary listing failure"): + self._pool.reap("abc123") + + self.assertTrue(os.path.exists(state_path)) + with mock.patch.object( + local_server_pool.os, "listdir", side_effect=OSError("failed during enter") + ): + with self.assertRaisesRegex(OSError, "failed during enter"): + with self._directory: + pass + self.assertIsNone(self._directory._lock_fd) + with self._directory: + pass + + def test_pool_directory_lock_blocks_another_process(self) -> None: + child = ( + "import errno\n" + "import fcntl\n" + "import os\n" + "import sys\n" + "fd = os.open(sys.argv[1], os.O_RDWR)\n" + "try:\n" + " fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)\n" + "except OSError as error:\n" + " if error.errno not in (errno.EACCES, errno.EAGAIN):\n" + " raise\n" + "else:\n" + " raise RuntimeError('acquired a held lock')\n" + "finally:\n" + " os.close(fd)\n" + ) + with self._directory: + result = subprocess.run( + [sys.executable, "-c", child, os.path.join(self._directory.path, ".lock")], + capture_output=True, + text=True, + timeout=10, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_parse_entry_grammar(self) -> None: + # uids match the acquisition layer's uuid4().hex[:12]: nonempty lowercase hex. + uid = "0123456789ab" + cases = [ + # Well-formed state entries of every kind. + (f"pending-{uid}.json", ("pending", uid)), + (f"conf-{uid}.json", ("conf", uid)), + (f"server-{uid}.json", ("server", uid)), + (f"retired-{uid}.json", ("retired", uid)), + (f"claimed-4321-{uid}.json", ("claimed", uid)), + (f"member-{uid}", ("member", uid)), + # Short but valid hex uids. + ("member-abc", ("member", "abc")), + ("server-abc.json", ("server", "abc")), + # The lock file and unrelated entries. + (".lock", (None, None)), + ("random.txt", (None, None)), + # Editor droppings: the finding-2 cases that used to slip through as phantom uids. + (f"member-{uid}.json.swp", (None, None)), + (f"server-{uid}.json.swp", (None, None)), + # Empty uids are rejected for every kind. + ("server-.json", (None, None)), + ("member-", (None, None)), + ("claimed-1234-.json", (None, None)), + # Non-hex uids (uppercase, out-of-range letters) are rejected. + ("server-ABCDEF.json", (None, None)), + ("server-ghij.json", (None, None)), + # Malformed claimed stems (non-numeric or missing pid, or a malformed uid). + (f"claimed-notapid-{uid}.json", (None, None)), + (f"claimed-{uid}.json", (None, None)), + ("claimed-4321-ABCDEF.json", (None, None)), + # A pid that is str.isdigit() but not int()-parsable (superscript two, U+00B2) must + # classify as "not claimed", never raise, since parse_entry runs over every entry. + (f"claimed-{chr(0xB2)}-{uid}.json", (None, None)), + ] + for name, expected in cases: + with self.subTest(name=name): + self.assertEqual(PoolDirectory.parse_entry(name), expected) + + def test_claiming_pid(self) -> None: + uid = "0123456789ab" + path = self._directory.claimed_path(4321, uid) + self.assertEqual(PoolDirectory.claiming_pid(path), 4321) + # Parses from the basename alone, independent of the directory prefix. + self.assertEqual(PoolDirectory.claiming_pid(f"claimed-7-{uid}.json"), 7) + + def test_locked_accessors_enumerate_state(self) -> None: + uid_a, uid_b = "aaaaaaaaaaaa", "bbbbbbbbbbbb" + with self._directory as directory: + directory.write_json(directory.server_path(uid_a), {"a": 1}) + directory.write_json(directory.pending_path(uid_a), {"a": 2}) + directory.write_json(directory.server_path(uid_b), {"b": 1}) + os.makedirs(directory.member_dir(uid_a), mode=0o700) + + self.assertEqual(sorted(directory.uids()), [uid_a, uid_b]) + + states_a = directory.states(uid_a) + self.assertEqual(set(states_a), {"server", "pending", "member"}) + self.assertEqual(states_a["server"], directory.server_path(uid_a)) + self.assertEqual(states_a["member"], directory.member_dir(uid_a)) + + servers = dict(directory.paths_of_kind("server")) + self.assertEqual(set(servers), {uid_a, uid_b}) + self.assertEqual(servers[uid_a], directory.server_path(uid_a)) + self.assertEqual(directory.paths_of_kind("retired"), []) + + def test_rename_remove_and_member_dir(self) -> None: + uid = "cccccccccccc" + with self._directory as directory: + src = directory.pending_path(uid) + dst = directory.server_path(uid) + directory.write_json(src, {"x": 1}) + directory.rename(src, dst) + self.assertFalse(os.path.exists(src)) + self.assertEqual(directory.read_json(dst), {"x": 1}) + + directory.remove(dst) + self.assertFalse(os.path.exists(dst)) + # Removing a missing path is a no-op. + directory.remove(dst) + + member = directory.member_dir(uid) + os.makedirs(member, mode=0o700) + with open(os.path.join(member, "inner"), "w") as f: + f.write("data") + directory.remove_member_dir(uid) + self.assertFalse(os.path.exists(member)) + # Removing a missing member directory is a no-op. + directory.remove_member_dir(uid) + + def test_states_rejects_duplicate_claimed(self) -> None: + uid = "dddddddddddd" + with self._directory as directory: + directory.write_json(directory.claimed_path(111, uid), {}) + directory.write_json(directory.claimed_path(222, uid), {}) + with self.assertRaisesRegex(AssertionError, "duplicate claimed"): + directory.states(uid) + + def test_accessors_require_the_lock(self) -> None: + # The locked accessors must refuse to run outside the context manager. + with self.assertRaisesRegex(AssertionError, "context manager"): + self._directory.uids() + + def test_not_reentrant(self) -> None: + with self._directory: + with self.assertRaisesRegex(AssertionError, "not reentrant"): + with self._directory: + pass + + @unittest.skipUnless( + sys.platform.startswith("linux") and os.path.isdir("/proc") and hasattr(os, "waitid"), + "requires Linux process state and waitid", + ) + def test_pid_alive_treats_zombie_as_dead(self) -> None: + proc = subprocess.Popen([sys.executable, "-c", "pass"]) + try: + # Wait for the child to exit but leave it waitable, which keeps it as a zombie + # until the finally block reaps it. + os.waitid(os.P_PID, proc.pid, os.WEXITED | os.WNOWAIT) + self.assertFalse(_pid_alive(proc.pid)) + finally: + proc.wait(timeout=10) + + def test_fingerprint_identity(self) -> None: + base = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + self.assertEqual(base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"})) + self.assertNotEqual( + base, pool_fingerprint("local[2]", {"spark.sql.shuffle.partitions": "4"}) + ) + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "8"}) + ) + self.assertNotEqual(base, pool_fingerprint("local[*]", {})) + # The working directory shapes the server (relative warehouse and metastore paths), + # so members are never shared across directories. + cwd = os.getcwd() + try: + os.chdir(self._tmpdir) + self.assertNotEqual(base, pool_fingerprint("local[*]", {"x": "4"})) + in_tmpdir = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + finally: + os.chdir(cwd) + self.assertNotEqual(base, in_tmpdir) + # So does the Python environment the server would run UDFs with. + os.environ["PYSPARK_PYTHON"] = "/some/other/python" + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + ) + # Match SparkConnectPlanner.pythonExec's PYSPARK_PYTHON -> PYSPARK_DRIVER_PYTHON -> + # python3 precedence for Connect UDFs, so clients never claim a server that resolved a + # different fallback interpreter. + os.environ.pop("PYSPARK_PYTHON") + os.environ["PYSPARK_DRIVER_PYTHON"] = "python3" + self.assertEqual(base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"})) + os.environ["PYSPARK_DRIVER_PYTHON"] = "/driver/python" + self.assertNotEqual( + base, pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + ) + # PYSPARK_DRIVER_PYTHON also feeds PythonUtils.defaultPythonExec (Python data sources), + # which prefers it over PYSPARK_PYTHON. So even with PYSPARK_PYTHON fixed, changing + # PYSPARK_DRIVER_PYTHON changes the server a run would boot and must change the identity. + os.environ["PYSPARK_PYTHON"] = "/worker/python" + worker_python = pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}) + os.environ["PYSPARK_DRIVER_PYTHON"] = "/other/driver/python" + self.assertNotEqual( + worker_python, + pool_fingerprint("local[*]", {"spark.sql.shuffle.partitions": "4"}), + ) + + def test_fingerprint_resolves_server_python_from_path(self) -> None: + # Relative worker commands are resolved through PATH by the server. Include that + # resolution so equal command strings cannot identify different Python environments. + executable_name = "pool-test-python" + executable_dirs = [os.path.join(self._tmpdir, name) for name in ("env-a", "env-b")] + for directory in executable_dirs: + os.makedirs(directory) + executable = os.path.join(directory, executable_name) + with open(executable, "w") as executable_file: + executable_file.write("#!/bin/sh\n") + os.chmod(executable, 0o700) + os.environ["PYSPARK_PYTHON"] = executable_name + with mock.patch.dict(os.environ, {"PATH": executable_dirs[0]}): + first_environment = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"PATH": executable_dirs[1]}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + + def test_fingerprint_includes_python_and_spark_paths(self) -> None: + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/a", "SPARK_HOME": "/spark/a"}): + first_environment = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/b", "SPARK_HOME": "/spark/a"}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + with mock.patch.dict(os.environ, {"PYTHONPATH": "/python/a", "SPARK_HOME": "/spark/b"}): + self.assertNotEqual(first_environment, pool_fingerprint("local[*]", {})) + + def test_fingerprint_conf_order_independent_and_string_keyed(self) -> None: + # sorted() over seed_conf makes the identity independent of dict insertion order, so a + # run that builds the same confs in a different order still matches. + self.assertEqual( + pool_fingerprint("local[*]", {"a": "1", "b": "2"}), + pool_fingerprint("local[*]", {"b": "2", "a": "1"}), + ) + # Confs serialize to a properties file, so values are compared as strings: 1 and "1" + # are the same seed and share an identity. + self.assertEqual( + pool_fingerprint("local[*]", {"k": 1}), pool_fingerprint("local[*]", {"k": "1"}) + ) + + def test_fingerprint_includes_python_executable(self) -> None: + # sys.executable is the client interpreter the server inherits; a packaging change that + # moves it must not silently reuse a server booted under the old one. + base = pool_fingerprint("local[*]", {}) + with mock.patch.object(sys, "executable", "/other/python"): + self.assertNotEqual(base, pool_fingerprint("local[*]", {})) + + def test_fingerprint_includes_jvm_env(self) -> None: + # Every JVM-shaping variable must change the identity. SPARK_CONF_DIR is the common CI + # case (it selects the spark-defaults.conf / spark-env.sh the server reads); the rest + # feed the classpath, heap, and JVM options. The expected list is spelled out here + # independently rather than derived from _JVM_ENV_VARS: the fingerprint reads that same + # tuple, so dropping a variable from it would silently leave the fingerprint AND a loop + # over it in agreement. The equality check catches such drift (a removal or an unlisted + # addition), and the loop proves each variable still changes the identity. + expected = ( + "SPARK_CONF_DIR", + "JAVA_HOME", + "SPARK_DIST_CLASSPATH", + "SPARK_DAEMON_MEMORY", + "SPARK_DRIVER_MEMORY", + "SPARK_SUBMIT_OPTS", + "SPARK_DAEMON_JAVA_OPTS", + ) + self.assertEqual(set(_JVM_ENV_VARS), set(expected)) + for var in expected: + with self.subTest(var=var): + with mock.patch.dict(os.environ, {var: "/value/a"}): + with_a = pool_fingerprint("local[*]", {}) + with mock.patch.dict(os.environ, {var: "/value/b"}): + self.assertNotEqual(with_a, pool_fingerprint("local[*]", {})) + + def test_pool_member_validation(self) -> None: + valid = self._server_data(12345, 123, created=1) + member = PoolMember.from_data(valid) + self.assertIsNotNone(member) + self.assertEqual(member.host, "localhost") + self.assertEqual(member.port, 12345) + self.assertEqual(member.pid, 123) + self.assertEqual(member.process_start_id, valid["process_start_id"]) + self.assertEqual(member.created, 1.0) + self.assertEqual(member.url, "sc://localhost:12345") + + invalid_records = { + "missing fields": {"fingerprint": "fp"}, + "empty token": self._server_data(12345, 123, token=""), + "non-string host": self._server_data(12345, 123, host=None), + "empty process start id": self._server_data(12345, 123, process_start_id=""), + "non-string process start id": self._server_data(12345, 123, process_start_id=None), + "boolean port": self._server_data(True, 123), + "string port": self._server_data("12345", 123), + "fractional port": self._server_data(12345.5, 123), + "zero port": self._server_data(0, 123), + "out-of-range port": self._server_data(65536, 123), + "boolean pid": self._server_data(12345, True), + "string pid": self._server_data(12345, "123"), + "fractional pid": self._server_data(12345, 123.5), + "zero pid": self._server_data(12345, 0), + "string created": self._server_data(12345, 123, created="1"), + "boolean created": self._server_data(12345, 123, created=True), + "negative created": self._server_data(12345, 123, created=-1), + "nan created": self._server_data(12345, 123, created=float("nan")), + "infinite created": self._server_data(12345, 123, created=float("inf")), + # A finite float, but far past any real clock: rejected so it cannot look + # perpetually fresh to age-based reaping. + "far-future created": self._server_data(12345, 123, created=2**100), + # Too large to convert to float at all -- the OverflowError guard in from_data keeps + # a corrupt state file (this round-trips through json) from crashing the caller. + "overflow created": self._server_data(12345, 123, created=10**400), + } + for name, data in invalid_records.items(): + with self.subTest(name=name): + self.assertIsNone(PoolMember.from_data(data)) + + def test_lifecycle_state_fields_and_validation(self) -> None: + pending = PendingState.from_data({"attendant_pid": 123, "created": 1, "fingerprint": "fp"}) + assert pending is not None + self.assertEqual(pending.attendant_pid, 123) + self.assertEqual(pending.created, 1.0) + self.assertEqual(pending.fingerprint, "fp") + self.assertIsNone( + PendingState.from_data({"attendant_pid": "123", "created": 1, "fingerprint": "fp"}) + ) + + retired = RetiredState.from_data( + {"pid": 456, "process_start_id": "process-1", "retired": 2} + ) + assert retired is not None + self.assertEqual(retired.pid, 456) + self.assertEqual(retired.process_start_id, "process-1") + self.assertEqual(retired.retired, 2.0) + self.assertFalse(retired.signalled) + self.assertEqual( + retired.as_data(), + { + "pid": 456, + "process_start_id": "process-1", + "retired": 2.0, + "signalled": False, + }, + ) + delivered = RetiredState.from_data( + { + "pid": 456, + "process_start_id": "process-1", + "retired": 2, + "signalled": True, + } + ) + assert delivered is not None + self.assertTrue(delivered.signalled) + self.assertIsNone( + RetiredState.from_data( + {"pid": 456, "process_start_id": "process-1", "retired": "not-a-time"} + ) + ) + self.assertIsNone( + RetiredState.from_data( + { + "pid": 456, + "process_start_id": "process-1", + "retired": 2, + "signalled": 1, + } + ) + ) + self.assertIsNone(RetiredState.from_data({"pid": 456, "retired": 2})) + + def test_claim_matches_fingerprint_and_renames(self) -> None: + with _listening_socket() as port: + server_process = self._live_process() + self._write_state( + self._directory.server_path("aaa"), + self._server_data(port, server_process.pid, fingerprint="other-fp"), + ) + self._write_state( + self._directory.server_path("bbb"), + self._server_data(port, server_process.pid, fingerprint="my-fp", token="t-bbb"), + ) + with self._directory: + member = self._pool.claim("my-fp") + self.assertIsNotNone(member) + self.assertEqual(member.token, "t-bbb") + claim_name = f"claimed-{os.getpid()}-bbb.json" + self.assertEqual(os.path.basename(member.claim_path), claim_name) + states = self._states("bbb") + self.assertEqual(set(states), {"claimed"}) + with self._directory as directory: + claimed = directory.read_json(states["claimed"]) + assert claimed is not None + self.assertEqual( + claimed["client_process_start_id"], ServerPool._process_start_id(os.getpid()) + ) + # The mismatched member is untouched, and a second claim finds nothing. + self.assertEqual(set(self._states("aaa")), {"server"}) + with self._directory: + self.assertFalse(self._pool.reap("bbb")) + self.assertIsNone(self._pool.claim("my-fp")) + + def test_claim_prefers_the_oldest_member(self) -> None: + with _listening_socket() as port: + server_process = self._live_process() + for uid, created in (("aaaa", time.time()), ("bbbb", time.time() - 100)): + self._write_state( + self._directory.server_path(uid), + self._server_data(port, server_process.pid, token="t-" + uid, created=created), + ) + with self._directory: + member = self._pool.claim("fp") + # Prefer the oldest ready member. Ordering is by wall-clock created, so this is + # approximate FIFO rather than a guarantee (see ServerPool.claim). + self.assertEqual(member.token, "t-bbbb") + + def test_claim_requires_the_lock(self) -> None: + # claim reaches the directory only through the locked accessors, so it inherits their + # assertion; pin the caller obligation directly so a future reordering that touches the + # directory before the first locked accessor is still caught. + with self.assertRaisesRegex(AssertionError, "context manager"): + self._pool.claim("fp") + + def test_claim_ignores_already_claimed_member(self) -> None: + # The kind filter is the exclusion invariant: a member already renamed to claimed-* is + # no longer of kind "server", so claim never hands it out a second time. A live, usable + # record under a claimed-* name must still be invisible to a new claimer. + with _listening_socket() as port: + server_process = self._live_process() + uid = "feed" + self._write_state( + self._directory.claimed_path(os.getpid() + 1, uid), + self._server_data(port, server_process.pid), + ) + with self._directory: + self.assertIsNone(self._pool.claim("fp")) + self.assertEqual(set(self._states(uid)), {"claimed"}) + + def test_concurrent_claimers_claim_one_member_once(self) -> None: + child = ( + "import sys\n" + "from pyspark.sql.connect.local_server_pool import PoolDirectory, ServerPool\n" + "directory = PoolDirectory(sys.argv[1])\n" + "with directory:\n" + " member = ServerPool(directory).claim('fp')\n" + "print(member.token if member is not None else 'NONE')\n" + ) + claimers = [] + results = [] + with _listening_socket() as port: + server_process = self._live_process() + uid = "cafe" + self._write_state( + self._directory.server_path(uid), + self._server_data(port, server_process.pid, token="claimed-once"), + ) + try: + with self._directory: + for _ in range(2): + claimers.append( + subprocess.Popen( + [sys.executable, "-c", child, self._directory.path], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + ) + for proc in claimers: + stdout, stderr = proc.communicate(timeout=20) + self.assertEqual(proc.returncode, 0, stderr) + results.append(stdout.strip()) + finally: + for proc in claimers: + if proc.poll() is None: + proc.kill() + proc.communicate(timeout=10) + + # Pins "claimed at most once": one child claims the member and the other sees none. It + # does not guarantee the two ever contend on the lock -- if the first child finishes + # before the second reaches flock(), the second simply finds no "server" entry -- so + # read this as a regression test for the exclusion invariant, not for lock contention. + self.assertEqual(sorted(results), ["NONE", "claimed-once"]) + states = self._states(uid) + self.assertEqual(set(states), {"claimed"}) + claiming_pid = PoolDirectory.claiming_pid(states["claimed"]) + self.assertIn(claiming_pid, [proc.pid for proc in claimers]) + + def test_claim_skips_unreachable_member(self) -> None: + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("ccc"), self._server_data(port, os.getpid()) + ) + with self._directory: + self.assertIsNone(self._pool.claim("fp")) + + def test_claim_skips_directory_with_state_filename(self) -> None: + os.makedirs(self._directory.path) + os.makedirs(self._directory.server_path("dead")) + with _listening_socket() as port: + server_process = self._live_process() + self._write_state( + self._directory.server_path("cafe"), self._server_data(port, server_process.pid) + ) + with self._directory: + member = self._pool.claim("fp") + + self.assertIsNotNone(member) + self.assertEqual(os.path.basename(member.claim_path), f"claimed-{os.getpid()}-cafe.json") + + def test_claim_skips_member_with_mismatched_process_identity(self) -> None: + with _listening_socket() as port: + server_process = self._live_process() + self._write_state( + self._directory.server_path("fade"), + self._server_data( + port, + server_process.pid, + process_start_id="a-different-process-generation", + ), + ) + with self._directory: + self.assertIsNone(self._pool.claim("fp")) + + self.assertEqual(set(self._states("fade")), {"server"}) + self.assertIsNone(server_process.poll()) + + def test_claim_skips_malformed_and_incompatible_members(self) -> None: + with _listening_socket() as port: + server_process = self._live_process() + bad_pid = self._server_data(port, server_process.pid) + bad_pid["pid"] = "not-a-pid" + bad_port = self._server_data(port, server_process.pid) + bad_port["port"] = "not-a-port" + bad_created = self._server_data(port, server_process.pid) + bad_created["created"] = "not-a-time" + non_finite_created = self._server_data(port, server_process.pid) + non_finite_created["created"] = float("nan") + out_of_range_port = self._server_data(port, server_process.pid) + out_of_range_port["port"] = 65536 + bad_host = self._server_data(port, server_process.pid) + bad_host["host"] = None + records = { + "a0": {"fingerprint": "fp"}, + "a1": bad_pid, + "a2": bad_port, + "a3": bad_created, + "a4": non_finite_created, + "a5": out_of_range_port, + "a6": bad_host, + "a7": self._server_data(port, server_process.pid, spark_version="not-this-version"), + "a8": self._server_data(port, 2**31 - 1), + "a9": self._server_data(port, 2**100), + } + for uid, data in records.items(): + self._write_state(self._directory.server_path(uid), data) + self._write_state( + self._directory.server_path("b0"), + self._server_data(port, server_process.pid, token="valid-token"), + ) + + with self._directory: + member = self._pool.claim("fp") + + self.assertIsInstance(member, PoolMember) + self.assertEqual(member.token, "valid-token") + for uid in records: + self.assertEqual(set(self._states(uid)), {"server"}) + + def test_reap_pending_of_dead_attendant(self) -> None: + # The attendant died mid-boot: its pending marker and conf seed are withdrawn, and + # the half-started server whose pid spark-daemon.sh recorded remains tracked. A daemon + # pid has no process-generation identity, so it cannot safely authorize a signal. + half_started = self._stubborn_process() + self._write_daemon_pid("b007", half_started.pid) + self._write_state( + self._directory.pending_path("b007"), + {"attendant_pid": 2**31 - 1, "created": time.time(), "fingerprint": "fp"}, + ) + self._write_state(self._directory.conf_path("b007"), {"spark.foo": "bar"}) + + with self._directory: + self._pool.reap("b007") + + states = self._states("b007") + self.assertNotIn("pending", states) + self.assertNotIn("conf", states) + self.assertEqual(set(states), {"member", "retired"}) + with self._directory as directory: + retired = directory.read_json(states["retired"]) + assert retired is not None + self.assertEqual(retired["pid"], half_started.pid) + self.assertNotIn("process_start_id", retired) + self.assertIsNone(half_started.poll()) + + half_started.kill() + self.assertTrue(_wait_proc_dead(half_started)) + with self._directory: + self.assertTrue(self._pool.reap("b007")) + + def test_reap_malformed_pending(self) -> None: + attendant = self._attendant("bad3") + self._write_state( + self._directory.pending_path("bad3"), + {"attendant_pid": attendant.pid, "created": "not-a-time"}, + ) + self._write_state(self._directory.conf_path("bad3"), {"spark.foo": "bar"}) + + with self._directory: + self.assertTrue(self._pool.reap("bad3")) + + self.assertTrue(_wait_proc_dead(attendant)) + + def test_reap_does_not_signal_reused_attendant_pid(self) -> None: + unrelated = self._live_process() + self._write_state( + self._directory.pending_path("bad8"), + { + "attendant_pid": unrelated.pid, + "created": time.time() - 181, + "fingerprint": "fp", + }, + ) + self._write_state(self._directory.conf_path("bad8"), {"spark.foo": "bar"}) + + with mock.patch.object(local_server_pool, "_process_command", return_value=None): + with self._directory: + self.assertFalse(self._pool.reap("bad8")) + self.assertEqual(set(self._states("bad8")), {"conf", "pending"}) + + with self._directory: + self.assertTrue(self._pool.reap("bad8")) + + self.assertIsNone(unrelated.poll()) + + def test_reap_timed_out_attendant_kills_its_launch_group(self) -> None: + attendant, launch_child_pid = self._attendant_with_launch_child("bad0") + self._write_state( + self._directory.pending_path("bad0"), + { + "attendant_pid": attendant.pid, + "created": time.time() - 181, + "fingerprint": "fp", + }, + ) + self._write_state(self._directory.conf_path("bad0"), {"spark.foo": "bar"}) + + with self._directory: + self.assertTrue(self._pool.reap("bad0")) + + self.assertTrue(_wait_proc_dead(attendant)) + self.assertTrue(_wait_pid_dead(launch_child_pid)) + + def test_reap_pending_with_published_server_retires_server(self) -> None: + # Publishing writes server-* before removing pending-*. If the attendant dies between + # those operations, the janitor must not leave that server available to claim. + server = self._stubborn_process() + self._write_daemon_pid("bad7", server.pid) + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("bad7"), self._server_data(port, server.pid) + ) + self._write_state( + self._directory.pending_path("bad7"), + {"attendant_pid": 2**31 - 1, "created": time.time(), "fingerprint": "fp"}, + ) + self._write_state(self._directory.conf_path("bad7"), {"spark.foo": "bar"}) + with self._directory: + self._pool.reap("bad7") + self.assertIsNone(self._pool.claim("fp")) + + self.assertEqual(set(self._states("bad7")), {"member", "retired"}) + self.assertIsNone(server.poll()) + + def test_reap_removes_conf_left_after_publication(self) -> None: + server = self._live_process() + with _listening_socket() as port: + self._write_state( + self._directory.server_path("c0f1"), self._server_data(port, server.pid) + ) + self._write_state(self._directory.conf_path("c0f1"), {"spark.foo": "bar"}) + + with self._directory: + self._pool.reap("c0f1") + + self.assertEqual(set(self._states("c0f1")), {"server"}) + self.assertIsNone(server.poll()) + + def test_reap_removes_stale_conf_without_an_attendant(self) -> None: + conf_path = self._write_state(self._directory.conf_path("c0f2"), {"spark.foo": "bar"}) + old = time.time() - 181 + os.utime(conf_path, (old, old)) + + with self._directory: + self.assertTrue(self._pool.reap("c0f2")) + + self.assertFalse(os.path.exists(conf_path)) + + def test_reap_keeps_live_pending(self) -> None: + attendant = self._live_process() + self._write_state( + self._directory.pending_path("11ce"), + {"attendant_pid": attendant.pid, "created": time.time(), "fingerprint": "fp"}, + ) + with self._directory: + self._pool.reap("11ce") + self.assertIn("pending", self._states("11ce")) + self.assertIsNone(attendant.poll()) + + def test_reap_server_unreachable_and_idle(self) -> None: + with self.subTest("unreachable member is retired"): + gone = self._live_process() + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("dead"), self._server_data(port, gone.pid) + ) + with self._directory: + self._pool.reap("dead") + self.assertEqual(set(self._states("dead")), {"retired"}) + self.assertTrue(_wait_proc_dead(gone)) + with self.subTest("member idle past the timeout is retired"): + os.environ["SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT"] = "10" + with _listening_socket() as port: + idle = self._live_process() + self._write_state( + self._directory.server_path("1d1e"), + self._server_data(port, idle.pid, created=time.time() - 60), + ) + fresh = self._live_process() + self._write_state( + self._directory.server_path("f2e5"), self._server_data(port, fresh.pid) + ) + with self._directory: + self._pool.janitor() + self.assertEqual(set(self._states("1d1e")), {"retired"}) + self.assertEqual(set(self._states("f2e5")), {"server"}) + self.assertTrue(_wait_proc_dead(idle)) + self.assertIsNone(fresh.poll()) + + def test_reap_does_not_signal_reused_server_pid(self) -> None: + unrelated = subprocess.Popen( + [sys.executable, "-c", "import sys; sys.stdin.buffer.read()"], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + self._procs.append(unrelated) + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("bad6"), self._server_data(port, unrelated.pid) + ) + with self._directory: + self._pool.reap("bad6") + + self.assertEqual(set(self._states("bad6")), {"retired"}) + self.assertIsNone(unrelated.poll()) + self.assertEqual(self._pool.purge(), 0) + self.assertIsNone(unrelated.poll()) + + def test_reap_malformed_member_recovers_server_pid(self) -> None: + # An out-of-range created value makes the full member invalid, but its independently + # valid pid must still be retired and tracked rather than leaving a live JVM orphaned. + invalid_created = self._live_process() + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("bad0"), + self._server_data(port, invalid_created.pid, created=2**100), + ) + with self._directory: + self._pool.reap("bad0") + states = self._states("bad0") + self.assertEqual(set(states), {"retired"}) + with self._directory as directory: + self.assertEqual(directory.read_json(states["retired"])["pid"], invalid_created.pid) + self.assertTrue(_wait_proc_dead(invalid_created)) + + # Claimed records use the same recovery when their client has disappeared. + claimed = self._live_process() + with _non_listening_socket() as port: + self._write_state( + self._directory.claimed_path(2**31 - 1, "bad2"), + self._server_data(port, claimed.pid, created=2**100), + ) + with self._directory: + self._pool.reap("bad2") + states = self._states("bad2") + self.assertEqual(set(states), {"retired"}) + with self._directory as directory: + self.assertEqual(directory.read_json(states["retired"])["pid"], claimed.pid) + self.assertTrue(_wait_proc_dead(claimed)) + + # If the record itself has no usable pid, fall back to spark-daemon.sh's pid file. The + # fallback keeps the shutdown tracked but lacks the process identity needed to signal it. + unreadable_pid = self._live_process() + self._write_daemon_pid("bad1", unreadable_pid.pid) + self._write_state(self._directory.server_path("bad1"), {"malformed": True}) + with self._directory: + self._pool.reap("bad1") + states = self._states("bad1") + self.assertEqual(set(states), {"member", "retired"}) + with self._directory as directory: + retired = directory.read_json(states["retired"]) + assert retired is not None + self.assertEqual(retired["pid"], unreadable_pid.pid) + self.assertNotIn("process_start_id", retired) + self.assertIsNone(unreadable_pid.poll()) + + def test_reap_claimed_of_dead_client(self) -> None: + orphan = self._live_process() + dead_client = 2**31 - 1 + ours = self._live_process() + with _non_listening_socket() as port: + self._write_state( + self._directory.claimed_path(dead_client, "0a0a"), + self._server_data(port, orphan.pid), + ) + self._write_state( + self._directory.claimed_path(os.getpid(), "0b0b"), + self._server_data(port, ours.pid), + ) + self._write_state( + self._directory.claimed_path(os.getpid(), "0c0c"), + self._server_data(port, 2**31 - 1), + ) + with self._directory: + self._pool.janitor() + # The orphaned claim and our dead server are retired; our live claim is untouched. + self.assertEqual(set(self._states("0a0a")), {"retired"}) + self.assertTrue(_wait_proc_dead(orphan), "the orphaned server was not stopped") + self.assertEqual(set(self._states("0b0b")), {"claimed"}) + self.assertIsNone(ours.poll()) + self.assertEqual(set(self._states("0c0c")), {"retired"}) + + def test_reap_claimed_detects_a_reused_client_pid(self) -> None: + server = self._live_process() + data = self._server_data(12345, server.pid) + data["client_process_start_id"] = "a-different-process-generation" + self._write_state(self._directory.claimed_path(os.getpid(), "bad5"), data) + + with self._directory: + self.assertFalse(self._pool.reap("bad5")) + + self.assertEqual(set(self._states("bad5")), {"retired"}) + self.assertTrue(_wait_proc_dead(server), "the orphaned server was not stopped") + + def test_reap_claimed_detects_a_reused_server_pid(self) -> None: + unrelated = self._stubborn_process() + data = self._server_data( + 12345, + unrelated.pid, + process_start_id="a-different-process-generation", + ) + self._write_state(self._directory.claimed_path(os.getpid(), "bad6"), data) + + with self._directory: + self.assertFalse(self._pool.reap("bad6")) + + self.assertEqual(set(self._states("bad6")), {"retired"}) + self.assertIsNone(unrelated.poll()) + + def test_reap_retired_escalates_to_sigkill(self) -> None: + with self.subTest("a fresh retirement is left to shut down gracefully"): + fresh = self._stubborn_process() + self._write_state( + self._directory.retired_path("f2e5"), + self._retired_data(fresh.pid, time.time()), + ) + with self._directory: + self._pool.reap("f2e5") + self.assertEqual(set(self._states("f2e5")), {"retired"}) + self.assertIsNone(fresh.poll()) + with self.subTest("a hung shutdown is hard-killed"): + stubborn = self._stubborn_process() + self._write_state( + self._directory.retired_path("a0a0"), + self._retired_data( + stubborn.pid, time.time() - ServerPool._RETIRE_KILL_AFTER_SECONDS - 1 + ), + ) + with self._directory: + self._pool.reap("a0a0") + self.assertTrue(_wait_proc_dead(stubborn), "SIGKILL escalation did not happen") + with self._directory: + self.assertTrue(self._pool.reap("a0a0")) + with self.subTest("a late first reaper hard-kills before giving up"): + abandoned = self._stubborn_process() + self._write_state( + self._directory.retired_path("ab4d"), + self._retired_data( + abandoned.pid, time.time() - ServerPool._RETIRE_GIVE_UP_AFTER_SECONDS - 1 + ), + ) + with self._directory: + self.assertTrue(self._pool.reap("ab4d")) + self.assertTrue(_wait_proc_dead(abandoned), "the abandoned server was not killed") + + def test_reap_repairs_malformed_retired_state(self) -> None: + server = self._stubborn_process() + malformed = self._retired_data(server.pid, "not-a-time") + malformed["signalled"] = True + path = self._write_state( + self._directory.retired_path("bad4"), + malformed, + ) + + with self._directory as directory: + self._pool.reap("bad4") + repaired = directory.read_json(path) + + assert repaired is not None + self.assertEqual(repaired["pid"], server.pid) + self.assertIsInstance(repaired["retired"], float) + self.assertTrue(repaired["signalled"]) + self.assertIsNone(server.poll()) + + def test_reap_keeps_retired_state_when_process_inspection_fails(self) -> None: + server = self._stubborn_process() + path = self._write_state( + self._directory.retired_path("bad9"), + self._retired_data( + server.pid, time.time() - ServerPool._RETIRE_GIVE_UP_AFTER_SECONDS - 1 + ), + ) + + with mock.patch.object(local_server_pool, "_is_local_connect_server", return_value=None): + with self._directory: + self.assertFalse(self._pool.reap("bad9")) + + self.assertEqual(set(self._states("bad9")), {"retired"}) + self.assertTrue(os.path.exists(path)) + self.assertIsNone(server.poll()) + + def test_reap_keeps_retired_state_without_a_process_identity(self) -> None: + server = self._stubborn_process() + original = { + "pid": server.pid, + "retired": time.time() - ServerPool._RETIRE_GIVE_UP_AFTER_SECONDS - 1, + } + path = self._write_state( + self._directory.retired_path("bad7"), + original, + ) + + with self._directory as directory: + self.assertFalse(self._pool.reap("bad7")) + self.assertFalse(self._pool.reap("bad7")) + retained = directory.read_json(path) + + self.assertEqual(retained, original) + self.assertIsNone(server.poll()) + + def test_reap_eventually_removes_state_without_a_process_handle(self) -> None: + uid = "fade" + self._write_state(self._directory.server_path(uid), {"malformed": True}) + + with self._directory as directory: + self.assertFalse(self._pool.reap(uid)) + states = self._directory.states(uid) + self.assertEqual(set(states), {"retired"}) + retired = directory.read_json(states["retired"]) + assert retired is not None + self.assertNotIn("pid", retired) + retired["retired"] = time.time() - ServerPool._RETIRE_GIVE_UP_AFTER_SECONDS - 1 + directory.write_json(states["retired"], retired) + self.assertTrue(self._pool.reap(uid)) + + self.assertFalse(self._states(uid)) + + def test_reap_keeps_a_daemon_pid_without_a_process_identity(self) -> None: + server = self._stubborn_process() + uid = "daed" + original = {"retired": time.time() - ServerPool._RETIRE_GIVE_UP_AFTER_SECONDS - 1} + path = self._write_state( + self._directory.retired_path(uid), + original, + ) + self._write_daemon_pid(uid, server.pid) + + with self._directory as directory: + self.assertFalse(self._pool.reap(uid)) + retained = directory.read_json(path) + + self.assertEqual(retained, original) + self.assertIsNone(server.poll()) + + server.kill() + self.assertTrue(_wait_proc_dead(server)) + with self._directory: + self.assertTrue(self._pool.reap(uid)) + self.assertFalse(os.path.exists(path)) + self.assertFalse(os.path.exists(self._directory.member_dir(uid))) + + def test_reap_prefers_the_record_pid_and_its_process_identity(self) -> None: + recorded_server = self._stubborn_process() + daemon_server = self._stubborn_process() + uid = "d00d" + path = self._write_state( + self._directory.retired_path(uid), + self._server_data(12345, recorded_server.pid), + ) + self._write_daemon_pid(uid, daemon_server.pid) + + with self._directory as directory: + self.assertFalse(self._pool.reap(uid)) + repaired = directory.read_json(path) + + assert repaired is not None + self.assertEqual(repaired["pid"], recorded_server.pid) + self.assertEqual( + repaired["process_start_id"], ServerPool._process_start_id(recorded_server.pid) + ) + self.assertIsNone(recorded_server.poll()) + self.assertIsNone(daemon_server.poll()) + + def test_reap_does_not_signal_a_reused_server_pid(self) -> None: + other_server = self._stubborn_process() + stale = self._retired_data( + other_server.pid, time.time() - ServerPool._RETIRE_KILL_AFTER_SECONDS - 1 + ) + stale["process_start_id"] = "a-different-process-generation" + self._write_state(self._directory.retired_path("bad8"), stale) + + with self._directory: + self.assertTrue(self._pool.reap("bad8")) + + self.assertIsNone(other_server.poll()) + + def test_reap_garbage_collects_old_unreferenced_member_directory(self) -> None: + old_dir = self._directory.member_dir("01d0") + fresh_dir = self._directory.member_dir("f2e5") + os.makedirs(old_dir) + os.makedirs(fresh_dir) + old = time.time() - 24 * 3600 - 1 + os.utime(old_dir, (old, old)) + + with self._directory: + self.assertTrue(self._pool.reap("01d0")) + self.assertFalse(self._pool.reap("f2e5")) + + self.assertFalse(os.path.exists(old_dir)) + self.assertTrue(os.path.isdir(fresh_dir)) + + def test_retire_survives_interrupted_state_rewrite(self) -> None: + server = self._stubborn_process() + with _non_listening_socket() as port: + server_path = self._write_state( + self._directory.server_path("c0de"), self._server_data(port, server.pid) + ) + with self._directory: + with mock.patch.object( + local_server_pool.os, + "replace", + side_effect=OSError("interrupted rewrite"), + ): + with self.assertRaisesRegex(OSError, "interrupted rewrite"): + process_start_id = ServerPool._process_start_id(server.pid) + assert process_start_id is not None + self._pool._retire(server_path, server.pid, process_start_id) + + states = self._states("c0de") + self.assertEqual(set(states), {"retired"}) + with self._directory as directory: + # The atomic rename preserved the old member payload. The next reaper recovers its + # pid, repairs the missing retirement timestamp, and keeps tracking the live JVM. + self._pool.reap("c0de") + repaired = directory.read_json(states["retired"]) + assert repaired is not None + self.assertEqual(repaired["pid"], server.pid) + self.assertIsInstance(repaired["retired"], float) + self.assertIsNone(server.poll()) + + def test_reap_retries_sigterm_only_until_delivered(self) -> None: + server = self._stubborn_process() + process_start_id = ServerPool._process_start_id(server.pid) + assert process_start_id is not None + state_path = self._write_state( + self._directory.server_path("fade"), + self._server_data(12345, server.pid), + ) + + with self._directory as directory: + with mock.patch.object( + ServerPool, + "_signal_server", + side_effect=[False, False, True], + ) as signal_server: + self._pool._retire(state_path, server.pid, process_start_id) + retired_path = self._directory.retired_path("fade") + initial = directory.read_json(retired_path) + assert initial is not None + self.assertFalse(initial["signalled"]) + + self.assertFalse(self._pool.reap("fade")) + failed_retry = directory.read_json(retired_path) + self.assertEqual(failed_retry, initial) + + self.assertFalse(self._pool.reap("fade")) + delivered = directory.read_json(retired_path) + assert delivered is not None + self.assertTrue(delivered["signalled"]) + self.assertEqual(delivered["retired"], initial["retired"]) + + self.assertFalse(self._pool.reap("fade")) + self.assertEqual(signal_server.call_count, 3) + + self.assertIsNone(server.poll()) + + def test_release_retires_the_claimed_member(self) -> None: + server = self._live_process() + with _non_listening_socket() as port: + server_data = self._server_data(port, server.pid) + claim_path = self._write_state( + self._directory.claimed_path(os.getpid(), "a1a1"), server_data + ) + member = PoolMember(server_data) + member.claim_path = claim_path + local_server_pool._claimed_member = member + + # Release must use the directory that owns the claim even if the override changes + # between acquisition and process-exit cleanup. + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = os.path.join(self._tmpdir, "other-pool") + local_server_pool.release_pooled_local_connect_server() + + self.assertIsNone(local_server_pool._claimed_member) + states = self._states("a1a1") + self.assertEqual(set(states), {"retired"}) + with self._directory as directory: + retired = directory.read_json(states["retired"]) + assert retired is not None + self.assertEqual(retired["pid"], server.pid) + self.assertTrue(retired["signalled"]) + self.assertTrue(_wait_proc_dead(server), "release did not stop the server") + # Releasing again is a no-op. + local_server_pool.release_pooled_local_connect_server() + + def test_release_does_not_signal_a_mismatched_process_identity(self) -> None: + server = self._live_process() + server_data = self._server_data( + 12345, + server.pid, + process_start_id="a-different-process-generation", + ) + claim_path = self._write_state( + self._directory.claimed_path(os.getpid(), "a1a4"), server_data + ) + member = PoolMember(server_data) + member.claim_path = claim_path + + self._pool.release(member) + + self.assertEqual(set(self._states("a1a4")), {"retired"}) + self.assertIsNone(server.poll()) + + def test_forked_child_does_not_release_its_parents_claim(self) -> None: + server = self._live_process() + with _non_listening_socket() as port: + server_data = self._server_data(port, server.pid) + parent_pid = os.getpid() + 1 + claim_path = self._write_state( + self._directory.claimed_path(parent_pid, "a1a3"), server_data + ) + member = PoolMember(server_data) + member.claim_path = claim_path + + self._pool.release(member) + + self.assertEqual(set(self._states("a1a3")), {"claimed"}) + self.assertIsNone(server.poll()) + + def test_release_retries_failures_and_tolerates_prior_retirement(self) -> None: + server = self._stubborn_process() + with _non_listening_socket() as port: + server_data = self._server_data(port, server.pid) + claim_path = self._write_state( + self._directory.claimed_path(os.getpid(), "a1a2"), server_data + ) + member = PoolMember(server_data) + member.claim_path = claim_path + local_server_pool._claimed_member = member + + with mock.patch.object( + PoolDirectory, "write_json", side_effect=OSError("interrupted retirement") + ): + with self.assertRaisesRegex(OSError, "interrupted retirement"): + local_server_pool.release_pooled_local_connect_server() + self.assertIs(local_server_pool._claimed_member, member) + self.assertEqual(set(self._states("a1a2")), {"retired"}) + + local_server_pool.release_pooled_local_connect_server() + + self.assertIsNone(local_server_pool._claimed_member) + self.assertEqual(set(self._states("a1a2")), {"retired"}) + # A concurrent janitor already completed the claim -> retired transition. + self._pool.release(member) + self.assertEqual(set(self._states("a1a2")), {"retired"}) + + def test_pool_size_parsing(self) -> None: + from pyspark.sql.connect.local_server_pool import _pool_size + + self.assertEqual(_pool_size({}), 2) + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "3"}), 3) + os.environ["SPARK_LOCAL_CONNECT_POOL_SIZE"] = "5" + self.assertEqual(_pool_size({}), 5) + # Junk falls back to the default; values below one are clamped up. + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "abc"}), 2) + self.assertEqual(_pool_size({"spark.local.connect.pool.size": float("inf")}), 2) + self.assertEqual(_pool_size({"spark.local.connect.pool.size": "0"}), 1) + + def test_refill_only_counts_matching_members(self) -> None: + from unittest import mock + + with self._directory as directory: + directory.write_json( + directory.server_path("a11"), self._server_data(1, os.getpid(), "my-fp") + ) + directory.write_json( + directory.pending_path("b22"), + {"attendant_pid": os.getpid(), "created": time.time(), "fingerprint": "other"}, + ) + with mock.patch.object(MemberAttendant, "spawn") as spawn: + self._pool.refill("local[2]", {}, "my-fp", target=2) + # One matching member exists (the other-fingerprint launch does not count), so one + # launch tops the pool up to the target of two. + self.assertEqual(spawn.call_count, 1) + + def test_acquire_returns_the_member_already_claimed_by_this_process(self) -> None: + member = PoolMember(self._server_data(15002, os.getpid())) + member.claim_path = "unused" + local_server_pool._claimed_member = member + url = local_server_pool.acquire_pooled_local_connect_server("local[2]", {}) + self.assertEqual(url, "sc://localhost:15002") + self.assertEqual(os.environ.get("SPARK_CONNECT_AUTHENTICATE_TOKEN"), "t") + + def test_acquire_requires_posix(self) -> None: + from unittest import mock + + from pyspark.errors import PySparkRuntimeError + + with mock.patch.object(os, "name", "nt"): + with self.assertRaises(PySparkRuntimeError) as ctx: + local_server_pool.acquire_pooled_local_connect_server("local[2]", {}) + self.assertIn("POSIX", str(ctx.exception)) + + + def test_purge_kills_everything_and_empties_the_directory(self) -> None: + warm = self._live_process() + attendant = self._attendant("b007") + half_started = self._stubborn_process() + duplicate_a = self._live_process() + duplicate_b = self._live_process() + retiring = self._stubborn_process() + with _non_listening_socket() as port: + self._write_state( + self._directory.server_path("a2a2"), self._server_data(port, warm.pid) + ) + self._write_state( + self._directory.pending_path("b007"), + {"attendant_pid": attendant.pid, "created": time.time(), "fingerprint": "fp"}, + ) + self._write_state(self._directory.conf_path("b007"), {"spark.foo": "bar"}) + os.makedirs(self._directory.member_dir("a2a2")) + self._write_daemon_pid("b007", half_started.pid) + # Purge is the corruption escape hatch, so malformed process metadata must not + # prevent it from clearing the rest of the pool. + self._write_state(self._directory.server_path("bad5"), {"pid": "not-a-pid"}) + self._write_state( + self._directory.claimed_path(111, "d00d"), + self._server_data(port, duplicate_a.pid), + ) + self._write_state( + self._directory.claimed_path(222, "d00d"), + self._server_data(port, duplicate_b.pid), + ) + self._write_state( + self._directory.retired_path("e7ed"), + {"pid": retiring.pid, "retired": time.time()}, + ) + + signalled = local_server_pool.purge_local_connect_pool() + + self.assertEqual(signalled, 6) + self.assertEqual(os.listdir(self._directory.path), [".lock"]) + self.assertTrue(_wait_proc_dead(warm)) + self.assertTrue(_wait_proc_dead(attendant)) + self.assertTrue(_wait_proc_dead(half_started)) + self.assertTrue(_wait_proc_dead(duplicate_a)) + self.assertTrue(_wait_proc_dead(duplicate_b)) + self.assertTrue(_wait_proc_dead(retiring)) + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires JVM access to start local Connect servers", +) +@unittest.skipUnless(os.name == "posix", "the pool relies on the POSIX sbin scripts") +class LocalConnectServerPoolE2ETests(unittest.TestCase): + """End-to-end tests that boot real pooled servers (slow).""" + + CLIENT = textwrap.dedent( + """ + import json + + from pyspark.sql import SparkSession + + spark = ( + SparkSession.builder.remote("local[2]") + .config("spark.local.connect.pool", "true") + .getOrCreate() + ) + from pyspark.sql.connect import local_server_pool + + try: + # Pool cleanup belongs only to the builder-created session that claimed the + # server. Stopping another session must leave the claimed server available. + spark.newSession().stop() + count = spark.range(1).count() + member = local_server_pool._claimed_member + print(json.dumps({"count": count, "server_pid": member.pid})) + finally: + spark.stop() + """ + ) + + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp() + self._pool_dir = os.path.join(self._tmpdir, "pool") + self._saved_env = {k: os.environ.get(k) for k in _SAVED_ENV_KEYS} + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = self._pool_dir + + def tearDown(self) -> None: + try: + local_server_pool.purge_local_connect_pool() + deadline = time.time() + 60 + while time.time() < deadline: + if set(os.listdir(self._pool_dir)) <= {".lock"}: + break + # Supervising attendants notice the purged directory and exit on their own; + # purge again in case one republished state in between. + local_server_pool.purge_local_connect_pool() + time.sleep(1) + finally: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _run_clients(self, n: int, pool_size: str) -> list: + env = dict(os.environ) + env["SPARK_LOCAL_CONNECT_POOL_DIR"] = self._pool_dir + env["SPARK_LOCAL_CONNECT_POOL"] = "1" + env["SPARK_LOCAL_CONNECT_POOL_SIZE"] = pool_size + env.pop("SPARK_CONNECT_AUTHENTICATE_TOKEN", None) + procs = [ + subprocess.Popen( + [sys.executable, "-c", self.CLIENT], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for _ in range(n) + ] + outputs = [] + try: + for proc in procs: + stdout, stderr = proc.communicate(timeout=300) + self.assertEqual(proc.returncode, 0, stderr) + lines = stdout.strip().splitlines() + self.assertTrue(lines, stderr) + outputs.append(json.loads(lines[-1])) + finally: + for proc in procs: + if proc.poll() is None: + proc.kill() + proc.communicate() + return outputs + + def test_sequential_runs_use_fresh_servers_and_tear_them_down(self) -> None: + first = self._run_clients(1, pool_size="1")[0] + self.assertEqual(first["count"], 1) + # The used server is torn down asynchronously after its run. + self.assertTrue( + _wait_pid_gone(first["server_pid"]), + "the server was not torn down after its run ended", + ) + second = self._run_clients(1, pool_size="1")[0] + self.assertEqual(second["count"], 1) + self.assertNotEqual( + first["server_pid"], second["server_pid"], "a pooled server was reused across runs" + ) + + def test_concurrent_cold_clients_get_distinct_servers(self) -> None: + outputs = self._run_clients(2, pool_size="2") + self.assertEqual({o["count"] for o in outputs}, {1}) + pids = [o["server_pid"] for o in outputs] + self.assertEqual(len(set(pids)), 2, "two concurrent runs shared a pooled server") + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/connect/test_connect_plan.py b/python/pyspark/sql/tests/connect/test_connect_plan.py index 5105a58b400b5..60f0add9a52a8 100644 --- a/python/pyspark/sql/tests/connect/test_connect_plan.py +++ b/python/pyspark/sql/tests/connect/test_connect_plan.py @@ -14,46 +14,55 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest -import uuid import datetime import decimal import math +import unittest +import uuid +from unittest.mock import MagicMock +from pyspark.errors import PySparkValueError from pyspark.testing.connectutils import ( PlanOnlyTestFixture, - should_test_connect, connect_requirement_message, + should_test_connect, ) -from pyspark.errors import PySparkValueError - -from unittest.mock import MagicMock if should_test_connect: import pyspark.sql.connect.proto as proto from pyspark.sql.connect.column import Column from pyspark.sql.connect.dataframe import DataFrame + from pyspark.sql.connect.expressions import LiteralExpression + from pyspark.sql.connect.functions import ( + bitmap_and, + bitmap_andnot, + bitmap_or, + bitmap_xor, + col, + lit, + max, + min, + sum, + ) + from pyspark.sql.connect.observation import Observation from pyspark.sql.connect.plan import ( - WriteOperation, - Read, - Join, - SetOperation, CollectMetrics, + Join, LogicalPlan, + Read, + SetOperation, + WriteOperation, ) - from pyspark.sql.connect.observation import Observation from pyspark.sql.connect.readwriter import DataFrameReader - from pyspark.sql.connect.expressions import LiteralExpression - from pyspark.sql.connect.functions import col, lit, max, min, sum from pyspark.sql.connect.types import pyspark_types_to_proto_types from pyspark.sql.types import ( - StringType, - StructType, - StructField, - IntegerType, - MapType, ArrayType, DoubleType, + IntegerType, + MapType, + StringType, + StructField, + StructType, ) @@ -71,6 +80,20 @@ def test_simple_project(self): self.assertIsNotNone(plan.root, "Root relation must be set") self.assertIsNotNone(plan.root.read) + def test_bitmap_scalar_functions(self): + df = self.connect.readTable(table_name=self.tbl_name) + plan = df.select( + bitmap_and(col("bytes"), col("bytes")), + bitmap_or(col("bytes"), col("bytes")), + bitmap_andnot(col("bytes"), col("bytes")), + bitmap_xor(col("bytes"), col("bytes")), + )._plan.to_proto(self.connect) + function_names = [ + expression.unresolved_function.function_name + for expression in plan.root.project.expressions + ] + self.assertEqual(function_names, ["bitmap_and", "bitmap_or", "bitmap_andnot", "bitmap_xor"]) + def test_join_using_columns(self): left_input = self.connect.readTable(table_name=self.tbl_name) right_input = self.connect.readTable(table_name=self.tbl_name) diff --git a/python/pyspark/sql/tests/connect/test_connect_readwriter.py b/python/pyspark/sql/tests/connect/test_connect_readwriter.py index 864ee822a3c16..5db76600a0dc5 100644 --- a/python/pyspark/sql/tests/connect/test_connect_readwriter.py +++ b/python/pyspark/sql/tests/connect/test_connect_readwriter.py @@ -21,25 +21,24 @@ import time from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql.tests.connect.test_connect_basic import SparkConnectSQLTestCase from pyspark.sql.types import ( - StructType, - StructField, - LongType, - StringType, - IntegerType, ArrayType, + IntegerType, + LongType, MapType, Row, + StringType, + StructField, + StructType, ) +from pyspark.testing.connectutils import should_test_connect from pyspark.testing.objects import ( - PythonOnlyUDT, ExamplePoint, PythonOnlyPoint, + PythonOnlyUDT, ) -from pyspark.testing.connectutils import should_test_connect -from pyspark.sql.tests.connect.test_connect_basic import SparkConnectSQLTestCase - if should_test_connect: from pyspark.sql.connect.readwriter import DataFrameWriterV2 @@ -390,7 +389,8 @@ def test_write_operations(self): def test_writeTo_operations(self): # SPARK-42002: Implement DataFrameWriterV2 import datetime - from pyspark.sql.connect.functions import col, years, months, days, hours, bucket + + from pyspark.sql.connect.functions import bucket, col, days, hours, months, years df = self.connect.createDataFrame( [(1, datetime.datetime(2000, 1, 1), "foo")], ("id", "ts", "value") diff --git a/python/pyspark/sql/tests/connect/test_connect_retry.py b/python/pyspark/sql/tests/connect/test_connect_retry.py index 9757e75798fd5..3bbc0ac55d751 100644 --- a/python/pyspark/sql/tests/connect/test_connect_retry.py +++ b/python/pyspark/sql/tests/connect/test_connect_retry.py @@ -19,12 +19,13 @@ from collections import defaultdict from pyspark.testing.connectutils import ( - should_test_connect, connect_requirement_message, + should_test_connect, ) if should_test_connect: import grpc + from pyspark.sql.connect.client.core import Retrying from pyspark.sql.connect.client.retries import RetryPolicy diff --git a/python/pyspark/sql/tests/connect/test_connect_session.py b/python/pyspark/sql/tests/connect/test_connect_session.py index 3b5a816b7d24f..e158c05443aa9 100644 --- a/python/pyspark/sql/tests/connect/test_connect_session.py +++ b/python/pyspark/sql/tests/connect/test_connect_session.py @@ -20,26 +20,27 @@ import uuid from typing import Optional -from pyspark.util import is_remote_only from pyspark.errors import PySparkException from pyspark.sql import SparkSession as PySparkSession from pyspark.testing.connectutils import ( - should_test_connect, ReusedConnectTestCase, connect_requirement_message, + should_test_connect, ) from pyspark.testing.utils import timeout +from pyspark.util import is_remote_only if should_test_connect: import grpc - from pyspark.sql.connect.session import SparkSession as RemoteSparkSession - from pyspark.sql.connect.client import ChannelBuilder, DefaultChannelBuilder + from pyspark.errors.exceptions.connect import ( AnalysisException, SparkConnectException, SparkConnectGrpcException, SparkUpgradeException, ) + from pyspark.sql.connect.client import ChannelBuilder, DefaultChannelBuilder + from pyspark.sql.connect.session import SparkSession as RemoteSparkSession class CustomChannelBuilder(DefaultChannelBuilder): @property @@ -82,6 +83,20 @@ def handler(**kwargs): self.spark.sql("select 1").collect() self.assertGreaterEqual(len(handler_called), 0) + @timeout(10) + def test_operation_id_in_execution_info_and_exception(self): + df = self.spark.sql("select 1") + df.collect() + self.assertIsNotNone(df.executionInfo) + operation_id = df.executionInfo.operation_id + self.assertIsNotNone(operation_id) + uuid.UUID(operation_id) + + with self.assertRaises(SparkConnectException) as error: + self.spark.sql("select raise_error('expected')").collect() + self.assertIsNotNone(error.exception.operation_id) + uuid.UUID(error.exception.operation_id) + def _check_no_active_session_error(self, e: PySparkException): self.check_error(exception=e, errorClass="NO_ACTIVE_SESSION", messageParameters=dict()) diff --git a/python/pyspark/sql/tests/connect/test_connect_stat.py b/python/pyspark/sql/tests/connect/test_connect_stat.py index 3d05cb7b4bb5c..827059653dc1f 100644 --- a/python/pyspark/sql/tests/connect/test_connect_stat.py +++ b/python/pyspark/sql/tests/connect/test_connect_stat.py @@ -17,16 +17,16 @@ from pyspark.errors import PySparkTypeError, PySparkValueError -from pyspark.testing.connectutils import should_test_connect from pyspark.sql.tests.connect.test_connect_basic import SparkConnectSQLTestCase +from pyspark.testing.connectutils import should_test_connect if should_test_connect: - from pyspark.sql import functions as SF - from pyspark.sql.connect import functions as CF from pyspark.errors.exceptions.connect import ( AnalysisException, SparkConnectException, ) + from pyspark.sql import functions as SF + from pyspark.sql.connect import functions as CF class SparkConnectStatTests(SparkConnectSQLTestCase): diff --git a/python/pyspark/sql/tests/connect/test_df_debug.py b/python/pyspark/sql/tests/connect/test_df_debug.py index 1e316831173c9..741c2d76f79fd 100644 --- a/python/pyspark/sql/tests/connect/test_df_debug.py +++ b/python/pyspark/sql/tests/connect/test_df_debug.py @@ -18,7 +18,7 @@ import unittest from pyspark.testing.connectutils import ReusedConnectTestCase -from pyspark.testing.utils import have_graphviz, graphviz_requirement_message +from pyspark.testing.utils import graphviz_requirement_message, have_graphviz class SparkConnectDataFrameDebug(ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py index 1f0369e19a6bf..4f25208711469 100644 --- a/python/pyspark/sql/tests/connect/test_parity_frame_plot.py +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.sql.tests.plot.test_frame_plot import DataFramePlotTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase class FramePlotParityTests(DataFramePlotTestsMixin, ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py index 10807e485d787..84e41c5a967ad 100644 --- a/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/connect/test_parity_frame_plot_plotly.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.sql.tests.plot.test_frame_plot_plotly import DataFramePlotPlotlyTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase class FramePlotPlotlyParityTests(DataFramePlotPlotlyTestsMixin, ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/connect/test_parity_memory_profiler.py b/python/pyspark/sql/tests/connect/test_parity_memory_profiler.py index 112af30cd1fdc..fde715a6e4f5f 100644 --- a/python/pyspark/sql/tests/connect/test_parity_memory_profiler.py +++ b/python/pyspark/sql/tests/connect/test_parity_memory_profiler.py @@ -17,8 +17,8 @@ import inspect import os -from pyspark.tests.test_memory_profiler import MemoryProfiler2TestsMixin, _do_computation from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.tests.test_memory_profiler import MemoryProfiler2TestsMixin, _do_computation class MemoryProfilerParityTests(MemoryProfiler2TestsMixin, ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/connect/test_parity_readwriter.py b/python/pyspark/sql/tests/connect/test_parity_readwriter.py index 08f12e392ec46..981b8ff35e695 100644 --- a/python/pyspark/sql/tests/connect/test_parity_readwriter.py +++ b/python/pyspark/sql/tests/connect/test_parity_readwriter.py @@ -16,7 +16,7 @@ # from pyspark.sql.tests.test_readwriter import ReadwriterTestsMixin, ReadwriterV2TestsMixin -from pyspark.testing.connectutils import should_test_connect, ReusedConnectTestCase +from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect if should_test_connect: from pyspark.sql.connect.readwriter import DataFrameWriterV2 diff --git a/python/pyspark/sql/tests/connect/test_parity_resources.py b/python/pyspark/sql/tests/connect/test_parity_resources.py index 0eb7c4b31338c..8d4c5f89fd1a2 100644 --- a/python/pyspark/sql/tests/connect/test_parity_resources.py +++ b/python/pyspark/sql/tests/connect/test_parity_resources.py @@ -16,8 +16,8 @@ # import os -from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.sql.tests.test_resources import ResourceProfileTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase class ResourceProfileParityTests(ResourceProfileTestsMixin, ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/connect/test_parity_udf.py b/python/pyspark/sql/tests/connect/test_parity_udf.py index 33ec11adcf747..d6e44759185de 100644 --- a/python/pyspark/sql/tests/connect/test_parity_udf.py +++ b/python/pyspark/sql/tests/connect/test_parity_udf.py @@ -17,19 +17,21 @@ import unittest -from pyspark.testing.connectutils import should_test_connect - -if should_test_connect: - from pyspark import sql - from pyspark.sql.connect.udf import UserDefinedFunction - - sql.udf.UserDefinedFunction = UserDefinedFunction - from pyspark.sql.tests.test_udf import BaseUDFTestsMixin from pyspark.testing.connectutils import ReusedConnectTestCase class UDFParityTests(BaseUDFTestsMixin, ReusedConnectTestCase): + @classmethod + def setUpClass(cls): + # test_udf uses UserDefinedFunction so we need to monkeypatch it + import pyspark.sql.tests.test_udf + from pyspark.sql.connect.udf import UserDefinedFunction + + pyspark.sql.tests.test_udf.UserDefinedFunction = UserDefinedFunction + + super().setUpClass() + @unittest.skip("Spark Connect does not support mapPartitions() but the test depends on it.") def test_worker_original_stdin_closed(self): super().test_worker_original_stdin_closed() diff --git a/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py b/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py new file mode 100644 index 0000000000000..67228b5c60fbf --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_parity_udf_in_higher_order_function.py @@ -0,0 +1,37 @@ +# +# 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. +# + +import unittest + +from pyspark.sql.tests.test_udf_in_higher_order_function import ( + UDFInHigherOrderFunctionTestsMixin, +) +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class UDFInHigherOrderFunctionParityTests( + UDFInHigherOrderFunctionTestsMixin, ReusedConnectTestCase +): + @unittest.skip("Asserts on the JVM optimized plan via _jdf, unavailable in Spark Connect.") + def test_lambda_without_udf_is_unchanged(self): + super().test_lambda_without_udf_is_unchanged() + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/connect/test_parity_udtf.py b/python/pyspark/sql/tests/connect/test_parity_udtf.py index 2ab0e733a4be3..62a91c4822aa8 100644 --- a/python/pyspark/sql/tests/connect/test_parity_udtf.py +++ b/python/pyspark/sql/tests/connect/test_parity_udtf.py @@ -17,24 +17,19 @@ import os import unittest -from pyspark.testing.connectutils import should_test_connect +from pyspark.sql.functions import lit, udtf from pyspark.sql.tests.test_udtf import ( BaseUDTFTestsMixin, - UDTFArrowTestsMixin, LegacyUDTFArrowTestsMixin, + UDTFArrowTestsMixin, ) -from pyspark.testing.connectutils import ReusedConnectTestCase +from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect if should_test_connect: - from pyspark import sql - from pyspark.sql.connect.udtf import UserDefinedTableFunction - - sql.udtf.UserDefinedTableFunction = UserDefinedTableFunction - from pyspark.sql.connect.functions import lit, udtf from pyspark.errors.exceptions.connect import ( + InvalidPlanInput, PickleException, PythonException, - InvalidPlanInput, ) diff --git a/python/pyspark/sql/tests/connect/test_parity_utils.py b/python/pyspark/sql/tests/connect/test_parity_utils.py index 521b6082cf222..692b735a5a59a 100644 --- a/python/pyspark/sql/tests/connect/test_parity_utils.py +++ b/python/pyspark/sql/tests/connect/test_parity_utils.py @@ -15,8 +15,8 @@ # limitations under the License. # -from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.sql.tests.test_utils import UtilsTestsMixin +from pyspark.testing.connectutils import ReusedConnectTestCase class UtilsParityTests(UtilsTestsMixin, ReusedConnectTestCase): diff --git a/python/pyspark/sql/tests/df_golden/df_golden.py b/python/pyspark/sql/tests/df_golden/df_golden.py index 69b9c77fc225c..7a0f3aa74e80e 100644 --- a/python/pyspark/sql/tests/df_golden/df_golden.py +++ b/python/pyspark/sql/tests/df_golden/df_golden.py @@ -92,7 +92,6 @@ import re from decimal import Decimal - _CASE_END = "!-- end" _SECTION_PREFIX = "--! " _FILE_METADATA_NAME = "__file_metadata__" diff --git a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py index c212f96113b21..864558c9e6bd0 100644 --- a/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py +++ b/python/pyspark/sql/tests/df_golden/scripts/group_by/bool_and_decimal_error.py @@ -1,6 +1,7 @@ # input type checking Decimal from decimal import Decimal + from pyspark.sql.functions import bool_and, lit df = spark.table("test_agg").select(bool_and(lit(Decimal("1.0")))) diff --git a/python/pyspark/sql/tests/df_golden/test_df_golden.py b/python/pyspark/sql/tests/df_golden/test_df_golden.py index 311491c5d7ed2..b71e945235aab 100644 --- a/python/pyspark/sql/tests/df_golden/test_df_golden.py +++ b/python/pyspark/sql/tests/df_golden/test_df_golden.py @@ -49,9 +49,8 @@ import os -from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.sql.tests.df_golden.df_golden import run_golden_test - +from pyspark.testing.connectutils import ReusedConnectTestCase _THIS_DIR = os.path.dirname(os.path.abspath(__file__)) diff --git a/python/pyspark/sql/tests/pandas/bench_arrow_columnar_udf.py b/python/pyspark/sql/tests/pandas/bench_arrow_columnar_udf.py index b788f90bec263..549300595e8ef 100644 --- a/python/pyspark/sql/tests/pandas/bench_arrow_columnar_udf.py +++ b/python/pyspark/sql/tests/pandas/bench_arrow_columnar_udf.py @@ -38,8 +38,8 @@ """ import argparse -import sys import os +import sys import time # Allow running from the Spark root directory. @@ -48,8 +48,7 @@ import pandas as pd from pyspark.sql import SparkSession -from pyspark.sql.functions import pandas_udf, col - +from pyspark.sql.functions import col, pandas_udf ARROW_SOURCE = "org.apache.spark.sql.execution.python.ArrowBackedDataSourceV2" diff --git a/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py b/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py index 5deed766cbce6..88c2c229d0571 100644 --- a/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py +++ b/python/pyspark/sql/tests/pandas/bench_pipelined_udf.py @@ -45,7 +45,6 @@ import subprocess import sys - SPARK_HOME = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../../..") PIPELINED_CONF = "spark.python.udf.pipelined.enabled" QUEUE_DEPTH_CONF = "spark.python.udf.pipelined.queueDepth" diff --git a/python/pyspark/sql/tests/pandas/helper/helper_pandas_transform_with_state.py b/python/pyspark/sql/tests/pandas/helper/helper_pandas_transform_with_state.py index 6a5b65ebf9586..aa8918ba7fa5d 100644 --- a/python/pyspark/sql/tests/pandas/helper/helper_pandas_transform_with_state.py +++ b/python/pyspark/sql/tests/pandas/helper/helper_pandas_transform_with_state.py @@ -15,29 +15,30 @@ # limitations under the License. # -from abc import abstractmethod import sys +import unittest +from abc import abstractmethod from typing import ( Iterator, NamedTuple, Optional, ) -import unittest + from pyspark.errors import PySparkRuntimeError from pyspark.sql.streaming import StatefulProcessor, StatefulProcessorHandle from pyspark.sql.types import ( - StringType, - StructType, - StructField, - Row, - IntegerType, - TimestampType, - LongType, + ArrayType, BooleanType, - FloatType, DoubleType, - ArrayType, + FloatType, + IntegerType, + LongType, MapType, + Row, + StringType, + StructField, + StructType, + TimestampType, ) from pyspark.testing.utils import have_pandas @@ -1808,6 +1809,7 @@ def handleInputRows(self, key, rows, timerValues) -> Iterator[pd.DataFrame]: attributes_map, confs_map = self._update_map_state(key, total_temperature) import json + import numpy as np def np_int64_to_int(x): diff --git a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state.py b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state.py index e1003ddfbddb9..1d9701e1cacdb 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state.py @@ -15,30 +15,48 @@ # limitations under the License. # -from abc import abstractmethod - import json import os -import time import tempfile -from pyspark.sql.streaming import StatefulProcessor - +import time import unittest +from abc import abstractmethod from typing import cast from pyspark import SparkConf from pyspark.sql.functions import split +from pyspark.sql.streaming import StatefulProcessor +from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import ( + AddFieldsProcessorFactory, + BasicProcessorFactory, + BasicProcessorNotNullableFactory, + ChunkCountProcessorFactory, + ChunkCountProcessorWithInitialStateFactory, + CompositeOutputProcessorFactory, + EventTimeStatefulProcessorFactory, + LargeValueStatefulProcessorFactory, + MapStateProcessorFactory, + MinEventTimeStatefulProcessorFactory, + ProcTimeStatefulProcessorFactory, + RemoveFieldsProcessorFactory, + ReorderedFieldsProcessorFactory, + SimpleStatefulProcessorFactory, + SimpleStatefulProcessorWithInitialStateFactory, + StatefulProcessorChainingOpsFactory, + StatefulProcessorCompositeTypeFactory, + UpcastProcessorFactory, +) from pyspark.sql.types import ( + ArrayType, + DecimalType, + DoubleType, + IntegerType, + MapType, + Row, StringType, - StructType, StructField, - Row, - IntegerType, + StructType, TimestampType, - DecimalType, - ArrayType, - MapType, - DoubleType, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( @@ -48,27 +66,6 @@ pyarrow_requirement_message, ) -from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import ( - SimpleStatefulProcessorWithInitialStateFactory, - EventTimeStatefulProcessorFactory, - ProcTimeStatefulProcessorFactory, - SimpleStatefulProcessorFactory, - StatefulProcessorChainingOpsFactory, - MapStateProcessorFactory, - LargeValueStatefulProcessorFactory, - BasicProcessorFactory, - BasicProcessorNotNullableFactory, - AddFieldsProcessorFactory, - RemoveFieldsProcessorFactory, - ReorderedFieldsProcessorFactory, - UpcastProcessorFactory, - MinEventTimeStatefulProcessorFactory, - StatefulProcessorCompositeTypeFactory, - ChunkCountProcessorFactory, - ChunkCountProcessorWithInitialStateFactory, - CompositeOutputProcessorFactory, -) - class TransformWithStateTestsMixin: @classmethod diff --git a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_checkpoint_v2.py b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_checkpoint_v2.py index 6b822f2bc664f..80b32b4cf444f 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_checkpoint_v2.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_checkpoint_v2.py @@ -16,10 +16,10 @@ # -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state import ( TransformWithStateInPandasTestsMixin, ) +from pyspark.testing.sqlutils import ReusedSQLTestCase class TransformWithStateInPandasWithCheckpointV2TestsMixin(TransformWithStateInPandasTestsMixin): diff --git a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py index 01959ed02eab3..a633727473426 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py @@ -15,25 +15,37 @@ # limitations under the License. # -from abc import abstractmethod - import json import os -import time import tempfile -from pyspark.sql.streaming import StatefulProcessor - +import time import unittest +from abc import abstractmethod from typing import cast from pyspark import SparkConf from pyspark.sql.functions import array_sort, col, explode, split +from pyspark.sql.streaming import StatefulProcessor +from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import ( + InvalidSimpleStatefulProcessorFactory, + ListStateLargeListProcessorFactory, + ListStateLargeTTLProcessorFactory, + ListStateProcessorFactory, + MapStateLargeTTLProcessorFactory, + MapStateProcessorFactory, + SimpleStatefulProcessorFactory, + SimpleStatefulProcessorWithInitialStateFactory, + SimpleTTLStatefulProcessorFactory, + StatefulProcessorWithInitialStateTimersFactory, + StatefulProcessorWithListStateInitialStateFactory, + TTLStatefulProcessorFactory, +) from pyspark.sql.types import ( + IntegerType, + Row, StringType, - StructType, StructField, - Row, - IntegerType, + StructType, ) from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -44,21 +56,6 @@ pyarrow_requirement_message, ) -from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import ( - SimpleStatefulProcessorWithInitialStateFactory, - StatefulProcessorWithInitialStateTimersFactory, - StatefulProcessorWithListStateInitialStateFactory, - SimpleStatefulProcessorFactory, - SimpleTTLStatefulProcessorFactory, - TTLStatefulProcessorFactory, - InvalidSimpleStatefulProcessorFactory, - ListStateProcessorFactory, - ListStateLargeListProcessorFactory, - ListStateLargeTTLProcessorFactory, - MapStateProcessorFactory, - MapStateLargeTTLProcessorFactory, -) - class TransformWithStateStateVariableTestsMixin: @classmethod @@ -1006,7 +1003,6 @@ class TransformWithStateInPandasStateVariableTests( if __name__ == "__main__": from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state import * # noqa: F403 - from pyspark.testing import main main() diff --git a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable_checkpoint_v2.py b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable_checkpoint_v2.py index 4ed1f63fd4051..d4c371fc3ee82 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable_checkpoint_v2.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable_checkpoint_v2.py @@ -16,10 +16,10 @@ # -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state_state_variable import ( TransformWithStateInPandasStateVariableTestsMixin, ) +from pyspark.testing.sqlutils import ReusedSQLTestCase class TransformWithStateInPandasStateVariableWithCheckpointV2TestsMixin( diff --git a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state.py b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state.py index 995e7f5b2fabd..1fb178e632edc 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state.py @@ -17,15 +17,16 @@ import os import unittest + from pyspark import SparkConf +from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state import ( + TransformWithStateTestsMixin, +) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pyarrow, pyarrow_requirement_message, ) -from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state import ( - TransformWithStateTestsMixin, -) @unittest.skipIf( diff --git a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_checkpoint_v2.py b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_checkpoint_v2.py index bec2eab2a1117..c8510e10a7b90 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_checkpoint_v2.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_checkpoint_v2.py @@ -16,10 +16,10 @@ # -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.tests.pandas.streaming.test_transform_with_state import ( TransformWithStateInPySparkTestsMixin, ) +from pyspark.testing.sqlutils import ReusedSQLTestCase class TransformWithStateInPySparkWithCheckpointV2TestsMixin(TransformWithStateInPySparkTestsMixin): diff --git a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable.py b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable.py index 16edb07c2f31e..1bd7bc3c23fa8 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable.py @@ -17,17 +17,17 @@ import os import unittest + from pyspark import SparkConf +from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state_state_variable import ( + TransformWithStateStateVariableTestsMixin, +) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pyarrow, pyarrow_requirement_message, ) -from pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state_state_variable import ( - TransformWithStateStateVariableTestsMixin, -) - @unittest.skipIf( not have_pyarrow or os.environ.get("PYTHON_GIL", "?") == "0", diff --git a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable_checkpoint_v2.py b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable_checkpoint_v2.py index 9208de303d2cc..ce7bd7e45e7f4 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable_checkpoint_v2.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable_checkpoint_v2.py @@ -16,10 +16,10 @@ # -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql.tests.pandas.streaming.test_transform_with_state_state_variable import ( TransformWithStateInPySparkStateVariableTestsMixin, ) +from pyspark.testing.sqlutils import ReusedSQLTestCase class TransformWithStateInPySparkStateVariableWithCheckpointV2TestsMixin( diff --git a/python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py b/python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py index 3f7c9bdc72488..631ffeff4ae44 100644 --- a/python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py +++ b/python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py @@ -24,6 +24,8 @@ import pandas.testing as pdt from pyspark import SparkConf +from pyspark.errors import PySparkAssertionError, PySparkValueError +from pyspark.errors.exceptions.base import IllegalArgumentException from pyspark.sql import DataFrame from pyspark.sql.functions import split from pyspark.sql.streaming import StatefulProcessor, TwsTester @@ -45,8 +47,6 @@ StructField, StructType, ) -from pyspark.errors import PySparkValueError, PySparkAssertionError -from pyspark.errors.exceptions.base import IllegalArgumentException from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, @@ -585,11 +585,12 @@ def test_pandas_dtype_preservation(self): self.assertEqual(result["count"].dtype, "int64") def test_processor_init_called_once(self): + from typing import Iterator + from pyspark.sql.streaming.stateful_processor import ( StatefulProcessor, StatefulProcessorHandle, ) - from typing import Iterator init_call_count = [0] diff --git a/python/pyspark/sql/tests/pandas/test_converter.py b/python/pyspark/sql/tests/pandas/test_converter.py index 3fbe4109fd159..fe0cdbfcd1560 100644 --- a/python/pyspark/sql/tests/pandas/test_converter.py +++ b/python/pyspark/sql/tests/pandas/test_converter.py @@ -16,13 +16,14 @@ # import unittest + from pyspark.sql.types import ( ArrayType, IntegerType, MapType, + Row, StringType, StructType, - Row, ) from pyspark.testing.utils import ( have_pandas, @@ -32,10 +33,10 @@ ) if have_pandas: - import pandas as pd import numpy as np - + import pandas as pd from pandas.testing import assert_series_equal + from pyspark.sql.pandas.types import _create_converter_from_pandas, _create_converter_to_pandas if have_pyarrow: diff --git a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py index b5a4daf3279cf..ef6d7de93c9bf 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map.py @@ -17,6 +17,7 @@ import unittest +from pyspark.errors import IllegalArgumentException, PythonException from pyspark.loose_version import LooseVersion from pyspark.sql import functions as sf from pyspark.sql.functions import pandas_udf, udf @@ -24,12 +25,11 @@ ArrayType, DoubleType, LongType, - StructType, + Row, StructField, + StructType, YearMonthIntervalType, - Row, ) -from pyspark.errors import IllegalArgumentException, PythonException from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, diff --git a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py index ed03cedb09716..c475f8084206e 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_cogrouped_map_misc.py @@ -15,8 +15,8 @@ # limitations under the License. # -import unittest import logging +import unittest from pyspark.sql import functions as sf from pyspark.sql.types import Row diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py index 92fb2d0cb06a8..833c5083be296 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py @@ -16,36 +16,36 @@ # import datetime -import unittest import logging import os - +import unittest from collections import OrderedDict from decimal import Decimal -from typing import Iterator, Tuple, Any +from typing import Any, Iterator, Tuple +from pyspark.errors import PySparkTypeError, PySparkValueError, PythonException from pyspark.loose_version import LooseVersion -from pyspark.sql import Row, functions as sf -from pyspark.sql.functions import udf, pandas_udf, PandasUDFType +from pyspark.sql import Row +from pyspark.sql import functions as sf +from pyspark.sql.functions import PandasUDFType, pandas_udf, udf from pyspark.sql.types import ( - IntegerType, - DoubleType, ArrayType, BinaryType, + BooleanType, ByteType, - LongType, DecimalType, - ShortType, + DoubleType, FloatType, + IntegerType, + LongType, + MapType, + NullType, + ShortType, StringType, - BooleanType, - StructType, StructField, - NullType, - MapType, + StructType, YearMonthIntervalType, ) -from pyspark.errors import PythonException, PySparkTypeError, PySparkValueError from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( assertDataFrameEqual, @@ -243,7 +243,8 @@ def test_register_grouped_map_udf(self): "SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, " "SQL_SCALAR_PANDAS_ITER_UDF, SQL_SCALAR_ARROW_ITER_UDF, " "SQL_GROUPED_AGG_PANDAS_UDF, SQL_GROUPED_AGG_ARROW_UDF, " - "SQL_GROUPED_AGG_PANDAS_ITER_UDF or SQL_GROUPED_AGG_ARROW_ITER_UDF" + "SQL_GROUPED_AGG_PANDAS_ITER_UDF, SQL_GROUPED_AGG_ARROW_ITER_UDF " + "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF" }, ) diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py index 7d536e9a91752..38f0933c35302 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py @@ -20,18 +20,17 @@ import string import sys import tempfile - import unittest from decimal import Decimal -from pyspark.sql.streaming.state import GroupStateTimeout, GroupState +from pyspark.sql.streaming.state import GroupState, GroupStateTimeout from pyspark.sql.types import ( + DecimalType, LongType, + Row, StringType, - StructType, StructField, - Row, - DecimalType, + StructType, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( diff --git a/python/pyspark/sql/tests/pandas/test_pandas_map.py b/python/pyspark/sql/tests/pandas/test_pandas_map.py index bfcedc6c8899f..b6cad5b7328e2 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_map.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_map.py @@ -14,17 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging import os import shutil import tempfile import time import unittest -import logging +from pyspark.errors import PythonException from pyspark.loose_version import LooseVersion from pyspark.sql import Row from pyspark.sql.functions import col, encode, lit -from pyspark.errors import PythonException from pyspark.sql.session import SparkSession from pyspark.sql.types import StructType from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -86,11 +86,39 @@ def test_map_in_pandas(self): expected = df.collect() self.assertEqual(actual, expected) - # test returning list of DataFrames - df = self.spark.range(10, numPartitions=3) - actual = df.mapInPandas(lambda it: [pdf for pdf in it], "id long").collect() - expected = df.collect() - self.assertEqual(actual, expected) + def test_map_in_pandas_legacy_accept_any_iterable(self): + # With the legacy flag enabled, returning a non-Iterator iterable (e.g. list) is accepted. + with self.sql_conf( + {"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled": True} + ): + df = self.spark.range(10, numPartitions=3) + actual = df.mapInPandas(lambda it: [pdf for pdf in it], "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) + + def test_map_in_pandas_legacy_accept_sequence_protocol(self): + # A sequence-protocol object (implements __getitem__ but not __iter__) is iterable via + # iter(...) even though it is not a collections.abc.Iterable, so the legacy flag must + # accept it too. + class SequenceOnly: + def __init__(self, items): + self._items = items + + def __getitem__(self, index): + return self._items[index] + + self.assertFalse(hasattr(SequenceOnly([]), "__iter__")) + + def returns_sequence(iterator): + return SequenceOnly([pdf for pdf in iterator]) + + with self.sql_conf( + {"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled": True} + ): + df = self.spark.range(10, numPartitions=3) + actual = df.mapInPandas(returns_sequence, "id long").collect() + expected = df.collect() + self.assertEqual(actual, expected) def test_multiple_columns(self): data = [(1, "foo"), (2, None), (3, "bar"), (4, "bar")] @@ -186,6 +214,10 @@ def no_iter(_): def bad_iter_elem(_): return iter([1]) + def list_not_iter(iterator): + # Iterable but not an Iterator: violates the Iterator[pandas.DataFrame] contract. + return [pdf for pdf in iterator] + with self.assertRaisesRegex( PythonException, "Return type of the user-defined function should be iterator of pandas.DataFrame, " @@ -200,6 +232,13 @@ def bad_iter_elem(_): ): (self.spark.range(10, numPartitions=3).mapInPandas(bad_iter_elem, "a int").count()) + with self.assertRaisesRegex( + PythonException, + "Return type of the user-defined function should be iterator of pandas.DataFrame, " + "but is list", + ): + (self.spark.range(10, numPartitions=3).mapInPandas(list_not_iter, "a int").count()) + def test_dataframes_with_other_column_names(self): with self.quiet(): self.check_dataframes_with_other_column_names() @@ -459,8 +498,18 @@ def func(iterator): def test_map_in_pandas_with_barrier_mode(self): df = self.spark.range(10) + def func0(iterator): + from pyspark import BarrierTaskContext + + BarrierTaskContext.get() + for batch in iterator: + yield batch + + with self.assertRaisesRegex(PythonException, "\\[NOT_IN_BARRIER_STAGE\\]"): + df.mapInPandas(func0, "id long", False).collect() + def func1(iterator): - from pyspark import TaskContext, BarrierTaskContext + from pyspark import BarrierTaskContext, TaskContext tc = TaskContext.get() assert tc is not None @@ -471,7 +520,7 @@ def func1(iterator): df.mapInPandas(func1, "id long", False).collect() def func2(iterator): - from pyspark import TaskContext, BarrierTaskContext + from pyspark import BarrierTaskContext, TaskContext tc = TaskContext.get() assert tc is not None diff --git a/python/pyspark/sql/tests/pandas/test_pandas_sqlmetrics.py b/python/pyspark/sql/tests/pandas/test_pandas_sqlmetrics.py index d4739c76bb9f1..212027b2033ba 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_sqlmetrics.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_sqlmetrics.py @@ -16,6 +16,7 @@ # import unittest + from pyspark.sql.functions import pandas_udf from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf.py b/python/pyspark/sql/tests/pandas/test_pandas_udf.py index ae9d231d91367..353b0b8ff49c7 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf.py @@ -15,20 +15,19 @@ # limitations under the License. # -import unittest import datetime +import unittest -from pyspark.sql.functions import udf, pandas_udf, PandasUDFType, assert_true, lit +from pyspark.errors import ParseException, PySparkTypeError, PythonException +from pyspark.sql.functions import PandasUDFType, assert_true, lit, pandas_udf, udf from pyspark.sql.types import ( + DayTimeIntervalType, DoubleType, - StructType, - StructField, LongType, - DayTimeIntervalType, + StructField, + StructType, VariantType, ) -from pyspark.errors import ParseException, PythonException, PySparkTypeError -from pyspark.util import PythonEvalType from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, @@ -36,6 +35,7 @@ pandas_requirement_message, pyarrow_requirement_message, ) +from pyspark.util import PythonEvalType @unittest.skipIf( @@ -313,8 +313,8 @@ def foo(x): ) def test_pandas_udf_detect_unsafe_type_conversion(self): - import pandas as pd import numpy as np + import pandas as pd values = [1.0] * 3 pdf = pd.DataFrame({"A": values}) @@ -352,9 +352,10 @@ def udf(column): df.withColumn("udf", udf("id")).collect() def test_pandas_udf_int_to_decimal_coercion(self): - import pandas as pd from decimal import Decimal + import pandas as pd + df = self.spark.range(0, 3) @pandas_udf(returnType="decimal(10,2)") diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py index 47decf731a77b..fa6200e6e2e17 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py @@ -15,24 +15,24 @@ # limitations under the License. # -import unittest import logging +import unittest from typing import Iterator, Tuple -from pyspark.util import PythonEvalType, is_remote_only -from pyspark.sql import Row, functions as sf +from pyspark.errors import AnalysisException, PySparkNotImplementedError, PythonException +from pyspark.sql import Row +from pyspark.sql import functions as sf from pyspark.sql.functions import ( + PandasUDFType, array, - explode, col, + explode, lit, mean, - udf, pandas_udf, - PandasUDFType, + udf, ) from pyspark.sql.types import ArrayType, YearMonthIntervalType -from pyspark.errors import AnalysisException, PySparkNotImplementedError, PythonException from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( assertDataFrameEqual, @@ -41,6 +41,7 @@ pandas_requirement_message, pyarrow_requirement_message, ) +from pyspark.util import PythonEvalType, is_remote_only if have_pandas: import pandas as pd @@ -752,9 +753,10 @@ def biased_sum(v, w=None): ) def test_arrow_cast_enabled_numeric_to_decimal(self): - import numpy as np from decimal import Decimal + import numpy as np + columns = [ "int8", "int16", diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py index 72d9fa566deef..2e0d46bedcee0 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_scalar.py @@ -14,53 +14,52 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging import os import random import shutil import tempfile import time import unittest -import logging from datetime import date, datetime from decimal import Decimal from pyspark import TaskContext -from pyspark.util import PythonEvalType, is_remote_only +from pyspark.errors import AnalysisException, PythonException from pyspark.sql import Column, Row from pyspark.sql.functions import ( + PandasUDFType, array, col, expr, lit, - sum, - struct, - udf, pandas_udf, + struct, + sum, to_json, - PandasUDFType, + udf, ) from pyspark.sql.types import ( - IntegerType, - ByteType, - StructType, - ShortType, + ArrayType, + BinaryType, BooleanType, - LongType, - FloatType, - DoubleType, + ByteType, + DateType, DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + MapType, + ShortType, StringType, - ArrayType, StructField, + StructType, TimestampType, - MapType, - DateType, - BinaryType, - YearMonthIntervalType, VariantType, VariantVal, + YearMonthIntervalType, ) -from pyspark.errors import AnalysisException, PythonException from pyspark.testing.sqlutils import ( ReusedSQLTestCase, test_compiled, @@ -73,6 +72,7 @@ pandas_requirement_message, pyarrow_requirement_message, ) +from pyspark.util import PythonEvalType, is_remote_only if have_pandas: import pandas as pd @@ -678,7 +678,7 @@ def check_vectorized_udf_invalid_length(self): df = self.spark.range(10) raise_exception = pandas_udf(lambda _: pd.Series(1), LongType()) with self.assertRaisesRegex( - Exception, "Result vector from pandas_udf was not the required length" + Exception, "The number of output rows.*must match the number of input rows" ): df.select(raise_exception(col("id"))).collect() @@ -704,6 +704,18 @@ def iter_udf_not_reading_all_input(it): with self.assertRaisesRegex(Exception, "The input iterator must be fully consumed"): df1.select(iter_udf_not_reading_all_input(col("id"))).collect() + @pandas_udf(LongType(), PandasUDFType.SCALAR_ITER) + def iter_udf_too_many_output_rows(it): + for batch in it: + yield pd.Series([1] * (len(batch) + 1)) + + with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": 3}): + df1 = self.spark.range(10).repartition(1) + with self.assertRaisesRegex( + Exception, "The number of output rows must not exceed the number of input rows" + ): + df1.select(iter_udf_too_many_output_rows(col("id"))).collect() + def test_vectorized_udf_chained(self): df = self.spark.range(10) scalar_f = pandas_udf(lambda x: x + 1, LongType()) diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py index cdeb12d5db4ed..fadaef3dac15d 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py @@ -16,9 +16,12 @@ # import unittest from inspect import signature -from typing import Union, Iterator, Tuple, get_type_hints +from typing import Iterator, Tuple, Union, get_type_hints -from pyspark.sql.functions import mean, lit +from pyspark.sql import Row +from pyspark.sql.functions import lit, mean +from pyspark.sql.pandas.functions import PandasUDFType, pandas_udf +from pyspark.sql.pandas.typehints import infer_eval_type, infer_group_pandas_eval_type from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, @@ -26,14 +29,11 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.sql.pandas.typehints import infer_eval_type, infer_group_pandas_eval_type -from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType -from pyspark.sql import Row from pyspark.util import PythonEvalType if have_pandas: - import pandas as pd import numpy as np + import pandas as pd from pandas.testing import assert_frame_equal diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py index 94dcaa3b6caf6..339367501c15e 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_typehints_with_future_annotations.py @@ -18,9 +18,12 @@ import unittest from inspect import signature -from typing import Union, Iterator, Tuple, get_type_hints +from typing import Iterator, Tuple, Union, get_type_hints -from pyspark.sql.functions import mean, lit +from pyspark.sql import Row +from pyspark.sql.functions import lit, mean +from pyspark.sql.pandas.functions import PandasUDFType, pandas_udf +from pyspark.sql.pandas.typehints import infer_eval_type from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, @@ -28,13 +31,10 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.sql.pandas.typehints import infer_eval_type -from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType -from pyspark.sql import Row if have_pandas: - import pandas as pd import numpy as np + import pandas as pd from pandas.testing import assert_frame_equal diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py b/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py index dda068f253052..e112378e35cc1 100644 --- a/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py +++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_window.py @@ -15,21 +15,22 @@ # limitations under the License. # -import unittest import logging +import unittest from decimal import Decimal from pyspark.errors import AnalysisException, PythonException +from pyspark.sql import Row from pyspark.sql import functions as sf -from pyspark.sql.functions import udf, pandas_udf, PandasUDFType -from pyspark.sql.window import Window +from pyspark.sql.functions import PandasUDFType, pandas_udf, udf from pyspark.sql.types import ( DecimalType, + DoubleType, + FloatType, IntegerType, LongType, - FloatType, - DoubleType, ) +from pyspark.sql.window import Window from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( assertDataFrameEqual, @@ -38,7 +39,6 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.sql import Row from pyspark.util import is_remote_only if have_pandas: diff --git a/python/pyspark/sql/tests/pandas/test_pipelined_udf.py b/python/pyspark/sql/tests/pandas/test_pipelined_udf.py index 1e133f6e219b9..ad194070b6d84 100644 --- a/python/pyspark/sql/tests/pandas/test_pipelined_udf.py +++ b/python/pyspark/sql/tests/pandas/test_pipelined_udf.py @@ -31,8 +31,8 @@ DoubleType, LongType, StringType, - StructType, StructField, + StructType, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( diff --git a/python/pyspark/sql/tests/plot/test_frame_plot.py b/python/pyspark/sql/tests/plot/test_frame_plot.py index a75ec50ec2910..f0aaf09b29acc 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot.py @@ -16,14 +16,15 @@ # import unittest + from pyspark.errors import PySparkValueError from pyspark.sql import Row from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_plotly, - plotly_requirement_message, have_pandas, + have_plotly, pandas_requirement_message, + plotly_requirement_message, ) if have_plotly and have_pandas: diff --git a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py index 232e242fb085e..af213787f68b5 100644 --- a/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py +++ b/python/pyspark/sql/tests/plot/test_frame_plot_plotly.py @@ -21,10 +21,10 @@ from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( - have_plotly, - plotly_requirement_message, have_pandas, + have_plotly, pandas_requirement_message, + plotly_requirement_message, ) if have_plotly and have_pandas: diff --git a/python/pyspark/sql/tests/streaming/kafka_utils.py b/python/pyspark/sql/tests/streaming/kafka_utils.py index cfa5ea739bab2..75f0791677c24 100644 --- a/python/pyspark/sql/tests/streaming/kafka_utils.py +++ b/python/pyspark/sql/tests/streaming/kafka_utils.py @@ -43,7 +43,7 @@ def test_kafka_streaming(self): import logging import time -from typing import List, Tuple, Callable, Any +from typing import Any, Callable, List, Tuple from pyspark.sql import SparkSession diff --git a/python/pyspark/sql/tests/streaming/test_streaming.py b/python/pyspark/sql/tests/streaming/test_streaming.py index 0ca6f6b4bb4f2..8d03e0a859b7e 100644 --- a/python/pyspark/sql/tests/streaming/test_streaming.py +++ b/python/pyspark/sql/tests/streaming/test_streaming.py @@ -20,11 +20,11 @@ import tempfile import time +from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql import Row from pyspark.sql.functions import lit -from pyspark.sql.types import StructType, StructField, IntegerType, StringType, TimestampType +from pyspark.sql.types import IntegerType, StringType, StructField, StructType, TimestampType from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.errors import PySparkTypeError, PySparkValueError class StreamingTestsMixin: @@ -364,8 +364,8 @@ def test_stream_exception(self): finally: sq.stop() - from pyspark.sql.functions import col, udf from pyspark.errors import StreamingQueryException + from pyspark.sql.functions import col, udf bad_udf = udf(lambda x: 1 / 0) sq = ( diff --git a/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py b/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py index 818d8361537d8..a12d71ce3f9b6 100644 --- a/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py +++ b/python/pyspark/sql/tests/streaming/test_streaming_foreach_batch.py @@ -16,6 +16,7 @@ # import time + from pyspark.sql.dataframe import DataFrame from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/streaming/test_streaming_kafka_rtm.py b/python/pyspark/sql/tests/streaming/test_streaming_kafka_rtm.py index 868481f946ff5..56ecf31db53bd 100644 --- a/python/pyspark/sql/tests/streaming/test_streaming_kafka_rtm.py +++ b/python/pyspark/sql/tests/streaming/test_streaming_kafka_rtm.py @@ -31,7 +31,7 @@ import uuid from pyspark.sql.tests.streaming.kafka_utils import KafkaUtils -from pyspark.testing.sqlutils import ReusedSQLTestCase, search_jar, read_classpath +from pyspark.testing.sqlutils import ReusedSQLTestCase, read_classpath, search_jar from pyspark.testing.utils import ( have_kafka, have_testcontainers, diff --git a/python/pyspark/sql/tests/streaming/test_streaming_listener.py b/python/pyspark/sql/tests/streaming/test_streaming_listener.py index b4922f54b2170..ea250ec0b5c9b 100644 --- a/python/pyspark/sql/tests/streaming/test_streaming_listener.py +++ b/python/pyspark/sql/tests/streaming/test_streaming_listener.py @@ -20,17 +20,17 @@ from datetime import datetime from pyspark import Row +from pyspark.sql.functions import col, count, lit from pyspark.sql.streaming import StreamingQueryListener from pyspark.sql.streaming.listener import ( - QueryStartedEvent, QueryProgressEvent, + QueryStartedEvent, QueryTerminatedEvent, SinkProgress, SourceProgress, StateOperatorProgress, StreamingQueryProgress, ) -from pyspark.sql.functions import count, col, lit from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py b/python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py index 7244afcf66592..fcffc8c846498 100644 --- a/python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py +++ b/python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py @@ -27,7 +27,7 @@ SimpleStatefulProcessorWithInitialStateFactory, StatefulProcessorCompositeTypeFactory, ) -from pyspark.sql.types import LongType, StringType, StructType, StructField +from pyspark.sql.types import LongType, StringType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, diff --git a/python/pyspark/sql/tests/test_artifact.py b/python/pyspark/sql/tests/test_artifact.py index 7351762dadcac..408b33f1ed743 100644 --- a/python/pyspark/sql/tests/test_artifact.py +++ b/python/pyspark/sql/tests/test_artifact.py @@ -17,10 +17,10 @@ import os import tempfile -from pyspark.sql.tests.connect.client.test_artifact import ArtifactTestsMixin -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.errors import PySparkRuntimeError from pyspark.sql.functions import assert_true, lit, udf +from pyspark.sql.tests.connect.client.test_artifact import ArtifactTestsMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase class ArtifactTests(ArtifactTestsMixin, ReusedSQLTestCase): diff --git a/python/pyspark/sql/tests/test_catalog.py b/python/pyspark/sql/tests/test_catalog.py index 66f7167df09f8..d813384a3a747 100644 --- a/python/pyspark/sql/tests/test_catalog.py +++ b/python/pyspark/sql/tests/test_catalog.py @@ -18,7 +18,7 @@ from pyspark import StorageLevel from pyspark.errors import AnalysisException, PySparkTypeError -from pyspark.sql.types import StructType, StructField, IntegerType +from pyspark.sql.types import IntegerType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_collection.py b/python/pyspark/sql/tests/test_collection.py index 5d0e48b73ce83..fd8b3edc6e3d3 100644 --- a/python/pyspark/sql/tests/test_collection.py +++ b/python/pyspark/sql/tests/test_collection.py @@ -20,24 +20,24 @@ from pyspark.loose_version import LooseVersion from pyspark.sql.types import ( - Row, ArrayType, - StringType, - IntegerType, - StructType, - StructField, BooleanType, DateType, - TimestampType, - TimestampNTZType, - FloatType, DayTimeIntervalType, + FloatType, + IntegerType, + Row, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( assertDataFrameEqual, - have_pyarrow, have_pandas, + have_pyarrow, pandas_requirement_message, pyarrow_requirement_message, ) @@ -45,7 +45,7 @@ class DataFrameCollectionTestsMixin: def _to_pandas(self): - from datetime import datetime, date, timedelta + from datetime import date, datetime, timedelta schema = ( StructType() @@ -96,8 +96,8 @@ def _to_pandas(self): @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_to_pandas(self): - import pandas as pd import numpy as np + import pandas as pd pdf = self._to_pandas() types = pdf.dtypes @@ -165,8 +165,8 @@ def test_to_pandas_required_pandas_not_found(self): @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_to_pandas_avoid_astype(self): - import pandas as pd import numpy as np + import pandas as pd schema = StructType().add("a", IntegerType()).add("b", StringType()).add("c", IntegerType()) data = [(1, "foo", 16777220), (None, "bar", None)] @@ -222,8 +222,8 @@ def check_to_pandas_from_null_dataframe(self): # SPARK-29188 test that toPandas() on a dataframe with only nulls has correct dtypes # SPARK-30537 test that toPandas() on a dataframe with only nulls has correct dtypes # using arrow - import pandas as pd import numpy as np + import pandas as pd sql = """ SELECT CAST(NULL AS TINYINT) AS tinyint, diff --git a/python/pyspark/sql/tests/test_column.py b/python/pyspark/sql/tests/test_column.py index 143e11a922db9..8056f3436f312 100644 --- a/python/pyspark/sql/tests/test_column.py +++ b/python/pyspark/sql/tests/test_column.py @@ -15,17 +15,17 @@ # limitations under the License. # -from enum import Enum -from itertools import chain import datetime import unittest import uuid +from enum import Enum +from itertools import chain +from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError from pyspark.sql import Column, Row from pyspark.sql import functions as sf +from pyspark.sql.types import IntegerType, LongType, StringType, StructField, StructType from pyspark.sql.window import Window -from pyspark.sql.types import StructType, StructField, IntegerType, LongType, StringType -from pyspark.errors import AnalysisException, PySparkTypeError, PySparkValueError from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import have_pandas, pandas_requirement_message @@ -47,8 +47,8 @@ def test_and_in_expression(self): self.assertRaises(ValueError, lambda: not self.df.key == 1) def test_validate_column_types(self): - from pyspark.sql.functions import udf, to_json from pyspark.sql.classic.column import _to_java_column + from pyspark.sql.functions import to_json, udf self.assertTrue("Column" in _to_java_column("a").getClass().toString()) self.assertTrue("Column" in _to_java_column("a").getClass().toString()) @@ -204,7 +204,7 @@ def test_bitwise_operations(self): self.assertEqual(~75, result["~b"]) def test_with_field(self): - from pyspark.sql.functions import lit, col + from pyspark.sql.functions import col, lit df = self.spark.createDataFrame([Row(a=Row(b=1, c=2))]) self.assertIsInstance(df["a"].withField("b", lit(3)), Column) diff --git a/python/pyspark/sql/tests/test_connect_compatibility.py b/python/pyspark/sql/tests/test_connect_compatibility.py index 6c1268f927851..97dad75501e1b 100644 --- a/python/pyspark/sql/tests/test_connect_compatibility.py +++ b/python/pyspark/sql/tests/test_connect_compatibility.py @@ -15,50 +15,50 @@ # limitations under the License. # -import unittest -import inspect import functools +import inspect +import unittest -from pyspark.testing.connectutils import should_test_connect, connect_requirement_message -from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.sql.classic.dataframe import DataFrame as ClassicDataFrame -from pyspark.sql.classic.column import Column as ClassicColumn -from pyspark.sql.session import SparkSession as ClassicSparkSession +import pyspark.sql.avro.functions as ClassicAvro +import pyspark.sql.functions as ClassicFunctions +import pyspark.sql.protobuf.functions as ClassicProtobuf from pyspark.sql.catalog import Catalog as ClassicCatalog +from pyspark.sql.classic.column import Column as ClassicColumn +from pyspark.sql.classic.dataframe import DataFrame as ClassicDataFrame +from pyspark.sql.group import GroupedData as ClassicGroupedData from pyspark.sql.readwriter import DataFrameReader as ClassicDataFrameReader from pyspark.sql.readwriter import DataFrameWriter as ClassicDataFrameWriter from pyspark.sql.readwriter import DataFrameWriterV2 as ClassicDataFrameWriterV2 -from pyspark.sql.window import Window as ClassicWindow -from pyspark.sql.window import WindowSpec as ClassicWindowSpec -import pyspark.sql.functions as ClassicFunctions -from pyspark.sql.group import GroupedData as ClassicGroupedData -import pyspark.sql.avro.functions as ClassicAvro -import pyspark.sql.protobuf.functions as ClassicProtobuf +from pyspark.sql.session import SparkSession as ClassicSparkSession from pyspark.sql.streaming.query import StreamingQuery as ClassicStreamingQuery from pyspark.sql.streaming.query import StreamingQueryManager as ClassicStreamingQueryManager from pyspark.sql.streaming.readwriter import DataStreamReader as ClassicDataStreamReader from pyspark.sql.streaming.readwriter import DataStreamWriter as ClassicDataStreamWriter +from pyspark.sql.window import Window as ClassicWindow +from pyspark.sql.window import WindowSpec as ClassicWindowSpec +from pyspark.testing.connectutils import connect_requirement_message, should_test_connect +from pyspark.testing.sqlutils import ReusedSQLTestCase if should_test_connect: - from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame - from pyspark.sql.connect.column import Column as ConnectColumn - from pyspark.sql.connect.session import SparkSession as ConnectSparkSession + import pyspark.sql.connect.avro.functions as ConnectAvro + import pyspark.sql.connect.functions as ConnectFunctions + import pyspark.sql.connect.protobuf.functions as ConnectProtobuf from pyspark.sql.connect.catalog import Catalog as ConnectCatalog + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + from pyspark.sql.connect.group import GroupedData as ConnectGroupedData from pyspark.sql.connect.readwriter import DataFrameReader as ConnectDataFrameReader from pyspark.sql.connect.readwriter import DataFrameWriter as ConnectDataFrameWriter from pyspark.sql.connect.readwriter import DataFrameWriterV2 as ConnectDataFrameWriterV2 - from pyspark.sql.connect.window import Window as ConnectWindow - from pyspark.sql.connect.window import WindowSpec as ConnectWindowSpec - import pyspark.sql.connect.functions as ConnectFunctions - from pyspark.sql.connect.group import GroupedData as ConnectGroupedData - import pyspark.sql.connect.avro.functions as ConnectAvro - import pyspark.sql.connect.protobuf.functions as ConnectProtobuf + from pyspark.sql.connect.session import SparkSession as ConnectSparkSession from pyspark.sql.connect.streaming.query import StreamingQuery as ConnectStreamingQuery from pyspark.sql.connect.streaming.query import ( StreamingQueryManager as ConnectStreamingQueryManager, ) from pyspark.sql.connect.streaming.readwriter import DataStreamReader as ConnectDataStreamReader from pyspark.sql.connect.streaming.readwriter import DataStreamWriter as ConnectDataStreamWriter + from pyspark.sql.connect.window import Window as ConnectWindow + from pyspark.sql.connect.window import WindowSpec as ConnectWindowSpec class ConnectCompatibilityTestsMixin: diff --git a/python/pyspark/sql/tests/test_context.py b/python/pyspark/sql/tests/test_context.py index 2a2dc0cd69ed7..540ee48d1d8be 100644 --- a/python/pyspark/sql/tests/test_context.py +++ b/python/pyspark/sql/tests/test_context.py @@ -25,7 +25,7 @@ from pyspark import SparkContext, SQLContext from pyspark.sql import Row, SparkSession -from pyspark.sql.types import StructType, StringType, StructField +from pyspark.sql.types import StringType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index d02dd8e734974..c5e9fb55aba20 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -22,11 +22,11 @@ from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.sql.conversion import ( + ArrowArrayConversion, ArrowArrayToPandasConversion, + ArrowBatchTransformer, ArrowTableToRowsConversion, LocalDataToArrowConversion, - ArrowArrayConversion, - ArrowBatchTransformer, PandasToArrowConversion, ) from pyspark.sql.types import ( @@ -351,9 +351,10 @@ def test_convert_arrow_cast(self): def test_convert_decimal(self): """Test int to decimal coercion.""" - import pandas as pd from decimal import Decimal + import pandas as pd + # DataFrame with integers, schema expects decimal df = pd.DataFrame({"a": [1, 2, 3]}) schema = StructType([StructField("a", DecimalType(10, 2))]) @@ -461,6 +462,21 @@ def test_convert_categorical(self): result = PandasToArrowConversion.convert([cat_series], schema) self.assertEqual(result.column(0).to_pylist(), ["a", "b", "a", "c"]) + def test_convert_chunked_array_backed(self): + """Test a chunked arrow-backed series is converted to a single Array.""" + import pandas as pd + import pyarrow as pa + + # pa.Array.from_pandas returns a ChunkedArray here, which + # pa.RecordBatch.from_arrays rejects. + chunked = pa.chunked_array([pa.array(["a", "b"]), pa.array(["c", "d", "e"])]) + series = pd.Series(chunked, dtype="string[pyarrow]") + schema = StructType([StructField("s", StringType())]) + + result = PandasToArrowConversion.convert([series], schema, arrow_cast=True) + self.assertIsInstance(result.column(0), pa.Array) + self.assertEqual(result.column(0).to_pylist(), ["a", "b", "c", "d", "e"]) + @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ConversionTests(unittest.TestCase): @@ -471,12 +487,19 @@ def test_conversion(self): (IntegerType(), (1,), (None,)), ((IntegerType(), {"nullable": False}), (1,)), (StringType(), ("a",)), + # bool coerced to string matches the JVM (EvaluatePython.makeFromJava). + (StringType(), (True, "true"), (False, "false")), (BinaryType(), (b"a",)), (GeographyType("ANY"), (None,)), (GeometryType("ANY"), (None,)), (ArrayType(IntegerType()), ([1, None],)), (ArrayType(IntegerType(), containsNull=False), ([1, 2],)), (ArrayType(BinaryType()), ([b"a", b"b"],)), + # array<string> with already-str, coerced (int/bool) and null elements. + ( + ArrayType(StringType()), + (["ok", 42, True, False, None], ["ok", "42", "true", "false", None]), + ), (MapType(StringType(), IntegerType()), ({"a": 1, "b": None},)), ( MapType(StringType(), IntegerType(), valueContainsNull=False), diff --git a/python/pyspark/sql/tests/test_creation.py b/python/pyspark/sql/tests/test_creation.py index 0c77faf9fe3bf..0b4542206001d 100644 --- a/python/pyspark/sql/tests/test_creation.py +++ b/python/pyspark/sql/tests/test_creation.py @@ -15,26 +15,27 @@ # limitations under the License. # -from decimal import Decimal import os import time import unittest -from pyspark.sql import Row +from decimal import Decimal + import pyspark.sql.functions as F +from pyspark.errors import ( + PySparkTypeError, + PySparkValueError, +) +from pyspark.sql import Row from pyspark.sql.types import ( + DateType, DecimalType, IntegerType, - StructType, - StructField, StringType, - DateType, - TimeType, - TimestampType, + StructField, + StructType, TimestampNTZType, -) -from pyspark.errors import ( - PySparkTypeError, - PySparkValueError, + TimestampType, + TimeType, ) from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -85,9 +86,10 @@ def test_create_dataframe_from_datetime_time(self): @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_create_dataframe_from_pandas_with_timestamp(self): - import pandas as pd from datetime import datetime + import pandas as pd + pdf = pd.DataFrame( {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]}, columns=["d", "ts"], @@ -110,9 +112,10 @@ def test_create_dataframe_required_pandas_not_found(self): with self.assertRaisesRegex( ImportError, "(Pandas >= .* must be installed|No module named '?pandas'?)" ): - import pandas as pd from datetime import datetime + import pandas as pd + pdf = pd.DataFrame( {"ts": [datetime(2017, 10, 31, 1, 1, 1)], "d": [pd.Timestamp.now().date()]} ) @@ -121,9 +124,10 @@ def test_create_dataframe_required_pandas_not_found(self): # Regression test for SPARK-23360 @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_create_dataframe_from_pandas_with_dst(self): + from datetime import datetime + import pandas as pd from pandas.testing import assert_frame_equal - from datetime import datetime pdf = pd.DataFrame({"time": [datetime(2015, 10, 31, 22, 30)]}) @@ -147,9 +151,10 @@ def test_create_dataframe_from_pandas_with_dst(self): @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_create_dataframe_from_pandas_with_day_time_interval(self): # SPARK-37277: Test DayTimeIntervalType in createDataFrame without Arrow. - import pandas as pd from datetime import timedelta + import pandas as pd + df = self.spark.createDataFrame(pd.DataFrame({"a": [timedelta(microseconds=123)]})) self.assertEqual(df.toPandas().a.iloc[0], timedelta(microseconds=123)) diff --git a/python/pyspark/sql/tests/test_dataframe.py b/python/pyspark/sql/tests/test_dataframe.py index e66750f78f278..70c52404fc4f7 100644 --- a/python/pyspark/sql/tests/test_dataframe.py +++ b/python/pyspark/sql/tests/test_dataframe.py @@ -16,48 +16,49 @@ # import glob +import io import os import pydoc import shutil import tempfile -import warnings import unittest -import io +import warnings from contextlib import redirect_stdout -from pyspark.sql import Row, functions, DataFrame +from pyspark.errors import ( + AnalysisException, + IllegalArgumentException, + PySparkTypeError, + PySparkValueError, + QueryContextType, +) +from pyspark.sql import DataFrame, Row, functions from pyspark.sql.functions import ( + array, col, - lit, count, - struct, date_format, - to_date, - array, explode, + lit, + struct, + to_date, ) from pyspark.sql.types import ( - StringType, IntegerType, LongType, - StructType, + StringType, StructField, + StructType, ) from pyspark.storagelevel import StorageLevel -from pyspark.errors import ( - AnalysisException, - IllegalArgumentException, - PySparkTypeError, - PySparkValueError, -) from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ( - ReusedSQLTestCase, SPARK_HOME, + ReusedSQLTestCase, ) from pyspark.testing.utils import ( - have_pyarrow, have_pandas, + have_pyarrow, pandas_requirement_message, pyarrow_requirement_message, ) @@ -261,6 +262,20 @@ def test_with_columns_renamed(self): messageParameters={"expected_type": "dict", "arg_name": "colsMap", "arg_type": "tuple"}, ) + def test_sort_ascending_invalid_type(self): + df = self.spark.createDataFrame([("Alice", 10)], ["name", "age"]) + with self.assertRaises(PySparkTypeError) as pe: + df.sort("age", ascending="asc") + self.check_error( + exception=pe.exception, + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "bool, int or list", + "arg_name": "ascending", + "arg_type": "str", + }, + ) + def test_with_columns_renamed_with_duplicated_names(self): df1 = self.spark.createDataFrame([(1, "v1")], ["id", "value"]) df2 = self.spark.createDataFrame([(1, "x", "v2")], ["id", "a", "value"]) @@ -411,6 +426,40 @@ def test_with_columns(self): self.assertRaises(TypeError, self.df.withColumns, ["key"]) self.assertRaises(Exception, self.df.withColumns) + def test_with_columns_with_dependencies(self): + df = self.spark.range(3).withColumns( + { + "a": col("id") + 1, + "b": col("a") + 1, + "c": col("a") + col("b"), + "d": col("a") + col("b") + col("c"), + } + ) + + assertDataFrameEqual( + df, + [ + Row(id=0, a=1, b=2, c=3, d=6), + Row(id=1, a=2, b=3, c=5, d=10), + Row(id=2, a=3, b=4, c=7, d=14), + ], + ) + + with self.assertRaises(AnalysisException) as pe: + self.spark.range(1).withColumns( + { + "a": col("b") - 1, + "b": col("id") + 1, + } + ).collect() + self.check_error( + exception=pe.exception, + errorClass="UNRESOLVED_COLUMN.WITH_SUGGESTION", + messageParameters={"objectName": "`b`", "proposal": "`id`"}, + query_context_type=QueryContextType.DataFrame, + fragment="col", + ) + def test_generic_hints(self): df1 = self.spark.range(10e10).toDF("id") df2 = self.spark.range(10e10).toDF("id") diff --git a/python/pyspark/sql/tests/test_dataframe_query_context.py b/python/pyspark/sql/tests/test_dataframe_query_context.py index d57dddf47af92..c72f2a533a962 100644 --- a/python/pyspark/sql/tests/test_dataframe_query_context.py +++ b/python/pyspark/sql/tests/test_dataframe_query_context.py @@ -18,8 +18,8 @@ from pyspark.errors import ( AnalysisException, ArithmeticException, - QueryContextType, NumberFormatException, + QueryContextType, ) from pyspark.sql import functions as sf from pyspark.testing.sqlutils import ( diff --git a/python/pyspark/sql/tests/test_datasources.py b/python/pyspark/sql/tests/test_datasources.py index e4f0bdfaf70b1..2ff5ffffbf93d 100644 --- a/python/pyspark/sql/tests/test_datasources.py +++ b/python/pyspark/sql/tests/test_datasources.py @@ -15,15 +15,15 @@ # limitations under the License. # +import os import shutil import tempfile import uuid -import os -from pyspark.sql import Row -from pyspark.sql.datasource import InputPartition, DataSource -from pyspark.sql.types import IntegerType, StructField, StructType, LongType, StringType from pyspark.errors import PySparkNotImplementedError +from pyspark.sql import Row +from pyspark.sql.datasource import DataSource, InputPartition +from pyspark.sql.types import IntegerType, LongType, StringType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_functions.py b/python/pyspark/sql/tests/test_functions.py index d3a4072c707e1..594e9d4cf8706 100644 --- a/python/pyspark/sql/tests/test_functions.py +++ b/python/pyspark/sql/tests/test_functions.py @@ -15,25 +15,26 @@ # limitations under the License. # -from contextlib import redirect_stdout import datetime -from enum import Enum -from inspect import getmembers, isfunction, isclass import io -from itertools import chain import math import re import unittest +from contextlib import redirect_stdout +from enum import Enum +from inspect import getmembers, isclass, isfunction +from itertools import chain from pyspark.errors import PySparkTypeError, PySparkValueError, SparkRuntimeException from pyspark.errors.exceptions.base import IllegalArgumentException -from pyspark.sql import Row, Window, functions as F, types +from pyspark.sql import Row, Window, types +from pyspark.sql import functions as F from pyspark.sql.avro.functions import from_avro, to_avro from pyspark.sql.column import Column from pyspark.sql.functions.builtin import nullifzero, randstr, uniform, zeroifnull -from pyspark.sql.types import StructType, StructField, StringType +from pyspark.sql.types import StringType, StructField, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase, SQLTestUtils -from pyspark.testing.utils import have_numpy, assertDataFrameEqual +from pyspark.testing.utils import assertDataFrameEqual, have_numpy class FunctionsTestsMixin: @@ -65,7 +66,6 @@ def test_function_parity(self): "not", # equivalent to python ~expression "any", # equivalent to python ~some "len", # equivalent to python ~length - "udaf", # used for creating UDAF's which are not supported in PySpark "partitioning$", # partitioning expressions for DSv2 ] @@ -939,6 +939,18 @@ def test_string_functions(self): df.select(getattr(F, name)(F.col("name"))), ) + def test_bitmap_scalar_functions(self): + df = self.spark.createDataFrame([("F00F", "70")], ["left", "right"]) + left = F.to_binary("left", F.lit("hex")) + right = F.to_binary("right", F.lit("hex")) + actual = df.select( + F.substring(F.hex(F.bitmap_and(left, right)), 0, 4), + F.substring(F.hex(F.bitmap_or(left, right)), 0, 4), + F.substring(F.hex(F.bitmap_andnot(left, right)), 0, 4), + F.substring(F.hex(F.bitmap_xor(left, right)), 0, 4), + ) + assertDataFrameEqual([Row("7000", "F00F", "800F", "800F")], actual) + def test_collation(self): df = self.spark.createDataFrame([("a",), ("b",)], ["name"]) actual = df.select(F.collation(F.collate("name", "UNICODE"))).distinct() @@ -974,7 +986,7 @@ def test_levenshtein_function(self): assertDataFrameEqual([Row(b=-1)], actual_with_threshold) def test_vector_functions(self): - from pyspark.sql.types import ArrayType, FloatType, StructType, StructField + from pyspark.sql.types import ArrayType, FloatType, StructField, StructType schema = StructType( [ @@ -1022,6 +1034,14 @@ def test_jaro_winkler_similarity_function(self): null_result = df.select(F.jaro_winkler_similarity(df.l, F.lit(None))).first()[0] self.assertIsNone(null_result) + def test_normalize_function(self): + df = self.spark.createDataFrame([("\ufb01",)], ["s"]) + result = df.select(F.normalize(df.s, F.lit("NFKC"))).first()[0] + self.assertEqual(result, "fi") + # Null handling + null_result = df.select(F.normalize(F.lit(None), F.lit("NFC"))).first()[0] + self.assertIsNone(null_result) + def test_between_function(self): df = self.spark.createDataFrame( [Row(a=1, b=2, c=3), Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)] @@ -2148,6 +2168,56 @@ def test_collect_functions(self): ["1", "2", "2", "2"], ) + def test_collect_union(self): + # array<int>: distinct union across rows; NULL arrays ignored. + df = self.spark.createDataFrame([([1, 2],), ([2, 3],), ([1],), (None,)], ["value"]) + self.assertEqual( + sorted(df.select(F.collect_union(df.value).alias("r")).collect()[0].r), + [1, 2, 3], + ) + + # array<string> + sdf = self.spark.createDataFrame([(["a", "b"],), (["b", "c"],), (["a"],)], ["value"]) + self.assertEqual( + sorted(sdf.select(F.collect_union("value").alias("r")).collect()[0].r), + ["a", "b", "c"], + ) + + # array<double> (buffer keyed by bit pattern; values round-trip) + ddf = self.spark.createDataFrame([([1.5, 2.5],), ([2.5, 3.5],)], ["value"]) + self.assertEqual( + sorted(ddf.select(F.collect_union("value").alias("r")).collect()[0].r), + [1.5, 2.5, 3.5], + ) + + # NULL elements inside a non-null array are dropped by default (IGNORE NULLS) ... + ndf = self.spark.createDataFrame([([1, None],), ([2],)], "value: array<int>") + self.assertEqual( + sorted(ndf.select(F.collect_union("value").alias("r")).collect()[0].r), + [1, 2], + ) + # ... and kept (a single null) with RESPECT NULLS (SQL clause via expr). + respect = ndf.select(F.expr("collect_union(value) RESPECT NULLS").alias("r")).collect()[0].r + self.assertEqual(sorted(respect, key=lambda x: (x is not None, x)), [None, 1, 2]) + + # array<struct>: the motivating case (dedups whole structs, stays element-wise). + struct_data = [ + ([{"id": 1, "flag": True}, {"id": 2, "flag": False}],), + ([{"id": 2, "flag": False}, {"id": 3, "flag": True}],), + ] + stdf = self.spark.createDataFrame(struct_data, "value: array<struct<id:int,flag:boolean>>") + struct_rows = stdf.select(F.collect_union("value").alias("r")).collect()[0].r + self.assertEqual( + sorted((row.id, row.flag) for row in struct_rows), + [(1, True), (2, False), (3, True)], + ) + + # Per-group union. + gdf = self.spark.createDataFrame([("a", [1, 2]), ("a", [2, 3]), ("b", [4])], ["k", "value"]) + rows = gdf.groupBy("k").agg(F.collect_union("value").alias("r")).orderBy("k").collect() + self.assertEqual(sorted(rows[0].r), [1, 2, 3]) + self.assertEqual(sorted(rows[1].r), [4]) + def test_listagg_functions(self): df = self.spark.createDataFrame( [(1, "1"), (2, "2"), (None, None), (1, "2")], ["key", "value"] @@ -3591,6 +3661,15 @@ def check(resultDf, expected): df.select(F.to_json(F.try_variant_array_append(arr, df.arrpath, F.lit(9)))), ["[1,2,9]", "[[3,9],4]"], ) + strip_v = F.parse_json(F.lit('{"a": 1, "b": null, "c": [1, null]}')) + check( + df.select(F.to_json(F.variant_strip_nulls(strip_v))), + ['{"a":1,"c":[1]}', '{"a":1,"c":[1]}'], + ) + check( + df.select(F.to_json(F.variant_strip_nulls(strip_v, False))), + ['{"a":1,"c":[1,null]}', '{"a":1,"c":[1,null]}'], + ) check(df.select(F.schema_of_variant(v)), ["OBJECT<a: BIGINT>", "OBJECT<b: BIGINT>"]) check(df.select(F.schema_of_variant_agg(v)), ["OBJECT<a: BIGINT, b: BIGINT>"]) @@ -3674,6 +3753,23 @@ def test_to_variant_object(self): ).collect() self.assertEqual("""{"a":1}""", actual[0]["var"]) + def test_variant_from_arrays_and_entries(self): + df = self.spark.createDataFrame( + [(["a", "b"], [1, 2])], "keys array<string>, values array<int>" + ) + actual = df.select( + F.to_json(F.variant_from_arrays("keys", "values")).alias("var"), + ).collect() + self.assertEqual("""{"a":1,"b":2}""", actual[0]["var"]) + + df2 = self.spark.createDataFrame( + [([("a", 1), ("b", 2)],)], "entries array<struct<k string, v int>>" + ) + actual2 = df2.select( + F.to_json(F.variant_from_entries("entries")).alias("var"), + ).collect() + self.assertEqual("""{"a":1,"b":2}""", actual2[0]["var"]) + def test_schema_of_csv(self): with self.assertRaises(PySparkTypeError) as pe: F.schema_of_csv(1) @@ -4152,6 +4248,54 @@ def test_max_by_min_by_with_k(self): self.assertEqual(result[0][1], ["Alice", "Carol"]) # Eng self.assertEqual(result[1][1], ["Frank", "Dave"]) # Sales + def test_xxh3_64(self): + """Test xxh3_64 hash function""" + # Test with string input + df = self.spark.createDataFrame([("Spark",), ("",), (None,)], ["data"]) + result = df.select(F.xxh3_64("data")).collect() + + # Verify against known values from Scala tests + self.assertEqual(result[0][0], 80997306238743657) # "Spark" + self.assertEqual(result[1][0], 0x2D06800538D394C2) # empty string + self.assertIsNone(result[2][0]) # null + + # Test with binary input + df_binary = self.spark.createDataFrame([(bytearray([1, 2, 3, 4, 5, 6]),)], ["data"]) + result_binary = df_binary.select(F.xxh3_64("data")).collect() + # Value from DataFrameFunctionsSuite.scala + self.assertEqual(result_binary[0][0], -4044731995552965649) + + def test_xxh3_128(self): + """Test xxh3_128 hash function""" + # Test with string input + df = self.spark.createDataFrame([("Spark",), ("",), (None,)], ["data"]) + result = df.select(F.xxh3_128("data")).collect() + + # Verify against known values from Scala tests + self.assertEqual(result[0][0], "7d57dd84c60c86ca1f4e82ab91a12b5e") # "Spark" + self.assertEqual(result[1][0], "99aa06d3014798d86001c324468d497f") # empty string + self.assertIsNone(result[2][0]) # null + + # Test with binary input + df_binary = self.spark.createDataFrame([(bytearray([1, 2, 3, 4, 5, 6]),)], ["data"]) + result_binary = df_binary.select(F.xxh3_128("data")).collect() + # Value from DataFrameFunctionsSuite.scala + self.assertEqual(result_binary[0][0], "866737830f560dbf3e1f439d2d785f44") + + def test_xxh3_with_cast(self): + """Test xxh3 functions with explicit cast to binary""" + df = self.spark.createDataFrame([("ABC",)], ["a"]) + + # Test xxh3_64 with cast + result_64 = df.select(F.xxh3_64(F.col("a").cast("binary"))).collect() + # Value from DataFrameFunctionsSuite.scala + self.assertEqual(result_64[0][0], 2615927343983396622) + + # Test xxh3_128 with cast + result_128 = df.select(F.xxh3_128(F.col("a").cast("binary"))).collect() + # Value from DataFrameFunctionsSuite.scala + self.assertEqual(result_128[0][0], "9e947f00ecd6acb2244da40f405c870e") + class FunctionsTests(FunctionsTestsMixin, ReusedSQLTestCase): pass diff --git a/python/pyspark/sql/tests/test_group.py b/python/pyspark/sql/tests/test_group.py index af2ee6801f4d5..534caf0928b7c 100644 --- a/python/pyspark/sql/tests/test_group.py +++ b/python/pyspark/sql/tests/test_group.py @@ -16,9 +16,10 @@ # import unittest +from pyspark.errors import AnalysisException from pyspark.sql import Row from pyspark.sql import functions as sf -from pyspark.errors import AnalysisException +from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, @@ -26,7 +27,6 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.testing import assertDataFrameEqual class GroupTestsMixin: diff --git a/python/pyspark/sql/tests/test_job_cancellation.py b/python/pyspark/sql/tests/test_job_cancellation.py index c5be232afd3b8..e2f0f30420539 100644 --- a/python/pyspark/sql/tests/test_job_cancellation.py +++ b/python/pyspark/sql/tests/test_job_cancellation.py @@ -15,11 +15,12 @@ # limitations under the License. # -import unittest import threading import time +import unittest from pyspark import InheritableThread, inheritable_thread_target +from pyspark.errors import IllegalArgumentException, PySparkValueError from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import eventually @@ -37,6 +38,27 @@ def test_tags(self): self.assertEqual(self.spark.getTags(), set()) self.spark.clearTags() + def test_invalid_tags(self): + # A tag cannot be an empty string or contain the ',' separator (documented on + # SparkSession.addTag / removeTag). Both the classic and Spark Connect paths reject + # such tags. The two paths raise different types -- a JVM-backed + # IllegalArgumentException in Classic and a PySparkValueError in Connect -- so accept + # either while still rejecting unrelated session/transport failures. + invalid_tag_errors = (IllegalArgumentException, PySparkValueError) + self.spark.clearTags() + for invalid_tag in ["", "a,b", ","]: + with self.assertRaises(invalid_tag_errors): + self.spark.addTag(invalid_tag) + with self.assertRaises(invalid_tag_errors): + self.spark.removeTag(invalid_tag) + # A rejected tag must not have been recorded. + self.assertEqual(self.spark.getTags(), set()) + + # Removing a valid but absent tag is a no-op in both modes (does not raise). + self.spark.removeTag("absent") + self.assertEqual(self.spark.getTags(), set()) + self.spark.clearTags() + def test_tags_multithread(self): output1 = None output2 = None diff --git a/python/pyspark/sql/tests/test_listener.py b/python/pyspark/sql/tests/test_listener.py index b83b4a96f1282..0fa88d06b2e10 100644 --- a/python/pyspark/sql/tests/test_listener.py +++ b/python/pyspark/sql/tests/test_listener.py @@ -17,11 +17,12 @@ import os import unittest + from pyspark.sql import SparkSession from pyspark.testing.sqlutils import SQLTestUtils from pyspark.testing.utils import ( - have_pyarrow, have_pandas, + have_pyarrow, pandas_requirement_message, pyarrow_requirement_message, ) @@ -37,6 +38,7 @@ class QueryExecutionListenerTests( @classmethod def setUpClass(cls): import glob + from pyspark.find_spark_home import _find_spark_home SPARK_HOME = _find_spark_home() diff --git a/python/pyspark/sql/tests/test_observation.py b/python/pyspark/sql/tests/test_observation.py index 440bc207a756e..aaa1335ba3152 100644 --- a/python/pyspark/sql/tests/test_observation.py +++ b/python/pyspark/sql/tests/test_observation.py @@ -15,8 +15,6 @@ # limitations under the License. # -from pyspark.sql import Row, Observation, functions as F -from pyspark.sql.types import StructType, LongType from pyspark.errors import ( AnalysisException, PySparkAssertionError, @@ -24,6 +22,9 @@ PySparkTypeError, PySparkValueError, ) +from pyspark.sql import Observation, Row +from pyspark.sql import functions as F +from pyspark.sql.types import LongType, StructType from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import assertDataFrameEqual, eventually diff --git a/python/pyspark/sql/tests/test_python_datasource.py b/python/pyspark/sql/tests/test_python_datasource.py index 675c4b19fad94..5bb8e9df1e3b2 100644 --- a/python/pyspark/sql/tests/test_python_datasource.py +++ b/python/pyspark/sql/tests/test_python_datasource.py @@ -16,15 +16,15 @@ # import contextlib import io +import json +import logging import os import tempfile import unittest -import logging -import json from dataclasses import dataclass from datetime import datetime from decimal import Decimal -from typing import Callable, Iterable, List, Union, Iterator, Tuple +from typing import Callable, Iterable, Iterator, List, Tuple, Union from pyspark.errors import AnalysisException, PythonException from pyspark.memory_profiler_ext import has_memory_profiler @@ -53,7 +53,7 @@ ) from pyspark.sql.functions import spark_partition_id from pyspark.sql.session import SparkSession -from pyspark.sql.types import Row, StructField, StructType, IntegerType, DecimalType, VariantVal +from pyspark.sql.types import DecimalType, IntegerType, Row, StructField, StructType, VariantVal from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ( SPARK_HOME, @@ -409,6 +409,641 @@ def reader(self, schema) -> "DataSourceReader": with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): df.show() + def test_limit_pushdown(self): + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.limit = None + + def pushLimit(self, limit: int) -> bool: + self.limit = limit + return True + + def partitions(self): + assert self.limit == 2, self.limit + return super().partitions() + + def read(self, partition): + assert self.limit == 2, self.limit + # Only produce as many rows as the query asked for. + for i in range(self.limit): + yield (i,) + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + def test_limit_pushdown_not_supported(self): + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int) -> bool: + # The reader cannot make use of the limit. + return False + + def read(self, partition): + yield from [(0,), (1,), (2,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + def test_limit_pushdown_rejected_does_not_plan_under_mutated_state(self): + # A reader may mutate itself while considering a limit and then reject it. partitions() + # and read() must not run under that rejected state; the scan must behave as if pushLimit + # was never called. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.considering_limit = False + + def pushLimit(self, limit: int) -> bool: + # Mutate self, then reject the limit. + self.considering_limit = True + return False + + def partitions(self): + assert not self.considering_limit, "partitions() planned under rejected limit" + return [InputPartition(0)] + + def read(self, partition): + assert not self.considering_limit, "read() planned under rejected limit" + yield from [(0,), (1,), (2,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + def test_limit_pushdown_partitions_planned_after_push_limit(self): + # With both filter and limit pushdown enabled, partitions() must be planned only after + # pushLimit -- never in the filter-pushdown pass before the limit is known -- so a reader + # that shapes partitions from the pushed limit sees it, and does not do (discarded) full + # partition discovery first. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.limit = None + + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + return [] # accept all filters + + def pushLimit(self, limit: int) -> bool: + self.limit = limit + return True + + def partitions(self): + # Fails if planned before pushLimit set the limit. + assert self.limit == 2, self.limit + return [InputPartition(0)] + + def read(self, partition): + assert self.limit == 2, self.limit + yield from [(1,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(2) + assertDataFrameEqual(df, [Row(x=1), Row(x=1)]) + + def test_limit_pushdown_over_delivering_reader(self): + # A reader that accepts the limit but ignores it must not change the query result: + # Spark always applies the limit again after the scan. + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int) -> bool: + return True + + def read(self, partition): + yield from [(i,) for i in range(100)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(3) + self.assertEqual(df.count(), 3) + + def test_limit_pushdown_with_filter(self): + # The limit is pushed after the filters, and the reader sees both. All filters are + # accepted here, so no post-scan filter remains to block limit pushdown. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.filters = [] + self.limit = None + + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + self.filters = list(filters) + return [] + + def pushLimit(self, limit: int) -> bool: + # pushFilters must have been called before pushLimit. + assert EqualTo(("x",), 1) in self.filters, self.filters + self.limit = limit + return True + + def read(self, partition): + assert EqualTo(("x",), 1) in self.filters, self.filters + assert self.limit == 5, self.limit + yield from [(1,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(5) + # All filters are reported as fully pushed, so Spark does not re-apply them. + assertDataFrameEqual(df, [Row(x=1), Row(x=1)]) + + def test_limit_pushdown_blocked_by_post_scan_filter(self): + # A filter that the reader does not accept stays as a post-scan filter, which prevents + # LIMIT from being pushed: applying the limit before that filter could drop rows the + # query needs. The result must still be correct. + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + # Accept nothing. + return filters + + def pushLimit(self, limit: int) -> bool: + raise AssertionError("pushLimit should not be called") + + def read(self, partition): + yield from [(1,), (2,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(5) + assertDataFrameEqual(df, [Row(x=1), Row(x=1)]) + + def test_limit_pushdown_rejected_with_accepted_filter(self): + # The reader accepts the filter but rejects the limit. Because every filter was pushed, + # no post-scan filter blocks the limit, so pushLimit is called -- and returns False. The + # worker then plans no read info (the rejected reader may be mutated), so build() re-plans + # the filters-only read from a fresh reader. That re-plan must reproduce the pushed-filter + # state, and the reader must return only the matching rows. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.filters = [] + + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + self.filters = list(filters) + return [] # accept all filters + + def pushLimit(self, limit: int) -> bool: + # pushLimit runs only after the filters were accepted. + assert EqualTo(("x",), 1) in self.filters, self.filters + return False # reject the limit + + def read(self, partition): + # Spark removed the accepted filter, so the reader must apply it itself. + assert EqualTo(("x",), 1) in self.filters, self.filters + yield from [(1,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(5) + assertDataFrameEqual(df, [Row(x=1), Row(x=1)]) + + def test_limit_pushdown_nondeterministic_push_filters(self): + # Pushing a limit replays pushFilters on a fresh reader. Spark has already committed to + # the first pass's filter decision, so a reader that reports a different supported set + # the second time must fail the query instead of silently returning wrong rows. + with tempfile.TemporaryDirectory(prefix="test_limit_pushdown_nondet") as d: + counter_path = os.path.join(d, "calls") + + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + # Accept everything on the first call, nothing on the replay. + n = 0 + if os.path.exists(counter_path): + with open(counter_path) as f: + n = int(f.read().strip() or 0) + with open(counter_path, "w") as f: + f.write(str(n + 1)) + return [] if n == 0 else list(filters) + + def pushLimit(self, limit: int) -> bool: + return True + + def read(self, partition): + yield from [(1,), (2,), (3,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(5) + with self.assertRaisesRegex(Exception, "must be deterministic"): + df.collect() + + def test_limit_pushdown_rejected_nondeterministic_push_filters(self): + # When a pushed limit is rejected, planning falls back to a fresh reader. That fallback + # replay of pushFilters must also be validated: a reader that agrees on the earlier passes + # but diverges on the fallback must fail the query, not silently read unfiltered rows. + with tempfile.TemporaryDirectory(prefix="test_limit_pushdown_reject_nondet") as d: + counter_path = os.path.join(d, "calls") + + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + n = 0 + if os.path.exists(counter_path): + with open(counter_path) as f: + n = int(f.read().strip() or 0) + with open(counter_path, "w") as f: + f.write(str(n + 1)) + # Accept everything on the first two passes, nothing on the fallback replay. + return [] if n < 2 else list(filters) + + def pushLimit(self, limit: int) -> bool: + return False # reject the limit, triggering the fresh-reader fallback + + def read(self, partition): + yield from [(1,), (2,), (3,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(5) + with self.assertRaisesRegex(Exception, "must be deterministic"): + df.collect() + + def test_limit_pushdown_zero(self): + # `LIMIT 0` never reaches the data source: EliminateLimits rewrites it to an empty + # relation before operator pushdown runs, so the scan is removed altogether and neither + # pushLimit nor read is called. + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int) -> bool: + raise AssertionError("pushLimit should not be called for LIMIT 0") + + def read(self, partition): + raise AssertionError("read should not be called for LIMIT 0") + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(0) + assertDataFrameEqual(df, []) + + def test_limit_pushdown_disabled_with_filter_pushdown_enabled(self): + # The reader implements pushLimit while limit pushdown is disabled. Filter pushdown runs + # a different planning worker, which must not silently ignore pushLimit either. + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + return [] + + def pushLimit(self, limit: int) -> bool: + return True + + def read(self, partition): + yield from [(1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": False, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1").limit(1) + with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): + df.show() + + def test_filter_pushdown_disabled_with_limit_pushdown_enabled(self): + # The reader implements pushFilters while filter pushdown is disabled. A limit-only scan + # runs the limit-pushdown worker, which caches the read info, so `plan_data_source_read` + # never runs. That worker must not silently ignore pushFilters either. No filter is used + # here on purpose: with filter pushdown disabled a filter would stay above the scan and + # block limit pushdown, so the limit worker would never run. + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + return [] + + def pushLimit(self, limit: int) -> bool: + return True + + def read(self, partition): + yield from [(1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": False, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(1) + with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): + df.show() + + def test_limit_pushdown_disabled(self): + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int) -> bool: + assert False + + def read(self, partition): + assert False + + class TestDataSource(DataSource): + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": False}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").schema("x int").load() + with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): + df.show() + + def test_pushdown_disabled_check_survives_reused_readinfo_cache(self): + # The provider caches the pushdown-free read info across scans of the same relation. That + # cache must not hide the DATA_SOURCE_PUSHDOWN_DISABLED check: warming it with pushdown + # enabled and then rescanning the same relation with pushdown disabled must still raise, + # because the reader implements a pushdown method the disabled config would silently + # ignore. The cache records the pushdown flags it was populated under and recomputes when + # they change, so the check runs for the flags currently in effect. + class TestDataSourceReader(DataSourceReader): + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + return [] + + def read(self, partition): + yield from [(0,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load() + + # Warm the provider's read-info cache with filter pushdown enabled. The query has no + # filters, so pushFilters is never called and the check passes. + with self.sql_conf({"spark.sql.python.filterPushdown.enabled": True}): + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + # Rescan the same relation with filter pushdown disabled. `df.select("*")` re-plans the + # scan on the same (provider-scoped) data source, so a stale cache would skip the check. + with self.sql_conf({"spark.sql.python.filterPushdown.enabled": False}): + with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): + df.select("*").collect() + + def test_pushdown_disabled_check_survives_reused_readinfo_cache_limit(self): + # Symmetric to the filter case above, for the limit flag. The read-info cache is keyed by + # both pushdown flags, so a regression that dropped only the limit flag from the key must + # be caught too. Warm the cache with limit pushdown enabled (no LIMIT in the query, so + # pushLimit is never called and the check passes), then rescan the same relation with + # limit pushdown disabled and assert the reader's pushLimit implementation is reported + # instead of silently ignored. + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int) -> bool: + return True + + def read(self, partition): + yield from [(0,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": False}): + with self.assertRaisesRegex(Exception, "DATA_SOURCE_PUSHDOWN_DISABLED"): + df.select("*").collect() + + def test_limit_pushdown_not_implemented(self): + # A reader that does not implement pushLimit is unaffected when the conf is on. + class TestDataSourceReader(DataSourceReader): + def read(self, partition): + yield from [(0,), (1,), (2,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + def test_limit_pushdown_only_does_not_call_push_filters(self): + # A limit-only query (no filters) must not trigger a spurious pushFilters([]) call, even + # when filter pushdown is also enabled: the worker skips pushFilters when there is nothing + # to push. A reader may therefore implement both and still expect pushFilters to be called + # only when the query has pushable filters. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.limit = None + + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + raise AssertionError("pushFilters should not be called without filters") + + def pushLimit(self, limit: int) -> bool: + self.limit = limit + return True + + def read(self, partition): + assert self.limit == 2, self.limit + yield from [(0,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + assertDataFrameEqual(df, [Row(x=0), Row(x=1)]) + + def test_limit_pushdown_invalid_return_type(self): + # pushLimit must return a bool. A non-bool return is rejected during planning. + class TestDataSourceReader(DataSourceReader): + def pushLimit(self, limit: int): + return "yes" + + def read(self, partition): + yield from [(0,), (1,), (2,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf({"spark.sql.python.limitPushdown.enabled": True}): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().limit(2) + with self.assertRaisesRegex(Exception, "DATA_SOURCE_INVALID_RETURN_TYPE"): + df.collect() + + def test_filter_pushdown_no_limit_with_limit_pushdown_enabled(self): + # Filter pushdown defers planning while limit pushdown is enabled, expecting a possible + # limit pass. When the query has no limit, that pass never comes, so build() plans the + # read with the filters only. The reader accepts the filter (so Spark removes it from the + # plan), which means the deferred re-plan must reproduce the pushed-filter state and the + # reader itself must return only the matching rows. + class TestDataSourceReader(DataSourceReader): + def __init__(self): + self.filters = [] + + def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]: + self.filters = list(filters) + return [] # accept all filters + + def pushLimit(self, limit: int) -> bool: + raise AssertionError("pushLimit should not be called without a limit") + + def read(self, partition): + # Spark removed the accepted filter, so the reader is responsible for it. + assert EqualTo(("x",), 1) in self.filters, self.filters + yield from [(1,), (1,)] + + class TestDataSource(DataSource): + def schema(self): + return "x int" + + def reader(self, schema) -> "DataSourceReader": + return TestDataSourceReader() + + with self.sql_conf( + { + "spark.sql.python.filterPushdown.enabled": True, + "spark.sql.python.limitPushdown.enabled": True, + } + ): + self.spark.dataSource.register(TestDataSource) + df = self.spark.read.format("TestDataSource").load().filter("x = 1") + assertDataFrameEqual(df, [Row(x=1), Row(x=1)]) + def _check_filters(self, sql_type, sql_filter, python_filters): """ Parameters diff --git a/python/pyspark/sql/tests/test_python_streaming_datasource.py b/python/pyspark/sql/tests/test_python_streaming_datasource.py index 704a2986ded88..47f381e249302 100644 --- a/python/pyspark/sql/tests/test_python_streaming_datasource.py +++ b/python/pyspark/sql/tests/test_python_streaming_datasource.py @@ -14,33 +14,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import json import os import tempfile import time import unittest -import json +from pyspark.errors import PySparkException from pyspark.sql.datasource import ( DataSource, + DataSourceStreamArrowWriter, DataSourceStreamReader, - InputPartition, DataSourceStreamWriter, - DataSourceStreamArrowWriter, + InputPartition, SimpleDataSourceStreamReader, WriterCommitMessage, ) +from pyspark.sql.streaming import StreamingQueryException from pyspark.sql.streaming.datasource import ( ReadAllAvailable, ReadLimit, ReadMaxRows, SupportsTriggerAvailableNow, ) -from pyspark.sql.streaming import StreamingQueryException from pyspark.sql.types import Row -from pyspark.errors import PySparkException from pyspark.testing import assertDataFrameEqual -from pyspark.testing.utils import eventually, have_pyarrow, pyarrow_requirement_message from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import eventually, have_pyarrow, pyarrow_requirement_message def wait_for_condition(query, condition_fn, timeout_sec=30): @@ -604,13 +604,14 @@ def test_stream_writer(self): def test_stream_arrow_writer(self): """Test DataSourceStreamArrowWriter with Arrow RecordBatch format.""" - import tempfile - import shutil import json import os - import pyarrow as pa + import shutil + import tempfile from dataclasses import dataclass + import pyarrow as pa + @dataclass class ArrowCommitMessage(WriterCommitMessage): partition_id: int diff --git a/python/pyspark/sql/tests/test_readwriter.py b/python/pyspark/sql/tests/test_readwriter.py index 5dd5b1ebff634..2f7ba28f9614c 100644 --- a/python/pyspark/sql/tests/test_readwriter.py +++ b/python/pyspark/sql/tests/test_readwriter.py @@ -24,12 +24,12 @@ from pyspark.sql.functions import col, lit from pyspark.sql.readwriter import DataFrameWriterV2 from pyspark.sql.types import ( - StructType, - StructField, - StringType, - BinaryType, ArrayType, + BinaryType, MapType, + StringType, + StructField, + StructType, ) from pyspark.testing import assertDataFrameEqual from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -356,7 +356,8 @@ def test_partitioning_functions(self): def check_partitioning_functions(self, tpe): import datetime - from pyspark.sql.functions.partitioning import years, months, days, hours, bucket + + from pyspark.sql.functions.partitioning import bucket, days, hours, months, years df = self.spark.createDataFrame( [(1, datetime.datetime(2000, 1, 1), "foo")], ("id", "ts", "value") @@ -374,7 +375,8 @@ def check_partitioning_functions(self, tpe): def partitioning_functions_user_error(self): import datetime - from pyspark.sql.functions.partitioning import years, months, days, hours, bucket + + from pyspark.sql.functions.partitioning import bucket, days, hours, months, years df = self.spark.createDataFrame( [(1, datetime.datetime(2000, 1, 1), "foo")], ("id", "ts", "value") diff --git a/python/pyspark/sql/tests/test_repartition.py b/python/pyspark/sql/tests/test_repartition.py index 707d58dba3312..c6991105f4abc 100644 --- a/python/pyspark/sql/tests/test_repartition.py +++ b/python/pyspark/sql/tests/test_repartition.py @@ -16,15 +16,15 @@ # -from pyspark.sql.functions import spark_partition_id, col, lit, when +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql.functions import col, lit, spark_partition_id, when from pyspark.sql.types import ( - StringType, - IntegerType, DoubleType, - StructType, + IntegerType, + StringType, StructField, + StructType, ) -from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_resources.py b/python/pyspark/sql/tests/test_resources.py index 1e9e9f796b5f1..b61cc515a9ac2 100644 --- a/python/pyspark/sql/tests/test_resources.py +++ b/python/pyspark/sql/tests/test_resources.py @@ -17,7 +17,7 @@ import unittest from pyspark import TaskContext -from pyspark.resource import TaskResourceRequests, ResourceProfileBuilder +from pyspark.resource import ResourceProfileBuilder, TaskResourceRequests from pyspark.sql import SparkSession from pyspark.testing.utils import ( ReusedPySparkTestCase, diff --git a/python/pyspark/sql/tests/test_serde.py b/python/pyspark/sql/tests/test_serde.py index 10a3385dce3fc..4a76be0e4535f 100644 --- a/python/pyspark/sql/tests/test_serde.py +++ b/python/pyspark/sql/tests/test_serde.py @@ -22,7 +22,7 @@ from pyspark.sql import Row from pyspark.sql.functions import lit -from pyspark.sql.types import StructType, StructField, DecimalType, BinaryType +from pyspark.sql.types import BinaryType, DecimalType, StructField, StructType from pyspark.testing.objects import UTCOffsetTimezone from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_session.py b/python/pyspark/sql/tests/test_session.py index fb86deb33a2da..a91f306cf9a0d 100644 --- a/python/pyspark/sql/tests/test_session.py +++ b/python/pyspark/sql/tests/test_session.py @@ -22,16 +22,16 @@ from pyspark import SparkConf, SparkContext from pyspark.errors import PySparkRuntimeError, PySparkValueError -from pyspark.sql import SparkSession, SQLContext, Row +from pyspark.errors.exceptions.captured import SparkNoSuchElementException +from pyspark.sql import Row, SparkSession, SQLContext from pyspark.sql.functions import col +from pyspark.sql.profiler import Profile from pyspark.testing.connectutils import ( - should_test_connect, connect_requirement_message, + should_test_connect, ) -from pyspark.errors.exceptions.captured import SparkNoSuchElementException -from pyspark.sql.profiler import Profile from pyspark.testing.sqlutils import ReusedSQLTestCase -from pyspark.testing.utils import PySparkTestCase, PySparkErrorTestUtils +from pyspark.testing.utils import PySparkErrorTestUtils, PySparkTestCase class SparkSessionTests(ReusedSQLTestCase): @@ -737,6 +737,7 @@ class SparkExtensionsTest(unittest.TestCase): @classmethod def setUpClass(cls): import glob + from pyspark.find_spark_home import _find_spark_home SPARK_HOME = _find_spark_home() diff --git a/python/pyspark/sql/tests/test_stat.py b/python/pyspark/sql/tests/test_stat.py index 5eee90526ce73..24a596f84d8de 100644 --- a/python/pyspark/sql/tests/test_stat.py +++ b/python/pyspark/sql/tests/test_stat.py @@ -16,20 +16,21 @@ # -from pyspark.sql import Row, functions as sf -from pyspark.sql.types import ( - StringType, - IntegerType, - DoubleType, - StructType, - StructField, - BooleanType, -) from pyspark.errors import ( AnalysisException, PySparkTypeError, PySparkValueError, ) +from pyspark.sql import Row +from pyspark.sql import functions as sf +from pyspark.sql.types import ( + BooleanType, + DoubleType, + IntegerType, + StringType, + StructField, + StructType, +) from pyspark.testing.sqlutils import ReusedSQLTestCase diff --git a/python/pyspark/sql/tests/test_types.py b/python/pyspark/sql/tests/test_types.py index c1e0ee6f096aa..e391fcb441ea9 100644 --- a/python/pyspark/sql/tests/test_types.py +++ b/python/pyspark/sql/tests/test_types.py @@ -22,69 +22,67 @@ import pickle import sys import unittest -from dataclasses import dataclass, asdict +from dataclasses import asdict, dataclass -from pyspark.sql import Row -from pyspark.sql import functions as F from pyspark.errors import ( AnalysisException, IllegalArgumentException, - SparkRuntimeException, ParseException, + PySparkNotImplementedError, + PySparkRuntimeError, PySparkTypeError, PySparkValueError, - PySparkRuntimeError, - PySparkNotImplementedError, + SparkRuntimeException, ) +from pyspark.sql import Row +from pyspark.sql import functions as F from pyspark.sql.types import ( - DataType, + ArrayType, + BinaryType, + BooleanType, ByteType, - ShortType, - IntegerType, - FloatType, - DateType, - TimeType, - TimestampType, - TimestampNTZType, - DayTimeIntervalType, - YearMonthIntervalType, CalendarIntervalType, - MapType, - StringType, CharType, - Geography, - Geometry, - VarcharType, - StructType, - StructField, - ArrayType, - DoubleType, - LongType, + DataType, + DateType, + DayTimeIntervalType, DecimalType, - BinaryType, - BooleanType, + DoubleType, + FloatType, + Geography, GeographyType, + Geometry, GeometryType, + IntegerType, + LongType, + MapType, NullType, + ShortType, + StringType, + StructField, + StructType, + TimestampNTZType, + TimestampType, + TimeType, UserDefinedType, + VarcharType, VariantType, VariantVal, - _create_row, -) -from pyspark.sql.types import ( + YearMonthIntervalType, _array_signed_int_typecode_ctype_mappings, _array_type_mappings, _array_unsigned_int_typecode_ctype_mappings, + _create_row, _infer_type, _make_type_verifier, _merge_type, ) from pyspark.testing.objects import ( - ExamplePointUDT, - PythonOnlyUDT, ExamplePoint, - PythonOnlyPoint, + ExamplePointUDT, MyObject, + PythonOnlyPoint, + PythonOnlyUDT, ) from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import PySparkErrorTestUtils @@ -540,7 +538,7 @@ def test_create_dataframe_from_objects(self): self.assertEqual(df.first(), Row(key=1, value="1")) def test_apply_schema(self): - from datetime import date, time, datetime, timedelta + from datetime import date, datetime, time, timedelta rdd = self.sc.parallelize( [ @@ -764,7 +762,7 @@ def test_schema_with_collations_json_ser_de(self): assert schema == _parse_datatype_json_string(schema.json()) def test_schema_with_collations_on_non_string_types(self): - from pyspark.sql.types import _parse_datatype_json_string, _COLLATIONS_METADATA_KEY + from pyspark.sql.types import _COLLATIONS_METADATA_KEY, _parse_datatype_json_string collations_on_int_col_json = f""" {{ @@ -910,7 +908,7 @@ def test_map_type_from_json(self): self.assertEqual(mapWithCollations, MapType.fromJson(map_json, collationsMap=collationsMap)) def test_schema_with_bad_collations_provider(self): - from pyspark.sql.types import _parse_datatype_json_string, _COLLATIONS_METADATA_KEY + from pyspark.sql.types import _COLLATIONS_METADATA_KEY, _parse_datatype_json_string schema_json = f""" {{ @@ -933,7 +931,7 @@ def test_schema_with_bad_collations_provider(self): self.assertRaises(PySparkValueError, lambda: _parse_datatype_json_string(schema_json)) def test_geography_json_serde(self): - from pyspark.sql.types import _parse_datatype_json_value, _parse_datatype_json_string + from pyspark.sql.types import _parse_datatype_json_string, _parse_datatype_json_value valid_test_cases = [ ("geography", GeographyType(4326)), @@ -981,7 +979,7 @@ def test_geography_json_serde(self): _parse_datatype_json_value(json) def test_geometry_json_serde(self): - from pyspark.sql.types import _parse_datatype_json_value, _parse_datatype_json_string + from pyspark.sql.types import _parse_datatype_json_string, _parse_datatype_json_value valid_test_cases = [ ("geometry", GeometryType(4326)), @@ -1025,7 +1023,7 @@ def test_geometry_json_serde(self): _parse_datatype_json_value(json) def test_udt(self): - from pyspark.sql.types import _parse_datatype_json_string, _infer_type, _make_type_verifier + from pyspark.sql.types import _infer_type, _make_type_verifier, _parse_datatype_json_string def check_datatype(datatype): pickled = pickle.loads(pickle.dumps(datatype)) @@ -2506,9 +2504,10 @@ def test_variant_type(self): self.spark.createDataFrame([VariantVal.parseJson("2")], "v variant") def test_variant_to_pandas(self): - import pandas as pd import json + import pandas as pd + expected_values = [ ("str", '"%s"' % ("0123456789" * 10), "0123456789" * 10), ("short_str", '"abc"', "abc"), diff --git a/python/pyspark/sql/tests/test_udf.py b/python/pyspark/sql/tests/test_udf.py index 5701fb6dcda97..b1a7f9214fbbf 100644 --- a/python/pyspark/sql/tests/test_udf.py +++ b/python/pyspark/sql/tests/test_udf.py @@ -15,46 +15,46 @@ # limitations under the License. # +import datetime import functools +import io +import logging import pydoc import shutil +import sys import tempfile -import unittest -import datetime -import io import time +import unittest from contextlib import redirect_stdout -import logging -import sys -from pyspark.sql import SparkSession, Column, Row -from pyspark.sql.functions import col, udf, assert_true, lit, rand -from pyspark.sql.udf import UserDefinedFunction +from pyspark.errors import AnalysisException, PySparkTypeError, PythonException +from pyspark.logger import PySparkLogger +from pyspark.sql import Column, Row, SparkSession +from pyspark.sql.functions import assert_true, col, lit, rand, udf from pyspark.sql.types import ( - StringType, - IntegerType, + ArrayType, BinaryType, BooleanType, + DayTimeIntervalType, DoubleType, + IntegerType, LongType, - ArrayType, MapType, - StructType, + StringType, StructField, + StructType, TimestampNTZType, - DayTimeIntervalType, VariantType, VariantVal, ) -from pyspark.errors import AnalysisException, PythonException, PySparkTypeError -from pyspark.logger import PySparkLogger +from pyspark.sql.udf import UserDefinedFunction from pyspark.testing.objects import ExamplePoint, ExamplePointUDT from pyspark.testing.sqlutils import ( ReusedSQLTestCase, test_compiled, test_not_compiled_message, ) -from pyspark.testing.utils import assertDataFrameEqual, timeout +from pyspark.testing.utils import assertDataFrameEqual, eventually, timeout from pyspark.util import is_remote_only @@ -190,9 +190,10 @@ def test_nondeterministic_udf_in_aggregate(self): self.check_nondeterministic_udf_in_aggregate() def check_nondeterministic_udf_in_aggregate(self): - from pyspark.sql.functions import sum import random + from pyspark.sql.functions import sum + udf_random_col = udf(lambda: int(100 * random.random()), "int").asNondeterministic() df = self.spark.range(10) @@ -836,7 +837,7 @@ def test_nonparam_udf_with_aggregate(self): # SPARK-24721 @unittest.skipIf(not test_compiled, test_not_compiled_message) def test_datasource_with_udf(self): - from pyspark.sql.functions import lit, col + from pyspark.sql.functions import col, lit path = tempfile.mkdtemp() shutil.rmtree(path) @@ -1706,20 +1707,23 @@ def my_udf(x): [Row(result=str(i)) for i in range(2)], ) - logs = self.spark.tvf.python_worker_logs() - - assertDataFrameEqual( - logs.select("level", "msg", "context", "logger"), - [ - Row( - level="WARNING", - msg="PySparkLogger test", - context={"func_name": my_udf.__name__, "x": str(i)}, - logger="PySparkLogger", - ) - for i in range(2) - ], - ) + @eventually(timeout=10, catch_assertions=True) + def check_logs(): + logs = self.spark.tvf.python_worker_logs() + assertDataFrameEqual( + logs.select("level", "msg", "context", "logger"), + [ + Row( + level="WARNING", + msg="PySparkLogger test", + context={"func_name": my_udf.__name__, "x": str(i)}, + logger="PySparkLogger", + ) + for i in range(2) + ], + ) + + check_logs() class UDFTests(BaseUDFTestsMixin, ReusedSQLTestCase): diff --git a/python/pyspark/sql/tests/test_udf_combinations.py b/python/pyspark/sql/tests/test_udf_combinations.py index c3ceb046d1df2..d4d68208ee80a 100644 --- a/python/pyspark/sql/tests/test_udf_combinations.py +++ b/python/pyspark/sql/tests/test_udf_combinations.py @@ -15,11 +15,11 @@ # limitations under the License. # -from typing import Iterator import itertools import unittest +from typing import Iterator -from pyspark.sql.functions import udf, arrow_udf, pandas_udf +from pyspark.sql.functions import arrow_udf, pandas_udf, udf from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, diff --git a/python/pyspark/sql/tests/test_udf_in_higher_order_function.py b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py new file mode 100644 index 0000000000000..b3479f4a26609 --- /dev/null +++ b/python/pyspark/sql/tests/test_udf_in_higher_order_function.py @@ -0,0 +1,1131 @@ +# +# 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. +# + +import unittest + +from pyspark.errors import AnalysisException +from pyspark.sql import functions as sf +from pyspark.sql.functions import udf +from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StringType +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import ( + assertDataFrameEqual, + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) + + +@unittest.skipIf( + not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message +) +class UDFInHigherOrderFunctionTestsMixin: + """Tests for scalar Python UDFs used inside higher-order function lambdas (SPARK-27052). + + ``ExtractPythonUDFFromLambda`` rewrites such a plan so the UDF is applied to the whole array + outside the lambda. Each test asserts the *result*, comparing against the equivalent native + expression wherever one exists, so that a rewrite that runs but computes the wrong thing + fails rather than passing quietly. + """ + + def test_transform(self): + df = self.spark.createDataFrame([([1, 2, 3],), ([],), ([10],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")), + df.select(sf.transform("values", lambda x: x + 1).alias("r")), + ) + + def test_transform_null_array_and_null_elements(self): + # A null array must stay null, and a null *element* must reach the UDF as None. + df = self.spark.createDataFrame([([1, None, 3],), (None,), ([],)], "values array<int>") + # Null-aware so the UDF itself can observe the null element. + f = udf(lambda x: -1 if x is None else x * 2, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: f(x)).alias("r")), + [([2, -1, 6],), (None,), ([],)], + ) + + def test_transform_udf_returning_null(self): + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + f = udf(lambda x: None if x == 2 else x, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: f(x)).alias("r")), + [([1, None, 3],)], + ) + + def test_transform_with_index(self): + df = self.spark.createDataFrame([([10, 20, 30],), ([],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + # The index parameter must still work once the element is read from the carrier struct. + assertDataFrameEqual( + df.select(sf.transform("values", lambda x, i: plus_one(x) + i).alias("r")), + df.select(sf.transform("values", lambda x, i: (x + 1) + i).alias("r")), + ) + + def test_composition_around_udf_result(self): + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + # Arithmetic, `when` and casts around the UDF result are ordinary JVM work. + assertDataFrameEqual( + df.select( + sf.transform("values", lambda x: plus_one(x) * 2).alias("mul"), + sf.transform( + "values", lambda x: sf.when(plus_one(x) > 2, sf.lit(1)).otherwise(sf.lit(0)) + ).alias("cond"), + sf.transform("values", lambda x: plus_one(x).cast("string")).alias("cast"), + ), + df.select( + sf.transform("values", lambda x: (x + 1) * 2).alias("mul"), + sf.transform( + "values", lambda x: sf.when((x + 1) > 2, sf.lit(1)).otherwise(sf.lit(0)) + ).alias("cond"), + sf.transform("values", lambda x: (x + 1).cast("string")).alias("cast"), + ), + ) + + def test_udf_argument_is_expression_over_element(self): + # `udf(x * 2)`: the argument is itself an expression over the element. + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x * 2)).alias("r")), + df.select(sf.transform("values", lambda x: x * 2 + 1).alias("r")), + ) + + def test_multiple_udfs_in_one_lambda(self): + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + times_ten = udf(lambda x: x * 10, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x) + times_ten(x)).alias("r")), + df.select(sf.transform("values", lambda x: (x + 1) + (x * 10)).alias("r")), + ) + + def test_nested_udfs(self): + # `f(g(x))`: both are lifted, and compose as array UDFs outside the lambda. + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + times_ten = udf(lambda x: x * 10, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: times_ten(plus_one(x))).alias("r")), + df.select(sf.transform("values", lambda x: (x + 1) * 10).alias("r")), + ) + + def test_nested_udf_inside_composite_argument(self): + # SPARK-27052: `f(g(x) + 1)` / `f(-g(x))`. The inner call is buried inside a composite + # argument, not a direct child. The inner result must be substituted before lifting `f`, + # or a raw `g` over the lambda variable would be left inside a lambda and mis-extracted + # (SPARK-48706). The result must match evaluating the composition element-wise. + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + times_ten = udf(lambda x: x * 10, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: times_ten(plus_one(x) + 1)).alias("r")), + df.select(sf.transform("values", lambda x: ((x + 1) + 1) * 10).alias("r")), + ) + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: times_ten(-plus_one(x))).alias("r")), + df.select(sf.transform("values", lambda x: (-(x + 1)) * 10).alias("r")), + ) + + def test_udf_with_outer_column_argument(self): + # A non-element argument must be broadcast to every element of its row. + df = self.spark.createDataFrame([([1, 2], 100), ([3], 200)], "values array<int>, base int") + add = udf(lambda x, b: x + b, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: add(x, sf.col("base"))).alias("r")), + [([101, 102],), ([203],)], + ) + + def test_udf_with_constant_argument_only(self): + # SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the element, but the UDF + # must still take the lambda's call domain: once per element, and zero times for an empty + # or null array (where the lambda never runs), rather than once per row. + df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)], "values array<int>") + const = udf(lambda v: v * 2, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: const(sf.lit(10))).alias("r")), + [([20, 20, 20],), ([],), (None,)], + ) + + def test_filter(self): + # `filter`'s result is built from the input elements, not the lambda's value. + df = self.spark.createDataFrame([([1, 2, 3, 4],), ([],), (None,)], "values array<int>") + is_even = udf(lambda x: x % 2 == 0, "boolean") + + assertDataFrameEqual( + df.select(sf.filter("values", lambda x: is_even(x)).alias("r")), + df.select(sf.filter("values", lambda x: (x % 2) == 0).alias("r")), + ) + + def test_exists_and_forall(self): + df = self.spark.createDataFrame([([1, 2, 3],), ([2, 4],), ([],)], "values array<int>") + is_even = udf(lambda x: x % 2 == 0, "boolean") + + assertDataFrameEqual( + df.select( + sf.exists("values", lambda x: is_even(x)).alias("e"), + sf.forall("values", lambda x: is_even(x)).alias("f"), + ), + df.select( + sf.exists("values", lambda x: (x % 2) == 0).alias("e"), + sf.forall("values", lambda x: (x % 2) == 0).alias("f"), + ), + ) + + def test_zip_with(self): + # Two arrays at once. `arrays_zip` pads the shorter side with nulls, which is what + # `zip_with` does itself, so differing lengths must agree with the native version. + df = self.spark.createDataFrame( + [([1, 2], [10, 20]), ([1, 2, 3], [10]), ([], []), (None, [1]), ([1], None)], + "l array<int>, r array<int>", + ) + add = udf(lambda a, b: (0 if a is None else a) + (0 if b is None else b), IntegerType()) + + # Compare against the equivalent native expression with the same null handling. + assertDataFrameEqual( + df.select(sf.zip_with("l", "r", lambda a, b: add(a, b)).alias("r")), + df.select( + sf.zip_with( + "l", + "r", + lambda a, b: sf.coalesce(a, sf.lit(0)) + sf.coalesce(b, sf.lit(0)), + ).alias("r") + ), + ) + + def test_zip_with_udf_on_one_side_only(self): + df = self.spark.createDataFrame([([1, 2], [10, 20])], "l array<int>, r array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + assertDataFrameEqual( + df.select(sf.zip_with("l", "r", lambda a, b: plus_one(a) + b).alias("r")), + [([12, 23],)], + ) + + def test_array_sort_with_udf_key(self): + # When the UDF applies per element, it is precomputed as a sort key that the JVM comparator + # compares (a UDF taking both elements is instead precomputed over the pairs; see the + # pairwise test). This must actually reorder, not be a no-op. + df = self.spark.createDataFrame([([3, 1, 2],), ([],), (None,)], "values array<int>") + negate = udf(lambda x: -x, IntegerType()) + + # Sorting by -x gives descending order. + assertDataFrameEqual( + df.select( + sf.array_sort( + "values", + lambda a, b: sf.when(negate(a) < negate(b), sf.lit(-1)) + .when(negate(a) > negate(b), sf.lit(1)) + .otherwise(sf.lit(0)), + ).alias("r") + ), + [([3, 2, 1],), ([],), (None,)], + ) + + def test_array_sort_pairwise_comparator(self): + # One UDF call receiving both elements has no per-element key, so the UDF is precomputed + # over every ordered pair and the comparator indexes that matrix. Assert an actual + # reordering, not merely that the query runs. + df = self.spark.createDataFrame( + [([3, 1, 2],), ([],), (None,), ([5],), ([2, 2, 1],)], "values array<int>" + ) + cmp_udf = udf(lambda a, b: (a > b) - (a < b), IntegerType()) + + assertDataFrameEqual( + df.select(sf.array_sort("values", lambda a, b: cmp_udf(a, b)).alias("r")), + [([1, 2, 3],), ([],), (None,), ([5],), ([1, 2, 2],)], + ) + + def test_array_sort_pairwise_comparator_descending(self): + # Reversing the comparator must reverse the order, which a no-op rewrite would not do. + df = self.spark.createDataFrame([([3, 1, 2],)], "values array<int>") + cmp_desc = udf(lambda a, b: (b > a) - (b < a), IntegerType()) + + assertDataFrameEqual( + df.select(sf.array_sort("values", lambda a, b: cmp_desc(a, b)).alias("r")), + [([3, 2, 1],)], + ) + + def test_transform_keys_and_values(self): + df = self.spark.createDataFrame( + [ + ({"a": 1, "b": 2},), + ], + "m map<string,int>", + ) + upper = udf(lambda s: s.upper(), StringType()) + plus_one = udf(lambda v: v + 1, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform_keys("m", lambda k, v: upper(k)).alias("r")), + [({"A": 1, "B": 2},)], + ) + assertDataFrameEqual( + df.select(sf.transform_values("m", lambda k, v: plus_one(v)).alias("r")), + [({"a": 2, "b": 3},)], + ) + # The lambda may read both the key and the value. + assertDataFrameEqual( + df.select(sf.transform_values("m", lambda k, v: plus_one(v) + sf.length(k)).alias("r")), + [({"a": 3, "b": 4},)], + ) + + def test_map_filter(self): + df = self.spark.createDataFrame([({"a": 1, "b": 2, "c": 3},)], "m map<string,int>") + is_odd = udf(lambda v: v % 2 == 1, "boolean") + + assertDataFrameEqual( + df.select(sf.map_filter("m", lambda k, v: is_odd(v)).alias("r")), + [({"a": 1, "c": 3},)], + ) + + def test_map_zip_with(self): + # The visited key set is the union of both maps' keys; a key missing from one side gives + # null on that side, matching map_zip_with's own semantics. + df = self.spark.createDataFrame( + [ + ( + {"a": 1, "b": 2}, + {"b": 20, "c": 30}, + ) + ], + "l map<string,int>, r map<string,int>", + ) + combine = udf( + lambda a, b: (0 if a is None else a) * 100 + (0 if b is None else b), IntegerType() + ) + + assertDataFrameEqual( + df.select(sf.map_zip_with("l", "r", lambda k, v1, v2: combine(v1, v2)).alias("r")), + [({"a": 100, "b": 220, "c": 30},)], + ) + + def test_transform_values_result_type_equals_key_type(self): + # SPARK-27052: transform_values on map<string,string> whose lambda also returns string. + # The rewrite must replace the values, not the keys (dispatch is by function, not type). + df = self.spark.createDataFrame([({"a": "x", "b": "y"},)], "m map<string,string>") + tag = udf(lambda v: v + "!", StringType()) + assertDataFrameEqual( + df.select(sf.transform_values("m", lambda k, v: tag(v)).alias("r")), + [({"a": "x!", "b": "y!"},)], + ) + + def test_nondeterministic_udf_calls_are_distinct(self): + # SPARK-27052: two calls to a nondeterministic UDF in one lambda must stay distinct, not be + # collapsed into one shared value. `rand_add` returns x plus a per-call random draw, so + # f(x) + f(x) equals 2x only if the two calls were (wrongly) deduplicated. + import random + + df = self.spark.createDataFrame([([10, 20, 30],)], "values array<int>") + rand_add = udf( + lambda x: x + random.randint(1, 1_000_000), IntegerType() + ).asNondeterministic() + row = df.select( + sf.transform("values", lambda x: rand_add(x) - rand_add(x)).alias("r") + ).collect()[0] + # If the two calls were deduplicated, every element would be exactly 0. + self.assertTrue(any(v != 0 for v in row["r"]), row["r"]) + + def test_fused_transforms_with_different_lengths(self): + # SPARK-27052: two transforms over arrays of different per-row lengths and null layouts get + # fused by ExtractPythonUDFs into one element-wise batch. Each UDF must be re-nested by its + # own array's shape, not a shared one. + df = self.spark.createDataFrame( + [([1, 2, 3], [10]), ([], None), (None, [7, 8])], + "a array<int>, b array<int>", + ) + f = udf(lambda v: v + 1, IntegerType()) + assertDataFrameEqual( + df.select( + sf.transform("a", lambda x: f(x)).alias("ra"), + sf.transform("b", lambda x: f(x)).alias("rb"), + ), + [([2, 3, 4], [11]), ([], None), (None, [8, 9])], + ) + + def test_element_and_return_types(self): + df = self.spark.createDataFrame([(["a", "bb"],)], "values array<string>") + upper_len = udf(lambda s: len(s), IntegerType()) + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: upper_len(x)).alias("r")), + [([1, 2],)], + ) + + df2 = self.spark.createDataFrame([([1.5, 2.5],)], "values array<double>") + half = udf(lambda v: v / 2, DoubleType()) + assertDataFrameEqual( + df2.select(sf.transform("values", lambda x: half(x)).alias("r")), + [([0.75, 1.25],)], + ) + + # A UDF returning a non-atomic type, so the re-nesting handles nested lists. + df3 = self.spark.createDataFrame([([1, 2],)], "values array<int>") + repeat = udf(lambda v: [v, v], ArrayType(IntegerType())) + assertDataFrameEqual( + df3.select(sf.transform("values", lambda x: repeat(x)).alias("r")), + [([[1, 1], [2, 2]],)], + ) + + def test_long_arrays_and_many_rows(self): + # Exercises batching: the wrapper evaluates all elements of a batch in one pass. + df = self.spark.range(0, 200).select( + sf.transform( + sf.sequence(sf.lit(1), sf.lit(20)), lambda x: (x + sf.col("id")).cast("int") + ).alias("values") + ) + plus_one = udf(lambda x: x + 1, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")), + df.select(sf.transform("values", lambda x: x + 1).alias("r")), + ) + + def test_all_null_rows(self): + df = self.spark.createDataFrame([(None,), (None,)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")), + [(None,), (None,)], + ) + + def test_empty_dataframe(self): + df = self.spark.createDataFrame([], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + self.assertEqual( + df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")).count(), 0 + ) + + def test_udf_inside_nested_lambda(self): + # A UDF in a *nested* lambda is lifted onto the fully flattened leaves and the nested + # structure is rebuilt around the result: `transform(matrix, row -> transform(row, x -> + # f(x)))` applies `f` to every leaf. The UDF runs once over all leaves (a depth-2 lift), and + # the result is compared against the equivalent native expression. + df = self.spark.createDataFrame( + [([[1, 2], [3]],), ([[], [4, 5]],), (None,), ([None, [6]],)], + "values array<array<int>>", + ) + plus_one = udf(lambda x: x + 1, IntegerType()) + is_even = udf(lambda x: x % 2 == 0, "boolean") + + # transform inside transform. + assertDataFrameEqual( + df.select( + sf.transform("values", lambda row: sf.transform(row, lambda x: plus_one(x))).alias( + "r" + ) + ), + df.select( + sf.transform("values", lambda row: sf.transform(row, lambda x: x + 1)).alias("r") + ), + ) + # filter inside transform: the inner result length differs from the input, but the UDF still + # runs over every leaf before the (JVM) filtering. + assertDataFrameEqual( + df.select( + sf.transform("values", lambda row: sf.filter(row, lambda x: is_even(x))).alias("r") + ), + df.select( + sf.transform("values", lambda row: sf.filter(row, lambda x: x % 2 == 0)).alias("r") + ), + ) + + def test_udf_inside_nested_lambda_capturing_outer_variable(self): + # The inner UDF reads both the inner element and the *enclosing* lambda's variable + # (`sf.size(row)`, where `row` is the outer element), so its argument depends on two nesting + # levels; the lift aligns both onto the leaves. Compared against the equivalent native + # expression, over null outer rows, null inner arrays, and empty inner arrays. + df = self.spark.createDataFrame( + [([[1, 2], [3, 4]],), (None,), ([[]],), ([None, [5]],)], + "values array<array<int>>", + ) + add = udf(lambda a, b: a + b, IntegerType()) + assertDataFrameEqual( + df.select( + sf.transform( + "values", lambda row: sf.transform(row, lambda x: add(x, sf.size(row))) + ).alias("r") + ), + df.select( + sf.transform( + "values", lambda row: sf.transform(row, lambda x: x + sf.size(row)) + ).alias("r") + ), + ) + + def test_udf_inside_three_level_nested_lambda(self): + # Three levels of nesting: `f` is lifted to a depth-3 element-wise UDF, so the worker + # flattens three array levels to the leaves and re-nests three levels. + df = self.spark.createDataFrame( + [([[[1, 2], [3]], [[4]]],), (None,), ([[[], None]],)], + "values array<array<array<int>>>", + ) + plus_one = udf(lambda x: x + 1, IntegerType()) + assertDataFrameEqual( + df.select( + sf.transform( + "values", + lambda a: sf.transform(a, lambda b: sf.transform(b, lambda x: plus_one(x))), + ).alias("r") + ), + df.select( + sf.transform( + "values", + lambda a: sf.transform(a, lambda b: sf.transform(b, lambda x: x + 1)), + ).alias("r") + ), + ) + + def test_vectorized_udf_inside_nested_lambda(self): + # All four vectorized flavors lifted out of a nested lambda (depth 2): the worker flattens + # two array levels to the leaves, runs the native-batch function once, and re-nests two + # levels. Includes null outer rows, null inner arrays, and empty inner arrays. + from typing import Iterator + + import pandas as pd + import pyarrow as pa + + from pyspark.sql.functions import arrow_udf, pandas_udf + + df = self.spark.createDataFrame( + [([[1, 2], [3]],), ([[], [4, 5]],), (None,), ([None, [6]],)], + "values array<array<int>>", + ) + + @pandas_udf(IntegerType()) + def plus_one_pandas(s: pd.Series) -> pd.Series: + return s + 1 + + @arrow_udf(IntegerType()) + def plus_one_arrow(a: pa.Array) -> pa.Array: + return pa.compute.add(a, 1) + + @pandas_udf(IntegerType()) + def plus_one_pandas_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield s + 1 + + @arrow_udf(IntegerType()) + def plus_one_arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]: + for a in it: + yield pa.compute.add(a, 1) + + native = df.select( + sf.transform("values", lambda row: sf.transform(row, lambda x: x + 1)).alias("r") + ) + for f in (plus_one_pandas, plus_one_arrow, plus_one_pandas_iter, plus_one_arrow_iter): + assertDataFrameEqual( + df.select( + sf.transform("values", lambda row: sf.transform(row, lambda x: f(x))).alias("r") + ), + native, + ) + + def test_nested_lambda_with_nondeterministic_inner_argument_fails(self): + # The inner iterated argument is nondeterministic, so the rewrite - which references it more + # than once - would evaluate it independently and misalign the results. It must keep failing + # analysis at the nest root, even though the UDF itself is liftable. + df = self.spark.createDataFrame([([[1, 2], [3]],)], "values array<array<int>>") + plus_one = udf(lambda x: x + 1, IntegerType()) + with self.assertRaises(AnalysisException) as ctx: + df.select( + sf.transform( + "values", lambda row: sf.transform(sf.shuffle(row), lambda x: plus_one(x)) + ) + ).collect() + self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception)) + + def test_udf_outside_inner_higher_order_function(self): + # The UDF applies to the outer array's element (itself an array), which *is* a real + # column, so this is rewritable even though a higher-order function is also present. + df = self.spark.createDataFrame([([[1, 2], [3]],)], "values array<array<int>>") + total = udf(lambda a: sum(a), IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda inner: total(inner)).alias("r")), + [([3, 3],)], + ) + + def test_rewritable_higher_order_function_inside_outer_lambda(self): + # A rewritable inner HOF over a *real* column, sitting inside an outer HOF's lambda: + # `transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)`. Analysis accepts it + # (the inner `transform` iterates the real column `arr`, not the outer lambda variable), + # the inner UDF is lifted out of the inner lambda, and the resulting element-wise UDF stays + # inside the outer lambda for `ExtractPythonUDFs` to extract per row. The result must match + # the native computation for a deterministic UDF. + df = self.spark.createDataFrame([([1, 2, 3], [10, 20])], "arr array<int>, arr2 array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + + assertDataFrameEqual( + df.select( + sf.transform( + "arr2", + lambda i: sf.array_max(sf.transform("arr", lambda x: plus_one(x))) + i, + ).alias("r") + ), + # array_max([2, 3, 4]) = 4; 4 + 10 = 14, 4 + 20 = 24 + [([14, 24],)], + ) + + def test_sql_string_syntax(self): + # The rewrite is on the analyzed plan, so it must fire for the SQL string syntax too, not + # only the DataFrame API. + self.spark.udf.register("py_plus_one", udf(lambda x: x + 1, IntegerType())) + with self.temp_view("t"): + self.spark.createDataFrame([([1, 2, 3],)], "values array<int>").createOrReplaceTempView( + "t" + ) + assertDataFrameEqual( + self.spark.sql("SELECT transform(values, x -> py_plus_one(x)) AS r FROM t"), + [([2, 3, 4],)], + ) + + def test_kwargs_call(self): + # A UDF called with a keyword argument is lifted too: the NamedArgumentExpression stays a + # direct child of the lifted UDF (only its value becomes an aligned array), so the runner + # still derives the kwargs mapping. add(x, y=10) = x + 10. + df = self.spark.createDataFrame([([1, 2],), (None,)], "values array<int>") + add = udf(lambda x, y: x + y, IntegerType()) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: add(x, y=sf.lit(10))).alias("r")), + [([11, 12],), (None,)], + ) + + def test_zero_argument_udf_still_fails(self): + # A zero-argument UDF has no argument to carry the iterated array's shape, so the rewrite + # cannot express it; it must keep failing analysis rather than crash the worker at runtime. + df = self.spark.createDataFrame([([1, 2],)], "values array<int>") + const = udf(lambda: 7, IntegerType()) + + with self.assertRaises(AnalysisException) as ctx: + df.select(sf.transform("values", lambda x: const())).collect() + self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception)) + + def test_nondeterministic_iterated_argument_still_fails(self): + # The rewrite references the iterated argument several times; a nondeterministic one (e.g. + # shuffle) would be evaluated independently per reference and misalign the results, so it + # must fail analysis instead of being rewritten. + df = self.spark.createDataFrame([([1, 2, 3, 4],)], "values array<int>") + is_even = udf(lambda x: x % 2 == 0, "boolean") + + with self.assertRaises(AnalysisException) as ctx: + df.select(sf.filter(sf.shuffle("values"), lambda x: is_even(x))).collect() + self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception)) + + def test_decimal_timestamp_and_struct_element_types(self): + # Arrow conversion edge cases beyond int/double/string: decimal, timestamp, and a struct + # return type. Each must round-trip through the element-wise wrapper correctly. + import datetime + from decimal import Decimal + + from pyspark.sql.types import ( + DecimalType, + StructField, + StructType, + TimestampType, + ) + + dec_df = self.spark.createDataFrame( + [([Decimal("1.50"), Decimal("2.25")],)], "values array<decimal(5,2)>" + ) + add_half = udf(lambda v: v + Decimal("0.50"), DecimalType(5, 2)) + assertDataFrameEqual( + dec_df.select(sf.transform("values", lambda x: add_half(x)).alias("r")), + [([Decimal("2.00"), Decimal("2.75")],)], + ) + + ts_df = self.spark.createDataFrame( + [([datetime.datetime(2020, 1, 1, 0, 0, 0)],)], "values array<timestamp>" + ) + add_day = udf(lambda t: t + datetime.timedelta(days=1), TimestampType()) + assertDataFrameEqual( + ts_df.select(sf.transform("values", lambda x: add_day(x)).alias("r")), + [([datetime.datetime(2020, 1, 2, 0, 0, 0)],)], + ) + + struct_type = StructType([StructField("a", IntegerType()), StructField("b", IntegerType())]) + int_df = self.spark.createDataFrame([([1, 2],)], "values array<int>") + to_struct = udf(lambda v: (v, v * 10), struct_type) + assertDataFrameEqual( + int_df.select(sf.transform("values", lambda x: to_struct(x)).alias("r")), + [([(1, 10), (2, 20)],)], + ) + + def test_integration_with_joins_grouping_and_caching(self): + left = self.spark.createDataFrame( + [ + (1, [1, 2]), + ( + 2, + [3], + ), + ], + "k int, values array<int>", + ) + right = self.spark.createDataFrame([(1,), (2,)], "k int") + plus_one = udf(lambda x: x + 1, IntegerType()) + + joined = left.join(right, "k").select( + "k", sf.transform("values", lambda x: plus_one(x)).alias("r") + ) + assertDataFrameEqual(joined, [(1, [2, 3]), (2, [4])]) + + cached = left.select(sf.transform("values", lambda x: plus_one(x)).alias("r")).cache() + try: + assertDataFrameEqual(cached, [([2, 3],), ([4],)]) + finally: + cached.unpersist() + + grouped = left.groupBy().agg( + sf.sum( + sf.aggregate( + sf.transform("values", lambda x: plus_one(x)), sf.lit(0), lambda a, x: a + x + ) + ).alias("s") + ) + # (1+1)+(2+1) + (3+1) = 9 + assertDataFrameEqual(grouped, [(9,)]) + + def test_mixed_with_plain_python_udf(self): + df = self.spark.createDataFrame([([1, 2],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + size_udf = udf(lambda a: len(a), IntegerType()) + + assertDataFrameEqual( + df.select( + sf.transform("values", lambda x: plus_one(x)).alias("r"), + size_udf("values").alias("n"), + ), + [([2, 3], 2)], + ) + + def test_lambda_without_udf_is_unchanged(self): + # The rewrite must be inert for plans that contain no Python UDF in a lambda. + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + native = df.select(sf.transform("values", lambda x: x + 1).alias("r")) + self.assertNotIn("pythonUDF", native._jdf.queryExecution().optimizedPlan().toString()) + assertDataFrameEqual(native, [([2, 3, 4],)]) + + def test_disabled_by_conf(self): + df = self.spark.createDataFrame([([1, 2],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + with self.sql_conf({"spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled": False}): + with self.assertRaises(AnalysisException) as ctx: + df.select(sf.transform("values", lambda x: plus_one(x))).collect() + self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception)) + + def test_udf_in_aggregate_fails(self): + # `aggregate` / `reduce` is a sequential fold: the values a UDF sees are outputs of earlier + # steps, not elements of a collection, so it cannot be applied once to the whole array. A + # UDF anywhere in `aggregate` / `reduce` - `merge` or `finish` - must fail analysis. + df = self.spark.createDataFrame([([1, 2],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType()) + double_it = udf(lambda acc: acc * 2, IntegerType()) + + aggregates = [ + sf.aggregate("values", sf.lit(0), lambda acc, x: acc + plus_one(x)), + sf.aggregate("values", sf.lit(0), lambda acc, x: plus_one(acc) + x), + sf.aggregate("values", sf.lit(0), lambda acc, x: acc + x, lambda acc: double_it(acc)), + # `reduce` is an alias of `aggregate`, so it is rejected the same way. + sf.reduce("values", sf.lit(0), lambda acc, x: acc + plus_one(x)), + ] + for agg in aggregates: + with self.assertRaises(AnalysisException) as ctx: + df.select(agg).collect() + self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception)) + + def test_scalar_pandas_udf_in_lambda(self): + # A vectorized scalar pandas UDF is lifted and applied over the flattened elements, so it + # still receives a pandas Series (its native contract) once per batch, not per element. + import pandas as pd + + from pyspark.sql.functions import pandas_udf + + df = self.spark.createDataFrame( + [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>" + ) + + @pandas_udf(IntegerType()) + def plus_one_pandas(s: pd.Series) -> pd.Series: + return s + 1 + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one_pandas(x)).alias("r")), + [([2, 3, 4],), ([],), (None,), ([11, None],)], + ) + # Arithmetic composition around the UDF result is ordinary JVM work. + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one_pandas(x) * 2).alias("r")), + [([4, 6, 8],), ([],), (None,), ([22, None],)], + ) + + def test_scalar_arrow_udf_in_lambda(self): + # A vectorized scalar Arrow UDF is lifted the same way; it takes and returns a pyarrow + # Array over the flattened elements. + import pyarrow as pa + + from pyspark.sql.functions import arrow_udf + + df = self.spark.createDataFrame( + [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>" + ) + + @arrow_udf(IntegerType()) + def plus_one_arrow(a: pa.Array) -> pa.Array: + return pa.compute.add(a, 1) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one_arrow(x)).alias("r")), + [([2, 3, 4],), ([],), (None,), ([11, None],)], + ) + assertDataFrameEqual( + df.select(sf.filter("values", lambda x: plus_one_arrow(x) > 2).alias("r")), + [([2, 3],), ([],), (None,), ([10],)], + ) + + def test_scalar_pandas_iter_udf_in_lambda(self): + # A scalar iterator pandas UDF keeps its iterator contract: it consumes and produces an + # iterator of Series. The worker feeds it the flattened elements and re-groups the streamed + # results back into arrays positionally, so output batch boundaries need not match input. + from typing import Iterator + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + + df = self.spark.createDataFrame( + [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values array<int>" + ) + + @pandas_udf(IntegerType()) + def plus_one_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield s + 1 + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one_iter(x)).alias("r")), + [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)], + ) + + def test_scalar_arrow_iter_udf_in_lambda(self): + # A scalar iterator Arrow UDF, lifted the same way as the pandas iterator variant. + from typing import Iterator + + import pyarrow as pa + + from pyspark.sql.functions import arrow_udf + + df = self.spark.createDataFrame( + [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values array<int>" + ) + + @arrow_udf(IntegerType()) + def plus_one_arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]: + for a in it: + yield pa.compute.add(a, 1) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one_arrow_iter(x)).alias("r")), + [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)], + ) + + def test_scalar_pandas_iter_udf_multiple_arguments_differ_in_type(self): + # A two-argument iterator UDF whose arguments have different element types: the array + # element (int) and an outer column (string) that the rewrite repeats into an aligned + # array. Each argument must be flattened with its own element type, not the first's. + from typing import Iterator, Tuple + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + + df = self.spark.createDataFrame( + [([1, 2, 3], "a"), ([], "b"), (None, "c"), ([10], "d")], + "values array<int>, tag string", + ) + + @pandas_udf(StringType()) + def tag_each(it: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]: + for x, t in it: + yield t + x.astype("string") + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: tag_each(x, sf.col("tag"))).alias("r")), + [(["a1", "a2", "a3"],), ([],), (None,), (["d10"],)], + ) + + def test_scalar_pandas_iter_udf_timestamp_return_type(self): + # A timestamp-returning pandas iterator UDF with a non-UTC session timezone: the result + # chunks are typed with the session timezone, so the streamed buffer must take its type + # from the first chunk rather than assuming UTC, or pa.concat_arrays would fail. Assert both + # against the equivalent non-iterator pandas UDF (isolates the concat fix) and against a + # native Spark expression computing the same instants (so a timezone bug common to both UDF + # paths would still be caught, while going through identical driver-collection semantics). + from typing import Iterator + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + from pyspark.sql.types import TimestampType + + with self.sql_conf({"spark.sql.session.timeZone": "America/Los_Angeles"}): + df = self.spark.createDataFrame([([1, 2],), (None,), ([3],)], "values array<int>") + + def compute(x): + return pd.to_datetime(x, unit="D", origin="2020-01-01") + + @pandas_udf(TimestampType()) + def to_ts(s: pd.Series) -> pd.Series: + return compute(s) + + @pandas_udf(TimestampType()) + def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield compute(s) + + iter_df = df.select(sf.transform("values", lambda x: to_ts_iter(x)).alias("r")) + # Native equivalent: pandas interprets the tz-naive origin in the session timezone, so + # `timestamp_add(DAY, x, TIMESTAMP '2020-01-01 00:00:00')` (also session-local) matches. + native_df = df.select( + sf.transform( + "values", + lambda x: sf.timestamp_add("DAY", x, sf.lit("2020-01-01").cast("timestamp")), + ).alias("r") + ) + # Consistent with the non-iterator pandas UDF, and with the native instants. + assertDataFrameEqual( + iter_df, + df.select(sf.transform("values", lambda x: to_ts(x)).alias("r")), + ) + assertDataFrameEqual(iter_df, native_df) + + def test_scalar_iter_udf_over_all_empty_and_null_partition(self): + # SPARK-58695: when a whole partition holds only empty/null arrays, the flattened inputs are + # all zero-length and a skip-empty iterator UDF yields no chunks. Those rows still need one + # (empty / null) output row each, or the positional JVM join drops them silently. Cover both + # the pandas and Arrow iterator flavors. + from typing import Iterator + + import pandas as pd + import pyarrow as pa + + from pyspark.sql.functions import arrow_udf, pandas_udf + + # Single partition so the whole batch is empty/null arrays. + df = self.spark.createDataFrame( + [([],), (None,), ([],), (None,)], "values array<int>" + ).coalesce(1) + + @pandas_udf(IntegerType()) + def skip_empty_pandas(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + if len(s) == 0: + continue + yield s + 1 + + @arrow_udf(IntegerType()) + def skip_empty_arrow(it: Iterator[pa.Array]) -> Iterator[pa.Array]: + for a in it: + if len(a) == 0: + continue + yield pa.compute.add(a, 1) + + for f in (skip_empty_pandas, skip_empty_arrow): + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: f(x)).alias("r")), + [([],), (None,), ([],), (None,)], + ) + + def test_scalar_pandas_iter_udf_timestamp_after_empty_batch(self): + # A zero-length result chunk (from an input batch holding only empty/null arrays) must not + # pin the output stream's timestamp type to the UTC-typed default: the rows it emits and + # the rows emitted from a later real chunk (typed with the session timezone) would then + # disagree, and the Arrow stream writer would reject the second output batch. Assert against + # the equivalent non-iterator pandas UDF (identical driver-collection semantics), so the + # check proves the schema fix without depending on absolute timezone offsets. + from typing import Iterator + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + from pyspark.sql.types import TimestampType + + with self.sql_conf( + { + "spark.sql.session.timeZone": "America/Los_Angeles", + # One row per Arrow batch so the empty-array row forms its own (first) batch. + "spark.sql.execution.arrow.maxRecordsPerBatch": "1", + } + ): + df = self.spark.createDataFrame([([],), ([1],)], "values array<int>").coalesce(1) + + def compute(x): + return pd.to_datetime(x, unit="D", origin="2020-01-01") + + @pandas_udf(TimestampType()) + def to_ts(s: pd.Series) -> pd.Series: + return compute(s) + + @pandas_udf(TimestampType()) + def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield compute(s) + + # Without the fix this raises ArrowInvalid ("different schema") writing the second + # output batch; with it the iterator result matches the non-iterator pandas UDF. + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: to_ts_iter(x)).alias("r")), + df.select(sf.transform("values", lambda x: to_ts(x)).alias("r")), + ) + + def test_scalar_pandas_udf_struct_element_return_type(self): + # A vectorized pandas UDF returning a struct element (a pandas.DataFrame per batch) inside a + # lambda. Covers the struct-DataFrame result path of the element-wise pandas branch. + import pandas as pd + + from pyspark.sql.functions import pandas_udf + from pyspark.sql.types import StructField, StructType + + df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)], "values array<int>") + ret = StructType([StructField("v", IntegerType()), StructField("neg", IntegerType())]) + + @pandas_udf(ret) + def to_struct(s: pd.Series) -> pd.DataFrame: + return pd.DataFrame({"v": s, "neg": -s}) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: to_struct(x)).alias("r")), + [ + ([(1, -1), (2, -2), (3, -3)],), + ([],), + (None,), + ], + ) + + def test_chained_vectorized_udfs_in_lambda(self): + # Nested calls f(g(x)) inside a lambda: g is lifted first, then f consumes g's array result. + # Cover both the non-iterator and iterator vectorized flavors. + from typing import Iterator + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + + df = self.spark.createDataFrame( + [([1, 2, 3],), ([],), (None,), ([10],)], "values array<int>" + ) + + @pandas_udf(IntegerType()) + def plus_one(s: pd.Series) -> pd.Series: + return s + 1 + + @pandas_udf(IntegerType()) + def times_two(s: pd.Series) -> pd.Series: + return s * 2 + + @pandas_udf(IntegerType()) + def plus_one_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield s + 1 + + @pandas_udf(IntegerType()) + def times_two_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]: + for s in it: + yield s * 2 + + # (x + 1) * 2, verified against the equivalent native expression. + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: times_two(plus_one(x))).alias("r")), + df.select(sf.transform("values", lambda x: (x + 1) * 2).alias("r")), + ) + assertDataFrameEqual( + df.select( + sf.transform("values", lambda x: times_two_iter(plus_one_iter(x))).alias("r") + ), + df.select(sf.transform("values", lambda x: (x + 1) * 2).alias("r")), + ) + + def test_scalar_iter_udf_struct_element_return_type(self): + # A scalar iterator pandas UDF returning a struct element (a pandas.DataFrame per batch) + # inside a lambda. Covers the struct-DataFrame result path of the iterator element-wise + # branch (Iterator[pd.DataFrame] contract). + from typing import Iterator + + import pandas as pd + + from pyspark.sql.functions import pandas_udf + from pyspark.sql.types import StructField, StructType + + df = self.spark.createDataFrame([([1, 2, 3],), ([],), (None,)], "values array<int>") + ret = StructType([StructField("v", IntegerType()), StructField("neg", IntegerType())]) + + @pandas_udf(ret) + def to_struct_iter(it: Iterator[pd.Series]) -> Iterator[pd.DataFrame]: + for s in it: + yield pd.DataFrame({"v": s, "neg": -s}) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: to_struct_iter(x)).alias("r")), + [ + ([(1, -1), (2, -2), (3, -3)],), + ([],), + (None,), + ], + ) + + def test_non_arrow_udf_is_also_supported(self): + # A UDF created with useArrow=False is still rewritable; the generated array wrapper + # is Arrow-based regardless of how the user's UDF was declared. + df = self.spark.createDataFrame([([1, 2, 3],)], "values array<int>") + plus_one = udf(lambda x: x + 1, IntegerType(), useArrow=False) + + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: plus_one(x)).alias("r")), + [([2, 3, 4],)], + ) + + def test_string_return_type(self): + df = self.spark.createDataFrame([([1, 2],), (None,)], "values array<int>") + to_str = udf(lambda x: f"v{x}", StringType()) + assertDataFrameEqual( + df.select(sf.transform("values", lambda x: to_str(x)).alias("r")), + [(["v1", "v2"],), (None,)], + ) + + +class UDFInHigherOrderFunctionTests(UDFInHigherOrderFunctionTestsMixin, ReusedSQLTestCase): + pass + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/test_udf_profiler.py b/python/pyspark/sql/tests/test_udf_profiler.py index feda3958224ca..b435574d8887d 100644 --- a/python/pyspark/sql/tests/test_udf_profiler.py +++ b/python/pyspark/sql/tests/test_udf_profiler.py @@ -15,27 +15,27 @@ # limitations under the License. # -from contextlib import contextmanager import inspect -import tempfile -import unittest import os import sys +import tempfile +import unittest import warnings +from contextlib import contextmanager from io import StringIO from typing import Iterator from pyspark import SparkConf from pyspark.errors import PySparkValueError +from pyspark.profiler import UDFBasicProfiler from pyspark.sql import SparkSession -from pyspark.sql.functions import col, arrow_udf, pandas_udf, udf +from pyspark.sql.functions import arrow_udf, col, pandas_udf, udf from pyspark.sql.window import Window -from pyspark.profiler import UDFBasicProfiler from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( + have_flameprof, have_pandas, have_pyarrow, - have_flameprof, pandas_requirement_message, pyarrow_requirement_message, ) @@ -526,9 +526,10 @@ def min_udf(v: pa.Array) -> float: @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) def test_perf_profiler_arrow_udf_grouped_agg_iter(self): - import pyarrow as pa from typing import Iterator + import pyarrow as pa + @arrow_udf("double") def arrow_mean_iter(it: Iterator[pa.Array]) -> float: sum_val = 0.0 diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py new file mode 100644 index 0000000000000..91d90106f271a --- /dev/null +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -0,0 +1,711 @@ +# +# 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. +# +""" +Differential / property-based tests for UDF transpilation (SPARK-54783). + +These tests run a fixed set of small Python UDFs twice -- once with +``spark.sql.experimental.optimizer.transpilePyUDFs`` enabled (so the catalyst +transpiler in :mod:`pyspark.sql.transpile` rewrites them into native +expressions) and once without -- and assert that the two runs produce the +same results for inputs generated by Hypothesis. + +The transpiler is intentionally minimal at this point so we expect this +suite to surface bugs (e.g. truthiness / NULL-handling mismatches between +Python's ``if x:`` semantics and SQL's ``CASE WHEN``). Failures here should +be treated as real correctness gaps in the transpiler, not as test bugs to +silence. + +The suite is gated on two things, both required: + +* the ``RUN_HYPOTHESIS`` env var must be set to a truthy value + (``1``, ``true``, or ``yes``, case-insensitive), and +* the ``hypothesis`` package must be installed. + +The gate is value-based rather than presence-based because CI always sets +``RUN_HYPOTHESIS`` (to ``"true"`` or ``"false"``) via the transpile +precondition in ``build_and_test.yml``; mere presence must not opt in, or +the slow suite would run on every PySpark job. + +If either gate is unmet the entire suite is skipped cleanly so it never +becomes a CI tax for folks who haven't opted in. In CI, this opt-in suite +is wired through ``.github/workflows/build_and_test.yml``, which flips both +gates on for the relevant job when PR changes touch the transpiler or this +test file. + +To run locally:: + + pip install hypothesis + RUN_HYPOTHESIS=1 RUN_HYPOTHESIS_MAX_EXAMPLES=1000 \ + python/run-tests --testnames pyspark.sql.tests.test_udf_transpile_hypothesis + +Set ``RUN_HYPOTHESIS_MAX_EXAMPLES`` to override the per-test example count +(default 1000). Each generated example runs two full Spark jobs (a +transpiled-vs-interpreted differential), so CI caps this at 50 via +``build_and_test.yml`` to stay under ``PYSPARK_TEST_TIMEOUT``; the explicit +``@example`` edge seeds always run on top of the generated ones regardless. +""" + +import os +import unittest +import warnings +from typing import Optional + +from pyspark.sql import Row +from pyspark.sql.types import ( + BooleanType, + LongType, + StructField, + StructType, +) +from pyspark.sql.udf import UserDefinedFunction +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import have_package +from pyspark.util import is_remote_only + +# Sentinel value used by ``_run`` to mark "this side raised". A unique +# object is sufficient because we only ever compare it against itself +# inside the helper. +_SENTINEL_RAISED = object() + + +_HYPOTHESIS_ENV = "RUN_HYPOTHESIS" +_have_hypothesis = have_package("hypothesis") + + +def _env_opts_in(value: Optional[str]) -> bool: + """Value-based opt-in: only 1/true/yes (case-insensitive) enable the suite. + + CI always sets ``RUN_HYPOTHESIS`` -- to ``"true"`` or ``"false"`` -- via the + transpile precondition in ``build_and_test.yml``, so a presence-based check + would run this very slow suite on every PySpark job regardless of the + gating decision. + """ + return value is not None and value.strip().lower() in ("1", "true", "yes") + + +_hypothesis_enabled = _env_opts_in(os.environ.get(_HYPOTHESIS_ENV)) +# Transpilation is only supported in regular (non-Connect) Spark for now, +# so the hypothesis suite skips cleanly under a pyspark-client-only install. +_regular_spark = not is_remote_only() +_skip_reason = ( + f"Set {_HYPOTHESIS_ENV}=1 (or true/yes) to run; hypothesis must also be installed, " + "and the suite only runs under regular (non-Connect) Spark." +) + + +if _have_hypothesis: + from hypothesis import HealthCheck, example, given, settings + from hypothesis import strategies as st + + _DEFAULT_MAX_EXAMPLES = int(os.environ.get("RUN_HYPOTHESIS_MAX_EXAMPLES", "1000")) + + # The ``function_scoped_fixture` health check is suppressed because we intentionally reuse the + # class-level SparkSession across examples; the per-example ``deadline`` is disabled because + # Spark task execution is much slower than hypothesis's default budget. + _hyp_settings = settings( + max_examples=_DEFAULT_MAX_EXAMPLES, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + + # Full 64-bit signed range -- used by comparison and equality tests where + # no arithmetic can overflow. + _LONG_BOUND = 2**63 - 1 + _long_strategy = st.one_of( + st.none(), st.integers(min_value=-_LONG_BOUND, max_value=_LONG_BOUND) + ) + + # Narrower range for arithmetic tests (+4, -2, *3, +7, x+y). Python's + # arithmetic never overflows, but Spark's ANSI mode raises on LongType + # overflow. Worse, the Python UDF runner silently wraps out-of-range + # return values even with ANSI=True, so "both raised" never fires for + # overflow values and the test sees a spurious mismatch instead. + # 2**61 is safe for all operations: (2**61)*3 ~= 6.9e18 < 9.2e18 = Long.MAX. + _LONG_ARITH_BOUND = 2**61 + _long_arith_strategy = st.one_of( + st.none(), + st.integers(min_value=-_LONG_ARITH_BOUND, max_value=_LONG_ARITH_BOUND), + ) + + _bool_strategy = st.one_of(st.none(), st.booleans()) + + # 32-bit signed boundaries. Values round-trip through LongType, but the + # int32 limits are where narrowing / off-by-one bugs in parameter-index + # plumbing and boundary handling tend to hide, so we always seed them. + _INT32_MAX = 2**31 - 1 + _INT32_MIN = -(2**31) + + # ---- Edge-case seeds (scalacheck-style) ----------------------------- + # + # Hypothesis already biases toward "interesting" boundary values, but + # explicit ``@example`` decorators make a regression on a specific + # value -- e.g. NULL, zero, the type's max -- deterministic across + # runs. These are the values we always want to try, before random + # generation kicks in. + _LONG_EDGES = (None, 0, 1, -1, 7, -7, _INT32_MAX, _INT32_MIN, _LONG_BOUND, -_LONG_BOUND) + _LONG_ARITH_EDGES = (None, 0, 1, -1, 7, -7, _LONG_ARITH_BOUND, -_LONG_ARITH_BOUND) + # Bool space is exhaustive (only three values) so the @example + # decorators here serve more as documentation of the NULL handling + # we care about than as new coverage on top of hypothesis's + # generator. + _BOOL_EDGES = (None, True, False) + # Multi-arg edges -- nulls plus the four sign-combos for non-zero + # values. Catches off-by-one errors in parameter-index plumbing + # better than random generation alone. + _LONG_PAIR_EDGES = ( + (None, None), + (None, 0), + (0, None), + (0, 0), + (1, -1), + (-1, 1), + (_INT32_MAX, _INT32_MIN), + (_INT32_MIN, _INT32_MAX), + (_LONG_BOUND, 1), + (1, -_LONG_BOUND), + ) + _LONG_ARITH_PAIR_EDGES = ( + (None, None), + (None, 0), + (0, None), + (0, 0), + (1, -1), + (-1, 1), + (_LONG_ARITH_BOUND, 1), + (1, -_LONG_ARITH_BOUND), + ) + # Sign-combo edges (plus NULL combinations) for the boolean tests. + # The bodies (``x > 0 and y > 0`` / ``x > 0 or y > 0``) raise in + # pure Python on a None input (``TypeError``), and the transpiler's + # NULL-guarded Compare also raises -- so the ``_run`` helper's "both + # raised" equivalence covers the NULL cases here. + _BOOLEAN_PAIR_EDGES = ( + (None, None), + (None, 0), + (0, None), + (0, 0), + (1, -1), + (-1, 1), + (1, 1), + (-1, -1), + (_LONG_BOUND, 1), + (1, -_LONG_BOUND), + ) + + def _seed_examples(values, key="value"): + """Stack one ``@example`` decorator per seed value.""" + + def wrapper(method): + for v in reversed(values): + method = example(**{key: v})(method) + return method + + return wrapper + + def _seed_pair_examples(pairs, keys=("x", "y")): + def wrapper(method): + for v0, v1 in reversed(pairs): + method = example(**{keys[0]: v0, keys[1]: v1})(method) + return method + + return wrapper + + +# ---- The UDF templates we exercise -------------------------------------- +# +# We keep these as module-level callables so that ``inspect.getsource`` works +# (the transpiler reads source via inspection). They are deliberately written +# the way a user would: idiomatic Python, including ``if x is not None`` +# guards and bare ``if x:`` truthiness checks. + + +def plus_four(x): + if x is not None: + return x + 4 + + +def plus_four_unsafe(x): + return x + 4 + + +def plus_four_with_else(x): + if x is not None: + return x + 4 + else: + return 0 + + +def is_none_branch(x): + if x is None: + return -1 + else: + return x + + +def truthy_bool_branch(x): + # The transpiler currently mishandles ``if x:`` for nullable bool inputs: + # Python treats ``None`` as falsy and takes the else branch, but a naive + # SQL lowering can produce NULL. This test is the canonical regression. + if x: + return 1 + else: + return 2 + + +def add_then_mod(x): + if x is not None: + return (x + 7) % 5 + + +def minus_two(x): + # Exercises ast.Sub. + if x is not None: + return x - 2 + + +def times_three(x): + # Exercises ast.Mult. + if x is not None: + return x * 3 + + +def negate_truthy(x): + # Exercises ast.UnaryOp(Not). Same NULL-as-falsy semantics as + # ``truthy_bool_branch`` above, just inverted. + if not x: + return 0 + else: + return 1 + + +def both_positive(x, y): + # Exercises ast.BoolOp(And) over Compare operands -- both operands + # are statically boolean, so the transpiler should lower to `&`. + # Kept as a single-statement body since the transpiler doesn't yet + # support multi-statement function bodies; NULL inputs flow through + # `>` to NULL on the Spark side and to a raise on the Python side, + # so the strategy below skips None. + return x > 0 and y > 0 + + +def either_positive(x, y): + # Exercises ast.BoolOp(Or) over Compare operands. + return x > 0 or y > 0 + + +def add_two(x, y): + # Multi-arg UDF -- exercises the parameter-index plumbing for + # functions with more than one positional argument. + if x is not None and y is not None: + return x + y + else: + return 0 + + +def eq_zero(x): + # Exercises ast.Compare(Eq) via _lower_eq with a non-None left and + # a literal 0 on the right. The None guard keeps the comparison + # itself away from NULL operands so the test stays single-path. + if x is not None: + return x == 0 + + +def neq_zero(x): + # Exercises ast.Compare(NotEq) through _lower_eq. + if x is not None: + return x != 0 + + +def eq_pair(x, y): + # Two-arg ``x == y`` exercising the full _lower_eq four-branch when + # chain that reproduces Python's None-equality semantics: + # None == None -> True; None == 0 -> False; 0 == None -> False. + # Note: no None guard, so every NULL combination runs through the + # transpiler's lowering. + return x == y + + +def neq_pair(x, y): + # Sister of ``eq_pair`` for ast.Compare(NotEq). + return x != y + + +# A lambda captured at module scope so ``inspect.getsource`` can read +# its definition. Exercises the ``ast.Lambda`` branch in +# ``_get_function_from_ast``. +lambda_plus_four = lambda x: x + 4 if x is not None else 0 # noqa: E731 + + +# ---------------------------------------------------------------------------- + + +@unittest.skipUnless(_have_hypothesis and _hypothesis_enabled and _regular_spark, _skip_reason) +class UDFTranspileHypothesisTests(ReusedSQLTestCase): + """Compare transpiled vs. interpreted Python UDF output on Hypothesis-generated inputs.""" + + # Markers we treat as "transpilation didn't actually happen" -- if any + # warning matches one of these, the differential comparison would + # collapse to interpreted-vs-interpreted and pass meaninglessly, so we + # fail loudly. + _BAD_TRANSPILE_WARNING_MARKERS = ( + "Unable to transpile", + "Errors encountered during transpilation", + "Exception transpiling", + "ANSI mode", + ) + + def _run(self, func, return_type, df, *udf_arg_columns, kwargs=None): + """Run ``func`` as a UDF with transpilation on and off, return both rows. + + ``udf_arg_columns`` and ``kwargs`` mirror what a caller would + write at the dataframe API: positional column names go in + ``udf_arg_columns`` and named-argument bindings go in ``kwargs`` + (e.g. ``kwargs={"y": "b", "x": "a"}`` to bind UDF parameter ``y`` + to column ``b`` and ``x`` to column ``a``). Use either, or both. + + Asserts the transpiled code path was actually exercised: + ``transpiled`` must be non-empty after construction, and no + transpilation-related warning may fire. Without these checks both + runs could silently fall back to interpreted Python and the + differential assertion would succeed for the wrong reason. + """ + func_name = getattr(func, "__name__", repr(func)) + kwargs = kwargs or {} + + transpile_on_conf = { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + # Transpilation requires ANSI; pin it on so the test result + # doesn't depend on the surrounding session default. + "spark.sql.ansi.enabled": True, + } + transpiled_error: Optional[Exception] = None + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with self.sql_conf(transpile_on_conf): + transpiled_udf = UserDefinedFunction(func, return_type) + self.assertTrue( + transpiled_udf.transpiled, + f"transpilation produced no Catalyst expression for " + f"{func_name!r} -- the differential comparison would be " + "meaningless without it", + ) + try: + transpiled_value = df.select( + transpiled_udf(*udf_arg_columns, **kwargs) + ).collect()[0][0] + except Exception as e: + transpiled_value = _SENTINEL_RAISED + transpiled_error = e + bad = [ + w + for w in caught + if any(marker in str(w.message) for marker in self._BAD_TRANSPILE_WARNING_MARKERS) + ] + self.assertFalse( + bad, + f"unexpected transpile warnings for {func_name!r}: {[str(w.message) for w in bad]}", + ) + + interpreted_error: Optional[Exception] = None + # Pin ANSI on for the interpreted path too so both sides see the same + # overflow semantics. Without this, the interpreted path would run with + # the ambient session default (likely False), causing LongType overflow + # to silently wrap in Python UDF results while ANSI raises on the + # transpiled path. + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": False, + "spark.sql.ansi.enabled": True, + } + ): + interpreted_udf = UserDefinedFunction(func, return_type) + try: + interpreted_value = df.select( + interpreted_udf(*udf_arg_columns, **kwargs) + ).collect()[0][0] + except Exception as e: + interpreted_value = _SENTINEL_RAISED + interpreted_error = e + + # If the transpiled path raises an exception we also need the interpreted path to raise one, + # however if the Python code (that in the interpreted path) raises an exception, the transpiled + # path may return a valid value. + if transpiled_error is not None: + self.assertIsNotNone( + interpreted_error, + f"{func_name!r}: transpiled raised {transpiled_error!r} but interpreted did not", + ) + elif interpreted_error is not None: + interpreted_value = transpiled_value + + return transpiled_value, interpreted_value + + def _single_arg_df(self, value, dtype): + schema = StructType([StructField("a", dtype, nullable=True)]) + return self.spark.createDataFrame([Row(a=value)], schema=schema) + + def _two_long_arg_df(self, x, y): + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + ] + ) + return self.spark.createDataFrame([Row(a=x, b=y)], schema=schema) + + if _have_hypothesis: + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_plus_four_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(plus_four, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"plus_four mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_plus_four_unsafe_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(plus_four_unsafe, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"plus_four mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_plus_four_with_else_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(plus_four_with_else, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"plus_four_with_else mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_is_none_branch_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(is_none_branch, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"is_none_branch mismatch on {value!r}") + + @_hyp_settings + @given(value=_bool_strategy) + @_seed_examples(_BOOL_EDGES) + def test_truthy_bool_branch_falls_back(self, value): + # `if x:` on a bare parameter name is a bare truthiness test whose + # type is unknown at transpile time. The transpiler must refuse to + # lower it (Spark's coalesce(x, false) is unsound for non-boolean + # columns) and fall back to interpreted Python instead. + df = self._single_arg_df(value, BooleanType()) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(truthy_bool_branch, LongType()) + self.assertEqual( + [], + pudf.transpiled, + "truthy_bool_branch: bare truthiness test must NOT transpile", + ) + interpreted = df.select(pudf("a")).collect()[0][0] + expected = 1 if value else 2 + self.assertEqual(interpreted, expected, f"truthy_bool_branch mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + # add_then_mod is the case that surfaced the Python-vs-SQL mod + # sign mismatch; the seed values cover the four sign combinations + # of `(x + 7) % 5` so we always re-prove the pmod fix on every + # run regardless of the random seed. + @_seed_examples((*_LONG_ARITH_EDGES, -2, -8, 8, 100, -100)) + def test_add_then_mod_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(add_then_mod, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"add_then_mod mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_minus_two_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(minus_two, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"minus_two mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_times_three_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(times_three, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"times_three mismatch on {value!r}") + + @_hyp_settings + @given(value=_bool_strategy) + @_seed_examples(_BOOL_EDGES) + def test_negate_truthy_falls_back(self, value): + # `if not x:` where x is a bare parameter name is unknown-type at + # transpile time. The transpiler must refuse (Spark's `~` is + # bitwise, not Python truthiness) and fall back to interpreted Python. + df = self._single_arg_df(value, BooleanType()) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(negate_truthy, LongType()) + self.assertEqual( + [], + pudf.transpiled, + "negate_truthy: bare `not x` must NOT transpile", + ) + interpreted = df.select(pudf("a")).collect()[0][0] + expected = 0 if not value else 1 + self.assertEqual(interpreted, expected, f"negate_truthy mismatch on {value!r}") + + @_hyp_settings + @given(x=_long_arith_strategy, y=_long_arith_strategy) + @_seed_pair_examples(_LONG_ARITH_PAIR_EDGES) + def test_add_two_matches_python(self, x, y): + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(add_two, LongType(), df, "a", "b") + self.assertEqual(transpiled, interpreted, f"add_two mismatch on (x={x!r}, y={y!r})") + + @_hyp_settings + @given(x=_long_arith_strategy, y=_long_arith_strategy) + @_seed_pair_examples(_LONG_ARITH_PAIR_EDGES) + def test_add_two_named_args_matches_python(self, x, y): + # Same UDF as above but called with kwargs (and intentionally + # in reversed order) to exercise the named-argument codepath + # documented in the udf() reference. The Python side resolves + # the kwargs to the function's positional params; the + # transpiled side has to align ``_udf_param_0`` / + # ``_udf_param_1`` with the same resolved positions, so any + # mistake here produces a swapped-argument bug. + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run( + add_two, + LongType(), + df, + kwargs={"y": "b", "x": "a"}, + ) + self.assertEqual( + transpiled, + interpreted, + f"add_two named-args mismatch on (x={x!r}, y={y!r})", + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy) + @_seed_pair_examples(_BOOLEAN_PAIR_EDGES) + def test_both_positive_matches_python(self, x, y): + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(both_positive, BooleanType(), df, "a", "b") + self.assertEqual( + transpiled, interpreted, f"both_positive mismatch on (x={x!r}, y={y!r})" + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy) + @_seed_pair_examples(_BOOLEAN_PAIR_EDGES) + def test_either_positive_matches_python(self, x, y): + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(either_positive, BooleanType(), df, "a", "b") + self.assertEqual( + transpiled, interpreted, f"either_positive mismatch on (x={x!r}, y={y!r})" + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy) + @_seed_pair_examples(_LONG_PAIR_EDGES) + def test_eq_pair_matches_python(self, x, y): + # Python's ``==`` has different NULL semantics from SQL ``=``: + # ``None == None`` is True, ``None == n`` is False. The + # transpiler's _lower_eq reproduces those semantics, so the + # transpiled and interpreted runs must agree on every NULL + # combination as well as the non-NULL cases. + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(eq_pair, BooleanType(), df, "a", "b") + self.assertEqual(transpiled, interpreted, f"eq_pair mismatch on (x={x!r}, y={y!r})") + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy) + @_seed_pair_examples(_LONG_PAIR_EDGES) + def test_neq_pair_matches_python(self, x, y): + # Sister of ``test_eq_pair_matches_python`` for the NotEq arm. + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(neq_pair, BooleanType(), df, "a", "b") + self.assertEqual(transpiled, interpreted, f"neq_pair mismatch on (x={x!r}, y={y!r})") + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_eq_zero_matches_python(self, value): + # Single-arg ``x == 0`` with a None guard, exercising _lower_eq's + # non-None-on-both-sides arm. + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(eq_zero, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"eq_zero mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_neq_zero_matches_python(self, value): + # Sister of ``test_eq_zero_matches_python`` for the NotEq arm. + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(neq_zero, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"neq_zero mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_arith_strategy) + @_seed_examples(_LONG_ARITH_EDGES) + def test_lambda_plus_four_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(lambda_plus_four, LongType(), df, "a") + self.assertEqual(transpiled, interpreted, f"lambda_plus_four mismatch on {value!r}") + + +class UDFTranspileHypothesisGatingTests(unittest.TestCase): + """Smoke tests that always run, regardless of the env gate. + + These don't talk to Spark; they just verify the gating / skipping logic + so a misconfigured environment doesn't silently skip everything forever. + """ + + def test_env_var_name_is_documented(self): + # If we ever rename the env var, the docstring needs to follow. + self.assertIn(_HYPOTHESIS_ENV, __doc__) + + def test_skip_reason_mentions_env_var(self): + self.assertIn(_HYPOTHESIS_ENV, _skip_reason) + + def test_env_gate_is_value_based(self): + # CI always sets RUN_HYPOTHESIS (to "true" or "false") via the + # transpile precondition in build_and_test.yml. If this gate ever + # regresses to presence-based, the very slow suite silently runs on + # every PySpark job. Pin the contract from both directions. + for opted_in in ("1", "true", "TRUE", "yes", " True "): + self.assertTrue(_env_opts_in(opted_in), f"{opted_in!r} must opt in") + for opted_out in (None, "", "0", "false", "FALSE", "no", "off"): + self.assertFalse(_env_opts_in(opted_out), f"{opted_out!r} must NOT opt in") + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/test_udf_transpile_parity.py b/python/pyspark/sql/tests/test_udf_transpile_parity.py new file mode 100644 index 0000000000000..b3b60c8cf431a --- /dev/null +++ b/python/pyspark/sql/tests/test_udf_transpile_parity.py @@ -0,0 +1,109 @@ +# +# 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. +# +""" +Parity tests that re-run the existing UDF test suites with Python-to-Catalyst +transpilation enabled. + +Transpilation is only attempted when both +``spark.sql.experimental.optimizer.transpilePyUDFs`` and +``spark.sql.ansi.enabled`` are true, and it is designed to fall back to +interpreted Python rather than risk semantic drift. These classes re-run the +shared UDF mixins under that configuration so we can confirm that turning on the +experimental feature does not change UDF results compared with the default +(transpilation off) runs covered by the original concrete classes +(``UDFTests``, ``UDFCombinationsTests``, ``UnifiedUDFTests``). + +Transpilation is currently only supported in regular (non-Connect) Spark, so +these classes are guarded with ``is_remote_only()`` and are intentionally not +inherited into the Spark Connect parity tests. The companion suites that test +the transpiler directly live in ``test_udf_transpile_unit.py`` and +``test_udf_transpile_hypothesis.py``. + +Note on configuration: enabling transpilation requires ANSI mode, so an "on" +run is unavoidably also an ANSI run. All inherited tests currently pass as-is +under this configuration, so no per-test overrides are defined here. If a future +change makes an inherited test diverge purely due to ANSI semantics or because +transpilation bypasses a Python-side effect (rather than a genuine result +change), override it here with a documented ``unittest.skip`` rather than +editing the inherited test body. +""" + +import unittest + +from pyspark.sql.tests.test_udf import BaseUDFTestsMixin +from pyspark.sql.tests.test_udf_combinations import UDFCombinationsTestsMixin +from pyspark.sql.tests.test_unified_udf import UnifiedUDFTestsMixin +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) +from pyspark.util import is_remote_only + +# Transpilation is gated on both of these being enabled, both at UDF +# construction time (python/pyspark/sql/udf.py) and again in the Catalyst +# optimizer (the ConvertToCatalyst rule). +# spark.conf.set requires strings, so we use "true" rather than Python True here. +_TRANSPILE_CONF = { + "spark.sql.experimental.optimizer.transpilePyUDFs": "true", + "spark.sql.ansi.enabled": "true", +} + +_NON_CONNECT_ONLY = "UDF transpilation is only supported in regular (non-Connect) Spark." + + +def _enable_transpilation(cls): + for key, value in _TRANSPILE_CONF.items(): + cls.spark.conf.set(key, value) + + +@unittest.skipIf(is_remote_only(), _NON_CONNECT_ONLY) +class TranspiledUDFParityTests(BaseUDFTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false") + _enable_transpilation(cls) + + +@unittest.skipIf(is_remote_only(), _NON_CONNECT_ONLY) +class TranspiledUDFCombinationsParityTests(UDFCombinationsTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false") + _enable_transpilation(cls) + + +@unittest.skipIf(is_remote_only(), _NON_CONNECT_ONLY) +@unittest.skipIf( + not have_pandas or not have_pyarrow, + pandas_requirement_message or pyarrow_requirement_message, +) +class TranspiledUnifiedUDFParityTests(UnifiedUDFTestsMixin, ReusedSQLTestCase): + @classmethod + def setUpClass(cls): + ReusedSQLTestCase.setUpClass() + _enable_transpilation(cls) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py new file mode 100644 index 0000000000000..fa294e3347cbb --- /dev/null +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -0,0 +1,2090 @@ +# +# 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. +# +""" +Unit tests for UDF transpilation. + +These were previously interleaved with the broader UDF mixin in +``test_udf.py``. They are split out because UDF transpilation is currently +only supported in regular (non-Connect) Spark, so they should not be +inherited into the Spark Connect parity test class. The companion +property-based suite lives in ``test_udf_transpile_hypothesis.py``. +""" + +import unittest + +from pyspark.sql import Row +from pyspark.sql.types import ( + BinaryType, + BooleanType, + DoubleType, + LongType, + StringType, +) +from pyspark.sql.udf import UserDefinedFunction +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.util import is_remote_only + +# Both flags must be on for the transpiler to attempt a rewrite (at UDF +# construction time and again in the optimizer); ANSI is required because +# transpilation targets ANSI semantics. +_TRANSPILE_ON = { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, +} + + +@unittest.skipIf( + is_remote_only(), + "UDF transpilation is only supported in regular (non-Connect) Spark.", +) +class UDFTranspileUnitTests(ReusedSQLTestCase): + def test_udf_transpile_basic(self): + # Test callable object + class PlusFour: + def __call__(self, col): + return col + 4 + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # Make sure we can transpile the object + call = PlusFour() + pudf = UserDefinedFunction(call, LongType()) + self.assertTrue(pudf.transpiled) + # Now make sure we can run the transpiled UDF* + input_df = self.spark.createDataFrame([Row(a=1)]) + transformed_df = input_df.select(pudf("a")) + [row] = transformed_df.collect() + self.assertEqual(row[0], 5) + + with self.sql_conf({"spark.sql.experimental.optimizer.transpilePyUDFs": False}): + call = PlusFour() + pudf = UserDefinedFunction(call, LongType()) + self.assertEqual([], pudf.transpiled) + # Now make sure we can run the UDF + input_df = self.spark.createDataFrame([Row(a=1)]) + transformed_df = input_df.select(pudf("a")) + [row] = transformed_df.collect() + self.assertEqual(row[0], 5) + + def test_udf_transpile_with_nones(self): + # Test callable object + class PlusFour: + def __call__(self, col): + if col is not None: + return col + 4 + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # Make sure we can transpile the object + call = PlusFour() + pudf = UserDefinedFunction(call, LongType()) + self.assertTrue(pudf.transpiled) + # Now make sure we can run the transpiled UDF* + input_df = self.spark.createDataFrame([Row(a=1)]) + transformed_df = input_df.select(pudf("a").alias("result")) + [row] = transformed_df.collect() + self.assertEqual(row[0], 5) + physical_plan = transformed_df._jdf.queryExecution().executedPlan().toString() + self.assertNotIn("UDF", physical_plan) + + with self.sql_conf({"spark.sql.experimental.optimizer.transpilePyUDFs": False}): + call = PlusFour() + pudf = UserDefinedFunction(call, LongType()) + self.assertEqual([], pudf.transpiled) + # Now make sure we can run the UDF + input_df = self.spark.createDataFrame([Row(a=1)]) + transformed_df = input_df.select(pudf("a").alias("result")) + [row] = transformed_df.collect() + self.assertEqual(row[0], 5) + physical_plan = transformed_df._jdf.queryExecution().executedPlan().toString() + self.assertIn("UDF", physical_plan) + + def test_udf_not_transpilable(self): + class UnsupportedEx: + def __call__(self, col): + if col is not None: + return col in "4" + + with self.sql_conf({"spark.sql.experimental.optimizer.transpilePyUDFs": True}): + call = UnsupportedEx() + pudf = UserDefinedFunction(call, BooleanType()) + self.assertEqual([], pudf.transpiled) + + def test_udf_transpile_requires_ansi(self): + # Transpilation targets ANSI semantics. With ANSI off the transpiler + # must skip rewriting (and warn the user) so we don't silently + # diverge from the Python interpretation; with ANSI on it should + # produce a Catalyst expression. + import warnings + + def plus_four(x): + if x is not None: + return x + 4 + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": False, + } + ): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + pudf = UserDefinedFunction(plus_four, LongType()) + self.assertEqual([], pudf.transpiled) + ansi_warnings = [w for w in caught if "ANSI mode" in str(w.message)] + self.assertTrue( + ansi_warnings, + "expected an 'ANSI mode' warning when transpilation is " + "requested but ANSI is disabled", + ) + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(plus_four, LongType()) + self.assertTrue( + pudf.transpiled, + "expected transpilation to produce a Catalyst expression " + "when both transpilePyUDFs and ANSI mode are enabled", + ) + + def test_udf_transpile_falls_back_for_unsupported_patterns(self): + # The transpiler intentionally only handles a small subset of + # Python AST today. Everything outside that subset must + # gracefully fall back to interpreted Python (with an empty + # `transpiled` list and a UserWarning) rather than break the + # UDF -- the "don't break people's Spark code" promise. This test + # walks the most common unsupported shapes, registers each as a + # UDF with transpilation on, and asserts (a) construction does + # not raise, (b) `transpiled == []`, (c) the UDF still produces + # the correct interpreted result. + + def divide_by_two(x): # `/` -- ast.Div, not handled. + if x is not None: + return x / 2 + + def floor_divide_by_two(x): # `//` -- ast.FloorDiv, not handled. + if x is not None: + return x // 2 + + def bit_and_one(x): # `&` -- ast.BitAnd, not handled. + if x is not None: + return x & 1 + + def bit_or_one(x): # `|` -- ast.BitOr, not handled. + if x is not None: + return x | 1 + + def left_shift(x): # `<<` -- ast.LShift, not handled. + if x is not None: + return x << 1 + + def multi_statement(x): # > 1 top-level statement, not handled. + y = 1 + return x + y if x is not None else 0 + + def func_closure_capture(x): + offset = 7 + if x is not None: + return x + offset + + cases = [ + ("divide_by_two", divide_by_two, DoubleType(), Row(a=4.0), 2.0), + ("floor_divide_by_two", floor_divide_by_two, LongType(), Row(a=5), 2), + ("bit_and_one", bit_and_one, LongType(), Row(a=5), 1), + ("bit_or_one", bit_or_one, LongType(), Row(a=4), 5), + ("left_shift", left_shift, LongType(), Row(a=3), 6), + ("multi_statement", multi_statement, LongType(), Row(a=5), 6), + ("func_closure_capture", func_closure_capture, LongType(), Row(a=10), 17), + ] + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + for label, func, return_type, row, expected in cases: + with self.subTest(case=label): + import warnings as _warnings + + with _warnings.catch_warnings(record=True) as caught_warnings: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, return_type) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: transpiler should not produce a Catalyst " + "expression for this AST shape", + ) + fallback = [ + w + for w in caught_warnings + if "Unable to transpile" in str(w.message) + or "Errors encountered" in str(w.message) + or "Exception transpiling" in str(w.message) + ] + self.assertTrue( + fallback, + f"{label}: expected a fallback warning when the " + "transpiler can't lower the function", + ) + df = self.spark.createDataFrame([row]) + [result] = df.select(pudf("a")).collect() + self.assertEqual( + result[0], + expected, + f"{label}: interpreted UDF result diverged from expected", + ) + + def test_udf_transpile_boolean_and_or_lowered(self): + # When `and`/`or` operands are syntactically boolean (Compare + # results in this case), the transpiler should lower to bitwise + # `&`/`|` and produce results matching the interpreted UDF. + # Each UDF is a single top-level statement (the transpiler + # doesn't support multi-statement bodies yet). + from pyspark.sql.types import StructField, StructType + + def both_positive(x, y): + return x > 0 and y > 0 + + def either_positive(x, y): + return x > 0 or y > 0 + + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + ] + ) + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # NULL inputs propagate through `>` to NULL, which then + # passes through `&` / `|` per SQL three-valued logic. We + # only assert on non-NULL inputs here since Python's + # interpreted `x > 0 and y > 0` would raise on None; the + # NULL handling itself is covered by the hypothesis suite. + for func, x, y, expected in [ + (both_positive, 1, 2, True), + (both_positive, 1, -1, False), + (both_positive, -1, -1, False), + (either_positive, -1, 2, True), + (either_positive, -1, -1, False), + (either_positive, 1, 1, True), + ]: + with self.subTest(func=func.__name__, x=x, y=y): + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue( + pudf.transpiled, + f"{func.__name__}: bool-typed and/or should transpile", + ) + df = self.spark.createDataFrame([Row(a=x, b=y)], schema=schema) + [row] = df.select(pudf("a", "b")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_less_than_zero(self): + # Restored from the unsupported-patterns matrix: now that the + # transpiler handles ast.Lt, `x < 0` should lower to a Catalyst + # expression and match interpreted Python. The ``is not None`` + # guard short-circuits None inputs through the else branch, so + # the comparison itself never sees a NULL in this UDF. + from pyspark.sql.types import StructField, StructType + + def less_than_zero(x): + if x is not None: + return x < 0 + + schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(less_than_zero, BooleanType()) + self.assertTrue(pudf.transpiled, "less_than_zero should now transpile") + for value, expected in [(-1, True), (0, False), (5, False), (None, None)]: + with self.subTest(value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_compare_with_none_raises(self): + # When a comparison's operand is NULL in Spark, Python would have + # raised TypeError ('>' not supported between NoneType and int). + # The transpiler wraps Compare ops with a raise_error guard so + # the rewritten plan fails loudly instead of silently producing + # NULL three-valued-logic results. + from pyspark.sql.types import StructField, StructType + + def gt_zero(x): + return x > 0 + + schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(gt_zero, BooleanType()) + self.assertTrue(pudf.transpiled, "gt_zero should transpile") + df = self.spark.createDataFrame([Row(a=None)], schema=schema) + with self.assertRaises(Exception) as ctx: + df.select(pudf("a")).collect() + self.assertIn("cannot compare NULL", str(ctx.exception)) + + def test_udf_transpile_eq_none_semantics(self): + # Python ``==``/``!=`` differ from Spark's three-valued NULL equality: + # in Python ``None == None`` is ``True`` and ``None == 0`` is ``False``, + # whereas SQL ``NULL = NULL`` and ``NULL = 0`` both yield ``NULL``. The + # transpiler's ``_lower_eq`` reproduces Python's semantics; this test + # exercises every arm of that logic. + from pyspark.sql.types import StructField, StructType + + def x_eq_zero(x): + if x is not None: + return x == 0 + else: + return None + + def x_neq_zero(x): + if x is not None: + return x != 0 + else: + return None + + def x_eq_y(x, y): + return x == y + + def x_neq_y(x, y): + return x != y + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + two_col_schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + ] + ) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # Single-arg ``x == 0`` / ``x != 0`` with a None guard. + pudf_eq = UserDefinedFunction(x_eq_zero, BooleanType()) + pudf_neq = UserDefinedFunction(x_neq_zero, BooleanType()) + self.assertTrue(pudf_eq.transpiled, "x == 0 should transpile") + self.assertTrue(pudf_neq.transpiled, "x != 0 should transpile") + for value, eq_expected, neq_expected in [ + (0, True, False), + (1, False, True), + (-3, False, True), + (None, None, None), + ]: + with self.subTest(value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=long_schema) + [row_eq] = df.select(pudf_eq("a")).collect() + [row_neq] = df.select(pudf_neq("a")).collect() + self.assertEqual(row_eq[0], eq_expected) + self.assertEqual(row_neq[0], neq_expected) + + # Two-arg ``x == y`` / ``x != y`` exercising every NULL combination. + pudf_eq_xy = UserDefinedFunction(x_eq_y, BooleanType()) + pudf_neq_xy = UserDefinedFunction(x_neq_y, BooleanType()) + self.assertTrue(pudf_eq_xy.transpiled, "x == y should transpile") + self.assertTrue(pudf_neq_xy.transpiled, "x != y should transpile") + # Python semantics: + # None == None -> True; None != None -> False + # None == 0 -> False; None != 0 -> True + # 0 == None -> False; 0 != None -> True + # 1 == 1 -> True; 1 != 1 -> False + # 1 == 2 -> False; 1 != 2 -> True + for x, y, eq_expected, neq_expected in [ + (None, None, True, False), + (None, 0, False, True), + (0, None, False, True), + (1, 1, True, False), + (1, 2, False, True), + ]: + with self.subTest(x=x, y=y): + df = self.spark.createDataFrame([Row(a=x, b=y)], schema=two_col_schema) + [row_eq] = df.select(pudf_eq_xy("a", "b")).collect() + [row_neq] = df.select(pudf_neq_xy("a", "b")).collect() + self.assertEqual(row_eq[0], eq_expected, f"({x} == {y})") + self.assertEqual(row_neq[0], neq_expected, f"({x} != {y})") + + def test_udf_transpile_lte_gte(self): + # ``<=`` and ``>=`` go through the same ``_lower_value_compare`` path + # as ``<`` / ``>`` (and so share the NULL-raises-TypeError guard), but + # the entry points are not exercised elsewhere. Cover both with a None + # guard so the comparison only sees non-NULL operands here. + from pyspark.sql.types import StructField, StructType + + def lte_zero(x): + if x is not None: + return x <= 0 + + def gte_zero(x): + if x is not None: + return x >= 0 + + schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf_lte = UserDefinedFunction(lte_zero, BooleanType()) + pudf_gte = UserDefinedFunction(gte_zero, BooleanType()) + self.assertTrue(pudf_lte.transpiled, "x <= 0 should transpile") + self.assertTrue(pudf_gte.transpiled, "x >= 0 should transpile") + for value, lte_expected, gte_expected in [ + (-1, True, False), + (0, True, True), + (1, False, True), + (None, None, None), + ]: + with self.subTest(value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=schema) + [row_lte] = df.select(pudf_lte("a")).collect() + [row_gte] = df.select(pudf_gte("a")).collect() + self.assertEqual(row_lte[0], lte_expected) + self.assertEqual(row_gte[0], gte_expected) + + def test_udf_transpile_chained_comparison_falls_back(self): + # ``a < b < c`` is a chained comparison: Python evaluates it as + # ``(a < b) and (b < c)``. The transpiler refuses chained Compare + # nodes (``len(ops) != 1``) and must fall back to interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def chained(x): + return 0 < x < 10 + + schema = StructType([StructField("a", LongType(), nullable=False)]) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(chained, BooleanType()) + self.assertEqual([], pudf.transpiled, "chained comparison must NOT transpile") + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, "expected a fallback warning") + for value, expected in [(5, True), (0, False), (10, False), (-3, False)]: + with self.subTest(value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_multi_row(self): + # Every other transpile test uses a 1-row DataFrame; this one runs + # the same arithmetic transpile on a multi-row input to catch any + # column-reference / batch-boundary bug that single-row tests can't. + from pyspark.sql.types import StructField, StructType + + def plus_four(x): + if x is not None: + return x + 4 + + schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + pudf = UserDefinedFunction(plus_four, LongType()) + self.assertTrue(pudf.transpiled) + inputs = [Row(a=v) for v in [-3, -1, 0, 1, 7, None, 100]] + df = self.spark.createDataFrame(inputs, schema=schema) + transformed_df = df.select(pudf("a").alias("result")) + rows = transformed_df.collect() + actual = [row[0] for row in rows] + expected = [None if v is None else v + 4 for v in [-3, -1, 0, 1, 7, None, 100]] + self.assertEqual(actual, expected) + # Plan should also have the UDF stripped under the rewrite. + physical_plan = transformed_df._jdf.queryExecution().executedPlan().toString() + self.assertNotIn("UDF", physical_plan) + + def test_udf_transpile_falls_back_for_non_boolean_short_circuit(self): + # Python's `x or 0` returns x if truthy else 0; Spark's `|` is + # bitwise, so we'd silently produce wrong results. The transpiler + # must refuse, fall back to interpreted Python, and still produce + # the correct result. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def or_zero(x): + return x or 0 + + def and_one(x): + return x and 1 + + def not_int(x): + return not 0 + x # operand is BinOp, statically non-boolean + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + + cases = [ + ("or_zero", or_zero, LongType(), long_schema, Row(a=5), 5), + ("or_zero_none", or_zero, LongType(), long_schema, Row(a=None), 0), + ("and_one", and_one, LongType(), long_schema, Row(a=5), 1), + ("and_one_zero", and_one, LongType(), long_schema, Row(a=0), 0), + ("not_int", not_int, BooleanType(), long_schema, Row(a=0), True), + ("not_int_nonzero", not_int, BooleanType(), long_schema, Row(a=3), False), + ] + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + for label, func, return_type, schema, row, expected in cases: + with self.subTest(case=label): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, return_type) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: non-boolean and/or/not must NOT be lowered", + ) + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) + or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, f"{label}: expected a fallback warning") + df = self.spark.createDataFrame([row], schema=schema) + [result] = df.select(pudf("a")).collect() + self.assertEqual(result[0], expected, f"{label}: interpreted mismatch") + + def test_udf_transpile_falls_back_for_bare_truthiness_test(self): + # A bare `if x:` applied to a non-boolean column cannot be soundly + # lowered: Python truthiness is type-dependent (0, "", [], None are + # falsy) and the transpiler has no input type information at this + # point. Emitting coalesce(x, false) either fails Spark analysis for + # non-boolean columns or silently produces wrong answers. The + # transpiler must refuse and fall back to interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def truthy_int(x): + if x: + return x + return -1 + + def truthy_string(x): + return x if x else "default" + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + str_schema = StructType([StructField("a", StringType(), nullable=True)]) + + cases = [ + ("truthy_int_zero", truthy_int, LongType(), long_schema, Row(a=0), -1), + ("truthy_int_nonzero", truthy_int, LongType(), long_schema, Row(a=3), 3), + ("truthy_string_empty", truthy_string, StringType(), str_schema, Row(a=""), "default"), + ("truthy_string_val", truthy_string, StringType(), str_schema, Row(a="hi"), "hi"), + ] + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + for label, func, return_type, schema, row, expected in cases: + with self.subTest(case=label): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, return_type) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: bare truthiness test must NOT be lowered to Catalyst", + ) + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) + or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, f"{label}: expected a fallback warning") + df = self.spark.createDataFrame([row], schema=schema) + [result] = df.select(pudf("a")).collect() + self.assertEqual(result[0], expected, f"{label}: interpreted mismatch") + + def test_udf_transpile_falls_back_for_mismatched_branch_types(self): + # An if/ternary whose two branches produce different Spark categories + # (e.g. numeric vs string) would lower to a CASE WHEN whose branch + # values share no common type under ANSI. That node is carried as a + # child of the TranspiledPythonUDF and is type-checked by CheckAnalysis + # before ConvertToCatalyst can drop it, so without a guard the whole + # query would fail analysis instead of falling back. The transpiler must + # refuse and run the UDF as interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def mixed_ternary(x): + return 1 if x > 0 else "neg" + + def mixed_if(x): + # Single top-level `if`/`else` so the If-statement lowering path + # (not the "more than one statement" fallback) exercises the guard. + if x > 0: + return "pos" + else: + return x + + # Positive control: matching-category branches must still transpile, so + # the guard does not over-refuse. + def homogeneous(x): + return x if x > 0 else 0 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # Inputs are chosen to take the string-returning branch so the + # interpreted result is unambiguous. + mismatch_cases = [ + ("mixed_ternary", mixed_ternary, Row(a=-3), "neg"), + ("mixed_if", mixed_if, Row(a=10), "pos"), + ] + for label, func, row, expected in mismatch_cases: + with self.subTest(case=label): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, StringType()) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: mismatched branch types must NOT be lowered to Catalyst", + ) + fallback = [w for w in caught if "Unable to transpile" in str(w.message)] + self.assertTrue(fallback, f"{label}: expected a fallback warning") + df = self.spark.createDataFrame([row], schema=long_schema) + # Must run without an analysis failure and match interpreted Python. + [result] = df.select(pudf("a")).collect() + self.assertEqual(result[0], expected, f"{label}: interpreted mismatch") + + with self.subTest(case="homogeneous"): + pudf = UserDefinedFunction(homogeneous, LongType()) + self.assertNotEqual( + [], + pudf.transpiled, + "matching-category branches must still transpile", + ) + df = self.spark.createDataFrame([Row(a=5), Row(a=-3)], schema=long_schema) + results = [r[0] for r in df.select(pudf("a")).collect()] + self.assertEqual(results, [5, 0], "homogeneous branch result mismatch") + + def test_udf_transpile_falls_back_for_cross_category_eq(self): + # `x == True` on a numeric column would lower to `x = true`, which + # fails ANSI analysis (BIGINT vs BOOLEAN) while the option is still a + # child of the TranspiledPythonUDF -- breaking a working UDF. The + # category gate must refuse so it runs as interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def eq_true(x): + return x == True # noqa: E712 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(eq_true, BooleanType()) + self.assertEqual([], pudf.transpiled, "cross-category == must not transpile") + df = self.spark.createDataFrame([Row(a=5), Row(a=1)], schema=long_schema) + results = [r[0] for r in df.select(pudf("a")).collect()] + self.assertEqual(results, [5 == True, 1 == True]) + + def test_udf_transpile_falls_back_for_nested_ternary_eq(self): + # A ternary operand used inside `==` must contribute its branches' + # category, not the old "numeric" catch-all: `("5" if c else "6") == 5` + # previously passed the equality guard as numeric-vs-numeric and + # Spark's string-number coercion silently returned True where Python's + # cross-type == is False. (Reported by Codex review on PR #34.) + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def nested_ternary_eq(x): + return ("5" if x > 0 else "6") == 5 + + def none_branch_ternary_eq(x): + return ("5" if x > 0 else None) == 5 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + df = self.spark.createDataFrame([Row(a=5), Row(a=-5)], schema=long_schema) + with self.sql_conf(_TRANSPILE_ON): + for func in [nested_ternary_eq, none_branch_ternary_eq]: + with self.subTest(func=func.__name__): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, BooleanType()) + self.assertEqual( + [], pudf.transpiled, "string-ternary == int must not transpile" + ) + results = [r[0] for r in df.select(pudf("a")).collect()] + self.assertEqual(results, [False, False], "must match Python's ==") + + def test_udf_transpile_str_int_compare_matches_python(self): + # Comparing a value against a string literal (``x == "5"`` / ``x < "5"``) + # under the untyped-parameter path produces two candidate options -- a + # numeric variant and a string variant. The numeric variant mixes + # categories (numeric column vs string literal): Python compares such + # values as unequal / raises TypeError, while a lowered ``x = '5'`` would + # coerce under ANSI and silently diverge -- so ``_lower_eq`` / + # ``_lower_value_compare`` refuse it. Only the string variant survives. + # When the UDF is applied to a numeric column, ResolveTranspiledPython- + # UDFOptions drops the string option (no category match) and the UDF + # falls back to interpreted Python. We verify the observable guarantee on + # both sides: ``==`` returns Python's result (no coerced ``True``), and + # ``<`` raises on the Python side directly and surfaces the same error + # through Spark rather than returning a coerced answer. + from pyspark.errors import PythonException + from pyspark.sql.types import StructField, StructType + + def eq_str(x): + return x == "5" + + def lt_str(x): + return x < "5" + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + df = self.spark.createDataFrame([Row(a=5), Row(a=1)], schema=long_schema) + + # ``==`` : numeric column -> string option dropped -> interpreted + # Python, so the result matches ``5 == "5"`` (False), not a coerced + # ``True`` from ``bigint = '5'``. + pudf_eq = UserDefinedFunction(eq_str, BooleanType()) + results = [r[0] for r in df.select(pudf_eq("a")).collect()] + self.assertEqual(results, [5 == "5", 1 == "5"], "must match Python's ==") + + # ``<`` : Python raises TypeError for ``int < str``; the transpiled + # string option is dropped for a numeric column, so Spark runs the + # interpreted UDF and surfaces the same error rather than coercing. + # Asserting PythonException (not a bare Exception) pins this to the + # interpreted-fallback path: a Catalyst AnalysisException here would + # instead mean the string option was wrongly kept and broke the + # query rather than falling back. + pudf_lt = UserDefinedFunction(lt_str, BooleanType()) + with self.assertRaises(TypeError): + lt_str(5) + with self.assertRaises(PythonException) as ctx: + df.select(pudf_lt("a")).collect() + self.assertIn("not supported between", str(ctx.exception)) + + def test_udf_transpile_falls_back_for_bool_arithmetic(self): + # `(x > 0) + 1` is valid Python (True + 1 == 2), but the lowered + # Add(boolean, int) fails ANSI analysis. The category of a + # boolean-producing operand is now "bool" (not the numeric catch-all), + # so this refuses and runs as interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def bool_plus_one(x): + return (x > 0) + 1 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(bool_plus_one, LongType()) + self.assertEqual([], pudf.transpiled, "bool arithmetic must not transpile") + df = self.spark.createDataFrame([Row(a=5), Row(a=-5)], schema=long_schema) + results = [r[0] for r in df.select(pudf("a")).collect()] + self.assertEqual(results, [2, 1]) + + def test_udf_transpile_falls_back_for_return_wrapped_bool_branch(self): + # If-statement branches arrive as ast.Return nodes; the branch-category + # guard must see through the wrapper. A boolean-returning branch vs a + # numeric one previously slipped past the guard and failed analysis + # (CASE WHEN [BOOLEAN, INT]) instead of falling back. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def mixed(x): + if x > 0: + return x > 5 + else: + return 1 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(mixed, LongType()) + self.assertEqual([], pudf.transpiled, "bool-vs-int branches must not transpile") + df = self.spark.createDataFrame([Row(a=-3)], schema=long_schema) + [result] = df.select(pudf("a")).collect() + self.assertEqual(result[0], 1) + + def test_udf_transpile_plain_self_param_is_an_ordinary_param(self): + # A plain function whose first parameter is literally named `self` is not a + # bound receiver -- the call site supplies it. Stripping it emitted + # `_udf_param_-1` and threw at call construction, so it used to be refused + # outright; the receiver is decided by dispatch now, so this lowers. + def weird(self, other): + return self + other + + self.assertEqual( + self._vals(weird, LongType(), "a long, b long", [(2, 3)]), + [5], + "both parameters are supplied at the call site, so both get placeholders", + ) + + def test_udf_transpile_falls_back_for_wraps_decorated_function(self): + # inspect.getsource follows __wrapped__, so a functools.wraps-decorated + # UDF previously transpiled the WRAPPED function's source while the + # interpreted path ran the wrapper -- a silent wrong result. + import functools + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def base(x): + return x + 1 + + @functools.wraps(base) + def wrapper(x): + return base(x) * 10 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(wrapper, LongType()) + self.assertEqual([], pudf.transpiled, "wraps-decorated UDF must not transpile") + df = self.spark.createDataFrame([Row(a=5)], schema=long_schema) + [result] = df.select(pudf("a")).collect() + self.assertEqual(result[0], 60, "must run the wrapper, not the wrapped source") + + def test_udf_transpile_falls_back_for_none_in_boolop(self): + # Python's `None and x` short-circuits to None; Spark's three-valued + # `null AND false` is false. A literal None operand must force a + # fallback rather than silently diverge. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def none_and(x): + return None and (x > 0) + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(none_and, BooleanType()) + self.assertEqual([], pudf.transpiled, "literal None in and/or must not transpile") + df = self.spark.createDataFrame([Row(a=-5)], schema=long_schema) + [result] = df.select(pudf("a")).collect() + self.assertIsNone(result[0]) + + def test_udf_transpile_falls_back_for_uncastable_return_type(self): + # The lowered expression is cast to the declared return type; a return + # type no atomic lowering can be cast to (arrays, maps, datetimes, ...) + # would make that Cast fail CheckAnalysis and break the whole query + # (the options are children of TranspiledPythonUDF), so such UDFs must + # fall back at construction instead. Interpreted execution still works + # (the pickled-UDF converter nulls the type-mismatched results). + import warnings as _warnings + + from pyspark.sql.types import ArrayType, TimestampType + + plus_one = lambda x: x + 1 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + for rt in (ArrayType(LongType()), TimestampType()): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(plus_one, rt) + self.assertEqual([], pudf.transpiled, f"return type {rt} must not transpile") + # Interpreted execution keeps working; an int result for an array + # return type is nulled by the pickled-UDF converter. (Timestamp + # is not exercised here: its converter accepts ints as micros.) + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + array_udf = UserDefinedFunction(plus_one, ArrayType(LongType())) + df = self.spark.createDataFrame([Row(a=1)]) + [result] = df.select(array_udf("a")).collect() + self.assertIsNone(result[0], "interpreted fallback nulls the mismatch") + + def test_udf_transpile_falls_back_for_cross_category_return_cast(self): + # Per-variant guard: the body category must MATCH the declared return + # type's category. Un-castable combos (binary body -> numeric return, + # boolean body -> binary return) would fail analysis outright, and + # analysis-valid cross-category casts (string -> long, numeric -> + # boolean, anything -> decimal) diverge from the interpreted path, + # which nulls type-mismatched results instead of casting -- e.g. + # `def f(s: str): return s` declared LongType() would return 123 for + # '123' (or raise CAST_INVALID_INPUT) where interpreted returns NULL. + import warnings as _warnings + + from pyspark.sql.types import DecimalType + + def bytes_to_long(x: bytes): + return x + + def bool_to_binary(x): + return (x > 0) if x is not None else None + + def str_ident(s: str): + return s + + def plus_one(x): + return x + 1 + + with self.sql_conf(_TRANSPILE_ON): + for func, rt, label in ( + (bytes_to_long, LongType(), "binary body -> numeric return"), + (bool_to_binary, BinaryType(), "boolean body -> binary return"), + (bytes_to_long, StringType(), "binary body -> string return"), + (str_ident, LongType(), "string body -> numeric return"), + (plus_one, BooleanType(), "numeric body -> boolean return"), + (plus_one, DecimalType(10, 2), "numeric body -> decimal return"), + ): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, rt) + self.assertEqual([], pudf.transpiled, f"{label} must not transpile") + # Interpreted execution of the Codex-flagged example: NULL, not a + # cast. (The transpiled cast would have returned 123.) + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + str_long = UserDefinedFunction(str_ident, LongType()) + df = self.spark.createDataFrame([("123",)], "s string") + self.assertIsNone(df.select(str_long("s")).first()[0]) + + def test_udf_transpile_falls_back_for_non_numeric_unary(self): + # Unary +/- only lower for numeric operands: Python raises TypeError + # on `+s`/`-s` for strings while Spark's ANSI string promotion would + # silently coerce (`-'5'` -> -5.0), and `-x` on a boolean would fail + # analysis outright rather than fall back. + import warnings as _warnings + + def neg_str(s: str): + return -s + + def pos_str(s: str): + return +s + + def neg_bool(x: bool): + return -x + + with self.sql_conf(_TRANSPILE_ON): + for func in (neg_str, pos_str, neg_bool): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, LongType()) + self.assertEqual([], pudf.transpiled, f"{func.__name__} must not transpile") + # Numeric unary still lowers and matches Python. + neg = lambda x: -x # noqa: E731 + self.assertEqual(self._vals(neg, LongType(), "a long", [(5,), (-3,)]), [-5, 3]) + + def test_udf_transpile_falls_back_for_self_reference(self): + # A __call__ body that references bare `self` has no column + # equivalent; the offset scheme previously emitted `_udf_param_-1`, + # which the JVM builder rejected with an AnalysisException at call + # construction instead of falling back to interpreted Python. + import warnings as _warnings + + class PickSelf: + def __call__(self, x): + return x if x is not None else self + + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(PickSelf(), LongType()) + self.assertEqual([], pudf.transpiled, "`self` reference must not transpile") + # Interpreted execution still works (previously the call itself + # raised). Only non-null rows are exercised: a row that RETURNS + # `self` would fail JVM-side unpickling of the instance, which is + # interpreted-UDF behavior unrelated to this guard. + df = self.spark.createDataFrame([(2,), (7,)], "a long") + results = [r[0] for r in df.select(pudf("a")).collect()] + self.assertEqual(results, [2, 7]) + + def test_udf_transpile_preserves_auto_column_name(self): + # The auto-generated column name must stay `f(a)` whether or not the + # rewrite engages; the TranspiledPythonUDF wrapper (and its option + # children) must not leak into user-visible schema names. + from pyspark.sql.types import StructField, StructType + + def plus_four(x): + return x + 4 + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + df = self.spark.createDataFrame([Row(a=1)], schema=long_schema) + with self.sql_conf(_TRANSPILE_ON): + pudf = UserDefinedFunction(plus_four, LongType()) + self.assertTrue(pudf.transpiled) + self.assertEqual(df.select(pudf("a")).columns, ["plus_four(a)"]) + + def test_udf_transpile_arity_mismatch_falls_back(self): + # Calling with the wrong number of arguments is a user error that must + # surface as the standard Python-side TypeError, not be silently + # absorbed by a transpiled constant (zero-param case) nor raise a + # misleading "internal error" AnalysisException (too-few-args case). + import warnings as _warnings + + from pyspark.errors import PythonException + from pyspark.sql.types import StructField, StructType + + def zero(): + return 42 + + def two(x, y): + return x + y + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + df = self.spark.createDataFrame([Row(a=5)], schema=long_schema) + with self.sql_conf(_TRANSPILE_ON): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf_zero = UserDefinedFunction(zero, LongType()) + pudf_two = UserDefinedFunction(two, LongType()) + with self.assertRaises(PythonException): + df.select(pudf_zero("a")).collect() + with self.assertRaises(PythonException): + df.select(pudf_two("a")).collect() + + def test_udf_transpile_decimal_input_falls_back(self): + # Python receives decimal.Decimal objects, which raise TypeError when + # mixed with float literals; the transpiled numeric lowering would + # silently succeed. Decimal columns must fall back to interpreted + # Python (pruned by input category at analysis time). + from pyspark.errors import PythonException + + def add_half(x): + return x + 1.5 + + with self.sql_conf(_TRANSPILE_ON): + # DoubleType: the return type must category-match the numeric body + # for the option to be emitted (a string return type would itself + # force a fallback before the decimal-input pruning under test). + pudf = UserDefinedFunction(add_half, DoubleType()) + self.assertTrue(pudf.transpiled, "numeric option should still be produced") + df = self.spark.sql("SELECT CAST(1.0 AS DECIMAL(10,2)) AS d") + with self.assertRaises(PythonException): + df.select(pudf("d")).collect() + + def test_udf_transpile_collated_string_falls_back(self): + # Under a non-binary collation Spark's `=` follows collation rules + # ('abc' = 'ABC' is true under UTF8_LCASE) while Python compares + # codepoints. Collated columns must fall back to interpreted Python. + def eq_abc(s): + return s == "ABC" + + with self.sql_conf(_TRANSPILE_ON): + pudf = UserDefinedFunction(eq_abc, BooleanType()) + self.assertTrue(pudf.transpiled, "string option should still be produced") + df = self.spark.sql("SELECT 'abc' COLLATE UTF8_LCASE AS s") + [result] = df.select(pudf("s")).collect() + self.assertIs(result[0], False, "must match Python, not collation semantics") + + def test_udf_transpile_is_none_semantics(self): + # `x is None` and `None is x` (and their `is not` variants) should + # transpile to isNull/isNotNull. Any other identity check (`x is 0`, + # `x is y`, `x is True`) must NOT transpile -- Python's `is` is an + # object-identity test with no SQL equivalent outside of None. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + + def x_is_none(x): + return x is None + + def x_is_not_none(x): + if x is not None: + return x + 1 + + def none_is_x(x): + return None is x + + def none_is_not_x(x): + if None is not x: + return x + 1 + + def x_is_zero(x): + return x is 0 # noqa: F632 identity vs equality + + def x_is_true(x): + return x is True + + def x_is_y(x, y): + return x is y + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + # `x is None` and `None is x` should transpile and produce + # identical results. + for func, label in [(x_is_none, "x_is_none"), (none_is_x, "none_is_x")]: + with self.subTest(case=label): + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue( + pudf.transpiled, + f"{label}: expected transpilation to succeed", + ) + df = self.spark.createDataFrame([Row(a=None)], schema=long_schema) + [row] = df.select(pudf("a")).collect() + self.assertTrue(row[0], f"{label}: None is None should be True") + df = self.spark.createDataFrame([Row(a=1)], schema=long_schema) + [row] = df.select(pudf("a")).collect() + self.assertFalse(row[0], f"{label}: 1 is None should be False") + + # `x is not None` and `None is not x` should transpile. + for func, label in [ + (x_is_not_none, "x_is_not_none"), + (none_is_not_x, "none_is_not_x"), + ]: + with self.subTest(case=label): + pudf = UserDefinedFunction(func, LongType()) + self.assertTrue( + pudf.transpiled, + f"{label}: expected transpilation to succeed", + ) + df = self.spark.createDataFrame([Row(a=2)], schema=long_schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], 3, f"{label}: non-None input should return x+1") + df = self.spark.createDataFrame([Row(a=None)], schema=long_schema) + [row] = df.select(pudf("a")).collect() + self.assertIsNone(row[0], f"{label}: None input should return None") + + # Non-None identity checks must NOT transpile and must still + # return correct results via interpreted Python. + bool_schema = StructType([StructField("a", BooleanType(), nullable=True)]) + two_col_schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + ] + ) + non_none_cases = [ + # CPython interns small ints so `0 is 0` happens to be True in CPython, + # but that is an implementation detail. The transpiler must still refuse + # to lower these to isNull/isNotNull. We just verify: (a) no transpile, + # (b) the interpreted result matches what Python actually produces. + ("x_is_zero", x_is_zero, BooleanType(), long_schema, Row(a=0), True), + # `True is True` is True because bool singletons are interned. + ("x_is_true", x_is_true, BooleanType(), bool_schema, Row(a=True), True), + ("x_is_y", x_is_y, BooleanType(), two_col_schema, Row(a=1, b=1), True), + ] + for label, func, return_type, schema, row, expected in non_none_cases: + with self.subTest(case=label): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, return_type) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: non-None identity check must NOT transpile", + ) + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) + or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, f"{label}: expected a fallback warning") + df = self.spark.createDataFrame([row], schema=schema) + args = ["a", "b"] if "b" in schema.fieldNames() else ["a"] + [result] = df.select(pudf(*args)).collect() + self.assertEqual(result[0], expected, f"{label}: interpreted result mismatch") + + def test_udf_transpile_not_bare_param_falls_back(self): + # `not x` where x is a bare UDF parameter (unknown type at + # transpile time) must NOT be lowered: Spark's `~` is bitwise, not + # Python truthiness, so `not 0` would produce True via Python but + # Spark's `~0L` is -1 (truthy). The transpiler must refuse and fall + # back to interpreted Python. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def not_x(x): + return not x + + long_schema = StructType([StructField("a", LongType(), nullable=True)]) + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(not_x, BooleanType()) + self.assertEqual([], pudf.transpiled, "not x on bare param must NOT transpile") + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, "expected a fallback warning for `not x`") + # Verify interpreted result is still correct. + for value, expected in [(0, True), (1, False), (None, True)]: + with self.subTest(value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=long_schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_and_or_bare_param_falls_back(self): + # `x and y` / `x or y` where x/y are bare UDF parameters (unknown + # type) must NOT be lowered: Python returns one of the operands + # (truthiness semantics) while Spark's `&`/`|` are bitwise. The + # transpiler must refuse and fall back. + import warnings as _warnings + + from pyspark.sql.types import StructField, StructType + + def x_and_y(x, y): + return x and y + + def x_or_y(x, y): + return x or y + + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + ] + ) + + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": True, + "spark.sql.ansi.enabled": True, + } + ): + for func, label, row, expected in [ + (x_and_y, "x_and_y_falsy", Row(a=0, b=5), 0), + (x_and_y, "x_and_y_truthy", Row(a=3, b=5), 5), + (x_or_y, "x_or_y_falsy_left", Row(a=0, b=5), 5), + (x_or_y, "x_or_y_truthy_left", Row(a=3, b=0), 3), + ]: + with self.subTest(case=label): + with _warnings.catch_warnings(record=True) as caught: + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, LongType()) + self.assertEqual( + [], + pudf.transpiled, + f"{label}: and/or on bare params must NOT transpile", + ) + fallback = [ + w + for w in caught + if "Unable to transpile" in str(w.message) + or "Errors encountered" in str(w.message) + ] + self.assertTrue(fallback, f"{label}: expected a fallback warning") + df = self.spark.createDataFrame([row], schema=schema) + [result] = df.select(pudf("a", "b")).collect() + self.assertEqual(result[0], expected, f"{label}: interpreted result mismatch") + + def test_cannot_convert_column_into_bool_includes_column_repr(self): + # The error fired by ``Column.__bool__`` should name the offending + # column so users can see which expression triggered the fallback. + from pyspark.errors import PySparkValueError + + df = self.spark.createDataFrame([Row(a=1, b=2)]) + col_a = df["a"] + with self.assertRaises(PySparkValueError) as ctx: + bool(col_a) + message = str(ctx.exception) + self.assertIn("Cannot convert column into bool", message) + # Column's stringification is JVM-side and may render the column + # as ``a`` (unresolved) or with a backtick variant, so we just + # require the column name appears somewhere in the message. + self.assertIn("a", message) + + # ------------------------------------------------------------------ + # Edge cases (SPARK-55206 follow-up). Helpers build a UDF with + # transpilation on; `_vals` runs it and returns outputs (asserting it + # transpiled), `_raises` asserts it raises. Arg columns come from the + # schema. Operator cases are table-driven. Plan-elision checks count + # `EvalPython` nodes because an ordering compare's `raise_error` message + # contains "UDF" (so the "UDF" substring is unreliable). + # ------------------------------------------------------------------ + + @staticmethod + def _udf_and_warnings(func, return_type): + """Build a UDF, returning it with the text of any warnings it emitted. + + ``udf.py`` reports WHY a UDF fell back only as a warning, so every test that + cares about a fallback needs them captured. The caller must already be inside + ``sql_conf(_TRANSPILE_ON)`` -- this does not set the conf. + """ + import warnings + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + u = UserDefinedFunction(func, return_type) + return u, " ".join(str(w.message) for w in caught) + + def _fallback_reason(self, func, return_type=LongType()): + """Assert ``func`` produced no options, and hand back the reason it reported.""" + u, reasons = self._udf_and_warnings(func, return_type) + self.assertEqual([], u.transpiled, f"{func} must fall back") + self.assertTrue(reasons, f"{func} fell back without saying why") + return u, reasons + + def _transpiled_udf(self, func, return_type): + """A UDF asserted to have produced options; naming the fallback reason if not. + + Without the captured warning an empty ``transpiled`` asserts as bare "[] is + not true", and from CI the text is only in a credentialed log artifact. The + caller must already be inside ``sql_conf(_TRANSPILE_ON)``. + """ + u, reasons = self._udf_and_warnings(func, return_type) + self.assertTrue(u.transpiled, f"{func} produced no transpiled options: {reasons}") + return u + + def _vals(self, func, return_type, schema, rows, require_lowered=True): + with self.sql_conf(_TRANSPILE_ON): + u = self._transpiled_udf(func, return_type) + df = self.spark.createDataFrame(rows, schema) + projected = df.select(u(*df.columns)) + # ``u.transpiled`` only says options were PRODUCED; the JVM may still + # discard them and run interpreted Python, returning the right value and + # hiding a wrong lowering. Pass ``require_lowered=False`` only where the + # JVM is EXPECTED to discard them. + if require_lowered: + self.assertEqual(0, self._eval_python_count(projected), str(func)) + return [r[0] for r in projected.collect()] + + def _raises(self, func, schema, rows, needle="numeric"): + with self.sql_conf(_TRANSPILE_ON): + u = self._transpiled_udf(func, LongType()) + df = self.spark.createDataFrame(rows, schema) + with self.assertRaises(Exception) as ctx: + df.select(u(*df.columns)).collect() + self.assertIn(needle, str(ctx.exception).lower(), str(func)) + + @staticmethod + def _eval_python_count(df): + return df._jdf.queryExecution().executedPlan().toString().count("EvalPython") + + def test_udf_transpile_lowers_operators(self): + # Operators lower to Catalyst and match Python: modulo sign-parity, + # non-commutative -/* (parameter order), unary nesting, constant + # body, not(compare), nested boolean, string ==/<, reversed-operand and + # column-to-column comparisons, if/elif/else, and assigned lambdas. + L, B = LongType(), BooleanType() + modulo = lambda x, y: x % y # noqa: E731 + subtract = lambda a, b: a - b # noqa: E731 + multiply = lambda a, b: a * b # noqa: E731 + double_neg = lambda x: --x # noqa: E731 + unary_pm = lambda x: +(-x) # noqa: E731 + constant = lambda x: 42 # noqa: E731 + not_pos = lambda x: (not (x > 0)) if x is not None else None # noqa: E731 + nested = lambda x, y, z: ((x > 0) and (y > 0)) or (z == 0) # noqa: E731 + str_eq = lambda x: (x == "foo") if x is not None else None # noqa: E731 + str_lt = lambda x: (x < "m") if x is not None else None # noqa: E731 + rev_lt = lambda x: (0 < x) if x is not None else None # noqa: E731 + rev_eq = lambda x: 5 == x # noqa: E731 + none_eq = lambda x: None == x # noqa: E711,E731 + col_lt = lambda a, b: (a < b) if a is not None and b is not None else None # noqa: E731 + assigned = lambda v: v + 1 # noqa: E731 + + def if_elif_else(x): + if x is None: + return -1 + elif x == 0: + return 0 + else: + return 1 + + # (func, return_type, schema, rows, expected); arg columns come from the schema. + cases = [ + (modulo, L, "a long, b long", [(7, 3), (7, -3), (-7, 3), (-7, -3)], [1, -2, 2, -1]), + (subtract, L, "a long, b long", [(5, 3), (3, 5)], [2, -2]), + (multiply, L, "a long, b long", [(4, 3), (-2, 5)], [12, -10]), + (double_neg, L, "a long", [(5,), (-3,)], [5, -3]), + (unary_pm, L, "a long", [(5,), (-3,)], [-5, 3]), + (constant, L, "a long", [(1,), (999,)], [42, 42]), + (not_pos, B, "a long", [(1,), (0,), (-1,), (None,)], [False, True, True, None]), + (str_eq, B, "a string", [("foo",), ("bar",), (None,)], [True, False, None]), + (str_lt, B, "a string", [("a",), ("z",), (None,)], [True, False, None]), + (rev_lt, B, "a long", [(1,), (0,), (-1,)], [True, False, False]), + (rev_eq, B, "a long", [(5,), (3,), (None,)], [True, False, False]), + (none_eq, B, "a long", [(None,), (5,)], [True, False]), + (col_lt, B, "a long, b long", [(1, 2), (2, 1), (1, 1)], [True, False, False]), + (if_elif_else, L, "a long", [(None,), (0,), (5,), (-3,)], [-1, 0, 1, 1]), + (assigned, L, "a long", [(1,), (10,)], [2, 11]), + ( + nested, + B, + "a long, b long, c long", + [(1, 1, 5), (-1, 1, 0), (-1, 1, 5)], + [True, True, False], + ), + ] + for i, (func, rt, schema, rows, expected) in enumerate(cases): + with self.subTest(case=i): + self.assertEqual(self._vals(func, rt, schema, rows), expected, f"case {i}: {rows}") + + def test_udf_transpile_callable_object_drops_its_receiver(self): + # A callable instance's `self` is dropped before anything indexes the param + # list, so a/b are _udf_param_0/_udf_param_1 with no offsetting anywhere + # downstream (the non-commutative body proves the order). + class SubAB: + def __call__(self, a, b): + return a - b + + self.assertEqual( + self._vals(SubAB(), LongType(), "a long, b long", [(5, 3), (3, 5)]), [2, -2] + ) + + def test_udf_transpile_plan_elision(self): + # Transpiled UDFs are elided in filter (not just select); a mixed + # non-convertible -> convertible -> non-convertible chain inlines only + # the middle UDF, leaving exactly two Python eval nodes. + offset = 3 + gt5 = lambda x: (x > 5) if x is not None else None # noqa: E731 + add_offset = lambda x: x + offset # noqa: E731 closure -> fallback + plus_one = lambda x: x + 1 # noqa: E731 convertible + div_two = lambda x: x / 2 # noqa: E731 `/` -> fallback + with self.sql_conf(_TRANSPILE_ON): + f = UserDefinedFunction(gt5, BooleanType()) + self.assertTrue(f.transpiled) + fdf = self.spark.createDataFrame([(3,), (7,), (1,), (None,)], "a long").filter(f("a")) + self.assertEqual([r[0] for r in fdf.collect()], [7]) + self.assertEqual(0, self._eval_python_count(fdf)) + + u1 = UserDefinedFunction(add_offset, LongType()) + u2 = UserDefinedFunction(plus_one, LongType()) + u3 = UserDefinedFunction(div_two, DoubleType()) + self.assertEqual(([], True, []), (u1.transpiled, bool(u2.transpiled), u3.transpiled)) + chained = ( + self.spark.createDataFrame([(10,)], "a long") + .select(u1("a").alias("x")) + .select(u2("x").alias("y")) + .select(u3("y").alias("z")) + ) + self.assertEqual(chained.first()[0], 7.0) # ((10 + 3) + 1) / 2 + self.assertEqual(2, self._eval_python_count(chained)) + + def test_udf_transpile_config_toggle_no_stale_nodes(self): + # Built with the flags on, executed with them off -> clean fallback to + # interpreted Python (the optimizer drops the transpiled node), no error. + plus_one = lambda x: x + 1 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + u = UserDefinedFunction(plus_one, LongType()) + self.assertTrue(u.transpiled) + with self.sql_conf( + { + "spark.sql.experimental.optimizer.transpilePyUDFs": False, + "spark.sql.ansi.enabled": False, + } + ): + df = self.spark.createDataFrame([(1,), (5,)], "a long") + self.assertEqual([r[0] for r in df.select(u("a")).collect()], [2, 6]) + + def test_udf_transpile_casts_to_return_type(self): + # The lowered expression is cast to the declared return type. + plus_one = lambda x: x + 1 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + d = UserDefinedFunction(plus_one, DoubleType()) + col = self.spark.createDataFrame([(1,)], "a long").select(d("a").alias("r")) + self.assertEqual(col.schema["r"].dataType, DoubleType()) + self.assertEqual(col.first()[0], 2.0) + self.assertEqual(self._vals(plus_one, LongType(), "a long", [(1,)]), [2]) + + def test_udf_transpile_falls_back(self): + # Shapes that must NOT transpile (and still compute via Python): + # inline/wrapped/partial lambdas, default/variadic/keyword-only args, and + # `%` string formatting. (String `+`/`*` now lower to concat/repeat -- see + # test_udf_transpile_string_operands -- but `%` as a format is not handled.) + import functools + + def wrapper(fn): + return fn + + def with_default(a, b=0): + return a + 10 * b + + def with_varargs(a, *rest): + return a + + def with_kwargs(a, **opts): + return a + + base = lambda v, w: v + w # noqa: E731 + percent_fmt = lambda x: "n=%d" % x # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + # An inline or wrapped lambda is a call ARGUMENT, so the source is + # read fine but parses as ``Call`` rather than a definition we can + # unwrap; ``functools.partial`` has no reachable source at all. + self.assertEqual([], UserDefinedFunction(lambda v: v + 1, LongType()).transpiled) + self.assertEqual( + [], UserDefinedFunction(wrapper(lambda v: v + 1), LongType()).transpiled + ) + self.assertEqual( + [], UserDefinedFunction(functools.partial(base, 1), LongType()).transpiled + ) + # default / variadic / keyword-only args, and `%` string formatting + for func, rt in [ + (with_default, LongType()), + (with_varargs, LongType()), + (with_kwargs, LongType()), + (percent_fmt, StringType()), + ]: + with self.subTest(func=func): + self.assertEqual([], UserDefinedFunction(func, rt).transpiled) + # Fell back -> interpreted Python still computes correctly. + wd = UserDefinedFunction(with_default, LongType()) + num = self.spark.createDataFrame([(5,)], "a long") + self.assertEqual( + [num.select(wd("a")).first()[0], num.select(wd("a", "a")).first()[0]], [5, 55] + ) + + def test_udf_transpile_falls_back_when_a_sibling_lambda_shares_the_line(self): + # ``inspect.getsource`` works in whole lines, so every lambda on a line hands + # back the same source and nothing in it says which one we hold. Taking + # whichever came first was right by position rather than by identity: + # ``minus_one`` lowered ``x + 1``, and the lambda in the decorator lowered + # the decorated ``def``'s body -- both silently wrong. Refusing costs + # ``plus_one``, which the first-match rule happened to get right; being right + # for 1-of-N by position is not something a caller can rely on (SPARK-58650). + # + # ``fmt: off`` keeps the two on one line; asserted below, since the formatter + # would otherwise split them and quietly make this test vacuous. + # fmt: off + plus_one = lambda x: x + 1; minus_one = lambda x: x - 1 # noqa: E702,E731 + # fmt: on + self.assertEqual( + plus_one.__code__.co_firstlineno, + minus_one.__code__.co_firstlineno, + "fixture must keep both lambdas on ONE line or it proves nothing", + ) + + captured = [] + + def capture(g): + captured.append(g) + return lambda fn: fn + + @capture(lambda x: x + 1) + def unrelated(a): + return a * 12345 + + (decorator_lambda,) = captured + + # A class body is the same hazard: master read the whole line and lowered + # ``helper``, computing 5*100 for a UDF that really returns 5*2. + class TwoOnALine: + helper, __call__ = lambda self, x: x * 100, lambda self, x: x * 2 + + self.assertEqual(10, TwoOnALine()(5)) + + num = self.spark.createDataFrame([(5,)], "a long") + sibling = "put each lambda on its own line" + not_a_statement = "does not define it as a statement of its own" + with self.sql_conf(_TRANSPILE_ON): + for label, func, expected, needle in [ + # Lowered the WRONG body before -- the bug. + ("second lambda on the line", minus_one, 4, sibling), + ("lambda inside a decorator", decorator_lambda, 6, not_a_statement), + ("__call__ sharing a class-body line", TwoOnALine(), 10, not_a_statement), + # Lowered correctly before; refused now as acknowledged collateral. + ("first lambda on the line", plus_one, 6, sibling), + ]: + with self.subTest(case=label): + u, reasons = self._fallback_reason(func) + # Assert the REASON, so relaxing the guard fails here loudly + # rather than passing for a new and unrelated cause. + self.assertIn(needle, reasons) + self.assertEqual(expected, num.select(u("a")).first()[0]) + + # The ``def`` itself is NOT collateral: we know we are holding it, so a + # lambda in its decorator cannot be the body and does not block lowering. + self.assertEqual(self._vals(unrelated, LongType(), "a long", [(5,)]), [61725]) + + def test_udf_transpile_refuses_a_lambda_whose_source_now_reads_as_a_def(self): + # ``inspect.getsource`` reads the file as it is NOW, so an edit after import + # can hand back source holding no lambda at all. Refusing only when a SIBLING + # lambda is in view failed open here, lowering an unrelated ``def`` as the + # body: 5 * 12345 for a lambda Python evaluates to 6. + import importlib.util + import os + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "drifted_lambda.py") + with open(path, "w") as handle: + handle.write("f = lambda x: x + 1\n") + spec = importlib.util.spec_from_file_location("drifted_lambda", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + with open(path, "w") as handle: + handle.write("def unrelated(a):\n return a * 12345\n") + + with self.sql_conf(_TRANSPILE_ON): + u, reasons = self._fallback_reason(module.f) + # Pin the REASON, or a source read that merely failed would pass too. + self.assertIn("which lambda to lower cannot be determined", reasons) + num = self.spark.createDataFrame([(5,)], "a long") + self.assertEqual(6, num.select(u("a")).first()[0]) + + def test_udf_transpile_ambiguity_check_sees_only_rival_lambdas(self): + # Two ways the check used to misfire or not fire at all. + from pyspark.sql.transpile import _held_code + + # A lambda nested in the held lambda's own body is not a rival -- it can + # never be the UDF, and the user cannot split it onto another line. It fell + # back with the sibling message, advice that could not be acted on. + nested = lambda x: (lambda y: y + 1)(x) # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + _, reasons = self._fallback_reason(nested) + self.assertIn("Call", reasons, "must fall back for the body, not for ambiguity") + self.assertNotIn("put each lambda on its own line", reasons) + + # The mirror: a lambda RETURNED by a one-line lambda. Here the outer one is + # what the source read locates and the inner one is what we hold, so treating + # nested lambdas as never-rivals let this through -- and it proceeded with the + # outer signature, reporting `n` as the UDF's parameter for a UDF whose only + # parameter is `x`. Matching the located lambda's parameters against the held + # code object's is what separates this from the case above. + # fmt: off + make_adder = lambda n: lambda x: x + n # noqa: E731 + # fmt: on + add_three = make_adder(3) + self.assertEqual(8, add_three(5)) + with self.sql_conf(_TRANSPILE_ON): + u, reasons = self._fallback_reason(add_three) + self.assertIn("takes different parameters", reasons) + self.assertEqual([], u._transpiled_param_names or []) + self.assertEqual( + 8, self.spark.createDataFrame([(5,)], "a long").select(u("a")).first()[0] + ) + + # ``staticmethod``/``classmethod`` hide ``__code__`` behind the descriptor, so + # reading it off them left the guard inactive -- skipping the check for a shape + # it exists to catch. Now that these lower at all, the skip would be a wrong + # answer: ``helper`` shares the line and takes the same parameter, so it is a + # true rival and would be lowered instead (5 * 9, not 5 * 2). + class Wrapped: + # fmt: off + helper = lambda x: x * 9; __call__ = staticmethod(lambda x: x * 2) # noqa: E702,E731 + # fmt: on + + self.assertEqual( + Wrapped.helper.__code__.co_firstlineno, + Wrapped.__call__.__code__.co_firstlineno, + "fixture must keep both lambdas on ONE line or it proves nothing", + ) + self.assertEqual( + "<lambda>", + getattr(_held_code(Wrapped()), "co_name", None), + "the held code must be found inside the descriptor", + ) + self.assertEqual(10, Wrapped()(5)) + with self.sql_conf(_TRANSPILE_ON): + u, reasons = self._fallback_reason(Wrapped()) + self.assertIn("put each lambda on its own line", reasons) + num = self.spark.createDataFrame([(5,)], "a long") + self.assertEqual(10, num.select(u("a")).first()[0]) + + def test_udf_transpile_lowers_an_annotated_lambda_binding(self): + # An annotated binding is the same shape as a plain one, and the form a typed + # codebase writes. Only ``ast.Assign`` was unwrapped, so this was refused as + # "not a statement of its own" -- while the module docstring told users to + # bind the lambda to a name and give it a line, which is exactly this. + from typing import Callable + + annotated: Callable[[int], int] = lambda x: x + 1 # noqa: E731 + self.assertEqual(self._vals(annotated, LongType(), "a long", [(5,)]), [6]) + + def test_udf_transpile_recovers_shapes_with_an_unheld_lambda_in_view(self): + # The ambiguity check applies only when the callable we hold IS a lambda. The + # lambdas below belong to a ``def`` we are not lowering, so refusing on their + # account would cost lowering for nothing. + from typing import Annotated + + def annotated(x: Annotated[int, lambda v: v > 0]) -> int: + return x + 1 + + def returns_annotated(x) -> Annotated[int, lambda v: v > 0]: + return x + 2 + + for label, func, expected in [ + ("lambda in a parameter annotation", annotated, 6), + ("lambda in the return annotation", returns_annotated, 7), + ]: + with self.subTest(case=label): + self.assertEqual(self._vals(func, LongType(), "a long", [(5,)]), [expected]) + + def test_udf_transpile_resolves_call_on_the_type_not_the_instance(self): + # Python's call protocol looks ``__call__`` up on the TYPE, so an instance + # attribute of that name is never what runs. ``getattr(obj, "__call__")`` + # finds it anyway, so the transpiler used to lower the shadowing body and + # return 5*99 where Python returns 5*4 -- silently wrong. The type's + # ``__call__`` must win, and it must still lower. + class Shadowed: + def __call__(self, x): + return x * 4 + + shadowed = Shadowed() + # Alone on its line: a leading statement on the same line would make the + # shadowing body unreachable for an unrelated reason and prove nothing. + shadowed.__call__ = lambda x: x * 99 + + self.assertEqual(20, shadowed(5), "the type's __call__ is what Python runs") + self.assertEqual(self._vals(shadowed, LongType(), "a long", [(5,)]), [20]) + + def test_udf_transpile_refuses_a_class_object(self): + # Calling a CLASS whose metaclass is ``type`` runs ``__init__`` and yields an + # instance, so its own ``__call__`` is never the body -- but that is the body + # the old ``getattr(func, "__call__")`` found. Pinned for the + # dynamically-created case too, where ``getsource`` cannot fall back to a + # ``ClassDef``. (A class with a custom metaclass IS callable through + # ``Meta.__call__``, and is resolved through it rather than refused.) + class Lexical: + def __init__(self, x): + self.v = x * 7 + + def __call__(self, x): + return x * 1000 + + def impl(self, x): + return x * 1000 + + Dynamic = type("Dynamic", (), {"__call__": impl}) + + with self.sql_conf(_TRANSPILE_ON): + for label, cls in [("lexical class", Lexical), ("type() class", Dynamic)]: + with self.subTest(case=label): + self.assertEqual([], UserDefinedFunction(cls, LongType()).transpiled) + + def test_udf_transpile_strips_a_bound_receiver_by_dispatch_not_by_name(self): + # The receiver used to be dropped only when literally named ``self``, so a + # bound ``__call__(this, x)`` or ``@classmethod f(cls, x)`` kept it in the + # public parameter list. That declares one parameter too many and shifts + # every ``_udf_param_N``: a two-column call returned column b's value where + # Python raises TypeError. + class Recv: + def __call__(this, x): + return x + 1 + + class Meth: + def act(this, x): + return x + 2 + + class Cls: + @classmethod + def act(kls, x): + return x + 3 + + # A ``__call__`` that is a classmethod, or one that is ALREADY a bound method, + # also has its receiver spoken for -- the class and the method's own + # ``__self__`` respectively. Both used to keep it in the public list, which + # shifts every placeholder. Each expectation below is what Python returns. + class ClsCall: + @classmethod + def __call__(kls, x): + return x + 4 + + class Helper: + def impl(self, x): + return x + 5 + + class BoundCall: + __call__ = Helper().impl + + # And the two compose: a ``staticmethod`` prepends nothing, but the method it + # wraps is already bound, so one parameter is still spoken for. Counting only + # what the descriptor prepends declared a parameter too many here. + class StaticBound: + __call__ = staticmethod(Helper().impl) + + for label, func, expected in [ + ("__call__ receiver not named self", Recv(), 6), + ("bound method receiver not named self", Meth().act, 7), + ("classmethod receiver not named self", Cls.act, 8), + ("classmethod as __call__", ClsCall(), 9), + ("already-bound method as __call__", BoundCall(), 10), + ("staticmethod over a bound method", StaticBound(), 10), + ]: + with self.subTest(case=label): + self.assertEqual(func(5), expected, "fixture must match Python's own answer") + self.assertEqual(self._vals(func, LongType(), "a long", [(5,)]), [expected]) + + # Both at once is a callable Python itself rejects: the classmethod prepends + # the class ON TOP of the method's own receiver, leaving no parameter for the + # call site. Counting one receiver returned a value where Python raises. + class ClassBound: + __call__ = classmethod(Helper().impl) + + with self.assertRaises(TypeError): + ClassBound()(5) + with self.sql_conf(_TRANSPILE_ON): + _, reasons = self._fallback_reason(ClassBound()) + self.assertIn("leaves no parameter for the call site", reasons) + + # The mirror: a ``staticmethod`` ``__call__`` prepends nothing, so its leading + # ``self`` IS supplied at the call site and both parameters are public. + # Resolving ``__call__`` on the type (rather than via getattr on the instance, + # which fires the descriptor) hands back the raw ``staticmethod``, which + # carries a ``__wrapped__`` of its own -- so for a while on this branch the + # wraps guard refused every one of them for a decorator that is not there. + class Static: + @staticmethod + def __call__(self, x): + return self + x + + self.assertEqual(14, Static()(5, 9), "staticmethod __call__ binds no receiver") + self.assertEqual( + self._vals(Static(), LongType(), "a long, b long", [(5, 9)]), + [14], + "both parameters come from the call site", + ) + + def test_udf_transpile_known_value_divergences(self): + # Transpile but DIVERGE from Python (documented in transpile.py; pinned so + # a future fix is noticed): unguarded arithmetic on NULL yields NULL + # (Python raises TypeError), and NaN > 0 is True (Python False; Spark + # orders NaN highest). Mixed str/numeric arithmetic is handled or falls + # back -- see test_udf_transpile_string_operands{,_fall_back}. + unguarded = lambda x: x + 1 # noqa: E731 + nan_gt = lambda x: (x > 0) if x is not None else None # noqa: E731 + eq_strlit = lambda x: (x == "5") if x is not None else None # noqa: E731 + self.assertEqual(self._vals(unguarded, LongType(), "a long", [(None,), (5,)]), [None, 6]) + self.assertEqual( + self._vals(nan_gt, BooleanType(), "a double", [(float("nan"),), (1.0,)]), [True, True] + ) + # `x == "5"` used to be pinned as a coercion divergence (int == "5" -> True). + # The eq category gate now drops the numeric variant, so on a long column the + # string option is pruned, nothing is left to lower (hence require_lowered + # =False), and the UDF falls back to interpreted Python -- matching Python's + # cross-type == (always False). + self.assertEqual( + self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)], require_lowered=False), + [False, False], + ) + + def test_udf_transpile_overflow_and_modulo_zero_raise(self): + # Transpiled arithmetic that raises at runtime: `*` overflow raises under + # ANSI where Python promotes to a big int (a real divergence, SPARK-55210), + # while `% 0` raises in both Spark and Python (compatible -- pinned here so + # it isn't mistaken for a divergence). + overflow = lambda x: x * x # noqa: E731 + modulo_zero = lambda x: x % 0 # noqa: E731 + self._raises(overflow, "a long", [(4000000000,)], "overflow") + self._raises(modulo_zero, "a long", [(5,)], "zero") + + def test_udf_transpile_string_operands(self): + # Textual `+`/`*` lower to Catalyst string ops and match Python: `str + + # str` -> concat, and `str * int` / `int * str` -> repeat (including a + # string column times a numeric literal). The transpiler emits a string- + # typed variant whose declared categories the JVM matches against the bound + # column types (see UserDefinedPythonFunction.builder). + S = StringType() + add = lambda a, b: a + b # noqa: E731 + mul = lambda a, b: a * b # noqa: E731 + mul3 = lambda a: a * 3 # noqa: E731 + concat_right = lambda a: a + "!" # noqa: E731 + concat_left = lambda a: "pre-" + a # noqa: E731 + repeat_lit = lambda x: "ab" * x # noqa: E731 + # (func, return_type, schema, rows, expected); arg columns come from schema. + cases = [ + (add, S, "a string, b string", [("x", "y"), ("a", "b")], ["xy", "ab"]), + (mul, S, "a string, b long", [("ab", 3)], ["ababab"]), + (mul, S, "a long, b string", [(3, "ab")], ["ababab"]), + (mul3, S, "a string", [("2",), ("ab",)], ["222", "ababab"]), + (concat_right, S, "a string", [("hi",)], ["hi!"]), + (concat_left, S, "a string", [("x",)], ["pre-x"]), + (repeat_lit, S, "a long", [(3,)], ["ababab"]), + ] + for i, (func, rt, schema, rows, expected) in enumerate(cases): + with self.subTest(case=i): + self.assertEqual(self._vals(func, rt, schema, rows), expected, f"case {i}") + + def test_udf_transpile_string_operands_fall_back(self): + # Operand/type combos with no valid string lowering for the bound column + # types fall back to the Python UDF, which raises the same way CPython does: + # `str + int` (and reversed), `str - int`, `str * str`, `str % int`, and a + # string column plus a numeric literal. The transpiler still emits numeric + # (and/or concat/repeat) variants, but none match the column types, so the + # JVM drops them and runs Python -- matching its TypeError. + add = lambda a, b: a + b # noqa: E731 + sub = lambda a, b: a - b # noqa: E731 + mul = lambda a, b: a * b # noqa: E731 + mod = lambda a, b: a % b # noqa: E731 + add5 = lambda a: a + 5 # noqa: E731 + # needle="" -> assert only that it raises (the message is CPython's). + for func, schema, rows in [ + (add, "a string, b long", [("10", 5)]), # str + int + (add, "a long, b string", [(5, "10")]), # int + str + (sub, "a string, b long", [("10", 5)]), # str - int + (mul, "a string, b string", [("a", "b")]), # str * str + (mod, "a string, b long", [("10", 3)]), # str % int + (add5, "a string", [("10",)]), # str column + numeric literal + ]: + with self.subTest(func=func, schema=schema): + self._raises(func, schema, rows, needle="") + + def test_udf_transpile_power_falls_back(self): + # `**` is intentionally not lowered (Spark's pow is DOUBLE and loses + # precision for large ints), so a UDF using it falls back to interpreted + # Python. TODO(SPARK-55210): revisit once an exact integer-power lowering + # exists. + square = lambda x: x**2 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + self.assertFalse(UserDefinedFunction(square, LongType()).transpiled) + + def test_udf_transpile_non_numeric_constant_falls_back(self): + # bool/None constants have no faithful numeric/string lowering, so + # arithmetic against them must fall back rather than emit an option that + # crashes analysis (`x * True`) or silently returns NULL (`x + None`). + mul_bool = lambda x: x * True # noqa: E731 + add_none = lambda x: x + None # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + self.assertFalse(UserDefinedFunction(mul_bool, LongType()).transpiled) + self.assertFalse(UserDefinedFunction(add_none, LongType()).transpiled) + + def test_udf_transpile_mixed_type_comparison_falls_back(self): + # Python forbids ordering across types (`a < b` for int/str -> TypeError); + # Spark would coerce and return a wrong boolean. A comparison whose + # operand categories differ is dropped (so int-vs-str `<` falls back), + # while a same-category comparison still transpiles. + def lt_mixed(a: int, b: str): + return (a < b) if a is not None and b is not None else None + + def lt_same(a: int, b: int): + return (a < b) if a is not None and b is not None else None + + with self.sql_conf(_TRANSPILE_ON): + self.assertFalse(UserDefinedFunction(lt_mixed, BooleanType()).transpiled) + self.assertTrue(UserDefinedFunction(lt_same, BooleanType()).transpiled) + + def test_udf_transpile_skips_nondeterministic(self): + # A nondeterministic UDF must not be transpiled: the optimizer could + # fold/reorder/duplicate the plain expression, dropping the barrier. + # Holds whether marked at construction or via asNondeterministic(). + plus_one = lambda x: x + 1 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + self.assertTrue(UserDefinedFunction(plus_one, LongType()).transpiled) + self.assertFalse( + UserDefinedFunction(plus_one, LongType()).asNondeterministic().transpiled + ) + self.assertFalse( + UserDefinedFunction(plus_one, LongType(), deterministic=False).transpiled + ) + + def test_udf_transpile_bool_and_binary_params(self): + # bool/bytes annotations map to the "bool"/"binary" categories and match + # Boolean/Binary columns. Identity and same-category comparison transpile + # (and match Python); boolean arithmetic has no lowering and falls back. + def bool_ident(x: bool): + return x + + def bool_lt(a: bool, b: bool): + return (a < b) if a is not None and b is not None else None + + def bool_add(x: bool): + return x + 1 # no boolean arithmetic lowering -> fall back + + def bytes_ident(x: bytes): + return x + + self.assertEqual( + self._vals(bool_ident, BooleanType(), "a boolean", [(True,), (False,), (None,)]), + [True, False, None], + ) + self.assertEqual( + self._vals( + bool_lt, + BooleanType(), + "a boolean, b boolean", + [(False, True), (True, False), (True, True)], + ), + [True, False, False], + ) + with self.sql_conf(_TRANSPILE_ON): + self.assertFalse(UserDefinedFunction(bool_add, LongType()).transpiled) + self.assertTrue(UserDefinedFunction(bytes_ident, BinaryType()).transpiled) + + def test_param_category_combos_caps_preserve_typed_pins(self): + # With more than three untyped params the cap collapses the untyped ones + # to numeric/string but keeps each typed param pinned (here a: str). + import ast as _ast + + from pyspark.sql.transpile import _param_category_combos + + fn = _ast.parse("def f(a: str, b, c, d, e): return a").body[0] + combos = _param_category_combos(fn, ["a", "b", "c", "d", "e"]) + self.assertEqual(len(combos), 2) + for combo in combos: + self.assertEqual(combo[0], "string") + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py index 2b30e41b30db2..10641b0769a0e 100644 --- a/python/pyspark/sql/tests/test_udtf.py +++ b/python/pyspark/sql/tests/test_udtf.py @@ -15,27 +15,33 @@ # limitations under the License. # -from decimal import Decimal import datetime +import logging import os import shutil import tempfile -import unittest -import logging import time +import unittest from dataclasses import dataclass +from decimal import Decimal from typing import Iterator, Optional from pyspark.errors import ( - PySparkAttributeError, - PythonException, - PySparkTypeError, AnalysisException, - PySparkPicklingError, IllegalArgumentException, + PySparkAttributeError, + PySparkPicklingError, + PySparkTypeError, + PythonException, ) -from pyspark.util import PythonEvalType +from pyspark.logger import PySparkLogger from pyspark.sql.functions import ( + AnalyzeArgument, + AnalyzeResult, + OrderingColumn, + PartitioningColumn, + SelectedColumn, + SkipRestOfInputTableException, array, col, create_map, @@ -43,12 +49,6 @@ named_struct, udf, udtf, - AnalyzeArgument, - AnalyzeResult, - OrderingColumn, - PartitioningColumn, - SelectedColumn, - SkipRestOfInputTableException, ) from pyspark.sql.types import ( ArrayType, @@ -64,7 +64,6 @@ StructType, VariantVal, ) -from pyspark.logger import PySparkLogger from pyspark.testing import assertDataFrameEqual, assertSchemaEqual from pyspark.testing.objects import ExamplePoint, ExamplePointUDT from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -74,7 +73,7 @@ pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.util import is_remote_only +from pyspark.util import PythonEvalType, is_remote_only class BaseUDTFTestsMixin: diff --git a/python/pyspark/sql/tests/test_unified_udf.py b/python/pyspark/sql/tests/test_unified_udf.py index bb1510c0dc9af..02d8f6ab2a7ca 100644 --- a/python/pyspark/sql/tests/test_unified_udf.py +++ b/python/pyspark/sql/tests/test_unified_udf.py @@ -19,16 +19,16 @@ from typing import Iterator, Tuple from pyspark.sql import functions as sf -from pyspark.sql.window import Window from pyspark.sql.functions import udf from pyspark.sql.types import LongType +from pyspark.sql.window import Window +from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import ( have_pandas, have_pyarrow, pandas_requirement_message, pyarrow_requirement_message, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.util import PythonEvalType diff --git a/python/pyspark/sql/tests/test_utils.py b/python/pyspark/sql/tests/test_utils.py index 6b7f6f3d3fbbe..e4f9420f3bc3d 100644 --- a/python/pyspark/sql/tests/test_utils.py +++ b/python/pyspark/sql/tests/test_utils.py @@ -14,44 +14,44 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import unittest import difflib +import unittest from itertools import zip_longest -from pyspark.errors import QueryContextType +import pyspark.sql.functions as F from pyspark.errors import ( AnalysisException, + IllegalArgumentException, ParseException, PySparkAssertionError, + PySparkTypeError, PySparkValueError, - IllegalArgumentException, + QueryContextType, SparkUpgradeException, - PySparkTypeError, -) -from pyspark.testing.utils import ( - assertDataFrameEqual, - assertSchemaEqual, - _context_diff, - have_numpy, - have_pandas, - have_pyarrow, ) -from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.sql import Row -import pyspark.sql.functions as F -from pyspark.sql.functions import to_date, unix_timestamp, from_unixtime +from pyspark.sql.functions import from_unixtime, to_date, unix_timestamp from pyspark.sql.types import ( - DecimalType, - StringType, ArrayType, + BooleanType, + DecimalType, + DoubleType, + FloatType, + IntegerType, LongType, - StructType, MapType, - FloatType, - DoubleType, + StringType, StructField, - IntegerType, - BooleanType, + StructType, +) +from pyspark.testing.sqlutils import ReusedSQLTestCase +from pyspark.testing.utils import ( + _context_diff, + assertDataFrameEqual, + assertSchemaEqual, + have_numpy, + have_pandas, + have_pyarrow, ) @@ -757,8 +757,8 @@ def test_assert_unequal_null_expected(self): "no pandas or numpy or pyarrow dependency", ) def test_assert_equal_exact_pandas_df(self): - import pandas as pd import numpy as np + import pandas as pd df1 = pd.DataFrame( data=np.array([(1, 2, 3), (4, 5, 6), (7, 8, 9)]), columns=["a", "b", "c"] @@ -775,8 +775,8 @@ def test_assert_equal_exact_pandas_df(self): "no pandas or numpy or pyarrow dependency", ) def test_assert_approx_equal_pandas_df(self): - import pandas as pd import numpy as np + import pandas as pd # test that asserts close enough equality for pandas df df1 = pd.DataFrame( @@ -794,8 +794,8 @@ def test_assert_approx_equal_pandas_df(self): "no pandas or numpy or pyarrow dependency", ) def test_assert_approx_equal_fail_exact_pandas_df(self): - import pandas as pd import numpy as np + import pandas as pd # test that asserts close enough equality for pandas df df1 = pd.DataFrame( @@ -838,8 +838,8 @@ def test_assert_approx_equal_fail_exact_pandas_df(self): "no pandas or numpy or pyarrow dependency", ) def test_assert_unequal_pandas_df(self): - import pandas as pd import numpy as np + import pandas as pd df1 = pd.DataFrame( data=np.array([(1, 2, 3), (4, 5, 6), (6, 5, 4)]), columns=["a", "b", "c"] @@ -881,9 +881,10 @@ def test_assert_unequal_pandas_df(self): "no pandas or numpy or pyarrow dependency", ) def test_assert_type_error_pandas_df(self): - import pyspark.pandas as ps - import pandas as pd import numpy as np + import pandas as pd + + import pyspark.pandas as ps df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) df2 = pd.DataFrame( @@ -949,9 +950,10 @@ def test_assert_equal_approx_pandas_on_spark_df(self): @unittest.skipIf(not have_pandas or not have_pyarrow, "no pandas or pyarrow dependency") def test_assert_error_pandas_pyspark_df(self): - import pyspark.pandas as ps import pandas as pd + import pyspark.pandas as ps + df1 = ps.DataFrame(data=[10, 20, 30], columns=["Numbers"]) df2 = self.spark.createDataFrame([(10,), (11,), (13,)], ["Numbers"]) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py new file mode 100644 index 0000000000000..f98910ea29f78 --- /dev/null +++ b/python/pyspark/sql/transpile.py @@ -0,0 +1,1261 @@ +# +# 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. +# +""" +Experimental tools for transpiling UDFS. + +Transpilation is only attempted when both +``spark.sql.experimental.optimizer.transpilePyUDFs=true`` and +``spark.sql.ansi.enabled=true``. The generated Catalyst expressions +target ANSI-mode SQL semantics (overflow raises, divide-by-zero raises, +etc.); running them under non-ANSI mode would silently diverge from the +Python interpretation in ways we don't currently track. If you flip +transpilation on with ANSI off the UDF will fall back to interpreted +Python execution and a warning is logged at UDF construction time. + +Python's ``+`` and ``*`` are overloaded for text (concat / repeat), so an +untyped parameter is transpiled into one option per input-type category +(numeric and string) and the JVM picks the one matching the bound column +types -- falling back to interpreted Python when none fit. Annotating the +UDF's parameters (e.g. ``def f(a: int, b: str)``) pins each category and +keeps the option matrix small; prefer doing so. To bound plan growth, +functions with more than three untyped parameters only emit the +all-numeric and all-string variants. + +A lambda is lowered only when its source names it directly and alone: bind it +to a name (``f = lambda x: x + 1``, annotated if you like) and give it a line +of its own. Passed straight to ``udf(...)``, wrapped in another call, returned +by another lambda, or sharing a line with a second lambda, nothing in the +source read back says which lambda is the UDF, so it falls back to interpreted +Python rather than risk the wrong body. +""" + +import ast +import contextlib +import inspect +import itertools +import sys +import textwrap +import threading +import warnings +from typing import TYPE_CHECKING, Any, Callable, Iterator, List, Optional, Tuple + +from pyspark.errors import UnsupportedOperationException +from pyspark.sql.column import Column +from pyspark.sql.functions import ( + abs as _abs, +) +from pyspark.sql.functions import ( + coalesce, + col, + concat, + lit, + pmod, + raise_error, + repeat, + when, +) +from pyspark.sql.types import ( + BinaryType, + BooleanType, + DataType, + DecimalType, + NumericType, + StringType, +) + +if TYPE_CHECKING: + from pyspark.sql import SparkSession + from pyspark.sql._typing import DataTypeOrString + + +class AbstractTranspiler(object): + """Base class for transpilers. All experimental.""" + + varieties: dict[str, type["AbstractTranspiler"]] = {} + # Specify the "friendly" name a user can add to spark.sql.experimental.optimizer.pyTranspilers + # to enable this transpiler. + variety: str = "" + + @classmethod + def register(cls) -> None: + AbstractTranspiler.varieties[cls.variety] = cls + + def _transpile_from_ast( + self, + src: Optional[str], + ast_info: ast.AST, + function_ast: ast.FunctionDef, + params: List[str], + returnType: "DataTypeOrString", + param_categories: Optional[dict] = None, + ) -> Optional[Column]: + """Lower ``function_ast`` to a :class:`Column`, or return ``None`` to decline. + + The override point for ``spark.sql.experimental.optimizer.pyTranspilers``. + + ``params`` is the CALLER-FACING parameter list: a receiver already bound, as + on a method or callable instance, has been removed, so ``params[i]`` is the + name bound to placeholder ``_udf_param_i`` with no offsetting needed. It is + also the list ``param_categories`` is keyed by. + """ + pass + + +def _is_definitely_basic_type(node: ast.AST) -> bool: + """ + Return True when ``node`` is statically guaranteed to produce a Python + basic/builtin type (int, float, str, bool, None, lists, etc.). + All ast.Name's are treated as basic types for now this will need to be updated + if/when we add free variables / closures to transpilation. + """ + match node: + case ast.Constant(): + return True + case ast.BinOp(left=left, right=right): + return _is_definitely_basic_type(left) and _is_definitely_basic_type(right) + case ast.UnaryOp(operand=operand): + return _is_definitely_basic_type(operand) + case ast.Name(): + return True + case _: + return False + + +def _is_definitely_boolean(node: ast.AST) -> bool: + """Return True when ``node`` is statically guaranteed to produce a Python + ``bool`` (or ``None``, which round-trips through ``coalesce``). + + Used to gate ``if``/ternary lowering: we only allow the test expression + into Catalyst's ``when(coalesce(test, false), ...)`` form when it provably + produces a boolean. Everything else (bare Name, arithmetic, function calls, + subscript, ...) must force a fallback to interpreted Python instead of + silently diverging. + """ + match node: + case ast.Constant(value=v): + return v is None or isinstance(v, bool) + case ast.Compare(left=left, comparators=comparators): + # All comparison operators of simple types bool + return all(_is_definitely_basic_type(v) for v in comparators + [left]) + case ast.BoolOp(values=values): + return all(_is_definitely_boolean(v) for v in values) + case ast.UnaryOp(op=ast.Not()): + # `not x` always produces bool. + return True + case ast.IfExp(body=body, orelse=orelse): + # Ternary is boolean only if both branches are. + return _is_definitely_boolean(body) and _is_definitely_boolean(orelse) + case _: + return False + + +class CatalystTranspiler(AbstractTranspiler): + """Transpiler that attempts to convert a Python UDF into native Spark SQL expressions.""" + + variety = "catalyst" + + # TODO (SPARK-55218): handle implicit-None return bodies like + # ``def f(x): x + x`` -- no return statement means return None; + # we should lower to lit(None) and optionally warn since it's + # likely a mistake. + def _convert_branch(self, params: List[str], statements: List[ast.stmt], slot: str) -> Column: + """Lower a single-statement if-body / if-else block. + + ``slot`` is just used to disambiguate the multi-statement error + message between the body and the else arm. + """ + if len(statements) > 1: + raise UnsupportedOperationException( + f"if statements with more than one expression in the {slot} " + "are not currently supported by the transpiler" + ) + if len(statements) == 0: + return lit(None) + return self._convert_chunk(params, statements[0]) + + def _safe_category(self, params: List[str], node: Optional[ast.AST]) -> Optional[str]: + """Best-effort input-type category for an if/else branch, or ``None`` when + it can't be pinned down statically. + + Used only to compare the two branches of an if/ternary. A ``None`` result + means "treat as compatible" (don't force a fallback): the node is absent, + is a bare ``None`` literal (which unifies with any branch type via + ``coalesce``/``Cast``), or its category can't be determined. + """ + if node is None: + return None + # If-statement branches arrive as ``Return`` statements; classify the + # returned value, not the statement wrapper (``_is_definitely_boolean`` + # has no ``Return`` case, so without this a boolean-returning branch + # would fall through to ``_category``'s numeric catch-all). + if isinstance(node, ast.Return): + return self._safe_category(params, node.value) + # An if-statement's category is its branches' common category (the + # ``_category`` catch-all would mislabel every ``ast.If`` "numeric"). + # Mismatched branches return None ("can't be pinned down"); the + # branch-compatibility check in ``_convert_if_like`` raises for them. + if isinstance(node, ast.If): + body_c = self._safe_category(params, node.body[0]) if node.body else None + else_c = self._safe_category(params, node.orelse[0]) if node.orelse else None + if body_c is not None and else_c is not None and body_c != else_c: + return None + return body_c if body_c is not None else else_c + if isinstance(node, ast.Constant) and node.value is None: + return None + # Comparisons / ``not`` / boolean ops produce a boolean column; classify + # them as "bool" (``_category``'s catch-all would mislabel them numeric). + if _is_definitely_boolean(node): + return "bool" + try: + return self._category(params, node) + except UnsupportedOperationException: + return None + + def _convert_if_like( + self, + params: List[str], + test_col: Column, + body_col: Column, + else_col: Column, + test_node: ast.AST, + body_node: Optional[ast.AST], + else_node: Optional[ast.AST], + ) -> Column: + # We cannot soundly lower a generic Python truthiness test here. + # Python truthiness depends on the runtime input type and value: + # for example, 0, 0.0, "", empty collections, and None are all + # falsy, while most other values are truthy. The transpiler does + # not have enough input type information at this point to decide + # whether ``test_col`` is a boolean expression or a bare value + # whose truthiness would need Python-specific handling. Emitting + # ``when(coalesce(test_col, false), ...)`` is therefore unsound: + # it can either fail Spark analysis for non-boolean columns or + # silently diverge from Python semantics. Fail closed so the UDF + # falls back to interpreted Python execution instead. + if not _is_definitely_boolean(test_node): + raise UnsupportedOperationException( + f"bare truthiness tests ({ast.dump(test_node)}) in if-expressions are " + " not currently supported by the transpiler" + ) + # When the two branches resolve to concrete but different categories + # (e.g. numeric vs string), the lowered ``when(...).otherwise(...)`` is a + # CASE WHEN whose branch values share no common type under ANSI. That node + # is carried as a child of the TranspiledPythonUDF and is type-checked by + # CheckAnalysis *before* ConvertToCatalyst can drop it, so it would fail + # the whole query rather than fall back. Refuse here so the UDF runs as + # interpreted Python instead. Branches whose category we can't pin down + # (e.g. a bare ``None``) are treated as compatible and don't force this. + body_cat = self._safe_category(params, body_node) + else_cat = self._safe_category(params, else_node) + if body_cat is not None and else_cat is not None and body_cat != else_cat: + raise UnsupportedOperationException( + f"if/else branches have incompatible categories ({body_cat} vs " + f"{else_cat}); the lowered CASE WHEN has no common type under ANSI, " + "so the transpiler falls back to interpreted Python" + ) + safe_test = coalesce(test_col, lit(False)) + return when(safe_test, body_col).otherwise(else_col) + + def _lower_eq( + self, + params: List[str], + left_node: ast.AST, + right_node: ast.AST, + equal: bool, + ) -> Column: + """Lower ``==`` / ``!=`` with Python's None-equality semantics. + + Unlike ordering operators, Python doesn't raise on ``None == x`` / + ``None != x``: ``None == None`` is True, ``None == 0`` is False, + and ``!=`` is the negation. Spark's ``==`` returns NULL on NULL + operands (three-valued logic), which would round-trip through + the UDF as ``None`` rather than the bool Python would have + produced. Hand-roll the four cases via ``when`` branches. + + When the two operands resolve to concrete but DIFFERENT categories + (e.g. ``x == True`` on a numeric column, or ``x == "5"`` under the + numeric variant), the lowered ``=`` either fails analysis under ANSI + (bool vs bigint) -- which would break a working UDF since the option + is type-checked before ConvertToCatalyst can drop it -- or coerces + where Python's ``==`` is simply False. Refuse those so the UDF falls + back to interpreted Python. A ``None`` literal operand stays allowed + (the four-branch NULL handling above reproduces Python exactly). + + One value-level difference remains (needs runtime values, so it is + documented, not guarded): Spark treats ``NaN = NaN`` as true, while + Python's ``nan == nan`` is False. + """ + lc = self._safe_category(params, left_node) + rc = self._safe_category(params, right_node) + if lc is not None and rc is not None and lc != rc: + raise UnsupportedOperationException( + f"`==`/`!=` operands have incompatible categories ({lc} vs {rc}); " + "Python compares across types as unequal while Spark would coerce " + "or fail analysis, so the transpiler falls back to interpreted Python" + ) + left_col = self._convert_chunk(params, left_node) + right_col = self._convert_chunk(params, right_node) + left_null = left_col.isNull() + right_null = right_col.isNull() + if equal: + both_null_val: Column = lit(True) + one_null_val: Column = lit(False) + value_cmp = left_col == right_col + else: + both_null_val = lit(False) + one_null_val = lit(True) + value_cmp = left_col != right_col + return ( + when(left_null & right_null, both_null_val) + .when(left_null | right_null, one_null_val) + .otherwise(value_cmp) + ) + + def _lower_value_compare( + self, + params: List[str], + left_node: ast.AST, + right_node: ast.AST, + op: Callable[[Column, Column], Column], + op_repr: str, + ) -> Column: + """Lower a value comparison (``<``, ``<=``, ``>``, ``>=``). + + Python raises ``TypeError`` when an operand of these operators is + ``None`` (e.g. ``None > 0``), whereas Spark's three-valued logic + returns ``NULL``. To stay faithful to the source UDF we guard the + comparison: if either operand is ``NULL`` we raise via + ``raise_error``, otherwise we evaluate ``left op right`` as usual. + Callers that have already proven the operand non-null (``if x is + not None: x > 0``) take the otherwise branch, so they never trip + the raise. + + Python also forbids ordering across types (``1 < "a"`` -> TypeError), + whereas Spark would coerce the operands and return a (wrong) boolean. + We therefore only lower when both operands share a category; a + mismatch raises so this variant is dropped and the UDF falls back to + interpreted Python rather than silently diverging. + + One value-level difference from Python remains (it needs runtime + value info, so it is documented, not guarded): Spark orders ``NaN`` + as greater than every value, whereas Python's ``NaN`` comparisons + are all ``False``. + """ + lc = self._category(params, left_node) + rc = self._category(params, right_node) + if lc != rc: + raise UnsupportedOperationException( + f"`{op_repr}` compares operands of different categories " + f"({lc} vs {rc}); Python would raise TypeError, so the " + "transpiler falls back to interpreted Python" + ) + left_col = self._convert_chunk(params, left_node) + right_col = self._convert_chunk(params, right_node) + null_guard = left_col.isNull() | right_col.isNull() + err = lit( + "Python UDF transpiler: cannot compare NULL with operator " + f"`{op_repr}`; Python would raise TypeError here. Add an " + "`is not None` guard or filter NULLs upstream." + ) + return when(null_guard, raise_error(err)).otherwise(op(left_col, right_col)) + + def _category(self, params: List[str], node: ast.AST) -> str: + """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current + ``self._param_categories`` assumption (set per input-type variant). + + Drives operator selection (``+`` -> add vs concat, ``*`` -> multiply vs + repeat) and raises ``UnsupportedOperationException`` when an operator's + operands are type-incompatible, so the caller drops that variant and the + JVM picks another option / falls back to the Python UDF. + """ + match node: + case ast.Constant(value=v): + # bool subclasses int, so classify it first: int/float -> numeric, + # str -> string, bool -> bool, bytes -> binary. None/complex/ + # Ellipsis have no usable Spark column type, so raise to drop this + # variant and fall back rather than emit an option that fails + # CheckAnalysis or silently diverges (e.g. `x + None` -> NULL where + # Python raises TypeError). + if isinstance(v, bool): + return "bool" + if isinstance(v, bytes): + return "binary" + if isinstance(v, (int, float)): + return "numeric" + if isinstance(v, str): + return "string" + raise UnsupportedOperationException( + f"constant {v!r} ({type(v).__name__}) has no usable column " + "category; falling back to interpreted Python" + ) + case ast.Name(id=name) if name in params: + # ``params`` is the caller-facing list, so its indexes are already + # the ``_udf_param_N`` / category indexes -- see ``_transpile_func``. + return self._param_categories.get(params.index(name), "numeric") + case ast.BinOp(left=left, op=op, right=right): + lc = self._category(params, left) + rc = self._category(params, right) + if isinstance(op, ast.Add) and lc == rc: + return lc # str + str -> str, num + num -> num + if isinstance(op, ast.Mult): + if {lc, rc} == {"numeric", "numeric"}: + return "numeric" + if {lc, rc} == {"numeric", "string"}: + return "string" # str * int / int * str -> repeat + if isinstance(op, (ast.Sub, ast.Mod)) and lc == rc == "numeric": + return "numeric" + raise UnsupportedOperationException( + f"operands of `{type(op).__name__}` are not type-compatible " + "for this input-type variant" + ) + case ast.Return(value=value) if value is not None: + return self._category(params, value) + case ast.IfExp(body=if_body, orelse=if_orelse): + # A ternary's category is its branches' common category. Without + # this arm the catch-all labeled every IfExp "numeric", so e.g. + # `("5" if c else "6") == 5` passed the equality guard as + # numeric-vs-numeric and Spark's string-number coercion silently + # diverged from Python's cross-type `==` (always False). A + # None-literal branch adopts the other branch's category (NULL + # unifies with any type in the lowered CASE WHEN); mismatched or + # all-None branches raise so the variant is dropped. + def branch_category(b: ast.AST) -> Optional[str]: + if isinstance(b, ast.Constant) and b.value is None: + return None + return self._category(params, b) + + body_cat = branch_category(if_body) + else_cat = branch_category(if_orelse) + if body_cat is not None and else_cat is not None and body_cat != else_cat: + raise UnsupportedOperationException( + f"ternary branches have mismatched categories ({body_cat} " + f"vs {else_cat}) and cannot drive operator selection" + ) + result_cat = body_cat if body_cat is not None else else_cat + if result_cat is None: + raise UnsupportedOperationException( + "ternary with all-None branches has no usable column category" + ) + return result_cat + case _ if _is_definitely_boolean(node): + # Comparisons, `not`, and boolean ops produce a boolean column. + # Labeling them "numeric" (the old catch-all) let booleans into + # arithmetic/equality lowerings where ANSI analysis fails (e.g. + # `(x > 0) + 1`, valid Python) instead of falling back. + return "bool" + case _: + # Remaining nodes (unsupported calls, subscripts, ...) don't + # drive concat/repeat selection and are rejected later by + # `_convert_chunk`; treat as numeric for category purposes. + return "numeric" + + def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: + match body: + case None: + # Special case literal None, the implicit return None + return lit(None) + case ast.UnaryOp(op=ast.Not(), operand=operand): + # Python's `not None` is `True` (None is falsy), but Spark's + # `~NULL` is `NULL`. Coalesce against `lit(True)` so a NULL + # operand mirrors Python's "None is falsy" rule. We only + # accept operands that are statically known to be boolean; + # for non-boolean operands (e.g. `not 0`, `not x` where x is + # a bare parameter name) Spark's `~` is bitwise, not Python + # truthiness, so we bail and let the caller fall back to + # interpreted Python rather than silently diverge. + if not _is_definitely_boolean(operand): + raise UnsupportedOperationException( + "`not` operand type is not statically known to be " + "boolean; Spark's `~` is bitwise, not Python " + "truthiness, so the transpiler refuses to lower this " + "and the UDF falls back to interpreted Python" + ) + return coalesce(self._convert_chunk(params, operand).__invert__(), lit(True)) + case ast.UnaryOp(op=(ast.USub() | ast.UAdd()) as op, operand=operand): + # `-x` / `+x` -- like the binary arithmetic operators, only + # lower for numeric operands. Python raises TypeError for + # unary +/- on strings, but Spark's ANSI string promotion + # would silently coerce the string to double (`-'5'` -> + # -5.0), and a boolean operand emits UnaryMinus(bool), which + # fails CheckAnalysis outright -- breaking the query instead + # of falling back, since the option is type-checked as a + # child of TranspiledPythonUDF before ConvertToCatalyst can + # drop it. Fail closed for every non-numeric category. + if self._category(params, operand) != "numeric": + raise UnsupportedOperationException( + "unary `+`/`-` is only supported for numeric operands " + "(Python raises TypeError on strings, and Spark would " + "coerce or fail analysis); the transpiler falls back " + "to interpreted Python" + ) + if isinstance(op, ast.USub): + # Handles both literal negative ints (USub on a Constant) + # and runtime negation of a column. + return self._convert_chunk(params, operand).__neg__() + # `+x` -- identity, kept for symmetry with USub. + return self._convert_chunk(params, operand) + case ast.BoolOp(op=op, values=values): + # Python `and` / `or` short-circuit and return one of the + # operands rather than a strict boolean. For the booleans + # produced by Compare / UnaryOp(Not) / nested BoolOps this + # maps cleanly onto Spark Column `&` / `|`. For + # non-boolean operands (including bare parameter names whose + # runtime type is unknown) the right semantics would require + # Python's truthiness rules (0 / "" / None / [] all + # falsy), which we can't faithfully reproduce without the + # input column types -- Spark's `&` / `|` would silently + # do bitwise instead. Require all operands to be statically + # known boolean so the caller falls back to interpreted + # Python rather than producing a plan whose results diverge. + if not all(_is_definitely_boolean(v) for v in values): + raise UnsupportedOperationException( + "`and` / `or` operand type is not statically known " + "to be boolean; Spark's `&` / `|` are bitwise, not " + "Python truthiness, so the transpiler refuses to " + "lower this and the UDF falls back to interpreted " + "Python" + ) + # A literal None operand short-circuits differently: Python's + # `None and (x > 0)` returns None regardless of x, but Spark's + # three-valued `null AND false` is false (and `null OR true` is + # true), so the lowered form diverges. `_is_definitely_boolean` + # accepts None for `not`/if-test contexts where coalesce handles + # it; here it must force a fallback instead. + if any(isinstance(v, ast.Constant) and v.value is None for v in values): + raise UnsupportedOperationException( + "literal None operand in `and` / `or` cannot be lowered: " + "Spark's three-valued logic diverges from Python's " + "short-circuit-return-operand semantics, so the UDF " + "falls back to interpreted Python" + ) + cols = [self._convert_chunk(params, v) for v in values] + if isinstance(op, ast.And): + result = cols[0] + for c in cols[1:]: + result = result & c + return result + if isinstance(op, ast.Or): + result = cols[0] + for c in cols[1:]: + result = result | c + return result + raise UnsupportedOperationException(f"BoolOp operator {op} is not supported") + case ast.IfExp(test=test, body=body_expr, orelse=orelse_expr): + # Ternary `body if test else orelse` -- shares the + # NULL-as-falsy lowering with the if-statement case. + return self._convert_if_like( + params, + self._convert_chunk(params, test), + self._convert_chunk(params, body_expr), + self._convert_chunk(params, orelse_expr), + test, + body_expr, + orelse_expr, + ) + case ast.If(test, success, orelse): + return self._convert_if_like( + params, + self._convert_chunk(params, test), + self._convert_branch(params, success, "body"), + self._convert_branch(params, orelse, "else body"), + test, + success[0] if success else None, + orelse[0] if orelse else None, + ) + case ast.Compare(left, ops, comps): + if len(ops) != 1 or len(comps) != 1: + raise UnsupportedOperationException( + "chained comparisons (e.g. `a < b < c`) are not supported by the transpiler" + ) + comp = comps[0] + match ops[0]: + case ast.Is() | ast.IsNot(): + # Only lower `x is None` / `None is x` (and their + # `is not` variants) to isNull/isNotNull. For any + # other comparator (e.g. `x is 0`, `x is y`) Python + # performs an object-identity check that has no SQL + # equivalent, so we must fall back to interpreted + # Python rather than silently emitting a null check. + is_none_left = isinstance(left, ast.Constant) and left.value is None + is_none_right = isinstance(comp, ast.Constant) and comp.value is None + if not (is_none_left or is_none_right): + raise UnsupportedOperationException( + "`is`/`is not` is only supported when one " + "operand is the literal None; other identity " + "checks (e.g. `x is 0`, `x is y`) cannot be " + "lowered to SQL and the UDF falls back to " + "interpreted Python" + ) + subject_node = comp if is_none_left else left + subject_col = self._convert_chunk(params, subject_node) + if isinstance(ops[0], ast.Is): + return subject_col.isNull() + else: + return subject_col.isNotNull() + case ast.Eq(): + return self._lower_eq(params, left, comp, equal=True) + case ast.NotEq(): + return self._lower_eq(params, left, comp, equal=False) + case ast.Lt(): + return self._lower_value_compare( + params, left, comp, lambda l, r: l < r, "<" + ) + case ast.LtE(): + return self._lower_value_compare( + params, left, comp, lambda l, r: l <= r, "<=" + ) + case ast.Gt(): + return self._lower_value_compare( + params, left, comp, lambda l, r: l > r, ">" + ) + case ast.GtE(): + return self._lower_value_compare( + params, left, comp, lambda l, r: l >= r, ">=" + ) + case _: + raise UnsupportedOperationException( + f"comparison operator {type(ops[0]).__name__} " + "is not supported by the transpiler" + ) + case ast.BinOp(left=left, op=op, right=right): + # Operator selection is driven by the operand *categories* under + # the current input-type variant (see ``_category``): Python's + # `+` / `*` are overloaded for text. `+` -> add (num,num) or + # concat (str,str); `*` -> multiply (num,num) or repeat (str,int + # / int,str); `-` / `%` are numeric-only. Combos that don't fit + # (str+int, str-str, ...) raise so this variant is dropped and + # the JVM picks another option or falls back to the Python UDF. + # + # `**` is intentionally NOT lowered: Spark's `pow` is DOUBLE and + # loses precision for large integers, so it would silently return + # wrong results. TODO (SPARK-55210): add an exact integer-power + # lowering and re-enable it. + # + # Value-level divergences remain documented (need runtime value + # info, not type): overflow raises ARITHMETIC_OVERFLOW under ANSI + # where Python promotes to a big int; arithmetic is not + # NULL-guarded (`x + 1` on NULL -> NULL vs Python TypeError). + # TODO (SPARK-55210): map overflow / divide-by-zero precisely. + lc = self._category(params, left) + rc = self._category(params, right) + left_col = self._convert_chunk(params, left) + right_col = self._convert_chunk(params, right) + match op: + case ast.Add(): + if lc == rc == "string": + return concat(left_col, right_col) + if lc == rc == "numeric": + return left_col.__add__(right_col) + case ast.Sub(): + if lc == rc == "numeric": + return left_col.__sub__(right_col) + case ast.Mult(): + if lc == "numeric" and rc == "numeric": + return left_col.__mul__(right_col) + if lc == "string" and rc == "numeric": + return repeat(left_col, right_col.cast("int")) + if lc == "numeric" and rc == "string": + return repeat(right_col, left_col.cast("int")) + case ast.Mod(): + if lc == rc == "numeric": + # Python's `%` takes the sign of the divisor; Spark's + # takes the dividend's. `sign(b) * pmod(sign(b) * a, + # abs(b))` reproduces Python for every non-zero divisor + # except at the LongType overflow boundaries -- `a = + # Long.MinValue` with `b < 0` (the `sign(b) * a` negate + # overflows) and `b = Long.MinValue` (the `abs(b)` + # overflows) -- where this raises ARITHMETIC_OVERFLOW + # under ANSI while Python returns a value. That matches + # the documented overflow caveat for `+`/`-`/`*` above. + # Use a CASE-based integer sign rather than sign() to + # avoid promoting operands to DoubleType, which loses + # precision near LongType boundaries. + sb = ( + when(right_col > 0, lit(1)) + .when(right_col < 0, lit(-1)) + .otherwise(lit(0)) + ) + return sb * pmod(sb * left_col, _abs(right_col)) + case _: + raise UnsupportedOperationException( + f"binary operator {type(op).__name__} is not " + "supported by the transpiler" + ) + raise UnsupportedOperationException( + f"`{type(op).__name__}` operands are not type-compatible for " + "this input-type variant" + ) + case ast.Return(value=value): + return self._convert_chunk(params, value) + case ast.Constant(value=value): + # Avoid circular import issue. + return lit(value) + case ast.Name(id=name, ctx=ast.Load()): + # Insert columns referencing the param indexes for children + if name in params: + # ``params`` excludes any bound receiver (see ``_transpile_func``), + # so its indexes ARE the placeholder indexes. A body referencing + # the receiver (``return self``) is not in this list and so takes + # the branch below, which refuses -- there is no column for it. + return col(f"_udf_param_{params.index(name)}") + else: + # TODO (SPARK-55207): Handle assignments, class vars, and closures + # via scope evaluation. + raise UnsupportedOperationException( + f"name {name!r} is not in the UDF's parameter list " + "and free variables / closures are not supported" + ) + case _: + raise UnsupportedOperationException( + f"AST node {type(body).__name__} is not supported by the " + f"transpiler ({ast.dump(body)[:120]})" + ) + + def _transpile_from_ast( + self, + src: Optional[str], + ast_info: ast.AST, + function_ast: ast.FunctionDef, + params: List[str], + returnType: "DataTypeOrString", + param_categories: Optional[dict] = None, + ) -> Optional[Column]: + # Short circuit on nothing to transpile. + if src == "" or ast_info is None: + return None + # Per-variant input-type assumption ({public_param_index -> category}), + # read by ``_category`` to choose str vs numeric operators. + self._param_categories = param_categories or {} + function_body = function_ast.body + if len(function_body) != 1: + raise UnsupportedOperationException( + "functions with more than one top-level statement are not " + "supported by the transpiler" + ) + # Refuse variants whose body category does not MATCH the declared + # return type's category. Two distinct failure modes hide here: + # + # * A cast that can never resolve (binary -> numeric, bool -> binary): + # the options are type-checked by CheckAnalysis as children of + # TranspiledPythonUDF before ConvertToCatalyst could drop them, so + # the whole query fails instead of falling back. + # * A cast that IS analysis-valid but that the interpreted + # SQL_BATCHED_UDF path never performs: EvaluatePython.makeFromJava + # accepts only the expected JVM types for the declared return type + # and nulls everything else. E.g. `def f(s: str): return s` declared + # LongType() returns NULL interpreted, but a lowered + # cast(string as bigint) would return 123 for '123' (or raise + # CAST_INVALID_INPUT for 'abc') -- a silent divergence. + # + # So require the strict match: numeric -> non-decimal NumericType + # (DecimalType is excluded like it is for inputs: the interpreted + # converter accepts only decimal.Decimal results there and nulls the + # ints/floats these lowerings produce), string -> StringType, bool -> + # BooleanType, binary -> BinaryType. An unknown category (e.g. a bare + # None body) lowers to NULL, which every return type accepts as NULL + # on both paths. Within-numeric conversions (e.g. a bigint body cast + # to a double return type) are intentionally still allowed and + # documented as the transpiled-cast behavior pinned by + # test_udf_transpile_casts_to_return_type. + if isinstance(returnType, DataType): + body_cat = self._safe_category(params, function_body[0]) + cast_ok = ( + body_cat is None + or ( + body_cat == "numeric" + and isinstance(returnType, NumericType) + and not isinstance(returnType, DecimalType) + ) + or (body_cat == "string" and isinstance(returnType, StringType)) + or (body_cat == "bool" and isinstance(returnType, BooleanType)) + or (body_cat == "binary" and isinstance(returnType, BinaryType)) + ) + if not cast_ok: + raise UnsupportedOperationException( + f"a {body_cat}-typed lowering does not match the declared " + f"return type {returnType.simpleString()}; the interpreted " + "path would return NULL where the lowered cast would " + "convert (or fail), so the transpiler falls back to " + "interpreted Python" + ) + converted = self._convert_chunk(params, function_body[0]) + # Cast to the declared return type so the rewritten plan reports a + # known data type to the optimizer's plan validator (otherwise it + # sees an UnresolvedFunction tree and reports VOID, which fails + # the schema-stability check on this rule). + return converted.cast(returnType) + + +CatalystTranspiler.register() + + +def _get_transpilers(session: "SparkSession") -> List[AbstractTranspiler]: + """Get the transpilers we should try.""" + configured_transpilers = session.conf.get("spark.sql.experimental.optimizer.pyTranspilers") + if not configured_transpilers: + return [] + transpiler_names = configured_transpilers.split(",") + return [ + AbstractTranspiler.varieties[name]() + for name in transpiler_names + if name in AbstractTranspiler.varieties + ] + + +def _annotation_category(annotation: Optional[ast.AST]) -> Optional[str]: + """Map a parameter's type annotation to a category + (``"numeric"``/``"string"``/``"bool"``/``"binary"``), or ``None`` when it's + absent or unrecognised (the caller then tries both numeric and string).""" + name: Optional[str] = None + if isinstance(annotation, ast.Name): + name = annotation.id + elif isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + name = annotation.value # stringized annotation, e.g. def f(a: "int") + # str -> "string", int/float -> "numeric", bool -> "bool", bytes -> "binary" + # (matching the constant handling in ``_category``). complex and anything + # unrecognised return None so the caller tries both numeric and string. + if name == "str": + return "string" + if name in ("int", "float"): + return "numeric" + if name == "bool": + return "bool" + if name == "bytes": + return "binary" + return None + + +def _param_category_combos(function_ast: ast.FunctionDef, public_params: List[str]) -> List[dict]: + """Per-variant maps ``{public_param_index -> category}`` where category is + one of ``"numeric"``/``"string"``/``"bool"``/``"binary"``. + + A typed param (``def f(a: str, b: int)``) is pinned to its category; an + untyped param is tried as both numeric and string. To cap plan growth, when + more than three params are untyped we collapse the untyped ones to the + all-numeric and all-string variants (encourage typing inputs to keep the + matrix small) while keeping every typed param pinned. + """ + n = len(public_params) + public_args = function_ast.args.args[len(function_ast.args.args) - n :] + candidates: List[List[str]] = [] + untyped = 0 + for arg in public_args: + cat = _annotation_category(arg.annotation) + if cat is None: + candidates.append(["numeric", "string"]) + untyped += 1 + else: + candidates.append([cat]) + if untyped > 3: + # Cap the 2**untyped blow-up, but keep each typed param pinned to its + # category (a single-element ``candidates`` entry); only the untyped + # params collapse to the all-numeric / all-string pair. + return [ + {i: c[0] if len(c) == 1 else fill for i, c in enumerate(candidates)} + for fill in ("numeric", "string") + ] + return [{i: choice[i] for i in range(n)} for choice in itertools.product(*candidates)] or [{}] + + +def _call_dunder(func: Callable) -> Any: + """The ``__call__`` entry from ``func``'s type. + + Not ``getattr(func, "__call__")``, which is wrong in two ways that both end with + lowering a body that never runs: an instance attribute ``obj.__call__ = f`` + shadows the type's for ``getattr`` but is ignored when ``obj`` is called, and on + a CLASS object it finds the ``__call__`` its instances use while calling the + class runs ``__init__``. + + ``getattr_static`` looks the name up without firing the descriptor protocol, so + deciding what to transpile never runs user code -- a custom descriptor used as + ``__call__`` would otherwise have its ``__get__`` called here. + + Everything comes back undisturbed, so a ``staticmethod`` or ``classmethod`` + arrives as the descriptor rather than the function inside it -- see + ``_call_impl``. There is always something to return: ``getattr_static`` on a type + falls through to the metatype, so the floor is ``type.__call__``. + """ + return inspect.getattr_static(type(func), "__call__") + + +def _call_impl(entry: Any) -> Any: + """The function inside a ``staticmethod`` / ``classmethod``, else ``entry`` itself. + + Both get in the way, in opposite directions: they do not forward the wrapped + function's ``__code__``, and they synthesize a ``__wrapped__`` pointing at it even + when no decorator is involved. So asking the descriptor directly finds no code + object and a wraps decorator that is not there -- unwrap before either question. + + Narrow on purpose: unwrapping any ``__func__`` would follow the attribute on + unrelated callables that expose one, and read the wrong code object. + """ + return entry.__func__ if isinstance(entry, (staticmethod, classmethod)) else entry + + +def _held_code(func: Callable) -> Any: + """The code object that runs when ``func`` is called, or ``None``. + + A function or method runs its own ``__code__``; anything else runs its type's + ``__call__``. Used only to ask whether we are holding a lambda. + """ + target = func if (inspect.isfunction(func) or inspect.ismethod(func)) else _call_dunder(func) + return getattr(_call_impl(target), "__code__", None) + + +_WARNINGS_LOCK = threading.Lock() + + +@contextlib.contextmanager +def _syntax_warnings_suppressed() -> Iterator[None]: + """Parse without re-emitting, or tripping over, the source's own SyntaxWarnings. + + The import already reported them. Without this, ``udf()`` repeats the warning, + and under warnings-as-errors the parse raises and lowering silently turns off. + Before 3.12 an invalid escape sequence was a DeprecationWarning, so ignore that + too rather than lose lowering on the oldest Python we support. + + The lock serializes our own use of ``warnings``, whose state is process-global. + It cannot serialize anyone else's: a thread entering ``catch_warnings`` while + this is open has its filters restored from our older snapshot on exit, and + entering at all bumps the filter version, so a "once"-filtered warning + elsewhere in the process can fire again. Both are inherent to the stdlib API, + and are why the parse is the only thing inside here. + """ + with _WARNINGS_LOCK: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=SyntaxWarning) + if sys.version_info < (3, 12): + warnings.filterwarnings( + "ignore", message="invalid escape sequence", category=DeprecationWarning + ) + yield + + +def _get_src_ast_from_func(func: Callable) -> Tuple[Optional[str], Optional[ast.AST]]: + """Try and get the AST from a given callable + + KNOWN LIMITATION: this is the source on disk NOW, not necessarily the source + ``func`` was compiled from. ``inspect.getsource`` reads through ``linecache``, + which re-reads an edited file while the code object stays as it was at import, + so editing a module in a long-lived driver and then building a UDF from a + function imported earlier lowers the NEW body while Python runs the old one -- + verified: rewriting ``lambda x: x + 1`` to ``x * 9`` gives Python 6, Spark 45. + Closing it needs the parsed node checked against the held code object; it is + not tracked separately, being part of the experimental transpiler + (SPARK-54783). Until then, transpilation assumes source files are not edited + underneath a running session. + """ + # Note: consider maybe dill? (see the JYTHON PR) + # inspect getsource does not work for functions defined in vanilla + # repl, but does for those in files or in ipython. + # It also fails when we give it an instance of a callable class. + try: + src = inspect.getsource(func) + src = textwrap.dedent(src).strip() + with _syntax_warnings_suppressed(): + ast_info = ast.parse(src) + except Exception: + try: + src = inspect.getsource(_call_dunder(func)) + src = textwrap.dedent(src).strip() + with _syntax_warnings_suppressed(): + ast_info = ast.parse(src) + except Exception: + # No usable source (REPL/stdin definition, builtin, ...) -- + # return cleanly so the caller reports "cannot transpile" + # instead of surfacing an UnboundLocalError as the reason. + return None, None + return src, ast_info + + +def _get_parameter_list(node: ast.FunctionDef) -> list[str]: + """Return the positional argument names in order.""" + return [arg.arg for arg in node.args.args] + + +def _get_function_from_ast(body: ast.AST, held_code: Any) -> Tuple[Optional[ast.FunctionDef], str]: + """ + Extract a :class:`ast.FunctionDef` node from an AST produced by + ``ast.parse(inspect.getsource(udf_func))``. + + Handles the following source patterns (in order): + + * ``f = lambda x: x + 1`` -- lambda bound to a name, annotated or not + * ``lambda x: x + 1`` -- bare expression (getsource on a raw lambda) + * ``def f(x): ... return x + 1`` + * a class with a ``__call__`` method + + ``held_code`` is the code object that runs when the callable is called; a + ``co_name`` of ``<lambda>`` is what makes the ambiguity checks below apply, and + its parameter names are what tell a located lambda apart from a rival. + + Returns the node and an empty reason, or ``None`` and why -- paired so no refusal + reaches the caller unexplained. + """ + if not hasattr(body, "body") or not body.body: + return None, "no statement was found in the source read for this callable" + + stmt = body.body[0] + + # Grab the value side of a top level assign (e.g. x = lambda ...). An annotated + # binding is the same shape, and the form a typed codebase writes. + if isinstance(stmt, ast.Assign): + stmt = stmt.value + elif isinstance(stmt, ast.AnnAssign) and stmt.value is not None: + stmt = stmt.value + + # Bare ``lambda x: ...`` (when ``inspect.getsource`` returns a raw + # lambda expression at module top level) parses as ``Expr(Lambda)``. + if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Lambda): + stmt = stmt.value + + # ``inspect.getsource`` works in whole lines, so refuse unless the lambda located + # here IS the one we hold: anything else lowers a body that never runs + # (SPARK-58650). + if getattr(held_code, "co_name", None) == "<lambda>": + if not isinstance(stmt, ast.Lambda): + return None, ( + "the source read for this lambda does not define it as a statement of " + "its own -- it is wrapped in a call or a tuple assignment, or a " + "surrounding definition, or the file has changed since import and no " + "longer holds it -- so which lambda to lower cannot be determined" + ) + # The located lambda must take the parameters the held one does, or it is a + # different lambda that merely sits where ours was read from. This is what + # separates a lambda nested in the body of the one we hold (fine: it can never + # be the UDF) from one that RETURNED the lambda we hold, as in the one-line + # ``make_adder = lambda n: lambda x: x + n`` -- there the outer lambda is + # located and the inner is held, and lowering the outer would be wrong. + located_args = [arg.arg for arg in stmt.args.args] + if located_args != list(held_code.co_varnames[: held_code.co_argcount]): + return None, ( + "the lambda defined in the source read for this one takes different " + f"parameters ({', '.join(located_args) or 'none'}), so it is not the " + "lambda being transpiled -- a lambda returning another lambda on one " + "line, or a file changed since import; put each lambda on its own line" + ) + # Only lambdas OUTSIDE ``stmt`` are rivals; one in its body cannot be the UDF, + # and the user could not split it onto another line. + own = set(map(id, ast.walk(stmt))) + if any(id(node) not in own for node in ast.walk(body) if isinstance(node, ast.Lambda)): + return None, ( + "more than one lambda is visible in the source line(s) this one was " + "read from, and nothing there says which is the UDF, so it is not " + "safe to lower; put each lambda on its own line to transpile it" + ) + + if isinstance(stmt, ast.Lambda): + # Synthesize a one-statement FunctionDef wrapping the lambda body so + # the rest of the transpiler can treat lambdas and ``def`` uniformly. + fn_ctor: Any = ast.FunctionDef + synthesized = fn_ctor( + name="<lambda>", + args=stmt.args, + body=[ast.Return(value=stmt.body)], + decorator_list=[], + ) + # A node without ``lineno`` cannot be unparsed or compiled; seed from the + # lambda so positions point at real source rather than line 1. + return ast.fix_missing_locations(ast.copy_location(synthesized, stmt)), "" + + if isinstance(stmt, ast.FunctionDef): + return stmt, "" + return None, ( + f"the source read for this callable is a {type(stmt).__name__}, which the " + "transpiler cannot reduce to a single function definition" + ) + + +def _transpile_func( + session: "SparkSession", + func: Callable[..., Any], + returnType: "DataTypeOrString", +) -> Tuple[List[Column], List[str], List[str], List[List[str]]]: + """ + An experimental internal function that attempts to transpile a callable function. + + Returns + ------- + list of transpiled options (one per backend x input-type variant) + list of errors as strings + list of positional parameter names (excluding a receiver already bound, as on a + method or callable instance) -- needed so the caller can resolve named-argument + invocations to positional order at call time, since the ``_udf_param_N`` + substitution in :class:`UserDefinedPythonFunction` is positional. + list of per-option input-type categories (``"numeric"`` / ``"string"`` per + public param) -- the JVM picks the option whose categories match the bound + column types, or falls back to the Python UDF when none match. + """ + try: + # The transpiler lowers to atomic (numeric/string/boolean/binary) + # expressions and casts the result to the declared return type. For a + # return type no lowering can even category-match (arrays, maps, + # structs, datetimes, ...), that Cast either never resolves -- and + # because the options ride along as children of TranspiledPythonUDF, + # an unresolvable Cast fails the WHOLE query at CheckAnalysis instead + # of falling back -- or diverges from the interpreted converter, which + # nulls type-mismatched results. Restrict transpilation to return + # types some lowering can match (the strict per-variant body-category + # check lives in ``_transpile_from_ast``); everything else falls back + # to interpreted Python. + if isinstance(returnType, str): + from pyspark.sql.types import _parse_datatype_string + + returnType = _parse_datatype_string(returnType) + if not isinstance(returnType, (NumericType, StringType, BooleanType, BinaryType)): + return ( + [], + [ + f"return type {returnType.simpleString()} is not supported by " + "the transpiler (no lowered expression can be cast to it " + "under ANSI rules); falling back to interpreted Python" + ], + [], + [], + ) + # A functools.wraps-style decorator makes ``inspect.getsource`` return + # the WRAPPED function's source (getsource follows ``__wrapped__``), + # while the UDF actually executes the wrapper. Transpiling would + # silently reproduce the wrong behavior, so refuse and fall back. + # ``_call_impl`` first: a ``staticmethod`` / ``classmethod`` exposes a + # ``__wrapped__`` of its own, so asking the descriptor refuses every one of + # them for a wraps decorator that is not there. + if ( + getattr(func, "__wrapped__", None) is not None + or getattr(_call_impl(_call_dunder(func)), "__wrapped__", None) is not None + ): + return ( + [], + [ + "decorated callables (functools.wraps) are not supported: " + "the visible source is the wrapped function's, not the " + "wrapper's, so transpilation would change behavior" + ], + [], + [], + ) + # Not ``ast``: that name would shadow the module for this whole function. + src, ast_info = _get_src_ast_from_func(func) + if ast_info is None: + return ([], ["Error getting ast for function, cannot transpile"], [], []) + # Get the lambda body and parameters + function_ast, extraction_error = _get_function_from_ast(ast_info, _held_code(func)) + if function_ast is None: + return ([], [extraction_error], [], []) + # Default, variadic (``*args`` / ``**kwargs``), keyword-only, and + # positional-only parameters can't be represented by the positional + # ``_udf_param_N`` placeholder scheme: a call site may omit a + # defaulted argument, leaving the placeholder referencing a position + # the call never bound, and ``_get_parameter_list`` only reads + # ``args``. Fall back to interpreted Python rather than emit an + # invalid plan. + fn_args = function_ast.args + if ( + fn_args.defaults + or any(d is not None for d in fn_args.kw_defaults) + or fn_args.kwonlyargs + or fn_args.vararg is not None + or fn_args.kwarg is not None + or fn_args.posonlyargs + ): + return ( + [], + [ + "functions with default, variadic, keyword-only, or " + "positional-only arguments are not supported by the transpiler" + ], + [], + [], + ) + params = _get_parameter_list(function_ast) + # Drop a receiver that is already bound, so what is left is what the call + # site supplies. Decided by HOW ``func`` dispatches, not by the parameter's + # name: a bound ``__call__(this, x)`` or ``@classmethod f(cls, x)`` has a + # receiver not named ``self``, while a plain ``def f(self, x)`` supplies its + # ``self`` at the call site. Going by the name misnumbered every + # ``_udf_param_N`` -- a two-column call on ``__call__(this, x)`` read column b + # for ``x`` and returned a value where Python raises TypeError. Asking + # ``inspect.signature`` is both weaker (a ``__signature__`` off by exactly one + # is undetectable) and worse behaved (it runs user code). + # + # For a callable instance, two things can consume a leading parameter, and + # they compose: what the descriptor prepends when Python looks ``__call__`` + # up -- the instance for a plain function, the class for a ``classmethod``, + # nothing for a ``staticmethod`` or for an already-bound method, whose + # ``__get__`` returns itself -- and whatever the callable already has bound. + # Each count below is checked against what Python returns for that shape. + if inspect.isfunction(func): + spoken_for = 0 + elif inspect.ismethod(func): + spoken_for = 1 + else: + call_entry = _call_dunder(func) + call_target = _call_impl(call_entry) + if not (inspect.isfunction(call_target) or inspect.ismethod(call_target)): + # A slot wrapper, property, partial, or other descriptor: what it + # prepends is not knowable from here. + return ( + [], + [ + f"a {type(call_entry).__name__} as __call__ does not say which " + "parameters the call site supplies, so the placeholder " + "positions cannot be assigned" + ], + [], + [], + ) + spoken_for = int( + inspect.isfunction(call_entry) or isinstance(call_entry, classmethod) + ) + int(inspect.ismethod(call_target)) + if spoken_for > 1 or spoken_for > len(params): + # Two receivers at once -- a ``classmethod`` over an already-bound method + # prepends the class ON TOP of the method's own ``__self__`` -- or one with + # no parameter to hold it. Python raises for whatever the call site passes, + # so there is nothing correct to lower. + return ([], ["callable leaves no parameter for the call site to bind"], [], []) + # Caller-facing params: callers match user-supplied kwargs against this, + # and the receiver is not named at the call site. Everything downstream + # indexes off THIS list, so the placeholder numbering needs no offset. + public_params = params[spoken_for:] + transpiled: list[Column] = [] + input_categories: list[list[str]] = [] + errors = [] + # One transpiled option per (backend x input-type variant). Untyped + # params are tried as both numeric and string so the JVM can pick the + # option matching the actual column types (or fall back if none match). + combos = _param_category_combos(function_ast, public_params) + # Maybe multiple transpilers (think CUDA, etc.). + transpilers = _get_transpilers(session) + for transpiler in transpilers: + for combo in combos: + try: + transpiled_column = transpiler._transpile_from_ast( + src, ast_info, function_ast, public_params, returnType, combo + ) + if transpiled_column is not None: + transpiled.append(transpiled_column) + input_categories.append( + [combo.get(i, "numeric") for i in range(len(public_params))] + ) + except Exception as e: + errors.append(str(e)) + return (transpiled, errors, public_params, input_categories) + except Exception as e: + # Don't re-raise: an inability to transpile must never break a + # working UDF. The caller treats an empty ``transpiled`` list as a + # silent fall-back to interpreted Python. + return ([], [str(e)], [], []) diff --git a/python/pyspark/sql/tvf.py b/python/pyspark/sql/tvf.py index 90874578a3562..b32f62ab61e5d 100644 --- a/python/pyspark/sql/tvf.py +++ b/python/pyspark/sql/tvf.py @@ -362,7 +362,7 @@ def json_tuple(self, input: Column, *fields: Column) -> DataFrame: |value1|value2| +------+------+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq if len(fields) == 0: raise PySparkValueError( @@ -507,7 +507,7 @@ def stack(self, n: Column, *fields: Column) -> DataFrame: | 3|NULL| +----+----+ """ - from pyspark.sql.classic.column import _to_seq, _to_java_column + from pyspark.sql.classic.column import _to_java_column, _to_seq sc = self._sparkSession.sparkContext return DataFrame( @@ -724,9 +724,10 @@ def _fn(self, functionName: str, *args: Column) -> DataFrame: def _test() -> None: - import os import doctest + import os import sys + import pyspark.sql.tvf os.chdir(os.environ["SPARK_HOME"]) diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index a57d045e47b77..e777325d5c356 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -15,23 +15,22 @@ # limitations under the License. # -import os -import sys -import decimal -import time -import math -import datetime +import base64 import calendar +import ctypes +import datetime +import decimal import json +import math +import os import re -import base64 +import sys +import time from array import array -import ctypes from collections.abc import Iterable from functools import reduce from typing import ( - cast, - overload, + TYPE_CHECKING, Any, Callable, ClassVar, @@ -39,39 +38,42 @@ Iterator, List, Optional, - Union, Tuple, Type, TypeVar, - TYPE_CHECKING, + Union, + cast, + overload, ) -from pyspark.util import is_remote_only, JVM_INT_MAX -from pyspark.serializers import CloudPickleSerializer -from pyspark.sql.utils import ( - get_active_spark_context, - escape_meta_characters, - IllegalArgumentException, - StringConcat, -) -from pyspark.sql.variant_utils import VariantUtils from pyspark.errors import ( + PySparkAttributeError, + PySparkIndexError, + PySparkKeyError, PySparkNotImplementedError, + PySparkRuntimeError, PySparkTypeError, PySparkValueError, - PySparkIndexError, - PySparkRuntimeError, - PySparkAttributeError, - PySparkKeyError, ) +from pyspark.serializers import CloudPickleSerializer from pyspark.sql.geo_utils import ( - GeographicSpatialReferenceSystemMapper as _GeographicSRSMapper, CartesianSpatialReferenceSystemMapper as _CartesianSRSMapper, ) +from pyspark.sql.geo_utils import ( + GeographicSpatialReferenceSystemMapper as _GeographicSRSMapper, +) +from pyspark.sql.utils import ( + IllegalArgumentException, + StringConcat, + escape_meta_characters, + get_active_spark_context, +) +from pyspark.sql.variant_utils import VariantUtils +from pyspark.util import JVM_INT_MAX, is_remote_only if TYPE_CHECKING: import numpy as np - from py4j.java_gateway import GatewayClient, JavaGateway, JavaClass + from py4j.java_gateway import GatewayClient, JavaClass, JavaGateway T = TypeVar("T") U = TypeVar("U") @@ -3849,6 +3851,7 @@ def convert(self, obj: "np.ndarray", gateway_client: "GatewayClient") -> "JavaGa def _test() -> None: import doctest + from pyspark.sql import SparkSession globs = globals() diff --git a/python/pyspark/sql/udf.py b/python/pyspark/sql/udf.py index e586c560deeb4..2618ecb7d4b04 100644 --- a/python/pyspark/sql/udf.py +++ b/python/pyspark/sql/udf.py @@ -22,11 +22,12 @@ import inspect import sys import warnings -from typing import Callable, Any, TYPE_CHECKING, Optional, cast, Union +from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast - -from pyspark.util import PythonEvalType +from pyspark.errors import PySparkNotImplementedError, PySparkRuntimeError, PySparkTypeError from pyspark.sql.column import Column +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version from pyspark.sql.types import ( DataType, StringType, @@ -34,14 +35,13 @@ _parse_datatype_string, ) from pyspark.sql.utils import get_active_spark_context -from pyspark.sql.pandas.types import to_arrow_type -from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version -from pyspark.errors import PySparkTypeError, PySparkNotImplementedError, PySparkRuntimeError +from pyspark.util import PythonEvalType if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.core.context import SparkContext - from pyspark.sql._typing import DataTypeOrString, ColumnOrName, UserDefinedFunctionLike + from pyspark.sql._typing import ColumnOrName, DataTypeOrString, UserDefinedFunctionLike from pyspark.sql.session import SparkSession __all__ = ["UDFRegistration"] @@ -76,11 +76,17 @@ def _create_udf( evalType: int, name: Optional[str] = None, deterministic: bool = True, + bufferSchema: Optional[StructType] = None, ) -> "UserDefinedFunctionLike": """Create a regular(non-Arrow-optimized) Python UDF.""" # Set the name of the UserDefinedFunction object to be the name of function f udf_obj = UserDefinedFunction( - f, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic + f, + returnType=returnType, + name=name, + evalType=evalType, + deterministic=deterministic, + bufferSchema=bufferSchema, ) return udf_obj._wrapped() @@ -165,6 +171,7 @@ def __init__( name: Optional[str] = None, evalType: int = PythonEvalType.SQL_BATCHED_UDF, deterministic: bool = True, + bufferSchema: Optional[StructType] = None, ): if not callable(func): raise PySparkTypeError( @@ -206,6 +213,104 @@ def __init__( ) self.evalType = evalType self.deterministic = deterministic + # Schema of the intermediate aggregation buffer, set only for an incremental Python + # aggregator (see :class:`pyspark.sql.aggregator.Aggregator`); ``None`` otherwise. It is a + # first-class field so it survives reconstruction paths such as ``_wrapped()``, + # ``asNondeterministic()`` and ``spark.udf.register``, and is threaded to the JVM in + # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage aggregation. + self.bufferSchema = bufferSchema + # Extract Python UDF details if transpilation is enabled. + self.transpiled: list = [] + self._transpiled_param_names: list[str] = [] + # Per-option input-type categories ("numeric"/"string" per public param), + # parallel to ``self.transpiled``; the JVM picks the option matching the + # actual column types or falls back to interpreted Python. + self._transpiled_input_categories: list = [] + # When we have a transpiled rewrite, ``__call__`` resolves any + # user-supplied kwargs against this positional parameter list so + # the JVM-side ``_udf_param_N`` substitution sees the inputs in + # the right order. Empty list when transpilation didn't happen. + from pyspark.sql import SparkSession + + session = SparkSession._instantiatedSession + + # A nondeterministic UDF must not be transpiled: replacing it with a plain + # Catalyst expression would let the optimizer fold/reorder/duplicate it, + # discarding the nondeterminism barrier. (asNondeterministic() also clears + # any options set here, for the udf(f).asNondeterministic() ordering.) + # Conf values are compared case-insensitively: `SET conf=True` stores + # the literal "True", which would otherwise silently disable + # transpilation (or mis-trigger the ANSI warning below). + # + # Each conf read is a JVM roundtrip, so keep the default construction + # path cheap: the experimental gate is only read for deterministic + # batched UDFs (the only shape we transpile), and the ANSI conf is only + # read once the gate is known to be on. When ``default`` is given it is + # passed through to ``RuntimeConfig.get`` so construction never depends + # on the JVM having the (experimental) conf registered -- e.g. a newer + # Python client against an older driver. No default is passed for + # ``spark.sql.ansi.enabled``: its registered default is dynamic + # (environment-driven) and must be respected when the key is unset. + def _conf_is_true(key: str, default: Optional[str] = None) -> bool: + if session is None: + return False + if default is None: + value = session.conf.get(key) + else: + value = session.conf.get(key, default) + return value is not None and value.lower() == "true" + + try: + transpile_enabled = ( + deterministic + and evalType == PythonEvalType.SQL_BATCHED_UDF + and _conf_is_true("spark.sql.experimental.optimizer.transpilePyUDFs", "false") + ) + # Transpilation only attempts to reproduce ANSI-mode Spark SQL + # semantics (no silent integer overflow, divide-by-zero raises, + # etc.). Running it against non-ANSI Spark would balloon the test + # matrix we'd have to maintain to verify Python-vs-SQL equivalence, + # so we gate on ANSI here and warn the user instead of trying to + # transpile in a mode we don't claim to support yet. + if transpile_enabled and not _conf_is_true("spark.sql.ansi.enabled"): + warnings.warn( + "Python UDF transpilation " + "(spark.sql.experimental.optimizer.transpilePyUDFs) is only " + "supported when ANSI mode is enabled " + "(spark.sql.ansi.enabled=true). Skipping transpilation for " + f"{func} -- enable ANSI mode or set transpilePyUDFs=false to " + "silence this warning.", + RuntimeWarning, + ) + transpile_enabled = False + if transpile_enabled and session: + # Import only if needed, also avoid circular import loops. + from pyspark.sql.transpile import _transpile_func + + # ``self.returnType`` parses (and caches) the declared return + # type; the transpiler needs the parsed form to decide whether + # the final Cast to it can resolve at all. The parse is reused + # later by ``_create_judf``, so this adds no extra JVM work. + ( + self.transpiled, + errors, + self._transpiled_param_names, + self._transpiled_input_categories, + ) = _transpile_func(session, func, self.returnType) + if not self.transpiled: + detail = f": {errors}" if errors else "" + warnings.warn(f"Unable to transpile UDF {func}{detail}") + except Exception as e: + # An inability to transpile must never break a working UDF -- fall + # back to interpreted Python execution and surface the failure as a + # warning so users can opt to investigate without losing their + # query. The conf reads above are included: a session whose JVM + # cannot answer them should degrade to "no transpilation", not + # break UDF definition. + warnings.warn(f"Exception transpiling UDF {func}: {e}") + self.transpiled = [] + self._transpiled_param_names = [] + self._transpiled_input_categories = [] @staticmethod def _check_return_type(returnType: DataType, evalType: int) -> None: @@ -411,8 +516,11 @@ def _judf(self) -> "JavaObject": self._judf_placeholder = self._create_judf(self.func) return self._judf_placeholder - def _create_judf(self, func: Callable[..., Any]) -> "JavaObject": + def _create_judf( + self, func: Callable[..., Any], include_transpiled: bool = True + ) -> "JavaObject": from pyspark.sql import SparkSession + from pyspark.sql.classic.column import _to_java_column_opt spark = SparkSession._getActiveSessionOrCreate() sc = spark.sparkContext @@ -420,8 +528,26 @@ def _create_judf(self, func: Callable[..., Any]) -> "JavaObject": wrapped_func = _wrap_function(sc, func, self.returnType) jdt = spark._jsparkSession.parseDataType(self.returnType.json()) assert sc._jvm is not None + transpiled = self.transpiled if include_transpiled else [] + input_categories = self._transpiled_input_categories if include_transpiled else [] + # Incremental Python aggregators additionally carry the intermediate buffer schema, which + # the JVM needs at planning time to build the two-stage aggregation (see PythonAggregate). + # Everyone else passes ``None`` here, which Py4J maps to the JVM ``null`` the ``bufferType`` + # parameter already defaults to. + jbuf = ( + spark._jsparkSession.parseDataType(self.bufferSchema.json()) + if self.bufferSchema is not None + else None + ) judf = getattr(sc._jvm, "org.apache.spark.sql.execution.python.UserDefinedPythonFunction")( - self._name, wrapped_func, jdt, self.evalType, self.deterministic + self._name, + wrapped_func, + jdt, + self.evalType, + self.deterministic, + map(_to_java_column_opt, transpiled), + input_categories, + jbuf, ) return judf @@ -430,6 +556,32 @@ def __call__(self, *args: "ColumnOrName", **kwargs: "ColumnOrName") -> Column: sc = get_active_spark_context() + # Transpilation rewrites the UDF into a Catalyst expression that + # references its inputs positionally via ``_udf_param_N`` (see + # ``UserDefinedPythonFunction.builder.resolveUDFParams``). If the + # caller used kwargs, the JVM-side substitution would otherwise + # splice ``NamedArgumentExpression`` wrappers into the rewritten + # tree (and into nested function calls like ``isnotnull``, which + # rejects named arguments). Resolve kwargs to positional here + # using the parameter list captured at transpilation time so the + # rewritten expression sees plain column refs in declared order. + if kwargs and self.transpiled and self._transpiled_param_names: + params = self._transpiled_param_names + ordered: list = list(args) + remaining_kwargs = dict(kwargs) + for pname in params[len(args) :]: + if pname in remaining_kwargs: + ordered.append(remaining_kwargs.pop(pname)) + else: + # Caller didn't supply this param positionally or by + # name -- bail out of the rewrite and let the regular + # JVM-side path raise a user-facing error. + break + else: + if not remaining_kwargs: + args = tuple(ordered) + kwargs = {} + assert sc._jvm is not None jcols = [_to_java_column(arg) for arg in args] + [ sc._jvm.PythonSQLUtils.namedArgumentExpression(key, _to_java_column(value)) @@ -440,6 +592,19 @@ def __call__(self, *args: "ColumnOrName", **kwargs: "ColumnOrName") -> Column: memory_profiler_enabled = sc._conf.get("spark.python.profile.memory", "false") == "true" if profiler_enabled or memory_profiler_enabled: + # Profiling is not supported for incremental Python aggregators. Their ``self.func`` is + # an ``Aggregator`` object, not a plain function: the profiler wrappers below would + # replace it with a function the worker cannot drive (it has no ``zero``/``reduce``/ + # ``bufferSchema``), and the memory profiler's ``inspect.getsourcelines(f.__code__)`` + # fails on the driver because an ``Aggregator`` instance has no ``__code__``. + if self.evalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF: + warnings.warn( + "Profiling incremental Python aggregators is not supported.", + UserWarning, + ) + judf = self._judf + return Column(judf.apply(_to_seq(sc, jcols))) + # Disable profiling Pandas UDFs with iterators as input/output. if self.evalType in [ PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, @@ -477,7 +642,11 @@ def func(*args: Any, **kwargs: Any) -> Any: return profiler.profile(f, *args, **kwargs) func.__signature__ = inspect.signature(f) # type: ignore[attr-defined] - judf = self._create_judf(func) + # Profiling requires the Python function to actually execute, + # and the transpiled path never runs it (it also produces a + # TranspiledPythonUDF, which has no resultId for the profiler + # to key on). Build this call's judf without transpiled options. + judf = self._create_judf(func, include_transpiled=False) jUDFExpr = judf.builderWithColumns(_to_seq(sc, jcols)) jPythonUDF = judf.fromUDFExpr(jUDFExpr) id = jUDFExpr.resultId().id() @@ -499,7 +668,9 @@ def func(*args: Any, **kwargs: Any) -> Any: ) func.__signature__ = inspect.signature(f) # type: ignore[attr-defined] - judf = self._create_judf(func) + # See the profiler branch above: no transpiled options while + # profiling, since only the interpreted path runs the function. + judf = self._create_judf(func, include_transpiled=False) jUDFExpr = judf.builderWithColumns(_to_seq(sc, jcols)) jPythonUDF = judf.fromUDFExpr(jUDFExpr) id = jUDFExpr.resultId().id() @@ -541,6 +712,7 @@ def wrapper(*args: "ColumnOrName", **kwargs: "ColumnOrName") -> Column: wrapper.returnType = self.returnType # type: ignore[attr-defined] wrapper.evalType = self.evalType # type: ignore[attr-defined] wrapper.deterministic = self.deterministic # type: ignore[attr-defined] + wrapper.bufferSchema = self.bufferSchema # type: ignore[attr-defined] wrapper.asNondeterministic = functools.wraps( # type: ignore[attr-defined] self.asNondeterministic )(lambda: self.asNondeterministic()._wrapped()) @@ -557,6 +729,14 @@ def asNondeterministic(self) -> "UserDefinedFunction": # with 'deterministic' updated. See SPARK-23233. self._judf_placeholder = None self.deterministic = False + # A transpiled rewrite replaces the (now nondeterministic) Python UDF + # with a plain Catalyst expression, which the optimizer is free to + # fold, reorder, or duplicate -- discarding the nondeterminism barrier + # the caller just asked for. Drop any transpiled options so a + # nondeterministic UDF always runs as interpreted Python. + self.transpiled = [] + self._transpiled_param_names = [] + self._transpiled_input_categories = [] return self @@ -696,6 +876,7 @@ def register( PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, ]: raise PySparkTypeError( errorClass="INVALID_UDF_EVAL_TYPE", @@ -704,7 +885,8 @@ def register( "SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, " "SQL_SCALAR_PANDAS_ITER_UDF, SQL_SCALAR_ARROW_ITER_UDF, " "SQL_GROUPED_AGG_PANDAS_UDF, SQL_GROUPED_AGG_ARROW_UDF, " - "SQL_GROUPED_AGG_PANDAS_ITER_UDF or SQL_GROUPED_AGG_ARROW_ITER_UDF" + "SQL_GROUPED_AGG_PANDAS_ITER_UDF, SQL_GROUPED_AGG_ARROW_ITER_UDF " + "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF" }, ) source_udf = _create_udf( @@ -713,6 +895,8 @@ def register( name=name, evalType=f.evalType, deterministic=f.deterministic, + # Preserve the incremental aggregator's buffer schema (None for other UDFs). + bufferSchema=getattr(f, "bufferSchema", None), ) register_udf = source_udf._unwrapped # type: ignore[attr-defined] return_udf = register_udf @@ -810,8 +994,9 @@ def registerJavaUDAF(self, name: str, javaClassName: str) -> None: def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.udf + from pyspark.sql import SparkSession from pyspark.testing.utils import have_pandas, have_pyarrow globs = pyspark.sql.udf.__dict__.copy() diff --git a/python/pyspark/sql/udtf.py b/python/pyspark/sql/udtf.py index 74400c3d2beae..e8dc25dc907ec 100644 --- a/python/pyspark/sql/udtf.py +++ b/python/pyspark/sql/udtf.py @@ -18,26 +18,27 @@ User-defined table function related classes and functions """ -import pickle -from dataclasses import dataclass, field import inspect +import pickle import sys import warnings -from typing import Any, Type, TYPE_CHECKING, Optional, Sequence, Union +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Optional, Sequence, Type, Union from pyspark.errors import ( PySparkAttributeError, + PySparkImportError, PySparkPicklingError, PySparkTypeError, - PySparkImportError, ) -from pyspark.util import PythonEvalType from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version from pyspark.sql.types import DataType, StructType, _parse_datatype_string from pyspark.sql.udf import _wrap_function +from pyspark.util import PythonEvalType if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql._typing import TVFArgumentOrName from pyspark.sql.dataframe import DataFrame from pyspark.sql.session import SparkSession @@ -419,9 +420,8 @@ def _create_judtf(self, func: Type) -> "JavaObject": return judtf def __call__(self, *args: "TVFArgumentOrName", **kwargs: "TVFArgumentOrName") -> "DataFrame": - from pyspark.sql.classic.column import _to_java_column, _to_seq - from pyspark.sql import DataFrame, SparkSession + from pyspark.sql.classic.column import _to_java_column, _to_seq from pyspark.sql.table_arg import TableArg spark = SparkSession._getActiveSessionOrCreate() @@ -556,8 +556,9 @@ def register( def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.sql.udf + from pyspark.sql import SparkSession globs = pyspark.sql.udtf.__dict__.copy() spark = SparkSession.builder.master("local[4]").appName("sql.udtf tests").getOrCreate() diff --git a/python/pyspark/sql/utils.py b/python/pyspark/sql/utils.py index 493eb6ee0abc3..876f2937758dd 100644 --- a/python/pyspark/sql/utils.py +++ b/python/pyspark/sql/utils.py @@ -14,41 +14,41 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from enum import Enum -import inspect import functools +import inspect import os +from enum import Enum from typing import ( + TYPE_CHECKING, Any, Callable, Dict, - Optional, List, - overload, + Optional, Sequence, - TYPE_CHECKING, - cast, TypeVar, Union, + cast, + overload, ) # For backward compatibility. from pyspark.errors import ( # noqa: F401 AnalysisException, - ParseException, IllegalArgumentException, - StreamingQueryException, - QueryExecutionException, - PythonException, - UnknownException, - SparkUpgradeException, + ParseException, PySparkImportError, PySparkNotImplementedError, PySparkRuntimeError, + PythonException, + QueryExecutionException, + SparkUpgradeException, + StreamingQueryException, + UnknownException, ) -from pyspark.util import is_remote_only, JVM_INT_MAX from pyspark.errors.exceptions.captured import CapturedException # noqa: F401 from pyspark.find_spark_home import _find_spark_home +from pyspark.util import JVM_INT_MAX, is_remote_only if TYPE_CHECKING: from py4j.java_collections import JavaArray @@ -58,10 +58,11 @@ JavaObject, JVMView, ) + from pyspark import SparkContext - from pyspark.sql.session import SparkSession - from pyspark.sql.dataframe import DataFrame from pyspark.pandas._typing import IndexOpsLike, SeriesOrIndex + from pyspark.sql.dataframe import DataFrame + from pyspark.sql.session import SparkSession FuncT = TypeVar("FuncT", bound=Callable[..., Any]) @@ -96,8 +97,8 @@ def to_scala_map(jvm: "JVMView", dic: Dict) -> "JavaObject": def require_test_compiled() -> None: """Raise Exception if test classes are not compiled""" - import os import glob + import os test_class_path = os.path.join(_find_spark_home(), "sql", "core", "target", "*", "test-classes") paths = glob.glob(test_class_path) @@ -462,8 +463,8 @@ def pyspark_column_op( Wrapper function for column_op to get proper Column class. """ from pyspark.pandas.base import column_op - from pyspark.sql.column import Column from pyspark.pandas.data_type_ops.base import _is_extension_dtypes + from pyspark.sql.column import Column result = column_op(getattr(Column, func_name))(left, right) # It works as expected on extension dtype, so we don't need to call `fillna` for this case. diff --git a/python/pyspark/sql/variant_utils.py b/python/pyspark/sql/variant_utils.py index 80b7efeaca100..41bed740d3da5 100644 --- a/python/pyspark/sql/variant_utils.py +++ b/python/pyspark/sql/variant_utils.py @@ -16,15 +16,16 @@ # import base64 -import decimal import datetime +import decimal import json import struct from array import array from typing import Any, Callable, Dict, List, NamedTuple, Tuple -from pyspark.errors import PySparkValueError from zoneinfo import ZoneInfo +from pyspark.errors import PySparkValueError + class VariantUtils: """ diff --git a/python/pyspark/sql/window.py b/python/pyspark/sql/window.py index 0b2687a038866..94bf719e19f32 100644 --- a/python/pyspark/sql/window.py +++ b/python/pyspark/sql/window.py @@ -18,16 +18,17 @@ # mypy: disable-error-code="empty-body" import sys -from typing import Sequence, TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Sequence, Union from pyspark.sql.utils import dispatch_window_method from pyspark.util import ( - JVM_LONG_MIN, JVM_LONG_MAX, + JVM_LONG_MIN, ) if TYPE_CHECKING: from py4j.java_gateway import JavaObject + from pyspark.sql._typing import ColumnOrName __all__ = ["Window", "WindowSpec"] diff --git a/python/pyspark/sql/worker/analyze_udtf.py b/python/pyspark/sql/worker/analyze_udtf.py index 8a9f4fc3666da..bd5639313d0fa 100644 --- a/python/pyspark/sql/worker/analyze_udtf.py +++ b/python/pyspark/sql/worker/analyze_udtf.py @@ -17,10 +17,11 @@ import inspect from textwrap import dedent -from typing import Any, Dict, List, IO, Protocol, Tuple +from typing import IO, Any, Dict, List, Protocol, Tuple from pyspark.errors import PySparkRuntimeError, PySparkValueError -from pyspark.logger.worker_io import capture_outputs, context_provider as default_context_provider +from pyspark.logger.worker_io import capture_outputs +from pyspark.logger.worker_io import context_provider as default_context_provider from pyspark.serializers import ( read_bool, read_int, @@ -28,13 +29,13 @@ write_with_length, ) from pyspark.sql.functions import OrderingColumn, PartitioningColumn, SelectedColumn -from pyspark.sql.types import _parse_datatype_json_string, StructType +from pyspark.sql.types import StructType, _parse_datatype_json_string from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult from pyspark.sql.worker.utils import worker_run from pyspark.worker_util import ( get_sock_file_to_executor, - read_command, pickleSer, + read_command, utf8_deserializer, ) diff --git a/python/pyspark/sql/worker/create_data_source.py b/python/pyspark/sql/worker/create_data_source.py index d9c1f799fa46f..00d3701ca9930 100644 --- a/python/pyspark/sql/worker/create_data_source.py +++ b/python/pyspark/sql/worker/create_data_source.py @@ -25,13 +25,13 @@ write_int, write_with_length, ) -from pyspark.sql.datasource import DataSource, CaseInsensitiveDict -from pyspark.sql.types import _parse_datatype_json_string, StructType +from pyspark.sql.datasource import CaseInsensitiveDict, DataSource +from pyspark.sql.types import StructType, _parse_datatype_json_string from pyspark.sql.worker.utils import worker_run from pyspark.worker_util import ( get_sock_file_to_executor, - read_command, pickleSer, + read_command, utf8_deserializer, ) diff --git a/python/pyspark/sql/worker/data_source_pushdown_filters.py b/python/pyspark/sql/worker/data_source_pushdown_filters.py index a649936422996..4be11e13a8bac 100644 --- a/python/pyspark/sql/worker/data_source_pushdown_filters.py +++ b/python/pyspark/sql/worker/data_source_pushdown_filters.py @@ -24,7 +24,7 @@ from pyspark.errors import PySparkAssertionError, PySparkValueError from pyspark.errors.exceptions.base import PySparkNotImplementedError from pyspark.logger.worker_io import capture_outputs -from pyspark.serializers import UTF8Deserializer, read_int, read_bool, write_int +from pyspark.serializers import UTF8Deserializer, read_bool, read_int, write_int from pyspark.sql.datasource import ( DataSource, DataSourceReader, @@ -45,7 +45,11 @@ ) from pyspark.sql.types import StructType, VariantVal, _parse_datatype_json_string from pyspark.sql.worker.plan_data_source_read import write_read_func_and_partitions -from pyspark.sql.worker.utils import worker_run +from pyspark.sql.worker.utils import ( + check_pushdown_not_disabled, + is_method_overridden, + worker_run, +) from pyspark.worker_util import ( get_sock_file_to_executor, pickleSer, @@ -113,24 +117,37 @@ def deserializeFilter(jsonDict: dict) -> Filter: def _main(infile: IO, outfile: IO) -> None: """ - Main method for planning a data source read with filter pushdown. + Main method for planning a data source read with filter and limit pushdown. - This process is invoked from the `UserDefinedPythonDataSourceReadRunner.runInPython` - method in the optimizer rule `PlanPythonDataSourceScan` in JVM. This process is responsible - for creating a `DataSourceReader` object, applying filter pushdown, and sending the - information needed back to the JVM. + This process is invoked from the `UserDefinedPythonDataSourceFilterPushdownRunner.runInPython` + method, which `PythonScanBuilder` runs on the JVM while pushing filters and a limit into the + scan. This process is responsible for creating a `DataSourceReader` object, applying filter + and limit pushdown, and sending the information needed back to the JVM. The infile and outfile are connected to the JVM via a socket. The JVM sends the following information to this process via the socket: - a `DataSource` instance representing the data source - a `StructType` instance representing the output schema of the data source - a list of filters to be pushed down + - the limit to be pushed down, or -1 if there is none - configuration values This process then creates a `DataSourceReader` instance by calling the `reader` method on the `DataSource` instance. It applies the filters by calling the `pushFilters` method - on the reader and determines which filters are supported. The indices of the supported - filters are sent back to the JVM, along with the list of partitions and the read function. + on the reader and determines which filters are supported. + + When a limit is sent, it is pushed down by calling the `pushLimit` method on the reader + after `pushFilters`, and whether the reader accepted it is sent back to the JVM. The JVM + replays the same filters when pushing down a limit, so that the reader reaches the same + state as it did during filter pushdown before `pushLimit` is called on it. + + The values sent back to the JVM, in order, are: + - a flag (1 or 0) for whether a read function and partitions follow. They do not when the + filter-pushdown pass defers planning (limit pushdown is enabled, so a limit pass may follow) + or when a pushed limit was rejected; the JVM then plans the read itself, once it knows + whether a limit is pushed. When the flag is 1, the read function and partitions follow. + - the indices of the supported filters. + - whether the limit was pushed down (1 or 0). """ # Receive the data source instance. data_source = read_command(pickleSer, infile) @@ -173,9 +190,12 @@ def _main(infile: IO, outfile: IO) -> None: filter_dicts = json.loads(json_str) filters = [FilterRef(deserializeFilter(f)) for f in filter_dicts] - # Push down the filters and get the indices of the unsupported filters. + # Push down the filters and get the indices of the unsupported filters. `pushFilters` is + # not called when there is nothing to push, so that a reader planning a limit-only scan + # does not observe a spurious empty pushFilters call. unsupported_filters = set( - FilterRef(f) for f in reader.pushFilters([ref.filter for ref in filters]) + FilterRef(f) + for f in (reader.pushFilters([ref.filter for ref in filters]) if filters else []) ) supported_filter_indices = [] for i, filter in enumerate(filters): @@ -195,30 +215,88 @@ def _main(infile: IO, outfile: IO) -> None: }, ) + # Receive the limit to push down. -1 means there is no limit. + limit = read_int(infile) + # Receive the max arrow batch size. max_arrow_batch_size = read_int(infile) assert max_arrow_batch_size > 0, ( "The maximum arrow batch size should be greater than 0, but got " f"'{max_arrow_batch_size}'" ) + enable_filter_pushdown = read_bool(infile) + enable_limit_pushdown = read_bool(infile) binary_as_bytes = read_bool(infile) + # Whether this worker should plan (and send) the read function and partitions. The JVM + # sets this to False for the filter-pushdown pass when limit pushdown is enabled, so that + # `partitions()`/`read()` are not planned before a possible `pushLimit`. A later limit + # pass (or build-time planning on the JVM) then plans once all pushdowns are known, + # matching the public contract that `pushLimit` runs before `partitions()`/`read()`. + plan_read_info = read_bool(infile) - # Return the read function and partitions. Doing this in the same worker - # as filter pushdown helps reduce the number of Python worker calls. - write_read_func_and_partitions( - outfile, - reader=reader, - data_source=data_source, - schema=schema, - max_arrow_batch_size=max_arrow_batch_size, - binary_as_bytes=binary_as_bytes, - ) + # Do not silently ignore a pushdown method that the reader implements while the + # corresponding pushdown is disabled. This worker also runs when only one of filter or + # limit pushdown is enabled -- e.g. a limit-only scan caches the read info here, so + # `plan_data_source_read` never runs and cannot validate the other method -- so both are + # validated here as well as in `plan_data_source_read`. + check_pushdown_not_disabled(reader, enable_filter_pushdown, enable_limit_pushdown) + + # Push down the limit, if any. This must happen after pushFilters, matching the + # operator order that DSv2 uses on the JVM side. + # + # Whether the read function and partitions follow. Start from what the JVM requested + # (`plan_read_info`), then turn it off if the reader rejects a pushed limit (below). + is_limit_pushed = False + send_read_info = plan_read_info + # Only call `pushLimit` if the reader actually overrides it: the inherited default is a + # no-op that returns False, so calling it would needlessly drive the rejection path (and + # an extra reader construction on the JVM) for readers that do not implement limit + # pushdown. + if limit >= 0 and is_method_overridden(reader, "pushLimit"): + is_limit_pushed = reader.pushLimit(limit) + if not isinstance(is_limit_pushed, bool): + raise PySparkValueError( + errorClass="DATA_SOURCE_INVALID_RETURN_TYPE", + messageParameters={ + "type": type(is_limit_pushed).__name__, + "name": type(reader).__name__ + ".pushLimit", + "supported_types": "bool", + }, + ) + if not is_limit_pushed: + # The reader considered the limit and rejected it, possibly mutating itself. Do + # not plan `partitions()`/`read()` from this reader; send no read info so the JVM + # plans the filters-only read from a fresh reader instead. That path also + # re-validates that the replayed `pushFilters` makes the same decision, keeping + # the rejected reader state out of the scan. + send_read_info = False + + # Send whether the read function and partitions follow, before them, so the JVM knows + # whether to read them (and so a planning exception is still surfaced first). None follow + # when the filter-pushdown pass defers planning (limit pushdown enabled) or when a pushed + # limit was rejected; the JVM then plans once it knows whether a limit is pushed, never on + # a reader whose plan would be discarded. + write_int(int(send_read_info), outfile) + if send_read_info: + # Planning here in the same worker as filter/limit pushdown avoids an extra Python + # worker call. + write_read_func_and_partitions( + outfile, + reader=reader, + data_source=data_source, + schema=schema, + max_arrow_batch_size=max_arrow_batch_size, + binary_as_bytes=binary_as_bytes, + ) # Return the supported filter indices. write_int(len(supported_filter_indices), outfile) for index in supported_filter_indices: write_int(index, outfile) + # Return whether the limit was pushed down, as 1 or 0. + write_int(int(is_limit_pushed), outfile) + def main(infile: IO, outfile: IO) -> None: worker_run(_main, infile, outfile) diff --git a/python/pyspark/sql/worker/plan_data_source_read.py b/python/pyspark/sql/worker/plan_data_source_read.py index 1de63bd74166d..e2e66bdc0bd02 100644 --- a/python/pyspark/sql/worker/plan_data_source_read.py +++ b/python/pyspark/sql/worker/plan_data_source_read.py @@ -16,9 +16,10 @@ # import functools +from itertools import chain, islice +from typing import IO, Iterable, Iterator, List, Tuple, Union + import pyarrow as pa -from itertools import islice, chain -from typing import IO, List, Iterator, Iterable, Tuple, Union from pyspark.errors import PySparkAssertionError, PySparkRuntimeError from pyspark.logger.worker_io import capture_outputs @@ -38,15 +39,15 @@ from pyspark.sql.datasource_internal import _streamReader from pyspark.sql.pandas.types import to_arrow_schema from pyspark.sql.types import ( - _parse_datatype_json_string, BinaryType, StructType, + _parse_datatype_json_string, ) -from pyspark.sql.worker.utils import worker_run +from pyspark.sql.worker.utils import check_pushdown_not_disabled, worker_run from pyspark.worker_util import ( get_sock_file_to_executor, - read_command, pickleSer, + read_command, utf8_deserializer, ) @@ -283,9 +284,15 @@ def _main(infile: IO, outfile: IO) -> None: for creating a `DataSourceReader` object and send the information needed back to the JVM. The infile and outfile are connected to the JVM via a socket. The JVM sends the following - information to this process via the socket: + information to this process via the socket, in this order (the protocol is positional): - a `DataSource` instance representing the data source + - a `StructType` instance representing the input schema from the child plan - a `StructType` instance representing the output schema of the data source + - the max Arrow batch size (int) + - whether filter pushdown is enabled (bool) + - whether limit pushdown is enabled (bool) + - whether this is a streaming read (bool) + - whether binary values are returned as `bytes` (bool) This process then creates a `DataSourceReader` instance by calling the `reader` method on the `DataSource` instance. Then it calls the `partitions()` method of the reader and @@ -339,6 +346,7 @@ def _main(infile: IO, outfile: IO) -> None: f"The maximum arrow batch size should be greater than 0, but got '{max_arrow_batch_size}'" ) enable_pushdown = read_bool(infile) + enable_limit_pushdown = read_bool(infile) is_streaming = read_bool(infile) binary_as_bytes = read_bool(infile) @@ -360,19 +368,9 @@ def _main(infile: IO, outfile: IO) -> None: "actual": f"'{type(reader).__name__}'", }, ) - is_pushdown_implemented = ( - getattr(reader.pushFilters, "__func__", None) is not DataSourceReader.pushFilters - ) - if is_pushdown_implemented and not enable_pushdown: - # Do not silently ignore pushFilters when pushdown is disabled. - # Raise an error to ask the user to enable pushdown. - raise PySparkAssertionError( - errorClass="DATA_SOURCE_PUSHDOWN_DISABLED", - messageParameters={ - "type": type(reader).__name__, - "conf": "spark.sql.python.filterPushdown.enabled", - }, - ) + # Do not silently ignore a pushdown method that the reader implements while the + # corresponding pushdown is disabled. Raise an error to ask the user to enable it. + check_pushdown_not_disabled(reader, enable_pushdown, enable_limit_pushdown) # Send the read function and partitions to the JVM. write_read_func_and_partitions( diff --git a/python/pyspark/sql/worker/python_streaming_sink_runner.py b/python/pyspark/sql/worker/python_streaming_sink_runner.py index 2a4ea0b95b287..d999b5580273c 100644 --- a/python/pyspark/sql/worker/python_streaming_sink_runner.py +++ b/python/pyspark/sql/worker/python_streaming_sink_runner.py @@ -27,14 +27,14 @@ ) from pyspark.sql.datasource import DataSource, WriterCommitMessage from pyspark.sql.types import ( - _parse_datatype_json_string, StructType, + _parse_datatype_json_string, ) from pyspark.sql.worker.utils import worker_run from pyspark.worker_util import ( get_sock_file_to_executor, - read_command, pickleSer, + read_command, utf8_deserializer, ) diff --git a/python/pyspark/sql/worker/utils.py b/python/pyspark/sql/worker/utils.py index 12bdb25e62529..e22745cc4fc72 100644 --- a/python/pyspark/sql/worker/utils.py +++ b/python/pyspark/sql/worker/utils.py @@ -17,36 +17,36 @@ import os import sys -from typing import Callable, IO, Optional +from typing import IO, Any, Callable, Optional from pyspark.accumulators import ( + SpecialAccumulatorIds, _accumulatorRegistry, _deserialize_accumulator, - SpecialAccumulatorIds, +) +from pyspark.serializers import ( + SpecialLengths, + read_int, + write_int, ) from pyspark.sql.profiler import ( ProfileResultsParam, ProfileResultsParamV2, - WorkerPerfProfiler, WorkerMemoryProfiler, -) -from pyspark.serializers import ( - read_int, - write_int, - SpecialLengths, + WorkerPerfProfiler, ) from pyspark.util import ( - start_faulthandler_periodic_traceback, handle_worker_exception, + start_faulthandler_periodic_traceback, with_faulthandler, ) from pyspark.worker_util import ( + Conf, check_python_version, send_accumulator_updates, + setup_broadcasts, setup_memory_limits, setup_spark_files, - setup_broadcasts, - Conf, ) @@ -56,6 +56,45 @@ def profiler(self) -> Optional[str]: return self.get("spark.sql.pyspark.dataSource.profiler", None) +def is_method_overridden(reader: Any, name: str) -> bool: + """ + Whether `reader` overrides the `DataSourceReader` method `name`, rather than inheriting the + default implementation. Used to detect pushdown methods that a reader implements while the + corresponding pushdown configuration is disabled, so that they are not silently ignored. + """ + from pyspark.sql.datasource import DataSourceReader + + return getattr(getattr(reader, name), "__func__", None) is not getattr(DataSourceReader, name) + + +def check_pushdown_not_disabled( + reader: Any, enable_filter_pushdown: bool, enable_limit_pushdown: bool +) -> None: + """ + Raise `DATA_SOURCE_PUSHDOWN_DISABLED` if `reader` implements a pushdown method while the + corresponding pushdown configuration is disabled, so that the method is not silently ignored. + + This is shared by both planning workers: `plan_data_source_read` runs it for a plain read, + and `data_source_pushdown_filters` runs it too, because a filter- or limit-only scan caches + the read info in that worker and `plan_data_source_read` never runs for such a scan. + """ + from pyspark.errors import PySparkAssertionError + + for method, conf, enabled in ( + ("pushFilters", "spark.sql.python.filterPushdown.enabled", enable_filter_pushdown), + ("pushLimit", "spark.sql.python.limitPushdown.enabled", enable_limit_pushdown), + ): + if not enabled and is_method_overridden(reader, method): + raise PySparkAssertionError( + errorClass="DATA_SOURCE_PUSHDOWN_DISABLED", + messageParameters={ + "type": type(reader).__name__, + "method": method, + "conf": conf, + }, + ) + + @with_faulthandler def worker_run(main: Callable, infile: IO, outfile: IO) -> None: try: diff --git a/python/pyspark/sql/worker/write_into_data_source.py b/python/pyspark/sql/worker/write_into_data_source.py index ca97090f004f6..4f104e4364723 100644 --- a/python/pyspark/sql/worker/write_into_data_source.py +++ b/python/pyspark/sql/worker/write_into_data_source.py @@ -17,7 +17,6 @@ import inspect from typing import IO, Iterator, Union -from pyspark.sql.conversion import ArrowTableToRowsConversion from pyspark.errors import PySparkAssertionError, PySparkRuntimeError, PySparkTypeError from pyspark.logger.worker_io import capture_outputs from pyspark.serializers import ( @@ -25,26 +24,27 @@ read_int, ) from pyspark.sql import Row +from pyspark.sql.conversion import ArrowTableToRowsConversion from pyspark.sql.datasource import ( + CaseInsensitiveDict, DataSource, - DataSourceWriter, DataSourceArrowWriter, - WriterCommitMessage, - CaseInsensitiveDict, - DataSourceStreamWriter, DataSourceStreamArrowWriter, + DataSourceStreamWriter, + DataSourceWriter, + WriterCommitMessage, ) from pyspark.sql.types import ( - _parse_datatype_json_string, - StructType, BinaryType, + StructType, _create_row, + _parse_datatype_json_string, ) from pyspark.sql.worker.utils import worker_run from pyspark.worker_util import ( get_sock_file_to_executor, - read_command, pickleSer, + read_command, utf8_deserializer, ) diff --git a/python/pyspark/streaming/context.py b/python/pyspark/streaming/context.py index bb0a659a6b33e..7b2203324c66a 100644 --- a/python/pyspark/streaming/context.py +++ b/python/pyspark/streaming/context.py @@ -14,20 +14,19 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import warnings from typing import Any, Callable, List, Optional, TypeVar -from py4j.java_gateway import java_import, is_instance_of, JavaObject +from py4j.java_gateway import JavaObject, is_instance_of, java_import from pyspark import RDD, SparkConf -from pyspark.serializers import NoOpSerializer, UTF8Deserializer, CloudPickleSerializer from pyspark.core.context import SparkContext +from pyspark.serializers import CloudPickleSerializer, NoOpSerializer, UTF8Deserializer from pyspark.storagelevel import StorageLevel from pyspark.streaming.dstream import DStream from pyspark.streaming.listener import StreamingListener from pyspark.streaming.util import TransformFunction, TransformFunctionSerializer -import warnings - __all__ = ["StreamingContext"] T = TypeVar("T") diff --git a/python/pyspark/streaming/dstream.py b/python/pyspark/streaming/dstream.py index bfc885bb5779d..493fd20ca8968 100644 --- a/python/pyspark/streaming/dstream.py +++ b/python/pyspark/streaming/dstream.py @@ -17,9 +17,10 @@ import operator import time -from itertools import chain from datetime import datetime +from itertools import chain from typing import ( + TYPE_CHECKING, Any, Callable, Generic, @@ -30,18 +31,17 @@ Tuple, TypeVar, Union, - TYPE_CHECKING, cast, overload, ) -from py4j.protocol import Py4JJavaError from py4j.java_gateway import JavaObject +from py4j.protocol import Py4JJavaError -from pyspark.storagelevel import StorageLevel -from pyspark.streaming.util import rddToFileName, TransformFunction -from pyspark.core.rdd import portable_hash, RDD +from pyspark.core.rdd import RDD, portable_hash from pyspark.resultiterable import ResultIterable +from pyspark.storagelevel import StorageLevel +from pyspark.streaming.util import TransformFunction, rddToFileName if TYPE_CHECKING: from pyspark.serializers import Serializer diff --git a/python/pyspark/streaming/kinesis.py b/python/pyspark/streaming/kinesis.py index 12a50f69b0b9f..a78a33611d2f0 100644 --- a/python/pyspark/streaming/kinesis.py +++ b/python/pyspark/streaming/kinesis.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import overload, Callable, Optional, TypeVar, Union +from typing import Callable, Optional, TypeVar, Union, overload from pyspark.serializers import NoOpSerializer from pyspark.storagelevel import StorageLevel diff --git a/python/pyspark/streaming/tests/test_kinesis.py b/python/pyspark/streaming/tests/test_kinesis.py index 8deaed4e1bae0..55c0f9306fd0f 100644 --- a/python/pyspark/streaming/tests/test_kinesis.py +++ b/python/pyspark/streaming/tests/test_kinesis.py @@ -18,11 +18,11 @@ import unittest from pyspark import StorageLevel -from pyspark.streaming.kinesis import KinesisUtils, InitialPositionInStream, MetricsLevel +from pyspark.streaming.kinesis import InitialPositionInStream, KinesisUtils, MetricsLevel from pyspark.testing.streamingutils import ( - should_test_kinesis, - kinesis_requirement_message, PySparkStreamingTestCase, + kinesis_requirement_message, + should_test_kinesis, ) diff --git a/python/pyspark/streaming/util.py b/python/pyspark/streaming/util.py index 7d69e69e13817..4ee97b6d63875 100644 --- a/python/pyspark/streaming/util.py +++ b/python/pyspark/streaming/util.py @@ -15,14 +15,14 @@ # limitations under the License. # +import sys import time -from datetime import datetime import traceback -import sys +from datetime import datetime from py4j.java_gateway import is_instance_of -from pyspark import SparkContext, RDD +from pyspark import RDD, SparkContext class TransformFunction: diff --git a/python/pyspark/taskcontext.py b/python/pyspark/taskcontext.py index 796e3b5ef5f67..6bb2102b47a8e 100644 --- a/python/pyspark/taskcontext.py +++ b/python/pyspark/taskcontext.py @@ -14,12 +14,13 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Any, ClassVar, Type, TypeVar, Dict, List, Optional, Union, cast +from decimal import Decimal +from typing import Any, ClassVar, Dict, List, Optional, Type, TypeVar, Union, cast -from pyspark.util import local_connect_and_auth -from pyspark.serializers import read_int, write_int, write_with_length, UTF8Deserializer from pyspark.errors import PySparkRuntimeError from pyspark.resource import ResourceInformation +from pyspark.serializers import UTF8Deserializer, read_int, write_int, write_with_length +from pyspark.util import local_connect_and_auth T = TypeVar("T", bound="TaskContext") @@ -128,7 +129,7 @@ class TaskContext: _taskAttemptId: Optional[int] = None _localProperties: Optional[Dict[str, str]] = None _cpus: Optional[int] = None - _cpuAmount: Optional[float] = None + _cpuAmount: Optional[Decimal] = None _resources: Optional[Dict[str, "ResourceInformation"]] = None def __new__(cls: Type["TaskContext"], **kwargs: Any) -> "TaskContext": @@ -192,7 +193,7 @@ def from_json(cls: Type[T], json: dict) -> T: attemptNumber=json["attemptNumber"], taskAttemptId=json["taskAttemptId"], cpus=json["cpus"], - cpuAmount=float(json["cpuAmount"]) if "cpuAmount" in json else None, + cpuAmount=Decimal(json["cpuAmount"]) if "cpuAmount" in json else None, resources={ k: ResourceInformation(v["name"], v["addresses"]) for k, v in json["resources"].items() @@ -276,7 +277,7 @@ def cpus(self) -> int: """ return cast(int, self._cpus) - def cpuAmount(self) -> float: + def cpuAmount(self) -> Decimal: """ The exact amount of CPUs allocated to the task. This can be fractional when ``spark.task.cpus`` or the task resource profile requests a fractional amount @@ -286,14 +287,14 @@ def cpuAmount(self) -> float: Returns ------- - float - the exact, possibly fractional, amount of CPUs. + decimal.Decimal + the exact, possibly fractional amount of CPUs. See Also -------- TaskContext.cpus """ - return cast(float, self._cpuAmount) + return cast(Decimal, self._cpuAmount) def resources(self) -> Dict[str, "ResourceInformation"]: """ @@ -547,6 +548,7 @@ def __init__(self, address: str) -> None: def _test() -> None: import doctest import sys + from pyspark.sql import SparkSession globs = globals().copy() diff --git a/python/pyspark/testing/connectutils.py b/python/pyspark/testing/connectutils.py index bc7125e821d01..75fe3e26a3d71 100644 --- a/python/pyspark/testing/connectutils.py +++ b/python/pyspark/testing/connectutils.py @@ -14,33 +14,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import contextlib +import functools +import os import shutil import tempfile -import os -import functools import unittest import uuid -import contextlib from typing import Callable, Optional from pyspark import Row, SparkConf from pyspark.loose_version import LooseVersion -from pyspark.util import is_remote_only +from pyspark.sql.session import SparkSession as PySparkSession +from pyspark.testing.sqlutils import SQLTestUtils from pyspark.testing.utils import ( - have_pandas, + PySparkBaseTestCase, + PySparkErrorTestUtils, connect_requirement_message, + have_pandas, should_test_connect, - PySparkErrorTestUtils, ) -from pyspark.testing.utils import PySparkBaseTestCase -from pyspark.testing.sqlutils import SQLTestUtils -from pyspark.sql.session import SparkSession as PySparkSession +from pyspark.util import is_remote_only if should_test_connect: + import pyspark.sql.connect.proto as pb2 from pyspark.sql.connect.dataframe import DataFrame - from pyspark.sql.connect.plan import Read, Range, SQL, LogicalPlan + from pyspark.sql.connect.plan import SQL, LogicalPlan, Range, Read from pyspark.sql.connect.session import SparkSession - import pyspark.sql.connect.proto as pb2 class MockRemoteSession: diff --git a/python/pyspark/testing/goldenutils.py b/python/pyspark/testing/goldenutils.py index d0a191e9fe9ad..f61cc21a0f84f 100644 --- a/python/pyspark/testing/goldenutils.py +++ b/python/pyspark/testing/goldenutils.py @@ -15,10 +15,10 @@ # limitations under the License. # -from typing import Any, Callable, List, Optional import inspect import os import time +from typing import Any, Callable, List, Optional, Union try: import numpy as np @@ -311,12 +311,56 @@ def repr_type(t: Any) -> str: # "halffloat" -> "float16", "float" -> "float32", "double" -> "float64" return _ARROW_FLOAT_ALIASES.get(s, s) + @staticmethod + def _scalar_str(scalar: Any) -> str: + """ + Render one PyArrow scalar for a golden cell via PyArrow's own ``str(scalar)``. + + Some values that Arrow stores are unrenderable via ``str``: + - a temporal value past Python's ``datetime`` range (e.g. date32 past year 9999) + raises ``OverflowError`` -> the marker ``temporal overflow``; + - an unsafe ``binary``->``string`` cast relabels bytes without UTF-8 validation, + so non-UTF-8 bytes raise ``UnicodeDecodeError`` -> the raw bytes, so the golden + still tracks what Arrow stored for them; + - a nanosecond ``time64`` holding INT64_MIN collides with pandas' NaT sentinel and + raises ``ValueError`` -> the marker ``NaT collision``. + An unexpected error (non-temporal overflow, non-string decode error, or a + ValueError from any other value) propagates. + """ + try: + return str(scalar).replace("\x00", "\\0") + except OverflowError: + # A valid Arrow temporal value can exceed Python's datetime range (e.g. a + # date32 past year 9999); record ``temporal overflow`` rather than fail. + # A non-temporal overflow is unexpected, so re-raise. + if pa.types.is_temporal(scalar.type): + return "temporal overflow" + raise + except UnicodeDecodeError: + # An unsafe binary->string cast relabels bytes as a string without UTF-8 + # validation, so str() cannot decode non-UTF-8 bytes; render the raw bytes + # so the golden still tracks what Arrow stored for them. A UnicodeDecodeError + # on a non-string type is unexpected, so re-raise. + if pa.types.is_string(scalar.type) or pa.types.is_large_string(scalar.type): + return repr(scalar.as_buffer().to_pybytes()) + raise + except ValueError: + # A nanosecond time64 renders through pandas (Python's ``time`` is microsecond + # resolution), and pandas reads INT64_MIN as its NaT sentinel, so it refuses + # that one value. Any other ValueError is unexpected, so re-raise. + if pa.types.is_time(scalar.type) and scalar.value == pd.NaT.value: + return "NaT collision" + raise + @classmethod - def repr_arrow_value(cls, value: Any, max_len: int = 32) -> str: + def repr_arrow_value( + cls, value: Union["pa.Array", "pa.ChunkedArray"], max_len: int = 32 + ) -> str: """ Format a PyArrow Array/ChunkedArray for golden file. - Each element uses str(scalar) from PyArrow's own scalar formatting. + Each element is rendered by ``_scalar_str`` (PyArrow's scalar formatting, with + fallbacks for values that are valid in Arrow but unrenderable via ``str``). Parameters ---------- @@ -331,17 +375,92 @@ def repr_arrow_value(cls, value: Any, max_len: int = 32) -> str: "[val1, val2, None]@arrow_type" """ # Escape NULL bytes so the value can be safely stored in CSV files. - elements = [str(scalar).replace("\x00", "\\0") for scalar in value] + elements = [cls._scalar_str(scalar) for scalar in value] v_str = "[" + ", ".join(elements) + "]" if max_len > 0: v_str = v_str[:max_len] return f"{v_str}@{cls.repr_type(value.type)}" @classmethod - def repr_pandas_value(cls, value: Any, max_len: int = 32) -> str: + def _repr_arrow_columns( + cls, value: Union["pa.Table", "pa.RecordBatch"], max_len: int + ) -> "tuple[str, str]": + """Render a Table/RecordBatch as a "{name: [scalars], ...}" body and schema string.""" + columns = [] + for name, column in zip(value.schema.names, value.columns): + # Escape NULL bytes so the value can be safely stored in CSV files. + elements = [cls._scalar_str(scalar) for scalar in column] + columns.append(f"{name}: [" + ", ".join(elements) + "]") + v_str = "{" + ", ".join(columns) + "}" + if max_len > 0: + v_str = v_str[:max_len] + schema = ", ".join(f"{f.name}: {cls.repr_type(f.type)}" for f in value.schema) + return v_str, schema + + @classmethod + def repr_arrow_table_value(cls, value: "pa.Table", max_len: int = 32) -> str: + """ + Format a PyArrow Table for golden file. + + Renders each column with PyArrow's scalar formatting (as ``repr_arrow_value`` + does for an Array), keyed by column name, plus the Arrow schema. + + Returns + ------- + str + "{col: [val1, val2, None], ...}@Table[name: type, ...]" + """ + v_str, schema = cls._repr_arrow_columns(value, max_len) + return f"{v_str}@Table[{schema}]" + + @classmethod + def repr_arrow_record_batch_value(cls, value: "pa.RecordBatch", max_len: int = 32) -> str: + """ + Format a PyArrow RecordBatch for golden file. + + Same shape as ``repr_arrow_table_value`` (a RecordBatch is a single batch of + columns), keyed by column name, plus the Arrow schema. + + Returns + ------- + str + "{col: [val1, val2, None], ...}@RecordBatch[name: type, ...]" + """ + v_str, schema = cls._repr_arrow_columns(value, max_len) + return f"{v_str}@RecordBatch[{schema}]" + + @classmethod + def repr_arrow_schema_value(cls, value: "pa.Schema", max_len: int = 32) -> str: + """ + Format a PyArrow Schema for golden file. + + Renders each field as "name: type nullable=...". A Schema carries no data, so + nullability is included (unlike the Table/RecordBatch schema string): Spark reads + ``field.nullable`` to build its StructType, so a change there silently alters the + inferred Spark schema. + + Returns + ------- + str + "[name: type nullable=True, ...]@Schema" + """ + fields = [f"{f.name}: {cls.repr_type(f.type)} nullable={f.nullable}" for f in value] + v_str = "[" + ", ".join(fields) + "]" + if max_len > 0: + v_str = v_str[:max_len] + return f"{v_str}@Schema" + + @classmethod + def repr_pandas_value(cls, value: "pd.DataFrame", max_len: int = 32) -> str: """ Format a pandas DataFrame for golden file. + Renders each column with tolist() (as ``repr_pandas_series_value`` does for a + Series), keyed by column name, plus the schema. tolist() gives a stable + Python-native representation and avoids ``DataFrame.to_json``'s epoch date + serialization, which overflows on out-of-nanosecond-range dates (year 9999 + with the default date_as_object=True) and misreads non-ns units on pandas 2. + Parameters ---------- value : pd.DataFrame @@ -352,16 +471,16 @@ def repr_pandas_value(cls, value: Any, max_len: int = 32) -> str: Returns ------- str - "value@Dataframe[schema]" + "{col: [val1, val2], ...}@Dataframe[schema]" """ - v_str = value.to_json().replace("\n", " ") + v_str = str({name: col.tolist() for name, col in value.items()}).replace("\n", " ") if max_len > 0: v_str = v_str[:max_len] simple_schema = ", ".join([f"{t} {d.name}" for t, d in value.dtypes.items()]) return f"{v_str}@Dataframe[{simple_schema}]" @classmethod - def repr_numpy_value(cls, value: Any, max_len: int = 32) -> str: + def repr_numpy_value(cls, value: "np.ndarray", max_len: int = 32) -> str: """ Format a numpy ndarray for golden file. @@ -404,6 +523,9 @@ def repr_value(cls, value: Any, max_len: int = 32) -> str: based on the value's type. - PyArrow Array/ChunkedArray -> repr_arrow_value + - PyArrow Table -> repr_arrow_table_value + - PyArrow RecordBatch -> repr_arrow_record_batch_value + - PyArrow Schema -> repr_arrow_schema_value - pandas DataFrame -> repr_pandas_value - numpy ndarray -> repr_numpy_value - Everything else -> repr_python_value @@ -422,6 +544,12 @@ def repr_value(cls, value: Any, max_len: int = 32) -> str: """ if have_pyarrow and isinstance(value, (pa.Array, pa.ChunkedArray)): return cls.repr_arrow_value(value, max_len) + if have_pyarrow and isinstance(value, pa.Table): + return cls.repr_arrow_table_value(value, max_len) + if have_pyarrow and isinstance(value, pa.RecordBatch): + return cls.repr_arrow_record_batch_value(value, max_len) + if have_pyarrow and isinstance(value, pa.Schema): + return cls.repr_arrow_schema_value(value, max_len) if have_pandas and isinstance(value, pd.DataFrame): return cls.repr_pandas_value(value, max_len) @@ -433,7 +561,7 @@ def repr_value(cls, value: Any, max_len: int = 32) -> str: return cls.repr_python_value(value, max_len) @classmethod - def repr_pandas_series_value(cls, value: Any, max_len: int = 32) -> str: + def repr_pandas_series_value(cls, value: "pd.Series", max_len: int = 32) -> str: """ Format a pandas Series for golden file. diff --git a/python/pyspark/testing/mlutils.py b/python/pyspark/testing/mlutils.py index e26a4cc83ee52..a5218c9f42c51 100644 --- a/python/pyspark/testing/mlutils.py +++ b/python/pyspark/testing/mlutils.py @@ -19,10 +19,10 @@ from pyspark import keyword_only from pyspark.ml import Estimator, Model, Transformer, UnaryTransformer +from pyspark.ml.classification import ClassificationModel, Classifier from pyspark.ml.evaluation import Evaluator from pyspark.ml.param import Param, Params, TypeConverters from pyspark.ml.param.shared import HasMaxIter, HasRegParam -from pyspark.ml.classification import Classifier, ClassificationModel from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWritable from pyspark.ml.wrapper import _java2py from pyspark.sql import DataFrame, SparkSession @@ -233,8 +233,8 @@ def __init__(self): def _transform(self, dataset): # A dummy transform impl which always predict label 1 - from pyspark.sql.functions import array, lit from pyspark.ml.functions import array_to_vector + from pyspark.sql.functions import array, lit rawPredCol = self.getRawPredictionCol() if rawPredCol: diff --git a/python/pyspark/testing/pandasutils.py b/python/pyspark/testing/pandasutils.py index 8483bfd75965e..e779ef0fb3ee3 100644 --- a/python/pyspark/testing/pandasutils.py +++ b/python/pyspark/testing/pandasutils.py @@ -15,12 +15,12 @@ # limitations under the License. # +import decimal import functools import shutil import tempfile import warnings from contextlib import contextmanager -import decimal from typing import Any, Union try: @@ -47,15 +47,15 @@ except ImportError: pass -from pyspark.loose_version import LooseVersion import pyspark.pandas as ps +from pyspark.errors import PySparkAssertionError +from pyspark.loose_version import LooseVersion from pyspark.pandas.frame import DataFrame from pyspark.pandas.indexes import Index from pyspark.pandas.series import Series from pyspark.pandas.utils import SPARK_CONF_ARROW_ENABLED from pyspark.testing.sqlutils import ReusedSQLTestCase from pyspark.testing.utils import is_ansi_mode_test -from pyspark.errors import PySparkAssertionError def _assert_pandas_equal( diff --git a/python/pyspark/testing/sqlutils.py b/python/pyspark/testing/sqlutils.py index 1760ad5ffd79c..cb8658a7a4a91 100644 --- a/python/pyspark/testing/sqlutils.py +++ b/python/pyspark/testing/sqlutils.py @@ -22,13 +22,13 @@ import tempfile from contextlib import contextmanager +from pyspark.find_spark_home import _find_spark_home from pyspark.sql import SparkSession from pyspark.sql.types import Row from pyspark.testing.utils import ( - ReusedPySparkTestCase, PySparkErrorTestUtils, + ReusedPySparkTestCase, ) -from pyspark.find_spark_home import _find_spark_home SPARK_HOME = _find_spark_home() @@ -69,8 +69,8 @@ def get_sbt_runtime_classpath(project_relative_path, project_name_map): Returns: Comma-separated string of JAR paths, or None if SBT command fails """ - import subprocess import re + import subprocess sbt_project = project_name_map.get(project_relative_path) if not sbt_project: diff --git a/python/pyspark/testing/streamingutils.py b/python/pyspark/testing/streamingutils.py index 57d459c04673f..8f3e8803fcb4e 100644 --- a/python/pyspark/testing/streamingutils.py +++ b/python/pyspark/testing/streamingutils.py @@ -19,7 +19,7 @@ import time import unittest -from pyspark import SparkConf, SparkContext, RDD +from pyspark import RDD, SparkConf, SparkContext from pyspark.streaming import StreamingContext from pyspark.testing.sqlutils import search_jar diff --git a/python/pyspark/testing/tests/test_changed_files.py b/python/pyspark/testing/tests/test_changed_files.py new file mode 100644 index 0000000000000..2413bc57ea30b --- /dev/null +++ b/python/pyspark/testing/tests/test_changed_files.py @@ -0,0 +1,98 @@ +# +# 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. +# +import os +import tempfile +import unittest + +from pyspark.testing.utils import ( + PySparkBaseTestCase, + grimp_requirement_message, + have_grimp, +) + +# A test module and modules it does / does not import (transitively). These relationships are +# stable parts of the pyspark import graph. +_TEST_MODULE = "pyspark.sql.tests.test_functions" +_RELEVANT_FILE = "python/pyspark/sql/functions/builtin.py" # imported by _TEST_MODULE +_IRRELEVANT_FILE = "python/pyspark/ml/classification.py" # not imported by _TEST_MODULE + + +@unittest.skipIf(not have_grimp, grimp_requirement_message) +class ChangedFilesSelectionTests(unittest.TestCase): + """Tests for the "smart" test selection driven by PYSPARK_CHANGED_FILES. + + ``PySparkBaseTestCase.skip_if_changed_files_irrelevant`` skips a test class when none of the + changed files' modules are reachable from the test's own module in the pyspark import graph. + These exercise that logic directly against real pyspark modules with known relationships. + """ + + def setUp(self): + # The relevance check is memoized with functools.cache; clear it so each case starts fresh. + PySparkBaseTestCase._is_module_relevant_to_changed_files.cache_clear() + + def _is_relevant(self, module, files): + with tempfile.NamedTemporaryFile("w") as f: + f.write("\n".join(files)) + f.flush() + return PySparkBaseTestCase._is_module_relevant_to_changed_files(module, f.name) + + def test_relevance(self): + own_file = "python/" + _TEST_MODULE.replace(".", os.path.sep) + ".py" + cases = [ + ("imported module is relevant", [_RELEVANT_FILE], True), + ("unimported module is irrelevant", [_IRRELEVANT_FILE], False), + ("relevant among irrelevant is relevant", [_IRRELEVANT_FILE, _RELEVANT_FILE], True), + ("the test's own file is relevant via the self-module short circuit", [own_file], True), + ( + "non-pyspark files are conservatively relevant", + ["sql/core/src/main/scala/Foo.scala"], + True, + ), + ( + "package __init__ with no graph node is conservatively relevant", + ["python/pyspark/sql/__init__.py"], + True, + ), + ] + for desc, files, expected in cases: + with self.subTest(desc): + self.assertEqual(self._is_relevant(_TEST_MODULE, files), expected) + + def test_skip_if_changed_files_irrelevant(self): + class _Dummy(PySparkBaseTestCase): + pass + + _Dummy.__module__ = _TEST_MODULE + + # Irrelevant changes raise SkipTest. + with tempfile.NamedTemporaryFile("w") as f: + f.write(_IRRELEVANT_FILE) + f.flush() + with self.assertRaises(unittest.SkipTest): + _Dummy.skip_if_changed_files_irrelevant(f.name) + + # Relevant changes do not. + with tempfile.NamedTemporaryFile("w") as f: + f.write(_RELEVANT_FILE) + f.flush() + _Dummy.skip_if_changed_files_irrelevant(f.name) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/testing/utils.py b/python/pyspark/testing/utils.py index cb549af75f3ca..e3af3fe40c904 100644 --- a/python/pyspark/testing/utils.py +++ b/python/pyspark/testing/utils.py @@ -15,33 +15,33 @@ # limitations under the License. # +import difflib +import faulthandler +import functools import os +import signal import struct import sys import unittest -import difflib -import faulthandler -import functools from decimal import Decimal -from time import time, sleep -import signal +from itertools import zip_longest +from time import sleep, time from typing import ( Any, - Optional, - Union, + Callable, Dict, List, - Callable, + Optional, + Union, ) -from itertools import zip_longest from pyspark import SparkConf from pyspark.errors import PySparkAssertionError, PySparkException, PySparkTypeError from pyspark.errors.exceptions.base import QueryContextType -from pyspark.sql.dataframe import DataFrame from pyspark.sql import Row -from pyspark.sql.types import StructType, StructField, VariantVal +from pyspark.sql.dataframe import DataFrame from pyspark.sql.functions import col, when +from pyspark.sql.types import StructField, StructType, VariantVal __all__ = ["assertDataFrameEqual", "assertSchemaEqual"] @@ -85,6 +85,9 @@ def have_package(name: str) -> bool: have_flameprof = have_package("flameprof") flameprof_requirement_message = "" if have_flameprof else "No module named 'flameprof'" +have_grimp = have_package("grimp") +grimp_requirement_message = "" if have_grimp else "No module named 'grimp'" + have_jinja2 = have_package("jinja2") jinja2_requirement_message = "" if have_jinja2 else "No module named 'jinja2'" @@ -294,9 +297,58 @@ def __exit__(self, exc_type, exc_val, exc_tb): class PySparkBaseTestCase(unittest.TestCase): @classmethod def setUpClass(cls): + if have_grimp and (path := os.environ.get("PYSPARK_CHANGED_FILES")): + # PYSPARK_CHANGED_FILES should only be used when ONLY pyspark files are changed. + # If other files (JVM for example) are changed, do NOT set this. + cls.skip_if_changed_files_irrelevant(path) + if os.environ.get("PYSPARK_TEST_TIMEOUT"): faulthandler.register(signal.SIGTERM, file=sys.__stderr__, all_threads=True) + @classmethod + def skip_if_changed_files_irrelevant(cls, path: str) -> None: + module = cls.__module__ + if module == "__main__": + mod = sys.modules["__main__"] + if mod.__spec__ and mod.__spec__.name: + module = mod.__spec__.name + else: + return + + if not cls._is_module_relevant_to_changed_files(module, path): + raise unittest.SkipTest("Skipping test because changed files are irrelevant") + + @staticmethod + @functools.cache + def _is_module_relevant_to_changed_files(module: str, path: str) -> bool: + import grimp + + with open(path, "r") as f: + changed_files = f.read().strip().splitlines() + + if not all(f.startswith("python/pyspark/") and f.endswith(".py") for f in changed_files): + # We have a wrong list of files, just run the test. + return True + + changed_modules = [ + f.removeprefix("python/").rsplit(".", 1)[0].replace(os.path.sep, ".") + for f in changed_files + ] + + graph = grimp.build_graph("pyspark") + + for changed_module in changed_modules: + if changed_module == module: + return True + try: + if graph.chain_exists(module, changed_module): + return True + except Exception: + # Any exception, we just be conservative and run the test. + return True + + return False + @classmethod def tearDownClass(cls): if os.environ.get("PYSPARK_TEST_TIMEOUT"): @@ -686,6 +738,7 @@ def compare_datatypes_ignore_nullable(dt1: Any, dt2: Any): if TYPE_CHECKING: import pandas + import pyspark.pandas @@ -1230,8 +1283,9 @@ def record_diff(r1, r2): def _test() -> None: import doctest - from pyspark.sql import SparkSession + import pyspark.testing.utils + from pyspark.sql import SparkSession globs = pyspark.testing.utils.__dict__.copy() spark = SparkSession.builder.master("local[4]").appName("testing.utils tests").getOrCreate() diff --git a/python/pyspark/tests/test_broadcast.py b/python/pyspark/tests/test_broadcast.py index 191616bc526a9..c08b08b544934 100644 --- a/python/pyspark/tests/test_broadcast.py +++ b/python/pyspark/tests/test_broadcast.py @@ -17,16 +17,16 @@ import os import pickle import random -import time import tempfile +import time import unittest from py4j.protocol import Py4JJavaError -from pyspark import SparkConf, SparkContext, Broadcast +from pyspark import Broadcast, SparkConf, SparkContext from pyspark.java_gateway import launch_gateway from pyspark.serializers import ChunkedStream -from pyspark.sql import SparkSession, Row +from pyspark.sql import Row, SparkSession class BroadcastTest(unittest.TestCase): diff --git a/python/pyspark/tests/test_conf.py b/python/pyspark/tests/test_conf.py index 5e501a0f3379f..bf6cbebaebb26 100644 --- a/python/pyspark/tests/test_conf.py +++ b/python/pyspark/tests/test_conf.py @@ -17,7 +17,7 @@ import random import unittest -from pyspark import SparkContext, SparkConf +from pyspark import SparkConf, SparkContext class ConfTests(unittest.TestCase): diff --git a/python/pyspark/tests/test_context.py b/python/pyspark/tests/test_context.py index 66bb0ebdd0973..648d9ef20a2c4 100644 --- a/python/pyspark/tests/test_context.py +++ b/python/pyspark/tests/test_context.py @@ -23,9 +23,9 @@ import unittest from collections import namedtuple -from pyspark import SparkConf, SparkFiles, SparkContext +from pyspark import SparkConf, SparkContext, SparkFiles from pyspark.testing.sqlutils import SPARK_HOME -from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest +from pyspark.testing.utils import PySparkTestCase, QuietTest, ReusedPySparkTestCase class CheckpointTests(ReusedPySparkTestCase): @@ -321,6 +321,33 @@ def create_spark_context(): with SparkContext("local-cluster[3, 1, 1024]") as sc: sc.range(2).foreach(lambda _: create_spark_context()) + def test_cancel_all_jobs_reason_reaches_the_job_failure(self): + # SPARK-58616: the reason must survive the trip to the JVM and land in the cancelled + # job's error, instead of the generic "as part of cancellation of all jobs". + with SparkContext() as sc: + errors = [] + + def run_job(): + try: + sc.parallelize(range(4), 4).map(lambda x: time.sleep(60)).collect() + except Exception as e: + errors.append(str(e)) + + job = threading.Thread(target=run_job) + job.start() + # Wait for the job to reach the scheduler before cancelling it. + deadline = time.time() + 60 + while not sc.statusTracker().getActiveJobsIds(): + self.assertLess(time.time(), deadline, "job never reached the scheduler") + time.sleep(0.1) + + sc.cancelAllJobs(reason="because the test asked for it") + job.join(60) + + self.assertEqual(len(errors), 1) + self.assertIn("because the test asked for it", errors[0]) + self.assertNotIn("as part of cancellation of all jobs", errors[0]) + class ContextTestsWithResources(unittest.TestCase): def setUp(self): diff --git a/python/pyspark/tests/test_daemon.py b/python/pyspark/tests/test_daemon.py index f0975cce85b98..09bcf19c9cff1 100644 --- a/python/pyspark/tests/test_daemon.py +++ b/python/pyspark/tests/test_daemon.py @@ -24,7 +24,7 @@ class DaemonTests(unittest.TestCase): def connect(self, port): - from socket import socket, AF_INET, AF_INET6, SOCK_STREAM + from socket import AF_INET, AF_INET6, SOCK_STREAM, socket family, host = AF_INET, "127.0.0.1" if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true": @@ -37,8 +37,8 @@ def connect(self, port): return True def do_termination_test(self, terminator): - from subprocess import Popen, PIPE from errno import ECONNREFUSED + from subprocess import PIPE, Popen # start daemon daemon_path = os.path.join(os.path.dirname(__file__), "..", "daemon.py") diff --git a/python/pyspark/tests/test_install_spark.py b/python/pyspark/tests/test_install_spark.py index 385e5a6a77844..315f90c0bc7d0 100644 --- a/python/pyspark/tests/test_install_spark.py +++ b/python/pyspark/tests/test_install_spark.py @@ -23,14 +23,14 @@ import urllib.request from pyspark.install import ( - get_preferred_mirrors, - install_spark, - _extract_tar, DEFAULT_HADOOP, DEFAULT_HIVE, UNSUPPORTED_COMBINATIONS, - checked_versions, + _extract_tar, checked_package_name, + checked_versions, + get_preferred_mirrors, + install_spark, ) diff --git a/python/pyspark/tests/test_memory_profiler.py b/python/pyspark/tests/test_memory_profiler.py index 5d9e41f16f06e..c6b7cc3a25c0f 100644 --- a/python/pyspark/tests/test_memory_profiler.py +++ b/python/pyspark/tests/test_memory_profiler.py @@ -15,9 +15,9 @@ # limitations under the License. # +import inspect import os import sys -import inspect import tempfile import unittest import warnings diff --git a/python/pyspark/tests/test_pin_thread.py b/python/pyspark/tests/test_pin_thread.py index a194c3e6b0fdd..894bb49c93bae 100644 --- a/python/pyspark/tests/test_pin_thread.py +++ b/python/pyspark/tests/test_pin_thread.py @@ -15,11 +15,11 @@ # limitations under the License. # import os -import time import threading +import time import unittest -from pyspark import SparkContext, SparkConf, InheritableThread +from pyspark import InheritableThread, SparkConf, SparkContext class PinThreadTests(unittest.TestCase): diff --git a/python/pyspark/tests/test_profiler.py b/python/pyspark/tests/test_profiler.py index 26834dbed1457..33be9dbd588b0 100644 --- a/python/pyspark/tests/test_profiler.py +++ b/python/pyspark/tests/test_profiler.py @@ -21,12 +21,12 @@ import unittest from io import StringIO -from pyspark import SparkConf, SparkContext, BasicProfiler +from pyspark import BasicProfiler, SparkConf, SparkContext +from pyspark.errors import PySparkRuntimeError, PythonException from pyspark.memory_profiler_ext import has_memory_profiler from pyspark.sql import SparkSession from pyspark.sql.functions import udf -from pyspark.errors import PythonException, PySparkRuntimeError -from pyspark.testing.utils import PySparkTestCase, PySparkErrorTestUtils +from pyspark.testing.utils import PySparkErrorTestUtils, PySparkTestCase class ProfilerTests(PySparkTestCase): diff --git a/python/pyspark/tests/test_rdd.py b/python/pyspark/tests/test_rdd.py index 7f277813338ea..d545b158831a0 100644 --- a/python/pyspark/tests/test_rdd.py +++ b/python/pyspark/tests/test_rdd.py @@ -14,30 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from datetime import datetime, timedelta import hashlib import os import random import tempfile import time import unittest +from datetime import datetime, timedelta from glob import glob from py4j.protocol import Py4JJavaError -from pyspark import shuffle, RDD +from pyspark import RDD, shuffle from pyspark.resource import ExecutorResourceRequests, ResourceProfileBuilder, TaskResourceRequests from pyspark.serializers import ( - CloudPickleSerializer, BatchedSerializer, + CloudPickleSerializer, CPickleSerializer, MarshalSerializer, - UTF8Deserializer, NoOpSerializer, + UTF8Deserializer, ) from pyspark.sql import SparkSession -from pyspark.testing.utils import ReusedPySparkTestCase, QuietTest, have_numpy, have_pandas from pyspark.testing.sqlutils import SPARK_HOME +from pyspark.testing.utils import QuietTest, ReusedPySparkTestCase, have_numpy, have_pandas global_func = lambda: "Hi" # noqa: E731 diff --git a/python/pyspark/tests/test_rddsampler.py b/python/pyspark/tests/test_rddsampler.py index f80372def2da3..b382410e623c5 100644 --- a/python/pyspark/tests/test_rddsampler.py +++ b/python/pyspark/tests/test_rddsampler.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from pyspark.testing.utils import ReusedPySparkTestCase from pyspark.rddsampler import RDDSampler, RDDStratifiedSampler +from pyspark.testing.utils import ReusedPySparkTestCase class RDDSamplerTests(ReusedPySparkTestCase): diff --git a/python/pyspark/tests/test_serializers.py b/python/pyspark/tests/test_serializers.py index f96ca9ab64be2..243ecc89fbb43 100644 --- a/python/pyspark/tests/test_serializers.py +++ b/python/pyspark/tests/test_serializers.py @@ -20,32 +20,32 @@ from pyspark import serializers from pyspark.serializers import ( - CloudPickleSerializer, - CompressedSerializer, AutoBatchedSerializer, BatchedSerializer, - AutoSerializer, - NoOpSerializer, - PairDeserializer, - FlattenedValuesSerializer, CartesianDeserializer, + CloudPickleSerializer, + CompressedSerializer, CPickleSerializer, - UTF8Deserializer, + FlattenedValuesSerializer, MarshalSerializer, + NoOpSerializer, + PairDeserializer, + UTF8Deserializer, ) from pyspark.testing.utils import ( - PySparkTestCase, - read_int, - write_int, ByteArrayOutput, + PySparkTestCase, have_numpy, have_scipy, + read_int, + write_int, ) class SerializationTestCase(unittest.TestCase): def test_namedtuple(self): from collections import namedtuple + from pyspark.cloudpickle import dumps, loads P = namedtuple("P", "x y") @@ -150,7 +150,6 @@ def test_hash_serializer(self): hash(UTF8Deserializer()) hash(CPickleSerializer()) hash(MarshalSerializer()) - hash(AutoSerializer()) hash(BatchedSerializer(CPickleSerializer())) hash(AutoBatchedSerializer(MarshalSerializer())) hash(PairDeserializer(NoOpSerializer(), UTF8Deserializer())) diff --git a/python/pyspark/tests/test_shuffle.py b/python/pyspark/tests/test_shuffle.py index 7edcf21a2e568..7d2b025c835bc 100644 --- a/python/pyspark/tests/test_shuffle.py +++ b/python/pyspark/tests/test_shuffle.py @@ -14,21 +14,21 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import os import random import unittest from tempfile import TemporaryDirectory -import os from py4j.protocol import Py4JJavaError -from pyspark import shuffle, CPickleSerializer, SparkConf, SparkContext +from pyspark import CPickleSerializer, SparkConf, SparkContext, shuffle from pyspark.shuffle import ( Aggregator, + ExternalGroupBy, ExternalMerger, ExternalSorter, - SimpleAggregator, Merger, - ExternalGroupBy, + SimpleAggregator, ) diff --git a/python/pyspark/tests/test_stage_sched.py b/python/pyspark/tests/test_stage_sched.py index c0e1e1539d791..1c0a8c787740d 100644 --- a/python/pyspark/tests/test_stage_sched.py +++ b/python/pyspark/tests/test_stage_sched.py @@ -15,12 +15,12 @@ # limitations under the License. # +import json import os +import shutil import tempfile -import unittest import time -import shutil -import json +import unittest from pyspark import SparkConf, SparkContext from pyspark.resource.profile import ResourceProfileBuilder diff --git a/python/pyspark/tests/test_taskcontext.py b/python/pyspark/tests/test_taskcontext.py index 2f39ff53bc9e4..59fa900c4ae01 100644 --- a/python/pyspark/tests/test_taskcontext.py +++ b/python/pyspark/tests/test_taskcontext.py @@ -22,8 +22,9 @@ import tempfile import time import unittest +from decimal import Decimal -from pyspark import SparkConf, SparkContext, TaskContext, BarrierTaskContext +from pyspark import BarrierTaskContext, SparkConf, SparkContext, TaskContext from pyspark.testing.sqlutils import SPARK_HOME from pyspark.testing.utils import PySparkTestCase, eventually @@ -355,7 +356,7 @@ def test_cpu_amount(self): """SPARK-58192: the exact cpu amount is available.""" rdd = self.sc.parallelize(range(10)) cpu_amount = rdd.map(lambda x: TaskContext.get().cpuAmount()).take(1)[0] - self.assertEqual(cpu_amount, 2.0) + self.assertEqual(cpu_amount, Decimal("2")) def test_resources(self): """Test that multiple resources are all available (SPARK-54929).""" @@ -379,7 +380,7 @@ class TaskContextTestsWithFractionalCpus(unittest.TestCase): def setUp(self): class_name = self.__class__.__name__ conf = SparkConf().set("spark.test.home", SPARK_HOME) - conf = conf.set("spark.task.cpus", "0.5") + conf = conf.set("spark.task.cpus", "0.123456789") self.sc = SparkContext("local-cluster[1,2,1024]", class_name, conf=conf) def test_fractional_cpu_amount(self): @@ -388,7 +389,7 @@ def test_fractional_cpu_amount(self): cpu_amount, cpus = rdd.map( lambda x: (TaskContext.get().cpuAmount(), TaskContext.get().cpus()) ).take(1)[0] - self.assertEqual(cpu_amount, 0.5) + self.assertEqual(cpu_amount, Decimal("0.123456789")) self.assertEqual(cpus, 1) def test_omp_num_threads_follows_task_cpus(self): diff --git a/python/pyspark/tests/test_util.py b/python/pyspark/tests/test_util.py index 686e946773b05..d387997a2e158 100644 --- a/python/pyspark/tests/test_util.py +++ b/python/pyspark/tests/test_util.py @@ -23,10 +23,10 @@ from py4j.protocol import Py4JJavaError from pyspark import keyword_only -from pyspark.util import _parse_memory, disable_gc +from pyspark.find_spark_home import _find_spark_home from pyspark.loose_version import LooseVersion from pyspark.testing.utils import PySparkTestCase, eventually, timeout -from pyspark.find_spark_home import _find_spark_home +from pyspark.util import _parse_memory, disable_gc class KeywordOnlyTests(unittest.TestCase): @@ -173,6 +173,7 @@ class HandleWorkerExceptionTests(unittest.TestCase): def run_handle_worker_exception(self, hide_traceback=None): import io + from pyspark.util import handle_worker_exception try: diff --git a/python/pyspark/tests/test_worker.py b/python/pyspark/tests/test_worker.py index 14a40dc233894..1fb25b7d19868 100644 --- a/python/pyspark/tests/test_worker.py +++ b/python/pyspark/tests/test_worker.py @@ -31,7 +31,7 @@ from py4j.protocol import Py4JJavaError from pyspark import SparkConf, SparkContext -from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest, eventually +from pyspark.testing.utils import PySparkTestCase, QuietTest, ReusedPySparkTestCase, eventually class WorkerTests(ReusedPySparkTestCase): diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.csv new file mode 100644 index 0000000000000..7403a6618187d --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.csv @@ -0,0 +1,104 @@ +test case pandas series arrow array +int8:standard [0, 1, -1, 127, -128]@Series[int8] [0, 1, -1, 127, -128]@int8 +int8:empty []@Series[int8] []@int8 +int16:standard [0, 1, -1, 32767, -32768]@Series[int16] [0, 1, -1, 32767, -32768]@int16 +int16:empty []@Series[int16] []@int16 +int32:standard [0, 1, -1, 2147483647, -2147483648]@Series[int32] [0, 1, -1, 2147483647, -2147483648]@int32 +int32:empty []@Series[int32] []@int32 +int64:standard [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 +int64:empty []@Series[int64] []@int64 +int64:nullable [0.0, 1.0, nan]@Series[float64] [0.0, 1.0, None]@float64 +uint8:standard [0, 1, 255]@Series[uint8] [0, 1, 255]@uint8 +uint16:standard [0, 1, 65535]@Series[uint16] [0, 1, 65535]@uint16 +uint32:standard [0, 1, 4294967295]@Series[uint32] [0, 1, 4294967295]@uint32 +uint64:standard [0, 1, 18446744073709551615]@Series[uint64] [0, 1, 18446744073709551615]@uint64 +float32:standard [0.0, 1.5, -1.5]@Series[float32] [0.0, 1.5, -1.5]@float32 +float32:nullable [0.0, nan, 1.5]@Series[float32] [0.0, None, 1.5]@float32 +float32:empty []@Series[float32] []@float32 +float64:standard [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@float64 +float64:nullable [0.0, nan, 1.5]@Series[float64] [0.0, None, 1.5]@float64 +float64:empty []@Series[float64] []@float64 +bool:standard [True, False, True]@Series[bool] [True, False, True]@bool +bool:empty []@Series[bool] []@bool +object:string ['hello', 'world', '']@Series[object] [hello, world, ]@string +object:string-nullable ['hello', None, 'world']@Series[object] [hello, None, world]@string +string:inferred ['hello', 'world']@Series[object] [hello, world]@string +object:bytes [b'hello', b'world']@Series[object] [b'hello', b'world']@binary +object:empty []@Series[object] []@null +object:all-null [None, None]@Series[object] [None, None]@null +object:decimal [Decimal('1.50'), Decimal('-2.25')]@Series[object] [1.50, -2.25]@decimal128(3, 2) +list<int64>:standard [[1, 2], [3]]@Series[object] [[1, 2], [3]]@list<item: int64> +list<int64>:nullable [[1, 2], None]@Series[object] [[1, 2], None]@list<item: int64> +list<int64>:null-element [[1, None], [3]]@Series[object] [[1, None], [3]]@list<item: int64> +list<string>:standard [['a', 'b'], ['c']]@Series[object] [['a', 'b'], ['c']]@list<item: string> +list<list<int64>>:standard [[[1, 2], [3]], [[4]]]@Series[object] [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> +list<struct>:standard [[{'a': 1}], [{'a': 2}]]@Series[object] [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> +struct:standard [{'a': 1, 'b': 'x'}]@Series[object] [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> +struct:nullable [{'a': 1, 'b': 'x'}, None]@Series[object] [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> +struct<struct>:standard [{'a': {'b': 1}}]@Series[object] [[('a', {'b': 1})]]@struct<a: struct<b: int64>> +struct<list<int64>>:standard [{'a': [1, 2]}]@Series[object] [[('a', [1, 2])]]@struct<a: list<item: int64>> +list<int64>:overflow [[300, 2], [3]]@Series[object] [[300, 2], [3]]@list<item: int64> +struct:overflow [{'a': 300, 'b': 'x'}]@Series[object] [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> +struct<int64>:standard [{'a': 1, 'b': 2}]@Series[object] [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> +struct<int64>:overflow [{'a': 300, 'b': 2}]@Series[object] [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> +datetime64[ns]:standard [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [2024-06-15 18:30:00]@timestamp[ns] +datetime64[ns]:nullable [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[ns]] [2024-06-15 18:30:00, None]@timestamp[ns] +datetime64[ns]:empty []@Series[datetime64[ns]] []@timestamp[ns] +datetime64[ns,tz]:standard [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] +datetime64[ns,tz]:nullable [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] +timedelta64[ns]:standard [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] +timedelta64[ns]:nullable [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [1 days 00:00:00, None]@duration[ns] +datetime64[us]:standard [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] [2024-06-15 18:30:00]@timestamp[us] +datetime64[us]:nullable [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[us]] [2024-06-15 18:30:00, None]@timestamp[us] +datetime64[us]:empty []@Series[datetime64[us]] []@timestamp[us] +datetime64[us,tz]:standard [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] +datetime64[us,tz]:nullable [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] +timedelta64[us]:standard [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[us]] [1 day, 0:00:00, 2:00:00]@duration[us] +timedelta64[us]:nullable [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] [1 day, 0:00:00, None]@duration[us] +datetime64:inferred [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [2024-06-15 18:30:00]@timestamp[ns] +timedelta64:inferred [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] +datetime64[us]:out-of-ns-range [Timestamp('1500-01-01 00:00:00')]@Series[datetime64[us]] [1500-01-01 00:00:00]@timestamp[us] +date:standard [datetime.date(2024, 6, 15)]@Series[object] [2024-06-15]@date32[day] +time:standard [datetime.time(18, 30, 45)]@Series[object] [18:30:45]@time64[us] +object:datetime [datetime.datetime(2024, 6, 15, 18, 30)]@Series[object] [2024-06-15 18:30:00]@timestamp[us] +object:datetime-sub-us [Timestamp('2024-01-01 00:00:00.000000123')]@Series[object] [2024-01-01 00:00:00]@timestamp[us] +object:timedelta [datetime.timedelta(days=1, seconds=7200)]@Series[object] [1 day, 2:00:00]@duration[us] +category:standard ['a', 'b', 'a']@Series[category] [a, b, a]@dictionary<values=string, indices=int8, ordered=0> +category:nullable ['a', nan, 'b']@Series[category] [a, None, b]@dictionary<values=string, indices=int8, ordered=0> +Int8:standard [0, 1, 127, -128]@Series[Int8] [0, 1, 127, -128]@int8 +Int8:nullable [0, 1, <NA>]@Series[Int8] [0, 1, None]@int8 +Int16:standard [0, 1, 32767, -32768]@Series[Int16] [0, 1, 32767, -32768]@int16 +Int16:nullable [0, 1, <NA>]@Series[Int16] [0, 1, None]@int16 +Int32:standard [0, 1, 2147483647, -2147483648]@Series[Int32] [0, 1, 2147483647, -2147483648]@int32 +Int32:nullable [0, 1, <NA>]@Series[Int32] [0, 1, None]@int32 +Int64:standard [0, 1, 9223372036854775807, -9223372036854775808]@Series[Int64] [0, 1, 9223372036854775807, -9223372036854775808]@int64 +Int64:nullable [0, 1, <NA>]@Series[Int64] [0, 1, None]@int64 +UInt64:standard [0, 1, 18446744073709551615]@Series[UInt64] [0, 1, 18446744073709551615]@uint64 +Int64:empty []@Series[Int64] []@int64 +Int64:all-null [<NA>, <NA>]@Series[Int64] [None, None]@int64 +Float64:standard [0.0, 1.5]@Series[Float64] [0.0, 1.5]@float64 +Float64:nullable [0.0, <NA>]@Series[Float64] [0.0, None]@float64 +boolean:standard [True, False]@Series[boolean] [True, False]@bool +boolean:nullable [True, <NA>]@Series[boolean] [True, None]@bool +string[python]:standard ['hello', 'world']@Series[string] [hello, world]@string +string[python]:nullable ['hello', <NA>]@Series[string] [hello, None]@string +string[python]:empty []@Series[string] []@string +int64[pyarrow]:standard [0, 1, -1]@Series[int64[pyarrow]] [0, 1, -1]@int64 +int64[pyarrow]:nullable [0, 1, <NA>]@Series[int64[pyarrow]] [0, 1, None]@int64 +int64[pyarrow]:empty []@Series[int64[pyarrow]] []@int64 +double[pyarrow]:nullable [0.0, <NA>]@Series[double[pyarrow]] [0.0, None]@float64 +bool[pyarrow]:nullable [True, <NA>]@Series[bool[pyarrow]] [True, None]@bool +string[pyarrow]:standard ['hello', 'world']@Series[string] [hello, world]@large_string +string[pyarrow]:nullable ['hello', <NA>]@Series[string] [hello, None]@large_string +string[pyarrow]:empty []@Series[string] []@large_string +large_binary[pyarrow]:standard [b'hello', b'world']@Series[large_binary[pyarrow]] [b'hello', b'world']@large_binary +timestamp[us][pyarrow]:standard [Timestamp('2024-01-01 12:00:00')]@Series[timestamp[us][pyarrow]] [2024-01-01 12:00:00]@timestamp[us] +int64[pyarrow]:single-chunk [1, 2]@Series[int64[pyarrow]] [1, 2]@int64 +int64[pyarrow]:multi-chunk [1, 2, 3]@Series[int64[pyarrow]] [1, 2, 3]@chunked<int64> +int64:overflow [300, 1]@Series[int64] [300, 1]@int64 +float64:fractional [1.5, 2.5]@Series[float64] [1.5, 2.5]@float64 +float64:infinity [inf, 1.0]@Series[float64] [inf, 1.0]@float64 +float64:precision [1.1234567890123]@Series[float64] [1.1234567890123]@float64 +datetime64[ns]:sub-us [Timestamp('2024-01-01 00:00:00.000000123')]@Series[datetime64[ns]] [2024-01-01 00:00:00.000000123]@timestamp[ns] +object:date-then-datetime [datetime.date(2024, 1, 1), datetime.datetime(2024, 1, 1, 5, 30)]@Series[object] [2024-01-01, 2024-01-01]@date32[day] +object:datetime-then-date [datetime.datetime(2024, 1, 1, 5, 30), datetime.date(2024, 1, 1)]@Series[object] ERR@ArrowTypeError diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.md new file mode 100644 index 0000000000000..a9acf237c1298 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_default.md @@ -0,0 +1,105 @@ +| test case | pandas series | arrow array | +|---------------------------------|--------------------------------------------------------------------------------------|-----------------------------------------------------------------| +| int8:standard | [0, 1, -1, 127, -128]@Series[int8] | [0, 1, -1, 127, -128]@int8 | +| int8:empty | []@Series[int8] | []@int8 | +| int16:standard | [0, 1, -1, 32767, -32768]@Series[int16] | [0, 1, -1, 32767, -32768]@int16 | +| int16:empty | []@Series[int16] | []@int16 | +| int32:standard | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | [0, 1, -1, 2147483647, -2147483648]@int32 | +| int32:empty | []@Series[int32] | []@int32 | +| int64:standard | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | +| int64:empty | []@Series[int64] | []@int64 | +| int64:nullable | [0.0, 1.0, nan]@Series[float64] | [0.0, 1.0, None]@float64 | +| uint8:standard | [0, 1, 255]@Series[uint8] | [0, 1, 255]@uint8 | +| uint16:standard | [0, 1, 65535]@Series[uint16] | [0, 1, 65535]@uint16 | +| uint32:standard | [0, 1, 4294967295]@Series[uint32] | [0, 1, 4294967295]@uint32 | +| uint64:standard | [0, 1, 18446744073709551615]@Series[uint64] | [0, 1, 18446744073709551615]@uint64 | +| float32:standard | [0.0, 1.5, -1.5]@Series[float32] | [0.0, 1.5, -1.5]@float32 | +| float32:nullable | [0.0, nan, 1.5]@Series[float32] | [0.0, None, 1.5]@float32 | +| float32:empty | []@Series[float32] | []@float32 | +| float64:standard | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@float64 | +| float64:nullable | [0.0, nan, 1.5]@Series[float64] | [0.0, None, 1.5]@float64 | +| float64:empty | []@Series[float64] | []@float64 | +| bool:standard | [True, False, True]@Series[bool] | [True, False, True]@bool | +| bool:empty | []@Series[bool] | []@bool | +| object:string | ['hello', 'world', '']@Series[object] | [hello, world, ]@string | +| object:string-nullable | ['hello', None, 'world']@Series[object] | [hello, None, world]@string | +| string:inferred | ['hello', 'world']@Series[object] | [hello, world]@string | +| object:bytes | [b'hello', b'world']@Series[object] | [b'hello', b'world']@binary | +| object:empty | []@Series[object] | []@null | +| object:all-null | [None, None]@Series[object] | [None, None]@null | +| object:decimal | [Decimal('1.50'), Decimal('-2.25')]@Series[object] | [1.50, -2.25]@decimal128(3, 2) | +| list<int64>:standard | [[1, 2], [3]]@Series[object] | [[1, 2], [3]]@list<item: int64> | +| list<int64>:nullable | [[1, 2], None]@Series[object] | [[1, 2], None]@list<item: int64> | +| list<int64>:null-element | [[1, None], [3]]@Series[object] | [[1, None], [3]]@list<item: int64> | +| list<string>:standard | [['a', 'b'], ['c']]@Series[object] | [['a', 'b'], ['c']]@list<item: string> | +| list<list<int64>>:standard | [[[1, 2], [3]], [[4]]]@Series[object] | [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> | +| list<struct>:standard | [[{'a': 1}], [{'a': 2}]]@Series[object] | [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> | +| struct:standard | [{'a': 1, 'b': 'x'}]@Series[object] | [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> | +| struct:nullable | [{'a': 1, 'b': 'x'}, None]@Series[object] | [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> | +| struct<struct>:standard | [{'a': {'b': 1}}]@Series[object] | [[('a', {'b': 1})]]@struct<a: struct<b: int64>> | +| struct<list<int64>>:standard | [{'a': [1, 2]}]@Series[object] | [[('a', [1, 2])]]@struct<a: list<item: int64>> | +| list<int64>:overflow | [[300, 2], [3]]@Series[object] | [[300, 2], [3]]@list<item: int64> | +| struct:overflow | [{'a': 300, 'b': 'x'}]@Series[object] | [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> | +| struct<int64>:standard | [{'a': 1, 'b': 2}]@Series[object] | [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> | +| struct<int64>:overflow | [{'a': 300, 'b': 2}]@Series[object] | [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> | +| datetime64[ns]:standard | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [2024-06-15 18:30:00]@timestamp[ns] | +| datetime64[ns]:nullable | [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[ns]] | [2024-06-15 18:30:00, None]@timestamp[ns] | +| datetime64[ns]:empty | []@Series[datetime64[ns]] | []@timestamp[ns] | +| datetime64[ns,tz]:standard | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] | [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] | +| datetime64[ns,tz]:nullable | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] | [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] | +| timedelta64[ns]:standard | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | +| timedelta64[ns]:nullable | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [1 days 00:00:00, None]@duration[ns] | +| datetime64[us]:standard | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | [2024-06-15 18:30:00]@timestamp[us] | +| datetime64[us]:nullable | [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[us]] | [2024-06-15 18:30:00, None]@timestamp[us] | +| datetime64[us]:empty | []@Series[datetime64[us]] | []@timestamp[us] | +| datetime64[us,tz]:standard | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | +| datetime64[us,tz]:nullable | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] | [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] | +| timedelta64[us]:standard | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[us]] | [1 day, 0:00:00, 2:00:00]@duration[us] | +| timedelta64[us]:nullable | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] | [1 day, 0:00:00, None]@duration[us] | +| datetime64:inferred | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [2024-06-15 18:30:00]@timestamp[ns] | +| timedelta64:inferred | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | +| datetime64[us]:out-of-ns-range | [Timestamp('1500-01-01 00:00:00')]@Series[datetime64[us]] | [1500-01-01 00:00:00]@timestamp[us] | +| date:standard | [datetime.date(2024, 6, 15)]@Series[object] | [2024-06-15]@date32[day] | +| time:standard | [datetime.time(18, 30, 45)]@Series[object] | [18:30:45]@time64[us] | +| object:datetime | [datetime.datetime(2024, 6, 15, 18, 30)]@Series[object] | [2024-06-15 18:30:00]@timestamp[us] | +| object:datetime-sub-us | [Timestamp('2024-01-01 00:00:00.000000123')]@Series[object] | [2024-01-01 00:00:00]@timestamp[us] | +| object:timedelta | [datetime.timedelta(days=1, seconds=7200)]@Series[object] | [1 day, 2:00:00]@duration[us] | +| category:standard | ['a', 'b', 'a']@Series[category] | [a, b, a]@dictionary<values=string, indices=int8, ordered=0> | +| category:nullable | ['a', nan, 'b']@Series[category] | [a, None, b]@dictionary<values=string, indices=int8, ordered=0> | +| Int8:standard | [0, 1, 127, -128]@Series[Int8] | [0, 1, 127, -128]@int8 | +| Int8:nullable | [0, 1, <NA>]@Series[Int8] | [0, 1, None]@int8 | +| Int16:standard | [0, 1, 32767, -32768]@Series[Int16] | [0, 1, 32767, -32768]@int16 | +| Int16:nullable | [0, 1, <NA>]@Series[Int16] | [0, 1, None]@int16 | +| Int32:standard | [0, 1, 2147483647, -2147483648]@Series[Int32] | [0, 1, 2147483647, -2147483648]@int32 | +| Int32:nullable | [0, 1, <NA>]@Series[Int32] | [0, 1, None]@int32 | +| Int64:standard | [0, 1, 9223372036854775807, -9223372036854775808]@Series[Int64] | [0, 1, 9223372036854775807, -9223372036854775808]@int64 | +| Int64:nullable | [0, 1, <NA>]@Series[Int64] | [0, 1, None]@int64 | +| UInt64:standard | [0, 1, 18446744073709551615]@Series[UInt64] | [0, 1, 18446744073709551615]@uint64 | +| Int64:empty | []@Series[Int64] | []@int64 | +| Int64:all-null | [<NA>, <NA>]@Series[Int64] | [None, None]@int64 | +| Float64:standard | [0.0, 1.5]@Series[Float64] | [0.0, 1.5]@float64 | +| Float64:nullable | [0.0, <NA>]@Series[Float64] | [0.0, None]@float64 | +| boolean:standard | [True, False]@Series[boolean] | [True, False]@bool | +| boolean:nullable | [True, <NA>]@Series[boolean] | [True, None]@bool | +| string[python]:standard | ['hello', 'world']@Series[string] | [hello, world]@string | +| string[python]:nullable | ['hello', <NA>]@Series[string] | [hello, None]@string | +| string[python]:empty | []@Series[string] | []@string | +| int64[pyarrow]:standard | [0, 1, -1]@Series[int64[pyarrow]] | [0, 1, -1]@int64 | +| int64[pyarrow]:nullable | [0, 1, <NA>]@Series[int64[pyarrow]] | [0, 1, None]@int64 | +| int64[pyarrow]:empty | []@Series[int64[pyarrow]] | []@int64 | +| double[pyarrow]:nullable | [0.0, <NA>]@Series[double[pyarrow]] | [0.0, None]@float64 | +| bool[pyarrow]:nullable | [True, <NA>]@Series[bool[pyarrow]] | [True, None]@bool | +| string[pyarrow]:standard | ['hello', 'world']@Series[string] | [hello, world]@large_string | +| string[pyarrow]:nullable | ['hello', <NA>]@Series[string] | [hello, None]@large_string | +| string[pyarrow]:empty | []@Series[string] | []@large_string | +| large_binary[pyarrow]:standard | [b'hello', b'world']@Series[large_binary[pyarrow]] | [b'hello', b'world']@large_binary | +| timestamp[us][pyarrow]:standard | [Timestamp('2024-01-01 12:00:00')]@Series[timestamp[us][pyarrow]] | [2024-01-01 12:00:00]@timestamp[us] | +| int64[pyarrow]:single-chunk | [1, 2]@Series[int64[pyarrow]] | [1, 2]@int64 | +| int64[pyarrow]:multi-chunk | [1, 2, 3]@Series[int64[pyarrow]] | [1, 2, 3]@chunked<int64> | +| int64:overflow | [300, 1]@Series[int64] | [300, 1]@int64 | +| float64:fractional | [1.5, 2.5]@Series[float64] | [1.5, 2.5]@float64 | +| float64:infinity | [inf, 1.0]@Series[float64] | [inf, 1.0]@float64 | +| float64:precision | [1.1234567890123]@Series[float64] | [1.1234567890123]@float64 | +| datetime64[ns]:sub-us | [Timestamp('2024-01-01 00:00:00.000000123')]@Series[datetime64[ns]] | [2024-01-01 00:00:00.000000123]@timestamp[ns] | +| object:date-then-datetime | [datetime.date(2024, 1, 1), datetime.datetime(2024, 1, 1, 5, 30)]@Series[object] | [2024-01-01, 2024-01-01]@date32[day] | +| object:datetime-then-date | [datetime.datetime(2024, 1, 1, 5, 30), datetime.date(2024, 1, 1)]@Series[object] | ERR@ArrowTypeError | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.csv new file mode 100644 index 0000000000000..69debccb7824f --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.csv @@ -0,0 +1,104 @@ +test case pandas series mask=None mask=isnull() +int8:standard [0, 1, -1, 127, -128]@Series[int8] [0, 1, -1, 127, -128]@int8 [0, 1, -1, 127, -128]@int8 +int8:empty []@Series[int8] []@int8 []@int8 +int16:standard [0, 1, -1, 32767, -32768]@Series[int16] [0, 1, -1, 32767, -32768]@int16 [0, 1, -1, 32767, -32768]@int16 +int16:empty []@Series[int16] []@int16 []@int16 +int32:standard [0, 1, -1, 2147483647, -2147483648]@Series[int32] [0, 1, -1, 2147483647, -2147483648]@int32 [0, 1, -1, 2147483647, -2147483648]@int32 +int32:empty []@Series[int32] []@int32 []@int32 +int64:standard [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 +int64:empty []@Series[int64] []@int64 []@int64 +int64:nullable [0.0, 1.0, nan]@Series[float64] [0.0, 1.0, None]@float64 [0.0, 1.0, None]@float64 +uint8:standard [0, 1, 255]@Series[uint8] [0, 1, 255]@uint8 [0, 1, 255]@uint8 +uint16:standard [0, 1, 65535]@Series[uint16] [0, 1, 65535]@uint16 [0, 1, 65535]@uint16 +uint32:standard [0, 1, 4294967295]@Series[uint32] [0, 1, 4294967295]@uint32 [0, 1, 4294967295]@uint32 +uint64:standard [0, 1, 18446744073709551615]@Series[uint64] [0, 1, 18446744073709551615]@uint64 [0, 1, 18446744073709551615]@uint64 +float32:standard [0.0, 1.5, -1.5]@Series[float32] [0.0, 1.5, -1.5]@float32 [0.0, 1.5, -1.5]@float32 +float32:nullable [0.0, nan, 1.5]@Series[float32] [0.0, None, 1.5]@float32 [0.0, None, 1.5]@float32 +float32:empty []@Series[float32] []@float32 []@float32 +float64:standard [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@float64 [0.0, 1.5, -1.5]@float64 +float64:nullable [0.0, nan, 1.5]@Series[float64] [0.0, None, 1.5]@float64 [0.0, None, 1.5]@float64 +float64:empty []@Series[float64] []@float64 []@float64 +bool:standard [True, False, True]@Series[bool] [True, False, True]@bool [True, False, True]@bool +bool:empty []@Series[bool] []@bool []@bool +object:string ['hello', 'world', '']@Series[object] [hello, world, ]@string [hello, world, ]@string +object:string-nullable ['hello', None, 'world']@Series[object] [hello, None, world]@string [hello, None, world]@string +string:inferred ['hello', 'world']@Series[object] [hello, world]@string [hello, world]@string +object:bytes [b'hello', b'world']@Series[object] [b'hello', b'world']@binary [b'hello', b'world']@binary +object:empty []@Series[object] []@null []@null +object:all-null [None, None]@Series[object] [None, None]@null [None, None]@null +object:decimal [Decimal('1.50'), Decimal('-2.25')]@Series[object] [1.50, -2.25]@decimal128(3, 2) [1.50, -2.25]@decimal128(3, 2) +list<int64>:standard [[1, 2], [3]]@Series[object] [[1, 2], [3]]@list<item: int64> [[1, 2], [3]]@list<item: int64> +list<int64>:nullable [[1, 2], None]@Series[object] [[1, 2], None]@list<item: int64> [[1, 2], None]@list<item: int64> +list<int64>:null-element [[1, None], [3]]@Series[object] [[1, None], [3]]@list<item: int64> [[1, None], [3]]@list<item: int64> +list<string>:standard [['a', 'b'], ['c']]@Series[object] [['a', 'b'], ['c']]@list<item: string> [['a', 'b'], ['c']]@list<item: string> +list<list<int64>>:standard [[[1, 2], [3]], [[4]]]@Series[object] [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> +list<struct>:standard [[{'a': 1}], [{'a': 2}]]@Series[object] [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> +struct:standard [{'a': 1, 'b': 'x'}]@Series[object] [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> +struct:nullable [{'a': 1, 'b': 'x'}, None]@Series[object] [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> +struct<struct>:standard [{'a': {'b': 1}}]@Series[object] [[('a', {'b': 1})]]@struct<a: struct<b: int64>> [[('a', {'b': 1})]]@struct<a: struct<b: int64>> +struct<list<int64>>:standard [{'a': [1, 2]}]@Series[object] [[('a', [1, 2])]]@struct<a: list<item: int64>> [[('a', [1, 2])]]@struct<a: list<item: int64>> +list<int64>:overflow [[300, 2], [3]]@Series[object] [[300, 2], [3]]@list<item: int64> [[300, 2], [3]]@list<item: int64> +struct:overflow [{'a': 300, 'b': 'x'}]@Series[object] [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> +struct<int64>:standard [{'a': 1, 'b': 2}]@Series[object] [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> +struct<int64>:overflow [{'a': 300, 'b': 2}]@Series[object] [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> +datetime64[ns]:standard [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [2024-06-15 18:30:00]@timestamp[ns] [2024-06-15 18:30:00]@timestamp[ns] +datetime64[ns]:nullable [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[ns]] [2024-06-15 18:30:00, None]@timestamp[ns] [2024-06-15 18:30:00, None]@timestamp[ns] +datetime64[ns]:empty []@Series[datetime64[ns]] []@timestamp[ns] []@timestamp[ns] +datetime64[ns,tz]:standard [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] +datetime64[ns,tz]:nullable [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] +timedelta64[ns]:standard [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] +timedelta64[ns]:nullable [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [1 days 00:00:00, None]@duration[ns] [1 days 00:00:00, None]@duration[ns] +datetime64[us]:standard [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] [2024-06-15 18:30:00]@timestamp[us] [2024-06-15 18:30:00]@timestamp[us] +datetime64[us]:nullable [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[us]] [2024-06-15 18:30:00, None]@timestamp[us] [2024-06-15 18:30:00, None]@timestamp[us] +datetime64[us]:empty []@Series[datetime64[us]] []@timestamp[us] []@timestamp[us] +datetime64[us,tz]:standard [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] +datetime64[us,tz]:nullable [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] +timedelta64[us]:standard [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[us]] [1 day, 0:00:00, 2:00:00]@duration[us] [1 day, 0:00:00, 2:00:00]@duration[us] +timedelta64[us]:nullable [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] [1 day, 0:00:00, None]@duration[us] [1 day, 0:00:00, None]@duration[us] +datetime64:inferred [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [2024-06-15 18:30:00]@timestamp[ns] [2024-06-15 18:30:00]@timestamp[ns] +timedelta64:inferred [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] [1 days 00:00:00, 0 days 02:00:00]@duration[ns] +datetime64[us]:out-of-ns-range [Timestamp('1500-01-01 00:00:00')]@Series[datetime64[us]] [1500-01-01 00:00:00]@timestamp[us] [1500-01-01 00:00:00]@timestamp[us] +date:standard [datetime.date(2024, 6, 15)]@Series[object] [2024-06-15]@date32[day] [2024-06-15]@date32[day] +time:standard [datetime.time(18, 30, 45)]@Series[object] [18:30:45]@time64[us] [18:30:45]@time64[us] +object:datetime [datetime.datetime(2024, 6, 15, 18, 30)]@Series[object] [2024-06-15 18:30:00]@timestamp[us] [2024-06-15 18:30:00]@timestamp[us] +object:datetime-sub-us [Timestamp('2024-01-01 00:00:00.000000123')]@Series[object] [2024-01-01 00:00:00]@timestamp[us] [2024-01-01 00:00:00]@timestamp[us] +object:timedelta [datetime.timedelta(days=1, seconds=7200)]@Series[object] [1 day, 2:00:00]@duration[us] [1 day, 2:00:00]@duration[us] +category:standard ['a', 'b', 'a']@Series[category] [a, b, a]@dictionary<values=string, indices=int8, ordered=0> [a, b, a]@dictionary<values=string, indices=int8, ordered=0> +category:nullable ['a', nan, 'b']@Series[category] [a, None, b]@dictionary<values=string, indices=int8, ordered=0> [a, None, b]@dictionary<values=string, indices=int8, ordered=0> +Int8:standard [0, 1, 127, -128]@Series[Int8] [0, 1, 127, -128]@int8 ERR@ValueError +Int8:nullable [0, 1, <NA>]@Series[Int8] [0, 1, None]@int8 ERR@ValueError +Int16:standard [0, 1, 32767, -32768]@Series[Int16] [0, 1, 32767, -32768]@int16 ERR@ValueError +Int16:nullable [0, 1, <NA>]@Series[Int16] [0, 1, None]@int16 ERR@ValueError +Int32:standard [0, 1, 2147483647, -2147483648]@Series[Int32] [0, 1, 2147483647, -2147483648]@int32 ERR@ValueError +Int32:nullable [0, 1, <NA>]@Series[Int32] [0, 1, None]@int32 ERR@ValueError +Int64:standard [0, 1, 9223372036854775807, -9223372036854775808]@Series[Int64] [0, 1, 9223372036854775807, -9223372036854775808]@int64 ERR@ValueError +Int64:nullable [0, 1, <NA>]@Series[Int64] [0, 1, None]@int64 ERR@ValueError +UInt64:standard [0, 1, 18446744073709551615]@Series[UInt64] [0, 1, 18446744073709551615]@uint64 ERR@ValueError +Int64:empty []@Series[Int64] []@int64 ERR@ValueError +Int64:all-null [<NA>, <NA>]@Series[Int64] [None, None]@int64 ERR@ValueError +Float64:standard [0.0, 1.5]@Series[Float64] [0.0, 1.5]@float64 ERR@ValueError +Float64:nullable [0.0, <NA>]@Series[Float64] [0.0, None]@float64 ERR@ValueError +boolean:standard [True, False]@Series[boolean] [True, False]@bool ERR@ValueError +boolean:nullable [True, <NA>]@Series[boolean] [True, None]@bool ERR@ValueError +string[python]:standard ['hello', 'world']@Series[string] [hello, world]@string ERR@ValueError +string[python]:nullable ['hello', <NA>]@Series[string] [hello, None]@string ERR@ValueError +string[python]:empty []@Series[string] []@string ERR@ValueError +int64[pyarrow]:standard [0, 1, -1]@Series[int64[pyarrow]] [0, 1, -1]@int64 ERR@ValueError +int64[pyarrow]:nullable [0, 1, <NA>]@Series[int64[pyarrow]] [0, 1, None]@int64 ERR@ValueError +int64[pyarrow]:empty []@Series[int64[pyarrow]] []@int64 ERR@ValueError +double[pyarrow]:nullable [0.0, <NA>]@Series[double[pyarrow]] [0.0, None]@float64 ERR@ValueError +bool[pyarrow]:nullable [True, <NA>]@Series[bool[pyarrow]] [True, None]@bool ERR@ValueError +string[pyarrow]:standard ['hello', 'world']@Series[string] [hello, world]@large_string ERR@ValueError +string[pyarrow]:nullable ['hello', <NA>]@Series[string] [hello, None]@large_string ERR@ValueError +string[pyarrow]:empty []@Series[string] []@large_string ERR@ValueError +large_binary[pyarrow]:standard [b'hello', b'world']@Series[large_binary[pyarrow]] [b'hello', b'world']@large_binary ERR@ValueError +timestamp[us][pyarrow]:standard [Timestamp('2024-01-01 12:00:00')]@Series[timestamp[us][pyarrow]] [2024-01-01 12:00:00]@timestamp[us] ERR@ValueError +int64[pyarrow]:single-chunk [1, 2]@Series[int64[pyarrow]] [1, 2]@int64 ERR@ValueError +int64[pyarrow]:multi-chunk [1, 2, 3]@Series[int64[pyarrow]] [1, 2, 3]@chunked<int64> ERR@ValueError +int64:overflow [300, 1]@Series[int64] [300, 1]@int64 [300, 1]@int64 +float64:fractional [1.5, 2.5]@Series[float64] [1.5, 2.5]@float64 [1.5, 2.5]@float64 +float64:infinity [inf, 1.0]@Series[float64] [inf, 1.0]@float64 [inf, 1.0]@float64 +float64:precision [1.1234567890123]@Series[float64] [1.1234567890123]@float64 [1.1234567890123]@float64 +datetime64[ns]:sub-us [Timestamp('2024-01-01 00:00:00.000000123')]@Series[datetime64[ns]] [2024-01-01 00:00:00.000000123]@timestamp[ns] [2024-01-01 00:00:00.000000123]@timestamp[ns] +object:date-then-datetime [datetime.date(2024, 1, 1), datetime.datetime(2024, 1, 1, 5, 30)]@Series[object] [2024-01-01, 2024-01-01]@date32[day] [2024-01-01, 2024-01-01]@date32[day] +object:datetime-then-date [datetime.datetime(2024, 1, 1, 5, 30), datetime.date(2024, 1, 1)]@Series[object] ERR@ArrowTypeError ERR@ArrowTypeError diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.md new file mode 100644 index 0000000000000..b8085d64629ef --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_mask.md @@ -0,0 +1,105 @@ +| test case | pandas series | mask=None | mask=isnull() | +|---------------------------------|--------------------------------------------------------------------------------------|-----------------------------------------------------------------|-----------------------------------------------------------------| +| int8:standard | [0, 1, -1, 127, -128]@Series[int8] | [0, 1, -1, 127, -128]@int8 | [0, 1, -1, 127, -128]@int8 | +| int8:empty | []@Series[int8] | []@int8 | []@int8 | +| int16:standard | [0, 1, -1, 32767, -32768]@Series[int16] | [0, 1, -1, 32767, -32768]@int16 | [0, 1, -1, 32767, -32768]@int16 | +| int16:empty | []@Series[int16] | []@int16 | []@int16 | +| int32:standard | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | [0, 1, -1, 2147483647, -2147483648]@int32 | [0, 1, -1, 2147483647, -2147483648]@int32 | +| int32:empty | []@Series[int32] | []@int32 | []@int32 | +| int64:standard | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | +| int64:empty | []@Series[int64] | []@int64 | []@int64 | +| int64:nullable | [0.0, 1.0, nan]@Series[float64] | [0.0, 1.0, None]@float64 | [0.0, 1.0, None]@float64 | +| uint8:standard | [0, 1, 255]@Series[uint8] | [0, 1, 255]@uint8 | [0, 1, 255]@uint8 | +| uint16:standard | [0, 1, 65535]@Series[uint16] | [0, 1, 65535]@uint16 | [0, 1, 65535]@uint16 | +| uint32:standard | [0, 1, 4294967295]@Series[uint32] | [0, 1, 4294967295]@uint32 | [0, 1, 4294967295]@uint32 | +| uint64:standard | [0, 1, 18446744073709551615]@Series[uint64] | [0, 1, 18446744073709551615]@uint64 | [0, 1, 18446744073709551615]@uint64 | +| float32:standard | [0.0, 1.5, -1.5]@Series[float32] | [0.0, 1.5, -1.5]@float32 | [0.0, 1.5, -1.5]@float32 | +| float32:nullable | [0.0, nan, 1.5]@Series[float32] | [0.0, None, 1.5]@float32 | [0.0, None, 1.5]@float32 | +| float32:empty | []@Series[float32] | []@float32 | []@float32 | +| float64:standard | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@float64 | [0.0, 1.5, -1.5]@float64 | +| float64:nullable | [0.0, nan, 1.5]@Series[float64] | [0.0, None, 1.5]@float64 | [0.0, None, 1.5]@float64 | +| float64:empty | []@Series[float64] | []@float64 | []@float64 | +| bool:standard | [True, False, True]@Series[bool] | [True, False, True]@bool | [True, False, True]@bool | +| bool:empty | []@Series[bool] | []@bool | []@bool | +| object:string | ['hello', 'world', '']@Series[object] | [hello, world, ]@string | [hello, world, ]@string | +| object:string-nullable | ['hello', None, 'world']@Series[object] | [hello, None, world]@string | [hello, None, world]@string | +| string:inferred | ['hello', 'world']@Series[object] | [hello, world]@string | [hello, world]@string | +| object:bytes | [b'hello', b'world']@Series[object] | [b'hello', b'world']@binary | [b'hello', b'world']@binary | +| object:empty | []@Series[object] | []@null | []@null | +| object:all-null | [None, None]@Series[object] | [None, None]@null | [None, None]@null | +| object:decimal | [Decimal('1.50'), Decimal('-2.25')]@Series[object] | [1.50, -2.25]@decimal128(3, 2) | [1.50, -2.25]@decimal128(3, 2) | +| list<int64>:standard | [[1, 2], [3]]@Series[object] | [[1, 2], [3]]@list<item: int64> | [[1, 2], [3]]@list<item: int64> | +| list<int64>:nullable | [[1, 2], None]@Series[object] | [[1, 2], None]@list<item: int64> | [[1, 2], None]@list<item: int64> | +| list<int64>:null-element | [[1, None], [3]]@Series[object] | [[1, None], [3]]@list<item: int64> | [[1, None], [3]]@list<item: int64> | +| list<string>:standard | [['a', 'b'], ['c']]@Series[object] | [['a', 'b'], ['c']]@list<item: string> | [['a', 'b'], ['c']]@list<item: string> | +| list<list<int64>>:standard | [[[1, 2], [3]], [[4]]]@Series[object] | [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> | [[[1, 2], [3]], [[4]]]@list<item: list<item: int64>> | +| list<struct>:standard | [[{'a': 1}], [{'a': 2}]]@Series[object] | [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> | [[{'a': 1}], [{'a': 2}]]@list<item: struct<a: int64>> | +| struct:standard | [{'a': 1, 'b': 'x'}]@Series[object] | [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> | +| struct:nullable | [{'a': 1, 'b': 'x'}, None]@Series[object] | [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> | +| struct<struct>:standard | [{'a': {'b': 1}}]@Series[object] | [[('a', {'b': 1})]]@struct<a: struct<b: int64>> | [[('a', {'b': 1})]]@struct<a: struct<b: int64>> | +| struct<list<int64>>:standard | [{'a': [1, 2]}]@Series[object] | [[('a', [1, 2])]]@struct<a: list<item: int64>> | [[('a', [1, 2])]]@struct<a: list<item: int64>> | +| list<int64>:overflow | [[300, 2], [3]]@Series[object] | [[300, 2], [3]]@list<item: int64> | [[300, 2], [3]]@list<item: int64> | +| struct:overflow | [{'a': 300, 'b': 'x'}]@Series[object] | [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> | [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> | +| struct<int64>:standard | [{'a': 1, 'b': 2}]@Series[object] | [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> | [[('a', 1), ('b', 2)]]@struct<a: int64, b: int64> | +| struct<int64>:overflow | [{'a': 300, 'b': 2}]@Series[object] | [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> | [[('a', 300), ('b', 2)]]@struct<a: int64, b: int64> | +| datetime64[ns]:standard | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [2024-06-15 18:30:00]@timestamp[ns] | [2024-06-15 18:30:00]@timestamp[ns] | +| datetime64[ns]:nullable | [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[ns]] | [2024-06-15 18:30:00, None]@timestamp[ns] | [2024-06-15 18:30:00, None]@timestamp[ns] | +| datetime64[ns]:empty | []@Series[datetime64[ns]] | []@timestamp[ns] | []@timestamp[ns] | +| datetime64[ns,tz]:standard | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] | [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] | [2024-06-15 18:30:00+00:00]@timestamp[ns, tz=UTC] | +| datetime64[ns,tz]:nullable | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] | [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] | [2024-06-15 18:30:00+00:00, None]@timestamp[ns, tz=UTC] | +| timedelta64[ns]:standard | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | +| timedelta64[ns]:nullable | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [1 days 00:00:00, None]@duration[ns] | [1 days 00:00:00, None]@duration[ns] | +| datetime64[us]:standard | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | [2024-06-15 18:30:00]@timestamp[us] | [2024-06-15 18:30:00]@timestamp[us] | +| datetime64[us]:nullable | [Timestamp('2024-06-15 18:30:00'), NaT]@Series[datetime64[us]] | [2024-06-15 18:30:00, None]@timestamp[us] | [2024-06-15 18:30:00, None]@timestamp[us] | +| datetime64[us]:empty | []@Series[datetime64[us]] | []@timestamp[us] | []@timestamp[us] | +| datetime64[us,tz]:standard | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | [2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | +| datetime64[us,tz]:nullable | [Timestamp('2024-06-15 18:30:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] | [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] | [2024-06-15 18:30:00+00:00, None]@timestamp[us, tz=UTC] | +| timedelta64[us]:standard | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[us]] | [1 day, 0:00:00, 2:00:00]@duration[us] | [1 day, 0:00:00, 2:00:00]@duration[us] | +| timedelta64[us]:nullable | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] | [1 day, 0:00:00, None]@duration[us] | [1 day, 0:00:00, None]@duration[us] | +| datetime64:inferred | [Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [2024-06-15 18:30:00]@timestamp[ns] | [2024-06-15 18:30:00]@timestamp[ns] | +| timedelta64:inferred | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:00:00')]@Series[timedelta64[ns]] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | [1 days 00:00:00, 0 days 02:00:00]@duration[ns] | +| datetime64[us]:out-of-ns-range | [Timestamp('1500-01-01 00:00:00')]@Series[datetime64[us]] | [1500-01-01 00:00:00]@timestamp[us] | [1500-01-01 00:00:00]@timestamp[us] | +| date:standard | [datetime.date(2024, 6, 15)]@Series[object] | [2024-06-15]@date32[day] | [2024-06-15]@date32[day] | +| time:standard | [datetime.time(18, 30, 45)]@Series[object] | [18:30:45]@time64[us] | [18:30:45]@time64[us] | +| object:datetime | [datetime.datetime(2024, 6, 15, 18, 30)]@Series[object] | [2024-06-15 18:30:00]@timestamp[us] | [2024-06-15 18:30:00]@timestamp[us] | +| object:datetime-sub-us | [Timestamp('2024-01-01 00:00:00.000000123')]@Series[object] | [2024-01-01 00:00:00]@timestamp[us] | [2024-01-01 00:00:00]@timestamp[us] | +| object:timedelta | [datetime.timedelta(days=1, seconds=7200)]@Series[object] | [1 day, 2:00:00]@duration[us] | [1 day, 2:00:00]@duration[us] | +| category:standard | ['a', 'b', 'a']@Series[category] | [a, b, a]@dictionary<values=string, indices=int8, ordered=0> | [a, b, a]@dictionary<values=string, indices=int8, ordered=0> | +| category:nullable | ['a', nan, 'b']@Series[category] | [a, None, b]@dictionary<values=string, indices=int8, ordered=0> | [a, None, b]@dictionary<values=string, indices=int8, ordered=0> | +| Int8:standard | [0, 1, 127, -128]@Series[Int8] | [0, 1, 127, -128]@int8 | ERR@ValueError | +| Int8:nullable | [0, 1, <NA>]@Series[Int8] | [0, 1, None]@int8 | ERR@ValueError | +| Int16:standard | [0, 1, 32767, -32768]@Series[Int16] | [0, 1, 32767, -32768]@int16 | ERR@ValueError | +| Int16:nullable | [0, 1, <NA>]@Series[Int16] | [0, 1, None]@int16 | ERR@ValueError | +| Int32:standard | [0, 1, 2147483647, -2147483648]@Series[Int32] | [0, 1, 2147483647, -2147483648]@int32 | ERR@ValueError | +| Int32:nullable | [0, 1, <NA>]@Series[Int32] | [0, 1, None]@int32 | ERR@ValueError | +| Int64:standard | [0, 1, 9223372036854775807, -9223372036854775808]@Series[Int64] | [0, 1, 9223372036854775807, -9223372036854775808]@int64 | ERR@ValueError | +| Int64:nullable | [0, 1, <NA>]@Series[Int64] | [0, 1, None]@int64 | ERR@ValueError | +| UInt64:standard | [0, 1, 18446744073709551615]@Series[UInt64] | [0, 1, 18446744073709551615]@uint64 | ERR@ValueError | +| Int64:empty | []@Series[Int64] | []@int64 | ERR@ValueError | +| Int64:all-null | [<NA>, <NA>]@Series[Int64] | [None, None]@int64 | ERR@ValueError | +| Float64:standard | [0.0, 1.5]@Series[Float64] | [0.0, 1.5]@float64 | ERR@ValueError | +| Float64:nullable | [0.0, <NA>]@Series[Float64] | [0.0, None]@float64 | ERR@ValueError | +| boolean:standard | [True, False]@Series[boolean] | [True, False]@bool | ERR@ValueError | +| boolean:nullable | [True, <NA>]@Series[boolean] | [True, None]@bool | ERR@ValueError | +| string[python]:standard | ['hello', 'world']@Series[string] | [hello, world]@string | ERR@ValueError | +| string[python]:nullable | ['hello', <NA>]@Series[string] | [hello, None]@string | ERR@ValueError | +| string[python]:empty | []@Series[string] | []@string | ERR@ValueError | +| int64[pyarrow]:standard | [0, 1, -1]@Series[int64[pyarrow]] | [0, 1, -1]@int64 | ERR@ValueError | +| int64[pyarrow]:nullable | [0, 1, <NA>]@Series[int64[pyarrow]] | [0, 1, None]@int64 | ERR@ValueError | +| int64[pyarrow]:empty | []@Series[int64[pyarrow]] | []@int64 | ERR@ValueError | +| double[pyarrow]:nullable | [0.0, <NA>]@Series[double[pyarrow]] | [0.0, None]@float64 | ERR@ValueError | +| bool[pyarrow]:nullable | [True, <NA>]@Series[bool[pyarrow]] | [True, None]@bool | ERR@ValueError | +| string[pyarrow]:standard | ['hello', 'world']@Series[string] | [hello, world]@large_string | ERR@ValueError | +| string[pyarrow]:nullable | ['hello', <NA>]@Series[string] | [hello, None]@large_string | ERR@ValueError | +| string[pyarrow]:empty | []@Series[string] | []@large_string | ERR@ValueError | +| large_binary[pyarrow]:standard | [b'hello', b'world']@Series[large_binary[pyarrow]] | [b'hello', b'world']@large_binary | ERR@ValueError | +| timestamp[us][pyarrow]:standard | [Timestamp('2024-01-01 12:00:00')]@Series[timestamp[us][pyarrow]] | [2024-01-01 12:00:00]@timestamp[us] | ERR@ValueError | +| int64[pyarrow]:single-chunk | [1, 2]@Series[int64[pyarrow]] | [1, 2]@int64 | ERR@ValueError | +| int64[pyarrow]:multi-chunk | [1, 2, 3]@Series[int64[pyarrow]] | [1, 2, 3]@chunked<int64> | ERR@ValueError | +| int64:overflow | [300, 1]@Series[int64] | [300, 1]@int64 | [300, 1]@int64 | +| float64:fractional | [1.5, 2.5]@Series[float64] | [1.5, 2.5]@float64 | [1.5, 2.5]@float64 | +| float64:infinity | [inf, 1.0]@Series[float64] | [inf, 1.0]@float64 | [inf, 1.0]@float64 | +| float64:precision | [1.1234567890123]@Series[float64] | [1.1234567890123]@float64 | [1.1234567890123]@float64 | +| datetime64[ns]:sub-us | [Timestamp('2024-01-01 00:00:00.000000123')]@Series[datetime64[ns]] | [2024-01-01 00:00:00.000000123]@timestamp[ns] | [2024-01-01 00:00:00.000000123]@timestamp[ns] | +| object:date-then-datetime | [datetime.date(2024, 1, 1), datetime.datetime(2024, 1, 1, 5, 30)]@Series[object] | [2024-01-01, 2024-01-01]@date32[day] | [2024-01-01, 2024-01-01]@date32[day] | +| object:datetime-then-date | [datetime.datetime(2024, 1, 1, 5, 30), datetime.date(2024, 1, 1)]@Series[object] | ERR@ArrowTypeError | ERR@ArrowTypeError | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.csv new file mode 100644 index 0000000000000..47e9f859dc1f0 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.csv @@ -0,0 +1,11 @@ +source \ target list<element: int64> list<element: int8> list<element: string> map<string, int64> map<string, int8> struct<a: int64, b: string> struct<a: int8, b: string> +list<int64>:standard [[1, 2], [3]]@list<element: int64> [[1, 2], [3]]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:overflow [[300, 2], [3]]@list<element: int64> ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:nullable [[1, 2], None]@list<element: int64> [[1, 2], None]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:null-element [[1, None], [3]]@list<element: int64> [[1, None], [3]]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<string>:standard ERR@ArrowInvalid ERR@ArrowInvalid [['a', 'b'], ['c']]@list<element: string> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +struct:standard ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')]]@struct<a: int8, b: string> +struct:overflow ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> ERR@ArrowInvalid +struct:nullable ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')], None]@struct<a: int8, b: string> +struct<int64>:standard ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [[('a', 1), ('b', 2)]]@map<string, int64> [[('a', 1), ('b', 2)]]@map<string, int8> ERR@ArrowTypeError ERR@ArrowTypeError +struct<int64>:overflow ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [[('a', 300), ('b', 2)]]@map<string, int64> ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowInvalid diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.md new file mode 100644 index 0000000000000..398ce8fd2c9e7 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_safe.md @@ -0,0 +1,12 @@ +| source \ target | list<element: int64> | list<element: int8> | list<element: string> | map<string, int64> | map<string, int8> | struct<a: int64, b: string> | struct<a: int8, b: string> | +|--------------------------|---------------------------------------|--------------------------------------|-------------------------------------------|---------------------------------------------|------------------------------------------|------------------------------------------------------------|-----------------------------------------------------------| +| list<int64>:standard | [[1, 2], [3]]@list<element: int64> | [[1, 2], [3]]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:overflow | [[300, 2], [3]]@list<element: int64> | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:nullable | [[1, 2], None]@list<element: int64> | [[1, 2], None]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:null-element | [[1, None], [3]]@list<element: int64> | [[1, None], [3]]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<string>:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | [['a', 'b'], ['c']]@list<element: string> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| struct:standard | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')]]@struct<a: int8, b: string> | +| struct:overflow | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> | ERR@ArrowInvalid | +| struct:nullable | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')], None]@struct<a: int8, b: string> | +| struct<int64>:standard | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [[('a', 1), ('b', 2)]]@map<string, int64> | [[('a', 1), ('b', 2)]]@map<string, int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | +| struct<int64>:overflow | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [[('a', 300), ('b', 2)]]@map<string, int64> | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowInvalid | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.csv new file mode 100644 index 0000000000000..47e9f859dc1f0 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.csv @@ -0,0 +1,11 @@ +source \ target list<element: int64> list<element: int8> list<element: string> map<string, int64> map<string, int8> struct<a: int64, b: string> struct<a: int8, b: string> +list<int64>:standard [[1, 2], [3]]@list<element: int64> [[1, 2], [3]]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:overflow [[300, 2], [3]]@list<element: int64> ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:nullable [[1, 2], None]@list<element: int64> [[1, 2], None]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<int64>:null-element [[1, None], [3]]@list<element: int64> [[1, None], [3]]@list<element: int8> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +list<string>:standard ERR@ArrowInvalid ERR@ArrowInvalid [['a', 'b'], ['c']]@list<element: string> ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +struct:standard ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')]]@struct<a: int8, b: string> +struct:overflow ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> ERR@ArrowInvalid +struct:nullable ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowInvalid ERR@ArrowInvalid [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> [[('a', 1), ('b', 'x')], None]@struct<a: int8, b: string> +struct<int64>:standard ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [[('a', 1), ('b', 2)]]@map<string, int64> [[('a', 1), ('b', 2)]]@map<string, int8> ERR@ArrowTypeError ERR@ArrowTypeError +struct<int64>:overflow ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [[('a', 300), ('b', 2)]]@map<string, int64> ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowInvalid diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.md new file mode 100644 index 0000000000000..398ce8fd2c9e7 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_nested_unsafe.md @@ -0,0 +1,12 @@ +| source \ target | list<element: int64> | list<element: int8> | list<element: string> | map<string, int64> | map<string, int8> | struct<a: int64, b: string> | struct<a: int8, b: string> | +|--------------------------|---------------------------------------|--------------------------------------|-------------------------------------------|---------------------------------------------|------------------------------------------|------------------------------------------------------------|-----------------------------------------------------------| +| list<int64>:standard | [[1, 2], [3]]@list<element: int64> | [[1, 2], [3]]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:overflow | [[300, 2], [3]]@list<element: int64> | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:nullable | [[1, 2], None]@list<element: int64> | [[1, 2], None]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<int64>:null-element | [[1, None], [3]]@list<element: int64> | [[1, None], [3]]@list<element: int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| list<string>:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | [['a', 'b'], ['c']]@list<element: string> | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| struct:standard | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 1), ('b', 'x')]]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')]]@struct<a: int8, b: string> | +| struct:overflow | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 300), ('b', 'x')]]@struct<a: int64, b: string> | ERR@ArrowInvalid | +| struct:nullable | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowInvalid | ERR@ArrowInvalid | [[('a', 1), ('b', 'x')], None]@struct<a: int64, b: string> | [[('a', 1), ('b', 'x')], None]@struct<a: int8, b: string> | +| struct<int64>:standard | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [[('a', 1), ('b', 2)]]@map<string, int64> | [[('a', 1), ('b', 2)]]@map<string, int8> | ERR@ArrowTypeError | ERR@ArrowTypeError | +| struct<int64>:overflow | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [[('a', 300), ('b', 2)]]@map<string, int64> | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowInvalid | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.csv new file mode 100644 index 0000000000000..0ef13693498db --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.csv @@ -0,0 +1,20 @@ +source \ target int8 int64 float32 timestamp[us] date32[day] duration[us] time64[ns] string binary +int64:standard ERR@ArrowInvalid [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 ERR@ArrowInvalid [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, 1969-12-31 23:59:59.999999, temporal overflow, temporal overflow]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00, 0:00:00.000001, -1 day, 23:59:59.999999, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] [00:00:00, 00:00:00, 23:59:59.999999, 23:47:16.854775, NaT collision]@time64[ns] ERR@ArrowTypeError [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\xff', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary +float64:standard ERR@ArrowInvalid ERR@ArrowInvalid [0.0, 1.5, -1.5]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'', b'']@binary +bool:standard [1, 0, 1]@int8 [1, 0, 1]@int64 [1.0, 0.0, 1.0]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'\x01', b'', b'\x01']@binary +object:string ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [hello, world, ]@string [b'hello', b'world', b'']@binary +object:bytes ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [hello, world]@string [b'hello', b'world']@binary +date:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-06-15]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:datetime ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [2024-06-15 18:30:00]@timestamp[us] [2024-06-15]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:timedelta ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError [1 day, 2:00:00]@duration[us] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +time:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [18:30:45]@time64[ns] ERR@ArrowTypeError ERR@ArrowTypeError +Int64:standard ERR@ArrowInvalid [0, 1, 9223372036854775807, -9223372036854775808]@int64 ERR@ArrowInvalid [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, temporal overflow, temporal overflow]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00, 0:00:00.000001, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] [00:00:00, 00:00:00, 23:47:16.854775, NaT collision]@time64[ns] ERR@ArrowTypeError [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary +string[pyarrow]:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [hello, world]@string [b'hello', b'world']@binary +large_binary[pyarrow]:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [hello, world]@string [b'hello', b'world']@binary +int64:overflow ERR@ArrowInvalid [300, 1]@int64 [300.0, 1.0]@float32 [1970-01-01 00:00:00.000300, 1970-01-01 00:00:00.000001]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00.000300, 0:00:00.000001]@duration[us] [00:00:00, 00:00:00]@time64[ns] ERR@ArrowTypeError [b',\x01', b'\x01']@binary +float64:fractional ERR@ArrowInvalid ERR@ArrowInvalid [1.5, 2.5]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'']@binary +float64:infinity ERR@ArrowInvalid ERR@ArrowInvalid [inf, 1.0]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'']@binary +float64:precision ERR@ArrowInvalid ERR@ArrowInvalid [1.1234568357467651]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'\x98nt\xd3\xad\xf9\xf1?']@binary +datetime64[ns]:sub-us ERR@ArrowNotImplementedError [1704067200000000123]@int64 ERR@ArrowNotImplementedError ERR@ArrowInvalid [2024-01-01]@date32[day] ERR@ArrowNotImplementedError [00:00:00]@time64[ns] ERR@ArrowTypeError [b'{']@binary +object:date-then-datetime ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-01-01, 2024-01-01]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:datetime-then-date ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-01-01, 2024-01-01]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.md new file mode 100644 index 0000000000000..2157f49343ac9 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_safe.md @@ -0,0 +1,21 @@ +| source \ target | int8 | int64 | float32 | timestamp[us] | date32[day] | duration[us] | time64[ns] | string | binary | +|--------------------------------|------------------------------|-------------------------------------------------------------|------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------|-------------------------|------------------------------------------------------------------------------------------------------| +| int64:standard | ERR@ArrowInvalid | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, 1969-12-31 23:59:59.999999, temporal overflow, temporal overflow]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00.000001, -1 day, 23:59:59.999999, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] | [00:00:00, 00:00:00, 23:59:59.999999, 23:47:16.854775, NaT collision]@time64[ns] | ERR@ArrowTypeError | [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\xff', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary | +| float64:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'', b'']@binary | +| bool:standard | [1, 0, 1]@int8 | [1, 0, 1]@int64 | [1.0, 0.0, 1.0]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'\x01', b'', b'\x01']@binary | +| object:string | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [hello, world, ]@string | [b'hello', b'world', b'']@binary | +| object:bytes | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [hello, world]@string | [b'hello', b'world']@binary | +| date:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-06-15]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:datetime | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [2024-06-15 18:30:00]@timestamp[us] | [2024-06-15]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:timedelta | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | [1 day, 2:00:00]@duration[us] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| time:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [18:30:45]@time64[ns] | ERR@ArrowTypeError | ERR@ArrowTypeError | +| Int64:standard | ERR@ArrowInvalid | [0, 1, 9223372036854775807, -9223372036854775808]@int64 | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, temporal overflow, temporal overflow]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00.000001, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] | [00:00:00, 00:00:00, 23:47:16.854775, NaT collision]@time64[ns] | ERR@ArrowTypeError | [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary | +| string[pyarrow]:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [hello, world]@string | [b'hello', b'world']@binary | +| large_binary[pyarrow]:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [hello, world]@string | [b'hello', b'world']@binary | +| int64:overflow | ERR@ArrowInvalid | [300, 1]@int64 | [300.0, 1.0]@float32 | [1970-01-01 00:00:00.000300, 1970-01-01 00:00:00.000001]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00.000300, 0:00:00.000001]@duration[us] | [00:00:00, 00:00:00]@time64[ns] | ERR@ArrowTypeError | [b',\x01', b'\x01']@binary | +| float64:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | [1.5, 2.5]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'']@binary | +| float64:infinity | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, 1.0]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'']@binary | +| float64:precision | ERR@ArrowInvalid | ERR@ArrowInvalid | [1.1234568357467651]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'\x98nt\xd3\xad\xf9\xf1?']@binary | +| datetime64[ns]:sub-us | ERR@ArrowNotImplementedError | [1704067200000000123]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [2024-01-01]@date32[day] | ERR@ArrowNotImplementedError | [00:00:00]@time64[ns] | ERR@ArrowTypeError | [b'{']@binary | +| object:date-then-datetime | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-01-01, 2024-01-01]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:datetime-then-date | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-01-01, 2024-01-01]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.csv new file mode 100644 index 0000000000000..32a24d6a661cf --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.csv @@ -0,0 +1,20 @@ +source \ target int8 int64 float32 timestamp[us] date32[day] duration[us] time64[ns] string binary +int64:standard [0, 1, -1, -1, 0]@int8 [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 [0.0, 1.0, -1.0, 9.223372036854776e+18, -9.223372036854776e+18]@float32 [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, 1969-12-31 23:59:59.999999, temporal overflow, temporal overflow]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00, 0:00:00.000001, -1 day, 23:59:59.999999, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] [00:00:00, 00:00:00, 23:59:59.999999, 23:47:16.854775, NaT collision]@time64[ns] ERR@ArrowTypeError [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\xff', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary +float64:standard [0, 1, -1]@int8 [0, 1, -1]@int64 [0.0, 1.5, -1.5]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'', b'']@binary +bool:standard [1, 0, 1]@int8 [1, 0, 1]@int64 [1.0, 0.0, 1.0]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'\x01', b'', b'\x01']@binary +object:string ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [hello, world, ]@string [b'hello', b'world', b'']@binary +object:bytes ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [hello, world]@string [b'hello', b'world']@binary +date:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-06-15]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:datetime ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [2024-06-15 18:30:00]@timestamp[us] [2024-06-15]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:timedelta ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError [1 day, 2:00:00]@duration[us] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +time:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError [18:30:45]@time64[ns] ERR@ArrowTypeError ERR@ArrowTypeError +Int64:standard ERR@ArrowInvalid [0, 1, 9223372036854775807, -9223372036854775808]@int64 ERR@ArrowInvalid [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, temporal overflow, temporal overflow]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00, 0:00:00.000001, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] [00:00:00, 00:00:00, 23:47:16.854775, NaT collision]@time64[ns] ERR@ArrowTypeError [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary +string[pyarrow]:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [hello, world]@string [b'hello', b'world']@binary +large_binary[pyarrow]:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [hello, world]@string [b'hello', b'world']@binary +int64:overflow [44, 1]@int8 [300, 1]@int64 [300.0, 1.0]@float32 [1970-01-01 00:00:00.000300, 1970-01-01 00:00:00.000001]@timestamp[us] ERR@ArrowNotImplementedError [0:00:00.000300, 0:00:00.000001]@duration[us] [00:00:00, 00:00:00]@time64[ns] ERR@ArrowTypeError [b',\x01', b'\x01']@binary +float64:fractional [1, 2]@int8 [1, 2]@int64 [1.5, 2.5]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'']@binary +float64:infinity [0, 1]@int8 [-9223372036854775808, 1]@int64 [inf, 1.0]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'', b'']@binary +float64:precision [1]@int8 [1]@int64 [1.1234568357467651]@float32 ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowTypeError [b'\x98nt\xd3\xad\xf9\xf1?']@binary +datetime64[ns]:sub-us ERR@ArrowNotImplementedError [1704067200000000123]@int64 ERR@ArrowNotImplementedError [2024-01-01 00:00:00]@timestamp[us] [2024-01-01]@date32[day] ERR@ArrowNotImplementedError [00:00:00]@time64[ns] ERR@ArrowTypeError [b'{']@binary +object:date-then-datetime ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-01-01, 2024-01-01]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError +object:datetime-then-date ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowTypeError [2024-01-01, 2024-01-01]@date32[day] ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError ERR@ArrowTypeError diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.md new file mode 100644 index 0000000000000..a6f12040dc1de --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_array_from_pandas_type_scalar_unsafe.md @@ -0,0 +1,21 @@ +| source \ target | int8 | int64 | float32 | timestamp[us] | date32[day] | duration[us] | time64[ns] | string | binary | +|--------------------------------|------------------------------|-------------------------------------------------------------|-------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------|-------------------------|------------------------------------------------------------------------------------------------------| +| int64:standard | [0, 1, -1, -1, 0]@int8 | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | [0.0, 1.0, -1.0, 9.223372036854776e+18, -9.223372036854776e+18]@float32 | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, 1969-12-31 23:59:59.999999, temporal overflow, temporal overflow]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00.000001, -1 day, 23:59:59.999999, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] | [00:00:00, 00:00:00, 23:59:59.999999, 23:47:16.854775, NaT collision]@time64[ns] | ERR@ArrowTypeError | [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\xff', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary | +| float64:standard | [0, 1, -1]@int8 | [0, 1, -1]@int64 | [0.0, 1.5, -1.5]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'', b'']@binary | +| bool:standard | [1, 0, 1]@int8 | [1, 0, 1]@int64 | [1.0, 0.0, 1.0]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'\x01', b'', b'\x01']@binary | +| object:string | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [hello, world, ]@string | [b'hello', b'world', b'']@binary | +| object:bytes | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [hello, world]@string | [b'hello', b'world']@binary | +| date:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-06-15]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:datetime | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [2024-06-15 18:30:00]@timestamp[us] | [2024-06-15]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:timedelta | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | [1 day, 2:00:00]@duration[us] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| time:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | [18:30:45]@time64[ns] | ERR@ArrowTypeError | ERR@ArrowTypeError | +| Int64:standard | ERR@ArrowInvalid | [0, 1, 9223372036854775807, -9223372036854775808]@int64 | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, temporal overflow, temporal overflow]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00.000001, 106751991 days, 4:00:54.775807, -106751992 days, 19:59:05.224192]@duration[us] | [00:00:00, 00:00:00, 23:47:16.854775, NaT collision]@time64[ns] | ERR@ArrowTypeError | [b'', b'\x01', b'\xff\xff\xff\xff\xff\xff\xff\x7f', b'']@binary | +| string[pyarrow]:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [hello, world]@string | [b'hello', b'world']@binary | +| large_binary[pyarrow]:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [hello, world]@string | [b'hello', b'world']@binary | +| int64:overflow | [44, 1]@int8 | [300, 1]@int64 | [300.0, 1.0]@float32 | [1970-01-01 00:00:00.000300, 1970-01-01 00:00:00.000001]@timestamp[us] | ERR@ArrowNotImplementedError | [0:00:00.000300, 0:00:00.000001]@duration[us] | [00:00:00, 00:00:00]@time64[ns] | ERR@ArrowTypeError | [b',\x01', b'\x01']@binary | +| float64:fractional | [1, 2]@int8 | [1, 2]@int64 | [1.5, 2.5]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'']@binary | +| float64:infinity | [0, 1]@int8 | [-9223372036854775808, 1]@int64 | [inf, 1.0]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'', b'']@binary | +| float64:precision | [1]@int8 | [1]@int64 | [1.1234568357467651]@float32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowTypeError | [b'\x98nt\xd3\xad\xf9\xf1?']@binary | +| datetime64[ns]:sub-us | ERR@ArrowNotImplementedError | [1704067200000000123]@int64 | ERR@ArrowNotImplementedError | [2024-01-01 00:00:00]@timestamp[us] | [2024-01-01]@date32[day] | ERR@ArrowNotImplementedError | [00:00:00]@time64[ns] | ERR@ArrowTypeError | [b'{']@binary | +| object:date-then-datetime | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-01-01, 2024-01-01]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | +| object:datetime-then-date | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowTypeError | [2024-01-01, 2024-01-01]@date32[day] | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | ERR@ArrowTypeError | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.csv index d4a8df1a2a627..afc22fdfac99e 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.csv +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.csv @@ -1,4 +1,10 @@ test case pyarrow array pandas series pandas series (date_as_object=False) +date32:standard [2024-01-01, 2024-06-15]@date32[day] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] +date32:nullable [2024-01-01, None]@date32[day] [datetime.date(2024, 1, 1), None]@Series[object] [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] +date32:empty []@date32[day] []@Series[object] []@Series[datetime64[ns]] +date64:standard [2024-01-01, 2024-06-15]@date64[ms] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] +date64:nullable [2024-01-01, None]@date64[ms] [datetime.date(2024, 1, 1), None]@Series[object] [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] +date64:empty []@date64[ms] []@Series[object] []@Series[datetime64[ns]] timestamp[s]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] timestamp[s]:nullable [2024-01-01 12:00:00, None]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] timestamp[s]:empty []@timestamp[s] []@Series[datetime64[ns]] []@Series[datetime64[ns]] @@ -14,7 +20,6 @@ timestamp[ns]:empty []@timestamp[ns] []@Series[datetime64[ns]] []@Series[datetim timestamp[us,tz=UTC]:standard [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] timestamp[us,tz=UTC]:nullable [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] timestamp[us,tz=UTC]:empty []@timestamp[us, tz=UTC] []@Series[datetime64[ns, UTC]] []@Series[datetime64[ns, UTC]] -timestamp[s]:overflow [2500-01-01 00:00:00]@timestamp[s] ERR@ArrowInvalid ERR@ArrowInvalid duration[s]:standard [1 day, 0:00:00, 2:30:00]@duration[s] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] duration[s]:nullable [1 day, 0:00:00, None]@duration[s] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] duration[s]:empty []@duration[s] []@Series[timedelta64[ns]] []@Series[timedelta64[ns]] @@ -27,13 +32,6 @@ duration[us]:empty []@duration[us] []@Series[timedelta64[ns]] []@Series[timedelt duration[ns]:standard [1 days 00:00:00, 0 days 02:30:00]@duration[ns] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] duration[ns]:nullable [1 days 00:00:00, None]@duration[ns] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] duration[ns]:empty []@duration[ns] []@Series[timedelta64[ns]] []@Series[timedelta64[ns]] -duration[s]:overflow [109500 days, 0:00:00]@duration[s] [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] -date32:standard [2024-01-01, 2024-06-15]@date32[day] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] -date32:nullable [2024-01-01, None]@date32[day] [datetime.date(2024, 1, 1), None]@Series[object] [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] -date32:empty []@date32[day] []@Series[object] []@Series[datetime64[ns]] -date64:standard [2024-01-01, 2024-06-15]@date64[ms] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] -date64:nullable [2024-01-01, None]@date64[ms] [datetime.date(2024, 1, 1), None]@Series[object] [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] -date64:empty []@date64[ms] []@Series[object] []@Series[datetime64[ns]] time32[s]:standard [12:30:00, 18:45:30]@time32[s] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] time32[s]:nullable [12:30:00, None]@time32[s] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] time32[s]:empty []@time32[s] []@Series[object] []@Series[object] @@ -46,6 +44,10 @@ time64[us]:empty []@time64[us] []@Series[object] []@Series[object] time64[ns]:standard [12:30:00, 18:45:30]@time64[ns] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] time64[ns]:nullable [12:30:00, None]@time64[ns] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] time64[ns]:empty []@time64[ns] []@Series[object] []@Series[object] +timestamp[s]:overflow [2500-01-01 00:00:00]@timestamp[s] ERR@ArrowInvalid ERR@ArrowInvalid +duration[s]:overflow [109500 days, 0:00:00]@duration[s] [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] +timestamp[us]:single-chunk [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] +timestamp[us]:multi-chunk [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] int64:standard [0, 1, -1]@int64 [0, 1, -1]@Series[int64] [0, 1, -1]@Series[int64] int64:nullable [0, 1, None]@int64 [0.0, 1.0, nan]@Series[float64] [0.0, 1.0, nan]@Series[float64] float64:standard [0.0, 1.5, -1.5]@float64 [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@Series[float64] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.md index 45eec337f0593..3e3e500fc52e3 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.md +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_coerce_temporal.md @@ -1,5 +1,11 @@ | test case | pyarrow array | pandas series | pandas series (date_as_object=False) | |-------------------------------|------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------| +| date32:standard | [2024-01-01, 2024-06-15]@date32[day] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] | +| date32:nullable | [2024-01-01, None]@date32[day] | [datetime.date(2024, 1, 1), None]@Series[object] | [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] | +| date32:empty | []@date32[day] | []@Series[object] | []@Series[datetime64[ns]] | +| date64:standard | [2024-01-01, 2024-06-15]@date64[ms] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] | +| date64:nullable | [2024-01-01, None]@date64[ms] | [datetime.date(2024, 1, 1), None]@Series[object] | [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] | +| date64:empty | []@date64[ms] | []@Series[object] | []@Series[datetime64[ns]] | | timestamp[s]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | | timestamp[s]:nullable | [2024-01-01 12:00:00, None]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | | timestamp[s]:empty | []@timestamp[s] | []@Series[datetime64[ns]] | []@Series[datetime64[ns]] | @@ -15,7 +21,6 @@ | timestamp[us,tz=UTC]:standard | [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] | | timestamp[us,tz=UTC]:nullable | [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] | | timestamp[us,tz=UTC]:empty | []@timestamp[us, tz=UTC] | []@Series[datetime64[ns, UTC]] | []@Series[datetime64[ns, UTC]] | -| timestamp[s]:overflow | [2500-01-01 00:00:00]@timestamp[s] | ERR@ArrowInvalid | ERR@ArrowInvalid | | duration[s]:standard | [1 day, 0:00:00, 2:30:00]@duration[s] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | | duration[s]:nullable | [1 day, 0:00:00, None]@duration[s] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | | duration[s]:empty | []@duration[s] | []@Series[timedelta64[ns]] | []@Series[timedelta64[ns]] | @@ -28,13 +33,6 @@ | duration[ns]:standard | [1 days 00:00:00, 0 days 02:30:00]@duration[ns] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | | duration[ns]:nullable | [1 days 00:00:00, None]@duration[ns] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | | duration[ns]:empty | []@duration[ns] | []@Series[timedelta64[ns]] | []@Series[timedelta64[ns]] | -| duration[s]:overflow | [109500 days, 0:00:00]@duration[s] | [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] | [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] | -| date32:standard | [2024-01-01, 2024-06-15]@date32[day] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] | -| date32:nullable | [2024-01-01, None]@date32[day] | [datetime.date(2024, 1, 1), None]@Series[object] | [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] | -| date32:empty | []@date32[day] | []@Series[object] | []@Series[datetime64[ns]] | -| date64:standard | [2024-01-01, 2024-06-15]@date64[ms] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [Timestamp('2024-01-01 00:00:00'), Timestamp('2024-06-15 00:00:00')]@Series[datetime64[ns]] | -| date64:nullable | [2024-01-01, None]@date64[ms] | [datetime.date(2024, 1, 1), None]@Series[object] | [Timestamp('2024-01-01 00:00:00'), NaT]@Series[datetime64[ns]] | -| date64:empty | []@date64[ms] | []@Series[object] | []@Series[datetime64[ns]] | | time32[s]:standard | [12:30:00, 18:45:30]@time32[s] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | | time32[s]:nullable | [12:30:00, None]@time32[s] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | | time32[s]:empty | []@time32[s] | []@Series[object] | []@Series[object] | @@ -47,6 +45,10 @@ | time64[ns]:standard | [12:30:00, 18:45:30]@time64[ns] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | | time64[ns]:nullable | [12:30:00, None]@time64[ns] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | | time64[ns]:empty | []@time64[ns] | []@Series[object] | []@Series[object] | +| timestamp[s]:overflow | [2500-01-01 00:00:00]@timestamp[s] | ERR@ArrowInvalid | ERR@ArrowInvalid | +| duration[s]:overflow | [109500 days, 0:00:00]@duration[s] | [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] | [Timedelta('-104004 days +00:25:26.290448384')]@Series[timedelta64[ns]] | +| timestamp[us]:single-chunk | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | +| timestamp[us]:multi-chunk | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | | int64:standard | [0, 1, -1]@int64 | [0, 1, -1]@Series[int64] | [0, 1, -1]@Series[int64] | | int64:nullable | [0, 1, None]@int64 | [0.0, 1.0, nan]@Series[float64] | [0.0, 1.0, nan]@Series[float64] | | float64:standard | [0.0, 1.5, -1.5]@float64 | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@Series[float64] | diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.csv index 13450e5634edf..5015d7c396f83 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.csv +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.csv @@ -120,3 +120,15 @@ map<string,map<string,int64>>:standard [[('a', [('x', 1)]), ('b', [('y', 2)])], dictionary<int32,string>:standard [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', 'a', 'b']@Series[category] dictionary<int32,string>:nullable [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', nan, 'a']@Series[category] dictionary<int32,string>:empty []@dictionary<values=string, indices=int32, ordered=0> []@Series[category] +int64:zero-chunk []@int64 []@Series[int64] +int64:empty-chunk []@int64 []@Series[int64] +int64:single-chunk [1, 2, 3]@int64 [1, 2, 3]@Series[int64] +int64:multi-chunk [1, 2, 3, 4]@int64 [1, 2, 3, 4]@Series[int64] +int64:multi-chunk-nullable [1, None, 2]@int64 [1.0, nan, 2.0]@Series[float64] +int64:multi-chunk-with-empty [1, 2, 3]@int64 [1, 2, 3]@Series[int64] +float64:multi-chunk [1.5, 2.5, 3.5]@float64 [1.5, 2.5, 3.5]@Series[float64] +string:single-chunk [a, b]@string ['a', 'b']@Series[object] +string:multi-chunk [a, b, c]@string ['a', 'b', 'c']@Series[object] +string:multi-chunk-nullable [a, None, c]@string ['a', None, 'c']@Series[object] +list<int64>:multi-chunk [[1, 2], [3], [4]]@list<item: int64> [array([1, 2]), array([3]), array([4])]@Series[object] +struct:multi-chunk [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.md index 04debdc77e03b..48d561f5963ab 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.md +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_default.md @@ -120,4 +120,16 @@ | map<string,map<string,int64>>:standard | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] | | dictionary<int32,string>:standard | [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', 'a', 'b']@Series[category] | | dictionary<int32,string>:nullable | [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', nan, 'a']@Series[category] | -| dictionary<int32,string>:empty | []@dictionary<values=string, indices=int32, ordered=0> | []@Series[category] | \ No newline at end of file +| dictionary<int32,string>:empty | []@dictionary<values=string, indices=int32, ordered=0> | []@Series[category] | +| int64:zero-chunk | []@int64 | []@Series[int64] | +| int64:empty-chunk | []@int64 | []@Series[int64] | +| int64:single-chunk | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | +| int64:multi-chunk | [1, 2, 3, 4]@int64 | [1, 2, 3, 4]@Series[int64] | +| int64:multi-chunk-nullable | [1, None, 2]@int64 | [1.0, nan, 2.0]@Series[float64] | +| int64:multi-chunk-with-empty | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | +| float64:multi-chunk | [1.5, 2.5, 3.5]@float64 | [1.5, 2.5, 3.5]@Series[float64] | +| string:single-chunk | [a, b]@string | ['a', 'b']@Series[object] | +| string:multi-chunk | [a, b, c]@string | ['a', 'b', 'c']@Series[object] | +| string:multi-chunk-nullable | [a, None, c]@string | ['a', None, 'c']@Series[object] | +| list<int64>:multi-chunk | [[1, 2], [3], [4]]@list<item: int64> | [array([1, 2]), array([3]), array([4])]@Series[object] | +| struct:multi-chunk | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.csv new file mode 100644 index 0000000000000..e9ef08af1bd0f --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.csv @@ -0,0 +1,151 @@ +test case pyarrow array integer_object_nulls=False integer_object_nulls=True spark pandas_options +int8:standard [0, 1, -1, 127, -128]@int8 [0, 1, -1, 127, -128]@Series[int8] [0, 1, -1, 127, -128]@Series[int8] [0, 1, -1, 127, -128]@Series[int8] +int8:nullable [0, 1, None]@int8 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +int8:empty []@int8 []@Series[int8] []@Series[int8] []@Series[int8] +int16:standard [0, 1, -1, 32767, -32768]@int16 [0, 1, -1, 32767, -32768]@Series[int16] [0, 1, -1, 32767, -32768]@Series[int16] [0, 1, -1, 32767, -32768]@Series[int16] +int16:nullable [0, 1, None]@int16 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +int16:empty []@int16 []@Series[int16] []@Series[int16] []@Series[int16] +int32:standard [0, 1, -1, 2147483647, -2147483648]@int32 [0, 1, -1, 2147483647, -2147483648]@Series[int32] [0, 1, -1, 2147483647, -2147483648]@Series[int32] [0, 1, -1, 2147483647, -2147483648]@Series[int32] +int32:nullable [0, 1, None]@int32 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +int32:empty []@int32 []@Series[int32] []@Series[int32] []@Series[int32] +int64:standard [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] +int64:nullable [0, 1, None]@int64 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +int64:empty []@int64 []@Series[int64] []@Series[int64] []@Series[int64] +uint8:standard [0, 1, 255]@uint8 [0, 1, 255]@Series[uint8] [0, 1, 255]@Series[uint8] [0, 1, 255]@Series[uint8] +uint8:nullable [0, 1, None]@uint8 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +uint8:empty []@uint8 []@Series[uint8] []@Series[uint8] []@Series[uint8] +uint16:standard [0, 1, 65535]@uint16 [0, 1, 65535]@Series[uint16] [0, 1, 65535]@Series[uint16] [0, 1, 65535]@Series[uint16] +uint16:nullable [0, 1, None]@uint16 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +uint16:empty []@uint16 []@Series[uint16] []@Series[uint16] []@Series[uint16] +uint32:standard [0, 1, 4294967295]@uint32 [0, 1, 4294967295]@Series[uint32] [0, 1, 4294967295]@Series[uint32] [0, 1, 4294967295]@Series[uint32] +uint32:nullable [0, 1, None]@uint32 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +uint32:empty []@uint32 []@Series[uint32] []@Series[uint32] []@Series[uint32] +uint64:standard [0, 1, 18446744073709551615]@uint64 [0, 1, 18446744073709551615]@Series[uint64] [0, 1, 18446744073709551615]@Series[uint64] [0, 1, 18446744073709551615]@Series[uint64] +uint64:nullable [0, 1, None]@uint64 [0.0, 1.0, nan]@Series[float64] [0, 1, None]@Series[object] [0, 1, None]@Series[object] +uint64:empty []@uint64 []@Series[uint64] []@Series[uint64] []@Series[uint64] +float32:standard [0.0, 1.5, -1.5]@float32 [0.0, 1.5, -1.5]@Series[float32] [0.0, 1.5, -1.5]@Series[float32] [0.0, 1.5, -1.5]@Series[float32] +float32:nullable [0.0, 1.5, None]@float32 [0.0, 1.5, nan]@Series[float32] [0.0, 1.5, nan]@Series[float32] [0.0, 1.5, nan]@Series[float32] +float32:empty []@float32 []@Series[float32] []@Series[float32] []@Series[float32] +float64:standard [0.0, 1.5, -1.5]@float64 [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@Series[float64] +float64:nullable [0.0, 1.5, None]@float64 [0.0, 1.5, nan]@Series[float64] [0.0, 1.5, nan]@Series[float64] [0.0, 1.5, nan]@Series[float64] +float64:special [nan, inf, -inf]@float64 [nan, inf, -inf]@Series[float64] [nan, inf, -inf]@Series[float64] [nan, inf, -inf]@Series[float64] +float64:empty []@float64 []@Series[float64] []@Series[float64] []@Series[float64] +bool:standard [True, False, True]@bool [True, False, True]@Series[bool] [True, False, True]@Series[bool] [True, False, True]@Series[bool] +bool:nullable [True, False, None]@bool [True, False, None]@Series[object] [True, False, None]@Series[object] [True, False, None]@Series[object] +bool:empty []@bool []@Series[bool] []@Series[bool] []@Series[bool] +string:standard [hello, world, ]@string ['hello', 'world', '']@Series[object] ['hello', 'world', '']@Series[object] ['hello', 'world', '']@Series[object] +string:nullable [hello, None, world]@string ['hello', None, 'world']@Series[object] ['hello', None, 'world']@Series[object] ['hello', None, 'world']@Series[object] +string:empty []@string []@Series[object] []@Series[object] []@Series[object] +large_string:standard [hello, world]@large_string ['hello', 'world']@Series[object] ['hello', 'world']@Series[object] ['hello', 'world']@Series[object] +large_string:nullable [hello, None]@large_string ['hello', None]@Series[object] ['hello', None]@Series[object] ['hello', None]@Series[object] +large_string:empty []@large_string []@Series[object] []@Series[object] []@Series[object] +binary:standard [b'hello', b'world']@binary [b'hello', b'world']@Series[object] [b'hello', b'world']@Series[object] [b'hello', b'world']@Series[object] +binary:nullable [b'hello', None]@binary [b'hello', None]@Series[object] [b'hello', None]@Series[object] [b'hello', None]@Series[object] +binary:empty []@binary []@Series[object] []@Series[object] []@Series[object] +large_binary:standard [b'hello', b'world']@large_binary [b'hello', b'world']@Series[object] [b'hello', b'world']@Series[object] [b'hello', b'world']@Series[object] +large_binary:nullable [b'hello', None]@large_binary [b'hello', None]@Series[object] [b'hello', None]@Series[object] [b'hello', None]@Series[object] +large_binary:empty []@large_binary []@Series[object] []@Series[object] []@Series[object] +decimal128:standard [1.23, 4.56, -7.89]@decimal128(5, 2) [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] +decimal128:nullable [1.23, None, 4.56]@decimal128(5, 2) [Decimal('1.23'), None, Decimal('4.56')]@Series[object] [Decimal('1.23'), None, Decimal('4.56')]@Series[object] [Decimal('1.23'), None, Decimal('4.56')]@Series[object] +decimal128:empty []@decimal128(5, 2) []@Series[object] []@Series[object] []@Series[object] +date32:standard [2024-01-01, 2024-06-15]@date32[day] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] +date32:nullable [2024-01-01, None]@date32[day] [datetime.date(2024, 1, 1), None]@Series[object] [datetime.date(2024, 1, 1), None]@Series[object] [datetime.date(2024, 1, 1), None]@Series[object] +date32:empty []@date32[day] []@Series[object] []@Series[object] []@Series[object] +date64:standard [2024-01-01, 2024-06-15]@date64[ms] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] +date64:nullable [2024-01-01, None]@date64[ms] [datetime.date(2024, 1, 1), None]@Series[object] [datetime.date(2024, 1, 1), None]@Series[object] [datetime.date(2024, 1, 1), None]@Series[object] +date64:empty []@date64[ms] []@Series[object] []@Series[object] []@Series[object] +timestamp[s]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] +timestamp[s]:nullable [2024-01-01 12:00:00, None]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] +timestamp[s]:empty []@timestamp[s] []@Series[datetime64[s]] []@Series[datetime64[s]] []@Series[datetime64[ns]] +timestamp[ms]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] +timestamp[ms]:nullable [2024-01-01 12:00:00, None]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] +timestamp[ms]:empty []@timestamp[ms] []@Series[datetime64[ms]] []@Series[datetime64[ms]] []@Series[datetime64[ns]] +timestamp[us]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] +timestamp[us]:nullable [2024-01-01 12:00:00, None]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] +timestamp[us]:empty []@timestamp[us] []@Series[datetime64[us]] []@Series[datetime64[us]] []@Series[datetime64[ns]] +timestamp[ns]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] +timestamp[ns]:nullable [2024-01-01 12:00:00, None]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] +timestamp[ns]:empty []@timestamp[ns] []@Series[datetime64[ns]] []@Series[datetime64[ns]] []@Series[datetime64[ns]] +timestamp[us,tz=UTC]:standard [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] +timestamp[us,tz=UTC]:nullable [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] +timestamp[us,tz=UTC]:empty []@timestamp[us, tz=UTC] []@Series[datetime64[us, UTC]] []@Series[datetime64[us, UTC]] []@Series[datetime64[ns, UTC]] +duration[s]:standard [1 day, 0:00:00, 2:30:00]@duration[s] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] +duration[s]:nullable [1 day, 0:00:00, None]@duration[s] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] +duration[s]:empty []@duration[s] []@Series[timedelta64[s]] []@Series[timedelta64[s]] []@Series[timedelta64[ns]] +duration[ms]:standard [1 day, 0:00:00, 2:30:00]@duration[ms] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] +duration[ms]:nullable [1 day, 0:00:00, None]@duration[ms] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] +duration[ms]:empty []@duration[ms] []@Series[timedelta64[ms]] []@Series[timedelta64[ms]] []@Series[timedelta64[ns]] +duration[us]:standard [1 day, 0:00:00, 2:30:00]@duration[us] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] +duration[us]:nullable [1 day, 0:00:00, None]@duration[us] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] +duration[us]:empty []@duration[us] []@Series[timedelta64[us]] []@Series[timedelta64[us]] []@Series[timedelta64[ns]] +duration[ns]:standard [1 days 00:00:00, 0 days 02:30:00]@duration[ns] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] +duration[ns]:nullable [1 days 00:00:00, None]@duration[ns] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] +duration[ns]:empty []@duration[ns] []@Series[timedelta64[ns]] []@Series[timedelta64[ns]] []@Series[timedelta64[ns]] +time32[s]:standard [12:30:00, 18:45:30]@time32[s] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] +time32[s]:nullable [12:30:00, None]@time32[s] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] +time32[s]:empty []@time32[s] []@Series[object] []@Series[object] []@Series[object] +time32[ms]:standard [12:30:00, 18:45:30]@time32[ms] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] +time32[ms]:nullable [12:30:00, None]@time32[ms] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] +time32[ms]:empty []@time32[ms] []@Series[object] []@Series[object] []@Series[object] +time64[us]:standard [12:30:00, 18:45:30]@time64[us] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] +time64[us]:nullable [12:30:00, None]@time64[us] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] +time64[us]:empty []@time64[us] []@Series[object] []@Series[object] []@Series[object] +time64[ns]:standard [12:30:00, 18:45:30]@time64[ns] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] +time64[ns]:nullable [12:30:00, None]@time64[ns] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] [datetime.time(12, 30), None]@Series[object] +time64[ns]:empty []@time64[ns] []@Series[object] []@Series[object] []@Series[object] +null:standard [None, None, None]@null [None, None, None]@Series[object] [None, None, None]@Series[object] [None, None, None]@Series[object] +null:empty []@null []@Series[object] []@Series[object] []@Series[object] +list<int64>:standard [[1, 2], [3, 4, 5]]@list<item: int64> [array([1, 2]), array([3, 4, 5])]@Series[object] [array([1, 2]), array([3, 4, 5])]@Series[object] [array([1, 2]), array([3, 4, 5])]@Series[object] +list<int64>:nullable [[1, 2], None, [3]]@list<item: int64> [array([1, 2]), None, array([3])]@Series[object] [array([1, 2]), None, array([3])]@Series[object] [array([1, 2]), None, array([3])]@Series[object] +list<int64>:empty []@list<item: int64> []@Series[object] []@Series[object] []@Series[object] +list<string>:standard [['a', 'b'], ['c']]@list<item: string> [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] +large_list<int64>:standard [[1, 2], [3, 4]]@large_list<item: int64> [array([1, 2]), array([3, 4])]@Series[object] [array([1, 2]), array([3, 4])]@Series[object] [array([1, 2]), array([3, 4])]@Series[object] +large_list<int64>:empty []@large_list<item: int64> []@Series[object] []@Series[object] []@Series[object] +fixed_size_list<int64>[3]:standard [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] [array([1, 2, 3]), array([4, 5, 6])]@Series[object] [array([1, 2, 3]), array([4, 5, 6])]@Series[object] [array([1, 2, 3]), array([4, 5, 6])]@Series[object] +fixed_size_list<int64>[3]:empty []@fixed_size_list<item: int64>[3] []@Series[object] []@Series[object] []@Series[object] +struct:standard [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] +struct:nullable [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, None]@Series[object] [{'x': 1, 'y': 'a'}, None]@Series[object] [{'x': 1, 'y': 'a'}, None]@Series[object] +struct:empty []@struct<x: int64, y: string> []@Series[object] []@Series[object] []@Series[object] +map<string,int64>:standard [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] +map<string,int64>:empty []@map<string, int64> []@Series[object] []@Series[object] []@Series[object] +list<list<int64>>:standard [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] +list<struct>:standard [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] +list<map<string,int64>>:standard [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] +struct<struct>:standard [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] +struct<list<int64>>:standard [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] +struct<map<string,int64>>:standard [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] +map<string,list<int64>>:standard [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] +map<string,struct>:standard [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] +map<string,map<string,int64>>:standard [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] +dictionary<int32,string>:standard [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', 'a', 'b']@Series[category] ['a', 'b', 'a', 'b']@Series[category] ['a', 'b', 'a', 'b']@Series[category] +dictionary<int32,string>:nullable [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', nan, 'a']@Series[category] ['a', 'b', nan, 'a']@Series[category] ['a', 'b', nan, 'a']@Series[category] +dictionary<int32,string>:empty []@dictionary<values=string, indices=int32, ordered=0> []@Series[category] []@Series[category] []@Series[category] +int64:zero-chunk []@int64 []@Series[int64] []@Series[int64] []@Series[int64] +int64:empty-chunk []@int64 []@Series[int64] []@Series[int64] []@Series[int64] +int64:single-chunk [1, 2, 3]@int64 [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] +int64:multi-chunk [1, 2, 3, 4]@int64 [1, 2, 3, 4]@Series[int64] [1, 2, 3, 4]@Series[int64] [1, 2, 3, 4]@Series[int64] +int64:multi-chunk-nullable [1, None, 2]@int64 [1.0, nan, 2.0]@Series[float64] [1, None, 2]@Series[object] [1, None, 2]@Series[object] +int64:multi-chunk-with-empty [1, 2, 3]@int64 [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] +float64:multi-chunk [1.5, 2.5, 3.5]@float64 [1.5, 2.5, 3.5]@Series[float64] [1.5, 2.5, 3.5]@Series[float64] [1.5, 2.5, 3.5]@Series[float64] +string:single-chunk [a, b]@string ['a', 'b']@Series[object] ['a', 'b']@Series[object] ['a', 'b']@Series[object] +string:multi-chunk [a, b, c]@string ['a', 'b', 'c']@Series[object] ['a', 'b', 'c']@Series[object] ['a', 'b', 'c']@Series[object] +string:multi-chunk-nullable [a, None, c]@string ['a', None, 'c']@Series[object] ['a', None, 'c']@Series[object] ['a', None, 'c']@Series[object] +list<int64>:multi-chunk [[1, 2], [3], [4]]@list<item: int64> [array([1, 2]), array([3]), array([4])]@Series[object] [array([1, 2]), array([3]), array([4])]@Series[object] [array([1, 2]), array([3]), array([4])]@Series[object] +struct:multi-chunk [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] +int8:extremes-nullable [127, -128, None]@int8 [127.0, -128.0, nan]@Series[float64] [127, -128, None]@Series[object] [127, -128, None]@Series[object] +int16:extremes-nullable [32767, -32768, None]@int16 [32767.0, -32768.0, nan]@Series[float64] [32767, -32768, None]@Series[object] [32767, -32768, None]@Series[object] +int32:extremes-nullable [2147483647, -2147483648, None]@int32 [2147483647.0, -2147483648.0, nan]@Series[float64] [2147483647, -2147483648, None]@Series[object] [2147483647, -2147483648, None]@Series[object] +int64:extremes-nullable [9223372036854775807, -9223372036854775808, None]@int64 [9.223372036854776e+18, -9.223372036854776e+18, nan]@Series[float64] [9223372036854775807, -9223372036854775808, None]@Series[object] [9223372036854775807, -9223372036854775808, None]@Series[object] +uint8:extremes-nullable [255, 0, None]@uint8 [255.0, 0.0, nan]@Series[float64] [255, 0, None]@Series[object] [255, 0, None]@Series[object] +uint16:extremes-nullable [65535, 0, None]@uint16 [65535.0, 0.0, nan]@Series[float64] [65535, 0, None]@Series[object] [65535, 0, None]@Series[object] +uint32:extremes-nullable [4294967295, 0, None]@uint32 [4294967295.0, 0.0, nan]@Series[float64] [4294967295, 0, None]@Series[object] [4294967295, 0, None]@Series[object] +uint64:extremes-nullable [18446744073709551615, 0, None]@uint64 [1.8446744073709552e+19, 0.0, nan]@Series[float64] [18446744073709551615, 0, None]@Series[object] [18446744073709551615, 0, None]@Series[object] +int64:all-null [None, None]@int64 [nan, nan]@Series[float64] [None, None]@Series[object] [None, None]@Series[object] +list<int64>:null-element [[1, None], [2, 3]]@list<item: int64> [array([ 1., nan]), array([2., 3.])]@Series[object] [array([1, None], dtype=object), array([2, 3], dtype=object)]@Series[object] [array([1, None], dtype=object), array([2, 3], dtype=object)]@Series[object] +list<int64>:null-element-extreme [[9223372036854775807, None]]@list<item: int64> [array([9.22337204e+18, nan])]@Series[object] [array([9223372036854775807, None], dtype=object)]@Series[object] [array([9223372036854775807, None], dtype=object)]@Series[object] +large_list<int64>:null-element [[1, None]]@large_list<item: int64> [array([ 1., nan])]@Series[object] [array([1, None], dtype=object)]@Series[object] [array([1, None], dtype=object)]@Series[object] +fixed_size_list<int64>[3]:null-element [[1, None, 3]]@fixed_size_list<item: int64>[3] [array([ 1., nan, 3.])]@Series[object] [array([1, None, 3], dtype=object)]@Series[object] [array([1, None, 3], dtype=object)]@Series[object] +list<list<int64>>:null-element [[[1, None], [2]]]@list<item: list<item: int64>> [array([array([ 1., nan]), array([2.])], dtype=object)]@Series[object] [array([array([1, None], dtype=object), array([2], dtype=object)], dtype=object)]@Series[object] [array([array([1, None], dtype=object), array([2], dtype=object)], dtype=object)]@Series[object] +struct:null-int-field [[('x', 1), ('y', 'a')], [('x', None), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1.0, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] +map<string,int64>:null-value [[('a', 1), ('b', None)]]@map<string, int64> [[('a', 1.0), ('b', None)]]@Series[object] [[('a', 1), ('b', None)]]@Series[object] [[('a', 1), ('b', None)]]@Series[object] +dictionary<int64>:nullable [1, None, 1]@dictionary<values=int64, indices=int32, ordered=0> [1, nan, 1]@Series[category] [1, nan, 1]@Series[category] [1, nan, 1]@Series[category] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.md new file mode 100644 index 0000000000000..661016dc3df10 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_integer_object_nulls.md @@ -0,0 +1,152 @@ +| test case | pyarrow array | integer_object_nulls=False | integer_object_nulls=True | spark pandas_options | +|----------------------------------------|-----------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------| +| int8:standard | [0, 1, -1, 127, -128]@int8 | [0, 1, -1, 127, -128]@Series[int8] | [0, 1, -1, 127, -128]@Series[int8] | [0, 1, -1, 127, -128]@Series[int8] | +| int8:nullable | [0, 1, None]@int8 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| int8:empty | []@int8 | []@Series[int8] | []@Series[int8] | []@Series[int8] | +| int16:standard | [0, 1, -1, 32767, -32768]@int16 | [0, 1, -1, 32767, -32768]@Series[int16] | [0, 1, -1, 32767, -32768]@Series[int16] | [0, 1, -1, 32767, -32768]@Series[int16] | +| int16:nullable | [0, 1, None]@int16 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| int16:empty | []@int16 | []@Series[int16] | []@Series[int16] | []@Series[int16] | +| int32:standard | [0, 1, -1, 2147483647, -2147483648]@int32 | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | +| int32:nullable | [0, 1, None]@int32 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| int32:empty | []@int32 | []@Series[int32] | []@Series[int32] | []@Series[int32] | +| int64:standard | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | +| int64:nullable | [0, 1, None]@int64 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| int64:empty | []@int64 | []@Series[int64] | []@Series[int64] | []@Series[int64] | +| uint8:standard | [0, 1, 255]@uint8 | [0, 1, 255]@Series[uint8] | [0, 1, 255]@Series[uint8] | [0, 1, 255]@Series[uint8] | +| uint8:nullable | [0, 1, None]@uint8 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| uint8:empty | []@uint8 | []@Series[uint8] | []@Series[uint8] | []@Series[uint8] | +| uint16:standard | [0, 1, 65535]@uint16 | [0, 1, 65535]@Series[uint16] | [0, 1, 65535]@Series[uint16] | [0, 1, 65535]@Series[uint16] | +| uint16:nullable | [0, 1, None]@uint16 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| uint16:empty | []@uint16 | []@Series[uint16] | []@Series[uint16] | []@Series[uint16] | +| uint32:standard | [0, 1, 4294967295]@uint32 | [0, 1, 4294967295]@Series[uint32] | [0, 1, 4294967295]@Series[uint32] | [0, 1, 4294967295]@Series[uint32] | +| uint32:nullable | [0, 1, None]@uint32 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| uint32:empty | []@uint32 | []@Series[uint32] | []@Series[uint32] | []@Series[uint32] | +| uint64:standard | [0, 1, 18446744073709551615]@uint64 | [0, 1, 18446744073709551615]@Series[uint64] | [0, 1, 18446744073709551615]@Series[uint64] | [0, 1, 18446744073709551615]@Series[uint64] | +| uint64:nullable | [0, 1, None]@uint64 | [0.0, 1.0, nan]@Series[float64] | [0, 1, None]@Series[object] | [0, 1, None]@Series[object] | +| uint64:empty | []@uint64 | []@Series[uint64] | []@Series[uint64] | []@Series[uint64] | +| float32:standard | [0.0, 1.5, -1.5]@float32 | [0.0, 1.5, -1.5]@Series[float32] | [0.0, 1.5, -1.5]@Series[float32] | [0.0, 1.5, -1.5]@Series[float32] | +| float32:nullable | [0.0, 1.5, None]@float32 | [0.0, 1.5, nan]@Series[float32] | [0.0, 1.5, nan]@Series[float32] | [0.0, 1.5, nan]@Series[float32] | +| float32:empty | []@float32 | []@Series[float32] | []@Series[float32] | []@Series[float32] | +| float64:standard | [0.0, 1.5, -1.5]@float64 | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@Series[float64] | +| float64:nullable | [0.0, 1.5, None]@float64 | [0.0, 1.5, nan]@Series[float64] | [0.0, 1.5, nan]@Series[float64] | [0.0, 1.5, nan]@Series[float64] | +| float64:special | [nan, inf, -inf]@float64 | [nan, inf, -inf]@Series[float64] | [nan, inf, -inf]@Series[float64] | [nan, inf, -inf]@Series[float64] | +| float64:empty | []@float64 | []@Series[float64] | []@Series[float64] | []@Series[float64] | +| bool:standard | [True, False, True]@bool | [True, False, True]@Series[bool] | [True, False, True]@Series[bool] | [True, False, True]@Series[bool] | +| bool:nullable | [True, False, None]@bool | [True, False, None]@Series[object] | [True, False, None]@Series[object] | [True, False, None]@Series[object] | +| bool:empty | []@bool | []@Series[bool] | []@Series[bool] | []@Series[bool] | +| string:standard | [hello, world, ]@string | ['hello', 'world', '']@Series[object] | ['hello', 'world', '']@Series[object] | ['hello', 'world', '']@Series[object] | +| string:nullable | [hello, None, world]@string | ['hello', None, 'world']@Series[object] | ['hello', None, 'world']@Series[object] | ['hello', None, 'world']@Series[object] | +| string:empty | []@string | []@Series[object] | []@Series[object] | []@Series[object] | +| large_string:standard | [hello, world]@large_string | ['hello', 'world']@Series[object] | ['hello', 'world']@Series[object] | ['hello', 'world']@Series[object] | +| large_string:nullable | [hello, None]@large_string | ['hello', None]@Series[object] | ['hello', None]@Series[object] | ['hello', None]@Series[object] | +| large_string:empty | []@large_string | []@Series[object] | []@Series[object] | []@Series[object] | +| binary:standard | [b'hello', b'world']@binary | [b'hello', b'world']@Series[object] | [b'hello', b'world']@Series[object] | [b'hello', b'world']@Series[object] | +| binary:nullable | [b'hello', None]@binary | [b'hello', None]@Series[object] | [b'hello', None]@Series[object] | [b'hello', None]@Series[object] | +| binary:empty | []@binary | []@Series[object] | []@Series[object] | []@Series[object] | +| large_binary:standard | [b'hello', b'world']@large_binary | [b'hello', b'world']@Series[object] | [b'hello', b'world']@Series[object] | [b'hello', b'world']@Series[object] | +| large_binary:nullable | [b'hello', None]@large_binary | [b'hello', None]@Series[object] | [b'hello', None]@Series[object] | [b'hello', None]@Series[object] | +| large_binary:empty | []@large_binary | []@Series[object] | []@Series[object] | []@Series[object] | +| decimal128:standard | [1.23, 4.56, -7.89]@decimal128(5, 2) | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] | +| decimal128:nullable | [1.23, None, 4.56]@decimal128(5, 2) | [Decimal('1.23'), None, Decimal('4.56')]@Series[object] | [Decimal('1.23'), None, Decimal('4.56')]@Series[object] | [Decimal('1.23'), None, Decimal('4.56')]@Series[object] | +| decimal128:empty | []@decimal128(5, 2) | []@Series[object] | []@Series[object] | []@Series[object] | +| date32:standard | [2024-01-01, 2024-06-15]@date32[day] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | +| date32:nullable | [2024-01-01, None]@date32[day] | [datetime.date(2024, 1, 1), None]@Series[object] | [datetime.date(2024, 1, 1), None]@Series[object] | [datetime.date(2024, 1, 1), None]@Series[object] | +| date32:empty | []@date32[day] | []@Series[object] | []@Series[object] | []@Series[object] | +| date64:standard | [2024-01-01, 2024-06-15]@date64[ms] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | +| date64:nullable | [2024-01-01, None]@date64[ms] | [datetime.date(2024, 1, 1), None]@Series[object] | [datetime.date(2024, 1, 1), None]@Series[object] | [datetime.date(2024, 1, 1), None]@Series[object] | +| date64:empty | []@date64[ms] | []@Series[object] | []@Series[object] | []@Series[object] | +| timestamp[s]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | +| timestamp[s]:nullable | [2024-01-01 12:00:00, None]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | +| timestamp[s]:empty | []@timestamp[s] | []@Series[datetime64[s]] | []@Series[datetime64[s]] | []@Series[datetime64[ns]] | +| timestamp[ms]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | +| timestamp[ms]:nullable | [2024-01-01 12:00:00, None]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | +| timestamp[ms]:empty | []@timestamp[ms] | []@Series[datetime64[ms]] | []@Series[datetime64[ms]] | []@Series[datetime64[ns]] | +| timestamp[us]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | +| timestamp[us]:nullable | [2024-01-01 12:00:00, None]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | +| timestamp[us]:empty | []@timestamp[us] | []@Series[datetime64[us]] | []@Series[datetime64[us]] | []@Series[datetime64[ns]] | +| timestamp[ns]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | +| timestamp[ns]:nullable | [2024-01-01 12:00:00, None]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | +| timestamp[ns]:empty | []@timestamp[ns] | []@Series[datetime64[ns]] | []@Series[datetime64[ns]] | []@Series[datetime64[ns]] | +| timestamp[us,tz=UTC]:standard | [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[ns, UTC]] | +| timestamp[us,tz=UTC]:nullable | [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[ns, UTC]] | +| timestamp[us,tz=UTC]:empty | []@timestamp[us, tz=UTC] | []@Series[datetime64[us, UTC]] | []@Series[datetime64[us, UTC]] | []@Series[datetime64[ns, UTC]] | +| duration[s]:standard | [1 day, 0:00:00, 2:30:00]@duration[s] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | +| duration[s]:nullable | [1 day, 0:00:00, None]@duration[s] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | +| duration[s]:empty | []@duration[s] | []@Series[timedelta64[s]] | []@Series[timedelta64[s]] | []@Series[timedelta64[ns]] | +| duration[ms]:standard | [1 day, 0:00:00, 2:30:00]@duration[ms] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | +| duration[ms]:nullable | [1 day, 0:00:00, None]@duration[ms] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | +| duration[ms]:empty | []@duration[ms] | []@Series[timedelta64[ms]] | []@Series[timedelta64[ms]] | []@Series[timedelta64[ns]] | +| duration[us]:standard | [1 day, 0:00:00, 2:30:00]@duration[us] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | +| duration[us]:nullable | [1 day, 0:00:00, None]@duration[us] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | +| duration[us]:empty | []@duration[us] | []@Series[timedelta64[us]] | []@Series[timedelta64[us]] | []@Series[timedelta64[ns]] | +| duration[ns]:standard | [1 days 00:00:00, 0 days 02:30:00]@duration[ns] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | +| duration[ns]:nullable | [1 days 00:00:00, None]@duration[ns] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | +| duration[ns]:empty | []@duration[ns] | []@Series[timedelta64[ns]] | []@Series[timedelta64[ns]] | []@Series[timedelta64[ns]] | +| time32[s]:standard | [12:30:00, 18:45:30]@time32[s] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | +| time32[s]:nullable | [12:30:00, None]@time32[s] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | +| time32[s]:empty | []@time32[s] | []@Series[object] | []@Series[object] | []@Series[object] | +| time32[ms]:standard | [12:30:00, 18:45:30]@time32[ms] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | +| time32[ms]:nullable | [12:30:00, None]@time32[ms] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | +| time32[ms]:empty | []@time32[ms] | []@Series[object] | []@Series[object] | []@Series[object] | +| time64[us]:standard | [12:30:00, 18:45:30]@time64[us] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | +| time64[us]:nullable | [12:30:00, None]@time64[us] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | +| time64[us]:empty | []@time64[us] | []@Series[object] | []@Series[object] | []@Series[object] | +| time64[ns]:standard | [12:30:00, 18:45:30]@time64[ns] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | +| time64[ns]:nullable | [12:30:00, None]@time64[ns] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | [datetime.time(12, 30), None]@Series[object] | +| time64[ns]:empty | []@time64[ns] | []@Series[object] | []@Series[object] | []@Series[object] | +| null:standard | [None, None, None]@null | [None, None, None]@Series[object] | [None, None, None]@Series[object] | [None, None, None]@Series[object] | +| null:empty | []@null | []@Series[object] | []@Series[object] | []@Series[object] | +| list<int64>:standard | [[1, 2], [3, 4, 5]]@list<item: int64> | [array([1, 2]), array([3, 4, 5])]@Series[object] | [array([1, 2]), array([3, 4, 5])]@Series[object] | [array([1, 2]), array([3, 4, 5])]@Series[object] | +| list<int64>:nullable | [[1, 2], None, [3]]@list<item: int64> | [array([1, 2]), None, array([3])]@Series[object] | [array([1, 2]), None, array([3])]@Series[object] | [array([1, 2]), None, array([3])]@Series[object] | +| list<int64>:empty | []@list<item: int64> | []@Series[object] | []@Series[object] | []@Series[object] | +| list<string>:standard | [['a', 'b'], ['c']]@list<item: string> | [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] | [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] | [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] | +| large_list<int64>:standard | [[1, 2], [3, 4]]@large_list<item: int64> | [array([1, 2]), array([3, 4])]@Series[object] | [array([1, 2]), array([3, 4])]@Series[object] | [array([1, 2]), array([3, 4])]@Series[object] | +| large_list<int64>:empty | []@large_list<item: int64> | []@Series[object] | []@Series[object] | []@Series[object] | +| fixed_size_list<int64>[3]:standard | [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] | [array([1, 2, 3]), array([4, 5, 6])]@Series[object] | [array([1, 2, 3]), array([4, 5, 6])]@Series[object] | [array([1, 2, 3]), array([4, 5, 6])]@Series[object] | +| fixed_size_list<int64>[3]:empty | []@fixed_size_list<item: int64>[3] | []@Series[object] | []@Series[object] | []@Series[object] | +| struct:standard | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | +| struct:nullable | [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, None]@Series[object] | [{'x': 1, 'y': 'a'}, None]@Series[object] | [{'x': 1, 'y': 'a'}, None]@Series[object] | +| struct:empty | []@struct<x: int64, y: string> | []@Series[object] | []@Series[object] | []@Series[object] | +| map<string,int64>:standard | [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] | +| map<string,int64>:empty | []@map<string, int64> | []@Series[object] | []@Series[object] | []@Series[object] | +| list<list<int64>>:standard | [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> | [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] | [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] | [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] | +| list<struct>:standard | [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> | [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] | [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] | [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] | +| list<map<string,int64>>:standard | [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> | [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] | [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] | [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] | +| struct<struct>:standard | [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] | +| struct<list<int64>>:standard | [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> | [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] | [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] | [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] | +| struct<map<string,int64>>:standard | [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] | +| map<string,list<int64>>:standard | [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> | [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] | [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] | [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] | +| map<string,struct>:standard | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] | +| map<string,map<string,int64>>:standard | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] | +| dictionary<int32,string>:standard | [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', 'a', 'b']@Series[category] | ['a', 'b', 'a', 'b']@Series[category] | ['a', 'b', 'a', 'b']@Series[category] | +| dictionary<int32,string>:nullable | [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', nan, 'a']@Series[category] | ['a', 'b', nan, 'a']@Series[category] | ['a', 'b', nan, 'a']@Series[category] | +| dictionary<int32,string>:empty | []@dictionary<values=string, indices=int32, ordered=0> | []@Series[category] | []@Series[category] | []@Series[category] | +| int64:zero-chunk | []@int64 | []@Series[int64] | []@Series[int64] | []@Series[int64] | +| int64:empty-chunk | []@int64 | []@Series[int64] | []@Series[int64] | []@Series[int64] | +| int64:single-chunk | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | +| int64:multi-chunk | [1, 2, 3, 4]@int64 | [1, 2, 3, 4]@Series[int64] | [1, 2, 3, 4]@Series[int64] | [1, 2, 3, 4]@Series[int64] | +| int64:multi-chunk-nullable | [1, None, 2]@int64 | [1.0, nan, 2.0]@Series[float64] | [1, None, 2]@Series[object] | [1, None, 2]@Series[object] | +| int64:multi-chunk-with-empty | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | +| float64:multi-chunk | [1.5, 2.5, 3.5]@float64 | [1.5, 2.5, 3.5]@Series[float64] | [1.5, 2.5, 3.5]@Series[float64] | [1.5, 2.5, 3.5]@Series[float64] | +| string:single-chunk | [a, b]@string | ['a', 'b']@Series[object] | ['a', 'b']@Series[object] | ['a', 'b']@Series[object] | +| string:multi-chunk | [a, b, c]@string | ['a', 'b', 'c']@Series[object] | ['a', 'b', 'c']@Series[object] | ['a', 'b', 'c']@Series[object] | +| string:multi-chunk-nullable | [a, None, c]@string | ['a', None, 'c']@Series[object] | ['a', None, 'c']@Series[object] | ['a', None, 'c']@Series[object] | +| list<int64>:multi-chunk | [[1, 2], [3], [4]]@list<item: int64> | [array([1, 2]), array([3]), array([4])]@Series[object] | [array([1, 2]), array([3]), array([4])]@Series[object] | [array([1, 2]), array([3]), array([4])]@Series[object] | +| struct:multi-chunk | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | +| int8:extremes-nullable | [127, -128, None]@int8 | [127.0, -128.0, nan]@Series[float64] | [127, -128, None]@Series[object] | [127, -128, None]@Series[object] | +| int16:extremes-nullable | [32767, -32768, None]@int16 | [32767.0, -32768.0, nan]@Series[float64] | [32767, -32768, None]@Series[object] | [32767, -32768, None]@Series[object] | +| int32:extremes-nullable | [2147483647, -2147483648, None]@int32 | [2147483647.0, -2147483648.0, nan]@Series[float64] | [2147483647, -2147483648, None]@Series[object] | [2147483647, -2147483648, None]@Series[object] | +| int64:extremes-nullable | [9223372036854775807, -9223372036854775808, None]@int64 | [9.223372036854776e+18, -9.223372036854776e+18, nan]@Series[float64] | [9223372036854775807, -9223372036854775808, None]@Series[object] | [9223372036854775807, -9223372036854775808, None]@Series[object] | +| uint8:extremes-nullable | [255, 0, None]@uint8 | [255.0, 0.0, nan]@Series[float64] | [255, 0, None]@Series[object] | [255, 0, None]@Series[object] | +| uint16:extremes-nullable | [65535, 0, None]@uint16 | [65535.0, 0.0, nan]@Series[float64] | [65535, 0, None]@Series[object] | [65535, 0, None]@Series[object] | +| uint32:extremes-nullable | [4294967295, 0, None]@uint32 | [4294967295.0, 0.0, nan]@Series[float64] | [4294967295, 0, None]@Series[object] | [4294967295, 0, None]@Series[object] | +| uint64:extremes-nullable | [18446744073709551615, 0, None]@uint64 | [1.8446744073709552e+19, 0.0, nan]@Series[float64] | [18446744073709551615, 0, None]@Series[object] | [18446744073709551615, 0, None]@Series[object] | +| int64:all-null | [None, None]@int64 | [nan, nan]@Series[float64] | [None, None]@Series[object] | [None, None]@Series[object] | +| list<int64>:null-element | [[1, None], [2, 3]]@list<item: int64> | [array([ 1., nan]), array([2., 3.])]@Series[object] | [array([1, None], dtype=object), array([2, 3], dtype=object)]@Series[object] | [array([1, None], dtype=object), array([2, 3], dtype=object)]@Series[object] | +| list<int64>:null-element-extreme | [[9223372036854775807, None]]@list<item: int64> | [array([9.22337204e+18, nan])]@Series[object] | [array([9223372036854775807, None], dtype=object)]@Series[object] | [array([9223372036854775807, None], dtype=object)]@Series[object] | +| large_list<int64>:null-element | [[1, None]]@large_list<item: int64> | [array([ 1., nan])]@Series[object] | [array([1, None], dtype=object)]@Series[object] | [array([1, None], dtype=object)]@Series[object] | +| fixed_size_list<int64>[3]:null-element | [[1, None, 3]]@fixed_size_list<item: int64>[3] | [array([ 1., nan, 3.])]@Series[object] | [array([1, None, 3], dtype=object)]@Series[object] | [array([1, None, 3], dtype=object)]@Series[object] | +| list<list<int64>>:null-element | [[[1, None], [2]]]@list<item: list<item: int64>> | [array([array([ 1., nan]), array([2.])], dtype=object)]@Series[object] | [array([array([1, None], dtype=object), array([2], dtype=object)], dtype=object)]@Series[object] | [array([array([1, None], dtype=object), array([2], dtype=object)], dtype=object)]@Series[object] | +| struct:null-int-field | [[('x', 1), ('y', 'a')], [('x', None), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1.0, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': None, 'y': 'b'}]@Series[object] | +| map<string,int64>:null-value | [[('a', 1), ('b', None)]]@map<string, int64> | [[('a', 1.0), ('b', None)]]@Series[object] | [[('a', 1), ('b', None)]]@Series[object] | [[('a', 1), ('b', None)]]@Series[object] | +| dictionary<int64>:nullable | [1, None, 1]@dictionary<values=int64, indices=int32, ordered=0> | [1, nan, 1]@Series[category] | [1, nan, 1]@Series[category] | [1, nan, 1]@Series[category] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.csv new file mode 100644 index 0000000000000..f587c3753a27b --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.csv @@ -0,0 +1,136 @@ +test case pyarrow array zero_copy_only=False zero_copy_only=True verified zero-copy +int8:standard [0, 1, -1, 127, -128]@int8 [0, 1, -1, 127, -128]@Series[int8] [0, 1, -1, 127, -128]@Series[int8] zero-copy +int8:nullable [0, 1, None]@int8 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +int8:empty []@int8 []@Series[int8] []@Series[int8] zero-copy +int16:standard [0, 1, -1, 32767, -32768]@int16 [0, 1, -1, 32767, -32768]@Series[int16] [0, 1, -1, 32767, -32768]@Series[int16] zero-copy +int16:nullable [0, 1, None]@int16 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +int16:empty []@int16 []@Series[int16] []@Series[int16] zero-copy +int32:standard [0, 1, -1, 2147483647, -2147483648]@int32 [0, 1, -1, 2147483647, -2147483648]@Series[int32] [0, 1, -1, 2147483647, -2147483648]@Series[int32] zero-copy +int32:nullable [0, 1, None]@int32 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +int32:empty []@int32 []@Series[int32] []@Series[int32] zero-copy +int64:standard [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] zero-copy +int64:nullable [0, 1, None]@int64 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +int64:empty []@int64 []@Series[int64] []@Series[int64] zero-copy +uint8:standard [0, 1, 255]@uint8 [0, 1, 255]@Series[uint8] [0, 1, 255]@Series[uint8] zero-copy +uint8:nullable [0, 1, None]@uint8 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +uint8:empty []@uint8 []@Series[uint8] []@Series[uint8] zero-copy +uint16:standard [0, 1, 65535]@uint16 [0, 1, 65535]@Series[uint16] [0, 1, 65535]@Series[uint16] zero-copy +uint16:nullable [0, 1, None]@uint16 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +uint16:empty []@uint16 []@Series[uint16] []@Series[uint16] zero-copy +uint32:standard [0, 1, 4294967295]@uint32 [0, 1, 4294967295]@Series[uint32] [0, 1, 4294967295]@Series[uint32] zero-copy +uint32:nullable [0, 1, None]@uint32 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +uint32:empty []@uint32 []@Series[uint32] []@Series[uint32] zero-copy +uint64:standard [0, 1, 18446744073709551615]@uint64 [0, 1, 18446744073709551615]@Series[uint64] [0, 1, 18446744073709551615]@Series[uint64] zero-copy +uint64:nullable [0, 1, None]@uint64 [0.0, 1.0, nan]@Series[float64] ERR@ArrowInvalid copied +uint64:empty []@uint64 []@Series[uint64] []@Series[uint64] zero-copy +float32:standard [0.0, 1.5, -1.5]@float32 [0.0, 1.5, -1.5]@Series[float32] [0.0, 1.5, -1.5]@Series[float32] zero-copy +float32:nullable [0.0, 1.5, None]@float32 [0.0, 1.5, nan]@Series[float32] ERR@ArrowInvalid copied +float32:empty []@float32 []@Series[float32] []@Series[float32] zero-copy +float64:standard [0.0, 1.5, -1.5]@float64 [0.0, 1.5, -1.5]@Series[float64] [0.0, 1.5, -1.5]@Series[float64] zero-copy +float64:nullable [0.0, 1.5, None]@float64 [0.0, 1.5, nan]@Series[float64] ERR@ArrowInvalid copied +float64:special [nan, inf, -inf]@float64 [nan, inf, -inf]@Series[float64] [nan, inf, -inf]@Series[float64] zero-copy +float64:empty []@float64 []@Series[float64] []@Series[float64] zero-copy +bool:standard [True, False, True]@bool [True, False, True]@Series[bool] ERR@ArrowInvalid copied +bool:nullable [True, False, None]@bool [True, False, None]@Series[object] ERR@ArrowInvalid copied +bool:empty []@bool []@Series[bool] ERR@ArrowInvalid copied +string:standard [hello, world, ]@string ['hello', 'world', '']@Series[object] ERR@ArrowInvalid copied +string:nullable [hello, None, world]@string ['hello', None, 'world']@Series[object] ERR@ArrowInvalid copied +string:empty []@string []@Series[object] ERR@ArrowInvalid copied +large_string:standard [hello, world]@large_string ['hello', 'world']@Series[object] ERR@ArrowInvalid copied +large_string:nullable [hello, None]@large_string ['hello', None]@Series[object] ERR@ArrowInvalid copied +large_string:empty []@large_string []@Series[object] ERR@ArrowInvalid copied +binary:standard [b'hello', b'world']@binary [b'hello', b'world']@Series[object] ERR@ArrowInvalid copied +binary:nullable [b'hello', None]@binary [b'hello', None]@Series[object] ERR@ArrowInvalid copied +binary:empty []@binary []@Series[object] ERR@ArrowInvalid copied +large_binary:standard [b'hello', b'world']@large_binary [b'hello', b'world']@Series[object] ERR@ArrowInvalid copied +large_binary:nullable [b'hello', None]@large_binary [b'hello', None]@Series[object] ERR@ArrowInvalid copied +large_binary:empty []@large_binary []@Series[object] ERR@ArrowInvalid copied +decimal128:standard [1.23, 4.56, -7.89]@decimal128(5, 2) [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] ERR@ArrowInvalid copied +decimal128:nullable [1.23, None, 4.56]@decimal128(5, 2) [Decimal('1.23'), None, Decimal('4.56')]@Series[object] ERR@ArrowInvalid copied +decimal128:empty []@decimal128(5, 2) []@Series[object] ERR@ArrowInvalid copied +date32:standard [2024-01-01, 2024-06-15]@date32[day] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] ERR@ArrowInvalid copied +date32:nullable [2024-01-01, None]@date32[day] [datetime.date(2024, 1, 1), None]@Series[object] ERR@ArrowInvalid copied +date32:empty []@date32[day] []@Series[object] ERR@ArrowInvalid copied +date64:standard [2024-01-01, 2024-06-15]@date64[ms] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] ERR@ArrowInvalid copied +date64:nullable [2024-01-01, None]@date64[ms] [datetime.date(2024, 1, 1), None]@Series[object] ERR@ArrowInvalid copied +date64:empty []@date64[ms] []@Series[object] ERR@ArrowInvalid copied +timestamp[s]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] zero-copy +timestamp[s]:nullable [2024-01-01 12:00:00, None]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] ERR@ArrowInvalid copied +timestamp[s]:empty []@timestamp[s] []@Series[datetime64[s]] []@Series[datetime64[s]] zero-copy +timestamp[ms]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] zero-copy +timestamp[ms]:nullable [2024-01-01 12:00:00, None]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] ERR@ArrowInvalid copied +timestamp[ms]:empty []@timestamp[ms] []@Series[datetime64[ms]] []@Series[datetime64[ms]] zero-copy +timestamp[us]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] zero-copy +timestamp[us]:nullable [2024-01-01 12:00:00, None]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] ERR@ArrowInvalid copied +timestamp[us]:empty []@timestamp[us] []@Series[datetime64[us]] []@Series[datetime64[us]] zero-copy +timestamp[ns]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] zero-copy +timestamp[ns]:nullable [2024-01-01 12:00:00, None]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] ERR@ArrowInvalid copied +timestamp[ns]:empty []@timestamp[ns] []@Series[datetime64[ns]] []@Series[datetime64[ns]] zero-copy +timestamp[us,tz=UTC]:standard [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] copied +timestamp[us,tz=UTC]:nullable [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] ERR@ArrowInvalid copied +timestamp[us,tz=UTC]:empty []@timestamp[us, tz=UTC] []@Series[datetime64[us, UTC]] []@Series[datetime64[us, UTC]] copied +duration[s]:standard [1 day, 0:00:00, 2:30:00]@duration[s] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] zero-copy +duration[s]:nullable [1 day, 0:00:00, None]@duration[s] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] ERR@ArrowInvalid copied +duration[s]:empty []@duration[s] []@Series[timedelta64[s]] []@Series[timedelta64[s]] zero-copy +duration[ms]:standard [1 day, 0:00:00, 2:30:00]@duration[ms] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] zero-copy +duration[ms]:nullable [1 day, 0:00:00, None]@duration[ms] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] ERR@ArrowInvalid copied +duration[ms]:empty []@duration[ms] []@Series[timedelta64[ms]] []@Series[timedelta64[ms]] zero-copy +duration[us]:standard [1 day, 0:00:00, 2:30:00]@duration[us] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] zero-copy +duration[us]:nullable [1 day, 0:00:00, None]@duration[us] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] ERR@ArrowInvalid copied +duration[us]:empty []@duration[us] []@Series[timedelta64[us]] []@Series[timedelta64[us]] zero-copy +duration[ns]:standard [1 days 00:00:00, 0 days 02:30:00]@duration[ns] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] zero-copy +duration[ns]:nullable [1 days 00:00:00, None]@duration[ns] [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] ERR@ArrowInvalid copied +duration[ns]:empty []@duration[ns] []@Series[timedelta64[ns]] []@Series[timedelta64[ns]] zero-copy +time32[s]:standard [12:30:00, 18:45:30]@time32[s] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] ERR@ArrowInvalid copied +time32[s]:nullable [12:30:00, None]@time32[s] [datetime.time(12, 30), None]@Series[object] ERR@ArrowInvalid copied +time32[s]:empty []@time32[s] []@Series[object] ERR@ArrowInvalid copied +time32[ms]:standard [12:30:00, 18:45:30]@time32[ms] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] ERR@ArrowInvalid copied +time32[ms]:nullable [12:30:00, None]@time32[ms] [datetime.time(12, 30), None]@Series[object] ERR@ArrowInvalid copied +time32[ms]:empty []@time32[ms] []@Series[object] ERR@ArrowInvalid copied +time64[us]:standard [12:30:00, 18:45:30]@time64[us] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] ERR@ArrowInvalid copied +time64[us]:nullable [12:30:00, None]@time64[us] [datetime.time(12, 30), None]@Series[object] ERR@ArrowInvalid copied +time64[us]:empty []@time64[us] []@Series[object] ERR@ArrowInvalid copied +time64[ns]:standard [12:30:00, 18:45:30]@time64[ns] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] ERR@ArrowInvalid copied +time64[ns]:nullable [12:30:00, None]@time64[ns] [datetime.time(12, 30), None]@Series[object] ERR@ArrowInvalid copied +time64[ns]:empty []@time64[ns] []@Series[object] ERR@ArrowInvalid copied +null:standard [None, None, None]@null [None, None, None]@Series[object] ERR@ArrowInvalid copied +null:empty []@null []@Series[object] ERR@ArrowInvalid copied +list<int64>:standard [[1, 2], [3, 4, 5]]@list<item: int64> [array([1, 2]), array([3, 4, 5])]@Series[object] ERR@ArrowInvalid copied +list<int64>:nullable [[1, 2], None, [3]]@list<item: int64> [array([1, 2]), None, array([3])]@Series[object] ERR@ArrowInvalid copied +list<int64>:empty []@list<item: int64> []@Series[object] ERR@ArrowInvalid copied +list<string>:standard [['a', 'b'], ['c']]@list<item: string> [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] ERR@ArrowInvalid copied +large_list<int64>:standard [[1, 2], [3, 4]]@large_list<item: int64> [array([1, 2]), array([3, 4])]@Series[object] ERR@ArrowInvalid copied +large_list<int64>:empty []@large_list<item: int64> []@Series[object] ERR@ArrowInvalid copied +fixed_size_list<int64>[3]:standard [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] [array([1, 2, 3]), array([4, 5, 6])]@Series[object] ERR@ArrowInvalid copied +fixed_size_list<int64>[3]:empty []@fixed_size_list<item: int64>[3] []@Series[object] ERR@ArrowInvalid copied +struct:standard [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] ERR@ArrowInvalid copied +struct:nullable [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, None]@Series[object] ERR@ArrowInvalid copied +struct:empty []@struct<x: int64, y: string> []@Series[object] ERR@ArrowInvalid copied +map<string,int64>:standard [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] ERR@ArrowInvalid copied +map<string,int64>:empty []@map<string, int64> []@Series[object] ERR@ArrowInvalid copied +list<list<int64>>:standard [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] ERR@ArrowInvalid copied +list<struct>:standard [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] ERR@ArrowInvalid copied +list<map<string,int64>>:standard [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] ERR@ArrowInvalid copied +struct<struct>:standard [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] ERR@ArrowInvalid copied +struct<list<int64>>:standard [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] ERR@ArrowInvalid copied +struct<map<string,int64>>:standard [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] ERR@ArrowInvalid copied +map<string,list<int64>>:standard [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] ERR@ArrowInvalid copied +map<string,struct>:standard [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] ERR@ArrowInvalid copied +map<string,map<string,int64>>:standard [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] ERR@ArrowInvalid copied +dictionary<int32,string>:standard [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', 'a', 'b']@Series[category] ERR@ArrowInvalid copied +dictionary<int32,string>:nullable [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', nan, 'a']@Series[category] ERR@ArrowInvalid copied +dictionary<int32,string>:empty []@dictionary<values=string, indices=int32, ordered=0> []@Series[category] ERR@ArrowInvalid copied +int64:zero-chunk []@int64 []@Series[int64] ERR@ArrowInvalid copied +int64:empty-chunk []@int64 []@Series[int64] []@Series[int64] zero-copy +int64:single-chunk [1, 2, 3]@int64 [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] zero-copy +int64:multi-chunk [1, 2, 3, 4]@int64 [1, 2, 3, 4]@Series[int64] ERR@ArrowInvalid copied +int64:multi-chunk-nullable [1, None, 2]@int64 [1.0, nan, 2.0]@Series[float64] ERR@ArrowInvalid copied +int64:multi-chunk-with-empty [1, 2, 3]@int64 [1, 2, 3]@Series[int64] ERR@ArrowInvalid copied +float64:multi-chunk [1.5, 2.5, 3.5]@float64 [1.5, 2.5, 3.5]@Series[float64] ERR@ArrowInvalid copied +string:single-chunk [a, b]@string ['a', 'b']@Series[object] ERR@ArrowInvalid copied +string:multi-chunk [a, b, c]@string ['a', 'b', 'c']@Series[object] ERR@ArrowInvalid copied +string:multi-chunk-nullable [a, None, c]@string ['a', None, 'c']@Series[object] ERR@ArrowInvalid copied +list<int64>:multi-chunk [[1, 2], [3], [4]]@list<item: int64> [array([1, 2]), array([3]), array([4])]@Series[object] ERR@ArrowInvalid copied +struct:multi-chunk [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] ERR@ArrowInvalid copied +int64:sliced [2, 3, 4]@int64 [2, 3, 4]@Series[int64] [2, 3, 4]@Series[int64] zero-copy +int64:sliced-with-null [2, None, 4]@int64 [2.0, nan, 4.0]@Series[float64] ERR@ArrowInvalid copied diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.md new file mode 100644 index 0000000000000..1091d1e5d3d64 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy.md @@ -0,0 +1,137 @@ +| test case | pyarrow array | zero_copy_only=False | zero_copy_only=True | verified zero-copy | +|----------------------------------------|-----------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------|----------------------| +| int8:standard | [0, 1, -1, 127, -128]@int8 | [0, 1, -1, 127, -128]@Series[int8] | [0, 1, -1, 127, -128]@Series[int8] | zero-copy | +| int8:nullable | [0, 1, None]@int8 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| int8:empty | []@int8 | []@Series[int8] | []@Series[int8] | zero-copy | +| int16:standard | [0, 1, -1, 32767, -32768]@int16 | [0, 1, -1, 32767, -32768]@Series[int16] | [0, 1, -1, 32767, -32768]@Series[int16] | zero-copy | +| int16:nullable | [0, 1, None]@int16 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| int16:empty | []@int16 | []@Series[int16] | []@Series[int16] | zero-copy | +| int32:standard | [0, 1, -1, 2147483647, -2147483648]@int32 | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | [0, 1, -1, 2147483647, -2147483648]@Series[int32] | zero-copy | +| int32:nullable | [0, 1, None]@int32 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| int32:empty | []@int32 | []@Series[int32] | []@Series[int32] | zero-copy | +| int64:standard | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64] | zero-copy | +| int64:nullable | [0, 1, None]@int64 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| int64:empty | []@int64 | []@Series[int64] | []@Series[int64] | zero-copy | +| uint8:standard | [0, 1, 255]@uint8 | [0, 1, 255]@Series[uint8] | [0, 1, 255]@Series[uint8] | zero-copy | +| uint8:nullable | [0, 1, None]@uint8 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| uint8:empty | []@uint8 | []@Series[uint8] | []@Series[uint8] | zero-copy | +| uint16:standard | [0, 1, 65535]@uint16 | [0, 1, 65535]@Series[uint16] | [0, 1, 65535]@Series[uint16] | zero-copy | +| uint16:nullable | [0, 1, None]@uint16 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| uint16:empty | []@uint16 | []@Series[uint16] | []@Series[uint16] | zero-copy | +| uint32:standard | [0, 1, 4294967295]@uint32 | [0, 1, 4294967295]@Series[uint32] | [0, 1, 4294967295]@Series[uint32] | zero-copy | +| uint32:nullable | [0, 1, None]@uint32 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| uint32:empty | []@uint32 | []@Series[uint32] | []@Series[uint32] | zero-copy | +| uint64:standard | [0, 1, 18446744073709551615]@uint64 | [0, 1, 18446744073709551615]@Series[uint64] | [0, 1, 18446744073709551615]@Series[uint64] | zero-copy | +| uint64:nullable | [0, 1, None]@uint64 | [0.0, 1.0, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| uint64:empty | []@uint64 | []@Series[uint64] | []@Series[uint64] | zero-copy | +| float32:standard | [0.0, 1.5, -1.5]@float32 | [0.0, 1.5, -1.5]@Series[float32] | [0.0, 1.5, -1.5]@Series[float32] | zero-copy | +| float32:nullable | [0.0, 1.5, None]@float32 | [0.0, 1.5, nan]@Series[float32] | ERR@ArrowInvalid | copied | +| float32:empty | []@float32 | []@Series[float32] | []@Series[float32] | zero-copy | +| float64:standard | [0.0, 1.5, -1.5]@float64 | [0.0, 1.5, -1.5]@Series[float64] | [0.0, 1.5, -1.5]@Series[float64] | zero-copy | +| float64:nullable | [0.0, 1.5, None]@float64 | [0.0, 1.5, nan]@Series[float64] | ERR@ArrowInvalid | copied | +| float64:special | [nan, inf, -inf]@float64 | [nan, inf, -inf]@Series[float64] | [nan, inf, -inf]@Series[float64] | zero-copy | +| float64:empty | []@float64 | []@Series[float64] | []@Series[float64] | zero-copy | +| bool:standard | [True, False, True]@bool | [True, False, True]@Series[bool] | ERR@ArrowInvalid | copied | +| bool:nullable | [True, False, None]@bool | [True, False, None]@Series[object] | ERR@ArrowInvalid | copied | +| bool:empty | []@bool | []@Series[bool] | ERR@ArrowInvalid | copied | +| string:standard | [hello, world, ]@string | ['hello', 'world', '']@Series[object] | ERR@ArrowInvalid | copied | +| string:nullable | [hello, None, world]@string | ['hello', None, 'world']@Series[object] | ERR@ArrowInvalid | copied | +| string:empty | []@string | []@Series[object] | ERR@ArrowInvalid | copied | +| large_string:standard | [hello, world]@large_string | ['hello', 'world']@Series[object] | ERR@ArrowInvalid | copied | +| large_string:nullable | [hello, None]@large_string | ['hello', None]@Series[object] | ERR@ArrowInvalid | copied | +| large_string:empty | []@large_string | []@Series[object] | ERR@ArrowInvalid | copied | +| binary:standard | [b'hello', b'world']@binary | [b'hello', b'world']@Series[object] | ERR@ArrowInvalid | copied | +| binary:nullable | [b'hello', None]@binary | [b'hello', None]@Series[object] | ERR@ArrowInvalid | copied | +| binary:empty | []@binary | []@Series[object] | ERR@ArrowInvalid | copied | +| large_binary:standard | [b'hello', b'world']@large_binary | [b'hello', b'world']@Series[object] | ERR@ArrowInvalid | copied | +| large_binary:nullable | [b'hello', None]@large_binary | [b'hello', None]@Series[object] | ERR@ArrowInvalid | copied | +| large_binary:empty | []@large_binary | []@Series[object] | ERR@ArrowInvalid | copied | +| decimal128:standard | [1.23, 4.56, -7.89]@decimal128(5, 2) | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[object] | ERR@ArrowInvalid | copied | +| decimal128:nullable | [1.23, None, 4.56]@decimal128(5, 2) | [Decimal('1.23'), None, Decimal('4.56')]@Series[object] | ERR@ArrowInvalid | copied | +| decimal128:empty | []@decimal128(5, 2) | []@Series[object] | ERR@ArrowInvalid | copied | +| date32:standard | [2024-01-01, 2024-06-15]@date32[day] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | ERR@ArrowInvalid | copied | +| date32:nullable | [2024-01-01, None]@date32[day] | [datetime.date(2024, 1, 1), None]@Series[object] | ERR@ArrowInvalid | copied | +| date32:empty | []@date32[day] | []@Series[object] | ERR@ArrowInvalid | copied | +| date64:standard | [2024-01-01, 2024-06-15]@date64[ms] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[object] | ERR@ArrowInvalid | copied | +| date64:nullable | [2024-01-01, None]@date64[ms] | [datetime.date(2024, 1, 1), None]@Series[object] | ERR@ArrowInvalid | copied | +| date64:empty | []@date64[ms] | []@Series[object] | ERR@ArrowInvalid | copied | +| timestamp[s]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[s]] | zero-copy | +| timestamp[s]:nullable | [2024-01-01 12:00:00, None]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[s]] | ERR@ArrowInvalid | copied | +| timestamp[s]:empty | []@timestamp[s] | []@Series[datetime64[s]] | []@Series[datetime64[s]] | zero-copy | +| timestamp[ms]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ms]] | zero-copy | +| timestamp[ms]:nullable | [2024-01-01 12:00:00, None]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ms]] | ERR@ArrowInvalid | copied | +| timestamp[ms]:empty | []@timestamp[ms] | []@Series[datetime64[ms]] | []@Series[datetime64[ms]] | zero-copy | +| timestamp[us]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]] | zero-copy | +| timestamp[us]:nullable | [2024-01-01 12:00:00, None]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[us]] | ERR@ArrowInvalid | copied | +| timestamp[us]:empty | []@timestamp[us] | []@Series[datetime64[us]] | []@Series[datetime64[us]] | zero-copy | +| timestamp[ns]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[datetime64[ns]] | zero-copy | +| timestamp[ns]:nullable | [2024-01-01 12:00:00, None]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), NaT]@Series[datetime64[ns]] | ERR@ArrowInvalid | copied | +| timestamp[ns]:empty | []@timestamp[ns] | []@Series[datetime64[ns]] | []@Series[datetime64[ns]] | zero-copy | +| timestamp[us,tz=UTC]:standard | [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[datetime64[us, UTC]] | copied | +| timestamp[us,tz=UTC]:nullable | [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), NaT]@Series[datetime64[us, UTC]] | ERR@ArrowInvalid | copied | +| timestamp[us,tz=UTC]:empty | []@timestamp[us, tz=UTC] | []@Series[datetime64[us, UTC]] | []@Series[datetime64[us, UTC]] | copied | +| duration[s]:standard | [1 day, 0:00:00, 2:30:00]@duration[s] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[s]] | zero-copy | +| duration[s]:nullable | [1 day, 0:00:00, None]@duration[s] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[s]] | ERR@ArrowInvalid | copied | +| duration[s]:empty | []@duration[s] | []@Series[timedelta64[s]] | []@Series[timedelta64[s]] | zero-copy | +| duration[ms]:standard | [1 day, 0:00:00, 2:30:00]@duration[ms] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ms]] | zero-copy | +| duration[ms]:nullable | [1 day, 0:00:00, None]@duration[ms] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ms]] | ERR@ArrowInvalid | copied | +| duration[ms]:empty | []@duration[ms] | []@Series[timedelta64[ms]] | []@Series[timedelta64[ms]] | zero-copy | +| duration[us]:standard | [1 day, 0:00:00, 2:30:00]@duration[us] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[us]] | zero-copy | +| duration[us]:nullable | [1 day, 0:00:00, None]@duration[us] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[us]] | ERR@ArrowInvalid | copied | +| duration[us]:empty | []@duration[us] | []@Series[timedelta64[us]] | []@Series[timedelta64[us]] | zero-copy | +| duration[ns]:standard | [1 days 00:00:00, 0 days 02:30:00]@duration[ns] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[timedelta64[ns]] | zero-copy | +| duration[ns]:nullable | [1 days 00:00:00, None]@duration[ns] | [Timedelta('1 days 00:00:00'), NaT]@Series[timedelta64[ns]] | ERR@ArrowInvalid | copied | +| duration[ns]:empty | []@duration[ns] | []@Series[timedelta64[ns]] | []@Series[timedelta64[ns]] | zero-copy | +| time32[s]:standard | [12:30:00, 18:45:30]@time32[s] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | ERR@ArrowInvalid | copied | +| time32[s]:nullable | [12:30:00, None]@time32[s] | [datetime.time(12, 30), None]@Series[object] | ERR@ArrowInvalid | copied | +| time32[s]:empty | []@time32[s] | []@Series[object] | ERR@ArrowInvalid | copied | +| time32[ms]:standard | [12:30:00, 18:45:30]@time32[ms] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | ERR@ArrowInvalid | copied | +| time32[ms]:nullable | [12:30:00, None]@time32[ms] | [datetime.time(12, 30), None]@Series[object] | ERR@ArrowInvalid | copied | +| time32[ms]:empty | []@time32[ms] | []@Series[object] | ERR@ArrowInvalid | copied | +| time64[us]:standard | [12:30:00, 18:45:30]@time64[us] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | ERR@ArrowInvalid | copied | +| time64[us]:nullable | [12:30:00, None]@time64[us] | [datetime.time(12, 30), None]@Series[object] | ERR@ArrowInvalid | copied | +| time64[us]:empty | []@time64[us] | []@Series[object] | ERR@ArrowInvalid | copied | +| time64[ns]:standard | [12:30:00, 18:45:30]@time64[ns] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[object] | ERR@ArrowInvalid | copied | +| time64[ns]:nullable | [12:30:00, None]@time64[ns] | [datetime.time(12, 30), None]@Series[object] | ERR@ArrowInvalid | copied | +| time64[ns]:empty | []@time64[ns] | []@Series[object] | ERR@ArrowInvalid | copied | +| null:standard | [None, None, None]@null | [None, None, None]@Series[object] | ERR@ArrowInvalid | copied | +| null:empty | []@null | []@Series[object] | ERR@ArrowInvalid | copied | +| list<int64>:standard | [[1, 2], [3, 4, 5]]@list<item: int64> | [array([1, 2]), array([3, 4, 5])]@Series[object] | ERR@ArrowInvalid | copied | +| list<int64>:nullable | [[1, 2], None, [3]]@list<item: int64> | [array([1, 2]), None, array([3])]@Series[object] | ERR@ArrowInvalid | copied | +| list<int64>:empty | []@list<item: int64> | []@Series[object] | ERR@ArrowInvalid | copied | +| list<string>:standard | [['a', 'b'], ['c']]@list<item: string> | [array(['a', 'b'], dtype=object), array(['c'], dtype=object)]@Series[object] | ERR@ArrowInvalid | copied | +| large_list<int64>:standard | [[1, 2], [3, 4]]@large_list<item: int64> | [array([1, 2]), array([3, 4])]@Series[object] | ERR@ArrowInvalid | copied | +| large_list<int64>:empty | []@large_list<item: int64> | []@Series[object] | ERR@ArrowInvalid | copied | +| fixed_size_list<int64>[3]:standard | [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] | [array([1, 2, 3]), array([4, 5, 6])]@Series[object] | ERR@ArrowInvalid | copied | +| fixed_size_list<int64>[3]:empty | []@fixed_size_list<item: int64>[3] | []@Series[object] | ERR@ArrowInvalid | copied | +| struct:standard | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | ERR@ArrowInvalid | copied | +| struct:nullable | [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, None]@Series[object] | ERR@ArrowInvalid | copied | +| struct:empty | []@struct<x: int64, y: string> | []@Series[object] | ERR@ArrowInvalid | copied | +| map<string,int64>:standard | [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[object] | ERR@ArrowInvalid | copied | +| map<string,int64>:empty | []@map<string, int64> | []@Series[object] | ERR@ArrowInvalid | copied | +| list<list<int64>>:standard | [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> | [array([array([1, 2]), array([3])], dtype=object), array([array([4, 5, 6])], dtype=object)]@Series[object] | ERR@ArrowInvalid | copied | +| list<struct>:standard | [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> | [array([{'x': 1}, {'x': 2}], dtype=object), array([{'x': 3}], dtype=object)]@Series[object] | ERR@ArrowInvalid | copied | +| list<map<string,int64>>:standard | [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> | [array([list([('a', 1)]), list([('b', 2)])], dtype=object), array([list([('c', 3)])], dtype=object)]@Series[object] | ERR@ArrowInvalid | copied | +| struct<struct>:standard | [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[object] | ERR@ArrowInvalid | copied | +| struct<list<int64>>:standard | [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> | [{'items': array([1, 2, 3])}, {'items': array([4, 5])}]@Series[object] | ERR@ArrowInvalid | copied | +| struct<map<string,int64>>:standard | [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[object] | ERR@ArrowInvalid | copied | +| map<string,list<int64>>:standard | [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> | [[('a', array([1, 2])), ('b', array([3]))], [('c', array([4, 5, 6]))]]@Series[object] | ERR@ArrowInvalid | copied | +| map<string,struct>:standard | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[object] | ERR@ArrowInvalid | copied | +| map<string,map<string,int64>>:standard | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[object] | ERR@ArrowInvalid | copied | +| dictionary<int32,string>:standard | [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', 'a', 'b']@Series[category] | ERR@ArrowInvalid | copied | +| dictionary<int32,string>:nullable | [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', nan, 'a']@Series[category] | ERR@ArrowInvalid | copied | +| dictionary<int32,string>:empty | []@dictionary<values=string, indices=int32, ordered=0> | []@Series[category] | ERR@ArrowInvalid | copied | +| int64:zero-chunk | []@int64 | []@Series[int64] | ERR@ArrowInvalid | copied | +| int64:empty-chunk | []@int64 | []@Series[int64] | []@Series[int64] | zero-copy | +| int64:single-chunk | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | zero-copy | +| int64:multi-chunk | [1, 2, 3, 4]@int64 | [1, 2, 3, 4]@Series[int64] | ERR@ArrowInvalid | copied | +| int64:multi-chunk-nullable | [1, None, 2]@int64 | [1.0, nan, 2.0]@Series[float64] | ERR@ArrowInvalid | copied | +| int64:multi-chunk-with-empty | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | ERR@ArrowInvalid | copied | +| float64:multi-chunk | [1.5, 2.5, 3.5]@float64 | [1.5, 2.5, 3.5]@Series[float64] | ERR@ArrowInvalid | copied | +| string:single-chunk | [a, b]@string | ['a', 'b']@Series[object] | ERR@ArrowInvalid | copied | +| string:multi-chunk | [a, b, c]@string | ['a', 'b', 'c']@Series[object] | ERR@ArrowInvalid | copied | +| string:multi-chunk-nullable | [a, None, c]@string | ['a', None, 'c']@Series[object] | ERR@ArrowInvalid | copied | +| list<int64>:multi-chunk | [[1, 2], [3], [4]]@list<item: int64> | [array([1, 2]), array([3]), array([4])]@Series[object] | ERR@ArrowInvalid | copied | +| struct:multi-chunk | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | ERR@ArrowInvalid | copied | +| int64:sliced | [2, 3, 4]@int64 | [2, 3, 4]@Series[int64] | [2, 3, 4]@Series[int64] | zero-copy | +| int64:sliced-with-null | [2, None, 4]@int64 | [2.0, nan, 4.0]@Series[float64] | ERR@ArrowInvalid | copied | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.csv new file mode 100644 index 0000000000000..55a84ce933e3b --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.csv @@ -0,0 +1,136 @@ +test case pyarrow array types_mapper=pd.ArrowDtype, zero_copy_only=False types_mapper=pd.ArrowDtype, zero_copy_only=True verified zero-copy +int8:standard [0, 1, -1, 127, -128]@int8 [0, 1, -1, 127, -128]@Series[int8[pyarrow]] [0, 1, -1, 127, -128]@Series[int8[pyarrow]] zero-copy +int8:nullable [0, 1, None]@int8 [0, 1, <NA>]@Series[int8[pyarrow]] [0, 1, <NA>]@Series[int8[pyarrow]] zero-copy +int8:empty []@int8 []@Series[int8[pyarrow]] []@Series[int8[pyarrow]] zero-copy +int16:standard [0, 1, -1, 32767, -32768]@int16 [0, 1, -1, 32767, -32768]@Series[int16[pyarrow]] [0, 1, -1, 32767, -32768]@Series[int16[pyarrow]] zero-copy +int16:nullable [0, 1, None]@int16 [0, 1, <NA>]@Series[int16[pyarrow]] [0, 1, <NA>]@Series[int16[pyarrow]] zero-copy +int16:empty []@int16 []@Series[int16[pyarrow]] []@Series[int16[pyarrow]] zero-copy +int32:standard [0, 1, -1, 2147483647, -2147483648]@int32 [0, 1, -1, 2147483647, -2147483648]@Series[int32[pyarrow]] [0, 1, -1, 2147483647, -2147483648]@Series[int32[pyarrow]] zero-copy +int32:nullable [0, 1, None]@int32 [0, 1, <NA>]@Series[int32[pyarrow]] [0, 1, <NA>]@Series[int32[pyarrow]] zero-copy +int32:empty []@int32 []@Series[int32[pyarrow]] []@Series[int32[pyarrow]] zero-copy +int64:standard [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64[pyarrow]] [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64[pyarrow]] zero-copy +int64:nullable [0, 1, None]@int64 [0, 1, <NA>]@Series[int64[pyarrow]] [0, 1, <NA>]@Series[int64[pyarrow]] zero-copy +int64:empty []@int64 []@Series[int64[pyarrow]] []@Series[int64[pyarrow]] zero-copy +uint8:standard [0, 1, 255]@uint8 [0, 1, 255]@Series[uint8[pyarrow]] [0, 1, 255]@Series[uint8[pyarrow]] zero-copy +uint8:nullable [0, 1, None]@uint8 [0, 1, <NA>]@Series[uint8[pyarrow]] [0, 1, <NA>]@Series[uint8[pyarrow]] zero-copy +uint8:empty []@uint8 []@Series[uint8[pyarrow]] []@Series[uint8[pyarrow]] zero-copy +uint16:standard [0, 1, 65535]@uint16 [0, 1, 65535]@Series[uint16[pyarrow]] [0, 1, 65535]@Series[uint16[pyarrow]] zero-copy +uint16:nullable [0, 1, None]@uint16 [0, 1, <NA>]@Series[uint16[pyarrow]] [0, 1, <NA>]@Series[uint16[pyarrow]] zero-copy +uint16:empty []@uint16 []@Series[uint16[pyarrow]] []@Series[uint16[pyarrow]] zero-copy +uint32:standard [0, 1, 4294967295]@uint32 [0, 1, 4294967295]@Series[uint32[pyarrow]] [0, 1, 4294967295]@Series[uint32[pyarrow]] zero-copy +uint32:nullable [0, 1, None]@uint32 [0, 1, <NA>]@Series[uint32[pyarrow]] [0, 1, <NA>]@Series[uint32[pyarrow]] zero-copy +uint32:empty []@uint32 []@Series[uint32[pyarrow]] []@Series[uint32[pyarrow]] zero-copy +uint64:standard [0, 1, 18446744073709551615]@uint64 [0, 1, 18446744073709551615]@Series[uint64[pyarrow]] [0, 1, 18446744073709551615]@Series[uint64[pyarrow]] zero-copy +uint64:nullable [0, 1, None]@uint64 [0, 1, <NA>]@Series[uint64[pyarrow]] [0, 1, <NA>]@Series[uint64[pyarrow]] zero-copy +uint64:empty []@uint64 []@Series[uint64[pyarrow]] []@Series[uint64[pyarrow]] zero-copy +float32:standard [0.0, 1.5, -1.5]@float32 [0.0, 1.5, -1.5]@Series[float[pyarrow]] [0.0, 1.5, -1.5]@Series[float[pyarrow]] zero-copy +float32:nullable [0.0, 1.5, None]@float32 [0.0, 1.5, <NA>]@Series[float[pyarrow]] [0.0, 1.5, <NA>]@Series[float[pyarrow]] zero-copy +float32:empty []@float32 []@Series[float[pyarrow]] []@Series[float[pyarrow]] zero-copy +float64:standard [0.0, 1.5, -1.5]@float64 [0.0, 1.5, -1.5]@Series[double[pyarrow]] [0.0, 1.5, -1.5]@Series[double[pyarrow]] zero-copy +float64:nullable [0.0, 1.5, None]@float64 [0.0, 1.5, <NA>]@Series[double[pyarrow]] [0.0, 1.5, <NA>]@Series[double[pyarrow]] zero-copy +float64:special [nan, inf, -inf]@float64 [nan, inf, -inf]@Series[double[pyarrow]] [nan, inf, -inf]@Series[double[pyarrow]] zero-copy +float64:empty []@float64 []@Series[double[pyarrow]] []@Series[double[pyarrow]] zero-copy +bool:standard [True, False, True]@bool [True, False, True]@Series[bool[pyarrow]] [True, False, True]@Series[bool[pyarrow]] zero-copy +bool:nullable [True, False, None]@bool [True, False, <NA>]@Series[bool[pyarrow]] [True, False, <NA>]@Series[bool[pyarrow]] zero-copy +bool:empty []@bool []@Series[bool[pyarrow]] []@Series[bool[pyarrow]] zero-copy +string:standard [hello, world, ]@string ['hello', 'world', '']@Series[string[pyarrow]] ['hello', 'world', '']@Series[string[pyarrow]] zero-copy +string:nullable [hello, None, world]@string ['hello', <NA>, 'world']@Series[string[pyarrow]] ['hello', <NA>, 'world']@Series[string[pyarrow]] zero-copy +string:empty []@string []@Series[string[pyarrow]] []@Series[string[pyarrow]] zero-copy +large_string:standard [hello, world]@large_string ['hello', 'world']@Series[large_string[pyarrow]] ['hello', 'world']@Series[large_string[pyarrow]] zero-copy +large_string:nullable [hello, None]@large_string ['hello', <NA>]@Series[large_string[pyarrow]] ['hello', <NA>]@Series[large_string[pyarrow]] zero-copy +large_string:empty []@large_string []@Series[large_string[pyarrow]] []@Series[large_string[pyarrow]] zero-copy +binary:standard [b'hello', b'world']@binary [b'hello', b'world']@Series[binary[pyarrow]] [b'hello', b'world']@Series[binary[pyarrow]] zero-copy +binary:nullable [b'hello', None]@binary [b'hello', <NA>]@Series[binary[pyarrow]] [b'hello', <NA>]@Series[binary[pyarrow]] zero-copy +binary:empty []@binary []@Series[binary[pyarrow]] []@Series[binary[pyarrow]] zero-copy +large_binary:standard [b'hello', b'world']@large_binary [b'hello', b'world']@Series[large_binary[pyarrow]] [b'hello', b'world']@Series[large_binary[pyarrow]] zero-copy +large_binary:nullable [b'hello', None]@large_binary [b'hello', <NA>]@Series[large_binary[pyarrow]] [b'hello', <NA>]@Series[large_binary[pyarrow]] zero-copy +large_binary:empty []@large_binary []@Series[large_binary[pyarrow]] []@Series[large_binary[pyarrow]] zero-copy +decimal128:standard [1.23, 4.56, -7.89]@decimal128(5, 2) [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[decimal128(5, 2)[pyarrow]] [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[decimal128(5, 2)[pyarrow]] zero-copy +decimal128:nullable [1.23, None, 4.56]@decimal128(5, 2) [Decimal('1.23'), <NA>, Decimal('4.56')]@Series[decimal128(5, 2)[pyarrow]] [Decimal('1.23'), <NA>, Decimal('4.56')]@Series[decimal128(5, 2)[pyarrow]] zero-copy +decimal128:empty []@decimal128(5, 2) []@Series[decimal128(5, 2)[pyarrow]] []@Series[decimal128(5, 2)[pyarrow]] zero-copy +date32:standard [2024-01-01, 2024-06-15]@date32[day] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date32[day][pyarrow]] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date32[day][pyarrow]] zero-copy +date32:nullable [2024-01-01, None]@date32[day] [datetime.date(2024, 1, 1), <NA>]@Series[date32[day][pyarrow]] [datetime.date(2024, 1, 1), <NA>]@Series[date32[day][pyarrow]] zero-copy +date32:empty []@date32[day] []@Series[date32[day][pyarrow]] []@Series[date32[day][pyarrow]] zero-copy +date64:standard [2024-01-01, 2024-06-15]@date64[ms] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date64[ms][pyarrow]] [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date64[ms][pyarrow]] zero-copy +date64:nullable [2024-01-01, None]@date64[ms] [datetime.date(2024, 1, 1), <NA>]@Series[date64[ms][pyarrow]] [datetime.date(2024, 1, 1), <NA>]@Series[date64[ms][pyarrow]] zero-copy +date64:empty []@date64[ms] []@Series[date64[ms][pyarrow]] []@Series[date64[ms][pyarrow]] zero-copy +timestamp[s]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[s][pyarrow]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[s][pyarrow]] zero-copy +timestamp[s]:nullable [2024-01-01 12:00:00, None]@timestamp[s] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[s][pyarrow]] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[s][pyarrow]] zero-copy +timestamp[s]:empty []@timestamp[s] []@Series[timestamp[s][pyarrow]] []@Series[timestamp[s][pyarrow]] zero-copy +timestamp[ms]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ms][pyarrow]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ms][pyarrow]] zero-copy +timestamp[ms]:nullable [2024-01-01 12:00:00, None]@timestamp[ms] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ms][pyarrow]] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ms][pyarrow]] zero-copy +timestamp[ms]:empty []@timestamp[ms] []@Series[timestamp[ms][pyarrow]] []@Series[timestamp[ms][pyarrow]] zero-copy +timestamp[us]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[us][pyarrow]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[us][pyarrow]] zero-copy +timestamp[us]:nullable [2024-01-01 12:00:00, None]@timestamp[us] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[us][pyarrow]] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[us][pyarrow]] zero-copy +timestamp[us]:empty []@timestamp[us] []@Series[timestamp[us][pyarrow]] []@Series[timestamp[us][pyarrow]] zero-copy +timestamp[ns]:standard [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ns][pyarrow]] [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ns][pyarrow]] zero-copy +timestamp[ns]:nullable [2024-01-01 12:00:00, None]@timestamp[ns] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ns][pyarrow]] [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ns][pyarrow]] zero-copy +timestamp[ns]:empty []@timestamp[ns] []@Series[timestamp[ns][pyarrow]] []@Series[timestamp[ns][pyarrow]] zero-copy +timestamp[us,tz=UTC]:standard [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[timestamp[us, tz=UTC][pyarrow]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[timestamp[us, tz=UTC][pyarrow]] zero-copy +timestamp[us,tz=UTC]:nullable [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), <NA>]@Series[timestamp[us, tz=UTC][pyarrow]] [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), <NA>]@Series[timestamp[us, tz=UTC][pyarrow]] zero-copy +timestamp[us,tz=UTC]:empty []@timestamp[us, tz=UTC] []@Series[timestamp[us, tz=UTC][pyarrow]] []@Series[timestamp[us, tz=UTC][pyarrow]] zero-copy +duration[s]:standard [1 day, 0:00:00, 2:30:00]@duration[s] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[s][pyarrow]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[s][pyarrow]] zero-copy +duration[s]:nullable [1 day, 0:00:00, None]@duration[s] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[s][pyarrow]] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[s][pyarrow]] zero-copy +duration[s]:empty []@duration[s] []@Series[duration[s][pyarrow]] []@Series[duration[s][pyarrow]] zero-copy +duration[ms]:standard [1 day, 0:00:00, 2:30:00]@duration[ms] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ms][pyarrow]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ms][pyarrow]] zero-copy +duration[ms]:nullable [1 day, 0:00:00, None]@duration[ms] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ms][pyarrow]] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ms][pyarrow]] zero-copy +duration[ms]:empty []@duration[ms] []@Series[duration[ms][pyarrow]] []@Series[duration[ms][pyarrow]] zero-copy +duration[us]:standard [1 day, 0:00:00, 2:30:00]@duration[us] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[us][pyarrow]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[us][pyarrow]] zero-copy +duration[us]:nullable [1 day, 0:00:00, None]@duration[us] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[us][pyarrow]] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[us][pyarrow]] zero-copy +duration[us]:empty []@duration[us] []@Series[duration[us][pyarrow]] []@Series[duration[us][pyarrow]] zero-copy +duration[ns]:standard [1 days 00:00:00, 0 days 02:30:00]@duration[ns] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ns][pyarrow]] [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ns][pyarrow]] zero-copy +duration[ns]:nullable [1 days 00:00:00, None]@duration[ns] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ns][pyarrow]] [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ns][pyarrow]] zero-copy +duration[ns]:empty []@duration[ns] []@Series[duration[ns][pyarrow]] []@Series[duration[ns][pyarrow]] zero-copy +time32[s]:standard [12:30:00, 18:45:30]@time32[s] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[s][pyarrow]] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[s][pyarrow]] zero-copy +time32[s]:nullable [12:30:00, None]@time32[s] [datetime.time(12, 30), <NA>]@Series[time32[s][pyarrow]] [datetime.time(12, 30), <NA>]@Series[time32[s][pyarrow]] zero-copy +time32[s]:empty []@time32[s] []@Series[time32[s][pyarrow]] []@Series[time32[s][pyarrow]] zero-copy +time32[ms]:standard [12:30:00, 18:45:30]@time32[ms] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[ms][pyarrow]] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[ms][pyarrow]] zero-copy +time32[ms]:nullable [12:30:00, None]@time32[ms] [datetime.time(12, 30), <NA>]@Series[time32[ms][pyarrow]] [datetime.time(12, 30), <NA>]@Series[time32[ms][pyarrow]] zero-copy +time32[ms]:empty []@time32[ms] []@Series[time32[ms][pyarrow]] []@Series[time32[ms][pyarrow]] zero-copy +time64[us]:standard [12:30:00, 18:45:30]@time64[us] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[us][pyarrow]] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[us][pyarrow]] zero-copy +time64[us]:nullable [12:30:00, None]@time64[us] [datetime.time(12, 30), <NA>]@Series[time64[us][pyarrow]] [datetime.time(12, 30), <NA>]@Series[time64[us][pyarrow]] zero-copy +time64[us]:empty []@time64[us] []@Series[time64[us][pyarrow]] []@Series[time64[us][pyarrow]] zero-copy +time64[ns]:standard [12:30:00, 18:45:30]@time64[ns] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[ns][pyarrow]] [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[ns][pyarrow]] zero-copy +time64[ns]:nullable [12:30:00, None]@time64[ns] [datetime.time(12, 30), <NA>]@Series[time64[ns][pyarrow]] [datetime.time(12, 30), <NA>]@Series[time64[ns][pyarrow]] zero-copy +time64[ns]:empty []@time64[ns] []@Series[time64[ns][pyarrow]] []@Series[time64[ns][pyarrow]] zero-copy +null:standard [None, None, None]@null [<NA>, <NA>, <NA>]@Series[null[pyarrow]] [<NA>, <NA>, <NA>]@Series[null[pyarrow]] zero-copy +null:empty []@null []@Series[null[pyarrow]] []@Series[null[pyarrow]] zero-copy +list<int64>:standard [[1, 2], [3, 4, 5]]@list<item: int64> [[1, 2], [3, 4, 5]]@Series[list<item: int64>[pyarrow]] [[1, 2], [3, 4, 5]]@Series[list<item: int64>[pyarrow]] zero-copy +list<int64>:nullable [[1, 2], None, [3]]@list<item: int64> [[1, 2], <NA>, [3]]@Series[list<item: int64>[pyarrow]] [[1, 2], <NA>, [3]]@Series[list<item: int64>[pyarrow]] zero-copy +list<int64>:empty []@list<item: int64> []@Series[list<item: int64>[pyarrow]] []@Series[list<item: int64>[pyarrow]] zero-copy +list<string>:standard [['a', 'b'], ['c']]@list<item: string> [['a', 'b'], ['c']]@Series[list<item: string>[pyarrow]] [['a', 'b'], ['c']]@Series[list<item: string>[pyarrow]] zero-copy +large_list<int64>:standard [[1, 2], [3, 4]]@large_list<item: int64> [[1, 2], [3, 4]]@Series[large_list<item: int64>[pyarrow]] [[1, 2], [3, 4]]@Series[large_list<item: int64>[pyarrow]] zero-copy +large_list<int64>:empty []@large_list<item: int64> []@Series[large_list<item: int64>[pyarrow]] []@Series[large_list<item: int64>[pyarrow]] zero-copy +fixed_size_list<int64>[3]:standard [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] [[1, 2, 3], [4, 5, 6]]@Series[fixed_size_list<item: int64>[3][pyarrow]] [[1, 2, 3], [4, 5, 6]]@Series[fixed_size_list<item: int64>[3][pyarrow]] zero-copy +fixed_size_list<int64>[3]:empty []@fixed_size_list<item: int64>[3] []@Series[fixed_size_list<item: int64>[3][pyarrow]] []@Series[fixed_size_list<item: int64>[3][pyarrow]] zero-copy +struct:standard [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] zero-copy +struct:nullable [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, <NA>]@Series[struct<x: int64, y: string>[pyarrow]] [{'x': 1, 'y': 'a'}, <NA>]@Series[struct<x: int64, y: string>[pyarrow]] zero-copy +struct:empty []@struct<x: int64, y: string> []@Series[struct<x: int64, y: string>[pyarrow]] []@Series[struct<x: int64, y: string>[pyarrow]] zero-copy +map<string,int64>:standard [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> [[('a', 1), ('b', 2)], [('c', 3)]]@Series[map<string, int64>[pyarrow]] [[('a', 1), ('b', 2)], [('c', 3)]]@Series[map<string, int64>[pyarrow]] zero-copy +map<string,int64>:empty []@map<string, int64> []@Series[map<string, int64>[pyarrow]] []@Series[map<string, int64>[pyarrow]] zero-copy +list<list<int64>>:standard [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> [[[1, 2], [3]], [[4, 5, 6]]]@Series[list<item: list<item: int64>>[pyarrow]] [[[1, 2], [3]], [[4, 5, 6]]]@Series[list<item: list<item: int64>>[pyarrow]] zero-copy +list<struct>:standard [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> [[{'x': 1}, {'x': 2}], [{'x': 3}]]@Series[list<item: struct<x: int64>>[pyarrow]] [[{'x': 1}, {'x': 2}], [{'x': 3}]]@Series[list<item: struct<x: int64>>[pyarrow]] zero-copy +list<map<string,int64>>:standard [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@Series[list<item: map<string, int64>>[pyarrow]] [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@Series[list<item: map<string, int64>>[pyarrow]] zero-copy +struct<struct>:standard [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[struct<outer: struct<inner: int64>>[pyarrow]] [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[struct<outer: struct<inner: int64>>[pyarrow]] zero-copy +struct<list<int64>>:standard [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> [{'items': [1, 2, 3]}, {'items': [4, 5]}]@Series[struct<items: list<item: int64>>[pyarrow]] [{'items': [1, 2, 3]}, {'items': [4, 5]}]@Series[struct<items: list<item: int64>>[pyarrow]] zero-copy +struct<map<string,int64>>:standard [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[struct<mapping: map<string, int64>>[pyarrow]] [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[struct<mapping: map<string, int64>>[pyarrow]] zero-copy +map<string,list<int64>>:standard [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@Series[map<string, list<item: int64>>[pyarrow]] [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@Series[map<string, list<item: int64>>[pyarrow]] zero-copy +map<string,struct>:standard [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[map<string, struct<v: int64>>[pyarrow]] [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[map<string, struct<v: int64>>[pyarrow]] zero-copy +map<string,map<string,int64>>:standard [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[map<string, map<string, int64>>[pyarrow]] [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[map<string, map<string, int64>>[pyarrow]] zero-copy +dictionary<int32,string>:standard [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', 'a', 'b']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] ['a', 'b', 'a', 'b']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] zero-copy +dictionary<int32,string>:nullable [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> ['a', 'b', <NA>, 'a']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] ['a', 'b', <NA>, 'a']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] zero-copy +dictionary<int32,string>:empty []@dictionary<values=string, indices=int32, ordered=0> []@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] []@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] zero-copy +int64:zero-chunk []@int64 []@Series[int64[pyarrow]] []@Series[int64[pyarrow]] zero-copy +int64:empty-chunk []@int64 []@Series[int64[pyarrow]] []@Series[int64[pyarrow]] zero-copy +int64:single-chunk [1, 2, 3]@int64 [1, 2, 3]@Series[int64[pyarrow]] [1, 2, 3]@Series[int64[pyarrow]] zero-copy +int64:multi-chunk [1, 2, 3, 4]@int64 [1, 2, 3, 4]@Series[int64[pyarrow]] [1, 2, 3, 4]@Series[int64[pyarrow]] zero-copy +int64:multi-chunk-nullable [1, None, 2]@int64 [1, <NA>, 2]@Series[int64[pyarrow]] [1, <NA>, 2]@Series[int64[pyarrow]] zero-copy +int64:multi-chunk-with-empty [1, 2, 3]@int64 [1, 2, 3]@Series[int64[pyarrow]] [1, 2, 3]@Series[int64[pyarrow]] zero-copy +float64:multi-chunk [1.5, 2.5, 3.5]@float64 [1.5, 2.5, 3.5]@Series[double[pyarrow]] [1.5, 2.5, 3.5]@Series[double[pyarrow]] zero-copy +string:single-chunk [a, b]@string ['a', 'b']@Series[string[pyarrow]] ['a', 'b']@Series[string[pyarrow]] zero-copy +string:multi-chunk [a, b, c]@string ['a', 'b', 'c']@Series[string[pyarrow]] ['a', 'b', 'c']@Series[string[pyarrow]] zero-copy +string:multi-chunk-nullable [a, None, c]@string ['a', <NA>, 'c']@Series[string[pyarrow]] ['a', <NA>, 'c']@Series[string[pyarrow]] zero-copy +list<int64>:multi-chunk [[1, 2], [3], [4]]@list<item: int64> [[1, 2], [3], [4]]@Series[list<item: int64>[pyarrow]] [[1, 2], [3], [4]]@Series[list<item: int64>[pyarrow]] zero-copy +struct:multi-chunk [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] zero-copy +int64:sliced [2, 3, 4]@int64 [2, 3, 4]@Series[int64[pyarrow]] [2, 3, 4]@Series[int64[pyarrow]] zero-copy +int64:sliced-with-null [2, None, 4]@int64 [2, <NA>, 4]@Series[int64[pyarrow]] [2, <NA>, 4]@Series[int64[pyarrow]] zero-copy diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.md new file mode 100644 index 0000000000000..daa25761c7b45 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed.md @@ -0,0 +1,137 @@ +| test case | pyarrow array | types_mapper=pd.ArrowDtype, zero_copy_only=False | types_mapper=pd.ArrowDtype, zero_copy_only=True | verified zero-copy | +|----------------------------------------|-----------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|----------------------| +| int8:standard | [0, 1, -1, 127, -128]@int8 | [0, 1, -1, 127, -128]@Series[int8[pyarrow]] | [0, 1, -1, 127, -128]@Series[int8[pyarrow]] | zero-copy | +| int8:nullable | [0, 1, None]@int8 | [0, 1, <NA>]@Series[int8[pyarrow]] | [0, 1, <NA>]@Series[int8[pyarrow]] | zero-copy | +| int8:empty | []@int8 | []@Series[int8[pyarrow]] | []@Series[int8[pyarrow]] | zero-copy | +| int16:standard | [0, 1, -1, 32767, -32768]@int16 | [0, 1, -1, 32767, -32768]@Series[int16[pyarrow]] | [0, 1, -1, 32767, -32768]@Series[int16[pyarrow]] | zero-copy | +| int16:nullable | [0, 1, None]@int16 | [0, 1, <NA>]@Series[int16[pyarrow]] | [0, 1, <NA>]@Series[int16[pyarrow]] | zero-copy | +| int16:empty | []@int16 | []@Series[int16[pyarrow]] | []@Series[int16[pyarrow]] | zero-copy | +| int32:standard | [0, 1, -1, 2147483647, -2147483648]@int32 | [0, 1, -1, 2147483647, -2147483648]@Series[int32[pyarrow]] | [0, 1, -1, 2147483647, -2147483648]@Series[int32[pyarrow]] | zero-copy | +| int32:nullable | [0, 1, None]@int32 | [0, 1, <NA>]@Series[int32[pyarrow]] | [0, 1, <NA>]@Series[int32[pyarrow]] | zero-copy | +| int32:empty | []@int32 | []@Series[int32[pyarrow]] | []@Series[int32[pyarrow]] | zero-copy | +| int64:standard | [0, 1, -1, 9223372036854775807, -9223372036854775808]@int64 | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64[pyarrow]] | [0, 1, -1, 9223372036854775807, -9223372036854775808]@Series[int64[pyarrow]] | zero-copy | +| int64:nullable | [0, 1, None]@int64 | [0, 1, <NA>]@Series[int64[pyarrow]] | [0, 1, <NA>]@Series[int64[pyarrow]] | zero-copy | +| int64:empty | []@int64 | []@Series[int64[pyarrow]] | []@Series[int64[pyarrow]] | zero-copy | +| uint8:standard | [0, 1, 255]@uint8 | [0, 1, 255]@Series[uint8[pyarrow]] | [0, 1, 255]@Series[uint8[pyarrow]] | zero-copy | +| uint8:nullable | [0, 1, None]@uint8 | [0, 1, <NA>]@Series[uint8[pyarrow]] | [0, 1, <NA>]@Series[uint8[pyarrow]] | zero-copy | +| uint8:empty | []@uint8 | []@Series[uint8[pyarrow]] | []@Series[uint8[pyarrow]] | zero-copy | +| uint16:standard | [0, 1, 65535]@uint16 | [0, 1, 65535]@Series[uint16[pyarrow]] | [0, 1, 65535]@Series[uint16[pyarrow]] | zero-copy | +| uint16:nullable | [0, 1, None]@uint16 | [0, 1, <NA>]@Series[uint16[pyarrow]] | [0, 1, <NA>]@Series[uint16[pyarrow]] | zero-copy | +| uint16:empty | []@uint16 | []@Series[uint16[pyarrow]] | []@Series[uint16[pyarrow]] | zero-copy | +| uint32:standard | [0, 1, 4294967295]@uint32 | [0, 1, 4294967295]@Series[uint32[pyarrow]] | [0, 1, 4294967295]@Series[uint32[pyarrow]] | zero-copy | +| uint32:nullable | [0, 1, None]@uint32 | [0, 1, <NA>]@Series[uint32[pyarrow]] | [0, 1, <NA>]@Series[uint32[pyarrow]] | zero-copy | +| uint32:empty | []@uint32 | []@Series[uint32[pyarrow]] | []@Series[uint32[pyarrow]] | zero-copy | +| uint64:standard | [0, 1, 18446744073709551615]@uint64 | [0, 1, 18446744073709551615]@Series[uint64[pyarrow]] | [0, 1, 18446744073709551615]@Series[uint64[pyarrow]] | zero-copy | +| uint64:nullable | [0, 1, None]@uint64 | [0, 1, <NA>]@Series[uint64[pyarrow]] | [0, 1, <NA>]@Series[uint64[pyarrow]] | zero-copy | +| uint64:empty | []@uint64 | []@Series[uint64[pyarrow]] | []@Series[uint64[pyarrow]] | zero-copy | +| float32:standard | [0.0, 1.5, -1.5]@float32 | [0.0, 1.5, -1.5]@Series[float[pyarrow]] | [0.0, 1.5, -1.5]@Series[float[pyarrow]] | zero-copy | +| float32:nullable | [0.0, 1.5, None]@float32 | [0.0, 1.5, <NA>]@Series[float[pyarrow]] | [0.0, 1.5, <NA>]@Series[float[pyarrow]] | zero-copy | +| float32:empty | []@float32 | []@Series[float[pyarrow]] | []@Series[float[pyarrow]] | zero-copy | +| float64:standard | [0.0, 1.5, -1.5]@float64 | [0.0, 1.5, -1.5]@Series[double[pyarrow]] | [0.0, 1.5, -1.5]@Series[double[pyarrow]] | zero-copy | +| float64:nullable | [0.0, 1.5, None]@float64 | [0.0, 1.5, <NA>]@Series[double[pyarrow]] | [0.0, 1.5, <NA>]@Series[double[pyarrow]] | zero-copy | +| float64:special | [nan, inf, -inf]@float64 | [nan, inf, -inf]@Series[double[pyarrow]] | [nan, inf, -inf]@Series[double[pyarrow]] | zero-copy | +| float64:empty | []@float64 | []@Series[double[pyarrow]] | []@Series[double[pyarrow]] | zero-copy | +| bool:standard | [True, False, True]@bool | [True, False, True]@Series[bool[pyarrow]] | [True, False, True]@Series[bool[pyarrow]] | zero-copy | +| bool:nullable | [True, False, None]@bool | [True, False, <NA>]@Series[bool[pyarrow]] | [True, False, <NA>]@Series[bool[pyarrow]] | zero-copy | +| bool:empty | []@bool | []@Series[bool[pyarrow]] | []@Series[bool[pyarrow]] | zero-copy | +| string:standard | [hello, world, ]@string | ['hello', 'world', '']@Series[string[pyarrow]] | ['hello', 'world', '']@Series[string[pyarrow]] | zero-copy | +| string:nullable | [hello, None, world]@string | ['hello', <NA>, 'world']@Series[string[pyarrow]] | ['hello', <NA>, 'world']@Series[string[pyarrow]] | zero-copy | +| string:empty | []@string | []@Series[string[pyarrow]] | []@Series[string[pyarrow]] | zero-copy | +| large_string:standard | [hello, world]@large_string | ['hello', 'world']@Series[large_string[pyarrow]] | ['hello', 'world']@Series[large_string[pyarrow]] | zero-copy | +| large_string:nullable | [hello, None]@large_string | ['hello', <NA>]@Series[large_string[pyarrow]] | ['hello', <NA>]@Series[large_string[pyarrow]] | zero-copy | +| large_string:empty | []@large_string | []@Series[large_string[pyarrow]] | []@Series[large_string[pyarrow]] | zero-copy | +| binary:standard | [b'hello', b'world']@binary | [b'hello', b'world']@Series[binary[pyarrow]] | [b'hello', b'world']@Series[binary[pyarrow]] | zero-copy | +| binary:nullable | [b'hello', None]@binary | [b'hello', <NA>]@Series[binary[pyarrow]] | [b'hello', <NA>]@Series[binary[pyarrow]] | zero-copy | +| binary:empty | []@binary | []@Series[binary[pyarrow]] | []@Series[binary[pyarrow]] | zero-copy | +| large_binary:standard | [b'hello', b'world']@large_binary | [b'hello', b'world']@Series[large_binary[pyarrow]] | [b'hello', b'world']@Series[large_binary[pyarrow]] | zero-copy | +| large_binary:nullable | [b'hello', None]@large_binary | [b'hello', <NA>]@Series[large_binary[pyarrow]] | [b'hello', <NA>]@Series[large_binary[pyarrow]] | zero-copy | +| large_binary:empty | []@large_binary | []@Series[large_binary[pyarrow]] | []@Series[large_binary[pyarrow]] | zero-copy | +| decimal128:standard | [1.23, 4.56, -7.89]@decimal128(5, 2) | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[decimal128(5, 2)[pyarrow]] | [Decimal('1.23'), Decimal('4.56'), Decimal('-7.89')]@Series[decimal128(5, 2)[pyarrow]] | zero-copy | +| decimal128:nullable | [1.23, None, 4.56]@decimal128(5, 2) | [Decimal('1.23'), <NA>, Decimal('4.56')]@Series[decimal128(5, 2)[pyarrow]] | [Decimal('1.23'), <NA>, Decimal('4.56')]@Series[decimal128(5, 2)[pyarrow]] | zero-copy | +| decimal128:empty | []@decimal128(5, 2) | []@Series[decimal128(5, 2)[pyarrow]] | []@Series[decimal128(5, 2)[pyarrow]] | zero-copy | +| date32:standard | [2024-01-01, 2024-06-15]@date32[day] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date32[day][pyarrow]] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date32[day][pyarrow]] | zero-copy | +| date32:nullable | [2024-01-01, None]@date32[day] | [datetime.date(2024, 1, 1), <NA>]@Series[date32[day][pyarrow]] | [datetime.date(2024, 1, 1), <NA>]@Series[date32[day][pyarrow]] | zero-copy | +| date32:empty | []@date32[day] | []@Series[date32[day][pyarrow]] | []@Series[date32[day][pyarrow]] | zero-copy | +| date64:standard | [2024-01-01, 2024-06-15]@date64[ms] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date64[ms][pyarrow]] | [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15)]@Series[date64[ms][pyarrow]] | zero-copy | +| date64:nullable | [2024-01-01, None]@date64[ms] | [datetime.date(2024, 1, 1), <NA>]@Series[date64[ms][pyarrow]] | [datetime.date(2024, 1, 1), <NA>]@Series[date64[ms][pyarrow]] | zero-copy | +| date64:empty | []@date64[ms] | []@Series[date64[ms][pyarrow]] | []@Series[date64[ms][pyarrow]] | zero-copy | +| timestamp[s]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[s][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[s][pyarrow]] | zero-copy | +| timestamp[s]:nullable | [2024-01-01 12:00:00, None]@timestamp[s] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[s][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[s][pyarrow]] | zero-copy | +| timestamp[s]:empty | []@timestamp[s] | []@Series[timestamp[s][pyarrow]] | []@Series[timestamp[s][pyarrow]] | zero-copy | +| timestamp[ms]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ms][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ms][pyarrow]] | zero-copy | +| timestamp[ms]:nullable | [2024-01-01 12:00:00, None]@timestamp[ms] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ms][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ms][pyarrow]] | zero-copy | +| timestamp[ms]:empty | []@timestamp[ms] | []@Series[timestamp[ms][pyarrow]] | []@Series[timestamp[ms][pyarrow]] | zero-copy | +| timestamp[us]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[us][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[us][pyarrow]] | zero-copy | +| timestamp[us]:nullable | [2024-01-01 12:00:00, None]@timestamp[us] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[us][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[us][pyarrow]] | zero-copy | +| timestamp[us]:empty | []@timestamp[us] | []@Series[timestamp[us][pyarrow]] | []@Series[timestamp[us][pyarrow]] | zero-copy | +| timestamp[ns]:standard | [2024-01-01 12:00:00, 2024-06-15 18:30:00]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ns][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), Timestamp('2024-06-15 18:30:00')]@Series[timestamp[ns][pyarrow]] | zero-copy | +| timestamp[ns]:nullable | [2024-01-01 12:00:00, None]@timestamp[ns] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ns][pyarrow]] | [Timestamp('2024-01-01 12:00:00'), <NA>]@Series[timestamp[ns][pyarrow]] | zero-copy | +| timestamp[ns]:empty | []@timestamp[ns] | []@Series[timestamp[ns][pyarrow]] | []@Series[timestamp[ns][pyarrow]] | zero-copy | +| timestamp[us,tz=UTC]:standard | [2024-01-01 12:00:00+00:00, 2024-06-15 18:30:00+00:00]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[timestamp[us, tz=UTC][pyarrow]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), Timestamp('2024-06-15 18:30:00+0000', tz='UTC')]@Series[timestamp[us, tz=UTC][pyarrow]] | zero-copy | +| timestamp[us,tz=UTC]:nullable | [2024-01-01 12:00:00+00:00, None]@timestamp[us, tz=UTC] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), <NA>]@Series[timestamp[us, tz=UTC][pyarrow]] | [Timestamp('2024-01-01 12:00:00+0000', tz='UTC'), <NA>]@Series[timestamp[us, tz=UTC][pyarrow]] | zero-copy | +| timestamp[us,tz=UTC]:empty | []@timestamp[us, tz=UTC] | []@Series[timestamp[us, tz=UTC][pyarrow]] | []@Series[timestamp[us, tz=UTC][pyarrow]] | zero-copy | +| duration[s]:standard | [1 day, 0:00:00, 2:30:00]@duration[s] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[s][pyarrow]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[s][pyarrow]] | zero-copy | +| duration[s]:nullable | [1 day, 0:00:00, None]@duration[s] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[s][pyarrow]] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[s][pyarrow]] | zero-copy | +| duration[s]:empty | []@duration[s] | []@Series[duration[s][pyarrow]] | []@Series[duration[s][pyarrow]] | zero-copy | +| duration[ms]:standard | [1 day, 0:00:00, 2:30:00]@duration[ms] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ms][pyarrow]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ms][pyarrow]] | zero-copy | +| duration[ms]:nullable | [1 day, 0:00:00, None]@duration[ms] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ms][pyarrow]] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ms][pyarrow]] | zero-copy | +| duration[ms]:empty | []@duration[ms] | []@Series[duration[ms][pyarrow]] | []@Series[duration[ms][pyarrow]] | zero-copy | +| duration[us]:standard | [1 day, 0:00:00, 2:30:00]@duration[us] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[us][pyarrow]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[us][pyarrow]] | zero-copy | +| duration[us]:nullable | [1 day, 0:00:00, None]@duration[us] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[us][pyarrow]] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[us][pyarrow]] | zero-copy | +| duration[us]:empty | []@duration[us] | []@Series[duration[us][pyarrow]] | []@Series[duration[us][pyarrow]] | zero-copy | +| duration[ns]:standard | [1 days 00:00:00, 0 days 02:30:00]@duration[ns] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ns][pyarrow]] | [Timedelta('1 days 00:00:00'), Timedelta('0 days 02:30:00')]@Series[duration[ns][pyarrow]] | zero-copy | +| duration[ns]:nullable | [1 days 00:00:00, None]@duration[ns] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ns][pyarrow]] | [Timedelta('1 days 00:00:00'), <NA>]@Series[duration[ns][pyarrow]] | zero-copy | +| duration[ns]:empty | []@duration[ns] | []@Series[duration[ns][pyarrow]] | []@Series[duration[ns][pyarrow]] | zero-copy | +| time32[s]:standard | [12:30:00, 18:45:30]@time32[s] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[s][pyarrow]] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[s][pyarrow]] | zero-copy | +| time32[s]:nullable | [12:30:00, None]@time32[s] | [datetime.time(12, 30), <NA>]@Series[time32[s][pyarrow]] | [datetime.time(12, 30), <NA>]@Series[time32[s][pyarrow]] | zero-copy | +| time32[s]:empty | []@time32[s] | []@Series[time32[s][pyarrow]] | []@Series[time32[s][pyarrow]] | zero-copy | +| time32[ms]:standard | [12:30:00, 18:45:30]@time32[ms] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[ms][pyarrow]] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time32[ms][pyarrow]] | zero-copy | +| time32[ms]:nullable | [12:30:00, None]@time32[ms] | [datetime.time(12, 30), <NA>]@Series[time32[ms][pyarrow]] | [datetime.time(12, 30), <NA>]@Series[time32[ms][pyarrow]] | zero-copy | +| time32[ms]:empty | []@time32[ms] | []@Series[time32[ms][pyarrow]] | []@Series[time32[ms][pyarrow]] | zero-copy | +| time64[us]:standard | [12:30:00, 18:45:30]@time64[us] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[us][pyarrow]] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[us][pyarrow]] | zero-copy | +| time64[us]:nullable | [12:30:00, None]@time64[us] | [datetime.time(12, 30), <NA>]@Series[time64[us][pyarrow]] | [datetime.time(12, 30), <NA>]@Series[time64[us][pyarrow]] | zero-copy | +| time64[us]:empty | []@time64[us] | []@Series[time64[us][pyarrow]] | []@Series[time64[us][pyarrow]] | zero-copy | +| time64[ns]:standard | [12:30:00, 18:45:30]@time64[ns] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[ns][pyarrow]] | [datetime.time(12, 30), datetime.time(18, 45, 30)]@Series[time64[ns][pyarrow]] | zero-copy | +| time64[ns]:nullable | [12:30:00, None]@time64[ns] | [datetime.time(12, 30), <NA>]@Series[time64[ns][pyarrow]] | [datetime.time(12, 30), <NA>]@Series[time64[ns][pyarrow]] | zero-copy | +| time64[ns]:empty | []@time64[ns] | []@Series[time64[ns][pyarrow]] | []@Series[time64[ns][pyarrow]] | zero-copy | +| null:standard | [None, None, None]@null | [<NA>, <NA>, <NA>]@Series[null[pyarrow]] | [<NA>, <NA>, <NA>]@Series[null[pyarrow]] | zero-copy | +| null:empty | []@null | []@Series[null[pyarrow]] | []@Series[null[pyarrow]] | zero-copy | +| list<int64>:standard | [[1, 2], [3, 4, 5]]@list<item: int64> | [[1, 2], [3, 4, 5]]@Series[list<item: int64>[pyarrow]] | [[1, 2], [3, 4, 5]]@Series[list<item: int64>[pyarrow]] | zero-copy | +| list<int64>:nullable | [[1, 2], None, [3]]@list<item: int64> | [[1, 2], <NA>, [3]]@Series[list<item: int64>[pyarrow]] | [[1, 2], <NA>, [3]]@Series[list<item: int64>[pyarrow]] | zero-copy | +| list<int64>:empty | []@list<item: int64> | []@Series[list<item: int64>[pyarrow]] | []@Series[list<item: int64>[pyarrow]] | zero-copy | +| list<string>:standard | [['a', 'b'], ['c']]@list<item: string> | [['a', 'b'], ['c']]@Series[list<item: string>[pyarrow]] | [['a', 'b'], ['c']]@Series[list<item: string>[pyarrow]] | zero-copy | +| large_list<int64>:standard | [[1, 2], [3, 4]]@large_list<item: int64> | [[1, 2], [3, 4]]@Series[large_list<item: int64>[pyarrow]] | [[1, 2], [3, 4]]@Series[large_list<item: int64>[pyarrow]] | zero-copy | +| large_list<int64>:empty | []@large_list<item: int64> | []@Series[large_list<item: int64>[pyarrow]] | []@Series[large_list<item: int64>[pyarrow]] | zero-copy | +| fixed_size_list<int64>[3]:standard | [[1, 2, 3], [4, 5, 6]]@fixed_size_list<item: int64>[3] | [[1, 2, 3], [4, 5, 6]]@Series[fixed_size_list<item: int64>[3][pyarrow]] | [[1, 2, 3], [4, 5, 6]]@Series[fixed_size_list<item: int64>[3][pyarrow]] | zero-copy | +| fixed_size_list<int64>[3]:empty | []@fixed_size_list<item: int64>[3] | []@Series[fixed_size_list<item: int64>[3][pyarrow]] | []@Series[fixed_size_list<item: int64>[3][pyarrow]] | zero-copy | +| struct:standard | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] | zero-copy | +| struct:nullable | [[('x', 1), ('y', 'a')], None]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, <NA>]@Series[struct<x: int64, y: string>[pyarrow]] | [{'x': 1, 'y': 'a'}, <NA>]@Series[struct<x: int64, y: string>[pyarrow]] | zero-copy | +| struct:empty | []@struct<x: int64, y: string> | []@Series[struct<x: int64, y: string>[pyarrow]] | []@Series[struct<x: int64, y: string>[pyarrow]] | zero-copy | +| map<string,int64>:standard | [[('a', 1), ('b', 2)], [('c', 3)]]@map<string, int64> | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[map<string, int64>[pyarrow]] | [[('a', 1), ('b', 2)], [('c', 3)]]@Series[map<string, int64>[pyarrow]] | zero-copy | +| map<string,int64>:empty | []@map<string, int64> | []@Series[map<string, int64>[pyarrow]] | []@Series[map<string, int64>[pyarrow]] | zero-copy | +| list<list<int64>>:standard | [[[1, 2], [3]], [[4, 5, 6]]]@list<item: list<item: int64>> | [[[1, 2], [3]], [[4, 5, 6]]]@Series[list<item: list<item: int64>>[pyarrow]] | [[[1, 2], [3]], [[4, 5, 6]]]@Series[list<item: list<item: int64>>[pyarrow]] | zero-copy | +| list<struct>:standard | [[{'x': 1}, {'x': 2}], [{'x': 3}]]@list<item: struct<x: int64>> | [[{'x': 1}, {'x': 2}], [{'x': 3}]]@Series[list<item: struct<x: int64>>[pyarrow]] | [[{'x': 1}, {'x': 2}], [{'x': 3}]]@Series[list<item: struct<x: int64>>[pyarrow]] | zero-copy | +| list<map<string,int64>>:standard | [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@list<item: map<string, int64>> | [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@Series[list<item: map<string, int64>>[pyarrow]] | [[[('a', 1)], [('b', 2)]], [[('c', 3)]]]@Series[list<item: map<string, int64>>[pyarrow]] | zero-copy | +| struct<struct>:standard | [[('outer', {'inner': 1})], [('outer', {'inner': 2})]]@struct<outer: struct<inner: int64>> | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[struct<outer: struct<inner: int64>>[pyarrow]] | [{'outer': {'inner': 1}}, {'outer': {'inner': 2}}]@Series[struct<outer: struct<inner: int64>>[pyarrow]] | zero-copy | +| struct<list<int64>>:standard | [[('items', [1, 2, 3])], [('items', [4, 5])]]@struct<items: list<item: int64>> | [{'items': [1, 2, 3]}, {'items': [4, 5]}]@Series[struct<items: list<item: int64>>[pyarrow]] | [{'items': [1, 2, 3]}, {'items': [4, 5]}]@Series[struct<items: list<item: int64>>[pyarrow]] | zero-copy | +| struct<map<string,int64>>:standard | [[('mapping', [('a', 1)])], [('mapping', [('b', 2)])]]@struct<mapping: map<string, int64>> | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[struct<mapping: map<string, int64>>[pyarrow]] | [{'mapping': [('a', 1)]}, {'mapping': [('b', 2)]}]@Series[struct<mapping: map<string, int64>>[pyarrow]] | zero-copy | +| map<string,list<int64>>:standard | [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@map<string, list<item: int64>> | [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@Series[map<string, list<item: int64>>[pyarrow]] | [[('a', [1, 2]), ('b', [3])], [('c', [4, 5, 6])]]@Series[map<string, list<item: int64>>[pyarrow]] | zero-copy | +| map<string,struct>:standard | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@map<string, struct<v: int64>> | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[map<string, struct<v: int64>>[pyarrow]] | [[('a', {'v': 1}), ('b', {'v': 2})], [('c', {'v': 3})]]@Series[map<string, struct<v: int64>>[pyarrow]] | zero-copy | +| map<string,map<string,int64>>:standard | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@map<string, map<string, int64>> | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[map<string, map<string, int64>>[pyarrow]] | [[('a', [('x', 1)]), ('b', [('y', 2)])], [('c', [('z', 3)])]]@Series[map<string, map<string, int64>>[pyarrow]] | zero-copy | +| dictionary<int32,string>:standard | [a, b, a, b]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', 'a', 'b']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | ['a', 'b', 'a', 'b']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | zero-copy | +| dictionary<int32,string>:nullable | [a, b, None, a]@dictionary<values=string, indices=int32, ordered=0> | ['a', 'b', <NA>, 'a']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | ['a', 'b', <NA>, 'a']@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | zero-copy | +| dictionary<int32,string>:empty | []@dictionary<values=string, indices=int32, ordered=0> | []@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | []@Series[dictionary<values=string, indices=int32, ordered=0>[pyarrow]] | zero-copy | +| int64:zero-chunk | []@int64 | []@Series[int64[pyarrow]] | []@Series[int64[pyarrow]] | zero-copy | +| int64:empty-chunk | []@int64 | []@Series[int64[pyarrow]] | []@Series[int64[pyarrow]] | zero-copy | +| int64:single-chunk | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64[pyarrow]] | [1, 2, 3]@Series[int64[pyarrow]] | zero-copy | +| int64:multi-chunk | [1, 2, 3, 4]@int64 | [1, 2, 3, 4]@Series[int64[pyarrow]] | [1, 2, 3, 4]@Series[int64[pyarrow]] | zero-copy | +| int64:multi-chunk-nullable | [1, None, 2]@int64 | [1, <NA>, 2]@Series[int64[pyarrow]] | [1, <NA>, 2]@Series[int64[pyarrow]] | zero-copy | +| int64:multi-chunk-with-empty | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64[pyarrow]] | [1, 2, 3]@Series[int64[pyarrow]] | zero-copy | +| float64:multi-chunk | [1.5, 2.5, 3.5]@float64 | [1.5, 2.5, 3.5]@Series[double[pyarrow]] | [1.5, 2.5, 3.5]@Series[double[pyarrow]] | zero-copy | +| string:single-chunk | [a, b]@string | ['a', 'b']@Series[string[pyarrow]] | ['a', 'b']@Series[string[pyarrow]] | zero-copy | +| string:multi-chunk | [a, b, c]@string | ['a', 'b', 'c']@Series[string[pyarrow]] | ['a', 'b', 'c']@Series[string[pyarrow]] | zero-copy | +| string:multi-chunk-nullable | [a, None, c]@string | ['a', <NA>, 'c']@Series[string[pyarrow]] | ['a', <NA>, 'c']@Series[string[pyarrow]] | zero-copy | +| list<int64>:multi-chunk | [[1, 2], [3], [4]]@list<item: int64> | [[1, 2], [3], [4]]@Series[list<item: int64>[pyarrow]] | [[1, 2], [3], [4]]@Series[list<item: int64>[pyarrow]] | zero-copy | +| struct:multi-chunk | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[struct<x: int64, y: string>[pyarrow]] | zero-copy | +| int64:sliced | [2, 3, 4]@int64 | [2, 3, 4]@Series[int64[pyarrow]] | [2, 3, 4]@Series[int64[pyarrow]] | zero-copy | +| int64:sliced-with-null | [2, None, 4]@int64 | [2, <NA>, 4]@Series[int64[pyarrow]] | [2, <NA>, 4]@Series[int64[pyarrow]] | zero-copy | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.csv new file mode 100644 index 0000000000000..7ab114cfbf897 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.csv @@ -0,0 +1,11 @@ +test case pyarrow array default spark memory options source readable after self_destruct +int64:single-chunk [1, 2, 3]@int64 [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] readable +int64:multi-chunk [1, 2, 3, 4]@int64 [1, 2, 3, 4]@Series[int64] [1, 2, 3, 4]@Series[int64] readable +int64:multi-chunk-nullable [1, None, 2]@int64 [1.0, nan, 2.0]@Series[float64] [1.0, nan, 2.0]@Series[float64] readable +int64:multi-chunk-with-empty [1, 2, 3]@int64 [1, 2, 3]@Series[int64] [1, 2, 3]@Series[int64] readable +int64:empty-chunk []@int64 []@Series[int64] []@Series[int64] readable +float64:multi-chunk [1.5, 2.5, 3.5]@float64 [1.5, 2.5, 3.5]@Series[float64] [1.5, 2.5, 3.5]@Series[float64] readable +string:multi-chunk [a, b, c]@string ['a', 'b', 'c']@Series[object] ['a', 'b', 'c']@Series[object] readable +string:multi-chunk-nullable [a, None, c]@string ['a', None, 'c']@Series[object] ['a', None, 'c']@Series[object] readable +list<int64>:multi-chunk [[1, 2], [3], [4]]@list<item: int64> [array([1, 2]), array([3]), array([4])]@Series[object] [array([1, 2]), array([3]), array([4])]@Series[object] readable +struct:multi-chunk [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] readable diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.md new file mode 100644 index 0000000000000..e41a51201015c --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_chunked_array_to_pandas_memory_flags.md @@ -0,0 +1,12 @@ +| test case | pyarrow array | default | spark memory options | source readable after self_destruct | +|------------------------------|------------------------------------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|---------------------------------------| +| int64:single-chunk | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | readable | +| int64:multi-chunk | [1, 2, 3, 4]@int64 | [1, 2, 3, 4]@Series[int64] | [1, 2, 3, 4]@Series[int64] | readable | +| int64:multi-chunk-nullable | [1, None, 2]@int64 | [1.0, nan, 2.0]@Series[float64] | [1.0, nan, 2.0]@Series[float64] | readable | +| int64:multi-chunk-with-empty | [1, 2, 3]@int64 | [1, 2, 3]@Series[int64] | [1, 2, 3]@Series[int64] | readable | +| int64:empty-chunk | []@int64 | []@Series[int64] | []@Series[int64] | readable | +| float64:multi-chunk | [1.5, 2.5, 3.5]@float64 | [1.5, 2.5, 3.5]@Series[float64] | [1.5, 2.5, 3.5]@Series[float64] | readable | +| string:multi-chunk | [a, b, c]@string | ['a', 'b', 'c']@Series[object] | ['a', 'b', 'c']@Series[object] | readable | +| string:multi-chunk-nullable | [a, None, c]@string | ['a', None, 'c']@Series[object] | ['a', None, 'c']@Series[object] | readable | +| list<int64>:multi-chunk | [[1, 2], [3], [4]]@list<item: int64> | [array([1, 2]), array([3]), array([4])]@Series[object] | [array([1, 2]), array([3]), array([4])]@Series[object] | readable | +| struct:multi-chunk | [[('x', 1), ('y', 'a')], [('x', 2), ('y', 'b')]]@struct<x: int64, y: string> | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | [{'x': 1, 'y': 'a'}, {'x': 2, 'y': 'b'}]@Series[object] | readable | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv new file mode 100644 index 0000000000000..b918763436dae --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.csv @@ -0,0 +1,11 @@ +test case pandas dataframe preserve_index=None preserve_index=False preserve_index=True +0-columns:range-index {}@Dataframe[][index=RangeIndex[0:3:1]] {}@RecordBatch[][num_rows=3] {}@RecordBatch[][num_rows=0] {__index_level_0__: [0, 1, 2]}@RecordBatch[__index_level_0__: int64][num_rows=3] +0-columns:named-index {}@Dataframe[][index='idx':[100, 200, 300]] {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] {}@RecordBatch[][num_rows=0] {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] +0-columns:unnamed-index {}@Dataframe[][index=None:[10, 20, 30]] {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] {}@RecordBatch[][num_rows=0] {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] +0-columns:empty {}@Dataframe[][index=RangeIndex[0:0:1]] {}@RecordBatch[][num_rows=0] {}@RecordBatch[][num_rows=0] {__index_level_0__: []}@RecordBatch[__index_level_0__: int64][num_rows=0] +single-column:range-index {'a': [1, 2, 3]}@Dataframe[a int64][index=RangeIndex[0:3:1]] {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3], __index_level_0__: [0, 1, 2]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] +single-column:named-index {'a': [1, 2, 3]}@Dataframe[a int64][index='idx':[100, 200, 300]] {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] +single-column:unnamed-index {'a': [1, 2, 3]}@Dataframe[a int64][index=None:[10, 20, 30]] {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] +multi-column:standard {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] +multi-column:nullable {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT, Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] +multi-column:no-rows {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64, f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] {i: [], f: [], b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0] {i: [], f: [], b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0] {i: [], f: [], b: [], t: [], __index_level_0__: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns], __index_level_0__: int64][num_rows=0] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md new file mode 100644 index 0000000000000..bc1bc210c5360 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_record_batch_from_pandas.md @@ -0,0 +1,12 @@ +| test case | pandas dataframe | preserve_index=None | preserve_index=False | preserve_index=True | +|-----------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| 0-columns:range-index | {}@Dataframe[][index=RangeIndex[0:3:1]] | {}@RecordBatch[][num_rows=3] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: [0, 1, 2]}@RecordBatch[__index_level_0__: int64][num_rows=3] | +| 0-columns:named-index | {}@Dataframe[][index='idx':[100, 200, 300]] | {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] | {}@RecordBatch[][num_rows=0] | {idx: [100, 200, 300]}@RecordBatch[idx: int64][num_rows=3] | +| 0-columns:unnamed-index | {}@Dataframe[][index=None:[10, 20, 30]] | {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: [10, 20, 30]}@RecordBatch[__index_level_0__: int64][num_rows=3] | +| 0-columns:empty | {}@Dataframe[][index=RangeIndex[0:0:1]] | {}@RecordBatch[][num_rows=0] | {}@RecordBatch[][num_rows=0] | {__index_level_0__: []}@RecordBatch[__index_level_0__: int64][num_rows=0] | +| single-column:range-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=RangeIndex[0:3:1]] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], __index_level_0__: [0, 1, 2]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] | +| single-column:named-index | {'a': [1, 2, 3]}@Dataframe[a int64][index='idx':[100, 200, 300]] | {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], idx: [100, 200, 300]}@RecordBatch[a: int64, idx: int64][num_rows=3] | +| single-column:unnamed-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=None:[10, 20, 30]] | {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] | {a: [1, 2, 3]}@RecordBatch[a: int64][num_rows=3] | {a: [1, 2, 3], __index_level_0__: [10, 20, 30]}@RecordBatch[a: int64, __index_level_0__: int64][num_rows=3] | +| multi-column:standard | {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | {i: [1, 2, 3], f: [1.5, 2.5, 3.5], b: [True, False, True], s: [a, b, c], t: [2020-01-01 05:30:00, 2020-01-01 05:30:00, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[i: int64, f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] | +| multi-column:nullable | {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT, Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns]][num_rows=3] | {f: [1.5, None, 3.5], b: [True, None, False], s: [a, None, c], t: [2020-01-01 05:30:00, None, 2020-01-01 05:30:00], __index_level_0__: [0, 1, 2]}@RecordBatch[f: float64, b: bool, s: string, t: timestamp[ns], __index_level_0__: int64][num_rows=3] | +| multi-column:no-rows | {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64, f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] | {i: [], f: [], b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0] | {i: [], f: [], b: [], t: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns]][num_rows=0] | {i: [], f: [], b: [], t: [], __index_level_0__: []}@RecordBatch[i: int64, f: float64, b: bool, t: timestamp[ns], __index_level_0__: int64][num_rows=0] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.csv index 129d4b535179f..a4b523de7a0d5 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.csv +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.csv @@ -7,7 +7,7 @@ int16:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int6 int16:max_min ERR@ArrowInvalid [32767, -32768, None]@int16 [32767, -32768, None]@int32 [32767, -32768, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [32768.0, -32768.0, None]@float16 [32767.0, -32768.0, None]@float32 [32767.0, -32768.0, None]@float64 [True, True, None]@bool [32767, -32768, None]@string [32767, -32768, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int32:standard [0, 1, None]@int8 [0, 1, None]@int16 [0, 1, None]@int32 [0, 1, None]@int64 [0, 1, None]@uint8 [0, 1, None]@uint16 [0, 1, None]@uint32 [0, 1, None]@uint64 [0.0, 1.0, None]@float16 [0.0, 1.0, None]@float32 [0.0, 1.0, None]@float64 [False, True, None]@bool [0, 1, None]@string [0, 1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0E-10, 1.0000000000, None]@decimal128(38, 10) [0E-10, 1.0000000000, None]@decimal256(76, 10) [1970-01-01, 1970-01-02, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:00:00, 00:00:01, None]@time32[s] [00:00:00, 00:00:00.001000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int32:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [-1.0, None]@float16 [-1.0, None]@float32 [-1.0, None]@float64 [True, None]@bool [-1, None]@string [-1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [-1.0000000000, None]@decimal128(38, 10) [-1.0000000000, None]@decimal256(76, 10) [1969-12-31, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [23:59:59, None]@time32[s] [23:59:59.999000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError -int32:max_min ERR@ArrowInvalid ERR@ArrowInvalid [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) ERR@OverflowError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [03:14:07, 20:45:52, None]@time32[s] [20:31:23.647000, 03:28:36.352000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError +int32:max_min ERR@ArrowInvalid ERR@ArrowInvalid [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) [temporal overflow, temporal overflow, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [03:14:07, 20:45:52, None]@time32[s] [20:31:23.647000, 03:28:36.352000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int64:standard [0, 1, None]@int8 [0, 1, None]@int16 [0, 1, None]@int32 [0, 1, None]@int64 [0, 1, None]@uint8 [0, 1, None]@uint16 [0, 1, None]@uint32 [0, 1, None]@uint64 [0.0, 1.0, None]@float16 [0.0, 1.0, None]@float32 [0.0, 1.0, None]@float64 [False, True, None]@bool [0, 1, None]@string [0, 1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0E-10, 1.0000000000, None]@decimal128(38, 10) [0E-10, 1.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1970-01-01, 1970-01-01, None]@date64[ms] [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] [0:00:00, 0:00:01, None]@duration[s] [0:00:00, 0:00:00.001000, None]@duration[ms] [0:00:00, 0:00:00.000001, None]@duration[us] [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:00:00, 00:00:00.000001, None]@time64[us] [00:00:00, 00:00:00, None]@time64[ns] int64:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [-1.0, None]@float16 [-1.0, None]@float32 [-1.0, None]@float64 [True, None]@bool [-1, None]@string [-1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [-1.0000000000, None]@decimal128(38, 10) [-1.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1969-12-31, None]@date64[ms] [1969-12-31 23:59:59, None]@timestamp[s] [1969-12-31 23:59:59.999000, None]@timestamp[ms] [1969-12-31 23:59:59.999999, None]@timestamp[us] [1969-12-31 23:59:59.999999999, None]@timestamp[ns] [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] [-1 day, 23:59:59, None]@duration[s] [-1 day, 23:59:59.999000, None]@duration[ms] [-1 day, 23:59:59.999999, None]@duration[us] [-1 days +23:59:59.999999999, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [23:59:59.999999, None]@time64[us] [23:59:59.999999, None]@time64[ns] int64:max_min ERR@ArrowInvalid ERR@ArrowInvalid [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [inf, -inf, None]@float16 ERR@ArrowInvalid [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1970-01-25, 1969-12-07, None]@date64[ms] [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:35:47.483647, 23:24:12.516352, None]@time64[us] [00:00:02.147483, 23:59:57.852516, None]@time64[ns] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.md index 1e2351c6a2847..9e27ec5dce3b6 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.md +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_safe.md @@ -1,85 +1,85 @@ -| source \ target | int8 | int16 | int32 | int64 | uint8 | uint16 | uint32 | uint64 | float16 | float32 | float64 | bool | string | large_string | binary | large_binary | fixed_size_binary[16] | decimal128(38, 10) | decimal256(76, 10) | date32[day] | date64[ms] | timestamp[s] | timestamp[ms] | timestamp[us] | timestamp[ns] | timestamp[s, tz=UTC] | timestamp[ms, tz=UTC] | timestamp[us, tz=UTC] | timestamp[ns, tz=UTC] | timestamp[s, tz=America/New_York] | timestamp[s, tz=Asia/Shanghai] | duration[s] | duration[ms] | duration[us] | duration[ns] | time32[s] | time32[ms] | time64[us] | time64[ns] | -|--------------------------------------------|------------------------------|------------------------------|---------------------------------------|---------------------------------------|------------------------------|------------------------------|------------------------------|------------------------------|------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|--------------------------------|-------------------------------------------------------------------------------|-------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------|-------------------------------------------|---------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------|----------------------------------------------------------------|-------------------------------------------------------------------------|--------------------------------------------------------------|-----------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------| -| int8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int8:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int8:max_min | [127, -128, None]@int8 | [127, -128, None]@int16 | [127, -128, None]@int32 | [127, -128, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [127.0, -128.0, None]@float16 | [127.0, -128.0, None]@float32 | [127.0, -128.0, None]@float64 | [True, True, None]@bool | [127, -128, None]@string | [127, -128, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [127.0000000000, -128.0000000000, None]@decimal128(38, 10) | [127.0000000000, -128.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:max_min | ERR@ArrowInvalid | [32767, -32768, None]@int16 | [32767, -32768, None]@int32 | [32767, -32768, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [32768.0, -32768.0, None]@float16 | [32767.0, -32768.0, None]@float32 | [32767.0, -32768.0, None]@float64 | [True, True, None]@bool | [32767, -32768, None]@string | [32767, -32768, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) | [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | [1970-01-01, 1970-01-02, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | [1969-12-31, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:max_min | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@OverflowError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [03:14:07, 20:45:52, None]@time32[s] | [20:31:23.647000, 03:28:36.352000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | -| int64:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| int64:max_min | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, None]@float16 | ERR@ArrowInvalid | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-25, 1969-12-07, None]@date64[ms] | [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] | [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] | [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] | [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] | [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] | [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] | [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] | [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] | [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] | [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] | [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:35:47.483647, 23:24:12.516352, None]@time64[us] | [00:00:02.147483, 23:59:57.852516, None]@time64[ns] | -| uint8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint8:max | ERR@ArrowInvalid | [255, None]@int16 | [255, None]@int32 | [255, None]@int64 | [255, None]@uint8 | [255, None]@uint16 | [255, None]@uint32 | [255, None]@uint64 | [255.0, None]@float16 | [255.0, None]@float32 | [255.0, None]@float64 | [True, None]@bool | [255, None]@string | [255, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [255.0000000000, None]@decimal128(38, 10) | [255.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint16:max | ERR@ArrowInvalid | ERR@ArrowInvalid | [65535, None]@int32 | [65535, None]@int64 | ERR@ArrowInvalid | [65535, None]@uint16 | [65535, None]@uint32 | [65535, None]@uint64 | [inf, None]@float16 | [65535.0, None]@float32 | [65535.0, None]@float64 | [True, None]@bool | [65535, None]@string | [65535, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [65535.0000000000, None]@decimal128(38, 10) | [65535.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint32:max | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@uint32 | [4294967295, None]@uint64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint64:max | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | ERR@ArrowInvalid | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, nan, None]@float16 | [inf, nan, None]@float32 | [inf, nan, None]@float64 | ERR@ArrowNotImplementedError | [inf, nan, None]@string | [inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.0999755859375, 0.89990234375, None]@float32 | [0.0999755859375, 0.89990234375, None]@float64 | ERR@ArrowNotImplementedError | [0.0999755859375, 0.89990234375, None]@string | [0.0999755859375, 0.89990234375, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.10000000149011612, 0.8999999761581421, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000015, 0.8999999762, None]@decimal128(38, 10) | [0.1000000015, 0.8999999762, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.1, 0.9, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000000, 0.9000000000, None]@decimal128(38, 10) | [0.1000000000, 0.9000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| bool:standard | [1, 0, None]@int8 | [1, 0, None]@int16 | [1, 0, None]@int32 | [1, 0, None]@int64 | [1, 0, None]@uint8 | [1, 0, None]@uint16 | [1, 0, None]@uint32 | [1, 0, None]@uint64 | ERR@ArrowNotImplementedError | [1.0, 0.0, None]@float32 | [1.0, 0.0, None]@float64 | [True, False, None]@bool | [true, false, None]@string | [true, false, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| fixed_size_binary[16]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@string | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@large_string | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@large_binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@fixed_size_binary[16] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal128(38, 10):standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal128(38, 10):large | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal256(76, 10):standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal256(76, 10):large | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date32[day]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date32[day]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date64[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date64[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| timestamp[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@string | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| timestamp[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59, None]@string | [1969-12-31 23:59:59, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | -| timestamp[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@string | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | -| timestamp[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999, None]@string | [1969-12-31 23:59:59.999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | -| timestamp[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@string | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | -| timestamp[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999, None]@string | [1969-12-31 23:59:59.999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@string | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00, None]@time64[ns] | -| timestamp[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999, None]@string | [1969-12-31 23:59:59.999999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[ns] | -| timestamp[s, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@string | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| timestamp[s, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59Z, None]@string | [1969-12-31 23:59:59Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | -| timestamp[ms, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@string | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | -| timestamp[ms, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999Z, None]@string | [1969-12-31 23:59:59.999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | -| timestamp[us, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@string | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | -| timestamp[us, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999Z, None]@string | [1969-12-31 23:59:59.999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[ns, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@string | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00, None]@time64[ns] | -| timestamp[ns, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999Z, None]@string | [1969-12-31 23:59:59.999999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[ns] | -| timestamp[s, tz=America/New_York]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@string | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, 1969-12-31, None]@date32[day] | [1969-12-31, 1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [19:00:00, 19:00:01, None]@time32[s] | [19:00:00, 19:00:01, None]@time32[ms] | [19:00:00, 19:00:01, None]@time64[us] | [19:00:00, 19:00:01, None]@time64[ns] | -| timestamp[s, tz=America/New_York]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 18:59:59-0500, None]@string | [1969-12-31 18:59:59-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [18:59:59, None]@time32[s] | [18:59:59, None]@time32[ms] | [18:59:59, None]@time64[us] | [18:59:59, None]@time64[ns] | -| timestamp[s, tz=Asia/Shanghai]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@string | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [08:00:00, 08:00:01, None]@time32[s] | [08:00:00, 08:00:01, None]@time32[ms] | [08:00:00, 08:00:01, None]@time64[us] | [08:00:00, 08:00:01, None]@time64[ns] | -| timestamp[s, tz=Asia/Shanghai]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 07:59:59+0800, None]@string | [1970-01-01 07:59:59+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, None]@date32[day] | [1970-01-01, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [07:59:59, None]@time32[s] | [07:59:59, None]@time32[ms] | [07:59:59, None]@time64[us] | [07:59:59, None]@time64[ns] | -| duration[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:01, None]@duration[ms] | [0:00:00, 0:00:01, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:01, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59, None]@duration[ms] | [-1 day, 23:59:59, None]@duration[us] | [-1 days +23:59:59, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.001000, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.001000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999000, None]@duration[us] | [-1 days +23:59:59.999000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| time32[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@string | [00:00:00, 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time32[s]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@string | [12:00:00, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time32[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000, 00:00:01.000, None]@string | [00:00:00.000, 00:00:01.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time32[ms]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000, None]@string | [12:00:00.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time64[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000, 00:00:01.000000, None]@string | [00:00:00.000000, 00:00:01.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time64[us]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000, None]@string | [12:00:00.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time64[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000000, 00:00:01.000000000, None]@string | [00:00:00.000000000, 00:00:01.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time64[ns]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000000, None]@string | [12:00:00.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | \ No newline at end of file +| source \ target | int8 | int16 | int32 | int64 | uint8 | uint16 | uint32 | uint64 | float16 | float32 | float64 | bool | string | large_string | binary | large_binary | fixed_size_binary[16] | decimal128(38, 10) | decimal256(76, 10) | date32[day] | date64[ms] | timestamp[s] | timestamp[ms] | timestamp[us] | timestamp[ns] | timestamp[s, tz=UTC] | timestamp[ms, tz=UTC] | timestamp[us, tz=UTC] | timestamp[ns, tz=UTC] | timestamp[s, tz=America/New_York] | timestamp[s, tz=Asia/Shanghai] | duration[s] | duration[ms] | duration[us] | duration[ns] | time32[s] | time32[ms] | time64[us] | time64[ns] | +|--------------------------------------------|------------------------------|------------------------------|---------------------------------------|---------------------------------------|------------------------------|------------------------------|------------------------------|------------------------------|------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|--------------------------------|-------------------------------------------------------------------------------|-------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------------------------------------|----------------------------------------------------------|-------------------------------------------|---------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------|----------------------------------------------------------------|-------------------------------------------------------------------------|--------------------------------------------------------------|-----------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------| +| int8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int8:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int8:max_min | [127, -128, None]@int8 | [127, -128, None]@int16 | [127, -128, None]@int32 | [127, -128, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [127.0, -128.0, None]@float16 | [127.0, -128.0, None]@float32 | [127.0, -128.0, None]@float64 | [True, True, None]@bool | [127, -128, None]@string | [127, -128, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [127.0000000000, -128.0000000000, None]@decimal128(38, 10) | [127.0000000000, -128.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:max_min | ERR@ArrowInvalid | [32767, -32768, None]@int16 | [32767, -32768, None]@int32 | [32767, -32768, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [32768.0, -32768.0, None]@float16 | [32767.0, -32768.0, None]@float32 | [32767.0, -32768.0, None]@float64 | [True, True, None]@bool | [32767, -32768, None]@string | [32767, -32768, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) | [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | [1970-01-01, 1970-01-02, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | [1969-12-31, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:max_min | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | [temporal overflow, temporal overflow, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [03:14:07, 20:45:52, None]@time32[s] | [20:31:23.647000, 03:28:36.352000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | +| int64:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| int64:max_min | ERR@ArrowInvalid | ERR@ArrowInvalid | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, None]@float16 | ERR@ArrowInvalid | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-25, 1969-12-07, None]@date64[ms] | [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] | [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] | [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] | [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] | [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] | [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] | [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] | [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] | [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] | [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] | [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:35:47.483647, 23:24:12.516352, None]@time64[us] | [00:00:02.147483, 23:59:57.852516, None]@time64[ns] | +| uint8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint8:max | ERR@ArrowInvalid | [255, None]@int16 | [255, None]@int32 | [255, None]@int64 | [255, None]@uint8 | [255, None]@uint16 | [255, None]@uint32 | [255, None]@uint64 | [255.0, None]@float16 | [255.0, None]@float32 | [255.0, None]@float64 | [True, None]@bool | [255, None]@string | [255, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [255.0000000000, None]@decimal128(38, 10) | [255.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint16:max | ERR@ArrowInvalid | ERR@ArrowInvalid | [65535, None]@int32 | [65535, None]@int64 | ERR@ArrowInvalid | [65535, None]@uint16 | [65535, None]@uint32 | [65535, None]@uint64 | [inf, None]@float16 | [65535.0, None]@float32 | [65535.0, None]@float64 | [True, None]@bool | [65535, None]@string | [65535, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [65535.0000000000, None]@decimal128(38, 10) | [65535.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint32:max | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@uint32 | [4294967295, None]@uint64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint64:max | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | ERR@ArrowInvalid | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, nan, None]@float16 | [inf, nan, None]@float32 | [inf, nan, None]@float64 | ERR@ArrowNotImplementedError | [inf, nan, None]@string | [inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.0999755859375, 0.89990234375, None]@float32 | [0.0999755859375, 0.89990234375, None]@float64 | ERR@ArrowNotImplementedError | [0.0999755859375, 0.89990234375, None]@string | [0.0999755859375, 0.89990234375, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.10000000149011612, 0.8999999761581421, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000015, 0.8999999762, None]@decimal128(38, 10) | [0.1000000015, 0.8999999762, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:special | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:fractional | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.1, 0.9, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000000, 0.9000000000, None]@decimal128(38, 10) | [0.1000000000, 0.9000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| bool:standard | [1, 0, None]@int8 | [1, 0, None]@int16 | [1, 0, None]@int32 | [1, 0, None]@int64 | [1, 0, None]@uint8 | [1, 0, None]@uint16 | [1, 0, None]@uint32 | [1, 0, None]@uint64 | ERR@ArrowNotImplementedError | [1.0, 0.0, None]@float32 | [1.0, 0.0, None]@float64 | [True, False, None]@bool | [true, false, None]@string | [true, false, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| fixed_size_binary[16]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@string | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@large_string | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@large_binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@fixed_size_binary[16] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal128(38, 10):standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal128(38, 10):large | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal256(76, 10):standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal256(76, 10):large | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date32[day]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date32[day]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date64[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date64[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| timestamp[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@string | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| timestamp[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59, None]@string | [1969-12-31 23:59:59, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | +| timestamp[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@string | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | +| timestamp[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999, None]@string | [1969-12-31 23:59:59.999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | +| timestamp[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@string | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | +| timestamp[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999, None]@string | [1969-12-31 23:59:59.999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@string | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00, None]@time64[ns] | +| timestamp[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999, None]@string | [1969-12-31 23:59:59.999999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[ns] | +| timestamp[s, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@string | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| timestamp[s, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59Z, None]@string | [1969-12-31 23:59:59Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | +| timestamp[ms, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@string | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | +| timestamp[ms, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999Z, None]@string | [1969-12-31 23:59:59.999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | ERR@ArrowInvalid | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | +| timestamp[us, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@string | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | +| timestamp[us, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999Z, None]@string | [1969-12-31 23:59:59.999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[ns, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@string | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [00:00:00, 00:00:00, None]@time64[ns] | +| timestamp[ns, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999Z, None]@string | [1969-12-31 23:59:59.999999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [23:59:59.999999, None]@time64[ns] | +| timestamp[s, tz=America/New_York]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@string | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, 1969-12-31, None]@date32[day] | [1969-12-31, 1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [19:00:00, 19:00:01, None]@time32[s] | [19:00:00, 19:00:01, None]@time32[ms] | [19:00:00, 19:00:01, None]@time64[us] | [19:00:00, 19:00:01, None]@time64[ns] | +| timestamp[s, tz=America/New_York]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 18:59:59-0500, None]@string | [1969-12-31 18:59:59-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [18:59:59, None]@time32[s] | [18:59:59, None]@time32[ms] | [18:59:59, None]@time64[us] | [18:59:59, None]@time64[ns] | +| timestamp[s, tz=Asia/Shanghai]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@string | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [08:00:00, 08:00:01, None]@time32[s] | [08:00:00, 08:00:01, None]@time32[ms] | [08:00:00, 08:00:01, None]@time64[us] | [08:00:00, 08:00:01, None]@time64[ns] | +| timestamp[s, tz=Asia/Shanghai]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 07:59:59+0800, None]@string | [1970-01-01 07:59:59+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, None]@date32[day] | [1970-01-01, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [07:59:59, None]@time32[s] | [07:59:59, None]@time32[ms] | [07:59:59, None]@time64[us] | [07:59:59, None]@time64[ns] | +| duration[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:01, None]@duration[ms] | [0:00:00, 0:00:01, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:01, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59, None]@duration[ms] | [-1 day, 23:59:59, None]@duration[us] | [-1 days +23:59:59, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.001000, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.001000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999000, None]@duration[us] | [-1 days +23:59:59.999000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| time32[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@string | [00:00:00, 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time32[s]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@string | [12:00:00, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time32[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000, 00:00:01.000, None]@string | [00:00:00.000, 00:00:01.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time32[ms]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000, None]@string | [12:00:00.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time64[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000, 00:00:01.000000, None]@string | [00:00:00.000000, 00:00:01.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time64[us]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000, None]@string | [12:00:00.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time64[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000000, 00:00:01.000000000, None]@string | [00:00:00.000000000, 00:00:01.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time64[ns]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000000, None]@string | [12:00:00.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.csv index ab5d671be452a..6e8bf7bf685c3 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.csv +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.csv @@ -7,7 +7,7 @@ int16:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int6 int16:max_min [-1, 0, None]@int8 [32767, -32768, None]@int16 [32767, -32768, None]@int32 [32767, -32768, None]@int64 [255, 0, None]@uint8 [32767, 32768, None]@uint16 [32767, 4294934528, None]@uint32 [32767, 18446744073709518848, None]@uint64 [32768.0, -32768.0, None]@float16 [32767.0, -32768.0, None]@float32 [32767.0, -32768.0, None]@float64 [True, True, None]@bool [32767, -32768, None]@string [32767, -32768, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int32:standard [0, 1, None]@int8 [0, 1, None]@int16 [0, 1, None]@int32 [0, 1, None]@int64 [0, 1, None]@uint8 [0, 1, None]@uint16 [0, 1, None]@uint32 [0, 1, None]@uint64 [0.0, 1.0, None]@float16 [0.0, 1.0, None]@float32 [0.0, 1.0, None]@float64 [False, True, None]@bool [0, 1, None]@string [0, 1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0E-10, 1.0000000000, None]@decimal128(38, 10) [0E-10, 1.0000000000, None]@decimal256(76, 10) [1970-01-01, 1970-01-02, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:00:00, 00:00:01, None]@time32[s] [00:00:00, 00:00:00.001000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int32:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int64 [255, None]@uint8 [65535, None]@uint16 [4294967295, None]@uint32 [18446744073709551615, None]@uint64 [-1.0, None]@float16 [-1.0, None]@float32 [-1.0, None]@float64 [True, None]@bool [-1, None]@string [-1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [-1.0000000000, None]@decimal128(38, 10) [-1.0000000000, None]@decimal256(76, 10) [1969-12-31, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [23:59:59, None]@time32[s] [23:59:59.999000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError -int32:max_min [-1, 0, None]@int8 [-1, 0, None]@int16 [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 [255, 0, None]@uint8 [65535, 0, None]@uint16 [2147483647, 2147483648, None]@uint32 [2147483647, 18446744071562067968, None]@uint64 [inf, -inf, None]@float16 [2147483648.0, -2147483648.0, None]@float32 [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) ERR@OverflowError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [03:14:07, 20:45:52, None]@time32[s] [20:31:23.647000, 03:28:36.352000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError +int32:max_min [-1, 0, None]@int8 [-1, 0, None]@int16 [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 [255, 0, None]@uint8 [65535, 0, None]@uint16 [2147483647, 2147483648, None]@uint32 [2147483647, 18446744071562067968, None]@uint64 [inf, -inf, None]@float16 [2147483648.0, -2147483648.0, None]@float32 [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) [temporal overflow, temporal overflow, None]@date32[day] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [03:14:07, 20:45:52, None]@time32[s] [20:31:23.647000, 03:28:36.352000, None]@time32[ms] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError int64:standard [0, 1, None]@int8 [0, 1, None]@int16 [0, 1, None]@int32 [0, 1, None]@int64 [0, 1, None]@uint8 [0, 1, None]@uint16 [0, 1, None]@uint32 [0, 1, None]@uint64 [0.0, 1.0, None]@float16 [0.0, 1.0, None]@float32 [0.0, 1.0, None]@float64 [False, True, None]@bool [0, 1, None]@string [0, 1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0E-10, 1.0000000000, None]@decimal128(38, 10) [0E-10, 1.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1970-01-01, 1970-01-01, None]@date64[ms] [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] [0:00:00, 0:00:01, None]@duration[s] [0:00:00, 0:00:00.001000, None]@duration[ms] [0:00:00, 0:00:00.000001, None]@duration[us] [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:00:00, 00:00:00.000001, None]@time64[us] [00:00:00, 00:00:00, None]@time64[ns] int64:negative [-1, None]@int8 [-1, None]@int16 [-1, None]@int32 [-1, None]@int64 [255, None]@uint8 [65535, None]@uint16 [4294967295, None]@uint32 [18446744073709551615, None]@uint64 [-1.0, None]@float16 [-1.0, None]@float32 [-1.0, None]@float64 [True, None]@bool [-1, None]@string [-1, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [-1.0000000000, None]@decimal128(38, 10) [-1.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1969-12-31, None]@date64[ms] [1969-12-31 23:59:59, None]@timestamp[s] [1969-12-31 23:59:59.999000, None]@timestamp[ms] [1969-12-31 23:59:59.999999, None]@timestamp[us] [1969-12-31 23:59:59.999999999, None]@timestamp[ns] [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] [-1 day, 23:59:59, None]@duration[s] [-1 day, 23:59:59.999000, None]@duration[ms] [-1 day, 23:59:59.999999, None]@duration[us] [-1 days +23:59:59.999999999, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [23:59:59.999999, None]@time64[us] [23:59:59.999999, None]@time64[ns] int64:max_min [-1, 0, None]@int8 [-1, 0, None]@int16 [2147483647, -2147483648, None]@int32 [2147483647, -2147483648, None]@int64 [255, 0, None]@uint8 [65535, 0, None]@uint16 [2147483647, 2147483648, None]@uint32 [2147483647, 18446744071562067968, None]@uint64 [inf, -inf, None]@float16 [2147483648.0, -2147483648.0, None]@float32 [2147483647.0, -2147483648.0, None]@float64 [True, True, None]@bool [2147483647, -2147483648, None]@string [2147483647, -2147483648, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError [1970-01-25, 1969-12-07, None]@date64[ms] [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [00:35:47.483647, 23:24:12.516352, None]@time64[us] [00:00:02.147483, 23:59:57.852516, None]@time64[ns] @@ -35,8 +35,8 @@ string:unicode ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInval large_string:numeric [0, 1, -1, None]@int8 [0, 1, -1, None]@int16 [0, 1, -1, None]@int32 [0, 1, -1, None]@int64 ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [0.0, 1.0, -1.0, None]@float16 [0.0, 1.0, -1.0, None]@float32 [0.0, 1.0, -1.0, None]@float64 ERR@ArrowInvalid [0, 1, -1, None]@string [0, 1, -1, None]@large_string [b'0', b'1', b'-1', None]@binary [b'0', b'1', b'-1', None]@large_binary ERR@ArrowInvalid [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError large_string:alpha ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [abc, , None]@string [abc, , None]@large_string [b'abc', b'', None]@binary [b'abc', b'', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError large_string:unicode ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [你好, مرحبا, 🎉, None]@string [你好, مرحبا, 🎉, None]@large_string [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError -binary:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@UnicodeDecodeError ERR@UnicodeDecodeError [b'\x00', b'\xff', b'hello', b'', None]@binary [b'\x00', b'\xff', b'hello', b'', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError -large_binary:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@UnicodeDecodeError ERR@UnicodeDecodeError [b'\x00', b'\xff', b'hello', b'', None]@binary [b'\x00', b'\xff', b'hello', b'', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError +binary:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [\0, b'\xff', hello, , None]@string [\0, b'\xff', hello, , None]@large_string [b'\x00', b'\xff', b'hello', b'', None]@binary [b'\x00', b'\xff', b'hello', b'', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError +large_binary:standard ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid [\0, b'\xff', hello, , None]@string [\0, b'\xff', hello, , None]@large_string [b'\x00', b'\xff', b'hello', b'', None]@binary [b'\x00', b'\xff', b'hello', b'', None]@large_binary ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowInvalid ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError fixed_size_binary[16]:standard ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@string [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@large_string [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@binary [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@large_binary [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@fixed_size_binary[16] ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError decimal128(38, 10):standard [0, 1, -1, None]@int8 [0, 1, -1, None]@int16 [0, 1, -1, None]@int32 [0, 1, -1, None]@int64 [0, 1, 255, None]@uint8 [0, 1, 65535, None]@uint16 [0, 1, 4294967295, None]@uint32 [0, 1, 18446744073709551615, None]@uint64 ERR@ArrowNotImplementedError [0.0, 1.5, -1.5, None]@float32 [0.0, 1.5, -1.5, None]@float64 ERR@ArrowNotImplementedError [0E-10, 1.5000000000, -1.5000000000, None]@string [0E-10, 1.5000000000, -1.5000000000, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError decimal128(38, 10):large [-1, None]@int8 [-7169, None]@int16 [1410065407, None]@int32 [9999999999, None]@int64 [255, None]@uint8 [58367, None]@uint16 [1410065407, None]@uint32 [9999999999, None]@uint64 ERR@ArrowNotImplementedError [10000000000.0, None]@float32 [9999999999.0, None]@float64 ERR@ArrowNotImplementedError [9999999999.0000000000, None]@string [9999999999.0000000000, None]@large_string ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError [9999999999.0000000000, None]@decimal128(38, 10) [9999999999.0000000000, None]@decimal256(76, 10) ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError ERR@ArrowNotImplementedError diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.md index f22c634c4a7db..dac777b15df4b 100644 --- a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.md +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_scalar_cast_unsafe.md @@ -1,85 +1,85 @@ -| source \ target | int8 | int16 | int32 | int64 | uint8 | uint16 | uint32 | uint64 | float16 | float32 | float64 | bool | string | large_string | binary | large_binary | fixed_size_binary[16] | decimal128(38, 10) | decimal256(76, 10) | date32[day] | date64[ms] | timestamp[s] | timestamp[ms] | timestamp[us] | timestamp[ns] | timestamp[s, tz=UTC] | timestamp[ms, tz=UTC] | timestamp[us, tz=UTC] | timestamp[ns, tz=UTC] | timestamp[s, tz=America/New_York] | timestamp[s, tz=Asia/Shanghai] | duration[s] | duration[ms] | duration[us] | duration[ns] | time32[s] | time32[ms] | time64[us] | time64[ns] | -|--------------------------------------------|------------------------------|------------------------------|-----------------------------------------------------|--------------------------------------------------------------------------------|------------------------------|------------------------------|------------------------------------------|------------------------------------------------------------|------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|--------------------------------|-------------------------------------------------------------------------------|-------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------|-------------------------------------------|---------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------|----------------------------------------------------------------|-------------------------------------------------------------------------|--------------------------------------------------------------|-----------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------| -| int8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int8:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int8:max_min | [127, -128, None]@int8 | [127, -128, None]@int16 | [127, -128, None]@int32 | [127, -128, None]@int64 | [127, 128, None]@uint8 | [127, 65408, None]@uint16 | [127, 4294967168, None]@uint32 | [127, 18446744073709551488, None]@uint64 | [127.0, -128.0, None]@float16 | [127.0, -128.0, None]@float32 | [127.0, -128.0, None]@float64 | [True, True, None]@bool | [127, -128, None]@string | [127, -128, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [127.0000000000, -128.0000000000, None]@decimal128(38, 10) | [127.0000000000, -128.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int16:max_min | [-1, 0, None]@int8 | [32767, -32768, None]@int16 | [32767, -32768, None]@int32 | [32767, -32768, None]@int64 | [255, 0, None]@uint8 | [32767, 32768, None]@uint16 | [32767, 4294934528, None]@uint32 | [32767, 18446744073709518848, None]@uint64 | [32768.0, -32768.0, None]@float16 | [32767.0, -32768.0, None]@float32 | [32767.0, -32768.0, None]@float64 | [True, True, None]@bool | [32767, -32768, None]@string | [32767, -32768, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) | [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | [1970-01-01, 1970-01-02, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | [1969-12-31, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int32:max_min | [-1, 0, None]@int8 | [-1, 0, None]@int16 | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | [255, 0, None]@uint8 | [65535, 0, None]@uint16 | [2147483647, 2147483648, None]@uint32 | [2147483647, 18446744071562067968, None]@uint64 | [inf, -inf, None]@float16 | [2147483648.0, -2147483648.0, None]@float32 | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@OverflowError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [03:14:07, 20:45:52, None]@time32[s] | [20:31:23.647000, 03:28:36.352000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| int64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | -| int64:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| int64:max_min | [-1, 0, None]@int8 | [-1, 0, None]@int16 | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | [255, 0, None]@uint8 | [65535, 0, None]@uint16 | [2147483647, 2147483648, None]@uint32 | [2147483647, 18446744071562067968, None]@uint64 | [inf, -inf, None]@float16 | [2147483648.0, -2147483648.0, None]@float32 | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-25, 1969-12-07, None]@date64[ms] | [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] | [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] | [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] | [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] | [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] | [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] | [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] | [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] | [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] | [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] | [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:35:47.483647, 23:24:12.516352, None]@time64[us] | [00:00:02.147483, 23:59:57.852516, None]@time64[ns] | -| uint8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint8:max | [-1, None]@int8 | [255, None]@int16 | [255, None]@int32 | [255, None]@int64 | [255, None]@uint8 | [255, None]@uint16 | [255, None]@uint32 | [255, None]@uint64 | [255.0, None]@float16 | [255.0, None]@float32 | [255.0, None]@float64 | [True, None]@bool | [255, None]@string | [255, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [255.0000000000, None]@decimal128(38, 10) | [255.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint16:max | [-1, None]@int8 | [-1, None]@int16 | [65535, None]@int32 | [65535, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [65535, None]@uint32 | [65535, None]@uint64 | [inf, None]@float16 | [65535.0, None]@float32 | [65535.0, None]@float64 | [True, None]@bool | [65535, None]@string | [65535, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [65535.0000000000, None]@decimal128(38, 10) | [65535.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint32:max | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [4294967295, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | [4294967296.0, None]@float32 | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| uint64:max | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [4294967295, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | [4294967296.0, None]@float32 | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:special | [0, 0, None]@int8 | [0, 0, None]@int16 | [-2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 9223372036854775808, None]@uint64 | [inf, nan, None]@float16 | [inf, nan, None]@float32 | [inf, nan, None]@float64 | ERR@ArrowNotImplementedError | [inf, nan, None]@string | [inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float16:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.0999755859375, 0.89990234375, None]@float32 | [0.0999755859375, 0.89990234375, None]@float64 | ERR@ArrowNotImplementedError | [0.0999755859375, 0.89990234375, None]@string | [0.0999755859375, 0.89990234375, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:special | [0, 0, 0, None]@int8 | [0, 0, 0, None]@int16 | [-2147483648, -2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, 0, None]@uint8 | [0, 0, 0, None]@uint16 | [0, 2147483648, 2147483648, None]@uint32 | [0, 9223372036854775808, 9223372036854775808, None]@uint64 | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 0E-10, 0E-10, None]@decimal128(38, 10) | [0E-10, 0E-10, 0E-10, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float32:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.10000000149011612, 0.8999999761581421, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000015, 0.8999999762, None]@decimal128(38, 10) | [0.1000000015, 0.8999999762, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:special | [0, 0, 0, None]@int8 | [0, 0, 0, None]@int16 | [-2147483648, -2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, 0, None]@uint8 | [0, 0, 0, None]@uint16 | [0, 2147483648, 2147483648, None]@uint32 | [0, 9223372036854775808, 9223372036854775808, None]@uint64 | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 0E-10, 0E-10, None]@decimal128(38, 10) | [0E-10, 0E-10, 0E-10, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| float64:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.1, 0.9, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000000, 0.9000000000, None]@decimal128(38, 10) | [0.1000000000, 0.9000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| bool:standard | [1, 0, None]@int8 | [1, 0, None]@int16 | [1, 0, None]@int32 | [1, 0, None]@int64 | [1, 0, None]@uint8 | [1, 0, None]@uint16 | [1, 0, None]@uint32 | [1, 0, None]@uint64 | ERR@ArrowNotImplementedError | [1.0, 0.0, None]@float32 | [1.0, 0.0, None]@float64 | [True, False, None]@bool | [true, false, None]@string | [true, false, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@UnicodeDecodeError | ERR@UnicodeDecodeError | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| large_binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@UnicodeDecodeError | ERR@UnicodeDecodeError | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| fixed_size_binary[16]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@string | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@large_string | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@large_binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@fixed_size_binary[16] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal128(38, 10):standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal128(38, 10):large | [-1, None]@int8 | [-7169, None]@int16 | [1410065407, None]@int32 | [9999999999, None]@int64 | [255, None]@uint8 | [58367, None]@uint16 | [1410065407, None]@uint32 | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal256(76, 10):standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| decimal256(76, 10):large | [-1, None]@int8 | [-7169, None]@int16 | [1410065407, None]@int32 | [9999999999, None]@int64 | [255, None]@uint8 | [58367, None]@uint16 | [1410065407, None]@uint32 | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date32[day]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date32[day]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date64[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| date64[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| timestamp[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@string | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| timestamp[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59, None]@string | [1969-12-31 23:59:59, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | -| timestamp[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@string | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | -| timestamp[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999, None]@string | [1969-12-31 23:59:59.999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | -| timestamp[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@string | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | -| timestamp[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999, None]@string | [1969-12-31 23:59:59.999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@string | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | -| timestamp[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999, None]@string | [1969-12-31 23:59:59.999999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[s, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@string | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| timestamp[s, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59Z, None]@string | [1969-12-31 23:59:59Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | -| timestamp[ms, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@string | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | -| timestamp[ms, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999Z, None]@string | [1969-12-31 23:59:59.999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | -| timestamp[us, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@string | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | -| timestamp[us, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999Z, None]@string | [1969-12-31 23:59:59.999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[ns, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@string | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | -| timestamp[ns, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999Z, None]@string | [1969-12-31 23:59:59.999999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | -| timestamp[s, tz=America/New_York]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@string | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, 1969-12-31, None]@date32[day] | [1969-12-31, 1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [19:00:00, 19:00:01, None]@time32[s] | [19:00:00, 19:00:01, None]@time32[ms] | [19:00:00, 19:00:01, None]@time64[us] | [19:00:00, 19:00:01, None]@time64[ns] | -| timestamp[s, tz=America/New_York]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 18:59:59-0500, None]@string | [1969-12-31 18:59:59-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [18:59:59, None]@time32[s] | [18:59:59, None]@time32[ms] | [18:59:59, None]@time64[us] | [18:59:59, None]@time64[ns] | -| timestamp[s, tz=Asia/Shanghai]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@string | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [08:00:00, 08:00:01, None]@time32[s] | [08:00:00, 08:00:01, None]@time32[ms] | [08:00:00, 08:00:01, None]@time64[us] | [08:00:00, 08:00:01, None]@time64[ns] | -| timestamp[s, tz=Asia/Shanghai]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 07:59:59+0800, None]@string | [1970-01-01 07:59:59+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, None]@date32[day] | [1970-01-01, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [07:59:59, None]@time32[s] | [07:59:59, None]@time32[ms] | [07:59:59, None]@time64[us] | [07:59:59, None]@time64[ns] | -| duration[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:01, None]@duration[ms] | [0:00:00, 0:00:01, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:01, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59, None]@duration[ms] | [-1 day, 23:59:59, None]@duration[us] | [-1 days +23:59:59, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.001000, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.001000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999000, None]@duration[us] | [-1 days +23:59:59.999000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [0:00:00, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00, None]@duration[ms] | [0:00:00, 0:00:00, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| duration[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [0:00:00, None]@duration[ms] | [0:00:00, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | -| time32[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@string | [00:00:00, 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time32[s]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@string | [12:00:00, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time32[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000, 00:00:01.000, None]@string | [00:00:00.000, 00:00:01.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time32[ms]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000, None]@string | [12:00:00.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time64[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000, 00:00:01.000000, None]@string | [00:00:00.000000, 00:00:01.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time64[us]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000, None]@string | [12:00:00.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | -| time64[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000000, 00:00:01.000000000, None]@string | [00:00:00.000000000, 00:00:01.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | -| time64[ns]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000000, None]@string | [12:00:00.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | \ No newline at end of file +| source \ target | int8 | int16 | int32 | int64 | uint8 | uint16 | uint32 | uint64 | float16 | float32 | float64 | bool | string | large_string | binary | large_binary | fixed_size_binary[16] | decimal128(38, 10) | decimal256(76, 10) | date32[day] | date64[ms] | timestamp[s] | timestamp[ms] | timestamp[us] | timestamp[ns] | timestamp[s, tz=UTC] | timestamp[ms, tz=UTC] | timestamp[us, tz=UTC] | timestamp[ns, tz=UTC] | timestamp[s, tz=America/New_York] | timestamp[s, tz=Asia/Shanghai] | duration[s] | duration[ms] | duration[us] | duration[ns] | time32[s] | time32[ms] | time64[us] | time64[ns] | +|--------------------------------------------|------------------------------|------------------------------|-----------------------------------------------------|--------------------------------------------------------------------------------|------------------------------|------------------------------|------------------------------------------|------------------------------------------------------------|------------------------------------------------|---------------------------------------------------------|---------------------------------------------------------|--------------------------------|-------------------------------------------------------------------------------|-------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|--------------------------------------------------------------------------|----------------------------------------------------------|-------------------------------------------|---------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------|----------------------------------------------------------------|-------------------------------------------------------------------------|--------------------------------------------------------------|-----------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------|-----------------------------------------------------|-----------------------------------------------------| +| int8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int8:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int8:max_min | [127, -128, None]@int8 | [127, -128, None]@int16 | [127, -128, None]@int32 | [127, -128, None]@int64 | [127, 128, None]@uint8 | [127, 65408, None]@uint16 | [127, 4294967168, None]@uint32 | [127, 18446744073709551488, None]@uint64 | [127.0, -128.0, None]@float16 | [127.0, -128.0, None]@float32 | [127.0, -128.0, None]@float64 | [True, True, None]@bool | [127, -128, None]@string | [127, -128, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [127.0000000000, -128.0000000000, None]@decimal128(38, 10) | [127.0000000000, -128.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int16:max_min | [-1, 0, None]@int8 | [32767, -32768, None]@int16 | [32767, -32768, None]@int32 | [32767, -32768, None]@int64 | [255, 0, None]@uint8 | [32767, 32768, None]@uint16 | [32767, 4294934528, None]@uint32 | [32767, 18446744073709518848, None]@uint64 | [32768.0, -32768.0, None]@float16 | [32767.0, -32768.0, None]@float32 | [32767.0, -32768.0, None]@float64 | [True, True, None]@bool | [32767, -32768, None]@string | [32767, -32768, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [32767.0000000000, -32768.0000000000, None]@decimal128(38, 10) | [32767.0000000000, -32768.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | [1970-01-01, 1970-01-02, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | [1969-12-31, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int32:max_min | [-1, 0, None]@int8 | [-1, 0, None]@int16 | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | [255, 0, None]@uint8 | [65535, 0, None]@uint16 | [2147483647, 2147483648, None]@uint32 | [2147483647, 18446744071562067968, None]@uint64 | [inf, -inf, None]@float16 | [2147483648.0, -2147483648.0, None]@float32 | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | [temporal overflow, temporal overflow, None]@date32[day] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [03:14:07, 20:45:52, None]@time32[s] | [20:31:23.647000, 03:28:36.352000, None]@time32[ms] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| int64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | +| int64:negative | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [-1, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [18446744073709551615, None]@uint64 | [-1.0, None]@float16 | [-1.0, None]@float32 | [-1.0, None]@float64 | [True, None]@bool | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1.0000000000, None]@decimal128(38, 10) | [-1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| int64:max_min | [-1, 0, None]@int8 | [-1, 0, None]@int16 | [2147483647, -2147483648, None]@int32 | [2147483647, -2147483648, None]@int64 | [255, 0, None]@uint8 | [65535, 0, None]@uint16 | [2147483647, 2147483648, None]@uint32 | [2147483647, 18446744071562067968, None]@uint64 | [inf, -inf, None]@float16 | [2147483648.0, -2147483648.0, None]@float32 | [2147483647.0, -2147483648.0, None]@float64 | [True, True, None]@bool | [2147483647, -2147483648, None]@string | [2147483647, -2147483648, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [2147483647.0000000000, -2147483648.0000000000, None]@decimal128(38, 10) | [2147483647.0000000000, -2147483648.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | [1970-01-25, 1969-12-07, None]@date64[ms] | [2038-01-19 03:14:07, 1901-12-13 20:45:52, None]@timestamp[s] | [1970-01-25 20:31:23.647000, 1969-12-07 03:28:36.352000, None]@timestamp[ms] | [1970-01-01 00:35:47.483647, 1969-12-31 23:24:12.516352, None]@timestamp[us] | [1970-01-01 00:00:02.147483647, 1969-12-31 23:59:57.852516352, None]@timestamp[ns] | [2038-01-19 03:14:07+00:00, 1901-12-13 20:45:52+00:00, None]@timestamp[s, tz=UTC] | [1970-01-25 20:31:23.647000+00:00, 1969-12-07 03:28:36.352000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:35:47.483647+00:00, 1969-12-31 23:24:12.516352+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:02.147483647+00:00, 1969-12-31 23:59:57.852516352+00:00, None]@timestamp[ns, tz=UTC] | [2038-01-18 22:14:07-05:00, 1901-12-13 15:45:52-05:00, None]@timestamp[s, tz=America/New_York] | [2038-01-19 11:14:07+08:00, 1901-12-14 04:45:52+08:00, None]@timestamp[s, tz=Asia/Shanghai] | [24855 days, 3:14:07, -24856 days, 20:45:52, None]@duration[s] | [24 days, 20:31:23.647000, -25 days, 3:28:36.352000, None]@duration[ms] | [0:35:47.483647, -1 day, 23:24:12.516352, None]@duration[us] | [0 days 00:00:02.147483647, -1 days +23:59:57.852516352, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:35:47.483647, 23:24:12.516352, None]@time64[us] | [00:00:02.147483, 23:59:57.852516, None]@time64[ns] | +| uint8:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint8:max | [-1, None]@int8 | [255, None]@int16 | [255, None]@int32 | [255, None]@int64 | [255, None]@uint8 | [255, None]@uint16 | [255, None]@uint32 | [255, None]@uint64 | [255.0, None]@float16 | [255.0, None]@float32 | [255.0, None]@float64 | [True, None]@bool | [255, None]@string | [255, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [255.0000000000, None]@decimal128(38, 10) | [255.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint16:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint16:max | [-1, None]@int8 | [-1, None]@int16 | [65535, None]@int32 | [65535, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [65535, None]@uint32 | [65535, None]@uint64 | [inf, None]@float16 | [65535.0, None]@float32 | [65535.0, None]@float64 | [True, None]@bool | [65535, None]@string | [65535, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [65535.0000000000, None]@decimal128(38, 10) | [65535.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint32:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint32:max | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [4294967295, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | [4294967296.0, None]@float32 | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint64:standard | [0, 1, None]@int8 | [0, 1, None]@int16 | [0, 1, None]@int32 | [0, 1, None]@int64 | [0, 1, None]@uint8 | [0, 1, None]@uint16 | [0, 1, None]@uint32 | [0, 1, None]@uint64 | [0.0, 1.0, None]@float16 | [0.0, 1.0, None]@float32 | [0.0, 1.0, None]@float64 | [False, True, None]@bool | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| uint64:max | [-1, None]@int8 | [-1, None]@int16 | [-1, None]@int32 | [4294967295, None]@int64 | [255, None]@uint8 | [65535, None]@uint16 | [4294967295, None]@uint32 | [4294967295, None]@uint64 | [inf, None]@float16 | [4294967296.0, None]@float32 | [4294967295.0, None]@float64 | [True, None]@bool | [4294967295, None]@string | [4294967295, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [4294967295.0000000000, None]@decimal128(38, 10) | [4294967295.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:special | [0, 0, None]@int8 | [0, 0, None]@int16 | [-2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 9223372036854775808, None]@uint64 | [inf, nan, None]@float16 | [inf, nan, None]@float32 | [inf, nan, None]@float64 | ERR@ArrowNotImplementedError | [inf, nan, None]@string | [inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float16:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.0999755859375, 0.89990234375, None]@float32 | [0.0999755859375, 0.89990234375, None]@float64 | ERR@ArrowNotImplementedError | [0.0999755859375, 0.89990234375, None]@string | [0.0999755859375, 0.89990234375, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:special | [0, 0, 0, None]@int8 | [0, 0, 0, None]@int16 | [-2147483648, -2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, 0, None]@uint8 | [0, 0, 0, None]@uint16 | [0, 2147483648, 2147483648, None]@uint32 | [0, 9223372036854775808, 9223372036854775808, None]@uint64 | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 0E-10, 0E-10, None]@decimal128(38, 10) | [0E-10, 0E-10, 0E-10, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float32:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.10000000149011612, 0.8999999761581421, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000015, 0.8999999762, None]@decimal128(38, 10) | [0.1000000015, 0.8999999762, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | [0.0, 1.5, -1.5, None]@float16 | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | [False, True, True, None]@bool | [0, 1.5, -1.5, None]@string | [0, 1.5, -1.5, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:special | [0, 0, 0, None]@int8 | [0, 0, 0, None]@int16 | [-2147483648, -2147483648, -2147483648, None]@int32 | [-9223372036854775808, -9223372036854775808, -9223372036854775808, None]@int64 | [0, 0, 0, None]@uint8 | [0, 0, 0, None]@uint16 | [0, 2147483648, 2147483648, None]@uint32 | [0, 9223372036854775808, 9223372036854775808, None]@uint64 | [inf, -inf, nan, None]@float16 | [inf, -inf, nan, None]@float32 | [inf, -inf, nan, None]@float64 | [True, True, True, None]@bool | [inf, -inf, nan, None]@string | [inf, -inf, nan, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 0E-10, 0E-10, None]@decimal128(38, 10) | [0E-10, 0E-10, 0E-10, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| float64:fractional | [0, 0, None]@int8 | [0, 0, None]@int16 | [0, 0, None]@int32 | [0, 0, None]@int64 | [0, 0, None]@uint8 | [0, 0, None]@uint16 | [0, 0, None]@uint32 | [0, 0, None]@uint64 | [0.0999755859375, 0.89990234375, None]@float16 | [0.10000000149011612, 0.8999999761581421, None]@float32 | [0.1, 0.9, None]@float64 | [True, True, None]@bool | [0.1, 0.9, None]@string | [0.1, 0.9, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0.1000000000, 0.9000000000, None]@decimal128(38, 10) | [0.1000000000, 0.9000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| bool:standard | [1, 0, None]@int8 | [1, 0, None]@int16 | [1, 0, None]@int32 | [1, 0, None]@int64 | [1, 0, None]@uint8 | [1, 0, None]@uint16 | [1, 0, None]@uint32 | [1, 0, None]@uint64 | ERR@ArrowNotImplementedError | [1.0, 0.0, None]@float32 | [1.0, 0.0, None]@float64 | [True, False, None]@bool | [true, false, None]@string | [true, false, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:numeric | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [0.0, 1.0, -1.0, None]@float16 | [0.0, 1.0, -1.0, None]@float32 | [0.0, 1.0, -1.0, None]@float64 | ERR@ArrowInvalid | [0, 1, -1, None]@string | [0, 1, -1, None]@large_string | [b'0', b'1', b'-1', None]@binary | [b'0', b'1', b'-1', None]@large_binary | ERR@ArrowInvalid | [0E-10, 1.0000000000, -1.0000000000, None]@decimal128(38, 10) | [0E-10, 1.0000000000, -1.0000000000, None]@decimal256(76, 10) | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:alpha | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [abc, , None]@string | [abc, , None]@large_string | [b'abc', b'', None]@binary | [b'abc', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_string:unicode | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [你好, مرحبا, 🎉, None]@string | [你好, مرحبا, 🎉, None]@large_string | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@binary | [b'\xe4\xbd\xa0\xe5\xa5\xbd', b'\xd9\x85\xd8\xb1\xd8\xad\xd8\xa8\xd8\xa7', b'\xf0\x9f\x8e\x89', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [\0, b'\xff', hello, , None]@string | [\0, b'\xff', hello, , None]@large_string | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| large_binary:standard | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | [\0, b'\xff', hello, , None]@string | [\0, b'\xff', hello, , None]@large_string | [b'\x00', b'\xff', b'hello', b'', None]@binary | [b'\x00', b'\xff', b'hello', b'', None]@large_binary | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowInvalid | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| fixed_size_binary[16]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@string | [0123456789abcdef, \0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0, None]@large_string | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@large_binary | [b'0123456789abcdef', b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', None]@fixed_size_binary[16] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal128(38, 10):standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal128(38, 10):large | [-1, None]@int8 | [-7169, None]@int16 | [1410065407, None]@int32 | [9999999999, None]@int64 | [255, None]@uint8 | [58367, None]@uint16 | [1410065407, None]@uint32 | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal256(76, 10):standard | [0, 1, -1, None]@int8 | [0, 1, -1, None]@int16 | [0, 1, -1, None]@int32 | [0, 1, -1, None]@int64 | [0, 1, 255, None]@uint8 | [0, 1, 65535, None]@uint16 | [0, 1, 4294967295, None]@uint32 | [0, 1, 18446744073709551615, None]@uint64 | ERR@ArrowNotImplementedError | [0.0, 1.5, -1.5, None]@float32 | [0.0, 1.5, -1.5, None]@float64 | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@string | [0E-10, 1.5000000000, -1.5000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0E-10, 1.5000000000, -1.5000000000, None]@decimal128(38, 10) | [0E-10, 1.5000000000, -1.5000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| decimal256(76, 10):large | [-1, None]@int8 | [-7169, None]@int16 | [1410065407, None]@int32 | [9999999999, None]@int64 | [255, None]@uint8 | [58367, None]@uint16 | [1410065407, None]@uint32 | [9999999999, None]@uint64 | ERR@ArrowNotImplementedError | [10000000000.0, None]@float32 | [9999999999.0, None]@float64 | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@string | [9999999999.0000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [9999999999.0000000000, None]@decimal128(38, 10) | [9999999999.0000000000, None]@decimal256(76, 10) | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date32[day]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date32[day]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date64[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@string | [1970-01-01, 1970-01-02, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-02, None]@date32[day] | [1970-01-01, 1970-01-02, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-02 00:00:00, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-02 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1970-01-01 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-02 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| date64[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-86400000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@string | [1969-12-31, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 00:00:00, None]@timestamp[s] | [1969-12-31 00:00:00, None]@timestamp[ms] | [1969-12-31 00:00:00, None]@timestamp[us] | [1969-12-31 00:00:00, None]@timestamp[ns] | [1969-12-31 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 00:00:00+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-30 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1969-12-31 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| timestamp[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@string | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| timestamp[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59, None]@string | [1969-12-31 23:59:59, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | +| timestamp[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@string | [1970-01-01 00:00:00.000, 1970-01-01 00:00:00.001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | +| timestamp[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999, None]@string | [1969-12-31 23:59:59.999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | +| timestamp[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@string | [1970-01-01 00:00:00.000000, 1970-01-01 00:00:00.000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | +| timestamp[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999, None]@string | [1969-12-31 23:59:59.999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@string | [1970-01-01 00:00:00.000000000, 1970-01-01 00:00:00.000000001, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | +| timestamp[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999, None]@string | [1969-12-31 23:59:59.999999999, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[s, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@string | [1970-01-01 00:00:00Z, 1970-01-01 00:00:01Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| timestamp[s, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59Z, None]@string | [1969-12-31 23:59:59Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59, None]@time32[ms] | [23:59:59, None]@time64[us] | [23:59:59, None]@time64[ns] | +| timestamp[ms, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@string | [1970-01-01 00:00:00.000Z, 1970-01-01 00:00:00.001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.001000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.001000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00.001000, None]@time32[ms] | [00:00:00, 00:00:00.001000, None]@time64[us] | [00:00:00, 00:00:00.001000, None]@time64[ns] | +| timestamp[ms, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999Z, None]@string | [1969-12-31 23:59:59.999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1969-12-31 23:59:59.999000, None]@timestamp[ms] | [1969-12-31 23:59:59.999000, None]@timestamp[us] | [1969-12-31 23:59:59.999000, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999000+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999000, None]@time64[us] | [23:59:59.999000, None]@time64[ns] | +| timestamp[us, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@string | [1970-01-01 00:00:00.000000Z, 1970-01-01 00:00:00.000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00.000001, None]@time64[us] | [00:00:00, 00:00:00.000001, None]@time64[ns] | +| timestamp[us, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999Z, None]@string | [1969-12-31 23:59:59.999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1969-12-31 23:59:59.999999, None]@timestamp[us] | [1969-12-31 23:59:59.999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[ns, tz=UTC]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@string | [1970-01-01 00:00:00.000000000Z, 1970-01-01 00:00:00.000000001Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:00, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:00.000000001, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:00.000000001+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:00, None]@time32[s] | [00:00:00, 00:00:00, None]@time32[ms] | [00:00:00, 00:00:00, None]@time64[us] | [00:00:00, 00:00:00, None]@time64[ns] | +| timestamp[ns, tz=UTC]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 23:59:59.999999999Z, None]@string | [1969-12-31 23:59:59.999999999Z, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, None]@timestamp[s] | [1970-01-01 00:00:00, None]@timestamp[ms] | [1970-01-01 00:00:00, None]@timestamp[us] | [1969-12-31 23:59:59.999999999, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59.999999999+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [23:59:59, None]@time32[s] | [23:59:59.999000, None]@time32[ms] | [23:59:59.999999, None]@time64[us] | [23:59:59.999999, None]@time64[ns] | +| timestamp[s, tz=America/New_York]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@string | [1969-12-31 19:00:00-0500, 1969-12-31 19:00:01-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, 1969-12-31, None]@date32[day] | [1969-12-31, 1969-12-31, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [19:00:00, 19:00:01, None]@time32[s] | [19:00:00, 19:00:01, None]@time32[ms] | [19:00:00, 19:00:01, None]@time64[us] | [19:00:00, 19:00:01, None]@time64[ns] | +| timestamp[s, tz=America/New_York]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31 18:59:59-0500, None]@string | [1969-12-31 18:59:59-0500, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1969-12-31, None]@date32[day] | [1969-12-31, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [18:59:59, None]@time32[s] | [18:59:59, None]@time32[ms] | [18:59:59, None]@time64[us] | [18:59:59, None]@time64[ns] | +| timestamp[s, tz=Asia/Shanghai]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@string | [1970-01-01 08:00:00+0800, 1970-01-01 08:00:01+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, 1970-01-01, None]@date32[day] | [1970-01-01, 1970-01-01, None]@date64[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[s] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ms] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[us] | [1970-01-01 00:00:00, 1970-01-01 00:00:01, None]@timestamp[ns] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[s, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ms, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[us, tz=UTC] | [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 19:00:00-05:00, 1969-12-31 19:00:01-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 08:00:00+08:00, 1970-01-01 08:00:01+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [08:00:00, 08:00:01, None]@time32[s] | [08:00:00, 08:00:01, None]@time32[ms] | [08:00:00, 08:00:01, None]@time64[us] | [08:00:00, 08:00:01, None]@time64[ns] | +| timestamp[s, tz=Asia/Shanghai]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01 07:59:59+0800, None]@string | [1970-01-01 07:59:59+0800, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [1970-01-01, None]@date32[day] | [1970-01-01, None]@date64[ms] | [1969-12-31 23:59:59, None]@timestamp[s] | [1969-12-31 23:59:59, None]@timestamp[ms] | [1969-12-31 23:59:59, None]@timestamp[us] | [1969-12-31 23:59:59, None]@timestamp[ns] | [1969-12-31 23:59:59+00:00, None]@timestamp[s, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ms, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[us, tz=UTC] | [1969-12-31 23:59:59+00:00, None]@timestamp[ns, tz=UTC] | [1969-12-31 18:59:59-05:00, None]@timestamp[s, tz=America/New_York] | [1970-01-01 07:59:59+08:00, None]@timestamp[s, tz=Asia/Shanghai] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [07:59:59, None]@time32[s] | [07:59:59, None]@time32[ms] | [07:59:59, None]@time64[us] | [07:59:59, None]@time64[ns] | +| duration[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:01, None]@duration[s] | [0:00:00, 0:00:01, None]@duration[ms] | [0:00:00, 0:00:01, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:01, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[s]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1 day, 23:59:59, None]@duration[s] | [-1 day, 23:59:59, None]@duration[ms] | [-1 day, 23:59:59, None]@duration[us] | [-1 days +23:59:59, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00.001000, None]@duration[ms] | [0:00:00, 0:00:00.001000, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.001000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ms]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [-1 day, 23:59:59.999000, None]@duration[ms] | [-1 day, 23:59:59.999000, None]@duration[us] | [-1 days +23:59:59.999000, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00, None]@duration[ms] | [0:00:00, 0:00:00.000001, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[us]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [0:00:00, None]@duration[ms] | [-1 day, 23:59:59.999999, None]@duration[us] | [-1 days +23:59:59.999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@string | [0, 1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, 0:00:00, None]@duration[s] | [0:00:00, 0:00:00, None]@duration[ms] | [0:00:00, 0:00:00, None]@duration[us] | [0 days 00:00:00, 0 days 00:00:00.000000001, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| duration[ns]:negative | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [-1, None]@string | [-1, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0:00:00, None]@duration[s] | [0:00:00, None]@duration[ms] | [0:00:00, None]@duration[us] | [-1 days +23:59:59.999999999, None]@duration[ns] | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | +| time32[s]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@string | [00:00:00, 00:00:01, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time32[s]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@string | [12:00:00, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time32[ms]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000, 00:00:01.000, None]@string | [00:00:00.000, 00:00:01.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time32[ms]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000, None]@int32 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000, None]@string | [12:00:00.000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time64[us]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000, 00:00:01.000000, None]@string | [00:00:00.000000, 00:00:01.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time64[us]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000, None]@string | [12:00:00.000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | +| time64[ns]:standard | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [0, 1000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00.000000000, 00:00:01.000000000, None]@string | [00:00:00.000000000, 00:00:01.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [00:00:00, 00:00:01, None]@time32[s] | [00:00:00, 00:00:01, None]@time32[ms] | [00:00:00, 00:00:01, None]@time64[us] | [00:00:00, 00:00:01, None]@time64[ns] | +| time64[ns]:noon | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [43200000000000, None]@int64 | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00.000000000, None]@string | [12:00:00.000000000, None]@large_string | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | ERR@ArrowNotImplementedError | [12:00:00, None]@time32[s] | [12:00:00, None]@time32[ms] | [12:00:00, None]@time64[us] | [12:00:00, None]@time64[ns] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.csv new file mode 100644 index 0000000000000..66f7bbaf691b7 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.csv @@ -0,0 +1,13 @@ +test case pandas dataframe preserve_index=None preserve_index=False preserve_index=True +0-columns:range-index {}@Dataframe[][index=RangeIndex[0:3:1]] []@Schema []@Schema [__index_level_0__: int64 nullable=True]@Schema +0-columns:named-index {}@Dataframe[][index='idx':[100, 200, 300]] [idx: int64 nullable=True]@Schema []@Schema [idx: int64 nullable=True]@Schema +0-columns:unnamed-index {}@Dataframe[][index=None:[10, 20, 30]] [__index_level_0__: int64 nullable=True]@Schema []@Schema [__index_level_0__: int64 nullable=True]@Schema +0-columns:empty {}@Dataframe[][index=RangeIndex[0:0:1]] []@Schema []@Schema [__index_level_0__: int64 nullable=True]@Schema +single-column:range-index {'a': [1, 2, 3]}@Dataframe[a int64][index=RangeIndex[0:3:1]] [a: int64 nullable=True]@Schema [a: int64 nullable=True]@Schema [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema +single-column:named-index {'a': [1, 2, 3]}@Dataframe[a int64][index='idx':[100, 200, 300]] [a: int64 nullable=True, idx: int64 nullable=True]@Schema [a: int64 nullable=True]@Schema [a: int64 nullable=True, idx: int64 nullable=True]@Schema +single-column:unnamed-index {'a': [1, 2, 3]}@Dataframe[a int64][index=None:[10, 20, 30]] [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema [a: int64 nullable=True]@Schema [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema +multi-column:standard {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema +multi-column:nullable {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT, Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema +multi-column:no-rows {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64, f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True]@Schema [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True]@Schema [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema +single-column:multiindex {'a': [1, 2, 3]}@Dataframe[a int64][index=MultiIndex[names=['g', 'n']]] [a: int64 nullable=True, g: int64 nullable=True, n: int64 nullable=True]@Schema [a: int64 nullable=True]@Schema [a: int64 nullable=True, g: int64 nullable=True, n: int64 nullable=True]@Schema +single-column:multiindex-partial-name {'a': [1, 2]}@Dataframe[a int64][index=MultiIndex[names=['g', None]]] [a: int64 nullable=True, g: int64 nullable=True, __index_level_1__: int64 nullable=True]@Schema [a: int64 nullable=True]@Schema [a: int64 nullable=True, g: int64 nullable=True, __index_level_1__: int64 nullable=True]@Schema diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.md new file mode 100644 index 0000000000000..bcf6286bdae4d --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_schema_from_pandas.md @@ -0,0 +1,14 @@ +| test case | pandas dataframe | preserve_index=None | preserve_index=False | preserve_index=True | +|---------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| 0-columns:range-index | {}@Dataframe[][index=RangeIndex[0:3:1]] | []@Schema | []@Schema | [__index_level_0__: int64 nullable=True]@Schema | +| 0-columns:named-index | {}@Dataframe[][index='idx':[100, 200, 300]] | [idx: int64 nullable=True]@Schema | []@Schema | [idx: int64 nullable=True]@Schema | +| 0-columns:unnamed-index | {}@Dataframe[][index=None:[10, 20, 30]] | [__index_level_0__: int64 nullable=True]@Schema | []@Schema | [__index_level_0__: int64 nullable=True]@Schema | +| 0-columns:empty | {}@Dataframe[][index=RangeIndex[0:0:1]] | []@Schema | []@Schema | [__index_level_0__: int64 nullable=True]@Schema | +| single-column:range-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=RangeIndex[0:3:1]] | [a: int64 nullable=True]@Schema | [a: int64 nullable=True]@Schema | [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema | +| single-column:named-index | {'a': [1, 2, 3]}@Dataframe[a int64][index='idx':[100, 200, 300]] | [a: int64 nullable=True, idx: int64 nullable=True]@Schema | [a: int64 nullable=True]@Schema | [a: int64 nullable=True, idx: int64 nullable=True]@Schema | +| single-column:unnamed-index | {'a': [1, 2, 3]}@Dataframe[a int64][index=None:[10, 20, 30]] | [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema | [a: int64 nullable=True]@Schema | [a: int64 nullable=True, __index_level_0__: int64 nullable=True]@Schema | +| multi-column:standard | {'i': [1, 2, 3], 'f': [1.5, 2.5, 3.5], 'b': [True, False, True], 's': ['a', 'b', 'c'], 't': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')]}@Dataframe[i int64, f float64, b bool, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema | +| multi-column:nullable | {'f': [1.5, nan, 3.5], 'b': [True, None, False], 's': ['a', None, 'c'], 't': [Timestamp('2020-01-01 05:30:00'), NaT, Timestamp('2020-01-01 05:30:00')]}@Dataframe[f float64, b object, s object, t datetime64[ns]][index=RangeIndex[0:3:1]] | [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema | [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True]@Schema | [f: float64 nullable=True, b: bool nullable=True, s: string nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema | +| multi-column:no-rows | {'i': [], 'f': [], 'b': [], 't': []}@Dataframe[i int64, f float64, b bool, t datetime64[ns]][index=RangeIndex[0:0:1]] | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True]@Schema | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True]@Schema | [i: int64 nullable=True, f: float64 nullable=True, b: bool nullable=True, t: timestamp[ns] nullable=True, __index_level_0__: int64 nullable=True]@Schema | +| single-column:multiindex | {'a': [1, 2, 3]}@Dataframe[a int64][index=MultiIndex[names=['g', 'n']]] | [a: int64 nullable=True, g: int64 nullable=True, n: int64 nullable=True]@Schema | [a: int64 nullable=True]@Schema | [a: int64 nullable=True, g: int64 nullable=True, n: int64 nullable=True]@Schema | +| single-column:multiindex-partial-name | {'a': [1, 2]}@Dataframe[a int64][index=MultiIndex[names=['g', None]]] | [a: int64 nullable=True, g: int64 nullable=True, __index_level_1__: int64 nullable=True]@Schema | [a: int64 nullable=True]@Schema | [a: int64 nullable=True, g: int64 nullable=True, __index_level_1__: int64 nullable=True]@Schema | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.csv new file mode 100644 index 0000000000000..1cdbb39f58d04 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.csv @@ -0,0 +1,21 @@ +test case pyarrow table cast result +types:downcast {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] +types:upcast {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] +overflow:int64->int32 {a: [1099511627776, 1]}@Table[a: int64] ERR@ArrowInvalid +truncate:float->int {a: [1.9, -2.1]}@Table[a: float64] ERR@ArrowInvalid +names:mismatch {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +names:reordered {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +names:field-count {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +nullable:false-with-nulls {a: [1, None]}@Table[a: int64] ERR@ValueError +nullable:false-no-nulls {a: [1, 2]}@Table[a: int64] {a: [1, 2]}@Table[a: int32] +timestamp:us->ns {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[ns]] +timestamp:attach-tz {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] {ts: [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00]}@Table[ts: timestamp[us, tz=UTC]] +string->large_string {s: [hello, world, None]}@Table[s: string] {s: [hello, world, None]}@Table[s: large_string] +binary->large_binary {b: [b'x', b'yz', None]}@Table[b: binary] {b: [b'x', b'yz', None]}@Table[b: large_binary] +nested:list {lst: [[1, 2], [3], None]}@Table[lst: list<item: int64>] {lst: [[1, 2], [3], None]}@Table[lst: list<item: int32>] +nested:list-overflow {lst: [[1099511627776, 1]]}@Table[lst: list<item: int64>] ERR@ArrowInvalid +nested:struct {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int64, y: string>] {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int32, y: large_string>] +nested:map {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int64>] {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int32>] +multi-chunk-column {a: [1, 2, 3, None]}@Table[a: int64] {a: [1, 2, 3, None]}@Table[a: int32] +empty:0-columns {}@Table[] {}@Table[] +empty:columns-no-rows {i: [], s: []}@Table[i: int64, s: string] {i: [], s: []}@Table[i: int32, s: large_string] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.md new file mode 100644 index 0000000000000..397f4ebb9799a --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_safe.md @@ -0,0 +1,22 @@ +| test case | pyarrow table | cast result | +|---------------------------|-----------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------| +| types:downcast | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] | +| types:upcast | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] | +| overflow:int64->int32 | {a: [1099511627776, 1]}@Table[a: int64] | ERR@ArrowInvalid | +| truncate:float->int | {a: [1.9, -2.1]}@Table[a: float64] | ERR@ArrowInvalid | +| names:mismatch | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| names:reordered | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| names:field-count | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| nullable:false-with-nulls | {a: [1, None]}@Table[a: int64] | ERR@ValueError | +| nullable:false-no-nulls | {a: [1, 2]}@Table[a: int64] | {a: [1, 2]}@Table[a: int32] | +| timestamp:us->ns | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[ns]] | +| timestamp:attach-tz | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] | {ts: [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00]}@Table[ts: timestamp[us, tz=UTC]] | +| string->large_string | {s: [hello, world, None]}@Table[s: string] | {s: [hello, world, None]}@Table[s: large_string] | +| binary->large_binary | {b: [b'x', b'yz', None]}@Table[b: binary] | {b: [b'x', b'yz', None]}@Table[b: large_binary] | +| nested:list | {lst: [[1, 2], [3], None]}@Table[lst: list<item: int64>] | {lst: [[1, 2], [3], None]}@Table[lst: list<item: int32>] | +| nested:list-overflow | {lst: [[1099511627776, 1]]}@Table[lst: list<item: int64>] | ERR@ArrowInvalid | +| nested:struct | {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int64, y: string>] | {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int32, y: large_string>] | +| nested:map | {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int64>] | {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int32>] | +| multi-chunk-column | {a: [1, 2, 3, None]}@Table[a: int64] | {a: [1, 2, 3, None]}@Table[a: int32] | +| empty:0-columns | {}@Table[] | {}@Table[] | +| empty:columns-no-rows | {i: [], s: []}@Table[i: int64, s: string] | {i: [], s: []}@Table[i: int32, s: large_string] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.csv new file mode 100644 index 0000000000000..f9065e14ff57b --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.csv @@ -0,0 +1,21 @@ +test case pyarrow table cast result +types:downcast {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] +types:upcast {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] +overflow:int64->int32 {a: [1099511627776, 1]}@Table[a: int64] {a: [0, 1]}@Table[a: int32] +truncate:float->int {a: [1.9, -2.1]}@Table[a: float64] {a: [1, -2]}@Table[a: int64] +names:mismatch {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +names:reordered {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +names:field-count {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] ERR@ValueError +nullable:false-with-nulls {a: [1, None]}@Table[a: int64] ERR@ValueError +nullable:false-no-nulls {a: [1, 2]}@Table[a: int64] {a: [1, 2]}@Table[a: int32] +timestamp:us->ns {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[ns]] +timestamp:attach-tz {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] {ts: [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00]}@Table[ts: timestamp[us, tz=UTC]] +string->large_string {s: [hello, world, None]}@Table[s: string] {s: [hello, world, None]}@Table[s: large_string] +binary->large_binary {b: [b'x', b'yz', None]}@Table[b: binary] {b: [b'x', b'yz', None]}@Table[b: large_binary] +nested:list {lst: [[1, 2], [3], None]}@Table[lst: list<item: int64>] {lst: [[1, 2], [3], None]}@Table[lst: list<item: int32>] +nested:list-overflow {lst: [[1099511627776, 1]]}@Table[lst: list<item: int64>] {lst: [[0, 1]]}@Table[lst: list<item: int32>] +nested:struct {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int64, y: string>] {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int32, y: large_string>] +nested:map {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int64>] {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int32>] +multi-chunk-column {a: [1, 2, 3, None]}@Table[a: int64] {a: [1, 2, 3, None]}@Table[a: int32] +empty:0-columns {}@Table[] {}@Table[] +empty:columns-no-rows {i: [], s: []}@Table[i: int64, s: string] {i: [], s: []}@Table[i: int32, s: large_string] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.md new file mode 100644 index 0000000000000..f1b1ddd381e15 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_cast_unsafe.md @@ -0,0 +1,22 @@ +| test case | pyarrow table | cast result | +|---------------------------|-----------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------| +| types:downcast | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] | +| types:upcast | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int32, b: float32] | {a: [1, 2, 3], b: [1.5, 2.5, 3.5]}@Table[a: int64, b: float64] | +| overflow:int64->int32 | {a: [1099511627776, 1]}@Table[a: int64] | {a: [0, 1]}@Table[a: int32] | +| truncate:float->int | {a: [1.9, -2.1]}@Table[a: float64] | {a: [1, -2]}@Table[a: int64] | +| names:mismatch | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| names:reordered | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| names:field-count | {a: [1, 2], b: [1.5, 2.5]}@Table[a: int64, b: float64] | ERR@ValueError | +| nullable:false-with-nulls | {a: [1, None]}@Table[a: int64] | ERR@ValueError | +| nullable:false-no-nulls | {a: [1, 2]}@Table[a: int64] | {a: [1, 2]}@Table[a: int32] | +| timestamp:us->ns | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[ns]] | +| timestamp:attach-tz | {ts: [1970-01-01 00:00:00, 1970-01-01 00:00:01]}@Table[ts: timestamp[us]] | {ts: [1970-01-01 00:00:00+00:00, 1970-01-01 00:00:01+00:00]}@Table[ts: timestamp[us, tz=UTC]] | +| string->large_string | {s: [hello, world, None]}@Table[s: string] | {s: [hello, world, None]}@Table[s: large_string] | +| binary->large_binary | {b: [b'x', b'yz', None]}@Table[b: binary] | {b: [b'x', b'yz', None]}@Table[b: large_binary] | +| nested:list | {lst: [[1, 2], [3], None]}@Table[lst: list<item: int64>] | {lst: [[1, 2], [3], None]}@Table[lst: list<item: int32>] | +| nested:list-overflow | {lst: [[1099511627776, 1]]}@Table[lst: list<item: int64>] | {lst: [[0, 1]]}@Table[lst: list<item: int32>] | +| nested:struct | {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int64, y: string>] | {st: [[('x', 1), ('y', 'a')], None]}@Table[st: struct<x: int32, y: large_string>] | +| nested:map | {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int64>] | {m: [[('k', 1), ('j', 2)], None]}@Table[m: map<string, int32>] | +| multi-chunk-column | {a: [1, 2, 3, None]}@Table[a: int64] | {a: [1, 2, 3, None]}@Table[a: int32] | +| empty:0-columns | {}@Table[] | {}@Table[] | +| empty:columns-no-rows | {i: [], s: []}@Table[i: int64, s: string] | {i: [], s: []}@Table[i: int32, s: large_string] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.csv new file mode 100644 index 0000000000000..bdb4eb03b4861 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.csv @@ -0,0 +1,13 @@ +test case pyarrow table pandas dataframe +empty:0-columns {}@Table[] {}@Dataframe[] +empty:columns-no-rows {i: [], s: []}@Table[i: int64, s: string] {'i': [], 's': []}@Dataframe[i int64, s object] +single-column {i: [1, 2, None]}@Table[i: int64] {'i': [1.0, 2.0, nan]}@Dataframe[i float64] +single-column:string {s: [hello, world, None]}@Table[s: string] {'s': ['hello', 'world', None]}@Dataframe[s object] +multi-column:mixed-scalar {i: [1, 2, None], s: [a, b, None], f: [1.5, 2.5, 3.5], b: [True, False, None]}@Table[i: int64, s: string, f: float64, b: bool] {'i': [1.0, 2.0, nan], 's': ['a', 'b', None], 'f': [1.5, 2.5, 3.5], 'b': [True, False, None]}@Dataframe[i float64, s object, f float64, b object] +multi-column:all-null {i: [None, None], s: [None, None]}@Table[i: int64, s: string] {'i': [nan, nan], 's': [None, None]}@Dataframe[i float64, s object] +multi-column:nested {lst: [[1, 2], [3], None], st: [[('x', 1)], None, [('x', 3)]]}@Table[lst: list<item: int64>, st: struct<x: int64>] {'lst': [array([1, 2]), array([3]), None], 'st': [{'x': 1}, None, {'x': 3}]}@Dataframe[lst object, st object] +temporal:timestamp {ts: [2020-01-01 05:30:00, 2021-06-15 23:59:00]}@Table[ts: timestamp[us]] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[us]] +temporal:date32 {d: [2020-01-01, 2021-06-15]}@Table[d: date32[day]] {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] +temporal:date64 {d: [2020-01-01, 2021-06-15]}@Table[d: date64[ms]] {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] +temporal:date32-far-future {d: [9999-12-31]}@Table[d: date32[day]] {'d': [datetime.date(9999, 12, 31)]}@Dataframe[d object] +temporal:multi-column-mix {ts: [2020-01-01 05:30:00, 2020-01-01 05:30:00], d: [2020-01-01, 9999-12-31], i: [1, 2]}@Table[ts: timestamp[us], d: date32[day], i: int64] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [datetime.date(2020, 1, 1), datetime.date(9999, 12, 31)], 'i': [1, 2]}@Dataframe[ts datetime64[us], d object, i int64] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.md new file mode 100644 index 0000000000000..cafeca290702d --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas.md @@ -0,0 +1,14 @@ +| test case | pyarrow table | pandas dataframe | +|----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| empty:0-columns | {}@Table[] | {}@Dataframe[] | +| empty:columns-no-rows | {i: [], s: []}@Table[i: int64, s: string] | {'i': [], 's': []}@Dataframe[i int64, s object] | +| single-column | {i: [1, 2, None]}@Table[i: int64] | {'i': [1.0, 2.0, nan]}@Dataframe[i float64] | +| single-column:string | {s: [hello, world, None]}@Table[s: string] | {'s': ['hello', 'world', None]}@Dataframe[s object] | +| multi-column:mixed-scalar | {i: [1, 2, None], s: [a, b, None], f: [1.5, 2.5, 3.5], b: [True, False, None]}@Table[i: int64, s: string, f: float64, b: bool] | {'i': [1.0, 2.0, nan], 's': ['a', 'b', None], 'f': [1.5, 2.5, 3.5], 'b': [True, False, None]}@Dataframe[i float64, s object, f float64, b object] | +| multi-column:all-null | {i: [None, None], s: [None, None]}@Table[i: int64, s: string] | {'i': [nan, nan], 's': [None, None]}@Dataframe[i float64, s object] | +| multi-column:nested | {lst: [[1, 2], [3], None], st: [[('x', 1)], None, [('x', 3)]]}@Table[lst: list<item: int64>, st: struct<x: int64>] | {'lst': [array([1, 2]), array([3]), None], 'st': [{'x': 1}, None, {'x': 3}]}@Dataframe[lst object, st object] | +| temporal:timestamp | {ts: [2020-01-01 05:30:00, 2021-06-15 23:59:00]}@Table[ts: timestamp[us]] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[us]] | +| temporal:date32 | {d: [2020-01-01, 2021-06-15]}@Table[d: date32[day]] | {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] | +| temporal:date64 | {d: [2020-01-01, 2021-06-15]}@Table[d: date64[ms]] | {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] | +| temporal:date32-far-future | {d: [9999-12-31]}@Table[d: date32[day]] | {'d': [datetime.date(9999, 12, 31)]}@Dataframe[d object] | +| temporal:multi-column-mix | {ts: [2020-01-01 05:30:00, 2020-01-01 05:30:00], d: [2020-01-01, 9999-12-31], i: [1, 2]}@Table[ts: timestamp[us], d: date32[day], i: int64] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [datetime.date(2020, 1, 1), datetime.date(9999, 12, 31)], 'i': [1, 2]}@Dataframe[ts datetime64[us], d object, i int64] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.csv b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.csv new file mode 100644 index 0000000000000..706bd3c5bff1f --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.csv @@ -0,0 +1,8 @@ +test case pyarrow table pandas dataframe pandas dataframe (date_as_object=False) +temporal:timestamp {ts: [2020-01-01 05:30:00, 2021-06-15 23:59:00]}@Table[ts: timestamp[us]] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[ns]] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[ns]] +temporal:date32 {d: [2020-01-01, 2021-06-15]}@Table[d: date32[day]] {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] {'d': [Timestamp('2020-01-01 00:00:00'), Timestamp('2021-06-15 00:00:00')]}@Dataframe[d datetime64[ns]] +temporal:date64 {d: [2020-01-01, 2021-06-15]}@Table[d: date64[ms]] {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] {'d': [Timestamp('2020-01-01 00:00:00'), Timestamp('2021-06-15 00:00:00')]}@Dataframe[d datetime64[ns]] +temporal:date32-far-future {d: [9999-12-31]}@Table[d: date32[day]] {'d': [datetime.date(9999, 12, 31)]}@Dataframe[d object] {'d': [Timestamp('1816-03-29 05:56:08.066277376')]}@Dataframe[d datetime64[ns]] +temporal:multi-column-mix {ts: [2020-01-01 05:30:00, 2020-01-01 05:30:00], d: [2020-01-01, 9999-12-31], i: [1, 2]}@Table[ts: timestamp[us], d: date32[day], i: int64] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [datetime.date(2020, 1, 1), datetime.date(9999, 12, 31)], 'i': [1, 2]}@Dataframe[ts datetime64[ns], d object, i int64] {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [Timestamp('2020-01-01 00:00:00'), Timestamp('1816-03-29 05:56:08.066277376')], 'i': [1, 2]}@Dataframe[ts datetime64[ns], d datetime64[ns], i int64] +temporal:timestamp-overflow {ts: [2500-01-01 00:00:00]}@Table[ts: timestamp[s]] ERR@ArrowInvalid ERR@ArrowInvalid +temporal:duration-overflow {dur: [109500 days, 0:00:00]}@Table[dur: duration[s]] {'dur': [Timedelta('-104004 days +00:25:26.290448384')]}@Dataframe[dur timedelta64[ns]] {'dur': [Timedelta('-104004 days +00:25:26.290448384')]}@Dataframe[dur timedelta64[ns]] diff --git a/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.md b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.md new file mode 100644 index 0000000000000..7730a2ec0ab79 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/golden_pyarrow_table_to_pandas_coerce_temporal.md @@ -0,0 +1,9 @@ +| test case | pyarrow table | pandas dataframe | pandas dataframe (date_as_object=False) | +|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| temporal:timestamp | {ts: [2020-01-01 05:30:00, 2021-06-15 23:59:00]}@Table[ts: timestamp[us]] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[ns]] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2021-06-15 23:59:00')]}@Dataframe[ts datetime64[ns]] | +| temporal:date32 | {d: [2020-01-01, 2021-06-15]}@Table[d: date32[day]] | {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] | {'d': [Timestamp('2020-01-01 00:00:00'), Timestamp('2021-06-15 00:00:00')]}@Dataframe[d datetime64[ns]] | +| temporal:date64 | {d: [2020-01-01, 2021-06-15]}@Table[d: date64[ms]] | {'d': [datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)]}@Dataframe[d object] | {'d': [Timestamp('2020-01-01 00:00:00'), Timestamp('2021-06-15 00:00:00')]}@Dataframe[d datetime64[ns]] | +| temporal:date32-far-future | {d: [9999-12-31]}@Table[d: date32[day]] | {'d': [datetime.date(9999, 12, 31)]}@Dataframe[d object] | {'d': [Timestamp('1816-03-29 05:56:08.066277376')]}@Dataframe[d datetime64[ns]] | +| temporal:multi-column-mix | {ts: [2020-01-01 05:30:00, 2020-01-01 05:30:00], d: [2020-01-01, 9999-12-31], i: [1, 2]}@Table[ts: timestamp[us], d: date32[day], i: int64] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [datetime.date(2020, 1, 1), datetime.date(9999, 12, 31)], 'i': [1, 2]}@Dataframe[ts datetime64[ns], d object, i int64] | {'ts': [Timestamp('2020-01-01 05:30:00'), Timestamp('2020-01-01 05:30:00')], 'd': [Timestamp('2020-01-01 00:00:00'), Timestamp('1816-03-29 05:56:08.066277376')], 'i': [1, 2]}@Dataframe[ts datetime64[ns], d datetime64[ns], i int64] | +| temporal:timestamp-overflow | {ts: [2500-01-01 00:00:00]}@Table[ts: timestamp[s]] | ERR@ArrowInvalid | ERR@ArrowInvalid | +| temporal:duration-overflow | {dur: [109500 days, 0:00:00]}@Table[dur: duration[s]] | {'dur': [Timedelta('-104004 days +00:25:26.290448384')]}@Dataframe[dur timedelta64[ns]] | {'dur': [Timedelta('-104004 days +00:25:26.290448384')]}@Dataframe[dur timedelta64[ns]] | \ No newline at end of file diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py index 128557fa6548b..fe585b207a809 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_cast.py @@ -60,15 +60,15 @@ from decimal import Decimal from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_pyarrow: import pyarrow as pa @@ -127,9 +127,9 @@ def _try_cast(self, src_arr, tgt_type, safe=True): """ try: result = src_arr.cast(tgt_type, safe=safe) - return self.repr_value(result, max_len=0) except Exception as e: return f"ERR@{type(e).__name__}" + return self.repr_value(result, max_len=0) # ============================================================ diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_default.py new file mode 100644 index 0000000000000..07a209641d546 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_default.py @@ -0,0 +1,442 @@ +# +# 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. +# + +""" +Tests for PyArrow Array.from_pandas() with default arguments using golden file comparison. + +This test monitors the pandas -> Arrow direction, which PySpark relies on for +``createDataFrame(pandas_df)`` and for every pandas UDF's return value. PySpark calls +``pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck)`` in +``pyspark/sql/conversion.py`` and ``pyspark/sql/pandas/conversion.py``; the non-default +arguments are covered by ``test_pyarrow_array_from_pandas_non_default.py``, which reuses +the source Series built here. + +Rows are grouped by how pandas stores the Series, because that decides which branch +PySpark takes when computing ``mask``: + +- numpy-backed dtypes, where ``mask=series.isnull()`` is passed; +- dtypes implementing the ``__arrow_array__`` protocol, where PySpark passes ``mask=None`` + because supplying a mask raises. The protocol means "can export Arrow", not "is stored + as Arrow": ``Int64`` is numpy values plus a byte mask and ``string[python]`` is an object + ndarray, yet both implement it alongside the genuinely Arrow-backed ``[pyarrow]`` dtypes; +- values chosen to be lossy or ambiguous when coerced, which the non-default type tests + need and whose unconverted baseline is recorded here. + +## Golden File Cell Format + +Each cell uses the value@type format: +- pandas Series: "python_list_repr@Series[dtype]" +- PyArrow Array: "python_list_repr@arrow_type" +- PyArrow ChunkedArray: "python_list_repr@chunked<arrow_type>" +- Error: "ERR@ExceptionClassName" + +``from_pandas`` returns a ChunkedArray rather than an Array when the input Series is backed +by a chunked Arrow array, and the two are otherwise indistinguishable because both report +the element type. The distinction is a real contract: ``create_arrow_table_from_pandas`` +builds a ``pa.Table`` (which accepts either) specifically because of it -- see SPARK-46776. +``chunked<...>`` follows Arrow's angle-bracket spelling for parameterized types, keeping it +distinct from the square-bracket ``Series[...]`` / ``ndarray[...]`` used for pandas and +numpy containers. The chunk count is deliberately not recorded: it varies with batching, +whereas chunked-versus-not is the behavior under test. + +## Regenerating Golden Files + +Set SPARK_GENERATE_GOLDEN_FILES=1 before running: + + SPARK_GENERATE_GOLDEN_FILES=1 python -m pytest \\ + python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_default.py + +## PyArrow and pandas Version Compatibility + +The golden files capture behavior for specific PyArrow and pandas versions. +Regenerate when upgrading either dependency, as from_pandas() behavior may change. +""" + +import datetime +import decimal +import unittest + +from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.utils import ( + have_numpy, + have_pandas, + have_pyarrow, + numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, +) + +if have_pandas: + import pandas as pd +if have_pyarrow: + import pyarrow as pa +if have_numpy: + import numpy as np + + +class _PyArrowFromPandasTestBase(GoldenFileTestMixin, unittest.TestCase): + """ + Shared machinery for pa.Array.from_pandas() golden file tests. + + Owns the source Series inventory: three disjoint group methods unioned by + ``_build_source_arrays``, whose order fixes the row order. The default test below and + the non-default tests in ``test_pyarrow_array_from_pandas_non_default.py`` subclass this + to reuse the whole inventory or one group, along with ``repr_from_pandas_result``. This + base defines no ``test_*`` methods, so it contributes no tests itself. + """ + + @staticmethod + def repr_from_pandas_result(value): + """ + Format a ``from_pandas`` return value, marking a ChunkedArray as ``chunked<type>``. + + ``repr_value`` reports the element type, which is the same for an Array and a + ChunkedArray, so without this the two are indistinguishable in a cell. + """ + rendered = GoldenFileTestMixin.repr_value(value, max_len=0) + if isinstance(value, pa.ChunkedArray): + values, _, arrow_type = rendered.rpartition("@") + return f"{values}@chunked<{arrow_type}>" + return rendered + + def _from_pandas_cell(self, series, **from_pandas_kwargs) -> str: + """ + Convert ``series`` via ``from_pandas(**from_pandas_kwargs)`` and format the result + as a golden-file cell, returning ``ERR@<ExceptionClass>`` if the conversion raises. + """ + try: + result = pa.Array.from_pandas(series, **from_pandas_kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + return self.repr_from_pandas_result(result) + + def _numpy_backed_sources(self): + """ + Series whose storage is numpy, so ``hasattr(series.array, "__arrow_array__")`` is + False and PySpark passes ``mask=series.isnull()``. + """ + sources = {} + + # ===================================================================== + # Integer types + # ===================================================================== + for dtype in ["int8", "int16", "int32", "int64"]: + info = np.iinfo(dtype) + sources[f"{dtype}:standard"] = pd.Series([0, 1, -1, info.max, info.min], dtype=dtype) + sources[f"{dtype}:empty"] = pd.Series([], dtype=dtype) + # A numpy integer Series cannot hold a null, so pandas widens it to float64. + sources["int64:nullable"] = pd.Series([0, 1, None]) + + for dtype in ["uint8", "uint16", "uint32", "uint64"]: + sources[f"{dtype}:standard"] = pd.Series([0, 1, np.iinfo(dtype).max], dtype=dtype) + + # ===================================================================== + # Float and boolean types + # ===================================================================== + for dtype in ["float32", "float64"]: + sources[f"{dtype}:standard"] = pd.Series([0.0, 1.5, -1.5], dtype=dtype) + sources[f"{dtype}:nullable"] = pd.Series([0.0, np.nan, 1.5], dtype=dtype) + sources[f"{dtype}:empty"] = pd.Series([], dtype=dtype) + + sources["bool:standard"] = pd.Series([True, False, True]) + sources["bool:empty"] = pd.Series([], dtype="bool") + + # ===================================================================== + # Object types (string, binary, decimal) + # ===================================================================== + # Only the string rows need an explicit dtype; pandas 3 would infer its ``str`` + # dtype and they would stop being object rows. + sources["object:string"] = pd.Series(["hello", "world", ""], dtype=object) + sources["object:string-nullable"] = pd.Series(["hello", None, "world"], dtype=object) + # Unpinned, so this records the dtype pandas infers from Python strings: object on + # pandas 2, its dedicated str dtype on pandas 3, which Arrow reads as large_string. + sources["string:inferred"] = pd.Series(["hello", "world"]) + sources["object:bytes"] = pd.Series([b"hello", b"world"]) + sources["object:empty"] = pd.Series([], dtype=object) + sources["object:all-null"] = pd.Series([None, None], dtype=object) + sources["object:decimal"] = pd.Series([decimal.Decimal("1.50"), decimal.Decimal("-2.25")]) + + # ===================================================================== + # Nested types + # ===================================================================== + # Inference builds these from Python containers. It cannot reach ``map`` (that + # needs an explicit type), nor ``large_list`` / ``fixed_size_list``, which + # ``to_arrow_type`` never requests. + sources["list<int64>:standard"] = pd.Series([[1, 2], [3]]) + sources["list<int64>:nullable"] = pd.Series([[1, 2], None]) + sources["list<int64>:null-element"] = pd.Series([[1, None], [3]]) + sources["list<string>:standard"] = pd.Series([["a", "b"], ["c"]]) + sources["list<list<int64>>:standard"] = pd.Series([[[1, 2], [3]], [[4]]]) + sources["list<struct>:standard"] = pd.Series([[{"a": 1}], [{"a": 2}]]) + sources["struct:standard"] = pd.Series([{"a": 1, "b": "x"}]) + sources["struct:nullable"] = pd.Series([{"a": 1, "b": "x"}, None]) + sources["struct<struct>:standard"] = pd.Series([{"a": {"b": 1}}]) + sources["struct<list<int64>>:standard"] = pd.Series([{"a": [1, 2]}]) + + # These feed the nested type tests (non_default): a child value that overflows a + # narrower element type, and homogeneous-value dicts, which infer as ``struct`` but + # convert to ``map`` when the requested type asks for one. + sources["list<int64>:overflow"] = pd.Series([[300, 2], [3]]) + sources["struct:overflow"] = pd.Series([{"a": 300, "b": "x"}]) + sources["struct<int64>:standard"] = pd.Series([{"a": 1, "b": 2}]) + sources["struct<int64>:overflow"] = pd.Series([{"a": 300, "b": 2}]) + + # ===================================================================== + # Temporal types + # ===================================================================== + # Each row pins its resolution so the row name stays accurate on both pandas + # majors: pandas 2 infers "ns" from a datetime, pandas 3 infers "us". + dt = datetime.datetime(2024, 6, 15, 18, 30, 0) + for unit in ["ns", "us"]: + sources[f"datetime64[{unit}]:standard"] = pd.Series([dt], dtype=f"datetime64[{unit}]") + sources[f"datetime64[{unit}]:nullable"] = pd.Series( + [dt, None], dtype=f"datetime64[{unit}]" + ) + sources[f"datetime64[{unit}]:empty"] = pd.Series([], dtype=f"datetime64[{unit}]") + sources[f"datetime64[{unit},tz]:standard"] = pd.Series( + [dt], dtype=f"datetime64[{unit}, UTC]" + ) + sources[f"datetime64[{unit},tz]:nullable"] = pd.Series( + [dt, None], dtype=f"datetime64[{unit}, UTC]" + ) + sources[f"timedelta64[{unit}]:standard"] = pd.Series( + pd.to_timedelta(["1 days", "2 hours"]), dtype=f"timedelta64[{unit}]" + ) + sources[f"timedelta64[{unit}]:nullable"] = pd.Series( + pd.to_timedelta(["1 days", None]), dtype=f"timedelta64[{unit}]" + ) + + # Unpinned, so the cell always duplicates one of the pinned rows above. Which one + # it matches is the behavior recorded: "ns" on pandas 2, "us" on pandas 3. + sources["datetime64:inferred"] = pd.Series([dt]) + sources["timedelta64:inferred"] = pd.Series(pd.to_timedelta(["1 days", "2 hours"])) + + # datetime64[ns] spans only 1677-2262, so a far-past value needs microseconds. + sources["datetime64[us]:out-of-ns-range"] = pd.Series( + [datetime.datetime(1500, 1, 1)], dtype="datetime64[us]" + ) + + # pandas has no native date or time dtype, so these stay object on their own. + sources["date:standard"] = pd.Series([datetime.date(2024, 6, 15)]) + sources["time:standard"] = pd.Series([datetime.time(18, 30, 45)]) + + # In object dtype pyarrow takes the unit from datetime's own microsecond + # resolution, not from a numpy dtype, so these give timestamp[us] where the + # datetime64[ns] rows above give timestamp[ns] -- and sub-microsecond digits are + # dropped silently, while an explicit type=timestamp("us") raises ArrowInvalid. + sources["object:datetime"] = pd.Series([dt], dtype=object) + sources["object:datetime-sub-us"] = pd.Series( + [pd.Timestamp("2024-01-01 00:00:00.000000123")], dtype=object + ) + sources["object:timedelta"] = pd.Series([datetime.timedelta(days=1, hours=2)], dtype=object) + + # ===================================================================== + # Categorical type + # ===================================================================== + # PySpark casts a categorical to its categories' dtype first, but the raw + # categorical is what a user hands to createDataFrame. + sources["category:standard"] = pd.Series(pd.Categorical(["a", "b", "a"])) + sources["category:nullable"] = pd.Series(pd.Categorical(["a", None, "b"])) + + return sources + + def _protocol_sources(self): + """ + Series implementing ``__arrow_array__``, so PySpark passes ``mask=None`` -- + supplying one raises, because the protocol returns an array whose validity bitmap + is already built. + + The protocol means "can export Arrow", not "is stored as Arrow", and the rows + below deliberately cover both: ``Int64`` is numpy values plus a byte mask and + ``string[python]`` is an object ndarray, while the ``[pyarrow]`` dtypes hold Arrow + buffers. PySpark does not distinguish them -- ``conversion.py:435`` branches only + on the protocol -- so they share one group. + """ + sources = {} + + # ===================================================================== + # Nullable extension dtypes (numpy values plus a byte mask) + # ===================================================================== + for dtype in ["Int8", "Int16", "Int32", "Int64"]: + info = np.iinfo(dtype.lower()) + sources[f"{dtype}:standard"] = pd.Series([0, 1, info.max, info.min], dtype=dtype) + sources[f"{dtype}:nullable"] = pd.Series([0, 1, None], dtype=dtype) + sources["UInt64:standard"] = pd.Series([0, 1, np.iinfo("uint64").max], dtype="UInt64") + sources["Int64:empty"] = pd.Series([], dtype="Int64") + sources["Int64:all-null"] = pd.Series([None, None], dtype="Int64") + + sources["Float64:standard"] = pd.Series([0.0, 1.5], dtype="Float64") + sources["Float64:nullable"] = pd.Series([0.0, None], dtype="Float64") + + sources["boolean:standard"] = pd.Series([True, False], dtype="boolean") + sources["boolean:nullable"] = pd.Series([True, None], dtype="boolean") + + # StringArray: an object ndarray of Python str, yet it exports Arrow. + sources["string[python]:standard"] = pd.Series(["hello", "world"], dtype="string[python]") + sources["string[python]:nullable"] = pd.Series(["hello", None], dtype="string[python]") + sources["string[python]:empty"] = pd.Series([], dtype="string[python]") + + # ===================================================================== + # PyArrow-backed dtypes (Arrow buffers, so the handoff is zero-copy) + # ===================================================================== + # Since pandas 2.2 a pyarrow-backed string stores large_string, and the protocol + # hands that back rather than the 32-bit string the name suggests. PySpark works + # around pyarrow < 19 ignoring a narrower request (SPARK-46776). + sources["int64[pyarrow]:standard"] = pd.Series([0, 1, -1], dtype="int64[pyarrow]") + sources["int64[pyarrow]:nullable"] = pd.Series([0, 1, None], dtype="int64[pyarrow]") + sources["int64[pyarrow]:empty"] = pd.Series([], dtype="int64[pyarrow]") + sources["double[pyarrow]:nullable"] = pd.Series([0.0, None], dtype="double[pyarrow]") + sources["bool[pyarrow]:nullable"] = pd.Series([True, None], dtype="bool[pyarrow]") + sources["string[pyarrow]:standard"] = pd.Series(["hello", "world"], dtype="string[pyarrow]") + sources["string[pyarrow]:nullable"] = pd.Series(["hello", None], dtype="string[pyarrow]") + sources["string[pyarrow]:empty"] = pd.Series([], dtype="string[pyarrow]") + # Binary is not auto-promoted to large_binary the way string is, so large_binary + # has to be requested explicitly. This is the binary counterpart of the string + # rows, and lets the type tests exercise the binary half of SPARK-46776. + sources["large_binary[pyarrow]:standard"] = pd.Series( + [b"hello", b"world"], dtype="large_binary[pyarrow]" + ) + sources["timestamp[us][pyarrow]:standard"] = pd.Series( + [datetime.datetime(2024, 1, 1, 12, 0, 0)], dtype="timestamp[us][pyarrow]" + ) + + # A chunked backing array makes from_pandas return a ChunkedArray, which + # pa.RecordBatch.from_arrays rejects -- hence PySpark builds a pa.Table. + sources["int64[pyarrow]:single-chunk"] = pd.Series( + pd.arrays.ArrowExtensionArray(pa.chunked_array([pa.array([1, 2], pa.int64())])) + ) + sources["int64[pyarrow]:multi-chunk"] = pd.Series( + pd.arrays.ArrowExtensionArray( + pa.chunked_array([pa.array([1, 2], pa.int64()), pa.array([3], pa.int64())]) + ) + ) + + return sources + + def _coercion_sources(self): + """ + Series whose values are lossy or ambiguous once a target type is requested. + + These exist for the non-default type tests, where ``safe=`` decides whether the + conversion raises or silently changes the value. Recording their unconverted + results here gives those cells a baseline to be read against. + """ + sources = {} + + # ===================================================================== + # Values that narrow lossily + # ===================================================================== + sources["int64:overflow"] = pd.Series([300, 1]) + sources["float64:fractional"] = pd.Series([1.5, 2.5]) + sources["float64:infinity"] = pd.Series([np.inf, 1.0]) + sources["float64:precision"] = pd.Series([1.1234567890123]) + + # The only temporal case where safe= flips. + sources["datetime64[ns]:sub-us"] = pd.Series( + pd.to_datetime(["2024-01-01 00:00:00.000000123"]) + ) + + # ===================================================================== + # Values whose inferred type depends on element order + # ===================================================================== + # Inference commits to the first element's type, so these same values either lose + # the time component or raise, depending on their order. + sources["object:date-then-datetime"] = pd.Series( + [datetime.date(2024, 1, 1), datetime.datetime(2024, 1, 1, 5, 30)], dtype=object + ) + sources["object:datetime-then-date"] = pd.Series( + [datetime.datetime(2024, 1, 1, 5, 30), datetime.date(2024, 1, 1)], dtype=object + ) + + return sources + + def _build_source_arrays(self): + """Build an ordered dict of named source pandas Series for testing.""" + sources = {} + for group in [ + self._numpy_backed_sources(), + self._protocol_sources(), + self._coercion_sources(), + ]: + sources.update(group) + return sources + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayFromPandasDefaultTests(_PyArrowFromPandasTestBase): + """Tests pa.Array.from_pandas() with default arguments via golden file comparison.""" + + def test_from_pandas_default(self): + """Test pa.Array.from_pandas() with default arguments against golden file.""" + sources = self._build_source_arrays() + row_names = list(sources.keys()) + col_names = ["pandas series", "arrow array"] + + overrides = {} + if LooseVersion(pd.__version__) >= LooseVersion("3.0.0"): + # Only the deliberately unpinned rows move: pandas 3 infers microseconds where + # pandas 2 inferred nanoseconds, and its dedicated str dtype is backed by + # large_string -- which a categorical's values inherit too. + overrides.update( + { + ("string:inferred", "pandas series"): "['hello', 'world']@Series[str]", + ("string:inferred", "arrow array"): "[hello, world]@large_string", + ("datetime64:inferred", "pandas series"): ( + "[Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]]" + ), + ("datetime64:inferred", "arrow array"): "[2024-06-15 18:30:00]@timestamp[us]", + ("timedelta64:inferred", "pandas series"): ( + "[Timedelta('1 days 00:00:00'), " + "Timedelta('0 days 02:00:00')]@Series[timedelta64[us]]" + ), + ("timedelta64:inferred", "arrow array"): ( + "[1 day, 0:00:00, 2:00:00]@duration[us]" + ), + ("category:standard", "arrow array"): ( + "[a, b, a]@dictionary<values=large_string, indices=int8, ordered=0>" + ), + ("category:nullable", "arrow array"): ( + "[a, None, b]@dictionary<values=large_string, indices=int8, ordered=0>" + ), + } + ) + + def compute_cell(row_name, col_name): + series = sources[row_name] + if col_name == "pandas series": + return self.repr_value(series, max_len=0) + else: + return self._from_pandas_cell(series) + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_array_from_pandas_default", + index_name="test case", + overrides=overrides, + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_non_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_non_default.py new file mode 100644 index 0000000000000..fa2833b078d16 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_non_default.py @@ -0,0 +1,390 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +Tests for PyArrow Array.from_pandas() with non-default arguments, using golden file +comparison. + +PySpark passes ``pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck)`` +to convert a pandas Series into an Arrow array. The bare-argument behavior is covered by +``test_pyarrow_array_from_pandas_default.py``, whose source Series inventory +(``_PyArrowFromPandasTestBase``) is reused here; this file records the non-default +arguments so CI fails loudly if they drift across pandas/PyArrow/NumPy upgrades. + +## The ``mask`` argument + +PySpark derives ``mask`` from how the Series is stored +(``conversion.py:435``, ``pandas/conversion.py:113``):: + + mask = None if hasattr(series.array, "__arrow_array__") else series.isnull() + +- **numpy-backed** dtypes take ``mask=series.isnull()``, which agrees with the nulls + ``from_pandas`` infers at ``mask=None``; this test pins that agreement. +- **protocol** dtypes (implementing ``__arrow_array__``) return a finished Arrow array with + its own validity bitmap, so PyArrow rejects any mask -- even an all-False no-op -- with + ``ValueError``; PySpark passes ``mask=None``. + +So ``mask`` is fixed by the row's dtype, hence a column pair (``mask=None`` vs +``mask=isnull()``) rather than a matrix dimension. + +## Golden File Cell Format + +Each cell uses the value@type format: +- pandas Series: "python_list_repr@Series[dtype]" +- PyArrow Array: "python_list_repr@arrow_type" +- PyArrow ChunkedArray: "python_list_repr@chunked<arrow_type>" +- Error: "ERR@ExceptionClassName" + +## Regenerating Golden Files + +Set SPARK_GENERATE_GOLDEN_FILES=1 before running: + + SPARK_GENERATE_GOLDEN_FILES=1 python -m pytest \\ + python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_from_pandas_non_default.py + +## PyArrow and pandas Version Compatibility + +The golden files capture behavior for specific PyArrow and pandas versions. +Regenerate when upgrading either dependency, as from_pandas() behavior may change. +The committed golden files were generated with pandas 2.3.3, pyarrow 24.0.0, and +numpy 2.4.1. +""" + +import platform +import unittest + +from pyspark.loose_version import LooseVersion +from pyspark.testing.utils import ( + have_numpy, + have_pandas, + have_pyarrow, + numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, +) + +# Imported as a module, not `from ... import _PyArrowFromPandasTestBase`, so that the +# non-default classes pick up only the shared base -- the default test class is not +# re-collected here. +from pyspark.tests.upstream.pyarrow import test_pyarrow_array_from_pandas_default + +if have_pandas: + import pandas as pd +if have_pyarrow: + import pyarrow as pa + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayFromPandasMaskTests( + test_pyarrow_array_from_pandas_default._PyArrowFromPandasTestBase +): + """ + Tests pa.Array.from_pandas(series, mask=...) via golden file comparison. + + Reuses the full source inventory. numpy-backed rows accept a mask and record equal + results for mask=None and mask=isnull(); protocol rows reject a mask and record + ERR@ValueError for mask=isnull(). + """ + + COL_MASK_NONE = "mask=None" + COL_MASK_ISNULL = "mask=isnull()" + + def test_from_pandas_mask(self): + """Test pa.Array.from_pandas() with the mask argument against golden file.""" + sources = self._build_source_arrays() + row_names = list(sources.keys()) + col_names = ["pandas series", self.COL_MASK_NONE, self.COL_MASK_ISNULL] + + # Version-specific expected values go here, keyed by (row, col), when a newer + # pandas/PyArrow/NumPy legitimately changes a cell's output. + overrides: dict[tuple[str, str], str] = {} + if LooseVersion(pd.__version__) >= LooseVersion("3.0.0"): + # The deliberately unpinned rows infer microseconds on pandas 3 where pandas 2 + # infers nanoseconds, and a categorical's values inherit the new str dtype's + # large_string backing. These stay numpy-backed, so both mask columns move + # together and accept the mask. + overrides[("datetime64:inferred", "pandas series")] = ( + "[Timestamp('2024-06-15 18:30:00')]@Series[datetime64[us]]" + ) + overrides[("timedelta64:inferred", "pandas series")] = ( + "[Timedelta('1 days 00:00:00'), " + "Timedelta('0 days 02:00:00')]@Series[timedelta64[us]]" + ) + for col in (self.COL_MASK_NONE, self.COL_MASK_ISNULL): + overrides[("datetime64:inferred", col)] = "[2024-06-15 18:30:00]@timestamp[us]" + overrides[("timedelta64:inferred", col)] = "[1 day, 0:00:00, 2:00:00]@duration[us]" + overrides[("category:standard", col)] = ( + "[a, b, a]@dictionary<values=large_string, indices=int8, ordered=0>" + ) + overrides[("category:nullable", col)] = ( + "[a, None, b]@dictionary<values=large_string, indices=int8, ordered=0>" + ) + # string:inferred is object (numpy-backed) on pandas 2 but the dedicated str + # dtype (large_string, __arrow_array__) on pandas 3, so it moves from a + # mask-accepting row to a protocol row: mask=isnull() flips to ERR@ValueError. + overrides[("string:inferred", "pandas series")] = "['hello', 'world']@Series[str]" + overrides[("string:inferred", self.COL_MASK_NONE)] = "[hello, world]@large_string" + overrides[("string:inferred", self.COL_MASK_ISNULL)] = "ERR@ValueError" + + def compute_cell(row_name, col_name): + series = sources[row_name] + if col_name == "pandas series": + return self.repr_value(series, max_len=0) + elif col_name == self.COL_MASK_NONE: + return self._from_pandas_cell(series, mask=None) + elif col_name == self.COL_MASK_ISNULL: + return self._from_pandas_cell(series, mask=series.isnull()) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_array_from_pandas_mask", + index_name="test case", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayFromPandasTypeScalarTests( + test_pyarrow_array_from_pandas_default._PyArrowFromPandasTestBase +): + """ + Tests pa.Array.from_pandas(series, type=..., safe=...) for SCALAR target types via golden + file comparison. Nested targets (list / map / struct) are covered by a separate class. + + Spark uses ``type=`` diagonally (source dtype and requested Arrow type share one + schema), so this pins a focused set of rows that make type=/safe= observable rather + than a dense source x target product: the off-diagonal cells (e.g. int -> large_binary) + are pyarrow-construction trivia Spark never hits, and general conversion is + pa.Array.cast's job. safe=True/False are two methods and two goldens, not doubled + columns. ``mask`` stays None (fixed by storage backend, covered by the mask tests), + isolating type=/safe=. + """ + + @staticmethod + def _get_target_types(): + """Scalar to_arrow_type targets that discriminate type=/safe=, plus duration/time64 + for the timedelta and time-of-day diagonals.""" + return [ + pa.int8(), + pa.int64(), + pa.float32(), + pa.timestamp("us"), + pa.date32(), + pa.duration("us"), + pa.time64("ns"), + pa.string(), + pa.binary(), + ] + + def _type_source_arrays(self): + """ + Clean family representatives, the protocol rows that expose the safe= drop and + SPARK-46776, and the shared coercion rows (reused so a changed cell is attributable + across the default/mask/type goldens). + """ + pool = {**self._numpy_backed_sources(), **self._protocol_sources()} + selected = [ + # Clean family reps. + "int64:standard", + "float64:standard", + "bool:standard", + "object:string", + "object:bytes", + "date:standard", + "object:datetime", + "object:timedelta", + "time:standard", + # Protocol rows: Int64 = safe-drop contrast; last two = SPARK-46776's pyarrow < 19 + # followup (a narrower type is ignored on the __arrow_array__ path -> stored type). + "Int64:standard", + "string[pyarrow]:standard", + "large_binary[pyarrow]:standard", + ] + sources = {name: pool[name] for name in selected} + # Coercion group reused in full; the numpy/protocol rows above are cherry-picked. + sources.update(self._coercion_sources()) + return sources + + def _compare_type_matrix(self, safe, golden_file_prefix, overrides): + sources = self._type_source_arrays() + target_types = self._get_target_types() + target_names = [self.repr_type(t) for t in target_types] + target_lookup = dict(zip(target_names, target_types)) + + self.compare_or_generate_golden_matrix( + row_names=list(sources.keys()), + col_names=target_names, + compute_cell=lambda src, tgt: self._from_pandas_cell( + sources[src], type=target_lookup[tgt], safe=safe + ), + golden_file_prefix=golden_file_prefix, + overrides=overrides, + ) + + def test_from_pandas_type_scalar_safe(self): + """Test pa.Array.from_pandas(type=<scalar>, safe=True) against golden file.""" + # pyarrow < 19 ignores the requested type on the protocol path and returns the stored + # type (the SPARK-46776 followup), for every target. Other sources are version-stable. + overrides: dict[tuple[str, str], str] = {} + if LooseVersion(pa.__version__) < LooseVersion("19.0.0"): + for col in [self.repr_type(t) for t in self._get_target_types()]: + overrides[("string[pyarrow]:standard", col)] = "[hello, world]@large_string" + overrides[("large_binary[pyarrow]:standard", col)] = ( + "[b'hello', b'world']@large_binary" + ) + self._compare_type_matrix( + safe=True, + golden_file_prefix="golden_pyarrow_array_from_pandas_type_scalar_safe", + overrides=overrides, + ) + + def test_from_pandas_type_scalar_unsafe(self): + """Test pa.Array.from_pandas(type=<scalar>, safe=False) against golden file.""" + # Same overrides as the safe method: the protocol drops the type request before any + # cast, so safe=False changes none of these cells (SPARK-46776, pyarrow < 19 followup). + overrides: dict[tuple[str, str], str] = {} + if LooseVersion(pa.__version__) < LooseVersion("19.0.0"): + for col in [self.repr_type(t) for t in self._get_target_types()]: + overrides[("string[pyarrow]:standard", col)] = "[hello, world]@large_string" + overrides[("large_binary[pyarrow]:standard", col)] = ( + "[b'hello', b'world']@large_binary" + ) + # inf -> int is undefined behavior in C++; x86 (the golden's platform) yields the + # "integer indefinite" INT_MIN, while ARM's FCVT saturates to INT_MAX. Linux + # aarch64 and macOS arm64 agree on these two cells (safe=True rejects inf -> int on + # every platform, so only this unsafe method needs the override). + if platform.machine() in ("aarch64", "arm64"): + overrides[("float64:infinity", "int8")] = "[-1, 1]@int8" + overrides[("float64:infinity", "int64")] = "[9223372036854775807, 1]@int64" + self._compare_type_matrix( + safe=False, + golden_file_prefix="golden_pyarrow_array_from_pandas_type_scalar_unsafe", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayFromPandasTypeNestedTests( + test_pyarrow_array_from_pandas_default._PyArrowFromPandasTestBase +): + """ + Tests pa.Array.from_pandas(series, type=..., safe=...) for NESTED target types + (list / map / struct) via golden file comparison; scalar targets are covered by + PyArrowArrayFromPandasTypeScalarTests. + + Targets are the container types ``to_arrow_type`` builds, spelled to match it + (``list<element: ...>``). The behavior recorded is that from_pandas does NOT propagate + safe= into a nested child: an overflowing child value raises for both safe settings, + whereas the same overflow at the top level (the scalar tests) raises only under + safe=True. The safe and unsafe goldens therefore match; both are kept so a future + pyarrow that honors safe= in children fails loudly. safe=True/False are two methods and + two goldens; ``mask`` stays None. + """ + + @staticmethod + def _get_target_types(): + """Nested to_arrow_type targets, each at a clean child type and a narrower one it can + overflow, so the safe= non-propagation is observable against the scalar goldens.""" + return [ + pa.list_(pa.field("element", pa.int64())), + pa.list_(pa.field("element", pa.int8())), + pa.list_(pa.field("element", pa.string())), + pa.map_(pa.string(), pa.int64()), + pa.map_(pa.string(), pa.int8()), + pa.struct([("a", pa.int64()), ("b", pa.string())]), + pa.struct([("a", pa.int8()), ("b", pa.string())]), + ] + + def _type_source_arrays(self): + """ + Nested source Series cherry-picked from the shared inventory: clean list/struct rows, + a child-overflow row per container, and homogeneous dicts for the map targets (a dict + Series infers as struct but converts to map when the type asks for one). + """ + pool = self._numpy_backed_sources() + selected = [ + "list<int64>:standard", + "list<int64>:overflow", + "list<int64>:nullable", + "list<int64>:null-element", + "list<string>:standard", + "struct:standard", + "struct:overflow", + "struct:nullable", + "struct<int64>:standard", + "struct<int64>:overflow", + ] + return {name: pool[name] for name in selected} + + def _compare_type_matrix(self, safe, golden_file_prefix, overrides): + sources = self._type_source_arrays() + target_types = self._get_target_types() + target_names = [self.repr_type(t) for t in target_types] + target_lookup = dict(zip(target_names, target_types)) + + self.compare_or_generate_golden_matrix( + row_names=list(sources.keys()), + col_names=target_names, + compute_cell=lambda src, tgt: self._from_pandas_cell( + sources[src], type=target_lookup[tgt], safe=safe + ), + golden_file_prefix=golden_file_prefix, + overrides=overrides, + ) + + def test_from_pandas_type_nested_safe(self): + """Test pa.Array.from_pandas(type=<nested>, safe=True) against golden file.""" + # No version-conditional cells: the nested sources are explicit list/dict object + # Series (not inference-sensitive), so the matrix holds across the pyarrow/pandas + # sweep. Entries would land here only if a future version changed a cell. + overrides: dict[tuple[str, str], str] = {} + self._compare_type_matrix( + safe=True, + golden_file_prefix="golden_pyarrow_array_from_pandas_type_nested_safe", + overrides=overrides, + ) + + def test_from_pandas_type_nested_unsafe(self): + """Test pa.Array.from_pandas(type=<nested>, safe=False) against golden file.""" + # Same empty overrides as the safe method, and version-stable for the same reason; + # since from_pandas does not propagate safe= into children, this golden also matches + # the safe one. + overrides: dict[tuple[str, str], str] = {} + self._compare_type_matrix( + safe=False, + golden_file_prefix="golden_pyarrow_array_from_pandas_type_nested_unsafe", + overrides=overrides, + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_type_inference.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_type_inference.py index dafb08646abbb..c2134279d0632 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_type_inference.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_array_type_inference.py @@ -97,6 +97,7 @@ def test_nullable_data(self): def test_plain_python_list(self): """Test type inference from Python lists.""" import math + import pyarrow as pa sg = ZoneInfo("Asia/Singapore") @@ -295,6 +296,7 @@ def test_pandas_series_numpy_backed(self): import numpy as np import pandas as pd import pyarrow as pa + from pyspark.loose_version import LooseVersion # pandas >= 3 infers large_string instead of string for object-dtype string Series @@ -365,6 +367,7 @@ def test_pandas_series_nullable_extension(self): import numpy as np import pandas as pd import pyarrow as pa + from pyspark.loose_version import LooseVersion # pandas >= 3 uses pyarrow-backed StringDtype, which infers large_string diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py index e3a9a0c0f5bdf..020a9581d9270 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_default.py @@ -27,6 +27,10 @@ - How null values are handled (NaN, None, NaT, etc.) - Whether values are preserved correctly after conversion +The shared ``_PyArrowToPandasTestBase`` holds the conversion helper and the source-array +inventory, split into type-family group methods that these and the non-default tests +reuse whole or by group. + ## Golden File Cell Format Each cell uses the value@type format: @@ -55,15 +59,15 @@ from decimal import Decimal from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, ) -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_pandas: import pandas as pd @@ -71,23 +75,30 @@ import pyarrow as pa -@unittest.skipIf( - not have_pyarrow or not have_pandas or not have_numpy, - pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, -) -class PyArrowArrayToPandasDefaultTests(GoldenFileTestMixin, unittest.TestCase): +class _PyArrowToPandasTestBase(GoldenFileTestMixin, unittest.TestCase): """ - Tests pa.Array.to_pandas() with default arguments via golden file comparison. + Shared machinery for pa.Array.to_pandas() golden file tests. - Covers all major Arrow types: integers, floats, bool, string, binary, - decimal, date, timestamp, duration, time, null, and nested types. - Each type is tested both without and with null values. + Holds the conversion helper and the source-array inventory, split into type-family + group methods reusable whole (via ``_build_source_arrays``) or one at a time, plus a + ``_chunked_sources`` layout group of ChunkedArrays (chunk count is orthogonal to + type). Defines no ``test_*`` of its own. """ - def _build_source_arrays(self): - """Build an ordered dict of named source PyArrow arrays for testing.""" - import pyarrow as pa - + def _to_pandas_cell(self, arr, **to_pandas_kwargs) -> str: + """ + Convert ``arr`` via ``to_pandas(**to_pandas_kwargs)`` and format the + result as a golden-file cell, returning ``ERR@<ExceptionClass>`` if the + conversion raises. + """ + try: + result = arr.to_pandas(**to_pandas_kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + return self.repr_value(result, max_len=0) + + def _numeric_sources(self): + """Integer, float, and boolean arrays.""" sources = {} # ===================================================================== @@ -136,6 +147,12 @@ def _build_source_arrays(self): sources["bool:nullable"] = pa.array([True, False, None], pa.bool_()) sources["bool:empty"] = pa.array([], pa.bool_()) + return sources + + def _string_binary_sources(self): + """String, binary, and decimal arrays.""" + sources = {} + # ===================================================================== # String types # ===================================================================== @@ -168,6 +185,12 @@ def _build_source_arrays(self): ) sources["decimal128:empty"] = pa.array([], pa.decimal128(5, 2)) + return sources + + def _temporal_sources(self): + """Date, timestamp, duration, and time arrays.""" + sources = {} + # ===================================================================== # Date types # ===================================================================== @@ -226,6 +249,12 @@ def _build_source_arrays(self): sources["time64[ns]:nullable"] = pa.array([t1, None], pa.time64("ns")) sources["time64[ns]:empty"] = pa.array([], pa.time64("ns")) + return sources + + def _nested_sources(self): + """Null, list, struct, map, and their nested combinations.""" + sources = {} + # ===================================================================== # Null type # ===================================================================== @@ -307,6 +336,12 @@ def _build_source_arrays(self): pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())), ) + return sources + + def _dictionary_sources(self): + """Dictionary-encoded arrays.""" + sources = {} + # ===================================================================== # Dictionary type # ===================================================================== @@ -325,6 +360,97 @@ def _build_source_arrays(self): return sources + def _chunked_sources(self): + """ + ChunkedArray inputs covering the chunk-count axis. + + pa.Array and pa.ChunkedArray share one ``to_pandas``, and Spark feeds a + ChunkedArray into it on the Table->columns path + (``python/pyspark/sql/pandas/conversion.py``): ``df.toPandas()`` builds one per + partition (no ``combine_chunks``), the empty-dataset branch uses + ``empty_table()`` (a single empty chunk), and the UDF paths ``combine_chunks()`` + first (one chunk). Chunk count is orthogonal to type, so these rows reuse a few + representative types and vary only how the values are split. + """ + sources = {} + + # ===================================================================== + # Numeric chunk layouts (int64: the type most Spark columns reduce to) + # ===================================================================== + sources["int64:zero-chunk"] = pa.chunked_array([], type=pa.int64()) + sources["int64:empty-chunk"] = pa.chunked_array([pa.array([], pa.int64())]) + sources["int64:single-chunk"] = pa.chunked_array([pa.array([1, 2, 3], pa.int64())]) + sources["int64:multi-chunk"] = pa.chunked_array( + [pa.array([1, 2], pa.int64()), pa.array([3, 4], pa.int64())] + ) + # Null in one chunk, not the other. Reachable via cogrouped applyInPandas, the + # one UDF path that does not combine_chunks() first. + sources["int64:multi-chunk-nullable"] = pa.chunked_array( + [pa.array([1, None], pa.int64()), pa.array([2], pa.int64())] + ) + # Empty chunk between two non-empty ones: concatenation must skip it cleanly. + sources["int64:multi-chunk-with-empty"] = pa.chunked_array( + [pa.array([1, 2], pa.int64()), pa.array([], pa.int64()), pa.array([3], pa.int64())] + ) + sources["float64:multi-chunk"] = pa.chunked_array( + [pa.array([1.5, 2.5], pa.float64()), pa.array([3.5], pa.float64())] + ) + + # ===================================================================== + # Variable-width chunk layouts (offset buffers, not just values) + # ===================================================================== + sources["string:single-chunk"] = pa.chunked_array([pa.array(["a", "b"], pa.string())]) + sources["string:multi-chunk"] = pa.chunked_array( + [pa.array(["a", "b"], pa.string()), pa.array(["c"], pa.string())] + ) + sources["string:multi-chunk-nullable"] = pa.chunked_array( + [pa.array(["a", None], pa.string()), pa.array(["c"], pa.string())] + ) + + # ===================================================================== + # Nested chunk layouts (child arrays split across chunks) + # ===================================================================== + sources["list<int64>:multi-chunk"] = pa.chunked_array( + [pa.array([[1, 2], [3]], pa.list_(pa.int64())), pa.array([[4]], pa.list_(pa.int64()))] + ) + struct_type = pa.struct([("x", pa.int64()), ("y", pa.string())]) + sources["struct:multi-chunk"] = pa.chunked_array( + [ + pa.array([{"x": 1, "y": "a"}], struct_type), + pa.array([{"x": 2, "y": "b"}], struct_type), + ] + ) + + return sources + + def _build_source_arrays(self): + """Build an ordered dict of named source PyArrow arrays for testing.""" + sources = {} + for group in [ + self._numeric_sources(), + self._string_binary_sources(), + self._temporal_sources(), + self._nested_sources(), + self._dictionary_sources(), + self._chunked_sources(), + ]: + sources.update(group) + return sources + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayToPandasDefaultTests(_PyArrowToPandasTestBase): + """ + Tests pa.Array.to_pandas() with default arguments via golden file comparison. + + Covers all major Arrow types: integers, floats, bool, string, binary, + decimal, date, timestamp, duration, time, null, and nested types. + Each type is tested both without and with null values. + """ + def test_to_pandas_default(self): """Test pa.Array.to_pandas() with default arguments against golden file.""" sources = self._build_source_arrays() @@ -340,6 +466,9 @@ def test_to_pandas_default(self): ("string:nullable", "pandas series"): ("['hello', nan, 'world']@Series[str]"), ("large_string:standard", "pandas series"): ("['hello', 'world']@Series[str]"), ("large_string:nullable", "pandas series"): "['hello', nan]@Series[str]", + ("string:single-chunk", "pandas series"): "['a', 'b']@Series[str]", + ("string:multi-chunk", "pandas series"): "['a', 'b', 'c']@Series[str]", + ("string:multi-chunk-nullable", "pandas series"): "['a', nan, 'c']@Series[str]", } ) # PyArrow 24 extends the pandas string dtype conversion to empty arrays. @@ -356,11 +485,7 @@ def compute_cell(row_name, col_name): if col_name == "pyarrow array": return self.repr_value(arr, max_len=0) else: - try: - result = arr.to_pandas() - return self.repr_value(result, max_len=0) - except Exception as e: - return f"ERR@{type(e).__name__}" + return self._to_pandas_cell(arr) self.compare_or_generate_golden_matrix( row_names=row_names, diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py index ddfc8f759e26f..ef35868bba08f 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py @@ -28,12 +28,10 @@ arguments so CI fails loudly if the behavior drifts across pandas/PyArrow/NumPy upgrades. -The shared ``_PyArrowToPandasTestBase`` holds the golden-file matrix driver and -the per-cell conversion helper. Each concrete test class supplies its own source -arrays (the rows relevant to the argument under test) and one test method per -argument combination (each producing its own golden file). New ``to_pandas`` -arguments (e.g. ``zero_copy_only``) can be added as additional classes without -touching the existing ones. +The shared ``_PyArrowToPandasTestBase`` (in ``test_pyarrow_arrow_to_pandas_default.py``) +holds the conversion helper and the source-array inventory. Each class reuses the +inventory whole or by group and adds a test method per argument combination (its own +golden file), so a new argument is a new class rather than an edit to the existing ones. ## Golden File Cell Format @@ -64,40 +62,26 @@ from pyspark.loose_version import LooseVersion from pyspark.testing.utils import ( - have_pyarrow, - have_pandas, have_numpy, - pyarrow_requirement_message, - pandas_requirement_message, + have_pandas, + have_pyarrow, numpy_requirement_message, + pandas_requirement_message, + pyarrow_requirement_message, +) + +# Import the shared base (which defines no test_* methods), not a concrete test class, +# so that unittest does not collect and re-run the default file's tests here. +from pyspark.tests.upstream.pyarrow.test_pyarrow_arrow_to_pandas_default import ( + _PyArrowToPandasTestBase, ) -from pyspark.testing.goldenutils import GoldenFileTestMixin if have_pandas: import pandas as pd if have_pyarrow: import pyarrow as pa - - -class _PyArrowToPandasTestBase(GoldenFileTestMixin, unittest.TestCase): - """ - Shared machinery for pa.Array.to_pandas() golden file tests. - - Concrete subclasses provide their own ``_build_source_arrays`` (the rows) and - one or more ``test_*`` methods that call ``compare_or_generate_golden_matrix``. - This base defines no ``test_*`` methods, so it contributes no tests itself. - """ - - def _to_pandas_cell(self, arr, **to_pandas_kwargs) -> str: - """ - Convert ``arr`` via ``to_pandas(**to_pandas_kwargs)`` and format the - result as a golden-file cell, returning ``ERR@<ExceptionClass>`` if the - conversion raises. - """ - try: - return self.repr_value(arr.to_pandas(**to_pandas_kwargs), max_len=0) - except Exception as e: - return f"ERR@{type(e).__name__}" +if have_numpy: + import numpy as np @unittest.skipIf( @@ -120,41 +104,17 @@ class PyArrowArrayToPandasCoerceTemporalTests(_PyArrowToPandasTestBase): """ def _build_source_arrays(self): - """Build an ordered dict of named source PyArrow arrays for testing.""" - sources = {} + """ + Reuse the base's temporal group (the types this argument targets), then add + coercion-specific overflow rows, chunked timestamps, and non-temporal controls. + """ + sources = self._temporal_sources() - # ===================================================================== - # Timestamp types (the primary target of coerce_temporal_nanoseconds) - # ===================================================================== - dt1 = datetime.datetime(2024, 1, 1, 12, 0, 0) - dt2 = datetime.datetime(2024, 6, 15, 18, 30, 0) - for unit in ["s", "ms", "us", "ns"]: - sources[f"timestamp[{unit}]:standard"] = pa.array([dt1, dt2], pa.timestamp(unit)) - sources[f"timestamp[{unit}]:nullable"] = pa.array([dt1, None], pa.timestamp(unit)) - sources[f"timestamp[{unit}]:empty"] = pa.array([], pa.timestamp(unit)) - # Timestamp with timezone - sources["timestamp[us,tz=UTC]:standard"] = pa.array( - [dt1, dt2], pa.timestamp("us", tz="UTC") - ) - sources["timestamp[us,tz=UTC]:nullable"] = pa.array( - [dt1, None], pa.timestamp("us", tz="UTC") - ) - sources["timestamp[us,tz=UTC]:empty"] = pa.array([], pa.timestamp("us", tz="UTC")) # Overflow: coercion to nanoseconds has a valid range (~1677-2262); a # far-future second-resolution timestamp cannot fit and should error. sources["timestamp[s]:overflow"] = pa.array( [datetime.datetime(2500, 1, 1)], pa.timestamp("s") ) - - # ===================================================================== - # Duration types (also coerced to nanoseconds) - # ===================================================================== - td1 = datetime.timedelta(days=1) - td2 = datetime.timedelta(hours=2, minutes=30) - for unit in ["s", "ms", "us", "ns"]: - sources[f"duration[{unit}]:standard"] = pa.array([td1, td2], pa.duration(unit)) - sources[f"duration[{unit}]:nullable"] = pa.array([td1, None], pa.duration(unit)) - sources[f"duration[{unit}]:empty"] = pa.array([], pa.duration(unit)) # Overflow: a duration beyond ~292 years exceeds the int64 nanosecond # range. Unlike timestamp overflow (which raises), coercing it silently # wraps around to a bogus value; this row pins that behavior. @@ -162,43 +122,19 @@ def _build_source_arrays(self): [datetime.timedelta(days=300 * 365)], pa.duration("s") ) - # ===================================================================== - # Date types. With the default date_as_object=True, pandas yields an - # object-dtype Series of datetime.date, so coerce_temporal_nanoseconds - # has nothing to coerce; the "date_as_object=False" column forces the - # numeric datetime64[ns] path where the argument actually takes effect. - # ===================================================================== - d1 = datetime.date(2024, 1, 1) - d2 = datetime.date(2024, 6, 15) - sources["date32:standard"] = pa.array([d1, d2], pa.date32()) - sources["date32:nullable"] = pa.array([d1, None], pa.date32()) - sources["date32:empty"] = pa.array([], pa.date32()) - sources["date64:standard"] = pa.array([d1, d2], pa.date64()) - sources["date64:nullable"] = pa.array([d1, None], pa.date64()) - sources["date64:empty"] = pa.array([], pa.date64()) - - # ===================================================================== - # Time types (control: pandas yields object-dtype datetime.time, no - # native time-of-day dtype, so coerce_temporal_nanoseconds is a no-op) - # ===================================================================== - t1 = datetime.time(12, 30, 0) - t2 = datetime.time(18, 45, 30) - sources["time32[s]:standard"] = pa.array([t1, t2], pa.time32("s")) - sources["time32[s]:nullable"] = pa.array([t1, None], pa.time32("s")) - sources["time32[s]:empty"] = pa.array([], pa.time32("s")) - sources["time32[ms]:standard"] = pa.array([t1, t2], pa.time32("ms")) - sources["time32[ms]:nullable"] = pa.array([t1, None], pa.time32("ms")) - sources["time32[ms]:empty"] = pa.array([], pa.time32("ms")) - sources["time64[us]:standard"] = pa.array([t1, t2], pa.time64("us")) - sources["time64[us]:nullable"] = pa.array([t1, None], pa.time64("us")) - sources["time64[us]:empty"] = pa.array([], pa.time64("us")) - sources["time64[ns]:standard"] = pa.array([t1, t2], pa.time64("ns")) - sources["time64[ns]:nullable"] = pa.array([t1, None], pa.time64("ns")) - sources["time64[ns]:empty"] = pa.array([], pa.time64("ns")) - - # ===================================================================== - # Non-temporal controls (unaffected by coerce_temporal_nanoseconds) - # ===================================================================== + # Chunked timestamps: the base's _chunked_sources has no temporal rows, so pin + # here that a chunked (non-ns) timestamp coerces to the same datetime64[ns] as + # the equivalent contiguous Array. + dt1 = datetime.datetime(2024, 1, 1, 12, 0, 0) + dt2 = datetime.datetime(2024, 6, 15, 18, 30, 0) + sources["timestamp[us]:single-chunk"] = pa.chunked_array( + [pa.array([dt1, dt2], pa.timestamp("us"))] + ) + sources["timestamp[us]:multi-chunk"] = pa.chunked_array( + [pa.array([dt1], pa.timestamp("us")), pa.array([dt2], pa.timestamp("us"))] + ) + + # Non-temporal controls (unaffected by coerce_temporal_nanoseconds). sources["int64:standard"] = pa.array([0, 1, -1], pa.int64()) sources["int64:nullable"] = pa.array([0, 1, None], pa.int64()) sources["float64:standard"] = pa.array([0.0, 1.5, -1.5], pa.float64()) @@ -258,6 +194,591 @@ def compute_cell(row_name, col_name): ) +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayToPandasZeroCopyTests(_PyArrowToPandasTestBase): + """ + Tests pa.Array.to_pandas(zero_copy_only=True) via golden file comparison. + + PySpark converts Arrow data to (numpy-backed) pandas objects throughout its + conversion layer (``python/pyspark/sql/pandas/conversion.py``); whether a + given Arrow type can make that conversion WITHOUT copying its buffers + directly affects the memory and latency of ``toPandas`` and pandas UDFs. + ``zero_copy_only=True`` makes PyArrow raise ``ArrowInvalid`` instead of + silently copying, so it is the natural probe for "is this conversion + zero-copy?". These tests record, per Arrow type, whether the conversion is + zero-copy so CI fails loudly if that changes across pandas/PyArrow/NumPy + upgrades. + + Three output columns are recorded for each source array: + + - ``zero_copy_only=False``: the default, where PyArrow silently copies when a + view is not possible. This always succeeds, and records the resulting dtype. + - ``zero_copy_only=True``: PyArrow's own verdict -- ``Series[dtype]`` when the + conversion is zero-copy, or ``ERR@ArrowInvalid`` when a copy is required. + - ``verified zero-copy``: an INDEPENDENT check of whether the conversion + actually reused the Arrow buffers, rather than trusting PyArrow's flag. + ``zero-copy`` when every buffer the result needs was borrowed, ``copied`` + when none were, and ``partial-copy`` when only some were (e.g. the values + reused but the offsets reallocated). The last two columns are expected to + agree, but are recorded separately rather than asserted equal, because they + do not always: a tz-aware timestamp reports ``zero_copy_only=True`` while + pandas materializes it into a ``DatetimeTZDtype`` array that shares nothing. + Pinning both as data makes such disagreements visible instead of hiding them + behind an assertion. + + The row set mirrors ``test_pyarrow_arrow_to_pandas_default.py``'s rows exactly + -- every Arrow type it covers, in its standard / nullable / empty variants, plus + the shared ChunkedArray layouts -- so both golden files pin the same types, and + appends the one layout variant that only matters for zero-copy: sliced (offset) + arrays. + + A second test method repeats those rows with ``types_mapper=pd.ArrowDtype``, which + asks pandas to keep pointing at the Arrow buffers instead of materializing them + into NumPy -- avoiding the copy is the point of that backend. Its own golden file + records the result, so the two can be read side by side. PySpark takes this path + in ``ArrowArrayToPandasConversion.convert_numpy`` + (``python/pyspark/sql/conversion.py``). + """ + + def _build_source_arrays(self): + """ + Reuse the base's full inventory, then add the sliced layout variants that only + matter for zero-copy: a slice views a contiguous no-null region, so it stays + zero-copy even though its data starts partway into the parent buffer. The + chunk-count variants live in the base's ``_chunked_sources`` so every golden + records their zero-copy behavior, not just this one. + """ + sources = super()._build_source_arrays() + + sources["int64:sliced"] = pa.array(list(range(10)), pa.int64()).slice(2, 3) + sources["int64:sliced-with-null"] = pa.array([1, 2, None, 4, 5], pa.int64()).slice(1, 3) + + return sources + + @staticmethod + def _arrow_buffers(arrow_obj): + """ + Every non-null buffer backing ``arrow_obj``. + + A ChunkedArray has no buffers of its own -- its data lives in its chunks -- + so it is expanded first. ``None`` entries (e.g. an absent validity bitmap) + are skipped. + """ + if isinstance(arrow_obj, pa.ChunkedArray): + chunks = [arrow_obj.chunk(i) for i in range(arrow_obj.num_chunks)] + else: + chunks = [arrow_obj] + + buffers = [] + for chunk in chunks: + for buffer in chunk.buffers(): + if buffer is not None: + buffers.append(buffer) + return buffers + + @classmethod + def _verify_zero_copy(cls, arr, **to_pandas_kwargs) -> str: + """ + Independently verify whether ``to_pandas`` reused ``arr``'s buffers, + instead of trusting PyArrow's own ``zero_copy_only`` verdict. + + The check inspects whatever storage pandas returned rather than assuming a + numpy-backed Series, so it stays correct as pandas moves more dtypes to + Arrow-backed storage. + + Returns ``"zero-copy"``, ``"partial-copy"``, ``"copied"``, or + ``"ERR@<ExceptionClass>"``. + """ + try: + series = arr.to_pandas(**to_pandas_kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + + backing_array = series.array + + # Arrow-backed result: compare buffer addresses, reading the stored data + # back through the public __arrow_array__ protocol. to_numpy() would + # materialize a copy here and wrongly report no sharing. Keying on the + # protocol rather than on ArrowDtype also covers dtypes that are + # Arrow-backed without being ArrowDtype, such as pandas 3's string. + # + # The result has several buffers (validity, offsets, values) and can borrow + # some while allocating others, so count how many it borrowed rather than + # stopping at the first match. The result's buffers are the denominator + # because the question is what pandas had to allocate. + if hasattr(backing_array, "__arrow_array__"): + source_addresses = {buffer.address for buffer in cls._arrow_buffers(arr)} + stored_buffers = cls._arrow_buffers(pa.array(backing_array)) + borrowed = 0 + for buffer in stored_buffers: + if buffer.address in source_addresses: + borrowed += 1 + if borrowed == len(stored_buffers): + return "zero-copy" + return "partial-copy" if borrowed else "copied" + + # numpy-backed result: np.shares_memory accounts for slice offsets, so it + # is robust where raw address equality is not. + pandas_values = series.to_numpy() + for buffer in cls._arrow_buffers(arr): + if np.shares_memory(np.frombuffer(buffer, dtype=np.uint8), pandas_values): + return "zero-copy" + + # Zero bytes cannot overlap, so an empty result is zero-copy if the numpy + # array still borrows the Arrow array (numpy's .base is its memory owner). + if len(arr) == 0: + owner = pandas_values + while (base := getattr(owner, "base", None)) is not None: + if base is arr: + return "zero-copy" + owner = base + + return "copied" + + # Output column for the default zero_copy_only=False: PyArrow silently copies + # when a view is not possible, so this always succeeds and records the dtype. + COL_ZERO_COPY_OFF = "zero_copy_only=False" + + # Output column for zero_copy_only=True: PyArrow raises ArrowInvalid instead + # of copying, so this records its verdict on whether a view was possible. + COL_ZERO_COPY_ON = "zero_copy_only=True" + + # Output column independently verifying that buffers were actually reused. + COL_VERIFIED = "verified zero-copy" + + def test_to_pandas_zero_copy_only(self): + """Test pa.Array.to_pandas(zero_copy_only=True/False) against golden file.""" + sources = self._build_source_arrays() + row_names = list(sources.keys()) + col_names = [ + "pyarrow array", + self.COL_ZERO_COPY_OFF, + self.COL_ZERO_COPY_ON, + self.COL_VERIFIED, + ] + + # Version-specific expected values go here, keyed by (row, col), when a + # newer pandas/PyArrow/NumPy legitimately changes a cell's output. + overrides: dict[tuple[str, str], str] = {} + # Pandas 3 renders non-empty Arrow string arrays with its dedicated string + # dtype, so the copying conversion reports Series[str] instead of object. + if LooseVersion(pd.__version__) >= LooseVersion("3.0.0"): + # Pandas stores that dtype as large_string, so `string` keeps its values + # but has its 32-bit offsets rebuilt as 64-bit, while `large_string` + # already matches and passes through untouched. + non_empty_strings = [ + ("string:standard", "['hello', 'world', '']@Series[str]", "partial-copy"), + ("string:nullable", "['hello', nan, 'world']@Series[str]", "partial-copy"), + ("large_string:standard", "['hello', 'world']@Series[str]", "zero-copy"), + ("large_string:nullable", "['hello', nan]@Series[str]", "zero-copy"), + ("string:single-chunk", "['a', 'b']@Series[str]", "partial-copy"), + ("string:multi-chunk", "['a', 'b', 'c']@Series[str]", "partial-copy"), + ("string:multi-chunk-nullable", "['a', nan, 'c']@Series[str]", "partial-copy"), + ] + for row, expected, _ in non_empty_strings: + overrides[(row, self.COL_ZERO_COPY_OFF)] = expected + + # Only from PyArrow 24 is that string conversion actually Arrow-backed, + # so zero_copy_only succeeds and the buffers are genuinely reused. On + # PyArrow < 24 it still materializes, so the pandas 2 expectations + # (ERR@ArrowInvalid / copied) remain correct and need no override. + if LooseVersion(pa.__version__) >= LooseVersion("24.0.0"): + for row, expected, verified in non_empty_strings: + overrides[(row, self.COL_ZERO_COPY_ON)] = expected + overrides[(row, self.COL_VERIFIED)] = verified + # Empty arrays also gain the string dtype in PyArrow 24. Offsets are + # the only buffer an empty result has, so rebuilding them shares + # nothing at all -- fully copied rather than partial. + for row, verified in [ + ("string:empty", "copied"), + ("large_string:empty", "zero-copy"), + ]: + overrides[(row, self.COL_ZERO_COPY_OFF)] = "[]@Series[str]" + overrides[(row, self.COL_ZERO_COPY_ON)] = "[]@Series[str]" + overrides[(row, self.COL_VERIFIED)] = verified + + def compute_cell(row_name, col_name): + arr = sources[row_name] + if col_name == "pyarrow array": + return self.repr_value(arr, max_len=0) + elif col_name == self.COL_ZERO_COPY_OFF: + return self._to_pandas_cell(arr, zero_copy_only=False) + elif col_name == self.COL_ZERO_COPY_ON: + return self._to_pandas_cell(arr, zero_copy_only=True) + elif col_name == self.COL_VERIFIED: + return self._verify_zero_copy(arr) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_arrow_to_pandas_zero_copy", + index_name="test case", + overrides=overrides, + ) + + # Output columns for the Arrow-backed conversion. Both zero_copy_only states are + # recorded even though the flag is expected to make no difference here + COL_ARROW_ZERO_COPY_OFF = "types_mapper=pd.ArrowDtype, zero_copy_only=False" + COL_ARROW_ZERO_COPY_ON = "types_mapper=pd.ArrowDtype, zero_copy_only=True" + COL_ARROW_VERIFIED = "verified zero-copy" + + def test_to_pandas_zero_copy_only_arrow_backed(self): + """Test pa.Array.to_pandas(types_mapper=pd.ArrowDtype) against golden file.""" + sources = self._build_source_arrays() + row_names = list(sources.keys()) + col_names = [ + "pyarrow array", + self.COL_ARROW_ZERO_COPY_OFF, + self.COL_ARROW_ZERO_COPY_ON, + self.COL_ARROW_VERIFIED, + ] + + # Version-specific expected values go here, keyed by (row, col), when a + # newer pandas/PyArrow/NumPy legitimately changes a cell's output. + overrides: dict[tuple[str, str], str] = {} + + def compute_cell(row_name, col_name): + arr = sources[row_name] + if col_name == "pyarrow array": + return self.repr_value(arr, max_len=0) + elif col_name == self.COL_ARROW_ZERO_COPY_OFF: + return self._to_pandas_cell(arr, types_mapper=pd.ArrowDtype, zero_copy_only=False) + elif col_name == self.COL_ARROW_ZERO_COPY_ON: + return self._to_pandas_cell(arr, types_mapper=pd.ArrowDtype, zero_copy_only=True) + elif col_name == self.COL_ARROW_VERIFIED: + return self._verify_zero_copy(arr, types_mapper=pd.ArrowDtype) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_arrow_to_pandas_zero_copy_arrow_backed", + index_name="test case", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowArrayToPandasIntegerObjectNullsTests(_PyArrowToPandasTestBase): + """ + Tests pa.Array.to_pandas(integer_object_nulls=True) via golden file comparison. + + numpy integers have no null, so by default PyArrow widens a null-bearing integer + array to ``float64`` with ``NaN`` -- which cannot hold every int64 exactly, so a + large value silently changes. ``integer_object_nulls=True`` keeps ``object`` + dtype (Python ``int`` and ``None``) instead, preserving the values. + + PySpark passes it in ``ArrowArrayToPandasConversion.convert_legacy`` + (``python/pyspark/sql/conversion.py``), bundled with ``date_as_object`` and + ``coerce_temporal_nanoseconds``, then narrows the object Series to a nullable + extension dtype (``Int8Dtype`` .. ``Int64Dtype``) -- the only bridge from Arrow to + those dtypes that avoids ``float64``. + + Three output columns are recorded per source array: the argument off, on, and the + full ``pandas_options`` dict ``convert_legacy`` passes. The last is not a + duplicate of the second -- ``coerce_temporal_nanoseconds`` shifts the temporal + rows to ``ns`` -- and pinning the call as Spark makes it also catches PyArrow + changing the ``date_as_object=True`` default it relies on. + + The row set mirrors ``test_pyarrow_arrow_to_pandas_default.py``'s rows so both + golden files pin the same types, and appends the integer variants those rows do + not reach (see ``_build_source_arrays``). + """ + + def _build_source_arrays(self): + """ + Reuse the base's full inventory, then add the integer variants it does not + reach: its nullable values are small enough to survive ``float64``, and its + nested rows are missing whole sub-lists rather than a single integer inside one. + """ + sources = super()._build_source_arrays() + + # Each width's min and max alongside a null. The shared rows use small values, + # which float64 represents exactly; only at 64 bits does the range exceed its + # 53-bit mantissa and the value itself change. + for pa_type in [ + pa.int8(), + pa.int16(), + pa.int32(), + pa.int64(), + pa.uint8(), + pa.uint16(), + pa.uint32(), + pa.uint64(), + ]: + if pa.types.is_signed_integer(pa_type): + bounds = [2 ** (pa_type.bit_width - 1) - 1, -(2 ** (pa_type.bit_width - 1))] + else: + bounds = [2**pa_type.bit_width - 1, 0] + sources[f"{pa_type}:extremes-nullable"] = pa.array(bounds + [None], pa_type) + + # Every value is null, so there is no integer left to convert. + sources["int64:all-null"] = pa.array([None, None], pa.int64()) + + # Nested types whose null is an integer ELEMENT, not a missing sub-list: the + # shared rows only cover the latter, which this argument does not affect. + # These are also the types convert_legacy still serves. + sources["list<int64>:null-element"] = pa.array([[1, None], [2, 3]], pa.list_(pa.int64())) + sources["list<int64>:null-element-extreme"] = pa.array( + [[2**63 - 1, None]], pa.list_(pa.int64()) + ) + sources["large_list<int64>:null-element"] = pa.array([[1, None]], pa.large_list(pa.int64())) + sources["fixed_size_list<int64>[3]:null-element"] = pa.array( + [[1, None, 3]], pa.list_(pa.int64(), 3) + ) + sources["list<list<int64>>:null-element"] = pa.array( + [[[1, None], [2]]], pa.list_(pa.list_(pa.int64())) + ) + sources["struct:null-int-field"] = pa.array( + [{"x": 1, "y": "a"}, {"x": None, "y": "b"}], + pa.struct([("x", pa.int64()), ("y", pa.string())]), + ) + sources["map<string,int64>:null-value"] = pa.array( + [[("a", 1), ("b", None)]], pa.map_(pa.string(), pa.int64()) + ) + # Control: dictionary encoding stores the distinct values once plus an index per + # row, so the null lives in the indices and the int64 values hold none. With no + # null to represent there, the argument has nothing to decide -- both columns + # stay category, whose own values also remain int64. + sources["dictionary<int64>:nullable"] = pa.array( + [1, None, 1], pa.int64() + ).dictionary_encode() + + return sources + + # Output column for the default: a null-bearing integer array widens to float64. + COL_INTEGER_OBJECT_NULLS_OFF = "integer_object_nulls=False" + # Output column for the argument on: the result stays object dtype. + COL_INTEGER_OBJECT_NULLS_ON = "integer_object_nulls=True" + # Output column for all three arguments as convert_legacy passes them together. + COL_SPARK_PANDAS_OPTIONS = "spark pandas_options" + + # Kept as one dict so the column cannot drift from the call site it mirrors. + SPARK_PANDAS_OPTIONS = { + "date_as_object": True, + "coerce_temporal_nanoseconds": True, + "integer_object_nulls": True, + } + + def test_to_pandas_integer_object_nulls(self): + """Test pa.Array.to_pandas(integer_object_nulls=True/False) against golden file.""" + sources = self._build_source_arrays() + row_names = list(sources.keys()) + col_names = [ + "pyarrow array", + self.COL_INTEGER_OBJECT_NULLS_OFF, + self.COL_INTEGER_OBJECT_NULLS_ON, + self.COL_SPARK_PANDAS_OPTIONS, + ] + + # Version-specific expected values go here, keyed by (row, col), when a newer + # pandas/PyArrow/NumPy legitimately changes a cell's output. This argument does + # not touch strings, so a string row shifts in every output column at once. + overrides: dict[tuple[str, str], str] = {} + + def override_outputs(row: str, expected: str) -> None: + for col in col_names[1:]: # every column but the "pyarrow array" input + overrides[(row, col)] = expected + + pandas_3_or_later = LooseVersion(pd.__version__) >= LooseVersion("3.0.0") + pyarrow_24_or_later = LooseVersion(pa.__version__) >= LooseVersion("24.0.0") + + # Pandas 3 renders Arrow string arrays with its dedicated string dtype. + if pandas_3_or_later: + override_outputs("string:standard", "['hello', 'world', '']@Series[str]") + override_outputs("string:nullable", "['hello', nan, 'world']@Series[str]") + override_outputs("large_string:standard", "['hello', 'world']@Series[str]") + override_outputs("large_string:nullable", "['hello', nan]@Series[str]") + override_outputs("string:single-chunk", "['a', 'b']@Series[str]") + override_outputs("string:multi-chunk", "['a', 'b', 'c']@Series[str]") + override_outputs("string:multi-chunk-nullable", "['a', nan, 'c']@Series[str]") + + # Empty ones stay object until PyArrow 24, so the baseline holds before that. + # Spark supports PyArrow 18+, so both branches are reachable. + if pandas_3_or_later and pyarrow_24_or_later: + override_outputs("string:empty", "[]@Series[str]") + override_outputs("large_string:empty", "[]@Series[str]") + + def compute_cell(row_name, col_name): + arr = sources[row_name] + if col_name == "pyarrow array": + return self.repr_value(arr, max_len=0) + elif col_name == self.COL_INTEGER_OBJECT_NULLS_OFF: + return self._to_pandas_cell(arr, integer_object_nulls=False) + elif col_name == self.COL_INTEGER_OBJECT_NULLS_ON: + return self._to_pandas_cell(arr, integer_object_nulls=True) + elif col_name == self.COL_SPARK_PANDAS_OPTIONS: + return self._to_pandas_cell(arr, **self.SPARK_PANDAS_OPTIONS) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_arrow_to_pandas_integer_object_nulls", + index_name="test case", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas or not have_numpy, + pyarrow_requirement_message or pandas_requirement_message or numpy_requirement_message, +) +class PyArrowChunkedArrayToPandasMemoryFlagsTests(_PyArrowToPandasTestBase): + """ + Tests pa.ChunkedArray.to_pandas() under the memory-tuning arguments + ``self_destruct`` / ``split_blocks`` / ``use_threads`` via golden file comparison. + + Spark sets all three together for ``df.toPandas()`` when Arrow self-destruct is + enabled (``python/pyspark/sql/pandas/conversion.py``), freeing each column's buffers + as it is converted to keep peak memory near one column. They form one unit -- + freeing (``self_destruct``) needs per-column blocks (``split_blocks``) and + single-threaded conversion (``use_threads=False``) -- so this golden records the + bundle (``spark memory options``) against ``default`` rather than a column per flag. + The path always converts a ``pa.ChunkedArray`` (``Table.column(i)``), hence the rows. + + The arguments tune HOW the conversion runs, not WHAT it produces, so the bundle is + expected to match the default. Pinning that equivalence catches drift: if these + arguments ever alter the output -- or ``self_destruct`` actually consumes the input, + which today it does not at this level -- a cell moves. A final column records whether + the source is still readable after ``self_destruct=True``, since Spark treats freeing + as an optional optimization, not a contract. Each cell rebuilds its source from a + factory so a destructive ``self_destruct`` cannot corrupt another cell. + """ + + def _memory_flag_source_factories(self): + """ + Named factories, each returning a FRESH ChunkedArray per call, since + ``self_destruct=True`` may consume its input. Rows span the chunk-count axis + across numeric, variable-width, and nested types (buffer layout differs by type). + """ + struct_type = pa.struct([("x", pa.int64()), ("y", pa.string())]) + return { + "int64:single-chunk": lambda: pa.chunked_array([pa.array([1, 2, 3], pa.int64())]), + "int64:multi-chunk": lambda: pa.chunked_array( + [pa.array([1, 2], pa.int64()), pa.array([3, 4], pa.int64())] + ), + "int64:multi-chunk-nullable": lambda: pa.chunked_array( + [pa.array([1, None], pa.int64()), pa.array([2], pa.int64())] + ), + "int64:multi-chunk-with-empty": lambda: pa.chunked_array( + [pa.array([1, 2], pa.int64()), pa.array([], pa.int64()), pa.array([3], pa.int64())] + ), + "int64:empty-chunk": lambda: pa.chunked_array([pa.array([], pa.int64())]), + "float64:multi-chunk": lambda: pa.chunked_array( + [pa.array([1.5, 2.5], pa.float64()), pa.array([3.5], pa.float64())] + ), + "string:multi-chunk": lambda: pa.chunked_array( + [pa.array(["a", "b"], pa.string()), pa.array(["c"], pa.string())] + ), + "string:multi-chunk-nullable": lambda: pa.chunked_array( + [pa.array(["a", None], pa.string()), pa.array(["c"], pa.string())] + ), + "list<int64>:multi-chunk": lambda: pa.chunked_array( + [ + pa.array([[1, 2], [3]], pa.list_(pa.int64())), + pa.array([[4]], pa.list_(pa.int64())), + ] + ), + "struct:multi-chunk": lambda: pa.chunked_array( + [ + pa.array([{"x": 1, "y": "a"}], struct_type), + pa.array([{"x": 2, "y": "b"}], struct_type), + ] + ), + } + + # Default (no memory arguments) -- the baseline the bundle should match. + COL_DEFAULT = "default" + # All three flags as convert_arrow_table_to_pandas passes them together. + COL_SPARK_MEMORY_OPTIONS = "spark memory options" + # Whether the source ChunkedArray is still readable after self_destruct=True. + COL_SOURCE_READABLE = "source readable after self_destruct" + + # Kept as one dict so the column cannot drift from the call site it mirrors + # (python/pyspark/sql/pandas/conversion.py). + SPARK_MEMORY_OPTIONS = { + "self_destruct": True, + "split_blocks": True, + "use_threads": False, + } + + def _source_readable_after_self_destruct(self, factory) -> str: + """ + Convert a fresh source with ``self_destruct=True``, then report whether it is + still readable: ``readable``, or ``unreadable@<ExceptionClass>`` if freed. + """ + arr = factory() + try: + arr.to_pandas(self_destruct=True) + except Exception as e: + return f"ERR@{type(e).__name__}" + try: + arr.to_pylist() + return "readable" + except Exception as e: + return f"unreadable@{type(e).__name__}" + + def test_to_pandas_memory_flags(self): + """Test pa.ChunkedArray.to_pandas() under the memory-tuning arguments.""" + factories = self._memory_flag_source_factories() + row_names = list(factories.keys()) + col_names = [ + "pyarrow array", + self.COL_DEFAULT, + self.COL_SPARK_MEMORY_OPTIONS, + self.COL_SOURCE_READABLE, + ] + + # Version-specific expected values go here, keyed by (row, col), when a newer + # pandas/PyArrow/NumPy legitimately changes a cell's output. These arguments do + # not touch strings, so a string row shifts in both value columns at once. + overrides: dict[tuple[str, str], str] = {} + if LooseVersion(pd.__version__) >= LooseVersion("3.0.0"): + value_cols = [self.COL_DEFAULT, self.COL_SPARK_MEMORY_OPTIONS] + for row, expected in [ + ("string:multi-chunk", "['a', 'b', 'c']@Series[str]"), + ("string:multi-chunk-nullable", "['a', nan, 'c']@Series[str]"), + ]: + for col in value_cols: + overrides[(row, col)] = expected + + def compute_cell(row_name, col_name): + factory = factories[row_name] + if col_name == "pyarrow array": + return self.repr_value(factory(), max_len=0) + elif col_name == self.COL_DEFAULT: + return self._to_pandas_cell(factory()) + elif col_name == self.COL_SPARK_MEMORY_OPTIONS: + return self._to_pandas_cell(factory(), **self.SPARK_MEMORY_OPTIONS) + elif col_name == self.COL_SOURCE_READABLE: + return self._source_readable_after_self_destruct(factory) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_chunked_array_to_pandas_memory_flags", + index_name="test case", + overrides=overrides, + ) + + if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py new file mode 100644 index 0000000000000..72c5fdcc097f9 --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_dataframe_from_pandas.py @@ -0,0 +1,253 @@ +# +# 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. +# + +""" +Golden-file tests for the PyArrow ``from_pandas`` constructors that take a whole pandas +DataFrame: ``pa.RecordBatch.from_pandas`` and ``pa.Schema.from_pandas`` (with +``pa.Table.from_pandas`` to follow in this file). These take a DataFrame, unlike +``pa.Array.from_pandas`` which takes a Series (covered by test_pyarrow_array_from_pandas_*). + +Per-column type inference matches the Array tests, so these pin the DataFrame-level +behavior instead: whole-frame assembly, the pandas index under ``preserve_index``, and -- +for RecordBatch -- num_rows preservation for a 0-column DataFrame. Spark calls +``RecordBatch.from_pandas`` bare at pandas/conversion.py:1026 and connect/session.py:632 +(the createDataFrame 0-column branch) and stateful_processor_api_client.py:557, relying on +the default ``preserve_index=None`` to carry num_rows via the index metadata -- otherwise a +0-column relation loses its rows. + +``Schema.from_pandas`` is inspected to build a Spark schema, and the two prod call sites +diverge on ``preserve_index``: classic pandas/conversion.py:971 passes ``False`` (index +dropped), Connect session.py:573 passes it bare/``None`` (a named or non-range index becomes +an extra field) -- so a named-index frame yields different field sets. Spark reads each +field's type AND nullability (conversion.py:989 / session.py:590), so the schema test pins +name/type/nullability across ``preserve_index``. + +Regenerate with SPARK_GENERATE_GOLDEN_FILES=1. +""" + +import datetime +import unittest + +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.utils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) + +if have_pandas: + import pandas as pd +if have_pyarrow: + import pyarrow as pa + + +class _PyArrowFromPandasFrameTestBase(GoldenFileTestMixin, unittest.TestCase): + """ + Shared machinery for the DataFrame-input ``from_pandas`` constructors (RecordBatch and + Schema here; Table as a followup). Owns the source-frame inventory and the index-aware + input-cell rendering (both constructors depend on the index under ``preserve_index``); + defines no ``test_*`` of its own. + """ + + @staticmethod + def _index_desc(index) -> str: + """Compact, deterministic description of a pandas index for the input cell.""" + if isinstance(index, pd.MultiIndex): + return f"MultiIndex[names={list(index.names)}]" + if isinstance(index, pd.RangeIndex): + return f"RangeIndex[{index.start}:{index.stop}:{index.step}]" + return f"{index.name!r}:{index.tolist()}" + + def _input_cell(self, df) -> str: + """Input DataFrame repr, extended with its index (repr_value drops it).""" + return f"{self.repr_value(df, max_len=0)}[index={self._index_desc(df.index)}]" + + def _build_source_frames(self): + """Named pandas DataFrames covering shape x index-kind, plus a dtype sample.""" + dt = datetime.datetime(2020, 1, 1, 5, 30) + named = pd.Index([100, 200, 300], name="idx") + unnamed = pd.Index([10, 20, 30]) + frames = {} + + # ===================================================================== + # 0-column frames -- only the index carries the row count + # ===================================================================== + frames["0-columns:range-index"] = pd.DataFrame(index=range(3)) + frames["0-columns:named-index"] = pd.DataFrame(index=named) + frames["0-columns:unnamed-index"] = pd.DataFrame(index=unnamed) + frames["0-columns:empty"] = pd.DataFrame(index=range(0)) + + # ===================================================================== + # Single column -- a non-RangeIndex becomes an extra column + # ===================================================================== + frames["single-column:range-index"] = pd.DataFrame({"a": [1, 2, 3]}) + frames["single-column:named-index"] = pd.DataFrame({"a": [1, 2, 3]}, index=named) + frames["single-column:unnamed-index"] = pd.DataFrame({"a": [1, 2, 3]}, index=unnamed) + + # ===================================================================== + # Multi-column assembly. Drift-prone columns (object strings, datetime64[ns]) are + # pinned so the Arrow output is stable across pandas 2/3; per-dtype inference itself + # is already covered by test_pyarrow_array_from_pandas_default. + # ===================================================================== + frames["multi-column:standard"] = pd.DataFrame( + { + "i": pd.Series([1, 2, 3], dtype="int64"), + "f": pd.Series([1.5, 2.5, 3.5], dtype="float64"), + "b": pd.Series([True, False, True], dtype=bool), + "s": pd.Series(["a", "b", "c"], dtype=object), + "t": pd.Series([dt, dt, dt], dtype="datetime64[ns]"), + } + ) + frames["multi-column:nullable"] = pd.DataFrame( + { + "f": pd.Series([1.5, None, 3.5], dtype="float64"), + "b": pd.Series([True, None, False], dtype=object), + "s": pd.Series(["a", None, "c"], dtype=object), + "t": pd.Series([dt, None, dt], dtype="datetime64[ns]"), + } + ) + # Multiple columns but zero rows (an empty object column is omitted -- it would + # infer to Arrow ``null`` rather than a concrete type). + frames["multi-column:no-rows"] = pd.DataFrame( + { + "i": pd.Series([], dtype="int64"), + "f": pd.Series([], dtype="float64"), + "b": pd.Series([], dtype=bool), + "t": pd.Series([], dtype="datetime64[ns]"), + } + ) + return frames + + +@unittest.skipIf( + not have_pyarrow or not have_pandas, + pyarrow_requirement_message or pandas_requirement_message, +) +class PyArrowRecordBatchFromPandasTests(_PyArrowFromPandasFrameTestBase): + """Tests pa.RecordBatch.from_pandas() across preserve_index via golden file comparison.""" + + def _from_pandas_cell(self, df, **kwargs) -> str: + """ + Convert ``df`` via RecordBatch.from_pandas(**kwargs) and append num_rows -- the + property this test pins, which a 0-column batch has no column to imply. Returns + ERR@<ExceptionClass> if the conversion raises; a formatting error is a test bug. + """ + try: + batch = pa.RecordBatch.from_pandas(df, **kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + return f"{self.repr_value(batch, max_len=0)}[num_rows={batch.num_rows}]" + + def test_from_pandas(self): + """Test pa.RecordBatch.from_pandas() across preserve_index against golden file.""" + sources = self._build_source_frames() + row_names = list(sources.keys()) + preserve = { + "preserve_index=None": None, + "preserve_index=False": False, + "preserve_index=True": True, + } + col_names = ["pandas dataframe", *preserve.keys()] + + # Version-specific expected values go here, keyed by (row, col), for known drift. + overrides: dict[tuple[str, str], str] = {} + + def compute_cell(row_name, col_name): + df = sources[row_name] + if col_name == "pandas dataframe": + return self._input_cell(df) + return self._from_pandas_cell(df, preserve_index=preserve[col_name]) + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_record_batch_from_pandas", + index_name="test case", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas, + pyarrow_requirement_message or pandas_requirement_message, +) +class PyArrowSchemaFromPandasTests(_PyArrowFromPandasFrameTestBase): + """Tests pa.Schema.from_pandas() across preserve_index via golden file comparison.""" + + def _schema_source_frames(self): + """Shared frames plus two MultiIndex rows -- a MultiIndex has several index levels, + each becoming its own field, so these pin multi-level index-to-field naming at the + schema. Level values are integers (stable int64 on pandas 2 and 3; strings drift).""" + frames = self._build_source_frames() + frames["single-column:multiindex"] = pd.DataFrame( + {"a": [1, 2, 3]}, + index=pd.MultiIndex.from_tuples([(1, 10), (1, 20), (2, 30)], names=["g", "n"]), + ) + frames["single-column:multiindex-partial-name"] = pd.DataFrame( + {"a": [1, 2]}, + index=pd.MultiIndex.from_tuples([(1, 10), (2, 20)], names=["g", None]), + ) + return frames + + def _from_pandas_cell(self, df, **kwargs) -> str: + """ + Infer the schema via Schema.from_pandas(**kwargs) and render its fields with + nullability -- the name/type/nullable Spark reads to build its StructType. Returns + ERR@<ExceptionClass> if inference raises; a formatting error is a test bug. + """ + try: + schema = pa.Schema.from_pandas(df, **kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + return self.repr_value(schema, max_len=0) + + def test_from_pandas(self): + """Test pa.Schema.from_pandas() across preserve_index against golden file.""" + sources = self._schema_source_frames() + row_names = list(sources.keys()) + preserve = { + "preserve_index=None": None, + "preserve_index=False": False, + "preserve_index=True": True, + } + col_names = ["pandas dataframe", *preserve.keys()] + + # Version-specific expected values go here, keyed by (row, col), for known drift. + overrides: dict[tuple[str, str], str] = {} + + def compute_cell(row_name, col_name): + df = sources[row_name] + if col_name == "pandas dataframe": + return self._input_cell(df) + return self._from_pandas_cell(df, preserve_index=preserve[col_name]) + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_schema_from_pandas", + index_name="test case", + overrides=overrides, + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_ignore_timezone.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_ignore_timezone.py index 08ed807e7dcd4..aa1b1a6002ceb 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_ignore_timezone.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_ignore_timezone.py @@ -15,10 +15,10 @@ # limitations under the License. # -import os import datetime -from zoneinfo import ZoneInfo +import os import unittest +from zoneinfo import ZoneInfo from pyspark.testing.utils import ( have_pandas, @@ -100,8 +100,8 @@ def test_timezone_with_python(self): @unittest.skipIf(not have_pandas, pandas_requirement_message) def test_timezone_with_pandas(self): - import pyarrow as pa import pandas as pd + import pyarrow as pa tz = "Asia/Singapore" ts1 = pd.Timestamp(2022, 1, 5, 15, 0, 1, tzinfo=ZoneInfo(tz)) diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_cast.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_cast.py new file mode 100644 index 0000000000000..01db33a1a41ca --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_cast.py @@ -0,0 +1,296 @@ +# +# 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. +# + +""" +Tests for PyArrow's pa.Table.cast() method using golden file comparison. + +Unlike pa.Array.cast() (covered by test_pyarrow_array_cast.py), Table.cast() casts a +whole table to a target *schema*. Per-column type conversion matches the Array tests, +so this file pins the genuinely Table-level behavior instead: multi-column casts, strict +field-name/order matching, target-field nullability enforcement, temporal coercion, and +the empty-table edges. Spark calls Table.cast() at +python/pyspark/sql/pandas/conversion.py:499 and :1108 (classic toArrow / createDataFrame) +and, in Spark Connect, at python/pyspark/sql/connect/dataframe.py:1995 and +python/pyspark/sql/connect/session.py:647,663. + +This suite covers both safe=True (default) and safe=False modes: +- safe=True: checks for overflow/truncation, raising on unsafe conversions. +- safe=False: allows unsafe conversions (overflow wrapping, truncation). +Each mode has its own golden file. Field-name/order and nullability errors raise in both +modes -- they are checked before the per-column value cast. + +## Golden File Cell Format + +Each cell uses the value@type format: +- pyarrow Table: "{col: [val1, val2, None], ...}@Table[name: type, ...]" +- Error: "ERR@ExceptionClassName" + +## Regenerating Golden Files + +Set SPARK_GENERATE_GOLDEN_FILES=1 before running: + + SPARK_GENERATE_GOLDEN_FILES=1 python -m pytest \\ + python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_cast.py + +If package tabulate (https://pypi.org/project/tabulate/) is installed, +it will also regenerate the Markdown files. + +## PyArrow Version Compatibility + +The golden files capture behavior for a specific PyArrow version. Regenerate when +upgrading PyArrow, as cast support may change between versions. Table.cast() is +pandas-free (PyArrow in, PyArrow out), so no pandas-version differences apply. +""" + +import unittest + +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.utils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) + +if have_pyarrow: + import pyarrow as pa + + +class _PyArrowTableCastTestBase(GoldenFileTestMixin, unittest.TestCase): + """Base class for pa.Table.cast() golden file tests. Defines no test_* of its own.""" + + def _try_cast(self, table, target_schema, safe=True) -> str: + """ + Cast ``table`` to ``target_schema`` and format the result as a golden cell, + returning ``ERR@<ExceptionClass>`` if the cast raises. Only the cast is guarded: + a formatting error is a test bug, not a cast signal, so it propagates. + """ + try: + result = table.cast(target_schema, safe=safe) + except Exception as e: + return f"ERR@{type(e).__name__}" + return self.repr_value(result, max_len=0) + + def _cast_scenarios(self): + """ + Ordered {name: (source_table, target_schema)} pairs, each isolating one + Table.cast contract. Shared by the safe and unsafe test methods. + """ + scenarios = {} + + # ===================================================================== + # Multi-column type cast (whole-schema assembly) + # ===================================================================== + scenarios["types:downcast"] = ( + pa.table( + { + "a": pa.array([1, 2, 3], pa.int64()), + "b": pa.array([1.5, 2.5, 3.5], pa.float64()), + } + ), + pa.schema([("a", pa.int32()), ("b", pa.float32())]), + ) + scenarios["types:upcast"] = ( + pa.table( + { + "a": pa.array([1, 2, 3], pa.int32()), + "b": pa.array([1.5, 2.5, 3.5], pa.float32()), + } + ), + pa.schema([("a", pa.int64()), ("b", pa.float64())]), + ) + + # ===================================================================== + # safe axis: these flip between the safe and unsafe goldens + # ===================================================================== + scenarios["overflow:int64->int32"] = ( + pa.table({"a": pa.array([2**40, 1], pa.int64())}), + pa.schema([("a", pa.int32())]), + ) + scenarios["truncate:float->int"] = ( + pa.table({"a": pa.array([1.9, -2.1], pa.float64())}), + pa.schema([("a", pa.int64())]), + ) + + # ===================================================================== + # Columns are matched positionally + name-equal (no match/reorder by name), so + # a name mismatch, a pure reorder, and a wrong field count all raise ValueError. + # ===================================================================== + base_ab = pa.table( + {"a": pa.array([1, 2], pa.int64()), "b": pa.array([1.5, 2.5], pa.float64())} + ) + scenarios["names:mismatch"] = ( + base_ab, + pa.schema([("x", pa.int32()), ("b", pa.float32())]), + ) + scenarios["names:reordered"] = ( + base_ab, + pa.schema([("b", pa.float64()), ("a", pa.int64())]), + ) + scenarios["names:field-count"] = ( + base_ab, + pa.schema([("a", pa.int32())]), + ) + + # ===================================================================== + # Target-field nullability enforcement + # ===================================================================== + scenarios["nullable:false-with-nulls"] = ( + pa.table({"a": pa.array([1, None], pa.int64())}), + pa.schema([pa.field("a", pa.int32(), nullable=False)]), + ) + scenarios["nullable:false-no-nulls"] = ( + pa.table({"a": pa.array([1, 2], pa.int64())}), + pa.schema([pa.field("a", pa.int32(), nullable=False)]), + ) + + # ===================================================================== + # Temporal unit / timezone coercion + # ===================================================================== + ts_us = pa.table({"ts": pa.array([0, 1_000_000], pa.timestamp("us"))}) + scenarios["timestamp:us->ns"] = ( + ts_us, + pa.schema([("ts", pa.timestamp("ns"))]), + ) + scenarios["timestamp:attach-tz"] = ( + ts_us, + pa.schema([("ts", pa.timestamp("us", "UTC"))]), + ) + + # ===================================================================== + # Variable-width widening + # ===================================================================== + scenarios["string->large_string"] = ( + pa.table({"s": pa.array(["hello", "world", None], pa.string())}), + pa.schema([("s", pa.large_string())]), + ) + scenarios["binary->large_binary"] = ( + pa.table({"b": pa.array([b"x", b"yz", None], pa.binary())}), + pa.schema([("b", pa.large_binary())]), + ) + + # ===================================================================== + # Nested column types: cast recurses into the container and casts each inner + # element, carrying safe= down (see nested:list-overflow). Kept name- and + # order-matched so these stay clean successes on every PyArrow version. + # ===================================================================== + scenarios["nested:list"] = ( + pa.table({"lst": pa.array([[1, 2], [3], None], pa.list_(pa.int64()))}), + pa.schema([("lst", pa.list_(pa.int32()))]), + ) + scenarios["nested:list-overflow"] = ( + pa.table({"lst": pa.array([[2**40, 1]], pa.list_(pa.int64()))}), + pa.schema([("lst", pa.list_(pa.int32()))]), + ) + scenarios["nested:struct"] = ( + pa.table( + { + "st": pa.array( + [{"x": 1, "y": "a"}, None], + pa.struct([("x", pa.int64()), ("y", pa.string())]), + ) + } + ), + pa.schema([("st", pa.struct([("x", pa.int32()), ("y", pa.large_string())]))]), + ) + scenarios["nested:map"] = ( + pa.table( + {"m": pa.array([[("k", 1), ("j", 2)], None], pa.map_(pa.string(), pa.int64()))} + ), + pa.schema([("m", pa.map_(pa.string(), pa.int32()))]), + ) + + # ===================================================================== + # Multi-chunk column: exercises pa.ChunkedArray.cast under Table.cast + # ===================================================================== + scenarios["multi-chunk-column"] = ( + pa.table({"a": pa.chunked_array([[1, 2], [3, None]], pa.int64())}), + pa.schema([("a", pa.int32())]), + ) + + # ===================================================================== + # Empty edges + # ===================================================================== + scenarios["empty:0-columns"] = (pa.table({}), pa.schema([])) + scenarios["empty:columns-no-rows"] = ( + pa.table({"i": pa.array([], pa.int64()), "s": pa.array([], pa.string())}), + pa.schema([("i", pa.int32()), ("s", pa.large_string())]), + ) + + return scenarios + + +@unittest.skipIf( + not have_pyarrow or not have_pandas, + pyarrow_requirement_message or pandas_requirement_message, +) +class PyArrowTableCastTests(_PyArrowTableCastTestBase): + """ + Tests pa.Table.cast(target_schema) with safe=True and safe=False via golden files. + + Pins Table-level cast behavior distinct from pa.Array.cast: whole-schema casts, + strict field-name/order matching, target-field nullability enforcement, temporal + coercion, and the empty-table edges. + """ + + def _run(self, safe, golden_file_prefix, overrides): + scenarios = self._cast_scenarios() + row_names = list(scenarios.keys()) + col_names = ["pyarrow table", "cast result"] + + def compute_cell(row_name, col_name): + source_table, target_schema = scenarios[row_name] + if col_name == "pyarrow table": + return self.repr_value(source_table, max_len=0) + elif col_name == "cast result": + return self._try_cast(source_table, target_schema, safe=safe) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix=golden_file_prefix, + index_name="test case", + overrides=overrides, + ) + + def test_table_cast_matrix(self): + """Test pa.Table.cast(target_schema) with safe=True (default).""" + # PyArrow-version-specific expected cells; empty at the pa24/pd2 baseline. + overrides: dict[tuple[str, str], str] = {} + self._run( + safe=True, + golden_file_prefix="golden_pyarrow_table_cast_safe", + overrides=overrides, + ) + + def test_table_cast_matrix_unsafe(self): + """Test pa.Table.cast(target_schema) with safe=False.""" + overrides: dict[tuple[str, str], str] = {} + self._run( + safe=False, + golden_file_prefix="golden_pyarrow_table_cast_unsafe", + overrides=overrides, + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_to_pandas.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_to_pandas.py new file mode 100644 index 0000000000000..f902281eee3eb --- /dev/null +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_to_pandas.py @@ -0,0 +1,325 @@ +# +# 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. +# + +""" +Tests for PyArrow Table.to_pandas() using golden file comparison. + +Unlike Array/ChunkedArray.to_pandas() (which returns a Series and is covered by +test_pyarrow_arrow_to_pandas_{default,non_default}.py), Table.to_pandas() returns a +DataFrame. Its per-column conversion matches the Array tests, so this file pins the +genuinely Table-specific behavior instead: multi-column DataFrame assembly and the +empty-table edges (0 columns / 0 rows). Spark calls Table.to_pandas() at +python/pyspark/sql/pandas/conversion.py:255 (the 0-column path) and, in Spark Connect, +at python/pyspark/sql/connect/client/core.py:1423 (a bare whole-Table conversion). + +## Golden File Cell Format + +Each cell uses the value@type format: +- pyarrow Table: "{col: [val1, val2, None], ...}@Table[name: type, ...]" +- pandas DataFrame: "{col: [values], ...}@Dataframe[name dtype, ...]" +- Error: "ERR@ExceptionClassName" + +Values are formatted via tolist() for stable, Python-native representation. + +## Regenerating Golden Files + +Set SPARK_GENERATE_GOLDEN_FILES=1 before running: + + SPARK_GENERATE_GOLDEN_FILES=1 python -m pytest \\ + python/pyspark/tests/upstream/pyarrow/test_pyarrow_table_to_pandas.py +""" + +import datetime +import unittest + +from pyspark.loose_version import LooseVersion +from pyspark.testing.goldenutils import GoldenFileTestMixin +from pyspark.testing.utils import ( + have_pandas, + have_pyarrow, + pandas_requirement_message, + pyarrow_requirement_message, +) + +if have_pandas: + import pandas as pd +if have_pyarrow: + import pyarrow as pa + + +class _PyArrowTableToPandasTestBase(GoldenFileTestMixin, unittest.TestCase): + """ + Shared machinery for pa.Table.to_pandas() golden file tests. + + Holds the conversion helper and the source-table inventory, split into group + methods that these and the (temporal-flag) tests reuse. Defines no ``test_*`` of + its own. + """ + + def _to_pandas_cell(self, table, **to_pandas_kwargs) -> str: + """ + Convert ``table`` via ``to_pandas(**to_pandas_kwargs)`` and format the result + as a golden-file cell, returning ``ERR@<ExceptionClass>`` if the conversion + raises. Only the conversion is guarded: a formatting error is a test bug, not + a conversion signal, so it propagates instead of masquerading as ``ERR@``. + """ + try: + pdf = table.to_pandas(**to_pandas_kwargs) + except Exception as e: + return f"ERR@{type(e).__name__}" + return self.repr_value(pdf, max_len=0) + + def _structural_tables(self): + """Assembly shapes and empty edges (no temporal coercion involved).""" + sources = {} + + # ===================================================================== + # Empty edges (shapes an Array/ChunkedArray cannot be) + # ===================================================================== + sources["empty:0-columns"] = pa.table({}) + sources["empty:columns-no-rows"] = pa.table( + {"i": pa.array([], pa.int64()), "s": pa.array([], pa.string())} + ) + + # ===================================================================== + # Multi-column assembly + # ===================================================================== + sources["single-column"] = pa.table({"i": pa.array([1, 2, None], pa.int64())}) + sources["single-column:string"] = pa.table( + {"s": pa.array(["hello", "world", None], pa.string())} + ) + sources["multi-column:mixed-scalar"] = pa.table( + { + "i": pa.array([1, 2, None], pa.int64()), + "s": pa.array(["a", "b", None], pa.string()), + "f": pa.array([1.5, 2.5, 3.5], pa.float64()), + "b": pa.array([True, False, None], pa.bool_()), + } + ) + sources["multi-column:all-null"] = pa.table( + { + "i": pa.array([None, None], pa.int64()), + "s": pa.array([None, None], pa.string()), + } + ) + sources["multi-column:nested"] = pa.table( + { + "lst": pa.array([[1, 2], [3], None], pa.list_(pa.int64())), + "st": pa.array([{"x": 1}, None, {"x": 3}], pa.struct([("x", pa.int64())])), + } + ) + + return sources + + def _temporal_tables(self): + """ + Tables with temporal columns. Shared with the temporal-flag tests, where + coerce_temporal_nanoseconds / date_as_object actually move cells. + """ + sources = {} + + sources["temporal:timestamp"] = pa.table( + { + "ts": pa.array( + [ + datetime.datetime(2020, 1, 1, 5, 30), + datetime.datetime(2021, 6, 15, 23, 59), + ], + pa.timestamp("us"), + ) + } + ) + sources["temporal:date32"] = pa.table( + {"d": pa.array([datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)], pa.date32())} + ) + sources["temporal:date64"] = pa.table( + {"d": pa.array([datetime.date(2020, 1, 1), datetime.date(2021, 6, 15)], pa.date64())} + ) + sources["temporal:date32-far-future"] = pa.table( + {"d": pa.array([datetime.date(9999, 12, 31)], pa.date32())} + ) + sources["temporal:multi-column-mix"] = pa.table( + { + "ts": pa.array([datetime.datetime(2020, 1, 1, 5, 30)] * 2, pa.timestamp("us")), + "d": pa.array( + [datetime.date(2020, 1, 1), datetime.date(9999, 12, 31)], pa.date32() + ), + "i": pa.array([1, 2], pa.int64()), + } + ) + + return sources + + def _build_all_tables(self): + """Build an ordered dict of named source PyArrow tables for testing.""" + sources = {} + for group in [ + self._structural_tables(), + self._temporal_tables(), + ]: + sources.update(group) + return sources + + +@unittest.skipIf( + not have_pyarrow or not have_pandas, + pyarrow_requirement_message or pandas_requirement_message, +) +class PyArrowTableToPandasDefaultTests(_PyArrowTableToPandasTestBase): + """ + Tests pa.Table.to_pandas() with default arguments via golden file comparison. + + Pins multi-column DataFrame assembly (combined schema, column names, coexisting + dtypes) and the empty-table edges. + """ + + def test_to_pandas_default(self): + """Test pa.Table.to_pandas() with default arguments against golden file.""" + sources = self._build_all_tables() + row_names = list(sources.keys()) + col_names = ["pyarrow table", "pandas dataframe"] + + overrides = {} + pandas_3_plus = LooseVersion(pd.__version__) >= LooseVersion("3.0.0") + pyarrow_19_plus = LooseVersion(pa.__version__) >= LooseVersion("19.0.0") + + # On pandas 3 with PyArrow 19+, Arrow string columns convert to the "str" + # dtype (missing values become nan) rather than object dtype (None). This + # boundary differs from the Array path (test_pyarrow_arrow_to_pandas_*), + # whose empty string columns switch only at PyArrow 24. + if pandas_3_plus and pyarrow_19_plus: + overrides.update( + { + ("single-column:string", "pandas dataframe"): ( + "{'s': ['hello', 'world', nan]}@Dataframe[s str]" + ) + } + ) + overrides.update( + { + ("multi-column:mixed-scalar", "pandas dataframe"): ( + "{'i': [1.0, 2.0, nan], 's': ['a', 'b', nan], " + "'f': [1.5, 2.5, 3.5], 'b': [True, False, None]}" + "@Dataframe[i float64, s str, f float64, b object]" + ) + } + ) + overrides.update( + { + ("multi-column:all-null", "pandas dataframe"): ( + "{'i': [nan, nan], 's': [nan, nan]}@Dataframe[i float64, s str]" + ) + } + ) + overrides.update( + { + ("empty:columns-no-rows", "pandas dataframe"): ( + "{'i': [], 's': []}@Dataframe[i int64, s str]" + ) + } + ) + + def compute_cell(row_name, col_name): + table = sources[row_name] + if col_name == "pyarrow table": + return self.repr_value(table, max_len=0) + else: + return self._to_pandas_cell(table) + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_table_to_pandas", + index_name="test case", + overrides=overrides, + ) + + +@unittest.skipIf( + not have_pyarrow or not have_pandas, + pyarrow_requirement_message or pandas_requirement_message, +) +class PyArrowTableToPandasCoerceTemporalTests(_PyArrowTableToPandasTestBase): + """ + Tests pa.Table.to_pandas(coerce_temporal_nanoseconds=True) via golden file comparison. + + Reuses the shared temporal tables plus a coercion overflow row, recorded under two + columns: the default date_as_object=True (dates stay object, unaffected by coercion) + and date_as_object=False (the datetime64[ns] path where coercion applies, including + the far-future overflow). + """ + + # to_pandas(coerce_temporal_nanoseconds=True); date_as_object at its default (True). + COL_PANDAS = "pandas dataframe" + + # to_pandas(coerce_temporal_nanoseconds=True, date_as_object=False) -- the only path + # on which coercion observably affects date columns. + COL_PANDAS_DATE_AS_OBJECT_FALSE = "pandas dataframe (date_as_object=False)" + + def test_to_pandas_coerce_temporal_nanoseconds(self): + """Test pa.Table.to_pandas(coerce_temporal_nanoseconds=True) against golden file.""" + sources = self._temporal_tables() + # Coercion to nanoseconds has a valid range (~1677-2262); a far-future + # second-resolution timestamp column cannot fit and raises. + sources["temporal:timestamp-overflow"] = pa.table( + {"ts": pa.array([datetime.datetime(2500, 1, 1)], pa.timestamp("s"))} + ) + # A duration beyond ~292 years also exceeds the int64 nanosecond range, but + # unlike timestamp overflow, coercion wraps it silently to a bogus value. + sources["temporal:duration-overflow"] = pa.table( + {"dur": pa.array([datetime.timedelta(days=300 * 365)], pa.duration("s"))} + ) + row_names = list(sources.keys()) + col_names = [ + "pyarrow table", + self.COL_PANDAS, + self.COL_PANDAS_DATE_AS_OBJECT_FALSE, + ] + + # Version-specific expected values go here, keyed by (row, col), when a newer + # pandas/PyArrow legitimately changes a cell. Add a LooseVersion-guarded block + # for each known drift. + overrides: dict[tuple[str, str], str] = {} + + def compute_cell(row_name, col_name): + table = sources[row_name] + if col_name == "pyarrow table": + return self.repr_value(table, max_len=0) + elif col_name == self.COL_PANDAS: + return self._to_pandas_cell(table, coerce_temporal_nanoseconds=True) + elif col_name == self.COL_PANDAS_DATE_AS_OBJECT_FALSE: + return self._to_pandas_cell( + table, coerce_temporal_nanoseconds=True, date_as_object=False + ) + else: + raise ValueError(f"unknown column: {col_name}") + + self.compare_or_generate_golden_matrix( + row_names=row_names, + col_names=col_names, + compute_cell=compute_cell, + golden_file_prefix="golden_pyarrow_table_to_pandas_coerce_temporal", + index_name="test case", + overrides=overrides, + ) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_type_coercion.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_type_coercion.py index aeb51c9020cd2..5c6d18f1a97af 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_type_coercion.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_type_coercion.py @@ -30,9 +30,9 @@ """ import datetime -from decimal import Decimal import math import unittest +from decimal import Decimal from typing import Any, List, Tuple from pyspark.loose_version import LooseVersion @@ -159,9 +159,10 @@ def test_pandas_na_coercion(self): def test_python_instances_coercion(self): """Test type coercion from Python list, tuple, generator with all data types.""" - import pyarrow as pa from zoneinfo import ZoneInfo + import pyarrow as pa + # ==== 2.1 Numeric Types ==== # (data, target_type, expected_values) @@ -447,10 +448,11 @@ def test_python_instances_coercion(self): ) def test_pandas_instances_coercion(self): """Test type coercion from pandas Series with various backend types.""" + from zoneinfo import ZoneInfo + import numpy as np import pandas as pd import pyarrow as pa - from zoneinfo import ZoneInfo # Constants int8_min, int8_max = np.iinfo(np.int8).min, np.iinfo(np.int8).max diff --git a/python/pyspark/traceback_utils.py b/python/pyspark/traceback_utils.py index fcc388ba080e0..8697842e182a3 100644 --- a/python/pyspark/traceback_utils.py +++ b/python/pyspark/traceback_utils.py @@ -15,9 +15,9 @@ # limitations under the License. # -from collections import namedtuple import os import traceback +from collections import namedtuple CallSite = namedtuple("CallSite", "function file linenum") diff --git a/python/pyspark/util.py b/python/pyspark/util.py index 858dc1013e94d..f28e17fe33cbb 100644 --- a/python/pyspark/util.py +++ b/python/pyspark/util.py @@ -17,25 +17,25 @@ import contextlib import copy -import functools import faulthandler +import functools import gc import itertools import os import re +import socket import sys import threading import traceback import typing -import socket import warnings from contextlib import contextmanager from types import TracebackType from typing import ( + IO, Any, Callable, Generator, - IO, Iterator, List, Optional, @@ -47,11 +47,11 @@ from pyspark.errors import PySparkRuntimeError from pyspark.serializers import ( - write_int, - read_int, - write_with_length, SpecialLengths, UTF8Deserializer, + read_int, + write_int, + write_with_length, ) __all__: List[str] = [] @@ -63,40 +63,48 @@ from py4j.java_gateway import JavaObject from pyspark._typing import NonUDFType - from pyspark.sql.pandas._typing import ( - PandasScalarUDFType, - PandasGroupedMapUDFType, - PandasGroupedAggUDFType, - PandasWindowAggUDFType, - PandasScalarIterUDFType, - PandasMapIterUDFType, - PandasCogroupedMapUDFType, - ArrowMapIterUDFType, - PandasGroupedMapUDFWithStateType, - ArrowGroupedMapUDFType, - ArrowGroupedMapIterUDFType, - ArrowCogroupedMapUDFType, - PandasGroupedMapIterUDFType, - PandasGroupedAggIterUDFType, - PandasGroupedMapUDFTransformWithStateType, - PandasGroupedMapUDFTransformWithStateInitStateType, - GroupedMapUDFTransformWithStateType, - GroupedMapUDFTransformWithStateInitStateType, - ArrowScalarUDFType, - ArrowScalarIterUDFType, - ArrowGroupedAggUDFType, - ArrowGroupedAggIterUDFType, - ArrowWindowAggUDFType, - ) + from pyspark.serializers import Serializer + from pyspark.sql import SparkSession from pyspark.sql._typing import ( SQLArrowBatchedUDFType, + SQLArrowElementwiseUDFType, SQLArrowTableUDFType, + SQLArrowUDTFType, SQLBatchedUDFType, + SQLScalarArrowElementwiseUDFType, + SQLScalarArrowIterElementwiseUDFType, + SQLScalarPandasElementwiseUDFType, + SQLScalarPandasIterElementwiseUDFType, SQLTableUDFType, - SQLArrowUDTFType, ) - from pyspark.serializers import Serializer - from pyspark.sql import SparkSession + from pyspark.sql.pandas._typing import ( + ArrowCogroupedMapUDFType, + ArrowGroupedAggIncrementalFinalUDFType, + ArrowGroupedAggIncrementalPartialUDFType, + ArrowGroupedAggIterUDFType, + ArrowGroupedAggUDFType, + ArrowGroupedMapIterUDFType, + ArrowGroupedMapUDFType, + ArrowMapIterUDFType, + ArrowScalarIterUDFType, + ArrowScalarUDFType, + ArrowWindowAggIncrementalUDFType, + ArrowWindowAggUDFType, + GroupedMapUDFTransformWithStateInitStateType, + GroupedMapUDFTransformWithStateType, + PandasCogroupedMapUDFType, + PandasGroupedAggIterUDFType, + PandasGroupedAggUDFType, + PandasGroupedMapIterUDFType, + PandasGroupedMapUDFTransformWithStateInitStateType, + PandasGroupedMapUDFTransformWithStateType, + PandasGroupedMapUDFType, + PandasGroupedMapUDFWithStateType, + PandasMapIterUDFType, + PandasScalarIterUDFType, + PandasScalarUDFType, + PandasWindowAggUDFType, + ) JVM_BYTE_MIN: int = -(1 << 7) @@ -442,10 +450,11 @@ def inner(*args: Any, **kwargs: Any) -> Any: return outer # Non Spark Connect with SparkSession or Callable - from pyspark.sql import SparkSession - from pyspark import SparkContext from py4j.clientserver import ClientServer + from pyspark import SparkContext + from pyspark.sql import SparkSession + if isinstance(SparkContext._gateway, ClientServer): # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. @@ -596,9 +605,10 @@ def copy_local_properties(*a: Any, **k: Any) -> Any: super().__init__(target=copy_local_properties, *args, **kwargs) # type: ignore[misc] else: # Non Spark Connect - from pyspark import SparkContext from py4j.clientserver import ClientServer + from pyspark import SparkContext + self._session = session # type: ignore[assignment] if isinstance(SparkContext._gateway, ClientServer): # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. @@ -632,9 +642,10 @@ def start(self) -> None: self._tags = set(thread_local.tags) else: # Non Spark Connect - from pyspark import SparkContext from py4j.clientserver import ClientServer + from pyspark import SparkContext + if isinstance(SparkContext._gateway, ClientServer): # Here's when the pinned-thread mode (PYSPARK_PIN_THREAD) is on. @@ -662,6 +673,11 @@ class PythonEvalType: SQL_BATCHED_UDF: "SQLBatchedUDFType" = 100 SQL_ARROW_BATCHED_UDF: "SQLArrowBatchedUDFType" = 101 + SQL_ARROW_ELEMENTWISE_UDF: "SQLArrowElementwiseUDFType" = 102 + SQL_SCALAR_PANDAS_ELEMENTWISE_UDF: "SQLScalarPandasElementwiseUDFType" = 103 + SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF: "SQLScalarPandasIterElementwiseUDFType" = 104 + SQL_SCALAR_ARROW_ELEMENTWISE_UDF: "SQLScalarArrowElementwiseUDFType" = 105 + SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF: "SQLScalarArrowIterElementwiseUDFType" = 106 SQL_SCALAR_PANDAS_UDF: "PandasScalarUDFType" = 200 SQL_GROUPED_MAP_PANDAS_UDF: "PandasGroupedMapUDFType" = 201 @@ -689,6 +705,17 @@ class PythonEvalType: SQL_WINDOW_AGG_ARROW_UDF: "ArrowWindowAggUDFType" = 253 SQL_GROUPED_AGG_ARROW_ITER_UDF: "ArrowGroupedAggIterUDFType" = 254 + # Incremental (partial + final) Arrow aggregator. See ``pyspark.sql.aggregator``. + # PARTIAL folds input rows into a per-group buffer via ``Aggregator.reduce`` on the map side; + # FINAL merges partial buffers via ``Aggregator.merge`` and produces output via + # ``Aggregator.finish`` after the shuffle. + SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF: "ArrowGroupedAggIncrementalPartialUDFType" = 255 + SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF: "ArrowGroupedAggIncrementalFinalUDFType" = 256 + + # Window aggregation with an incremental ``Aggregator``. A window has no shuffle, so each + # frame's rows are folded with ``reduce`` (from ``zero``) and finished with ``finish``. + SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF: "ArrowWindowAggIncrementalUDFType" = 257 + SQL_TABLE_UDF: "SQLTableUDFType" = 300 SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301 SQL_ARROW_UDTF: "SQLArrowUDTFType" = 302 @@ -1047,6 +1074,7 @@ def enable_faulthandler(self, start_periodic_traceback: bool = True) -> Iterator if __name__ == "__main__": import doctest + import pyspark.util from pyspark.core.context import SparkContext diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 2db77d73148ea..a37c81893e036 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -19,28 +19,28 @@ Worker that receives input from Piped RDD. """ -import os -import sys import dataclasses -import time import inspect import itertools import json +import os +import sys +import time import warnings from collections.abc import Iterator from typing import ( + TYPE_CHECKING, Any, + BinaryIO, Callable, Iterable, Optional, Tuple, Type, TypeVar, - TYPE_CHECKING, Union, get_args, get_origin, - BinaryIO, ) T = TypeVar("T") @@ -51,35 +51,40 @@ from pyspark.sql.pandas._typing import GroupedBatch +from pyspark import _NoValue, shuffle from pyspark.accumulators import ( SpecialAccumulatorIds, _accumulatorRegistry, _deserialize_accumulator, ) -from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient -from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode -from pyspark.taskcontext import BarrierTaskContext, TaskContext -from pyspark.util import PythonEvalType +from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError +from pyspark.logger.worker_io import capture_outputs +from pyspark.messages import ( + SparkMessageReceiver, + SparkSocketMessageReceiver, +) from pyspark.serializers import ( + BatchedSerializer, + CPickleSerializer, + SpecialLengths, write_int, write_long, - SpecialLengths, - CPickleSerializer, - BatchedSerializer, ) from pyspark.sql.conversion import ( - LocalDataToArrowConversion, - ArrowTableToRowsConversion, ArrowBatchTransformer, + ArrowTableToRowsConversion, + LocalDataToArrowConversion, PandasToArrowConversion, ) from pyspark.sql.functions import SkipRestOfInputTableException from pyspark.sql.pandas.serializers import ( - ArrowStreamSerializer, - ArrowStreamGroupSerializer, ArrowStreamCoGroupSerializer, + ArrowStreamGroupSerializer, + ArrowStreamSerializer, ) from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type +from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient +from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode from pyspark.sql.types import ( ArrayType, BinaryType, @@ -94,30 +99,25 @@ _create_row, _parse_datatype_json_string, ) +from pyspark.taskcontext import BarrierTaskContext, TaskContext from pyspark.util import ( + PythonEvalType, fail_on_stopiteration, handle_worker_exception, - with_faulthandler, start_faulthandler_periodic_traceback, + with_faulthandler, ) -from pyspark import _NoValue, shuffle -from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.worker_message import WorkerInitInfo from pyspark.worker_util import ( + Conf, check_python_version, get_sock_file_to_executor, - read_command, pickleSer, + read_command, send_accumulator_updates, setup_broadcasts, setup_memory_limits, setup_spark_files, - Conf, -) -from pyspark.logger.worker_io import capture_outputs -from pyspark.messages import ( - SparkMessageReceiver, - SparkSocketMessageReceiver, ) @@ -147,6 +147,16 @@ def use_legacy_pandas_udtf_conversion(self) -> bool: == "true" ) + @property + def map_in_batch_legacy_accept_any_iterable(self) -> bool: + return ( + self.get( + "spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled", + "false", + ) + == "true" + ) + @property def binary_as_bytes(self) -> bool: return self.get("spark.sql.execution.pyspark.binaryAsBytes", "true") == "true" @@ -224,6 +234,17 @@ def input_type(self) -> Optional[DataType]: return None return _parse_datatype_json_string(input_type) + @property + def elementwise_nesting(self) -> Optional[list]: + # Per-UDF nesting depth (parallel to the UDF list) for the element-wise lift: how many + # ``array`` levels the worker flattens off each argument and re-nests onto the result. A UDF + # in a single lambda is depth 1; one lifted out of nested lambdas is deeper. Absent/empty + # means depth 1 for every UDF. See ExtractPythonUDFFromLambda. + raw = self.get("elementwise_nesting", None) + if raw is None or raw == "": + return None + return [int(x) for x in raw.split(",")] + @property def table_arg_offsets(self) -> Optional[list[int]]: offsets = self.get("table_arg_offsets", None) @@ -245,6 +266,42 @@ def chain(f, g): return lambda *a: g(f(*a)) +# Sentinel standing in for NaN grouping-key values in the map-side incremental-aggregate combine. +# ``float('nan') != float('nan')``, so distinct NaN objects would never collide in a dict; mapping +# them to one sentinel lets the map-side combine collapse them (correctness does not depend on it, +# as the FINAL stage re-groups authoritatively). +_NAN_GROUPING_KEY = object() + + +def _hashable_grouping_key(key_values: Tuple[Any, ...]) -> Any: + """ + Canonicalize a grouping-key value tuple (extracted from Arrow via ``to_pylist``) into a + hashable form usable as a ``dict`` key for the map-side PARTIAL combine of incremental Python + aggregators. + + Lists (from ``array`` columns) become tuples and dicts (from ``struct`` columns) become tuples + of ``(name, value)`` pairs, so nested complex keys hash by value; NaN floats map to a sentinel. + This is a best-effort combine only: an exotic unhashable value falls back to a unique object so + the row forms its own group, and the FINAL stage merges any keys left uncollapsed here. + """ + + def canon(v: Any) -> Any: + if isinstance(v, float) and v != v: + return _NAN_GROUPING_KEY + if isinstance(v, list): + return tuple(canon(x) for x in v) + if isinstance(v, dict): + return tuple((k, canon(x)) for k, x in v.items()) + return v + + try: + canonical = tuple(canon(v) for v in key_values) + hash(canonical) + return canonical + except TypeError: + return object() + + def verify_return_type(result: T, expected_type: Type[T]) -> T: """ Verify a UDF return value against an expected type. @@ -325,73 +382,49 @@ def verify_scalar_result(result: Any, num_rows: int) -> Any: "actual": type(result).__name__, }, ) - if result_length != num_rows: - # TODO: change error class to RESULT_ROWS_MISMATCH - raise PySparkRuntimeError( - errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF", - messageParameters={ - "udf_type": "arrow_udf", - "expected": str(num_rows), - "actual": str(result_length), - }, - ) + verify_result_row_count(result_length, num_rows) return result -def verify_iterator_exhausted(iterator: Iterator, error_class: str) -> None: +def verify_iterator_exhausted(iterator: Iterator) -> None: """Verify that an iterator has been fully consumed.""" try: next(iterator) except StopIteration: pass else: - raise PySparkRuntimeError(errorClass=error_class, messageParameters={}) + raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) def verify_output_row_limit( iterator: Iterator, max_rows: Union[int, Callable[[], int]], - error_class: str, ) -> Iterator: """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" total_rows = 0 for element in iterator: total_rows += len(element) if total_rows > (max_rows() if callable(max_rows) else max_rows): - raise PySparkRuntimeError(errorClass=error_class, messageParameters={}) + raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) yield element -def verify_output_row_count( +def verify_iter_result_row_count( iterator: Iterator, - expected_rows: Union[int, Callable[[], int]], - error_class: str, + expected_rows: Callable[[], int], ) -> Iterator: - """Yield elements and verify final row count matches expected exactly.""" + """Yield elements and verify final row count matches expected exactly. + + ``expected_rows`` is a callable because the expected count is only known once + the iterator is fully consumed (input rows are counted lazily as a side effect + of pulling batches), so it must be read after this generator is exhausted. + """ actual_rows = 0 for element in iterator: actual_rows += len(element) yield element - expected = expected_rows() if callable(expected_rows) else expected_rows - if actual_rows != expected: - raise PySparkRuntimeError( - errorClass=error_class, - messageParameters={ - "output_length": str(actual_rows), - "input_length": str(expected), - }, - ) - - -def wrap_udf(f, args_offsets, kwargs_offsets, return_type): - func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets, kwargs_offsets) - - if return_type.needConversion(): - toInternal = return_type.toInternal - return args_kwargs_offsets, lambda *a: toInternal(func(*a)) - else: - return args_kwargs_offsets, lambda *a: func(*a) + verify_result_row_count(actual_rows, expected_rows()) def _verify_column_schema( @@ -522,6 +555,11 @@ def _is_iter_based(eval_type: int) -> bool: return eval_type in ( PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, + # Iterator UDFs lifted out of a higher-order function lambda keep the iterator contract: + # the user function still consumes and produces an iterator of batches; the worker only + # feeds it the flattened elements and re-nests the results. See ExtractPythonUDFFromLambda. + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, @@ -563,14 +601,13 @@ def profiling_func(*args, **kwargs): def wrap_memory_profiler(f, eval_type, result_id): + import pyspark.memory_profiler_ext from pyspark.sql.profiler import ( ProfileResultsParam, ProfileResultsParamV2, WorkerMemoryProfiler, ) - import pyspark.memory_profiler_ext - if not pyspark.memory_profiler_ext.has_memory_profiler: return f @@ -631,6 +668,7 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index): if eval_type in ( PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_ARROW_BATCHED_UDF, ): func = profiling_func @@ -644,9 +682,24 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index): # The last returnType will be the return type of UDF. Eval types are grouped below by the # shape of the value they return. + # Incremental Python aggregators: the pickled "function" is the Aggregator object itself, whose + # zero/reduce/merge/finish methods the worker calls directly. Return it unwrapped (not through + # fail_on_stopiteration, which would treat it as a plain callable). + if eval_type in ( + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, + PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF, + ): + return chained_func, args_offsets, kwargs_offsets, return_type + # Scalar, aggregation and window UDFs: (func, args_offsets, kwargs_offsets, return_type). if eval_type in ( PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, @@ -685,8 +738,15 @@ def read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index): return func, None, None, return_type elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF: return func, None, None, None + # Batched (plain Python) UDFs: (args_kwargs_offsets, eval func); apply kwargs binding and + # convert each result to the internal representation only when the return type requires it. elif eval_type == PythonEvalType.SQL_BATCHED_UDF: - return wrap_udf(func, args_offsets, kwargs_offsets, return_type) + func, args_kwargs_offsets = wrap_kwargs_support(func, args_offsets, kwargs_offsets) + if return_type.needConversion(): + toInternal = return_type.toInternal + return args_kwargs_offsets, lambda *a: toInternal(func(*a)) + else: + return args_kwargs_offsets, lambda *a: func(*a) else: raise ValueError("Unknown eval type: {}".format(eval_type)) @@ -1835,9 +1895,162 @@ def mapper(_, it): return mapper, None, ser, ser +def _elementwise_renest(flat_values, shape_lengths, is_large): + """Re-nest a flat Array of per-element results into an ``array<R>`` column. + + ``flat_values`` holds the results for every non-null element in order; ``shape_lengths`` is + the per-array element count of the iterated argument (``None`` for a null array, which stays + null and consumes no elements). ``is_large`` preserves the input's list width (``ListArray`` + with int32 offsets vs. ``LargeListArray`` with int64). + + Shared by the vectorized element-wise worker paths (scalar pandas / Arrow and their iterator + variants) that back Python UDFs inside higher-order function lambdas. See + ``ExtractPythonUDFFromLambda``. + """ + import pyarrow as pa + + offsets = [0] + running = 0 + mask = [] + for n in shape_lengths: + mask.append(n is None) + if n is not None: + running += n + offsets.append(running) + list_cls = pa.LargeListArray if is_large else pa.ListArray + offsets_arr = pa.array(offsets, type=pa.int64() if is_large else pa.int32()) + null_mask = pa.array(mask, type=pa.bool_()) + return list_cls.from_arrays(offsets_arr, flat_values, mask=null_mask) + + +def _elementwise_leaf_type(data_type, depth): + """The element type ``depth`` ``ArrayType`` levels below ``data_type``. + + A lifted UDF's argument arrives as ``array^depth<T>`` (one ``array`` level per enclosing higher- + order function lambda); this peels them off to the scalar leaf ``T`` the user function sees. See + ``ExtractPythonUDFFromLambda``. + """ + for _ in range(depth): + data_type = data_type.elementType + return data_type + + +def _elementwise_flatten_deep(col, depth): + """Flatten ``depth`` list levels off ``col``, keeping each level's shape for re-nesting. + + Returns ``(leaf, shape_levels, is_large_levels)``: ``leaf`` is the fully flattened element + ``pa.Array`` (the leaves of the ``depth``-deep nesting), ``shape_levels[k]`` is the per-slot + length (``None`` for a null slot) at level ``k`` (0 = outermost), and ``is_large_levels[k]`` + whether that level is a ``LargeListArray``. ``depth`` is 1 for a UDF in a single lambda and more + for one lifted out of nested lambdas. Shared by the element-wise worker paths. See + ``ExtractPythonUDFFromLambda``. + """ + import pyarrow as pa + import pyarrow.compute as pc + + shape_levels = [] + is_large_levels = [] + cur = col + for _ in range(depth): + shape_levels.append(pc.list_value_length(cur).to_pylist()) + is_large_levels.append(pa.types.is_large_list(cur.type)) + cur = cur.flatten() + return cur, shape_levels, is_large_levels + + +def _elementwise_flatten_leaf(col, depth): + """Flatten ``depth`` list levels off ``col`` to its leaf ``pa.Array``, without capturing shape. + + A lifted UDF re-nests its result by the *first* argument's per-level shapes only, so the other + arguments need just their leaves. This skips the ``pc.list_value_length(...).to_pylist()`` and + ``is_large`` bookkeeping ``_elementwise_flatten_deep`` does for the first argument. See + ``ExtractPythonUDFFromLambda``. + """ + cur = col + for _ in range(depth): + cur = cur.flatten() + return cur + + +def _elementwise_renest_deep(flat_values, shape_levels, is_large_levels): + """Re-nest a flat leaf Array back through ``len(shape_levels)`` list levels, innermost first. + + Inverse of ``_elementwise_flatten_deep``: rebuilds the ``array^depth<R>`` result from the flat + per-leaf results and the per-level shapes captured while flattening the input. For ``depth`` 1 + this is a single ``_elementwise_renest``. + """ + result = flat_values + for lengths, is_large in zip(reversed(shape_levels), reversed(is_large_levels)): + result = _elementwise_renest(result, lengths, is_large) + return result + + +def _elementwise_flatten_column(flat, element_type, is_pandas, runner_conf): + """Adapt one already-flattened ``array<T>`` element column to the vectorized fn's input. + + ``flat`` is the flattened element ``pa.Array`` (the caller flattens once per batch and shares it + across fused UDFs). Returns it unchanged for the Arrow flavor, or converted to a pandas Series / + DataFrame with the element type ``T`` for the pandas flavor. Shared by the vectorized + element-wise worker paths that back Python UDFs inside higher-order function lambdas. See + ``ExtractPythonUDFFromLambda``. + """ + if not is_pandas: + return flat + from pyspark.sql.conversion import ArrowArrayToPandasConversion + + return ArrowArrayToPandasConversion.convert( + flat, + element_type, + timezone=runner_conf.timezone, + struct_in_pandas="dict", + ndarray_as_list=False, + prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, + df_for_struct=True, + ) + + +def _elementwise_result_to_arrow(result, return_type, arrow_element_type, is_pandas, runner_conf): + """Convert one vectorized UDF result over the flat elements to a single flat Arrow Array. + + ``result`` is a pandas Series / DataFrame (pandas flavor) or a ``pa.Array`` (Arrow flavor); the + returned array holds one element per input element. The Arrow flavor is coerced to + ``arrow_element_type`` (UTC-typed); the pandas flavor is typed by ``PandasToArrowConversion`` + using the session timezone, so its timestamp type may differ from ``arrow_element_type`` - + callers that concatenate results must take the type from the returned array, not assume UTC. + Shared by the vectorized element-wise worker paths. See ``ExtractPythonUDFFromLambda``. + """ + import pyarrow as pa + + if is_pandas: + batch = PandasToArrowConversion.convert( + [result], + StructType([StructField("_0", return_type)]), + timezone=runner_conf.timezone, + safecheck=runner_conf.safecheck, + arrow_cast=True, + prefers_large_types=runner_conf.use_large_var_types, + assign_cols_by_name=runner_conf.assign_cols_by_name, + int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled, + ) + else: + batch = ArrowBatchTransformer.enforce_schema( + pa.RecordBatch.from_arrays([result], ["_0"]), + pa.schema([pa.field("_0", arrow_element_type)]), + safecheck=runner_conf.safecheck, + ) + # PandasToArrowConversion / enforce_schema both return a pa.RecordBatch, so column(0) is a + # single pa.Array (never a ChunkedArray). + return batch.column(0) + + def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): if eval_type in ( PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_SCALAR_PANDAS_UDF, PythonEvalType.SQL_SCALAR_ARROW_UDF, PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, @@ -1861,6 +2074,9 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, + PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF, ): # NOTE: if timezone is set here, that implies respectSessionTimeZone is True if eval_type in ( @@ -1868,12 +2084,17 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, + # The map-side PARTIAL stage streams ordinary (multi-group) batches and hash-combines + # inside the worker, so it uses the plain (non-grouped) stream serializer below. Only + # the post-shuffle FINAL stage receives one Arrow stream per group. + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, + PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF, ): ser = ArrowStreamGroupSerializer(write_start_stream=True) elif eval_type in ( @@ -1940,6 +2161,20 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # invoke the UDF output_batches = udf_func(input_batches) + # The declared signature is Iterator[...], so a strict iterator is required by + # default. With the legacy flag, accept any object Python can iterate over -- via + # iter(...), which honors both __iter__ and the sequence protocol (__getitem__) -- + # by adapting it into an iterator before the shared element-type verification. + if runner_conf.map_in_batch_legacy_accept_any_iterable and not isinstance( + output_batches, Iterator + ): + try: + output_batches = iter(output_batches) + except TypeError: + # Not iterable at all; leave it so verify_return_type below raises the + # standard UDF_RETURN_TYPE error. + pass + # Post-processing verified_iter = verify_return_type( output_batches, @@ -2026,24 +2261,19 @@ def process_results(): limited = verify_output_row_limit( process_results(), lambda: num_input_rows, - error_class="OUTPUT_EXCEEDS_INPUT_ROWS", ) # Apply row count match check (final) - matched = verify_output_row_count( + matched = verify_iter_result_row_count( limited, lambda: num_input_rows, - error_class="RESULT_ROWS_MISMATCH", ) # Yield batches yield from matched # Verify iterator consumed - verify_iterator_exhausted( - args_iter, - error_class="INPUT_NOT_FULLY_CONSUMED", - ) + verify_iterator_exhausted(args_iter) # profiling is not supported for UDF return func, None, ser, ser @@ -2119,9 +2349,175 @@ def grouped_func( # profiling is not supported for UDF return grouped_func, None, ser, ser - if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF: + if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF: + import pyarrow as pa + + # Map-side PARTIAL stage: hash-combine input rows into a per-group buffer via the + # aggregator's `reduce`. Ordinary (multi-group) batches are streamed in; the worker keeps + # one running buffer per distinct grouping key -- never whole groups of rows -- and, at end + # of partition, emits one row per key: the grouping key columns followed by one buffer + # struct column per aggregator. Because the FINAL stage re-groups these authoritatively + # after the shuffle, the worker's grouping only needs to be a best-effort combine: any keys + # it fails to collapse (e.g. NaN, which compares unequal to itself) are merged downstream. + # + # The leading `num_grouping_keys` input columns are the grouping keys (see the operator); + # `grouping_key_schema` carries their names/types so the emitted key columns round-trip. + grouping_key_schema = eval_conf.grouping_key_schema + num_grouping_keys = ( + len(grouping_key_schema.fields) if grouping_key_schema is not None else 0 + ) + + buffer_col_names = ["_%d" % i for i in range(len(udfs))] + buffer_arrow_types = [ + to_arrow_type( + agg.bufferSchema, + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ) + for agg, _, _, _ in udfs + ] + # Buffer field names are invariant across groups; compute them once per aggregator. + field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs] + + # The aggregator's `reduce` receives a single positional tuple, so any named arguments at + # the call site are appended after the positional ones, in call order (kwargs_offsets + # preserves that order). This mirrors how a Python call `f(*args, **kwargs)` would order + # them into one value tuple. + input_offsets_by_udf = [ + list(args_offsets) + list(kwargs_offsets.values()) + for _, args_offsets, kwargs_offsets, _ in udfs + ] + # Input columns actually consumed per batch: the leading grouping keys plus every + # aggregator input, deduplicated. A UDF input may reuse a grouping-key column (the operator + # dedups its projection), so converting by distinct offset avoids repeated `to_pylist()`. + needed_offsets = sorted( + set(range(num_grouping_keys)) | {o for offsets in input_offsets_by_udf for o in offsets} + ) + + # Cap the map-side buffer so a high-cardinality partition -- exactly where partial + # aggregation degenerates -- cannot grow the per-key dict without bound and OOM the worker. + # When the distinct-key count reaches the cap we flush the whole map as one batch and start + # fresh; end-of-partition buffers are emitted in equally bounded chunks. This is safe + # because the FINAL stage re-groups the emitted partial buffers authoritatively after the + # shuffle and merges any duplicate keys that early flushes produce. A non-positive + # maxRecordsPerBatch means "unbounded" (mirroring the reader side). + max_records = runner_conf.arrow_max_records_per_batch + cap = max_records if max_records > 0 else None + + def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + # hashable key -> (representative key value tuple, list of per-aggregator buffers) + groups: "dict[Any, tuple]" = {} + key_field_types: Optional[list] = None + + def make_batch(entries: list) -> pa.RecordBatch: + arrays = [] + names = [] + for j in range(num_grouping_keys): + arrays.append( + pa.array( + [e[0][j] for e in entries], + type=key_field_types[j], # type: ignore + ) + ) + names.append("k_%d" % j) + for i, (agg, _, _, _) in enumerate(udfs): + field_names = field_names_by_udf[i] + structs = [ + {name: e[1][i][t] for t, name in enumerate(field_names)} for e in entries + ] + arrays.append(pa.array(structs, type=buffer_arrow_types[i])) + names.append(buffer_col_names[i]) + return pa.RecordBatch.from_arrays(arrays, names) + + for batch in data: + if key_field_types is None: + key_field_types = [batch.schema.field(j).type for j in range(num_grouping_keys)] + pylist_by_offset = {o: batch.column(o).to_pylist() for o in needed_offsets} + key_cols = [pylist_by_offset[j] for j in range(num_grouping_keys)] + udf_cols = [ + [pylist_by_offset[o] for o in input_offsets_by_udf[i]] for i in range(len(udfs)) + ] + for r in range(batch.num_rows): + key_values = tuple(key_cols[j][r] for j in range(num_grouping_keys)) + hashable_key = _hashable_grouping_key(key_values) + entry = groups.get(hashable_key) + if entry is None: + buffers = [agg.zero() for agg, _, _, _ in udfs] + groups[hashable_key] = (key_values, buffers) + else: + buffers = entry[1] + for i, (agg, _, _, _) in enumerate(udfs): + cols_i = udf_cols[i] + buffers[i] = agg.reduce(buffers[i], tuple(c[r] for c in cols_i)) + if cap is not None and len(groups) >= cap: + yield make_batch(list(groups.values())) + groups = {} + + if not groups: + # Empty partition (or fully flushed above): emit nothing more. The FINAL stage + # supplies the global-aggregation identity row when there is no input at all. + return + + entries = list(groups.values()) + if cap is None: + yield make_batch(entries) + else: + for start in range(0, len(entries), cap): + yield make_batch(entries[start : start + cap]) + + # profiling is not supported for UDF + return func, None, ser, ser + + if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF: import pyarrow as pa + + # Post-shuffle FINAL stage: merge each group's partial buffers via the aggregator's `merge` + # and produce the output via `finish`. Buffers are streamed and merged one batch at a time. + # Every group the JVM sends yields exactly one output row; null partial-buffer rows are + # skipped, so an empty global aggregation (a single all-null buffer row injected by the + # operator) still produces `finish(zero)`. + col_names = ["_%d" % i for i in range(len(udfs))] + return_schema = to_arrow_schema( + StructType([StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]), + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ) + # Buffer field names are invariant across groups and batches; compute once per aggregator. + field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs] + + def grouped_func( + split_index: int, data: Iterator["GroupedBatch"] + ) -> Iterator[pa.RecordBatch]: + for group in data: + merged: list = [None] * len(udfs) + for batch in group: + for i, (agg, args_offsets, _, _) in enumerate(udfs): + field_names = field_names_by_udf[i] + m = merged[i] + for row in batch.column(args_offsets[0]).to_pylist(): + if row is None: + continue + partial = tuple(row[name] for name in field_names) + m = partial if m is None else agg.merge(m, partial) + merged[i] = m + results = [] + for i, (agg, _, _, _) in enumerate(udfs): + m = merged[i] if merged[i] is not None else agg.zero() + results.append(agg.finish(m)) + # Type each output array explicitly (mirroring the PARTIAL stage) so a non-trivial + # outputType or an all-None column does not depend on Arrow type inference. + result_arrays = [ + pa.array([r], type=return_schema.field(i).type) for i, r in enumerate(results) + ] + batch = pa.RecordBatch.from_arrays(result_arrays, col_names) + yield ArrowBatchTransformer.enforce_schema(batch, return_schema) + + # profiling is not supported for UDF + return grouped_func, None, ser, ser + + if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF: import pandas as pd + import pyarrow as pa col_names = ["_%d" % i for i in range(len(udfs))] output_schema = StructType( @@ -2164,8 +2560,8 @@ def grouped_func( return grouped_func, None, ser, ser if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One GROUPED_AGG_PANDAS_ITER UDF expected here." udf_func, args_offsets, _, return_type = udfs[0] @@ -2277,9 +2673,104 @@ def grouped_func( # profiling is not supported for UDF return grouped_func, None, ser, ser - if eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF: + if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF: import pyarrow as pa + + # Window aggregation with an incremental ``Aggregator``. The operator sends each frame -- + # the whole partition for an unbounded frame, or per-row ``[begin, end)`` slices for a + # bounded one -- and the worker folds the frame's rows with ``reduce`` (from a fresh + # ``zero``) and produces the value with ``finish``, one output value per input row. A window + # has no shuffle, so the intermediate buffer never leaves the worker (unlike the two-stage + # groupBy path); ``merge`` is not used here. + window_bound_types_str = runner_conf.get("window_bound_types") + window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] + + col_names = ["_%d" % i for i in range(len(udfs))] + return_schema = to_arrow_schema( + StructType([StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]), + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ) + + def fold(agg: Any, buffer: Any, value_cols: list, start: int, end: int) -> Any: + # Fold rows ``[start, end)`` (each a tuple across ``value_cols``, matching the call-site + # argument order) into ``buffer`` via the aggregator's ``reduce``. + for r in range(start, end): + buffer = agg.reduce(buffer, tuple(c[r] for c in value_cols)) + return buffer + + def grouped_func( + split_index: int, data: Iterator["GroupedBatch"] + ) -> Iterator[pa.RecordBatch]: + for group in data: + batch_list = list(group) + if not batch_list: + continue + if hasattr(pa, "concat_batches"): + concatenated = pa.concat_batches(batch_list) + else: + # pyarrow.concat_batches not supported before 19.0.0 + # remove this once we drop support for old versions + concatenated = pa.RecordBatch.from_struct_array( + pa.concat_arrays([b.to_struct_array() for b in batch_list]) + ) + num_rows = concatenated.num_rows + + result_arrays = [] + for udf_index, (agg, args_offsets, kwargs_offsets, _) in enumerate(udfs): + bound_type = window_bound_types[udf_index] + result_type = return_schema.field(udf_index).type + if bound_type == "unbounded": + # One frame spanning the whole partition: compute once, repeat per row. + value_cols = [concatenated.column(o).to_pylist() for o in args_offsets] + [ + concatenated.column(v).to_pylist() for v in kwargs_offsets.values() + ] + result = agg.finish(fold(agg, agg.zero(), value_cols, 0, num_rows)) + result_arrays.append(pa.array([result] * num_rows, type=result_type)) + elif bound_type == "bounded": + # Per-row frame ``[begin, end)``. Materialize the aggregator's input columns + # once; frames index into them by row. + begin_col = concatenated.column(args_offsets[0]) + end_col = concatenated.column(args_offsets[1]) + data_offsets = list(args_offsets[2:]) + list(kwargs_offsets.values()) + value_cols = [concatenated.column(o).to_pylist() for o in data_offsets] + # When consecutive frames share the same lower bound and only grow on the + # right (e.g. rowsBetween(unboundedPreceding, currentRow)), extend the + # running buffer by the newly-included rows instead of refolding from + # ``zero`` -- O(n) overall rather than O(n^2). Otherwise -- the lower bound + # advanced (a row left the window, which ``reduce`` cannot subtract) or the + # frame shrank -- refold the frame from ``zero``. + results = [] + have_running = False + running: Any = None + prev_begin = -1 + prev_end = 0 + for i in range(num_rows): + begin = begin_col[i].as_py() + end = end_col[i].as_py() + if have_running and begin == prev_begin and end >= prev_end: + running = fold(agg, running, value_cols, prev_end, end) + else: + running = fold(agg, agg.zero(), value_cols, begin, end) + have_running = True + prev_begin, prev_end = begin, end + results.append(agg.finish(running)) + result_arrays.append(pa.array(results, type=result_type)) + else: + raise PySparkRuntimeError( + errorClass="INVALID_WINDOW_BOUND_TYPE", + messageParameters={"window_bound_type": bound_type}, + ) + + batch = pa.RecordBatch.from_arrays(result_arrays, col_names) + yield ArrowBatchTransformer.enforce_schema(batch, return_schema) + + # profiling is not supported for UDF + return grouped_func, None, ser, ser + + if eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF: import pandas as pd + import pyarrow as pa window_bound_types_str = runner_conf.get("window_bound_types") window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] @@ -2491,8 +2982,8 @@ def grouped_func( return grouped_func, None, ser, ser if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One GROUPED_MAP_PANDAS UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] @@ -2562,8 +3053,8 @@ def grouped_func( return grouped_func, None, ser, ser if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One GROUPED_MAP_PANDAS_ITER UDF expected here." grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] @@ -2691,8 +3182,8 @@ def cogrouped_func( return cogrouped_func, None, ser, ser if eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One MAP_PANDAS_ITER UDF expected here." map_udf, _, _, return_type = udfs[0] @@ -2720,11 +3211,19 @@ def dataframe_iter(): df_for_struct=True, )[0] - # mapInPandas accepts any iterable (e.g. a list), not just an - # iterator, so the standard verify_return_type (which requires an - # Iterator) is intentionally not reused here. result = map_udf(dataframe_iter()) - if not isinstance(result, Iterator) and not hasattr(result, "__iter__"): + # The declared signature is Iterator[...], so a strict iterator is required by + # default. With the legacy flag, accept any object Python can iterate over -- via + # iter(...), which honors both __iter__ and the sequence protocol (__getitem__) -- + # by adapting it into an iterator. + if runner_conf.map_in_batch_legacy_accept_any_iterable and not isinstance( + result, Iterator + ): + try: + result = iter(result) + except TypeError: + pass # Not iterable; fall through to the UDF_RETURN_TYPE error below. + if not isinstance(result, Iterator): raise PySparkTypeError( errorClass="UDF_RETURN_TYPE", messageParameters={ @@ -2760,8 +3259,8 @@ def dataframe_iter(): return func, None, ser, ser if eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One COGROUPED_MAP_PANDAS UDF expected here." cogrouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] @@ -3011,6 +3510,427 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record # profiling is not supported for UDF return func, None, ser, ser + if eval_type == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF: + # This path exchanges data with the JVM over Arrow, so PyArrow is required. Fail with a + # clear message rather than a bare ImportError from `import pyarrow` below. + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + + require_minimum_pyarrow_version() + + import pyarrow as pa + + # Element-wise UDFs back higher-order lambdas like transform(arr, x -> udf(x)). + # ExtractPythonUDFFromLambda rewrites them so the UDF receives *all* array elements + # at once (as ``array<T>``) rather than per-element. Flatten each argument down to its + # leaves, evaluate once over the batch, then re-nest with the input offsets. A UDF lifted + # out of nested lambdas (e.g. transform(arr, i -> transform(i, x -> udf(x)))) flattens more + # than one ``array`` level - its per-UDF depth comes from ``elementwise_nesting``. + # Example: array<array<int>> -> udf(depth 2) -> array<array<int>>. + + # UDF preparation + input_fields = list(eval_conf.input_type) + nesting = eval_conf.elementwise_nesting + udf_infos = [] + for udf_index, udf in enumerate(udfs): + udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf + wrapped_func, args_kwargs_offsets = wrap_kwargs_support( + udf_func, udf_args_offsets, udf_kwargs_offsets + ) + depth = nesting[udf_index] if nesting is not None else 1 + # Each argument arrives as ``array^depth<T>``; convert its leaves with the element type + # ``T`` reached by peeling ``depth`` array levels. + arg_converters = [ + ArrowTableToRowsConversion._create_converter( + _elementwise_leaf_type(input_fields[o].dataType, depth), + none_on_identity=True, + binary_as_bytes=runner_conf.binary_as_bytes, + ) + for o in args_kwargs_offsets + ] + udf_infos.append( + ( + wrapped_func, + args_kwargs_offsets, + depth, + arg_converters, + # UDF returns one value per element; return type was pickled, unchanged. This is + # per-element, so element type equals the declared return type. + to_arrow_type( + udf_return_type, + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ), + LocalDataToArrowConversion._create_converter( + udf_return_type, + none_on_identity=True, + int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled, + ), + ) + ) + col_names = [f"_{i}" for i in range(len(udfs))] + + @fail_on_stopiteration + def _evaluate_elementwise_udf(udf_func, rows): + if runner_conf.arrow_concurrency_level <= 0: + return [udf_func(*row) for row in rows] + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool: + return list(pool.map(lambda row: udf_func(*row), rows)) + + def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + for input_batch in data: + # Each UDF is re-nested by *its own* first argument's shape. ExtractPythonUDFs can + # fuse UDFs over differently shaped/nested arrays into one batch, so a single shared + # shape would misalign every UDF but the first. The rewrite always passes at least + # one array argument, so `offsets` is non-empty. + output_arrays = [] + for info in udf_infos: + ( + wrapped_func, + offsets, + depth, + arg_converters, + arrow_element_type, + result_conv, + ) = info + # Flatten each argument `depth` list levels to its leaves; the first argument's + # per-level shapes drive the re-nest. + leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep( + input_batch.column(offsets[0]), depth + ) + columns = [] + for i, (o, conv) in enumerate(zip(offsets, arg_converters)): + leaf = ( + leaf0 + if i == 0 + else _elementwise_flatten_leaf(input_batch.column(o), depth) + ) + values = ArrowTableToRowsConversion._to_pylist(leaf) + if conv is not None: + values = [conv(v) for v in values] + columns.append(values) + + total_elements = len(columns[0]) + # Stream the argument tuples rather than materializing a batch-sized list. + rows = zip(*columns) + results = _evaluate_elementwise_udf(wrapped_func, rows) + verify_result_row_count(len(results), total_elements) + + # Convert results and re-nest to array<R> using that UDF's offsets. + converted = ( + [result_conv(r) for r in results] if result_conv is not None else results + ) + try: + flat_arr = pa.array(converted, type=arrow_element_type) + # Broader than the SQL_ARROW_BATCHED_UDF path above (which catches only + # ArrowInvalid): the element-wise wrapper commonly returns list/struct-typed + # elements, whose type mismatches surface as ArrowTypeError, so both are caught + # before falling back to an explicit cast. + except (pa.lib.ArrowInvalid, pa.lib.ArrowTypeError): + flat_arr = pa.array(converted).cast( + target_type=arrow_element_type, safe=runner_conf.safecheck + ) + output_arrays.append( + _elementwise_renest_deep(flat_arr, shape_levels, is_large_levels) + ) + + yield pa.RecordBatch.from_arrays(output_arrays, col_names) + + # profiling is not supported for UDF + return func, None, ser, ser + + if eval_type in ( + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + ): + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + + require_minimum_pyarrow_version() + + import pyarrow as pa + + # A scalar pandas or Arrow UDF lifted out of a higher-order function's lambda by + # ExtractPythonUDFFromLambda. Each argument arrives as ``array<T>`` aligned with the + # iterated array. We flatten each list column to its element column, run the *vectorized* + # function once over that flat column (so it still receives a pandas Series / DataFrame or a + # pa.Array, its native contract), then re-nest the flat result to ``array<R>`` using the + # input's offsets - one row in, one row out, one Python round trip per batch. + is_pandas = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF + + input_fields = list(eval_conf.input_type) + nesting = eval_conf.elementwise_nesting + udf_infos = [] + for udf_index, udf in enumerate(udfs): + udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf + wrapped_func, args_kwargs_offsets = wrap_kwargs_support( + udf_func, udf_args_offsets, udf_kwargs_offsets + ) + # The UDF returns one value per element, so its declared return type is the element + # type of the ``array<R>`` this operator produces. + arrow_element_type = to_arrow_type( + udf_return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types + ) + depth = nesting[udf_index] if nesting is not None else 1 + # Each argument arrives as ``array^depth<T>``; the vectorized function must see the leaf + # element type ``T`` reached by peeling ``depth`` array levels. + arg_leaf_types = [ + _elementwise_leaf_type(input_fields[o].dataType, depth) for o in args_kwargs_offsets + ] + udf_infos.append( + ( + wrapped_func, + args_kwargs_offsets, + udf_return_type, + arrow_element_type, + depth, + arg_leaf_types, + ) + ) + col_names = [f"_{i}" for i in range(len(udfs))] + + if is_pandas: + import pandas as pd + + def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + for input_batch in data: + output_arrays = [] + for ( + wrapped_func, + offsets, + return_type, + arrow_element_type, + depth, + arg_leaf_types, + ) in udf_infos: + # Flatten each argument `depth` list levels to its leaves and adapt to the + # vectorized fn's input. Different UDFs in one operator may iterate differently + # shaped or differently nested arrays, so each flattens and re-nests by its own + # argument (the first argument's per-level shapes drive the re-nest). + leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep( + input_batch.column(offsets[0]), depth + ) + total_elements = len(leaf0) + flat_columns = [ + _elementwise_flatten_column( + leaf0 + if i == 0 + else _elementwise_flatten_leaf(input_batch.column(o), depth), + t, + is_pandas, + runner_conf, + ) + for i, (o, t) in enumerate(zip(offsets, arg_leaf_types)) + ] + + result = wrapped_func(*flat_columns) + if is_pandas: + if not hasattr(result, "__len__"): + pd_type = ( + "pandas.DataFrame" + if isinstance(return_type, StructType) + else "pandas.Series" + ) + raise PySparkTypeError( + errorClass="UDF_RETURN_TYPE", + messageParameters={ + "expected": pd_type, + "actual": type(result).__name__, + }, + ) + # struct return type must be a DataFrame (matches the base pandas path). + if isinstance(return_type, StructType) and not isinstance( + result, pd.DataFrame + ): + raise PySparkValueError( + "Invalid return type. Please make sure that the UDF returns a " + "pandas.DataFrame when the specified return type is StructType." + ) + # Verify the flat length before re-nesting so a wrong-length result raises + # the friendly RESULT_ROWS_MISMATCH rather than an opaque pyarrow error. + verify_result_row_count(len(result), total_elements) + else: + # Arrow flavor: a non-array-like result (e.g. a bare int) raises the + # friendly UDF_RETURN_TYPE rather than a bare TypeError from len(), matching + # the base SQL_SCALAR_ARROW_UDF path, and also checks the flat length. + verify_scalar_result(result, total_elements) + + flat_arr = _elementwise_result_to_arrow( + result, return_type, arrow_element_type, is_pandas, runner_conf + ) + nested = _elementwise_renest_deep(flat_arr, shape_levels, is_large_levels) + output_arrays.append(nested) + + yield pa.RecordBatch.from_arrays(output_arrays, col_names) + + # profiling is not supported for UDF + return func, None, ser, ser + + if eval_type in ( + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, + ): + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version + + require_minimum_pyarrow_version() + + import collections + + import pyarrow as pa + + is_pandas = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF + if is_pandas: + import pandas as pd + + assert num_udfs == 1, "One SCALAR_*_ITER_ELEMENTWISE UDF expected here." + udf_func, args_offsets, kwargs_offsets, return_type = udfs[0] + assert not kwargs_offsets, "Iterator UDFs do not take keyword arguments." + + # A scalar iterator UDF (pandas or Arrow) lifted out of a higher-order function's lambda. + # The user function keeps its iterator contract: it consumes an iterator of batches and + # yields an iterator of batches, one output value per input value. We preserve that by + # feeding it the *flattened* elements of each input batch and, since the JVM joins UDF + # output to input positionally by row (one ``array<R>`` per input ``array<T>`` row, in + # order), buffering a FIFO of the per-row element counts to re-group the streamed flat + # results back into arrays. Output batch boundaries need not match input ones. + # Each argument arrives as ``array^depth<T>``; the vectorized function must see the leaf + # element type ``T`` (arguments may differ, e.g. an outer column repeated into an aligned + # array). ``depth`` > 1 for a UDF lifted out of nested lambdas. + input_fields = list(eval_conf.input_type) + nesting = eval_conf.elementwise_nesting + depth = nesting[0] if nesting is not None else 1 + arg_leaf_types = [ + _elementwise_leaf_type(input_fields[o].dataType, depth) for o in args_offsets + ] + arrow_element_type = to_arrow_type( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types + ) + + def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + # FIFO of per-input-batch nesting shapes awaiting their flat leaf results. Each entry + # (per-level shapes, per-level list width, leaf count) re-nests one input batch's worth + # of rows back to ``array^depth<R>`` once that many leaf results have streamed in. + pending_shapes: "collections.deque" = collections.deque() + num_input_elements = 0 + + def extract_flat(batch: pa.RecordBatch): + nonlocal num_input_elements + # Flatten each argument `depth` levels to its leaves; the first argument's per-level + # shapes re-nest this batch's rows. The user function sees the flat leaves as a + # pandas Series / DataFrame (pandas) or a pa.Array (Arrow), each with its leaf type. + leaf0, shape_levels, is_large_levels = _elementwise_flatten_deep( + batch.column(args_offsets[0]), depth + ) + pending_shapes.append((shape_levels, is_large_levels, len(leaf0))) + num_input_elements += len(leaf0) + flat_cols = [ + _elementwise_flatten_column( + leaf0 if i == 0 else _elementwise_flatten_leaf(batch.column(o), depth), + arg_leaf_types[i], + is_pandas, + runner_conf, + ) + for i, o in enumerate(args_offsets) + ] + return flat_cols[0] if len(flat_cols) == 1 else tuple(flat_cols) + + flat_args_iter = map(extract_flat, data) + + if not is_pandas: + verified_iter = verify_return_type( + udf_func(flat_args_iter), + Iterator[pa.Array], # type: ignore[type-abstract] + ) + else: + pandas_iter_type = ( + Iterator[pd.DataFrame] + if isinstance(return_type, StructType) + else Iterator[pd.Series] + ) + verified_iter = verify_return_type(udf_func(flat_args_iter), pandas_iter_type) + + # Buffer the streamed flat element results and emit an ``array<R>`` row as soon as the + # shape at the head of the FIFO is fully covered. A row whose length is 0 (an empty + # array) or None (a null array) needs no elements, so it is emitted immediately even + # before any chunk arrives - this matters when a whole partition is empty/null arrays + # and the UDF yields nothing, otherwise those rows would be dropped by the positional + # JVM join. Chunks are held in a list and concatenated only when a shape spans more than + # one, so a UDF that yields once per input batch (the common case) never re-copies the + # buffer. ``empty_type`` supplies the element type for a zero-length emit; it tracks the + # most recent chunk's type (even a zero-length chunk carries the flavor's type - the + # pandas flavor types timestamps with the session timezone), falling back to the + # UTC-typed ``arrow_element_type`` only before any chunk arrives, so all emitted batches + # share one schema. + pending_chunks: "list" = [] + pending_len = 0 + empty_type = arrow_element_type + num_output_elements = 0 + + def emit_ready(): + nonlocal pending_chunks, pending_len + while pending_shapes: + shape_levels, is_large_levels, needed = pending_shapes[0] + if needed > pending_len: + break + pending_shapes.popleft() + if needed == 0: + flat = pa.nulls(0, type=empty_type) + else: + combined = ( + pending_chunks[0] + if len(pending_chunks) == 1 + else pa.concat_arrays(pending_chunks) + ) + flat = combined.slice(0, needed) + remainder = combined.slice(needed) + pending_chunks = [remainder] if len(remainder) else [] + pending_len -= needed + nested = _elementwise_renest_deep(flat, shape_levels, is_large_levels) + yield pa.RecordBatch.from_arrays([nested], ["_0"]) + + def process_results(): + nonlocal pending_chunks, pending_len, empty_type, num_output_elements + for result in verified_iter: + if is_pandas: + verify_pandas_result( + result, + return_type, + assign_cols_by_name=True, + truncate_return_schema=True, + ) + chunk = _elementwise_result_to_arrow( + result, return_type, arrow_element_type, is_pandas, runner_conf + ) + num_output_elements += len(chunk) + # Fail fast if the UDF over-produces, before the buffer grows unbounded (the + # base iterator paths do the same via verify_output_row_limit). + if num_output_elements > num_input_elements: + raise PySparkRuntimeError( + errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={} + ) + # Even a zero-length chunk carries the flavor's element type (the pandas flavor + # types timestamps with the session timezone), so always take it: otherwise + # rows emitted for an all-empty batch before the first non-empty chunk would use + # the UTC-typed default and disagree with later batches, breaking the output + # stream's single-schema contract. + empty_type = chunk.type + if len(chunk): + pending_chunks.append(chunk) + pending_len += len(chunk) + yield from emit_ready() + + # The iterator is exhausted: every input row's flat elements must have arrived. + verify_result_row_count(num_output_elements, num_input_elements) + # Flush any residual all-empty / all-null rows (they consume no elements). + if pending_shapes: + yield from emit_ready() + verify_iterator_exhausted(flat_args_iter) + + yield from process_results() + + # profiling is not supported for UDF + return func, None, ser, ser + if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF: import pandas as pd import pyarrow as pa @@ -3057,15 +3977,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record "actual": type(result).__name__, }, ) - if len(result) != num_rows: - raise PySparkRuntimeError( - errorClass="SCHEMA_MISMATCH_FOR_PANDAS_UDF", - messageParameters={ - "udf_type": "pandas_udf", - "expected": str(num_rows), - "actual": str(len(result)), - }, - ) + verify_result_row_count(len(result), num_rows) # struct_in_pandas="dict": UDF must return DataFrame for struct types if isinstance(udf_return_type, StructType) and not isinstance( result, pd.DataFrame @@ -3151,31 +4063,26 @@ def process_results(): limited = verify_output_row_limit( process_results(), lambda: num_input_rows, - error_class="OUTPUT_EXCEEDS_INPUT_ROWS", ) # Apply row count match check (final) - matched = verify_output_row_count( + matched = verify_iter_result_row_count( limited, lambda: num_input_rows, - error_class="RESULT_ROWS_MISMATCH", ) # Yield batches yield from matched # Verify iterator consumed - verify_iterator_exhausted( - args_iter, - error_class="INPUT_NOT_FULLY_CONSUMED", - ) + verify_iterator_exhausted(args_iter) # profiling is not supported for UDF return func, None, ser, ser if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PANDAS UDF expected here." udf, arg_offsets, return_type = udfs[0] @@ -3315,8 +4222,8 @@ def convert_results(result_iter): return transform_with_state_func, None, ser, ser if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF: - import pyarrow as pa import pandas as pd + import pyarrow as pa assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PANDAS_INIT_STATE UDF expected here." udf, arg_offsets, return_type = udfs[0] @@ -3526,8 +4433,9 @@ def convert_results( return func, None, ser, ser if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: - import pyarrow as pa import pandas as pd + import pyarrow as pa + from pyspark.sql.streaming.state import GroupState assert num_udfs == 1, "One GROUPED_MAP_PANDAS_UDF_WITH_STATE UDF expected here." @@ -4137,22 +5045,26 @@ def convert_results(result_rows: Iterable[Any]) -> Iterator["pa.RecordBatch"]: # profiling is not supported for UDF return func, None, ser, ser - else: - - def mapper(a): - result = tuple(f(*[a[o] for o in arg_offsets]) for arg_offsets, f in udfs) - # In the special case of a single UDF this will return a single result rather - # than a tuple of results; this is the format that the JVM side expects. - if len(result) == 1: - return result[0] - else: - return result + elif eval_type == PythonEvalType.SQL_BATCHED_UDF: + # Plain Python (pickle) UDFs, the only eval type reaching this branch. read_single_udf + # prepared each UDF as an (arg_offsets, eval_func) pair. Apply every one to each input + # row: a single UDF yields its bare result, multiple UDFs yield a tuple of results, + # which is the shape the JVM side expects. num_udfs is fixed, so the single-result + # case is handled once here rather than by unwrapping a one-element tuple per row. + def func(split_index: int, data: Iterator[Any]) -> Iterator[Any]: + if num_udfs == 1: + arg_offsets, f = udfs[0] + return (f(*[row[offset] for offset in arg_offsets]) for row in data) + return ( + tuple(f(*[row[offset] for offset in arg_offsets]) for arg_offsets, f in udfs) + for row in data + ) - def func(_, it): - return map(mapper, it) + # profiling is not supported for UDF + return func, None, ser, ser - # profiling is not supported for UDF - return func, None, ser, ser + else: + raise ValueError("Unknown eval type: {}".format(eval_type)) def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO): diff --git a/python/pyspark/worker_message.py b/python/pyspark/worker_message.py index 5be806882c977..72db43c6aad2f 100644 --- a/python/pyspark/worker_message.py +++ b/python/pyspark/worker_message.py @@ -18,14 +18,15 @@ import dataclasses import json import sys -from typing import Optional, TypeAlias, Union, IO, Any +from decimal import Decimal +from typing import IO, Any, Optional, TypeAlias, Union from pyspark.errors import PySparkValueError -from pyspark.serializers import read_bool, read_int, read_long, SpecialLengths +from pyspark.messages import ZeroCopyByteStream +from pyspark.serializers import SpecialLengths, read_bool, read_int, read_long from pyspark.taskcontext import BarrierTaskContext, ResourceInformation, TaskContext from pyspark.util import PythonEvalType from pyspark.worker_util import utf8_deserializer -from pyspark.messages import ZeroCopyByteStream @dataclasses.dataclass @@ -43,7 +44,7 @@ class ResourceInfo: attempt_number: int task_attempt_id: int cpus: int - cpu_amount: float + cpu_amount: Decimal resources: dict[str, ResourceInfo] local_properties: dict[str, str] @@ -59,7 +60,7 @@ def from_stream(cls, stream: ZeroCopyByteStream) -> "TaskContextInfo": attempt_number=task_context_json["attemptNumber"], task_attempt_id=task_context_json["taskAttemptId"], cpus=task_context_json["cpus"], - cpu_amount=float(task_context_json["cpuAmount"]), + cpu_amount=Decimal(task_context_json["cpuAmount"]), resources={ k: cls.ResourceInfo(name=v["name"], addresses=v["addresses"]) for k, v in task_context_json["resources"].items() diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 13e08449de22d..44c693ee70fa7 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -19,13 +19,13 @@ Util functions for workers. """ -from contextlib import contextmanager import importlib -from inspect import currentframe, getframeinfo import os import sys -from typing import Any, Generator, IO, Optional, Union, overload import warnings +from contextlib import contextmanager +from inspect import currentframe, getframeinfo +from typing import IO, Any, Generator, Optional, Union, overload from pyspark.messages import ZeroCopyByteStream @@ -42,17 +42,16 @@ has_resource_module = False from pyspark.accumulators import _accumulatorRegistry -from pyspark.util import is_remote_only from pyspark.errors import PySparkRuntimeError -from pyspark.util import local_connect_and_auth from pyspark.serializers import ( + CPickleSerializer, + FramedSerializer, + UTF8Deserializer, read_int, read_long, write_int, - FramedSerializer, - UTF8Deserializer, - CPickleSerializer, ) +from pyspark.util import is_remote_only, local_connect_and_auth pickleSer = CPickleSerializer() utf8_deserializer = UTF8Deserializer() diff --git a/python/run-tests.py b/python/run-tests.py index df9af0d4f3fe5..e4413cfe9bf44 100755 --- a/python/run-tests.py +++ b/python/run-tests.py @@ -18,32 +18,30 @@ # import asyncio +import io import logging -from argparse import ArgumentParser import os -import io import platform import pty +import queue as Queue import re import shutil import subprocess import sys import tempfile -from threading import Thread, Lock import time import uuid -import queue as Queue +from argparse import ArgumentParser from multiprocessing import Manager - +from threading import Lock, Thread # Append `SPARK_HOME/dev` to the Python path so that we can import the sparktestsupport module sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../dev/")) from sparktestsupport import SPARK_HOME -from sparktestsupport.shellutils import which, subprocess_check_output from sparktestsupport.modules import all_modules, pyspark_sql # noqa - +from sparktestsupport.shellutils import subprocess_check_output, which python_modules = dict((m.name, m) for m in all_modules if m.python_test_goals if m.name != "root") @@ -420,6 +418,12 @@ def parse_opts(): default=4, help="The number of suites to test in parallel (default %(default)d)", ) + parser.add_argument( + "--changed-files", + type=str, + default=None, + help="A file containing a list of changed files (default: %(default)s)", + ) parser.add_argument("--verbose", action="store_true", help="Enable additional debug logging") group = parser.add_argument_group("Developer Options") @@ -485,6 +489,12 @@ def main(): python_execs = opts.python_executables.split(",") LOGGER.info("Will test against the following Python executables: %s", python_execs) + if opts.changed_files: + with open(opts.changed_files, "r") as f: + changed_files = f.read().splitlines() + os.environ["PYSPARK_CHANGED_FILES"] = opts.changed_files + LOGGER.info("Will select tests based on the following changed files: %s", changed_files) + if should_test_modules: modules_to_test = [] for module_name in opts.modules.split(","): diff --git a/python/test_support/pytorch_training_test_file.py b/python/test_support/pytorch_training_test_file.py index 150246563f094..528f497713376 100644 --- a/python/test_support/pytorch_training_test_file.py +++ b/python/test_support/pytorch_training_test_file.py @@ -21,13 +21,14 @@ momentum = 0.5 log_interval = 100 +import shutil +import tempfile + import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms -import tempfile -import shutil class Net(nn.Module): diff --git a/repl/src/main/scala/org/apache/spark/repl/Signaling.scala b/repl/src/main/scala/org/apache/spark/repl/Signaling.scala index 9577e0ecaa2ef..81cd0794edf82 100644 --- a/repl/src/main/scala/org/apache/spark/repl/Signaling.scala +++ b/repl/src/main/scala/org/apache/spark/repl/Signaling.scala @@ -33,7 +33,7 @@ private[repl] object Signaling extends Logging { if (!ctx.statusTracker.getActiveJobIds().isEmpty) { logWarning("Cancelling all active jobs, this can take a while. " + "Press Ctrl+C again to exit now.") - ctx.cancelAllJobs() + ctx.cancelAllJobs("because the driver process received an interrupt signal (SIGINT)") true } else { false diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStep.scala index 579e5baeffdac..3201ab48e86e3 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStep.scala @@ -100,30 +100,24 @@ private[spark] class BasicDriverFeatureStep(conf: KubernetesDriverConf) val driverUIPort = SparkUI.getUIPort(conf.sparkConf) val driverSparkConnectServerPort = conf.sparkConf.getInt(CONNECT_GRPC_BINDING_PORT, DEFAULT_SPARK_CONNECT_SERVER_PORT) + // 0 is invalid as kubernetes containerPort request, we shall leave it unmounted. + val driverContainerPorts = Seq( + DRIVER_PORT_NAME -> driverPort, + BLOCK_MANAGER_PORT_NAME -> driverBlockManagerPort, + UI_PORT_NAME -> driverUIPort, + SPARK_CONNECT_SERVER_PORT_NAME -> driverSparkConnectServerPort + ).collect { case (name, port) if port != 0 => + new ContainerPortBuilder() + .withName(name) + .withContainerPort(port) + .withProtocol("TCP") + .build() + } val driverContainer = new ContainerBuilder(pod.container) .withName(Option(pod.container.getName).getOrElse(DEFAULT_DRIVER_CONTAINER_NAME)) .withImage(driverContainerImage) .withImagePullPolicy(conf.imagePullPolicy) - .addNewPort() - .withName(DRIVER_PORT_NAME) - .withContainerPort(driverPort) - .withProtocol("TCP") - .endPort() - .addNewPort() - .withName(BLOCK_MANAGER_PORT_NAME) - .withContainerPort(driverBlockManagerPort) - .withProtocol("TCP") - .endPort() - .addNewPort() - .withName(UI_PORT_NAME) - .withContainerPort(driverUIPort) - .withProtocol("TCP") - .endPort() - .addNewPort() - .withName(SPARK_CONNECT_SERVER_PORT_NAME) - .withContainerPort(driverSparkConnectServerPort) - .withProtocol("TCP") - .endPort() + .addAllToPorts(driverContainerPorts.asJava) .addNewEnv() .withName(ENV_SPARK_USER) .withValue(Utils.getCurrentUserName()) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStep.scala index 9eb9303fa0ddf..af6273c131d4b 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStep.scala @@ -53,7 +53,8 @@ private[spark] class BasicExecutorFeatureStep( private val executorPodNamePrefix = kubernetesConf.resourceNamePrefix - private val driverAddress = if (kubernetesConf.get(KUBERNETES_EXECUTOR_USE_DRIVER_POD_IP)) { + private val driverAddress = if (kubernetesConf.get(KUBERNETES_EXECUTOR_USE_DRIVER_POD_IP) && + !Utils.isAnyLocalAddress(kubernetesConf.get(DRIVER_BIND_ADDRESS))) { kubernetesConf.get(DRIVER_BIND_ADDRESS) } else { kubernetesConf.get(DRIVER_HOST_ADDRESS) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStep.scala index 39dfec95176c2..00680c3e319cc 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStep.scala @@ -29,9 +29,11 @@ import org.apache.spark.deploy.k8s.{KubernetesConf, SparkPod} import org.apache.spark.deploy.k8s.Config._ import org.apache.spark.deploy.k8s.Constants._ import org.apache.spark.deploy.k8s.KubernetesUtils.buildPodWithServiceAccount +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{CONFIG, CONFIGS, PREFIX, SERVICE_ACCOUNT_NAME, VALUE} private[spark] class DriverKubernetesCredentialsFeatureStep(kubernetesConf: KubernetesConf) - extends KubernetesFeatureConfigStep { + extends KubernetesFeatureConfigStep with Logging { private val maybeMountedOAuthTokenFile = kubernetesConf.getOption( s"$KUBERNETES_AUTH_DRIVER_MOUNTED_CONF_PREFIX.$OAUTH_TOKEN_FILE_CONF_SUFFIX") @@ -59,10 +61,16 @@ private[spark] class DriverKubernetesCredentialsFeatureStep(kubernetesConf: Kube s"$KUBERNETES_AUTH_DRIVER_CONF_PREFIX.$CLIENT_CERT_FILE_CONF_SUFFIX", "Driver client cert file") - private val shouldMountSecret = oauthTokenBase64.isDefined || - caCertDataBase64.isDefined || - clientKeyDataBase64.isDefined || - clientCertDataBase64.isDefined + private val submittedCredentialConfs = Seq( + OAUTH_TOKEN_CONF_SUFFIX -> oauthTokenBase64, + CA_CERT_FILE_CONF_SUFFIX -> caCertDataBase64, + CLIENT_KEY_FILE_CONF_SUFFIX -> clientKeyDataBase64, + CLIENT_CERT_FILE_CONF_SUFFIX -> clientCertDataBase64) + .collect { case (suffix, credential) if credential.isDefined => + s"$KUBERNETES_AUTH_DRIVER_CONF_PREFIX.$suffix" + } + + private val shouldMountSecret = submittedCredentialConfs.nonEmpty private val driverCredentialsSecretName = s"${kubernetesConf.resourceNamePrefix}-kubernetes-credentials" @@ -71,6 +79,28 @@ private[spark] class DriverKubernetesCredentialsFeatureStep(kubernetesConf: Kube if (!shouldMountSecret) { pod.copy(pod = buildPodWithServiceAccount(driverServiceAccount, pod).getOrElse(pod.pod)) } else { + // The credentials secret takes precedence over the driver service account: this branch never + // applies the account, so warn that the pod keeps whatever its spec names, or the namespace + // default. Stay quiet when the spec already names the same account. Both spec fields are + // read, `serviceAccountName` winning, matching Kubernetes' SetDefaults_PodSpec: a pod + // template is deserialized with no API-server defaulting, so it can leave either one null. + val podSpec = Option(pod.pod.getSpec) + val podServiceAccount = podSpec.flatMap(s => Option(s.getServiceAccountName)) + .filter(_.nonEmpty) + .orElse(podSpec.flatMap(s => Option(s.getServiceAccount)).filter(_.nonEmpty)) + driverServiceAccount.filterNot(podServiceAccount.contains).foreach { account => + val keptAccount = podServiceAccount + .map(name => log"the pod keeps ${MDC(SERVICE_ACCOUNT_NAME, name)}, named by its spec") + .getOrElse(log"the pod falls back to the namespace's default account") + logWarning(log"Not applying " + + log"${MDC(CONFIG, KUBERNETES_DRIVER_SERVICE_ACCOUNT_NAME.key)}=${MDC(VALUE, account)} " + + log"to the driver pod, because the driver credentials given by " + + log"${MDC(CONFIGS, submittedCredentialConfs.mkString(", "))} take precedence: " + + keptAccount + log". To have Spark apply that configuration anyway, put the credentials " + + log"inside the driver pod and point the " + + log"${MDC(PREFIX, KUBERNETES_AUTH_DRIVER_MOUNTED_CONF_PREFIX)}.* configurations at " + + log"them instead, which does not mount a secret.") + } val driverPodWithMountedKubernetesCredentials = new PodBuilder(pod.pod) .editOrNewSpec() diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStep.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStep.scala index b52ad732dfba4..c80948343b23c 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStep.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStep.scala @@ -24,7 +24,6 @@ import io.fabric8.kubernetes.api.model._ import org.apache.spark.deploy.k8s.{KubernetesConf, SparkPod} import org.apache.spark.deploy.k8s.Config._ -import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils.randomize private[spark] class LocalDirsFeatureStep( @@ -47,12 +46,11 @@ private[spark] class LocalDirsFeatureStep( // exist in the image. // We could make utils.getConfiguredLocalDirs opinionated about Kubernetes, as it is already // a bit opinionated about YARN. - val resolvedLocalDirs = Option(conf.sparkConf.getenv("SPARK_LOCAL_DIRS")) + val resolvedLocalDirs = randomize(Option(conf.sparkConf.getenv("SPARK_LOCAL_DIRS")) .orElse(conf.getOption("spark.local.dir")) .getOrElse(defaultLocalDir) - .split(",") - randomize(resolvedLocalDirs) - localDirs = resolvedLocalDirs.toImmutableArraySeq + .split(",")) + localDirs = resolvedLocalDirs localDirVolumes = resolvedLocalDirs .zipWithIndex .map { case (_, index) => @@ -62,7 +60,7 @@ private[spark] class LocalDirsFeatureStep( .withMedium(if (useLocalDirTmpFs) "Memory" else null) .endEmptyDir() .build() - }.toImmutableArraySeq + } localDirVolumeMounts = localDirVolumes .zip(resolvedLocalDirs) diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOps.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOps.scala index bd8e0f97132dd..217b236af6733 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOps.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOps.scala @@ -50,7 +50,11 @@ private class KillApplication extends K8sSubmitOp { (implicit client: KubernetesClient): Unit = { val podToDelete = getPod(namespace, pName) - if (Option(podToDelete).isDefined) { + // `getPod` returns a request handle, which is never null; only resolving it reports whether + // the pod exists. Without the `get()` the check below is always true, so a name that is not + // in the cluster would issue a delete that the API server answers with a swallowed 404 and + // report nothing to the user. `ListStatus.executeOnPod` resolves it the same way. + if (Option(podToDelete.get()).isDefined) { getGracePeriod(sparkConf) match { case Some(period) => podToDelete.withGracePeriod(period).delete() case _ => podToDelete.delete() diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala index 84cfd0d72b3bb..b914c3054d209 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManager.scala @@ -228,8 +228,10 @@ private[spark] class ExecutorPodsLifecycleManager( .inNamespace(namespace) .withName(updatedPod.getMetadata.getName) - if (podToDelete.get() != null && - podToDelete.get.getMetadata.getDeletionTimestamp == null) { + // Fetch once: the pod can be removed between two `get` calls, making the second + // return null and NPE while dereferencing its metadata. + val fetchedPod = podToDelete.get() + if (fetchedPod != null && fetchedPod.getMetadata.getDeletionTimestamp == null) { podToDelete.delete() } } else if (!inactivatedPods.contains(execId) && !isPodInactive(updatedPod)) { diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala index 0784b82a85de2..27a0b320cbf9a 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackend.scala @@ -111,15 +111,18 @@ private[spark] class KubernetesClusterSchedulerBackend( super.start() // Must be called before setting the executors podAllocator.start(applicationId(), this) + // SPARK-38794: Create the executor ConfigMap before requesting executors. Executor + // allocation is asynchronous (background thread pool), so requesting executors first + // can race with ConfigMap creation, causing transient "configmap ... not found" mounts. + if (!conf.get(KUBERNETES_EXECUTOR_DISABLE_CONFIGMAP)) { + setUpExecutorConfigMap(podAllocator.driverPod) + } val defaultProfile = scheduler.sc.resourceProfileManager.defaultResourceProfile val initExecs = Map(defaultProfile -> initialExecutors) podAllocator.setTotalExpectedExecutors(initExecs) lifecycleManager.start(this) watchEvents.start(applicationId()) pollEvents.start(applicationId()) - if (!conf.get(KUBERNETES_EXECUTOR_DISABLE_CONFIGMAP)) { - setUpExecutorConfigMap(podAllocator.driverPod) - } } override def stop(): Unit = { @@ -192,6 +195,14 @@ private[spark] class KubernetesClusterSchedulerBackend( Future.successful(true) } + // The Deployment/StatefulSet allocators (and unknown custom ones) scale their controller + // down on a zero requirement, which deletes running executor pods after the termination + // grace period instead of letting them finish their tasks, so holding is supported only + // with the direct allocator. + private[spark] override def supportsExecutorHold: Boolean = { + conf.get(KUBERNETES_ALLOCATION_PODS_ALLOCATOR) == "direct" + } + override def sufficientResourcesRegistered(): Boolean = { totalRegisteredExecutors.get() >= minRegisteredExecutors } diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala index e44d7e29ef606..385bc82a35e21 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala @@ -128,6 +128,17 @@ private[spark] object KubernetesExecutorBackend extends Logging { val env = SparkEnv.createExecutorEnv(driverConf, execId, arguments.bindAddress, arguments.hostname, arguments.cores, cfg.ioEncryptionKey, isLocal = false) + // Apply initial user credentials to the executor credential store. + // Uses unconditional set() rather than updateIfNewer() because the store is guaranteed + // null at this point (executor startup, before any RPC or task is received). + // Note: there is a narrow window where a renewal broadcast (vN+1) could arrive between + // the SparkAppConfig reply (vN) and executor registration in executorDataMap, leaving + // this executor on vN until vN+2. The TaskDescription path covers this case since + // every dispatched task carries the latest credentials. + cfg.userCredentials.foreach { case (version, credentials) => + env.userCredentials.set(VersionedCredentials(version, credentials)) + } + val backend = backendCreateFn(env.rpcEnv, arguments, env, cfg.resourceProfile, execId) env.rpcEnv.setupEndpoint("Executor", backend) arguments.workerUrl.foreach { url => diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleExecutorComponents.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleExecutorComponents.scala index cbe215c3f218b..3d3b10bcecbaa 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleExecutorComponents.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleExecutorComponents.scala @@ -71,6 +71,42 @@ class KubernetesLocalDiskShuffleExecutorComponents(sparkConf: SparkConf) } object KubernetesLocalDiskShuffleExecutorComponents extends Logging { + /** + * `File.listFiles` returns null when the directory cannot be read (an IO error, a permission + * denial, or the directory being removed during the walk). Skip such a directory with a warning + * instead of failing the whole recovery, so one unreadable entry cannot cost us every other + * recoverable file. This is the single-level counterpart of `SparkFileUtils.recursiveList`, and + * unlike `JavaUtils.listFilesSafely` it does not throw on an unlistable directory. + */ + private def listEntriesOrEmpty(dir: File): Array[File] = { + val entries = dir.listFiles() + if (entries != null) { + entries + } else { + logWarning(log"Failed to list ${MDC(LogKeys.FILE_ABSOLUTE_PATH, dir.getAbsolutePath)}; " + + log"skipping it during shuffle data recovery.") + Array.empty[File] + } + } + + /** + * Walks two levels up from a local directory to reach the volume root shared by all executor + * generations. Returns None for a path with fewer than three components, where `getParent` + * yields null. + */ + private def volumeRootOf(localDir: String): Option[File] = { + val root = Option(new File(localDir).getParentFile).flatMap(p => Option(p.getParentFile)) + if (root.isEmpty) { + logWarning(log"Cannot locate the volume root of local directory " + + log"${MDC(LogKeys.PATH, localDir)}; skipping it during shuffle data recovery. Recovery " + + log"requires a local directory nested at least two levels below the volume root, such as " + + log"/data/spark-x/executor-y. Set the mount path of the spark-local-dir-* volume, or " + + log"${MDC(LogKeys.CONFIG, "spark.local.dir")} when no such volume is mounted, " + + log"accordingly.") + } + root + } + /** * This tries to recover shuffle data of dead executors' local dirs if exists. * Since the executors are already dead, we cannot use `getHostLocalDirs`. @@ -79,18 +115,18 @@ object KubernetesLocalDiskShuffleExecutorComponents extends Logging { def recoverDiskStore(conf: SparkConf, bm: BlockManager): Unit = { // Find All files val (checksumFiles, files) = Utils.getConfiguredLocalDirs(conf) - .filter(_ != null) - .map(s => new File(new File(new File(s).getParent).getParent)) + .filter(s => s != null && s.nonEmpty) + .flatMap(volumeRootOf) .flatMap { dir => - val oldDirs = dir.listFiles().filter { f => + val oldDirs = listEntriesOrEmpty(dir).filter { f => f.isDirectory && f.getName.startsWith("spark-") } - val files = oldDirs - .flatMap(_.listFiles).filter(_.isDirectory) // executor-xxx - .flatMap(_.listFiles).filter(_.isDirectory) // blockmgr-xxx - .flatMap(_.listFiles).filter(_.isDirectory) // 00 - .flatMap(_.listFiles) - if (files != null) files.toImmutableArraySeq else Seq.empty + oldDirs + .flatMap(listEntriesOrEmpty).filter(_.isDirectory) // executor-xxx + .flatMap(listEntriesOrEmpty).filter(_.isDirectory) // blockmgr-xxx + .flatMap(listEntriesOrEmpty).filter(_.isDirectory) // 00 + .flatMap(listEntriesOrEmpty) + .toImmutableArraySeq } .partition(_.getName.contains(".checksum")) val (indexFiles, dataFiles) = files.partition(_.getName.endsWith(".index")) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStepSuite.scala index 70e2239205aef..fce22a28c6d0f 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicDriverFeatureStepSuite.scala @@ -475,6 +475,44 @@ class BasicDriverFeatureStepSuite extends SparkFunSuite { assert(amountAndFormat(limits("memory")) === "5500Mi") } + test("SPARK-58202: containerPort entries are skipped when the corresponding Spark port is 0") { + val sparkConf = new SparkConf() + .set(CONTAINER_IMAGE, "spark-driver:latest") + .set(KUBERNETES_DRIVER_POD_NAME, "spark-driver-pod") + .set(DRIVER_PORT, 0) + .set(DRIVER_BLOCK_MANAGER_PORT, 0) + .set(UI_PORT, 0) + .set(CONNECT_GRPC_BINDING_PORT, "0") + val kubernetesConf: KubernetesDriverConf = KubernetesTestConf.createDriverConf( + sparkConf = sparkConf, + environment = DRIVER_ENVS, + annotations = DRIVER_ANNOTATIONS) + + val featureStep = new BasicDriverFeatureStep(kubernetesConf) + val configuredPod = featureStep.configurePod(SparkPod.initialPod()) + assert(configuredPod.container.getPorts.isEmpty) + } + + test("SPARK-58202: containerPort entries include only ports with non-zero values") { + val sparkConf = new SparkConf() + .set(CONTAINER_IMAGE, "spark-driver:latest") + .set(KUBERNETES_DRIVER_POD_NAME, "spark-driver-pod") + .set(DRIVER_PORT, 9000) + .set(DRIVER_BLOCK_MANAGER_PORT, 0) + .set(UI_PORT, 4040) + .set(CONNECT_GRPC_BINDING_PORT, "0") + val kubernetesConf: KubernetesDriverConf = KubernetesTestConf.createDriverConf( + sparkConf = sparkConf, + environment = DRIVER_ENVS, + annotations = DRIVER_ANNOTATIONS) + + val featureStep = new BasicDriverFeatureStep(kubernetesConf) + val configuredPod = featureStep.configurePod(SparkPod.initialPod()) + val portsByName = configuredPod.container.getPorts.asScala + .map(cp => cp.getName -> cp.getContainerPort).toMap + assert(portsByName === Map(DRIVER_PORT_NAME -> 9000, UI_PORT_NAME -> 4040)) + } + def containerPort(name: String, portNumber: Int): ContainerPort = new ContainerPortBuilder() diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStepSuite.scala index ba5aa35985dae..70d5f63398927 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/BasicExecutorFeatureStepSuite.scala @@ -352,11 +352,19 @@ class BasicExecutorFeatureStepSuite extends SparkFunSuite with BeforeAndAfter { ENV_EXECUTOR_ATTRIBUTE_EXECUTOR_ID -> KubernetesTestConf.EXECUTOR_ID)) } - test("SPARK-53944: Support spark.kubernetes.executor.useDriverPodIP") { - Seq((false, "localhost"), (true, "bindAddress")).foreach { - case (flag, address) => + test("SPARK-53944, SPARK-58748: Support spark.kubernetes.executor.useDriverPodIP") { + Seq( + (false, "10.138.148.230", "localhost"), + (true, "10.138.148.230", "10.138.148.230"), + (true, "bindAddress", "bindAddress"), + (true, "2001:DB8:0:0::BEEF", "[2001:db8::beef]"), + (true, "0.0.0.0", "localhost"), + (true, "::", "localhost"), + (true, "[::]", "localhost"), + (true, "0:0:0:0:0:0:0:0", "localhost")).foreach { + case (flag, bindAddress, address) => val conf = baseConf.clone() - .set(DRIVER_BIND_ADDRESS, "bindAddress") + .set(DRIVER_BIND_ADDRESS, bindAddress) .set(KUBERNETES_EXECUTOR_USE_DRIVER_POD_IP, flag) val kconf = KubernetesTestConf.createExecutorConf(sparkConf = conf) val step = new BasicExecutorFeatureStep(kconf, new SecurityManager(conf), defaultProfile) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala index 33982e9298887..eff0f2fb77a78 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/DriverKubernetesCredentialsFeatureStepSuite.scala @@ -23,7 +23,8 @@ import java.util.Base64 import scala.jdk.CollectionConverters._ -import io.fabric8.kubernetes.api.model.Secret +import io.fabric8.kubernetes.api.model.{PodBuilder, Secret} +import org.apache.logging.log4j.Level import org.apache.spark.{SparkConf, SparkFunSuite} import org.apache.spark.deploy.k8s.{KubernetesTestConf, SparkPod} @@ -35,6 +36,8 @@ class DriverKubernetesCredentialsFeatureStepSuite extends SparkFunSuite { private val credentialsTempDirectory = Utils.createTempDir() private val BASE_DRIVER_POD = SparkPod.initialPod() + private val SERVICE_ACCOUNT_CONF = KUBERNETES_DRIVER_SERVICE_ACCOUNT_NAME.key + private val STEP_LOGGER = classOf[DriverKubernetesCredentialsFeatureStep].getName test("Don't set any credentials") { val kubernetesConf = KubernetesTestConf.createDriverConf() @@ -127,6 +130,138 @@ class DriverKubernetesCredentialsFeatureStepSuite extends SparkFunSuite { assert(driverContainerVolumeMount.head.getMountPath === DRIVER_CREDENTIALS_SECRETS_BASE_DIR) } + test("SPARK-58872: warn when driver credentials drop the driver service account") { + val caCertConf = s"$KUBERNETES_AUTH_DRIVER_CONF_PREFIX.$CA_CERT_FILE_CONF_SUFFIX" + val caCertFile = writeCredentials("sa-ca.pem", "ca-cert") + val stepUnderTest = stepWith( + SERVICE_ACCOUNT_CONF -> "spark", caCertConf -> caCertFile.getAbsolutePath) + val logAppender = new LogAppender + val configuredPod = withLogAppenderReturning(logAppender) { + stepUnderTest.configurePod(BASE_DRIVER_POD) + } + // The documented behavior the warning describes: the credentials win, so the account is + // never applied, and with no account on the spec the pod falls back to the namespace default. + assert(configuredPod.pod.getSpec.getServiceAccount === null) + assert(configuredPod.pod.getSpec.getServiceAccountName === null) + val warnings = warningsFrom(logAppender) + val named = warnings.filter(w => w.contains(SERVICE_ACCOUNT_CONF) && w.contains(caCertConf)) + assert(named.size === 1, s"expected one warning naming both $SERVICE_ACCOUNT_CONF and " + + s"$caCertConf, got: $warnings") + assert(named.head.contains("namespace's default"), + s"warning does not say what the pod falls back to: ${named.head}") + // Only the credentials actually submitted are named, so a message that lists all four fails. + Seq(OAUTH_TOKEN_CONF_SUFFIX, CLIENT_KEY_FILE_CONF_SUFFIX, CLIENT_CERT_FILE_CONF_SUFFIX) + .map(suffix => s"$KUBERNETES_AUTH_DRIVER_CONF_PREFIX.$suffix") + .foreach(conf => assert(!named.head.contains(conf), + s"warning names $conf, which was not set: ${named.head}")) + // The way out of the conflict is worth spelling out, so require it stays in the message. + assert(named.head.contains(s"$KUBERNETES_AUTH_DRIVER_MOUNTED_CONF_PREFIX.*"), + s"warning does not point at the mounted configs: ${named.head}") + + // A template naming a different account does lose the configured one, so it must warn and say + // which account the pod keeps. The second case pins `serviceAccountName` beating the alias. + Seq( + podWithAccount(serviceAccountName = Some("other")), + podWithAccount(serviceAccount = Some("spark"), serviceAccountName = Some("other")) + ).foreach { otherAccountPod => + val otherAppender = new LogAppender + withLogAppender(otherAppender, loggerNames = Seq(STEP_LOGGER)) { + stepUnderTest.configurePod(otherAccountPod) + } + val otherWarnings = warningsFrom(otherAppender).filter(_.contains(SERVICE_ACCOUNT_CONF)) + assert(otherWarnings.size === 1, + s"expected one warning for a pod running as " + + s"${otherAccountPod.pod.getSpec.getServiceAccountName}/" + + s"${otherAccountPod.pod.getSpec.getServiceAccount}, got: $otherWarnings") + assert(otherWarnings.head.contains("other"), + s"warning does not name the account the pod keeps: ${otherWarnings.head}") + } + } + + test("SPARK-58872: stay quiet when the driver service account survives") { + // With the account alone there is nothing to mount, so it is applied and nothing is warned. + val saOnlyAppender = new LogAppender + val saOnlyPod = withLogAppenderReturning(saOnlyAppender) { + stepWith(SERVICE_ACCOUNT_CONF -> "spark").configurePod(BASE_DRIVER_POD) + } + assert(saOnlyPod.pod.getSpec.getServiceAccount === "spark") + assert(saOnlyPod.pod.getSpec.getServiceAccountName === "spark") + assert(warningsFrom(saOnlyAppender).isEmpty, + s"nothing was dropped, so nothing to warn about: ${warningsFrom(saOnlyAppender)}") + + // The mounted configs are the escape hatch the message and the docs point at, so setting one + // alongside the account must neither drop it nor warn. The mounted keys never feed + // `shouldMountSecret`, so this takes the first branch and never reaches the guard. + val mountedAppender = new LogAppender + val mountedPod = withLogAppenderReturning(mountedAppender) { + stepWith(SERVICE_ACCOUNT_CONF -> "spark", + s"$KUBERNETES_AUTH_DRIVER_MOUNTED_CONF_PREFIX.$CA_CERT_FILE_CONF_SUFFIX" -> "/etc/ca.pem") + .configurePod(BASE_DRIVER_POD) + } + assert(mountedPod.pod.getSpec.getServiceAccount === "spark") + assert(mountedPod.pod.getSpec.getServiceAccountName === "spark") + assert(warningsFrom(mountedAppender).isEmpty, + s"mounted configs must not drop the account: ${warningsFrom(mountedAppender)}") + + // A pod template that already names the same account loses nothing, so there is nothing to say. + // An explicitly empty `serviceAccountName` means unset, so Kubernetes copies the alias up. + val caCertFile = writeCredentials("quiet-ca.pem", "ca-cert") + val stepUnderTest = stepWith(SERVICE_ACCOUNT_CONF -> "spark", + s"$KUBERNETES_AUTH_DRIVER_CONF_PREFIX.$CA_CERT_FILE_CONF_SUFFIX" -> + caCertFile.getAbsolutePath) + Seq( + podWithAccount(serviceAccountName = Some("spark")), + podWithAccount(serviceAccount = Some("spark")), + podWithAccount(serviceAccount = Some("spark"), serviceAccountName = Some("spark")), + podWithAccount(serviceAccount = Some("spark"), serviceAccountName = Some("")) + ).foreach { templatePod => + val sameAccountAppender = new LogAppender + withLogAppender(sameAccountAppender, loggerNames = Seq(STEP_LOGGER)) { + stepUnderTest.configurePod(templatePod) + } + assert(warningsFrom(sameAccountAppender).isEmpty, + s"warned although the pod already runs as spark via " + + s"${templatePod.pod.getSpec.getServiceAccountName}/" + + s"${templatePod.pod.getSpec.getServiceAccount}: " + + s"${warningsFrom(sameAccountAppender)}") + } + } + + private def stepWith(confs: (String, String)*): DriverKubernetesCredentialsFeatureStep = { + val sparkConf = new SparkConf(false) + confs.foreach { case (k, v) => sparkConf.set(k, v) } + new DriverKubernetesCredentialsFeatureStep( + KubernetesTestConf.createDriverConf(sparkConf = sparkConf)) + } + + /** + * A driver pod whose spec names a service account, standing in for a user pod template. A + * template is parsed with no API-server defaulting, so it can name either field on its own. + */ + private def podWithAccount( + serviceAccount: Option[String] = None, + serviceAccountName: Option[String] = None): SparkPod = { + val spec = new PodBuilder(BASE_DRIVER_POD.pod).editOrNewSpec() + serviceAccount.foreach(spec.withServiceAccount(_)) + serviceAccountName.foreach(spec.withServiceAccountName(_)) + SparkPod(spec.endSpec().build(), BASE_DRIVER_POD.container) + } + + private def warningsFrom(appender: LogAppender): Seq[String] = + appender.loggingEvents + .filter(_.getLevel === Level.WARN) + .map(_.getMessage.getFormattedMessage) + .toSeq + + /** `withLogAppender` returns Unit, so carry the block's value out of it. */ + private def withLogAppenderReturning[T](appender: LogAppender)(f: => T): T = { + var result: Option[T] = None + withLogAppender(appender, loggerNames = Seq(STEP_LOGGER)) { + result = Some(f) + } + result.get + } + private def writeCredentials(credentialsFileName: String, credentialsContents: String): File = { val credentialsFile = new File(credentialsTempDirectory, credentialsFileName) Files.writeString(credentialsFile.toPath, credentialsContents) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala index 3a9561051a894..452a2cb7922c1 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/features/LocalDirsFeatureStepSuite.scala @@ -16,6 +16,8 @@ */ package org.apache.spark.deploy.k8s.features +import scala.jdk.CollectionConverters._ + import io.fabric8.kubernetes.api.model.{EnvVarBuilder, VolumeBuilder, VolumeMountBuilder} import org.apache.spark.{SparkConf, SparkFunSuite} @@ -105,6 +107,34 @@ class LocalDirsFeatureStepSuite extends SparkFunSuite { .build()) } + test("SPARK-58857: randomize the local dirs resolved from configuration") { + // SPARK-39755 added randomization to both branches of configurePod, but the emptyDir branch + // called Utils.randomize as a statement and dropped its result, so the order stayed as + // configured. Run the step repeatedly and require that not every run agrees. + val dirs = (1 to 4).map(i => s"/var/data/my-local-dir-$i") + val sparkConf = new SparkConfWithEnv(Map("SPARK_LOCAL_DIRS" -> dirs.mkString(","))) + val kubernetesConf = KubernetesTestConf.createDriverConf(sparkConf = sparkConf) + + val orders = (1 to 10).map { _ => + val configuredPod = + new LocalDirsFeatureStep(kubernetesConf, defaultLocalDir).configurePod( + SparkPod.initialPod()) + val env = configuredPod.container.getEnv.get(0) + assert(env.getName === "SPARK_LOCAL_DIRS") + // Whatever the order, the set of dirs is preserved and the mounts agree with the env var. + assert(env.getValue.split(",").sorted === dirs.sorted) + assert(configuredPod.pod.getSpec.getVolumes.size === dirs.size) + assert(configuredPod.container.getVolumeMounts.asScala.map(_.getName) === + (1 to dirs.size).map(i => s"spark-local-dir-$i")) + assert(configuredPod.container.getVolumeMounts.asScala.map(_.getMountPath).mkString(",") === + env.getValue) + env.getValue + }.toSet + + // 10 runs of 4 dirs: a false failure needs the same permutation every time, (1/24)^9. + assert(orders.size > 1, s"local dirs were never reordered across 10 runs: $orders") + } + test("Use tmpfs to back default local dir") { val sparkConf = new SparkConf(false).set(KUBERNETES_LOCAL_DIRS_TMPFS, true) val kubernetesConf = KubernetesTestConf.createDriverConf(sparkConf = sparkConf) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOpSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOpSuite.scala index 95a76b98b227b..5cb9d30d24c9f 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOpSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/submit/K8sSubmitOpSuite.scala @@ -25,7 +25,7 @@ import io.fabric8.kubernetes.api.model._ import io.fabric8.kubernetes.client.{KubernetesClient, PropagationPolicyConfigurable} import io.fabric8.kubernetes.client.dsl.{Deletable, NamespaceListVisitFromServerGetDeleteRecreateWaitApplicable, PodResource} import org.mockito.{ArgumentMatchers, Mock, MockitoAnnotations} -import org.mockito.Mockito.{times, verify, when} +import org.mockito.Mockito.{never, times, verify, when} import org.scalatest.BeforeAndAfter import org.apache.spark.{SparkConf, SparkFunSuite} @@ -38,6 +38,7 @@ import org.apache.spark.scheduler.cluster.k8s.ExecutorLifecycleTestUtils.TEST_SP class K8sSubmitOpSuite extends SparkFunSuite with BeforeAndAfter { private val driverPodName1 = "driver1" private val driverPodName2 = "driver2" + private val missingPodName = "driver3" private val driverPod1 = buildDriverPod(driverPodName1, "1") private val driverPod2 = buildDriverPod(driverPodName2, "2") private val podList = List(driverPod1, driverPod2) @@ -55,6 +56,11 @@ class K8sSubmitOpSuite extends SparkFunSuite with BeforeAndAfter { @Mock private var driverPodOperations2: PodResource = _ + // The missing pod needs a real mock whose `get` returns null. `executeOnPod` resolves the + // handle, so a bare unstubbed name would make `withName` return null and NPE instead. + @Mock + private var missingPodOperations: PodResource = _ + @Mock private var kubernetesClient: KubernetesClient = _ @@ -76,10 +82,12 @@ class K8sSubmitOpSuite extends SparkFunSuite with BeforeAndAfter { when(podOperations.inNamespace(namespace)).thenReturn(podsWithNamespace) when(podsWithNamespace.withName(driverPodName1)).thenReturn(driverPodOperations1) when(podsWithNamespace.withName(driverPodName2)).thenReturn(driverPodOperations2) + when(podsWithNamespace.withName(missingPodName)).thenReturn(missingPodOperations) when(driverPodOperations1.get).thenReturn(driverPod1) when(driverPodOperations1.delete()).thenReturn(Arrays.asList(new StatusDetails)) when(driverPodOperations2.get).thenReturn(driverPod2) when(driverPodOperations2.delete()).thenReturn(Arrays.asList(new StatusDetails)) + doReturn(null).when(missingPodOperations).get } test("List app status") { @@ -120,6 +128,17 @@ class K8sSubmitOpSuite extends SparkFunSuite with BeforeAndAfter { verify(deletable, times(1)).delete() } + test("SPARK-58725: Kill app that does not exist") { + implicit val kubeClient: KubernetesClient = kubernetesClient + val killApp = new KillApplication + killApp.printStream = err + killApp.executeOnPod(missingPodName, Option(namespace), new SparkConf()) + // scalastyle:off + verify(err).println(ArgumentMatchers.eq("Application not found.")) + // scalastyle:on + verify(missingPodOperations, never()).delete() + } + test("Kill multiple apps with glob without gracePeriod") { implicit val kubeClient: KubernetesClient = kubernetesClient val killApp = new KillApplication diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala index a0398f4f95143..00f7fb227bbd0 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/ExecutorPodsLifecycleManagerSuite.scala @@ -168,6 +168,22 @@ class ExecutorPodsLifecycleManagerSuite extends SparkFunSuite with BeforeAndAfte verify(namedExecutorPods(failedPod.getMetadata.getName), times(1)).delete() } + test("SPARK-59008: Pod deleted between the two get calls doesn't throw NPE.") { + val failedPod = failedExecutorWithoutDeletion(1) + val mockPodResource = mock(classOf[PodResource]) + namedExecutorPods.put("spark-executor-1", mockPodResource) + // The pod is present on the first lookup and gone right after, as when the API server + // removes it concurrently with the driver's deletion attempt. Re-reading it would NPE. + when(mockPodResource.get()).thenReturn(failedPod, null.asInstanceOf[Pod]) + snapshotsStore.updatePod(failedPod) + snapshotsStore.notifySubscribers() + + val msg = exitReasonMessage(1, failedPod, 1) + val expectedLossReason = ExecutorExited(1, exitCausedByApp = true, msg) + verify(schedulerBackend, times(1)).doRemoveExecutor("1", expectedLossReason) + verify(namedExecutorPods(failedPod.getMetadata.getName), times(1)).delete() + } + test("When the scheduler backend lists executor ids that aren't present in the cluster," + " remove those executors from Spark.") { when(schedulerBackend.getExecutorsWithRegistrationTs()).thenReturn(Map("1" -> 7L)) diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterManagerSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterManagerSuite.scala index 78e0942cfb82c..55eb692d8ec77 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterManagerSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterManagerSuite.scala @@ -16,6 +16,11 @@ */ package org.apache.spark.scheduler.cluster.k8s +import java.net.{InetSocketAddress, StandardProtocolFamily} +import java.nio.channels.ServerSocketChannel + +import scala.util.Using + import io.fabric8.kubernetes.client.KubernetesClient import org.mockito.{Mock, MockitoAnnotations} import org.mockito.Mockito.when @@ -28,6 +33,7 @@ import org.apache.spark.internal.config._ import org.apache.spark.scheduler.TaskSchedulerImpl import org.apache.spark.scheduler.cluster.k8s.ExecutorLifecycleTestUtils.TEST_SPARK_APP_ID import org.apache.spark.scheduler.local.LocalSchedulerBackend +import org.apache.spark.util.RpcUtils class KubernetesClusterManagerSuite extends SparkFunSuite with BeforeAndAfter { @@ -88,6 +94,32 @@ class KubernetesClusterManagerSuite extends SparkFunSuite with BeforeAndAfter { assert(backend2.applicationId() === "user-app-id") } + test("SPARK-58719: normalize IPv6 driver host when using the driver pod IP") { + assume( + Using(ServerSocketChannel.open(StandardProtocolFamily.INET6)) { channel => + channel.bind(new InetSocketAddress("::1", 0)) + }.isSuccess, + "IPv6 loopback is unavailable") + + val rawAddress = "0:0:0:0:0:0:0:1" + val conf = new SparkConf(false) + .setAppName("ipv6-driver-host") + .setMaster("k8s://test") + .set(KUBERNETES_DRIVER_MASTER_URL, "local[2]") + .set(KUBERNETES_EXECUTOR_USE_DRIVER_POD_IP, true) + .set(DRIVER_BIND_ADDRESS, rawAddress) + .set("spark.ui.enabled", "false") + + LocalSparkContext.withSpark(new SparkContext(conf)) { context => + assert(context.conf.get(DRIVER_BIND_ADDRESS) === rawAddress) + assert(context.conf.get(DRIVER_HOST_ADDRESS) === "[::1]") + assert(context.env.blockManager.blockManagerId.host === "[::1]") + val driverRef = RpcUtils.makeDriverRef( + HeartbeatReceiver.ENDPOINT_NAME, context.conf, context.env.rpcEnv) + assert(driverRef.address.host === "[::1]") + } + } + test("deployment allocator with dynamic allocation requires deletion cost") { val manager = new KubernetesClusterManager() sparkConf.set(KUBERNETES_ALLOCATION_PODS_ALLOCATOR, "deployment") diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala index cf172eb096d47..64734cc6c612a 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala @@ -27,24 +27,26 @@ import io.fabric8.kubernetes.client.dsl.PodResource import io.fabric8.kubernetes.client.dsl.base.PatchContext import org.jmock.lib.concurrent.DeterministicScheduler import org.mockito.{ArgumentCaptor, Mock, MockitoAnnotations} -import org.mockito.ArgumentMatchers.{any, eq => mockitoEq} -import org.mockito.Mockito.{atLeastOnce, mock, never, spy, verify, when} +import org.mockito.ArgumentMatchers.{any, anyBoolean, eq => mockitoEq} +import org.mockito.Mockito.{atLeastOnce, inOrder, mock, never, spy, times, verify, when} import org.scalatest.BeforeAndAfter -import org.apache.spark.{SparkConf, SparkContext, SparkEnv, SparkFunSuite} +import org.apache.spark.{SparkConf, SparkContext, SparkEnv, SparkException, SparkFunSuite} import org.apache.spark.deploy.k8s.Config._ import org.apache.spark.deploy.k8s.Constants._ import org.apache.spark.deploy.k8s.Fabric8Aliases._ +import org.apache.spark.internal.config.SCHEDULER_MAX_RETAINED_UNKNOWN_EXECUTORS import org.apache.spark.resource.{ResourceProfile, ResourceProfileManager} -import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv} -import org.apache.spark.scheduler.{ExecutorKilled, ExecutorLossReason, LiveListenerBus, TaskSchedulerImpl} -import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{RegisterExecutor, RemoveExecutor, StopDriver} +import org.apache.spark.rpc.{RpcAddress, RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv} +import org.apache.spark.scheduler.{ExecutorDecommissionInfo, ExecutorKilled, ExecutorLossReason, LiveListenerBus, TaskSchedulerImpl} +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{DecommissionExecutor, RegisterExecutor, RemoveExecutor, StopDriver, StopExecutors} import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend import org.apache.spark.scheduler.cluster.k8s.ExecutorLifecycleTestUtils.TEST_SPARK_APP_ID +import org.apache.spark.storage.{BlockManager, BlockManagerMaster} class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAndAfter { - private val schedulerExecutorService = new DeterministicScheduler() + private var schedulerExecutorService: DeterministicScheduler = _ private val sparkConf = new SparkConf(false) .set("spark.executor.instances", "3") .set("spark.app.id", TEST_SPARK_APP_ID) @@ -57,6 +59,12 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn @Mock private var env: SparkEnv = _ + @Mock + private var blockManager: BlockManager = _ + + @Mock + private var blockManagerMaster: BlockManagerMaster = _ + @Mock private var rpcEnv: RpcEnv = _ @@ -116,12 +124,17 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn private val defaultProfile = ResourceProfile.getOrCreateDefaultProfile(sparkConf) before { + schedulerExecutorService = new DeterministicScheduler() MockitoAnnotations.openMocks(this).close() when(taskScheduler.sc).thenReturn(sc) + when(taskScheduler.excludedNodes()).thenReturn(Set.empty[String]) when(sc.conf).thenReturn(sparkConf) + when(sc.listenerBus).thenReturn(listenerBus) when(sc.resourceProfileManager).thenReturn(resourceProfileManager) when(sc.env).thenReturn(env) when(env.rpcEnv).thenReturn(rpcEnv) + when(env.blockManager).thenReturn(blockManager) + when(blockManager.master).thenReturn(blockManagerMaster) driverEndpoint = ArgumentCaptor.forClass(classOf[RpcEndpoint]) when( rpcEnv.setupEndpoint( @@ -134,7 +147,15 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn when(configMapsOperations.inNamespace("default")).thenReturn(configMapsWithNamespace) when(configMapsWithNamespace.resource(any[ConfigMap]())).thenReturn(configMapResource) when(podAllocator.driverPod).thenReturn(None) - schedulerBackendUnderTest = new KubernetesClusterSchedulerBackend( + schedulerBackendUnderTest = createSchedulerBackend() + } + + after { + ResourceProfile.clearDefaultProfile() + } + + private def createSchedulerBackend(): KubernetesClusterSchedulerBackend = { + new KubernetesClusterSchedulerBackend( taskScheduler, sc, kubernetesClient, @@ -146,8 +167,35 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn pollEvents) } - after { - ResourceProfile.clearDefaultProfile() + private def registerExecutor( + backend: KubernetesClusterSchedulerBackend, + executorId: String): RpcEndpointRef = { + val executorEndpoint = mock(classOf[RpcEndpointRef]) + when(executorEndpoint.address).thenReturn(RpcAddress("localhost", 10000 + executorId.toInt)) + backend.createDriverEndpoint().receiveAndReply(mock(classOf[RpcCallContext])).apply( + RegisterExecutor(executorId, executorEndpoint, s"host-$executorId", 1, + Map.empty, Map.empty, Map.empty, defaultProfile.id)) + assert(backend.isExecutorActive(executorId)) + executorEndpoint + } + + private def withDecommissionMetadata( + f: KubernetesClusterSchedulerBackend => Unit): Unit = { + val keys = Seq(KUBERNETES_ALLOCATION_PODS_ALLOCATOR.key, + KUBERNETES_EXECUTOR_POD_DELETION_COST.key, SCHEDULER_MAX_RETAINED_UNKNOWN_EXECUTORS.key) + val originalValues = keys.map(key => key -> sparkConf.getOption(key)) + sparkConf.set(KUBERNETES_ALLOCATION_PODS_ALLOCATOR, "deployment") + sparkConf.set(KUBERNETES_EXECUTOR_POD_DELETION_COST, 7) + sparkConf.set(SCHEDULER_MAX_RETAINED_UNKNOWN_EXECUTORS, 10) + try { + // The backend captures the unknown-executor cache size when it is constructed. + f(createSchedulerBackend()) + } finally { + originalValues.foreach { + case (key, Some(value)) => sparkConf.set(key, value) + case (key, None) => sparkConf.remove(key) + } + } } test("Start all components") { @@ -160,6 +208,13 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn verify(configMapResource).create() } + test("SPARK-38794: executor ConfigMap is created before executors are requested") { + schedulerBackendUnderTest.start() + val ordered = inOrder(configMapResource, podAllocator) + ordered.verify(configMapResource).create() + ordered.verify(podAllocator).setTotalExpectedExecutors(Map(defaultProfile -> 3)) + } + test("SPARK-56684: kubernetesClient is exposed within the k8s package") { assert(schedulerBackendUnderTest.kubernetesClient eq kubernetesClient) } @@ -296,6 +351,83 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn sparkConf.remove(KUBERNETES_EXECUTOR_POD_DELETION_COST.key) } + test("SPARK-58879: idle decommission passes only selected executors to Kubernetes") { + withDecommissionMetadata { schedulerBackend => + val backend = spy[KubernetesClusterSchedulerBackend](schedulerBackend) + val idleExecutor = registerExecutor(backend, "1") + val busyExecutor = registerExecutor(backend, "2") + when(taskScheduler.isExecutorBusy("2")).thenReturn(true) + when(podsWithNamespace.withLabel(SPARK_APP_ID_LABEL, TEST_SPARK_APP_ID)) + .thenReturn(labeledPods) + when(labeledPods.withLabel(SPARK_ROLE_LABEL, SPARK_POD_EXECUTOR_ROLE)) + .thenReturn(labeledPods) + when(labeledPods.withLabelIn(SPARK_EXECUTOR_ID_LABEL, "1")).thenReturn(labeledPods) + val podResource = mock(classOf[PodResource]) + when(labeledPods.resources()) + .thenAnswer(_ => java.util.stream.Stream.of[PodResource](podResource)) + val decomInfo = ExecutorDecommissionInfo("test") + + val accepted = backend.decommissionExecutorsIfIdle( + Array("1" -> decomInfo, "2" -> decomInfo, "3" -> decomInfo, "1" -> decomInfo), + adjustTargetNumExecutors = false) + + assert(accepted === Seq("1")) + val requests = ArgumentCaptor.forClass( + classOf[Array[(String, ExecutorDecommissionInfo)]]) + verify(backend).decommissionExecutors( + requests.capture(), mockitoEq(false), mockitoEq(false)) + assert(requests.getValue.toSeq === Seq("1" -> decomInfo)) + assert(!backend.isExecutorActive("1")) + assert(backend.isExecutorActive("2")) + verify(blockManagerMaster).decommissionBlockManagers(Seq("1")) + verify(idleExecutor).send(DecommissionExecutor) + verify(busyExecutor, never()).send(DecommissionExecutor) + verify(kubernetesClient, never()).pods() + + schedulerExecutorService.runUntilIdle() + verify(labeledPods, times(2)).withLabel(SPARK_ROLE_LABEL, SPARK_POD_EXECUTOR_ROLE) + val executorIds = ArgumentCaptor.forClass(classOf[Array[String]]) + verify(labeledPods, atLeastOnce()).withLabelIn( + mockitoEq(SPARK_EXECUTOR_ID_LABEL), executorIds.capture(): _*) + assert(executorIds.getAllValues.asScala.forall(_.toSeq == Seq("1"))) + verify(labeledPods, times(2)).resources() + val patches = ArgumentCaptor.forClass(classOf[Pod]) + verify(podResource, times(2)).patch(any(classOf[PatchContext]), patches.capture()) + val appliedPods = patches.getAllValues.asScala + assert(appliedPods.exists { pod => + Option(pod.getMetadata.getLabels).exists(_.get("soLong") == "cruelWorld") + }) + assert(appliedPods.exists { pod => + Option(pod.getMetadata.getAnnotations).exists(_.get(POD_DELETION_COST) == "7") + }) + } + } + + test("SPARK-58879: rejected idle decommission has no pod updates or replay") { + withDecommissionMetadata { schedulerBackend => + val backend = spy[KubernetesClusterSchedulerBackend](schedulerBackend) + val busyExecutor = registerExecutor(backend, "1") + when(taskScheduler.isExecutorBusy("1")).thenReturn(true) + val decomInfo = ExecutorDecommissionInfo("test") + + assert(backend.decommissionExecutorsIfIdle( + Array.empty[(String, ExecutorDecommissionInfo)], + adjustTargetNumExecutors = false).isEmpty) + assert(backend.decommissionExecutorsIfIdle( + Array("1" -> decomInfo, "2" -> decomInfo, "1" -> decomInfo), + adjustTargetNumExecutors = false).isEmpty) + + val laterExecutor = registerExecutor(backend, "2") + schedulerExecutorService.runUntilIdle() + verify(backend, never()).decommissionExecutors( + any[Array[(String, ExecutorDecommissionInfo)]](), anyBoolean(), anyBoolean()) + verify(blockManagerMaster, never()).decommissionBlockManagers(any[Seq[String]]()) + verify(busyExecutor, never()).send(DecommissionExecutor) + verify(laterExecutor, never()).send(DecommissionExecutor) + verify(kubernetesClient, never()).pods() + } + } + test("SPARK-34407: CoarseGrainedSchedulerBackend.stop may throw SparkException") { schedulerBackendUnderTest.start() @@ -306,6 +438,20 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn verify(kubernetesClient).close() } + test("stopExecutors() reports a failed StopExecutors RPC as SCHEDULER_BACKEND_SHUTDOWN_FAILED") { + val rpcFailure = new RuntimeException("StopExecutors timed out") + when(driverEndpointRef.askSync[Boolean](StopExecutors)).thenThrow(rpcFailure) + val e = intercept[SparkException] { + schedulerBackendUnderTest.stopExecutors() + } + checkError( + exception = e, + condition = "SCHEDULER_BACKEND_SHUTDOWN_FAILED.EXECUTORS", + sqlState = Some("58030"), + parameters = Map.empty[String, String]) + assert(e.getCause === rpcFailure) + } + test("SPARK-34469: Ignore RegisterExecutor when SparkContext is stopped") { when(sc.isStopped).thenReturn(true) val endpoint = schedulerBackendUnderTest.createDriverEndpoint() @@ -350,4 +496,18 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn assert(id1 === id2, "applicationId() must return the same value on repeated calls") assert(id1.startsWith("spark-"), "generated app ID should have the spark- prefix") } + + test("SPARK-58915: the executors can be held only with the direct pods allocator") { + assert(schedulerBackendUnderTest.supportsExecutorHold) + Seq("statefulset", "deployment", "com.example.CustomAllocator").foreach { allocator => + sparkConf.set(KUBERNETES_ALLOCATION_PODS_ALLOCATOR, allocator) + try { + assert(!schedulerBackendUnderTest.supportsExecutorHold, + s"holding must not be supported with the $allocator allocator") + } finally { + sparkConf.remove(KUBERNETES_ALLOCATION_PODS_ALLOCATOR.key) + } + } + assert(schedulerBackendUnderTest.supportsExecutorHold) + } } diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleDataIOSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleDataIOSuite.scala index 25b2cad9ddd28..ba1777719afb7 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleDataIOSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/shuffle/KubernetesLocalDiskShuffleDataIOSuite.scala @@ -252,4 +252,56 @@ class KubernetesLocalDiskShuffleDataIOSuite extends SparkFunSuite with LocalRoot when(bm.TempFileBasedBlockStoreUpdater).thenAnswer(_ => throw new Exception()) KubernetesLocalDiskShuffleExecutorComponents.recoverDiskStore(sparkConf, bm) } + + // A sentinel thrown from the block-store updater. Reaching it proves the scan actually + // walked down to a recoverable shuffle file, so the test cannot pass vacuously. + private val sentinelMessage = "reached a shuffle file" + + private def sentinelBlockManager(): BlockManager = { + val bm = mock(classOf[BlockManager]) + when(bm.TempFileBasedBlockStoreUpdater) + .thenAnswer(_ => throw new IllegalStateException(sentinelMessage)) + bm + } + + private def createRecoverableShuffleFile(localDir: String): Unit = { + val dir = new File(localDir, "blockmgr-z/00") + Files.createDirectories(dir.toPath()) + Files.write(new File(dir, "shuffle_0_0_0.index").toPath(), Array[Byte](0, 0, 0, 0)) + } + + test("SPARK-58693: a local dir with fewer than three path components is skipped") { + val deepDir = conf.get("spark.local.dir") + "/spark-x/executor-y" + createRecoverableShuffleFile(deepDir) + // "/data" is what the Local Storage example in running-on-kubernetes.md configures, and + // walking two levels up from it yields null. The deep dir must still be recovered. + val sparkConf = conf.clone.set("spark.local.dir", s"/data,$deepDir") + + val m = intercept[IllegalStateException] { + KubernetesLocalDiskShuffleExecutorComponents + .recoverDiskStore(sparkConf, sentinelBlockManager()) + }.getMessage + assert(m.contains(sentinelMessage)) + } + + test("SPARK-58693: an unlistable directory in the scan does not abort recovery") { + val deepDir = conf.get("spark.local.dir") + "/spark-x/executor-y" + createRecoverableShuffleFile(deepDir) + // Shaped like the ext4 lost+found that sits at the root of a freshly formatted PVC: a + // directory the executor can see but not list, so listFiles() returns null. + val unlistable = new File(conf.get("spark.local.dir") + "/spark-x/lost+found") + Files.createDirectories(unlistable.toPath()) + assert(unlistable.setReadable(false, false)) + assume(!unlistable.canRead, "requires an unreadable directory; skipped when run as root") + + try { + val m = intercept[IllegalStateException] { + KubernetesLocalDiskShuffleExecutorComponents + .recoverDiskStore(conf.clone.set("spark.local.dir", deepDir), sentinelBlockManager()) + }.getMessage + assert(m.contains(sentinelMessage)) + } finally { + unlistable.setReadable(true, false) + } + } } diff --git a/resource-managers/kubernetes/integration-tests/README.md b/resource-managers/kubernetes/integration-tests/README.md index 35c0356acddfd..f1373a26f05c5 100644 --- a/resource-managers/kubernetes/integration-tests/README.md +++ b/resource-managers/kubernetes/integration-tests/README.md @@ -413,13 +413,13 @@ The suite is tagged with `YuniKornTag` which is excluded by default via ## Requirements -- Apache YuniKorn 1.8.0. +- Apache YuniKorn 1.9.0. ## Installation helm repo add yunikorn https://apache.github.io/yunikorn-release helm repo update - helm install yunikorn yunikorn/yunikorn --namespace yunikorn --version 1.8.0 \ + helm install yunikorn yunikorn/yunikorn --namespace yunikorn --version 1.9.0 \ --create-namespace --set embedAdmissionController=false ## Run tests diff --git a/resource-managers/kubernetes/integration-tests/tests/autoscale.py b/resource-managers/kubernetes/integration-tests/tests/autoscale.py index 809b698fcdd8c..ce8990ea1c072 100644 --- a/resource-managers/kubernetes/integration-tests/tests/autoscale.py +++ b/resource-managers/kubernetes/integration-tests/tests/autoscale.py @@ -20,7 +20,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": """ Usage: autoscale diff --git a/resource-managers/kubernetes/integration-tests/tests/decommissioning.py b/resource-managers/kubernetes/integration-tests/tests/decommissioning.py index 0880e8ab275b3..735e2b0e31504 100644 --- a/resource-managers/kubernetes/integration-tests/tests/decommissioning.py +++ b/resource-managers/kubernetes/integration-tests/tests/decommissioning.py @@ -20,7 +20,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": """ Usage: decommissioning diff --git a/resource-managers/kubernetes/integration-tests/tests/decommissioning_cleanup.py b/resource-managers/kubernetes/integration-tests/tests/decommissioning_cleanup.py index 8af558ee5214e..fe7e1747fb19a 100644 --- a/resource-managers/kubernetes/integration-tests/tests/decommissioning_cleanup.py +++ b/resource-managers/kubernetes/integration-tests/tests/decommissioning_cleanup.py @@ -20,7 +20,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": """ Usage: decommissioning diff --git a/resource-managers/kubernetes/integration-tests/tests/pyfiles.py b/resource-managers/kubernetes/integration-tests/tests/pyfiles.py index 73c53be482c03..6ab323282d7c7 100644 --- a/resource-managers/kubernetes/integration-tests/tests/pyfiles.py +++ b/resource-managers/kubernetes/integration-tests/tests/pyfiles.py @@ -19,7 +19,6 @@ from pyspark.sql import SparkSession from pyspark.sql.types import StringType - if __name__ == "__main__": """ Usage: pyfiles [major_python_version] diff --git a/resource-managers/kubernetes/integration-tests/tests/pyfiles_connect.py b/resource-managers/kubernetes/integration-tests/tests/pyfiles_connect.py index 4a30d89de9d6a..13c9b98c89286 100644 --- a/resource-managers/kubernetes/integration-tests/tests/pyfiles_connect.py +++ b/resource-managers/kubernetes/integration-tests/tests/pyfiles_connect.py @@ -19,7 +19,6 @@ from pyspark.sql import SparkSession from pyspark.sql.types import StringType - if __name__ == "__main__": """ Usage: pyfiles diff --git a/resource-managers/kubernetes/integration-tests/tests/python_executable_check.py b/resource-managers/kubernetes/integration-tests/tests/python_executable_check.py index 89fd2aacab1a3..9156c8a3bafd1 100644 --- a/resource-managers/kubernetes/integration-tests/tests/python_executable_check.py +++ b/resource-managers/kubernetes/integration-tests/tests/python_executable_check.py @@ -18,7 +18,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": spark = SparkSession \ .builder \ diff --git a/resource-managers/kubernetes/integration-tests/tests/worker_memory_check.py b/resource-managers/kubernetes/integration-tests/tests/worker_memory_check.py index 74559a0b54402..d1e2c6be2f8bf 100644 --- a/resource-managers/kubernetes/integration-tests/tests/worker_memory_check.py +++ b/resource-managers/kubernetes/integration-tests/tests/worker_memory_check.py @@ -20,7 +20,6 @@ from pyspark.sql import SparkSession - if __name__ == "__main__": """ Usage: worker_memory_check [Memory_in_Mi] diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala index b51c52e0e1697..efe63895efebd 100644 --- a/resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala +++ b/resource-managers/yarn/src/main/scala/org/apache/spark/scheduler/cluster/YarnSchedulerBackend.scala @@ -163,6 +163,10 @@ private[spark] abstract class YarnSchedulerBackend( yarnSchedulerEndpointRef.ask[Boolean](prepareRequestExecutors(resourceProfileToTotalExecs)) } + // The AM honors a zero executor target without killing the running executors, so the + // executors can be held gracefully. + private[spark] override def supportsExecutorHold: Boolean = true + /** * Request that the ApplicationMaster kill the specified executors. */ diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 index 1445e4015beb6..dcccba9edc041 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseLexer.g4 @@ -201,6 +201,7 @@ COMPENSATION: 'COMPENSATION'; COMPUTE: 'COMPUTE'; CONCATENATE: 'CONCATENATE'; CONDITION: 'CONDITION'; +CONDITIONAL: 'CONDITIONAL'; CONSTRAINT: 'CONSTRAINT'; CONTAINS: 'CONTAINS'; CONTINUE: 'CONTINUE'; @@ -254,8 +255,10 @@ DOUBLE: 'DOUBLE'; DROP: 'DROP'; ELSE: 'ELSE'; ELSEIF: 'ELSEIF'; +EMPTY: 'EMPTY'; END: 'END'; ENFORCED: 'ENFORCED'; +ERROR: 'ERROR'; ESCAPE: 'ESCAPE'; ESCAPED: 'ESCAPED'; EVOLUTION: 'EVOLUTION'; @@ -333,6 +336,11 @@ ITEMS: 'ITEMS'; ITERATE: 'ITERATE'; JOIN: 'JOIN'; JSON: 'JSON'; +JSON_EXISTS: 'JSON_EXISTS'; +JSON_QUERY: 'JSON_QUERY'; +JSON_TABLE: 'JSON_TABLE'; +JSON_VALUE: 'JSON_VALUE'; +KEEP: 'KEEP'; KEY: 'KEY'; KEYS: 'KEYS'; LANGUAGE: 'LANGUAGE'; @@ -391,8 +399,10 @@ NULL: 'NULL'; NULLS: 'NULLS'; NUMERIC: 'NUMERIC'; NORELY: 'NORELY'; +OBJECT: 'OBJECT'; OF: 'OF'; OFFSET: 'OFFSET'; +OMIT: 'OMIT'; ON: 'ON'; ONLY: 'ONLY'; OPEN: 'OPEN'; @@ -400,6 +410,7 @@ OPTION: 'OPTION'; OPTIONS: 'OPTIONS'; OR: 'OR'; ORDER: 'ORDER'; +ORDINALITY: 'ORDINALITY'; OUT: 'OUT'; OUTER: 'OUTER'; OUTPUTFORMAT: 'OUTPUTFORMAT'; @@ -425,6 +436,7 @@ PURGE: 'PURGE'; QUALIFY: 'QUALIFY'; QUARTER: 'QUARTER'; QUERY: 'QUERY'; +QUOTES: 'QUOTES'; RANGE: 'RANGE'; READ: 'READ'; READS: 'READS'; @@ -447,6 +459,7 @@ RESET: 'RESET'; RESPECT: 'RESPECT'; RESTRICT: 'RESTRICT'; RETURN: 'RETURN'; +RETURNING: 'RETURNING'; RETURNS: 'RETURNS'; REVOKE: 'REVOKE'; RIGHT: 'RIGHT'; @@ -534,11 +547,13 @@ TYPE: 'TYPE'; UNARCHIVE: 'UNARCHIVE'; UNBOUNDED: 'UNBOUNDED'; UNCACHE: 'UNCACHE'; +UNCONDITIONAL: 'UNCONDITIONAL'; UNIFORM: 'UNIFORM'; UNION: 'UNION'; UNIQUE: 'UNIQUE'; UNKNOWN: 'UNKNOWN'; UNLOCK: 'UNLOCK'; +UNNEST: 'UNNEST'; UNPIVOT: 'UNPIVOT'; UNSET: 'UNSET'; UNTIL: 'UNTIL'; @@ -567,6 +582,7 @@ WINDOW: 'WINDOW'; WITH: 'WITH'; WITHIN: 'WITHIN'; WITHOUT: 'WITHOUT'; +WRAPPER: 'WRAPPER'; YEAR: 'YEAR'; YEARS: 'YEARS'; ZONE: 'ZONE'; diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 index bc6c437a42e26..3003182471884 100644 --- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 +++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 @@ -1208,9 +1208,40 @@ relationPrimary | LEFT_PAREN relation RIGHT_PAREN sample? watermarkClause? tableAlias #aliasedRelation | inlineTable #inlineTableDefault2 + | unnest #unnestTable + | jsonTable #jsonTableRelation | tableFunctionCallWithTrailingClauses #tableValuedFunction ; +// ANSI SQL UNNEST of one or more arrays in the FROM clause, with an optional +// trailing ordinality column (WITH ORDINALITY). Multiple arrays are expanded in +// parallel, padded with NULLs to the length of the longest array. +unnest + : UNNEST LEFT_PAREN expression (COMMA expression)* RIGHT_PAREN + (WITH ORDINALITY)? tableAlias + ; + +// The ANSI SQL:2016 JSON_TABLE table-valued function. Because the COLUMNS clause is not a normal +// function-argument list, it has a dedicated production rather than going through +// tableFunctionCall. Only the flat (non-NESTED) subset is currently supported. +jsonTable + : JSON_TABLE LEFT_PAREN jsonExpr=expression COMMA rowPath=stringLit + COLUMNS LEFT_PAREN jsonTableColumn (COMMA jsonTableColumn)* RIGHT_PAREN + jsonTableOnErrorClause? + RIGHT_PAREN tableAlias + ; + +jsonTableColumn + : colName=errorCapturingIdentifier FOR ORDINALITY #jsonTableOrdinalityColumn + | colName=errorCapturingIdentifier dataType + EXISTS (PATH path=stringLit)? #jsonTableExistsColumn + | colName=errorCapturingIdentifier dataType (PATH path=stringLit)? #jsonTableValueColumn + ; + +jsonTableOnErrorClause + : (NULL | ERROR) ON ERROR + ; + optionsClause : WITH options=propertyList ; @@ -1431,6 +1462,18 @@ primaryExpression | ANY_VALUE LEFT_PAREN expression (IGNORE NULLS)? RIGHT_PAREN #any_value | LAST LEFT_PAREN expression (IGNORE NULLS)? RIGHT_PAREN #last | POSITION LEFT_PAREN substr=valueExpression IN str=valueExpression RIGHT_PAREN #position + | JSON_VALUE LEFT_PAREN jsonExpr=valueExpression COMMA path=stringLit + (RETURNING returning=dataType)? + (emptyBehavior=jsonValueBehavior ON EMPTY)? + (errorBehavior=jsonValueBehavior ON ERROR)? RIGHT_PAREN #jsonValue + | JSON_EXISTS LEFT_PAREN jsonExpr=valueExpression COMMA path=stringLit + (errorBehavior=jsonExistsErrorBehavior ON ERROR)? RIGHT_PAREN #jsonExists + | JSON_QUERY LEFT_PAREN jsonExpr=valueExpression COMMA path=stringLit + (RETURNING returning=dataType)? + wrapper=jsonQueryArrayWrapper? + quotes=jsonQueryQuotes? + (emptyBehavior=jsonQueryBehavior ON EMPTY)? + (errorBehavior=jsonQueryBehavior ON ERROR)? RIGHT_PAREN #jsonQuery | constant #constantDefault | ASTERISK exceptClause? #star | qualifiedName DOT ASTERISK exceptClause? #star @@ -1457,6 +1500,46 @@ primaryExpression FROM position=valueExpression (FOR length=valueExpression)? RIGHT_PAREN #overlay ; +// The behavior selected by a JSON_VALUE `... ON EMPTY` / `... ON ERROR` clause. NULL and ERROR are +// keywords; DEFAULT carries an expression evaluated in place of the missing/erroring value. +jsonValueBehavior + : NULL #jsonValueBehaviorNull + | ERROR #jsonValueBehaviorError + | DEFAULT defaultExpr=expression #jsonValueBehaviorDefault + ; + +// The behavior selected by a JSON_EXISTS `... ON ERROR` clause: the boolean (or UNKNOWN, i.e. a +// BOOLEAN NULL) to produce when the input is not a single well-formed JSON value. +jsonExistsErrorBehavior + : TRUE + | FALSE + | UNKNOWN + | ERROR + ; + +// The JSON_QUERY array-wrapper clause. `WITH [UNCONDITIONAL]` always wraps the result in `[...]`; +// `WITH CONDITIONAL` wraps only a non-array/object (scalar) result; `WITHOUT` (default) never wraps. +// The `ARRAY` word is optional, matching the SQL standard (`WITH WRAPPER` == `WITH ARRAY WRAPPER`). +jsonQueryArrayWrapper + : WITHOUT ARRAY? WRAPPER #jsonQueryWrapperWithout + | WITH wrapperType=(CONDITIONAL | UNCONDITIONAL)? ARRAY? WRAPPER #jsonQueryWrapperWith + ; + +// The JSON_QUERY quotes clause: `OMIT QUOTES` strips the surrounding quotes from a scalar string +// result; `KEEP QUOTES` (default) leaves them. +jsonQueryQuotes + : KEEP QUOTES #jsonQueryQuotesKeep + | OMIT QUOTES #jsonQueryQuotesOmit + ; + +// The behavior selected by a JSON_QUERY `... ON EMPTY` / `... ON ERROR` clause. +jsonQueryBehavior + : NULL #jsonQueryBehaviorNull + | ERROR #jsonQueryBehaviorError + | EMPTY ARRAY #jsonQueryBehaviorEmptyArray + | EMPTY OBJECT #jsonQueryBehaviorEmptyObject + ; + semiStructuredExtractionPath : jsonPathFirstPart (jsonPathParts)* ; @@ -2098,6 +2181,7 @@ ansiNonReserved | COMPUTE | CONCATENATE | CONDITION + | CONDITIONAL | CONTAINS | CONTINUE | COST @@ -2140,7 +2224,9 @@ ansiNonReserved | DOUBLE | DROP | ELSEIF + | EMPTY | ENFORCED + | ERROR | ESCAPED | EVOLUTION | EXACT @@ -2199,6 +2285,11 @@ ansiNonReserved | ITEMS | ITERATE | JSON + | JSON_EXISTS + | JSON_QUERY + | JSON_TABLE + | JSON_VALUE + | KEEP | KEY | KEYS | LANGUAGE @@ -2250,10 +2341,13 @@ ansiNonReserved | NORELY | NULLS | NUMERIC + | OBJECT | OF + | OMIT | OPEN | OPTION | OPTIONS + | ORDINALITY | OUT | OUTPUTFORMAT | OVER @@ -2276,6 +2370,7 @@ ansiNonReserved | QUALIFY | QUARTER | QUERY + | QUOTES | RANGE | READ | READS @@ -2296,6 +2391,7 @@ ansiNonReserved | RESPECT | RESTRICT | RETURN + | RETURNING | RETURNS | REVOKE | RLIKE @@ -2372,8 +2468,10 @@ ansiNonReserved | UNARCHIVE | UNBOUNDED | UNCACHE + | UNCONDITIONAL | UNIFORM | UNLOCK + | UNNEST | UNPIVOT | UNSET | UNTIL @@ -2396,6 +2494,7 @@ ansiNonReserved | WIDTH | WINDOW | WITHOUT + | WRAPPER | YEAR | YEARS | ZONE @@ -2505,6 +2604,7 @@ nonReserved | COMPUTE | CONCATENATE | CONDITION + | CONDITIONAL | CONSTRAINT | CONTAINS | CONTINUE @@ -2557,8 +2657,10 @@ nonReserved | DROP | ELSE | ELSEIF + | EMPTY | END | ENFORCED + | ERROR | ESCAPE | ESCAPED | EVOLUTION @@ -2631,6 +2733,11 @@ nonReserved | ITEMS | ITERATE | JSON + | JSON_EXISTS + | JSON_QUERY + | JSON_TABLE + | JSON_VALUE + | KEEP | KEY | KEYS | LANGUAGE @@ -2687,14 +2794,17 @@ nonReserved | NULL | NULLS | NUMERIC + | OBJECT | OF | OFFSET + | OMIT | ONLY | OPEN | OPTION | OPTIONS | OR | ORDER + | ORDINALITY | OUT | OUTER | OUTPUTFORMAT @@ -2720,6 +2830,7 @@ nonReserved | QUALIFY | QUARTER | QUERY + | QUOTES | RANGE | READ | READS @@ -2742,6 +2853,7 @@ nonReserved | RESPECT | RESTRICT | RETURN + | RETURNING | RETURNS | REVOKE | RLIKE @@ -2825,10 +2937,12 @@ nonReserved | UNARCHIVE | UNBOUNDED | UNCACHE + | UNCONDITIONAL | UNIFORM | UNIQUE | UNKNOWN | UNLOCK + | UNNEST | UNPIVOT | UNSET | UNTIL @@ -2856,6 +2970,7 @@ nonReserved | WITH | WITHIN | WITHOUT + | WRAPPER | YEAR | YEARS | ZONE diff --git a/sql/api/src/main/scala/org/apache/spark/sql/Column.scala b/sql/api/src/main/scala/org/apache/spark/sql/Column.scala index 56a9787db092e..4befebd8418b5 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/Column.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/Column.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql import scala.jdk.CollectionConverters._ -import org.apache.spark.annotation.Stable +import org.apache.spark.annotation.{DeveloperApi, Stable} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{LEFT_EXPR, RIGHT_EXPR} import org.apache.spark.sql.catalyst.parser.DataTypeParser @@ -30,11 +30,19 @@ import org.apache.spark.sql.internal.{ColumnNode, TableValuedFunctionArgument} import org.apache.spark.sql.types._ import org.apache.spark.util.ArrayImplicits._ -private[spark] object Column { +/** + * The companion object is public so that the `Column` type can be referenced as a value. This + * allows an implementation to add a `Column(expression)` factory through an extension method. All + * of its members are internal to Spark. + * + * @since 4.4.0 + */ +@DeveloperApi +object Column { - def apply(colName: String): Column = new Column(colName) + private[spark] def apply(colName: String): Column = new Column(colName) - def apply(node: => ColumnNode): Column = withOrigin(new Column(node)) + private[spark] def apply(node: => ColumnNode): Column = withOrigin(new Column(node)) /** * Invoke a function with an options map as its last argument. If there are no options, its diff --git a/sql/api/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala b/sql/api/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala index d6d3a19edc168..8d13ffea55c15 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/DataFrameStatFunctions.scala @@ -489,7 +489,7 @@ abstract class DataFrameStatFunctions { * @param seed * random seed * @return - * a `CountMinSketch` over column `colName` + * a `CountMinSketch` over column `col` * @since 2.0.0 */ def countMinSketch(col: Column, depth: Int, width: Int, seed: Int): CountMinSketch = { @@ -510,7 +510,7 @@ abstract class DataFrameStatFunctions { * @param seed * random seed * @return - * a `CountMinSketch` over column `colName` + * a `CountMinSketch` over column `col` * @since 2.0.0 */ def countMinSketch(col: Column, eps: Double, confidence: Double, seed: Int): CountMinSketch = @@ -529,6 +529,8 @@ abstract class DataFrameStatFunctions { * expected number of items which will be put into the filter. * @param fpp * expected false positive probability of the filter. + * @return + * a `BloomFilter` over column `colName` * @since 2.0.0 */ def bloomFilter(colName: String, expectedNumItems: Long, fpp: Double): BloomFilter = { @@ -544,6 +546,8 @@ abstract class DataFrameStatFunctions { * expected number of items which will be put into the filter. * @param fpp * expected false positive probability of the filter. + * @return + * a `BloomFilter` over column `col` * @since 2.0.0 */ def bloomFilter(col: Column, expectedNumItems: Long, fpp: Double): BloomFilter = { @@ -560,6 +564,8 @@ abstract class DataFrameStatFunctions { * expected number of items which will be put into the filter. * @param numBits * expected number of bits of the filter. + * @return + * a `BloomFilter` over column `colName` * @since 2.0.0 */ def bloomFilter(colName: String, expectedNumItems: Long, numBits: Long): BloomFilter = { @@ -575,6 +581,8 @@ abstract class DataFrameStatFunctions { * expected number of items which will be put into the filter. * @param numBits * expected number of bits of the filter. + * @return + * a `BloomFilter` over column `col` * @since 2.0.0 */ def bloomFilter(col: Column, expectedNumItems: Long, numBits: Long): BloomFilter = withOrigin { diff --git a/sql/api/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala b/sql/api/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala index 71daf86c94042..d9a86424f4198 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/DataFrameWriter.scala @@ -267,7 +267,9 @@ abstract class DataFrameWriter[T] { * +---+---+ * }}} * - * Because it inserts data to an existing table, format or options will be ignored. + * Because it inserts data to an existing table, the format is ignored. For data source V2 + * tables, catalog-declared table-state options are forwarded to the table load and all options + * are forwarded to the write; for V1 tables the options are ignored. * @since 1.4.0 */ def insertInto(tableName: String): Unit diff --git a/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala b/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala index 72cd1190ba40c..7c97d223a2a84 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/Encoders.scala @@ -86,14 +86,14 @@ object Encoders { * * @since 4.0.0 */ - def CHAR(length: Int): Encoder[java.lang.String] = CharEncoder(length) + def CHAR(length: Int): Encoder[java.lang.String] = CharEncoder(CharType(length)) /** * An encoder for nullable varchar type. * * @since 4.0.0 */ - def VARCHAR(length: Int): Encoder[java.lang.String] = VarcharEncoder(length) + def VARCHAR(length: Int): Encoder[java.lang.String] = VarcharEncoder(VarcharType(length)) /** * An encoder for nullable string type. diff --git a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala index a17a5dcac9e82..d6444feffc0c6 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/SparkSession.scala @@ -829,10 +829,7 @@ object SparkSession extends SparkSessionCompanion { } private[this] def lookupCompanion(name: String): SparkSessionCompanion = { - val cls = SparkClassUtils.classForName(name) - val mirror = scala.reflect.runtime.currentMirror - val module = mirror.classSymbol(cls).companion.asModule - mirror.reflectModule(module).instance.asInstanceOf[SparkSessionCompanion] + SparkClassUtils.getCompanionObject(name).asInstanceOf[SparkSessionCompanion] } /** @inheritdoc */ diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala index 6f5c4be42bbd4..60a5a597d337f 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/ScalaReflection.scala @@ -40,7 +40,7 @@ private[catalyst] object ScalaSubtypeLock * A default version of ScalaReflection that uses the runtime universe. */ object ScalaReflection extends ScalaReflection { - val universe: scala.reflect.runtime.universe.type = scala.reflect.runtime.universe + lazy val universe: scala.reflect.runtime.universe.type = scala.reflect.runtime.universe // Since we are creating a runtime mirror using the class loader of current thread, // we need to use def at here. So, every time we call mirror, it is using the // class loader of the current thread. @@ -218,7 +218,7 @@ object ScalaReflection extends ScalaReflection { } def encodeFieldNameToIdentifier(fieldName: String): String = { - TermName(fieldName).encodedName.toString + scala.reflect.NameTransformer.encode(fieldName) } /** diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala index 57c15de4f0db4..56dc001b227d5 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/AgnosticEncoder.scala @@ -237,8 +237,10 @@ object AgnosticEncoders { // Nullable leaf encoders case object NullEncoder extends LeafEncoder[java.lang.Void](NullType) case object StringEncoder extends LeafEncoder[String](StringType) - case class CharEncoder(length: Int) extends LeafEncoder[String](CharType(length)) - case class VarcharEncoder(length: Int) extends LeafEncoder[String](VarcharType(length)) + // Carry the full constrained type (length + collation), matching GeographyEncoder / + // GeometryEncoder. Reconstructing from length alone would drop a declared collation. + case class CharEncoder(dt: CharType) extends LeafEncoder[String](dt) + case class VarcharEncoder(dt: VarcharType) extends LeafEncoder[String](dt) case object BinaryEncoder extends LeafEncoder[Array[Byte]](BinaryType) case object ScalaBigIntEncoder extends LeafEncoder[BigInt](DecimalType.BigIntDecimal) case object JavaBigIntEncoder extends LeafEncoder[JBigInt](DecimalType.BigIntDecimal) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala index 5fce0d1491ba9..7627998f5a6df 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/encoders/RowEncoder.scala @@ -72,14 +72,31 @@ object RowEncoder extends DataTypeErrorsBase { encoderForDataType(schema, lenient).asInstanceOf[AgnosticEncoder[Row]] } + /** + * Builds an encoder for a schema that the engine produced, such as the result schema of a Spark + * Connect query. Whether CHAR/VARCHAR are first-class types is decided by the session that + * produced the schema, so they are always accepted here. A client cannot read that session's + * configuration, and refusing the type would make the result undecodable. + */ + private[sql] def encoderForResultSchema(schema: StructType): AgnosticEncoder[Row] = + encoderForDataType(schema, lenient = false, charVarcharFirstClassTypes = true) + .asInstanceOf[AgnosticEncoder[Row]] + private[sql] def encoderForDataType(dataType: DataType, lenient: Boolean): AgnosticEncoder[_] = + encoderForDataType(dataType, lenient, SqlApiConf.get.charVarcharFirstClassTypes) + + private def encoderForDataType( + dataType: DataType, + lenient: Boolean, + charVarcharFirstClassTypes: Boolean): AgnosticEncoder[_] = TypeApiOps(dataType) .map(_.getEncoder) - .getOrElse(encoderForDataTypeDefault(dataType, lenient)) + .getOrElse(encoderForDataTypeDefault(dataType, lenient, charVarcharFirstClassTypes)) private def encoderForDataTypeDefault( dataType: DataType, - lenient: Boolean): AgnosticEncoder[_] = + lenient: Boolean, + charVarcharFirstClassTypes: Boolean): AgnosticEncoder[_] = dataType match { case NullType => NullEncoder case BooleanType => BoxedBooleanEncoder @@ -91,10 +108,10 @@ object RowEncoder extends DataTypeErrorsBase { case DoubleType => BoxedDoubleEncoder case dt: DecimalType => JavaDecimalEncoder(dt, lenientSerialization = true) case BinaryType => BinaryEncoder - case c: CharType if SqlApiConf.get.preserveCharVarcharTypeInfo => - CharEncoder(c.length) - case v: VarcharType if SqlApiConf.get.preserveCharVarcharTypeInfo => - VarcharEncoder(v.length) + case c: CharType if charVarcharFirstClassTypes => + CharEncoder(c) + case v: VarcharType if charVarcharFirstClassTypes => + VarcharEncoder(v) case s: StringType if StringHelper.isPlainString(s) => StringEncoder case TimestampType if SqlApiConf.get.datetimeJava8ApiEnabled => InstantEncoder(lenient) case TimestampType => TimestampEncoder(lenient) @@ -107,25 +124,25 @@ object RowEncoder extends DataTypeErrorsBase { case _: VariantType => VariantEncoder case p: PythonUserDefinedType => // TODO check if this works. - encoderForDataType(p.sqlType, lenient) + encoderForDataType(p.sqlType, lenient, charVarcharFirstClassTypes) case udt: UserDefinedType[_] => UDTEncoder(udt, udt.getClass) case ArrayType(elementType, containsNull) => IterableEncoder( classTag[mutable.ArraySeq[_]], - encoderForDataType(elementType, lenient), + encoderForDataType(elementType, lenient, charVarcharFirstClassTypes), containsNull, lenientSerialization = true) case MapType(keyType, valueType, valueContainsNull) => MapEncoder( classTag[scala.collection.Map[_, _]], - encoderForDataType(keyType, lenient), - encoderForDataType(valueType, lenient), + encoderForDataType(keyType, lenient, charVarcharFirstClassTypes), + encoderForDataType(valueType, lenient, charVarcharFirstClassTypes), valueContainsNull) case StructType(fields) => AgnosticRowEncoder(fields.map { field => EncoderField( field.name, - encoderForDataType(field.dataType, lenient), + encoderForDataType(field.dataType, lenient, charVarcharFirstClassTypes), field.nullable, field.metadata) }.toImmutableArraySeq) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkCharVarcharUtils.scala b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkCharVarcharUtils.scala index 2e7806294ec0b..70684d3920425 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkCharVarcharUtils.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkCharVarcharUtils.scala @@ -31,11 +31,23 @@ trait SparkCharVarcharUtils { } /** - * Validate the given [[DataType]] to fail if it is char or varchar types or contains nested - * ones + * Logical type name for CHAR/VARCHAR, used when stamping `spark.sql.catalyst.type` on STRING + * storage (ORC/Avro). None for other types, including unbounded STRING. + */ + def charVarcharTypeName(dt: DataType): Option[String] = dt match { + case c: CharType => Some(c.typeName) + case v: VarcharType => Some(v.typeName) + case _ => None + } + + /** + * Fail if the type contains CHAR/VARCHAR unless legacy-as-string or first-class CHAR/VARCHAR is + * enabled (standard semantics or preserveCharVarcharTypeInfo). */ def failIfHasCharVarchar(dt: DataType): DataType = { - if (!SqlApiConf.get.charVarcharAsString && hasCharVarchar(dt)) { + if (SqlApiConf.get.charVarcharFirstClassTypes) { + dt + } else if (!SqlApiConf.get.charVarcharAsString && hasCharVarchar(dt)) { throw DataTypeErrors.charOrVarcharTypeAsStringUnsupportedError() } else { replaceCharVarcharWithString(dt) @@ -54,8 +66,8 @@ trait SparkCharVarcharUtils { StructType(fields.map { field => field.copy(dataType = replaceCharVarcharWithString(field.dataType)) }) - case c: CharType if !SqlApiConf.get.preserveCharVarcharTypeInfo => c.toStringType - case v: VarcharType if !SqlApiConf.get.preserveCharVarcharTypeInfo => v.toStringType + case c: CharType if !SqlApiConf.get.charVarcharFirstClassTypes => c.toStringType + case v: VarcharType if !SqlApiConf.get.charVarcharFirstClassTypes => v.toStringType case _ => dt } } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/errors/DataTypeErrors.scala b/sql/api/src/main/scala/org/apache/spark/sql/errors/DataTypeErrors.scala index 326d918f24865..f59470c1f1b33 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/errors/DataTypeErrors.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/errors/DataTypeErrors.scala @@ -139,7 +139,7 @@ private[sql] object DataTypeErrors extends DataTypeErrorsBase { def decimalCannotGreaterThanPrecisionError(scale: Int, precision: Int): Throwable = { new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1228", + errorClass = "DECIMAL_SCALE_EXCEEDS_PRECISION", messageParameters = Map("scale" -> scale.toString, "precision" -> precision.toString)) } diff --git a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala index 7cf205e75ac43..98af9876ccdab 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala @@ -138,6 +138,13 @@ private[sql] object QueryParsingErrors extends DataTypeErrorsBase { ctx) } + def duplicateJsonTableColumnError(columnName: String, ctx: JsonTableContext): Throwable = { + new ParseException( + errorClass = "INVALID_SQL_SYNTAX.DUPLICATE_JSON_TABLE_COLUMN", + messageParameters = Map("columnName" -> toSQLId(columnName)), + ctx) + } + def clausesWithPipeOperatorsUnsupportedError( ctx: QueryOrganizationContext, clauses: String): Throwable = { diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala index 5f9e02ae9b057..307a82e5a44a6 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala @@ -432,6 +432,47 @@ object functions { */ def collect_set(columnName: String): Column = collect_set(Column(columnName)) + /** + * Aggregate function: returns the distinct union of the elements of an array-typed column + * across rows. + * + * The aggregation buffer holds only the distinct elements, so its size is bounded by the + * element universe rather than by the number of input rows. Null elements are dropped by + * default (IGNORE NULLS), matching `collect_set`. With `RESPECT NULLS`, a single null element + * is kept, in which case this is equivalent to `array_distinct(flatten(collect_list(e)))`. The + * `RESPECT NULLS` clause is only available through SQL (e.g. + * `expr("collect_union(col) RESPECT NULLS")`). + * + * @param e + * The array column to collect the union of. A column of type array. + * @note + * The function is non-deterministic because the order of collected results depends on the + * order of the rows which may be non-deterministic after a shuffle. + * + * @group agg_funcs + * @since 4.3.0 + * @return + * Returns a column that evaluates to an array. + */ + def collect_union(e: Column): Column = Column.fn("collect_union", e) + + /** + * Aggregate function: returns the distinct union of the elements of an array-typed column + * across rows. + * + * @param columnName + * The name of the array column to collect the union of. A column of type array. + * @note + * The function is non-deterministic because the order of collected results depends on the + * order of the rows which may be non-deterministic after a shuffle. + * + * @group agg_funcs + * @since 4.3.0 + * @return + * Returns a column that evaluates to an array. + */ + def collect_union(columnName: String): Column = collect_union(Column(columnName)) + /** * Returns a count-min sketch of a column with the given esp, confidence and seed. The result is * an array of bytes, which can be deserialized to a `CountMinSketch` before usage. Count-min @@ -4863,6 +4904,7 @@ object functions { * }}} * * @group normal_funcs + * @since 1.5.0 */ def expr(expr: String): Column = Column(internal.SqlExpression(expr)) @@ -6027,6 +6069,53 @@ object functions { */ def round(e: Column, scale: Column): Column = Column.fn("round", e, scale) + /** + * Truncates the value of `e` toward zero to 0 decimal places. + * + * @param e + * the value to truncate. A column that evaluates to a numeric. + * @return + * Returns a column of the same type as the input, except that a decimal input may return a + * decimal of different precision and scale. + * @group math_funcs + * @since 4.4.0 + */ + def truncate(e: Column): Column = truncate(e, 0) + + /** + * Truncates the value of `e` toward zero to `scale` decimal places when `scale` is greater than + * or equal to 0, or to the left of the decimal point when `scale` is less than 0. + * + * @param e + * the value to truncate. A column that evaluates to a numeric. + * @param scale + * the number of decimal places to keep. A column that evaluates to an integral. Must be a + * constant. + * @return + * Returns a column of the same type as the input, except that a decimal input may return a + * decimal of different precision and scale. + * @group math_funcs + * @since 4.4.0 + */ + def truncate(e: Column, scale: Int): Column = Column.fn("truncate", e, lit(scale)) + + /** + * Truncates the value of `e` toward zero to `scale` decimal places when `scale` is greater than + * or equal to 0, or to the left of the decimal point when `scale` is less than 0. + * + * @param e + * the value to truncate. A column that evaluates to a numeric. + * @param scale + * the number of decimal places to keep. A column that evaluates to an integral. Must be a + * constant. + * @return + * Returns a column of the same type as the input, except that a decimal input may return a + * decimal of different precision and scale. + * @group math_funcs + * @since 4.4.0 + */ + def truncate(e: Column, scale: Column): Column = Column.fn("truncate", e, scale) + /** * Returns the value of the column `e` rounded to 0 decimal places with HALF_EVEN round mode. * @@ -6551,6 +6640,30 @@ object functions { @scala.annotation.varargs def xxhash64(cols: Column*): Column = Column.fn("xxhash64", cols: _*) + /** + * Returns a 64-bit hash value of the argument using the XXH3 algorithm. + * + * @param col + * the column to hash, which must have string or binary type. + * @group hash_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a long. + */ + def xxh3_64(col: Column): Column = Column.fn("xxh3_64", col) + + /** + * Returns a 128-bit XXH3 hash of the argument as a 32-character hex string. + * + * @param col + * the column to hash, which must have string or binary type. + * @group hash_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a string. + */ + def xxh3_128(col: Column): Column = Column.fn("xxh3_128", col) + /** * Returns null if the condition is true, and throws an exception otherwise. * @@ -7225,6 +7338,91 @@ object functions { */ def bitmap_count(col: Column): Column = Column.fn("bitmap_count", col) + /** + * Returns a bitmap that is the bitwise AND of two input bitmaps. The result is always a + * 4096-byte Spark Binary bitmap. If either input is NULL, the result is NULL. Missing bytes in + * shorter inputs are treated as zero, and inputs longer than 4096 bytes raise + * `BITMAP_INPUT_TOO_LARGE`. Both inputs must use the same bit-position mapping. If they were + * constructed by grouping `bitmap_bit_position` values by `bitmap_bucket_number`, they must + * represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + * function operates on two bitmaps from the same row; use `bitmap_*_agg` to combine bitmaps + * across rows. The representation is not a RoaringBitmap serialization. + * + * @param left + * A column that evaluates to a binary bitmap. + * @param right + * A column that evaluates to a binary bitmap. + * @group misc_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a binary bitmap. + */ + def bitmap_and(left: Column, right: Column): Column = Column.fn("bitmap_and", left, right) + + /** + * Returns a bitmap that is the bitwise OR of two input bitmaps. The result is always a + * 4096-byte Spark Binary bitmap. If either input is NULL, the result is NULL. Missing bytes in + * shorter inputs are treated as zero, and inputs longer than 4096 bytes raise + * `BITMAP_INPUT_TOO_LARGE`. Both inputs must use the same bit-position mapping. If they were + * constructed by grouping `bitmap_bit_position` values by `bitmap_bucket_number`, they must + * represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + * function operates on two bitmaps from the same row; use `bitmap_*_agg` to combine bitmaps + * across rows. The representation is not a RoaringBitmap serialization. + * + * @param left + * A column that evaluates to a binary bitmap. + * @param right + * A column that evaluates to a binary bitmap. + * @group misc_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a binary bitmap. + */ + def bitmap_or(left: Column, right: Column): Column = Column.fn("bitmap_or", left, right) + + /** + * Returns a bitmap that is the bitwise AND NOT of two input bitmaps. The result is always a + * 4096-byte Spark Binary bitmap. If either input is NULL, the result is NULL. Missing bytes in + * shorter inputs are treated as zero, and inputs longer than 4096 bytes raise + * `BITMAP_INPUT_TOO_LARGE`. Both inputs must use the same bit-position mapping. If they were + * constructed by grouping `bitmap_bit_position` values by `bitmap_bucket_number`, they must + * represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + * function operates on two bitmaps from the same row; use `bitmap_*_agg` to combine bitmaps + * across rows. The representation is not a RoaringBitmap serialization. + * + * @param left + * A column that evaluates to a binary bitmap. + * @param right + * A column that evaluates to a binary bitmap. + * @group misc_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a binary bitmap. + */ + def bitmap_andnot(left: Column, right: Column): Column = + Column.fn("bitmap_andnot", left, right) + + /** + * Returns a bitmap that is the bitwise XOR of two input bitmaps. The result is always a + * 4096-byte Spark Binary bitmap. If either input is NULL, the result is NULL. Missing bytes in + * shorter inputs are treated as zero, and inputs longer than 4096 bytes raise + * `BITMAP_INPUT_TOO_LARGE`. Both inputs must use the same bit-position mapping. If they were + * constructed by grouping `bitmap_bit_position` values by `bitmap_bucket_number`, they must + * represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + * function operates on two bitmaps from the same row; use `bitmap_*_agg` to combine bitmaps + * across rows. The representation is not a RoaringBitmap serialization. + * + * @param left + * A column that evaluates to a binary bitmap. + * @param right + * A column that evaluates to a binary bitmap. + * @group misc_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a binary bitmap. + */ + def bitmap_xor(left: Column, right: Column): Column = Column.fn("bitmap_xor", left, right) + /** * Returns a bitmap that is the bitwise OR of all of the bitmaps from the input column. The * input column should be bitmaps created from bitmap_construct_agg(). @@ -7253,6 +7451,20 @@ object functions { */ def bitmap_and_agg(col: Column): Column = Column.fn("bitmap_and_agg", col) + /** + * Returns a bitmap that is the bitwise XOR of all of the bitmaps from the input column. The + * input column should be bitmaps created from bitmap_construct_agg(). + * + * @param col + * A column containing bitmaps created by bitmap_construct_agg() and evaluating to binary + * data. + * @group agg_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a binary. + */ + def bitmap_xor_agg(col: Column): Column = Column.fn("bitmap_xor_agg", col) + ////////////////////////////////////////////////////////////////////////////////////////////// // String functions ////////////////////////////////////////////////////////////////////////////////////////////// @@ -7283,6 +7495,19 @@ object functions { */ def base64(e: Column): Column = Column.fn("base64", e) + /** + * Computes the BASE32 (RFC 4648) encoding of a binary column and returns it as a string column. + * This is the reverse of from_base32. + * + * @param e + * The target column to work on. A column that evaluates to a binary. + * @group string_funcs + * @since 4.3.0 + * @return + * Returns a column that evaluates to a string. + */ + def to_base32(e: Column): Column = Column.fn("to_base32", e) + /** * Calculates the bit length for the specified string column. * @@ -7407,6 +7632,34 @@ object functions { def try_validate_utf8(str: Column): Column = Column.fn("try_validate_utf8", str) + /** + * Returns the Unicode normalization of `str` using the given normalization `form`. Valid forms + * are 'NFC', 'NFD', 'NFKC', and 'NFKD', as defined by Unicode Standard Annex #15. The form name + * is case-insensitive. Normalization is backed by Spark's bundled ICU4J library rather than the + * JVM's own Unicode data, so results are stable across JVM vendors and versions. + * + * @param str + * the input string to normalize. + * @param form + * the normalization form: 'NFC', 'NFD', 'NFKC', or 'NFKD'. + * @group string_funcs + * @since 4.4.0 + */ + def normalize(str: Column, form: Column): Column = + Column.fn("normalize", str, form) + + /** + * Returns the Unicode normalization of `str` using the default form 'NFC'. To use a different + * form, call the two-argument overload. + * + * @param str + * the input string to normalize. + * @group string_funcs + * @since 4.4.0 + */ + def normalize(str: Column): Column = + Column.fn("normalize", str) + /** * Formats numeric column x to a format like '#,###,###.##', rounded to d decimal places with * HALF_EVEN round mode, and returns the result as a string column. @@ -8107,6 +8360,19 @@ object functions { */ def unbase64(e: Column): Column = Column.fn("unbase64", e) + /** + * Decodes a BASE32 (RFC 4648) encoded string column and returns it as a binary column. This is + * the reverse of to_base32. + * + * @param e + * target column to work on. A column that evaluates to a string. + * @group string_funcs + * @since 4.3.0 + * @return + * Returns a column that evaluates to a binary. + */ + def from_base32(e: Column): Column = Column.fn("from_base32", e) + /** * Right-pad the string column with pad to a length of len. If the string column is longer than * len, the return value is shortened to len characters. @@ -8371,6 +8637,7 @@ object functions { * @param count * number of occurrences. A column that evaluates to an integral. Must be a constant. * @group string_funcs + * @since 1.5.0 * @return * Returns a column that evaluates to a string. */ @@ -13052,6 +13319,40 @@ object functions { def slice(x: Column, start: Column, length: Column): Column = Column.fn("slice", x, start, length) + /** + * Returns the given array `x` with the last `n` elements removed. Raises an error if `n` is + * negative or greater than the number of elements in the array. + * + * @param x + * the array column to be trimmed. A column that evaluates to an array. + * @param n + * the number of elements to remove from the end of the array. Must be between 0 and the + * number of elements in the array (inclusive). + * + * @group array_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to an array. + */ + def trim_array(x: Column, n: Int): Column = trim_array(x, lit(n)) + + /** + * Returns the given array `x` with the last `n` elements removed. Raises an error if `n` is + * negative or greater than the number of elements in the array. + * + * @param x + * the array column to be trimmed. A column that evaluates to an array. + * @param n + * the number of elements to remove from the end of the array. Must be between 0 and the + * number of elements in the array (inclusive). + * + * @group array_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to an array. + */ + def trim_array(x: Column, n: Column): Column = Column.fn("trim_array", x, n) + /** * Concatenates the elements of `column` using the `delimiter`. Null values are replaced with * `nullReplacement`. @@ -14082,6 +14383,36 @@ object functions { */ def to_variant_object(col: Column): Column = Column.fn("to_variant_object", col) + /** + * Creates a variant object from the given arrays of keys and values. The keys must be non-null + * strings and the two arrays must have the same length. + * + * @param keys + * a column that evaluates to an array of string keys. + * @param values + * a column that evaluates to an array of values. + * @group variant_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a variant. + */ + def variant_from_arrays(keys: Column, values: Column): Column = + Column.fn("variant_from_arrays", keys, values) + + /** + * Creates a variant object from an array of key/value struct entries. The keys must be non-null + * strings. + * + * @param entries + * a column that evaluates to an array of key/value structs. + * @group variant_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a variant. + */ + def variant_from_entries(entries: Column): Column = + Column.fn("variant_from_entries", entries) + /** * Check if a variant value is a variant null. Returns true if and only if the input is a * variant null and false otherwise (including in the case of SQL NULL). @@ -14486,6 +14817,32 @@ object functions { def try_variant_array_append(v: Column, path: String, value: Column): Column = Column.fn("try_variant_array_append", v, lit(path), value) + /** + * Recursively removes object fields and array elements whose value is a variant null. Returns + * NULL if `v` is NULL. + * + * @param v + * a variant column. + * @group variant_funcs + * @since 4.3.0 + */ + def variant_strip_nulls(v: Column): Column = Column.fn("variant_strip_nulls", v) + + /** + * Recursively removes object fields and array elements whose value is a variant null, unless + * `includeArrays` is false, in which case null array elements are kept. Returns NULL if any + * argument is NULL. + * + * @param v + * a variant column. + * @param includeArrays + * whether null elements are also removed from arrays. + * @group variant_funcs + * @since 4.3.0 + */ + def variant_strip_nulls(v: Column, includeArrays: Boolean): Column = + Column.fn("variant_strip_nulls", v, lit(includeArrays)) + /** * Extracts a sub-variant from `v` according to `path` string, and then cast the sub-variant to * `targetType`. Returns null if the path does not exist. Throws an exception if the cast fails. @@ -14663,6 +15020,19 @@ object functions { */ def json_object_keys(e: Column): Column = Column.fn("json_object_keys", e) + /** + * Returns the type of the outermost JSON value as a string: one of 'object', 'array', 'string', + * 'number', 'boolean', or 'null'. Returns null for invalid or empty input. + * + * @param e + * the JSON string column. A column that evaluates to a string. + * @group json_funcs + * @since 4.4.0 + * @return + * Returns a column that evaluates to a string. + */ + def json_typeof(e: Column): Column = Column.fn("json_typeof", e) + // scalastyle:off line.size.limit /** * (Scala-specific) Converts a column containing a `StructType`, `ArrayType` or a `MapType` into @@ -16937,6 +17307,7 @@ object functions { * a UserDefinedFunction that can be used as an aggregating expression. * * @group udf_funcs + * @since 3.0.0 * @note * The input encoder is inferred from the input type IN. */ @@ -16975,6 +17346,7 @@ object functions { * a UserDefinedFunction that can be used as an aggregating expression * * @group udf_funcs + * @since 3.0.0 * @note * This overloading takes an explicit input encoder, to support UDAF declarations in Java. */ @@ -17484,6 +17856,32 @@ object functions { */ def unwrap_udt(column: Column): Column = Column.internalFn("unwrap_udt", column) + /** + * Wrap a column as a user-defined type. + * @param column + * the column to wrap. The column data type must match the UDT's underlying SQL type. + * @param udt + * the target user-defined type. + * @group udf_funcs + * @since 4.4.0 + */ + def wrap_udt(column: Column, udt: UserDefinedType[_]): Column = { + wrap_udt(column, lit(udt.json)) + } + + /** + * Wrap a column as a user-defined type. + * @param column + * the column to wrap. The column data type must match the UDT's underlying SQL type. + * @param udt + * the target user-defined type as a constant JSON string column. + * @group udf_funcs + * @since 4.4.0 + */ + def wrap_udt(column: Column, udt: Column): Column = { + Column.internalFn("wrap_udt", column, udt) + } + // ---------------------- Vector Functions ---------------------- /** diff --git a/sql/api/src/main/scala/org/apache/spark/sql/internal/SqlApiConf.scala b/sql/api/src/main/scala/org/apache/spark/sql/internal/SqlApiConf.scala index 7538580fb234a..854545088c8b8 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/internal/SqlApiConf.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/internal/SqlApiConf.scala @@ -42,6 +42,14 @@ private[sql] trait SqlApiConf { def allowNegativeScaleOfDecimalEnabled: Boolean def charVarcharAsString: Boolean def preserveCharVarcharTypeInfo: Boolean + def charVarcharStandardSemantics: Boolean + + /** + * True when CHAR/VARCHAR may appear as first-class types in schemas and plans (either the + * legacy preserve path or SQL standard semantics). + */ + def charVarcharFirstClassTypes: Boolean = + preserveCharVarcharTypeInfo || charVarcharStandardSemantics def datetimeJava8ApiEnabled: Boolean def sessionLocalTimeZone: String def legacyTimeParserPolicy: LegacyBehaviorPolicy.Value @@ -104,6 +112,7 @@ private[sql] object DefaultSqlApiConf extends SqlApiConf { override def allowNegativeScaleOfDecimalEnabled: Boolean = false override def charVarcharAsString: Boolean = false override def preserveCharVarcharTypeInfo: Boolean = false + override def charVarcharStandardSemantics: Boolean = false override def datetimeJava8ApiEnabled: Boolean = false override def sessionLocalTimeZone: String = TimeZone.getDefault.getID override def legacyTimeParserPolicy: LegacyBehaviorPolicy.Value = LegacyBehaviorPolicy.CORRECTED diff --git a/sql/api/src/main/scala/org/apache/spark/sql/streaming/progress.scala b/sql/api/src/main/scala/org/apache/spark/sql/streaming/progress.scala index 3411e962a5712..aabb09c5741f8 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/streaming/progress.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/streaming/progress.scala @@ -42,6 +42,32 @@ import org.apache.spark.sql.streaming.SinkProgress.DEFAULT_NUM_OUTPUT_ROWS /** * Information about updates made to stateful operators in a [[StreamingQuery]] during a trigger. + * + * @param operatorName + * Name of the stateful operator this progress describes. + * @param numRowsTotal + * Number of state rows held by the operator after this trigger. + * @param numRowsUpdated + * Number of state rows updated during this trigger. + * @param allUpdatesTimeMs + * Time taken, in milliseconds, to apply all state updates in this trigger. + * @param numRowsRemoved + * Number of state rows removed during this trigger. + * @param allRemovalsTimeMs + * Time taken, in milliseconds, to remove all evicted state rows in this trigger. + * @param commitTimeMs + * Time taken, in milliseconds, to commit the state changes of this trigger. + * @param memoryUsedBytes + * Memory used, in bytes, by the operator's state store. + * @param numRowsDroppedByWatermark + * Number of input rows dropped because their event time was older than the watermark. + * @param numShufflePartitions + * Number of shuffle partitions the operator ran with. + * @param numStateStoreInstances + * Number of state store instances backing the operator. + * @param customMetrics + * Custom metrics specific to the stateful operator or state store implementation, keyed by + * metric name. */ @Evolving class StateOperatorProgress private[spark] ( @@ -172,6 +198,8 @@ class StateOperatorProgress private[spark] ( * Information about operators in the query that store state. * @param sources * detailed statistics on data being read from each of the streaming sources. + * @param sink + * detailed statistics on data being written to the sink. * @since 2.1.0 */ @Evolving @@ -354,6 +382,8 @@ class SourceProgress protected[spark] ( * @param numOutputRows * Number of rows written to the sink or -1 for Continuous Mode (temporarily) or Sink V1 (until * decommissioned). + * @param metrics + * Sink-specific metrics reported for this trigger, keyed by metric name. * @since 2.1.0 */ @Evolving diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/DataType.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/DataType.scala index dc79267e120c0..91e6eade4eae2 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/DataType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/DataType.scala @@ -285,11 +285,18 @@ object DataType { TimestampType } else if (p < TimestampLTZNanosType.MIN_PRECISION || p > TimestampLTZNanosType.MAX_PRECISION) { - // Reject out-of-range precisions before the feature-flag check so the error is always - // INVALID_TIMESTAMP_PRECISION, not FEATURE_NOT_ENABLED. + // Reject out-of-range precisions so the error is always INVALID_TIMESTAMP_PRECISION. throw DataTypeErrors.invalidTimestampPrecisionError(precision, "TIMESTAMP_LTZ") } else { - DataTypeErrors.checkTimestampNanosTypesEnabled() + // The nanos preview flag is intentionally NOT enforced here (SPARK-57835). This JSON + // path is how a persisted schema is reconstructed from the catalog (e.g. + // HiveExternalCatalog.getTable -> DataType.fromJson), so gating it would make a table + // written with the flag on completely inaccessible once it is off -- DESCRIBE, SHOW + // CREATE TABLE, and even DROP would fail. Instead we mirror TIME (also flag-gated but + // reconstructed unconditionally): metadata reads succeed, and the flag is enforced at + // analysis/execution time via TypeUtils.failUnsupportedDataType when the data is + // actually read, written, or queried. The user-facing SQL parser path + // (DataTypeAstBuilder) stays gated, so DDL like TIMESTAMP_LTZ(9) still fails fast. TimestampLTZNanosType(p) } case TIMESTAMP_NTZ_NANOS_TYPE(precision) => @@ -305,11 +312,12 @@ object DataType { TimestampNTZType } else if (p < TimestampNTZNanosType.MIN_PRECISION || p > TimestampNTZNanosType.MAX_PRECISION) { - // Reject out-of-range precisions before the feature-flag check so the error is always - // INVALID_TIMESTAMP_PRECISION, not FEATURE_NOT_ENABLED. + // Reject out-of-range precisions so the error is always INVALID_TIMESTAMP_PRECISION. throw DataTypeErrors.invalidTimestampPrecisionError(precision, "TIMESTAMP_NTZ") } else { - DataTypeErrors.checkTimestampNanosTypesEnabled() + // Not flag-gated on purpose (SPARK-57835); see the TIMESTAMP_LTZ branch above for the + // rationale (catalog restoration must be able to reconstruct persisted nanos schemas + // regardless of the preview flag). TimestampNTZNanosType(p) } case "timestamp_ltz" => TimestampType diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/Decimal.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/Decimal.scala index a41847f02c2a2..a8de8e9d908a9 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/Decimal.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/Decimal.scala @@ -407,7 +407,7 @@ final class Decimal extends Ordered[Decimal] with Serializable { lv = roundMode match { case ROUND_FLOOR => if (lv < 0) -1L else 0L case ROUND_CEILING => if (lv > 0) 1L else 0L - case ROUND_HALF_UP | ROUND_HALF_EVEN => 0L + case ROUND_HALF_UP | ROUND_HALF_EVEN | ROUND_DOWN => 0L case _ => throw DataTypeErrors.unsupportedRoundingMode(roundMode) } } else { @@ -433,6 +433,8 @@ final class Decimal extends Ordered[Decimal] with Serializable { if (doubled > pow10diff || doubled == pow10diff && lv % 2 != 0) { lv += (if (droppedDigits < 0) -1L else 1L) } + case ROUND_DOWN => + // Truncation toward zero: `lv /= pow10diff` already dropped the fractional part. case _ => throw DataTypeErrors.unsupportedRoundingMode(roundMode) } @@ -578,6 +580,7 @@ object Decimal { val ROUND_HALF_EVEN = BigDecimal.RoundingMode.HALF_EVEN val ROUND_CEILING = BigDecimal.RoundingMode.CEILING val ROUND_FLOOR = BigDecimal.RoundingMode.FLOOR + val ROUND_DOWN = BigDecimal.RoundingMode.DOWN /** Maximum number of decimal digits an Int can represent */ val MAX_INT_DIGITS = 9 diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala index 34467c258d6c5..736fbfead3d77 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/StringType.scala @@ -159,25 +159,62 @@ case object StringHelper extends PartialOrdering[StringConstraint] { def isPlainString(s: StringType): Boolean = s.constraint == NoConstraint + /** + * Strip CHAR/VARCHAR length constraints, preserving collation. + * + * Used by transforming string expressions (upper, substr, concat, ...) so their result type is + * plain STRING even when inputs are CharType/VarcharType, when standard semantics are on. + */ + def plainStringType(dt: DataType): DataType = dt match { + case c: CharType => c.toStringType + case v: VarcharType => v.toStringType + case other => other + } + + def plainStringType(s: StringType): StringType = s match { + case c: CharType => c.toStringType + case v: VarcharType => v.toStringType + case other => other + } + def isMoreConstrained(a: StringType, b: StringType): Boolean = gteq(a.constraint, b.constraint) + /** + * Least common string type: CHAR -> VARCHAR -> STRING, with length max(n, m) when the result + * remains CHAR or VARCHAR. + * + * When first-class CHAR/VARCHAR are off, always widens to unbounded STRING (legacy + * annotated-STRING path where Char/Varchar do not appear in plans). + */ def tightestCommonString(s1: StringType, s2: StringType): Option[StringType] = { if (s1.collationId != s2.collationId) { return None } - if (!SqlApiConf.get.preserveCharVarcharTypeInfo) { + if (!SqlApiConf.get.charVarcharFirstClassTypes) { return Some(StringType(s1.collationId)) } + // Carry the declared collation onto CHAR/VARCHAR results. This propagates the Option rather + // than the id: None means "not explicitly declared" and renders as char(n) instead of + // char(n) collate UTF8_BINARY, so an all-default LCT keeps printing (and comparing) as before. + // The two collation ids are already known to be equal here. + val collation = declaredCollation(s1).orElse(declaredCollation(s2)) Some((s1.constraint, s2.constraint) match { - case (FixedLength(l1), FixedLength(l2)) => CharType(l1.max(l2)) - case (MaxLength(l1), FixedLength(l2)) => VarcharType(l1.max(l2)) - case (FixedLength(l1), MaxLength(l2)) => VarcharType(l1.max(l2)) - case (MaxLength(l1), MaxLength(l2)) => VarcharType(l1.max(l2)) + case (FixedLength(l1), FixedLength(l2)) => new CharType(l1.max(l2), collation) + case (MaxLength(l1), FixedLength(l2)) => new VarcharType(l1.max(l2), collation) + case (FixedLength(l1), MaxLength(l2)) => new VarcharType(l1.max(l2), collation) + case (MaxLength(l1), MaxLength(l2)) => new VarcharType(l1.max(l2), collation) case _ => StringType(s1.collationId) }) } + /** The explicitly declared collation of a CHAR/VARCHAR type, if any. */ + private def declaredCollation(s: StringType): Option[Int] = s match { + case c: CharType => c.collation + case v: VarcharType => v.collation + case _ => None + } + def removeCollation(s: StringType): StringType = s match { case c: CharType => CharType(c.length) case v: VarcharType => VarcharType(v.length) diff --git a/sql/api/src/main/scala/org/apache/spark/sql/types/StructType.scala b/sql/api/src/main/scala/org/apache/spark/sql/types/StructType.scala index 66c20330d4275..0d0cd5d2821b3 100644 --- a/sql/api/src/main/scala/org/apache/spark/sql/types/StructType.scala +++ b/sql/api/src/main/scala/org/apache/spark/sql/types/StructType.scala @@ -617,7 +617,7 @@ object StructType extends AbstractDataType { .map { case rightField @ StructField(rightName, rightType, rightNullable, _) => try { leftField.copy( - dataType = merge(leftType, rightType), + dataType = merge(leftType, rightType, caseSensitive), nullable = leftNullable || rightNullable) } catch { case NonFatal(e) => diff --git a/sql/catalyst/pom.xml b/sql/catalyst/pom.xml index 949843af062fe..a2b38d0b72841 100644 --- a/sql/catalyst/pom.xml +++ b/sql/catalyst/pom.xml @@ -162,7 +162,7 @@ so that the tests classes of external modules can use them. The two execution profiles are necessary - first one for 'mvn package', second one for 'mvn test-compile'. Ideally, 'mvn compile' should not compile test classes and therefore should not need this. - However, a closed due to "Cannot Reproduce" Maven bug (https://issues.apache.org/jira/browse/MNG-3559) + However, a Maven bug closed as "Cannot Reproduce" (https://issues.apache.org/jira/browse/MNG-3559) causes the compilation to fail if catalyst test-jar is not generated. Hence, the second execution profile for 'mvn test-compile'. --> diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java index 543cc3b5a7271..0e9ad2d61f9c6 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ArrayOfCollatedStringsSerDe.java @@ -45,8 +45,8 @@ public ArrayOfCollatedStringsSerDe(int collationId) { } private CollatedString wrap(String original) { - String key = CollationFactory.getCollationKey( - UTF8String.fromString(original), collationId).toString(); + byte[] key = CollationFactory.getCollationKeyBytes( + UTF8String.fromString(original), collationId); return new CollatedString(key, original); } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtils.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtils.java index a505b8d7d7566..714a3c28b4769 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtils.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtils.java @@ -57,6 +57,15 @@ public static void bitmapMerge(byte[] bitmap1, byte[] bitmap2) { } } + /** Performs bitwise XOR on both bitmaps and writes the result into bitmap1. */ + public static void bitmapXorMerge(byte[] bitmap1, byte[] bitmap2) { + // For XOR operation, bytes beyond the shorter bitmap's length are unchanged + // since XOR with 0 (absent bits) is identity: X ^ 0 = X. + for (int i = 0; i < java.lang.Math.min(bitmap1.length, bitmap2.length); ++i) { + bitmap1[i] = (byte) ((bitmap1[i] & 0x0FF) ^ (bitmap2[i] & 0x0FF)); + } + } + /** Performs bitwise AND on both bitmaps and writes the result into bitmap1. */ public static void bitmapAndMerge(byte[] bitmap1, byte[] bitmap2) { int minLen = java.lang.Math.min(bitmap1.length, bitmap2.length); @@ -69,4 +78,47 @@ public static void bitmapAndMerge(byte[] bitmap1, byte[] bitmap2) { bitmap1[i] = 0; } } + + /** Performs bitwise AND on both bitmaps and returns a new fixed-size bitmap. */ + public static byte[] bitmapAnd(byte[] bitmap1, byte[] bitmap2) { + byte[] result = new byte[NUM_BYTES]; + int numBytes = java.lang.Math.min( + NUM_BYTES, java.lang.Math.min(bitmap1.length, bitmap2.length)); + for (int i = 0; i < numBytes; ++i) { + result[i] = (byte) ((bitmap1[i] & 0x0FF) & (bitmap2[i] & 0x0FF)); + } + return result; + } + + /** Performs bitwise OR on both bitmaps and returns a new fixed-size bitmap. */ + public static byte[] bitmapOr(byte[] bitmap1, byte[] bitmap2) { + byte[] longer = bitmap1.length >= bitmap2.length ? bitmap1 : bitmap2; + byte[] shorter = bitmap1.length >= bitmap2.length ? bitmap2 : bitmap1; + byte[] result = java.util.Arrays.copyOf(longer, NUM_BYTES); + bitmapMerge(result, shorter); + return result; + } + + /** Performs bitwise AND NOT on both bitmaps and returns a new fixed-size bitmap. */ + public static byte[] bitmapAndNot(byte[] bitmap1, byte[] bitmap2) { + byte[] result = java.util.Arrays.copyOf(bitmap1, NUM_BYTES); + int numBytes = java.lang.Math.min( + NUM_BYTES, java.lang.Math.min(bitmap1.length, bitmap2.length)); + for (int i = 0; i < numBytes; ++i) { + result[i] = (byte) ((result[i] & 0x0FF) & ~(bitmap2[i] & 0x0FF)); + } + return result; + } + + /** Performs bitwise XOR on both bitmaps and returns a new fixed-size bitmap. */ + public static byte[] bitmapXor(byte[] bitmap1, byte[] bitmap2) { + byte[] longer = bitmap1.length >= bitmap2.length ? bitmap1 : bitmap2; + byte[] shorter = bitmap1.length >= bitmap2.length ? bitmap2 : bitmap1; + byte[] result = java.util.Arrays.copyOf(longer, NUM_BYTES); + int numBytes = java.lang.Math.min(NUM_BYTES, shorter.length); + for (int i = 0; i < numBytes; ++i) { + result[i] = (byte) ((result[i] & 0x0FF) ^ (shorter[i] & 0x0FF)); + } + return result; + } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java index 27dc7ab5cf9fe..0c82869062fa5 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/CollatedString.java @@ -17,6 +17,8 @@ package org.apache.spark.sql.catalyst.expressions; +import java.util.Arrays; + /** * A DataSketches ItemsSketch item for non-binary collated strings (SPARK-58069). * <p> @@ -24,18 +26,21 @@ * strings (e.g. {@code 'HELLO'} and {@code 'hello'} under {@code UTF8_LCASE}) are counted as a * single item. The {@code original} field retains an actual input value to return in the result, * mirroring how {@code mode()} returns a real value rather than the normalized collation key. + * <p> + * The {@code key} is the raw collation sort-key bytes (SPARK-58096). ICU sort keys are arbitrary + * bytes, not valid UTF-8, so decoding them to a {@code String} is lossy: two collation-distinct + * values whose keys differ only within invalid-byte regions would decode to the same {@code String} + * and be incorrectly merged. Keying on the bytes directly avoids that over-merge. */ public class CollatedString { - private final String key; + private final byte[] key; private final String original; + private final int hash; - public CollatedString(String key, String original) { + public CollatedString(byte[] key, String original) { this.key = key; this.original = original; - } - - public String key() { - return key; + this.hash = Arrays.hashCode(key); } public String original() { @@ -44,7 +49,7 @@ public String original() { @Override public int hashCode() { - return key.hashCode(); + return hash; } @Override @@ -55,6 +60,6 @@ public boolean equals(Object obj) { if (!(obj instanceof CollatedString)) { return false; } - return key.equals(((CollatedString) obj).key); + return Arrays.equals(key, ((CollatedString) obj).key); } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java index 3f2c7e1fc5d41..4f7446d0cd9c7 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java @@ -17,6 +17,8 @@ package org.apache.spark.sql.catalyst.expressions; +import com.ibm.icu.text.Normalizer2; + import java.nio.ByteBuffer; import java.security.GeneralSecurityException; import java.security.SecureRandom; @@ -149,6 +151,28 @@ public static UTF8String tryValidateUTF8String(UTF8String utf8String) { else return null; } + /** + * Normalizes the given string using the given Unicode normalization form, per the + * decomposition/composition algorithm in Unicode Standard Annex #15. Uses ICU4J, the same + * library backing Spark's collation support, instead of the JDK's {@code java.text.Normalizer} + * so the result is pinned to Spark's bundled ICU4J/Unicode data (see {@code icu4j.version} in + * pom.xml) and does not vary across JVM vendors or versions. + * + * @param input the input string to normalize. + * @param form the normalization form, one of NFC, NFD, NFKC, NFKD (case-insensitive). + * @return the normalized string. + */ + public static UTF8String normalize(UTF8String input, UTF8String form) { + Normalizer2 normalizer = switch (form.toString().toUpperCase(Locale.ROOT)) { + case "NFC" -> Normalizer2.getNFCInstance(); + case "NFD" -> Normalizer2.getNFDInstance(); + case "NFKC" -> Normalizer2.getNFKCInstance(); + case "NFKD" -> Normalizer2.getNFKDInstance(); + default -> throw QueryExecutionErrors.invalidNormalizeFormError(form.toString()); + }; + return UTF8String.fromString(normalizer.normalize(input.toString())); + } + public static byte[] aesEncrypt(byte[] input, byte[] key, UTF8String mode, diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionInfo.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionInfo.java index 325462f82c69f..4fd6e543cf09b 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionInfo.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionInfo.java @@ -150,36 +150,39 @@ public ExpressionInfo( if (!note.isEmpty()) { if (!note.contains(" ") || !note.endsWith(" ")) { throw new SparkIllegalArgumentException( - "_LEGACY_ERROR_TEMP_3201", Map.of("exprName", this.name, "note", note)); + "MALFORMED_EXPRESSION_INFO.NOTE", + Map.of("fieldName", "note", "exprName", this.name, "note", note)); } this.extended += "\n Note:\n " + note.trim() + "\n"; } if (!group.isEmpty() && !validGroups.contains(group)) { throw new SparkIllegalArgumentException( - "_LEGACY_ERROR_TEMP_3202", - Map.of("exprName", this.name, + "MALFORMED_EXPRESSION_INFO.GROUP", + Map.of("fieldName", "group", "exprName", this.name, "validGroups", String.valueOf(validGroups.stream().sorted().toList()), "group", group)); } if (!source.isEmpty() && !validSources.contains(source)) { throw new SparkIllegalArgumentException( - "_LEGACY_ERROR_TEMP_3203", - Map.of("exprName", this.name, + "MALFORMED_EXPRESSION_INFO.SOURCE", + Map.of("fieldName", "source", "exprName", this.name, "validSources", String.valueOf(validSources.stream().sorted().toList()), "source", source)); } if (!since.isEmpty()) { if (Integer.parseInt(since.split("\\.")[0]) < 0) { throw new SparkIllegalArgumentException( - "_LEGACY_ERROR_TEMP_3204", Map.of("exprName", this.name, "since", since)); + "MALFORMED_EXPRESSION_INFO.SINCE", + Map.of("fieldName", "since", "exprName", this.name, "since", since)); } this.extended += "\n Since: " + since + "\n"; } if (!deprecated.isEmpty()) { if (!deprecated.contains(" ") || !deprecated.endsWith(" ")) { throw new SparkIllegalArgumentException( - "_LEGACY_ERROR_TEMP_3205", - Map.of("exprName", this.name, "deprecated", deprecated)); + "MALFORMED_EXPRESSION_INFO.DEPRECATED", + Map.of("fieldName", "deprecated", "exprName", this.name, + "deprecated", deprecated)); } this.extended += "\n Deprecated:\n " + deprecated.trim() + "\n"; } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/VectorFunctionImplUtils.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/VectorFunctionImplUtils.java index 8a3223c588fb1..c288c58376686 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/VectorFunctionImplUtils.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/VectorFunctionImplUtils.java @@ -31,6 +31,9 @@ public class VectorFunctionImplUtils { * Returns NULL if either vector contains NULL elements, has zero magnitude, or is empty. * Throws an exception if vectors have different dimensions. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. + * The dot product and the squared norms are accumulated in double precision: their magnitudes + * are quadratic in the input values, so single precision would overflow to infinity (or + * underflow to zero) for vectors whose cosine similarity is perfectly representable as a float. */ public static Float vectorCosineSimilarity(ArrayData left, ArrayData right, UTF8String funcName) { int leftLen = left.numElements(); @@ -45,9 +48,9 @@ public static Float vectorCosineSimilarity(ArrayData left, ArrayData right, UTF8 return null; } - float dotProduct = 0.0f; - float norm1Sq = 0.0f; - float norm2Sq = 0.0f; + double dotProduct = 0.0d; + double norm1Sq = 0.0d; + double norm2Sq = 0.0d; int i = 0; int simdLimit = (leftLen / 8) * 8; @@ -66,15 +69,15 @@ public static Float vectorCosineSimilarity(ArrayData left, ArrayData right, UTF8 return null; } - float a0 = left.getFloat(i), a1 = left.getFloat(i + 1); - float a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); - float a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); - float a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); + double a0 = left.getFloat(i), a1 = left.getFloat(i + 1); + double a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); + double a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); + double a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); - float b0 = right.getFloat(i), b1 = right.getFloat(i + 1); - float b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); - float b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); - float b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); + double b0 = right.getFloat(i), b1 = right.getFloat(i + 1); + double b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); + double b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); + double b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); dotProduct += a0 * b0 + a1 * b1 + a2 * b2 + a3 * b3 + a4 * b4 + a5 * b5 + a6 * b6 + a7 * b7; @@ -90,19 +93,23 @@ public static Float vectorCosineSimilarity(ArrayData left, ArrayData right, UTF8 if (left.isNullAt(i) || right.isNullAt(i)) { return null; } - float a = left.getFloat(i); - float b = right.getFloat(i); + double a = left.getFloat(i); + double b = right.getFloat(i); dotProduct += a * b; norm1Sq += a * a; norm2Sq += b * b; i++; } - float normProduct = (float) Math.sqrt(norm1Sq * norm2Sq); - if (normProduct < Float.MIN_NORMAL) { + // For vectors of finite elements, `norm1Sq * norm2Sq` cannot overflow in double precision: + // both factors are bounded by MAX_ROUNDED_ARRAY_LENGTH * Float.MAX_VALUE^2, so their product + // stays well below Double.MAX_VALUE. An element that is already infinite makes the product + // infinite and the result NaN, exactly as it did before the accumulators were widened. + double normProduct = Math.sqrt(norm1Sq * norm2Sq); + if (normProduct == 0.0d) { return null; } - return dotProduct / normProduct; + return (float) (dotProduct / normProduct); } /** @@ -111,6 +118,8 @@ public static Float vectorCosineSimilarity(ArrayData left, ArrayData right, UTF8 * Returns 0.0 for empty vectors. * Throws an exception if vectors have different dimensions. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. + * The dot product is accumulated in double precision so that intermediate terms do not + * overflow to infinity when the final result is representable as a float. */ public static Float vectorInnerProduct(ArrayData left, ArrayData right, UTF8String funcName) { int leftLen = left.numElements(); @@ -125,7 +134,7 @@ public static Float vectorInnerProduct(ArrayData left, ArrayData right, UTF8Stri return 0.0f; } - float dotProduct = 0.0f; + double dotProduct = 0.0d; int i = 0; int simdLimit = (leftLen / 8) * 8; @@ -144,15 +153,15 @@ public static Float vectorInnerProduct(ArrayData left, ArrayData right, UTF8Stri return null; } - float a0 = left.getFloat(i), a1 = left.getFloat(i + 1); - float a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); - float a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); - float a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); + double a0 = left.getFloat(i), a1 = left.getFloat(i + 1); + double a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); + double a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); + double a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); - float b0 = right.getFloat(i), b1 = right.getFloat(i + 1); - float b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); - float b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); - float b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); + double b0 = right.getFloat(i), b1 = right.getFloat(i + 1); + double b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); + double b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); + double b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); dotProduct += a0 * b0 + a1 * b1 + a2 * b2 + a3 * b3 + a4 * b4 + a5 * b5 + a6 * b6 + a7 * b7; @@ -164,13 +173,13 @@ public static Float vectorInnerProduct(ArrayData left, ArrayData right, UTF8Stri if (left.isNullAt(i) || right.isNullAt(i)) { return null; } - float a = left.getFloat(i); - float b = right.getFloat(i); + double a = left.getFloat(i); + double b = right.getFloat(i); dotProduct += a * b; i++; } - return dotProduct; + return (float) dotProduct; } /** @@ -179,6 +188,8 @@ public static Float vectorInnerProduct(ArrayData left, ArrayData right, UTF8Stri * Returns 0.0 for empty vectors. * Throws an exception if vectors have different dimensions. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. + * The sum of squares is accumulated in double precision: it is quadratic in the input values, + * so single precision would overflow to infinity for distances representable as a float. */ public static Float vectorL2Distance(ArrayData left, ArrayData right, UTF8String funcName) { int leftLen = left.numElements(); @@ -193,7 +204,7 @@ public static Float vectorL2Distance(ArrayData left, ArrayData right, UTF8String return 0.0f; } - float sumSq = 0.0f; + double sumSq = 0.0d; int i = 0; int simdLimit = (leftLen / 8) * 8; @@ -212,18 +223,18 @@ public static Float vectorL2Distance(ArrayData left, ArrayData right, UTF8String return null; } - float a0 = left.getFloat(i), a1 = left.getFloat(i + 1); - float a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); - float a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); - float a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); + double a0 = left.getFloat(i), a1 = left.getFloat(i + 1); + double a2 = left.getFloat(i + 2), a3 = left.getFloat(i + 3); + double a4 = left.getFloat(i + 4), a5 = left.getFloat(i + 5); + double a6 = left.getFloat(i + 6), a7 = left.getFloat(i + 7); - float b0 = right.getFloat(i), b1 = right.getFloat(i + 1); - float b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); - float b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); - float b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); + double b0 = right.getFloat(i), b1 = right.getFloat(i + 1); + double b2 = right.getFloat(i + 2), b3 = right.getFloat(i + 3); + double b4 = right.getFloat(i + 4), b5 = right.getFloat(i + 5); + double b6 = right.getFloat(i + 6), b7 = right.getFloat(i + 7); - float d0 = a0 - b0, d1 = a1 - b1, d2 = a2 - b2, d3 = a3 - b3; - float d4 = a4 - b4, d5 = a5 - b5, d6 = a6 - b6, d7 = a7 - b7; + double d0 = a0 - b0, d1 = a1 - b1, d2 = a2 - b2, d3 = a3 - b3; + double d4 = a4 - b4, d5 = a5 - b5, d6 = a6 - b6, d7 = a7 - b7; sumSq += d0 * d0 + d1 * d1 + d2 * d2 + d3 * d3 + d4 * d4 + d5 * d5 + d6 * d6 + d7 * d7; @@ -235,9 +246,9 @@ public static Float vectorL2Distance(ArrayData left, ArrayData right, UTF8String if (left.isNullAt(i) || right.isNullAt(i)) { return null; } - float a = left.getFloat(i); - float b = right.getFloat(i); - float diff = a - b; + double a = left.getFloat(i); + double b = right.getFloat(i); + double diff = a - b; sumSq += diff * diff; i++; } @@ -246,19 +257,19 @@ public static Float vectorL2Distance(ArrayData left, ArrayData right, UTF8String } /** - * Computes the L1 norm (Manhattan norm) of a float vector. + * Computes the L1 norm (Manhattan norm) of a float vector, in double precision. * Returns NULL if the vector contains NULL elements. * Returns 0.0 for empty vectors. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. */ - public static Float vectorL1Norm(ArrayData vec) { + public static Double vectorL1Norm(ArrayData vec) { int len = vec.numElements(); if (len == 0) { - return 0.0f; + return 0.0d; } - float sum = 0.0f; + double sum = 0.0d; int i = 0; int simdLimit = (len / 8) * 8; @@ -273,10 +284,10 @@ public static Float vectorL1Norm(ArrayData vec) { return null; } - float a0 = vec.getFloat(i), a1 = vec.getFloat(i + 1); - float a2 = vec.getFloat(i + 2), a3 = vec.getFloat(i + 3); - float a4 = vec.getFloat(i + 4), a5 = vec.getFloat(i + 5); - float a6 = vec.getFloat(i + 6), a7 = vec.getFloat(i + 7); + double a0 = vec.getFloat(i), a1 = vec.getFloat(i + 1); + double a2 = vec.getFloat(i + 2), a3 = vec.getFloat(i + 3); + double a4 = vec.getFloat(i + 4), a5 = vec.getFloat(i + 5); + double a6 = vec.getFloat(i + 6), a7 = vec.getFloat(i + 7); sum += Math.abs(a0) + Math.abs(a1) + Math.abs(a2) + Math.abs(a3) + Math.abs(a4) + Math.abs(a5) + Math.abs(a6) + Math.abs(a7); @@ -288,7 +299,7 @@ public static Float vectorL1Norm(ArrayData vec) { if (vec.isNullAt(i)) { return null; } - float a = vec.getFloat(i); + double a = vec.getFloat(i); sum += Math.abs(a); i++; } @@ -297,19 +308,19 @@ public static Float vectorL1Norm(ArrayData vec) { } /** - * Computes the L2 norm (Euclidean norm) of a float vector. + * Computes the L2 norm (Euclidean norm) of a float vector, in double precision. * Returns NULL if the vector contains NULL elements. * Returns 0.0 for empty vectors. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. */ - public static Float vectorL2Norm(ArrayData vec) { + public static Double vectorL2Norm(ArrayData vec) { int len = vec.numElements(); if (len == 0) { - return 0.0f; + return 0.0d; } - float sumSq = 0.0f; + double sumSq = 0.0d; int i = 0; int simdLimit = (len / 8) * 8; @@ -324,10 +335,10 @@ public static Float vectorL2Norm(ArrayData vec) { return null; } - float a0 = vec.getFloat(i), a1 = vec.getFloat(i + 1); - float a2 = vec.getFloat(i + 2), a3 = vec.getFloat(i + 3); - float a4 = vec.getFloat(i + 4), a5 = vec.getFloat(i + 5); - float a6 = vec.getFloat(i + 6), a7 = vec.getFloat(i + 7); + double a0 = vec.getFloat(i), a1 = vec.getFloat(i + 1); + double a2 = vec.getFloat(i + 2), a3 = vec.getFloat(i + 3); + double a4 = vec.getFloat(i + 4), a5 = vec.getFloat(i + 5); + double a6 = vec.getFloat(i + 6), a7 = vec.getFloat(i + 7); sumSq += a0 * a0 + a1 * a1 + a2 * a2 + a3 * a3 + a4 * a4 + a5 * a5 + a6 * a6 + a7 * a7; @@ -339,24 +350,24 @@ public static Float vectorL2Norm(ArrayData vec) { if (vec.isNullAt(i)) { return null; } - float a = vec.getFloat(i); + double a = vec.getFloat(i); sumSq += a * a; i++; } - return (float) Math.sqrt(sumSq); + return Math.sqrt(sumSq); } /** - * Computes the infinity norm (maximum absolute value) of a float vector. + * Computes the infinity norm (maximum absolute value) of a float vector, in double precision. * Returns NULL if the vector contains NULL elements. * Returns 0.0 for empty vectors. */ - public static Float vectorInfNorm(ArrayData vec) { + public static Double vectorInfNorm(ArrayData vec) { int len = vec.numElements(); if (len == 0) { - return 0.0f; + return 0.0d; } float maxAbs = 0.0f; @@ -370,7 +381,7 @@ public static Float vectorInfNorm(ArrayData vec) { } } - return maxAbs; + return (double) maxAbs; } /** @@ -378,15 +389,17 @@ public static Float vectorInfNorm(ArrayData vec) { * Returns NULL if the vector contains NULL elements or if the norm is zero. * Returns an empty array for empty vectors. * Uses manual loop unrolling (8 elements at a time) for speculative SIMD optimization. + * The norm is taken in double precision so that vectors whose norm is not representable as a + * float (or is only representable as a subnormal float) are still normalized correctly. */ - public static ArrayData vectorNormalizeWithNorm(ArrayData vec, float norm) { + public static ArrayData vectorNormalizeWithNorm(ArrayData vec, double norm) { int len = vec.numElements(); if (len == 0) { return vec; } - if (norm < Float.MIN_NORMAL) { + if (norm == 0.0d) { return null; } @@ -405,14 +418,14 @@ public static ArrayData vectorNormalizeWithNorm(ArrayData vec, float norm) { return null; } - result[i] = vec.getFloat(i) / norm; - result[i + 1] = vec.getFloat(i + 1) / norm; - result[i + 2] = vec.getFloat(i + 2) / norm; - result[i + 3] = vec.getFloat(i + 3) / norm; - result[i + 4] = vec.getFloat(i + 4) / norm; - result[i + 5] = vec.getFloat(i + 5) / norm; - result[i + 6] = vec.getFloat(i + 6) / norm; - result[i + 7] = vec.getFloat(i + 7) / norm; + result[i] = (float) (vec.getFloat(i) / norm); + result[i + 1] = (float) (vec.getFloat(i + 1) / norm); + result[i + 2] = (float) (vec.getFloat(i + 2) / norm); + result[i + 3] = (float) (vec.getFloat(i + 3) / norm); + result[i + 4] = (float) (vec.getFloat(i + 4) / norm); + result[i + 5] = (float) (vec.getFloat(i + 5) / norm); + result[i + 6] = (float) (vec.getFloat(i + 6) / norm); + result[i + 7] = (float) (vec.getFloat(i + 7) / norm); i += 8; } @@ -421,7 +434,7 @@ public static ArrayData vectorNormalizeWithNorm(ArrayData vec, float norm) { if (vec.isNullAt(i)) { return null; } - result[i] = vec.getFloat(i) / norm; + result[i] = (float) (vec.getFloat(i) / norm); i++; } @@ -429,13 +442,13 @@ public static ArrayData vectorNormalizeWithNorm(ArrayData vec, float norm) { } /** - * Computes the Lp norm of a float vector using the specified degree. - * Supported degrees: 1.0 (L1), 2.0 (L2), Float.POSITIVE_INFINITY (L∞). + * Computes the Lp norm of a float vector using the specified degree, in double precision. + * Supported degrees: 1.0 (L1), 2.0 (L2), Float.POSITIVE_INFINITY (infinity norm). * Returns NULL if the vector contains NULL elements. * Returns 0.0 for empty vectors. * Throws INVALID_VECTOR_NORM_DEGREE if degree is not supported. */ - public static Float vectorNorm(ArrayData vec, float degree, UTF8String funcName) { + private static Double vectorNormAsDouble(ArrayData vec, float degree, UTF8String funcName) { // exact floating point comparison for degree since this is direct user input if (degree == 1.0f) { return vectorL1Norm(vec); @@ -448,15 +461,33 @@ public static Float vectorNorm(ArrayData vec, float degree, UTF8String funcName) } } + /** + * Computes the Lp norm of a float vector using the specified degree. + * Supported degrees: 1.0 (L1), 2.0 (L2), Float.POSITIVE_INFINITY (infinity norm). + * Returns NULL if the vector contains NULL elements. + * Returns 0.0 for empty vectors. + * Throws INVALID_VECTOR_NORM_DEGREE if degree is not supported. + */ + public static Float vectorNorm(ArrayData vec, float degree, UTF8String funcName) { + Double norm = vectorNormAsDouble(vec, degree, funcName); + if (norm == null) { + return null; + } + return (float) norm.doubleValue(); + } + /** * Normalizes a float vector to unit length using the specified norm degree. - * Supported degrees: 1.0 (L1), 2.0 (L2), Float.POSITIVE_INFINITY (L∞). + * Supported degrees: 1.0 (L1), 2.0 (L2), Float.POSITIVE_INFINITY (infinity norm). * Returns NULL if the vector contains NULL elements or has zero norm. * Returns an empty array for empty vectors. * Throws INVALID_VECTOR_NORM_DEGREE if degree is not supported. */ public static ArrayData vectorNormalize(ArrayData vec, float degree, UTF8String funcName) { - Float norm = vectorNorm(vec, degree, funcName); + // The norm is kept in double precision here: rounding it to a float first would turn a norm + // that overflows (or underflows) the float range into infinity (or zero) and produce an + // all-zero (or NULL) result for a vector that is perfectly normalizable. + Double norm = vectorNormAsDouble(vec, degree, funcName); if (norm == null) { return null; } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/XXH3.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/XXH3.java new file mode 100644 index 0000000000000..fe3fe6b5dad19 --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/XXH3.java @@ -0,0 +1,572 @@ +/* + * 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.catalyst.expressions; + +import java.util.HexFormat; + +import org.apache.spark.unsafe.types.UTF8String; + +/** + * A Java port of the XXH3 64-bit and 128-bit hash functions (from the reference implementation at + * https://github.com/Cyan4973/xxHash, as specified in doc/xxhash_spec.md). The output is byte + * compatible with the reference implementation, so it matches `xxhsum` and other XXH3 tools. + * + * <p>The multiplication and mixing machinery operates on 64-bit lanes; Java's signed + * {@code long} is used as an unsigned 64-bit integer throughout ({@code >>>} for logical shift, + * {@link Long#rotateLeft}, and {@link #unsignedMultiplyHigh} for the high half of 64x64 + * products). The 128-bit hash's 1-3 byte inputs are the exception: they are first composed as + * 32-bit {@code int}s (see {@link #len1to3128}) before being widened. + */ +public final class XXH3 { + + private XXH3() {} + + private static final long PRIME32_1 = 0x9E3779B1L; + private static final long PRIME32_2 = 0x85EBCA77L; + private static final long PRIME32_3 = 0xC2B2AE3DL; + private static final long PRIME64_1 = 0x9E3779B185EBCA87L; + private static final long PRIME64_2 = 0xC2B2AE3D27D4EB4FL; + private static final long PRIME64_3 = 0x165667B19E3779F9L; + private static final long PRIME64_4 = 0x85EBCA77C2B2AE63L; + private static final long PRIME64_5 = 0x27D4EB2F165667C5L; + private static final long PRIME_MX1 = 0x165667919E3779F9L; + private static final long PRIME_MX2 = 0x9FB21C651E98DF25L; + + private static final int SECRET_SIZE = 192; + private static final int SECRET_SIZE_MIN = 136; + private static final int STRIPE_LEN = 64; + private static final int SECRET_MERGEACCS_START = 11; + private static final int SECRET_LASTACC_START = 7; + private static final int NB_STRIPES_PER_BLOCK = (SECRET_SIZE - STRIPE_LEN) / 8; + private static final int BLOCK_LEN = STRIPE_LEN * NB_STRIPES_PER_BLOCK; + + // The default 192-byte secret from the reference implementation. + private static final byte[] SECRET = HexFormat.of().parseHex( + "b8fe6c3923a44bbe7c01812cf721ad1c" + + "ded46de9839097db7240a4a4b7b3671f" + + "cb79e64eccc0e578825ad07dccff7221" + + "b8084674f743248ee03590e6813a264c" + + "3c2852bb91c300cb88d0658b1b532ea3" + + "71644897a20df94e3819ef46a9deacd8" + + "a8fa763fe39c343ff9dcbbc7c70b4f1d" + + "8a51e04bcdb45931c89f7ec9d9787364" + + "eac5ac8334d3ebc3c581a0fffa1363eb" + + "170ddd51b7f0da49d316552629d4689e" + + "2b16be587d47a1fc8ff8b8d17ad031ce" + + "45cb3a8f95160428afd7fbcabb4b407e"); + + // ---- little-endian reads / writes ---- + + private static long readLE64(byte[] data, int offset) { + return (data[offset] & 0xFFL) + | ((data[offset + 1] & 0xFFL) << 8) + | ((data[offset + 2] & 0xFFL) << 16) + | ((data[offset + 3] & 0xFFL) << 24) + | ((data[offset + 4] & 0xFFL) << 32) + | ((data[offset + 5] & 0xFFL) << 40) + | ((data[offset + 6] & 0xFFL) << 48) + | ((data[offset + 7] & 0xFFL) << 56); + } + + // Reads 4 little-endian bytes as an unsigned 32-bit value (in the low 32 bits of the result). + private static long readLE32(byte[] data, int offset) { + return (data[offset] & 0xFFL) + | ((data[offset + 1] & 0xFFL) << 8) + | ((data[offset + 2] & 0xFFL) << 16) + | ((data[offset + 3] & 0xFFL) << 24); + } + + private static void writeLE64(byte[] data, int offset, long value) { + for (int i = 0; i < 8; i++) { + data[offset + i] = (byte) (value >>> (8 * i)); + } + } + + // ---- mixing helpers ---- + + private static long xxh64Avalanche(long h) { + h ^= h >>> 33; + h *= PRIME64_2; + h ^= h >>> 29; + h *= PRIME64_3; + h ^= h >>> 32; + return h; + } + + private static long xxh3Avalanche(long h) { + h ^= h >>> 37; + h *= PRIME_MX1; + h ^= h >>> 32; + return h; + } + + // Unsigned high 64 bits of the 128-bit product a*b (Math.multiplyHigh is signed). + private static long unsignedMultiplyHigh(long a, long b) { + return Math.multiplyHigh(a, b) + ((a >> 63) & b) + ((b >> 63) & a); + } + + // Low 64 bits XOR high 64 bits of the 128-bit product a*b. + private static long mul128Fold64(long a, long b) { + return (a * b) ^ unsignedMultiplyHigh(a, b); + } + + private static long mix16B(byte[] input, int inOff, int secOff, long seed) { + long lo = readLE64(input, inOff) ^ (readLE64(SECRET, secOff) + seed); + long hi = readLE64(input, inOff + 8) ^ (readLE64(SECRET, secOff + 8) - seed); + return mul128Fold64(lo, hi); + } + + // ---- XXH3 64-bit ---- + + public static long hash64(byte[] input, long seed) { + int len = input.length; + if (len <= 16) { + if (len > 8) { + return len9to16(input, len, seed); + } else if (len >= 4) { + return len4to8(input, len, seed); + } else if (len > 0) { + return len1to3(input, len, seed); + } + return xxh64Avalanche(seed ^ readLE64(SECRET, 56) ^ readLE64(SECRET, 64)); + } else if (len <= 128) { + return len17to128(input, len, seed); + } else if (len <= 240) { + return len129to240(input, len, seed); + } + return hashLong(input, len, seed); + } + + private static long len1to3(byte[] input, int len, long seed) { + long combined = ((input[0] & 0xFFL) << 16) + | ((input[len >> 1] & 0xFFL) << 24) + | (input[len - 1] & 0xFFL) + | ((long) len << 8); + long flip = (readLE32(SECRET, 0) ^ readLE32(SECRET, 4)) + seed; + return xxh64Avalanche(combined ^ flip); + } + + private static long len4to8(byte[] input, int len, long seed) { + seed ^= ((long) Integer.reverseBytes((int) seed) & 0xFFFFFFFFL) << 32; + long in1 = readLE32(input, 0); + long in2 = readLE32(input, len - 4); + long combined = in2 | (in1 << 32); + long flip = (readLE64(SECRET, 8) ^ readLE64(SECRET, 16)) - seed; + long x = combined ^ flip; + x ^= Long.rotateLeft(x, 49) ^ Long.rotateLeft(x, 24); + x *= PRIME_MX2; + x ^= (x >>> 35) + len; + x *= PRIME_MX2; + x ^= x >>> 28; + return x; + } + + private static long len9to16(byte[] input, int len, long seed) { + long flip1 = (readLE64(SECRET, 24) ^ readLE64(SECRET, 32)) + seed; + long flip2 = (readLE64(SECRET, 40) ^ readLE64(SECRET, 48)) - seed; + long in1 = readLE64(input, 0) ^ flip1; + long in2 = readLE64(input, len - 8) ^ flip2; + long acc = len + Long.reverseBytes(in1) + in2 + mul128Fold64(in1, in2); + return xxh3Avalanche(acc); + } + + private static long len17to128(byte[] input, int len, long seed) { + long acc = (long) len * PRIME64_1; + for (int i = (len - 1) >> 5; i >= 0; i--) { + acc += mix16B(input, 16 * i, 32 * i, seed); + acc += mix16B(input, len - 16 * (i + 1), 32 * i + 16, seed); + } + return xxh3Avalanche(acc); + } + + private static long len129to240(byte[] input, int len, long seed) { + long acc = (long) len * PRIME64_1; + int nbRounds = len / 16; + for (int i = 0; i < 8; i++) { + acc += mix16B(input, 16 * i, 16 * i, seed); + } + acc = xxh3Avalanche(acc); + for (int i = 8; i < nbRounds; i++) { + acc += mix16B(input, 16 * i, 16 * (i - 8) + 3, seed); + } + acc += mix16B(input, len - 16, SECRET_SIZE_MIN - 17, seed); + return xxh3Avalanche(acc); + } + + // ---- long input (> 240 bytes) ---- + + private static byte[] customSecret(long seed) { + if (seed == 0) { + return SECRET; + } + byte[] secret = new byte[SECRET_SIZE]; + for (int i = 0; i < SECRET_SIZE / 16; i++) { + writeLE64(secret, 16 * i, readLE64(SECRET, 16 * i) + seed); + writeLE64(secret, 16 * i + 8, readLE64(SECRET, 16 * i + 8) - seed); + } + return secret; + } + + private static void accumulate512( + long[] acc, byte[] input, int inOff, byte[] secret, int secOff) { + for (int i = 0; i < 8; i++) { + long data = readLE64(input, inOff + 8 * i); + long key = data ^ readLE64(secret, secOff + 8 * i); + acc[i ^ 1] += data; + acc[i] += (key & 0xFFFFFFFFL) * (key >>> 32); + } + } + + private static void scrambleAcc(long[] acc, byte[] secret, int secOff) { + for (int i = 0; i < 8; i++) { + acc[i] ^= acc[i] >>> 47; + acc[i] ^= readLE64(secret, secOff + 8 * i); + acc[i] *= PRIME32_1; + } + } + + private static long mergeAccs(long[] acc, byte[] secret, int secOff, long start) { + long result = start; + for (int i = 0; i < 4; i++) { + long a0 = acc[2 * i] ^ readLE64(secret, secOff + 16 * i); + long a1 = acc[2 * i + 1] ^ readLE64(secret, secOff + 16 * i + 8); + result += mul128Fold64(a0, a1); + } + return xxh3Avalanche(result); + } + + private static long[] hashLongAccumulate(byte[] input, int len, byte[] secret) { + long[] acc = {PRIME32_3, PRIME64_1, PRIME64_2, PRIME64_3, PRIME64_4, PRIME32_2, PRIME64_5, + PRIME32_1}; + int nbBlocks = (len - 1) / BLOCK_LEN; + for (int n = 0; n < nbBlocks; n++) { + for (int s = 0; s < NB_STRIPES_PER_BLOCK; s++) { + accumulate512(acc, input, n * BLOCK_LEN + s * STRIPE_LEN, secret, s * 8); + } + scrambleAcc(acc, secret, SECRET_SIZE - STRIPE_LEN); + } + int nbStripes = ((len - 1) - BLOCK_LEN * nbBlocks) / STRIPE_LEN; + for (int s = 0; s < nbStripes; s++) { + accumulate512(acc, input, nbBlocks * BLOCK_LEN + s * STRIPE_LEN, secret, s * 8); + } + accumulate512( + acc, input, len - STRIPE_LEN, secret, SECRET_SIZE - STRIPE_LEN - SECRET_LASTACC_START); + return acc; + } + + private static long hashLong(byte[] input, int len, long seed) { + byte[] secret = customSecret(seed); + long[] acc = hashLongAccumulate(input, len, secret); + return mergeAccs(acc, secret, SECRET_MERGEACCS_START, (long) len * PRIME64_1); + } + + /** Hashes the input with the default seed 0. */ + public static long hash64(byte[] input) { + return hash64(input, 0L); + } + + // ---- XXH3 128-bit ---- + + private static long mult32to64(long a, long b) { + return (a & 0xFFFFFFFFL) * (b & 0xFFFFFFFFL); + } + + private static void mix32B(long[] acc, byte[] input, int in1, int in2, int secOff, long seed) { + acc[0] += mix16B(input, in1, secOff, seed); + acc[0] ^= readLE64(input, in2) + readLE64(input, in2 + 8); + acc[1] += mix16B(input, in2, secOff + 16, seed); + acc[1] ^= readLE64(input, in1) + readLE64(input, in1 + 8); + } + + private static long[] finalize128(long[] acc, int len, long seed) { + long low = xxh3Avalanche(acc[0] + acc[1]); + long high = -xxh3Avalanche( + acc[0] * PRIME64_1 + acc[1] * PRIME64_4 + ((long) len - seed) * PRIME64_2); + return new long[] {low, high}; + } + + /** Returns the XXH3 128-bit hash as {@code {low64, high64}}. */ + public static long[] hash128(byte[] input, long seed) { + int len = input.length; + if (len <= 16) { + if (len > 8) { + return len9to16128(input, len, seed); + } else if (len >= 4) { + return len4to8128(input, len, seed); + } else if (len > 0) { + return len1to3128(input, len, seed); + } + long low = xxh64Avalanche(seed ^ readLE64(SECRET, 64) ^ readLE64(SECRET, 72)); + long high = xxh64Avalanche(seed ^ readLE64(SECRET, 80) ^ readLE64(SECRET, 88)); + return new long[] {low, high}; + } else if (len <= 128) { + return len17to128128(input, len, seed); + } else if (len <= 240) { + return len129to240128(input, len, seed); + } + return hashLong128(input, len, seed); + } + + private static long[] len1to3128(byte[] input, int len, long seed) { + int c1 = input[0] & 0xFF; + int c2 = input[len >> 1] & 0xFF; + int c3 = input[len - 1] & 0xFF; + int combinedl = (c1 << 16) | (c2 << 24) | c3 | (len << 8); + int combinedh = Integer.rotateLeft(Integer.reverseBytes(combinedl), 13); + long bitflipl = (readLE32(SECRET, 0) ^ readLE32(SECRET, 4)) + seed; + long bitfliph = (readLE32(SECRET, 8) ^ readLE32(SECRET, 12)) - seed; + long low = xxh64Avalanche((combinedl & 0xFFFFFFFFL) ^ bitflipl); + long high = xxh64Avalanche((combinedh & 0xFFFFFFFFL) ^ bitfliph); + return new long[] {low, high}; + } + + private static long[] len4to8128(byte[] input, int len, long seed) { + seed ^= ((long) Integer.reverseBytes((int) seed) & 0xFFFFFFFFL) << 32; + long inputLo = readLE32(input, 0); + long inputHi = readLE32(input, len - 4); + long input64 = inputLo | (inputHi << 32); + long keyed = input64 ^ ((readLE64(SECRET, 16) ^ readLE64(SECRET, 24)) + seed); + long mul = PRIME64_1 + ((long) len << 2); + long lo = keyed * mul; + long hi = unsignedMultiplyHigh(keyed, mul); + hi += lo << 1; + lo ^= hi >>> 3; + lo ^= lo >>> 35; + lo *= PRIME_MX2; + lo ^= lo >>> 28; + return new long[] {lo, xxh3Avalanche(hi)}; + } + + private static long[] len9to16128(byte[] input, int len, long seed) { + long bitflipl = (readLE64(SECRET, 32) ^ readLE64(SECRET, 40)) - seed; + long bitfliph = (readLE64(SECRET, 48) ^ readLE64(SECRET, 56)) + seed; + long inputLo = readLE64(input, 0); + long inputHi = readLE64(input, len - 8); + long m0 = inputLo ^ inputHi ^ bitflipl; + long lo = m0 * PRIME64_1; + long hi = unsignedMultiplyHigh(m0, PRIME64_1); + lo += (long) (len - 1) << 54; + inputHi ^= bitfliph; + hi += inputHi + mult32to64(inputHi, PRIME32_2 - 1); + lo ^= Long.reverseBytes(hi); + long h2lo = lo * PRIME64_2; + long h2hi = unsignedMultiplyHigh(lo, PRIME64_2) + hi * PRIME64_2; + return new long[] {xxh3Avalanche(h2lo), xxh3Avalanche(h2hi)}; + } + + private static long[] len17to128128(byte[] input, int len, long seed) { + long[] acc = {(long) len * PRIME64_1, 0L}; + int i = (len - 1) / 32; + do { + mix32B(acc, input, 16 * i, len - 16 * (i + 1), 32 * i, seed); + } while (i-- != 0); + return finalize128(acc, len, seed); + } + + private static long[] len129to240128(byte[] input, int len, long seed) { + long[] acc = {(long) len * PRIME64_1, 0L}; + for (int i = 0; i < 4; i++) { + mix32B(acc, input, 32 * i, 32 * i + 16, 32 * i, seed); + } + acc[0] = xxh3Avalanche(acc[0]); + acc[1] = xxh3Avalanche(acc[1]); + for (int i = 4; i < (len >> 5); i++) { + mix32B(acc, input, 32 * i, 32 * i + 16, (i - 4) * 32 + 3, seed); + } + mix32B(acc, input, len - 16, len - 32, 103, -seed); + return finalize128(acc, len, seed); + } + + private static long[] hashLong128(byte[] input, int len, long seed) { + byte[] secret = customSecret(seed); + long[] acc = hashLongAccumulate(input, len, secret); + long low = mergeAccs(acc, secret, SECRET_MERGEACCS_START, (long) len * PRIME64_1); + long high = mergeAccs(acc, secret, SECRET_SIZE - STRIPE_LEN - SECRET_MERGEACCS_START, + ~((long) len * PRIME64_2)); + return new long[] {low, high}; + } + + private static final byte[] HEX_DIGITS = { + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' + }; + + // Hex-writing siblings of the length-branch methods above: same arithmetic, but they encode + // the two lanes directly into the caller's hex buffer instead of returning a long[]. This + // avoids the result-pair array for these Into variants; hashLongAccumulate (>240 bytes, + // used by hashLong128Into below) still allocates its own accumulator. hash128 above is kept + // array-returning for callers that need the raw pair. + private static void len1to3128Into(byte[] input, int len, long seed, byte[] hex) { + int c1 = input[0] & 0xFF; + int c2 = input[len >> 1] & 0xFF; + int c3 = input[len - 1] & 0xFF; + int combinedl = (c1 << 16) | (c2 << 24) | c3 | (len << 8); + int combinedh = Integer.rotateLeft(Integer.reverseBytes(combinedl), 13); + long bitflipl = (readLE32(SECRET, 0) ^ readLE32(SECRET, 4)) + seed; + long bitfliph = (readLE32(SECRET, 8) ^ readLE32(SECRET, 12)) - seed; + long low = xxh64Avalanche((combinedl & 0xFFFFFFFFL) ^ bitflipl); + long high = xxh64Avalanche((combinedh & 0xFFFFFFFFL) ^ bitfliph); + writeHex64(hex, 0, high); + writeHex64(hex, 16, low); + } + + private static void len4to8128Into(byte[] input, int len, long seed, byte[] hex) { + seed ^= ((long) Integer.reverseBytes((int) seed) & 0xFFFFFFFFL) << 32; + long inputLo = readLE32(input, 0); + long inputHi = readLE32(input, len - 4); + long input64 = inputLo | (inputHi << 32); + long keyed = input64 ^ ((readLE64(SECRET, 16) ^ readLE64(SECRET, 24)) + seed); + long mul = PRIME64_1 + ((long) len << 2); + long lo = keyed * mul; + long hi = unsignedMultiplyHigh(keyed, mul); + hi += lo << 1; + lo ^= hi >>> 3; + lo ^= lo >>> 35; + lo *= PRIME_MX2; + lo ^= lo >>> 28; + writeHex64(hex, 0, xxh3Avalanche(hi)); + writeHex64(hex, 16, lo); + } + + private static void len9to16128Into(byte[] input, int len, long seed, byte[] hex) { + long bitflipl = (readLE64(SECRET, 32) ^ readLE64(SECRET, 40)) - seed; + long bitfliph = (readLE64(SECRET, 48) ^ readLE64(SECRET, 56)) + seed; + long inputLo = readLE64(input, 0); + long inputHi = readLE64(input, len - 8); + long m0 = inputLo ^ inputHi ^ bitflipl; + long lo = m0 * PRIME64_1; + long hi = unsignedMultiplyHigh(m0, PRIME64_1); + lo += (long) (len - 1) << 54; + inputHi ^= bitfliph; + hi += inputHi + mult32to64(inputHi, PRIME32_2 - 1); + lo ^= Long.reverseBytes(hi); + long h2lo = lo * PRIME64_2; + long h2hi = unsignedMultiplyHigh(lo, PRIME64_2) + hi * PRIME64_2; + writeHex64(hex, 0, xxh3Avalanche(h2hi)); + writeHex64(hex, 16, xxh3Avalanche(h2lo)); + } + + // Scalar-returning siblings of mix32B's two lanes, used so the Into variants below thread the + // accumulator through local variables instead of a long[2] array. + private static long mixLowLane(long acc0, byte[] input, int in1, int in2, int secOff, long seed) { + acc0 += mix16B(input, in1, secOff, seed); + acc0 ^= readLE64(input, in2) + readLE64(input, in2 + 8); + return acc0; + } + + private static long mixHighLane( + long acc1, byte[] input, int in1, int in2, int secOff, long seed) { + acc1 += mix16B(input, in2, secOff + 16, seed); + acc1 ^= readLE64(input, in1) + readLE64(input, in1 + 8); + return acc1; + } + + private static void len17to128128Into(byte[] input, int len, long seed, byte[] hex) { + long acc0 = (long) len * PRIME64_1; + long acc1 = 0L; + int i = (len - 1) / 32; + do { + acc0 = mixLowLane(acc0, input, 16 * i, len - 16 * (i + 1), 32 * i, seed); + acc1 = mixHighLane(acc1, input, 16 * i, len - 16 * (i + 1), 32 * i, seed); + } while (i-- != 0); + finalize128Into(acc0, acc1, len, seed, hex); + } + + private static void len129to240128Into(byte[] input, int len, long seed, byte[] hex) { + long acc0 = (long) len * PRIME64_1; + long acc1 = 0L; + for (int i = 0; i < 4; i++) { + acc0 = mixLowLane(acc0, input, 32 * i, 32 * i + 16, 32 * i, seed); + acc1 = mixHighLane(acc1, input, 32 * i, 32 * i + 16, 32 * i, seed); + } + acc0 = xxh3Avalanche(acc0); + acc1 = xxh3Avalanche(acc1); + for (int i = 4; i < (len >> 5); i++) { + acc0 = mixLowLane(acc0, input, 32 * i, 32 * i + 16, (i - 4) * 32 + 3, seed); + acc1 = mixHighLane(acc1, input, 32 * i, 32 * i + 16, (i - 4) * 32 + 3, seed); + } + acc0 = mixLowLane(acc0, input, len - 16, len - 32, 103, -seed); + acc1 = mixHighLane(acc1, input, len - 16, len - 32, 103, -seed); + finalize128Into(acc0, acc1, len, seed, hex); + } + + private static void finalize128Into(long acc0, long acc1, int len, long seed, byte[] hex) { + long low = xxh3Avalanche(acc0 + acc1); + long high = -xxh3Avalanche( + acc0 * PRIME64_1 + acc1 * PRIME64_4 + ((long) len - seed) * PRIME64_2); + writeHex64(hex, 0, high); + writeHex64(hex, 16, low); + } + + private static void hashLong128Into(byte[] input, int len, long seed, byte[] hex) { + byte[] secret = customSecret(seed); + long[] acc = hashLongAccumulate(input, len, secret); + long low = mergeAccs(acc, secret, SECRET_MERGEACCS_START, (long) len * PRIME64_1); + long high = mergeAccs(acc, secret, SECRET_SIZE - STRIPE_LEN - SECRET_MERGEACCS_START, + ~((long) len * PRIME64_2)); + writeHex64(hex, 0, high); + writeHex64(hex, 16, low); + } + + private static void hashHex128(byte[] input, long seed, byte[] hex) { + int len = input.length; + if (len <= 16) { + if (len > 8) { + len9to16128Into(input, len, seed, hex); + return; + } else if (len >= 4) { + len4to8128Into(input, len, seed, hex); + return; + } else if (len > 0) { + len1to3128Into(input, len, seed, hex); + return; + } + long low = xxh64Avalanche(seed ^ readLE64(SECRET, 64) ^ readLE64(SECRET, 72)); + long high = xxh64Avalanche(seed ^ readLE64(SECRET, 80) ^ readLE64(SECRET, 88)); + writeHex64(hex, 0, high); + writeHex64(hex, 16, low); + } else if (len <= 128) { + len17to128128Into(input, len, seed, hex); + } else if (len <= 240) { + len129to240128Into(input, len, seed, hex); + } else { + hashLong128Into(input, len, seed, hex); + } + } + + /** Returns the XXH3 128-bit hash as a 32-character lowercase hex string (default seed 0). */ + public static UTF8String hash128Hex(byte[] input) { + return hash128Hex(input, 0L); + } + + /** + * Returns the XXH3 128-bit hash as a 32-character lowercase hex string (canonical big-endian). + */ + public static UTF8String hash128Hex(byte[] input, long seed) { + byte[] hex = new byte[32]; + hashHex128(input, seed, hex); + return UTF8String.fromBytes(hex); + } + + private static void writeHex64(byte[] hex, int offset, long value) { + for (int i = 0; i < 8; i++) { + int b = (int) (value >>> (56 - 8 * i)) & 0xFF; + hex[offset + i * 2] = HEX_DIGITS[b >>> 4]; + hex[offset + i * 2 + 1] = HEX_DIGITS[b & 0x0F]; + } + } +} diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionUtils.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionUtils.java index 38bdcbec2069d..a3cdc355d75a3 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionUtils.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionUtils.java @@ -77,4 +77,44 @@ public static GenericArrayData jsonObjectKeys(UTF8String json) { return null; } } + + private static final UTF8String JSON_TYPE_OBJECT = UTF8String.fromString("object"); + private static final UTF8String JSON_TYPE_ARRAY = UTF8String.fromString("array"); + private static final UTF8String JSON_TYPE_STRING = UTF8String.fromString("string"); + private static final UTF8String JSON_TYPE_NUMBER = UTF8String.fromString("number"); + private static final UTF8String JSON_TYPE_BOOLEAN = UTF8String.fromString("boolean"); + private static final UTF8String JSON_TYPE_NULL = UTF8String.fromString("null"); + + public static UTF8String jsonTypeof(UTF8String json) { + try (JsonParser jsonParser = + CreateJacksonParser.utf8String(SharedFactory.jsonFactory(), json)) { + JsonToken token = jsonParser.nextToken(); + if (token == null) { + return null; + } + UTF8String type = switch (token) { + case START_OBJECT -> JSON_TYPE_OBJECT; + case START_ARRAY -> JSON_TYPE_ARRAY; + case VALUE_STRING -> JSON_TYPE_STRING; + case VALUE_NUMBER_INT, VALUE_NUMBER_FLOAT -> JSON_TYPE_NUMBER; + case VALUE_TRUE, VALUE_FALSE -> JSON_TYPE_BOOLEAN; + case VALUE_NULL -> JSON_TYPE_NULL; + default -> null; + }; + if (type == null) { + return null; + } + // Consume the value so malformed input surfaces as a parse error and returns null, + // matching json_object_keys and json_array_length. + jsonParser.skipChildren(); + // Reject trailing content after the first value, e.g. `123 true`, so only a single + // well-formed JSON value is accepted. + if (jsonParser.nextToken() != null) { + return null; + } + return type; + } catch (IOException e) { + return null; + } + } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/util/CharVarcharCodegenUtils.java b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/util/CharVarcharCodegenUtils.java index 1e183b9be5977..4c1d1014f5bb3 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/util/CharVarcharCodegenUtils.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/util/CharVarcharCodegenUtils.java @@ -54,6 +54,30 @@ public static UTF8String varcharTypeWriteSideCheck(UTF8String inputStr, int limi } } + /** + * Applies the SQL explicit-cast rules for a character string source and CHAR target. + * + * Unlike store assignment, an explicit character-to-character cast truncates non-space + * characters instead of raising a right-truncation exception. + */ + public static UTF8String charTypeCast(UTF8String inputStr, int limit) { + int numChars = inputStr.numChars(); + if (numChars == limit) { + return inputStr; + } else if (numChars < limit) { + return inputStr.rpad(limit, SPACE); + } else { + return inputStr.substring(0, limit); + } + } + + /** + * Applies the SQL explicit-cast rules for a character string source and VARCHAR target. + */ + public static UTF8String varcharTypeCast(UTF8String inputStr, int limit) { + return inputStr.numChars() > limit ? inputStr.substring(0, limit) : inputStr; + } + public static UTF8String readSidePadding(UTF8String inputStr, int limit) { int numChars = inputStr.numChars(); if (numChars == limit) { @@ -64,4 +88,27 @@ public static UTF8String readSidePadding(UTF8String inputStr, int limit) { return inputStr; } } + + /** + * Read-side CHAR check under standard semantics: pad to limit, or trim trailing + * spaces then error if still longer than limit. + * + * Standard semantics require a read to observe the same value a write would have + * produced, so this is deliberately the write-side check rather than + * {@link #readSidePadding}, which tolerates over-long values. Keep the two sides + * identical: a fix to one is a fix to both. + */ + public static UTF8String charTypeReadSideCheck(UTF8String inputStr, int limit) { + return charTypeWriteSideCheck(inputStr, limit); + } + + /** + * Read-side VARCHAR check under standard semantics: allow up to limit characters, + * or trim trailing spaces then error if still longer than limit. + * + * Identical to the write-side check by design; see {@link #varcharTypeWriteSideCheck}. + */ + public static UTF8String varcharTypeReadSideCheck(UTF8String inputStr, int limit) { + return varcharTypeWriteSideCheck(inputStr, limit); + } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/StagingTableCatalog.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/StagingTableCatalog.java index 6811ea380b3ae..bd8a34cb758e8 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/StagingTableCatalog.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/StagingTableCatalog.java @@ -99,7 +99,9 @@ default StagedTable stageCreate( * @param ident a table identifier * @param tableInfo information about the table * @return metadata for the new table. This can be null if the catalog does not support atomic - * creation for this table. Spark will call {@link #loadTable(Identifier)} later. + * creation for this table. Spark will call + * {@link #loadTable(Identifier, TableContext, CaseInsensitiveStringMap)} later, + * forwarding the catalog-declared table-state options and required privileges. * @throws TableAlreadyExistsException If a table or view already exists for the identifier * @throws UnsupportedOperationException If a requested partition transform is not supported * @throws NoSuchNamespaceException If the identifier namespace does not exist (optional) @@ -163,7 +165,9 @@ default StagedTable stageReplace( * @param ident a table identifier * @param tableInfo information about the table * @return metadata for the new table. This can be null if the catalog does not support atomic - * creation for this table. Spark will call {@link #loadTable(Identifier)} later. + * creation for this table. Spark will call + * {@link #loadTable(Identifier, TableContext, CaseInsensitiveStringMap)} later, + * forwarding the catalog-declared table-state options and required privileges. * @throws UnsupportedOperationException If a requested partition transform is not supported * @throws NoSuchNamespaceException If the identifier namespace does not exist (optional) * @throws NoSuchTableException If the table does not exist @@ -226,7 +230,9 @@ default StagedTable stageCreateOrReplace( * @param ident a table identifier * @param tableInfo information about the table * @return metadata for the new table. This can be null if the catalog does not support atomic - * creation for this table. Spark will call {@link #loadTable(Identifier)} later. + * creation for this table. Spark will call + * {@link #loadTable(Identifier, TableContext, CaseInsensitiveStringMap)} later, + * forwarding the catalog-declared table-state options and required privileges. * @throws UnsupportedOperationException If a requested partition transform is not supported * @throws NoSuchNamespaceException If the identifier namespace does not exist (optional) */ diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCapability.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCapability.java index 6a7c9b704e429..599e801f2a07f 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCapability.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCapability.java @@ -126,5 +126,33 @@ public enum TableCapability { * * @since 4.3.0 */ - GENERATE_COLUMN_VALUES_ON_WRITE + GENERATE_COLUMN_VALUES_ON_WRITE, + + /** + * Signals that Spark may fuse two batch scans of this table that differ only in their projected + * columns and/or pushed filters into a single scan (Spark-side scan merging). + * <p> + * By returning this capability a table declares a determinism contract: holding the scan options + * constant, the rows and columns a scan reads are fully determined by the filters pushed via + * {@link org.apache.spark.sql.connector.read.SupportsPushDownV2Filters} and the columns pruned + * via {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}. Equivalently, + * obtaining a fresh {@link org.apache.spark.sql.connector.read.ScanBuilder} with the same options + * and re-applying the same pushed filters and pruned columns yields an equivalent scan. + * <p> + * Given that contract, Spark builds the merged scan itself: it prunes a fresh ScanBuilder to the + * union of both read schemas, re-pushes the (possibly OR-widened) filters, and builds. The merged + * scan reads the union of the two scans' columns and a superset of their rows; each original + * scan's result is recovered by a projection and filter applied above it. The connector supplies + * no merge logic of its own. + * <p> + * This capability lives on the table rather than on the scan so that a source using the V1 scan + * fallback (whose scan Spark wraps in an internal wrapper) can still opt in. A table need not + * reason about pushdowns that are not reproducible this way (a pushed aggregate, join, variant + * extraction, limit, offset, top-N, or table sample): Spark tracks those on its own side while + * building the scan and never merges a scan that carries one, whether or not the table returns + * this capability. + * + * @since 4.3.0 + */ + SCAN_MERGING } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalog.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalog.java index 23e9499932c16..36916f907160a 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalog.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableCatalog.java @@ -17,6 +17,7 @@ package org.apache.spark.sql.connector.catalog; +import org.apache.spark.SparkIllegalArgumentException; import org.apache.spark.annotation.Evolving; import org.apache.spark.sql.connector.expressions.Transform; import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException; @@ -103,6 +104,26 @@ public interface TableCatalog extends CatalogPlugin { */ default Set<TableCatalogCapability> capabilities() { return Set.of(); } + /** + * Returns the connector-specific option keys that select the table state (such as a branch, tag, + * snapshot, or version) and therefore must be known when the table is loaded. Keys that Spark + * parses and handles itself, such as time travel, must not be listed here. + * <p> + * Spark may need to resolve the same table more than once while analyzing or refreshing a query. + * Spark reuses one table instance only for references whose table-state options match and passes + * only the declared options to {@code loadTable}. The complete user option map remains on each + * resolved relation for subsequent scan and write planning. + * <p> + * The default implementation returns an empty set, treating all options as unable to select a + * different table state. Option key matching is case-insensitive, while option values remain + * case-sensitive. + * + * @return a non-null set of case-insensitive option keys + * + * @since 4.3.0 + */ + default Set<String> tableStateOptionKeys() { return Set.of(); } + /** * List the tables in a namespace from the catalog. * @@ -194,6 +215,53 @@ default Table loadTable(Identifier ident, long timestamp) throws NoSuchTableExce throw QueryCompilationErrors.noSuchTableError(name(), ident); } + /** + * Load table metadata by {@link Identifier identifier} from the catalog, forwarding the + * user-specified options that may affect table state. + * <p> + * The default implementation ignores {@code stateOptions} and delegates to the existing + * {@code loadTable} overloads based on {@code context}. Catalogs that want to receive the user + * options while loading a table for a read or write must override + * {@link #tableStateOptionKeys()} and this method. + * Spark passes only the options declared by {@link #tableStateOptionKeys()}. Spark retains the + * complete user option map on the resolved relation for subsequent scan and write planning. + * <p> + * An override replaces that dispatch and must honor {@code context} itself: apply the time + * travel in {@link TableContext#timeTravel()}, and authorize the requested + * {@link TableContext#writePrivileges()} as it would in {@link #loadTable(Identifier, Set)}. + * Spark does not re-check either afterwards. + * + * @param ident a table identifier + * @param context the parsed load parameters (time travel, write privileges) + * @param stateOptions options declared to affect table state; Spark-parsed state such as time + * travel is provided through {@code context} instead + * @return the table's metadata + * @throws NoSuchTableException If the table doesn't exist + * + * @since 4.3.0 + */ + default Table loadTable( + Identifier ident, + TableContext context, + CaseInsensitiveStringMap stateOptions) throws NoSuchTableException { + if (context.timeTravel().isPresent()) { + TimeTravel timeTravel = context.timeTravel().get(); + if (timeTravel instanceof TimeTravel.AsOfVersion v) { + return loadTable(ident, v.version()); + } else if (timeTravel instanceof TimeTravel.AsOfTimestamp ts) { + return loadTable(ident, ts.micros()); + } else { + throw new SparkIllegalArgumentException( + "INTERNAL_ERROR", + Map.of("message", "Unsupported time travel spec: " + timeTravel)); + } + } else if (!context.writePrivileges().isEmpty()) { + return loadTable(ident, context.writePrivileges()); + } else { + return loadTable(ident); + } + } + /** * Load a {@link Changelog} for the given table, representing the row-level changes within the * range specified by {@code context}. @@ -278,7 +346,10 @@ default Table createTable( * @param ident a table identifier * @param tableInfo information about the table * @return metadata for the new table. This can be null if getting the metadata for the new table - * is expensive. Spark will call {@link #loadTable(Identifier)} if needed (e.g. CTAS). + * is expensive. Spark will call + * {@link #loadTable(Identifier, TableContext, CaseInsensitiveStringMap)} if needed + * (e.g. CTAS), forwarding the catalog-declared table-state options and required + * privileges. * * @throws TableAlreadyExistsException If a table already exists for the identifier * @throws UnsupportedOperationException If a requested partition transform is not supported diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableContext.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableContext.java new file mode 100644 index 0000000000000..9437149895d27 --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TableContext.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.catalog; + +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +import org.apache.spark.SparkIllegalArgumentException; +import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; + +/** + * Encapsulates the parsed, Spark-recognized parameters of a table load request, passed from the + * analyzer / DataFrame API to the catalog's + * {@link TableCatalog#loadTable(Identifier, TableContext, CaseInsensitiveStringMap)} method. + * <p> + * A load is either a read (optionally with time travel) or a write (carrying write privileges); + * time travel and write privileges are mutually exclusive. + * + * @since 4.3.0 + */ +@Evolving +public class TableContext { + + // null means no time travel. + private final TimeTravel timeTravel; + // Never null; an empty set means no write privileges (i.e. a read). + private final Set<TableWritePrivilege> writePrivileges; + + public TableContext(TimeTravel timeTravel, Set<TableWritePrivilege> privileges) { + this.timeTravel = timeTravel; + this.writePrivileges = privileges == null ? Set.of() : Set.copyOf(privileges); + if (timeTravel != null && !writePrivileges.isEmpty()) { + throw new SparkIllegalArgumentException( + "INTERNAL_ERROR", + Map.of("message", "Cannot set both time travel and write privileges")); + } + } + + /** Returns the time-travel spec, or empty for a current-version read. */ + public Optional<TimeTravel> timeTravel() { + return Optional.ofNullable(timeTravel); + } + + /** Returns the requested write privileges; empty for a read. */ + public Set<TableWritePrivilege> writePrivileges() { + return writePrivileges; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof TableContext that)) return false; + return Objects.equals(timeTravel, that.timeTravel) + && writePrivileges.equals(that.writePrivileges); + } + + @Override + public int hashCode() { + return Objects.hash(timeTravel, writePrivileges); + } + + @Override + public String toString() { + return "TableContext{timeTravel=" + timeTravel + + ", writePrivileges=" + writePrivileges + "}"; + } +} diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TimeTravel.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TimeTravel.java new file mode 100644 index 0000000000000..c9ecc80d65d5e --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/TimeTravel.java @@ -0,0 +1,43 @@ +/* + * 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.connector.catalog; + +import org.apache.spark.annotation.Evolving; + +/** + * A time-travel specification for reading a table as of a specific version or point in time. + * + * @since 4.3.0 + */ +@Evolving +public sealed interface TimeTravel permits TimeTravel.AsOfVersion, TimeTravel.AsOfTimestamp { + + /** + * Time travel to a specific version of the table. + * + * @param version the version identifier (connector-defined) + */ + record AsOfVersion(String version) implements TimeTravel {} + + /** + * Time travel to a specific point in time. + * + * @param micros microseconds since 1970-01-01 00:00:00 UTC + */ + record AsOfTimestamp(long micros) implements TimeTravel {} +} diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/BoundFunction.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/BoundFunction.java index 53a1beb9c4b14..2c8c33f1653fa 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/BoundFunction.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/BoundFunction.java @@ -90,6 +90,14 @@ default boolean isDeterministic() { * functions in other catalogs. For example, many catalogs may define a "bucket" function with a * different implementation. Adding context, like "com.mycompany.bucket(string)", is recommended * to avoid unintentional collisions. + * <p> + * Two functions that partition data differently must not return the same name; Spark may + * otherwise treat unrelated data as co-partitioned. An override should return a stable name + * across {@code bind} calls, since Spark may bind a function multiple times and has only this + * name to relate two bound instances by. Equal functions must share the same canonical name; + * the reverse is not true, as this name is deliberately coarser and says nothing about the + * rest of a function's state. For whether two transform expressions are the same expression, + * see {@link #equals(Object)}. * * @return a canonical name for this function */ @@ -100,4 +108,38 @@ default String canonicalName() { // bugs if not replaced before release. return UUID.randomUUID().toString(); } + + /** + * Implementations SHOULD override {@link Object#equals(Object)} and {@link Object#hashCode()}. + * <p> + * Spark may bind a function multiple times, so the same transform can be represented by two + * different bound instances. Without a semantic {@code equals} it cannot tell they are the + * same, and misses optimizations such as: + * <ul> + * <li>keeping a union's keyed partitioning</li> + * <li>retaining a reported ordering that matches the partitioning</li> + * <li>reusing identical bucketed scans</li> + * <li>recognizing two identical subplans</li> + * </ul> + * Missed matches cost performance only, never correctness. + * <p> + * Compare whatever state affects behaviour, and keep it stable across {@code bind} calls. + * {@link #canonicalName()} alone is not always enough: two functions can share a name and still + * differ in, say, {@link #resultType()} or a {@link ReducibleFunction}'s reducers. Two functions + * that can produce different values must not compare equal -- the same comparison decides whether + * two ordinary calls to a {@link ScalarFunction} or an {@link AggregateFunction} are the same + * expression, so Spark may otherwise evaluate one where the query asked for the other. + * {@code hashCode} must agree with {@code equals}. + * <p> + * If this method is overridden, {@link #canonicalName()} should be overridden as well, so that + * equal functions share the same name. + */ + @Override + boolean equals(Object other); + + /** + * Must agree with {@link #equals(Object)}. See that method. + */ + @Override + int hashCode(); } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportStatistics.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportStatistics.java index 031749dee0350..fa91f2c6117be 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportStatistics.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsReportStatistics.java @@ -17,6 +17,8 @@ package org.apache.spark.sql.connector.read; +import java.util.OptionalLong; + import org.apache.spark.annotation.Evolving; /** @@ -36,4 +38,51 @@ public interface SupportsReportStatistics extends Scan { * Returns the estimated statistics of this data source scan. */ Statistics estimateStatistics(); + + /** + * Returns the estimated size in bytes of this scan without computing full statistics. + * <p> + * When cost-based optimization or plan statistics are disabled, Spark primarily needs the scan's + * size in bytes (for example, for broadcast-join thresholding). This method lets connectors that + * can produce a size estimate cheaply serve it directly and avoid computing the full statistics. + * <p> + * The default implementation returns {@code OptionalLong.empty()}, signalling that the connector + * does not offer a cheap size estimate. In that case Spark falls back to + * {@link #estimateStatistics()}, so a connector that only implements + * {@link #estimateStatistics()} keeps the same size-estimation behavior it had before this method + * existed. Connectors override this method only when they have a genuinely cheaper size estimate + * than {@link #estimateStatistics()}. + * + * @since 4.3.0 + */ + default OptionalLong estimateSizeInBytes() { + return OptionalLong.empty(); + } + + /** + * Returns whether the statistics reported by this scan already reflect all filters that were + * fully pushed down to the data source. + * <p> + * When {@code true} (the default), the reported statistics describe exactly the data the scan + * will produce. When {@code false}, they do <em>not</em> account for the fully pushed filters + * (for example, they describe the whole table), so Spark may use those fully pushed filters to + * adjust stats. Re-applying those fully pushed filters in Spark should be redundant for query + * results because the data source already evaluates them. + * <p> + * The adjustment Spark performs when this returns {@code false} is best-effort: Spark re-applies + * a fully pushed filter for stats adjustment only when every column the filter references is + * still present in {@link Scan#readSchema()}. If {@code pruneColumns} removes a pushed-filter + * column, Spark drops that filter from the adjustment, so the reported statistics will not + * reflect it. + * <p> + * A connector that wants a fully pushed filter to participate in this stats adjustment must, when + * Spark requests schema pruning through {@link SupportsPushDownRequiredColumns}, retain the + * columns referenced only by that fully pushed filter; otherwise those columns are pruned and the + * filter is dropped from the adjustment. + * + * @since 4.3.0 + */ + default boolean reflectsFullyPushedDownFilters() { + return true; + } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java index 927d4a53e22fc..067202a362705 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java @@ -38,6 +38,10 @@ public interface SupportsRuntimeFiltering extends SupportsRuntimeV2Filtering { * <p> * Spark will call {@link #filter(Filter[])} if it can derive a runtime * predicate for any of the filter attributes. + * <p> + * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java index 94dbc3865958a..6b286f041b01e 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java @@ -28,8 +28,9 @@ * filter initially planned {@link InputPartition}s using predicates Spark infers at runtime. * This interface is very similar to {@link SupportsRuntimeFiltering} except it uses * data source V2 {@link Predicate} instead of data source V1 {@link Filter}. - * {@link SupportsRuntimeV2Filtering} is preferred over {@link SupportsRuntimeFiltering} - * and only one of them should be implemented by the data sources. + * {@link SupportsRuntimeV2Filtering} is preferred over {@link SupportsRuntimeFiltering}. + * A scan must not implement SupportsRuntimeCatalystFiltering together with this interface; + * Spark rejects such a scan. * <p> * <b>Iterative filtering:</b> When {@link #supportsIterativePushdown()} returns true, * {@link #filter(Predicate[])} may be called <i>multiple times</i> on the same @@ -50,6 +51,10 @@ public interface SupportsRuntimeV2Filtering extends Scan { * <p> * Spark will call {@link #filter(Predicate[])} if it can derive a runtime * predicate for any of the filter attributes. + * <p> + * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); @@ -72,6 +77,11 @@ public interface SupportsRuntimeV2Filtering extends Scan { * The implementation must accumulate state across all calls so that * {@link #pushedPredicates()} can return predicates from all of them. * <p> + * Independently of {@link #supportsIterativePushdown()}, this method may also be called once + * per scan node when a plan holds several scan nodes sharing one {@link Scan} instance (e.g. + * the two branches of a group-based UPDATE). Implementations must accumulate state across + * those calls as well. + * <p> * Note that Spark will call {@link Scan#toBatch()} again after filtering the scan at runtime. * * @param predicates data source V2 predicates used to filter the scan at runtime @@ -82,6 +92,10 @@ public interface SupportsRuntimeV2Filtering extends Scan { * Returns the predicates that are pushed to the data source via * {@link #filter(Predicate[])}. * <p> + * These are not fully pushed predicates: Spark may still evaluate them after the scan. + * They are predicates that fully or partially help the data source prune initially planned + * {@link InputPartition}s. + * <p> * When iterative filtering is supported and {@link #filter(Predicate[])} was called * multiple times, this method must return predicates from <i>all</i> calls. * <p> diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsRealTimeRead.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsRealTimeRead.java index 5542781f333a8..50150fc11feee 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsRealTimeRead.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsRealTimeRead.java @@ -25,7 +25,22 @@ /** * A variation on {@link PartitionReader} for use with low latency streaming processing. - * + * <p> + * <b>Which method to implement.</b> A source reader must provide a {@code nextWithTimeout} that + * proceeds to the next record, blocking until one is available or a timeout elapses. There are two + * overloads, and a source should override exactly one: + * <ul> + * <li>{@link #nextWithTimeout(Long)} -- implement this in the common case. It receives only the + * timeout and is all a third-party source needs.</li> + * <li>{@link #nextWithTimeout(Long, Long)} -- the overload the execution engine actually invokes. + * Its default implementation ignores the extra parameter and delegates to + * {@link #nextWithTimeout(Long)}. Override it only for engine-internal sources that must + * observe the engine's reference start time (used together with the engine-internal low + * latency clock); that clock is not part of the third-party contract, so most sources should + * not override this overload.</li> + * </ul> + * Because both overloads are {@code default} methods, overriding neither is not caught at compile + * time -- it fails at runtime when {@link #nextWithTimeout(Long)} throws. */ @Evolving public interface SupportsRealTimeRead<T> extends PartitionReader<T> { @@ -78,12 +93,39 @@ public Optional<Long> recArrivalTime() { * Alternative function to be called than next(), that proceed to the next record. The different * from next() is that, if there is no more records, the call needs to keep waiting until * the timeout. - * @param startTimeMs the base time (milliseconds) the was used to calculate the timeout. - * Sources should use it as the reference time to start waiting for the next - * record instead of getting the latest time from LowLatencyClock. + * <p> + * This is the recommended method to implement for a source. It is enough for any source that + * does not need the engine's reference start time (see {@link #nextWithTimeout(Long, Long)}). + * The engine always invokes {@link #nextWithTimeout(Long, Long)}, whose default implementation + * delegates here, so overriding only this method is sufficient. If a source overrides neither + * overload, this default implementation throws to surface the mistake. * @param timeoutMs if no result is available after this timeout (milliseconds), return * @return {@link RecordStatus} describing whether a record is available and its arrival time * @throws IOException */ - RecordStatus nextWithTimeout(Long startTimeMs, Long timeoutMs) throws IOException; + default RecordStatus nextWithTimeout(Long timeoutMs) throws IOException { + throw new UnsupportedOperationException( + "A SupportsRealTimeRead implementation must override either " + + "nextWithTimeout(Long) or nextWithTimeout(Long, Long)."); + } + + /** + * The overload of {@link #nextWithTimeout(Long)} that the execution engine actually invokes. In + * addition to the timeout it receives {@code startTimeMs}, the reference time the engine used + * to compute that timeout, so the engine and the source agree on when the wait started (this + * matters when the engine runs against its internal low latency clock, e.g. a manual clock in + * tests). The default implementation ignores {@code startTimeMs} and delegates to + * {@link #nextWithTimeout(Long)}. + * <p> + * This is intended for engine-internal sources. The reference clock is not part of the + * third-party contract, so third-party sources should implement {@link #nextWithTimeout(Long)} + * and leave this overload to its default. + * @param startTimeMs the base time (milliseconds) that was used to calculate the timeout + * @param timeoutMs if no result is available after this timeout (milliseconds), return + * @return {@link RecordStatus} describing whether a record is available and its arrival time + * @throws IOException + */ + default RecordStatus nextWithTimeout(Long startTimeMs, Long timeoutMs) throws IOException { + return nextWithTimeout(timeoutMs); + } } diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatchRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatchRow.java index c689d9faf4b31..89b0b7c0a466e 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatchRow.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatchRow.java @@ -86,6 +86,11 @@ public InternalRow copy() { row.update(i, getMap(i).copy()); } else if (pdt instanceof PhysicalVariantType) { row.update(i, getVariant(i)); + } else if (pdt instanceof PhysicalTimestampNTZNanosType) { + // TimestampNanosVal is immutable, so it can be shared without copying. + row.update(i, getTimestampNTZNanos(i)); + } else if (pdt instanceof PhysicalTimestampLTZNanosType) { + row.update(i, getTimestampLTZNanos(i)); } else { throw new RuntimeException("Not implemented. " + dt); } @@ -208,6 +213,10 @@ public Object get(int ordinal, DataType dataType) { return getLong(ordinal); } else if (dataType instanceof TimestampNTZType) { return getLong(ordinal); + } else if (dataType instanceof TimestampNTZNanosType) { + return getTimestampNTZNanos(ordinal); + } else if (dataType instanceof TimestampLTZNanosType) { + return getTimestampLTZNanos(ordinal); } else if (dataType instanceof ArrayType) { return getArray(ordinal); } else if (dataType instanceof StructType structType) { diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarRow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarRow.java index 9086726c5ee5b..2db280c53de61 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarRow.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarRow.java @@ -91,6 +91,11 @@ public InternalRow copy() { row.update(i, getMap(i).copy()); } else if (pdt instanceof PhysicalVariantType) { row.update(i, getVariant(i)); + } else if (pdt instanceof PhysicalTimestampNTZNanosType) { + // TimestampNanosVal is immutable, so it can be shared without copying. + row.update(i, getTimestampNTZNanos(i)); + } else if (pdt instanceof PhysicalTimestampLTZNanosType) { + row.update(i, getTimestampLTZNanos(i)); } else { throw new RuntimeException("Not implemented. " + dt); } @@ -212,6 +217,10 @@ public Object get(int ordinal, DataType dataType) { return getLong(ordinal); } else if (dataType instanceof TimestampNTZType) { return getLong(ordinal); + } else if (dataType instanceof TimestampNTZNanosType) { + return getTimestampNTZNanos(ordinal); + } else if (dataType instanceof TimestampLTZNanosType) { + return getTimestampLTZNanos(ordinal); } else if (dataType instanceof ArrayType) { return getArray(ordinal); } else if (dataType instanceof StructType) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystTypeConverters.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystTypeConverters.scala index 97804d1e84633..cab367b4badf4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystTypeConverters.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/CatalystTypeConverters.scala @@ -61,6 +61,15 @@ object CatalystTypeConverters { } } + private def throwInvalidExternalValue(value: Any, dataType: String): Nothing = { + throw new SparkIllegalArgumentException( + errorClass = "INVALID_EXTERNAL_VALUE", + messageParameters = scala.collection.immutable.Map( + "other" -> value.toString, + "otherClass" -> value.getClass.getCanonicalName, + "dataType" -> dataType)) + } + private def getConverterForType(dataType: DataType): CatalystTypeConverter[Any, Any, Any] = { TypeUtils.failUnsupportedDataType(dataType, SQLConf.get) TypeOps(dataType) @@ -296,12 +305,7 @@ object CatalystTypeConverters { idx += 1 } new GenericInternalRow(ar) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> structType.catalogString)) + case other => throwInvalidExternalValue(other, structType.catalogString) } override def toScala(row: InternalRow): Row = { @@ -355,12 +359,7 @@ object CatalystTypeConverters { case utf8: UTF8String => utf8 case chr: Char => UTF8String.fromString(chr.toString) case ac: Array[Char] => UTF8String.fromString(String.valueOf(ac)) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> StringType.sql)) + case other => throwInvalidExternalValue(other, StringType.sql) } override def toScala(catalystValue: UTF8String): String = if (catalystValue == null) null else catalystValue.toString @@ -381,12 +380,7 @@ object CatalystTypeConverters { override def toCatalystImpl(scalaValue: Any): BinaryView = scalaValue match { case g: org.apache.spark.sql.types.Geometry if SQLConf.get.geospatialEnabled => STUtils.serializeGeomFromWKB(g, dataType) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> StringType.sql)) + case other => throwInvalidExternalValue(other, StringType.sql) } override def toScala(catalystValue: BinaryView): org.apache.spark.sql.types.Geometry = { assertGeospatialEnabled() @@ -406,12 +400,7 @@ object CatalystTypeConverters { override def toCatalystImpl(scalaValue: Any): BinaryView = scalaValue match { case g: org.apache.spark.sql.types.Geography if SQLConf.get.geospatialEnabled => STUtils.serializeGeogFromWKB(g, dataType) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> StringType.sql)) + case other => throwInvalidExternalValue(other, StringType.sql) } override def toScala(catalystValue: BinaryView): org.apache.spark.sql.types.Geography = { assertGeospatialEnabled() @@ -430,12 +419,7 @@ object CatalystTypeConverters { override def toCatalystImpl(scalaValue: Any): Int = scalaValue match { case d: Date => DateTimeUtils.fromJavaDate(d) case l: LocalDate => DateTimeUtils.localDateToDays(l) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> DateType.sql)) + case other => throwInvalidExternalValue(other, DateType.sql) } override def toScala(catalystValue: Any): Date = if (catalystValue == null) null else DateTimeUtils.toJavaDate(catalystValue.asInstanceOf[Int]) @@ -470,12 +454,7 @@ object CatalystTypeConverters { override def toCatalystImpl(scalaValue: Any): Long = scalaValue match { case t: Timestamp => DateTimeUtils.fromJavaTimestamp(t) case i: Instant => DateTimeUtils.instantToMicros(i) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> TimestampType.sql)) + case other => throwInvalidExternalValue(other, TimestampType.sql) } override def toScala(catalystValue: Any): Timestamp = if (catalystValue == null) null @@ -498,12 +477,7 @@ object CatalystTypeConverters { extends CatalystTypeConverter[Any, LocalDateTime, Any] { override def toCatalystImpl(scalaValue: Any): Any = scalaValue match { case l: LocalDateTime => DateTimeUtils.localDateTimeToMicros(l) - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> TimestampNTZType.sql)) + case other => throwInvalidExternalValue(other, TimestampNTZType.sql) } override def toScala(catalystValue: Any): LocalDateTime = @@ -525,12 +499,7 @@ object CatalystTypeConverters { case d: JavaBigDecimal => Decimal(d) case d: JavaBigInteger => Decimal(d) case d: Decimal => d - case other => throw new SparkIllegalArgumentException( - errorClass = "INVALID_EXTERNAL_VALUE", - messageParameters = scala.collection.immutable.Map( - "other" -> other.toString, - "otherClass" -> other.getClass.getCanonicalName, - "dataType" -> dataType.catalogString)) + case other => throwInvalidExternalValue(other, dataType.catalogString) } decimal.toPrecision(dataType.precision, dataType.scale, Decimal.ROUND_HALF_UP, nullOnOverflow) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/DeserializerBuildHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/DeserializerBuildHelper.scala index b08804dccc673..d4c37f7eb5fa9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/DeserializerBuildHelper.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/DeserializerBuildHelper.scala @@ -327,10 +327,10 @@ object DeserializerBuildHelper { createDeserializerForGeographyType(path, g.dt) case g: GeometryEncoder => createDeserializerForGeometryType(path, g.dt) - case CharEncoder(length) => - createDeserializerForChar(path, returnNullable = false, length) - case VarcharEncoder(length) => - createDeserializerForVarchar(path, returnNullable = false, length) + case CharEncoder(dt) => + createDeserializerForChar(path, returnNullable = false, dt.length) + case VarcharEncoder(dt) => + createDeserializerForVarchar(path, returnNullable = false, dt.length) case StringEncoder => createDeserializerForString(path, returnNullable = false) case _: ScalaDecimalEncoder => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/FileSourceOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/FileSourceOptions.scala index 62db2d90ec8fd..1040bb0312875 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/FileSourceOptions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/FileSourceOptions.scala @@ -18,7 +18,11 @@ package org.apache.spark.sql.catalyst import java.util.regex.{Pattern, PatternSyntaxException} -import org.apache.spark.sql.catalyst.FileSourceOptions.{IGNORE_CORRUPT_FILES, IGNORE_MISSING_FILES, IGNORED_PATH_SEGMENT_REGEX} +import scala.util.control.NonFatal + +import org.apache.hadoop.fs.GlobPattern + +import org.apache.spark.sql.catalyst.FileSourceOptions.{ARCHIVE_PATH_FILTER, IGNORE_CORRUPT_FILES, IGNORE_MISSING_FILES, IGNORED_PATH_SEGMENT_REGEX} import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateFormatter} import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf} @@ -66,12 +70,37 @@ class FileSourceOptions( */ lazy val ignoredPathSegmentRegexPattern: Pattern = FileSourceOptions.compileIgnoredPathSegmentRegex(ignoredPathSegmentRegex) + + /** + * Glob selecting which inner archive entries to read, matched against each entry's full path + * within the archive (e.g. `subdir/*`, `*/*.csv`). An empty value disables the filter, matching + * how an empty [[ignoredPathSegmentRegex]] is treated. Validated here so an invalid glob fails on + * the driver. + */ + val archivePathFilter: Option[String] = { + val glob = parameters.get(ARCHIVE_PATH_FILTER).filter(_.nonEmpty) + glob.foreach(FileSourceOptions.compileArchivePathFilter) + glob + } + + /** + * The effective [[archivePathFilter]] is compiled once per instance of this class, so the archive + * reads sharing one options object reuse a single matcher rather than re-compiling per archive. + * `transient` because Hadoop's `GlobPattern` is not serializable, so an executor recompiles from + * [[archivePathFilter]] on first use. Paths that cannot reach this value -- the schema-inference + * RDDs, which would have to capture the matcher in a closure, and the parallel footer/schema + * readers, whose signatures are fixed -- carry [[archivePathFilter]] instead and compile it + * themselves. + */ + @transient lazy val archivePathFilterPattern: Option[GlobPattern] = + archivePathFilter.map(FileSourceOptions.compileArchivePathFilter) } object FileSourceOptions { val IGNORE_CORRUPT_FILES = "ignoreCorruptFiles" val IGNORE_MISSING_FILES = "ignoreMissingFiles" val IGNORED_PATH_SEGMENT_REGEX = "ignoredPathSegmentRegex" + val ARCHIVE_PATH_FILTER = "archivePathFilter" // A regex that never matches any name, used when the filter is disabled by an empty value. private val DISABLED_FILTER_PATTERN = Pattern.compile("(?!)") @@ -97,4 +126,15 @@ object FileSourceOptions { } } } + + /** Compiles `archivePathFilter` into a glob matcher, reporting an invalid glob clearly. */ + def compileArchivePathFilter(glob: String): GlobPattern = { + try { + new GlobPattern(glob) + } catch { + case NonFatal(e) => + throw new IllegalArgumentException( + s"The '$ARCHIVE_PATH_FILTER' value '$glob' is not a valid glob.", e) + } + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala index 3531ebc77c98c..370749b1d9bf9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SerializerBuildHelper.scala @@ -358,8 +358,8 @@ object SerializerBuildHelper { messageParameters = scala.collection.immutable.Map.empty) case g: GeographyEncoder => createSerializerForGeographyType(input, g.dt) case g: GeometryEncoder => createSerializerForGeometryType(input, g.dt) - case CharEncoder(length) => createSerializerForChar(input, length) - case VarcharEncoder(length) => createSerializerForVarchar(input, length) + case CharEncoder(dt) => createSerializerForChar(input, dt.length) + case VarcharEncoder(dt) => createSerializerForVarchar(input, dt.length) case StringEncoder => createSerializerForString(input) case ScalaDecimalEncoder(dt) => createSerializerForBigDecimal(input, dt) case JavaDecimalEncoder(dt, false) => createSerializerForBigDecimal(input, dt) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index d6eac52ebbb28..22a887ae12ce8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala @@ -132,9 +132,12 @@ object FakeV2SessionCatalog extends TableCatalog with FunctionCatalog with Suppo * @param nestedViewDepth The nested depth in the view resolution, this enables us to limit the * depth of nested views. * @param maxNestedViewDepth The maximum allowed depth of nested view resolution. - * @param relationCache A mapping from qualified table names and time travel spec to resolved - * relations. This can ensure that the table is resolved only once if a table - * is used multiple times in a query. + * @param relationCache A mapping from (qualified table name, time travel spec, options) to + * resolved relations. This can ensure that the table is resolved only once if + * a table is used multiple times in a query with the same options. + * @param tableCache A mapping from (catalog, identifier, time travel spec, table-state options) to + * concrete tables. This pins one table state while allowing references to keep + * different read-specific options. * @param referredTempViewNames All the temp view names referred by the current view we are * resolving. It's used to make sure the relation resolution is * consistent between view creation and view resolution. For example, @@ -154,8 +157,8 @@ case class AnalysisContext( resolutionPathEntries: Option[Seq[Seq[String]]] = None, nestedViewDepth: Int = 0, maxNestedViewDepth: Int = -1, - relationCache: mutable.Map[(Seq[String], Option[TimeTravelSpec]), LogicalPlan] = - mutable.Map.empty, + relationCache: mutable.Map[RelationCacheKey, LogicalPlan] = mutable.Map.empty, + tableCache: mutable.Map[TableCacheKey, Table] = mutable.Map.empty, referredTempViewNames: Seq[Seq[String]] = Seq.empty, // 1. If we are resolving a view, this field will be restored from the view metadata, // by calling `AnalysisContext.withAnalysisContext(viewDesc)`. @@ -250,6 +253,7 @@ object AnalysisContext { nestedViewDepth = originContext.nestedViewDepth + 1, maxNestedViewDepth = maxNestedViewDepth, relationCache = originContext.relationCache, + tableCache = originContext.tableCache, referredTempViewNames = viewDesc.viewReferredTempViewNames, referredTempFunctionNames = mutable.Set(viewDesc.viewReferredTempFunctionNames: _*), referredTempVariableNames = viewDesc.viewReferredTempVariableNames, @@ -638,7 +642,8 @@ class Analyzer( Seq( ResolveWithCTE, ExtractDistributedSequenceID, - ResolveAsOfJoin) ++ + ResolveAsOfJoin, + ResolveTranspiledPythonUDFOptions) ++ Seq(ResolveUpdateEventTimeWatermarkColumn) ++ extendedResolutionRules ++ Seq(NameStreamingSources) : _*), @@ -958,7 +963,8 @@ class Analyzer( // TODO: Support Pandas UDF. private def checkValidAggregateExpression(expr: Expression): Unit = expr match { case a: AggregateExpression => - if (a.aggregateFunction.isInstanceOf[PythonUDAF]) { + if (a.aggregateFunction.isInstanceOf[PythonUDAF] || + a.aggregateFunction.isInstanceOf[PythonAggregate]) { throw QueryCompilationErrors.pandasUDFAggregateNotSupportedInPivotError() } else { // OK and leave the argument check to CheckAnalysis. @@ -3930,12 +3936,11 @@ class Analyzer( val defaultValueFillMode = if (conf.coerceInsertNestedTypes && v2Write.schemaEvolutionEnabled) RECURSE else FILL - // Only let TableOutputResolver see generation expression metadata if the table - // supports auto-filling generated columns on write. + // Generation expressions live on the table's columns, so attach them to the expected + // output for TableOutputResolver, which auto-fills the generated columns the query is + // missing. val expected = v2Write.table match { - case r: DataSourceV2Relation - if !GeneratedColumn.supportsGeneratedColumnsOnWrite(r.table) => - r.output.map(GeneratedColumn.removeGenerationExpressionMetadata) + case r: DataSourceV2Relation => GeneratedColumn.attachGenerationExpressions(r) case _ => v2Write.table.output } val (projection, autoFilledGenCols) = @@ -3945,20 +3950,9 @@ class Analyzer( if (projection != v2Write.query) { val cleanedTable = v2Write.table match { case r: DataSourceV2Relation => - r.copy(output = r.output.map { attr => - val cleaned = CharVarcharUtils.cleanAttrMetadata(attr) - // Strip the generation expression metadata from columns Spark auto-filled, so - // ResolveTableConstraints does not add a (redundant) CheckInvariant for them: - // their values were computed from the generation expression and are correct by - // construction. User-provided generated columns keep the metadata so their - // values are still validated. - if (autoFilledGenCols.contains(attr.name)) { - GeneratedColumn.removeGenerationExpressionMetadata(cleaned) - .asInstanceOf[AttributeReference] - } else { - cleaned - } - }) + val cleaned = r.output.map(CharVarcharUtils.cleanAttrMetadata) + r.copy(output = + GeneratedColumn.markAutoFilledGeneratedColumns(cleaned, autoFilledGenCols)) case other => other } v2Write.withNewQuery(projection).withNewTable(cleanedTable) @@ -4718,7 +4712,50 @@ object ResolveUnresolvedHaving extends Rule[LogicalPlan] { plan.resolveOperatorsWithPruning(_.containsPattern(UNRESOLVED_HAVING), ruleId) { case u @ UnresolvedHaving(havingCondition, child) if havingCondition.resolved && child.resolved => - Filter(condition = havingCondition, child = child) + val filter = Filter(condition = havingCondition, child = child) + insertFilterBeforeWindow(filter).getOrElse(filter) + } + } + + /** + * Searches through Project and Generate nodes for a Window chain and places HAVING below every + * Window in that chain. This restores SQL clause order for plans produced by queries such as: + * + * {{{ + * SELECT explode(array(a)), count(*) OVER () + * FROM VALUES (1), (2), (NULL) AS t(a) + * GROUP BY a + * HAVING a IS NOT NULL + * }}} + * + * Returns None unless the condition can be evaluated below every Window in the chain. + */ + private def insertFilterBeforeWindow(filter: Filter): Option[LogicalPlan] = filter.child match { + case project: Project => + insertFilterBeforeWindow(filter.copy(child = project.child)) + .map(child => project.withNewChildren(Seq(child))) + case generate: Generate => + insertFilterBeforeWindow(filter.copy(child = generate.child)) + .map(child => generate.withNewChildren(Seq(child))) + case window: Window => + insertFilterBeforeWindowChain(filter, window) + case _ => + None + } + + private def insertFilterBeforeWindowChain( + filter: Filter, + window: Window): Option[LogicalPlan] = { + if (!filter.condition.references.subsetOf(window.child.outputSet)) { + None + } else { + val child = window.child match { + case childWindow: Window => + insertFilterBeforeWindowChain(filter, childWindow) + case child => + Some(filter.copy(child = child)) + } + child.map(child => window.withNewChildren(Seq(child))) } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala index 23c416dd4b383..d64f9d48bee2c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercion.scala @@ -168,6 +168,10 @@ object AnsiTypeCoercion extends TypeCoercionBase { private def implicitCast( inType: DataType, expectedType: AbstractDataType): Option[DataType] = { + // CHAR/VARCHAR promotion is checked first: the acceptsType case below would otherwise + // accept the constrained type unchanged, since CharType and VarcharType extend StringType. + charVarcharToPlainString(inType, expectedType).foreach(dt => return Some(dt)) + (inType, expectedType) match { // If the expected type equals the input type, no need to cast. case _ if expectedType.acceptsType(inType) => Some(inType) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/BinaryArithmeticWithDatetimeResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/BinaryArithmeticWithDatetimeResolver.scala index cdfd942ca09af..3b254f297242f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/BinaryArithmeticWithDatetimeResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/BinaryArithmeticWithDatetimeResolver.scala @@ -47,6 +47,7 @@ import org.apache.spark.sql.catalyst.expressions.{ } import org.apache.spark.sql.types.{ AnsiIntervalType, + AnyTimestampNanoType, AnyTimestampTypeExpression, CalendarIntervalType, DatetimeType, @@ -75,6 +76,10 @@ object BinaryArithmeticWithDatetimeResolver { TimestampAddYMInterval(l, r) case (_: YearMonthIntervalType, TimestampType | TimestampNTZType) => TimestampAddYMInterval(r, l) + case (_: AnyTimestampNanoType, _: YearMonthIntervalType) => + TimestampAddYMInterval(l, r) + case (_: YearMonthIntervalType, _: AnyTimestampNanoType) => + TimestampAddYMInterval(r, l) case (CalendarIntervalType, CalendarIntervalType) | (_: DayTimeIntervalType, _: DayTimeIntervalType) => a @@ -113,6 +118,9 @@ object BinaryArithmeticWithDatetimeResolver { case (TimestampType | TimestampNTZType, _: YearMonthIntervalType) => DatetimeSub(l, r, TimestampAddYMInterval(l, UnaryMinus(r, context.evalMode == EvalMode.ANSI))) + case (_: AnyTimestampNanoType, _: YearMonthIntervalType) => + DatetimeSub(l, r, TimestampAddYMInterval(l, + UnaryMinus(r, context.evalMode == EvalMode.ANSI))) case (CalendarIntervalType, CalendarIntervalType) | (_: DayTimeIntervalType, _: DayTimeIntervalType) => s @@ -137,7 +145,9 @@ object BinaryArithmeticWithDatetimeResolver { TimestampAddInterval(l, UnaryMinus(r, context.evalMode == EvalMode.ANSI))), l.dataType) case _ if AnyTimestampTypeExpression.unapply(l) || - AnyTimestampTypeExpression.unapply(r) => + AnyTimestampTypeExpression.unapply(r) || + AnyTimestampNanoType.acceptsType(l.dataType) || + AnyTimestampNanoType.acceptsType(r.dataType) => SubtractTimestamps(l, r) case (_, DateType) => SubtractDates(l, r) case (DateType, dt) if dt != StringType => DateSub(l, r) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala index 1d78d455084a6..53b5780847f2f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala @@ -307,7 +307,7 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString // We should inline all CTE relations to restore the original plan shape, as the analysis check // may need to match certain plan shapes. For dangling CTE relations, they will still be kept // in the original `WithCTE` node, as we need to perform analysis check for them as well. - val inlineCTE = InlineCTE(alwaysInline = true, keepDanglingRelations = true) + val inlineCTE = InlineCTE(alwaysInline = true, keepDanglingRelations = true, isAnalysis = true) val inlinedPlan: LogicalPlan = try { inlineCTE(plan) } catch { @@ -363,11 +363,13 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString plan.foreachUp { case p if p.analyzed => // Skip already analyzed sub-plans - case leaf: LeafNode if !SQLConf.get.preserveCharVarcharTypeInfo && - leaf.output.map(_.dataType).exists(CharVarcharUtils.hasCharVarchar) => + case leaf: LeafNode + if !SQLConf.get.charVarcharFirstClassTypes && + leaf.output.exists(attr => CharVarcharUtils.hasCharVarchar(attr.dataType)) => throw SparkException.internalError( s"Logical plan should not have output of char/varchar type when " + - s"${SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key} is false: " + leaf) + s"${SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key} and " + + s"${SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key} are both false: " + leaf) case u: UnresolvedNamespace => u.schemaNotFound(u.multipartIdentifier) @@ -501,10 +503,27 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString hof.invalidFormat(checkRes) } + // A Python UDF cannot be evaluated inside a lambda, because it needs a separate + // physical operator that cannot see the lambda's variables. For some shapes + // `ExtractPythonUDFFromLambda` rewrites the plan in the optimizer so the UDF is + // applied to the whole array outside the lambda instead; those are allowed through + // here. Everything else must still fail, or nothing downstream can evaluate it. + // + // Judge only at a *nest root* - a HOF that iterates real columns, not a free lambda + // variable. `canRewritePythonUDFInLambda` validates the whole nest below the root + // (a UDF in a nested lambda is lifted out one level at a time), and a root's + // `functions.exists` sees UDFs at every depth, so firing on inner HOFs too would + // double-report and reject nests the rule actually handles. case hof: HigherOrderFunction if hof.resolved && hof.functions - .exists(_.exists(_.isInstanceOf[PythonUDF])) => - val u = hof.functions.flatMap(_.find(_.isInstanceOf[PythonUDF])).head + .exists(_.exists(_.isInstanceOf[PythonUDF])) && + !PythonUDF.hasFreeLambdaVariable(hof) && + !(conf.pythonUDFInHigherOrderFunctionEnabled && + PythonUDF.canRewritePythonUDFInLambda(hof)) => + // Name the offending UDF: the first one of an unsupported eval type if any (e.g. a + // pandas UDF), otherwise the first Python UDF in the lambdas. + val udfs = hof.functions.flatMap(_.collect { case u: PythonUDF => u }) + val u = udfs.find(!PythonUDF.isElementwiseRewritableUDF(_)).getOrElse(udfs.head) hof.failAnalysis( errorClass = "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF", messageParameters = Map("funcName" -> toSQLExpr(u))) @@ -974,33 +993,17 @@ trait CheckAnalysis extends LookupCatalog with QueryErrorsBase with PlanToString operator match { case o if o.children.nonEmpty && o.missingInput.nonEmpty => - val missingAttributes = o.missingInput.map(attr => toSQLExpr(attr)).mkString(", ") - val input = o.inputSet.map(attr => toSQLExpr(attr)).mkString(", ") - val resolver = plan.conf.resolver val attrsWithSameName = o.missingInput.filter { missing => o.inputSet.exists(input => resolver(missing.name, input.name)) } - if (attrsWithSameName.nonEmpty) { - val sameNames = attrsWithSameName.map(attr => toSQLExpr(attr)).mkString(", ") - o.failAnalysis( - errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION", - messageParameters = Map( - "missingAttributes" -> missingAttributes, - "input" -> input, - "operator" -> operator.simpleString(SQLConf.get.maxToStringFields), - "operation" -> sameNames - )) - } else { - o.failAnalysis( - errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", - messageParameters = Map( - "missingAttributes" -> missingAttributes, - "input" -> input, - "operator" -> operator.simpleString(SQLConf.get.maxToStringFields) - )) - } + throw QueryCompilationErrors.missingAttributesError( + operator = o, + missingInput = o.missingInput, + input = o.inputSet, + attributesWithSameName = attrsWithSameName + ) case p @ Project(projectList, _) => checkForUnspecifiedWindow(projectList) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala index f43dfb77a9c00..ea1a1917d8f55 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CollationTypeCoercion.scala @@ -25,7 +25,10 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Proj import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLExpr import org.apache.spark.sql.errors.QueryCompilationErrors -import org.apache.spark.sql.types.{ArrayType, DataType, IndeterminateStringType, MapType, NullType, StringType, StructType} +import org.apache.spark.sql.types.{ + ArrayType, DataType, IndeterminateStringType, MapType, NullType, StringHelper, + StringType, StructType +} import org.apache.spark.sql.util.SchemaUtils /** @@ -107,6 +110,13 @@ object CollationTypeCoercion extends SQLConfHelper { /** * Changes the data type of the expression to the given `newType`. + * + * Never retarget an existing Cast (`cast.copy(dataType = ...)`). Explicit CAST + * truncation / overflow (ISO 6.13) and CHAR padding must stay on the inner node; + * LCT is an outer Cast. Uncollated TypeCoercion already nests. + * + * Literals: `copy(dataType)` is enough when only collation changes. When a string + * constraint changes (CHAR(2) to CHAR(4)), wrap in Cast so padding is re-applied. */ private def changeType(expr: Expression, newType: DataType): Expression = { mergeTypes(expr.dataType, newType) match { @@ -114,8 +124,11 @@ object CollationTypeCoercion extends SQLConfHelper { assert(!newDataType.existsRecursively(_.isInstanceOf[StringTypeWithContext])) expr match { + case lit: Literal if stringConstraintChanged(lit.dataType, newDataType) => + Cast(lit, newDataType, timeZoneId = Some(conf.sessionLocalTimeZone)) case lit: Literal => lit.copy(dataType = newDataType) - case cast: Cast => cast.copy(dataType = newDataType) + case cast: Cast => + Cast(cast, newDataType, timeZoneId = Some(conf.sessionLocalTimeZone)) case subquery: SubqueryExpression => changeTypeInSubquery(subquery, newType) @@ -127,6 +140,24 @@ object CollationTypeCoercion extends SQLConfHelper { } } + /** + * True when CHAR/VARCHAR length (or nested length) differs between `from` and `to`. + * Collation-only differences are not a constraint change. + */ + private def stringConstraintChanged(from: DataType, to: DataType): Boolean = { + (from, to) match { + case (f: StringType, t: StringType) => f.constraint != t.constraint + case (ArrayType(fe, _), ArrayType(te, _)) => stringConstraintChanged(fe, te) + case (MapType(fk, fv, _), MapType(tk, tv, _)) => + stringConstraintChanged(fk, tk) || stringConstraintChanged(fv, tv) + case (fs: StructType, ts: StructType) if fs.length == ts.length => + fs.fields.indices.exists { i => + stringConstraintChanged(fs.fields(i).dataType, ts.fields(i).dataType) + } + case _ => false + } + } + /** * Changes the data type of the expression in the subquery to the given `newType`. * Currently only supports subqueries with [[Project]] and [[Aggregate]] plan. @@ -414,7 +445,23 @@ object CollationTypeCoercion extends SQLConfHelper { } } - /** Determines the winning StringTypeWithContext based on the strength of the collation. */ + /** + * Resolves collation strength independently of CHAR/VARCHAR length. + * + * This rule always runs. First-class CHAR/VARCHAR appear whenever + * `charVarcharFirstClassTypes` is true (`standardSemantics` or + * `preserveCharVarcharTypeInfo`), not only under `standardSemantics`. + * + * Same collation, including mixed strength: take the string-family LCT `max(n, m)` + * (pads, never truncates) and attach the stronger strength. Example: + * `coalesce(CAST('a' AS CHAR(2) COLLATE UTF8_LCASE), + * CAST(1 AS CHAR(4) COLLATE UTF8_LCASE))` is CHAR(4) COLLATE UTF8_LCASE + * (Implicit CHAR(2) from a string CAST vs Default CHAR(4) from a non-string CAST). + * + * Different collations at equal strength: mismatch (error if Explicit, else + * indeterminate). Different collations at unequal strength: the stronger operand + * wins in full, including its length (SQL collation precedence). + */ private def getWinningStringType( left: StringTypeWithContext, right: StringTypeWithContext): StringTypeWithContext = { @@ -427,14 +474,13 @@ object CollationTypeCoercion extends SQLConfHelper { } } - (left.strength.priority, right.strength.priority) match { - case (leftPriority, rightPriority) if leftPriority == rightPriority => - if (left.sameType(right)) left - else handleMismatch() + val winner = + if (left.strength.priority <= right.strength.priority) left else right - case (leftPriority, rightPriority) => - if (leftPriority < rightPriority) left - else right + StringHelper.tightestCommonString(left.stringType, right.stringType) match { + case Some(lct) => StringTypeWithContext(lct, winner.strength) + case None if left.strength.priority == right.strength.priority => handleMismatch() + case None => winner } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala index 965330db58280..ca4225b645fb1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ +import org.apache.spark.sql.catalyst.expressions.ml._ import org.apache.spark.sql.catalyst.expressions.st._ import org.apache.spark.sql.catalyst.expressions.variant._ import org.apache.spark.sql.catalyst.expressions.xml._ @@ -473,6 +474,7 @@ object FunctionRegistry { expression[ToRadians]("radians"), expression[Rint]("rint"), expression[Round]("round"), + expression[Truncate]("truncate"), expression[ShiftLeft]("shiftleft"), expression[ShiftRight]("shiftright"), expression[ShiftRightUnsigned]("shiftrightunsigned"), @@ -551,6 +553,7 @@ object FunctionRegistry { expression[CollectList]("collect_list"), expression[CollectList]("array_agg", true, Some("3.3.0")), expression[CollectSet]("collect_set"), + expression[CollectUnion]("collect_union"), expression[ListAgg]("listagg"), expression[ListAgg]("string_agg", setAlias = true), expressionBuilder("count_min_sketch", CountMinSketchAggExpressionBuilder), @@ -609,6 +612,7 @@ object FunctionRegistry { expressionBuilder("startswith", StartsWithExpressionBuilder), expressionBuilder("endswith", EndsWithExpressionBuilder), expression[Base64]("base64"), + expression[Base32]("to_base32"), expression[BitLength]("bit_length"), expression[Length]("char_length", true, Some("2.3.0")), expression[Length]("character_length", true, Some("2.3.0")), @@ -668,6 +672,7 @@ object FunctionRegistry { expression[StringTrimBoth]("btrim"), expression[Upper]("ucase", true), expression[UnBase64]("unbase64"), + expression[UnBase32]("from_base32"), expression[Unhex]("unhex"), expression[Upper]("upper"), expression[XPathList]("xpath"), @@ -687,6 +692,7 @@ object FunctionRegistry { expression[ValidateUTF8]("validate_utf8"), expression[TryValidateUTF8]("try_validate_utf8"), expression[Quote]("quote"), + expression[Normalize]("normalize"), // url functions expression[UrlEncode]("url_encode"), @@ -698,11 +704,11 @@ object FunctionRegistry { expression[AddMonths]("add_months"), expression[CurrentDate]("current_date"), expressionBuilder("curdate", CurDateExpressionBuilder, setAlias = true), - expression[CurrentTimestamp]("current_timestamp"), + expressionBuilder("current_timestamp", CurrentTimestampExpressionBuilder), expression[CurrentTime]("current_time"), expression[CurrentTime]("localtime", since = Some("4.3.0")), expression[CurrentTimeZone]("current_timezone"), - expression[LocalTimestamp]("localtimestamp"), + expressionBuilder("localtimestamp", LocalTimestampExpressionBuilder), expression[DateDiff]("datediff"), expression[DateDiff]("date_diff", setAlias = true, Some("3.4.0")), expression[DateAdd]("date_add"), @@ -720,7 +726,7 @@ object FunctionRegistry { expression[Month]("month"), expression[MonthsBetween]("months_between"), expression[NextDay]("next_day"), - expression[Now]("now"), + expressionBuilder("now", NowExpressionBuilder), expression[Quarter]("quarter"), expressionBuilder("second", SecondExpressionBuilder), expression[ParseToTimestamp]("to_timestamp"), @@ -808,6 +814,7 @@ object FunctionRegistry { expression[MapConcat]("map_concat"), expression[Size]("size"), expression[Slice]("slice"), + expression[TrimArray]("trim_array"), expression[Size]("cardinality", true, Some("2.4.0")), expression[ArraysZip]("arrays_zip"), expression[SortArray]("sort_array"), @@ -846,6 +853,8 @@ object FunctionRegistry { expression[Uuid]("uuid"), expression[Murmur3Hash]("hash"), expression[XxHash64]("xxhash64"), + expression[Xxh364]("xxh3_64"), + expression[Xxh3128]("xxh3_128"), expression[Sha1]("sha", true), expression[Sha1]("sha1"), expression[Sha2]("sha2"), @@ -966,8 +975,13 @@ object FunctionRegistry { expression[BitmapBitPosition]("bitmap_bit_position"), expression[BitmapConstructAgg]("bitmap_construct_agg"), expression[BitmapCount]("bitmap_count"), + expression[BitmapAnd]("bitmap_and"), + expression[BitmapOr]("bitmap_or"), + expression[BitmapAndNot]("bitmap_andnot"), + expression[BitmapXor]("bitmap_xor"), expression[BitmapOrAgg]("bitmap_or_agg"), expression[BitmapAndAgg]("bitmap_and_agg"), + expression[BitmapXorAgg]("bitmap_xor_agg"), // json expression[StructsToJson]("to_json"), @@ -975,6 +989,7 @@ object FunctionRegistry { expression[SchemaOfJson]("schema_of_json"), expression[LengthOfJsonArray]("json_array_length"), expression[JsonObjectKeys]("json_object_keys"), + expression[JsonTypeof]("json_typeof"), // Variant expressionBuilder("parse_json", ParseJsonExpressionBuilder), @@ -985,6 +1000,8 @@ object FunctionRegistry { expression[SchemaOfVariant]("schema_of_variant"), expression[SchemaOfVariantAgg]("schema_of_variant_agg"), expression[ToVariantObject]("to_variant_object"), + expression[VariantFromArrays]("variant_from_arrays"), + expression[VariantFromEntries]("variant_from_entries"), expression[IsValidVariant]("is_valid_variant"), expression[VariantDelete]("variant_delete"), expressionBuilder("variant_insert", VariantInsertExpressionBuilder), @@ -993,6 +1010,7 @@ object FunctionRegistry { expressionBuilder("try_variant_set", TryVariantSetExpressionBuilder), expressionBuilder("variant_array_append", VariantArrayAppendExpressionBuilder), expressionBuilder("try_variant_array_append", TryVariantArrayAppendExpressionBuilder), + expressionBuilder("variant_strip_nulls", VariantStripNullsExpressionBuilder), // Spatial expression[ST_AsBinary]("st_asbinary"), @@ -1051,7 +1069,28 @@ object FunctionRegistry { fr } - val functionSet: Set[FunctionIdentifier] = builtin.listFunction().toSet + /** + * Builtin function identifiers known to SHOW FUNCTIONS / SessionCatalog. + * Starts as the catalyst [[builtin]] set; sql/core-only builtins (e.g. parse_sql) + * are added later via [[registerExtraBuiltin]]. + */ + @volatile private var _functionSet: Set[FunctionIdentifier] = builtin.listFunction().toSet + + def functionSet: Set[FunctionIdentifier] = _functionSet + + /** + * Registers a builtin that cannot live in catalyst (e.g. depends on SparkSqlParser + * in sql/core). Updates both [[builtin]] (so session clones / catalog reset see it) + * and [[functionSet]] (so SHOW FUNCTIONS classifies it as SYSTEM, not USER). + */ + private[sql] def registerExtraBuiltin( + name: String, + info: ExpressionInfo, + builder: FunctionBuilder): Unit = synchronized { + val id = builtinFunctionIdentifier(name) + builtin.registerFunction(id, info, builder) + _functionSet = _functionSet + id + } /** Registry for internal functions used by Connect and the Column API. */ private[sql] val internal: SimpleFunctionRegistry = @@ -1085,6 +1124,7 @@ object FunctionRegistry { registerInternalExpression[Days]("days") registerInternalExpression[Hours]("hours") registerInternalExpression[UnwrapUDT]("unwrap_udt") + registerInternalExpression[WrapUDT]("wrap_udt") registerInternalExpression[MonotonicallyIncreasingID]("distributed_id", setAlias = true) registerInternalExpression[DistributedSequenceID]("distributed_sequence_id") registerInternalExpression[PandasProduct]("pandas_product") @@ -1098,6 +1138,7 @@ object FunctionRegistry { registerInternalExpression[NullIndex]("null_index") registerInternalExpression[CastTimestampNTZToLong]("timestamp_ntz_to_long") registerInternalExpression[ArrayBinarySearch]("array_binary_search") + registerInternalExpression[VectorPosExplode]("vector_posexplode") private def makeExprInfoForVirtualOperator(name: String, usage: String): ExpressionInfo = { new ExpressionInfo( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala index 543e9da670a6e..8dac7858724da 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala @@ -530,8 +530,12 @@ class FunctionResolution( } // We get an aggregate function, we need to wrap it in an AggregateExpression. case agg: AggregateFunction => - // Note: PythonUDAF does not support these advanced clauses. - if (agg.isInstanceOf[PythonUDAF]) checkUnsupportedAggregateClause(agg, unresolvedFunc) + // Note: neither PythonUDAF nor the incremental PythonAggregate support these advanced + // clauses (DISTINCT / FILTER / ORDER BY / IGNORE NULLS). They have dedicated physical + // operators that do not honor them, so reject rather than silently drop the clause. + if (agg.isInstanceOf[PythonUDAF] || agg.isInstanceOf[PythonAggregate]) { + checkUnsupportedAggregateClause(agg, unresolvedFunc) + } // After parse, the functions not set the ordering within group yet. val newAgg = agg match { case owg: SupportsOrderingWithinGroup @@ -622,6 +626,7 @@ class FunctionResolution( case anyValue: AnyValue => anyValue.copy(ignoreNulls = ignoreNulls) case collectList: CollectList => collectList.copy(ignoreNulls = ignoreNulls) case collectSet: CollectSet => collectSet.copy(ignoreNulls = ignoreNulls) + case collectUnion: CollectUnion => collectUnion.copy(ignoreNulls = ignoreNulls) case _ if ignoreNulls => // Only fail for IGNORE NULLS; RESPECT NULLS is the default behavior throw QueryCompilationErrors.functionWithUnsupportedSyntaxError( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCache.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCache.scala index 770a5e780b24a..7c34ad8911f60 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCache.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCache.scala @@ -18,11 +18,18 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Identifier} +import org.apache.spark.sql.util.CaseInsensitiveStringMap private[sql] trait RelationCache { - def lookup(nameParts: Seq[String], resolver: Resolver): Option[LogicalPlan] + def lookup( + catalog: CatalogPlugin, + ident: Identifier, + tableId: Option[String], + stateOptions: CaseInsensitiveStringMap, + resolver: Resolver): Option[LogicalPlan] } private[sql] object RelationCache { - val empty: RelationCache = (_, _) => None + val empty: RelationCache = (_, _, _, _, _) => None } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCacheKey.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCacheKey.scala new file mode 100644 index 0000000000000..fc4217cdf701d --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationCacheKey.scala @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.analysis + +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Key for the per-query relation cache in [[AnalysisContext]], shared by [[RelationResolution]]. + * + * The complete option map is part of the key because the resolved relation carries it into scan + * and write planning. References may reuse the same `Table` when their table-state options match, + * but references with different options must not reuse the same cached relation. + */ +private[sql] case class RelationCacheKey( + nameParts: Seq[String], + timeTravelSpec: Option[TimeTravelSpec], + options: CaseInsensitiveStringMap) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala index 0a085fcc2971b..9a61138cd1807 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala @@ -17,8 +17,6 @@ package org.apache.spark.sql.catalyst.analysis -import scala.collection.mutable - import org.apache.spark.internal.Logging import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.SQLConfHelper @@ -61,11 +59,10 @@ class RelationResolution( with LookupCatalog with SQLConfHelper { - type CacheKey = (Seq[String], Option[TimeTravelSpec]) - val v1SessionCatalog = catalogManager.v1SessionCatalog - private def relationCache: mutable.Map[CacheKey, LogicalPlan] = AnalysisContext.get.relationCache + private def relationCache = AnalysisContext.get.relationCache + private def tableCache = AnalysisContext.get.tableCache /** * If we are resolving database objects (relations, functions, etc.) inside views, we may need to @@ -164,16 +161,26 @@ class RelationResolution( def resolveRelation( u: UnresolvedRelation, timeTravelSpec: Option[TimeTravelSpec] = None): Option[LogicalPlan] = { - val timeTravelSpecFromOptions = TimeTravelSpec.fromOptions( - u.options, - conf.getConf(SQLConf.TIME_TRAVEL_TIMESTAMP_KEY), - conf.getConf(SQLConf.TIME_TRAVEL_VERSION_KEY), - conf.sessionLocalTimeZone - ) + val isWriteTarget = u.options.containsKey(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) + val hasTimeTravelWriteOptions = + isWriteTarget && CatalogV2Util.containsTimeTravelOptions(u.options) + val timeTravelSpecFromOptions = if (hasTimeTravelWriteOptions) { + // Time travel applies to reads only. Defer the option check until the identifier is resolved + // so every write API reports the same qualified relation ID, without parsing the options as + // a read time-travel specification first. + None + } else { + TimeTravelSpec.fromOptions( + u.options, + conf.getConf(SQLConf.TIME_TRAVEL_TIMESTAMP_KEY), + conf.getConf(SQLConf.TIME_TRAVEL_VERSION_KEY), + conf.sessionLocalTimeZone) + } if (timeTravelSpec.nonEmpty && timeTravelSpecFromOptions.nonEmpty) { throw new AnalysisException("MULTIPLE_TIME_TRAVEL_SPEC", Map.empty[String, String]) } val finalTimeTravelSpec = timeTravelSpec.orElse(timeTravelSpecFromOptions) + val isTimeTravel = finalTimeTravelSpec.isDefined || hasTimeTravelWriteOptions val identifier = u.multipartIdentifier // system.session.v (3 parts): only local temp view by name; same as SessionCatalog matching. @@ -182,7 +189,7 @@ class RelationResolution( return resolveTempView( normalized, u.isStreaming, - finalTimeTravelSpec.isDefined + isTimeTravel ) } @@ -192,7 +199,7 @@ class RelationResolution( identifier.head.equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE)) { val viewNameOnly = Seq(identifier.last) val tempSession = () => - resolveTempView(viewNameOnly, u.isStreaming, finalTimeTravelSpec.isDefined) + resolveTempView(viewNameOnly, u.isStreaming, isTimeTravel) val persistentSessionDb = () => tryResolvePersistent(u, identifier, finalTimeTravelSpec) return if (conf.prioritizeSystemCatalog) { @@ -208,7 +215,7 @@ class RelationResolution( return resolveTempView( identifier, u.isStreaming, - finalTimeTravelSpec.isDefined + isTimeTravel ).orElse(tryResolvePersistent(u, identifier, finalTimeTravelSpec)) } @@ -218,7 +225,7 @@ class RelationResolution( for (step <- steps) { val result = step match { case SessionScopeStep => - resolveTempView(identifier, u.isStreaming, finalTimeTravelSpec.isDefined) + resolveTempView(identifier, u.isStreaming, isTimeTravel) case PersistentCatalogStep(prefix) => tryResolvePersistent(u, prefix ++ identifier, finalTimeTravelSpec) } @@ -236,26 +243,40 @@ class RelationResolution( finalTimeTravelSpec: Option[TimeTravelSpec]): Option[LogicalPlan] = { expandIdentifier(identifier) match { case CatalogAndIdentifier(catalog, ident) => - val key = toCacheKey(catalog, ident, finalTimeTravelSpec) val planId = u.getTagValue(LogicalPlan.PLAN_ID_TAG) val writePrivileges = u.options.get(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) val finalOptions = u.clearWritePrivileges.options + if (writePrivileges != null) { + CatalogV2Util.rejectTimeTravelOptionsForWrite(catalog, ident, finalOptions) + } + // Time travel applies to reads only; reject an explicit time-travel specification on a + // write target with a user-facing error. + if (finalTimeTravelSpec.nonEmpty && writePrivileges != null) { + throw QueryCompilationErrors.timeTravelUnsupportedError( + toSQLId(ident.toQualifiedNameParts(catalog))) + } + val key = toCacheKey(catalog, ident, finalTimeTravelSpec, finalOptions) // A reference that requires write privileges is never served from the per-query relation - // cache. The catalog authorizes the write in `loadTable(ident, writePrivileges)` below, and - // a cache hit would skip that call entirely. The hit happens whenever the write target is + // cache. The catalog authorizes the write during the uncached `loadTable` below, and a + // cache hit would skip that call entirely. The hit happens whenever the write target is // also read in the same statement -- the target is resolved after its query (see // `ResolveRelations`), so it finds the relation the query already put in the cache, e.g. // for `INSERT INTO t SELECT * FROM t`. + // + // The cache key includes the options, so a hit means the options already match and each + // reference's own bag is honored without re-applying it here. val cached = if (writePrivileges == null) relationCache.get(key) else None cached - // The per-query relation cache is not keyed by options. When the same table is referenced - // more than once in a single statement with different dynamic options (e.g. a self-join, - // or a second reference sharing the target's cache entry), a cache hit would otherwise - // reuse the first reference's options and silently drop this reference's. Re-apply this - // reference's options to the cached relation so each reference honors its own bag. - .map(applyOptions(_, finalOptions)) .map(adaptCachedRelation(_, planId)) .orElse { + lazy val tableKey = + toTableCacheKey(catalog, ident, finalTimeTravelSpec, finalOptions) + val pinnedTable = if (writePrivileges == null && catalog.isInstanceOf[TableCatalog]) { + tableCache.get(tableKey) + } else { + None + } + // For a `RelationCatalog` with no time-travel / write privileges, the single-RPC // `loadRelation` answers both "is there a table?" and "is there a view?" in one // call. Time-travel and write privileges apply to tables only, so for those the @@ -265,57 +286,66 @@ class RelationResolution( // Skip the table-side lookup entirely for view-only catalogs (no `TableCatalog` // mixin): `CatalogV2Util.loadTable` would call `asTableCatalog` and throw // MISSING_CATALOG_ABILITY.TABLES, masking the legitimate view-resolution path. - val relation: Option[Relation] = catalog match { - case mc: RelationCatalog if finalTimeTravelSpec.isEmpty && writePrivileges == null => - try { - Some(mc.loadRelation(ident)) - } catch { - case _: NoSuchTableException => None - } - case _ => - val tableSide: Option[Table] = if ( - CatalogV2Util.isSessionCatalog(catalog) || catalog.isInstanceOf[TableCatalog] - ) { - CatalogV2Util.loadTable( - catalog, - ident, - finalTimeTravelSpec, - Option(writePrivileges)) - } else { - None - } - // Fallback to ViewCatalog for catalogs that host views but where loadTable - // returned None (or was skipped because there's no TableCatalog mixin). - // Time-travel / write privileges only apply to tables, not views, so the - // fallback only fires when both are absent. - tableSide.orElse { - if (finalTimeTravelSpec.isEmpty && writePrivileges == null) { - catalog match { - case vc: ViewCatalog => - try { - Some(vc.loadView(ident)) - } catch { - case _: NoSuchViewException => None - } - case _ => None - } + val relation: Option[Relation] = pinnedTable.orElse { + catalog match { + case mc: RelationCatalog + if finalTimeTravelSpec.isEmpty && writePrivileges == null => + try { + Some(mc.loadRelation(ident)) + } catch { + case _: NoSuchTableException => None + } + case _ => + val tableSide: Option[Table] = if ( + CatalogV2Util.isSessionCatalog(catalog) || catalog.isInstanceOf[TableCatalog] + ) { + CatalogV2Util.loadTable( + catalog, + ident, + finalTimeTravelSpec, + Option(writePrivileges), + finalOptions) } else { None } - } + // Fallback to ViewCatalog for catalogs that host views but where loadTable + // returned None (or was skipped because there's no TableCatalog mixin). + // Time-travel / write privileges only apply to tables, not views, so the + // fallback only fires when both are absent. + tableSide.orElse { + if (finalTimeTravelSpec.isEmpty && writePrivileges == null) { + catalog match { + case vc: ViewCatalog => + try { + Some(vc.loadView(ident)) + } catch { + case _: NoSuchViewException => None + } + case _ => None + } + } else { + None + } + } + } } // `table` is `relation` filtered to tables only -- used for cache lookup since // we don't share-cache views. val table: Option[Table] = relation.collect { case t: Table => t } + // Reuse a cached relation only when its table identity and state options match. The + // returned relation still carries this read's complete option map. val sharedRelationCacheMatch = for { t <- table - if finalTimeTravelSpec.isEmpty && writePrivileges == null && !u.isStreaming - cached <- lookupSharedRelationCache(catalog, ident, t) + if pinnedTable.isEmpty && finalTimeTravelSpec.isEmpty && + writePrivileges == null && !u.isStreaming + cached <- lookupSharedRelationCache(catalog, ident, t, finalOptions) } yield { val updatedRelation = cached.copy(options = finalOptions) + updatedRelation.copyTagsFrom(cached) val nameParts = ident.toQualifiedNameParts(catalog) val aliasedRelation = SubqueryAlias(nameParts, updatedRelation) + tableCache.update(tableKey, cached.table) relationCache.update(key, aliasedRelation) adaptCachedRelation(aliasedRelation, planId) } @@ -328,6 +358,12 @@ class RelationResolution( finalOptions, u.isStreaming, finalTimeTravelSpec) + // A write target skips cache lookup above so authorization always runs, but its + // freshly loaded Table is still published for subsequent reads, matching the + // relation cache behavior below. + if (pinnedTable.isEmpty) { + table.foreach(tableCache.update(tableKey, _)) + } loaded.foreach(relationCache.update(key, _)) loaded.map(cloneWithPlanId(_, planId)) } @@ -354,7 +390,7 @@ class RelationResolution( val relation = if (u.isStreaming) { StreamingRelationV2( None, changelogTable.name, changelogTable, u.options, - changelogTable.columns.toAttributes, Some(catalog), Some(ident), None) + changelogTable.columns.toOutputAttributes, Some(catalog), Some(ident), None) } else { DataSourceV2Relation.create(changelogTable, Some(catalog), Some(ident), u.options) } @@ -366,21 +402,9 @@ class RelationResolution( private def lookupSharedRelationCache( catalog: CatalogPlugin, ident: Identifier, - table: Table): Option[DataSourceV2Relation] = { - CatalogV2Util.lookupCachedRelation(sharedRelationCache, catalog, ident, table, conf) - } - - /** - * Re-applies `options` to the relation in a cached plan. Every `relationCache` entry holds a - * single relation for its own identifier (a view's body is still unresolved when it is cached), - * so this cannot reach another table's relation. - */ - private def applyOptions( - cached: LogicalPlan, - options: CaseInsensitiveStringMap): LogicalPlan = cached transform { - case r: DataSourceV2Relation => r.copy(options = options) - case r: UnresolvedCatalogRelation => r.copy(options = options) - case r: StreamingRelationV2 => r.copy(extraOptions = options) + table: Table, + options: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { + CatalogV2Util.lookupCachedRelation(sharedRelationCache, catalog, ident, table, options, conf) } private def adaptCachedRelation(cached: LogicalPlan, planId: Option[Long]): LogicalPlan = { @@ -454,7 +478,7 @@ class RelationResolution( table.name, table, options, - table.columns.toAttributes, + table.columns.toOutputAttributes, Some(catalog), Some(ident), v1Fallback @@ -496,12 +520,32 @@ class RelationResolution( } private def getOrLoadRelation(ref: V2TableReference): LogicalPlan = { - val key = toCacheKey(ref.catalog, ref.identifier) + val key = toCacheKey(ref.catalog, ref.identifier, None, ref.options) relationCache.get(key) match { case Some(cached) => adaptCachedRelation(cached, ref) case None => - val relation = loadRelation(ref) + val catalog = catalogManager.catalog(ref.catalog.name).asTableCatalog + val tableKey = toTableCacheKey(catalog, ref.identifier, None, ref.options) + val relation = tableCache.get(tableKey) match { + case Some(pinnedTable) => + createRelation(ref, catalog, pinnedTable) + case None => + val table = CatalogV2Util.getTable(catalog, ref.identifier, options = ref.options) + val sharedCacheMatch = if (ref.context.sharedCacheable) { + lookupSharedRelationCache(catalog, ref.identifier, table, ref.options) + } else { + None + } + sharedCacheMatch match { + case Some(cached) => + tableCache.update(tableKey, cached.table) + adaptCachedRelation(cached, ref) + case None => + tableCache.update(tableKey, table) + createRelation(ref, catalog, table) + } + } relationCache.update(key, relation) relation } @@ -518,7 +562,17 @@ class RelationResolution( */ private def loadRelation(ref: V2TableReference): LogicalPlan = { val resolvedCatalog = catalogManager.catalog(ref.catalog.name).asTableCatalog + // Only WriteTargetContext gets here (the sole non-cacheable context); it is currently used by + // transactional streaming writes and does not retain required write privileges. Keep the + // legacy load unchanged; options and privileges must be handled together in a follow-up. val table = resolvedCatalog.loadTable(ref.identifier) + createRelation(ref, resolvedCatalog, table) + } + + private def createRelation( + ref: V2TableReference, + resolvedCatalog: TableCatalog, + table: Table): DataSourceV2Relation = { V2TableReferenceUtils.validateLoadedTable(table, ref) DataSourceV2Relation( table = table, @@ -556,8 +610,22 @@ class RelationResolution( private def toCacheKey( catalog: CatalogPlugin, ident: Identifier, - timeTravelSpec: Option[TimeTravelSpec] = None): CacheKey = { - ((catalog.name +: ident.namespace :+ ident.name).toImmutableArraySeq, timeTravelSpec) + timeTravelSpec: Option[TimeTravelSpec], + options: CaseInsensitiveStringMap): RelationCacheKey = { + val nameParts = (catalog.name +: ident.namespace :+ ident.name).toImmutableArraySeq + RelationCacheKey(nameParts, timeTravelSpec, options) + } + + private def toTableCacheKey( + catalog: CatalogPlugin, + ident: Identifier, + timeTravelSpec: Option[TimeTravelSpec], + options: CaseInsensitiveStringMap): TableCacheKey = { + TableCacheKey( + catalog, + ident, + timeTravelSpec, + CatalogV2Util.extractTableStateOptions(catalog, options)) } private def cloneWithPlanId(plan: LogicalPlan, planId: Option[Long]): LogicalPlan = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveAsOfJoin.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveAsOfJoin.scala index f09612d5902bd..f5f828fc6ed6f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveAsOfJoin.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveAsOfJoin.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.{ } import org.apache.spark.sql.catalyst.expressions.AttributeSet import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.MatchComparisonOperator import org.apache.spark.sql.catalyst.plans.logical.{AsOfJoin, LogicalPlan, Project} import org.apache.spark.sql.catalyst.plans.logical.AsOfJoin.MatchConditionTypes import org.apache.spark.sql.catalyst.rules.Rule @@ -73,24 +74,23 @@ object ResolveAsOfJoin extends Rule[LogicalPlan] with SQLConfHelper { } val resolvedJoin = (matchLeft, matchOp, matchRight) match { case (Some(leftExpr), Some(operator), Some(rightExpr)) => - AsOfJoinValidation.validateMatchConditionTableReferences( - joinBase, left, right, leftExpr, rightExpr) - if (leftExpr.resolved && rightExpr.resolved) { - AsOfJoinValidation.validateMatchConditionOperands(joinBase, leftExpr, rightExpr) - val (leftOperand, rightOperand, normalizedOp) = - AsOfJoin.normalizeMatchOperands(left, right, leftExpr, operator, rightExpr) - val (asOfCondition, orderExpression, leftSortExprs, rightSortExprs) = - AsOfJoin.materializeMatchComparison(leftOperand, rightOperand, normalizedOp) - joinBase.copy( - asOfCondition = asOfCondition, - orderExpression = orderExpression, - leftSortExprs = leftSortExprs, - rightSortExprs = rightSortExprs, - matchLeftOperand = None, - matchOperator = None, - matchRightOperand = None) - } else { - joinBase + AsOfJoinMatchConditionResolution.materialize( + join = joinBase, + leftSet = left.outputSet, + rightSet = right.outputSet, + leftOperand = leftExpr, + operator = operator, + rightOperand = rightExpr) match { + case Some(materialized) => + joinBase.copy( + asOfCondition = materialized.asOfCondition, + orderExpression = materialized.orderExpression, + leftSortExprs = materialized.leftSortExpressions, + rightSortExprs = materialized.rightSortExpressions, + matchLeftOperand = None, + matchOperator = None, + matchRightOperand = None) + case None => joinBase } case (None, None, None) => joinBase case _ => joinBase @@ -107,17 +107,75 @@ object ResolveAsOfJoin extends Rule[LogicalPlan] with SQLConfHelper { } } +/** + * The executable [[AsOfJoin]] fields that a SQL `MATCH_CONDITION` clause materializes into. + */ +private[analysis] case class MaterializedMatchCondition( + asOfCondition: Expression, + orderExpression: Expression, + leftSortExpressions: Seq[Expression], + rightSortExpressions: Seq[Expression]) + +/** + * Validates a SQL `MATCH_CONDITION` clause and materializes it into the executable [[AsOfJoin]] + * fields, shared by the fixed-point [[ResolveAsOfJoin]] rule and the single-pass + * [[org.apache.spark.sql.catalyst.analysis.resolver.AsOfJoinResolver]] so the two analyzers + * cannot diverge on match-condition validation or normalization. + */ +private[analysis] object AsOfJoinMatchConditionResolution { + + /** + * Returns [[None]] while either operand is still unresolved, which the fixed-point analyzer + * reaches on iterations before the operands resolve. Table-reference validation runs + * regardless, as it only needs the operands' attribute references. + */ + def materialize( + join: AsOfJoin, + leftSet: AttributeSet, + rightSet: AttributeSet, + leftOperand: Expression, + operator: MatchComparisonOperator, + rightOperand: Expression): Option[MaterializedMatchCondition] = { + AsOfJoinValidation.validateMatchConditionTableReferences( + join = join, + leftSet = leftSet, + rightSet = rightSet, + leftExpr = leftOperand, + rightExpr = rightOperand) + + if (leftOperand.resolved && rightOperand.resolved) { + AsOfJoinValidation.validateMatchConditionOperands(join, leftOperand, rightOperand) + val (normalizedLeft, normalizedRight, normalizedOperator) = + AsOfJoin.normalizeMatchOperands( + leftSet = leftSet, + rightSet = rightSet, + expr1 = leftOperand, + operator = operator, + expr2 = rightOperand) + val (asOfCondition, orderExpression, leftSortExpressions, rightSortExpressions) = + AsOfJoin.materializeMatchComparison( + leftOperand = normalizedLeft, + rightOperand = normalizedRight, + normalizedOp = normalizedOperator) + Some(MaterializedMatchCondition( + asOfCondition = asOfCondition, + orderExpression = orderExpression, + leftSortExpressions = leftSortExpressions, + rightSortExpressions = rightSortExpressions)) + } else { + None + } + } +} + private[analysis] object AsOfJoinValidation extends QueryErrorsBase { def validateMatchConditionTableReferences( join: AsOfJoin, - left: LogicalPlan, - right: LogicalPlan, + leftSet: AttributeSet, + rightSet: AttributeSet, leftExpr: Expression, rightExpr: Expression): Unit = { - val leftSet = left.outputSet - val rightSet = right.outputSet - def referencesBothJoinSides(refs: AttributeSet): Boolean = { refs.nonEmpty && refs.intersect(leftSet).nonEmpty && diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala index 185a5503b1107..6fc196774a048 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.catalyst.util.SparkCharVarcharUtils.replaceCharVarch import org.apache.spark.sql.connector.catalog._ import org.apache.spark.sql.errors.DataTypeErrors.toSQLId import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ /** @@ -127,6 +128,18 @@ class ResolveCatalogs(val catalogManager: CatalogManager) val resolvedIdentifier = resolveIdentifier(nameParts, allowTemp, Nil) c.copy(name = resolvedIdentifier) + case c @ CreateTableAsSelect( + UnresolvedIdentifier(nameParts, allowTemp), _, _, _, writeOptions, _, _) => + val resolvedIdentifier = resolveIdentifier(nameParts, allowTemp, Nil) + rejectTimeTravelOptionsForWrite(resolvedIdentifier, writeOptions) + c.copy(name = resolvedIdentifier) + + case r @ ReplaceTableAsSelect( + UnresolvedIdentifier(nameParts, allowTemp), _, _, _, writeOptions, _, _) => + val resolvedIdentifier = resolveIdentifier(nameParts, allowTemp, Nil) + rejectTimeTravelOptionsForWrite(resolvedIdentifier, writeOptions) + r.copy(name = resolvedIdentifier) + case UnresolvedIdentifier(nameParts, allowTemp) => resolveIdentifier(nameParts, allowTemp, Nil) @@ -143,7 +156,7 @@ class ResolveCatalogs(val catalogManager: CatalogManager) allowTemp: Boolean, columns: Seq[ColumnDefinition]): ResolvedIdentifier = { val columnOutput = columns.map { col => - val dataType = if (conf.preserveCharVarcharTypeInfo) { + val dataType = if (conf.charVarcharFirstClassTypes) { col.dataType } else { replaceCharVarcharWithString(col.dataType) @@ -159,6 +172,15 @@ class ResolveCatalogs(val catalogManager: CatalogManager) } } + private def rejectTimeTravelOptionsForWrite( + resolvedIdentifier: ResolvedIdentifier, + writeOptions: Map[String, String]): Unit = { + CatalogV2Util.rejectTimeTravelOptionsForWrite( + resolvedIdentifier.catalog, + resolvedIdentifier.identifier, + new CaseInsensitiveStringMap(writeOptions.asJava)) + } + private def isSystemBuiltinName(nameParts: Seq[String]): Boolean = { nameParts.length == 3 && nameParts(0).equalsIgnoreCase(CatalogManager.SYSTEM_CATALOG_NAME) && diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveInsertionBase.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveInsertionBase.scala index ad89005a093e9..fa8be3aa1f304 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveInsertionBase.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveInsertionBase.scala @@ -21,7 +21,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Cast} import org.apache.spark.sql.catalyst.plans.logical.{InsertIntoStatement, LogicalPlan, Project} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.errors.QueryCompilationErrors -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} import org.apache.spark.sql.util.SchemaUtils abstract class ResolveInsertionBase extends Rule[LogicalPlan] { @@ -52,6 +52,10 @@ abstract class ResolveInsertionBase extends Rule[LogicalPlan] { case (input: StructType, expected: StructType) => // Rename inner fields of the input column to pass the by-name INSERT analysis. Alias(Cast(queryOutputCol, renameFieldsInStruct(input, expected)), resolvedCol.name)() + case (input: ArrayType, expected: ArrayType) => + Alias(Cast(queryOutputCol, renameFieldsInType(input, expected)), resolvedCol.name)() + case (input: MapType, expected: MapType) => + Alias(Cast(queryOutputCol, renameFieldsInType(input, expected)), resolvedCol.name)() case _ => Alias(queryOutputCol, resolvedCol.name)() } @@ -62,16 +66,25 @@ abstract class ResolveInsertionBase extends Rule[LogicalPlan] { private def renameFieldsInStruct(input: StructType, expected: StructType): StructType = { if (input.length == expected.length) { val newFields = input.zip(expected).map { case (f1, f2) => - (f1.dataType, f2.dataType) match { - case (s1: StructType, s2: StructType) => - f1.copy(name = f2.name, dataType = renameFieldsInStruct(s1, s2)) - case _ => - f1.copy(name = f2.name) - } + f1.copy(name = f2.name, dataType = renameFieldsInType(f1.dataType, f2.dataType)) } StructType(newFields) } else { input } } + + // Recursively rename fields so that positional INSERT analysis applies at every nesting level, + // including structs inside arrays and maps. See SPARK-58816. + private def renameFieldsInType(input: DataType, expected: DataType): DataType = + (input, expected) match { + case (s1: StructType, s2: StructType) => + renameFieldsInStruct(s1, s2) + case (ArrayType(e1, n1), ArrayType(e2, _)) => + ArrayType(renameFieldsInType(e1, e2), n1) + case (MapType(k1, v1, n1), MapType(k2, v2, _)) => + MapType(renameFieldsInType(k1, k2), renameFieldsInType(v1, v2), n1) + case _ => + input + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSchemaEvolution.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSchemaEvolution.scala index 09f22c16bfd14..08237128358a2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSchemaEvolution.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSchemaEvolution.scala @@ -17,7 +17,6 @@ package org.apache.spark.sql.catalyst.analysis -import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal import org.apache.spark.SparkThrowable @@ -27,7 +26,7 @@ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.COMMAND import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap -import org.apache.spark.sql.connector.catalog.{Identifier, SupportsSchemaEvolution, Table, TableCatalog, TableChange} +import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, SupportsSchemaEvolution, Table, TableCatalog, TableChange} import org.apache.spark.sql.connector.catalog.TableChange.ColumnChange import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, ExtractV2CatalogAndIdentifier, ExtractV2Table} @@ -51,7 +50,8 @@ object ResolveSchemaEvolution extends Rule[LogicalPlan] { write.table match { case relation @ ExtractV2CatalogAndIdentifier(catalog, ident) => evolveSchema(catalog, ident, write.pendingSchemaChanges) - val newTable = catalog.loadTable(ident, write.writePrivileges.asJava) + val newTable = CatalogV2Util.loadTableForV2Write( + catalog, ident, write.writePrivileges, relation.options) val writeWithNewTarget = replaceWriteTarget(write, relation, newTable) val remainingChanges = writeWithNewTarget.pendingSchemaChanges diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala index 2a4b5693c705d..580af46533f30 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTableConstraints.scala @@ -80,10 +80,10 @@ class ResolveTableConstraints(val catalogManager: CatalogManager) extends Rule[L } /** - * For each user-provided generated column, add a CheckInvariant that validates the column - * value matches the generation expression. Auto-filled generated columns are excluded: - * ResolveOutputRelation strips their GENERATION_EXPRESSION metadata from the table output, - * so they won't appear here. + * For each generated column whose value was provided by the write, add a CheckInvariant that + * validates the column value matches the generation expression. Generated columns Spark + * auto-filled are excluded: ResolveOutputRelation marks them in the table output and their + * values are correct by construction. */ private def buildGeneratedColumnConstraints( r: DataSourceV2Relation, @@ -92,20 +92,16 @@ class ResolveTableConstraints(val catalogManager: CatalogManager) extends Rule[L return Seq.empty } - // Use V2 columns from the table to access both V2 expressions and SQL strings. - // Only add constraints for generated columns whose GENERATION_EXPRESSION metadata - // is still present in the table output -- ResolveOutputRelation strips the metadata - // from auto-filled columns so they are excluded here. - val v2Columns = r.table.columns() val resolver = catalogManager.v1SessionCatalog.conf.resolver - val userProvidedGenCols = r.output - .filter(attr => GeneratedColumn.isGeneratedColumn(attr.metadata)) + val autoFilledGenCols = r.output + .filter(GeneratedColumn.isAutoFilledGeneratedColumn) .map(_.name) .toSet - v2Columns.flatMap { col => + // Use V2 columns from the table to access both V2 expressions and SQL strings. + r.table.columns().flatMap { col => Option(col.columnGenerationExpression()) - .filter(_ => userProvidedGenCols.exists(n => resolver(n, col.name))) + .filter(_ => !autoFilledGenCols.exists(n => resolver(n, col.name))) .map { genExpr => val catalystExpr = buildGenerationCatalystExpression(genExpr) val colRef = UnresolvedAttribute.quoted(col.name) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala new file mode 100644 index 0000000000000..188b3e5ae22d0 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala @@ -0,0 +1,97 @@ +/* + * 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.catalyst.analysis + +import org.apache.spark.sql.catalyst.expressions.TranspiledPythonUDF +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.TRANSPILED_PYTHON_UDF +import org.apache.spark.sql.types.{BinaryType, BooleanType, DataType, DecimalType, NumericType, StringType} + +/** + * Prunes the per-input-type options carried by a [[TranspiledPythonUDF]] down to those whose + * declared categories match the resolved argument types. + * + * A Python operator such as `a + b` is overloaded for text, so the transpiler emits one option + * per input-type variant -- a numeric `Add` and a string `concat`, say -- each tagged with the + * input-type categories it expects. Those options are children of the node, so leaving a + * type-incompatible one in place (a numeric `Add` over string columns) would make `CheckAnalysis` + * reject the whole plan. We can only choose once the argument types are known, which is after + * reference resolution -- hence a rule here rather than in the builder, which runs at + * call-construction time before the columns are bound -- and we must run before `CheckAnalysis`. + * + * Matching is strict by category (a numeric option only for numeric columns, a string option only + * for string columns). We deliberately do not lean on implicit type coercion, which would, e.g., + * make a numeric `Add` "valid" over a string column and silently diverge from Python's + * `TypeError`. When no option matches, the list is emptied and `ConvertToCatalyst` falls back to + * the original Python UDF. + */ +object ResolveTranspiledPythonUDFOptions extends Rule[LogicalPlan] { + def apply(plan: LogicalPlan): LogicalPlan = { + if (!plan.containsPattern(TRANSPILED_PYTHON_UDF)) { + plan + } else { + plan.resolveOperatorsWithPruning(_.containsPattern(TRANSPILED_PYTHON_UDF)) { + case op if op.containsPattern(TRANSPILED_PYTHON_UDF) => + // Bottom-up so a nested TranspiledPythonUDF (a transpiled UDF feeding another) is pruned + // -- and thus resolved -- before its parent's input types are inspected. + op.transformExpressionsUpWithPruning(_.containsPattern(TRANSPILED_PYTHON_UDF)) { + case t: TranspiledPythonUDF + if t.optionInputCategories.nonEmpty && t.pythonUDFExpr.childrenResolved => + val argTypes = t.pythonUDFExpr.children.map(_.dataType) + val kept = t.transpiledOptions.zip(t.optionInputCategories).collect { + case (option, categories) if optionMatchesTypes(categories, argTypes) => option + } + t.copy(transpiledOptions = kept, optionInputCategories = Nil) + } + } + } + } + + // True when each declared category matches the corresponding argument type: + // "numeric" -> NumericType, "string" -> StringType, "bool" -> BooleanType, + // "binary" -> BinaryType. "string" matches only StringType (not BinaryType): a + // bytes/BinaryType column is tagged "binary" instead, so the string lowerings + // (e.g. `repeat`) never see it. Empty categories means "no restriction", so the + // option is kept. + // + // Two deliberate exclusions keep the transpiled semantics faithful to Python: + // - DecimalType is NOT "numeric": Python receives decimal.Decimal objects, + // which raise TypeError when mixed with float literals and carry different + // precision semantics than Spark's decimal arithmetic, so decimal columns + // fall back to interpreted Python. + // - "string" requires the default UTF8_BINARY collation: under a non-binary + // collation (e.g. UTF8_LCASE) Spark's `=`/`<`/`concat` follow collation + // rules while Python compares codepoints, so `'abc' == 'ABC'` would return + // true where Python returns False. + private def optionMatchesTypes(categories: Seq[String], argTypes: Seq[DataType]): Boolean = { + if (categories.isEmpty) { + true + } else if (categories.length != argTypes.length) { + false + } else { + categories.zip(argTypes).forall { + case ("numeric", dt) => dt.isInstanceOf[NumericType] && !dt.isInstanceOf[DecimalType] + case ("string", st: StringType) => st.isUTF8BinaryCollation + case ("bool", dt) => dt.isInstanceOf[BooleanType] + case ("binary", dt) => dt.isInstanceOf[BinaryType] + case _ => false + } + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala index c8795b8ad9b1f..8bb55b1822b16 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.sql.catalyst.expressions.{Alias, EqualNullSafe, Expression, Literal, Not} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, Filter, LogicalPlan, Project, ReplaceData, WriteDelta} +import org.apache.spark.sql.catalyst.trees.TreePattern.DELETE_FROM_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ import org.apache.spark.sql.connector.catalog.{SupportsDeleteV2, SupportsRowLevelOperations, TruncatableTable} import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -36,7 +37,8 @@ import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, Extr */ object RewriteDeleteFromTable extends RewriteRowLevelCommand { - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(DELETE_FROM_TABLE)) { case d @ DeleteFromTable(aliasedTable, cond) if d.resolved => EliminateSubqueryAliases(aliasedTable) match { case ExtractV2Table(_: TruncatableTable) if cond == TrueLiteral => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala index 714f815161aeb..1f8e9df4a1103 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala @@ -24,6 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, JoinType, LeftAnti, LeftOuter, RightOuter} import org.apache.spark.sql.catalyst.plans.logical.{DeleteAction, Filter, HintInfo, InsertAction, InsertOnlyMerge, Join, JoinHint, LogicalPlan, MergeAction, MergeIntoTable, MergeRows, NO_BROADCAST_AND_REPLICATION, Project, ReplaceData, UpdateAction, WriteDelta} import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Copy, Delete, Discard, Insert, Instruction, Keep, ROW_ID, Split, Update} +import org.apache.spark.sql.catalyst.trees.TreePattern.MERGE_INTO_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils.{COPY_OPERATION, INSERT_OPERATION, OPERATION_COLUMN, UPDATE_OPERATION} import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -42,7 +43,8 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper private final val ROW_FROM_SOURCE = "__row_from_source" private final val ROW_FROM_TARGET = "__row_from_target" - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(MERGE_INTO_TABLE)) { // aligned is false when schema evolution is pending (see ResolveRowLevelCommandAssignments) case m @ MergeIntoTable(aliasedTable, source, cond, matchedActions, notMatchedActions, notMatchedBySourceActions, _) if m.resolved && m.rewritable && m.aligned && diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala index 207b3e7217040..352ddc0b0acb1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, MetadataAttribute, Not, SubqueryExpression} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{Assignment, Expand, Filter, LogicalPlan, Project, ReplaceData, Union, UpdateTable, WriteDelta} +import org.apache.spark.sql.catalyst.trees.TreePattern.UPDATE_TABLE import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations import org.apache.spark.sql.connector.write.{RowLevelOperationTable, SupportsDelta} @@ -34,7 +35,8 @@ import org.apache.spark.sql.types.IntegerType */ object RewriteUpdateTable extends RewriteRowLevelCommand { - override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators { + override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsWithPruning( + _.containsPattern(UPDATE_TABLE)) { case u @ UpdateTable(aliasedTable, assignments, cond) if u.resolved && u.rewritable && u.aligned => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/StringPromotionTypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/StringPromotionTypeCoercion.scala index 5ca3736d2352b..d6e3df4442454 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/StringPromotionTypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/StringPromotionTypeCoercion.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.types.{ DoubleType, NullType, StringTypeExpression, + TimestampLTZNanosTypeExpression, TimestampType, TimestampTypeExpression } @@ -55,10 +56,20 @@ object StringPromotionTypeCoercion { // For equality between string and timestamp we cast the string to a timestamp // so that things like rounding of subsecond precision does not affect the comparison. + // Both micros and nanos are scoped to the LTZ family (TimestampType / TimestampLTZNanosType): + // the NTZ families need no arm here, because their equality reaches the general + // BinaryComparison arm below and findCommonTypeForBinaryComparison returns the config-blind + // nanos/micros common type for them -- the same Cast this arm would add. Only the LTZ arm is + // load-bearing, since without it LTZ equality would take the range path and be promoted to + // string under legacy castDatetimeToString. case p @ Equality(left @ StringTypeExpression(), right @ TimestampTypeExpression()) => p.withNewChildren(Seq(Cast(left, TimestampType), right)) case p @ Equality(left @ TimestampTypeExpression(), right @ StringTypeExpression()) => p.withNewChildren(Seq(left, Cast(right, TimestampType))) + case p @ Equality(left @ StringTypeExpression(), right @ TimestampLTZNanosTypeExpression()) => + p.withNewChildren(Seq(Cast(left, right.dataType), right)) + case p @ Equality(left @ TimestampLTZNanosTypeExpression(), right @ StringTypeExpression()) => + p.withNewChildren(Seq(left, Cast(right, left.dataType))) case p @ BinaryComparison(left, right) if findCommonTypeForBinaryComparison(left.dataType, right.dataType, conf).isDefined => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala new file mode 100644 index 0000000000000..186c4b9855a43 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.analysis + +import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Identifier} +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Key for the per-query table-state cache in [[AnalysisContext]]. + * + * Unlike [[RelationCacheKey]], this key contains only options declared to affect table state. This + * lets references retain different scan options while sharing one concrete table state. + */ +private[sql] case class TableCacheKey( + catalog: CatalogPlugin, + identifier: Identifier, + timeTravelSpec: Option[TimeTravelSpec], + stateOptions: CaseInsensitiveStringMap) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableOutputResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableOutputResolver.scala index b50c77eee4795..933513654dc15 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableOutputResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableOutputResolver.scala @@ -236,7 +236,8 @@ object TableOutputResolver extends SQLConfHelper with Logging { if (canWriteValue) { val nullCheckedValue = checkNullability(value, attr, conf, colPath) val casted = cast(nullCheckedValue, attrTypeWithoutCharVarchar, conf, colPath.quoted) - val exprWithStrLenCheck = if (conf.charVarcharAsString || !attrTypeHasCharVarchar) { + val exprWithStrLenCheck = if (!attrTypeHasCharVarchar || + !CharVarcharUtils.shouldApplyWriteSideLengthCheck(conf)) { casted } else { CharVarcharUtils.stringLengthCheck(casted, attr.dataType) @@ -333,7 +334,8 @@ object TableOutputResolver extends SQLConfHelper with Logging { expectedCol: Attribute, conf: SQLConf): NamedExpression = { val rawType = CharVarcharUtils.getRawType(expectedCol.metadata).getOrElse(expectedCol.dataType) - val checked = if (!conf.charVarcharAsString && CharVarcharUtils.hasCharVarchar(rawType)) { + val checked = if (CharVarcharUtils.hasCharVarchar(rawType) && + CharVarcharUtils.shouldApplyWriteSideLengthCheck(conf)) { val value = defaultExpr match { case a: Alias => a.child case other => other @@ -854,7 +856,8 @@ object TableOutputResolver extends SQLConfHelper with Logging { } else { val udtUnwrapped = unwrapUDT(queryExpr) val casted = cast(udtUnwrapped, attrTypeWithoutCharVarchar, conf, colPath.quoted) - if (conf.charVarcharAsString || !attrTypeHasCharVarchar) { + if (!attrTypeHasCharVarchar || + !CharVarcharUtils.shouldApplyWriteSideLengthCheck(conf)) { casted } else { CharVarcharUtils.stringLengthCheck(casted, tableAttr.dataType) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala index 53de166e69edf..755882325939f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercion.scala @@ -145,6 +145,14 @@ object TypeCoercion extends TypeCoercionBase { => if (conf.castDatetimeToString) Some(st) else Some(TimestampType) case (TimestampType, st: StringType) => if (conf.castDatetimeToString) Some(st) else Some(TimestampType) + // Mirror the micros TimestampType (LTZ) arms above: nanos TIMESTAMP_LTZ(p) honors + // castDatetimeToString. Nanos TIMESTAMP_NTZ(p) intentionally has no arm here and, exactly like + // micros TimestampNTZType, falls through to the config-blind canPromoteAsInBinaryComparison + // line below -- so the LTZ and NTZ families stay consistent across micros and nanos. + case (st: StringType, tsNanos: TimestampLTZNanosType) + => if (conf.castDatetimeToString) Some(st) else Some(tsNanos) + case (tsNanos: TimestampLTZNanosType, st: StringType) + => if (conf.castDatetimeToString) Some(st) else Some(tsNanos) case (st: StringType, NullType) => Some(st) case (NullType, st: StringType) => Some(st) @@ -194,6 +202,10 @@ object TypeCoercion extends TypeCoercionBase { private def implicitCast(inType: DataType, expectedType: AbstractDataType): Option[DataType] = { // Note that ret is nullable to avoid typing a lot of Some(...) in this local scope. // We wrap immediately an Option after this. + // CHAR/VARCHAR promotion is checked first: the acceptsType case below would otherwise + // accept the constrained type unchanged, since CharType and VarcharType extend StringType. + charVarcharToPlainString(inType, expectedType).foreach(dt => return Some(dt)) + @Nullable val ret: DataType = (inType, expectedType) match { // If the expected type is already a parent of the input type, no need to cast. case _ if expectedType.acceptsType(inType) => inType diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionHelper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionHelper.scala index 942a6be948d8e..f64f9a09a5e45 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionHelper.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionHelper.scala @@ -39,6 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.{ ImplicitCastInputTypes, In, InSubquery, + JsonTuple, Least, ListQuery, Literal, @@ -60,7 +61,11 @@ import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.types.{AbstractArrayType, StringTypeWithCollation} +import org.apache.spark.sql.internal.types.{ + AbstractArrayType, + AbstractStringType, + StringTypeWithCollation +} import org.apache.spark.sql.types.{ AbstractDataType, AnyDataType, @@ -80,6 +85,7 @@ import org.apache.spark.sql.types.{ IntegralType, MapType, NullType, + StringHelper, StringType, StringTypeExpression, StructType, @@ -89,7 +95,8 @@ import org.apache.spark.sql.types.{ TimestampNTZType, TimestampType, TimestampTypeExpression, - TimeType + TimeType, + TypeCollection } abstract class TypeCoercionHelper { @@ -133,6 +140,69 @@ abstract class TypeCoercionHelper { */ def implicitCast(e: Expression, expectedType: AbstractDataType): Option[Expression] + /** + * Where a plain string is expected, promote CHAR(n)/VARCHAR(n) to unbounded STRING the same way + * SHORT promotes to INT, and return the promoted type. + * + * CharType and VarcharType extend StringType, so an expectation such as + * `StringTypeWithCollation` accepts them as-is and the implicit cast rules leave the length + * constraint in place. Expressions that then require all their string inputs to share a single + * type (`overlay`, `string_agg`, ...) cannot unify CHAR(n) with STRING, and RuntimeReplaceable + * ones (`right`) build literals from the constrained type that no longer match their other + * branches. + * + * The expectation must actually mention a string type (or an array of strings). Promoting at an + * `AnyDataType` site would strip the length from pass-through expressions such as `max`, `lag`, + * and `element_at`, which are required to preserve CHAR/VARCHAR. + */ + protected def charVarcharToPlainString( + inType: DataType, + expectedType: AbstractDataType): Option[DataType] = { + if (!conf.charVarcharStandardSemantics) { + return None + } + inType match { + case st: StringType if !StringHelper.isPlainString(st) => + val plain = StringHelper.plainStringType(st) + if (expectsStringType(expectedType) && expectedType.acceptsType(plain)) { + Some(plain) + } else { + None + } + case ArrayType(et, containsNull) => + arrayElementExpectation(expectedType).flatMap { elemExpected => + charVarcharToPlainString(et, elemExpected).map(ArrayType(_, containsNull)) + } + case _ => None + } + } + + private def expectsStringType(expectedType: AbstractDataType): Boolean = expectedType match { + case _: StringType => true + case _: AbstractStringType => true + case TypeCollection(types) => types.exists(expectsStringType) + case _ => false + } + + private def arrayElementExpectation(expectedType: AbstractDataType): Option[AbstractDataType] = + expectedType match { + case AbstractArrayType(elem) => Some(elem) + case ArrayType(elem, _) => Some(elem) + case TypeCollection(types) => types.view.flatMap(arrayElementExpectation).headOption + case _ => None + } + + /** + * Concat/Elt stringify non-binary inputs. CHAR/VARCHAR go through + * [[charVarcharToPlainString]] so a collated constrained type promotes to unbounded STRING + * with the same collation. Non-strings still target the default UTF8_BINARY StringType. + */ + protected def implicitCastToString(e: Expression): Expression = { + charVarcharToPlainString(e.dataType, StringTypeWithCollation(supportsTrimCollation = true)) + .map(dt => if (dt == e.dataType) e else Cast(e, dt)) + .getOrElse(implicitCast(e, StringType).getOrElse(e)) + } + /** * Whether casting `from` as `to` is valid. */ @@ -244,6 +314,11 @@ abstract class TypeCoercionHelper { } } + // Fractional-seconds precision of the microsecond timestamp family. DATE has no time component + // and is treated as this precision when widening (so DATE <-> micro widens to a micro type and + // DATE <-> nanos to a nanos type). The nanos types carry their own precision in [7, 9]. + private final val MicrosPrecision = 6 + protected def findWiderDateTimeType(d1: DatetimeType, d2: DatetimeType): Option[DatetimeType] = (d1, d2) match { // Two TIME operands of differing fractional-seconds precision widen to the larger precision @@ -275,7 +350,6 @@ abstract class TypeCoercionHelper { // Fractional-seconds precision of the timestamp family (micros: 6, nanos: 7-9). DATE has no // time component and is treated as the micro precision (getOrElse) so that DATE <-> micro // widens to the micro type and DATE <-> nanos to the nanos type. - val MicrosPrecision = 6 def isLtz(d: DatetimeType): Boolean = TimestampFamily.isLtz(d) def isNtz(d: DatetimeType): Boolean = TimestampFamily.isNtz(d) def precisionOf(d: DatetimeType): Int = @@ -302,6 +376,36 @@ abstract class TypeCoercionHelper { } } + /** Whether `dt` is on the LTZ/NTZ timestamp fractional-precision axis (a micro or nanos type). */ + private def isTimestampFamily(dt: DataType): Boolean = + TimestampFamily.fractionalPrecision(dt).isDefined + + /** + * Common operand type for [[SubtractTimestamps]] over two differing timestamp-family operands. + * The operands are widened to the larger of the two fractional-second precisions (the micro types + * count as 6, the nanos types carry their own precision `p` in [7, 9]) and unified in one + * time-zone family: + * - a cross-family pair unifies in the no-time-zone (NTZ) family, mirroring the microsecond + * precedent where TIMESTAMP - TIMESTAMP_NTZ coerces both operands to TIMESTAMP_NTZ; + * - a same-family pair keeps that family, so a TIMESTAMP - TIMESTAMP_LTZ(p) style pair still + * subtracts in the session time zone (DST-aware) exactly as a pure LTZ pair does. + * The subtraction reads only each operand's epochMicros and always reports the difference on the + * microsecond grid (a DayTimeIntervalType in the default mode, a CalendarIntervalType when + * spark.sql.legacy.interval.enabled is set), so widening the precision never changes the numeric + * result -- it only keeps the two operands the same concrete type. For a pure-micro cross-family + * pair this returns TimestampNTZType, identical to the pre-nanos behavior. + */ + private def subtractTimestampsCommonType(dt1: DataType, dt2: DataType): DataType = { + val p = math.max( + TimestampFamily.fractionalPrecision(dt1).getOrElse(MicrosPrecision), + TimestampFamily.fractionalPrecision(dt2).getOrElse(MicrosPrecision)) + if (TimestampFamily.isLtz(dt1) && TimestampFamily.isLtz(dt2)) { + if (p <= MicrosPrecision) TimestampType else TimestampLTZNanosType(p) + } else { + if (p <= MicrosPrecision) TimestampNTZType else TimestampNTZNanosType(p) + } + } + /** * Type coercion helper that matches agaist [[In]] and [[InSubquery]] expressions in order to * type coerce LHS and RHS to expected types. @@ -478,9 +582,7 @@ abstract class TypeCoercionHelper { case c @ Concat(children) if conf.concatBinaryAsString || !children.map(_.dataType).forall(_ == BinaryType) => - val newChildren = c.children.map { e => - implicitCast(e, StringType).getOrElse(e) - } + val newChildren = c.children.map(implicitCastToString) c.copy(children = newChildren) case other => other } @@ -528,9 +630,7 @@ abstract class TypeCoercionHelper { val newInputs = if (conf.eltOutputAsString || !children.tail.map(_.dataType).forall(_ == BinaryType)) { - children.tail.map { e => - implicitCast(e, StringType).getOrElse(e) - } + children.tail.map(implicitCastToString) } else { children.tail } @@ -615,16 +715,35 @@ abstract class TypeCoercionHelper { } e.withNewChildren(children) + // JsonTuple validates its own input types and rejects non-string children with + // NON_STRING_TYPE, so it only takes the CHAR/VARCHAR promotion here. Do not fold this into + // the ExpectsInputTypes arm below: that would also apply the NullType rewrite and turn + // json_tuple(json, null) from an analysis error into a typed STRING null. + case j: JsonTuple => + val expected = StringTypeWithCollation(supportsTrimCollation = true) + val children = j.children.map { child => + charVarcharToPlainString(child.dataType, expected) + .map(dt => if (dt == child.dataType) child else Cast(child, dt)) + .getOrElse(child) + } + j.withNewChildren(children) + case e: ExpectsInputTypes if e.inputTypes.nonEmpty => // Convert NullType into some specific target type for ExpectsInputTypes that don't do - // general implicit casting. + // general implicit casting. Also promote CHAR/VARCHAR to STRING here: these + // expressions skip ImplicitCastInputTypes, so without this the length constraint would + // remain on the child. val children: Seq[Expression] = e.children.zip(e.inputTypes).map { case (in, expected) => - if (in.dataType == NullType && !expected.acceptsType(NullType)) { - Literal.create(null, expected.defaultConcreteType) - } else { - in - } + charVarcharToPlainString(in.dataType, expected) + .map(dt => if (dt == in.dataType) in else Cast(in, dt)) + .getOrElse { + if (in.dataType == NullType && !expected.acceptsType(NullType)) { + Literal.create(null, expected.defaultConcreteType) + } else { + in + } + } } e.withNewChildren(children) @@ -738,14 +857,18 @@ abstract class TypeCoercionHelper { d.copy(startDate = Cast(d.startDate, DateType)) case d @ DateSub(StringTypeExpression(), _) => d.copy(startDate = Cast(d.startDate, DateType)) - case s @ SubtractTimestamps(DateTypeExpression(), AnyTimestampTypeExpression(), _, _) => + case s @ SubtractTimestamps(DateTypeExpression(), r, _, _) + if isTimestampFamily(r.dataType) => s.copy(left = Cast(s.left, s.right.dataType)) - case s @ SubtractTimestamps(AnyTimestampTypeExpression(), DateTypeExpression(), _, _) => + case s @ SubtractTimestamps(l, DateTypeExpression(), _, _) + if isTimestampFamily(l.dataType) => s.copy(right = Cast(s.right, s.left.dataType)) - case s @ SubtractTimestamps(AnyTimestampTypeExpression(), AnyTimestampTypeExpression(), _, _) - if s.left.dataType != s.right.dataType => - val newLeft = castIfNotSameType(s.left, TimestampNTZType) - val newRight = castIfNotSameType(s.right, TimestampNTZType) + case s @ SubtractTimestamps(l, r, _, _) + if isTimestampFamily(l.dataType) && isTimestampFamily(r.dataType) && + l.dataType != r.dataType => + val commonType = subtractTimestampsCommonType(l.dataType, r.dataType) + val newLeft = castIfNotSameType(s.left, commonType) + val newRight = castIfNotSameType(s.right, commonType) s.copy(left = newLeft, right = newRight) case t @ TimestampAddInterval(StringTypeExpression(), _, _) => @@ -766,14 +889,18 @@ abstract class TypeCoercionHelper { case d @ DateSub(AnyTimestampTypeExpression(), _) => d.copy(startDate = Cast(d.startDate, DateType)) - case s @ SubtractTimestamps(DateTypeExpression(), AnyTimestampTypeExpression(), _, _) => + case s @ SubtractTimestamps(DateTypeExpression(), r, _, _) + if isTimestampFamily(r.dataType) => s.copy(left = Cast(s.left, s.right.dataType)) - case s @ SubtractTimestamps(AnyTimestampTypeExpression(), DateTypeExpression(), _, _) => + case s @ SubtractTimestamps(l, DateTypeExpression(), _, _) + if isTimestampFamily(l.dataType) => s.copy(right = Cast(s.right, s.left.dataType)) - case s @ SubtractTimestamps(AnyTimestampTypeExpression(), AnyTimestampTypeExpression(), _, _) - if s.left.dataType != s.right.dataType => - val newLeft = castIfNotSameType(s.left, TimestampNTZType) - val newRight = castIfNotSameType(s.right, TimestampNTZType) + case s @ SubtractTimestamps(l, r, _, _) + if isTimestampFamily(l.dataType) && isTimestampFamily(r.dataType) && + l.dataType != r.dataType => + val commonType = subtractTimestampsCommonType(l.dataType, r.dataType) + val newLeft = castIfNotSameType(s.left, commonType) + val newRight = castIfNotSameType(s.right, commonType) s.copy(left = newLeft, right = newRight) case other => other diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala index eddcee169e377..27f7599fa6536 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala @@ -23,7 +23,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{ANALYSIS_ERROR, QUERY_PLAN} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.ExtendedAnalysisException -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, CurrentDate, CurrentTimestampLike, Expression, GroupingSets, LocalTimestamp, MonotonicallyIncreasingID, NamedExpression, SessionWindow, WindowExpression} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, CurrentDate, CurrentTimestampLike, Expression, GroupingSets, LocalTimestamp, LocalTimestampNanos, MonotonicallyIncreasingID, NamedExpression, SessionWindow, WindowExpression} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ @@ -637,7 +637,8 @@ object UnsupportedOperationChecker extends Logging { subPlan.expressions.foreach { e => if (e.collectLeaves().exists { - case (_: CurrentTimestampLike | _: CurrentDate | _: LocalTimestamp) => true + case (_: CurrentTimestampLike | _: CurrentDate | _: LocalTimestamp | + _: LocalTimestampNanos) => true case _ => false }) { throwError(s"Continuous processing does not support current time operations.") @@ -656,6 +657,18 @@ object UnsupportedOperationChecker extends Logging { if (outputMode != InternalOutputModes.Update) { throwRealTimeError("OUTPUT_MODE_NOT_SUPPORTED", Map("outputMode" -> outputMode.toString)) } + + plan.foreachUp { + case u: Union => + // Block stateful operators before union + u.foreachUp { + case statefulOp @ (_: Aggregate | _: TransformWithState | + _: TransformWithStateInPySpark | _: Deduplicate) if statefulOp.isStateful => + throwRealTimeError("STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED", Map.empty) + case _ => + } + case _ => + } } private def throwRealTimeError(subClass: String, args: Map[String, String]): Unit = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/V2TableReference.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/V2TableReference.scala index 7bad6d149c602..2d4ef18b5f600 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/V2TableReference.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/V2TableReference.scala @@ -86,22 +86,29 @@ private[sql] object V2TableReference { metadataColumns: Seq[MetadataColumn]) sealed trait Context { + /** Whether re-resolution may reuse the per-query relation cache. */ def cacheable: Boolean + + /** Whether re-resolution may reuse the shared (CACHE TABLE) relation cache. */ + def sharedCacheable: Boolean } /** Context for relations that are re-resolved on access of a dataframe temp view. */ case class TemporaryViewContext(viewName: Seq[String]) extends Context { val cacheable = true + val sharedCacheable = true } /** Context for relations that are re-resolved through a transaction catalog. */ case object TransactionContext extends Context { val cacheable = true + val sharedCacheable = false } /** Context for write targets. */ case object WriteTargetContext extends Context { val cacheable = false + val sharedCacheable = false } def createForTempView(relation: DataSourceV2Relation, viewName: Seq[String]): V2TableReference = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AggregateResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AggregateResolver.scala index cd51f17c9766e..da9083f99fd20 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AggregateResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AggregateResolver.scala @@ -28,6 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{ Alias, AliasHelper, AttributeReference, + BaseGroupingSets, Expression, ExprId, IntegerLiteral, @@ -51,6 +52,8 @@ class AggregateResolver( private val operatorResolutionContextStack = operatorResolver.getOperatorResolutionContextStack private val lcaResolver = expressionResolver.getLcaResolver private val ordinalResolver = expressionResolver.getOrdinalResolver + private val groupingAnalyticsResolver = + new GroupingAnalyticsResolver(operatorResolver, expressionResolver) /** * Resolve [[Aggregate]] operator. @@ -144,25 +147,36 @@ class AggregateResolver( ) if (resolvedAggregateExpressions.hasLateralColumnAlias) { + // LCA + grouping analytics (CUBE/ROLLUP/GROUPING SETS) is not yet supported in the + // single-pass resolver because Expand mints new attribute IDs that get out of sync + // with the Project operator. Fall back to legacy analyzer for correct results. + if (finalAggregate.groupingExpressions.exists(_.isInstanceOf[BaseGroupingSets])) { + throw new ExplicitlyUnsupportedResolverFeature( + "lateral column alias with grouping analytics") + } val aggregateWithLcaResolutionResult = lcaResolver.handleLcaInAggregate(finalAggregate) + val lcaBaseAggregate = aggregateWithLcaResolutionResult.baseAggregate AggregateResolutionResult( operator = aggregateWithLcaResolutionResult.resolvedOperator, outputList = aggregateWithLcaResolutionResult.outputList, groupingAttributeIds = - getGroupingAttributeIds(aggregateWithLcaResolutionResult.baseAggregate), + getGroupingAttributeIds(lcaBaseAggregate), aggregateListAliases = aggregateWithLcaResolutionResult.aggregateListAliases, - baseAggregate = aggregateWithLcaResolutionResult.baseAggregate + baseAggregate = lcaBaseAggregate ) } else { - AggregationValidator(finalAggregate) + // Grouping analytics (CUBE/ROLLUP/GROUPING SETS) are expanded unconditionally. + // groupingAnalyticsResolver.resolve no-ops when no BaseGroupingSets are present. + val expandedAggregate = groupingAnalyticsResolver.resolve(finalAggregate) + AggregationValidator(expandedAggregate) AggregateResolutionResult( - operator = finalAggregate, - outputList = finalAggregate.aggregateExpressions, - groupingAttributeIds = getGroupingAttributeIds(finalAggregate), + operator = expandedAggregate, + outputList = expandedAggregate.aggregateExpressions, + groupingAttributeIds = getGroupingAttributeIds(expandedAggregate), aggregateListAliases = scopes.current.getTopAggregateExpressionAliases, - baseAggregate = finalAggregate + baseAggregate = expandedAggregate ) } } finally { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AsOfJoinResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AsOfJoinResolver.scala new file mode 100644 index 0000000000000..d7216e4f304fa --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/AsOfJoinResolver.scala @@ -0,0 +1,367 @@ +/* + * 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.catalyst.analysis.resolver + +import org.apache.spark.sql.catalyst.analysis.{ + AnalysisErrorAt, + AsOfJoinMatchConditionResolution, + MaterializedMatchCondition, + NaturalAndUsingJoinResolution +} +import org.apache.spark.sql.catalyst.expressions.{ + Attribute, + AttributeSet, + Expression, + LambdaFunction, + NamedExpression +} +import org.apache.spark.sql.catalyst.plans.JoinType +import org.apache.spark.sql.catalyst.plans.logical.{AsOfJoin, LogicalPlan, Project} +import org.apache.spark.sql.catalyst.util._ + +/** + * Resolves [[AsOfJoin]] operators, including SQL `MATCH_CONDITION` and `USING` clauses. + * + * An ASOF JOIN matches each left-side row to the closest right-side row whose order key is + * less than or equal to (or, for a reversed match condition, greater than or equal to) the + * left row's order key, optionally within a tolerance and equality keys (`USING` / `ON`). + * Unlike a regular join, at most one right row is chosen per left row based on temporal (or + * other ordered) proximity rather than arbitrary key equality. + * + * {{{ + * SELECT * FROM trades ASOF JOIN quotes + * MATCH_CONDITION (trades.time >= quotes.time) USING (symbol) + * + * Project [symbol, time, ...] + * +- AsOfJoin (trades.time >= quotes.time), (trades.symbol = quotes.symbol), Inner + * }}} + */ +class AsOfJoinResolver( + override val resolver: Resolver, + override val expressionResolver: ExpressionResolver) + extends TreeNodeResolver[AsOfJoin, LogicalPlan] + with JoinLikeResolver { + + override def resolve(unresolvedAsOfJoin: AsOfJoin): LogicalPlan = { + val (resolvedLeft, leftNameScope) = resolveJoinChild( + unresolvedOperator = unresolvedAsOfJoin, + child = unresolvedAsOfJoin.left + ) + + val (resolvedRight, rightNameScope) = resolveJoinChild( + unresolvedOperator = unresolvedAsOfJoin, + child = unresolvedAsOfJoin.right + ) + + ExpressionIdAssigner.assertOutputsHaveNoConflictingExpressionIds( + Seq(leftNameScope.output, rightNameScope.output) + ) + + expressionIdAssigner.createMappingFromChildMappings( + newOutputIds = leftNameScope.getOutputIds ++ rightNameScope.getOutputIds + ) + + val partiallyResolved = unresolvedAsOfJoin.copy( + left = resolvedLeft, + right = resolvedRight + ) + + val resolvedCondition = resolveJoinCondition( + unresolvedJoin = unresolvedAsOfJoin, + unresolvedCondition = partiallyResolved.condition, + leftNameScope = leftNameScope, + rightNameScope = rightNameScope, + collectInvalidExpressions = true + ) + + val (usingOutput, conditionWithUsingColumns) = resolveUsingColumns( + unresolvedAsOfJoin = unresolvedAsOfJoin, + partiallyResolved = partiallyResolved, + resolvedCondition = resolvedCondition, + leftNameScope = leftNameScope, + rightNameScope = rightNameScope + ) + + val resolvedJoin = resolveAsOfExpressions( + unresolvedAsOfJoin = unresolvedAsOfJoin, + partiallyResolved = partiallyResolved.copy( + condition = conditionWithUsingColumns, + usingColumns = None + ), + leftNameScope = leftNameScope, + rightNameScope = rightNameScope + ) + + usingOutput match { + case Some((outputList, hiddenList)) => + buildUsingProject( + unresolvedAsOfJoin = unresolvedAsOfJoin, + resolvedJoin = resolvedJoin, + outputList = outputList, + hiddenList = hiddenList, + rightNameScope = rightNameScope + ) + case None => + overwriteJoinOutputScope( + joinType = unresolvedAsOfJoin.joinType, + leftNameScope = leftNameScope, + rightNameScope = rightNameScope + ) + cteRegistry.currentScope.tryPutWithCTE( + unresolvedOperator = unresolvedAsOfJoin, + resolvedOperator = resolvedJoin + ) + } + } + + private def resolveUsingColumns( + unresolvedAsOfJoin: AsOfJoin, + partiallyResolved: AsOfJoin, + resolvedCondition: Option[Expression], + leftNameScope: NameScope, + rightNameScope: NameScope) + : (Option[(Seq[NamedExpression], Seq[Attribute])], Option[Expression]) = { + partiallyResolved.usingColumns match { + case Some(columns) if resolvedCondition.isEmpty => + val (outputList, hiddenList, newCondition) = + NaturalAndUsingJoinResolution.computeJoinOutputsAndNewCondition( + left = partiallyResolved.left, + leftOutput = leftNameScope.output, + right = partiallyResolved.right, + rightOutput = rightNameScope.output, + joinType = partiallyResolved.joinType, + joinNames = columns, + condition = None, + resolveName = conf.resolver + ) + val resolvedUsingCondition = resolveJoinCondition( + unresolvedJoin = unresolvedAsOfJoin, + unresolvedCondition = newCondition, + leftNameScope = leftNameScope, + rightNameScope = rightNameScope, + collectInvalidExpressions = true + ) + (Some((outputList, hiddenList)), resolvedUsingCondition) + case _ => + (None, resolvedCondition) + } + } + + private def buildUsingProject( + unresolvedAsOfJoin: AsOfJoin, + resolvedJoin: AsOfJoin, + outputList: Seq[NamedExpression], + hiddenList: Seq[Attribute], + rightNameScope: NameScope): Project = { + val resolvedOutputList = outputList.map { expression => + resolveExpressionInJoin(unresolvedAsOfJoin, expression).asInstanceOf[NamedExpression] + } + val outputAttributes = resolvedOutputList.map(_.toAttribute) + val filteredHiddenOutput = filterHiddenOutputMetadataForJoin( + joinType = unresolvedAsOfJoin.joinType, + oldHiddenOutput = scopes.current.hiddenOutput, + rightHiddenOutput = rightNameScope.hiddenOutput + ) + val newHiddenOutput = computeHiddenOutputForJoin( + mainOutput = outputAttributes, + oldHiddenOutput = filteredHiddenOutput, + extraHiddenOutput = hiddenList + ) + scopes.overwriteCurrent( + output = Some(outputAttributes), + hiddenOutput = Some(newHiddenOutput) + ) + + val qualifiedAccessOnlyColumns = newHiddenOutput.filter(_.qualifiedAccessOnly) + val projectList = + if (unresolvedAsOfJoin.containsTag(ResolverTag.TOP_LEVEL_OPERATOR)) { + resolvedOutputList + } else { + resolvedOutputList ++ qualifiedAccessOnlyColumns + } + + operatorResolutionContextStack.current.baseOperator = Some(resolvedJoin) + val project = Project(projectList, resolvedJoin) + project.setTagValue(Project.hiddenOutputTag, qualifiedAccessOnlyColumns) + project + } + + private def resolveAsOfExpressions( + unresolvedAsOfJoin: AsOfJoin, + partiallyResolved: AsOfJoin, + leftNameScope: NameScope, + rightNameScope: NameScope): AsOfJoin = { + ( + partiallyResolved.matchLeftOperand, + partiallyResolved.matchOperator, + partiallyResolved.matchRightOperand + ) match { + case (Some(unresolvedLeftOperand), Some(operator), Some(unresolvedRightOperand)) => + val leftOperand = resolveExpressionInJoin(unresolvedAsOfJoin, unresolvedLeftOperand) + val rightOperand = resolveExpressionInJoin(unresolvedAsOfJoin, unresolvedRightOperand) + AsOfJoinMatchConditionResolution.materialize( + join = partiallyResolved, + leftSet = AttributeSet(leftNameScope.output), + rightSet = AttributeSet(rightNameScope.output), + leftOperand = leftOperand, + operator = operator, + rightOperand = rightOperand + ) match { + case Some(materialized) => + materializedMatchConditionToJoin( + unresolvedAsOfJoin = unresolvedAsOfJoin, + partiallyResolved = partiallyResolved, + materialized = materialized + ) + case None => + partiallyResolved + } + case (None, None, None) => + resolvePreMaterializedExpressions(unresolvedAsOfJoin, partiallyResolved) + case _ => + partiallyResolved + } + } + + /** + * Places the materialized `MATCH_CONDITION` expressions on the join. They are freshly + * constructed trees, so each one still has to go through expression resolution before it can be + * placed in the resolved plan. + */ + private def materializedMatchConditionToJoin( + unresolvedAsOfJoin: AsOfJoin, + partiallyResolved: AsOfJoin, + materialized: MaterializedMatchCondition): AsOfJoin = { + throwIfLambdaBasedOrdering(materialized.orderExpression) + + val asOfCondition = resolveExpressionInJoin(unresolvedAsOfJoin, materialized.asOfCondition) + val orderExpression = + resolveExpressionInJoin(unresolvedAsOfJoin, materialized.orderExpression) + val leftSortExpressions = + materialized.leftSortExpressions.map(resolveExpressionInJoin(unresolvedAsOfJoin, _)) + val rightSortExpressions = + materialized.rightSortExpressions.map(resolveExpressionInJoin(unresolvedAsOfJoin, _)) + + partiallyResolved.copy( + asOfCondition = asOfCondition, + orderExpression = orderExpression, + matchLeftOperand = None, + matchOperator = None, + matchRightOperand = None, + leftSortExprs = leftSortExpressions, + rightSortExprs = rightSortExpressions + ) + } + + private def resolvePreMaterializedExpressions( + unresolvedAsOfJoin: AsOfJoin, + partiallyResolved: AsOfJoin): AsOfJoin = { + throwIfLambdaBasedOrdering(partiallyResolved.orderExpression) + + val resolvedTolerance = + partiallyResolved.toleranceAssertion.map(resolveExpressionInJoin(unresolvedAsOfJoin, _)) + validateTolerance(unresolvedAsOfJoin, resolvedTolerance) + + val asOfCondition = + resolveExpressionInJoin(unresolvedAsOfJoin, partiallyResolved.asOfCondition) + val orderExpression = + resolveExpressionInJoin(unresolvedAsOfJoin, partiallyResolved.orderExpression) + val leftSortExpressions = + partiallyResolved.leftSortExprs.map(resolveExpressionInJoin(unresolvedAsOfJoin, _)) + val rightSortExpressions = + partiallyResolved.rightSortExprs.map(resolveExpressionInJoin(unresolvedAsOfJoin, _)) + + partiallyResolved.copy( + asOfCondition = asOfCondition, + orderExpression = orderExpression, + toleranceAssertion = resolvedTolerance, + leftSortExprs = leftSortExpressions, + rightSortExprs = rightSortExpressions + ) + } + + /** + * `MATCH_CONDITION` operands that are ordered element-wise (`ARRAY` operands) materialize into + * an ordering expression built on top of a [[LambdaFunction]]: + * + * {{{ + * -- MATCH_CONDITION (t.a >= r.a) with ARRAY<INT> operands + * zip_with(t.a, r.a, lambdafunction((lambda left_elem - lambda right_elem), ...)) + * }}} + * + * The single-pass [[ExpressionResolver]] doesn't support lambda expressions, so bail out and + * let the fixed-point analyzer resolve those queries. + */ + private def throwIfLambdaBasedOrdering(orderExpression: Expression): Unit = { + if (orderExpression.exists(_.isInstanceOf[LambdaFunction])) { + throw new ExplicitlyUnsupportedResolverFeature( + "MATCH_CONDITION with a lambda-based ordering expression" + ) + } + } + + private def resolveExpressionInJoin( + unresolvedAsOfJoin: AsOfJoin, + unresolvedExpression: Expression): Expression = { + expressionResolver.resolveExpressionTreeInOperator( + unresolvedExpression, + unresolvedAsOfJoin + ) + } + + private def validateTolerance( + unresolvedAsOfJoin: AsOfJoin, + toleranceAssertion: Option[Expression]): Unit = { + toleranceAssertion.foreach { assertion => + if (!assertion.foldable) { + unresolvedAsOfJoin.failAnalysis( + errorClass = "AS_OF_JOIN.TOLERANCE_IS_UNFOLDABLE", + messageParameters = Map.empty + ) + } + if (!assertion.eval().asInstanceOf[Boolean]) { + unresolvedAsOfJoin.failAnalysis( + errorClass = "AS_OF_JOIN.TOLERANCE_IS_NON_NEGATIVE", + messageParameters = Map.empty + ) + } + } + } + + private def overwriteJoinOutputScope( + joinType: JoinType, + leftNameScope: NameScope, + rightNameScope: NameScope): Unit = { + val newOutput = AsOfJoin.computeOutput( + joinType = joinType, + leftOutput = leftNameScope.output, + rightOutput = rightNameScope.output + ) + + val filteredHiddenOutput = filterHiddenOutputMetadataForJoin( + joinType = joinType, + oldHiddenOutput = scopes.current.hiddenOutput, + rightHiddenOutput = rightNameScope.hiddenOutput + ) + + scopes.overwriteCurrent( + output = Some(newOutput), + hiddenOutput = Some(computeHiddenOutputForJoin(newOutput, filteredHiddenOutput)) + ) + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HavingResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HavingResolver.scala index 93c0eabb94873..ddf930ac7b1ed 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HavingResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HavingResolver.scala @@ -49,6 +49,14 @@ class HavingResolver(resolver: Resolver, expressionResolver: ExpressionResolver) override def resolve(unresolvedHaving: UnresolvedHaving): LogicalPlan = { val resolvedChild = resolver.resolve(unresolvedHaving.child) + // HAVING over grouping analytics (CUBE/ROLLUP/GROUPING SETS) is not yet supported in + // single-pass because the expanded Aggregate's ExprIds get out of sync when HAVING + // inserts missing expressions. Fall back to legacy for correct results (SPARK-57346). + if (operatorResolutionContextStack.current.hasGroupingAnalytics) { + throw new ExplicitlyUnsupportedResolverFeature( + "HAVING with grouping analytics (SPARK-57346)") + } + resolvedChild match { case window: Window if scopes.current.baseAggregate.isDefined => resolveHavingAboveWindow( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HybridAnalyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HybridAnalyzer.scala index 3fc6438597300..ad3d88e42afcc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HybridAnalyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/HybridAnalyzer.scala @@ -354,7 +354,8 @@ object HybridAnalyzer { SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_RELATION_BRIDGING_ENABLED.key, SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_RUN_EXTENDED_RESOLUTION_CHECKS.key, SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_THROW_FROM_RESOLVER_GUARD.key, - SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_VALIDATION_ENABLED.key + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_VALIDATION_ENABLED.key, + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLE_ASOF_JOIN_RESOLUTION.key ) /** diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinLikeResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinLikeResolver.scala new file mode 100644 index 0000000000000..9cc87bf7f00d9 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinLikeResolver.scala @@ -0,0 +1,275 @@ +/* + * 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.catalyst.analysis.resolver + +import java.util.HashSet + +import org.apache.spark.sql.catalyst.SQLConfHelper +import org.apache.spark.sql.catalyst.analysis.{withPosition, AnalysisErrorAt} +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, ExprId} +import org.apache.spark.sql.catalyst.plans.{ + ExistenceJoin, + FullOuter, + JoinType, + LeftAnti, + LeftOuter, + LeftSemi, + LeftSingle, + RightOuter +} +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.util._ +import org.apache.spark.sql.errors.QueryErrorsBase +import org.apache.spark.sql.types.BooleanType + +/** + * Shared resolution mechanics for join-like binary operators, mixed into [[JoinResolver]] and + * other join-like resolvers: resolving two children in isolated multi-child scopes, computing a + * combined hidden output, filtering hidden metadata columns by join type, and resolving a boolean + * join condition. These routines carry subtle scope, CTE, hidden-output, metadata-nullability and + * [[ExprId]] behavior; keeping them here ensures join-like resolvers stay aligned and fixes cannot + * drift between them. + * + * For example, in the following query: + * + * {{{ + * SELECT * FROM t1 JOIN t2 ON t1.key = t2.key; + * }}} + * + * the plan is: + * + * {{{ + * Project [key#1, key#2] + * +- Join Inner, (key#1 = key#2) + * :- SubqueryAlias t1 + * : +- Relation t1[key#1] + * +- SubqueryAlias t2 + * +- Relation t2[key#2] + * }}} + * + * `t1` and `t2` are each resolved by [[resolveJoinChild]] in their own [[NameScope]], and the + * condition `key#1 = key#2` is resolved by [[resolveJoinCondition]] against the union of the two + * child scopes. + */ +trait JoinLikeResolver extends SQLConfHelper with QueryErrorsBase { + + protected val resolver: Resolver + protected val expressionResolver: ExpressionResolver + + protected def scopes: NameScopeStack = resolver.getNameScopes + protected def cteRegistry: CteRegistry = resolver.getCteRegistry + protected def operatorResolutionContextStack: OperatorResolutionContextStack = + resolver.getOperatorResolutionContextStack + protected def expressionIdAssigner: ExpressionIdAssigner = + expressionResolver.getExpressionIdAssigner + + /** + * Resolves a single join child in the context of a) new [[NameScope]] b) new + * [[ExpressionIdAssigner]] mapping c) new [[CteScope]] for the multi-child operator. Returns the + * resolved child together with its [[NameScope]], which the caller uses to compute the join + * output. + */ + protected def resolveJoinChild( + unresolvedOperator: LogicalPlan, + child: LogicalPlan): (LogicalPlan, NameScope) = { + expressionIdAssigner.pushMapping() + scopes.pushScope() + cteRegistry.pushScopeForMultiChildOperator( + unresolvedOperator = unresolvedOperator, + unresolvedChild = child + ) + + try { + val resolvedChild = resolver.resolve(child) + (resolvedChild, scopes.current) + } finally { + cteRegistry.popScope() + scopes.popScope() + expressionIdAssigner.popMapping(collectChildMapping = true) + } + } + + /** + * Resolves the join condition against __all__ attributes from child scopes. We overwrite the + * current scope first to prepare for + * [[ExpressionResolver.resolveExpressionTreeInOperator]]. The join will actually produce a + * different output than the one set here, so an additional overwrite with the correct values is + * needed afterwards. Two overwrites are necessary because the condition is resolved from + * original children outputs, whereas the join output will either not contain all attributes or + * their nullabilities will be different. + * + * `collectInvalidExpressions` controls whether unsupported expressions (aggregate / window / + * generator, etc.) found in the just-resolved condition are thrown immediately as + * `UNSUPPORTED_EXPR_FOR_OPERATOR`. [[JoinResolver]] leaves this off and relies on the generic + * post-resolution check in [[Resolver.validateResolvedOperatorGenerically]], which inspects + * [[ExpressionResolver.getLastInvalidExpressionsInTheContextOfOperator]] once, after the last + * expression tree of the operator is resolved. Callers that resolve further expression trees + * after the condition must turn this on, or those later trees overwrite the "last invalid + * expressions" snapshot before the generic check runs. + */ + protected def resolveJoinCondition( + unresolvedJoin: LogicalPlan, + unresolvedCondition: Option[Expression], + leftNameScope: NameScope, + rightNameScope: NameScope, + collectInvalidExpressions: Boolean = false): Option[Expression] = { + scopes.overwriteCurrent( + output = Some(leftNameScope.output ++ rightNameScope.output), + hiddenOutput = Some(leftNameScope.hiddenOutput ++ rightNameScope.hiddenOutput) + ) + + val resolvedCondition = unresolvedCondition.map { condition => + expressionResolver.resolveExpressionTreeInOperator( + condition, + unresolvedJoin + ) + } + + validateJoinConditionDataType(resolvedCondition, unresolvedJoin) + + if (collectInvalidExpressions) { + val invalidExpressions = + expressionResolver.getLastInvalidExpressionsInTheContextOfOperator + if (invalidExpressions.nonEmpty) { + withPosition(unresolvedJoin) { + resolver.throwUnsupportedExprForOperator( + operator = unresolvedJoin, + invalidExpressions = invalidExpressions + ) + } + } + } + + resolvedCondition + } + + private def validateJoinConditionDataType( + condition: Option[Expression], + unresolvedJoin: LogicalPlan): Unit = { + condition match { + case Some(condition) => + if (condition.dataType != BooleanType) { + unresolvedJoin.failAnalysis( + errorClass = "JOIN_CONDITION_IS_NOT_BOOLEAN_TYPE", + messageParameters = Map( + "joinCondition" -> toSQLExpr(condition), + "conditionType" -> toSQLType(condition.dataType) + ) + ) + } + case None => + } + } + + /** + * Computes the new hidden output for a join. The result contains attributes from `mainOutput`, + * followed by `extraHiddenOutput` (marked as qualified access only), followed by qualified access + * only attributes from `oldHiddenOutput`. All attributes must be unique: `mainOutput` takes + * precedence over hidden output, and `extraHiddenOutput` takes precedence over `oldHiddenOutput`. + * + * For regular joins, `extraHiddenOutput` is empty and the result is simply `mainOutput` plus the + * qualified access only portion of `oldHiddenOutput`. For NATURAL / USING joins, + * `extraHiddenOutput` contains additional attributes from + * [[org.apache.spark.sql.catalyst.analysis.NaturalAndUsingJoinResolution]] that must be marked as + * qualified access only and take precedence over `oldHiddenOutput` so that name resolution in + * downstream operators (e.g. `Sort`, `Filter`) disambiguates correctly. + */ + protected def computeHiddenOutputForJoin( + mainOutput: Seq[Attribute], + oldHiddenOutput: Seq[Attribute], + extraHiddenOutput: Seq[Attribute] = Seq.empty): Seq[Attribute] = { + val mainOutputLookup = new HashSet[ExprId](mainOutput.size) + mainOutput.foreach { attribute => + mainOutputLookup.add(attribute.exprId) + } + + val extraHiddenOutputLookup = new HashSet[ExprId](extraHiddenOutput.size) + extraHiddenOutput.foreach { attribute => + extraHiddenOutputLookup.add(attribute.exprId) + } + + val filteredExtraHiddenOutput = extraHiddenOutput.collect { + case attribute if !mainOutputLookup.contains(attribute.exprId) => + attribute.markAsQualifiedAccessOnly() + } + + val filteredOldHiddenOutput = oldHiddenOutput.filter { attribute => + !mainOutputLookup.contains(attribute.exprId) && + !extraHiddenOutputLookup.contains(attribute.exprId) && + (attribute.qualifiedAccessOnly || attribute.isMetadataCol) + } + + mainOutput ++ filteredExtraHiddenOutput ++ filteredOldHiddenOutput + } + + /** + * Filters metadata columns from hidden output based on join type. + * + * For [[ExistenceJoin]] and left-existence joins ([[LeftSemi]], [[LeftAnti]]), right-side + * metadata columns are removed, matching [[org.apache.spark.sql.catalyst.plans.logical.Join]]'s + * `metadataOutput`, which propagates only the left side's metadata output for these join types. + * For outer joins, metadata columns on the nullable side have their nullability set to `true`, + * mirroring the adjustment applied to the main output -- metadata columns bypass it because they + * live in hidden output instead. For all other join types, metadata columns from both sides are + * kept as-is. + */ + protected def filterHiddenOutputMetadataForJoin( + joinType: JoinType, + oldHiddenOutput: Seq[Attribute], + rightHiddenOutput: Seq[Attribute]): Seq[Attribute] = { + val rightMetadataIds = new HashSet[ExprId]() + rightHiddenOutput.foreach { attribute => + if (attribute.isMetadataCol) { + rightMetadataIds.add(attribute.exprId) + } + } + + joinType match { + case _: ExistenceJoin | LeftSemi | LeftAnti => + oldHiddenOutput.filter(attribute => + !attribute.isMetadataCol || !rightMetadataIds.contains(attribute.exprId) + ) + case LeftOuter | LeftSingle => + oldHiddenOutput.map { attribute => + if (attribute.isMetadataCol && rightMetadataIds.contains(attribute.exprId)) { + attribute.withNullability(true) + } else { + attribute + } + } + case RightOuter => + oldHiddenOutput.map { attribute => + if (attribute.isMetadataCol && !rightMetadataIds.contains(attribute.exprId)) { + attribute.withNullability(true) + } else { + attribute + } + } + case FullOuter => + oldHiddenOutput.map { attribute => + if (attribute.isMetadataCol) { + attribute.withNullability(true) + } else { + attribute + } + } + case _ => + oldHiddenOutput + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinResolver.scala index 27a55c7f6b9f8..1fbd43074154d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/JoinResolver.scala @@ -19,24 +19,22 @@ package org.apache.spark.sql.catalyst.analysis.resolver import java.util.HashSet -import org.apache.spark.sql.catalyst.analysis.{AnalysisErrorAt, NaturalAndUsingJoinResolution} +import org.apache.spark.sql.catalyst.analysis.NaturalAndUsingJoinResolution import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, ExprId, NamedExpression} import org.apache.spark.sql.catalyst.plans.{JoinType, NaturalJoin, UsingJoin} import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, LogicalPlan, Project} import org.apache.spark.sql.catalyst.util._ -import org.apache.spark.sql.types.BooleanType /** * Resolves [[Join]] operator by resolving its left and right children and its join condition. If * the unresolved join is [[NaturalJoin]] or [[UsingJoin]], the resulting operator will be * [[Project]], otherwise it will be [[Join]]. */ -class JoinResolver(resolver: Resolver, expressionResolver: ExpressionResolver) - extends TreeNodeResolver[Join, LogicalPlan] { - private val scopes = resolver.getNameScopes - private val expressionIdAssigner = expressionResolver.getExpressionIdAssigner - private val cteRegistry = resolver.getCteRegistry - private val operatorResolutionContextStack = resolver.getOperatorResolutionContextStack +class JoinResolver( + override val resolver: Resolver, + override val expressionResolver: ExpressionResolver) + extends TreeNodeResolver[Join, LogicalPlan] + with JoinLikeResolver { /** * Resolves [[Join]] operator: @@ -52,12 +50,12 @@ class JoinResolver(resolver: Resolver, expressionResolver: ExpressionResolver) */ override def resolve(unresolvedJoin: Join): LogicalPlan = { val (resolvedLeftOperator: LogicalPlan, leftNameScope: NameScope) = resolveJoinChild( - unresolvedJoin = unresolvedJoin, + unresolvedOperator = unresolvedJoin, child = unresolvedJoin.left ) val (resolvedRightOperator: LogicalPlan, rightNameScope: NameScope) = resolveJoinChild( - unresolvedJoin = unresolvedJoin, + unresolvedOperator = unresolvedJoin, child = unresolvedJoin.right ) @@ -82,26 +80,6 @@ class JoinResolver(resolver: Resolver, expressionResolver: ExpressionResolver) ) } - private def resolveJoinChild( - unresolvedJoin: Join, - child: LogicalPlan): (LogicalPlan, NameScope) = { - expressionIdAssigner.pushMapping() - scopes.pushScope() - cteRegistry.pushScopeForMultiChildOperator( - unresolvedOperator = unresolvedJoin, - unresolvedChild = child - ) - - try { - val resolvedChild = resolver.resolve(child) - (resolvedChild, scopes.current) - } finally { - cteRegistry.popScope() - scopes.popScope() - expressionIdAssigner.popMapping(collectChildMapping = true) - } - } - /** * If the type of join is [[NaturalJoin]] or [[UsingJoin]], perform additional transformations in * [[commonNaturalJoinProcessing]]. Otherwise, overwrite current name scope output with the @@ -375,52 +353,4 @@ class JoinResolver(resolver: Resolver, expressionResolver: ExpressionResolver) rightNameScope.output.map(_.name) ) } - - /** - * Resolves join condition by __all__ attributes from child scopes. We need to overwrite current - * scope first to prepare for [[resolveExpressionTreeInOperator]]. [[Join]] will actually produce - * different output than the one we are setting here, so additional overwrite with correct values - * will be needed. Two overwrites are necessary because condition is resolved from original - * children outputs, whereas output of [[Join]] will either not contain all attributes or their - * nullabilities will be different. - */ - private def resolveJoinCondition( - unresolvedJoin: Join, - unresolvedCondition: Option[Expression], - leftNameScope: NameScope, - rightNameScope: NameScope) = { - scopes.overwriteCurrent( - output = Some(leftNameScope.output ++ rightNameScope.output), - hiddenOutput = Some(leftNameScope.hiddenOutput ++ rightNameScope.hiddenOutput) - ) - - val resolvedCondition = unresolvedCondition.map { condition => - expressionResolver.resolveExpressionTreeInOperator( - condition, - unresolvedJoin - ) - } - - validateJoinConditionDataType(resolvedCondition, unresolvedJoin) - - resolvedCondition - } - - private def validateJoinConditionDataType( - condition: Option[Expression], - unresolvedJoin: Join): Unit = { - condition match { - case Some(condition) => - if (condition.dataType != BooleanType) { - unresolvedJoin.failAnalysis( - errorClass = "JOIN_CONDITION_IS_NOT_BOOLEAN_TYPE", - messageParameters = Map( - "joinCondition" -> toSQLExpr(condition), - "conditionType" -> toSQLType(condition.dataType) - ) - ) - } - case None => - } - } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/LateralColumnAliasResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/LateralColumnAliasResolver.scala index 3d19c4cbde8c7..6859c98f1f999 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/LateralColumnAliasResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/LateralColumnAliasResolver.scala @@ -70,6 +70,9 @@ class LateralColumnAliasResolver(expressionResolver: ExpressionResolver, operato case _ @Project(projectList: Seq[_], aggregate: Aggregate) => operatorResolutionContextStack.current.baseOperator = Some(aggregate) + // Note: LCA + grouping analytics (CUBE/ROLLUP/GROUPING SETS) is intercepted in + // AggregateResolver which throws ExplicitlyUnsupportedResolverFeature before reaching + // here. This validation is retained as defense-in-depth in case the call path changes. AggregationValidator(aggregate) val remappedAliases = new HashMap[ExprId, Alias](projectList.size) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/OperatorResolutionContext.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/OperatorResolutionContext.scala index 5df31a9c09fdf..880fce67ba677 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/OperatorResolutionContext.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/OperatorResolutionContext.scala @@ -51,7 +51,8 @@ class OperatorResolutionContext( var baseOperator: Option[LogicalPlan] = None, val isResolvingTreeUnderHaving: Boolean = false, var parameterNamesToValues: Option[LinkedHashMap[String, Expression]] = None, - var hasGroupingAnalytics: Boolean = false) { + var hasGroupingAnalytics: Boolean = false, + val isSubqueryRoot: Boolean = false) { /** * @param subqueryExpressionsToValidate List of [[SubqueryExpression]]s in the context of current @@ -130,7 +131,8 @@ class OperatorResolutionContextStack { new OperatorResolutionContext( unresolvedPlan = Some(unresolvedPlan), isResolvingTreeUnderHaving = isResolvingTreeUnderHaving, - parameterNamesToValues = current.parameterNamesToValues + parameterNamesToValues = current.parameterNamesToValues, + isSubqueryRoot = isSubqueryRoot ) ) } @@ -138,11 +140,20 @@ class OperatorResolutionContextStack { /** * Pops the top resolution context from the stack. Before popping, propagates the * `hasGroupingAnalytics` from the child context to the parent context if it was set. + * The flag is NOT propagated across subquery boundaries (isSubqueryRoot) or derived-table + * boundaries (SubqueryAlias), because grouping analytics inside a subquery or derived table + * are unrelated to operators in the outer query. */ def pop(): Unit = { val childContext = current stack.pop() - if (childContext.hasGroupingAnalytics) { + val isDerivedTableBoundary = childContext.unresolvedPlan match { + case Some(_: SubqueryAlias) => true + case _ => false + } + if (childContext.hasGroupingAnalytics && + !childContext.isSubqueryRoot && + !isDerivedTableBoundary) { current.hasGroupingAnalytics = childContext.hasGroupingAnalytics } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolutionValidator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolutionValidator.scala index 6bdf2d4b0615b..452f6d9555aeb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolutionValidator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolutionValidator.scala @@ -101,6 +101,8 @@ class ResolutionValidator { validateSort(sort) case join: Join => validateJoin(join) + case asOfJoin: AsOfJoin => + validateAsOfJoin(asOfJoin) case repartition: Repartition => validateRepartition(repartition) case repartitionByExpression: RepartitionByExpression => @@ -349,6 +351,36 @@ class ResolutionValidator { handleOperatorOutput(join) } + private def validateAsOfJoin(asOfJoin: AsOfJoin): Unit = { + // The inner scope keeps the per-child output overwrites done by `handleOperatorOutput` out + // of the outer scope, which holds the combined join output the join expressions resolve + // against. + attributeScopeStack.pushScope() + try { + attributeScopeStack.pushScope() + try { + validate(asOfJoin.left) + validate(asOfJoin.right) + assert(asOfJoin.left.outputSet.intersect(asOfJoin.right.outputSet).isEmpty) + } finally { + attributeScopeStack.popScope() + } + + attributeScopeStack.overwriteCurrent(asOfJoin.left.output ++ asOfJoin.right.output) + + expressionResolutionValidator.validate(asOfJoin.asOfCondition) + expressionResolutionValidator.validate(asOfJoin.orderExpression) + asOfJoin.condition.foreach(expressionResolutionValidator.validate) + asOfJoin.toleranceAssertion.foreach(expressionResolutionValidator.validate) + asOfJoin.leftSortExprs.foreach(expressionResolutionValidator.validate) + asOfJoin.rightSortExprs.foreach(expressionResolutionValidator.validate) + } finally { + attributeScopeStack.popScope() + } + + handleOperatorOutput(asOfJoin) + } + private def validateSupervisingCommand(supervisingCommand: SupervisingCommand): Unit = {} private def handleOperatorOutput(operator: LogicalPlan): Unit = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/Resolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/Resolver.scala index 9aac3cb726ec0..d46227719f85b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/Resolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/Resolver.scala @@ -55,7 +55,6 @@ import org.apache.spark.sql.catalyst.trees.CurrentOrigin import org.apache.spark.sql.catalyst.util.EvaluateUnresolvedInlineTable import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} -import org.apache.spark.sql.internal.SQLConf /** * The Resolver implements a single-pass bottom-up analysis algorithm in the Catalyst. @@ -118,6 +117,7 @@ class Resolver( private val filterResolver = new FilterResolver(this, expressionResolver) private val sortResolver = new SortResolver(this, expressionResolver) private val joinResolver = new JoinResolver(this, expressionResolver) + private val asOfJoinResolver = new AsOfJoinResolver(this, expressionResolver) private val havingResolver = new HavingResolver(this, expressionResolver) /** @@ -265,6 +265,8 @@ class Resolver( unresolvedPlan match { case unresolvedJoin: Join => joinResolver.resolve(unresolvedJoin) + case unresolvedAsOfJoin: AsOfJoin => + asOfJoinResolver.resolve(unresolvedAsOfJoin) case unresolvedWith: UnresolvedWith => resolveWith(unresolvedWith) case withCte: WithCTE => @@ -807,7 +809,7 @@ class Resolver( if (operator.children.nonEmpty) { val missingInput = operator.missingInput if (missingInput.nonEmpty) { - throwMissingAttributesError(operator, missingInput) + Resolver.throwMissingAttributesError(operator, missingInput) } } } @@ -828,42 +830,6 @@ class Resolver( } } - private def throwMissingAttributesError( - operator: LogicalPlan, - missingInput: AttributeSet): Nothing = { - val inputSet = operator.inputSet - - val inputAttributesByName = new IdentifierMap[Attribute] - for (attribute <- inputSet) { - inputAttributesByName.put(attribute.name, attribute) - } - - val attributesWithSameName = missingInput.filter { missingAttribute => - inputAttributesByName.contains(missingAttribute.name) - } - - if (attributesWithSameName.nonEmpty) { - operator.failAnalysis( - errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION", - messageParameters = Map( - "missingAttributes" -> makeCommaSeparatedExpressionString(missingInput.toSeq), - "input" -> makeCommaSeparatedExpressionString(inputSet.toSeq), - "operator" -> operator.simpleString(SQLConf.get.maxToStringFields), - "operation" -> makeCommaSeparatedExpressionString(attributesWithSameName.toSeq) - ) - ) - } else { - operator.failAnalysis( - errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", - messageParameters = Map( - "missingAttributes" -> makeCommaSeparatedExpressionString(missingInput.toSeq), - "input" -> makeCommaSeparatedExpressionString(inputSet.toSeq), - "operator" -> operator.simpleString(SQLConf.get.maxToStringFields) - ) - ) - } - } - private def throwSinglePassFailedToResolveOperator(operator: LogicalPlan): Nothing = throw SparkException.internalError( msg = s"Failed to resolve operator in single-pass: $operator", @@ -871,7 +837,7 @@ class Resolver( summary = operator.origin.context.summary() ) - private def throwUnsupportedExprForOperator( + private[resolver] def throwUnsupportedExprForOperator( operator: LogicalPlan, invalidExpressions: Seq[Expression]): Nothing = { throw new AnalysisException( @@ -890,6 +856,30 @@ class Resolver( object Resolver { + /** + * Fails with `MISSING_ATTRIBUTES` because `operator` references attributes its child does not + * produce. + */ + def throwMissingAttributesError( + operator: LogicalPlan, + missingInput: AttributeSet): Nothing = { + val inputAttributesByName = new IdentifierMap[Attribute] + for (attribute <- operator.inputSet) { + inputAttributesByName.put(attribute.name, attribute) + } + + val attributesWithSameName = missingInput.filter { missingAttribute => + inputAttributesByName.contains(missingAttribute.name) + }.toSeq + + throw QueryCompilationErrors.missingAttributesError( + operator = operator, + missingInput = missingInput.toSeq, + input = operator.inputSet.toSeq, + attributesWithSameName = attributesWithSameName + ) + } + /** * Create a new instance of the [[RelationResolution]]. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolverGuard.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolverGuard.scala index 99c79aaee3f67..80270cd4653d6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolverGuard.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolverGuard.scala @@ -111,6 +111,11 @@ class ResolverGuard( checkFilter(filter) case join: Join => checkJoin(join) + case asOfJoin: AsOfJoin + if conf.getConf( + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLE_ASOF_JOIN_RESOLUTION + ) => + checkAsOfJoin(asOfJoin) case unresolvedSubqueryColumnAliases: UnresolvedSubqueryColumnAliases => checkUnresolvedSubqueryColumnAliases(unresolvedSubqueryColumnAliases) case subqueryAlias: SubqueryAlias => @@ -284,6 +289,52 @@ class ResolverGuard( } } + private def checkAsOfJoin(asOfJoin: AsOfJoin) = { + checkOperator(asOfJoin.left) + .orElse { + checkOperator(asOfJoin.right) + } + .orElse { + asOfJoin.condition match { + case Some(condition) => checkExpression(condition) + case None => None + } + } + .orElse { + asOfJoin.matchLeftOperand match { + case Some(expression) => checkExpression(expression) + case None => None + } + } + .orElse { + asOfJoin.matchRightOperand match { + case Some(expression) => checkExpression(expression) + case None => None + } + } + .orElse { + checkExpression(asOfJoin.asOfCondition) + } + .orElse { + checkExpression(asOfJoin.orderExpression) + } + .orElse { + asOfJoin.toleranceAssertion match { + case Some(expression) => checkExpression(expression) + case None => None + } + } + .orElse { + checkExpressions(asOfJoin.leftSortExprs) + } + .orElse { + checkExpressions(asOfJoin.rightSortExprs) + } + } + + private def checkExpressions(expressions: Seq[Expression]): Option[String] = + expressions.iterator.map(checkExpression).collectFirst { case Some(reason) => reason } + private def checkFilter(unresolvedFilter: Filter) = checkOperator(unresolvedFilter.child).orElse(checkExpression(unresolvedFilter.condition)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolvesNameByHiddenOutput.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolvesNameByHiddenOutput.scala index 6faecbe832813..0b06eaa9be458 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolvesNameByHiddenOutput.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/ResolvesNameByHiddenOutput.scala @@ -284,7 +284,7 @@ trait ResolvesNameByHiddenOutput extends SQLConfHelper { val (metadataCols, nonMetadataCols) = operatorOutput.partition(_.toAttribute.qualifiedAccessOnly) - operator match { + val expandedOperator = operator match { case aggregate: Aggregate => val newAggregateList = nonMetadataCols ++ filteredMissingExpressions ++ metadataCols aggregate.copy(aggregateExpressions = newAggregateList) @@ -306,11 +306,26 @@ trait ResolvesNameByHiddenOutput extends SQLConfHelper { project.copy(projectList = newProjectList, child = expandedChild) } + + checkMissingInput(expandedOperator) + + expandedOperator } else { operator } } + /** + * Rejects an expanded `operator` whose child does not produce an appended hidden-output + * expression, e.g. an ORDER BY on a column an operator re-outputs under a fresh id. + */ + private def checkMissingInput(operator: LogicalPlan): Unit = { + val missingInput = operator.missingInput + if (missingInput.nonEmpty) { + Resolver.throwMissingAttributesError(operator, missingInput) + } + } + private def filterMissingExpressions( operatorOutput: Seq[NamedExpression], missingExpressions: Seq[NamedExpression]): Seq[NamedExpression] = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/SortResolver.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/SortResolver.scala index ed0edf90969fe..1219d16e4c831 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/SortResolver.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/resolver/SortResolver.scala @@ -135,6 +135,14 @@ class SortResolver(operatorResolver: Resolver, expressionResolver: ExpressionRes val resolvedChild = operatorResolver.resolve(unresolvedSort.child) + // ORDER BY over grouping analytics (CUBE/ROLLUP/GROUPING SETS) is not yet supported in + // single-pass because the expanded Aggregate's ExprIds get out of sync when ORDER BY + // inserts missing expressions. Fall back to legacy for correct results (SPARK-57346). + if (operatorResolutionContextStack.current.hasGroupingAnalytics) { + throw new ExplicitlyUnsupportedResolverFeature( + "ORDER BY with grouping analytics (SPARK-57346)") + } + operatorResolutionContextStack.current.ordinalReplacementExpressions = Some( OrdinalReplacementSortOrderExpressions( expressions = scopes.current.output.toIndexedSeq, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala index fc51dd8c72d45..50b2970544dcf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala @@ -144,7 +144,9 @@ case class UnresolvedRelation( def requireWritePrivileges(privileges: Set[TableWritePrivilege]): UnresolvedRelation = { if (privileges.nonEmpty) { val newOptions = new java.util.HashMap[String, String] - newOptions.putAll(options) + // CaseInsensitiveStringMap's Map view exposes lowercase keys. Copy the original map to + // preserve user-provided key casing when adding the internal marker. + newOptions.putAll(options.asCaseSensitiveMap()) newOptions.put(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES, privileges.mkString(",")) copy(options = new CaseInsensitiveStringMap(newOptions)) } else { @@ -155,7 +157,8 @@ case class UnresolvedRelation( def clearWritePrivileges: UnresolvedRelation = { if (options.containsKey(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES)) { val newOptions = new java.util.HashMap[String, String] - newOptions.putAll(options) + // Preserve user-provided key casing while removing the internal marker as well. + newOptions.putAll(options.asCaseSensitiveMap()) newOptions.remove(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) copy(options = new CaseInsensitiveStringMap(newOptions)) } else { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/interface.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/interface.scala index 247d4124dae4a..377b649aaca04 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/interface.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/catalog/interface.scala @@ -972,17 +972,17 @@ case class CatalogColumnStat( */ def toMap(colName: String): Map[String, String] = { val map = new scala.collection.mutable.HashMap[String, String] - map.put(s"${colName}.${CatalogColumnStat.KEY_VERSION}", CatalogColumnStat.VERSION.toString) + map.put(s"$colName.${CatalogColumnStat.KEY_VERSION}", CatalogColumnStat.VERSION.toString) distinctCount.foreach { v => - map.put(s"${colName}.${CatalogColumnStat.KEY_DISTINCT_COUNT}", v.toString) + map.put(s"$colName.${CatalogColumnStat.KEY_DISTINCT_COUNT}", v.toString) } nullCount.foreach { v => - map.put(s"${colName}.${CatalogColumnStat.KEY_NULL_COUNT}", v.toString) + map.put(s"$colName.${CatalogColumnStat.KEY_NULL_COUNT}", v.toString) } - avgLen.foreach { v => map.put(s"${colName}.${CatalogColumnStat.KEY_AVG_LEN}", v.toString) } - maxLen.foreach { v => map.put(s"${colName}.${CatalogColumnStat.KEY_MAX_LEN}", v.toString) } - min.foreach { v => map.put(s"${colName}.${CatalogColumnStat.KEY_MIN_VALUE}", v) } - max.foreach { v => map.put(s"${colName}.${CatalogColumnStat.KEY_MAX_VALUE}", v) } + avgLen.foreach { v => map.put(s"$colName.${CatalogColumnStat.KEY_AVG_LEN}", v.toString) } + maxLen.foreach { v => map.put(s"$colName.${CatalogColumnStat.KEY_MAX_LEN}", v.toString) } + min.foreach { v => map.put(s"$colName.${CatalogColumnStat.KEY_MIN_VALUE}", v) } + max.foreach { v => map.put(s"$colName.${CatalogColumnStat.KEY_MAX_VALUE}", v) } histogram.foreach { h => CatalogTable.splitLargeTableProp( s"$colName.${CatalogColumnStat.KEY_HISTOGRAM}", @@ -1096,15 +1096,15 @@ object CatalogColumnStat extends Logging { try { Some(CatalogColumnStat( - distinctCount = map.get(s"${colName}.${KEY_DISTINCT_COUNT}").map(v => BigInt(v.toLong)), - min = map.get(s"${colName}.${KEY_MIN_VALUE}"), - max = map.get(s"${colName}.${KEY_MAX_VALUE}"), - nullCount = map.get(s"${colName}.${KEY_NULL_COUNT}").map(v => BigInt(v.toLong)), - avgLen = map.get(s"${colName}.${KEY_AVG_LEN}").map(_.toLong), - maxLen = map.get(s"${colName}.${KEY_MAX_LEN}").map(_.toLong), + distinctCount = map.get(s"$colName.${KEY_DISTINCT_COUNT}").map(v => BigInt(v.toLong)), + min = map.get(s"$colName.${KEY_MIN_VALUE}"), + max = map.get(s"$colName.${KEY_MAX_VALUE}"), + nullCount = map.get(s"$colName.${KEY_NULL_COUNT}").map(v => BigInt(v.toLong)), + avgLen = map.get(s"$colName.${KEY_AVG_LEN}").map(_.toLong), + maxLen = map.get(s"$colName.${KEY_MAX_LEN}").map(_.toLong), histogram = CatalogTable.readLargeTableProp(map, s"$colName.$KEY_HISTOGRAM") .map(HistogramSerializer.deserialize), - version = map(s"${colName}.${KEY_VERSION}").toInt + version = map(s"$colName.${KEY_VERSION}").toInt )) } catch { case NonFatal(e) => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVOptions.scala index 30b6877d3a82c..7db03a8a23231 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVOptions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/csv/CSVOptions.scala @@ -107,15 +107,15 @@ class CSVOptions( } private def getBool(paramName: String, default: Boolean = false): Boolean = { - val param = parameters.getOrElse(paramName, default.toString) - if (param == null) { - default - } else if (param.toLowerCase(Locale.ROOT) == "true") { - true - } else if (param.toLowerCase(Locale.ROOT) == "false") { - false - } else { - throw QueryExecutionErrors.paramIsNotBooleanValueError(paramName) + val paramValue = parameters.get(paramName) + paramValue match { + case None => default + case Some(null) => default + case Some(value) => value.toLowerCase(Locale.ROOT) match { + case "true" => true + case "false" => false + case _ => throw QueryExecutionErrors.paramIsNotBooleanValueError(paramName) + } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpression.scala index 2cee8303dc57b..78f0f8a7af613 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpression.scala @@ -35,6 +35,7 @@ case class ApplyFunctionExpression( override lazy val deterministic: Boolean = function.isDeterministic && children.forall(_.deterministic) override def foldable: Boolean = deterministic && children.forall(_.foldable) + override def stateful: Boolean = true private lazy val reusedRow = new SpecificInternalRow(function.inputTypes().toImmutableArraySeq) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApproxTopKExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApproxTopKExpressions.scala index 8d4f0a18d5708..fc647c8602086 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApproxTopKExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ApproxTopKExpressions.scala @@ -42,6 +42,13 @@ import org.apache.spark.sql.types._ _FUNC_(state, k) - Returns top k items with their frequency. `k` An optional INTEGER literal greater than 0. If k is not specified, it defaults to 5. """, + arguments = """ + Arguments: + * state - The sketch state produced by `approx_top_k_accumulate` or + `approx_top_k_combine`. + * k - Optional. A constant INTEGER literal greater than 0 giving the number + of top items to return. If omitted, it defaults to 5. + """, examples = """ Examples: > SELECT _FUNC_(approx_top_k_accumulate(expr)) FROM VALUES (0), (0), (1), (1), (2), (3), (4), (4) AS tab(expr); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/BroadcastValueProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/BroadcastValueProjection.scala new file mode 100644 index 0000000000000..0058414987a9f --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/BroadcastValueProjection.scala @@ -0,0 +1,31 @@ +/* + * 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.catalyst.expressions + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan + +/** + * Describes a value projection from the rows of an existing standard hash broadcast. + * + * The source hash keys are the complete, ordered keys of the broadcast join. The value expression + * is evaluated against the rows stored in the resulting hashed relation. + */ +case class BroadcastValueProjection( + sourcePlan: LogicalPlan, + sourceHashKeys: Seq[Expression], + valueExpression: Expression) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala index 90b9f19cf8d04..841648912df8b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/CallMethodViaReflection.scala @@ -51,6 +51,13 @@ import org.apache.spark.util.Utils */ @ExpressionDescription( usage = "_FUNC_(class, method[, arg1[, arg2 ..]]) - Calls a method with reflection.", + arguments = """ + Arguments: + * class - A literal string with the fully qualified name of the class. + * method - A literal string with the name of the static method to call. + * argN - Optional arguments passed to the method. Only primitive and string + types are supported, and each argument is matched to the method signature. + """, examples = """ Examples: > SELECT _FUNC_('java.util.UUID', 'randomUUID'); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala index eb7b1d269cdbe..5e5bb00df6c26 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Cast.scala @@ -660,6 +660,11 @@ object Cast extends QueryErrorsBase { @ExpressionDescription( usage = "_FUNC_(expr AS type) - Casts the value `expr` to the target data type `type`." + " `expr` :: `type` alternative casting syntax is also supported.", + arguments = """ + Arguments: + * expr - An expression whose value is converted to the target data type. + * type - The target data type to cast the value to. + """, examples = """ Examples: > SELECT _FUNC_('10' as int); @@ -687,6 +692,10 @@ case class Cast( override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = copy(timeZoneId = Option(timeZoneId)) + // Parser, Column.cast, and Connect set USER_SPECIFIED_CAST. Analyzer-inserted Casts do not. + override protected def truncateCharVarcharOnCast: Boolean = + containsTag(Cast.USER_SPECIFIED_CAST) + override protected def withNewChildInternal(newChild: Expression): Cast = copy(child = newChild) // CAST_TO_TIMESTAMP must be set on a superset of the targets accepted by @@ -1741,7 +1750,7 @@ case class Cast( """ } else { code""" - scala.Option<Integer> $intOpt = $dateTimeUtilsCls.stringToDate($c); + scala.Option $intOpt = $dateTimeUtilsCls.stringToDate($c); if ($intOpt.isDefined()) { $evPrim = ((Integer) $intOpt.get()).intValue(); } else { @@ -1786,7 +1795,7 @@ case class Cast( """ } else { code""" - scala.Option<Long> $longOpt = $dateTimeUtilsCls.stringToTime($c); + scala.Option $longOpt = $dateTimeUtilsCls.stringToTime($c); if ($longOpt.isDefined()) { $evPrim = $dateTimeUtilsCls.truncateTimeToPrecision( ((Long) $longOpt.get()).longValue(), ${to.precision}); @@ -1976,7 +1985,7 @@ case class Cast( """ } else { code""" - scala.Option<Long> $longOpt = $dateTimeUtilsCls.stringToTimestamp($c, $zid); + scala.Option $longOpt = $dateTimeUtilsCls.stringToTimestamp($c, $zid); if ($longOpt.isDefined()) { $evPrim = ((Long) $longOpt.get()).longValue(); } else { @@ -2052,7 +2061,7 @@ case class Cast( """ } else { code""" - scala.Option<Long> $longOpt = $dateTimeUtilsCls.stringToTimestampWithoutTimeZone($c); + scala.Option $longOpt = $dateTimeUtilsCls.stringToTimestampWithoutTimeZone($c); if ($longOpt.isDefined()) { $evPrim = ((Long) $longOpt.get()).longValue(); } else { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruning.scala index 1c33ed65c1df1..959acbc762b4e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruning.scala @@ -39,6 +39,7 @@ trait DynamicPruning extends Predicate * broadcast through ReuseExchange; otherwise, it will use the filter only if it * can reuse the results of the broadcast through ReuseExchange * @param broadcastKeyIndices the indices of the filtering keys collected from the broadcast + * @param broadcastValueProjection an optional value projection from an existing broadcast */ case class DynamicPruningSubquery( pruningKey: Expression, @@ -47,28 +48,49 @@ case class DynamicPruningSubquery( broadcastKeyIndices: Seq[Int], onlyInBroadcast: Boolean, exprId: ExprId = NamedExpression.newExprId, - hint: Option[HintInfo] = None) + hint: Option[HintInfo] = None)( + @transient private[sql] val broadcastValueProjection: Option[BroadcastValueProjection] = None) extends SubqueryExpression(buildQuery, Seq(pruningKey), exprId, Seq.empty, hint) with DynamicPruning with Unevaluable with UnaryLike[Expression] { + override protected def otherCopyArgs: Seq[AnyRef] = Seq(broadcastValueProjection) + override def child: Expression = pruningKey override def plan: LogicalPlan = buildQuery override def nullable: Boolean = false - override def withNewPlan(plan: LogicalPlan): DynamicPruningSubquery = copy(buildQuery = plan) + override def withNewPlan(plan: LogicalPlan): DynamicPruningSubquery = + copy(buildQuery = plan)(broadcastValueProjection) override def withNewOuterAttrs(outerAttrs: Seq[Expression]): DynamicPruningSubquery = { // Updating outer attrs of DynamicPruningSubquery is unsupported; assert that they match // pruningKey and return a copy without any changes. assert(outerAttrs.size == 1 && outerAttrs.head.semanticEquals(pruningKey)) - copy() + copy()(broadcastValueProjection) } - override def withNewHint(hint: Option[HintInfo]): SubqueryExpression = copy(hint = hint) + override def withNewHint(hint: Option[HintInfo]): SubqueryExpression = + copy(hint = hint)(broadcastValueProjection) + + private[sql] def usableBroadcastValueProjection: Option[BroadcastValueProjection] = { + broadcastValueProjection.filter { projection => + projection.sourcePlan.resolved && + projection.sourcePlan.deterministic && + projection.sourceHashKeys.nonEmpty && + projection.sourceHashKeys.forall(_.resolved) && + projection.sourceHashKeys.forall(_.deterministic) && + projection.sourceHashKeys.forall( + _.references.subsetOf(projection.sourcePlan.outputSet)) && + projection.valueExpression.resolved && + projection.valueExpression.deterministic && + projection.valueExpression.references.subsetOf(projection.sourcePlan.outputSet) && + projection.valueExpression.dataType == pruningKey.dataType + } + } override lazy val resolved: Boolean = { pruningKey.resolved && @@ -92,11 +114,11 @@ case class DynamicPruningSubquery( pruningKey = pruningKey.canonicalized, buildQuery = buildQuery.canonicalized, buildKeys = buildKeys.map(QueryPlan.normalizeExpressions(_, buildQuery.output)), - exprId = ExprId(0)) + exprId = ExprId(0))(None) } override protected def withNewChildInternal(newChild: Expression): DynamicPruningSubquery = - copy(pruningKey = newChild) + copy(pruningKey = newChild)(broadcastValueProjection) } /** diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala index 6c9ea6d656b2d..f3f2988abddcf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ExprUtils.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.catalyst.plans.logical.Aggregate +import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, CharVarcharUtils} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors} import org.apache.spark.sql.internal.types.{AbstractMapType, StringTypeWithCollation} @@ -220,4 +221,41 @@ object ExprUtils extends EvalHelper with QueryErrorsBase { a.groupingExpressions.foreach(checkValidGroupingExprs) a.aggregateExpressions.foreach(checkValidAggregateExpression) } + + /** + * Returns true if `e` is safe to evaluate unconditionally, i.e. on rows where the + * original plan would not have evaluated it: evaluating it must not raise an error and + * must not change the query result. This is the check to use when relocating an + * expression out of a short-circuited position, e.g. moving the base out of the taken + * branch of an If/CaseWhen, or hoisting a streamed-side join conjunct above the probe + * so it also runs for streamed rows that have no match. + * + * Only a whitelist of total, deterministic expressions qualifies: + * - leaves: attribute references and literals; + * - total accessors: GetStructField, GetArrayStructFields and GetMapValue never throw. + * Note that GetArrayItem/ElementAt are NOT included: they throw on invalid ordinals + * when ANSI mode is on; + * - logic/predicates: And, Or, Not, comparisons, IsNull, IsNotNull, IsNaN, NullIf, + * Coalesce, In and InSet are total boolean functions. + * Anything else (arithmetic, casts, string functions, UDFs, nested IF/CASE WHEN, ...) + * conservatively returns false. On top of the whitelist, the expression must be + * deterministic and must not contain subqueries. + * + * Note: this deliberately does not rely on [[Expression.throwable]], which is opt-in + * metadata that most expressions do not override. A throwing ScalaUDF with + * non-throwing children, for example, reports non-throwable. + */ + def canEvaluateUnconditionally(e: Expression): Boolean = + e.deterministic && !e.containsPattern(PLAN_EXPRESSION) && + canEvaluateUnconditionallyInternal(e) + + private def canEvaluateUnconditionallyInternal(e: Expression): Boolean = e match { + case _: AttributeReference | _: Literal => true + case _: GetStructField | _: GetArrayStructFields | _: GetMapValue => + e.children.forall(canEvaluateUnconditionallyInternal) + case _: And | _: Or | _: Not | _: BinaryComparison | _: IsNull | _: IsNotNull | + _: IsNaN | _: NullIf | _: Coalesce | _: In | _: InSet => + e.children.forall(canEvaluateUnconditionallyInternal) + case _ => false + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala index 834f3b0debd08..4959df90a53ea 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala @@ -519,6 +519,9 @@ trait NonSQLExpression extends Expression { case a: Attribute => new PrettyAttribute(a) case a: Alias => PrettyAttribute(a.sql, a.dataType) case p: PythonFuncExpression => PrettyPythonUDF(p.name, p.dataType, p.children) + // Render a transpiled UDF like the UDF it wraps (options must not leak + // into user-visible strings). + case t: TranspiledPythonUDF => PrettyPythonUDF(t.name, t.dataType, t.pythonUDFExpr.children) }.toString } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala index 99619d9de96cc..b820172bc2b53 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala @@ -23,7 +23,8 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.UnresolvedException import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} -import org.apache.spark.sql.catalyst.trees.TreePattern.{PYTHON_UDF, TreePattern} +import org.apache.spark.sql.catalyst.trees.TreePattern.{PYTHON_UDF, TRANSPILED_PYTHON_UDF, + TreePattern} import org.apache.spark.sql.catalyst.util.toPrettySQL import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.types._ @@ -35,6 +36,15 @@ object PythonUDF { private[this] val SCALAR_TYPES = Set( PythonEvalType.SQL_BATCHED_UDF, PythonEvalType.SQL_ARROW_BATCHED_UDF, + // Element-wise UDFs are row-shaped from the plan's point of view: one array column in, one + // array column out per row. They are extracted by `ExtractPythonUDFs` like any other scalar + // UDF; only the Python worker treats them element-wise. One eval type per lifted flavor keeps + // the worker's pandas- vs. Arrow-shaped batching and the iterator contract distinct. + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_SCALAR_PANDAS_UDF, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, PythonEvalType.SQL_SCALAR_ARROW_UDF, @@ -45,11 +55,196 @@ object PythonUDF { e.isInstanceOf[PythonUDF] && SCALAR_TYPES.contains(e.asInstanceOf[PythonUDF].evalType) } + /** + * Whether `e` is a Python UDF that can be lifted out of a higher-order function's lambda by + * `ExtractPythonUDFFromLambda`, which applies it to the whole array outside the lambda. + * + * Both the row-at-a-time eval types (plain and Arrow batched) and the vectorized scalar eval + * types (scalar pandas / Arrow and their iterator variants) qualify: the rule lifts the UDF + * structurally over `array<T>` arguments, and the Python worker flattens each array, invokes the + * function on the flat element column with its own batching contract, and re-nests. See + * [[liftedElementwiseEvalType]] for the mapping to the eval type the lifted UDF runs under. + * + * Otherwise-eligible shapes are excluded because the rewrite cannot preserve them: + * - a zero-argument call, `f()`: the lift turns each argument into an aligned array, so with no + * argument there is no array to carry the iterated shape, and the element-wise UDF would + * reach the worker with no input column and crash there instead of failing analysis; + * - a call with named arguments on an *iterator* UDF (scalar pandas / Arrow iterator): iterator + * UDFs do not take keyword arguments (the worker ignores their kwargs offsets), so such a + * call is invalid regardless of the lift. Named arguments on the non-iterator flavors are + * supported: the lift keeps each `NamedArgumentExpression` as a direct child of the lifted + * UDF (only its value becomes an aligned array), so the runner still derives the kwargs map; + * - a UDF whose argument or return type involves a UDT: the lift forces an Arrow element-wise + * eval type, which has no UDT fallback (unlike `correctEvalType`'s Arrow -> pickle path), so + * it would fail at runtime instead of at analysis. + * All keep the previous behavior (an analysis error) rather than being rewritten. + * + * This is shared with `CheckAnalysis` so that the shapes analysis accepts are exactly those the + * optimizer rule can rewrite. + */ + def isElementwiseRewritableUDF(e: Expression): Boolean = e match { + case udf: PythonUDF => + isElementwiseRewritableEvalType(udf.evalType) && + udf.children.nonEmpty && + (supportsNamedArgumentsWhenLifted(udf.evalType) || + !udf.children.exists(_.isInstanceOf[NamedArgumentExpression])) && + !containsUDT(udf.dataType) && + !udf.children.exists(c => containsUDT(c.dataType)) + case _ => false + } + + /** + * Whether a lifted UDF of this eval type can carry keyword arguments. Iterator UDFs (scalar + * pandas / Arrow iterator, and their lifted element-wise forms) cannot - the worker binds only + * positional arguments for them - so a named-argument call on an iterator UDF is not rewritable. + */ + private def supportsNamedArgumentsWhenLifted(evalType: Int): Boolean = evalType match { + case PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF | + PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF | + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF | + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF => false + case _ => true + } + + private def isElementwiseRewritableEvalType(evalType: Int): Boolean = evalType match { + case PythonEvalType.SQL_BATCHED_UDF | + PythonEvalType.SQL_ARROW_BATCHED_UDF | + PythonEvalType.SQL_SCALAR_PANDAS_UDF | + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF | + PythonEvalType.SQL_SCALAR_ARROW_UDF | + PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF => true + // The already-lifted element-wise types are rewritable again: a UDF inside a nested lambda is + // lifted once onto the inner lambda's variable (producing an element-wise UDF over that + // variable) and then re-lifted onto the enclosing array, incrementing its nesting depth. Users + // cannot create these eval types directly, so they only appear mid-rewrite - `CheckAnalysis`, + // which runs before this rule, never sees them. + case _ => PythonEvalType.isElementwiseUDF(evalType) + } + + /** + * The eval type a rewritable UDF runs under once lifted out of the lambda. Each maps to the + * element-wise flavor that preserves its worker contract: the row-at-a-time types share the one + * pickle-based element-wise path, while each vectorized scalar type keeps its own pandas- vs. + * Arrow-shaped batching and iterator behavior. An already-lifted element-wise type maps to itself + * (re-lifting for a nested lambda keeps the flavor and only bumps the nesting depth). `evalType` + * must satisfy [[isElementwiseRewritableEvalType]]. + */ + def liftedElementwiseEvalType(evalType: Int): Int = evalType match { + case PythonEvalType.SQL_BATCHED_UDF | PythonEvalType.SQL_ARROW_BATCHED_UDF => + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF + case PythonEvalType.SQL_SCALAR_PANDAS_UDF => + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF + case PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF => + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF + case PythonEvalType.SQL_SCALAR_ARROW_UDF => + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF + case PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF => + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF + case elementwise if PythonEvalType.isElementwiseUDF(elementwise) => elementwise + case other => + throw internalError(s"Not a rewritable elementwise UDF eval type: $other") + } + + /** + * Whether every Python UDF in `hof`'s lambdas can be lifted out by `ExtractPythonUDFFromLambda`. + * Used by `CheckAnalysis` to decide whether to reject the plan; `hof` must be a *nest root* - one + * that iterates real columns, not a free lambda variable - because `CheckAnalysis` fires only at + * nest roots (see its guard) and this predicate validates the whole nest below the root. + * + * Both the row-at-a-time and the vectorized scalar eval types are liftable, in a single lambda or + * nested lambdas: an inner lambda's UDF is lifted onto its (enclosing-variable) argument and then + * re-lifted outward one array level at a time, so `transform(arr, i -> transform(i, x -> f(x)))` + * works (`f` lifts to a depth-2 element-wise UDF over `arr`). A UDF in a nested *argument*, + * `transform(arr, x -> transform(udf(x), y -> y))`, lifts onto `arr` the same way. + * + * These shapes still cannot be rewritten and are rejected: + * - a UDF in `aggregate` / `reduce`: the fold is sequential, so the UDF sees earlier steps' + * outputs, not array elements, and cannot be applied once to the whole array (see + * [[isRewritableShape]]). + * - a *nondeterministic iterated argument*, `filter(shuffle(arr), x -> f(x))` (at the root or + * any nested level): the rewrite references that argument several times (the carrier's `c0`, + * each lifted UDF's argument, the `map_keys`/`map_values` desugar, the pairwise `array_sort` + * path), and nondeterministic expressions are not subexpression-eliminated, so the copies + * would evaluate independently and disagree - keeping the results misaligned. (This is + * distinct from a nondeterministic UDF *call*, which `ExtractPythonUDFFromLambda.liftKey` + * keeps distinct but well-defined.) + * - a HOF (at any level) whose shape the rewrite does not model (see [[isRewritableShape]]). + */ + def canRewritePythonUDFInLambda(hof: HigherOrderFunction): Boolean = { + // Every Python UDF anywhere in the lambdas must be a rewritable flavor. `collect` is recursive, + // so this also covers UDFs in nested lambdas. + val allUDFsRewritable = hof.functions.forall { f => + f.collect { case udf: PythonUDF => udf }.forall(isElementwiseRewritableUDF) + } + // Reading a free lambda variable means `hof` is itself nested in an enclosing lambda, so the + // array it iterates is not a real column. Such an inner HOF is validated as part of its + // enclosing root's nest by `everyHofInNestRewritable`, never on its own. + val iteratesRealColumns = !hasFreeLambdaVariable(hof) + iteratesRealColumns && allUDFsRewritable && everyHofInNestRewritable(hof) + } + + /** + * Whether `hof` and every higher-order function nested within its lambda bodies is a rewritable + * shape (see [[isRewritableShape]]) with deterministic arguments. This is the recursive core that + * supports UDFs in *nested* lambdas: every HOF on the path from the root down to a UDF must be + * rewritable, because the rule lifts the UDF out one lambda level at a time. A nested HOF's + * iterated argument is legitimately an enclosing lambda variable, which is why the free-variable + * check in [[canRewritePythonUDFInLambda]] applies only at the root, not here. + */ + private def everyHofInNestRewritable(hof: HigherOrderFunction): Boolean = { + val nestedHofs = hof.functions.flatMap(_.collect { case h: HigherOrderFunction => h }) + (hof +: nestedHofs).forall { h => + isRewritableShape(h) && h.arguments.forall(_.deterministic) + } + } + + /** + * The structural assumption the rewrite makes: one lambda with plain-variable parameters, over at + * least one array- or map-valued argument (`transform`, `filter`, the map family, ...). + * `aggregate` / `reduce` fail this - they have two lambdas (`merge`, `finish`) - so a UDF in a + * fold is rejected: the fold is sequential, so the UDF sees earlier steps' outputs, not array + * elements. Checking the shape rather than listing classes means a new function of a familiar + * shape needs no change here. + * + * The function must also carry one of the result-type marker traits the rewrite dispatches on + * ([[ResultTypeFromArgument]] or [[ResultTypeFromFunction]]). Every built-in single-lambda HOF is + * marked today, but requiring it here keeps "analysis accepts exactly what the rule rewrites" + * structural: a future HOF missing both traits is rejected at analysis rather than slipping + * through and leaving the UDF inside the lambda at runtime. + */ + private def isRewritableShape(hof: HigherOrderFunction): Boolean = + hof.functions.length == 1 && + hof.functions.head.isInstanceOf[LambdaFunction] && + hof.functions.head.asInstanceOf[LambdaFunction].arguments + .forall(_.isInstanceOf[NamedLambdaVariable]) && + (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) && + hof.arguments.exists { a => + a.dataType.isInstanceOf[ArrayType] || a.dataType.isInstanceOf[MapType] + } + + /** + * Whether `e` references a [[NamedLambdaVariable]] that it does not itself bind, i.e. one bound + * by an enclosing lambda. Such an expression cannot be evaluated outside that lambda. + * + * Shared with `ExtractPythonUDFFromLambda` so the rule can re-check the nested-lambda guard on + * its own, rather than relying only on `CheckAnalysis` having already rejected such plans. + */ + def hasFreeLambdaVariable(e: Expression): Boolean = { + def check(expr: Expression, bound: Set[ExprId]): Boolean = expr match { + case LambdaFunction(function, arguments, _) => + check(function, bound ++ arguments.map(_.exprId)) + case v: NamedLambdaVariable => !bound.contains(v.exprId) + case other => other.children.exists(check(_, bound)) + } + check(e, Set.empty) + } + def isWindowPandasUDF(e: PythonFuncExpression): Boolean = { - // This is currently only `PythonUDAF` (which means SQL_GROUPED_AGG_PANDAS_UDF or - // SQL_GROUPED_AGG_ARROW_UDF), but we might - // support new types in the future, e.g, N -> N transform. - e.isInstanceOf[PythonUDAF] + // `PythonUDAF` (SQL_GROUPED_AGG_PANDAS_UDF or SQL_GROUPED_AGG_ARROW_UDF) and the incremental + // `PythonAggregate` are the Python aggregate functions that run over a window through the + // Python window operator, rather than the JVM SQL window path. We might support new types in + // the future, e.g. N -> N transform. + e.isInstanceOf[PythonUDAF] || e.isInstanceOf[PythonAggregate] } def correctEvalType(udf: PythonUDF, pythonUDFArrowFallbackOnUDT: Boolean): Int = { @@ -89,6 +284,41 @@ trait PythonFuncExpression extends NonSQLExpression with UserDefinedExpression { override def nullable: Boolean = true } + +case class TranspiledPythonUDF( + name: String, + pythonUDFExpr: Expression, + transpiledOptions: List[Expression], + // Per-option input-type categories ("numeric"/"string" per public param), + // parallel to `transpiledOptions`. ResolveTranspiledPythonUDFOptions prunes the + // options to those whose categories match the resolved input types (before + // CheckAnalysis can reject a type-incompatible option) and clears this field; + // ConvertToCatalyst then picks the first survivor or falls back to the Python + // UDF. Empty means "no restriction" (kept as-is). + optionInputCategories: List[List[String]] = Nil) extends Expression with Unevaluable { + require( + optionInputCategories.isEmpty || optionInputCategories.length == transpiledOptions.length, + s"optionInputCategories (${optionInputCategories.length}) must be parallel to " + + s"transpiledOptions (${transpiledOptions.length}) or empty" + ) + override def children: Seq[Expression] = pythonUDFExpr +: transpiledOptions + override def dataType: DataType = pythonUDFExpr.dataType + override def nullable: Boolean = pythonUDFExpr.nullable + override protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): + TranspiledPythonUDF = + copy(pythonUDFExpr = newChildren.head, transpiledOptions = newChildren.tail.toList) + final override val nodePatterns: Seq[TreePattern] = Seq(TRANSPILED_PYTHON_UDF) + + // True when every direct input to pythonUDFExpr is a plain PythonUDF (not a + // TranspiledPythonUDF). Used to decide whether to preserve the UDF batch pipeline + // rather than inserting a Catalyst node in the middle of a Python UDF chain. + def hasOnlyPythonUDFInputs: Boolean = + pythonUDFExpr.children.nonEmpty && + pythonUDFExpr.children.forall { + _.isInstanceOf[PythonUDF] + } +} + /** * A serialized version of a Python lambda function. This is a special expression, which needs a * dedicated physical operator to execute it, and thus can't be pushed down to data sources. @@ -100,7 +330,14 @@ case class PythonUDF( children: Seq[Expression], evalType: Int, udfDeterministic: Boolean, - resultId: ExprId = NamedExpression.newExprId) + resultId: ExprId = NamedExpression.newExprId, + // For an element-wise UDF lifted out of a higher-order function's lambda (see + // `ExtractPythonUDFFromLambda`), the number of `array` levels the Python worker flattens off + // each argument before invoking the function, and re-nests onto the result: 1 for a UDF in a + // single lambda, and one more for each enclosing lambda when the UDF is lifted out of a nested + // lambda (e.g. `transform(arr, i -> transform(i, x -> f(x)))` lifts `f` to depth 2). Ignored + // for every non-element-wise eval type, where it stays at its default of 1. + elementwiseNestingDepth: Int = 1) extends Expression with PythonFuncExpression with Unevaluable { lazy val resultAttribute: Attribute = AttributeReference(toPrettySQL(this), dataType, nullable)( @@ -168,6 +405,55 @@ case class PythonUDAF( copy(children = newChildren) } +/** + * A serialized Python aggregator that supports true incremental (partial) aggregation, the + * analog of the Scala typed `org.apache.spark.sql.expressions.Aggregator[IN, BUF, OUT]`. Unlike + * [[PythonUDAF]] (which materializes the whole group and calls Python once), this is planned as a + * two-stage aggregation by + * [[org.apache.spark.sql.execution.python.PythonIncrementalAggregateExec]]: a map-side PARTIAL + * stage folds input rows into a per-group buffer via the aggregator's `reduce`, and a post-shuffle + * FINAL stage + * merges the partial buffers via `merge` and produces the output via `finish`. + * + * `bufferSchema` is the schema of the intermediate buffer that crosses the shuffle between the two + * stages (the analog of the Scala aggregator's `bufferEncoder`). It is exposed here rather than via + * [[aggBufferAttributes]] because, like [[PythonUDAF]], this expression is unevaluable in the JVM; + * the physical operator derives the buffer attributes from `bufferSchema` directly. + */ +case class PythonAggregate( + name: String, + func: PythonFunction, + dataType: DataType, + children: Seq[Expression], + udfDeterministic: Boolean, + bufferSchema: StructType, + evalType: Int = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, + resultId: ExprId = NamedExpression.newExprId) + extends UnevaluableAggregateFunc with PythonFuncExpression { + + override def sql(isDistinct: Boolean): String = { + val distinct = if (isDistinct) "DISTINCT " else "" + s"$name($distinct${children.mkString(", ")})" + } + + override def toAggString(isDistinct: Boolean): String = { + val start = if (isDistinct) "(distinct " else "(" + name + children.mkString(start, ", ", ")") + s"#${resultId.id}$typeSuffix" + } + + override lazy val canonicalized: Expression = { + val canonicalizedChildren = children.map(_.canonicalized) + // `resultId` can be seen as cosmetic variation, as it doesn't affect the result. + this.copy(resultId = ExprId(-1)).withNewChildren(canonicalizedChildren) + } + + final override val nodePatterns: Seq[TreePattern] = Seq(PYTHON_UDF) + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): PythonAggregate = + copy(children = newChildren) +} + abstract class UnevaluableGenerator extends Generator { final override def eval(input: InternalRow): IterableOnce[InternalRow] = throw QueryExecutionErrors.cannotEvaluateExpressionError(this) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala index b4dd41092871b..f9a62309463a1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ScalaUDF.scala @@ -1163,14 +1163,19 @@ case class ScalaUDF( } else { s"$resultTerm = ($boxedType)$resultConverter.apply($getFuncResult)" } + // `failedExecuteUserDefinedFunctionError` returns a checked SparkException; throwing + // it directly compiles under Janino but not under the JDK compiler ("unreported + // exception"). Re-throw via Platform.throwException (no declared `throws`), the same + // sneaky-throw used elsewhere in codegen, to stay compatible with both backends. val callFunc = s""" |$boxedType $resultTerm = null; |try { | $funcInvocation; |} catch (Throwable e) { - | throw QueryExecutionErrors.failedExecuteUserDefinedFunctionError( - | "$functionName", "$inputTypesString", "$outputType", e); + | org.apache.spark.unsafe.Platform.throwException( + | QueryExecutionErrors.failedExecuteUserDefinedFunctionError( + | "$functionName", "$inputTypesString", "$outputType", e)); |} """.stripMargin diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala index 69dfdbfc9a08e..808083e932084 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SelectedField.scala @@ -85,7 +85,7 @@ object SelectedField { expr match { case a: Attribute => dataTypeOpt.map { dt => - StructField(a.name, dt, a.nullable) + StructField(a.name, dt, a.nullable, a.metadata) } case c: GetStructField => val field = c.childSchema(c.ordinal) @@ -109,7 +109,7 @@ object SelectedField { // This should not happen. throw QueryCompilationErrors.dataTypeUnsupportedByClassError(x, "GetArrayStructFields") } - val newField = StructField(field.name, newFieldDataType, field.nullable) + val newField = StructField(field.name, newFieldDataType, field.nullable, field.metadata) selectField(child, Option(ArrayType(struct(newField), containsNull))) case GetMapValue(child, key) if key.foldable => // GetMapValue does not select a field from a struct (i.e. prune the struct) so it can't be diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SortOrder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SortOrder.scala index 166866c90b877..60a0671a8affc 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SortOrder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/SortOrder.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.expressions import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TypeCheckResult -import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode} import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.util.TypeUtils import org.apache.spark.sql.types._ @@ -194,9 +194,13 @@ case class SortPrefix(child: SortOrder) extends UnaryExpression { override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { val childCode = child.child.genCode(ctx) val input = childCode.value - val BinaryPrefixCmp = classOf[BinaryPrefixComparator].getName - val DoublePrefixCmp = classOf[DoublePrefixComparator].getName - val StringPrefixCmp = classOf[StringPrefixComparator].getName + // Use javaSourceName to emit the binary name form (e.g. + // `PrefixComparators$DoublePrefixComparator`); Janino loads that name directly, + // and the JDK backend's rewriteInnerClassRefs converts it to the dotted source + // form javac requires. + val BinaryPrefixCmp = CodeGenerator.javaSourceName(classOf[BinaryPrefixComparator]) + val DoublePrefixCmp = CodeGenerator.javaSourceName(classOf[DoublePrefixComparator]) + val StringPrefixCmp = CodeGenerator.javaSourceName(classOf[StringPrefixComparator]) val prefixCode = child.child.dataType match { case BooleanType => s"$input ? 1L : 0L" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala index 0fec0bd3e00e5..096c1788df63f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ToStringBase.scala @@ -35,6 +35,13 @@ import org.apache.spark.util.SparkStringUtils trait ToStringBase { self: UnaryExpression with TimeZoneAwareExpression => + /** + * ISO 6.13 truncation applies only to user-written CAST / TRY_CAST. Implicit and + * store-assignment Casts keep the write-side length check. [[ToPrettyString]] is never + * a CAST, so it stays false. + */ + protected def truncateCharVarcharOnCast: Boolean = false + private lazy val dateFormatter = DateFormatter() private lazy val timeFormatter = new FractionTimeFormatter() private lazy val timestampFormatter = TimestampFormatter.getFractionFormatter(zoneId) @@ -57,14 +64,22 @@ trait ToStringBase { self: UnaryExpression with TimeZoneAwareExpression => // Returns a function to convert a value to pretty string. The function assumes input is not null. protected final def castToString( - from: DataType, to: StringConstraint = NoConstraint): Any => UTF8String = - to match { - case FixedLength(length) => - s => CharVarcharCodegenUtils.charTypeWriteSideCheck(castToString(from)(s), length) - case MaxLength(length) => - s => CharVarcharCodegenUtils.varcharTypeWriteSideCheck(castToString(from)(s), length) - case NoConstraint => castToString(from) + from: DataType, to: StringConstraint = NoConstraint): Any => UTF8String = { + val toUTF8String = castToString(from) + (to, from) match { + case (FixedLength(length), _: StringType) + if SQLConf.get.charVarcharStandardSemantics && truncateCharVarcharOnCast => + s => CharVarcharCodegenUtils.charTypeCast(toUTF8String(s), length) + case (MaxLength(length), _: StringType) + if SQLConf.get.charVarcharStandardSemantics && truncateCharVarcharOnCast => + s => CharVarcharCodegenUtils.varcharTypeCast(toUTF8String(s), length) + case (FixedLength(length), _) => + s => CharVarcharCodegenUtils.charTypeWriteSideCheck(toUTF8String(s), length) + case (MaxLength(length), _) => + s => CharVarcharCodegenUtils.varcharTypeWriteSideCheck(toUTF8String(s), length) + case (NoConstraint, _) => toUTF8String } + } // The Types Framework is the single integration point for framework types' cast-to-string, via // the zone-less formatUTF8. The cast's session zone is threaded into the lookup so TIMESTAMP_LTZ @@ -196,14 +211,22 @@ trait ToStringBase { self: UnaryExpression with TimeZoneAwareExpression => (c, evPrim) => { val tmpVar = ctx.freshVariable("tmp", classOf[UTF8String]) val castToString = castToStringCode(from, ctx)(c, tmpVar) - val maintainConstraint = to match { - case FixedLength(length) => + val maintainConstraint = (to, from) match { + case (FixedLength(length), _: StringType) + if SQLConf.get.charVarcharStandardSemantics && truncateCharVarcharOnCast => + code"""$evPrim = org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils + .charTypeCast($tmpVar, $length);""".stripMargin + case (MaxLength(length), _: StringType) + if SQLConf.get.charVarcharStandardSemantics && truncateCharVarcharOnCast => + code"""$evPrim = org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils + .varcharTypeCast($tmpVar, $length);""".stripMargin + case (FixedLength(length), _) => code"""$evPrim = org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils .charTypeWriteSideCheck($tmpVar, $length);""".stripMargin - case MaxLength(length) => + case (MaxLength(length), _) => code"""$evPrim = org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils .varcharTypeWriteSideCheck($tmpVar, $length);""".stripMargin - case NoConstraint => code"$evPrim = $tmpVar;" + case (NoConstraint, _) => code"$evPrim = $tmpVar;" } code""" UTF8String $tmpVar; @@ -217,10 +240,19 @@ trait ToStringBase { self: UnaryExpression with TimeZoneAwareExpression => from: DataType, ctx: CodegenContext): (ExprValue, ExprValue) => Block = { from match { case BinaryType => + // Pass the public BinaryFormatter trait as the reference's cast type. + // `binaryFormatter` is a lambda (UTF8String.fromBytes); its runtime class is a + // non-nameable synthetic (e.g. ToStringBase$$anonfun$binaryFormatter$N), which + // the JDK compiler cannot reference ("cannot find symbol"); Janino tolerates it. val bf = JavaCode.global( - ctx.addReferenceObj("binaryFormatter", binaryFormatter), + ctx.addReferenceObj("binaryFormatter", binaryFormatter, classOf[BinaryFormatter].getName), classOf[BinaryFormatter]) - (c, evPrim) => code"$evPrim = $bf.apply($c);" + // `BinaryFormatter` extends `Array[Byte] => UTF8String` (a Function1). The JDK + // compiler resolves `bf.apply(c)` through the parameterised signature and infers + // `UTF8String`, but Janino binds it to the erased `apply(Object)` and infers + // `Object`, so the assignment to the `UTF8String` result needs an explicit cast + // to compile under both backends. + (c, evPrim) => code"$evPrim = (UTF8String) $bf.apply($c);" case DateType => val df = JavaCode.global( ctx.addReferenceObj("dateFormatter", dateFormatter), diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TryEval.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TryEval.scala index 289b102c1f6f9..b41f690f86900 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TryEval.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TryEval.scala @@ -329,6 +329,12 @@ case class TryToBinary( @ExpressionDescription( usage = "_FUNC_(class, method[, arg1[, arg2 ..]]) - This is a special version of `reflect` that" + " performs the same operation, but returns a NULL value instead of raising an error if the invoke method thrown exception.", + arguments = """ + Arguments: + * class - A string literal with the fully qualified name of the class. + * method - A string literal with the name of the static method to invoke. + * arg1, arg2, ... - Optional arguments passed to the invoked method. + """, examples = """ Examples: > SELECT _FUNC_('java.util.UUID', 'randomUUID'); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala index 249e3955a81f4..8183e6a447bd6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/UnwrapUDT.scala @@ -25,6 +25,8 @@ import org.apache.spark.sql.types.{DataType, UserDefinedType} /** * Unwrap UDT data type column into its underlying type. + * + * @see [[WrapUDT]] for converting an underlying SQL type column to a UDT. */ case class UnwrapUDT(child: Expression) extends UnaryExpression with NonSQLExpression { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala new file mode 100644 index 0000000000000..770dd1f8fa470 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/WrapUDT.scala @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.expressions + +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.expressions.Cast._ +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber +import org.apache.spark.sql.types.{DataType, UserDefinedType} + +/** + * Wrap a column with a UDT whose underlying SQL type matches the column data type. + * + * @see [[UnwrapUDT]] for converting a UDT column to its underlying SQL type. + */ +case class WrapUDT(child: Expression, udt: UserDefinedType[_]) + extends UnaryExpression with NonSQLExpression { + + def this(child: Expression, udt: Expression) = { + this(child, WrapUDT.parseUDT(udt)) + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + child.genCode(ctx) + } + + override def checkInputDataTypes(): TypeCheckResult = { + if (DataTypeUtils.sameType(child.dataType, udt.sqlType)) { + TypeCheckResult.TypeCheckSuccess + } else { + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> toSQLType(udt.sqlType), + "inputSql" -> toSQLExpr(child), + "inputType" -> toSQLType(child.dataType))) + } + } + + override def dataType: DataType = udt + + override def nullSafeEval(input: Any): Any = input + + override def prettyName: String = "wrap_udt" + + override protected def withNewChildInternal(newChild: Expression): WrapUDT = { + copy(child = newChild) + } +} + +object WrapUDT { + private def parseUDT(expression: Expression): UserDefinedType[_] = { + ExprUtils.evalTypeExpr(expression) match { + case udt: UserDefinedType[_] => udt + case dataType => + throw new AnalysisException( + errorClass = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "sqlExpr" -> toSQLExpr(expression), + "paramIndex" -> ordinalNumber(1), + "requiredType" -> toSQLType("UserDefinedType"), + "inputSql" -> toSQLExpr(expression), + "inputType" -> toSQLType(dataType))) + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala index 186056ab7e6a9..8ed6590793daa 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala @@ -404,7 +404,7 @@ class ApproxTopKAggregateBuffer[T](val sketch: ItemsSketch[T], private var nullC if (UnsafeRowUtils.isBinaryStable(st)) { sketch.asInstanceOf[ItemsSketch[String]].update(orig.toString) } else { - val cKey = CollationFactory.getCollationKey(orig, st.collationId).toString + val cKey = CollationFactory.getCollationKeyBytes(orig, st.collationId) sketch.asInstanceOf[ItemsSketch[CollatedString]] .update(new CollatedString(cKey, orig.toString)) } @@ -798,6 +798,15 @@ object CombineInternal { _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch. `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is specified, it will be set for the combined sketch. If maxItemsTracked is not specified, the input sketches must have the same maxItemsTracked value, otherwise an error will be thrown. The output sketch will use the same value from the input sketches. """, + arguments = """ + Arguments: + * state - The sketch state to combine, as produced by approx_top_k_accumulate. + An expression that evaluates to the sketch state struct. + * maxItemsTracked - Optional. The maximum number of items to track in the combined + sketch, with an upper limit of 1000000. An expression that evaluates to an integer. + Must be a constant. If not specified, the input sketches must share the same + maxItemsTracked value, which is used for the output sketch. + """, examples = """ Examples: > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr)); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala index 19e30c53b436d..b9caae5467ce8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproximatePercentile.scala @@ -66,6 +66,16 @@ import org.apache.spark.util.ArrayImplicits._ In this case, returns the approximate percentile array of column `col` at the given percentage array. """, + arguments = """ + Arguments: + * col - The numeric, ANSI interval or TIME column whose percentile is + computed. + * percentage - A value (or array of values) between 0.0 and 1.0 specifying + the percentile(s) to compute. + * accuracy - Optional. A positive numeric literal (default: 10000) that + controls approximation accuracy at the cost of memory. Higher values + yield better accuracy; `1.0/accuracy` is the relative error. + """, examples = """ Examples: > SELECT _FUNC_(col, array(0.5, 0.4, 0.1), 100) FROM VALUES (0), (1), (2), (10) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala index bc78dfdf8cec1..0502cf38f8492 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala @@ -136,7 +136,12 @@ case class Corr( override val evaluateExpression: Expression = { If(n === 0.0, Literal.create(null, DoubleType), - If(n === 1.0, divideByZeroEvalResult, ck / sqrt(xMk * yMk))) + If(n === 1.0, divideByZeroEvalResult, + // For an exactly zero variance accumulator, corr is undefined. Keep this check after + // n == 1.0 to preserve the legacy single-pair result, and before division to avoid the + // ANSI divide-by-zero error for this case. + If(xMk === 0.0 || yMk === 0.0, Literal.create(null, DoubleType), + ck / sqrt(xMk * yMk)))) } override def prettyName: String = "corr" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala index 758ef22f0a2c2..f5f248c48387a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Count.scala @@ -34,6 +34,11 @@ import org.apache.spark.sql.types._ _FUNC_(DISTINCT expr[, expr...]) - Returns the number of rows for which the supplied expression(s) are unique and non-null. """, + arguments = """ + Arguments: + * expr - One or more expressions. A row is counted only when all supplied + expressions are non-null. Use `*` to count all rows, including rows with nulls. + """, examples = """ Examples: > SELECT _FUNC_(*) FROM VALUES (NULL), (5), (5), (20) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala index f0a27677628dc..a1eb024488bd6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CountMinSketchAgg.scala @@ -206,6 +206,13 @@ case class CountMinSketchAgg( `CountMinSketch` before usage. Count-min sketch is a probabilistic data structure used for cardinality estimation using sub-linear space. """, + arguments = """ + Arguments: + * col - The column to build the count-min sketch from. + * eps - A double literal for the relative error of the sketch. + * confidence - A double literal for the confidence of the sketch. + * seed - An integer literal used as the random seed. + """, examples = """ Examples: > SELECT hex(_FUNC_(col, 0.5d, 0.5d, 1)) FROM VALUES (1), (2), (1) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala index dadf7ac53c596..fe89a535f264e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HistogramNumeric.scala @@ -51,6 +51,12 @@ import org.apache.spark.sql.util.NumericHistogram statistical computing packages. Note: the output type of the 'x' field in the return value is propagated from the input value consumed in the aggregate function. """, + arguments = """ + Arguments: + * expr - A numeric, date, timestamp, or interval expression whose values are aggregated + into the histogram. + * nb - A foldable integer expression (at least 2) giving the number of histogram bins. + """, examples = """ Examples: > SELECT _FUNC_(col, 5) FROM VALUES (0), (1), (2), (10) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala index f304b43358ad1..56f203aa9c4e1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/HyperLogLogPlusPlus.scala @@ -50,6 +50,12 @@ import org.apache.spark.sql.types._ usage = """ _FUNC_(expr[, relativeSD]) - Returns the estimated cardinality by HyperLogLog++. `relativeSD` defines the maximum relative standard deviation allowed.""", + arguments = """ + Arguments: + * expr - An expression of any type whose distinct values are counted. + * relativeSD - An optional double literal for the maximum relative standard + deviation allowed. Defaults to 0.05. + """, examples = """ Examples: > SELECT _FUNC_(col1) FROM VALUES (1), (1), (2), (2), (3) tab(col1); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala index f49297eba88bd..f2a36e571cd8a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Max.scala @@ -27,6 +27,11 @@ import org.apache.spark.sql.types._ @ExpressionDescription( usage = "_FUNC_(expr) - Returns the maximum value of `expr`.", + arguments = """ + Arguments: + * expr - An expression of any orderable type whose maximum value across the group is + returned. NULL values are ignored. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (10), (50), (20) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MaxMinByK.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MaxMinByK.scala index a0a7acd930974..567beb9a74846 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MaxMinByK.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MaxMinByK.scala @@ -240,6 +240,13 @@ case class MaxMinByK( maximum values of `y`, sorted in descending order by `y`. Returns NULL if there are no non-NULL ordering values. """, + arguments = """ + Arguments: + * x - The value expression to return. + * y - The ordering expression whose maximum selects the value of `x`. + * k - An optional positive integer. When present, returns an array of the `k` values + of `x` associated with the largest values of `y`. + """, examples = """ Examples: > SELECT _FUNC_(x, y) FROM VALUES ('a', 10), ('b', 50), ('c', 20) AS tab(x, y); @@ -275,6 +282,13 @@ object MaxByBuilder extends ExpressionBuilder { minimum values of `y`, sorted in ascending order by `y`. Returns NULL if there are no non-NULL ordering values. """, + arguments = """ + Arguments: + * x - The value expression to return. + * y - The ordering expression whose minimum selects the value of `x`. + * k - An optional positive integer. When present, returns an array of the `k` values + of `x` associated with the smallest values of `y`. + """, examples = """ Examples: > SELECT _FUNC_(x, y) FROM VALUES ('a', 10), ('b', 50), ('c', 20) AS tab(x, y); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala index eaef7b6bec113..7585b1e80b2c2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Min.scala @@ -27,6 +27,11 @@ import org.apache.spark.sql.types._ @ExpressionDescription( usage = "_FUNC_(expr) - Returns the minimum value of `expr`.", + arguments = """ + Arguments: + * expr - An expression of any orderable type whose minimum value across the group is + returned. NULL values are ignored. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (10), (-1), (20) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala index bdc5248ffa736..0ed76e0bfa7de 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/collect.scala @@ -104,6 +104,10 @@ abstract class Collect[T <: Growable[Any] with Iterable[Any]] extends TypedImper */ @ExpressionDescription( usage = "_FUNC_(expr) - Collects and returns a list of non-unique elements.", + arguments = """ + Arguments: + * expr - An expression of any type whose values are collected into a list. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (1), (2), (1) AS tab(col); @@ -181,6 +185,10 @@ case class CollectList( */ @ExpressionDescription( usage = "_FUNC_(expr) - Collects and returns a set of unique elements.", + arguments = """ + Arguments: + * expr - An expression of any type whose values are collected into a set. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (1), (2), (1) AS tab(col); @@ -316,6 +324,179 @@ case class CollectSet( copy(child = newChild) } +/** + * Collect the distinct union of the elements of an array-typed input across rows. + * + * Unlike collect_set, whose input is a scalar and whose output is the set of those scalars, + * collect_union's input is itself an array and its output is the set of the array's + * *elements* unioned across all rows. The aggregation buffer holds only the distinct + * elements (a set), so its size is bounded by the element universe rather than by the + * number of input rows. + * + * Null handling mirrors collect_set: by default (IGNORE NULLS) null elements are dropped. + * With RESPECT NULLS, one null element is kept, in which case collect_union is equivalent to + * `array_distinct(flatten(collect_list(arr)))`. + * + * @param ignoreNulls when true (IGNORE NULLS, the default), null elements are excluded from + * the result array. When false (RESPECT NULLS), a single null element is + * kept. + */ +@ExpressionDescription( + usage = + "_FUNC_(expr) - Collects and returns the distinct union of the elements of array `expr`.", + arguments = """ + Arguments: + * expr - An array expression whose elements are collected into a set across rows. + """, + examples = """ + Examples: + > SELECT _FUNC_(col) FROM VALUES (array(1, 2)), (array(2, 3)), (array(1)) AS tab(col); + [1,2,3] + """, + note = """ + The function is non-deterministic because the order of collected results depends + on the order of the rows which may be non-deterministic after a shuffle. + """, + group = "agg_funcs", + since = "4.3.0") +case class CollectUnion( + child: Expression, + mutableAggBufferOffset: Int = 0, + inputAggBufferOffset: Int = 0, + ignoreNulls: Boolean = true) + extends Collect[mutable.HashSet[Any]] with QueryErrorsBase with UnaryLike[Expression] { + + def this(child: Expression) = this(child, 0, 0, true) + + // The input is guarded by checkInputDataTypes to be an ArrayType; this is its element type. + private lazy val elementType: DataType = child.dataType match { + case ArrayType(et, _) => et + case other => other + } + + // The result array contains a null only when null elements are respected (RESPECT NULLS). + override protected def bufferContainsNull: Boolean = !ignoreNulls + + // Result is array<elementType>; containsNull is true iff null elements are kept. + override def dataType: DataType = ArrayType(elementType, containsNull = bufferContainsNull) + + // The buffer stores distinct elements. Mirror CollectSet's keying so equality is correct + // for float/double (bit pattern) and binary (byte array) element types. + override lazy val bufferElementType: DataType = elementType match { + case BinaryType => ArrayType(ByteType) + case DoubleType => LongType + case FloatType => IntegerType + case other => other + } + + @transient private lazy val complexNormalizer: Any => Any = { + val ref = BoundReference(0, elementType, nullable = true) + val proj = UnsafeProjection.create(NormalizeFloatingNumbers.normalize(ref)) + (value: Any) => InternalRow.copyValue(proj(InternalRow(value)).get(0, elementType)) + } + + override def convertToBufferElement(value: Any): Any = elementType match { + // See CollectSet.convertToBufferElement for why binary/float/double are keyed specially. + case BinaryType => UnsafeArrayData.fromPrimitiveArray(value.asInstanceOf[Array[Byte]]) + case DoubleType => + java.lang.Double.doubleToLongBits( + NormalizeFloatingNumbers.DOUBLE_NORMALIZER(value).asInstanceOf[Double]) + case FloatType => + java.lang.Float.floatToIntBits( + NormalizeFloatingNumbers.FLOAT_NORMALIZER(value).asInstanceOf[Float]) + case dt if NormalizeFloatingNumbers.needNormalize(dt) => complexNormalizer(value) + case _ => InternalRow.copyValue(value) + } + + // Iterate the input array and add each element to the set. NULL input arrays are skipped; + // a NULL element is dropped when ignoreNulls (IGNORE NULLS) and kept otherwise (RESPECT + // NULLS), where the HashSet naturally dedups it to a single null. + override def update( + buffer: mutable.HashSet[Any], + input: InternalRow): mutable.HashSet[Any] = { + val arr = child.eval(input) + if (arr != null) { + arr.asInstanceOf[ArrayData].foreach(elementType, (_, element: Any) => + if (element != null) { + buffer += convertToBufferElement(element) + } else if (!ignoreNulls) { + buffer += null + }) + } + buffer + } + + override def eval(buffer: mutable.HashSet[Any]): Any = { + val array = elementType match { + case BinaryType => + buffer.iterator.map { + case null => null + case v => v.asInstanceOf[ArrayData].toByteArray() + }.toArray[Any] + case DoubleType => + buffer.iterator.map { + case null => null + case v => java.lang.Double.longBitsToDouble(v.asInstanceOf[Long]) + }.toArray[Any] + case FloatType => + buffer.iterator.map { + case null => null + case v => java.lang.Float.intBitsToFloat(v.asInstanceOf[Int]) + }.toArray[Any] + case _ => buffer.toArray + } + new GenericArrayData(array) + } + + override def checkInputDataTypes(): TypeCheckResult = child.dataType match { + case ArrayType(et, _) + if !et.existsRecursively(_.isInstanceOf[MapType]) && UnsafeRowUtils.isBinaryStable(et) => + TypeCheckResult.TypeCheckSuccess + case ArrayType(_, _) => + DataTypeMismatch( + errorSubClass = "UNSUPPORTED_INPUT_TYPE", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), + "dataType" -> (s"${toSQLType(MapType)} " + "or \"COLLATED STRING\"") + ) + ) + case _ => + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> toSQLType(ArrayType), + "inputSql" -> toSQLExpr(child), + "inputType" -> toSQLType(child.dataType) + ) + ) + } + + override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): ImperativeAggregate = + copy(mutableAggBufferOffset = newMutableAggBufferOffset) + + override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate = + copy(inputAggBufferOffset = newInputAggBufferOffset) + + override def prettyName: String = "collect_union" + + override def createAggregationBuffer(): mutable.HashSet[Any] = mutable.HashSet.empty + + override def toString: String = { + val ignoreNullsStr = if (ignoreNulls) "" else " respect nulls" + s"$prettyName($child)$ignoreNullsStr" + } + + override def sql(isDistinct: Boolean): String = { + val distinct = if (isDistinct) "DISTINCT " else "" + val nullsStr = if (ignoreNulls) "" else " RESPECT NULLS" + s"$prettyName($distinct${child.sql})$nullsStr" + } + + override protected def withNewChildInternal(newChild: Expression): CollectUnion = + copy(child = newChild) +} + /** * Collect the top-k elements. This expression is dedicated only for Spark-ML. * @param reverse when true, returns the smallest k elements. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala index 952b331b7227d..0a06c54409c08 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala @@ -243,6 +243,14 @@ object HllSketchAgg { _FUNC_(expr, allowDifferentLgConfigK) - Returns the merged HllSketch's updatable binary representation. `allowDifferentLgConfigK` (optional) Allow sketches with different lgConfigK values to be unioned (defaults to false).""", + arguments = """ + Arguments: + * expr - The binary representation of an HllSketch to merge. + An expression that evaluates to binary. + * allowDifferentLgConfigK - Optional. Whether to allow sketches with different + lgConfigK values to be unioned. An expression that evaluates to a boolean. + Defaults to false. + """, examples = """ Examples: > SELECT hll_sketch_estimate(_FUNC_(sketch, true)) FROM (SELECT hll_sketch_agg(col) as sketch FROM VALUES (1) tab(col) UNION ALL SELECT hll_sketch_agg(col, 20) as sketch FROM VALUES (1) tab(col)); @@ -327,6 +335,30 @@ case class HllUnionAgg( } } + /** + * Merges `sketch` into the Union acting as the aggregation buffer, instantiating it if absent. + * + * An empty sketch holds no coupons, so it carries no precision: it is exempt from the lgConfigK + * check, and an empty Union is re-seeded at the lgConfigK of the first non-empty sketch. This + * keeps the empty sketch that `eval` emits for an all-NULL group mergeable, and makes the + * result independent of the order in which rows reach the aggregate. + * + * @param unionOption A previously initialized Union instance, or None + * @param sketch The sketch to merge in + */ + private def mergeSketch(unionOption: Option[Union], sketch: HllSketch): Option[Union] = { + val union = unionOption match { + case Some(buffer) if buffer.isEmpty && !sketch.isEmpty => new Union(sketch.getLgConfigK) + case Some(buffer) => buffer + case None => new Union(sketch.getLgConfigK) + } + if (!union.isEmpty && !sketch.isEmpty) { + compareLgConfigK(union.getLgConfigK, sketch.getLgConfigK) + } + union.update(sketch) + Some(union) + } + /** * Update the Union instance with the HllSketch byte array obtained from the row. * @@ -340,10 +372,7 @@ case class HllUnionAgg( case BinaryType => try { val sketch = HllSketch.wrap(Memory.wrap(v.asInstanceOf[Array[Byte]])) - val union = unionOption.getOrElse(new Union(sketch.getLgConfigK)) - compareLgConfigK(union.getLgConfigK, sketch.getLgConfigK) - union.update(sketch) - Some(union) + mergeSketch(unionOption, sketch) } catch { case _: SketchesArgumentException | _: java.lang.Error | _: ArrayIndexOutOfBoundsException => @@ -365,10 +394,8 @@ case class HllUnionAgg( */ override def merge(unionOption: Option[Union], inputOption: Option[Union]): Option[Union] = { (unionOption, inputOption) match { - case (Some(union), Some(input)) => - compareLgConfigK(union.getLgConfigK, input.getLgConfigK) - union.update(input.getResult(targetType)) - Some(union) + case (Some(_), Some(input)) => + mergeSketch(unionOption, input.getResult(targetType)) // unclear if these scenarios can ever occur case (Some(_), None) => unionOption diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/kllAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/kllAggregates.scala index a811736c5ee23..f13bad8291c66 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/kllAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/kllAggregates.scala @@ -57,9 +57,10 @@ import org.apache.spark.sql.types.{AbstractDataType, BinaryType, ByteType, DataT arguments = """ Arguments: * expr - The expression to aggregate into the KLL sketch. - An expression that evaluates to an integral. - * k - The parameter controlling the size and accuracy of the sketch. - An expression that evaluates to an integer. Must be a constant. + An expression that evaluates to an integral. + * k - Optional. The parameter controlling the size and accuracy of the sketch. + An expression that evaluates to an integer between 8 and 65535. Must be a + constant. Defaults to 200. """, examples = """ Examples: @@ -208,6 +209,14 @@ case class KllSketchAggBigint( The optional k parameter controls the size and accuracy of the sketch (default 200, range 8-65535). Larger k values provide more accurate quantile estimates but result in larger, slower sketches. """, + arguments = """ + Arguments: + * expr - The expression to aggregate into the KLL sketch. + An expression that evaluates to a float. + * k - Optional. The parameter controlling the size and accuracy of the sketch. + An expression that evaluates to an integer between 8 and 65535. Must be a + constant. Defaults to 200. + """, examples = """ Examples: > SELECT LENGTH(kll_sketch_to_string_float(_FUNC_(col))) > 0 FROM VALUES (CAST(1.0 AS FLOAT)), (CAST(2.0 AS FLOAT)), (CAST(3.0 AS FLOAT)), (CAST(4.0 AS FLOAT)), (CAST(5.0 AS FLOAT)) tab(col); @@ -347,9 +356,10 @@ case class KllSketchAggFloat( arguments = """ Arguments: * expr - The expression to aggregate into the KLL sketch. - An expression that evaluates to a float or double. - * k - The parameter controlling the size and accuracy of the sketch. - An expression that evaluates to an integer. Must be a constant. + An expression that evaluates to a float or double. + * k - Optional. The parameter controlling the size and accuracy of the sketch. + An expression that evaluates to an integer between 8 and 65535. Must be a + constant. Defaults to 200. """, examples = """ Examples: @@ -493,6 +503,14 @@ case class KllSketchAggDouble( The optional k parameter controls the size and accuracy of the merged sketch (range 8-65535). If k is not specified, the merged sketch adopts the k value from the first input sketch. """, + arguments = """ + Arguments: + * expr - The expression to merge into the KLL sketch. + An expression that evaluates to a binary KLL sketch representation. + * k - Optional. The parameter controlling the size and accuracy of the merged + sketch. An expression that evaluates to an integer between 8 and 65535. Must be + a constant. Defaults to the k value of the first input sketch. + """, examples = """ Examples: > SELECT kll_sketch_get_n_bigint(_FUNC_(sketch)) FROM (SELECT kll_sketch_agg_bigint(col) as sketch FROM VALUES (1), (2), (3) tab(col) UNION ALL SELECT kll_sketch_agg_bigint(col) as sketch FROM VALUES (4), (5), (6) tab(col)) t; @@ -566,6 +584,14 @@ case class KllMergeAggBigint( The optional k parameter controls the size and accuracy of the merged sketch (range 8-65535). If k is not specified, the merged sketch adopts the k value from the first input sketch. """, + arguments = """ + Arguments: + * expr - The expression to merge into the KLL sketch. + An expression that evaluates to a binary KLL sketch representation. + * k - Optional. The parameter controlling the size and accuracy of the merged + sketch. An expression that evaluates to an integer between 8 and 65535. Must be + a constant. Defaults to the k value of the first input sketch. + """, examples = """ Examples: > SELECT kll_sketch_get_n_float(_FUNC_(sketch)) FROM (SELECT kll_sketch_agg_float(col) as sketch FROM VALUES (CAST(1.0 AS FLOAT)), (CAST(2.0 AS FLOAT)), (CAST(3.0 AS FLOAT)) tab(col) UNION ALL SELECT kll_sketch_agg_float(col) as sketch FROM VALUES (CAST(4.0 AS FLOAT)), (CAST(5.0 AS FLOAT)), (CAST(6.0 AS FLOAT)) tab(col)) t; @@ -639,6 +665,14 @@ case class KllMergeAggFloat( The optional k parameter controls the size and accuracy of the merged sketch (range 8-65535). If k is not specified, the merged sketch adopts the k value from the first input sketch. """, + arguments = """ + Arguments: + * expr - The expression to merge into the KLL sketch. + An expression that evaluates to a binary KLL sketch representation. + * k - Optional. The parameter controlling the size and accuracy of the merged + sketch. An expression that evaluates to an integer between 8 and 65535. Must be + a constant. Defaults to the k value of the first input sketch. + """, examples = """ Examples: > SELECT kll_sketch_get_n_double(_FUNC_(sketch)) FROM (SELECT kll_sketch_agg_double(col) as sketch FROM VALUES (CAST(1.0 AS DOUBLE)), (CAST(2.0 AS DOUBLE)), (CAST(3.0 AS DOUBLE)) tab(col) UNION ALL SELECT kll_sketch_agg_double(col) as sketch FROM VALUES (CAST(4.0 AS DOUBLE)), (CAST(5.0 AS DOUBLE)), (CAST(6.0 AS DOUBLE)) tab(col)) t; diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/percentiles.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/percentiles.scala index cff1c6b9750f5..4427dc25097e1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/percentiles.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/percentiles.scala @@ -271,6 +271,15 @@ abstract class PercentileBase positive integral """, + arguments = """ + Arguments: + * col - The column to compute the percentile of. + An expression that evaluates to a numeric, ANSI interval, or time. + * percentage - The percentile(s) to compute, each between 0.0 and 1.0. Either a + single numeric value or an array of numeric values. Must be foldable. + * frequency - Optional. The number of times each value should be counted. + An expression that evaluates to a positive integral value. Defaults to 1. + """, examples = """ Examples: > SELECT _FUNC_(col, 0.3) FROM VALUES (0), (10) AS tab(col); @@ -493,6 +502,12 @@ case class PercentileDisc( usage = "_FUNC_(percentage) WITHIN GROUP (ORDER BY col) - Return a percentile value based on " + "a continuous distribution of numeric, ANSI interval or TIME column `col` at the given " + "`percentage` (specified in ORDER BY clause).", + arguments = """ + Arguments: + * percentage - The percentile to compute, between 0.0 and 1.0. Must be foldable. + * col - The column to compute the percentile of, specified in the ORDER BY clause. + An expression that evaluates to a numeric, ANSI interval, or time. + """, examples = """ Examples: > SELECT _FUNC_(0.25) WITHIN GROUP (ORDER BY col) FROM VALUES (0), (10) AS tab(col); @@ -519,6 +534,12 @@ object PercentileContBuilder extends ExpressionBuilder { usage = "_FUNC_(percentage) WITHIN GROUP (ORDER BY col) - Return a percentile value based on " + "a discrete distribution of numeric, ANSI interval or TIME column `col` at the given " + "`percentage` (specified in ORDER BY clause).", + arguments = """ + Arguments: + * percentage - The percentile to compute, between 0.0 and 1.0. Must be foldable. + * col - The column to compute the percentile of, specified in the ORDER BY clause. + An expression that evaluates to a numeric, ANSI interval, or time. + """, examples = """ Examples: > SELECT _FUNC_(0.25) WITHIN GROUP (ORDER BY col) FROM VALUES (0), (10) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala index d8cd93129ec81..a5326ae76c9f6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala @@ -323,6 +323,13 @@ case class ThetaSketchAgg( _FUNC_(expr, lgNomEntries) - Returns the ThetaSketch's Compact binary representation. `lgNomEntries` (optional) the log-base-2 of Nominal Entries, with Nominal Entries deciding the number buckets or slots for the ThetaSketch.""", + arguments = """ + Arguments: + * expr - A binary expression of Compact ThetaSketch representations to union. + * lgNomEntries - An optional integer expression, the log-base-2 of nominal entries which + sets the number of buckets for the resulting ThetaSketch. When omitted, it defaults to + the default log nominal entries. + """, examples = """ Examples: > SELECT theta_sketch_estimate(_FUNC_(sketch)) FROM (SELECT theta_sketch_agg(col) as sketch FROM VALUES (1) tab(col) UNION ALL SELECT theta_sketch_agg(col, 20) as sketch FROM VALUES (1) tab(col)); @@ -507,6 +514,10 @@ case class ThetaUnionAgg( usage = """ _FUNC_(expr) - Returns the ThetaSketch's Compact binary representation by intersecting all the Theta sketches in the input column.""", + arguments = """ + Arguments: + * expr - A binary expression of Compact ThetaSketch representations to intersect. + """, examples = """ Examples: > SELECT theta_sketch_estimate(_FUNC_(sketch)) FROM (SELECT theta_sketch_agg(col) as sketch FROM VALUES (1) tab(col) UNION ALL SELECT theta_sketch_agg(col, 20) as sketch FROM VALUES (1) tab(col)); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/arithmetic.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/arithmetic.scala index 03bb84c2623ea..d0d7507e3250f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/arithmetic.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/arithmetic.scala @@ -502,6 +502,13 @@ object Add { @ExpressionDescription( usage = "expr1 _FUNC_ expr2 - Returns `expr1`-`expr2`.", + arguments = """ + Arguments: + * expr1 - The minuend. + An expression that evaluates to a numeric, interval, date, timestamp, or time. + * expr2 - The subtrahend. + An expression that evaluates to a numeric, interval, date, timestamp, or time. + """, examples = """ Examples: > SELECT 2 _FUNC_ 1; @@ -1313,6 +1320,12 @@ object Pmod { */ @ExpressionDescription( usage = "_FUNC_(expr, ...) - Returns the least value of all parameters, skipping null values.", + arguments = """ + Arguments: + * expr - An expression of any orderable type. At least two expressions must be given, + and all arguments must share a common type. Null values are skipped; the result is + null only when all arguments are null. + """, examples = """ Examples: > SELECT _FUNC_(10, 9, 2, 4, 3); @@ -1402,6 +1415,12 @@ case class Least(children: Seq[Expression]) extends ComplexTypeMergingExpression */ @ExpressionDescription( usage = "_FUNC_(expr, ...) - Returns the greatest value of all parameters, skipping null values.", + arguments = """ + Arguments: + * expr - An expression of any orderable type. At least two expressions must be given, + and all arguments must share a common type. Null values are skipped; the result is + null only when all arguments are null. + """, examples = """ Examples: > SELECT _FUNC_(10, 9, 2, 4, 3); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/avroSqlFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/avroSqlFunctions.scala index aea04913e468d..cb553aeb96ffe 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/avroSqlFunctions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/avroSqlFunctions.scala @@ -40,6 +40,13 @@ import org.apache.spark.util.Utils usage = """ _FUNC_(child, jsonFormatSchema, options) - Converts a binary Avro value into a Catalyst value. """, + arguments = """ + Arguments: + * child - A binary value containing Avro-encoded data. + * jsonFormatSchema - A constant string with the Avro schema in JSON format used to + interpret the data. + * options - A constant map of string keys and values holding conversion options. + """, examples = """ Examples: > SELECT _FUNC_(s, '{"type": "record", "name": "struct", "fields": [{ "name": "u", "type": ["int","string"] }]}', map()) IS NULL AS result FROM (SELECT NAMED_STRUCT('u', NAMED_STRUCT('member0', member0, 'member1', member1)) AS s FROM VALUES (1, NULL), (NULL, 'a') tab(member0, member1)); @@ -141,6 +148,12 @@ case class FromAvro(child: Expression, jsonFormatSchema: Expression, options: Ex _FUNC_(child[, jsonFormatSchema]) - Converts a Catalyst binary input value into its corresponding Avro format result. """, + arguments = """ + Arguments: + * child - A value to encode into Avro binary format. + * jsonFormatSchema - An optional constant string with the Avro schema in JSON format + used to encode the value. If omitted, the schema is inferred from the input type. + """, examples = """ Examples: > SELECT _FUNC_(s, '{"type": "record", "name": "struct", "fields": [{ "name": "u", "type": ["int","string"] }]}') IS NULL FROM (SELECT NULL AS s); @@ -215,6 +228,11 @@ case class ToAvro(child: Expression, jsonFormatSchema: Expression) usage = """ _FUNC_(jsonFormatSchema, options) - Returns schema in the DDL format of the avro schema in JSON string format. """, + arguments = """ + Arguments: + * jsonFormatSchema - A constant string with the Avro schema in JSON format. + * options - A constant map of string keys and values holding conversion options. + """, examples = """ Examples: > SELECT _FUNC_('{"type": "record", "name": "struct", "fields": [{"name": "u", "type": ["int", "string"]}]}', map()); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala index 1f943ed7edb08..806d8bf510103 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/bitmapExpressions.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TypeCheckResult import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess} import org.apache.spark.sql.catalyst.expressions.aggregate.ImperativeAggregate +import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke import org.apache.spark.sql.catalyst.trees.UnaryLike import org.apache.spark.sql.catalyst.types.DataTypeUtils @@ -106,6 +107,11 @@ case class BitmapBitPosition(child: Expression) @ExpressionDescription( usage = "_FUNC_(child) - Returns the number of set bits in the child bitmap.", + arguments = """ + Arguments: + * child - The bitmap whose set bits are counted. An expression that evaluates to a binary + bitmap, typically produced by bitmap_construct_agg(). + """, examples = """ Examples: > SELECT _FUNC_(X '1010'); @@ -153,6 +159,222 @@ case class BitmapCount(child: Expression) copy(child = newChild) } +/** Base class for scalar bitmap binary operations. */ +abstract class BitmapBinaryExpression extends BinaryExpression with ExpectsInputTypes { + + override def inputTypes: Seq[AbstractDataType] = Seq(BinaryType, BinaryType) + + override def dataType: DataType = BinaryType + + override def nullIntolerant: Boolean = true + + protected def applyOperation(bitmap1: Array[Byte], bitmap2: Array[Byte]): Array[Byte] + + protected def genCodeOperation( + bitmapUtils: String, bitmap1: String, bitmap2: String): String + + private def checkBitmapLength(bitmap: Array[Byte]): Unit = { + if (bitmap.length > BitmapExpressionUtils.NUM_BYTES) { + throw QueryExecutionErrors.bitmapInputTooLargeError( + bitmap.length, BitmapExpressionUtils.NUM_BYTES) + } + } + + override protected def nullSafeEval(input1: Any, input2: Any): Any = { + val bitmap1 = input1.asInstanceOf[Array[Byte]] + val bitmap2 = input2.asInstanceOf[Array[Byte]] + checkBitmapLength(bitmap1) + checkBitmapLength(bitmap2) + applyOperation(bitmap1, bitmap2) + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val bitmapUtils = classOf[BitmapExpressionUtils].getName + val errors = QueryExecutionErrors.getClass.getName.stripSuffix("$") + nullSafeCodeGen(ctx, ev, (bitmap1, bitmap2) => { + s""" + |if ($bitmap1.length > ${BitmapExpressionUtils.NUM_BYTES}) { + | throw $errors.bitmapInputTooLargeError( + | $bitmap1.length, ${BitmapExpressionUtils.NUM_BYTES}); + |} + |if ($bitmap2.length > ${BitmapExpressionUtils.NUM_BYTES}) { + | throw $errors.bitmapInputTooLargeError( + | $bitmap2.length, ${BitmapExpressionUtils.NUM_BYTES}); + |} + |${ev.value} = ${genCodeOperation(bitmapUtils, bitmap1, bitmap2)}; + |""".stripMargin + }) + } +} + +@ExpressionDescription( + usage = "_FUNC_(left, right) - Returns a bitmap that is the bitwise AND of two input bitmaps.", + arguments = """ + Arguments: + * left - A binary bitmap. + * right - A binary bitmap. + """, + examples = """ + Examples: + > SELECT substring(hex(_FUNC_(X 'F0', X '70')), 0, 2); + 70 + """, + note = """ + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + BITMAP_INPUT_TOO_LARGE. Both inputs must use the same bit-position mapping. If they were + constructed by grouping bitmap_bit_position values by bitmap_bucket_number, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use bitmap_*_agg to combine bitmaps across + rows. + """, + since = "4.4.0", + group = "misc_funcs" +) +case class BitmapAnd(left: Expression, right: Expression) extends BitmapBinaryExpression { + + override def prettyName: String = "bitmap_and" + + override protected def applyOperation( + bitmap1: Array[Byte], bitmap2: Array[Byte]): Array[Byte] = + BitmapExpressionUtils.bitmapAnd(bitmap1, bitmap2) + + override protected def genCodeOperation( + bitmapUtils: String, bitmap1: String, bitmap2: String): String = + s"$bitmapUtils.bitmapAnd($bitmap1, $bitmap2)" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): BitmapAnd = + copy(left = newLeft, right = newRight) +} + +@ExpressionDescription( + usage = "_FUNC_(left, right) - Returns a bitmap that is the bitwise OR of two input bitmaps.", + arguments = """ + Arguments: + * left - A binary bitmap. + * right - A binary bitmap. + """, + examples = """ + Examples: + > SELECT substring(hex(_FUNC_(X '10', X '20')), 0, 2); + 30 + """, + note = """ + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + BITMAP_INPUT_TOO_LARGE. Both inputs must use the same bit-position mapping. If they were + constructed by grouping bitmap_bit_position values by bitmap_bucket_number, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use bitmap_*_agg to combine bitmaps across + rows. + """, + since = "4.4.0", + group = "misc_funcs" +) +case class BitmapOr(left: Expression, right: Expression) extends BitmapBinaryExpression { + + override def prettyName: String = "bitmap_or" + + override protected def applyOperation( + bitmap1: Array[Byte], bitmap2: Array[Byte]): Array[Byte] = + BitmapExpressionUtils.bitmapOr(bitmap1, bitmap2) + + override protected def genCodeOperation( + bitmapUtils: String, bitmap1: String, bitmap2: String): String = + s"$bitmapUtils.bitmapOr($bitmap1, $bitmap2)" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): BitmapOr = + copy(left = newLeft, right = newRight) +} + +@ExpressionDescription( + usage = "_FUNC_(left, right) - Returns a bitmap that is the bitwise AND NOT of two bitmaps.", + arguments = """ + Arguments: + * left - A binary bitmap. + * right - A binary bitmap. + """, + examples = """ + Examples: + > SELECT substring(hex(_FUNC_(X 'F0', X '70')), 0, 2); + 80 + """, + note = """ + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + BITMAP_INPUT_TOO_LARGE. Both inputs must use the same bit-position mapping. If they were + constructed by grouping bitmap_bit_position values by bitmap_bucket_number, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use bitmap_*_agg to combine bitmaps across + rows. + """, + since = "4.4.0", + group = "misc_funcs" +) +case class BitmapAndNot(left: Expression, right: Expression) extends BitmapBinaryExpression { + + override def prettyName: String = "bitmap_andnot" + + override protected def applyOperation( + bitmap1: Array[Byte], bitmap2: Array[Byte]): Array[Byte] = + BitmapExpressionUtils.bitmapAndNot(bitmap1, bitmap2) + + override protected def genCodeOperation( + bitmapUtils: String, bitmap1: String, bitmap2: String): String = + s"$bitmapUtils.bitmapAndNot($bitmap1, $bitmap2)" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): BitmapAndNot = + copy(left = newLeft, right = newRight) +} + +@ExpressionDescription( + usage = "_FUNC_(left, right) - Returns a bitmap that is the bitwise XOR of two input bitmaps.", + arguments = """ + Arguments: + * left - A binary bitmap. + * right - A binary bitmap. + """, + examples = """ + Examples: + > SELECT substring(hex(_FUNC_(X 'F0', X '70')), 0, 2); + 80 + """, + note = """ + Inputs use Spark's Binary bitmap representation, not a RoaringBitmap serialization. Each + input may contain 0 to 4096 bytes; missing bytes are treated as zero. The result is always a + 4096-byte Binary value. NULL input returns NULL, and inputs longer than 4096 bytes raise + BITMAP_INPUT_TOO_LARGE. Both inputs must use the same bit-position mapping. If they were + constructed by grouping bitmap_bit_position values by bitmap_bucket_number, they must + represent the same bucket because the bitmap bytes do not retain bucket metadata. This scalar + function combines two bitmaps from the same row; use bitmap_*_agg to combine bitmaps across + rows. + """, + since = "4.4.0", + group = "misc_funcs" +) +case class BitmapXor(left: Expression, right: Expression) extends BitmapBinaryExpression { + + override def prettyName: String = "bitmap_xor" + + override protected def applyOperation( + bitmap1: Array[Byte], bitmap2: Array[Byte]): Array[Byte] = + BitmapExpressionUtils.bitmapXor(bitmap1, bitmap2) + + override protected def genCodeOperation( + bitmapUtils: String, bitmap1: String, bitmap2: String): String = + s"$bitmapUtils.bitmapXor($bitmap1, $bitmap2)" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): BitmapXor = + copy(left = newLeft, right = newRight) +} + @ExpressionDescription( usage = """ _FUNC_(child) - Returns a bitmap with the positions of the bits set from all the values from @@ -250,6 +472,11 @@ case class BitmapConstructAgg(child: Expression, _FUNC_(child) - Returns a bitmap that is the bitwise OR of all of the bitmaps from the child expression. The input should be bitmaps created from bitmap_construct_agg(). """, + arguments = """ + Arguments: + * child - The expression whose bitmap values are combined with a bitwise OR. + An expression that evaluates to a binary bitmap created from bitmap_construct_agg(). + """, // scalastyle:off line.size.limit examples = """ Examples: @@ -343,6 +570,11 @@ case class BitmapOrAgg(child: Expression, _FUNC_(child) - Returns a bitmap that is the bitwise AND of all of the bitmaps from the child expression. The input should be bitmaps created from bitmap_construct_agg(). """, + arguments = """ + Arguments: + * child - The expression whose bitmap values are combined with a bitwise AND. + An expression that evaluates to a binary bitmap created from bitmap_construct_agg(). + """, // scalastyle:off line.size.limit examples = """ Examples: @@ -430,3 +662,101 @@ case class BitmapAndAgg( buffer.getBinary(mutableAggBufferOffset) } } + +@ExpressionDescription( + usage = """ + _FUNC_(child) - Returns a bitmap that is the bitwise XOR of all of the bitmaps from the child + expression. The input should be bitmaps created from bitmap_construct_agg(). + """, + arguments = """ + Arguments: + * child - The expression whose bitmap values are combined with a bitwise XOR. + An expression that evaluates to a binary bitmap created from bitmap_construct_agg(). + """, + // scalastyle:off line.size.limit + examples = """ + Examples: + > SELECT substring(hex(_FUNC_(col)), 0, 6) FROM VALUES (X'10'), (X'30'), (X'40') AS tab(col); + 600000 + > SELECT substring(hex(_FUNC_(col)), 0, 6) FROM VALUES (X'10'), (X'10') AS tab(col); + 000000 + """, + // scalastyle:on line.size.limit + since = "4.4.0", + group = "agg_funcs") +case class BitmapXorAgg( + child: Expression, + mutableAggBufferOffset: Int = 0, + inputAggBufferOffset: Int = 0) + extends ImperativeAggregate + with UnaryLike[Expression] { + + def this(child: Expression) = { + this(child = child, mutableAggBufferOffset = 0, inputAggBufferOffset = 0) + } + + override def checkInputDataTypes(): TypeCheckResult = { + if (child.dataType != BinaryType) { + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> toSQLType(BinaryType), + "inputSql" -> toSQLExpr(child), + "inputType" -> toSQLType(child.dataType))) + } else { + TypeCheckSuccess + } + } + + override def dataType: DataType = BinaryType + + override def prettyName: String = "bitmap_xor_agg" + + override protected def withNewChildInternal(newChild: Expression): BitmapXorAgg = + copy(child = newChild) + + override def withNewMutableAggBufferOffset( + newMutableAggBufferOffset: Int): ImperativeAggregate = + copy(mutableAggBufferOffset = newMutableAggBufferOffset) + + override def withNewInputAggBufferOffset(newInputAggBufferOffset: Int): ImperativeAggregate = + copy(inputAggBufferOffset = newInputAggBufferOffset) + + override def nullable: Boolean = false + + override def aggBufferSchema: StructType = DataTypeUtils.fromAttributes(aggBufferAttributes) + + // The aggregation buffer is a fixed size binary. + private val bitmapAttr = AttributeReference("bitmap", BinaryType, false)() + + override def aggBufferAttributes: Seq[AttributeReference] = bitmapAttr :: Nil + + override def defaultResult: Option[Literal] = + Option(Literal(Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(0))) + + override val inputAggBufferAttributes: Seq[AttributeReference] = + aggBufferAttributes.map(_.newInstance()) + + override def initialize(buffer: InternalRow): Unit = { + buffer.update(mutableAggBufferOffset, Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(0)) + } + + override def update(buffer: InternalRow, input: InternalRow): Unit = { + val input_bitmap = child.eval(input).asInstanceOf[Array[Byte]] + if (input_bitmap != null) { + val bitmap = buffer.getBinary(mutableAggBufferOffset) + BitmapExpressionUtils.bitmapXorMerge(bitmap, input_bitmap) + } + } + + override def merge(buffer1: InternalRow, buffer2: InternalRow): Unit = { + val bitmap1 = buffer1.getBinary(mutableAggBufferOffset) + val bitmap2 = buffer2.getBinary(inputAggBufferOffset) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + } + + override def eval(buffer: InternalRow): Any = { + buffer.getBinary(mutableAggBufferOffset) + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala new file mode 100644 index 0000000000000..2a64f5d007133 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompiler.scala @@ -0,0 +1,1666 @@ +/* + * 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.catalyst.expressions.codegen + +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, InputStream, IOException, StringWriter} +import java.lang.reflect.{Member, Method, Modifier} +import java.net.{JarURLConnection, URI, URL} +import java.util.Locale +import java.util.concurrent.{Callable, ExecutionException, ExecutorService} +import javax.tools.{Diagnostic, DiagnosticCollector, FileObject, ForwardingJavaFileManager, JavaCompiler, JavaFileManager, JavaFileObject, SimpleJavaFileObject, StandardJavaFileManager, StandardLocation, ToolProvider} + +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import com.google.common.cache.{Cache, CacheBuilder} +import com.google.common.util.concurrent.Uninterruptibles +import org.codehaus.commons.compiler.{CompileException, InternalCompilerException} +import org.codehaus.janino.ClassBodyEvaluator +import org.codehaus.janino.util.ClassFile +import org.codehaus.janino.util.ClassFile.CodeAttribute + +import org.apache.spark.{JobArtifactSet, SparkEnv, TaskContext, TaskKilledException} +import org.apache.spark.executor.{ExecutorClassLoader, InputMetrics} +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.metrics.source.CodegenMetrics +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, UnsafeArrayData, UnsafeMapData, UnsafeRow} +import org.apache.spark.sql.catalyst.util.{ArrayData, CollationAwareUTF8String, CollationFactory, CollationSupport, MapData} +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.Decimal +import org.apache.spark.unsafe.Platform +import org.apache.spark.unsafe.types.{BinaryView, CalendarInterval, TimestampNanosVal, UTF8String, VariantVal} +import org.apache.spark.util.{ParentClassLoader, ThreadUtils, Utils} + +/** + * Backend used to compile generator-produced Java source into a [[GeneratedClass]]. + * + * Two implementations are provided: + * - [[JaninoCodeCompiler]]: the default, uses Janino's `ClassBodyEvaluator`. Very fast. + * - [[JdkCodeCompiler]]: uses `javax.tools.JavaCompiler` from the JDK. Slower (~5x for + * large generated units, 30-300x for small ones), but maintained on the JDK + * release cadence and not subject to Janino's unmaintained-upstream risk. + * + * The backend is selected at compile time via [[SQLConf.CODEGEN_COMPILER]]. + */ +trait CodeCompiler { + /** Backend name as used in `spark.sql.codegen.compiler`. */ + def name: String + + /** + * Compile a generator-produced class body into an instance of the + * [[GeneratedClass]] subclass it defines. + * + * @return the instantiated generated class along with bytecode statistics. + */ + def compile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) +} + +object CodeCompiler extends Logging { + + // Emit log messages under CodeGenerator's logger name: operators and tests + // (SPARK-25113 / SPARK-51527) subscribe to that exact logger for codegen + // compilation events, and the backends are implementation details of + // `CodeGenerator.compile`, so their logs belong under its name. + override protected def logName: String = classOf[CodeGenerator[_, _]].getName + + val JANINO: String = "janino" + val JDK: String = "jdk" + + /** + * Fully-qualified imports made available to generated code by both backends. + * + * For Janino these are passed to `ClassBodyEvaluator.setDefaultImports`. + * For the JDK backend they are rendered into `import` statements inside the + * synthesized compilation unit. + * + * This is the single shared list - anything added here automatically applies to + * both backends. It intentionally excludes + * `org.apache.spark.sql.catalyst.expressions.codegen.GeneratedClass` to avoid + * a name collision with the generated subclass `GeneratedClass`; the extends + * clause uses the fully-qualified name instead. + * + * When adding an entry, keep its SIMPLE name distinct from any `import` line a + * generator emits at the top of a class body (currently only GenerateColumnAccessor + * does this): javac rejects two single-type imports sharing a simple name + * (JLS 7.5.1) while Janino resolves them leniently, so a collision would fail only + * under the JDK backend. + */ + val DefaultImports: Seq[String] = Seq( + classOf[Platform].getName, + classOf[InternalRow].getName, + classOf[UnsafeRow].getName, + classOf[BinaryView].getName, + classOf[UTF8String].getName, + classOf[Decimal].getName, + classOf[CalendarInterval].getName, + classOf[TimestampNanosVal].getName, + classOf[VariantVal].getName, + classOf[ArrayData].getName, + classOf[UnsafeArrayData].getName, + classOf[MapData].getName, + classOf[UnsafeMapData].getName, + classOf[Expression].getName, + classOf[TaskContext].getName, + classOf[TaskKilledException].getName, + classOf[InputMetrics].getName, + classOf[CollationAwareUTF8String].getName, + classOf[CollationFactory].getName, + classOf[CollationSupport].getName, + QueryExecutionErrors.getClass.getName.stripSuffix("$") + ) + + /** + * FQN of the generated class. Must NOT be under the `codegen` package or Janino + * fails with `java.lang.InstantiationException`. The same name is used for both + * backends so generated source, logs, and diagnostics name the same class + * whichever backend compiles it. (Compiled results are NOT shared across + * backends: the compile cache key includes the backend.) + */ + val GeneratedClassName: String = + "org.apache.spark.sql.catalyst.expressions.GeneratedClass" + + def active(): CodeCompiler = active(null) + + /** + * Resolve the active backend for the given generated unit. + * + * The configured backend ([[SQLConf.CODEGEN_COMPILER]]) governs ordinary codegen. The + * exception is codegen the JDK compiler is fundamentally *incapable* of compiling - not + * merely slower at compiling - which is always routed to Janino regardless of the + * configured backend. This is deterministic routing decided up front from the execution + * context and the generated source; it is never a fallback after a failed compile. Three + * such cases, all involving classes the JDK compiler cannot name that Janino's lenient + * loader/lexer accepts: + * + * - REPL / interactive sessions (spark-shell `$line*` wrappers, Spark Connect / + * Ammonite session artifacts): reachable only through a runtime class loader and + * carrying self-inconsistent reflection metadata the JDK compiler cannot resolve. + * This arm is context-wide by design: ALL codegen in such a session routes to + * Janino, whether or not the unit references a REPL class, because the reference + * cannot be told from the source text up front. See [[isReplContext]]. + * - A reference to a class nested in a Scala `package object` (binary name + * `a.b.package$Inner`): `package` is a Java reserved word that cannot be spelled as + * an identifier in any form - Java has no backtick/escape, unlike Scala - so javac + * can neither parse `a.b.package.Inner` nor resolve the flat `a.b.package$Inner`. + * See [[requiresJaninoSource]]. + * - A reference to a class the Java language forbids naming, i.e. an anonymous or local + * class (`a.b.Outer$1`, `a.b.Outer$$anon$1`) or a class nested inside one + * (`a.b.Outer$1$Inner`), that is shown to be unnarrowable. The JDK backend rewrites + * such a reference to the nearest nameable supertype, which works only while that + * supertype is itself referenceable and offers every member the generated code could + * access; this arm covers the classes for which reflection positively reports it does + * not. A class whose verdict reflection cannot determine is not routed: a reference to + * it in code position keeps its binary name, which javac rejects for these shapes. + * See [[JdkCodeCompiler.referencesUnnarrowableClass]]. + */ + def active(code: CodeAndComment): CodeCompiler = { + // Resolve the configured backend first: when it is already Janino - by configuration or + // because javac is absent - none of the overrides below can change the outcome, and + // their source scans would be pure overhead. + val configured = forBackend(SQLConf.get.codegenCompiler) + if (configured eq JaninoCodeCompiler) { + configured + } else if (isReplContext) { + logReplRoutingOnce() + JaninoCodeCompiler + } else if (requiresJaninoSource(code)) { + logPackageObjectRoutingOnce() + JaninoCodeCompiler + } else if (code != null && JdkCodeCompiler.referencesUnnarrowableClass(code.body)) { + logUnnarrowableClassRoutingOnce() + JaninoCodeCompiler + } else { + configured + } + } + + // One-time visibility for the deterministic routing above: an operator who set + // `jdk` should be able to tell from the logs why Janino still shows up. + private val replRoutingLogged = new java.util.concurrent.atomic.AtomicBoolean(false) + private def logReplRoutingOnce(): Unit = { + if (replRoutingLogged.compareAndSet(false, true)) { + logInfo(log"REPL / interactive session context detected; codegen is routed to " + + log"Janino although ${MDC(LogKeys.CONFIG, SQLConf.CODEGEN_COMPILER.key)} " + + log"requests another backend (the JDK compiler cannot resolve REPL-defined " + + log"classes). This notice is logged once per JVM.") + } + } + private val packageObjectRoutingLogged = new java.util.concurrent.atomic.AtomicBoolean(false) + private def logPackageObjectRoutingOnce(): Unit = { + if (packageObjectRoutingLogged.compareAndSet(false, true)) { + logInfo(log"Generated code references a Scala package-object class; that unit is " + + log"routed to Janino although ${MDC(LogKeys.CONFIG, SQLConf.CODEGEN_COMPILER.key)} " + + log"requests another backend (`package` is a Java reserved word the JDK compiler " + + log"cannot name). This notice is logged once per JVM.") + } + } + private val unnarrowableClassRoutingLogged = new java.util.concurrent.atomic.AtomicBoolean(false) + private def logUnnarrowableClassRoutingOnce(): Unit = { + if (unnarrowableClassRoutingLogged.compareAndSet(false, true)) { + logInfo(log"Generated code references a class Java cannot name (anonymous, local, or " + + log"nested in one) that cannot be narrowed to a nameable supertype; that unit is " + + log"routed to Janino although " + + log"${MDC(LogKeys.CONFIG, SQLConf.CODEGEN_COMPILER.key)} requests another backend. " + + log"This notice is logged once per JVM.") + } + } + + // A `package` segment in a qualified/binary class name - a Scala `package object`'s nested + // class such as `a.b.package$Inner`. `package` is a Java reserved word the JDK compiler can + // name in no form (parse error as `package.Inner`; unresolvable as the flat `package$Inner`), + // whereas Janino's lexer scans `package$Inner` as one identifier. + // + // `package` is the only keyword scanned for, by design. It is the only Java keyword the + // Scala compiler ever produces in a generated name (from `package object`); a class named + // after any other keyword (`class int`) requires pathological user code. It is also the only + // keyword that is *safe* to scan for: the rest (`int`, `new`, `this`, `return`, `switch`, ...) + // occur as legitimate tokens throughout the generated Java, so matching them would route + // almost all codegen to Janino, whereas a `package` token never appears in a generated class + // body except as such a class reference. (A fully general check would inspect the resolved + // class names rather than the source text, but that information is only available during + // compilation, i.e. after the backend is already chosen.) The lookbehind keeps a legal + // identifier like `mypackage$Inner` from matching; a false positive (e.g. text inside a string + // literal) is harmless - it only picks Janino, a superset of what javac accepts. + private val UnnameablePackageObjectClass = """(?<![\w$])package[.$]""".r + private def requiresJaninoSource(code: CodeAndComment): Boolean = { + // This runs on every compile() call (the result is part of the cache key), so gate + // the regex scan behind an intrinsified substring search: generated bodies almost + // never contain the literal `package` at all, and the regex runs only when they do. + code != null && code.body.contains("package") && + UnnameablePackageObjectClass.findFirstIn(code.body).isDefined + } + + /** + * True when codegen is running in a REPL / interactive context, detected via the three + * mechanisms Spark uses to ship such classes: + * - the active job/session carries a REPL or artifact class-dir URI + * ([[JobArtifactSet.getCurrentJobArtifactState]]'s `replClassDirUri`). This is the + * canonical signal: Spark Connect sets it per session and spark-shell falls back to + * it from `spark.repl.class.uri`. It is a thread-local set around both driver-side + * and executor-side work, so it catches driver-side codegen where no + * `ExecutorClassLoader` is in the loader chain (e.g. a Connect UDF over a local + * relation referencing an Ammonite `$sess` class); or + * - `spark.repl.class.uri` set in the active conf (spark-shell sets this globally); or + * - an [[org.apache.spark.executor.ExecutorClassLoader]] somewhere in the active + * class loader chain (created on executors when a session has such a class URI). + * + * The default (non-REPL) job state has no `replClassDirUri`, so ordinary codegen is not + * affected. Any lookup failure conservatively reports `false`, which preserves the + * configured backend. + */ + private def isReplContext: Boolean = { + def hasArtifactReplUri = + try JobArtifactSet.getCurrentJobArtifactState.exists(_.replClassDirUri.isDefined) + catch { case NonFatal(_) => false } + def confHasReplUri = + try Option(SparkEnv.get).exists(_.conf.contains("spark.repl.class.uri")) + catch { case NonFatal(_) => false } + def eclInChain = + try { + var loader = Utils.getContextOrSparkClassLoader + var found = false + while (loader != null && !found) { + if (loader.isInstanceOf[ExecutorClassLoader]) found = true + loader = loader.getParent + } + found + } catch { + case NonFatal(_) => false + } + hasArtifactReplUri || confHasReplUri || eclInChain + } + + /** + * Get the backend by name. SQLConf already validates the value via `checkValues` + * at config-set time, so unknown names should not reach here in normal use; + * tests may call this directly. When `jdk` is requested but the JDK compiler is + * not present at runtime (a JRE-only image), this logs a warning once and falls + * back to Janino so the query does not fail. + */ + private[codegen] def forBackend(requested: String): CodeCompiler = { + requested.toLowerCase(Locale.ROOT) match { + case JANINO => JaninoCodeCompiler + case JDK if JdkCodeCompiler.isAvailable => JdkCodeCompiler + case JDK => + logJdkUnavailableOnce() + JaninoCodeCompiler + case other => + throw new IllegalArgumentException( + s"Unknown ${SQLConf.CODEGEN_COMPILER.key} backend: $other " + + s"(supported: ${Seq(JANINO, JDK).mkString(", ")})") + } + } + + private val jdkUnavailableWarned = new java.util.concurrent.atomic.AtomicBoolean(false) + private def logJdkUnavailableOnce(): Unit = { + if (jdkUnavailableWarned.compareAndSet(false, true)) { + logWarning(log"${MDC(LogKeys.CONFIG, SQLConf.CODEGEN_COMPILER.key)}=jdk requested " + + log"but javax.tools.JavaCompiler is not available on this runtime " + + log"(JRE-only image?). Falling back to Janino for this JVM.") + } + } + + /** + * Compute bytecode statistics for a set of compiled classes. Both backends + * produce the same map shape (className -> classfile bytes), so the analysis is + * shared. This is the only piece of code that depends on Janino's + * `ClassFile` parser (from the `janino` artifact); it can be swapped for ASM + * later without touching either backend. + */ + private[codegen] def computeByteCodeStats( + classBytecodes: Iterable[(String, Array[Byte])]): ByteCodeStats = { + val perClass = classBytecodes.map { case (_, classBytes) => + val classCodeSize = classBytes.length + CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classCodeSize) + try { + val cf = new ClassFile(new ByteArrayInputStream(classBytes)) + val constPoolSize = cf.getConstantPoolSize + val methodCodeSizes = cf.methodInfos.asScala.flatMap { method => + method.getAttributes.collect { case attr: CodeAttribute => + val byteCodeSize = attr.code.length + CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize) + if (byteCodeSize > CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT) { + logInfo(log"Generated method too long to be JIT compiled: " + + log"${MDC(LogKeys.CLASS_NAME, cf.getThisClassName)}." + + log"${MDC(LogKeys.METHOD_NAME, method.getName)} is " + + log"${MDC(LogKeys.BYTECODE_SIZE, byteCodeSize)} bytes") + } + byteCodeSize + } + } + // Use `maxOption` to handle classes with no methods (e.g., a synthetic + // module-info-style class). The original Janino-only code would have raised + // UnsupportedOperationException there; we degrade gracefully to -1 instead. + (methodCodeSizes.maxOption.getOrElse(-1), constPoolSize) + } catch { + case NonFatal(e) => + logWarning("Error calculating stats of compiled class.", e) + (-1, -1) + } + } + + val (maxMethodSizes, constPoolSize) = perClass.unzip + ByteCodeStats( + maxMethodCodeSize = maxMethodSizes.maxOption.getOrElse(-1), + maxConstPoolSize = constPoolSize.maxOption.getOrElse(-1), + // Minus 2 for `GeneratedClass` and an outer-most generated class. + // Both backends wrap the class body in a single outer declaration, so the + // emitted class count has the same shape (1 outer wrapper + K user-declared + // classes) and the offset yields the same value under either backend. + // `max(0, ...)` keeps an unexpected emit shape from going negative. + numInnerClasses = math.max(0, classBytecodes.size - 2)) + } + + /** + * Log the generated source on a compilation failure. Behaviour matches the + * original [[CodeGenerator]] implementation. `maxLines` (the session's + * `loggingMaxLinesForCodegen`) is captured by the CALLER: the JDK backend invokes + * this from its compile worker thread, where `SQLConf.get` would silently return + * the default conf instead of the calling session's. + */ + private[codegen] def logGeneratedCodeOnFailure(code: CodeAndComment, maxLines: Int): Unit = { + val formatted = s"\n${CodeFormatter.format(code, maxLines)}" + if (Utils.isTesting) { + logError(formatted) + } else { + logInfo(formatted) + } + } +} + +/** + * The default backend using Janino's `ClassBodyEvaluator`. + * + * This lifts the previous body of `CodeGenerator.doCompile`, with the only + * changes being: imports moved to `CodeCompiler.DefaultImports`, stats + * computation moved to `CodeCompiler.computeByteCodeStats` (which degrades to + * `-1` for a class with no methods instead of throwing; see its comment). + * Behaviour is otherwise preserved, including the [[ParentClassLoader]] wrapping + * (workaround for the Janino `findIClass` behaviour described in SPARK-15622 / + * SPARK-11636). + */ +object JaninoCodeCompiler extends CodeCompiler with Logging { + + override val name: String = CodeCompiler.JANINO + + // Route source-code/debug log emissions under CodeGenerator's logger name. + override protected def logName: String = classOf[CodeGenerator[_, _]].getName + + override def compile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = { + val evaluator = new ClassBodyEvaluator() + + // See SPARK-15622 / SPARK-11636 for why this wrapping is required. + val parentClassLoader = new ParentClassLoader(Utils.getContextOrSparkClassLoader) + evaluator.setParentClassLoader(parentClassLoader) + evaluator.setClassName(CodeCompiler.GeneratedClassName) + evaluator.setDefaultImports(CodeCompiler.DefaultImports: _*) + evaluator.setExtendedClass(classOf[GeneratedClass]) + + logBasedOnLevel(SQLConf.get.codegenLogLevel) { + // Only add extra debugging info to byte code when we are going to print the source code. + evaluator.setDebuggingInformation(true, true, false) + log"\n${MDC(LogKeys.CODE, CodeFormatter.format(code))}" + } + + val codeStats = + try { + evaluator.cook("generated.java", code.body) + CodeCompiler.computeByteCodeStats(evaluator.getBytecodes.asScala) + } catch { + case e: InternalCompilerException => + logError("Failed to compile the generated Java code.", e) + CodeCompiler.logGeneratedCodeOnFailure(code, SQLConf.get.loggingMaxLinesForCodegen) + throw QueryExecutionErrors.internalCompilerError(e) + case e: CompileException => + logError("Failed to compile the generated Java code.", e) + CodeCompiler.logGeneratedCodeOnFailure(code, SQLConf.get.loggingMaxLinesForCodegen) + throw QueryExecutionErrors.compilerError(e) + } + + (evaluator.getClazz().getConstructor().newInstance().asInstanceOf[GeneratedClass], codeStats) + } +} + +/** + * Alternative backend using the JDK's `javax.tools.JavaCompiler`. + * + * Wraps the generator-produced class body in a synthesized compilation unit + * (package + imports + `public class GeneratedClass extends ...`) before + * handing it to the compiler. Compiled classes are captured in memory and + * loaded through a private [[ClassLoader]] that mirrors the Janino backend's + * [[ParentClassLoader]] wrapping (SPARK-15622 / SPARK-11636) so behaviour on + * containerised deployments stays consistent. + * + * Class resolution: referenced classes are resolved through the task's context + * [[ClassLoader]] (see [[ClassLoaderFileManager]]) rather than a file-based + * `-classpath`, mirroring how Janino resolves them. This lets the JDK compiler see + * classes that exist only on a runtime loader - REPL-generated, Spark Connect + * session artifacts - and avoids handing javac a giant classpath to index. + * + * Threading: the actual javac invocation runs on a dedicated single-threaded + * executor (see `compileExecutor`). This is required for correctness on Spark + * task threads (jar reads through interruptible NIO channels vs. task + * interruption), and it also confines the shared, not-thread-safe + * [[StandardJavaFileManager]] (used for platform-class lookups) to one + * thread, so no extra locking is needed. Caller threads only build the source and + * capture the context classloader, then await the result. + * + * Resource lifecycle: the shared [[StandardJavaFileManager]] is intentionally + * never closed. Like the compile executor and the per-jar package index, it is + * JVM-lifetime state rather than a per-compile resource, so this is not a leak. + * + * Performance is roughly 5x slower than Janino for large generated units and + * 30-300x slower for small ones. + * The benefit is decoupling Spark from Janino's release cadence. + */ +object JdkCodeCompiler extends CodeCompiler with Logging { + + override val name: String = CodeCompiler.JDK + + // Route source-code/debug log emissions under CodeGenerator's logger name. + override protected def logName: String = classOf[CodeGenerator[_, _]].getName + + /** True if `javax.tools.JavaCompiler` is available on this runtime. */ + lazy val isAvailable: Boolean = ToolProvider.getSystemJavaCompiler != null + + private lazy val compiler: JavaCompiler = { + val c = ToolProvider.getSystemJavaCompiler + require(c != null, + "javax.tools.JavaCompiler is not available; check isAvailable before use") + c + } + + /** + * Shared file manager. `StandardJavaFileManager` is not thread-safe; reuse is safe + * here because every javac invocation runs on the single-threaded [[compileExecutor]]. + */ + private lazy val sharedFileManager: StandardJavaFileManager = + compiler.getStandardFileManager(null, null, null) + + /** + * Compiler options applied to every compilation. + * + * There is deliberately no `--release`/`-source`/`-target`: the compiled class is + * loaded only into the same JVM that produced it (the cache is in-memory; generated + * code travels between JVMs as SOURCE), so the emitted class-file version is always + * consistent with the running runtime and there is no cross-compilation target to + * pin. Pinning `--release` would only reroute the delegate file manager's + * platform-class lookups through `ct.sym`, adding overhead without benefit. + * + * Note there is no `-classpath`: the [[ClassLoaderFileManager]] resolves the + * `CLASS_PATH` location through the compile's parent [[ClassLoader]] (the task's + * context classloader) rather than a file-based classpath. This mirrors how + * Janino resolves referenced classes, so the JDK backend sees exactly what + * Janino would - including classes that exist only on a runtime classloader + * (REPL-generated, Spark Connect session artifacts) and never on + * `java.class.path`. It also avoids handing javac a giant `-classpath` to index, + * which both bloats compiler memory and is brittle to harvest correctly across + * driver / executor / Connect deployments. + */ + private val compileOptions: java.util.List[String] = Seq( + "-proc:none", // skip annotation processing + "-g:none", // skip debug info + "-nowarn", // suppress warnings + "-implicit:none", // do not compile referenced source files + "-Xlint:none" // disable lints + ).asJava + + /** Source-position package name and simple class name derived once. */ + private val packageName: String = + CodeCompiler.GeneratedClassName.substring(0, CodeCompiler.GeneratedClassName.lastIndexOf('.')) + private val simpleName: String = + CodeCompiler.GeneratedClassName.substring(CodeCompiler.GeneratedClassName.lastIndexOf('.') + 1) + + /** Rendered import block, computed once. */ + private val importBlock: String = + CodeCompiler.DefaultImports.map(i => s"import $i;").mkString("\n") + + /** FQN of the abstract base used in the extends clause (avoids import collision). */ + private val extendsFqn: String = classOf[GeneratedClass].getName + + /** + * Wrap a generator-produced class body in a full compilation unit. Class name + * matches Janino's output so logs and diagnostics name the same class whichever + * backend compiled it. `classLoader` resolves the candidate inner-class + * references for the `$`-rewrite (see [[rewriteInnerClassRefs]]); it must be the + * same loader the compile resolves classes through. + * + * Hoisted imports are deliberately left un-rewritten: an `import` requires a canonical + * name, so a binary inner-class name there is a javac error whether or not it is + * narrowed. + */ + private[codegen] def wrapAsCompilationUnit(body: String, classLoader: ClassLoader): String = { + val (extraImports, cleanedBody) = extractLeadingImports(body) + val javacBody = stripFunction1ApplyBridges(cleanedBody) + // Built by plain concatenation, NOT a stripMargin template: stripMargin + // post-processes the final interpolated string, so a generated-body line whose + // first non-blank character is `|` (e.g. a line-wrapped `||` condition) would + // lose that character. No current generator emits such a line, but the unit + // must not depend on that. + s"package $packageName;\n" + + s"$importBlock\n" + + s"$extraImports\n" + + s"public class $simpleName extends $extendsFqn {\n" + + s"${rewriteInnerClassRefs(javacBody, classLoader)}\n" + + "}\n" + } + + // The explicit `scala.Function1` `apply(Object)` bridge that projection codegen emits + // for the Janino backend (see `CodeGenerator.function1ApplyBridge`). javac synthesizes + // this bridge itself for the typed `apply(InternalRow)` override and rejects an explicit + // duplicate with a "name clash" error, so it must be removed before compiling with the + // JDK backend. This pattern matches exactly the shape `function1ApplyBridge` emits + // (whitespace tolerant, `\1` ties the cast operand to the parameter); keep them in sync. + private val Function1ApplyBridgePattern = + ("""(?s)public\s+java\.lang\.Object\s+apply\(\s*java\.lang\.Object\s+(\w+)\s*\)\s*""" + + """\{\s*return\s+apply\(\(\s*InternalRow\s*\)\s*\1\s*\)\s*;\s*\}""").r + + /** Remove the Janino-only Function1 `apply(Object)` bridges so javac does not clash. */ + private[codegen] def stripFunction1ApplyBridges(body: String): String = + if (body.contains("apply(java.lang.Object")) { + Function1ApplyBridgePattern.replaceAllIn(body, "") + } else { + body + } + + /** + * Some generators (e.g. GenerateColumnAccessor) emit `import` statements at + * the top of the class body. Janino's ClassBodyEvaluator treats those as + * imports for the synthesized class, but the JDK compiler rejects imports + * inside a class declaration ("illegal start of type"). Extract any leading + * `import` lines from the body so they can be hoisted into the compilation + * unit header. + */ + private[codegen] def extractLeadingImports(body: String): (String, String) = { + // Fast path: no leading `import` line (every generator but GenerateColumnAccessor). + // Skips the full line-split allocation. + var p = 0 + while (p < body.length && Character.isWhitespace(body.charAt(p))) p += 1 + if (!body.startsWith("import ", p)) return ("", body) + // The `-1` limit keeps trailing empty lines so the reconstruction is faithful. + val lines = body.split("\n", -1) + val imports = new StringBuilder + var i = 0 + var scanning = true + while (scanning && i < lines.length) { + val trimmed = lines(i).trim + if (trimmed.startsWith("import ")) { + imports.append(lines(i)).append('\n') + i += 1 + } else if (trimmed.isEmpty) { + i += 1 + } else { + scanning = false + } + } + if (i == 0) { + ("", body) + } else { + (imports.toString, lines.drop(i).mkString("\n")) + } + } + + /** + * Rewrite JVM-binary inner-class references into the Java-source form the JDK + * compiler accepts. Spark generators emit class names via `Class#getName` in + * many places; for nested classes that yields the binary form (`Outer$Inner`), + * which Janino accepts as a source-level identifier but the JDK compiler does + * not. + * + * The correct source form depends on HOW the class is nested, and that cannot + * be told from the text alone: + * - a regular nested class `Outer$Inner` must be written `Outer.Inner`; + * - a class nested inside a Scala `object` has a binary name whose `$` + * separators include the module suffix (e.g. `Model$SaveLoad$Leaf` where + * `SaveLoad` is an object), and the JDK compiler resolves it ONLY via the + * raw binary name - the dotted canonical form `Model.SaveLoad$.Leaf` makes + * javac reconstruct a non-existent `Model$SaveLoad$$Leaf`. + * Textually `Model$SaveLoad$Leaf` (object-nested) is indistinguishable from + * `A$B$C` (three regular classes) yet they need opposite treatment, so the + * decision is made by resolving each candidate against the compile classpath + * and consulting reflection: `getCanonicalName` is the right source form when + * it is free of `$`, otherwise the binary name is. + * + * For each maximal qualified-name token that contains `$`, we find the longest + * dot-delimited prefix that loads as a class and replace it with that + * reflection-derived name, leaving any trailing member access untouched + * (so `Foo$.MODULE$.apply` and `List$.MODULE$.newBuilder()` resolve correctly). + * Tokens whose prefixes do not resolve - notably references to the + * not-yet-compiled inner classes of the generated unit itself - fall back to a + * conservative regex that dots `$`-before-uppercase, matching the historical + * behaviour for those. + * + * The rewrite is applied only to actual code spans: string literals, char + * literals, and `//` / block comments are copied verbatim so that a `$Upper` + * sequence inside generated string data (e.g. a column name or error message) + * is never corrupted. + */ + private[codegen] def rewriteInnerClassRefs(body: String, classLoader: ClassLoader): String = { + val out = new java.lang.StringBuilder(body.length + 16) + val code = new java.lang.StringBuilder() + val n = body.length + // A token's rewritten form is stable for a given classloader; memoize within + // this call so repeated type references resolve at most once. + val memo = mutable.HashMap.empty[String, String] + + def flushCode(): Unit = { + if (code.length > 0) { + out.append(rewriteCodeSpan(code.toString, classLoader, memo)) + code.setLength(0) + } + } + + // Copy a quoted literal (string or char) verbatim, honoring backslash escapes. + def copyQuoted(start: Int, quote: Char): Int = { + out.append(quote) + var j = start + 1 + var closed = false + while (j < n && !closed) { + val ch = body.charAt(j) + if (ch == '\\' && j + 1 < n) { + out.append(ch).append(body.charAt(j + 1)) + j += 2 + } else { + out.append(ch) + j += 1 + if (ch == quote) closed = true + } + } + j + } + + var i = 0 + while (i < n) { + val c = body.charAt(i) + if (c == '"' || c == '\'') { + flushCode() + i = copyQuoted(i, c) + } else if (c == '/' && i + 1 < n && body.charAt(i + 1) == '/') { + flushCode() + while (i < n && body.charAt(i) != '\n') { out.append(body.charAt(i)); i += 1 } + } else if (c == '/' && i + 1 < n && body.charAt(i + 1) == '*') { + flushCode() + out.append("/*") + i += 2 + while (i < n && !(body.charAt(i) == '*' && i + 1 < n && body.charAt(i + 1) == '/')) { + out.append(body.charAt(i)); i += 1 + } + // The scan exits either at the `*/` terminator (then i + 1 < n holds by the + // loop condition) or at end-of-body for an unterminated comment, whose + // characters the loop already copied verbatim. + if (i + 1 < n) { out.append("*/"); i += 2 } + } else { + code.append(c) + i += 1 + } + } + flushCode() + out.toString + } + + /** + * Rewrite the qualified-name tokens of a code span (no literals or comments). + * Runs of `[A-Za-z0-9_$.]` are treated as candidate qualified names; only + * those containing `$` are resolved (others cannot be binary inner-class + * references), and everything else is copied verbatim so whitespace and + * punctuation are preserved exactly. + */ + private def rewriteCodeSpan( + span: String, + classLoader: ClassLoader, + memo: mutable.Map[String, String]): String = { + val sb = new java.lang.StringBuilder(span.length + 16) + val n = span.length + var i = 0 + while (i < n) { + val c = span.charAt(i) + if (isNameStart(c)) { + val start = i + i += 1 + while (i < n && isNamePart(span.charAt(i))) i += 1 + val token = span.substring(start, i) + if (token.indexOf('$') < 0) { + sb.append(token) + } else { + sb.append(memo.getOrElseUpdate(token, rewriteQualifiedName(token, classLoader))) + } + } else { + sb.append(c) + i += 1 + } + } + sb.toString + } + + private def isNameStart(c: Char): Boolean = + (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || c == '_' || c == '$' + + private def isNamePart(c: Char): Boolean = + isNameStart(c) || (c >= '0' && c <= '9') || c == '.' + + /** + * Resolve a `$`-containing qualified token to a JDK-compiler-acceptable form by + * finding the longest dot-delimited prefix that loads as a class and replacing + * it with its reflection-derived source name, keeping any trailing member + * access. Falls back to the conservative `$`-before-uppercase regex when no + * prefix resolves (e.g. inner classes of the not-yet-compiled generated unit). + */ + private def rewriteQualifiedName(token: String, classLoader: ClassLoader): String = { + loadLongestPrefix(token, classLoader) match { + case Some((cls, rest)) => + val sourceName = narrowedSourceName(cls) + val resolved = if (rest.isEmpty) sourceName else s"$sourceName.$rest" + // `split('.')` drops a trailing empty segment, so a token that ends in `.` + // (member access wrapped onto the next line) must get its dot restored. + if (token.endsWith(".")) resolved + "." else resolved + case None => + InnerClassRefPattern.replaceAllIn(token, ".") + } + } + + /** + * The source name to emit for `cls`: the name of the nearest class Java can name, in the + * form javac accepts. + * + * Narrowing happens only on a positive [[Narrowable]] verdict; anything else keeps the + * binary name. Consulting the verdict is the load-bearing part, because the reflective + * failures are not symmetric: the supertype climb can succeed while member enumeration + * throws, so climbing here without it would narrow a class whose members were never + * checked, to a public supertype javac happily accepts. Taking the target from the verdict + * rather than re-deriving it costs one climb instead of two. Keeping the binary name of a + * class Java cannot name yields a name javac rejects, so the unit fails to compile rather + * than compiling into a wrong answer. + */ + private def narrowedSourceName(cls: Class[_]): String = narrowingVerdict(cls) match { + case Narrowable(target) => sourceNameOf(target) + case Unnarrowable => cls.getName + case Unknown(cause) => + // Logged from here, not from the verdict: this is the one call site that knows the + // token really is a type reference in code position, so the message's claim that the + // unit will fail to compile actually holds. The routing scan reads the raw body and + // can evaluate a token inside a literal, where the same verdict costs nothing. + logUnevaluableClassOnce(cls, cause) + cls.getName + } + + /** + * Find the longest dot-delimited prefix of `token` that loads as a class, and return it + * with the remaining member access. Only a prefix containing `$` is tried, since only + * such a prefix can be a binary inner-class name. Returns None when nothing loads (e.g. + * an inner class of the not-yet-compiled generated unit). + */ + private def loadLongestPrefix( + token: String, + classLoader: ClassLoader): Option[(Class[_], String)] = { + val parts = token.split('.') + var k = parts.length + while (k >= 1) { + val prefix = parts.iterator.take(k).mkString(".") + if (prefix.indexOf('$') >= 0) { + loadWithoutInit(prefix, classLoader) match { + case Some(cls) => return Some((cls, parts.iterator.drop(k).mkString("."))) + case None => // try a shorter prefix + } + } + k -= 1 + } + None + } + + /** Load `binaryName` without initializing it; None when it is not a loadable class. */ + private def loadWithoutInit(binaryName: String, classLoader: ClassLoader): Option[Class[_]] = { + try { + // scalastyle:off classforname + // Load with the exact loader passed in (the task's context loader), not the Spark + // class loader, so the JDK compiler sees what the runtime would; Utils.classForName + // cannot target an arbitrary loader. + Some(Class.forName(binaryName, false, classLoader)) + // scalastyle:on classforname + } catch { + case _: ClassNotFoundException | _: LinkageError => None + case NonFatal(_) => None + } + } + + /** + * The source name the JDK compiler accepts for `cls`: the canonical name when it is a + * plain dotted identifier name (regular nesting, e.g. `java.util.Map.Entry`), otherwise + * the binary name. The binary name is required for classes nested in Scala objects and + * for Scala companion-object classes, whose canonical form carries a module `$` that + * javac cannot resolve, and for Scala operator-named classes (e.g. + * `scala.collection.immutable.::`) whose canonical form is not a valid Java identifier - + * in all those cases the binary name is itself a legal Java type reference. + * + * The canonical name is also rejected when it is not a faithful rename of the binary + * name - it must keep the same package. Scala REPL classes (e.g. + * `$line21.$read$$iw$TestCaseClass`) report a misleading `getCanonicalName` that drops + * the package and returns just the simple name (`TestCaseClass`); using it would corrupt + * the reference into an unqualified one javac cannot resolve. + * + * Reflection here can raise a `LinkageError` when the class loaded but its enclosing + * class did not (a partial or shaded jar); `NonFatal` does not cover that, so it is + * caught explicitly and the binary name used, matching what an unresolvable prefix + * yields in [[loadLongestPrefix]]. + */ + private def sourceNameOf(cls: Class[_]): String = { + val canonical = + try cls.getCanonicalName + catch { + case _: LinkageError => null + case NonFatal(_) => null + } + val pkg = cls.getPackageName + val usableCanonical = canonical != null && isPlainDottedName(canonical) && + (pkg.isEmpty || canonical.startsWith(pkg + ".")) + if (usableCanonical) canonical else cls.getName + } + + /** + * Climb to the nearest class that can be named in Java source. A class whose + * `getCanonicalName` is null has no source-referenceable name: the JDK compiler rejects + * a qualified reference to it even when the `.class` file is on the classpath, because + * the Java language forbids naming it. That covers anonymous classes (a Scala + * `new HashMap[..]() {...}` compiles to `Outer$$anon$1`), local classes, and classes + * nested inside either of those (`Outer$1$Inner`), which are themselves neither + * anonymous nor local. Janino does not care - it resolves any class by its runtime + * binary name - so this only matters for the JDK backend. For an anonymous class + * implementing an interface (`new Comparator() {...}`, whose superclass is `Object`), + * the implemented interface is preferred over `Object`. + * + * Narrowing a reference this way is only sound while every member the generated code + * could access remains reachable through the replacement type; a unit referencing a class + * for which reflection shows that does not hold is routed to Janino instead of being + * rewritten. When reflection cannot tell, the unit stays on the JDK backend and the + * reference keeps its binary name, which javac rejects, so the unit fails to compile + * rather than compiling into a wrong answer (see [[referencesUnnarrowableClass]] and + * [[CodeCompiler.active]]). + */ + private def nameableSupertype(start: Class[_]): Class[_] = { + var c: Class[_] = start + while (c != null && c.getCanonicalName == null) { + val sup: Class[_] = c.getSuperclass + c = + if (sup != null && (sup ne classOf[Object])) sup + else c.getInterfaces.headOption.getOrElse(sup) + } + if (c == null) classOf[Object] else c + } + + /** + * True when `body` references a class that [[nameableSupertype]] cannot narrow soundly, + * i.e. the class carries a public member that the replacement type does not offer, or + * the replacement type is itself one javac cannot reference. Such a unit must go to + * Janino: rewriting the reference would either drop the member or emit a type name javac + * rejects. + * + * Called from [[CodeCompiler.active]] on every compile, so it is gated behind a scan for + * `$` followed by a digit. Every class the Java language forbids naming carries that in + * its binary name (`Outer$1`, `Outer$1Local`, Scala's `Outer$$anon$1`, and their nested + * members `Outer$1$Inner`). The other `$` forms the rewrite handles do not: regular + * nesting (`Map$Entry`), Scala modules (`Foo$`, `Model$Load$Leaf`), package objects + * (`pkg$Inner`), specialized (`Function1$mcII$sp`) and operator-named (`$colon$colon`) + * classes. Runtime-synthesized classes need not carry it at all: a lambda is + * `Outer$$Lambda$14/0x...` on JDK 17 but `Outer$$Lambda/0x...` on 21 and later, and a + * hidden class is `Host$Named/0x...`. Neither routes on its own name: the tokenizer stops + * at `/`, and the truncated remainder either resolves to nothing (a lambda) or to the + * ordinary class the hidden class was defined from, whose own verdict is then the one that + * answers. On JDK 17 the lambda shape does pass the gate, at the cost of one failed load. + * + * The scan adds one linear pass over the body ahead of the compile-cache lookup, behind + * the intrinsified `$`-digit gate that ordinary generated code fails immediately. + * + * The scan reads the raw body, so a `$`-digit sequence outside code position (inside a + * string literal, say a regex the optimizer folded in, or inside a comment) can trigger + * the resolution attempt. That is harmless: an unloadable token is ignored, and a loadable + * one only ever picks Janino, which accepts a superset of what javac does. + */ + private[codegen] def referencesUnnarrowableClass(body: String): Boolean = { + if (!containsDollarDigit(body)) return false + val classLoader = Utils.getContextOrSparkClassLoader + val checked = mutable.HashSet.empty[String] + var i = 0 + val n = body.length + while (i < n) { + if (isNameStart(body.charAt(i))) { + val start = i + i += 1 + while (i < n && isNamePart(body.charAt(i))) i += 1 + val token = body.substring(start, i) + if (containsDollarDigit(token) && checked.add(token) && + routesOnToken(token, classLoader)) { + return true + } + } else { + i += 1 + } + } + false + } + + /** + * True when `token` resolves to a class whose reference has to go to Janino, i.e. one + * [[narrowingVerdict]] positively reports as [[Unnarrowable]]. + * + * An [[Unknown]] verdict does not route: it is no evidence that narrowing is unsafe, only + * that reflection could not tell. The token then contributes nothing to the decision, so + * unless some other token in the same body is [[Unnarrowable]] the unit stays on the + * configured backend, where a reference in code position keeps its binary name and javac + * rejects it, so the unit fails to compile rather than compiling into a wrong answer. + * Routing it to Janino would compile it instead, which is the better outcome for this one + * unit, but it would also let a truncated classpath quietly move work off the configured + * backend. The routing arms exist for source Spark knows javac cannot express, not for a + * broken deployment. + */ + private def routesOnToken(token: String, classLoader: ClassLoader): Boolean = { + loadLongestPrefix(token, classLoader).exists { + case (cls, _) => narrowingVerdict(cls) match { + case Unnarrowable => true + case Narrowable(_) | Unknown(_) => false + } + } + } + + /** True iff `s` holds a `$` immediately followed by an ASCII digit. */ + private[codegen] def containsDollarDigit(s: String): Boolean = { + var i = s.indexOf('$') + while (i >= 0 && i < s.length - 1) { + val next = s.charAt(i + 1) + if (next >= '0' && next <= '9') return true + i = s.indexOf('$', i + 1) + } + false + } + + /** + * Whether a reference to a class can be replaced by the nearest class Java can name. + * + * [[Narrowable]] carries the replacement type, so the decision and the name emitted for it + * come from one evaluation. The other two both keep the JDK backend from rewriting a + * reference, but they differ in what a caller may conclude: [[Unnarrowable]] is positive + * evidence that narrowing loses access, so the unit is routed to Janino, while [[Unknown]] + * means reflection could not answer (a partial or shaded jar) and is no evidence either + * way. + */ + private sealed trait NarrowingVerdict + private case class Narrowable(target: Class[_]) extends NarrowingVerdict + private case object Unnarrowable extends NarrowingVerdict + private case class Unknown(cause: Throwable) extends NarrowingVerdict + + /** + * The verdict for replacing a reference to `cls` with [[nameableSupertype]]. A class that + * is already nameable needs no narrowing, and only has to pass the accessibility check + * below on itself. + * + * Otherwise two things must hold. First, the replacement type must be one the generated + * unit can reference: it and every enclosing class must be public. Same-package is NOT + * sufficient even though javac would accept it - the generated class is defined into + * `org.apache.spark.sql.catalyst.expressions` but loaded by [[InMemoryClassLoader]], so + * its runtime package differs from the same-named package on the app loader and a + * package-private access would fail with `IllegalAccessError` at execution time instead + * of at compile time. Second, every public member of the concrete class - including + * inherited ones, since the generated code may access any of them - must be reachable on + * the replacement type. + * + * A member is matched by its exact erased signature, with one allowance for bridges: an + * override of a generic method has a narrower erasure than the supertype declaration it + * implements (`compare(String, String)` against `Comparator.compare(Object, Object)`), + * and the compiler emits a bridge carrying the supertype's signature. Such a method is + * safe to narrow because `invokevirtual` on the supertype signature still dispatches to + * the override. An overload has no bridge of its own, so it is rejected - and it must be, + * because narrowing binds the call to the supertype's method instead: `Invoke` codegen + * always wraps the call in an explicit cast, which would hide the type mismatch from + * javac and silently produce the wrong result rather than fail to compile. + * + * Telling the two apart needs the bridge to be matched to the specific method it forwards + * to, not merely to some method of the same name and arity. A class can hold both shapes + * at once - an anonymous `Comparator[String]` that also declares `compare(int, int)` has + * a bridge for the override and an unrelated overload sharing its name and arity - and + * excusing the whole name/arity group would let the overload through. So a bridge covers + * a method only when their parameter types line up (boxing counts as equivalent, since a + * specialized primitive override such as `apply(int)` is reached through an + * `apply(Object)` bridge), and only when it is the group's sole non-bridge method: with + * two of them the bridge cannot forward to both, so at least one becomes unreachable + * after narrowing. + * + * Fields and static methods get no such allowance. Both are bound statically, so they are + * matched by declaring class rather than by signature: a class that redeclares an inherited + * public field, or hides an inherited public static method, exposes it under a name the + * replacement type also answers to, and the narrowed reference would silently reach the + * replacement type's member instead. `getFields`/`getMethods` report the hiding and the + * hidden member both, so a name- or signature-based check would accept the pair. There is + * no counterpart to `invokevirtual`'s re-dispatch here, so the only safe such member is one + * the replacement type declares or inherits itself. + * + * Reflection over either class can raise a `LinkageError` when a signature or an + * enclosing class names something the loader cannot find (a partial or shaded jar), and + * the failure is not symmetric: the climb can succeed while member enumeration throws, so + * a verdict of [[Narrowable]] would let the rewrite emit a supertype name javac accepts + * for a class whose members were never checked. Such a class is reported [[Unknown]] + * instead, which keeps both the routing decision and the rewrite from acting on it. + * + * A non-`LinkageError` failure is caught the same way, because this file already treats + * this reflection as fallible in both directions: [[sourceNameOf]] guards + * `getCanonicalName` with `NonFatal`, as did the `resolveSourceName` that it and + * [[loadWithoutInit]] were split out of, over its whole load-climb-name sequence. Catching + * both is also what lets the routing scan read the raw body safely. A token there need not + * be a type reference in this unit at all: a class name embedded in a generated + * error-message literal reaches this method exactly as a real reference does, and an + * escaping exception would cost the unit its compile over one, at the routing step before + * any compile is attempted (`CodeGenerator.compile` only unwraps cache exceptions). + */ + private def narrowingVerdict(cls: Class[_]): NarrowingVerdict = { + try { + // An array class cannot arrive from the tokenizer: `javaType` renders an array as + // `component[]` and neither `[` nor `;` is a name character, so `foo.Bar$1[]` yields + // the component token, which is the one that needs the verdict. A future caller could + // still pass one, and it would narrow silently: an array of an unnameable component + // type has a null canonical name and only Object's members, so the climb lands on + // `Cloneable` and the member check passes. Reporting it unnarrowable routes the unit + // to Janino, which needs no rewrite for the component name it would find there. + if (cls.isArray) return Unnarrowable + val target = nameableSupertype(cls) + // A nameable class needs no narrowing, only a check that the generated unit can reach + // it at run time. That is the class's own modifier, NOT [[isPubliclyNameable]]'s walk + // over the enclosing chain: for a nested class `getModifiers` reports the source-level + // modifier from the InnerClasses attribute, while the JVM checks the class file's own + // `ACC_PUBLIC`, and a public class nested in a package-private one has it and is + // reachable across loaders. The walk is the right test for the climbed target below, + // where rejecting means routing to Janino. Here it would mean emitting this class's own + // binary name, which javac resolves for no nested class at all, turning a reference + // that compiled and ran into a compile error. The two arms answer different questions + // and disagree on this shape by design. + if (cls eq target) { + return if (Modifier.isPublic(cls.getModifiers)) Narrowable(cls) else Unnarrowable + } + // The climbed target is emitted by canonical name, so here nameability is what matters: + // a target javac cannot name from the generated unit's package must not be narrowed to. + if (!isPubliclyNameable(target)) return Unnarrowable + val reachable: Seq[Class[_]] = Seq(target, classOf[Object]) + // Statics are excluded: they are matched by declaring class below, and letting one + // satisfy the signature set would accept an INSTANCE method of `cls` whose only + // counterpart on the target is static, a call the narrowed reference would bind + // statically to the target's method. scalac produces that pair, since it does not + // treat a Java static as an inherited member. + val targetSignatures = reachable.flatMap(_.getMethods) + .filterNot(m => Modifier.isStatic(m.getModifiers)).map(erasedSignature).toSet + val targetStaticSignatures = reachable.flatMap(_.getMethods) + .filter(m => Modifier.isStatic(m.getModifiers)).map(erasedSignature).toSet + val methods = cls.getMethods + val bridges = methods.filter(m => m.isBridge && targetSignatures.contains(erasedSignature(m))) + val nonBridgeCount = methods.iterator.filterNot(_.isBridge) + .foldLeft(Map.empty[(String, Int), Int]) { (counts, m) => + val key = (m.getName, m.getParameterCount) + counts.updated(key, counts.getOrElse(key, 0) + 1) + } + def coveredByBridge(m: Method): Boolean = + nonBridgeCount.getOrElse((m.getName, m.getParameterCount), 0) <= 1 && + bridges.exists(b => forwardsTo(b, m)) + def declaredOnTarget(m: Member): Boolean = m.getDeclaringClass.isAssignableFrom(target) + val targetFieldNames = Seq(target, classOf[Object]).flatMap(_.getFields).map(_.getName).toSet + // A synthetic member is invisible to javac's source-level lookup: referencing one is + // "cannot find symbol" even though `getMethods`/`getFields` report it, so no generated + // reference can reach it and its loss to narrowing costs nothing. scalac emits several: + // a lambda body in the class becomes a `public static final $anonfun$...`, and an inner + // class carries a public `$outer` field and accessor. The exemption is withheld when the + // target answers to the same name and kind, because the JVM ignores `ACC_SYNTHETIC` when + // it resolves a statically-bound member. There the synthetic one shadows the target's, + // and narrowing would change which is read. + def exemptSynthetic(m: Method): Boolean = m.isSynthetic && { + val shadowed = + if (Modifier.isStatic(m.getModifiers)) targetStaticSignatures else targetSignatures + !shadowed.contains(erasedSignature(m)) + } + val membersLineUp = methods.forall { m => + if (Modifier.isStatic(m.getModifiers)) declaredOnTarget(m) || exemptSynthetic(m) + else targetSignatures.contains(erasedSignature(m)) || coveredByBridge(m) || + exemptSynthetic(m) + } && cls.getFields.forall { f => + declaredOnTarget(f) || (f.isSynthetic && !targetFieldNames.contains(f.getName)) + } + if (membersLineUp) Narrowable(target) else Unnarrowable + } catch { + case e: LinkageError => Unknown(e) + case NonFatal(e) => Unknown(e) + } + } + + private val unevaluableClassLogged = new java.util.concurrent.atomic.AtomicBoolean(false) + // The one outcome an operator cannot diagnose from the routing logs: reflection over a + // referenced class threw, so Spark cannot tell whether narrowing it is safe. The unit stays + // on the configured backend with the reference spelled as a binary name javac rejects, which + // surfaces as a compile error that looks like a codegen bug rather than a classpath one, and + // repeats, since the compile cache does not retain failures. Called only from + // [[narrowedSourceName]], so the once-per-JVM budget is spent on a real type reference. + private def logUnevaluableClassOnce(cls: Class[_], e: Throwable): Unit = { + if (unevaluableClassLogged.compareAndSet(false, true)) { + logWarning(log"Reflection over ${MDC(LogKeys.CLASS_NAME, cls.getName)} failed, so " + + log"Spark cannot tell whether a reference to it can be narrowed to a name the JDK " + + log"compiler accepts; the reference keeps its binary name, which that compiler " + + log"rejects. This usually means a partial or shaded jar on the classpath. Setting " + + log"${MDC(LogKeys.CONFIG, SQLConf.CODEGEN_COMPILER.key)} to 'janino' compiles such " + + log"units anyway. This notice is logged once per JVM.", e) + } + } + + /** + * True when `bridge` can be the bridge the compiler emitted for `impl`: same name and + * arity, and every bridge parameter accepts the corresponding one of `impl`. Boxing is + * treated as equivalence so that a specialized primitive override (`apply(int)` reached + * through an `apply(Object)` bridge) matches. + */ + private def forwardsTo(bridge: Method, impl: Method): Boolean = + bridge.getName == impl.getName && + bridge.getParameterCount == impl.getParameterCount && + bridge.getParameterTypes.zip(impl.getParameterTypes).forall { + case (bridgeParam, implParam) => boxed(bridgeParam).isAssignableFrom(boxed(implParam)) + } + + /** The wrapper type for a primitive, or `cls` itself when it is already a reference type. */ + private def boxed(cls: Class[_]): Class[_] = cls match { + case Integer.TYPE => classOf[java.lang.Integer] + case java.lang.Long.TYPE => classOf[java.lang.Long] + case java.lang.Double.TYPE => classOf[java.lang.Double] + case java.lang.Float.TYPE => classOf[java.lang.Float] + case java.lang.Short.TYPE => classOf[java.lang.Short] + case java.lang.Byte.TYPE => classOf[java.lang.Byte] + case Character.TYPE => classOf[java.lang.Character] + case java.lang.Boolean.TYPE => classOf[java.lang.Boolean] + case other => other + } + + private def erasedSignature(m: Method): (String, Seq[Class[_]]) = + (m.getName, m.getParameterTypes.toSeq) + + /** True iff `cls` and every class enclosing it are public, i.e. javac can name it. */ + private def isPubliclyNameable(cls: Class[_]): Boolean = { + var c = cls + while (c != null) { + if (!Modifier.isPublic(c.getModifiers)) return false + c = c.getEnclosingClass + } + true + } + + /** True iff `s` contains only `[A-Za-z0-9_.]` - a dotted Java identifier path. */ + private def isPlainDottedName(s: String): Boolean = { + var i = 0 + while (i < s.length) { + val c = s.charAt(i) + val ok = (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || + (c >= '0' && c <= '9') || c == '_' || c == '.' + if (!ok) return false + i += 1 + } + true + } + + // `$` preceded by an identifier char and followed by an uppercase letter. + private val InnerClassRefPattern = """(?<=[A-Za-z0-9_])\$(?=[A-Z])""".r + + // A dotted, qualified Java name (at least one `.`), used to recover the class + // references in a generated unit when resolving classes from a non-enumerable + // classloader (see `ClassLoaderFileManager.resolveReferencedClasses`). + private val QualifiedNamePattern = """[A-Za-z_$][\w$]*(?:\.[A-Za-z_$][\w$]*)+""".r + + /** + * Dedicated single-threaded executor on which the actual javac invocation runs. + * + * Reading jars can go through interruptible NIO file channels (both javac's + * platform-class reads and the JDK's jar handling). Spark runs codegen on task + * threads whose interrupt status may be set (task cancellation / cleanup); an + * interrupt during a jar read raises `ClosedByInterruptException` (surfaced as + * "bad class file ... unable to access file") and poisons the JDK's process-wide + * cached zip filesystem for that jar, so later compiles fail too. Janino is immune + * because it resolves classes via `ClassLoader.loadClass`, which reads jar bytes + * through `ZipFile`'s native path rather than an interruptible NIO channel. + * + * Running every compile on this never-interrupted thread keeps those channels + * open and valid. The single worker also confines the shared + * [[StandardJavaFileManager]] to one thread, so no additional locking is needed. + * + * The single worker serializes JDK-backend compiles JVM-wide: when many sessions or + * streaming queries trigger first-time codegen at once, their javac runs queue + * behind one another - a throughput cliff the Janino backend (which compiles on the + * calling threads) does not have. This is an accepted trade-off: each distinct unit + * compiles once per JVM and is then served from the cache, and the worker must be a + * dedicated never-interrupted thread regardless; a small pool with per-thread file + * managers could lift the limit later if profiles ever demand it. Note that the + * compilation-time metric measures caller-observed wall clock, so under contention + * it includes time spent queued behind other compiles. + */ + private lazy val compileExecutor: ExecutorService = + ThreadUtils.newDaemonSingleThreadExecutor("jdk-code-compiler") + + override def compile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = { + // Capture the CALLING thread's context classloader: it differs from the compile + // worker's, and both class resolution (ClassLoaderFileManager) and loading of the + // compiled output must see the task's classes - executor jars, and classes that + // live only on a runtime loader such as REPL-generated or Spark Connect session + // artifacts. `wrapAsCompilationUnit` is pure given the loader. The failure-log + // line budget is likewise session-bound, so it too is read on the calling thread + // (the compile worker has no session SQLConf attached). + val resolveLoader = Utils.getContextOrSparkClassLoader + val source = wrapAsCompilationUnit(code.body, resolveLoader) + val failureLogMaxLines = SQLConf.get.loggingMaxLinesForCodegen + + logBasedOnLevel(SQLConf.get.codegenLogLevel) { + log"\n${MDC(LogKeys.CODE, CodeFormatter.format(code))}" + } + + // Wrap the parent in ParentClassLoader for the same SPARK-15622 / SPARK-11636 + // reason the Janino backend does: a ClassNotFoundException from the parent must + // not carry a cause, or downstream resolution can fail in non-local deployments. + val parentLoader = new ParentClassLoader(resolveLoader) + + val future = compileExecutor.submit(new Callable[(GeneratedClass, ByteCodeStats)] { + override def call(): (GeneratedClass, ByteCodeStats) = + doCompile(code, source, resolveLoader, parentLoader, failureLogMaxLines) + }) + try { + // Await uninterruptibly: the result is cached and the worker must finish its + // jar reads without interruption (see compileExecutor). The caller's interrupt + // status is preserved for Spark to act on after this returns. Worst case a + // killed/speculated task holds its slot for one javac run before Spark observes + // the kill; the work is never wasted because the result lands in the cache. + Uninterruptibles.getUninterruptibly(future) + } catch { + case e: ExecutionException => + // Surface the root cause: a QueryExecutionErrors throwable for compile + // failures, or a raw reflection/linkage exception when loading or + // instantiating the compiled class fails - the same exceptions the Janino + // path's unguarded instantiation surfaces. + throw Option(e.getCause).getOrElse(e) + } + } + + /** The actual compilation; always runs on [[compileExecutor]]. */ + private def doCompile( + code: CodeAndComment, + source: String, + resolveLoader: ClassLoader, + parentLoader: ClassLoader, + failureLogMaxLines: Int): (GeneratedClass, ByteCodeStats) = { + val fileObject = new InMemorySourceFile(CodeCompiler.GeneratedClassName, source) + val diagnostics = new DiagnosticCollector[JavaFileObject]() + val fileManager = new ClassLoaderFileManager(sharedFileManager, resolveLoader, source) + // Captures anything javac writes outside the diagnostics listener (internal + // failures can bypass it); folded into the error message below instead of being + // silently dropped on `System.err`. + val compilerOut = new StringWriter() + + val task = compiler.getTask( + /* out = */ compilerOut, + /* fileManager = */ fileManager, + /* diagnostics = */ diagnostics, + /* options = */ compileOptions, + /* classes = */ null, + /* compilationUnits = */ java.util.Collections.singletonList(fileObject)) + + // On failure, dump the full compilation unit rather than the raw class body: the + // line numbers in javac diagnostics refer to the wrapped unit (header + adapted + // body), so this keeps the dump's `/* NNN */` markers aligned with `line NNN`. + def logSourceOnFailure(): Unit = CodeCompiler.logGeneratedCodeOnFailure( + new CodeAndComment(source, code.comment), failureLogMaxLines) + + val success = try { + task.call().booleanValue() + } catch { + case NonFatal(e) => + logError("Failed to compile the generated Java code.", e) + logSourceOnFailure() + throw QueryExecutionErrors.internalCompilerError( + new InternalCompilerException(e.getMessage, e)) + } + + if (!success) { + val errors = diagnostics.getDiagnostics.asScala + .filter(_.getKind == Diagnostic.Kind.ERROR) + .map(formatDiagnostic) + .mkString("\n") + // javac can report failure without ERROR diagnostics (internal conditions may + // surface through other kinds or the output writer); never raise an + // empty-message exception. + val message = Seq(errors, compilerOut.toString.trim).filter(_.nonEmpty) match { + case Seq() => "the JDK compiler returned failure without diagnostics" + case parts => parts.mkString("\n") + } + val ex = new CompileException(message, null) + logError("Failed to compile the generated Java code.", ex) + logSourceOnFailure() + throw QueryExecutionErrors.compilerError(ex) + } + + val classBytecodes = fileManager.snapshot() + val codeStats = CodeCompiler.computeByteCodeStats(classBytecodes) + val loader = new InMemoryClassLoader(classBytecodes.toMap, parentLoader) + + // `getConstructor` (public-only), matching the Janino path's instantiation. + val generated = loader.loadClass(CodeCompiler.GeneratedClassName) + .getConstructor() + .newInstance() + .asInstanceOf[GeneratedClass] + + (generated, codeStats) + } + + private def formatDiagnostic(d: Diagnostic[_ <: JavaFileObject]): String = { + val line = if (d.getLineNumber > 0) s"line ${d.getLineNumber}: " else "" + s"$line${d.getMessage(Locale.ROOT)}" + } + + // --- in-memory plumbing --- + + private class InMemorySourceFile(className: String, code: String) + extends SimpleJavaFileObject( + URI.create("string:///" + className.replace('.', '/') + + JavaFileObject.Kind.SOURCE.extension), + JavaFileObject.Kind.SOURCE) { + override def getCharContent(ignoreEncodingErrors: Boolean): CharSequence = code + } + + private class InMemoryClassFile(className: String) + extends SimpleJavaFileObject( + URI.create("bytes:///" + className.replace('.', '/') + + JavaFileObject.Kind.CLASS.extension), + JavaFileObject.Kind.CLASS) { + val bytes = new ByteArrayOutputStream() + override def openOutputStream(): java.io.OutputStream = bytes + def toBytes: Array[Byte] = bytes.toByteArray + } + + /** + * Per-jar index of `package path -> class binary names`, built once per jar and + * shared across compilations. Jars are immutable, so caching their contents is + * safe and avoids re-scanning a jar's full entry list on every `list()` call + * (the expensive part of classloader-based enumeration). Directories are NOT + * cached - they may gain classes at runtime (Spark Connect session artifacts, + * REPL output) - so those are always enumerated fresh. A jar path cannot serve + * changed content within a JVM either: Spark refuses to overwrite an added jar, + * and the JDK's own JarURLConnection cache (`setUseCaches(true)` below) already + * assumes that immutability. + * + * The cache is size-bounded so a long-running driver (e.g. a Spark Connect server + * whose sessions add distinct jars for years) cannot grow it without limit; the + * bound is far above the handful of jars generated code actually references, so + * eviction is rare and costs only a one-off re-index of that jar. + */ + private val jarPackageIndex: Cache[String, Map[String, Seq[String]]] = + CacheBuilder.newBuilder() + .maximumSize(2048) + .build[String, Map[String, Seq[String]]]() + + /** + * A [[JavaFileManager]] that resolves the `CLASS_PATH` location through a runtime + * [[ClassLoader]] (the task's context classloader) instead of a file-based + * classpath, mirroring how Janino resolves referenced classes. This lets the JDK + * compiler see exactly what Janino would - including classes that exist only on a + * runtime loader (REPL-generated, Spark Connect session artifacts) and never on + * `java.class.path`. Platform-class lookups (`java.*`, `jdk.*`) still go through + * the wrapped [[StandardJavaFileManager]]; compiled output is captured in + * [[InMemoryClassFile]] objects and never reaches it. + */ + private class ClassLoaderFileManager( + delegate: StandardJavaFileManager, + classLoader: ClassLoader, + source: String) + extends ForwardingJavaFileManager[StandardJavaFileManager](delegate) { + + // Qualified names referenced by the generated source, used only to resolve + // classes in packages the classloader cannot enumerate (see `list`). Computed + // lazily, so a normal file-classpath compile (every package enumerable) never + // pays for it. False positives from string/comment text are harmless - they + // simply fail the class-file probe in `resolveReferencedClasses`. + private lazy val referencedNames: Set[String] = + QualifiedNamePattern.findAllIn(source).toSet + + private val classFiles = new mutable.HashMap[String, InMemoryClassFile]() + + override def getJavaFileForOutput( + location: JavaFileManager.Location, + className: String, + kind: JavaFileObject.Kind, + sibling: FileObject): JavaFileObject = { + val out = new InMemoryClassFile(className) + classFiles.put(className, out) + out + } + + def snapshot(): Iterable[(String, Array[Byte])] = + classFiles.iterator.map { case (name, file) => (name, file.toBytes) }.toVector + + override def list( + location: JavaFileManager.Location, + packageName: String, + kinds: java.util.Set[JavaFileObject.Kind], + recurse: Boolean): java.lang.Iterable[JavaFileObject] = { + if (location == StandardLocation.CLASS_PATH && + kinds.contains(JavaFileObject.Kind.CLASS) && packageName.nonEmpty) { + val out = mutable.ArrayBuffer.empty[JavaFileObject] + val seen = mutable.HashSet.empty[String] + // First, enumerate the package via getResources (file dirs / jar entries). + listFromClassLoader(packageName, recurse, out, seen) + // Then add any classes the source references that enumeration missed. A + // classloader can serve a class by name (getResourceAsStream / loadClass) yet + // not enumerate it (getResources) - the Scala REPL and Spark Connect session + // loaders hold generated classes only in memory. This also covers split + // packages (e.g. `org.apache.spark.sql.connect`), where some classes are on + // the classpath and others - test/session artifacts - are only in memory. + // Janino resolves all of these by name; mirror that here. + resolveReferencedClasses(packageName, out, seen) + out.asJava + } else if (location == StandardLocation.CLASS_PATH && + kinds.contains(JavaFileObject.Kind.CLASS)) { + // Root package: never needed (the generated unit and its references live in + // named packages), and enumerating it would scan every classpath root. + java.util.Collections.emptyList[JavaFileObject]() + } else { + super.list(location, packageName, kinds, recurse) + } + } + + override def inferBinaryName( + location: JavaFileManager.Location, file: JavaFileObject): String = file match { + case f: ClassLoaderFileObject => f.binaryName + case _ => super.inferBinaryName(location, file) + } + + /** Enumerate the classes of `packageName` visible to the parent classloader. */ + private def listFromClassLoader( + packageName: String, + recurse: Boolean, + out: mutable.ArrayBuffer[JavaFileObject], + seen: mutable.HashSet[String]): Unit = { + val path = packageName.replace('.', '/') + val urls = + try classLoader.getResources(path).asScala.toSeq + catch { case NonFatal(_) => Seq.empty[URL] } + urls.foreach { url => + try { + url.getProtocol match { + case "file" => + collectFromDir(new java.io.File(url.toURI), packageName, recurse, out, seen) + case "jar" => collectFromJar(url, path, packageName, recurse, out, seen) + case _ => // custom protocols cannot be enumerated; skip + } + } catch { + case NonFatal(_) => () // a bad classpath entry must not fail the compile + } + } + } + + /** + * Add classes of `packageName` that the generated source references but that + * enumeration did not surface, resolving each referenced child by name. The + * existence probe uses `getResourceAsStream` (not `getResource`): in-memory + * loaders such as the Scala REPL and Spark Connect session loaders override + * `getResourceAsStream` to return bytes but have no resource URL, so `getResource` + * would return null even though the class is loadable - exactly the path Janino + * uses. Already-seen classes (from enumeration) are skipped, so split packages add + * only their in-memory extras. + * + * When the referenced child is itself a nested class (e.g. the Scala REPL's + * `$read$$iw$TestCaseClass`), javac needs not only that class file but its + * enclosing classes (`$read$$iw`, `$read$`) to resolve the inner-class chain. + * Those enclosing classes live in the same package but are not referenced by the + * source, so each `$`-boundary prefix of the child is probed and added too. + * + * Only the first dot-component after the package is probed: a dotted member-class + * reference (`pkg.Outer.Inner`) relies on `Outer$Inner` surfacing through package + * enumeration, which holds for every enumerable loader. The known non-enumerable + * loaders (REPL / Connect session artifacts) are deterministically routed to + * Janino up front (see `CodeCompiler.active`), and their class references are flat + * binary names rather than dotted member forms in any case. + */ + private def resolveReferencedClasses( + packageName: String, + out: mutable.ArrayBuffer[JavaFileObject], + seen: mutable.HashSet[String]): Unit = { + val prefix = packageName + "." + referencedNames.foreach { name => + if (name.startsWith(prefix) && name.length > prefix.length) { + val rest = name.substring(prefix.length) + val end = rest.indexOf('.') + val child = if (end < 0) rest else rest.substring(0, end) + // Probe the child and every enclosing-class prefix (cut at each non-leading + // `$`). The immediate child can also be a sub-package rather than a class + // (e.g. `expressions` in `o.a.s.sql.catalyst.expressions.UnsafeRow`), and + // some loaders return a non-null stream for a package-shaped path, so the + // class-file magic is verified before adding (avoids a phantom class that + // would clash with the package). + var i = 1 + while (i <= child.length) { + if (i == child.length || child.charAt(i) == '$') { + val binary = prefix + child.substring(0, i) + if (seen.add(binary) && isClassFileResource(binary.replace('.', '/') + ".class")) { + out += new ClassLoaderFileObject(binary, classLoader) + } + } + i += 1 + } + } + } + } + + /** True iff `resource` resolves to bytes beginning with the `0xCAFEBABE` magic. */ + private def isClassFileResource(resource: String): Boolean = { + val stream = classLoader.getResourceAsStream(resource) + if (stream == null) return false + try { + val head = stream.readNBytes(4) + head.length == 4 && + (head(0) & 0xFF) == 0xCA && (head(1) & 0xFF) == 0xFE && + (head(2) & 0xFF) == 0xBA && (head(3) & 0xFF) == 0xBE + } catch { + case NonFatal(_) => false + } finally { + try stream.close() catch { case NonFatal(_) => () } + } + } + + private def collectFromDir( + dir: java.io.File, + pkg: String, + recurse: Boolean, + out: mutable.ArrayBuffer[JavaFileObject], + seen: mutable.HashSet[String]): Unit = { + val children = dir.listFiles() + if (children == null) return + children.foreach { f => + val name = f.getName + if (f.isFile && name.endsWith(".class")) { + val binary = if (pkg.isEmpty) name.dropRight(6) else s"$pkg.${name.dropRight(6)}" + if (seen.add(binary)) out += new ClassLoaderFileObject(binary, classLoader) + } else if (recurse && f.isDirectory) { + collectFromDir(f, if (pkg.isEmpty) name else s"$pkg.$name", recurse, out, seen) + } + } + } + + private def collectFromJar( + url: URL, + path: String, + pkg: String, + recurse: Boolean, + out: mutable.ArrayBuffer[JavaFileObject], + seen: mutable.HashSet[String]): Unit = { + val conn = url.openConnection().asInstanceOf[JarURLConnection] + conn.setUseCaches(true) + val jarFile = conn.getJarFile + val jarKey = jarFile.getName + // Build (and cache) this jar's full package -> class-names index once. + val index = jarPackageIndex.get(jarKey, () => buildJarIndex(jarFile)) + val direct = index.getOrElse(path, Seq.empty) + val classes = + if (!recurse) direct + else index.iterator.collect { + case (p, names) if p == path || p.startsWith(path + "/") => names + }.flatten.toSeq + classes.foreach { binary => + if (seen.add(binary)) out += new ClassLoaderFileObject(binary, classLoader) + } + } + + private def buildJarIndex(jarFile: java.util.jar.JarFile): Map[String, Seq[String]] = { + val acc = mutable.HashMap.empty[String, mutable.ArrayBuffer[String]] + val entries = jarFile.entries() + while (entries.hasMoreElements) { + val e = entries.nextElement() + val n = e.getName + if (!e.isDirectory && n.endsWith(".class")) { + val slash = n.lastIndexOf('/') + val pkgPath = if (slash < 0) "" else n.substring(0, slash) + val binary = n.dropRight(6).replace('/', '.') + acc.getOrElseUpdate(pkgPath, mutable.ArrayBuffer.empty) += binary + } + } + acc.iterator.map { case (k, v) => (k, v.toSeq) }.toMap + } + } + + /** + * A `CLASS`-kind [[JavaFileObject]] whose bytes are read from a [[ClassLoader]] + * resource on demand, so the compiler reads exactly the bytecode the runtime + * loader would serve for that class. + */ + private class ClassLoaderFileObject(val binaryName: String, loader: ClassLoader) + extends SimpleJavaFileObject( + URI.create("classloader:///" + binaryName.replace('.', '/') + + JavaFileObject.Kind.CLASS.extension), + JavaFileObject.Kind.CLASS) { + override def openInputStream(): InputStream = { + val resource = binaryName.replace('.', '/') + ".class" + val is = loader.getResourceAsStream(resource) + if (is == null) throw new IOException(s"class resource not found: $resource") + is + } + } + + private class InMemoryClassLoader( + classBytes: Map[String, Array[Byte]], + parent: ClassLoader) extends ClassLoader(parent) { + override def findClass(name: String): Class[_] = { + classBytes.get(name) match { + case Some(bytes) => defineClass(name, bytes, 0, bytes.length) + case None => throw new ClassNotFoundException(name) + } + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala index def7ec39b5714..57864957bc13a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala @@ -17,23 +17,14 @@ package org.apache.spark.sql.catalyst.expressions.codegen -import java.io.ByteArrayInputStream - import scala.annotation.tailrec import scala.collection.mutable import scala.collection.mutable.ArrayBuffer -import scala.jdk.CollectionConverters._ -import scala.util.control.NonFatal import com.google.common.util.concurrent.{ExecutionError, UncheckedExecutionException} -import org.codehaus.commons.compiler.{CompileException, InternalCompilerException} -import org.codehaus.janino.ClassBodyEvaluator -import org.codehaus.janino.util.ClassFile -import org.codehaus.janino.util.ClassFile.CodeAttribute - -import org.apache.spark.{SparkException, SparkIllegalArgumentException, TaskContext, TaskKilledException} -import org.apache.spark.executor.InputMetrics -import org.apache.spark.internal.{Logging, LogKeys} + +import org.apache.spark.{SparkException, SparkIllegalArgumentException} +import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys._ import org.apache.spark.metrics.source.CodegenMetrics import org.apache.spark.sql.catalyst.InternalRow @@ -42,14 +33,13 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.types._ import org.apache.spark.sql.catalyst.types.ops.TypeOps -import org.apache.spark.sql.catalyst.util.{ArrayData, CollationAwareUTF8String, CollationFactory, CollationSupport, MapData, SQLOrderingUtil, UnsafeRowUtils} +import org.apache.spark.sql.catalyst.util.{ArrayData, MapData, SQLOrderingUtil, UnsafeRowUtils} import org.apache.spark.sql.catalyst.util.DateTimeConstants.NANOS_PER_MILLIS import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.Platform import org.apache.spark.unsafe.types._ -import org.apache.spark.util.{LongAccumulator, NonFateSharingCache, ParentClassLoader, Utils} +import org.apache.spark.util.{LongAccumulator, NonFateSharingCache, Utils} /** * Java source for evaluating an [[Expression]] given a [[InternalRow]] of input. @@ -301,9 +291,10 @@ class CodegenContext extends Logging { * * @param javaType Java type of the field. Note that short names can be used for some types, * e.g. InternalRow, UnsafeRow, UnsafeArrayData, etc. Other types will have to - * specify the fully-qualified Java type name. See the code in doCompile() for - * the list of default imports available. - * Also, generic type arguments are accepted but ignored. + * specify the fully-qualified Java type name. See + * `CodeCompiler.DefaultImports` for the list of default imports available. + * Also, generic type arguments are kept in field declarations, but stripped + * from array creation expressions because Java forbids generic array creation. * @param variableName Name of the field. * @param initFunc Function includes statement(s) to put into the init() method to initialize * this field. The argument is the name of the mutable state variable. @@ -402,6 +393,18 @@ class CodegenContext extends Logging { } def declareMutableStates(): String = { + def rawJavaType(javaType: String): String = { + val builder = new StringBuilder + var depth = 0 + javaType.foreach { + case '<' => depth += 1 + case '>' => depth -= 1 + case c if depth == 0 => builder.append(c) + case _ => + } + builder.toString + } + // It's possible that we add same mutable state twice, e.g. the `mergeExpressions` in // `TypedAggregateExpression`, we should call `distinct` here to remove the duplicated ones. val inlinedStates = inlinedMutableStates.distinct.map { case (javaType, variableName) => @@ -419,10 +422,10 @@ class CodegenContext extends Logging { if (javaType.contains("[]")) { // initializer had an one-dimensional array variable val baseType = javaType.substring(0, javaType.length - 2) - s"private $javaType[] $arrayName = new $baseType[$length][];" + s"private $javaType[] $arrayName = new ${rawJavaType(baseType)}[$length][];" } else { // initializer had a scalar variable - s"private $javaType[] $arrayName = new $javaType[$length];" + s"private $javaType[] $arrayName = new ${rawJavaType(javaType)}[$length];" } } } @@ -794,7 +797,7 @@ class CodegenContext extends Logging { val isNullA = freshName("isNullA") val elementB = freshName("elementB") val isNullB = freshName("isNullB") - val jt = javaType(elementType); + val jt = javaType(elementType) s""" |boolean $isNullA = $arrayA.isNullAt($i); |boolean $isNullB = $arrayB.isNullAt($i); @@ -1478,7 +1481,8 @@ abstract class CodeGenerator[InType <: AnyRef, OutType <: AnyRef] extends Loggin } /** - * Java bytecode statistics of a compiled class by Janino. + * Java bytecode statistics of a compiled generated class. Populated by whichever + * [[CodeCompiler]] backend produced the class. */ case class ByteCodeStats(maxMethodCodeSize: Int, maxConstPoolSize: Int, numInnerClasses: Int) @@ -1508,6 +1512,25 @@ object CodeGenerator extends Logging { // class. final val GENERATED_CLASS_SIZE_THRESHOLD = 1000000 + /** + * The `scala.Function1` `apply(Object)` bridge that projection codegen must emit for + * the Janino backend but hide from the JDK backend. + * + * Generated projections extend a Scala `Projection` (an `InternalRow => *`) and define + * the typed `apply(InternalRow)`. Janino does not synthesize bridge methods for the + * inherited generic supertype method, so it reports the class as not implementing + * `scala.Function1.apply(Object)` unless this explicit bridge is present. The JDK + * compiler, on the other hand, synthesizes the bridge itself and rejects an explicit + * one as a name clash. The two backends therefore need different source, so the bridge + * is emitted here in exactly the shape [[JdkCodeCompiler]] strips before invoking javac + * (keep the two in sync). `argName` can be any valid Java identifier - the body just + * casts and delegates to the typed overload. + */ + def function1ApplyBridge(argName: String): String = + s"""public java.lang.Object apply(java.lang.Object $argName) { + | return apply((InternalRow) $argName); + |}""".stripMargin + // This is the threshold for the number of global variables, whose types are primitive type or // complex type (e.g. more than one-dimensional array), that will be placed at the outer class final val OUTER_CLASS_VARIABLES_THRESHOLD = 10000 @@ -1535,13 +1558,22 @@ object CodeGenerator extends Logging { def resetCompileTime(): Unit = _compileTime.reset() /** - * Compile the Java source code into a Java class, using Janino. + * Compile the Java source code into a Java class via the active [[CodeCompiler]] + * backend (normally the one [[SQLConf.CODEGEN_COMPILER]] selects; see + * [[CodeCompiler.active]] for the deterministic routing overrides). * * @return a pair of a generated class and the bytecode statistics of generated functions. */ def compile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = try { val classLoaderRef = new HashableWeakReference(Utils.getContextOrSparkClassLoader) - cache.get((classLoaderRef, code)) + // The active backend is part of the cache key: flipping + // `spark.sql.codegen.compiler` mid-session must not silently reuse a class + // (and `ByteCodeStats`) compiled by the previously selected backend. The key + // holds the CodeCompiler singleton itself rather than its name: identity + // equality/hashCode is stable for an in-memory cache and immune to case + // variants of the name. + val backend = CodeCompiler.active(code) + cache.get((classLoaderRef, backend, code)) } catch { // Cache.get() may wrap the original exception. See the following URL // https://guava.dev/releases/14.0.1/api/docs/com/google/common/cache/ @@ -1550,128 +1582,6 @@ object CodeGenerator extends Logging { throw e.getCause } - /** - * Compile the Java source code into a Java class, using Janino. - */ - private[this] def doCompile(code: CodeAndComment): (GeneratedClass, ByteCodeStats) = { - val evaluator = new ClassBodyEvaluator() - - // A special classloader used to wrap the actual parent classloader of - // [[org.codehaus.janino.ClassBodyEvaluator]] (see CodeGenerator.doCompile). This classloader - // does not throw a ClassNotFoundException with a cause set (i.e. exception.getCause returns - // a null). This classloader is needed because janino will throw the exception directly if - // the parent classloader throws a ClassNotFoundException with cause set instead of trying to - // find other possible classes (see org.codehaus.janinoClassLoaderIClassLoader's - // findIClass method). Please also see https://issues.apache.org/jira/browse/SPARK-15622 and - // https://issues.apache.org/jira/browse/SPARK-11636. - val parentClassLoader = new ParentClassLoader(Utils.getContextOrSparkClassLoader) - evaluator.setParentClassLoader(parentClassLoader) - // Cannot be under package codegen, or fail with java.lang.InstantiationException - evaluator.setClassName("org.apache.spark.sql.catalyst.expressions.GeneratedClass") - evaluator.setDefaultImports( - classOf[Platform].getName, - classOf[InternalRow].getName, - classOf[UnsafeRow].getName, - classOf[BinaryView].getName, - classOf[UTF8String].getName, - classOf[Decimal].getName, - classOf[CalendarInterval].getName, - classOf[org.apache.spark.unsafe.types.TimestampNanosVal].getName, - classOf[VariantVal].getName, - classOf[ArrayData].getName, - classOf[UnsafeArrayData].getName, - classOf[MapData].getName, - classOf[UnsafeMapData].getName, - classOf[Expression].getName, - classOf[TaskContext].getName, - classOf[TaskKilledException].getName, - classOf[InputMetrics].getName, - classOf[CollationAwareUTF8String].getName, - classOf[CollationFactory].getName, - classOf[CollationSupport].getName, - QueryExecutionErrors.getClass.getName.stripSuffix("$") - ) - evaluator.setExtendedClass(classOf[GeneratedClass]) - - logBasedOnLevel(SQLConf.get.codegenLogLevel) { - // Only add extra debugging info to byte code when we are going to print the source code. - evaluator.setDebuggingInformation(true, true, false) - log"\n${MDC(LogKeys.CODE, CodeFormatter.format(code))}" - } - - val codeStats = try { - evaluator.cook("generated.java", code.body) - updateAndGetCompilationStats(evaluator) - } catch { - case e: InternalCompilerException => - logError("Failed to compile the generated Java code.", e) - logGeneratedCode(code) - throw QueryExecutionErrors.internalCompilerError(e) - case e: CompileException => - logError("Failed to compile the generated Java code.", e) - logGeneratedCode(code) - throw QueryExecutionErrors.compilerError(e) - } - - (evaluator.getClazz().getConstructor().newInstance().asInstanceOf[GeneratedClass], codeStats) - } - - private def logGeneratedCode(code: CodeAndComment): Unit = { - val maxLines = SQLConf.get.loggingMaxLinesForCodegen - if (Utils.isTesting) { - logError(s"\n${CodeFormatter.format(code, maxLines)}") - } else { - logInfo(s"\n${CodeFormatter.format(code, maxLines)}") - } - } - - /** - * Returns the bytecode statistics (max method bytecode size, max constant pool size, and - * # of inner classes) of generated classes by inspecting Janino classes. - * Also, this method updates the metrics information. - */ - private def updateAndGetCompilationStats(evaluator: ClassBodyEvaluator): ByteCodeStats = { - // First retrieve the generated classes. - val classes = evaluator.getBytecodes.asScala - - // Then walk the classes to get at the method bytecode. - val codeStats = classes.map { case (_, classBytes) => - val classCodeSize = classBytes.length - CodegenMetrics.METRIC_GENERATED_CLASS_BYTECODE_SIZE.update(classCodeSize) - try { - val cf = new ClassFile(new ByteArrayInputStream(classBytes)) - val constPoolSize = cf.getConstantPoolSize - val methodCodeSizes = cf.methodInfos.asScala.flatMap { method => - method.getAttributes.collect { case attr: CodeAttribute => - val byteCodeSize = attr.code.length - CodegenMetrics.METRIC_GENERATED_METHOD_BYTECODE_SIZE.update(byteCodeSize) - - if (byteCodeSize > DEFAULT_JVM_HUGE_METHOD_LIMIT) { - logInfo(log"Generated method too long to be JIT compiled: " + - log"${MDC(LogKeys.CLASS_NAME, cf.getThisClassName)}." + - log"${MDC(LogKeys.METHOD_NAME, method.getName)} is " + - log"${MDC(LogKeys.BYTECODE_SIZE, byteCodeSize)} bytes") - } - - byteCodeSize - } - } - (methodCodeSizes.max, constPoolSize) - } catch { - case NonFatal(e) => - logWarning("Error calculating stats of compiled class.", e) - (-1, -1) - } - } - - val (maxMethodSizes, constPoolSize) = codeStats.unzip - ByteCodeStats( - maxMethodCodeSize = maxMethodSizes.max, - maxConstPoolSize = constPoolSize.max, - // Minus 2 for `GeneratedClass` and an outer-most generated class - numInnerClasses = classes.size - 2) - } - /** * A cache of generated classes. * @@ -1686,10 +1596,11 @@ object CodeGenerator extends Logging { * aborted. See [[NonFateSharingCache]] for more details. */ private val cache = { - val loadFunc: ((HashableWeakReference, CodeAndComment)) => (GeneratedClass, ByteCodeStats) = { - case (_, code) => + val loadFunc: ((HashableWeakReference, CodeCompiler, CodeAndComment)) + => (GeneratedClass, ByteCodeStats) = { + case (_, backend, code) => val startTime = System.nanoTime() - val result = doCompile(code) + val result = backend.compile(code) val endTime = System.nanoTime() val duration = endTime - startTime val timeMs: Double = duration.toDouble / NANOS_PER_MILLIS @@ -1719,6 +1630,20 @@ object CodeGenerator extends Logging { val primitiveTypes = Seq(JAVA_BOOLEAN, JAVA_BYTE, JAVA_SHORT, JAVA_INT, JAVA_LONG, JAVA_FLOAT, JAVA_DOUBLE) + /** + * Returns the class name to embed as a type reference in generated code: the + * JVM binary name from `Class#getName` (e.g. `Outer$Inner` for nested classes). + * + * Janino accepts binary names directly, so this matches the historical + * behaviour. The JDK backend cannot, and adapts the name to the source form it + * requires in `JdkCodeCompiler.rewriteInnerClassRefs`: a regular nested class + * becomes `Outer.Inner`, while a class nested in a Scala `object` keeps its + * binary name (its canonical form carries a module `$` that javac cannot + * resolve). Feeding both backends the binary name from one place keeps that + * single adaptation correct for every reference. + */ + def javaSourceName(cls: Class[_]): String = cls.getName + /** * Returns true if a Java type is Java primitive primitive type */ @@ -2026,7 +1951,7 @@ object CodeGenerator extends Logging { case _: GeographyType | _: GeometryType => "BinaryView" case udt: UserDefinedType[_] => javaType(udt.sqlType) case ObjectType(cls) if cls.isArray => s"${javaType(ObjectType(cls.getComponentType))}[]" - case ObjectType(cls) => cls.getName + case ObjectType(cls) => javaSourceName(cls) case _ => PhysicalDataType(dt) match { case _: PhysicalArrayType => "ArrayData" case PhysicalBinaryType => "byte[]" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala index 2e018de07101e..ed18469ba0495 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateMutableProjection.scala @@ -128,8 +128,9 @@ object GenerateMutableProjection extends CodeGenerator[Seq[Expression], MutableP return (InternalRow) mutableRow; } - public java.lang.Object apply(java.lang.Object _i) { - InternalRow ${ctx.INPUT_ROW} = (InternalRow) _i; + ${CodeGenerator.function1ApplyBridge(ctx.INPUT_ROW)} + + public InternalRow apply(InternalRow ${ctx.INPUT_ROW}) { $evalSubexpr $allProjections // copy all the results into MutableRow diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala index d4f3a6100d522..68ae30e42e09b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateSafeProjection.scala @@ -197,8 +197,9 @@ object GenerateSafeProjection extends CodeGenerator[Seq[Expression], Projection] ${ctx.initPartition()} } - public java.lang.Object apply(java.lang.Object _i) { - InternalRow ${ctx.INPUT_ROW} = (InternalRow) _i; + ${CodeGenerator.function1ApplyBridge(ctx.INPUT_ROW)} + + public InternalRow apply(InternalRow ${ctx.INPUT_ROW}) { $allExpressions return mutableRow; } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala index f729ecacb4154..c3543c06750b2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/GenerateUnsafeProjection.scala @@ -418,10 +418,7 @@ object GenerateUnsafeProjection extends CodeGenerator[Seq[Expression], UnsafePro | ${ctx.initPartition()} | } | - | // Scala.Function1 need this - | public java.lang.Object apply(java.lang.Object row) { - | return apply((InternalRow) row); - | } + | ${CodeGenerator.function1ApplyBridge(ctx.INPUT_ROW)} | | public UnsafeRow apply(InternalRow ${ctx.INPUT_ROW}) { | ${eval.code} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala index f9dabd38664a5..80db7c88c6e30 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala @@ -191,6 +191,10 @@ case class ArraySize(child: Expression) */ @ExpressionDescription( usage = "_FUNC_(map) - Returns an unordered array containing the keys of the map.", + arguments = """ + Arguments: + * map - A map expression whose keys are returned as an array. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b')); @@ -226,6 +230,11 @@ case class MapKeys(child: Expression) */ @ExpressionDescription( usage = "_FUNC_(map, key) - Returns true if the map contains the key.", + arguments = """ + Arguments: + * map - A map expression to search. + * key - A key to look for. Its type must match, or be coercible to, the map's key type. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b'), 1); @@ -488,6 +497,10 @@ object ArraysZip { */ @ExpressionDescription( usage = "_FUNC_(map) - Returns an unordered array containing the values of the map.", + arguments = """ + Arguments: + * map - A map expression whose values are returned as an array. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b')); @@ -522,6 +535,10 @@ case class MapValues(child: Expression) */ @ExpressionDescription( usage = "_FUNC_(map) - Returns an unordered array of all entries in the given map.", + arguments = """ + Arguments: + * map - A map expression whose entries are returned as an array of key-value structs. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b')); @@ -697,6 +714,11 @@ case class MapEntries(child: Expression) */ @ExpressionDescription( usage = "_FUNC_(map, ...) - Returns the union of all the given maps", + arguments = """ + Arguments: + * map - A map expression. There can be one or more of them, and all must share + compatible key and value types. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b'), map(3, 'c')); @@ -825,6 +847,10 @@ case class MapConcat(children: Seq[Expression]) */ @ExpressionDescription( usage = "_FUNC_(arrayOfEntries) - Returns a map created from the given array of entries.", + arguments = """ + Arguments: + * arrayOfEntries - An array of two-field key-value structs from which the map is built. + """, examples = """ Examples: > SELECT _FUNC_(array(struct(1, 'a'), struct(2, 'b'))); @@ -1404,6 +1430,8 @@ case class Reverse(child: Expression) BinaryType, ArrayType)) + // Reversing a string transforms its content, so ImplicitTypeCasts promotes CHAR/VARCHAR to + // STRING. Array and binary inputs are unaffected. override def dataType: DataType = child.dataType private def resultArrayElementNullable = dataType.asInstanceOf[ArrayType].containsNull @@ -1621,6 +1649,12 @@ case class ArrayContains(left: Expression, right: Expression) @ExpressionDescription( usage = "_FUNC_(array, value) - Return index (0-based) of the search value, " + "if it is contained in the array; otherwise, (-<insertion point> - 1).", + arguments = """ + Arguments: + * array - A sorted array expression to search in. + * value - The value to search for. Its type must match, or be coercible to, the + array's element type. + """, examples = """ Examples: > SELECT _FUNC_(array(1, 2, 3), 2); @@ -2128,16 +2162,22 @@ case class Slice(x: Expression, start: Expression, length: Expression) val lengthInt = lengthVal.asInstanceOf[Int] val arr = xVal.asInstanceOf[ArrayData] val startIndex = ArrayExpressionUtils.sliceStartIndex(startInt, arr.numElements(), prettyName) - if (lengthInt < 0) { - throw QueryExecutionErrors.unexpectedValueForLengthInFunctionError(prettyName, lengthInt) - } + // Resolve (and validate) the result length via the shared helper, mirroring the codegen path. + // Besides rejecting a negative length, this clamps the length to the elements remaining after + // `startIndex`. For an in-range `startIndex`, this clamp keeps `startIndex + resLength` from + // overflowing `Int` -- the unclamped `startIndex + lengthInt` could wrap negative and make + // `slice` drop all elements. An out-of-range `startIndex` (a large negative `start`) can + // still wrap the helper's own `numElements - startIndex`, but the guard below returns before + // `resLength` is used. + val resLength = + ArrayExpressionUtils.sliceLength(lengthInt, arr.numElements(), startIndex, prettyName) // startIndex can be negative if start is negative and its absolute value is greater than the // number of elements in the array if (startIndex < 0 || startIndex >= arr.numElements()) { return new GenericArrayData(Array.empty[AnyRef]) } val data = arr.toSeq[AnyRef](elementType) - new GenericArrayData(data.slice(startIndex, startIndex + lengthInt)) + new GenericArrayData(data.slice(startIndex, startIndex + resLength)) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { @@ -2188,6 +2228,91 @@ case class Slice(x: Expression, start: Expression, length: Expression) copy(x = newFirst, start = newSecond, length = newThird) } +/** + * Removes the last `n` elements from the given array, per the ANSI SQL `TRIM_ARRAY` function. + */ +@ExpressionDescription( + usage = """ + _FUNC_(array, n) - Returns the given array with the last `n` elements removed. Raises an error + if `n` is negative or greater than the number of elements in the array.""", + arguments = """ + Arguments: + * array - the array to trim. + * n - the number of elements to remove from the end of the array. Must be between 0 and the + number of elements in the array (inclusive). + """, + examples = """ + Examples: + > SELECT _FUNC_(array(1, 2, 3, 4, 5), 2); + [1,2,3] + > SELECT _FUNC_(array('a', 'b', 'c'), 0); + ["a","b","c"] + > SELECT _FUNC_(array(1, 2, 3), 3); + [] + """, + group = "array_funcs", + since = "4.4.0") +case class TrimArray(left: Expression, right: Expression) + extends BinaryExpression with ImplicitCastInputTypes { + override def nullIntolerant: Boolean = true + + override def prettyName: String = "trim_array" + + override def dataType: DataType = left.dataType + + private def resultArrayElementNullable = dataType.asInstanceOf[ArrayType].containsNull + + override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, IntegerType) + + @transient private lazy val elementType: DataType = + left.dataType.asInstanceOf[ArrayType].elementType + + override def nullSafeEval(arrayVal: Any, nVal: Any): Any = { + val arr = arrayVal.asInstanceOf[ArrayData] + val n = nVal.asInstanceOf[Int] + val numElements = arr.numElements() + if (n < 0 || n > numElements) { + throw QueryExecutionErrors.invalidElementCountForTrimArrayError(prettyName, numElements, n) + } + val retainCount = numElements - n + val values = new Array[Any](retainCount) + for (i <- 0 until retainCount) { + if (!arr.isNullAt(i)) values(i) = arr.get(i, elementType) + } + new GenericArrayData(values) + } + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + nullSafeCodeGen(ctx, ev, (array, n) => { + val numElements = ctx.freshName("numElements") + val resLength = ctx.freshName("resLength") + val values = ctx.freshName("values") + val i = ctx.freshName("i") + val allocation = CodeGenerator.createArrayData( + values, elementType, resLength, s" $prettyName failed.") + val assignment = CodeGenerator.createArrayAssignment( + values, elementType, array, i, i, resultArrayElementNullable) + s""" + |${CodeGenerator.JAVA_INT} $numElements = $array.numElements(); + |if ($n < 0 || $n > $numElements) { + | throw QueryExecutionErrors.invalidElementCountForTrimArrayError( + | "$prettyName", $numElements, $n); + |} + |${CodeGenerator.JAVA_INT} $resLength = $numElements - $n; + |$allocation + |for (int $i = 0; $i < $resLength; $i ++) { + | $assignment + |} + |${ev.value} = $values; + """.stripMargin + }) + } + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): TrimArray = + copy(left = newLeft, right = newRight) +} + /** * Creates a String containing all the elements of the input array separated by the delimiter. */ @@ -2386,7 +2511,9 @@ case class ArrayJoin( } } - override def dataType: DataType = array.dataType.asInstanceOf[ArrayType].elementType + // After ImplicitTypeCasts, array elements that were CHAR/VARCHAR are STRING. + override def dataType: DataType = + array.dataType.asInstanceOf[ArrayType].elementType override def prettyName: String = "array_join" @@ -3022,6 +3149,11 @@ case class TryElementAt(left: Expression, right: Expression, replacement: Expres */ @ExpressionDescription( usage = "_FUNC_(col1, col2, ..., colN) - Returns the concatenation of col1, col2, ..., colN.", + arguments = """ + Arguments: + * colN - An expression to concatenate. There can be one or more of them, and all must be + of the same type: strings, binaries, or arrays. + """, examples = """ Examples: > SELECT _FUNC_('Spark', 'SQL'); @@ -3243,6 +3375,11 @@ case class Concat(children: Seq[Expression]) extends ComplexTypeMergingExpressio */ @ExpressionDescription( usage = "_FUNC_(arrayOfArrays) - Transforms an array of arrays into a single array.", + arguments = """ + Arguments: + * arrayOfArrays - An array whose elements are themselves arrays; they are concatenated + in order into one array. + """, examples = """ Examples: > SELECT _FUNC_(array(array(1, 2), array(3, 4))); @@ -3290,7 +3427,7 @@ case class Flatten(child: Expression) extends UnaryExpression throw QueryExecutionErrors.arrayFunctionWithElementsExceedLimitError( prettyName, numberOfElements) } - val flattenedData = new Array(numberOfElements.toInt) + val flattenedData = new Array[Any](numberOfElements.toInt) var position = 0 for (ad <- arrayData) { val arr = ad.toObjectArray(elementType) @@ -3433,8 +3570,8 @@ case class Sequence( override def nullable: Boolean = children.exists(_.nullable) - // If step is defined, then an error will be thrown if the start and stop do not satisfy the step. - override lazy val throwable: Boolean = stepOpt.isDefined + // Can throw if step is defined and start and stop don't match or any of the children can throw. + override lazy val throwable: Boolean = stepOpt.isDefined || children.exists(_.throwable) override def dataType: ArrayType = ArrayType(start.dataType, containsNull = false) @@ -3491,7 +3628,7 @@ case class Sequence( val physicalDataType = PhysicalDataType(iType) type T = physicalDataType.InternalType val integral = PhysicalIntegralType.integral(iType) - val ct = ClassTag[T](physicalDataType.tag.mirror.runtimeClass(physicalDataType.tag.tpe)) + val ct = physicalDataType.tag new IntegralSequenceImpl[T](iType)(ct, integral.asInstanceOf[Integral[T]]) case TimestampType | TimestampNTZType => @@ -3535,7 +3672,7 @@ case class Sequence( val arr = ctx.freshName("arr") val arrElemType = CodeGenerator.javaType(dataType.elementType) s""" - |final $arrElemType[] $arr = null; + |$arrElemType[] $arr = null; |${impl.genCode(ctx, startGen.value, stopGen.value, stepGen.value, arr, arrElemType)} |${ev.value} = UnsafeArrayData.fromPrimitiveArray($arr); """.stripMargin @@ -3987,14 +4124,20 @@ object Sequence { estimatedStep: String, len: String): String = { val calcFn = classOf[Sequence].getName + ".sequenceLength" + // `$start` and `$stop` are numeric expressions and `$step` is numeric or a + // CalendarInterval reference, so they have to be converted before going into a + // `Map<String, String>`. Janino, which compiles the generated code, erases the type + // arguments and binds `put` to `put(Object, Object)`, which lets the raw values through + // and leaves the parameter map holding non-String values that + // `SparkThrowable.getMessageParameters` then hands out as Strings. s""" |if (!(($estimatedStep > 0 && $start <= $stop) || | ($estimatedStep < 0 && $start >= $stop) || | ($estimatedStep == 0 && $start == $stop))) { | java.util.Map<String, String> params = new java.util.HashMap<String, String>(); - | params.put("start", $start); - | params.put("stop", $stop); - | params.put("step", $step); + | params.put("start", String.valueOf($start)); + | params.put("stop", String.valueOf($stop)); + | params.put("step", String.valueOf($step)); | throw new org.apache.spark.SparkIllegalArgumentException( | "_LEGACY_ERROR_TEMP_3243", params); |} @@ -4467,7 +4610,7 @@ case class ArrayDistinct(child: Expression) val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()" val hashSet = ctx.freshName("hashSet") val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName - val arrayBuilderClass = s"$arrayBuilder$$of$ptName" + val arrayBuilderClass = s"$arrayBuilder.of$ptName" // Only need to track null element index when array's element is nullable. val declareNullTrackVariables = if (resultArrayElementNullable) { @@ -4668,7 +4811,7 @@ case class ArrayUnion(left: Expression, right: Expression) extends ArrayBinaryLi val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()" val hashSet = ctx.freshName("hashSet") val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName - val arrayBuilderClass = s"$arrayBuilder$$of$ptName" + val arrayBuilderClass = s"$arrayBuilder.of$ptName" val body = s""" @@ -4893,7 +5036,7 @@ case class ArrayIntersect(left: Expression, right: Expression) extends ArrayBina val hashSet = ctx.freshName("hashSet") val hashSetResult = ctx.freshName("hashSetResult") val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName - val arrayBuilderClass = s"$arrayBuilder$$of$ptName" + val arrayBuilderClass = s"$arrayBuilder.of$ptName" val withArray2NaNCheckCodeGenerator = (array: String, index: String) => @@ -5118,7 +5261,7 @@ case class ArrayExcept(left: Expression, right: Expression) extends ArrayBinaryL val classTag = s"scala.reflect.ClassTag$$.MODULE$$.$hsTypeName()" val hashSet = ctx.freshName("hashSet") val arrayBuilder = classOf[mutable.ArrayBuilder[_]].getName - val arrayBuilderClass = s"$arrayBuilder$$of$ptName" + val arrayBuilderClass = s"$arrayBuilder.of$ptName" val withArray2NaNCheckCodeGenerator = (array: String, index: String) => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala index 122272ba43339..2f5b05af043d3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala @@ -55,6 +55,11 @@ trait NoThrow */ @ExpressionDescription( usage = "_FUNC_(expr, ...) - Returns an array with the given elements.", + arguments = """ + Arguments: + * expr - An expression of any type to include as an array element. Zero or more + expressions can be given, and they must share a common type. + """, examples = """ Examples: > SELECT _FUNC_(1, 2, 3); @@ -179,6 +184,14 @@ private [sql] object GenArrayData { */ @ExpressionDescription( usage = "_FUNC_(key0, value0, key1, value1, ...) - Creates a map with the given key/value pairs.", + arguments = """ + Arguments: + * keyN - An expression for the N-th map key. Keys must not be NULL and all keys must + share a common type. + * valueN - An expression for the value paired with keyN. All values must share a common + type. Keys and values are supplied as alternating arguments, so the total number of + arguments must be even. + """, examples = """ Examples: > SELECT _FUNC_(1.0, '2', 3.0, '4'); @@ -423,7 +436,12 @@ object CreateStruct { null, "struct", "_FUNC_(col1, col2, col3, ...) - Creates a struct with the given field values.", - "", + """ + | Arguments: + | * colN - The field values of the struct. There can be zero or more of them, + | each an expression of any type. Field names are assigned as `colN` by + | default unless the value is a named expression. + | """.stripMargin, """ | Examples: | > SELECT _FUNC_(1, 2, 3); @@ -446,6 +464,13 @@ object CreateStruct { // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(name1, val1, name2, val2, ...) - Creates a struct with the given field names and values.", + arguments = """ + Arguments: + * nameN - A STRING literal giving the name of the N-th struct field. It must not be NULL. + * valN - An expression of any type providing the value for the field named nameN. Names + and values are supplied as alternating arguments, so the total number of arguments + must be even. + """, examples = """ Examples: > SELECT _FUNC_("a", 1, "b", 2, "c", 3); @@ -615,9 +640,14 @@ case class StringToMap(text: Expression, pairDelim: Expression, keyValueDelim: E override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeNonCSAICollation, StringTypeNonCSAICollation, StringTypeNonCSAICollation) - override def dataType: DataType = MapType(first.dataType, first.dataType) + // The entries are split out of the input, so they do not carry its CHAR(n)/VARCHAR(n) length + // constraint. ImplicitTypeCasts promotes CHAR/VARCHAR to STRING at this ExpectsInputTypes site. + private lazy val entryType: DataType = + first.dataType + + override def dataType: DataType = MapType(entryType, entryType) - private lazy val mapBuilder = new ArrayBasedMapBuilder(first.dataType, first.dataType) + private lazy val mapBuilder = new ArrayBasedMapBuilder(entryType, entryType) private final lazy val collationId: Int = text.dataType.asInstanceOf[StringType].collationId diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala index 621f02ca18b86..24dd3623cd137 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/conditionalExpressions.scala @@ -33,6 +33,12 @@ import org.apache.spark.util.ArrayImplicits._ // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr1, expr2, expr3) - If `expr1` evaluates to true, then returns `expr2`; otherwise returns `expr3`.", + arguments = """ + Arguments: + * expr1 - A boolean expression evaluated as the condition. + * expr2 - The expression returned when `expr1` evaluates to true. + * expr3 - The expression returned when `expr1` evaluates to false or null. + """, examples = """ Examples: > SELECT _FUNC_(1 < 2, 'a', 'b'); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala index 5110aab9c6f03..15d4c15dcbdee 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/csvExpressions.scala @@ -40,6 +40,14 @@ import org.apache.spark.unsafe.types.UTF8String // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(csvStr, schema[, options]) - Returns a struct value with the given `csvStr` and `schema`.", + arguments = """ + Arguments: + * csvStr - A string expression of a single CSV record. + * schema - A string literal or invocation of `schema_of_csv` describing the schema. + * options - An optional map literal of string key-value pairs specifying CSV parsing + options controlling how `csvStr` is parsed. Accepts the same options as the CSV data + source. + """, examples = """ Examples: > SELECT _FUNC_('1, 0.8', 'a INT, b DOUBLE'); @@ -102,6 +110,7 @@ case class CsvToStructs( @transient private lazy val evaluator: CsvToStructsEvaluator = CsvToStructsEvaluator( options, nullableSchema, nameOfCorruptRecord, timeZoneId, requiredSchema) + override def stateful: Boolean = true override def nullSafeEval(input: Any): Any = { evaluator.evaluate(input.asInstanceOf[UTF8String]) @@ -133,6 +142,12 @@ case class CsvToStructs( */ @ExpressionDescription( usage = "_FUNC_(csv[, options]) - Returns schema in the DDL format of CSV string.", + arguments = """ + Arguments: + * csv - A foldable string expression of a single CSV record. + * options - An optional map literal of string key-value pairs specifying CSV parsing + options that control schema inference. Accepts the same options as the CSV data source. + """, examples = """ Examples: > SELECT _FUNC_('1,abc'); @@ -205,6 +220,13 @@ case class SchemaOfCsv( // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr[, options]) - Returns a CSV string with a given struct value", + arguments = """ + Arguments: + * expr - A struct expression to convert into a CSV string. + * options - An optional map literal of string key-value pairs specifying CSV generation + options controlling how the struct is rendered. Accepts the same options as the CSV data + source. + """, examples = """ Examples: > SELECT _FUNC_(named_struct('a', 1, 'b', 2)); @@ -275,6 +297,7 @@ case class StructsToCsv( lazy val converter: Any => UTF8String = { (row: Any) => UTF8String.fromString(gen.writeToString(row.asInstanceOf[InternalRow])) } + override def stateful: Boolean = true override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = copy(timeZoneId = Option(timeZoneId)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala index a9baf473c822e..46f9e494d2a88 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datasketchesExpressions.scala @@ -29,6 +29,11 @@ import org.apache.spark.sql.types.{AbstractDataType, BinaryType, BooleanType, Da usage = """ _FUNC_(expr) - Returns the estimated number of unique values given the binary representation of a Datasketches HllSketch. """, + arguments = """ + Arguments: + * expr - A binary expression holding the serialized representation of a Datasketches + HllSketch. + """, examples = """ Examples: > SELECT _FUNC_(hll_sketch_agg(col)) FROM VALUES (1), (1), (2), (2), (3) tab(col); @@ -70,6 +75,15 @@ case class HllSketchEstimate(child: Expression) Datasketches HllSketch objects, using a Datasketches Union object. Set allowDifferentLgConfigK to true to allow unions of sketches with different lgConfigK values (defaults to false). """, + arguments = """ + Arguments: + * first - A binary expression holding the serialized representation of a Datasketches + HllSketch. + * second - A binary expression holding the serialized representation of a Datasketches + HllSketch. + * allowDifferentLgConfigK - A boolean. Set to true to allow unions of sketches with + different lgConfigK values. Defaults to false. + """, examples = """ Examples: > SELECT hll_sketch_estimate(_FUNC_(hll_sketch_agg(col1), hll_sketch_agg(col2))) FROM VALUES (1, 4), (1, 4), (2, 5), (2, 5), (3, 6) tab(col1, col2); @@ -121,11 +135,19 @@ case class HllUnion(first: Expression, second: Expression, third: Expression) throw QueryExecutionErrors.hllInvalidInputSketchBuffer(prettyName) } val allowDifferentLgConfigK = value3.asInstanceOf[Boolean] - if (!allowDifferentLgConfigK && sketch1.getLgConfigK != sketch2.getLgConfigK) { + if (!allowDifferentLgConfigK && !sketch1.isEmpty && !sketch2.isEmpty && + sketch1.getLgConfigK != sketch2.getLgConfigK) { throw QueryExecutionErrors.hllUnionDifferentLgK( sketch1.getLgConfigK, sketch2.getLgConfigK, function = prettyName) } - val union = new Union(Math.min(sketch1.getLgConfigK, sketch2.getLgConfigK)) + // An empty sketch holds no coupons, so it carries no precision: it is exempt from the + // lgConfigK check and does not drag the result down to its own lgConfigK. + val lgConfigK = (sketch1.isEmpty, sketch2.isEmpty) match { + case (true, false) => sketch2.getLgConfigK + case (false, true) => sketch1.getLgConfigK + case _ => Math.min(sketch1.getLgConfigK, sketch2.getLgConfigK) + } + val union = new Union(lgConfigK) union.update(sketch1) union.update(sketch2) union.getResult(targetType).toUpdatableByteArray diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala index 0a73945bd4f61..fea6a813c2923 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.catalyst.util.{DateTimeUtils, LegacyDateFormats, Tim import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.catalyst.util.DateTimeUtils._ import org.apache.spark.sql.catalyst.util.LegacyDateFormats.SIMPLE_DATE_FORMAT -import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.errors.{DataTypeErrors, QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.types.StringTypeWithCollation import org.apache.spark.sql.types._ @@ -197,15 +197,141 @@ abstract class CurrentTimestampLike() extends LeafExpression with CodegenFallbac } /** - * Returns the current timestamp at the start of query evaluation. + * Returns the current timestamp at the start of query evaluation. The no-argument micro form + * registered as `current_timestamp` / `now`; see [[CurrentTimestampExpressionBuilder]] for the + * `current_timestamp(precision)` / `now(precision)` variants. * There is no code generation since this expression should get constant folded by the optimizer. */ -// scalastyle:off line.size.limit +case class CurrentTimestamp() extends CurrentTimestampLike { + override def prettyName: String = "current_timestamp" +} + +case class Now() extends CurrentTimestampLike { + override def prettyName: String = "now" +} + +/** + * Returns the current timestamp without time zone at the start of query evaluation. The + * no-argument micro form registered as `localtimestamp`; see + * [[LocalTimestampExpressionBuilder]] for the `localtimestamp(precision)` variant. + * There is no code generation since this expression should get constant folded by the optimizer. + */ +case class LocalTimestamp(timeZoneId: Option[String] = None) extends LeafExpression + with TimeZoneAwareExpression with CodegenFallback { + def this() = this(None) + override def foldable: Boolean = true + override def nullable: Boolean = false + override def dataType: DataType = TimestampNTZType + final override def nodePatternsInternal(): Seq[TreePattern] = Seq(CURRENT_LIKE) + override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = + copy(timeZoneId = Option(timeZoneId)) + override def eval(input: InternalRow): Any = localDateTimeToMicros(LocalDateTime.now(zoneId)) + override def prettyName: String = "localtimestamp" +} + +/** + * Returns the current timestamp with local time zone at the start of query evaluation, as a + * nanosecond-precision `TIMESTAMP_LTZ(precision)` (`precision` in `[7, 9]`). This is the + * nanosecond counterpart of [[CurrentTimestamp]] / [[Now]] and is produced by + * [[CurrentTimestampExpressionBuilder]] when `current_timestamp(p)` / `now(p)` is called with a + * nanosecond precision. Like the microsecond current-timestamp expressions it is foldable and + * gets constant folded by [[org.apache.spark.sql.catalyst.optimizer.ComputeCurrentTime]]; there + * is no code generation. + */ +case class CurrentTimestampNanos(precision: Int) extends CurrentTimestampLike { + override def dataType: DataType = TimestampLTZNanosType(precision) + override def eval(input: InternalRow): Any = + instantToTimestampNanos(java.time.Instant.now(), precision) + override def prettyName: String = "current_timestamp" +} + +/** + * Returns the current timestamp without time zone at the start of query evaluation, as a + * nanosecond-precision `TIMESTAMP_NTZ(precision)` (`precision` in `[7, 9]`). This is the + * nanosecond counterpart of [[LocalTimestamp]] and is produced by + * [[LocalTimestampExpressionBuilder]] when `localtimestamp(p)` is called with a nanosecond + * precision. Like [[LocalTimestamp]] it is time-zone aware (the session time zone determines the + * wall-clock value), foldable, and gets constant folded by + * [[org.apache.spark.sql.catalyst.optimizer.ComputeCurrentTime]]; there is no code generation. + */ +case class LocalTimestampNanos(precision: Int, timeZoneId: Option[String] = None) + extends LeafExpression with TimeZoneAwareExpression with CodegenFallback { + def this(precision: Int) = this(precision, None) + override def foldable: Boolean = true + override def nullable: Boolean = false + override def dataType: DataType = TimestampNTZNanosType(precision) + final override def nodePatternsInternal(): Seq[TreePattern] = Seq(CURRENT_LIKE) + override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = + copy(timeZoneId = Option(timeZoneId)) + override def eval(input: InternalRow): Any = + localDateTimeToTimestampNanos(LocalDateTime.now(zoneId), precision) + override def prettyName: String = "localtimestamp" +} + +/** + * Shared precision handling for the `current_timestamp(p)` / `now(p)` / `localtimestamp(p)` + * expression builders. + */ +private[expressions] object CurrentTimestampPrecision { + /** + * Validates the foldable integer precision argument `p` of a current-timestamp function and + * returns it. `p == 6` selects the historical microsecond type; `p` in `[7, 9]` selects the + * nanosecond type (which requires `spark.sql.timestampNanosTypes.enabled`). Any other value is + * rejected with `INVALID_TIMESTAMP_PRECISION`. The `typeName` (`TIMESTAMP_LTZ` / `TIMESTAMP_NTZ`) + * is used in the precision / feature-flag error messages, matching the `TIMESTAMP(p)` type + * parser. + */ + def validate(funcName: String, precision: Expression, typeName: String): Int = { + if (!precision.foldable) { + throw QueryCompilationErrors.nonFoldableArgumentError( + funcName, "precision", precision.dataType) + } + if (!precision.dataType.isInstanceOf[IntegralType]) { + throw QueryCompilationErrors.unexpectedInputDataTypeError( + funcName, 1, IntegerType, precision) + } + val value = precision.eval() + if (value == null) { + throw QueryCompilationErrors.unexpectedNullError("precision", precision) + } + val p = value.asInstanceOf[Number].intValue() + // Reject out-of-range precisions before the feature-flag check so the error is always + // INVALID_TIMESTAMP_PRECISION, not FEATURE_NOT_ENABLED (mirrors the TIMESTAMP(p) type parser). + if (p != 6 && + (p < TimestampLTZNanosType.MIN_PRECISION || p > TimestampLTZNanosType.MAX_PRECISION)) { + throw DataTypeErrors.invalidTimestampPrecisionError(p.toString, typeName) + } + if (p != 6) { + DataTypeErrors.checkTimestampNanosTypesEnabled() + } + p + } +} + +/** + * Builds `current_timestamp` / `now`. The no-argument form keeps the historical microsecond + * `TIMESTAMP` type ([[CurrentTimestamp]] / [[Now]]); the single-argument form accepts a foldable + * integer precision `p`, returning the microsecond `TIMESTAMP` for `p == 6` and a nanosecond + * `TIMESTAMP_LTZ(p)` for `p` in `[7, 9]` (gated behind `spark.sql.timestampNanosTypes.enabled`). + * The precision handling mirrors the `TIMESTAMP_LTZ(p)` type parser and `current_time(p)`. + * + * `now` is registered through the separate [[NowExpressionBuilder]] so it can document its own + * examples (`now` has no braceless form, unlike the `current_timestamp` keyword), but delegates + * its build logic here. + */ +// scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = """ _FUNC_() - Returns the current timestamp at the start of query evaluation. All calls of current_timestamp within the same query return the same value. _FUNC_ - Returns the current timestamp at the start of query evaluation. + + _FUNC_(precision) - Returns the current timestamp at the start of query evaluation, with the given fractional-seconds precision. A precision in [7, 9] returns a nanosecond-precision TIMESTAMP_LTZ(precision) and requires spark.sql.timestampNanosTypes.enabled to be true; precision 6 returns the standard microsecond TIMESTAMP. + """, + arguments = """ + Arguments: + * precision - An optional integer literal. Either 6 (the microsecond TIMESTAMP default) or a + value in [7, 9] selecting a nanosecond-precision TIMESTAMP_LTZ(precision). """, examples = """ Examples: @@ -213,59 +339,116 @@ abstract class CurrentTimestampLike() extends LeafExpression with CodegenFallbac 2020-04-25 15:49:11.914 > SELECT _FUNC_; 2020-04-25 15:49:11.914 + > SET spark.sql.timestampNanosTypes.enabled=true; + spark.sql.timestampNanosTypes.enabled true + > SELECT _FUNC_(9); + 2020-04-25 15:49:11.914120463 """, note = """ The syntax without braces has been supported since 2.0.1. """, group = "datetime_funcs", since = "1.5.0") -// scalastyle:on line.size.limit -case class CurrentTimestamp() extends CurrentTimestampLike { - override def prettyName: String = "current_timestamp" +// scalastyle:on line.size.limit line.contains.tab +object CurrentTimestampExpressionBuilder extends ExpressionBuilder { + override def build(funcName: String, expressions: Seq[Expression]): Expression = { + expressions.length match { + case 0 => + // Preserve the exact historical expression so `now` still maps to `Now` and rendering / + // pattern matching are unchanged. + if (funcName.equalsIgnoreCase("now")) Now() else CurrentTimestamp() + case 1 => + val p = CurrentTimestampPrecision.validate(funcName, expressions.head, "TIMESTAMP_LTZ") + if (p == 6) { + if (funcName.equalsIgnoreCase("now")) Now() else CurrentTimestamp() + } else { + CurrentTimestampNanos(p) + } + case n => + throw QueryCompilationErrors.wrongNumArgsError(funcName, Seq(0, 1), n) + } + } } +/** + * Builds `now` / `now(precision)`. Shares [[CurrentTimestampExpressionBuilder]]'s build logic but + * carries its own `@ExpressionDescription`: `now` is a regular function with no braceless form, + * so its examples must not include the `SELECT now` case that the `current_timestamp` keyword + * documents. + */ +// scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( - usage = "_FUNC_() - Returns the current timestamp at the start of query evaluation.", + usage = """ + _FUNC_() - Returns the current timestamp at the start of query evaluation. All calls of now within the same query return the same value. + + _FUNC_(precision) - Returns the current timestamp at the start of query evaluation, with the given fractional-seconds precision. A precision in [7, 9] returns a nanosecond-precision TIMESTAMP_LTZ(precision) and requires spark.sql.timestampNanosTypes.enabled to be true; precision 6 returns the standard microsecond TIMESTAMP. + """, + arguments = """ + Arguments: + * precision - An optional integer literal. Either 6 (the microsecond TIMESTAMP default) or a + value in [7, 9] selecting a nanosecond-precision TIMESTAMP_LTZ(precision). + """, examples = """ Examples: > SELECT _FUNC_(); 2020-04-25 15:49:11.914 + > SET spark.sql.timestampNanosTypes.enabled=true; + spark.sql.timestampNanosTypes.enabled true + > SELECT _FUNC_(9); + 2020-04-25 15:49:11.914120463 """, group = "datetime_funcs", since = "1.6.0") -case class Now() extends CurrentTimestampLike { - override def prettyName: String = "now" +// scalastyle:on line.size.limit line.contains.tab +object NowExpressionBuilder extends ExpressionBuilder { + override def build(funcName: String, expressions: Seq[Expression]): Expression = + CurrentTimestampExpressionBuilder.build(funcName, expressions) } /** - * Returns the current timestamp without time zone at the start of query evaluation. - * There is no code generation since this expression should get constant folded by the optimizer. + * Builds `localtimestamp`. The no-argument form keeps the historical microsecond `TIMESTAMP_NTZ` + * type ([[LocalTimestamp]]); the single-argument form accepts a foldable integer precision `p`, + * returning the microsecond `TIMESTAMP_NTZ` for `p == 6` and a nanosecond `TIMESTAMP_NTZ(p)` for + * `p` in `[7, 9]` (gated behind `spark.sql.timestampNanosTypes.enabled`). The precision handling + * mirrors the `TIMESTAMP_NTZ(p)` type parser and `current_time(p)`. */ -// scalastyle:off line.size.limit +// scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = """ _FUNC_() - Returns the current timestamp without time zone at the start of query evaluation. All calls of localtimestamp within the same query return the same value. _FUNC_ - Returns the current local date-time at the session time zone at the start of query evaluation. + + _FUNC_(precision) - Returns the current local date-time with the given fractional-seconds precision. A precision in [7, 9] returns a nanosecond-precision TIMESTAMP_NTZ(precision) and requires spark.sql.timestampNanosTypes.enabled to be true; precision 6 returns the standard microsecond TIMESTAMP_NTZ. + """, + arguments = """ + Arguments: + * precision - An optional integer literal. Either 6 (the microsecond TIMESTAMP_NTZ default) + or a value in [7, 9] selecting a nanosecond-precision TIMESTAMP_NTZ(precision). """, examples = """ Examples: > SELECT _FUNC_(); 2020-04-25 15:49:11.914 + > SET spark.sql.timestampNanosTypes.enabled=true; + spark.sql.timestampNanosTypes.enabled true + > SELECT _FUNC_(9); + 2020-04-25 15:49:11.914120463 """, group = "datetime_funcs", since = "3.4.0") -case class LocalTimestamp(timeZoneId: Option[String] = None) extends LeafExpression - with TimeZoneAwareExpression with CodegenFallback { - def this() = this(None) - override def foldable: Boolean = true - override def nullable: Boolean = false - override def dataType: DataType = TimestampNTZType - final override def nodePatternsInternal(): Seq[TreePattern] = Seq(CURRENT_LIKE) - override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = - copy(timeZoneId = Option(timeZoneId)) - override def eval(input: InternalRow): Any = localDateTimeToMicros(LocalDateTime.now(zoneId)) - override def prettyName: String = "localtimestamp" +// scalastyle:on line.size.limit line.contains.tab +object LocalTimestampExpressionBuilder extends ExpressionBuilder { + override def build(funcName: String, expressions: Seq[Expression]): Expression = { + expressions.length match { + case 0 => LocalTimestamp() + case 1 => + val p = CurrentTimestampPrecision.validate(funcName, expressions.head, "TIMESTAMP_NTZ") + if (p == 6) LocalTimestamp() else LocalTimestampNanos(p) + case n => + throw QueryCompilationErrors.wrongNumArgsError(funcName, Seq(0, 1), n) + } + } } /** @@ -306,6 +489,14 @@ case class CurrentBatchTimestamp( case _: TimestampType => Literal(timestampUs, TimestampType) case _: TimestampNTZType => Literal(convertTz(timestampUs, ZoneOffset.UTC, zoneId), TimestampNTZType) + // The batch timestamp is millisecond resolution, so nanos-within-micro is always 0. + // TIMESTAMP_LTZ is instant-based (zone-independent); TIMESTAMP_NTZ takes the wall clock + // in the session zone, mirroring the micro TimestampNTZType branch above. + case ltz: TimestampLTZNanosType => + Literal(TimestampNanosVal.fromParts(timestampUs, 0), ltz) + case ntz: TimestampNTZNanosType => + Literal( + TimestampNanosVal.fromParts(convertTz(timestampUs, ZoneOffset.UTC, zoneId), 0), ntz) case _: DateType => Literal(microsToDays(timestampUs, zoneId), DateType) } } @@ -1126,6 +1317,10 @@ case class UnixMicros(child: Expression) extends TimestampToLongBase { // scalastyle:off line.contains.tab @ExpressionDescription( usage = "_FUNC_(timestamp) - Returns the number of nanoseconds since 1970-01-01 00:00:00 UTC.", + arguments = """ + Arguments: + * timestamp - A nanosecond-precision timestamp value (TIMESTAMP_LTZ or TIMESTAMP_NTZ). + """, examples = """ Examples: > SET spark.sql.timestampNanosTypes.enabled=true; @@ -2749,7 +2944,8 @@ case class TimestampAddYMInterval( override def toString: String = s"$left + $right" override def sql: String = s"${left.sql} + ${right.sql}" - override def inputTypes: Seq[AbstractDataType] = Seq(AnyTimestampType, YearMonthIntervalType) + override def inputTypes: Seq[AbstractDataType] = + Seq(TypeCollection(AnyTimestampType, AnyTimestampNanoType), YearMonthIntervalType) override def dataType: DataType = timestamp.dataType @@ -2758,16 +2954,25 @@ case class TimestampAddYMInterval( @transient private lazy val zoneIdInEval: ZoneId = zoneIdForType(left.dataType) - override def nullSafeEval(micros: Any, months: Any): Any = { - timestampAddMonths(micros.asInstanceOf[Long], months.asInstanceOf[Int], zoneIdInEval) + override def nullSafeEval(start: Any, months: Any): Any = left.dataType match { + case _: AnyTimestampNanoType => + timestampNanosAddMonths( + start.asInstanceOf[TimestampNanosVal], months.asInstanceOf[Int], zoneIdInEval) + case _ => + timestampAddMonths(start.asInstanceOf[Long], months.asInstanceOf[Int], zoneIdInEval) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName) val dtu = DateTimeUtils.getClass.getName.stripSuffix("$") - defineCodeGen(ctx, ev, (micros, months) => { - s"""$dtu.timestampAddMonths($micros, $months, $zid)""" - }) + left.dataType match { + case _: AnyTimestampNanoType => + defineCodeGen(ctx, ev, (sd, months) => + s"""$dtu.timestampNanosAddMonths($sd, $months, $zid)""") + case _ => + defineCodeGen(ctx, ev, (micros, months) => + s"""$dtu.timestampAddMonths($micros, $months, $zid)""") + } } override protected def withNewChildrenInternal( @@ -4499,6 +4704,10 @@ object Extract { * the given timestamps. * - Otherwise the expression returns `DayTimeIntervalType` with the difference in microseconds * between given timestamps. + * + * Both microsecond and nanosecond-precision timestamp types are accepted as operands. Because the + * difference is reported on the microsecond grid in either mode, a nanosecond operand contributes + * only its `epochMicros`; the sub-microsecond remainder is truncated. */ case class SubtractTimestamps( left: Expression, @@ -4513,7 +4722,14 @@ case class SubtractTimestamps( def this(endTimestamp: Expression, startTimestamp: Expression) = this(endTimestamp, startTimestamp, SQLConf.get.legacyIntervalEnabled) - override def inputTypes: Seq[AbstractDataType] = Seq(AnyTimestampType, AnyTimestampType) + // Nanosecond-precision timestamps are accepted alongside the microsecond types. The difference is + // always reported on the microsecond grid (both result types -- DayTimeIntervalType and, in + // legacy mode, CalendarIntervalType -- carry microsecond resolution), so each operand contributes + // only its epochMicros; the sub-microsecond remainder is truncated. + override def inputTypes: Seq[AbstractDataType] = + Seq( + TypeCollection(AnyTimestampType, AnyTimestampNanoType), + TypeCollection(AnyTimestampType, AnyTimestampNanoType)) override def dataType: DataType = if (legacyInterval) CalendarIntervalType else DayTimeIntervalType() @@ -4522,6 +4738,13 @@ case class SubtractTimestamps( @transient private lazy val zoneIdInEval: ZoneId = zoneIdForType(left.dataType) + // For the nanosecond carrier the child value is a boxed TimestampNanosVal, so read its + // epochMicros; for the microsecond timestamp types it is already a boxed Long. + private def toMicros(value: Any): Long = value match { + case v: TimestampNanosVal => v.epochMicros + case n => n.asInstanceOf[Long] + } + @transient private lazy val evalFunc: (Long, Long) => Any = if (legacyInterval) { (leftMicros, rightMicros) => @@ -4531,17 +4754,29 @@ case class SubtractTimestamps( subtractTimestamps(leftMicros, rightMicros, zoneIdInEval) } - override def nullSafeEval(leftMicros: Any, rightMicros: Any): Any = { - evalFunc(leftMicros.asInstanceOf[Long], rightMicros.asInstanceOf[Long]) + override def nullSafeEval(leftTs: Any, rightTs: Any): Any = { + evalFunc(toMicros(leftTs), toMicros(rightTs)) } - override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = if (legacyInterval) { - defineCodeGen(ctx, ev, (end, start) => - s"new org.apache.spark.unsafe.types.CalendarInterval(0, 0, $end - $start)") - } else { - val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName) - val dtu = DateTimeUtils.getClass.getName.stripSuffix("$") - defineCodeGen(ctx, ev, (l, r) => s"""$dtu.subtractTimestamps($l, $r, $zid)""") + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + // The nanosecond carrier exposes epochMicros as a public field; the microsecond types are + // already primitive longs. Reduce each operand to microseconds before subtracting. + def toMicrosCode(e: Expression): String => String = e.dataType match { + case _: AnyTimestampNanoType => c => s"$c.epochMicros" + case _ => c => c + } + val leftMicros = toMicrosCode(left) + val rightMicros = toMicrosCode(right) + if (legacyInterval) { + defineCodeGen(ctx, ev, (end, start) => + s"new org.apache.spark.unsafe.types.CalendarInterval(0, 0, " + + s"${leftMicros(end)} - ${rightMicros(start)})") + } else { + val zid = ctx.addReferenceObj("zoneId", zoneIdInEval, classOf[ZoneId].getName) + val dtu = DateTimeUtils.getClass.getName.stripSuffix("$") + defineCodeGen(ctx, ev, (l, r) => + s"""$dtu.subtractTimestamps(${leftMicros(l)}, ${rightMicros(r)}, $zid)""") + } } override def toString: String = s"($left - $right)" @@ -4659,20 +4894,64 @@ case class ConvertTimezone( Seq( StringTypeWithCollation(supportsTrimCollation = true), StringTypeWithCollation(supportsTrimCollation = true), - TimestampNTZType) - override def dataType: DataType = TimestampNTZType - - override def nullSafeEval(srcTz: Any, tgtTz: Any, micros: Any): Any = { - DateTimeUtils.convertTimestampNtzToAnotherTz( + TypeCollection(TimestampNTZType, AnyTimestampNanoType)) + + // sourceTs's requiredType, as actually enforced by this method: TypeCollection includes + // AnyTimestampNanoType (rather than an NTZ-only nanos type) only so that an LTZ(p) input is + // accepted here and rejected below with the friendly message, instead of being silently + // widened to TimestampNTZType by the generic datetime-to-datetime implicit cast rule. That + // makes AnyTimestampNanoType.simpleString (which lists timestamp_ltz(p)) leak into the + // generic super.checkInputDataTypes() mismatch message for this param when sourceTs is some + // unrelated type (e.g. an int), even though LTZ(p) is never actually accepted. Route both the + // generic mismatch and the explicit LTZ rejection through the same message so they agree. + private def wrongSourceTsType: DataTypeMismatch = DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(2), + "requiredType" -> toSQLType("(timestamp_ntz or timestamp_ntz(p) with p in [7, 9])"), + "inputSql" -> toSQLExpr(sourceTs), + "inputType" -> toSQLType(sourceTs.dataType))) + + override def checkInputDataTypes(): TypeCheckResult = super.checkInputDataTypes() match { + case TypeCheckSuccess if sourceTs.dataType.isInstanceOf[TimestampLTZNanosType] => + wrongSourceTsType + case DataTypeMismatch("UNEXPECTED_INPUT_TYPE", params) + if params.get("paramIndex").contains(ordinalNumber(2)) => + wrongSourceTsType + case result => result + } + + private def isTsNanos: Boolean = sourceTs.dataType.isInstanceOf[AnyTimestampNanoType] + + // Preserves the exact source precision (7/8/9); AnyTimestampNanoType.defaultConcreteType would + // always widen the result to precision 9. + override def dataType: DataType = if (isTsNanos) sourceTs.dataType else TimestampNTZType + + override def nullSafeEval(srcTz: Any, tgtTz: Any, ts: Any): Any = { + val micros = if (isTsNanos) ts.asInstanceOf[TimestampNanosVal].epochMicros else ts + val convertedTs = DateTimeUtils.convertTimestampNtzToAnotherTz( srcTz.asInstanceOf[UTF8String].toString, tgtTz.asInstanceOf[UTF8String].toString, micros.asInstanceOf[Long]) + if (isTsNanos) { + TimestampNanosVal.fromParts( + convertedTs, ts.asInstanceOf[TimestampNanosVal].nanosWithinMicro) + } else convertedTs } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { val dtu = DateTimeUtils.getClass.getName.stripSuffix("$") - defineCodeGen(ctx, ev, (srcTz, tgtTz, micros) => - s"""$dtu.convertTimestampNtzToAnotherTz($srcTz.toString(), $tgtTz.toString(), $micros)""") + if (isTsNanos) { + defineCodeGen(ctx, ev, (srcTz, tgtTz, ts) => { + val convertedMicros = s"$dtu.convertTimestampNtzToAnotherTz(" + + s"$srcTz.toString(), $tgtTz.toString(), $ts.epochMicros)" + s"org.apache.spark.unsafe.types.TimestampNanosVal.fromParts(" + + s"$convertedMicros, $ts.nanosWithinMicro)" + }) + } else { + defineCodeGen(ctx, ev, (srcTz, tgtTz, micros) => + s"""$dtu.convertTimestampNtzToAnotherTz($srcTz.toString(), $tgtTz.toString(), $micros)""") + } } override def prettyName: String = "convert_timezone" diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala index b513b3858bbdd..3f2e838c19037 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/generators.scala @@ -143,6 +143,13 @@ case class UserDefinedGenerator( // scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = "_FUNC_(n, expr1, ..., exprk) - Separates `expr1`, ..., `exprk` into `n` rows. Uses column names col0, col1, etc. by default unless specified otherwise.", + arguments = """ + Arguments: + * n - The number of rows to separate the expressions into. Must be a + positive integer literal. + * exprN - The expressions to separate into rows. Their values are laid out + row-major across the `n` output rows. + """, examples = """ Examples: > SELECT _FUNC_(2, 1, 2, 3); @@ -437,6 +444,10 @@ trait ExplodeGeneratorBuilderBase extends GeneratorBuilder { // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr) - Separates the elements of array `expr` into multiple rows, or the elements of map `expr` into multiple rows and columns. Unless specified otherwise, uses the default column name `col` for elements of the array or `key` and `value` for the elements of the map.", + arguments = """ + Arguments: + * expr - An array or map expression whose elements are separated into rows. + """, examples = """ Examples: > SELECT _FUNC_(array(10, 20)); @@ -522,6 +533,11 @@ trait PosExplodeGeneratorBuilderBase extends GeneratorBuilder { // scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = "_FUNC_(expr) - Separates the elements of array `expr` into multiple rows with positions, or the elements of map `expr` into multiple rows and columns with positions. Unless specified otherwise, uses the column name `pos` for position, `col` for elements of the array or `key` and `value` for elements of the map.", + arguments = """ + Arguments: + * expr - An array or map expression whose elements are separated into rows, + each paired with its position. + """, examples = """ Examples: > SELECT _FUNC_(array(10,20)); @@ -643,6 +659,11 @@ trait InlineGeneratorBuilderBase extends GeneratorBuilder { // scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = "_FUNC_(expr) - Explodes an array of structs into a table. Uses column names col1, col2, etc. by default unless specified otherwise.", + arguments = """ + Arguments: + * expr - An array-of-structs expression whose struct fields become the + columns of each generated row. + """, examples = """ Examples: > SELECT _FUNC_(array(struct(1, 'a'), struct(2, 'b'))); @@ -701,6 +722,127 @@ object InlineOuterGeneratorBuilder extends InlineGeneratorBuilderBase { override def isOuter: Boolean = true } +/** + * Expands one or more arrays into a table, one row per element, implementing the ANSI SQL + * `UNNEST` collection derived table used in the FROM clause. + * + * When several arrays are supplied they are expanded in parallel: the number of output rows is + * the length of the longest array, and shorter arrays are padded with NULLs. A NULL array is + * treated as an empty array (contributes no elements). This matches the multi-array semantics of + * PostgreSQL and Trino. + * + * Each array contributes exactly one output column holding its element as-is; unlike `inline`, + * arrays of structs are not expanded into one column per field. When `withOrdinality` is set, a + * trailing 1-based `BIGINT` ordinality column is appended, matching `WITH ORDINALITY` in + * PostgreSQL and Trino (BigQuery's 0-based `WITH OFFSET` is intentionally not adopted). + * + * {{{ + * SELECT * FROM UNNEST(array(10, 20), array(30)) WITH ORDINALITY -> + * 10 30 1 + * 20 NULL 2 + * }}} + * + * This generator uses interpreted evaluation ([[CodegenFallback]]); `GenerateExec` therefore + * disables whole-stage codegen for the enclosing `Generate`. The per-array/per-ordinality zip with + * NULL padding does not map onto the existing [[CollectionGenerator]] codegen path (which emits a + * single `ArrayData`/`MapData`), and its interpreted `eval` is already lazy (one row built per + * pull). Interpreted generation is the same choice made by other non-`CollectionGenerator` + * generators such as [[ReplicateRows]]. A dedicated codegen path (analogous to `arrays_zip`) is + * possible future work if UNNEST becomes hot in whole-stage-codegen pipelines. + */ +case class Unnest(children: Seq[Expression], withOrdinality: Boolean) + extends Generator with CodegenFallback { + + private lazy val arrayElementTypes: Seq[ArrayType] = + children.map(_.dataType.asInstanceOf[ArrayType]) + + override def checkInputDataTypes(): TypeCheckResult = { + if (children.isEmpty) { + throw QueryCompilationErrors.wrongNumArgsError( + toSQLId(prettyName), Seq("> 0"), children.length) + } + val nonArray = children.zipWithIndex.collectFirst { + case (e, idx) if !e.dataType.isInstanceOf[ArrayType] => (e, idx) + } + nonArray match { + case Some((e, idx)) => + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(idx), + "requiredType" -> toSQLType(ArrayType), + "inputSql" -> toSQLExpr(e), + "inputType" -> toSQLType(e.dataType))) + case None => + TypeCheckResult.TypeCheckSuccess + } + } + + override def elementSchema: StructType = { + // With a single array keep the `explode`-compatible default name `col`; with several arrays + // use positional names `col0`, `col1`, and so on. A shorter array yields NULLs in the + // padded rows, so a padded column is always nullable regardless of the array's own + // `containsNull`. + val padded = children.length > 1 + val arrayFields = arrayElementTypes.zipWithIndex.map { case (at, idx) => + val name = if (children.length == 1) "col" else s"col$idx" + StructField(name, at.elementType, nullable = at.containsNull || padded) + } + val fields = if (withOrdinality) { + arrayFields :+ StructField("ordinality", LongType, nullable = false) + } else { + arrayFields + } + StructType(fields) + } + + override def eval(input: InternalRow): IterableOnce[InternalRow] = { + val arrays = children.map(_.eval(input).asInstanceOf[ArrayData]) + // The number of output rows is the length of the longest array; a null array counts as empty. + val numRows = arrays.iterator + .map(a => if (a == null) 0 else a.numElements()) + .foldLeft(0)(math.max) + val numArrays = arrays.length + val numFields = if (withOrdinality) numArrays + 1 else numArrays + // Build each output row lazily so that neither a single wide array nor a correlated LATERAL + // UNNEST over many input rows materializes the whole expansion up front; GenerateExec pulls + // rows one at a time. The iterator is single-pass but only ever traversed once per input row + // (GenerateExec's `outer` emptiness check uses `hasNext`, which does not build a row). + (0 until numRows).iterator.map { row => + val fields = new Array[Any](numFields) + var col = 0 + while (col < numArrays) { + val arr = arrays(col) + fields(col) = if (arr != null && row < arr.numElements() && !arr.isNullAt(row)) { + arr.get(row, arrayElementTypes(col).elementType) + } else { + null + } + col += 1 + } + if (withOrdinality) { + // Widen before incrementing so the 1-based ordinality cannot overflow Int near the top of + // the range (the ordinality column is BIGINT). + fields(numArrays) = row.toLong + 1L + } + // Wrap the array by reference (no copy) rather than InternalRow.fromSeq, which re-copies. + new GenericInternalRow(fields) + } + } + + override def prettyName: String = "unnest" + + // Keep the `withOrdinality` flag out of the default `toString`/EXPLAIN rendering (where it would + // otherwise appear as a bare `true` argument) and instead surface it as a readable suffix. + override def stringArgs: Iterator[Any] = { + val childArgs = children.iterator + if (withOrdinality) childArgs ++ Iterator("WITH ORDINALITY") else childArgs + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): Unnest = copy(children = newChildren) +} + @ExpressionDescription( usage = """_FUNC_() - Get Spark SQL keywords""", examples = """ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/grouping.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/grouping.scala index 0a12735b38da8..adb26487b1bf1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/grouping.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/grouping.scala @@ -196,6 +196,11 @@ object GroupingSets { _FUNC_(col) - indicates whether a specified column in a GROUP BY is aggregated or not, returns 1 for aggregated or 0 for not aggregated in the result set.", """, + arguments = """ + Arguments: + * col - A grouping column referenced in the GROUP BY clause. It must exactly match one + of the grouping expressions of the query. + """, examples = """ Examples: > SELECT name, _FUNC_(name), sum(age) FROM VALUES (2, 'Alice'), (5, 'Bob') people(age, name) GROUP BY cube(name); @@ -229,6 +234,12 @@ case class Grouping(child: Expression) extends Expression with Unevaluable _FUNC_([col1[, col2 ..]]) - returns the level of grouping, equals to `(grouping(c1) << (n-1)) + (grouping(c2) << (n-2)) + ... + grouping(cn)` """, + arguments = """ + Arguments: + * colN - An optional grouping column referenced in the GROUP BY clause. Zero or more + columns can be given; when provided, they must match the grouping columns exactly, + and when omitted all grouping columns are used. + """, examples = """ Examples: > SELECT name, _FUNC_(), sum(age), avg(height) FROM VALUES (2, 'Alice', 165), (5, 'Bob', 180) people(age, name, height) GROUP BY cube(name, height); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala index 085586041b07e..eb4004f34f29a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/hash.scala @@ -248,6 +248,78 @@ case class Crc32(child: Expression) override protected def withNewChildInternal(newChild: Expression): Crc32 = copy(child = newChild) } +@ExpressionDescription( + usage = "_FUNC_(expr) - Returns a 64-bit hash value of the argument using the XXH3 algorithm.", + arguments = """ + Arguments: + * expr - The expression to compute the XXH3 hash of. + An expression that evaluates to a binary. + """, + examples = """ + Examples: + > SELECT _FUNC_('Spark'); + 80997306238743657 + """, + since = "4.4.0", + group = "hash_funcs") +case class Xxh364(child: Expression) + extends UnaryExpression with ImplicitCastInputTypes { + override def nullIntolerant: Boolean = true + + override def dataType: DataType = LongType + + override def inputTypes: Seq[DataType] = Seq(BinaryType) + + override def contextIndependentFoldable: Boolean = child.contextIndependentFoldable + + protected override def nullSafeEval(input: Any): Any = + XXH3.hash64(input.asInstanceOf[Array[Byte]]) + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val cls = classOf[XXH3].getName + nullSafeCodeGen(ctx, ev, value => s"${ev.value} = $cls.hash64($value);") + } + + override def prettyName: String = "xxh3_64" + + override protected def withNewChildInternal(newChild: Expression): Xxh364 = copy(child = newChild) +} + +@ExpressionDescription( + usage = "_FUNC_(expr) - Returns a 128-bit XXH3 hash of the argument as a hex string.", + arguments = """ + Arguments: + * expr - The expression to compute the XXH3 hash of. + An expression that evaluates to a binary. + """, + examples = """ + Examples: + > SELECT _FUNC_('Spark'); + 7d57dd84c60c86ca1f4e82ab91a12b5e + """, + since = "4.4.0", + group = "hash_funcs") +case class Xxh3128(child: Expression) + extends UnaryExpression with ImplicitCastInputTypes with DefaultStringProducingExpression { + override def nullIntolerant: Boolean = true + + override def inputTypes: Seq[DataType] = Seq(BinaryType) + + override def contextIndependentFoldable: Boolean = child.contextIndependentFoldable + + protected override def nullSafeEval(input: Any): Any = + XXH3.hash128Hex(input.asInstanceOf[Array[Byte]]) + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val cls = classOf[XXH3].getName + defineCodeGen(ctx, ev, c => s"$cls.hash128Hex($c)") + } + + override def prettyName: String = "xxh3_128" + + override protected def withNewChildInternal(newChild: Expression): Xxh3128 = + copy(child = newChild) +} /** * A function that calculates hash value for a group of expressions. Note that the `seed` argument @@ -762,6 +834,11 @@ abstract class InterpretedHashFunction { */ @ExpressionDescription( usage = "_FUNC_(expr1, expr2, ...) - Returns a hash value of the arguments.", + arguments = """ + Arguments: + * exprN - The values to hash. There can be one or more of them, each an + expression of any data type. + """, examples = """ Examples: > SELECT _FUNC_('Spark', array(123), 2); @@ -833,6 +910,11 @@ case class CollationAwareMurmur3Hash(children: Seq[Expression], seed: Int) @ExpressionDescription( usage = "_FUNC_(expr1, expr2, ...) - Returns a 64-bit hash value of the arguments. " + "Hash seed is 42.", + arguments = """ + Arguments: + * exprN - The values to hash. There can be one or more of them, each an + expression of any data type. + """, examples = """ Examples: > SELECT _FUNC_('Spark', array(123), 2); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala index 524e27321f4fa..2f90d718294b1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/higherOrderFunctions.scala @@ -86,9 +86,15 @@ case class NamedLambdaVariable( override def qualifier: Seq[String] = Seq.empty + override def stateful: Boolean = true + override def newInstance(): NamedExpression = copy(exprId = NamedExpression.newExprId, value = new AtomicReference()) + override def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): NamedLambdaVariable = + copy(value = new AtomicReference()) + override def toAttribute: Attribute = { AttributeReference(name, dataType, nullable, Metadata.empty)(exprId, Seq.empty) } @@ -238,8 +244,16 @@ trait HigherOrderFunction extends Expression with ExpectsInputTypes { override lazy val canonicalized: Expression = { var currExprId = -1 + // Number the lambda variables of this higher-order function, but only those not already + // canonicalized. A canonical `NamedLambdaVariable` carries `value = null` (see the rename + // below); an original one carries an `AtomicReference`. When this HOF is nested inside another, + // the enclosing HOF's `canonicalized` renames every variable in the whole subtree first, then + // canonicalizes the children - which re-enters this method on the nested HOF. Re-numbering the + // already-renamed variables here would give a reference to an *enclosing* lambda's variable a + // different id than its binding, leaking it into `references`; skipping them keeps a variable + // and its references in agreement across nesting levels. val argumentMap = functions.flatMap(_.collect { - case l: NamedLambdaVariable => + case l: NamedLambdaVariable if l.value != null => currExprId += 1 l.exprId -> currExprId }).toMap @@ -387,6 +401,29 @@ trait MapBasedSimpleHigherOrderFunction extends SimpleHigherOrderFunction { override def argumentType: AbstractDataType = MapType } +/** + * A higher-order function whose result type is its argument's type, because it returns a subset or + * reordering of the input rather than the lambda's values. Members: `filter`, `map_filter`, + * `array_sort` (e.g. `filter(array<int>, ...) => array<int>`). Provides the shared `dataType`. + * + * The counterpart is [[ResultTypeFromFunction]]; the split is a real property of the expression, + * not of any evaluation: `filter` keeps its input's type, `transform`'s type follows the lambda. + */ +trait ResultTypeFromArgument extends SimpleHigherOrderFunction { + override def dataType: DataType = argument.dataType +} + +/** + * A higher-order function whose result type follows its lambda, not its argument. Members: + * `transform`, `transform_keys`, `transform_values`, `zip_with`, `map_zip_with`, `aggregate`, and + * the predicates `exists` / `forall` (whose boolean lambda gives a boolean result). + * + * Each member computes its own `dataType` (array of the element type, a re-keyed/re-valued map, the + * fold result, ...), so this is a marker with no shared implementation - the counterpart of + * [[ResultTypeFromArgument]]. + */ +trait ResultTypeFromFunction extends HigherOrderFunction + /** * Transform elements in an array using the transform function. This is similar to * a `map` in functional programming. @@ -412,7 +449,7 @@ trait MapBasedSimpleHigherOrderFunction extends SimpleHigherOrderFunction { case class ArrayTransform( argument: Expression, function: Expression) - extends ArrayBasedSimpleHigherOrderFunction { + extends ArrayBasedSimpleHigherOrderFunction with ResultTypeFromFunction { override def dataType: ArrayType = ArrayType(function.dataType, function.nullable) @@ -543,7 +580,7 @@ case class ArraySort( argument: Expression, function: Expression, allowNullComparisonResult: Boolean) - extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback { + extends ArrayBasedSimpleHigherOrderFunction with CodegenFallback with ResultTypeFromArgument { def this(argument: Expression, function: Expression) = { this( @@ -557,7 +594,6 @@ case class ArraySort( @transient lazy val elementType: DataType = argument.dataType.asInstanceOf[ArrayType].elementType - override def dataType: ArrayType = argument.dataType.asInstanceOf[ArrayType] override def checkInputDataTypes(): TypeCheckResult = { checkArgumentDataTypes() match { case TypeCheckResult.TypeCheckSuccess => @@ -658,6 +694,12 @@ object ArraySort { */ @ExpressionDescription( usage = "_FUNC_(expr, func) - Filters entries in a map using the function.", + arguments = """ + Arguments: + * expr - A map expression. + * func - A lambda function `(k, v) -> boolean` that returns whether the entry with + key `k` and value `v` should be kept. + """, examples = """ Examples: > SELECT _FUNC_(map(1, 0, 2, 2, 3, -1), (k, v) -> k > v); @@ -668,7 +710,7 @@ object ArraySort { case class MapFilter( argument: Expression, function: Expression) - extends MapBasedSimpleHigherOrderFunction with CodegenFallback { + extends MapBasedSimpleHigherOrderFunction with CodegenFallback with ResultTypeFromArgument { @transient lazy val (keyVar, valueVar) = { val args = function.asInstanceOf[LambdaFunction].arguments @@ -698,8 +740,6 @@ case class MapFilter( ArrayBasedMapData(retKeys.toArray, retValues.toArray) } - override def dataType: DataType = argument.dataType - override def functionType: AbstractDataType = BooleanType override def nodeName: String = "map_filter" @@ -738,9 +778,7 @@ case class MapFilter( case class ArrayFilter( argument: Expression, function: Expression) - extends ArrayBasedSimpleHigherOrderFunction { - - override def dataType: DataType = argument.dataType + extends ArrayBasedSimpleHigherOrderFunction with ResultTypeFromArgument { override def functionType: AbstractDataType = BooleanType @@ -885,7 +923,7 @@ case class ArrayExists( argument: Expression, function: Expression, followThreeValuedLogic: Boolean) - extends ArrayBasedSimpleHigherOrderFunction with Predicate { + extends ArrayBasedSimpleHigherOrderFunction with Predicate with ResultTypeFromFunction { def this(argument: Expression, function: Expression) = { this( @@ -1023,7 +1061,7 @@ object ArrayExists { case class ArrayForAll( argument: Expression, function: Expression) - extends ArrayBasedSimpleHigherOrderFunction with Predicate { + extends ArrayBasedSimpleHigherOrderFunction with Predicate with ResultTypeFromFunction { override def nullable: Boolean = super.nullable || function.nullable @@ -1153,7 +1191,7 @@ case class ArrayAggregate( zero: Expression, merge: Expression, finish: Expression) - extends HigherOrderFunction with QuaternaryLike[Expression] { + extends HigherOrderFunction with QuaternaryLike[Expression] with ResultTypeFromFunction { def this(argument: Expression, zero: Expression, merge: Expression) = { this(argument, zero, merge, LambdaFunction.identity) @@ -1352,6 +1390,12 @@ case class ArrayAggregate( */ @ExpressionDescription( usage = "_FUNC_(expr, func) - Transforms elements in a map using the function.", + arguments = """ + Arguments: + * expr - A map expression. + * func - A lambda function `(k, v) -> newKey` producing the transformed key from the + entry with key `k` and value `v`. The values are kept unchanged. + """, examples = """ Examples: > SELECT _FUNC_(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> k + 1); @@ -1364,7 +1408,7 @@ case class ArrayAggregate( case class TransformKeys( argument: Expression, function: Expression) - extends MapBasedSimpleHigherOrderFunction with CodegenFallback { + extends MapBasedSimpleHigherOrderFunction with CodegenFallback with ResultTypeFromFunction { @transient lazy val MapType(keyType, valueType, valueContainsNull) = argument.dataType @@ -1412,6 +1456,12 @@ case class TransformKeys( */ @ExpressionDescription( usage = "_FUNC_(expr, func) - Transforms values in the map using the function.", + arguments = """ + Arguments: + * expr - A map expression. + * func - A lambda function `(k, v) -> newValue` producing the transformed value from + the entry with key `k` and value `v`. The keys are kept unchanged. + """, examples = """ Examples: > SELECT _FUNC_(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> v + 1); @@ -1424,7 +1474,7 @@ case class TransformKeys( case class TransformValues( argument: Expression, function: Expression) - extends MapBasedSimpleHigherOrderFunction with CodegenFallback { + extends MapBasedSimpleHigherOrderFunction with CodegenFallback with ResultTypeFromFunction { @transient lazy val MapType(keyType, valueType, valueContainsNull) = argument.dataType @@ -1472,6 +1522,14 @@ case class TransformValues( NULL will be passed as the value for the missing key. If an input map contains duplicated keys, only the first entry of the duplicated key is passed into the lambda function. """, + arguments = """ + Arguments: + * map1 - The first map expression. + * map2 - The second map expression. + * function - A lambda function `(k, v1, v2) -> newValue` that produces the merged value + for key `k`, where `v1` and `v2` are the values from `map1` and `map2` respectively + (NULL when the key is missing from that map). + """, examples = """ Examples: > SELECT _FUNC_(map(1, 'a', 2, 'b'), map(1, 'x', 2, 'y'), (k, v1, v2) -> concat(v1, v2)); @@ -1482,7 +1540,8 @@ case class TransformValues( since = "3.0.0", group = "lambda_funcs") case class MapZipWith(left: Expression, right: Expression, function: Expression) - extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression] { + extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression] + with ResultTypeFromFunction { def functionForEval: Expression = functionsForEval.head @@ -1720,7 +1779,8 @@ case class MapZipWith(left: Expression, right: Expression, function: Expression) group = "lambda_funcs") // scalastyle:on line.size.limit case class ZipWith(left: Expression, right: Expression, function: Expression) - extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression] { + extends HigherOrderFunction with CodegenFallback with TernaryLike[Expression] + with ResultTypeFromFunction { def functionForEval: Expression = functionsForEval.head diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala index e80b4a355b726..c4eca772b1783 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/json/JsonExpressionEvalUtils.scala @@ -24,7 +24,7 @@ import scala.util.parsing.combinator.RegexParsers import com.fasterxml.jackson.core._ import com.fasterxml.jackson.core.json.JsonReadFeature -import org.apache.spark.SparkException +import org.apache.spark.{SparkException, TaskContext} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{ExprUtils, GenericInternalRow, GetJsonObject} import org.apache.spark.sql.catalyst.json.{CreateJacksonParser, JacksonGenerator, JacksonParser, JsonInferSchema, JSONOptions} @@ -55,9 +55,9 @@ object JsonPathParser extends RegexParsers { def root: Parser[Char] = '$' - def long: Parser[Long] = "\\d+".r ^? { - case x => x.toLong - } + // Guard the conversion so an oversized index (e.g. `[999999999999999999999999]`) makes the path + // fail to parse rather than throwing NumberFormatException out of the parser. + def long: Parser[Long] = "\\d+".r ^? { case x if x.toLongOption.isDefined => x.toLong } // parse `[*]` and `[123]` subscripts def subscript: Parser[List[PathInstruction]] = @@ -97,6 +97,13 @@ object JsonPathParser extends RegexParsers { None } } + + /** + * Returns `Some(true)` if the path parses and contains a wildcard, `Some(false)` if it parses + * without a wildcard, and `None` if it does not parse. Used by `JSON_TABLE` to validate that + * column and (container) row paths are simple, wildcard-free paths. + */ + def hasWildcard(str: String): Option[Boolean] = parse(str).map(_.contains(Wildcard)) } private[this] object SharedFactory { @@ -350,6 +357,831 @@ case class JsonTupleEvaluator(foldableFieldNames: Array[Option[String]]) { } } +/** + * The three-state result of navigating a JSON path for `JSON_TABLE`. `get_json_object` collapses + * "the path is absent" and "the value is JSON null" into a single `null`, which is wrong for + * `JSON_TABLE`: `EXISTS` must treat a present-but-null value as existing, and a value column must + * distinguish SQL `NULL` from the literal string `"null"`. This ADT keeps the two cases distinct. + */ +sealed trait JsonPathResult +object JsonPathResult { + /** The path did not match (the key/index is absent). */ + case object Missing extends JsonPathResult + /** The path matched a JSON `null` literal. */ + case object NullValue extends JsonPathResult + /** The path matched a value; `raw` is its verbatim JSON text (including quoted strings). */ + case class Found(raw: UTF8String) extends JsonPathResult +} + +/** + * The result of a single-value [[JsonTableEvaluator.lookup]] for `JSON_VALUE`. Unlike + * [[JsonPathResult]] -- whose `Found` carries verbatim JSON text for `JSON_TABLE` to unquote later + * -- a `Scalar` here already holds the value's cast-ready text (a string's unquoted/unescaped + * content, a number's or boolean's source text), and an object/array match is reported as + * `NonScalar` without being serialized at all, since `JSON_VALUE` routes non-scalars to ON ERROR. + */ +sealed trait JsonValueLookup +object JsonValueLookup { + /** The path did not match (routes to ON EMPTY). */ + case object Missing extends JsonValueLookup + /** The path matched a JSON `null` literal (yields SQL NULL). */ + case object NullValue extends JsonValueLookup + /** The path matched an object or array, i.e. not a scalar (routes to ON ERROR). */ + case object NonScalar extends JsonValueLookup + /** The path matched a scalar; `text` is its cast-ready value (strings already unquoted). */ + case class Scalar(text: UTF8String) extends JsonValueLookup +} + +/** + * The result of a single-value [[JsonTableEvaluator.queryLookup]] for `JSON_QUERY`. Unlike + * [[JsonValueLookup]] -- which reports an object/array match as `NonScalar` without serializing it, + * since `JSON_VALUE` never returns a non-scalar -- `Found` here always carries the matched value's + * serialized JSON text (`JSON_QUERY` returns objects, arrays, and scalars alike). `structural` + * distinguishes an object/array match from a scalar match, which the caller needs for the array + * wrapper (`WITH CONDITIONAL`) and quotes (`OMIT QUOTES`) behaviors. A JSON `null` literal is a + * scalar match whose text is `null`, not a distinct case. + */ +sealed trait JsonQueryLookup +object JsonQueryLookup { + /** The path did not match (routes to ON EMPTY). */ + case object Missing extends JsonQueryLookup + /** + * The path matched a value; `raw` is its serialized JSON text (a string is still quoted) and + * `structural` is true iff the value is an object or array. `unquoted` is the `OMIT QUOTES` + * form -- a matched JSON string's decoded content (read straight from the parser), and `raw` + * itself for every other value (objects, arrays, numbers, booleans, and JSON `null`, for which + * `OMIT QUOTES` is a no-op). Carrying it here lets the caller apply `OMIT QUOTES` without + * re-parsing the serialized fragment. It differs from `raw` only for a matched string, and only + * when [[JsonTableEvaluator.queryLookup]] is called with `omitQuotes = true`; the default + * `KEEP QUOTES` path discards `unquoted` and leaves it equal to `raw`, so the string decode is + * not paid for. As an optimization, a matched string under `OMIT QUOTES` -- whose `raw` the + * caller would discard, since `OMIT QUOTES` cannot be combined with a wrapper -- skips + * serialization entirely, so `raw` then also holds the decoded content. + */ + case class Found(raw: UTF8String, structural: Boolean, unquoted: UTF8String) + extends JsonQueryLookup +} + +/** + * A prefix trie over the (wildcard-free) column paths of a single `JSON_TABLE` invocation, built + * once via [[JsonTableEvaluator.buildPathTrie]] and reused for every row. It lets + * [[JsonTableEvaluator.navigateAll]] resolve all columns in a single traversal of a row item + * instead of re-parsing the item once per column. + * + * Each node groups the paths that share a common prefix: `named`/`indexed` hold the object-key and + * array-index steps to child nodes, and `terminals` lists the result-slot indices of the columns + * whose path ends exactly at this node. + */ +private[expressions] final class JsonTablePathTrie { + // Result-slot indices of columns whose path terminates at this node. + var terminals: List[Int] = Nil + // Object-key children, keyed by field name. + val named: mutable.HashMap[String, JsonTablePathTrie] = mutable.HashMap.empty + // Array-index children, keyed by index. + val indexed: mutable.HashMap[Long, JsonTablePathTrie] = mutable.HashMap.empty + + def hasChildren: Boolean = named.nonEmpty || indexed.nonEmpty + + /** True if no column path was inserted (e.g. an ordinality-only table): nothing to resolve. */ + def isEmpty: Boolean = terminals.isEmpty && !hasChildren +} + +/** + * The result of positioning a parser at a JSON path for the `JSON_TABLE` row source (see + * `positionAt`). Like [[JsonPathResult]] it distinguishes a missing path from a JSON `null`, but + * `AtValue` leaves the parser on the matched value's first token (rather than serializing it) so + * the row source can be streamed. + */ +sealed trait PositionResult +object PositionResult { + /** The path did not match. */ + case object Missing extends PositionResult + /** The path matched a JSON `null` literal. */ + case object NullValue extends PositionResult + /** The path matched a value; the parser is positioned at its first token. */ + case object AtValue extends PositionResult +} + +/** + * Token-aware navigation of a `containerPath` shared by the SQL/JSON functions. Three entry + * points, each with its own path constraints -- so `containerPath` is NOT wildcard-free in general: + * + * - [[evaluate]] -- `JSON_TABLE` row source: given the input, a wildcard-free container path, and + * whether the row path ended in `[*]`, produces the per-row JSON documents that the + * [[org.apache.spark.sql.catalyst.expressions.JsonTable]] generator projects into columns via + * [[navigateColumns]]. `$.items[*]` (containerPath `$.items`, `explodeRoot` = true) explodes an + * array into rows; `$` or `$.x` (`explodeRoot` = false) yields exactly one row. + * - [[lookup]] -- `JSON_VALUE` single-scalar extraction, over a wildcard-free path + * (`explodeRoot` = false). + * - [[pathExists]] -- `JSON_EXISTS` existence test. Here `containerPath` MAY contain wildcards + * (`[*]`, `.*`, `['*']`) and is evaluated in SQL/JSON *lax* mode (auto-wrap/unwrap, see + * [[anyMatch]]); `explodeRoot` is unused, so construct with `explodeRoot` = false. + * + * Unlike `get_json_object`, navigation here is token-aware and distinguishes missing keys from + * JSON `null` values (see [[JsonPathResult]] / [[PositionResult]]). + * + * Every entry point requires the input to be exactly one well-formed JSON value (no trailing + * garbage, not empty); anything else is treated as malformed, so the caller applies the ON ERROR + * behavior consistently in both modes. + */ +case class JsonTableEvaluator(containerPath: Seq[PathInstruction], explodeRoot: Boolean) { + import PathInstruction._ + import SharedFactory._ + + /** + * Returns the per-row JSON documents selected by the row path as an iterator, or `None` if the + * JSON is null or malformed, or if `[*]` was applied to a non-array (the caller maps `None` to + * the configured ON ERROR behavior). A well-formed input whose row path matches nothing returns + * `Some(empty iterator)`. + * + * The input is first scanned once to validate it is a single well-formed JSON value (so trailing + * garbage is rejected consistently in both ON ERROR modes -- this pass is O(n) tokens and does + * not materialize values). For the array (`[*]`) case the elements are then serialized one at a + * time from a second parser, so the whole expanded payload is never held in memory at once. + */ + final def evaluate(json: UTF8String): Option[Iterator[UTF8String]] = { + if (json == null || !isSingleWellFormedValue(json)) return None + // The parser is positioned at the matched value and, for the array case, handed to a lazy + // iterator that reads elements directly from it -- the container is never serialized whole. + // Ownership of `parser` transfers to that iterator (which closes it on exhaustion); in every + // other branch we close it before returning. + val parser = CreateJacksonParser.utf8String(jsonFactory, json) + var transferred = false + try { + parser.nextToken() + positionAt(parser, containerPath) match { + case PositionResult.Missing => + // Well-formed JSON, but the row path matched nothing: no rows. + Some(Iterator.empty) + case PositionResult.NullValue => + // The container is JSON null. `[*]` over a non-array is an error; otherwise one row. + if (explodeRoot) None else Some(Iterator.single(UTF8String.fromString("null"))) + case PositionResult.AtValue => + if (explodeRoot) { + // `[*]` requires an array; a non-array match is an error. + if (parser.currentToken != JsonToken.START_ARRAY) { + None + } else { + val it = arrayElementIterator(parser) // owns and eventually closes `parser` + transferred = true + Some(it) + } + } else { + Some(Iterator.single(serializeCurrentValue(parser))) + } + } + } catch { + case _: JsonProcessingException => None + } finally { + if (!transferred) parser.close() + } + } + + /** + * Resolves `containerPath` against a single JSON value for `JSON_VALUE`, preserving the missing / + * JSON-null / found distinction that [[evaluate]] collapses. Returns: + * + * - `None` if the input is not a single well-formed JSON value (malformed / trailing garbage / + * empty); + * - `Some(Missing)` if the path matches nothing; + * - `Some(NullValue)` if the path matches an explicit JSON `null`; + * - `Some(NonScalar)` if the path matches an object or array; + * - `Some(Scalar(text))` if the path matches a scalar, where `text` is its cast-ready value -- + * a string's unquoted/unescaped content, a number's or boolean's source text (see + * [[scalarAt]]). + * + * A `null` input is the caller's responsibility. `explodeRoot` is ignored: this is a single-value + * lookup, so construct the evaluator with `explodeRoot = false`. + * + * A single parser both navigates the path and validates well-formedness: after the matched value + * is captured (or the path is found missing), [[drainToRootEnd]] consumes the rest of the root + * value and rejects any trailing content, so a valid prefix followed by garbage (or a second + * root value) is rejected exactly as a fully malformed document is. This avoids the extra + * O(document size) validation pass a separate `isSingleWellFormedValue` scan would add per row. + */ + final def lookup(json: UTF8String): Option[JsonValueLookup] = { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, json)) { parser => + try { + if (parser.nextToken() == null) { + None // empty or whitespace-only + } else { + val result = positionAt(parser, containerPath) match { + case PositionResult.Missing => JsonValueLookup.Missing + case PositionResult.NullValue => JsonValueLookup.NullValue + case PositionResult.AtValue => scalarAt(parser) + } + // Reject a valid prefix trailed by extra content, keeping malformed-input semantics + // identical to the array row source (which validates the whole document up front). + if (drainToRootEnd(parser)) Some(result) else None + } + } catch { + case _: JsonProcessingException => None + } + } + } + + /** + * Resolves `containerPath` against a single JSON value for `JSON_QUERY`, serializing the matched + * value to JSON text. Returns: + * + * - `None` if the input is not a single well-formed JSON value (malformed / trailing garbage / + * empty), which the caller maps to ON ERROR; + * - `Some(Missing)` if the path matches nothing (ON EMPTY); + * - `Some(Found(raw, structural, unquoted))` if the path matches, where `raw` is the value's + * serialized JSON text, `structural` is true for an object or array (as opposed to a scalar, + * including a JSON `null`, whose text is `null`), and `unquoted` is the `OMIT QUOTES` form + * (a matched JSON string's decoded content; `raw` for every other value). + * + * `omitQuotes` mirrors the caller's `OMIT QUOTES` clause: only a matched JSON string's `unquoted` + * form differs from `raw`, and only `OMIT QUOTES` consumes it, so the decode-and-allocate is done + * for a string only when `omitQuotes` is true -- the default `KEEP QUOTES` path leaves `unquoted` + * equal to `raw` and skips the work. + * + * A `null` input is the caller's responsibility. Like [[lookup]] this navigates and validates + * with a single parser: after the matched value is serialized (which consumes it), + * [[drainToRootEnd]] + * walks out of the enclosing containers and rejects any trailing content, so a valid prefix + * followed by garbage is rejected exactly as a fully malformed document is. + */ + final def queryLookup(json: UTF8String, omitQuotes: Boolean): Option[JsonQueryLookup] = { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, json)) { parser => + try { + if (parser.nextToken() == null) { + None // empty or whitespace-only + } else { + val result = positionAt(parser, containerPath) match { + case PositionResult.Missing => JsonQueryLookup.Missing + // A JSON `null` literal is a scalar value for JSON_QUERY: serialize it to the text + // `null` rather than reporting it specially. The parser is positioned on the token. + case PositionResult.NullValue => + val raw = serializeCurrentValue(parser) + JsonQueryLookup.Found(raw, structural = false, unquoted = raw) + case PositionResult.AtValue => + parser.currentToken match { + case JsonToken.START_OBJECT | JsonToken.START_ARRAY => + val raw = serializeCurrentValue(parser) + JsonQueryLookup.Found(raw, structural = true, unquoted = raw) + case JsonToken.VALUE_STRING if omitQuotes => + // `OMIT QUOTES` returns the string's decoded content and cannot be combined with + // an array wrapper (rejected at analysis time), so the serialized (re-quoted) + // `raw` form would be discarded on this branch. Decode straight from the parser + // and carry it as both fields, skipping the wasted `serializeCurrentValue`. Only + // `OMIT QUOTES` reaches here; a `KEEP QUOTES` string falls through below. + val unquoted = UTF8String.fromString(parser.getText) + JsonQueryLookup.Found(unquoted, structural = false, unquoted) + case _ => + // A non-string scalar (number/boolean), for which `OMIT QUOTES` is a no-op, or a + // string under the default `KEEP QUOTES` (its decoded form is never used): the + // `unquoted` form is just `raw`, so no separate decode is needed. + val raw = serializeCurrentValue(parser) + JsonQueryLookup.Found(raw, structural = false, unquoted = raw) + } + } + if (drainToRootEnd(parser)) Some(result) else None + } + } catch { + case _: JsonProcessingException => None + } + } + } + + /** + * Classifies the value the parser is positioned at (the `AtValue` case of [[positionAt]]) for a + * `JSON_VALUE` [[lookup]], avoiding the serialize-then-reparse round trip that + * [[serializeCurrentValue]] followed by [[unquotedString]] would incur: + * + * - an object or array is consumed with `skipChildren` -- so [[drainToRootEnd]] can still walk + * back out and validate the document -- and reported as `NonScalar`, never serialized, since + * `JSON_VALUE` routes non-scalars to ON ERROR and so never needs the value; + * - a scalar's cast-ready text is read straight from the parser via `getText`, which returns a + * string's unquoted, unescaped content and a number's or boolean's verbatim source characters + * (so a high-precision fraction reaches a DECIMAL/STRING cast intact, as with + * `copyCurrentStructureExact`). + * + * A JSON `null` never reaches here -- [[positionAt]] reports it as `NullValue`. + */ + private def scalarAt(parser: JsonParser): JsonValueLookup = parser.currentToken match { + case JsonToken.START_OBJECT | JsonToken.START_ARRAY => + parser.skipChildren() + JsonValueLookup.NonScalar + case _ => + JsonValueLookup.Scalar(UTF8String.fromString(parser.getText)) + } + + /** + * Finishes consuming the root JSON value the `parser` is partway through and verifies nothing + * follows it, returning false if the input is truncated or has trailing content. Callers navigate + * to (and serialize) a matched value with the same parser, which can leave it positioned inside + * the enclosing containers; this walks back out to the root and confirms the document held + * exactly one well-formed value. + */ + private def drainToRootEnd(parser: JsonParser): Boolean = { + // Walk out of any still-open containers, consuming the remainder of the root value. + while (!parser.getParsingContext.inRoot) { + if (parser.nextToken() == null) return false // truncated mid-value + } + // At the root now: exactly one value was present iff nothing remains. + parser.nextToken() == null + } + + /** + * Tests whether `containerPath` matches at least one item in a single JSON document, evaluated in + * SQL/JSON *lax* mode (see [[anyMatch]]): wildcards are supported and a structural mismatch is a + * non-match rather than an error. Returns `Some(true)` if the path matches (including a match + * whose value is an explicit JSON `null`), `Some(false)` if it matches nothing, and `None` if the + * input is not a single well-formed JSON value (malformed / trailing garbage / empty). Does not + * serialize the matched value. A `null` input is the caller's responsibility; `explodeRoot` is + * ignored (construct the evaluator with `explodeRoot = false`). + * + * A single parser both navigates the path and validates well-formedness: after the match + * decision, [[drainToRootEnd]] consumes the rest of the root value and rejects any trailing + * content, so a valid prefix followed by garbage (or a second root value) is malformed exactly as + * a fully bad document is. This avoids the extra O(document size) pass a separate + * `isSingleWellFormedValue` scan would add per row. + */ + final def pathExists(json: UTF8String): Option[Boolean] = { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, json)) { parser => + try { + if (parser.nextToken() == null) { + None // empty or whitespace-only + } else { + val exists = anyMatch(parser, containerPath) + if (drainToRootEnd(parser)) Some(exists) else None + } + } catch { + case _: JsonProcessingException => None + } + } + } + + /** + * Returns whether `path` matches at least one item within the JSON value at the parser's current + * token, evaluated in SQL/JSON *lax* mode (the `JSON_EXISTS` default, matching Oracle and + * PostgreSQL): + * - wildcards are supported: `[*]` (any array element) and `.*` / `['*']` (any object member); + * - arrays are auto-unwrapped, i.e. a member/index/wildcard step applied to an array is applied + * to each element (so `$.a.b` matches when `a` is an array of objects each having `b`); + * - a non-array is auto-wrapped as a single-element array, so `[*]` / `[0]` match the value; + * - a structural mismatch (e.g. a member step on a scalar) is a non-match, never an error. + * + * Invariant: the parser enters positioned on the first token of the current value and leaves + * positioned on that value's last token -- the value is always fully consumed, even after a match + * is found -- so wildcard branches compose and [[drainToRootEnd]] can validate trailing content. + */ + private def anyMatch(parser: JsonParser, path: Seq[PathInstruction]): Boolean = { + path match { + case Nil => + // End of path: the current value is present (including a JSON null) -> a match. Consume it. + parser.skipChildren() + true + + case Key :: Named(name) :: rest => + parser.currentToken match { + case JsonToken.START_OBJECT => + // First-match semantics for a named key: only the first member with this name is + // followed, so a duplicate key later in the object is ignored. This matches the + // first-match `positionAt` / `navigateAll` used by `JSON_VALUE` / `JSON_TABLE`, so a + // path resolves consistently across the JSON functions. The whole object is still + // drained (subsequent members skipped) to keep the parser-position invariant. + var found = false + var matched = false + var token = parser.nextToken() + while (token != null && token != JsonToken.END_OBJECT) { + val matches = !matched && parser.currentName == name + parser.nextToken() // move onto the value; each branch consumes it + if (matches) { + matched = true + if (anyMatch(parser, rest)) found = true + } else { + parser.skipChildren() + } + token = parser.nextToken() + } + found + case JsonToken.START_ARRAY => + forEachElement(parser)(anyMatch(parser, path)) // lax auto-unwrap: apply to each element + case _ => + false // member accessor on a scalar: no match (the scalar is already fully consumed) + } + + case Subscript :: Index(index) :: rest => + parser.currentToken match { + case JsonToken.START_ARRAY => + var found = false + var i = 0L + var token = parser.nextToken() + while (token != null && token != JsonToken.END_ARRAY) { + if (i == index) { + if (anyMatch(parser, rest)) found = true + } else { + parser.skipChildren() + } + i += 1 + token = parser.nextToken() + } + found + case _ => + // lax auto-wrap: a non-array is a single-element array; [0] matches, [i>0] does not. + if (index == 0) { + anyMatch(parser, rest) + } else { + parser.skipChildren() // consume the wrapped value to keep the invariant + false + } + } + + case Subscript :: Wildcard :: rest => + parser.currentToken match { + case JsonToken.START_ARRAY => + forEachElement(parser)(anyMatch(parser, rest)) + case _ => + anyMatch(parser, rest) // lax auto-wrap: a non-array is a single-element array + } + + case Wildcard :: rest => + parser.currentToken match { + case JsonToken.START_OBJECT => + var found = false + var token = parser.nextToken() + while (token != null && token != JsonToken.END_OBJECT) { + parser.nextToken() // move onto the member value + if (found) parser.skipChildren() else if (anyMatch(parser, rest)) found = true + token = parser.nextToken() + } + found + case JsonToken.START_ARRAY => + forEachElement(parser)(anyMatch(parser, path)) // lax auto-unwrap: apply to each element + case _ => + false // member wildcard on a scalar: no members + } + + case _ => + // Unreachable: JsonPathParser only produces the instruction pairs handled above. + parser.skipChildren() + false + } + } + + /** + * Iterates the array the parser is positioned on (its current token is `START_ARRAY`), evaluating + * `matchElement` once per element with the parser positioned on that element's first token; each + * call must fully consume its element. Returns whether any element matched, and always drains the + * whole array, leaving the parser on the closing `END_ARRAY`. Once a match is found the remaining + * elements are skipped, not matched (existence short-circuits, but the array is still drained). + */ + private def forEachElement(parser: JsonParser)(matchElement: => Boolean): Boolean = { + var found = false + var token = parser.nextToken() + while (token != null && token != JsonToken.END_ARRAY) { + if (found) parser.skipChildren() else if (matchElement) found = true + token = parser.nextToken() + } + found + } + + /** + * Navigates `path` and leaves the parser positioned at the first token of the matched value + * (returning `AtValue`), or returns `Missing`/`NullValue`. Unlike the column projection traversal + * ([[navigateColumns]]), this does not serialize the value or finish consuming the enclosing + * containers -- the caller either streams from the current position (array row source) or + * serializes the single matched value. + */ + private def positionAt(parser: JsonParser, path: Seq[PathInstruction]): PositionResult = { + path match { + case Nil => + if (parser.currentToken == JsonToken.VALUE_NULL) PositionResult.NullValue + else PositionResult.AtValue + + case Key :: Named(name) :: rest => + if (parser.currentToken != JsonToken.START_OBJECT) { + skipRest(parser) + PositionResult.Missing + } else { + var token = parser.nextToken() + while (token != null && token != JsonToken.END_OBJECT) { + if (parser.currentName == name) { + parser.nextToken() // move onto the value; stop here (first match wins) + return positionAt(parser, rest) + } + parser.nextToken() + parser.skipChildren() + token = parser.nextToken() + } + PositionResult.Missing + } + + case Subscript :: Index(index) :: rest => + if (parser.currentToken != JsonToken.START_ARRAY) { + skipRest(parser) + PositionResult.Missing + } else { + var i = 0L + var token = parser.nextToken() + while (token != null && token != JsonToken.END_ARRAY) { + if (i == index) { + return positionAt(parser, rest) + } + parser.skipChildren() + i += 1 + token = parser.nextToken() + } + PositionResult.Missing + } + + case _ => + // Should not happen: JSON_TABLE paths are validated to be simple and wildcard-free. + skipRest(parser) + PositionResult.Missing + } + } + + /** + * Returns true if the input is exactly one well-formed JSON value with no trailing content, so a + * valid prefix followed by garbage, or an empty document, is treated as malformed (consistently + * in both ON ERROR modes). + */ + private def isSingleWellFormedValue(json: UTF8String): Boolean = { + try { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, json)) { parser => + if (parser.nextToken() == null) { + false // empty or whitespace-only + } else { + parser.skipChildren() // consume the first value in full + parser.nextToken() == null // nothing must remain after it + } + } + } catch { + case _: JsonProcessingException => false + } + } + + /** + * Resolves every column of `trie` against the value at the parser's current token in a single + * traversal, writing each matched terminal's [[JsonPathResult]] into `out` at its slot index. + * Only the simple wildcard-free instruction set produced for `JSON_TABLE` paths is modeled by the + * trie (`Key`/`Named` object steps and `Subscript`/`Index` array steps). + * + * Slots left untouched keep their initial `Missing`. A matched value is stored as its raw JSON + * text (`Found.raw`), i.e. strings keep their enclosing quotes so the fragment stays + * re-parseable; value columns unquote scalar strings afterwards via [[JsonTable]]'s extraction. + */ + private def navigateAll( + parser: JsonParser, + trie: JsonTablePathTrie, + out: Array[JsonPathResult]): Unit = { + val isNull = parser.currentToken == JsonToken.VALUE_NULL + + if (!trie.hasChildren) { + // Leaf node: every column terminates here, so just record the current value (or null) and + // consume it. This is the common case for disjoint column paths. + if (trie.terminals.nonEmpty) { + val result = if (isNull) JsonPathResult.NullValue + else JsonPathResult.Found(serializeCurrentValue(parser)) + trie.terminals.foreach(out(_) = result) + } else { + skipRest(parser) + } + } else if (trie.terminals.nonEmpty && !isNull) { + // A column path both ends here and extends deeper (e.g. `$.a` alongside `$.a.b`). Serialize + // the value once for the terminals, then re-parse that fragment to resolve the deeper + // columns -- this rare prefix overlap is the only place *within a single traversal* that a + // value is parsed more than once, and even then the descendant columns are still resolved in + // a single sub-traversal. (Separately, an array row item is serialized by + // `arrayElementIterator` and parsed again here, once per row.) + val raw = serializeCurrentValue(parser) + val result = JsonPathResult.Found(raw) + trie.terminals.foreach(out(_) = result) + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, raw)) { sub => + sub.nextToken() + descendInto(sub, trie, out) + } + } else { + // Descend for the deeper columns. Any terminal ending at this node resolves to `NullValue`: + // the earlier `!isNull` branch already handled non-null terminals, so a terminal reaching + // here means the value is a JSON null. Descendant-only slots that do not match stay `Missing` + // (a null has no children, so `descendInto` skips it and leaves them untouched). + if (trie.terminals.nonEmpty) trie.terminals.foreach(out(_) = JsonPathResult.NullValue) + descendInto(parser, trie, out) + } + } + + /** + * Descends into the object or array at the parser's current token, dispatching each matching + * field/element to the corresponding child trie node via [[navigateAll]] and skipping the rest. + * A scalar (or JSON null) has no children, so the whole value is skipped and the deeper columns + * are left as `Missing`. + */ + private def descendInto( + parser: JsonParser, + trie: JsonTablePathTrie, + out: Array[JsonPathResult]): Unit = { + parser.currentToken match { + case JsonToken.START_OBJECT if trie.named.isEmpty => + // No object-key columns descend here (only array-index paths): skip the whole object. + skipRest(parser) + + case JsonToken.START_OBJECT => + // First match wins for duplicate keys: once a trie key has been dispatched, later fields + // with the same name are skipped. + val consumed = mutable.HashSet.empty[String] + var token = parser.nextToken() + while (token != null && token != JsonToken.END_OBJECT) { + val name = parser.currentName + val child = trie.named.get(name) + parser.nextToken() // move onto the field value + if (child.isDefined && consumed.add(name)) { + navigateAll(parser, child.get, out) + } else { + parser.skipChildren() + } + token = parser.nextToken() + } + + case JsonToken.START_ARRAY if trie.indexed.isEmpty => + // No array-index columns descend here (only object-key paths): skip the whole array. + skipRest(parser) + + case JsonToken.START_ARRAY => + var i = 0L + var token = parser.nextToken() + while (token != null && token != JsonToken.END_ARRAY) { + val child = trie.indexed.get(i) + if (child.isDefined) { + navigateAll(parser, child.get, out) + } else { + parser.skipChildren() + } + i += 1 + token = parser.nextToken() + } + + case _ => + // A scalar where some columns expected to descend: those stay Missing. + skipRest(parser) + } + } + + /** Skips the remainder of the value at the parser's current token. */ + private def skipRest(parser: JsonParser): Unit = parser.skipChildren() + + /** + * Serializes the value at the parser's current token to its raw JSON text. Strings keep their + * enclosing quotes, so the result is always a re-parseable JSON fragment (this matters because a + * matched value may be re-parsed as a row item). Value columns unquote scalar strings afterwards + * via [[JsonTableEvaluator.unquotedString]]. + */ + private def serializeCurrentValue(parser: JsonParser): UTF8String = { + val output = new ByteArrayOutputStream() + Utils.tryWithResource(jsonFactory.createGenerator(output, JsonEncoding.UTF8)) { + // `copyCurrentStructureExact` preserves floating-point tokens byte-for-byte; the plain + // `copyCurrentStructure` may round them for textual formats, which would corrupt a + // high-precision fraction before JSON_TABLE casts the reserialized text to DECIMAL/STRING. + generator => generator.copyCurrentStructureExact(parser) + } + UTF8String.fromBytes(output.toByteArray) + } + + // The array parser currently owned by an outstanding `arrayElementIterator`, or null. Since + // `GenerateExec` evaluates rows sequentially and fully drains each row's iterator before the + // next `eval`, at most one such parser is open at a time per task. Tracked so a single + // task-completion listener (registered once below) can close it on early termination. + @transient private var openArrayParser: JsonParser = _ + @transient private var completionListenerRegistered = false + + /** + * Streams the elements of the array the `parser` is currently positioned at (`START_ARRAY`), + * serializing one element at a time straight from the source parser -- the enclosing array is + * never materialized as a whole. The iterator owns `parser`: it closes it on exhaustion (the + * fast path). To also close it when the consumer stops early (e.g. a downstream `LIMIT`, or a + * per-column cast failure) -- the `Generator` API has no close hook -- a *single* task-completion + * listener is registered per evaluator (i.e. per task) that closes whichever parser is currently + * open, rather than one listener per input row, so processing many JSON rows does not accumulate + * an unbounded listener list. `Generate` can thus emit rows for a large array without holding the + * full expanded payload in memory. + */ + private def arrayElementIterator(parser: JsonParser): Iterator[UTF8String] = { + openArrayParser = parser + if (!completionListenerRegistered) { + Option(TaskContext.get()).foreach { tc => + tc.addTaskCompletionListener[Unit] { _ => + val p = openArrayParser + if (p != null && !p.isClosed) p.close() + } + completionListenerRegistered = true + } + } + new Iterator[UTF8String] { + private var nextToken = parser.nextToken() + + override def hasNext: Boolean = { + val more = nextToken != null && nextToken != JsonToken.END_ARRAY + if (!more) close() + more + } + + override def next(): UTF8String = { + val element = serializeCurrentValue(parser) + nextToken = parser.nextToken() + element + } + + // Close the parser and drop the evaluator's reference to it so the listener does not retain + // it (and does not double-close) after this iterator is exhausted. + private def close(): Unit = { + if (!parser.isClosed) parser.close() + if (openArrayParser eq parser) openArrayParser = null + } + } + } + + /** + * If `raw` is a JSON string literal (e.g. `"hi"`), returns its unquoted, unescaped value; + * otherwise (numbers, booleans, objects, arrays) returns the raw JSON text unchanged. Used to + * give a value column the string's content rather than its quoted JSON form. + * + * `raw` is a Jackson-serialized fragment with no leading whitespace, so only a fragment whose + * first byte is `"` can be a string literal. Non-string values (the common case) are returned + * without constructing a parser at all. + */ + def unquotedString(raw: UTF8String): UTF8String = { + if (raw.numBytes() == 0 || raw.getByte(0) != '"') return raw + try { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, raw)) { parser => + if (parser.nextToken() == JsonToken.VALUE_STRING) { + UTF8String.fromString(parser.getText) + } else { + raw + } + } + } catch { + case _: JsonProcessingException => raw + } + } + + /** + * Builds the prefix trie that lets [[navigateColumns]] resolve every column path in one pass. The + * `paths` are indexed by result slot; a slot whose `include` is false (an ordinality column, + * which has no JSON path) contributes nothing to the trie -- note this is distinct from a root + * path `$`, which is an *empty but included* path that must resolve to the whole item. Call once + * per `JSON_TABLE` invocation and reuse the result for every row. + */ + def buildPathTrie( + paths: Array[Seq[PathInstruction]], + include: Array[Boolean]): JsonTablePathTrie = { + val root = new JsonTablePathTrie + var slot = 0 + while (slot < paths.length) { + if (include(slot)) insertPath(root, paths(slot), slot) + slot += 1 + } + root + } + + private def insertPath(root: JsonTablePathTrie, path: Seq[PathInstruction], slot: Int): Unit = { + var node = root + var rest = path + var valid = true + while (rest.nonEmpty && valid) { + rest match { + case Key :: Named(name) :: tail => + node = node.named.getOrElseUpdate(name, new JsonTablePathTrie) + rest = tail + case Subscript :: Index(index) :: tail => + node = node.indexed.getOrElseUpdate(index, new JsonTablePathTrie) + rest = tail + case _ => + // Should not happen: JSON_TABLE column paths are validated to be simple and + // wildcard-free. Drop the slot rather than mis-resolve it (it will read as Missing). + valid = false + } + } + if (valid) node.terminals ::= slot + } + + /** + * Resolves every column path (as built by [[buildPathTrie]]) within a single row item in one + * traversal, returning the per-slot results. Slots for ordinality columns (empty paths) are not + * in the trie and stay `Missing`; the caller fills them directly. Used to extract value and + * EXISTS columns with correct missing-vs-null semantics. + */ + def navigateColumns(item: UTF8String, trie: JsonTablePathTrie, numColumns: Int) + : Array[JsonPathResult] = { + val out = Array.fill[JsonPathResult](numColumns)(JsonPathResult.Missing) + // No path columns (e.g. an ordinality-only table): skip parsing the item entirely. + if (trie.isEmpty) return out + try { + Utils.tryWithResource(CreateJacksonParser.utf8String(jsonFactory, item)) { parser => + parser.nextToken() + navigateAll(parser, trie, out) + } + } catch { + // A malformed item leaves already-resolved slots in place and the rest as Missing. + case _: JsonProcessingException => + } + out + } +} + /** * The expression `GetJsonObject` will utilize it to support codegen. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala index d8131c3afb422..a5f07d8eaa036 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala @@ -17,20 +17,22 @@ package org.apache.spark.sql.catalyst.expressions +import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TypeCheckResult import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, CodegenFallback, ExprCode} import org.apache.spark.sql.catalyst.expressions.codegen.Block.BlockHelper import org.apache.spark.sql.catalyst.expressions.json.{GetJsonObjectEvaluator, JsonExpressionUtils, - JsonPathParser, JsonToStructsEvaluator, JsonTupleEvaluator, MultiGetJsonObjectEvaluator, + JsonPathParser, JsonPathResult, JsonQueryLookup, JsonTableEvaluator, JsonTablePathTrie, + JsonToStructsEvaluator, JsonTupleEvaluator, JsonValueLookup, MultiGetJsonObjectEvaluator, PathInstruction, SchemaOfJsonEvaluator, StructsToJsonEvaluator} import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.json._ import org.apache.spark.sql.catalyst.trees.TreePattern.{GET_JSON_OBJECT, JSON_TO_STRUCT, RUNTIME_REPLACEABLE, TreePattern} import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap -import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.types.StringTypeWithCollation import org.apache.spark.sql.types._ @@ -82,6 +84,7 @@ case class GetJsonObject(json: Expression, path: Expression) } else { new GetJsonObjectEvaluator() } + override def stateful: Boolean = true override def eval(input: InternalRow): Any = { evaluator.setJson(json.eval(input).asInstanceOf[UTF8String]) @@ -213,6 +216,8 @@ case class MultiGetJsonObject( } } + override def stateful: Boolean = true + @transient private lazy val evaluator = MultiGetJsonObjectEvaluator( fallbackPaths.map(UTF8String.fromString), @@ -244,6 +249,12 @@ case class MultiGetJsonObject( // scalastyle:off line.size.limit line.contains.tab @ExpressionDescription( usage = "_FUNC_(jsonStr, p1, p2, ..., pn) - Returns a tuple like the function get_json_object, but it takes multiple names. All the input parameters and output column types are string.", + arguments = """ + Arguments: + * jsonStr - A JSON string to extract fields from. + * pN - The field names to extract. Each name yields one output column with + the corresponding field value. + """, examples = """ Examples: > SELECT _FUNC_('{"a":1, "b":2}', 'a', 'b'); @@ -275,8 +286,14 @@ case class JsonTuple(children: Seq[Expression]) }.toArray } + // The extracted fields are values from inside the JSON document, so they do not carry the + // CHAR(n)/VARCHAR(n) length of the document itself. ImplicitTypeCoercion promotes CHAR/VARCHAR + // children to STRING without applying general implicit casts or rewriting untyped NULL. + private lazy val fieldType: DataType = + children.head.dataType + override def elementSchema: StructType = StructType(fieldExpressions.zipWithIndex.map { - case (_, idx) => StructField(s"c$idx", children.head.dataType, nullable = true) + case (_, idx) => StructField(s"c$idx", fieldType, nullable = true) }) override def prettyName: String = "json_tuple" @@ -300,6 +317,7 @@ case class JsonTuple(children: Seq[Expression]) @transient private lazy val evaluator: JsonTupleEvaluator = JsonTupleEvaluator(foldableFieldNames) + override def stateful: Boolean = true override def eval(input: InternalRow): IterableOnce[InternalRow] = { val json = jsonExpr.eval(input).asInstanceOf[UTF8String] @@ -339,6 +357,850 @@ case class JsonTuple(children: Seq[Expression]) copy(children = newChildren) } +/** + * The kind of a single `JSON_TABLE` column. + */ +sealed trait JsonTableColumnKind +object JsonTableColumnKind { + /** A `FOR ORDINALITY` column: a 1-based sequential row counter. */ + case object Ordinality extends JsonTableColumnKind + /** A regular value column: extracts the value at `path` and casts it to `dataType`. */ + case object Value extends JsonTableColumnKind + /** An `EXISTS` column: true when `path` matches, cast to `dataType`. */ + case object Exists extends JsonTableColumnKind +} + +/** + * A single column definition of a `JSON_TABLE` invocation. + * + * @param name the output column name + * @param dataType the declared Spark type of the column (LongType for ORDINALITY columns) + * @param path the SQL/JSON path relative to a row item; None for ORDINALITY columns + * @param kind the column kind (ordinality / value / exists) + */ +case class JsonTableColumn( + name: String, + dataType: DataType, + path: Option[String], + kind: JsonTableColumnKind) + +/** + * Behavior when the JSON input is malformed. + */ +sealed trait JsonTableErrorMode +object JsonTableErrorMode { + /** Produce no rows on malformed input (the SQL-standard default). */ + case object NullOnError extends JsonTableErrorMode + /** Raise an error on malformed input. */ + case object ErrorOnError extends JsonTableErrorMode +} + +// scalastyle:off line.size.limit +/** + * The SQL:2016 `JSON_TABLE` table-valued function. Shreds a JSON document into a relational table: + * the `rowPath` selects a sequence of row items and each [[JsonTableColumn]] projects a value out + * of each item. Implemented as a [[Generator]] so it plugs into the existing, well-tested + * [[org.apache.spark.sql.catalyst.plans.logical.Generate]] operator; no new execution operator is + * introduced. + * + * Only the flat (non-`NESTED PATH`) subset of the standard is supported. Row-source and value + * extraction use the token-aware [[JsonTableEvaluator]], which (unlike `get_json_object`) + * distinguishes a missing path from a JSON `null` value; type coercion reuses [[Cast]]. + * + * {{{ + * SELECT t.* FROM json_table( + * '{"items":[{"id":1,"n":"a"},{"id":2,"n":"b"}]}', + * '$.items[*]' + * COLUMNS (seq FOR ORDINALITY, id INT PATH '$.id', name STRING PATH '$.n') + * ) AS t; + * }}} + */ +// scalastyle:on line.size.limit +case class JsonTable( + child: Expression, + rowPath: String, + columns: Seq[JsonTableColumn], + errorMode: JsonTableErrorMode, + timeZoneId: Option[String] = None, + // Captured at plan-construction time so column casts do not change behavior if the session's + // ANSI mode is flipped between building the plan and executing it (matching `Cast`, which + // fixes its eval mode when the expression is constructed). + ansiEnabled: Boolean = SQLConf.get.ansiEnabled) + extends UnaryExpression + with Generator + with TimeZoneAwareExpression + with CodegenFallback + with ImplicitCastInputTypes + with QueryErrorsBase { + + // Declared via ImplicitCastInputTypes so the analyzer coerces the JSON input to STRING. In + // particular an untyped SQL NULL (NullType) is cast to STRING rather than rejected, so + // `JSON_TABLE(NULL, ...)` reaches the runtime and applies the NULL ON ERROR behavior. + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + + // ORDINALITY columns always hold a non-null counter; value/EXISTS columns may be null. + override def elementSchema: StructType = + StructType(columns.map { c => + val nullable = c.kind != JsonTableColumnKind.Ordinality + StructField(c.name, c.dataType, nullable = nullable) + }) + + override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = + copy(timeZoneId = Option(timeZoneId)) + + override def checkInputDataTypes(): TypeCheckResult = { + // First the standard input-type check (STRING for the JSON input, with NULL coerced). + val inputCheck = super.checkInputDataTypes() + if (inputCheck.isFailure) { + inputCheck + } else { + // Validate the row path and every column path. A path is valid here iff it parses and is + // free of wildcards -- except the row path may end in a single `[*]`, which is stripped into + // `containerInstructions`, so the row path is checked on that already-stripped list. + val rowPathValid = JsonPathParser.parse(rowPath).isDefined && + !containerInstructions.contains(PathInstruction.Wildcard) + val invalid: Option[(String, String)] = if (!rowPathValid) { + Some(("row path", rowPath)) + } else { + columns.iterator.collect { case c if c.path.isDefined => (c.name, c.path.get) } + .collectFirst { + // Valid column path: parses and is wildcard-free, i.e. hasWildcard == Some(false). + case (name, path) if !JsonPathParser.hasWildcard(path).contains(false) => + (s"column '$name'", path) + } + } + invalid match { + case Some((location, path)) => + DataTypeMismatch( + errorSubClass = "INVALID_JSON_TABLE_PATH", + messageParameters = Map("location" -> location, "path" -> toSQLValue(path))) + case None => + // Every value/EXISTS column is produced by casting from a source type (StringType for + // value columns, BooleanType for EXISTS columns) to the declared column type. Reject a + // non-castable declared type (e.g. a value column declared STRUCT/ARRAY/MAP) here rather + // than failing at runtime. Ordinality columns are always LongType and need no check. + // The castability rules differ between ANSI and non-ANSI mode (e.g. BOOLEAN -> TIMESTAMP + // is allowed by non-ANSI casts but not ANSI casts), so this must use the same eval mode + // as the actual per-column `Cast` built in `columnCasts`. + def sourceType(c: JsonTableColumn): Option[DataType] = c.kind match { + case JsonTableColumnKind.Value => Some(StringType) + case JsonTableColumnKind.Exists => Some(BooleanType) + case JsonTableColumnKind.Ordinality => None + } + def castable(src: DataType, target: DataType): Boolean = + if (ansiEnabled) Cast.canAnsiCast(src, target) else Cast.canCast(src, target) + columns.iterator.flatMap(c => sourceType(c).map((c, _))) + .collectFirst { case (c, src) if !castable(src, c.dataType) => (c, src) } match { + case Some((c, srcType)) => + DataTypeMismatch( + errorSubClass = "CAST_WITHOUT_SUGGESTION", + messageParameters = Map( + "srcType" -> toSQLType(srcType), + "targetType" -> toSQLType(c.dataType))) + case None => + TypeCheckResult.TypeCheckSuccess + } + } + } + } + + // The row path is `containerRowPath` plus an optional trailing `[*]`. Splitting on the parsed + // instruction list (rather than the raw string) is whitespace-insensitive and unambiguous. + // `checkInputDataTypes` guarantees the path parses and is wildcard-free at this point. + @transient private lazy val (containerInstructions, explodeRoot) + : (Seq[PathInstruction], Boolean) = { + val parsed = JsonPathParser.parse(rowPath).getOrElse(Nil) + parsed match { + case init :+ PathInstruction.Subscript :+ PathInstruction.Wildcard => + (init, true) + case other => + (other, false) + } + } + + @transient private lazy val rowEvaluator: JsonTableEvaluator = + JsonTableEvaluator(containerInstructions, explodeRoot) + + // Parsed instruction list per column (empty for ordinality columns, which have no path). + @transient private lazy val columnPaths: Array[Seq[PathInstruction]] = + columns.map(c => c.path.flatMap(JsonPathParser.parse).getOrElse(Nil)).toArray + + // Column kinds snapshotted into an array, like `columnPaths` and `columnCasts`: `columns` is a + // `List` (the parser builds it with `.map(...).toSeq`), so `columns(i)` is O(i) and indexing it + // in the per-row projection loop would make `projectRow` O(n^2) in the column count. + @transient private lazy val columnKinds: Array[JsonTableColumnKind] = columns.map(_.kind).toArray + + // Prefix trie over the column paths, built once so every row's value/EXISTS columns are resolved + // in a single traversal of the item rather than one re-parse per column. Ordinality columns have + // no path and are excluded (their empty path must not be confused with a root path `$`, which is + // an included column reading the whole item). + @transient private lazy val columnTrie: JsonTablePathTrie = { + val include = columnKinds.map(_ != JsonTableColumnKind.Ordinality) + rowEvaluator.buildPathTrie(columnPaths, include) + } + + // One reusable Cast per non-ordinality column, evaluated against a single-slot mutable input + // row. Building the Cast once (over a BoundReference) avoids allocating an expression tree per + // row/column on the hot path. The source type is BooleanType for EXISTS, StringType otherwise. + @transient private lazy val columnCasts: Array[Expression] = { + val evalMode = EvalMode.fromBoolean(ansiEnabled) + columns.map { c => + c.kind match { + case JsonTableColumnKind.Ordinality => null + case JsonTableColumnKind.Exists => + Cast(BoundReference(0, BooleanType, nullable = false), c.dataType, timeZoneId, evalMode) + case JsonTableColumnKind.Value => + Cast(BoundReference(0, StringType, nullable = true), c.dataType, timeZoneId, evalMode) + } + }.toArray + } + + // Reusable single-slot input row for the per-column casts above. + @transient private lazy val castInput: GenericInternalRow = new GenericInternalRow(1) + + private def castColumn(i: Int, value: Any): Any = { + castInput.update(0, value) + columnCasts(i).eval(castInput) + } + + private def projectRow(item: UTF8String, ordinal: Long): InternalRow = { + val numColumns = columnKinds.length + // Resolve every value/EXISTS column in a single traversal of the item; ordinality slots are + // not in the trie and come back as Missing (filled below). + val resolved = rowEvaluator.navigateColumns(item, columnTrie, numColumns) + val values = new Array[Any](numColumns) + var i = 0 + while (i < numColumns) { + values(i) = columnKinds(i) match { + case JsonTableColumnKind.Ordinality => + ordinal + case JsonTableColumnKind.Exists => + // Present (including an explicit JSON null) counts as existing; only Missing is false. + val exists = resolved(i) != JsonPathResult.Missing + castColumn(i, exists) + case JsonTableColumnKind.Value => + resolved(i) match { + // `raw` is a re-parseable JSON fragment; unquote a scalar string so the column gets + // its content (e.g. `"hi"` -> `hi`), then cast to the declared type. + case JsonPathResult.Found(raw) => castColumn(i, rowEvaluator.unquotedString(raw)) + // A missing path and an explicit JSON null both yield SQL NULL for a value column. + case _ => null + } + } + i += 1 + } + new GenericInternalRow(values) + } + + override def eval(input: InternalRow): IterableOnce[InternalRow] = { + val json = child.eval(input).asInstanceOf[UTF8String] + rowEvaluator.evaluate(json) match { + case Some(items) => + // A manual Long counter for FOR ORDINALITY: `zipWithIndex` is Int-based and would wrap + // past Int.MaxValue for a very large array, whereas ordinality is a BIGINT. + var ordinal = 0L + items.map { item => + ordinal += 1L + projectRow(item, ordinal) + } + case None => + // `errorMode` governs the row-source JSON only (null / malformed input, or `[*]` over a + // non-array). Per-column value extraction follows normal `Cast` semantics: a bad cast is + // NULL in non-ANSI mode and raises in ANSI mode, independent of ON ERROR. + errorMode match { + case JsonTableErrorMode.NullOnError => Iterator.empty + case JsonTableErrorMode.ErrorOnError => + throw QueryExecutionErrors.malformedRecordsDetectedInRecordParsingError( + if (json == null) "null" else json.toString, + SparkException.internalError("JSON_TABLE encountered malformed JSON input.")) + } + } + } + + override def prettyName: String = "json_table" + + // The default `Expression.sql` renders only children, i.e. `json_table(<json_expr>)`, dropping + // the row path, columns, and ON ERROR mode. Render the full `JSON_TABLE(...)` syntax so + // analysis/type-check diagnostics (e.g. INVALID_JSON_TABLE_PATH) point at the whole invocation. + override def sql: String = { + val columnsSQL = columns.map { c => + val pathSQL = c.path.map(p => s" PATH '$p'").getOrElse("") + c.kind match { + case JsonTableColumnKind.Ordinality => s"${c.name} FOR ORDINALITY" + case JsonTableColumnKind.Exists => s"${c.name} ${c.dataType.sql} EXISTS$pathSQL" + case JsonTableColumnKind.Value => s"${c.name} ${c.dataType.sql}$pathSQL" + } + }.mkString(", ") + val errorSQL = errorMode match { + case JsonTableErrorMode.NullOnError => "NULL ON ERROR" + case JsonTableErrorMode.ErrorOnError => "ERROR ON ERROR" + } + s"JSON_TABLE(${child.sql}, '$rowPath' COLUMNS ($columnsSQL) $errorSQL)" + } + + override protected def withNewChildInternal(newChild: Expression): JsonTable = + copy(child = newChild) +} + +/** + * Behavior of `JSON_VALUE`'s `ON EMPTY` / `ON ERROR` clause: what to produce when the path matches + * nothing, or when the input/extraction fails. + */ +sealed trait JsonValueBehavior +object JsonValueBehavior { + /** Produce SQL NULL (the SQL-standard default for both ON EMPTY and ON ERROR). */ + case object Null extends JsonValueBehavior + /** Raise an error. */ + case object Error extends JsonValueBehavior + /** Produce the value of a `DEFAULT` expression, cast to the RETURNING type. */ + case object Default extends JsonValueBehavior +} + +// scalastyle:off line.size.limit +/** + * The SQL:2016 `JSON_VALUE` scalar function (feature T821): extracts a single scalar located by a + * SQL/JSON `path` from a JSON input, casts it to the `RETURNING` type (default STRING), and applies + * the `ON EMPTY` / `ON ERROR` behavior when the path matches nothing or the extraction/cast fails: + * + * - missing path -> ON EMPTY behavior + * - explicit JSON `null` -> SQL NULL + * - non-scalar (object/array) match -> ON ERROR behavior + * - malformed / non-single-value input -> ON ERROR behavior + * - scalar match, cast fails -> ON ERROR behavior + * - scalar match, cast succeeds -> the cast value + * + * Both clauses default to NULL per the standard. A `null` JSON input yields SQL NULL directly, not + * the ON EMPTY/ERROR path. + * + * `emptyDefault` / `errorDefault` hold the `DEFAULT <expr>` expressions, present only for the + * corresponding `Default` behavior. The child list is variable (0-2 defaults), so this extends + * `Expression` directly rather than `UnaryExpression`. + * + * {{{ + * JSON_VALUE('{"id":7}', '$.id' RETURNING INT) -- 7 + * JSON_VALUE('{"id":7}', '$.missing' DEFAULT -1 ON EMPTY) -- -1 + * JSON_VALUE('{"a":{}}', '$.a' ERROR ON ERROR) -- raises (non-scalar) + * }}} + */ +// scalastyle:on line.size.limit +case class JsonValue( + child: Expression, + path: String, + returning: DataType, + onEmpty: JsonValueBehavior, + onError: JsonValueBehavior, + emptyDefault: Option[Expression], + errorDefault: Option[Expression], + timeZoneId: Option[String] = None, + ansiEnabled: Boolean = SQLConf.get.ansiEnabled) + extends Expression + with TimeZoneAwareExpression + with CodegenFallback + with ExpectsInputTypes + with QueryErrorsBase { + + override def nullable: Boolean = true + + // Reuses the mutable `castInput` row across rows (see `castScalar`), so it holds evaluation + // state. Interpreted execution must fresh-copy the expression before use, or a shared instance + // could cast another concurrent evaluation's value; matches the neighboring JSON expressions. + override def stateful: Boolean = true + + // Children: the JSON input first, then whichever DEFAULT expressions are present. The two + // defaults are resolved/coerced through the normal child machinery; their cast to `returning` + // happens at eval time via `emptyDefaultCast` / `errorDefaultCast`. + override def children: Seq[Expression] = + child +: (emptyDefault.toSeq ++ errorDefault.toSeq) + + // One entry per child: the JSON input must be STRING; the DEFAULT children accept anything (they + // are cast to `returning` explicitly at eval). One entry per child is required because the + // coercion rule zips `children` against `inputTypes` and rebuilds via `withNewChildren`; a + // shorter list would truncate the zip and pass the wrong child count. + override def inputTypes: Seq[AbstractDataType] = + StringTypeWithCollation(supportsTrimCollation = true) +: + children.tail.map(_ => AnyDataType) + + override def dataType: DataType = returning + + override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = + copy(timeZoneId = Option(timeZoneId)) + + override def checkInputDataTypes(): TypeCheckResult = { + val inputCheck = super.checkInputDataTypes() + if (inputCheck.isFailure) { + inputCheck + } else if (!JsonPathParser.hasWildcard(path).contains(false)) { + // The path must parse and be wildcard-free (JSON_VALUE returns a single scalar). The + // `INVALID_JSON_PATH` message is shared with `JSON_EXISTS`, which does accept wildcards, so + // it must stay generic -- do not re-add wildcard-specific wording here. + DataTypeMismatch( + errorSubClass = "INVALID_JSON_PATH", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "path" -> toSQLValue(path))) + } else if (!JsonValue.isValidReturningType(returning)) { + // RETURNING is restricted to scalar (atomic) types per ANSI 9075-2 6.28. + DataTypeMismatch( + errorSubClass = "INVALID_JSON_SCALAR_RETURNING_TYPE", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "returningType" -> toSQLType(returning))) + } else { + // Each DEFAULT expression is cast to the RETURNING type at eval time (see + // `emptyDefaultCast` / `errorDefaultCast`). Those casts are not analyzed children, so an + // uncastable default (e.g. `DEFAULT array(1)` with `RETURNING INT`) would otherwise slip + // past analysis and fail late only when its branch is taken. Surface the cast's own type + // check here so it is rejected up front with the standard CAST_* message. + (emptyDefaultCast ++ errorDefaultCast) + .map(_.checkInputDataTypes()) + .find(_.isFailure) + .getOrElse(TypeCheckResult.TypeCheckSuccess) + } + } + + // Eval mode for the user-provided DEFAULT expression casts: follows the session ANSI setting like + // any ordinary value cast. The extracted-scalar cast is separate (see `valueCast`). + @transient private lazy val defaultEvalMode = EvalMode.fromBoolean(ansiEnabled) + + // Path parsed once (the grammar makes it a string literal). `checkInputDataTypes` guarantees it + // parses and is wildcard-free, so the evaluator is only built for a valid path. + @transient private lazy val evaluator: JsonTableEvaluator = + JsonTableEvaluator(JsonPathParser.parse(path).getOrElse(Nil), explodeRoot = false) + + // Cast from the extracted scalar's STRING form to the RETURNING type, built once over a reused + // input slot to avoid per-row allocation. Always an ANSI (throwing) cast, independent of the + // session's ANSI setting, so a failed conversion always routes to ON ERROR (see `eval`) rather + // than being silently turned into NULL by a non-ANSI session. + @transient private lazy val valueCast: Expression = + Cast(BoundReference(0, StringType, nullable = true), returning, timeZoneId, EvalMode.ANSI) + @transient private lazy val castInput: GenericInternalRow = new GenericInternalRow(1) + + // Casts for the DEFAULT expressions to the RETURNING type (only built when present). + @transient private lazy val emptyDefaultCast: Option[Expression] = + emptyDefault.map(e => Cast(e, returning, timeZoneId, defaultEvalMode)) + @transient private lazy val errorDefaultCast: Option[Expression] = + errorDefault.map(e => Cast(e, returning, timeZoneId, defaultEvalMode)) + + private def castScalar(text: UTF8String): Any = { + castInput.update(0, text) + valueCast.eval(castInput) + } + + // Handle the ON EMPTY case per the configured behavior. + private def onEmptyResult(input: InternalRow): Any = onEmpty match { + case JsonValueBehavior.Null => null + case JsonValueBehavior.Default => emptyDefaultCast.get.eval(input) + case JsonValueBehavior.Error => + throw QueryExecutionErrors.jsonValueOnEmptyError(prettyName, path, cause = null) + } + + // Handle the ON ERROR case per the configured behavior. `cause` (if any) is attached for context. + private def onErrorResult(input: InternalRow, cause: Throwable): Any = onError match { + case JsonValueBehavior.Null => null + case JsonValueBehavior.Default => errorDefaultCast.get.eval(input) + case JsonValueBehavior.Error => + throw QueryExecutionErrors.jsonValueOnErrorError(prettyName, path, cause) + } + + override def eval(input: InternalRow): Any = { + val json = child.eval(input).asInstanceOf[UTF8String] + // NULL input propagates to NULL (not ON EMPTY / ON ERROR), matching ANSI and the other engines. + if (json == null) return null + evaluator.lookup(json) match { + // Malformed / non-single-value input. + case None => onErrorResult(input, cause = null) + // Path matched nothing. + case Some(JsonValueLookup.Missing) => onEmptyResult(input) + // Matched an explicit JSON null: a present, scalar null value -> SQL NULL. + case Some(JsonValueLookup.NullValue) => null + // Matched an object or array: not a scalar -> ON ERROR. + case Some(JsonValueLookup.NonScalar) => onErrorResult(input, cause = null) + case Some(JsonValueLookup.Scalar(text)) => + // `valueCast` throws on a failed conversion, which routes to ON ERROR. + try castScalar(text) catch { case e: Exception => onErrorResult(input, e) } + } + } + + override def prettyName: String = "json_value" + + override def sql: String = { + val returningSQL = if (returning == StringType) "" else s" RETURNING ${returning.sql}" + def behaviorSQL(b: JsonValueBehavior, default: Option[Expression]): String = b match { + case JsonValueBehavior.Null => "NULL" + case JsonValueBehavior.Error => "ERROR" + case JsonValueBehavior.Default => s"DEFAULT ${default.get.sql}" + } + val emptySQL = if (onEmpty == JsonValueBehavior.Null) "" + else s" ${behaviorSQL(onEmpty, emptyDefault)} ON EMPTY" + val errorSQL = if (onError == JsonValueBehavior.Null) "" + else s" ${behaviorSQL(onError, errorDefault)} ON ERROR" + // Render the path as a properly escaped string literal so bracket-quoted paths such as + // `$['a']` (and any path containing a quote or backslash) round-trip as valid SQL. + val pathSQL = Literal(UTF8String.fromString(path), StringType).sql + s"JSON_VALUE(${child.sql}, $pathSQL$returningSQL$emptySQL$errorSQL)" + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): JsonValue = { + // Rebuild the child list in the same order `children` produced it: json, then the present + // defaults. `copy(child = ...)` alone would drop coercion applied to the DEFAULT children. + var i = 1 + val newEmpty = emptyDefault.map { _ => val e = newChildren(i); i += 1; e } + val newError = errorDefault.map { _ => val e = newChildren(i); i += 1; e } + copy(child = newChildren(0), emptyDefault = newEmpty, errorDefault = newError) + } +} + +object JsonValue { + /** + * ANSI (9075-2 6.28) restricts JSON_VALUE RETURNING to predefined scalar types: string, numeric, + * boolean, and datetime. We allow exactly those families. Note this deliberately excludes VARIANT + * (a Spark extension, deferred per the design's open question) and BINARY, even though both are + * `AtomicType`s -- so an `AtomicType` check is not sufficient. STRUCT/ARRAY/MAP are excluded as + * non-atomic. CHAR/VARCHAR are normalized to STRING by the parser before reaching here. + */ + def isValidReturningType(dt: DataType): Boolean = dt match { + case _: StringType => true + case _: NumericType => true + case BooleanType => true + case _: DatetimeType => true + case _ => false + } +} + +/** + * Behavior of `JSON_EXISTS`'s `ON ERROR` clause: the value produced when the input is not a single + * well-formed JSON value (malformed / trailing garbage). `Unknown` is a BOOLEAN NULL. + */ +sealed trait JsonExistsBehavior +object JsonExistsBehavior { + case object True extends JsonExistsBehavior + case object False extends JsonExistsBehavior + case object Unknown extends JsonExistsBehavior + case object Error extends JsonExistsBehavior +} + +/** + * The SQL:2016 `JSON_EXISTS` predicate (feature T821): returns whether a SQL/JSON `path` matches at + * least one item in a JSON input. + * + * - path matches (including an explicit JSON `null`) -> true + * - path matches nothing -> false + * - malformed / non-single-value input -> ON ERROR behavior (default FALSE) + * - SQL NULL input -> SQL NULL (Unknown, per 9075-2 8.23) + * + * This distinguishes "present but null" from "absent" (unlike `get_json_object(...) IS NOT NULL`). + * The `ON ERROR` clause chooses TRUE / FALSE / UNKNOWN (a BOOLEAN NULL) / ERROR; it defaults to + * FALSE ON ERROR. + * + * Paths are evaluated in SQL/JSON lax mode (matching Oracle / PostgreSQL): wildcards are supported + * and arrays are auto-wrapped/unwrapped, while a structural mismatch is a non-match, not an error. + * + * {{{ + * JSON_EXISTS('{"a":{"b":1}}', '$.a.b') -- true + * JSON_EXISTS('{"a":null}', '$.a') -- true (present, value is null) + * JSON_EXISTS('{"a":1}', '$.b') -- false (absent) + * JSON_EXISTS('{"a":[1,2]}', '$.a[*]') -- true (array has elements) + * JSON_EXISTS('not json', '$.a' TRUE ON ERROR) -- true + * }}} + */ +case class JsonExists( + child: Expression, + path: String, + onError: JsonExistsBehavior) + extends UnaryExpression + with CodegenFallback + with ExpectsInputTypes + with QueryErrorsBase { + + // The result is NULL only when the input is SQL NULL, or when `UNKNOWN ON ERROR` turns malformed + // input into a BOOLEAN NULL. With a non-nullable input and any other ON ERROR behavior the result + // is always a concrete boolean, which lets the optimizer treat e.g. a WHERE predicate as such. + override def nullable: Boolean = + child.nullable || onError == JsonExistsBehavior.Unknown + + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + + override def dataType: DataType = BooleanType + + // The path is a constant (the grammar makes it a string literal), so it is parsed once and shared + // by `checkInputDataTypes` and `evaluator` rather than reparsed. `None` means it did not parse. + @transient private lazy val parsedPath: Option[Seq[PathInstruction]] = JsonPathParser.parse(path) + + override def checkInputDataTypes(): TypeCheckResult = { + val inputCheck = super.checkInputDataTypes() + if (inputCheck.isFailure) { + inputCheck + } else if (parsedPath.isDefined) { + // A valid SQL/JSON path. Wildcards are allowed and evaluated in lax mode at runtime. + TypeCheckResult.TypeCheckSuccess + } else { + // The path is not a valid SQL/JSON path. `JSON_EXISTS` reaches this branch only for a + // syntactically malformed path -- wildcards parse and are accepted above. The shared + // `INVALID_JSON_PATH` error is also raised by `JSON_VALUE` (which additionally rejects + // wildcards, as it returns a single scalar); its message is worded generically for both. + DataTypeMismatch( + errorSubClass = "INVALID_JSON_PATH", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "path" -> toSQLValue(path))) + } + } + + // `checkInputDataTypes` guarantees the path parses before this is forced; `Nil` is an unreachable + // fallback that would match the document root. + @transient private lazy val evaluator: JsonTableEvaluator = + JsonTableEvaluator(parsedPath.getOrElse(Nil), explodeRoot = false) + + private def onErrorResult(): Any = onError match { + case JsonExistsBehavior.True => true + case JsonExistsBehavior.False => false + case JsonExistsBehavior.Unknown => null + case JsonExistsBehavior.Error => + throw QueryExecutionErrors.jsonExistsOnError(prettyName, path) + } + + override def eval(input: InternalRow): Any = { + val json = child.eval(input).asInstanceOf[UTF8String] + // SQL NULL input yields Unknown (BOOLEAN NULL), not the ON ERROR path, per 9075-2 8.23. + if (json == null) return null + evaluator.pathExists(json) match { + case Some(exists) => exists + case None => onErrorResult() // malformed / non-single-value input + } + } + + override def prettyName: String = "json_exists" + + override def sql: String = { + val errorSQL = onError match { + case JsonExistsBehavior.False => "" // the default + case JsonExistsBehavior.True => " TRUE ON ERROR" + case JsonExistsBehavior.Unknown => " UNKNOWN ON ERROR" + case JsonExistsBehavior.Error => " ERROR ON ERROR" + } + s"JSON_EXISTS(${child.sql}, ${toSQLValue(path)}$errorSQL)" + } + + override protected def withNewChildInternal(newChild: Expression): JsonExists = + copy(child = newChild) +} + +/** + * Behavior of `JSON_QUERY`'s `ON EMPTY` / `ON ERROR` clause: what to produce when the path matches + * nothing (`ON EMPTY`) or the input is not valid JSON (`ON ERROR`). + */ +sealed trait JsonQueryBehavior +object JsonQueryBehavior { + /** Produce SQL NULL (the SQL-standard default for both clauses). */ + case object Null extends JsonQueryBehavior + /** Raise an error. */ + case object Error extends JsonQueryBehavior + /** Produce an empty JSON array `[]`. */ + case object EmptyArray extends JsonQueryBehavior + /** Produce an empty JSON object `{}`. */ + case object EmptyObject extends JsonQueryBehavior +} + +/** + * The array-wrapper behavior of `JSON_QUERY` (SQL:2016 `... ARRAY WRAPPER`). This implementation + * resolves a single value per path (wildcard-free paths only), so a wrapper wraps that value in a + * one-element array: + * - `Without` (default): return the value unwrapped; + * - `Unconditional` (`WITH [UNCONDITIONAL] ARRAY WRAPPER`): always wrap; + * - `Conditional` (`WITH CONDITIONAL ARRAY WRAPPER`): wrap only a scalar; leave an object or + * array as is. + */ +sealed trait JsonQueryWrapper +object JsonQueryWrapper { + case object Without extends JsonQueryWrapper + case object Conditional extends JsonQueryWrapper + case object Unconditional extends JsonQueryWrapper +} + +/** The quotes behavior of `JSON_QUERY`: `KEEP QUOTES` (default) or `OMIT QUOTES`. */ +sealed trait JsonQueryQuotes +object JsonQueryQuotes { + case object Keep extends JsonQueryQuotes + case object Omit extends JsonQueryQuotes +} + +// scalastyle:off line.size.limit +/** + * The SQL:2016 `JSON_QUERY` function (feature T828): extracts the JSON value located by `path` from + * a JSON input and returns it as JSON text (STRING): + * + * - missing path -> ON EMPTY behavior + * - malformed / non-single-value input -> ON ERROR behavior + * - matched object / array / scalar -> its serialized JSON text, after applying the array + * wrapper and quotes clauses + * + * A matched scalar (including a JSON `null`) is not an error under the default `WITHOUT ARRAY + * WRAPPER`; it is emitted as JSON text (`JSON_QUERY('{"id":7}', '$.id')` -> `7`). `OMIT QUOTES` + * strips the surrounding quotes from a scalar string result (and cannot be combined with a wrapper). + * Both `ON EMPTY` and `ON ERROR` default to NULL per the standard, and a `null` JSON input yields + * SQL NULL directly. `RETURNING` is restricted to string types here (VARIANT is deferred); the + * result is always JSON text. + * + * {{{ + * JSON_QUERY('{"a":{"x":1}}', '$.a') -- '{"x":1}' + * JSON_QUERY('{"t":["x","y"]}', '$.t') -- '["x","y"]' + * JSON_QUERY('{"t":["x","y"]}', '$.t[0]' WITH ARRAY WRAPPER) -- '["x"]' + * JSON_QUERY('{"n":"Ada"}', '$.n' OMIT QUOTES) -- 'Ada' + * }}} + */ +// scalastyle:on line.size.limit +case class JsonQuery( + child: Expression, + path: String, + returning: DataType, + wrapper: JsonQueryWrapper, + quotes: JsonQueryQuotes, + onEmpty: JsonQueryBehavior, + onError: JsonQueryBehavior) + extends UnaryExpression + with CodegenFallback + with ExpectsInputTypes + with QueryErrorsBase { + + override def nullable: Boolean = true + + // The JSON input must be a STRING; the result is JSON text. + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + + override def dataType: DataType = returning + + override def checkInputDataTypes(): TypeCheckResult = { + val inputCheck = super.checkInputDataTypes() + if (inputCheck.isFailure) { + inputCheck + } else if (!JsonPathParser.hasWildcard(path).contains(false)) { + // The path must parse and be wildcard-free (a single value is resolved). + DataTypeMismatch( + errorSubClass = "INVALID_JSON_PATH", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "path" -> toSQLValue(path))) + } else if (!JsonQuery.isValidReturningType(returning)) { + // RETURNING is restricted to string types (the result is JSON text; VARIANT is deferred). + DataTypeMismatch( + errorSubClass = "INVALID_JSON_QUERY_RETURNING_TYPE", + messageParameters = Map( + "functionName" -> toSQLId(prettyName), "returningType" -> toSQLType(returning))) + } else if (quotes == JsonQueryQuotes.Omit && wrapper != JsonQueryWrapper.Without) { + // OMIT QUOTES applies only to an unwrapped scalar; the SQL standard forbids pairing it with + // an array wrapper. Enforced here (not only in the parser) so a directly-constructed + // expression cannot silently ignore the quotes clause. + DataTypeMismatch( + errorSubClass = "INVALID_JSON_QUERY_WRAPPER_AND_QUOTES", + messageParameters = Map("functionName" -> toSQLId(prettyName))) + } else { + TypeCheckResult.TypeCheckSuccess + } + } + + // Path parsed once (the grammar makes it a string literal). `checkInputDataTypes` guarantees it + // parses and is wildcard-free, so the evaluator is only built for a valid path. + @transient private lazy val evaluator: JsonTableEvaluator = + JsonTableEvaluator(JsonPathParser.parse(path).getOrElse(Nil), explodeRoot = false) + + // Handle the ON EMPTY / ON ERROR case per the configured behavior. + private def onEmptyResult(): Any = behaviorResult(onEmpty, isEmpty = true) + private def onErrorResult(): Any = behaviorResult(onError, isEmpty = false) + + private def behaviorResult(behavior: JsonQueryBehavior, isEmpty: Boolean): Any = behavior match { + case JsonQueryBehavior.Null => null + case JsonQueryBehavior.EmptyArray => JsonQuery.EmptyArrayText + case JsonQueryBehavior.EmptyObject => JsonQuery.EmptyObjectText + case JsonQueryBehavior.Error => + if (isEmpty) throw QueryExecutionErrors.jsonQueryOnEmptyError(prettyName, path, cause = null) + else throw QueryExecutionErrors.jsonQueryOnErrorError(prettyName, path, cause = null) + } + + // Apply the array-wrapper and quotes clauses to a matched value. `raw` is its serialized JSON + // text, `unquoted` is the OMIT QUOTES form (a string's decoded content; `raw` otherwise, so OMIT + // QUOTES is a no-op for objects, arrays, and non-string scalars), and `structural` is true for an + // object or array match (rather than a scalar, incl. JSON null). + private def wrapAndQuote(raw: UTF8String, unquoted: UTF8String, structural: Boolean): UTF8String = + wrapper match { + case JsonQueryWrapper.Without => + // OMIT QUOTES reuses the string decoded during the lookup rather than re-parsing the + // serialized fragment; OMIT QUOTES combined with a wrapper is rejected at analysis time. + if (quotes == JsonQueryQuotes.Omit) unquoted else raw + case JsonQueryWrapper.Unconditional => JsonQuery.wrapInArray(raw) + // CONDITIONAL wraps only a scalar; an object or array is already a structural result. + case JsonQueryWrapper.Conditional => if (structural) raw else JsonQuery.wrapInArray(raw) + } + + override def eval(input: InternalRow): Any = { + val json = child.eval(input).asInstanceOf[UTF8String] + // NULL input propagates to NULL (not ON EMPTY / ON ERROR), matching ANSI and the other engines. + if (json == null) return null + evaluator.queryLookup(json, omitQuotes = quotes == JsonQueryQuotes.Omit) match { + // Malformed / non-single-value input. + case None => onErrorResult() + // Path matched nothing. + case Some(JsonQueryLookup.Missing) => onEmptyResult() + // Matched a value: serialize it, applying the wrapper and quotes clauses. + case Some(JsonQueryLookup.Found(raw, structural, unquoted)) => + wrapAndQuote(raw, unquoted, structural) + } + } + + override def prettyName: String = "json_query" + + override def sql: String = { + val returningSQL = if (returning == StringType) "" else s" RETURNING ${returning.sql}" + val wrapperSQL = wrapper match { + case JsonQueryWrapper.Without => "" + case JsonQueryWrapper.Unconditional => " WITH UNCONDITIONAL ARRAY WRAPPER" + case JsonQueryWrapper.Conditional => " WITH CONDITIONAL ARRAY WRAPPER" + } + val quotesSQL = quotes match { + case JsonQueryQuotes.Keep => "" + case JsonQueryQuotes.Omit => " OMIT QUOTES" + } + def behaviorSQL(b: JsonQueryBehavior): String = b match { + case JsonQueryBehavior.Null => "NULL" + case JsonQueryBehavior.Error => "ERROR" + case JsonQueryBehavior.EmptyArray => "EMPTY ARRAY" + case JsonQueryBehavior.EmptyObject => "EMPTY OBJECT" + } + val emptySQL = + if (onEmpty == JsonQueryBehavior.Null) "" else s" ${behaviorSQL(onEmpty)} ON EMPTY" + val errorSQL = + if (onError == JsonQueryBehavior.Null) "" else s" ${behaviorSQL(onError)} ON ERROR" + // Render the path as a properly escaped string literal so bracket-quoted paths round-trip. + val pathSQL = Literal(UTF8String.fromString(path), StringType).sql + s"JSON_QUERY(${child.sql}, $pathSQL$returningSQL$wrapperSQL$quotesSQL$emptySQL$errorSQL)" + } + + override protected def withNewChildInternal(newChild: Expression): JsonQuery = + copy(child = newChild) +} + +object JsonQuery { + private val EmptyArrayText: UTF8String = UTF8String.fromString("[]") + private val EmptyObjectText: UTF8String = UTF8String.fromString("{}") + private val ArrayOpen: UTF8String = UTF8String.fromString("[") + private val ArrayClose: UTF8String = UTF8String.fromString("]") + + private def wrapInArray(raw: UTF8String): UTF8String = + UTF8String.concat(ArrayOpen, raw, ArrayClose) + + /** + * `JSON_QUERY` returns a JSON fragment as text, so RETURNING is restricted to a plain STRING here + * (VARIANT is deferred). `CharType` / `VarcharType` extend `StringType` but carry a length that + * `JSON_QUERY` does not enforce -- it returns the fragment verbatim without a cast -- so they are + * rejected: the parser normalizes a SQL `CHAR`/`VARCHAR` RETURNING to STRING before construction, + * and this guards a raw `CharType`/`VarcharType` supplied by direct Catalyst construction. + */ + def isValidReturningType(dt: DataType): Boolean = dt match { + case _: CharType | _: VarcharType => false + case _: StringType => true + case _ => false + } +} + /** * Converts an json input string to a [[StructType]], [[ArrayType]] or [[MapType]] * with the specified schema. @@ -346,6 +1208,14 @@ case class JsonTuple(children: Seq[Expression]) // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(jsonStr, schema[, options]) - Returns a struct value with the given `jsonStr` and `schema`.", + arguments = """ + Arguments: + * jsonStr - A JSON string to parse. + * schema - The schema to use when parsing the JSON string, given as a DDL + formatted string or a schema expression. + * options - Optional. A map of string key-value pairs that control how the + JSON is parsed. By default no options are set. + """, examples = """ Examples: > SELECT _FUNC_('{"a":1, "b":0.8}', 'a INT, b DOUBLE'); @@ -427,6 +1297,7 @@ case class JsonToStructs( @transient private lazy val evaluator = new JsonToStructsEvaluator( options, nullableSchema, nameOfCorruptRecord, timeZoneId, variantAllowDuplicateKeys) + override def stateful: Boolean = true override def nullSafeEval(json: Any): Any = evaluator.evaluate(json.asInstanceOf[UTF8String]) @@ -549,6 +1420,12 @@ case class StructsToJson( */ @ExpressionDescription( usage = "_FUNC_(json[, options]) - Returns schema in the DDL format of JSON string.", + arguments = """ + Arguments: + * json - A JSON string whose schema is inferred. + * options - Optional. A map of string key-value pairs that control how the + JSON is parsed. By default no options are set. + """, examples = """ Examples: > SELECT _FUNC_('[{"col":0}]'); @@ -709,3 +1586,49 @@ case class JsonObjectKeys(child: Expression) override protected def withNewChildInternal(newChild: Expression): JsonObjectKeys = copy(child = newChild) } + +/** + * A function which returns the type of the outermost JSON value as a string. + */ +@ExpressionDescription( + usage = "_FUNC_(json) - Returns the type of the outermost JSON value, or null if invalid.", + arguments = """ + Arguments: + * json - A JSON string. Returns the type of the outermost value ('object', 'array', + 'string', 'number', 'boolean', 'null'), or null for an invalid or empty string. + An expression that evaluates to a string. + """, + examples = """ + Examples: + > SELECT _FUNC_('{"a": 1}'); + object + > SELECT _FUNC_('[1, 2, 3]'); + array + > SELECT _FUNC_('123'); + number + """, + group = "json_funcs", + since = "4.4.0" +) +case class JsonTypeof(child: Expression) + extends UnaryExpression + with ExpectsInputTypes + with RuntimeReplaceable + with DefaultStringProducingExpression { + + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + override def nullable: Boolean = true + override def prettyName: String = "json_typeof" + + override def replacement: Expression = StaticInvoke( + classOf[JsonExpressionUtils], + dataType, + "jsonTypeof", + Seq(child), + inputTypes + ) + + override protected def withNewChildInternal(newChild: Expression): JsonTypeof = + copy(child = newChild) +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/kllExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/kllExpressions.scala index ba6c68ab95fcc..218146155837b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/kllExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/kllExpressions.scala @@ -564,7 +564,7 @@ abstract class KllSketchGetQuantileBase override def dataType: DataType = { right.dataType match { - case ArrayType(_, _) => ArrayType(outputDataType, false) + case ArrayType(_, _) => ArrayType(outputDataType, containsNull = false) case _ => outputDataType } } @@ -750,7 +750,7 @@ abstract class KllSketchGetRankBase } override def dataType: DataType = { right.dataType match { - case ArrayType(_, _) => ArrayType(DoubleType, false) + case ArrayType(_, _) => ArrayType(DoubleType, containsNull = false) case _ => DoubleType } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala index 546d82546d034..b0cdbcc321c1d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala @@ -171,7 +171,7 @@ object Literal { case _: DayTimeIntervalType if v.isInstanceOf[Duration] => Literal(CatalystTypeConverters.createToCatalystConverter(dataType)(v), dataType) case _: ObjectType => Literal(v, dataType) - case _: CharType | _: VarcharType if SQLConf.get.preserveCharVarcharTypeInfo => + case _: CharType | _: VarcharType if SQLConf.get.charVarcharFirstClassTypes => Literal(CatalystTypeConverters.createToCatalystConverter(dataType)(v), dataType) case _ if requiresSchemaAwareNanosConversion(dataType, v) => Literal(CatalystTypeConverters.createToCatalystConverter(dataType)(v), dataType) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala index 9613bc25d65f0..9d6f2c037df6b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/maskExpressions.scala @@ -290,7 +290,8 @@ case class Mask( * Returns the [[DataType]] of the result of evaluating this expression. It is invalid to query * the dataType of an unresolved expression (i.e., when `resolved` == false). */ - override def dataType: DataType = input.dataType + override def dataType: DataType = + input.dataType /** * Returns a Seq of the children of this node. Children should not change. Immutability required @@ -319,6 +320,7 @@ object Mask { val MASKED_DIGIT = 'n' // This value helps to retain original value in the input by ignoring the replacement rules val MASKED_IGNORE = null + private val IGNORE_CODE_POINT = -1 def transformInput( input: Any, @@ -330,29 +332,47 @@ object Mask { val transformedString = if (input == null) { null } else { - input.toString.map { - transformChar(_, maskUpper, maskLower, maskDigit, maskOther).toChar + val upper = replacementCodePoint(maskUpper) + val lower = replacementCodePoint(maskLower) + val digit = replacementCodePoint(maskDigit) + val other = replacementCodePoint(maskOther) + val str = input.toString + val sb = new java.lang.StringBuilder(str.length) + var i = 0 + while (i < str.length) { + val codePoint = str.codePointAt(i) + sb.appendCodePoint(transformCodePoint(codePoint, upper, lower, digit, other)) + i += Character.charCount(codePoint) } + sb.toString } org.apache.spark.unsafe.types.UTF8String.fromString(transformedString) } - private def transformChar( - c: Char, - maskUpper: Any, - maskLower: Any, - maskDigit: Any, - maskOther: Any): Int = { + private def replacementCodePoint(option: Any): Int = { + if (option != MASKED_IGNORE) { + option.asInstanceOf[UTF8String].toString.codePointAt(0) + } else { + IGNORE_CODE_POINT + } + } + + private def transformCodePoint( + codePoint: Int, + maskUpper: Int, + maskLower: Int, + maskDigit: Int, + maskOther: Int): Int = { - def maskedChar(c: Char, option: Any): Char = { - if (option != MASKED_IGNORE) option.asInstanceOf[UTF8String].toString.charAt(0) else c + def maskedCodePoint(replacement: Int): Int = { + if (replacement != IGNORE_CODE_POINT) replacement else codePoint } - Character.getType(c) match { - case Character.UPPERCASE_LETTER => maskedChar(c, maskUpper) - case Character.LOWERCASE_LETTER => maskedChar(c, maskLower) - case Character.DECIMAL_DIGIT_NUMBER => maskedChar(c, maskDigit) - case _ => maskedChar(c, maskOther) + Character.getType(codePoint) match { + case Character.UPPERCASE_LETTER => maskedCodePoint(maskUpper) + case Character.LOWERCASE_LETTER => maskedCodePoint(maskLower) + case Character.DECIMAL_DIGIT_NUMBER => maskedCodePoint(maskDigit) + case _ => maskedCodePoint(maskOther) } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala index aed6b34716a75..988b1f73a6e66 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/mathExpressions.scala @@ -1282,6 +1282,8 @@ case class Hex(child: Expression) Seq(TypeCollection(LongType, BinaryType, StringTypeWithCollation(supportsTrimCollation = true))) override def dataType: DataType = child.dataType match { + // After ImplicitTypeCasts, a CHAR/VARCHAR input is STRING. Keep collation from the + // promoted child rather than DefaultStringProducingExpression's UTF8_BINARY StringType. case st: StringType => st case _ => super.dataType } @@ -1971,6 +1973,51 @@ case class BRound( newLeft: Expression, newRight: Expression): BRound = copy(child = newLeft, scale = newRight) } +/** + * Truncate an expression toward zero to `scale` decimal places. + * A negative `scale` truncates digits to the left of the decimal point. + * truncate(1234.5678, 2) = 1234.56, truncate(-1234.5678, 2) = -1234.56. + */ +// scalastyle:off line.size.limit +@ExpressionDescription( + usage = "_FUNC_(expr[, scale]) - Returns `expr` truncated toward zero to `scale` decimal places. `scale` defaults to 0. A negative `scale` truncates digits to the left of the decimal point.", + arguments = """ + Arguments: + * expr - The expression to truncate. An expression that evaluates to a numeric. + * scale - The number of decimal places to keep. It must be a constant integer expression and defaults to 0. A negative value truncates digits to the left of the decimal point. + """, + examples = """ + Examples: + > SELECT _FUNC_(1234.5678, 2); + 1234.56 + > SELECT _FUNC_(1234.5678, -2); + 1200 + > SELECT _FUNC_(-1234.5678, 2); + -1234.56 + """, + since = "4.4.0", + group = "math_funcs") +// scalastyle:on line.size.limit +case class Truncate( + child: Expression, + scale: Expression, + // Kept for symmetry with Round/BRound, which do need it. Truncation toward zero can never + // increase magnitude, so the ANSI overflow checks inherited from RoundBase never trigger. + override val ansiEnabled: Boolean = SQLConf.get.ansiEnabled) + // Also inherits RoundBase's one-digit decimal precision widening, needed for rounding modes + // that can carry (e.g. ceil(9.9, 0) = 10) but never exercised here since truncation cannot. + extends RoundBase(child, scale, BigDecimal.RoundingMode.DOWN, "ROUND_DOWN") { + def this(child: Expression) = this(child, Literal(0), SQLConf.get.ansiEnabled) + + def this(child: Expression, scale: Expression) = this(child, scale, SQLConf.get.ansiEnabled) + + override def flatArguments: Iterator[Any] = Iterator(child, scale) + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): Truncate = + copy(child = newLeft, scale = newRight) +} + object WidthBucket { /** Shared by interpreted eval and generated Java code; must stay public for codegen. */ def computeBucketNumber(value: Double, min: Double, max: Double, numBucket: Long): jl.Long = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala index e802776661090..61e04b9f761ce 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/misc.scala @@ -84,6 +84,7 @@ case class RaiseError(errorClass: Expression, errorParms: Expression, dataType: override def foldable: Boolean = false override def nullable: Boolean = true + override lazy val throwable: Boolean = true override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeWithCollation(supportsTrimCollation = true), @@ -364,6 +365,10 @@ case class SparkVersion() @ExpressionDescription( usage = """_FUNC_(expr) - Return DDL-formatted type string for the data type of the input.""", + arguments = """ + Arguments: + * expr - An expression of any type whose data type is returned as a DDL-formatted string. + """, examples = """ Examples: > SELECT _FUNC_(1); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala new file mode 100644 index 0000000000000..5d5e8f7829bc6 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala @@ -0,0 +1,291 @@ +/* + * 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.catalyst.expressions.ml + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckFailure} +import org.apache.spark.sql.catalyst.expressions.{Expression, Generator, Literal} +import org.apache.spark.sql.catalyst.expressions.Cast._ +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.util.ArrayData +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String + +/** + * Explodes the SQL struct representation of an MLlib vector into index-value pairs. This + * expression is dedicated only for Spark ML and should be used together with `unwrap_udt`. + * The mode controls whether it emits all entries or nonzero entries. It always emits a marker + * row before each vector for ML computations that need a per-vector row. The marker index is + * `-1 - vector.size`. + * + * Sparse vector examples: + * {{{ + * // v = {type: 0, size: 4, indices: [1, 3], values: [2.0, 4.0]} + * vector_posexplode(v) + * index value + * -5 NaN + * 1 2.0 + * 3 4.0 + * + * vector_posexplode(v, mode = "dense") + * index value + * -5 NaN + * 0 0.0 + * 1 2.0 + * 2 0.0 + * 3 4.0 + * }}} + * + * Dense vector examples: + * {{{ + * // v = {type: 1, size: null, indices: null, values: [1.0, 0.0, 3.0]} + * vector_posexplode(v) + * index value + * -4 NaN + * 0 1.0 + * 2 3.0 + * + * vector_posexplode(v, mode = "dense") + * index value + * -4 NaN + * 0 1.0 + * 1 0.0 + * 2 3.0 + * }}} + */ +case class VectorPosExplode(child: Expression, mode: Expression) + extends Generator with CodegenFallback { + + def this(child: Expression) = this(child, Literal("sparse")) + + override def children: Seq[Expression] = Seq(child, mode) + + @transient private lazy val vectorMode: VectorPosExplode.VectorMode.Value = + VectorPosExplode.toMode(mode.eval().asInstanceOf[UTF8String].toString) + + override def checkInputDataTypes(): TypeCheckResult = { + if (!VectorPosExplode.isVectorType(child.dataType)) { + return DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> + toSQLType(s"STRUCT with SQL type ${VectorPosExplode.vectorSqlType.sql}"), + "inputSql" -> toSQLExpr(child), + "inputType" -> toSQLType(child.dataType))) + } + if (!mode.foldable || !mode.dataType.isInstanceOf[StringType]) { + return DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(1), + "requiredType" -> toSQLType("foldable STRING"), + "inputSql" -> toSQLExpr(mode), + "inputType" -> toSQLType(mode.dataType))) + } + val modeValue = mode.eval() + if (modeValue == null) { + return TypeCheckFailure("The second argument of vector_posexplode cannot be null.") + } + VectorPosExplode.toModeOption(modeValue.asInstanceOf[UTF8String].toString) match { + case Some(_) => + case None => + return TypeCheckFailure( + "The second argument of vector_posexplode must be one of: dense, sparse.") + } + TypeCheckResult.TypeCheckSuccess + } + + override def elementSchema: StructType = VectorPosExplode.elementSchema + + override def eval(input: InternalRow): IterableOnce[InternalRow] = { + val vector = child.eval(input).asInstanceOf[InternalRow] + if (vector == null) { + Iterator.empty + } else { + val values = vector.getArray(3) + val (size, rows) = vector.getByte(0) match { + case VectorPosExplode.SparseVectorType => + val indices = vector.getArray(2) + if (indices == null || values == null || vector.isNullAt(1)) { + return Iterator.empty + } + val size = vector.getInt(1) + (size, VectorPosExplode.explodeSparse(vectorMode, size, indices, values)) + case VectorPosExplode.DenseVectorType => + if (values == null) { + return Iterator.empty + } + (values.numElements(), VectorPosExplode.explodeDense(vectorMode, values)) + case vectorType => + throw new IllegalArgumentException(s"Unknown vector type $vectorType.") + } + Iterator.single(VectorPosExplode.markerRow(size)) ++ rows + } + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): VectorPosExplode = { + copy(child = newChildren(0), mode = newChildren(1)) + } +} + +object VectorPosExplode { + object VectorMode extends Enumeration { + val Dense, Sparse = Value + } + + private val SparseVectorType: Byte = 0 + private val DenseVectorType: Byte = 1 + + private val vectorSqlType = StructType(Array( + StructField("type", ByteType, nullable = false), + StructField("size", IntegerType, nullable = true), + StructField("indices", ArrayType(IntegerType, containsNull = false), nullable = true), + StructField("values", ArrayType(DoubleType, containsNull = false), nullable = true))) + + private val elementSchema = new StructType() + .add("index", IntegerType, nullable = false) + .add("value", DoubleType, nullable = false) + + private def isVectorType(dataType: DataType): Boolean = dataType match { + case struct: StructType => struct == vectorSqlType + case _ => false + } + + private def toModeOption(mode: String): Option[VectorMode.Value] = mode match { + case "dense" => Some(VectorMode.Dense) + case "sparse" => Some(VectorMode.Sparse) + case _ => None + } + + private def toMode(mode: String): VectorMode.Value = toModeOption(mode).get + + private def markerRow(size: Int): InternalRow = InternalRow(-1 - size, Double.NaN) + + private def explodeSparse( + mode: VectorMode.Value, + size: Int, + indices: ArrayData, + values: ArrayData): Iterator[InternalRow] = mode match { + case VectorMode.Dense => + explodeSparseAsDense(size, indices, values) + case VectorMode.Sparse => + explodeSparseNonzero(indices, values) + } + + private def explodeSparseAsDense( + vectorSize: Int, + indices: ArrayData, + values: ArrayData): Iterator[InternalRow] = { + val numActives = values.numElements() + // Mirrors SparseVector.iterator without depending on MLlib from Catalyst. + new Iterator[InternalRow] { + private var index = 0 + private var activeIndex = 0 + private var nextActiveIndex = if (numActives > 0) indices.getInt(0) else -1 + + override def hasNext: Boolean = index < vectorSize + + override def next(): InternalRow = { + if (!hasNext) { + throw new NoSuchElementException("next on empty iterator") + } + val value = if (index == nextActiveIndex) { + val activeValue = values.getDouble(activeIndex) + activeIndex += 1 + nextActiveIndex = if (activeIndex < numActives) indices.getInt(activeIndex) else -1 + activeValue + } else { + 0.0 + } + val row = InternalRow(index, value) + index += 1 + row + } + } + } + + private def explodeSparseNonzero( + indices: ArrayData, + values: ArrayData): Iterator[InternalRow] = { + val numElements = values.numElements() + new Iterator[InternalRow] { + private var index = 0 + private var nextRow: InternalRow = _ + + override def hasNext: Boolean = { + while (nextRow == null && index < numElements) { + val value = values.getDouble(index) + if (value != 0.0) { + nextRow = InternalRow(indices.getInt(index), value) + } + index += 1 + } + nextRow != null + } + + override def next(): InternalRow = { + if (!hasNext) { + throw new NoSuchElementException("next on empty iterator") + } + val row = nextRow + nextRow = null + row + } + } + } + + private def explodeDense(mode: VectorMode.Value, values: ArrayData): Iterator[InternalRow] = { + mode match { + case VectorMode.Dense => + explodeDense(values, skipZero = false) + case VectorMode.Sparse => + explodeDense(values, skipZero = true) + } + } + + private def explodeDense(values: ArrayData, skipZero: Boolean): Iterator[InternalRow] = { + val numElements = values.numElements() + new Iterator[InternalRow] { + private var index = 0 + private var nextRow: InternalRow = _ + + override def hasNext: Boolean = { + while (nextRow == null && index < numElements) { + val value = values.getDouble(index) + if (!skipZero || value != 0.0) { + nextRow = InternalRow(index, value) + } + index += 1 + } + nextRow != null + } + + override def next(): InternalRow = { + if (!hasNext) { + throw new NoSuchElementException("next on empty iterator") + } + val row = nextRow + nextRow = null + row + } + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala index eb8e67eca3f14..db77d6b82b9f3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/nullExpressions.scala @@ -41,6 +41,11 @@ import org.apache.spark.sql.types._ // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr1, expr2, ...) - Returns the first non-null argument if exists. Otherwise, null.", + arguments = """ + Arguments: + * exprN - An expression of any type. All arguments must share a common type. + Arguments are evaluated in order and the first non-null value is returned. + """, examples = """ Examples: > SELECT _FUNC_(NULL, 1, NULL); @@ -252,6 +257,10 @@ case class NullIfZero(input: Expression, replacement: Expression) @ExpressionDescription( usage = "_FUNC_(expr) - Returns zero if `expr` is equal to null, or `expr` otherwise.", + arguments = """ + Arguments: + * expr - An expression. Zero is returned when it is null, otherwise `expr` is returned. + """, examples = """ Examples: > SELECT _FUNC_(NULL); @@ -273,6 +282,11 @@ case class ZeroIfNull(input: Expression, replacement: Expression) @ExpressionDescription( usage = "_FUNC_(expr1, expr2) - Returns `expr2` if `expr1` is null, or `expr1` otherwise.", + arguments = """ + Arguments: + * expr1 - An expression. Returned when it is not null. + * expr2 - The value returned when `expr1` is null. + """, examples = """ Examples: > SELECT _FUNC_(NULL, array('2')); @@ -297,6 +311,12 @@ case class Nvl(left: Expression, right: Expression, replacement: Expression) // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr1, expr2, expr3) - Returns `expr2` if `expr1` is not null, or `expr3` otherwise.", + arguments = """ + Arguments: + * expr1 - An expression tested for nullability. + * expr2 - The value returned when `expr1` is not null. + * expr3 - The value returned when `expr1` is null. + """, examples = """ Examples: > SELECT _FUNC_(NULL, 2, 1); @@ -464,6 +484,10 @@ case class NaNvl(left: Expression, right: Expression) */ @ExpressionDescription( usage = "_FUNC_(expr) - Returns true if `expr` is null, or false otherwise.", + arguments = """ + Arguments: + * expr - An expression of any type to test for nullability. + """, examples = """ Examples: > SELECT _FUNC_(1); @@ -498,6 +522,10 @@ case class IsNull(child: Expression) extends UnaryExpression with Predicate { */ @ExpressionDescription( usage = "_FUNC_(expr) - Returns true if `expr` is not null, or false otherwise.", + arguments = """ + Arguments: + * expr - An expression of any type to test for nullability. + """, examples = """ Examples: > SELECT _FUNC_(1); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala index 28e934b5df951..4e9d272973f09 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/objects/objects.scala @@ -219,10 +219,17 @@ trait InvokeLike extends Expression with NonSQLExpression with ImplicitCastInput returnType: Option[String] = None): String = { val castFuncCall = if (returnType.isEmpty) funcCall else s"(${returnType.get}) $funcCall" if (needTryCatch) { + // Catch Throwable rather than Exception: invoked methods may declare + // `throws Throwable` (common in reflection-derived signatures), and the + // JDK compiler requires the catch type to cover the declared throw. Janino + // erases checked-exception checks; javac enforces them. Platform.throwException + // rethrows the caught value unchanged (an Unsafe sneaky-throw, no wrapping), so + // an Error passing through here is observably identical to one propagating + // uncaught - the wider catch type must not change runtime behavior. s""" try { $resultVal = $castFuncCall; - } catch (Exception e) { + } catch (Throwable e) { org.apache.spark.unsafe.Platform.throwException(e); } """ @@ -595,7 +602,13 @@ case class NewInstance( propagateNull: Boolean, dataType: DataType, outerPointer: Option[() => AnyRef]) extends InvokeLike { + // JVM-form binary name (with `$` for inner classes); used where literal `$` + // is intentional (e.g., Scala companion access `Foo$.MODULE$`). private val className = cls.getName + // The same binary name as `className`, routed through `javaSourceName` to mark + // it as a type reference the JDK backend rewrites to a javac-legal form at + // compile time (see JdkCodeCompiler.rewriteInnerClassRefs). + private val javaSourceClassName = CodeGenerator.javaSourceName(cls) override def nullable: Boolean = needNullCheck @@ -678,7 +691,7 @@ case class NewInstance( case _ => outer.map { gen => s"${gen.value}.new ${Utils.getSimpleName(cls)}($argString)" }.getOrElse { - s"new $className($argString)" + s"new $javaSourceClassName($argString)" } } @@ -1182,7 +1195,11 @@ case class MapObjects private( val (initCollection, addElement, getResult): (String, String => String, String) = customCollectionCls match { case Some(cls) if classOf[mutable.ArraySeq[_]].isAssignableFrom(cls) => - val tag = ctx.addReferenceObj("tag", elementClassTag()) + // Cast the reference to the public ClassTag interface rather than the + // ClassTag's concrete runtime class, which can be a non-public inner class + // (e.g. scala.reflect.ClassTag$GenericClassTag). The JDK compiler rejects a + // cast to an inaccessible type ("has private access"); Janino does not. + val tag = ctx.addReferenceObj("tag", elementClassTag(), classOf[ClassTag[_]].getName) val builderClassName = classOf[mutable.ArrayBuilder[_]].getName val getBuilder = s"$builderClassName$$.MODULE$$.make($tag)" val builder = ctx.freshName("collectionBuilder") @@ -1196,7 +1213,11 @@ case class MapObjects private( s"MODULE$$.make($builder.result());" ) case Some(cls) if classOf[immutable.ArraySeq[_]].isAssignableFrom(cls) => - val tag = ctx.addReferenceObj("tag", elementClassTag()) + // Cast the reference to the public ClassTag interface rather than the + // ClassTag's concrete runtime class, which can be a non-public inner class + // (e.g. scala.reflect.ClassTag$GenericClassTag). The JDK compiler rejects a + // cast to an inaccessible type ("has private access"); Janino does not. + val tag = ctx.addReferenceObj("tag", elementClassTag(), classOf[ClassTag[_]].getName) val builderClassName = classOf[mutable.ArrayBuilder[_]].getName val getBuilder = s"$builderClassName$$.MODULE$$.make($tag)" val builder = ctx.freshName("collectionBuilder") diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala index b52d09fc9c709..61fb573dc39f0 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala @@ -21,7 +21,7 @@ import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, UnresolvedSeed} import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch -import org.apache.spark.sql.catalyst.expressions.ExpectsInputTypes.{toSQLExpr, toSQLId} +import org.apache.spark.sql.catalyst.expressions.ExpectsInputTypes.{toSQLExpr, toSQLId, toSQLValue} import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral} import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.trees.{BinaryLike, TernaryLike, UnaryLike} @@ -301,14 +301,14 @@ case class Uniform( override def third: Expression = seedExpression override def withNewSeed(newSeed: Long): Expression = - Uniform(min, max, Literal(newSeed, LongType), hideSeed) + Uniform(min, max, Literal(newSeed, LongType), hideSeed, timeZoneId) override def withShiftedSeed(shift: Long): Expression = - Uniform(min, max, Literal(seed + shift, LongType), hideSeed) + Uniform(min, max, Literal(seed + shift, LongType), hideSeed, timeZoneId) override def withNewChildrenInternal( newFirst: Expression, newSecond: Expression, newThird: Expression): Expression = - Uniform(newFirst, newSecond, newThird, hideSeed) + copy(min = newFirst, max = newSecond, seedExpression = newThird) override def replacement: Expression = { if (Seq(min, max, seedExpression).exists(_.dataType == NullType)) { @@ -345,13 +345,13 @@ object Uniform { usage = """ _FUNC_(length[, seed]) - Returns a string of the specified length whose characters are chosen uniformly at random from the following pool of characters: 0-9, a-z, A-Z. The random seed is - optional. The string length must be a constant two-byte or four-byte integer (SMALLINT or INT, - respectively). + optional. The string length must be a non-negative constant two-byte or four-byte integer + (SMALLINT or INT, respectively). """, arguments = """ Arguments: * length - The length of the random string to generate. - An expression that evaluates to an integer. Must be a constant. + An expression that evaluates to a non-negative integer. Must be a constant. * seed - The seed used to produce reproducible random results. An expression that evaluates to an integer or long. Must be a constant. """, @@ -425,6 +425,18 @@ case class RandStr( "inputExpr" -> toSQLExpr(expr))) } } + if (result == TypeCheckResult.TypeCheckSuccess) { + val lengthValue = length.eval() + // randstr(NULL, 0) is valid (treated as 0), so only reject a negative length. + if (lengthValue != null && lengthValue.asInstanceOf[Int] < 0) { + result = DataTypeMismatch( + errorSubClass = "VALUE_OUT_OF_RANGE", + messageParameters = Map( + "exprName" -> toSQLId("length"), + "valueRange" -> s"[0, ${Int.MaxValue}]", + "currentValue" -> toSQLValue(lengthValue, IntegerType))) + } + } result } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala index 7800b72d44aae..5c1fb73cf1a80 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/regexpExpressions.scala @@ -380,7 +380,12 @@ sealed abstract class LikeAllBase extends MultiLikeBase { val javaDataType = CodeGenerator.javaType(child.dataType) val pattern = ctx.freshName("pattern") val valueArg = ctx.freshName("valueArg") - val patternCache = ctx.addReferenceObj("patternCache", cache.asJava) + // Cast to a parameterized List so the generated for-each binds elements to Pattern + // (javac rejects iterating a raw collection into a typed loop var; Janino allows it), + // and so the cast targets the public List interface rather than the non-public + // Scala collection wrapper's runtime class. + val patternCache = ctx.addReferenceObj( + "patternCache", cache.asJava, "java.util.List<java.util.regex.Pattern>") val checkNotMatchCode = if (isNotSpecified) { s"$pattern.matcher($valueArg.toString()).matches()" @@ -440,7 +445,12 @@ sealed abstract class LikeAnyBase extends MultiLikeBase { val javaDataType = CodeGenerator.javaType(child.dataType) val pattern = ctx.freshName("pattern") val valueArg = ctx.freshName("valueArg") - val patternCache = ctx.addReferenceObj("patternCache", cache.asJava) + // Cast to a parameterized List so the generated for-each binds elements to Pattern + // (javac rejects iterating a raw collection into a typed loop var; Janino allows it), + // and so the cast targets the public List interface rather than the non-public + // Scala collection wrapper's runtime class. + val patternCache = ctx.addReferenceObj( + "patternCache", cache.asJava, "java.util.List<java.util.regex.Pattern>") val checkMatchCode = if (isNotSpecified) { s"!$pattern.matcher($valueArg.toString()).matches()" @@ -607,7 +617,8 @@ case class RLike(left: Expression, right: Expression) extends StringRegexExpress case class StringSplit(str: Expression, regex: Expression, limit: Expression) extends TernaryExpression with ImplicitCastInputTypes { override def nullIntolerant: Boolean = true - override def dataType: DataType = ArrayType(str.dataType, containsNull = false) + override def dataType: DataType = + ArrayType(str.dataType, containsNull = false) override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeBinaryLcase, StringTypeWithCollation, IntegerType) override def first: Expression = str @@ -738,6 +749,7 @@ case class RegExpReplace(subject: Expression, regexp: Expression, rep: Expressio // last replacement string, we don't want to convert a UTF8String => java.langString every time. @transient private var lastReplacement: String = _ @transient private var lastReplacementInUTF8: UTF8String = _ + override def stateful: Boolean = true final override val nodePatterns: Seq[TreePattern] = Seq(REGEXP_REPLACE) override def nullSafeEval(s: Any, p: Any, r: Any, i: Any): Any = { @@ -754,7 +766,8 @@ case class RegExpReplace(subject: Expression, regexp: Expression, rep: Expressio RegExpUtils.replace(pattern, s.toString, lastReplacement, i.asInstanceOf[Int]) } - override def dataType: DataType = subject.dataType + override def dataType: DataType = + subject.dataType override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeBinaryLcase, StringTypeWithCollation, StringTypeBinaryLcase, IntegerType) @@ -855,6 +868,7 @@ abstract class RegExpExtractBase @transient private var lastRegex: UTF8String = _ // last regex pattern, we cache it for performance concern @transient private var pattern: Pattern = _ + override def stateful: Boolean = true final override val nodePatterns: Seq[TreePattern] = Seq(REGEXP_EXTRACT_FAMILY) @@ -928,7 +942,8 @@ case class RegExpExtract(subject: Expression, regexp: Expression, idx: Expressio RegExpExtractBase.extract(getLastMatcher(s, p), r.asInstanceOf[Int], prettyName) } - override def dataType: DataType = subject.dataType + override def dataType: DataType = + subject.dataType override def prettyName: String = "regexp_extract" override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { @@ -1004,7 +1019,8 @@ case class RegExpExtractAll(subject: Expression, regexp: Expression, idx: Expres RegExpExtractBase.extractAll(getLastMatcher(s, p), r.asInstanceOf[Int], prettyName) } - override def dataType: DataType = ArrayType(subject.dataType) + override def dataType: DataType = + ArrayType(subject.dataType) override def prettyName: String = "regexp_extract_all" override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala index 4ea1fa03173ba..afd894167ac97 100755 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala @@ -24,6 +24,8 @@ import java.util.{Base64 => JBase64, HashMap, Locale, Map => JMap} import scala.collection.mutable.ArrayBuffer +import org.apache.commons.codec.binary.{Base32 => CommonsBase32} + import org.apache.spark.QueryContext import org.apache.spark.network.util.JavaUtils import org.apache.spark.sql.catalyst.InternalRow @@ -97,7 +99,8 @@ case class ConcatWs(children: Seq[Expression]) Seq.fill(children.size - 1)(arrayOrStr) } - override def dataType: DataType = children.head.dataType + override def dataType: DataType = + children.head.dataType override def nullable: Boolean = children.head.nullable override def foldable: Boolean = children.forall(_.foldable) @@ -280,6 +283,12 @@ case class ConcatWs(children: Seq[Expression]) and `spark.sql.ansi.enabled` is set to false. If `spark.sql.ansi.enabled` is set to true, it throws ArrayIndexOutOfBoundsException for invalid indices. """, + arguments = """ + Arguments: + * n - An integer expression giving the 1-based index of the input to return. + * input1, input2, ... - The input expressions to select from. They can be strings + or binary values. + """, examples = """ Examples: > SELECT _FUNC_(1, 'scala', 'java'); @@ -445,7 +454,8 @@ trait String2StringExpression extends ImplicitCastInputTypes { def convert(v: UTF8String): UTF8String - override def dataType: DataType = child.dataType + override def dataType: DataType = + child.dataType override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCollation(supportsTrimCollation = true)) override def contextIndependentFoldable: Boolean = child.contextIndependentFoldable @@ -961,6 +971,66 @@ case class TryValidateUTF8(input: Expression) extends RuntimeReplaceable with Im } +/** + * A function that returns the Unicode normalization of a string. + */ +// scalastyle:off +@ExpressionDescription( + usage = """ + _FUNC_(str[, form]) - Returns the Unicode normalization of `str` using the normalization `form`. + Valid forms are 'NFC' (default), 'NFD', 'NFKC', and 'NFKD', as defined by Unicode Standard + Annex #15: 'NFD'/'NFKD' apply canonical/compatibility decomposition; 'NFC'/'NFKC' apply the + same decomposition followed by canonical composition. The form name is case-insensitive. + Normalization is backed by Spark's bundled ICU4J library rather than the JVM's own Unicode + data, so results are stable across JVM vendors and versions. + """, + arguments = """ + Arguments: + * str - a string expression to normalize. + * form - a string expression giving the normalization form: 'NFC', 'NFD', 'NFKC', or 'NFKD'. + If omitted, 'NFC' is used. + """, + examples = """ + Examples: + > SELECT _FUNC_('fi', 'NFKC'); + fi + """, + since = "4.4.0", + group = "string_funcs") +// scalastyle:on +case class Normalize(input: Expression, form: Expression) + extends RuntimeReplaceable with ImplicitCastInputTypes with BinaryLike[Expression] { + override def nullIntolerant: Boolean = true + + override lazy val replacement: Expression = + StaticInvoke( + classOf[ExpressionImplUtils], + input.dataType, + "normalize", + Seq(input, form), + inputTypes) + + def this(input: Expression) = this(input, Literal("NFC")) + + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true), + StringTypeWithCollation(supportsTrimCollation = true)) + + override def nodeName: String = "normalize" + + override def nullable: Boolean = true + + override def left: Expression = input + + override def right: Expression = form + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): Normalize = { + copy(input = newLeft, form = newRight) + } + +} + /** * Replace all occurrences with string. */ @@ -997,7 +1067,7 @@ case class StringReplace(srcExpr: Expression, searchExpr: Expression, replaceExp override def nullSafeEval(srcEval: Any, searchEval: Any, replaceEval: Any): Any = { CollationSupport.StringReplace.exec(srcEval.asInstanceOf[UTF8String], - searchEval.asInstanceOf[UTF8String], replaceEval.asInstanceOf[UTF8String], collationId); + searchEval.asInstanceOf[UTF8String], replaceEval.asInstanceOf[UTF8String], collationId) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { @@ -1005,7 +1075,8 @@ case class StringReplace(srcExpr: Expression, searchExpr: Expression, replaceExp CollationSupport.StringReplace.genCode(src, search, replace, collationId)) } - override def dataType: DataType = srcExpr.dataType + override def dataType: DataType = + srcExpr.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeNonCSAICollation(supportsTrimCollation = true), @@ -1099,7 +1170,8 @@ case class Overlay(input: Expression, replace: Expression, pos: Expression, len: this(str, replace, pos, Literal.create(-1, IntegerType)) } - override def dataType: DataType = input.dataType + override def dataType: DataType = + input.dataType override def inputTypes: Seq[AbstractDataType] = Seq( TypeCollection( @@ -1249,6 +1321,7 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac @transient private var lastMatching: UTF8String = _ @transient private var lastReplace: UTF8String = _ @transient private var dict: JMap[String, String] = _ + override def stateful: Boolean = true final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId @@ -1287,7 +1360,8 @@ case class StringTranslate(srcExpr: Expression, matchingExpr: Expression, replac }) } - override def dataType: DataType = srcExpr.dataType + override def dataType: DataType = + srcExpr.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeNonCSAICollation(supportsTrimCollation = true), @@ -1367,7 +1441,8 @@ trait String2TrimExpression extends Expression with ImplicitCastInputTypes { protected def direction: String override def children: Seq[Expression] = srcStr +: trimStr.toSeq - override def dataType: DataType = srcStr.dataType + override def dataType: DataType = + srcStr.dataType override def inputTypes: Seq[AbstractDataType] = Seq.fill(children.size)(StringTypeWithCollation(supportsTrimCollation = true)) @@ -1924,7 +1999,8 @@ case class SubstringIndex(strExpr: Expression, delimExpr: Expression, countExpr: override def nullIntolerant: Boolean = true final lazy val collationId: Int = first.dataType.asInstanceOf[StringType].collationId - override def dataType: DataType = strExpr.dataType + override def dataType: DataType = + strExpr.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeNonCSAICollation(supportsTrimCollation = true), @@ -1939,7 +2015,7 @@ case class SubstringIndex(strExpr: Expression, delimExpr: Expression, countExpr: override def nullSafeEval(str: Any, delim: Any, count: Any): Any = { CollationSupport.SubstringIndex.exec(str.asInstanceOf[UTF8String], - delim.asInstanceOf[UTF8String], count.asInstanceOf[Int], collationId); + delim.asInstanceOf[UTF8String], count.asInstanceOf[Int], collationId) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { @@ -2023,7 +2099,7 @@ case class StringLocate(substr: Expression, str: Expression, start: Expression) 0 } else { CollationSupport.StringLocate.exec(l.asInstanceOf[UTF8String], - r.asInstanceOf[UTF8String], s.asInstanceOf[Int] - 1, collationId) + 1; + r.asInstanceOf[UTF8String], s.asInstanceOf[Int] - 1, collationId) + 1 } } } @@ -2137,7 +2213,8 @@ case class StringLPad(str: Expression, len: Expression, pad: Expression) override def second: Expression = len override def third: Expression = pad - override def dataType: DataType = str.dataType + override def dataType: DataType = + str.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeWithCollation(supportsTrimCollation = true), @@ -2233,7 +2310,8 @@ case class StringRPad( override def second: Expression = len override def third: Expression = pad - override def dataType: DataType = str.dataType + override def dataType: DataType = + str.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeWithCollation(supportsTrimCollation = true), @@ -2289,7 +2367,8 @@ case class FormatString(children: Expression*) extends Expression with ImplicitC override def foldable: Boolean = children.forall(_.foldable) override def contextIndependentFoldable: Boolean = children.forall(_.contextIndependentFoldable) override def nullable: Boolean = children(0).nullable - override def dataType: DataType = children(0).dataType + override def dataType: DataType = + children(0).dataType override def inputTypes: Seq[AbstractDataType] = StringTypeWithCollation(supportsTrimCollation = true) :: @@ -2411,7 +2490,8 @@ case class InitCap(child: Expression) override def inputTypes: Seq[AbstractDataType] = Seq(StringTypeWithCollation(supportsTrimCollation = true)) - override def dataType: DataType = child.dataType + override def dataType: DataType = + child.dataType override def nullSafeEval(string: Any): Any = { CollationSupport.InitCap.exec(string.asInstanceOf[UTF8String], collationId, useICU) @@ -2448,7 +2528,8 @@ case class StringRepeat(str: Expression, times: Expression) override def nullIntolerant: Boolean = true override def left: Expression = str override def right: Expression = times - override def dataType: DataType = str.dataType + override def dataType: DataType = + str.dataType override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeWithCollation(supportsTrimCollation = true), @@ -2560,7 +2641,8 @@ case class Substring(str: Expression, pos: Expression, len: Expression) this(str, pos, Literal(Integer.MAX_VALUE)) } - override def dataType: DataType = str.dataType + override def dataType: DataType = + str.dataType override def inputTypes: Seq[AbstractDataType] = Seq( @@ -2624,12 +2706,16 @@ case class Substring(str: Expression, pos: Expression, len: Expression) case class Right(str: Expression, len: Expression) extends RuntimeReplaceable with ImplicitCastInputTypes with BinaryLike[Expression] { + // Type the literal branches after ImplicitTypeCasts promotes CHAR/VARCHAR to STRING. + // Substring then returns STRING, so the If branches match. + private lazy val resultType: DataType = str.dataType + override lazy val replacement: Expression = If( IsNull(str), - Literal(null, str.dataType), + Literal(null, resultType), If( LessThanOrEqual(len, Literal(0)), - Literal(UTF8String.EMPTY_UTF8, str.dataType), + Literal(UTF8String.EMPTY_UTF8, resultType), new Substring(str, UnaryMinus(len, failOnError = false)) ) ) @@ -3349,6 +3435,110 @@ object UnBase64 { } } +/** + * Converts the argument from binary to a base 32 string. + */ +@ExpressionDescription( + usage = "_FUNC_(bin) - Converts the argument from a binary `bin` to a base 32 string.", + arguments = """ + Arguments: + * bin - The binary value to encode as a base 32 string. + An expression that evaluates to a binary. + """, + examples = """ + Examples: + > SELECT _FUNC_('foobar'); + MZXW6YTBOI====== + > SELECT _FUNC_(x'666f6f626172'); + MZXW6YTBOI====== + """, + since = "4.3.0", + group = "string_funcs") +case class Base32(child: Expression) + extends UnaryExpression + with RuntimeReplaceable + with ImplicitCastInputTypes + with DefaultStringProducingExpression { + + override def inputTypes: Seq[DataType] = Seq(BinaryType) + + override def contextIndependentFoldable: Boolean = child.contextIndependentFoldable + + override lazy val replacement: Expression = StaticInvoke( + classOf[Base32], + dataType, + "encode", + Seq(child), + Seq(BinaryType), + returnNullable = false) + + override def toString: String = s"$prettyName($child)" + + override def prettyName: String = "to_base32" + + override protected def withNewChildInternal(newChild: Expression): Expression = + copy(child = newChild) +} + +object Base32 { + private lazy val codec = new CommonsBase32() + + def encode(input: Array[Byte]): UTF8String = { + UTF8String.fromBytes(codec.encode(input)) + } +} + +/** + * Converts the argument from a base 32 string to BINARY. + */ +@ExpressionDescription( + usage = "_FUNC_(str) - Converts the argument from a base 32 string `str` to a binary.", + arguments = """ + Arguments: + * str - The base 32 string to decode to binary. + An expression that evaluates to a string. + """, + examples = """ + Examples: + > SELECT _FUNC_('MZXW6YTBOI======'); + foobar + """, + since = "4.3.0", + group = "string_funcs") +case class UnBase32(child: Expression) + extends UnaryExpression + with RuntimeReplaceable + with ImplicitCastInputTypes { + + override def dataType: DataType = BinaryType + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + override def contextIndependentFoldable: Boolean = child.contextIndependentFoldable + + override lazy val replacement: Expression = StaticInvoke( + classOf[UnBase32], + dataType, + "decode", + Seq(child), + inputTypes, + returnNullable = false) + + override def toString: String = s"$prettyName($child)" + + override def prettyName: String = "from_base32" + + override protected def withNewChildInternal(newChild: Expression): Expression = + copy(child = newChild) +} + +object UnBase32 { + private lazy val codec = new CommonsBase32() + + def decode(input: UTF8String): Array[Byte] = { + codec.decode(input.getBytes) + } +} + object Decode { def createExpr(params: Seq[Expression]): Expression = { params.length match { @@ -3737,6 +3927,7 @@ case class FormatNumber(x: Expression, d: Expression) // as a decimal separator. @transient private lazy val numberFormat = new DecimalFormat("", new DecimalFormatSymbols(Locale.US)) + override def stateful: Boolean = true override protected def nullSafeEval(xObject: Any, dObject: Any): Any = { right.dataType match { @@ -3892,8 +4083,10 @@ case class Sentences( def this(str: Expression, language: Expression) = this(str, language, Literal("")) override def nullable: Boolean = true - override def dataType: DataType = - ArrayType(ArrayType(str.dataType, containsNull = false), containsNull = false) + override def dataType: DataType = { + val elementType = str.dataType + ArrayType(ArrayType(elementType, containsNull = false), containsNull = false) + } override def inputTypes: Seq[AbstractDataType] = Seq( StringTypeWithCollation(supportsTrimCollation = true), @@ -3927,8 +4120,13 @@ case class Sentences( */ case class StringSplitSQL( str: Expression, - delimiter: Expression) extends BinaryExpression { - override def dataType: DataType = ArrayType(str.dataType, containsNull = false) + delimiter: Expression) extends BinaryExpression with ExpectsInputTypes { + override def dataType: DataType = + ArrayType(str.dataType, containsNull = false) + override def inputTypes: Seq[AbstractDataType] = + Seq( + StringTypeWithCollation(supportsTrimCollation = true), + StringTypeWithCollation(supportsTrimCollation = true)) final lazy val collationId: Int = left.dataType.asInstanceOf[StringType].collationId override def left: Expression = str override def right: Expression = delimiter @@ -4016,6 +4214,11 @@ case class SplitPart ( case class Empty2Null(child: Expression) extends UnaryExpression with String2StringExpression { override def convert(v: UTF8String): UTF8String = if (v.numBytes() == 0) null else v + // Not a transforming function: every non-empty value is returned unchanged, so this keeps the + // child's type rather than taking the plain-STRING result that String2StringExpression gives + // its transforming implementations. + override def dataType: DataType = child.dataType + override def nullable: Boolean = true override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala index 8ac40a3fe2a57..1aa668a9b8940 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala @@ -28,6 +28,11 @@ import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType, Integ usage = """ _FUNC_(expr) - Returns the estimated number of unique values given the binary representation of a Datasketches ThetaSketch. """, + arguments = """ + Arguments: + * expr - The binary representation of a Datasketches ThetaSketch. + An expression that evaluates to binary. + """, examples = """ Examples: > SELECT _FUNC_(theta_sketch_agg(col)) FROM VALUES (1), (1), (2), (2), (3) tab(col); @@ -66,6 +71,16 @@ case class ThetaSketchEstimate(child: Expression) Datasketches ThetaSketch objects using a ThetaSketch Union object. Users can set lgNomEntries to a value between 4 and 26 to find the union of sketches with different union buffer size values (defaults to 12). """, + arguments = """ + Arguments: + * first - The binary representation of the first Datasketches ThetaSketch. + An expression that evaluates to binary. + * second - The binary representation of the second Datasketches ThetaSketch. + An expression that evaluates to binary. + * lgNomEntries - Optional. The log-base-2 of the nominal entries used to size the + union buffer, between 4 and 26. An expression that evaluates to an integer. + Defaults to 12. + """, examples = """ Examples: > SELECT theta_sketch_estimate(_FUNC_(theta_sketch_agg(col1), theta_sketch_agg(col2))) FROM VALUES (1, 4), (1, 4), (2, 5), (2, 5), (3, 6) tab(col1, col2); @@ -125,6 +140,13 @@ case class ThetaUnion(first: Expression, second: Expression, third: Expression) _FUNC_(first, second) - Subtracts two binary representations of Datasketches ThetaSketch objects from two input columns using a ThetaSketch AnotB object. """, + arguments = """ + Arguments: + * first - The binary representation of the first Datasketches ThetaSketch. + An expression that evaluates to binary. + * second - The binary representation of the second Datasketches ThetaSketch to + subtract from the first. An expression that evaluates to binary. + """, examples = """ Examples: > SELECT theta_sketch_estimate(_FUNC_(theta_sketch_agg(col1), theta_sketch_agg(col2))) FROM VALUES (5, 4), (1, 4), (2, 5), (2, 5), (3, 1) tab(col1, col2); @@ -173,6 +195,13 @@ case class ThetaDifference(first: Expression, second: Expression) _FUNC_(first, second) - Intersects two binary representations of Datasketches ThetaSketch objects from two input columns using a ThetaSketch Intersect object. """, + arguments = """ + Arguments: + * first - The binary representation of the first Datasketches ThetaSketch. + An expression that evaluates to binary. + * second - The binary representation of the second Datasketches ThetaSketch. + An expression that evaluates to binary. + """, examples = """ Examples: > SELECT theta_sketch_estimate(_FUNC_(theta_sketch_agg(col1), theta_sketch_agg(col2))) FROM VALUES (5, 4), (1, 4), (2, 5), (2, 5), (3, 1) tab(col1, col2); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromProtobufSqlFunctions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromProtobufSqlFunctions.scala index 7d0fa35b3f0f7..644695dfabbfd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromProtobufSqlFunctions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/toFromProtobufSqlFunctions.scala @@ -46,6 +46,18 @@ import org.apache.spark.util.Utils usage = """ _FUNC_(data, messageName, descFilePath, options) - Converts a binary Protobuf value into a Catalyst value. """, + arguments = """ + Arguments: + * data - The binary Protobuf value to convert. + * messageName - A constant string naming the Protobuf message to look for in + the descriptor file. + * descFilePath - Optional. A constant string or binary value with the + Protobuf descriptor file, created by `protoc` with `--descriptor_set_out` + and `--include_imports`. If omitted, the message must be resolvable + otherwise. + * options - Optional. A constant map of string key-value pairs controlling + the conversion. By default no options are set. + """, examples = """ Examples: > SELECT _FUNC_(s, 'Person', '/path/to/descriptor.desc', map()) IS NULL AS result FROM (SELECT NAMED_STRUCT('name', name, 'id', id) AS s FROM VALUES ('John Doe', 1), (NULL, 2) tab(name, id)); @@ -190,6 +202,17 @@ case class FromProtobuf( _FUNC_(child, messageName, descFilePath, options) - Converts a Catalyst binary input value into its corresponding Protobuf format result. """, + arguments = """ + Arguments: + * child - The Catalyst input value to convert to Protobuf binary. + * messageName - A constant string naming the Protobuf message to serialize to. + * descFilePath - Optional. A constant string or binary value with the + Protobuf descriptor file, created by `protoc` with `--descriptor_set_out` + and `--include_imports`. If omitted, the message must be resolvable + otherwise. + * options - Optional. A constant map of string key-value pairs controlling + the conversion. By default no options are set. + """, examples = """ Examples: > SELECT _FUNC_(s, 'Person', '/path/to/descriptor.desc', map('emitDefaultValues', 'true')) IS NULL FROM (SELECT NULL AS s); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleDifference.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleDifference.scala index e26354b54f6c9..539af821e8e80 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleDifference.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleDifference.scala @@ -31,6 +31,12 @@ import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType} _FUNC_(tupleSketch1, tupleSketch2) - Subtracts two binary representations of Datasketches TupleSketch objects with double summary data type using a TupleSketch AnotB object. Returns elements in the first sketch that are not in the second sketch. """, + arguments = """ + Arguments: + * tupleSketch1 - A binary representation of a TupleSketch with double summary. + * tupleSketch2 - A binary representation of a TupleSketch with double summary + to subtract from the first sketch. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), tuple_sketch_agg_double(col2, val2))) FROM VALUES (5, 5.0D, 4, 4.0D), (1, 1.0D, 4, 4.0D), (2, 2.0D, 5, 5.0D), (3, 3.0D, 1, 1.0D) tab(col1, val1, col2, val2); @@ -67,6 +73,12 @@ case class TupleDifferenceDouble(left: Expression, right: Expression) _FUNC_(tupleSketch1, tupleSketch2) - Subtracts two binary representations of Datasketches TupleSketch objects with integer summary data type using a TupleSketch AnotB object. Returns elements in the first sketch that are not in the second sketch. """, + arguments = """ + Arguments: + * tupleSketch1 - A binary representation of a TupleSketch with integer summary. + * tupleSketch2 - A binary representation of a TupleSketch with integer summary + to subtract from the first sketch. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), tuple_sketch_agg_integer(col2, val2))) FROM VALUES (5, 5, 4, 4), (1, 1, 4, 4), (2, 2, 5, 5), (3, 3, 1, 1) tab(col1, val1, col2, val2); @@ -103,6 +115,11 @@ case class TupleDifferenceInteger(left: Expression, right: Expression) _FUNC_(tupleSketch, thetaSketch) - Subtracts the binary representation of a Datasketches ThetaSketch from a TupleSketch with double summary data type using a TupleSketch AnotB object. Returns elements in the TupleSketch that are not in the ThetaSketch. """, + arguments = """ + Arguments: + * tupleSketch - A binary representation of a TupleSketch with double summary. + * thetaSketch - A binary representation of a ThetaSketch to subtract from the TupleSketch. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), theta_sketch_agg(col2))) FROM VALUES (5, 5.0D, 4), (1, 1.0D, 4), (2, 2.0D, 5), (3, 3.0D, 1) tab(col1, val1, col2); @@ -139,6 +156,11 @@ case class TupleDifferenceThetaDouble(left: Expression, right: Expression) _FUNC_(tupleSketch, thetaSketch) - Subtracts the binary representation of a Datasketches ThetaSketch from a TupleSketch with integer summary data type using a TupleSketch AnotB object. Returns elements in the TupleSketch that are not in the ThetaSketch. """, + arguments = """ + Arguments: + * tupleSketch - A binary representation of a TupleSketch with integer summary. + * thetaSketch - A binary representation of a ThetaSketch to subtract from the TupleSketch. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), theta_sketch_agg(col2))) FROM VALUES (5, 5, 4), (1, 1, 4), (2, 2, 5), (3, 3, 1) tab(col1, val1, col2); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleIntersection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleIntersection.scala index 2085b759af09e..b6a5cb1d38ce1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleIntersection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleIntersection.scala @@ -34,6 +34,15 @@ import org.apache.spark.unsafe.types.UTF8String _FUNC_(tupleSketch1, tupleSketch2, mode) - Intersects two binary representations of Datasketches TupleSketch objects with double summary data type using a TupleSketch Intersection object. Users can set mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch1 - The binary representation of a Datasketches TupleSketch with a double + summary data type. + * tupleSketch2 - The binary representation of a Datasketches TupleSketch with a double + summary data type. + * mode - The summary combination mode: 'sum', 'min', 'max', or 'alwaysone' + (optional, defaults to 'sum'). + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), tuple_sketch_agg_double(col2, val2))) FROM VALUES (1, 1.0D, 1, 4.0D), (2, 2.0D, 2, 5.0D), (3, 3.0D, 4, 6.0D) tab(col1, val1, col2, val2); @@ -80,6 +89,15 @@ case class TupleIntersectionDouble(first: Expression, second: Expression, third: _FUNC_(tupleSketch1, tupleSketch2, mode) - Intersects two binary representations of Datasketches TupleSketch objects with integer summary data type using a TupleSketch Intersection object. Users can set mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch1 - The binary representation of a Datasketches TupleSketch with an integer + summary data type. + * tupleSketch2 - The binary representation of a Datasketches TupleSketch with an integer + summary data type. + * mode - The summary combination mode: 'sum', 'min', 'max', or 'alwaysone' + (optional, defaults to 'sum'). + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), tuple_sketch_agg_integer(col2, val2))) FROM VALUES (1, 1, 1, 4), (2, 2, 2, 5), (3, 3, 4, 6) tab(col1, val1, col2, val2); @@ -131,6 +149,14 @@ case class TupleIntersectionInteger(first: Expression, second: Expression, third assigned a default double summary value based on the mode: 0.0 for 'sum' mode, +Infinity for 'min' mode, -Infinity for 'max' mode, or 1.0 for 'alwaysone' mode. Users can set mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch - The binary representation of a Datasketches TupleSketch with a double + summary data type. + * thetaSketch - The binary representation of a Datasketches ThetaSketch. + * mode - The summary combination mode: 'sum', 'min', 'max', or 'alwaysone' + (optional, defaults to 'sum'). + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1.0D, 1), (2, 2.0D, 2), (3, 3.0D, 4) tab(col1, val1, col2); @@ -182,6 +208,14 @@ case class TupleIntersectionThetaDouble(first: Expression, second: Expression, t assigned a default integer summary value based on the mode: 0 for 'sum' mode, Integer.MAX_VALUE for 'min' mode, Integer.MIN_VALUE for 'max' mode, or 1 for 'alwaysone' mode. Users can set mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch - The binary representation of a Datasketches TupleSketch with an integer + summary data type. + * thetaSketch - The binary representation of a Datasketches ThetaSketch. + * mode - The summary combination mode: 'sum', 'min', 'max', or 'alwaysone' + (optional, defaults to 'sum'). + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1, 1), (2, 2, 2), (3, 3, 4) tab(col1, val1, col2); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchEstimate.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchEstimate.scala index fcd0a048479b0..5ac9504bfa929 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchEstimate.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchEstimate.scala @@ -27,6 +27,11 @@ import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType, Doubl _FUNC_(child) - Returns the estimated number of unique values given the binary representation of a Datasketches TupleSketch. The sketch's summary type must be a double. """, + arguments = """ + Arguments: + * child - A binary value holding a serialized Datasketches TupleSketch + with a double summary. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_double(key, summary)) FROM VALUES (1, 1.0D), (1, 2.0D), (2, 3.0D) tab(key, summary); @@ -64,6 +69,11 @@ case class TupleSketchEstimateDouble(child: Expression) _FUNC_(child) - Returns the estimated number of unique values given the binary representation of a Datasketches TupleSketch. The sketch's summary type must be an integer. """, + arguments = """ + Arguments: + * child - A binary value holding a serialized Datasketches TupleSketch + with an integer summary. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_integer(key, summary)) FROM VALUES (1, 1), (1, 2), (2, 3) tab(key, summary); @@ -101,6 +111,11 @@ case class TupleSketchEstimateInteger(child: Expression) _FUNC_(child) - Returns the theta value (sampling rate) from a Datasketches TupleSketch. The theta value represents the effective sampling rate of the sketch, between 0.0 and 1.0. The sketch's summary type must be a double. """, + arguments = """ + Arguments: + * child - A binary value holding a serialized Datasketches TupleSketch + with a double summary. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_double(key, summary)) FROM VALUES (1, 1.0D), (2, 2.0D), (3, 3.0D) tab(key, summary); @@ -138,6 +153,11 @@ case class TupleSketchThetaDouble(child: Expression) _FUNC_(child) - Returns the theta value (sampling rate) from a Datasketches TupleSketch. The theta value represents the effective sampling rate of the sketch, between 0.0 and 1.0. The sketch's summary type must be an integer. """, + arguments = """ + Arguments: + * child - A binary value holding a serialized Datasketches TupleSketch + with an integer summary. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_integer(key, summary)) FROM VALUES (1, 1), (2, 2), (3, 3) tab(key, summary); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchSummary.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchSummary.scala index a37a03e41e534..625ab10c31197 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchSummary.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleSketchSummary.scala @@ -34,6 +34,13 @@ import org.apache.spark.unsafe.types.UTF8String _FUNC_(child, mode) - Aggregates the summary values from a double summary type Datasketches TupleSketch. The mode can be 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * child - A binary expression holding a serialized double summary type Datasketches + TupleSketch to aggregate. + * mode - An optional string aggregation mode: 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_double(key, summary)) FROM VALUES (1, 1.0D), (1, 2.0D), (2, 3.0D) tab(key, summary); @@ -98,6 +105,13 @@ case class TupleSketchSummaryDouble(left: Expression, right: Expression) _FUNC_(child, mode) - Aggregates the summary values from a integer summary type Datasketches TupleSketch. The mode can be 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * child - A binary expression holding a serialized integer summary type Datasketches + TupleSketch to aggregate. + * mode - An optional string aggregation mode: 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT _FUNC_(tuple_sketch_agg_integer(key, summary)) FROM VALUES (1, 1), (1, 2), (2, 3) tab(key, summary); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleUnion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleUnion.scala index be04fb20dd30b..d174e392bb939 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleUnion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/tupleUnion.scala @@ -377,6 +377,15 @@ abstract class TupleUnionBase[S <: Summary] TupleSketch objects with double summary data type using a TupleSketch Union object. Users can set lgNomEntries to a value between 4 and 26 (defaults to 12) and mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch1 - The first TupleSketch (double summary) as a binary value. + * tupleSketch2 - The second TupleSketch (double summary) as a binary value. + * lgNomEntries - Optional integer between 4 and 26 setting the log base 2 of the + number of nominal entries. Defaults to 12. + * mode - Optional summary combining mode, one of 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), tuple_sketch_agg_double(col2, val2))) FROM VALUES (1, 1.0D, 4, 4.0D), (2, 2.0D, 5, 5.0D), (3, 3.0D, 6, 6.0D) tab(col1, val1, col2, val2); @@ -407,6 +416,15 @@ object TupleUnionDoubleExpressionBuilder extends ExpressionBuilder { TupleSketch objects with integer summary data type using a TupleSketch Union object. Users can set lgNomEntries to a value between 4 and 26 (defaults to 12) and mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch1 - The first TupleSketch (integer summary) as a binary value. + * tupleSketch2 - The second TupleSketch (integer summary) as a binary value. + * lgNomEntries - Optional integer between 4 and 26 setting the log base 2 of the + number of nominal entries. Defaults to 12. + * mode - Optional summary combining mode, one of 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), tuple_sketch_agg_integer(col2, val2))) FROM VALUES (1, 1, 4, 4), (2, 2, 5, 5), (3, 3, 6, 6) tab(col1, val1, col2, val2); @@ -440,6 +458,15 @@ object TupleUnionIntegerExpressionBuilder extends ExpressionBuilder { -Infinity for 'max' mode, or 1.0 for 'alwaysone' mode. Users can set lgNomEntries to a value between 4 and 26 (defaults to 12) and mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch - The TupleSketch (double summary) as a binary value. + * thetaSketch - The ThetaSketch as a binary value. + * lgNomEntries - Optional integer between 4 and 26 setting the log base 2 of the + number of nominal entries. Defaults to 12. + * mode - Optional summary combining mode, one of 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_double(_FUNC_(tuple_sketch_agg_double(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1.0D, 4), (2, 2.0D, 5), (3, 3.0D, 6) tab(col1, val1, col2); @@ -473,6 +500,15 @@ object TupleUnionThetaDoubleExpressionBuilder extends ExpressionBuilder { mode, Integer.MIN_VALUE for 'max' mode, or 1 for 'alwaysone' mode. Users can set lgNomEntries to a value between 4 and 26 (defaults to 12) and mode to 'sum', 'min', 'max', or 'alwaysone' (defaults to 'sum'). """, + arguments = """ + Arguments: + * tupleSketch - The TupleSketch (integer summary) as a binary value. + * thetaSketch - The ThetaSketch as a binary value. + * lgNomEntries - Optional integer between 4 and 26 setting the log base 2 of the + number of nominal entries. Defaults to 12. + * mode - Optional summary combining mode, one of 'sum', 'min', 'max', or 'alwaysone'. + Defaults to 'sum'. + """, examples = """ Examples: > SELECT tuple_sketch_estimate_integer(_FUNC_(tuple_sketch_agg_integer(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1, 4), (2, 2, 5), (3, 3, 6) tab(col1, val1, col2); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala index da1bfbcea8294..b4ba27b79860d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala @@ -270,6 +270,12 @@ object VariantExpressionEvalUtils { arrayAppendAtPath(input, javaSegments, pathStr, value, valueDataType, functionName, failOnError) } + def stripNulls(input: VariantVal, includeArrays: Boolean): VariantVal = { + val v = new Variant(input.getValue, input.getMetadata) + val out = VariantBuilder.stripNulls(v, includeArrays) + new VariantVal(out.getValue, out.getMetadata) + } + /** Cast a Spark value from `dataType` into the variant type. */ def castToVariant(input: Any, dataType: DataType): VariantVal = { // Enforce strict check because it is illegal for input struct/map/variant to contain duplicate @@ -280,6 +286,77 @@ object VariantExpressionEvalUtils { new VariantVal(v.getValue, v.getMetadata) } + /** + * Build a variant object directly from a keys array and a values array, without materializing an + * intermediate map. Keys must be non-null strings and the two arrays must have equal length. A + * null key raises `NULL_MAP_KEY`, a duplicate key raises `VARIANT_DUPLICATE_KEY` (matching + * to_variant_object), and null values are kept as variant null. + */ + def variantFromArrays(keys: ArrayData, values: ArrayData, valueType: DataType): VariantVal = { + if (keys.numElements() != values.numElements()) { + // Reuse the same error map_from_arrays raises for a keys/values length mismatch. + throw QueryExecutionErrors.mapDataKeyArrayLengthDiffersFromValueArrayLengthError() + } + val builder = new VariantBuilder(false) + val start = builder.getWritePos + val numElements = keys.numElements() + val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](numElements) + var i = 0 + while (i < numElements) { + if (keys.isNullAt(i)) { + throw QueryExecutionErrors.nullAsMapKeyNotAllowedError() + } + val key = keys.getUTF8String(i).toString + val id = builder.addKey(key) + fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start)) + val value = if (values.isNullAt(i)) null else values.get(i, valueType) + buildVariant(builder, value, valueType) + i += 1 + } + builder.finishWritingObject(start, fields) + val v = builder.result() + new VariantVal(v.getValue, v.getMetadata) + } + + /** + * Build a variant object directly from an array of key/value struct entries, without an + * intermediate map. Keys must be non-null strings. A null entry makes the whole result null, + * and this is checked for every entry before any value is converted, so a null entry always + * dominates a conversion failure in an earlier entry (matching `map_from_entries`). A null key + * raises `NULL_MAP_KEY`, a duplicate key raises `VARIANT_DUPLICATE_KEY`, and null values are + * kept as variant null. + */ + def variantFromEntries(entries: ArrayData, valueType: DataType): VariantVal = { + val numElements = entries.numElements() + var i = 0 + while (i < numElements) { + if (entries.isNullAt(i)) { + return null + } + i += 1 + } + + val builder = new VariantBuilder(false) + val start = builder.getWritePos + val fields = new java.util.ArrayList[VariantBuilder.FieldEntry](numElements) + i = 0 + while (i < numElements) { + val entry = entries.getStruct(i, 2) + if (entry.isNullAt(0)) { + throw QueryExecutionErrors.nullAsMapKeyNotAllowedError() + } + val key = entry.getUTF8String(0).toString + val id = builder.addKey(key) + fields.add(new VariantBuilder.FieldEntry(key, id, builder.getWritePos - start)) + val value = if (entry.isNullAt(1)) null else entry.get(1, valueType) + buildVariant(builder, value, valueType) + i += 1 + } + builder.finishWritingObject(start, fields) + val v = builder.result() + new VariantVal(v.getValue, v.getMetadata) + } + /** Returns `true` if a data type is or has a child variant type. */ def typeContainsVariant(dt: DataType): Boolean = dt match { case _: VariantType => true diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala index e9ce291b655c3..c58634cfd8d47 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala @@ -32,9 +32,9 @@ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke import org.apache.spark.sql.catalyst.json.JsonInferSchema import org.apache.spark.sql.catalyst.plans.logical.{FunctionSignature, InputParameter} +import org.apache.spark.sql.catalyst.trees.{BinaryLike, UnaryLike} import org.apache.spark.sql.catalyst.trees.TreePattern.{TreePattern, VARIANT_GET} -import org.apache.spark.sql.catalyst.trees.UnaryLike -import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData, QuotingUtils} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, GenericArrayData, QuotingUtils} import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors} import org.apache.spark.sql.internal.SQLConf @@ -81,6 +81,10 @@ case class ParseJson(child: Expression, failOnError: Boolean = true) // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr) - Check if a variant value is a variant null. Returns true if and only if the input is a variant null and false otherwise (including in the case of SQL NULL).", + arguments = """ + Arguments: + * expr - A variant value to check. + """, examples = """ Examples: > SELECT _FUNC_(parse_json('null')); @@ -122,6 +126,11 @@ case class IsVariantNull(child: Expression) extends UnaryExpression // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr) - Convert a nested input (array/map/struct) into a variant where maps and structs are converted to variant objects which are unordered unlike SQL structs. Input maps can only have string keys.", + arguments = """ + Arguments: + * expr - A nested value of array, map, or struct type. Maps must have string keys. Nested + values of any type are allowed. + """, examples = """ Examples: > SELECT _FUNC_(named_struct('a', 1, 'b', 2)); @@ -185,6 +194,159 @@ case class ToVariantObject(child: Expression) } } +// scalastyle:off line.size.limit +@ExpressionDescription( + usage = "_FUNC_(keys, values) - Creates a variant object from the given arrays of keys and values. The keys must be non-null strings and the two arrays must have the same length.", + arguments = """ + Arguments: + * keys - An array of non-null strings used as the object keys. + * values - An array of values, with the same length as the keys array. + """, + examples = """ + Examples: + > SELECT _FUNC_(array('a', 'b'), array(1, 2)); + {"a":1,"b":2} + """, + since = "4.4.0", + group = "variant_funcs") +// scalastyle:on line.size.limit +case class VariantFromArrays(left: Expression, right: Expression) + extends BinaryExpression + with ExpectsInputTypes + with QueryErrorsBase { + override def nullIntolerant: Boolean = true + override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, ArrayType) + override def dataType: DataType = VariantType + + private lazy val valueType: DataType = right.dataType.asInstanceOf[ArrayType].elementType + + override def checkInputDataTypes(): TypeCheckResult = { + val defaultCheck = super.checkInputDataTypes() + if (defaultCheck.isFailure) { + defaultCheck + } else { + left.dataType.asInstanceOf[ArrayType].elementType match { + case _: StringType if VariantGet.checkDataType(valueType, allowStructsAndMaps = true) => + TypeCheckResult.TypeCheckSuccess + case _: StringType => + DataTypeMismatch( + errorSubClass = "CAST_WITHOUT_SUGGESTION", + messageParameters = + Map("srcType" -> toSQLType(valueType), "targetType" -> toSQLType(VariantType))) + case _ => + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> toSQLType(ArrayType(StringType)), + "inputSql" -> toSQLExpr(left), + "inputType" -> toSQLType(left.dataType))) + } + } + } + + override def prettyName: String = "variant_from_arrays" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): VariantFromArrays = + copy(left = newLeft, right = newRight) + + override protected def nullSafeEval(keyArray: Any, valueArray: Any): Any = + VariantExpressionEvalUtils.variantFromArrays( + keyArray.asInstanceOf[ArrayData], valueArray.asInstanceOf[ArrayData], valueType) + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + nullSafeCodeGen(ctx, ev, (keyArray, valueArray) => { + val cls = variant.VariantExpressionEvalUtils.getClass.getName.stripSuffix("$") + val valueTypeArg = ctx.addReferenceObj("valueType", valueType) + s"${ev.value} = $cls.variantFromArrays($keyArray, $valueArray, $valueTypeArg);" + }) + } +} + +// scalastyle:off line.size.limit +@ExpressionDescription( + usage = "_FUNC_(entries) - Creates a variant object from an array of key/value struct entries. The keys must be non-null strings.", + arguments = """ + Arguments: + * entries - An array of key/value structs, where the first field is a non-null string key. + """, + examples = """ + Examples: + > SELECT _FUNC_(array(struct('a', 1), struct('b', 2))); + {"a":1,"b":2} + """, + since = "4.4.0", + group = "variant_funcs") +// scalastyle:on line.size.limit +case class VariantFromEntries(child: Expression) + extends UnaryExpression + with QueryErrorsBase { + override def nullIntolerant: Boolean = true + + @transient + private lazy val dataTypeDetails: Option[(DataType, Boolean)] = child.dataType match { + case ArrayType( + StructType(Array(StructField(_, _, _, _), StructField(_, valueType, _, _))), + containsNull) => + Some((valueType, containsNull)) + case _ => None + } + + @transient private lazy val valueType: DataType = dataTypeDetails.get._1 + @transient private lazy val nullEntries: Boolean = dataTypeDetails.get._2 + + override def nullable: Boolean = child.nullable || nullEntries + override def dataType: DataType = VariantType + + override def checkInputDataTypes(): TypeCheckResult = child.dataType match { + case ArrayType( + StructType(Array(StructField(_, _: StringType, _, _), StructField(_, vt, _, _))), _) => + if (VariantGet.checkDataType(vt, allowStructsAndMaps = true)) { + TypeCheckResult.TypeCheckSuccess + } else { + DataTypeMismatch( + errorSubClass = "CAST_WITHOUT_SUGGESTION", + messageParameters = + Map("srcType" -> toSQLType(vt), "targetType" -> toSQLType(VariantType))) + } + case _ => + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> s"${toSQLType(ArrayType)} of pair ${toSQLType(StructType)}", + "inputSql" -> toSQLExpr(child), + "inputType" -> toSQLType(child.dataType))) + } + + override def prettyName: String = "variant_from_entries" + + override protected def withNewChildInternal(newChild: Expression): VariantFromEntries = + copy(child = newChild) + + override protected def nullSafeEval(input: Any): Any = { + val entries = input.asInstanceOf[ArrayData] + VariantExpressionEvalUtils.variantFromEntries(entries, valueType) + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + nullSafeCodeGen(ctx, ev, c => { + val cls = variant.VariantExpressionEvalUtils.getClass.getName.stripSuffix("$") + val valueTypeArg = ctx.addReferenceObj("valueType", valueType) + // nullSafeCodeGen only declares `ev.isNull` as a local when this expression is + // nullable; when it isn't, entries can never contain a null and the helper can never + // return null, so the reassignment must be skipped rather than referencing an + // undeclared variable. + val markNull = if (nullable) s"${ev.isNull} = ${ev.value} == null;" else "" + s""" + |${ev.value} = $cls.variantFromEntries($c, $valueTypeArg); + |$markNull + """.stripMargin + }) + } +} + // A path segment in the `VariantGet` expression represents either an object key access or an array // index access. sealed abstract class VariantPathSegment extends Serializable @@ -289,6 +451,11 @@ case class VariantGet( timeZoneId, zoneId) + override def eval(input: InternalRow): Any = { + val _ = parsedPath + super.eval(input) + } + protected override def nullSafeEval(input: Any, path: Any): Any = parsedPath match { case Some(pp) => VariantGet.variantGet(input.asInstanceOf[VariantVal], pp, dataType, castArgs) @@ -359,6 +526,7 @@ case object VariantGet { VariantType => true case ArrayType(elementType, _) => checkDataType(elementType, allowStructsAndMaps) + case MapType(_: CharType | _: VarcharType, _, _) => false case MapType(_: StringType, valueType, _) if allowStructsAndMaps => checkDataType(valueType, allowStructsAndMaps) case StructType(fields) if allowStructsAndMaps => @@ -633,6 +801,13 @@ abstract class VariantGetExpressionBuilderBase(failOnError: Boolean) extends Exp // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(v, path[, type]) - Extracts a sub-variant from `v` according to `path`, and then cast the sub-variant to `type`. When `type` is omitted, it is default to `variant`. Returns null if the path does not exist. Throws an exception if the cast fails.", + arguments = """ + Arguments: + * v - A variant value to extract from. + * path - A string literal in JSONPath format that identifies the sub-variant to extract. + * type - An optional string literal naming the SQL type to cast the extracted sub-variant to. + When omitted, it defaults to `variant`. + """, examples = """ Examples: > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a', 'int'); @@ -655,6 +830,13 @@ object VariantGetExpressionBuilder extends VariantGetExpressionBuilderBase(true) // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(v, path[, type]) - Extracts a sub-variant from `v` according to `path`, and then cast the sub-variant to `type`. When `type` is omitted, it is default to `variant`. Returns null if the path does not exist or the cast fails.", + arguments = """ + Arguments: + * v - A variant value to extract from. + * path - A string literal in JSONPath format that identifies the sub-variant to extract. + * type - An optional string literal naming the SQL type to cast the extracted sub-variant to. + When omitted, it defaults to `variant`. + """, examples = """ Examples: > SELECT _FUNC_(parse_json('{"a": 1}'), '$.a', 'int'); @@ -873,6 +1055,11 @@ case class VariantInsert( } } + override def eval(input: InternalRow): Any = { + val _ = foldablePath + super.eval(input) + } + override protected def nullSafeEval(v: Any, p: Any, valValue: Any): Any = { val inputVariant = v.asInstanceOf[VariantVal] foldablePath match { @@ -1051,6 +1238,13 @@ case class VariantSet( val result = super.checkInputDataTypes() if (result.isFailure) { result + } else if (!createIfMissing.foldable) { + DataTypeMismatch( + errorSubClass = "NON_FOLDABLE_INPUT", + messageParameters = Map( + "inputName" -> toSQLId("create_if_missing"), + "inputType" -> toSQLType(createIfMissing.dataType), + "inputExpr" -> toSQLExpr(createIfMissing))) } else if (value.dataType == NullType) { TypeCheckResult.TypeCheckSuccess } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps = false)) { @@ -1186,7 +1380,7 @@ abstract class VariantSetExpressionBuilderBase(failOnError: Boolean) extends Exp path should start with `$` and is followed by one or more segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root path `$` is not allowed. * val - Any expression castable to variant. - * create_if_missing - An optional boolean (default true). + * create_if_missing - An optional boolean (default true). Must be a constant. """, examples = """ Examples: @@ -1227,7 +1421,7 @@ object VariantSetExpressionBuilder extends VariantSetExpressionBuilderBase(true) path should start with `$` and is followed by one or more segments like `[123]`, `.name`, `['name']`, or `["name"]`. The root path `$` is not allowed. * val - Any expression castable to variant. - * create_if_missing - An optional boolean (default true). + * create_if_missing - An optional boolean (default true). Must be a constant. """, examples = """ Examples: @@ -1308,6 +1502,11 @@ case class VariantArrayAppend( } } + override def eval(input: InternalRow): Any = { + val _ = foldablePath + super.eval(input) + } + override protected def nullSafeEval(v: Any, p: Any, valValue: Any): Any = { val inputVariant = v.asInstanceOf[VariantVal] foldablePath match { @@ -1451,7 +1650,94 @@ object VariantArrayAppendExpressionBuilder extends VariantArrayAppendExpressionB ) // scalastyle:on line.size.limit object TryVariantArrayAppendExpressionBuilder - extends VariantArrayAppendExpressionBuilderBase(false) +extends VariantArrayAppendExpressionBuilderBase(false) + +case class VariantStripNulls(child: Expression, includeArrays: Expression) + extends RuntimeReplaceable + with ExpectsInputTypes + with BinaryLike[Expression] + with QueryErrorsBase { + + override def left: Expression = child + override def right: Expression = includeArrays + + override def inputTypes: Seq[AbstractDataType] = Seq(VariantType, BooleanType) + + override def checkInputDataTypes(): TypeCheckResult = { + val result = super.checkInputDataTypes() + if (result.isFailure) { + result + } else if (!includeArrays.foldable) { + DataTypeMismatch( + errorSubClass = "NON_FOLDABLE_INPUT", + messageParameters = Map( + "inputName" -> toSQLId("include_arrays"), + "inputType" -> toSQLType(includeArrays.dataType), + "inputExpr" -> toSQLExpr(includeArrays))) + } else { + TypeCheckResult.TypeCheckSuccess + } + } + + override lazy val replacement: Expression = StaticInvoke( + VariantExpressionEvalUtils.getClass, + VariantType, + "stripNulls", + Seq(child, includeArrays), + inputTypes, + returnNullable = false) + + override def prettyName: String = "variant_strip_nulls" + + override protected def withNewChildrenInternal( + newLeft: Expression, newRight: Expression): VariantStripNulls = + copy(child = newLeft, includeArrays = newRight) +} + +// scalastyle:off line.size.limit +@ExpressionDescription( + usage = "_FUNC_(v[, include_arrays]) - Recursively removes object fields and array elements " + + "whose value is a variant null, unless `include_arrays` is false, in which case null array " + + "elements are kept. Returns NULL if any argument is NULL.", + arguments = """ + Arguments: + * v - The variant value to strip. + * include_arrays - An optional boolean (default true). Must be a constant. + """, + examples = """ + Examples: + > SELECT _FUNC_(parse_json('{"a": 1, "b": null, "c": 3}')); + {"a":1,"c":3} + > SELECT _FUNC_(parse_json('[1, null, 3]')); + [1,3] + > SELECT _FUNC_(parse_json('{"a": {"b": null, "c": [1, null]}}')); + {"a":{"c":[1]}} + > SELECT _FUNC_(parse_json('{"a": [1, null], "b": null}'), false); + {"a":[1,null]} + > SELECT _FUNC_(parse_json('{"a": null}')); + {} + > SELECT _FUNC_(parse_json('null')); + null + > SELECT _FUNC_(NULL); + NULL + """, + since = "4.3.0", + group = "variant_funcs" +) +// scalastyle:on line.size.limit +object VariantStripNullsExpressionBuilder extends ExpressionBuilder { + override def functionSignature: Option[FunctionSignature] = { + val vArg = InputParameter("v") + val includeArraysArg = + InputParameter("include_arrays", Some(Literal.create(true, BooleanType))) + Some(FunctionSignature(Seq(vArg, includeArraysArg))) + } + + override def build(funcName: String, expressions: Seq[Expression]): Expression = { + assert(expressions.size == 2) + VariantStripNulls(expressions(0), expressions(1)) + } +} case class VariantExplode(child: Expression) extends UnaryExpression with Generator with ExpectsInputTypes { @@ -1569,6 +1855,10 @@ object VariantExplode { @ExpressionDescription( usage = "_FUNC_(v) - Returns schema in the SQL format of a variant.", + arguments = """ + Arguments: + * v - A variant value whose schema is returned. + """, examples = """ Examples: > SELECT _FUNC_(parse_json('null')); @@ -1643,13 +1933,24 @@ object SchemaOfVariant { val field = v.getFieldAtIndex(i) fields(i) = StructField(field.key, schemaOf(field.value)) } - // According to the variant spec, object fields must be sorted alphabetically. So we don't - // have to sort, but just need to validate they are sorted. + var utf8Sorted = true + var utf16Sorted = true + var previousKey = if (size > 0) fields(0).name else null + var previousKeyBytes = if (size > 0) VariantUtil.encodeKey(previousKey) else null for (i <- 1 until size) { - if (fields(i - 1).name >= fields(i).name) { - throw new SparkRuntimeException("MALFORMED_VARIANT", Map.empty) - } + val currentKey = fields(i).name + val currentKeyBytes = VariantUtil.encodeKey(currentKey) + utf8Sorted &&= VariantUtil.compareKeys(previousKeyBytes, currentKeyBytes) < 0 + utf16Sorted &&= previousKey.compareTo(currentKey) < 0 + previousKey = currentKey + previousKeyBytes = currentKeyBytes + } + if (!utf8Sorted && !utf16Sorted) { + throw new SparkRuntimeException("MALFORMED_VARIANT", Map.empty) } + // `mergeSchema` expects StructType fields in Java String order. Older Spark values already + // use that order, while spec-compliant values need to be reordered after validation. + java.util.Arrays.sort(fields, JsonInferSchema.structFieldComparator) StructType(fields) case Type.ARRAY => var elementType: DataType = NullType @@ -1684,6 +1985,10 @@ object SchemaOfVariant { // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(v) - Returns the merged schema in the SQL format of a variant column.", + arguments = """ + Arguments: + * v - A variant column whose per-row schemas are merged into a single schema. + """, examples = """ Examples: > SELECT _FUNC_(parse_json(j)) FROM VALUES ('1'), ('2'), ('3') AS tab(j); @@ -1748,6 +2053,10 @@ case class SchemaOfVariantAgg( @ExpressionDescription( usage = "_FUNC_(v) - Returns true if the variant is valid, false if it is malformed, " + "NULL if `v` is NULL.", + arguments = """ + Arguments: + * v - A variant value to validate. + """, examples = """ Examples: > SELECT _FUNC_(parse_json('null')); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/vectorExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/vectorExpressions.scala index e65fae3a2bc2c..ebda73e42a3b3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/vectorExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/vectorExpressions.scala @@ -43,10 +43,16 @@ import org.apache.spark.unsafe.Platform _FUNC_(array1, array2) - Returns the cosine similarity between two float vectors. The vectors must have the same dimension. """, + arguments = """ + Arguments: + * array1 - An ARRAY of FLOAT values representing the first vector. + * array2 - An ARRAY of FLOAT values representing the second vector. It must have the + same dimension as array1. + """, examples = """ Examples: > SELECT _FUNC_(array(1.0F, 2.0F, 3.0F), array(4.0F, 5.0F, 6.0F)); - 0.9746319 + 0.97463185 """, since = "4.2.0", group = "vector_funcs" @@ -101,6 +107,12 @@ case class VectorCosineSimilarity(left: Expression, right: Expression) _FUNC_(array1, array2) - Returns the inner product (dot product) between two float vectors. The vectors must have the same dimension. """, + arguments = """ + Arguments: + * array1 - An ARRAY of FLOAT values representing the first vector. + * array2 - An ARRAY of FLOAT values representing the second vector. It must have the + same dimension as array1. + """, examples = """ Examples: > SELECT _FUNC_(array(1.0F, 2.0F, 3.0F), array(4.0F, 5.0F, 6.0F)); @@ -159,6 +171,12 @@ case class VectorInnerProduct(left: Expression, right: Expression) _FUNC_(array1, array2) - Returns the Euclidean (L2) distance between two float vectors. The vectors must have the same dimension. """, + arguments = """ + Arguments: + * array1 - An ARRAY of FLOAT values representing the first vector. + * array2 - An ARRAY of FLOAT values representing the second vector. It must have the + same dimension as array1. + """, examples = """ Examples: > SELECT _FUNC_(array(1.0F, 2.0F, 3.0F), array(4.0F, 5.0F, 6.0F)); @@ -218,6 +236,12 @@ case class VectorL2Distance(left: Expression, right: Expression) Degree defaults to 2.0 (Euclidean norm) if unspecified. Supported values: 1.0 (L1 norm), 2.0 (L2 norm), float('inf') (infinity norm). """, + arguments = """ + Arguments: + * vector - An ARRAY of FLOAT values representing the vector. + * degree - A FLOAT specifying the norm degree. Defaults to 2.0 when omitted. Supported + values are 1.0 (L1 norm), 2.0 (L2 norm), and float('inf') (infinity norm). + """, examples = """ Examples: > SELECT _FUNC_(array(3.0F, 4.0F), 2.0F); @@ -284,6 +308,13 @@ case class VectorNorm(vector: Expression, degree: Expression) Degree defaults to 2.0 (Euclidean norm) if unspecified. Supported values: 1.0 (L1 norm), 2.0 (L2 norm), float('inf') (infinity norm). """, + arguments = """ + Arguments: + * vector - An ARRAY of FLOAT values representing the vector to normalize. + * degree - A FLOAT specifying the norm degree used for normalization. Defaults to 2.0 + when omitted. Supported values are 1.0 (L1 norm), 2.0 (L2 norm), and float('inf') + (infinity norm). + """, examples = """ Examples: > SELECT _FUNC_(array(3.0F, 4.0F), 2.0F); @@ -555,6 +586,11 @@ trait VectorAggregateBase extends ImperativeAggregate _FUNC_(array) - Returns the element-wise mean of float vectors in a group. All vectors must have the same dimension. """, + arguments = """ + Arguments: + * array - An ARRAY of FLOAT values representing a vector. All vectors aggregated in the + group must have the same dimension. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (array(1.0F, 2.0F)), (array(3.0F, 4.0F)) AS tab(col); @@ -635,6 +671,11 @@ case class VectorAvg( _FUNC_(array) - Returns the element-wise sum of float vectors in a group. All vectors must have the same dimension. """, + arguments = """ + Arguments: + * array - An ARRAY of FLOAT values representing a vector. All vectors aggregated in the + group must have the same dimension. + """, examples = """ Examples: > SELECT _FUNC_(col) FROM VALUES (array(1.0F, 2.0F)), (array(3.0F, 4.0F)) AS tab(col); diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala index f26a4dc9a0c0c..e6c52c08e3b1f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala @@ -410,6 +410,9 @@ object WindowFunctionType { case PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF => PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF case PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF => PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF } + // The incremental aggregator has a single window eval type: the window operator sends each + // frame's rows to the worker, which folds them with `reduce` and produces `finish`. + case _: PythonAggregate => PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala index 5a91d37203754..31aea4910543b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/xmlExpressions.scala @@ -38,6 +38,13 @@ import org.apache.spark.unsafe.types.UTF8String // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(xmlStr, schema[, options]) - Returns a struct value with the given `xmlStr` and `schema`.", + arguments = """ + Arguments: + * xmlStr - A string expression containing a single XML record to parse. + * schema - The schema of the output struct, as a DDL-formatted string or a schema + expression. + * options - Optional. A map of key-value pairs controlling how the XML is parsed. + """, examples = """ Examples: > SELECT _FUNC_('<p><a>1</a><b>0.8</b></p>', 'a INT, b DOUBLE'); @@ -96,6 +103,7 @@ case class XmlToStructs( @transient private lazy val evaluator: XmlToStructsEvaluator = XmlToStructsEvaluator(options, nullableSchema, nameOfCorruptRecord, timeZoneId, child) + override def stateful: Boolean = true private val nameOfCorruptRecord = SQLConf.get.getConf(SQLConf.COLUMN_NAME_OF_CORRUPT_RECORD) @@ -139,6 +147,11 @@ case class XmlToStructs( */ @ExpressionDescription( usage = "_FUNC_(xml[, options]) - Returns schema in the DDL format of XML string.", + arguments = """ + Arguments: + * xml - A foldable string expression containing an XML record whose schema is inferred. + * options - Optional. A map of key-value pairs controlling how the XML is parsed. + """, examples = """ Examples: > SELECT _FUNC_('<p><a>1</a></p>'); @@ -223,6 +236,11 @@ case class SchemaOfXml( // scalastyle:off line.size.limit @ExpressionDescription( usage = "_FUNC_(expr[, options]) - Returns a XML string with a given struct value", + arguments = """ + Arguments: + * expr - A struct-valued expression to convert to an XML string. + * options - Optional. A map of key-value pairs controlling how the XML is generated. + """, examples = """ Examples: > SELECT _FUNC_(named_struct('a', 1, 'b', 2)); @@ -278,6 +296,7 @@ case class StructsToXml( @transient private lazy val evaluator = StructsToXmlEvaluator(options, child.dataType, timeZoneId) + override def stateful: Boolean = true override def withTimeZone(timeZoneId: String): TimeZoneAwareExpression = copy(timeZoneId = Option(timeZoneId)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala new file mode 100644 index 0000000000000..fb6f50627840a --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentiles.scala @@ -0,0 +1,225 @@ +/* + * 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.catalyst.optimizer + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression, ExprId, GetArrayItem, LeafExpression, Literal, NamedExpression} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, ApproximatePercentile} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE +import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, DoubleType} + +private[optimizer] case class PercentileFusionIdentity( + aggregateFunctions: Seq[Expression], + mode: AggregateMode, + isDistinct: Boolean, + filter: Option[Expression], + percentageBits: Seq[Long]) + +/** + * Foldable percentage array that retains the original scalar aggregate structures in equality. + * + * Fusion removes those structures from the physical aggregate. Keeping them here prevents + * subquery or exchange reuse from equating plans that were distinct before fusion. + */ +private[optimizer] case class PercentileFusionArray(identity: PercentileFusionIdentity) + extends LeafExpression with CodegenFallback { + override def foldable: Boolean = true + override def nullable: Boolean = false + override def dataType: ArrayType = ArrayType(DoubleType, containsNull = false) + + private lazy val value = new GenericArrayData( + identity.percentageBits.map(java.lang.Double.longBitsToDouble)) + private lazy val literal = Literal(value, dataType) + + override def eval(input: InternalRow): Any = value + override def toString: String = literal.toString + override def sql: String = literal.sql +} + +/** + * Combines scalar approximate percentiles that can share the same percentile digest. + * + * An approximate percentile digest depends on its input, accuracy, filter, distinctness, and + * aggregate mode, but not on the percentile requested from the completed digest. Consequently, + * compatible scalar percentiles can be calculated by one array-valued aggregate and projected + * back to their original scalar outputs. + * + * Inputs and filters must retain their original expression structure so that floating-point + * evaluation and ANSI overflow behavior are preserved. Streaming aggregates are left unchanged + * to preserve the value schemas of existing checkpoints. + */ +object CombineApproximatePercentiles extends Rule[LogicalPlan] { + + private case class CompatibilityKey( + child: Expression, + accuracy: Long, + mode: AggregateMode, + isDistinct: Boolean, + filter: Option[Expression]) + + private case class PhysicalCompatibilityKey( + child: Expression, + percentage: Expression, + accuracy: Expression, + mode: AggregateMode, + isDistinct: Boolean, + filter: Option[Expression]) + + private def structurallyNormalize( + expression: Expression, + inputOrdinals: scala.collection.Map[ExprId, Int]): Expression = expression.transformUp { + case attribute: AttributeReference => + inputOrdinals.get(attribute.exprId) match { + case Some(ordinal) => AttributeReference("none", attribute.dataType)(ExprId(ordinal)) + case None => attribute + } + } + + private def physicalCompatibilityKey( + key: CompatibilityKey, + percentile: ApproximatePercentile): PhysicalCompatibilityKey = PhysicalCompatibilityKey( + key.child.canonicalized, + percentile.percentageExpression.canonicalized, + percentile.accuracyExpression.canonicalized, + key.mode, + key.isDistinct, + key.filter.map(_.canonicalized)) + + private def hasSafePhysicalFusion( + expressions: scala.collection.Iterable[AggregateExpression]): Boolean = { + val physicalGroups = expressions.groupBy(_.canonicalized) + // PhysicalAggregation already shares a digest within each canonical group. Fusion must both + // remove a digest and preserve cases where canonical percentages evaluate differently. + physicalGroups.sizeCompare(1) > 0 && physicalGroups.values.forall { group => + val percentages = group.iterator.map { expression => + expression.aggregateFunction + .asInstanceOf[ApproximatePercentile] + .percentageExpression + .eval() + } + val first = percentages.next() + percentages.forall(_ == first) + } + } + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED)) return plan + + plan.transformUpWithPruning(_.containsPattern(AGGREGATE), ruleId) { + case aggregate: Aggregate if aggregate.resolved && !aggregate.isStreaming => + combine(aggregate) + } + } + + private def combine(aggregate: Aggregate): Aggregate = { + val compatible = mutable.LinkedHashMap.empty[ + CompatibilityKey, mutable.ArrayBuffer[AggregateExpression]] + // PhysicalAggregation deduplicates semantically equivalent aggregates. Track every logical + // key that shares a physical key so fusion does not change that existing deduplication. + val physicalCompatibilityKeys = mutable.HashMap.empty[ + PhysicalCompatibilityKey, mutable.HashSet[CompatibilityKey]] + + aggregate.aggregateExpressions.foreach(_.foreach { + case expression @ AggregateExpression( + percentile: ApproximatePercentile, mode, isDistinct, filter, _) + if percentile.child.deterministic && + filter.forall(_.deterministic) => + val key = CompatibilityKey( + percentile.child, + // Analysis already validates that accuracy is foldable, non-null, and in range. + percentile.accuracyExpression.eval().asInstanceOf[Number].longValue, + mode, + isDistinct, + filter) + physicalCompatibilityKeys.getOrElseUpdate( + physicalCompatibilityKey(key, percentile), + mutable.HashSet.empty) += key + if (percentile.percentageExpression.dataType == DoubleType) { + compatible.getOrElseUpdate(key, mutable.ArrayBuffer.empty) += expression + } + case _ => + }) + + val replacements = mutable.HashMap.empty[ExprId, (AggregateExpression, Int)] + lazy val inputOrdinals = { + val ordinals = mutable.HashMap.empty[ExprId, Int] + aggregate.child.output.zipWithIndex.foreach { case (attribute, ordinal) => + ordinals.getOrElseUpdate(attribute.exprId, ordinal) + } + ordinals + } + compatible.iterator.map { case (key, expressions) => + key -> expressions.distinctBy(_.resultId) + }.filter { case (key, expressions) => + hasSafePhysicalFusion(expressions) && expressions.forall { expression => + val percentile = expression.aggregateFunction.asInstanceOf[ApproximatePercentile] + val physicalKey = physicalCompatibilityKey(key, percentile) + // OptimizeOneRowPlan can erase DISTINCT after fusion. Across distinctness boundaries, + // canonical matches are safe only when their original inputs and filters also match. + physicalCompatibilityKeys(physicalKey).sizeCompare(1) == 0 && + physicalCompatibilityKeys + .get(physicalKey.copy(isDistinct = !physicalKey.isDistinct)) + .forall(_.forall(other => other.child == key.child && other.filter == key.filter)) + } + }.foreach { case (key, expressions) => + val first = expressions.head + val percentile = first.aggregateFunction.asInstanceOf[ApproximatePercentile] + val percentages = expressions.map { expression => + expression.aggregateFunction + .asInstanceOf[ApproximatePercentile] + .percentageExpression + } + val percentageValues = percentages.map(_.eval().asInstanceOf[Double]).toSeq + val identity = PercentileFusionIdentity( + expressions.map { expression => + structurallyNormalize(expression.aggregateFunction, inputOrdinals) + }.toSeq, + key.mode, + key.isDistinct, + key.filter.map(structurallyNormalize(_, inputOrdinals)), + percentageValues.map(java.lang.Double.doubleToRawLongBits)) + val combinedFunction = percentile.copy(percentageExpression = PercentileFusionArray(identity)) + combinedFunction.copyTagsFrom(percentile) + val combined = first.copy( + aggregateFunction = combinedFunction, resultId = NamedExpression.newExprId) + expressions.zipWithIndex.foreach { case (expression, index) => + replacements.put(expression.resultId, (combined, index)) + } + } + + if (replacements.isEmpty) { + aggregate + } else { + val rewrittenExpressions = aggregate.aggregateExpressions.map { expression => + expression.transformUp { + case original: AggregateExpression if replacements.contains(original.resultId) => + val (combined, index) = replacements(original.resultId) + GetArrayItem(combined, Literal(index), failOnError = false) + }.asInstanceOf[NamedExpression] + } + aggregate.copy(aggregateExpressions = rewrittenExpressions) + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala index 5c1967c094ffa..34c7aa442d4ee 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala @@ -59,6 +59,9 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] { if (idx >= 0 && idx < elems.size) { // valid index elems(idx) + } else if (ga.failOnError) { + // Keep ANSI runtime behavior, which raises on out-of-bounds access. + ga } else { // out of bounds, mimic the runtime behavior and return null Literal(null, ga.dataType) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala index c6a7940b93aaf..f71c4934945b1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuery.scala @@ -725,9 +725,21 @@ object DecorrelateInnerQuery extends PredicateHelper { if (partitionFields.isEmpty) { // Underlying subquery has no predicates connecting inner and outer query. // In this case, limit can be computed over the inner query directly. + // The ORDER BY was peeled off the Sort above; re-apply it as a global Sort below + // the limit so that ORDER BY ... LIMIT (and ORDER BY ... LIMIT ... OFFSET) is + // order-preserving. Otherwise the ordering is dropped and the limit returns an + // arbitrary (non-deterministic) rows. + val orderedChild = + if (ordering.nonEmpty && !SQLConf.get.getConf( + SQLConf.DECORRELATE_LIMIT_OFFSET_LEGACY_INCORRECT_ORDER_HANDLING_ENABLED)) { + Sort(replaceOuterReferences(ordering, outerReferenceMap), global = true, newChild) + } else { + newChild + } offsetExpr match { - case IntegerLiteral(0) => (Limit(limit, newChild), joinCond, outerReferenceMap) - case _ => (Limit(limit, Offset(offsetExpr, newChild)), joinCond, outerReferenceMap) + case IntegerLiteral(0) => (Limit(limit, orderedChild), joinCond, outerReferenceMap) + case _ => + (Limit(limit, Offset(offsetExpr, orderedChild)), joinCond, outerReferenceMap) } } else { val orderByFields = replaceOuterReferences(ordering, outerReferenceMap) @@ -1006,7 +1018,7 @@ object DecorrelateInnerQuery extends PredicateHelper { // predicate does, and the correlations can not be replaced via equivalences. // Introduce a domain join on the left side of the join // (chosen arbitrarily) to provide values for the correlated attribute reference. - shouldPushToLeft = true; + shouldPushToLeft = true } val (newLeft, leftJoinCond, leftOuterReferenceMap) = if (shouldPushToLeft) { decorrelate(left, newOuterReferences, aggregated, underSetOp) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala index a8f5b3ca67ed0..aef521e8d0a7d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InjectRuntimeFilter.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.{INVOKE, JSON_TO_STRUCT, LIKE_FAMLIY, PYTHON_UDF, REGEXP_EXTRACT_FAMILY, REGEXP_REPLACE, SCALA_UDF} +import org.apache.spark.sql.catalyst.util.UnsafeRowUtils import org.apache.spark.sql.internal.SQLConf /** @@ -36,29 +37,42 @@ import org.apache.spark.sql.internal.SQLConf */ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with JoinSelectionHelper { + private case class FilterCreationSide( + key: Expression, + plan: LogicalPlan, + useMaterializedThreshold: Boolean, + materializedRowCount: Option[BigInt] = None, + materializedSizeInBytes: Option[BigInt] = None) + private def injectFilter( filterApplicationSideKey: Expression, filterApplicationSidePlan: LogicalPlan, - filterCreationSideKey: Expression, - filterCreationSidePlan: LogicalPlan): LogicalPlan = { + filterCreationSide: FilterCreationSide): LogicalPlan = { injectBloomFilter( filterApplicationSideKey, filterApplicationSidePlan, - filterCreationSideKey, - filterCreationSidePlan + filterCreationSide ) } private def injectBloomFilter( filterApplicationSideKey: Expression, filterApplicationSidePlan: LogicalPlan, - filterCreationSideKey: Expression, - filterCreationSidePlan: LogicalPlan): LogicalPlan = { + filterCreationSide: FilterCreationSide): LogicalPlan = { + val filterCreationSideKey = filterCreationSide.key + val filterCreationSidePlan = filterCreationSide.plan + val creationSideThreshold = if (filterCreationSide.useMaterializedThreshold) { + conf.runtimeFilterMaterializedCreationSideThreshold + } else { + conf.runtimeFilterCreationSideThreshold + } // Skip if the filter creation side is too big - if (filterCreationSidePlan.stats.sizeInBytes > conf.runtimeFilterCreationSideThreshold) { + if (filterCreationSide.materializedSizeInBytes + .getOrElse(filterCreationSidePlan.stats.sizeInBytes) > creationSideThreshold) { return filterApplicationSidePlan } - val rowCount = filterCreationSidePlan.stats.rowCount + val rowCount = filterCreationSide.materializedRowCount + .orElse(filterCreationSidePlan.stats.rowCount) val bloomFilterAgg = if (rowCount.isDefined && rowCount.get.longValue > 0L) { new BloomFilterAggregate(new XxHash64(Seq(filterCreationSideKey)), rowCount.get.longValue) @@ -81,23 +95,24 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } /** - * Extracts a sub-plan which is a simple filter over scan from the input plan. The simple - * filter should be selective and the filter condition (including expressions in the child - * plan referenced by the filter condition) should be a simple expression, so that we do - * not add a subquery that might have an expensive computation. The extracted sub-plan should - * produce a superset of the entire creation side output data, so that it's still correct to - * use the sub-plan to build the runtime filter to prune the application side. + * Extracts either a safely materialized leaf with accurate statistics or a simple selective + * filter over a scan. Filter conditions and the expressions they reference must remain simple, + * so the runtime-filter subquery does not introduce expensive computation. The extracted plan + * must produce a superset of the creation side's join keys. */ private def extractSelectiveFilterOverScan( plan: LogicalPlan, - filterCreationSideKey: Expression): Option[(Expression, LogicalPlan)] = { + filterCreationSideKey: Expression, + allowMaterializedCache: Boolean, + applicationDistinctCount: => Option[BigInt], + onMaterializedLeaf: => Unit = ()): Option[FilterCreationSide] = { def extract( p: LogicalPlan, predicateReference: AttributeSet, hasHitFilter: Boolean, hasHitSelectiveFilter: Boolean, currentPlan: LogicalPlan, - targetKey: Expression): Option[(Expression, LogicalPlan)] = p match { + targetKey: Expression): Option[FilterCreationSide] = p match { case Project(projectList, child) if hasHitFilter => // We need to make sure all expressions referenced by filter predicates are simple // expressions. @@ -175,8 +190,55 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J } else { None } + case leaf: MaterializedLeafNode => + onMaterializedLeaf + val safeLineage = currentPlan.deterministic && + findExpressionAndTrackLineageDown(targetKey, currentPlan).exists { + case (trackedKey, _) => isSimpleExpression(trackedKey) + } + val materializedMetadata = if (allowMaterializedCache && safeLineage && + leaf.mayHaveUsableMaterializedStats && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.isDefined)) { + leaf.materializedMetadata.filter(_.statsAvailable).flatMap { metadata => + val creationSize = if (currentPlan eq leaf) { + metadata.sizeInBytes + } else { + currentPlan.stats.sizeInBytes + } + Option.when(creationSize <= conf.runtimeFilterMaterializedCreationSideThreshold) { + metadata -> creationSize + } + } + } else { + None + } + materializedMetadata match { + case Some((metadata, creationSize)) => + val rowCount = metadata.rowCount + Option.when( + rowCount <= conf.getConf(SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS) && + (hasHitSelectiveFilter || leaf.hasSelectivePredicate || + applicationDistinctCount.exists(_ > rowCount))) { + FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = true, + materializedRowCount = Some(rowCount), + materializedSizeInBytes = Some(creationSize)) + } + case None if hasHitSelectiveFilter => + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false)) + case _ => None + } case _: LeafNode if hasHitSelectiveFilter => - Some((targetKey, currentPlan)) + Some(FilterCreationSide( + targetKey, + currentPlan, + useMaterializedThreshold = false)) case _ => None } @@ -237,18 +299,81 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J * Extracts the beneficial filter creation plan with check show below: * - The filterApplicationSideKey can be pushed down through joins, aggregates and windows * (ie the expression references originate from a single leaf node) - * - The filter creation side has a selective predicate + * - The filter creation side has a selective predicate, or its exact materialized row count + * is smaller than the application side's distinct join-key count * - The max filterApplicationSide scan size is greater than a configurable threshold */ private def extractBeneficialFilterCreatePlan( filterApplicationSide: LogicalPlan, filterCreationSide: LogicalPlan, filterApplicationSideKey: Expression, - filterCreationSideKey: Expression): Option[(Expression, LogicalPlan)] = { + filterCreationSideKey: Expression): Option[FilterCreationSide] = { if (findExpressionAndTrackLineageDown( filterApplicationSideKey, filterApplicationSide).isDefined && satisfyByteSizeRequirement(filterApplicationSide)) { - extractSelectiveFilterOverScan(filterCreationSide, filterCreationSideKey) + val allowMaterializedCache = UnsafeRowUtils.isBinaryStable(filterCreationSideKey.dataType) && + UnsafeRowUtils.isBinaryStable(filterApplicationSideKey.dataType) + def distinctCount(key: Expression, plan: LogicalPlan): Option[BigInt] = key match { + case attribute: Attribute => + plan.stats.attributeStats.get(attribute).flatMap(_.distinctCount) + case _ => None + } + def hasOnlyJoinKeyNullChecksOverScan( + plan: LogicalPlan, + targetKey: Expression): Boolean = plan match { + case project: Project => + hasOnlyJoinKeyNullChecksOverScan( + project.child, replaceAlias(targetKey, getAliasMap(project))) + case Filter(condition, child) => + splitConjunctivePredicates(condition).forall { + case IsNotNull(expression) => expression.semanticEquals(targetKey) + case _ => false + } && hasOnlyJoinKeyNullChecksOverScan(child, targetKey) + case _: LeafNode => true + case _ => false + } + lazy val currentDistinctCount = + distinctCount(filterApplicationSideKey, filterApplicationSide) + lazy val lineageDistinctCount = findExpressionAndTrackLineageDown( + filterApplicationSideKey, filterApplicationSide).flatMap { + case (trackedKey, origin) => distinctCount(trackedKey, origin) + } + lazy val applicationDistinctCount = { + if (hasOnlyJoinKeyNullChecksOverScan( + filterApplicationSide, filterApplicationSideKey)) { + lineageDistinctCount.orElse(currentDistinctCount) + } else { + currentDistinctCount + } + } + if (allowMaterializedCache) { + var sawMaterializedLeaf = false + val selectiveCreationSide = extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None, + onMaterializedLeaf = { sawMaterializedLeaf = true }) + selectiveCreationSide + .filter(_.plan.stats.sizeInBytes <= conf.runtimeFilterCreationSideThreshold) + .orElse { + if (sawMaterializedLeaf) { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = true, + applicationDistinctCount = applicationDistinctCount) + } else { + selectiveCreationSide + } + } + } else { + extractSelectiveFilterOverScan( + filterCreationSide, + filterCreationSideKey, + allowMaterializedCache = false, + applicationDistinctCount = None) + } } else { None } @@ -307,9 +432,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J val hasShuffle = isProbablyShuffleJoin(left, right, hint) if (canPruneLeft(joinType) && (hasShuffle || probablyHasShuffle(left)) && !hasBloomFilter(newLeft, l)) { - extractBeneficialFilterCreatePlan(left, right, l, r).foreach { - case (filterCreationSideKey, filterCreationSidePlan) => - newLeft = injectFilter(l, newLeft, filterCreationSideKey, filterCreationSidePlan) + extractBeneficialFilterCreatePlan(left, right, l, r).foreach { creationSide => + newLeft = injectFilter(l, newLeft, creationSide) } } // Did we actually inject on the left? If not, try on the right @@ -320,10 +444,8 @@ object InjectRuntimeFilter extends Rule[LogicalPlan] with PredicateHelper with J // 3. There is no bloom filter on the right key yet if (newLeft.fastEquals(oldLeft) && canPruneRight(joinType) && (hasShuffle || probablyHasShuffle(right)) && !hasBloomFilter(newRight, r)) { - extractBeneficialFilterCreatePlan(right, left, r, l).foreach { - case (filterCreationSideKey, filterCreationSidePlan) => - newRight = injectFilter( - r, newRight, filterCreationSideKey, filterCreationSidePlan) + extractBeneficialFilterCreatePlan(right, left, r, l).foreach { creationSide => + newRight = injectFilter(r, newRight, creationSide) } } if (!newLeft.fastEquals(oldLeft) || !newRight.fastEquals(oldRight)) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala index c8bd01ed8849f..1b36ba04cfd79 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTE.scala @@ -40,10 +40,19 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, PLAN_EXPRESSION} * @param alwaysInline if true, inline all CTEs in the query plan. * @param keepDanglingRelations if true, dangling CTE relations will be kept in the original * `WithCTE` node. + * @param isAnalysis if true, this rule runs during analysis (e.g. from `CheckAnalysis`), where + * the plan may be a subplan that references `CTERelationDef`s owned by a + * surrounding scope; such out-of-scope references are tolerated and left for + * the owning scope to resolve. Defaults to false: every other caller (the + * optimizer, `ProgressReporter`) operates on a complete plan, so a + * `CTERelationRef` with no definition in the plan indicates corruption and + * raises an error rather than being silently dropped. Only the scoped + * `CheckAnalysis` call opts into the tolerant behavior. */ case class InlineCTE( alwaysInline: Boolean = false, - keepDanglingRelations: Boolean = false) extends Rule[LogicalPlan] { + keepDanglingRelations: Boolean = false, + isAnalysis: Boolean = false) extends Rule[LogicalPlan] { override def apply(plan: LogicalPlan): LogicalPlan = { if (!plan.isInstanceOf[Subquery] && plan.containsPattern(CTE)) { @@ -148,18 +157,36 @@ case class InlineCTE( buildCTEMap(child, cteMap, outerCTEId) } - case ref: CTERelationRef => - cteMap(ref.cteId) = cteMap(ref.cteId).withRefCountIncreased(1) + case ref: CTERelationRef => cteMap.get(ref.cteId) match { + case Some(refInfo) => + cteMap(ref.cteId) = refInfo.withRefCountIncreased(1) - // The `outerCTEId` CTE definition can either reference `cteId` definition if `cteId` is in - // the same or in an outer `WithCTE` node, or `outerCTEId` can contain `cteId` definition if - // `cteId` is an inner `WithCTE` node inside `outerCTEId`. - // In both cases we can track the relations in `outgoingRefs` when we see a definition the - // first time. But if we encounter a conflicting duplicated contains relation later, then we - // will remove the references of the first contains relation. - outerCTEId.foreach { cteId => - cteMap(cteId).increaseOutgoingRefCount(ref.cteId, 1) - } + // The `outerCTEId` CTE definition can either reference `cteId` definition if `cteId` is + // in the same or in an outer `WithCTE` node, or `outerCTEId` can contain `cteId` + // definition if `cteId` is an inner `WithCTE` node inside `outerCTEId`. + // In both cases we can track the relations in `outgoingRefs` when we see a definition the + // first time. But if we encounter a conflicting duplicated contains relation later, then + // we will remove the references of the first contains relation. + outerCTEId.foreach { cteId => + cteMap(cteId).increaseOutgoingRefCount(ref.cteId, 1) + } + + case None => + // The referenced CTE definition is not present in this plan. During analysis + // (`isAnalysis` = true) this legitimately happens when a check runs on a subplan that + // contains the reference but not its enclosing `WithCTE` -- e.g. + // `ResolveSQLTableFunctions` calls `checkAnalysis` (which runs `InlineCTE`) on a resolved + // SQL table function whose argument is a scalar subquery referencing an outer CTE. The + // reference is resolved by the scope that owns the definition, so there is nothing to + // count here. Otherwise (`isAnalysis` = false) the plan is complete, so a missing + // definition indicates corruption -- fail loudly rather than silently dropping the + // reference. + if (!isAnalysis) { + throw SparkException.internalError( + "No CTERelationDef found for CTERelationRef with id " + + s"${ref.cteId} while building the CTE map.") + } + } case _ => if (plan.containsPattern(CTE)) { @@ -229,39 +256,50 @@ case class InlineCTE( WithCTE(inlined, notInlined) } - case ref: CTERelationRef => - val refInfo = cteMap(ref.cteId) + case ref: CTERelationRef => cteMap.get(ref.cteId) match { + case None => + // Out-of-scope reference whose definition is not in this plan (mirrors the guard in + // `buildCTEMap`). During analysis it is left unchanged for the scope that owns the + // definition; otherwise a missing definition is corruption. + if (!isAnalysis) { + throw SparkException.internalError( + "No CTERelationDef found for CTERelationRef with id " + + s"${ref.cteId} while inlining CTEs.") + } + ref - val cteBody = if (ref.isUnlimitedRecursion) { - setUnlimitedRecursion(refInfo.cteDef.child, ref.cteId) - } else { - refInfo.cteDef.child - } - if (refInfo.shouldInline) { - if (ref.outputSet == refInfo.cteDef.outputSet) { - cteBody + case Some(refInfo) => + val cteBody = if (ref.isUnlimitedRecursion) { + setUnlimitedRecursion(refInfo.cteDef.child, ref.cteId) } else { - val ctePlan = DeduplicateRelations( - Join( - cteBody, - cteBody, - Inner, - None, - JoinHint(None, None) - ) - ).children(1) - val projectList = ref.output.zip(ctePlan.output).map { case (tgtAttr, srcAttr) => - if (srcAttr.semanticEquals(tgtAttr)) { - tgtAttr - } else { - Alias(srcAttr, tgtAttr.name)(exprId = tgtAttr.exprId) + refInfo.cteDef.child + } + if (refInfo.shouldInline) { + if (ref.outputSet == refInfo.cteDef.outputSet) { + cteBody + } else { + val ctePlan = DeduplicateRelations( + Join( + cteBody, + cteBody, + Inner, + None, + JoinHint(None, None) + ) + ).children(1) + val projectList = ref.output.zip(ctePlan.output).map { case (tgtAttr, srcAttr) => + if (srcAttr.semanticEquals(tgtAttr)) { + tgtAttr + } else { + Alias(srcAttr, tgtAttr.name)(exprId = tgtAttr.exprId) + } } + Project(projectList, ctePlan) } - Project(projectList, ctePlan) + } else { + ref } - } else { - ref - } + } case _ if plan.containsPattern(CTE) => plan diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala index 9e613c54a49bd..e2043e06186f4 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortExpression.scala @@ -17,33 +17,49 @@ package org.apache.spark.sql.catalyst.optimizer +import scala.collection.immutable.VectorMap import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{Alias, ArrayTransform, CreateNamedStruct, Expression, GetStructField, If, IsNull, LambdaFunction, Literal, MapFromArrays, MapKeys, MapSort, MapValues, NamedExpression, NamedLambdaVariable} +import org.apache.spark.sql.catalyst.expressions.{Alias, ArrayTransform, Attribute, CreateNamedStruct, Expression, GetStructField, If, IsNull, LambdaFunction, Literal, MapFromArrays, MapKeys, MapSort, MapValues, NamedExpression, NamedLambdaVariable} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project, RepartitionByExpression} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, REPARTITION_OPERATION} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, MapType, StructType} import org.apache.spark.util.ArrayImplicits.SparkArrayOps /** - * Adds [[MapSort]] to [[Aggregate]] expressions containing map columns, - * as the key/value pairs need to be in the correct order before grouping: + * Adds [[MapSort]] to grouping expressions and distinct aggregate arguments that contain maps, + * ensuring key/value pairs have a consistent order before aggregation: * * SELECT map_column, COUNT(*) FROM TABLE GROUP BY map_column => * SELECT _groupingmapsort as map_column, COUNT(*) FROM ( * SELECT map_sort(map_column) as _groupingmapsort FROM TABLE * ) GROUP BY _groupingmapsort + * + * SELECT COUNT(DISTINCT map_column) FROM TABLE => + * SELECT COUNT(DISTINCT _distinctmapsort) FROM ( + * SELECT map_sort(map_column) as _distinctmapsort FROM TABLE + * ) + * + * Distinct arguments are normalized here instead of in [[RewriteDistinctAggregates]] because a + * single distinct group without a filter bypasses that rule and is handled by the physical planner. */ -object InsertMapSortInGroupingExpressions extends Rule[LogicalPlan] { +object InsertMapSortInAggregate extends Rule[LogicalPlan] { import InsertMapSortExpression._ override def apply(plan: LogicalPlan): LogicalPlan = { if (!plan.containsPattern(AGGREGATE)) { return plan } + val normalizeDistinctAggregates = + conf.getConf(SQLConf.INSERT_MAP_SORT_IN_DISTINCT_AGGREGATES_ENABLED) val shouldRewrite = plan.exists { - case agg: Aggregate if agg.groupingExpressions.exists(mapTypeExistsRecursively) => true + case agg: Aggregate => + agg.groupingExpressions.exists(mapTypeExistsRecursively) || + (normalizeDistinctAggregates && + distinctAggregateChildren(agg.aggregateExpressions).exists(mapTypeExistsRecursively)) case _ => false } if (!shouldRewrite) { @@ -52,32 +68,104 @@ object InsertMapSortInGroupingExpressions extends Rule[LogicalPlan] { plan transformUpWithNewOutput { case agg @ Aggregate(groupingExprs, aggregateExpressions, child, hint) => - val exprToMapSort = new mutable.HashMap[Expression, NamedExpression] - val newGroupingKeys = groupingExprs.map { expr => - val inserted = insertMapSortRecursively(expr) - if (expr.ne(inserted)) { - exprToMapSort.getOrElseUpdate( - expr.canonicalized, - Alias(inserted, "_groupingmapsort")() - ).toAttribute + val distinctExpressions = if (normalizeDistinctAggregates) { + distinctAggregateChildren(aggregateExpressions) + } else { + Seq.empty + } + val expressionsToNormalize = groupingExprs ++ distinctExpressions + if (!expressionsToNormalize.exists(mapTypeExistsRecursively)) { + agg -> Nil + } else { + val groupingMapSortAliases = insertMapSortInExpressions( + expressions = groupingExprs, + aliasName = "_groupingmapsort") + val distinctInputAliases = createDistinctInputAliases(distinctExpressions) + val distinctMapSortAliases = insertMapSortInExpressions( + expressions = distinctExpressions, + aliasName = "_distinctmapsort", + reusableAliases = groupingMapSortAliases, + inputAliases = distinctInputAliases) + val newGroupingKeys = groupingExprs.map { expr => + groupingMapSortAliases.get(expr.canonicalized).map(_.toAttribute).getOrElse(expr) + } + val newAggregateExprs = aggregateExpressions.map { + case named if groupingMapSortAliases.contains(named.canonicalized) => + // If we replace the top-level named expr, then should add back the original name + groupingMapSortAliases(named.canonicalized).toAttribute.withName(named.name) + case other => + // This must be top-down so distinct arguments are normalized before their children. + // Continuing downward also substitutes grouping aliases in other arguments. + other.transformDown { + case ae: AggregateExpression if ae.isDistinct => + ae.copy(aggregateFunction = ae.aggregateFunction.withNewChildren( + ae.aggregateFunction.children.map { child => + distinctMapSortAliases.get(child.canonicalized) + .map(_.toAttribute).getOrElse(child) + }).asInstanceOf[AggregateFunction]) + case e => + groupingMapSortAliases.get(e.canonicalized).map(_.toAttribute).getOrElse(e) + }.asInstanceOf[NamedExpression] + } + val distinctInput = if (distinctInputAliases.nonEmpty) { + // Project complex inputs once because recursive normalization can reference them + // repeatedly. The operator optimization batch later removes these temporary projection + // layers with CollapseProject and ColumnPruning. + Project(child.output ++ distinctInputAliases.values, child) } else { - expr + child } + val newChild = Project( + child.output ++ (groupingMapSortAliases ++ distinctMapSortAliases).values, + distinctInput) + val newAgg = Aggregate(newGroupingKeys, newAggregateExprs, newChild, hint) + newAgg -> agg.output.zip(newAgg.output) } - val newAggregateExprs = aggregateExpressions.map { - case named if exprToMapSort.contains(named.canonicalized) => - // If we replace the top-level named expr, then should add back the original name - exprToMapSort(named.canonicalized).toAttribute.withName(named.name) - case other => - other.transformUp { - case e => exprToMapSort.get(e.canonicalized).map(_.toAttribute).getOrElse(e) - }.asInstanceOf[NamedExpression] - } - val newChild = Project(child.output ++ exprToMapSort.values, child) - val newAgg = Aggregate(newGroupingKeys, newAggregateExprs, newChild, hint) - newAgg -> agg.output.zip(newAgg.output) } } + + private def distinctAggregateChildren( + aggregateExpressions: Seq[NamedExpression]): Seq[Expression] = { + aggregateExpressions + .flatMap(_.collect { + case ae: AggregateExpression if ae.isDistinct => ae + }) + .flatMap(_.aggregateFunction.children) + } + + private def createDistinctInputAliases( + expressions: Seq[Expression]): VectorMap[Expression, NamedExpression] = { + val aliases = new mutable.LinkedHashMap[Expression, NamedExpression] + expressions.foreach { + case _: Attribute => + case expr if mapTypeExistsRecursively(expr) => + aliases.getOrElseUpdate( + expr.canonicalized, + Alias(expr, "_distinctaggregateexpression")()) + case _ => + } + VectorMap.from(aliases) + } + + private def insertMapSortInExpressions( + expressions: Seq[Expression], + aliasName: String, + reusableAliases: Map[Expression, NamedExpression] = Map.empty, + inputAliases: Map[Expression, NamedExpression] = Map.empty) + : VectorMap[Expression, NamedExpression] = { + val aliases = new mutable.LinkedHashMap[Expression, NamedExpression] + expressions.foreach { expr => + val canonicalized = expr.canonicalized + val input = inputAliases.get(canonicalized).map(_.toAttribute).getOrElse(expr) + val inserted = insertMapSortRecursively(input) + if (input.ne(inserted)) { + aliases.getOrElseUpdate( + canonicalized, + reusableAliases.getOrElse(canonicalized, Alias(inserted, aliasName)())) + } + } + VectorMap.from(aliases) + } } /** diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala new file mode 100644 index 0000000000000..9ec2227f969cc --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansReferences.scala @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.optimizer + +import org.apache.spark.sql.catalyst.expressions.{Attribute, ExprId, LeafExpression, Unevaluable} +import org.apache.spark.sql.catalyst.plans.logical.LeafNode +import org.apache.spark.sql.catalyst.trees.TreePattern.{NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY_REFERENCE, TreePattern} +import org.apache.spark.sql.types.DataType + +// The temporary reference placeholders below are produced by the `MergeSubplans` rule (now in +// sql/core) but must remain in catalyst: `ScalarSubqueryReference` is referenced by the catalyst +// expression `BloomFilterMightContain`, and catalyst cannot depend on sql/core. + +/** + * Temporary reference to a subquery which is added to a `PlanMerger`. + * + * @param level The level of the replaced subquery. It defines the `PlanMerger` instance into which + * the subquery is merged. + * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`. + * @param outputIndex The index of the output attribute of the merged plan. + * @param dataType The data type of the original scalar subquery. + * @param exprId The expression id of the original scalar subquery. + */ +case class ScalarSubqueryReference( + level: Int, + mergedPlanIndex: Int, + outputIndex: Int, + override val dataType: DataType, + exprId: ExprId) extends LeafExpression with Unevaluable { + override def nullable: Boolean = true + + final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE) +} + +/** + * Temporary reference to a non-grouping aggregate which is added to a `PlanMerger`. + * + * @param level The level of the replaced aggregate. It defines the `PlanMerger` instance into which + * the aggregate is merged. + * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`. + * @param outputIndices The indices of the output attributes of the merged plan. + * @param output The output of the original aggregate. + */ +case class NonGroupingAggregateReference( + level: Int, + mergedPlanIndex: Int, + outputIndices: Seq[Int], + override val output: Seq[Attribute]) extends LeafNode { + final override val nodePatterns: Seq[TreePattern] = Seq(NO_GROUPING_AGGREGATE_REFERENCE) +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala index ec4ac257b3049..f04ff990251d7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeCsvJsonExprs.scala @@ -44,6 +44,10 @@ import org.apache.spark.unsafe.types.UTF8String object OptimizeCsvJsonExprs extends Rule[LogicalPlan] { private def nameOfCorruptRecord = conf.columnNameOfCorruptRecord + /** Whether the field this extracts is the corrupt record column of its `JsonToStructs`. */ + private def selectsCorruptRecord(g: GetStructField): Boolean = + g.childSchema(g.ordinal).name == nameOfCorruptRecord + private type SimpleJsonPath = Seq[GetJsonObject.SimpleJsonPathSegment] private type SharedJsonCandidate = (GetJsonObject, SimpleJsonPath, String) private type SharedJsonCandidateUnit = Seq[SharedJsonCandidate] @@ -393,10 +397,14 @@ object OptimizeCsvJsonExprs extends Rule[LogicalPlan] { private val jsonOptimization: PartialFunction[Expression, Expression] = { case c: CreateNamedStruct - // If we create struct from various fields of the same `JsonToStructs`. + // If we create struct from various fields of the same `JsonToStructs`. Pruning the + // schema stops the parser from converting the dropped fields, so a malformed value in + // one of them no longer fails under a parse mode such as failfast. To be more + // conservative, it does not optimize when any option is set, like the cases below. if c.valExprs.forall { v => v.isInstanceOf[GetStructField] && v.asInstanceOf[GetStructField].child.isInstanceOf[JsonToStructs] && + v.asInstanceOf[GetStructField].child.asInstanceOf[JsonToStructs].options.isEmpty && v.children.head.semanticEquals(c.valExprs.head.children.head) } => val jsonToStructs = c.valExprs.map(_.children.head) @@ -410,9 +418,18 @@ object OptimizeCsvJsonExprs extends Rule[LogicalPlan] { // `JsonToStructs` does not support parsing json with duplicated field names. val duplicateFields = c.names.map(_.toString).distinct.length != c.names.length + // Dropping a field stops the parser from converting it, so the parser never records the + // row when a dropped field contains a malformed value, and the corrupt record column comes + // back null. Selecting as many fields as the schema has drops nothing: with `sameFieldName` + // and no duplicates the selected names are distinct names of the schema, so the counts + // match only when every field is selected. + val fields = c.valExprs.map(_.asInstanceOf[GetStructField]) + val prunesCorruptRecord = fields.exists(selectsCorruptRecord) && + fields.length != fields.head.childSchema.length + // If we create struct from various fields of the same `JsonToStructs` and we don't // alias field names and there is no duplicated field in the struct. - if (sameFieldName && !duplicateFields) { + if (sameFieldName && !duplicateFields && !prunesCorruptRecord) { val fromJson = jsonToStructs.head.asInstanceOf[JsonToStructs].copy(schema = c.dataType) val nullFields = c.children.grouped(2).flatMap { case Seq(name, value) => Seq(name, Literal(null, value.dataType)) @@ -436,12 +453,15 @@ object OptimizeCsvJsonExprs extends Rule[LogicalPlan] { child case g @ GetStructField(j @ JsonToStructs(schema: StructType, _, _, _), ordinal, _) - if schema.length > 1 && j.options.isEmpty => + if schema.length > 1 && j.options.isEmpty && !selectsCorruptRecord(g) => // Options here should be empty because the optimization should not be enabled // for some options. For example, when the parse mode is failfast it should not // optimize, and should force to parse the whole input JSON with failing fast for // an invalid input. // To be more conservative, it does not optimize when any option is set for now. + // Pruning to the corrupt record column alone is excluded as well: it leaves the + // parser nothing to convert, so the parser never records the row when a dropped + // field contains a malformed value, and the column comes back null. val prunedSchema = StructType(Array(schema(ordinal))) g.copy(child = j.copy(schema = prunedSchema), ordinal = 0) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinCondition.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinCondition.scala index 7c41ebea050be..86248cf30828c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinCondition.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinCondition.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.catalyst.optimizer -import org.apache.spark.sql.catalyst.expressions.{And, EqualNullSafe, EqualTo, IsNull, Or, PredicateHelper} +import org.apache.spark.sql.catalyst.expressions.{And, EqualNullSafe, EqualTo, Expression, IsNull, Or, PredicateHelper} import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.{JOIN, OR} @@ -30,7 +30,16 @@ object OptimizeJoinCondition extends Rule[LogicalPlan] with PredicateHelper { override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( _.containsPattern(JOIN), ruleId) { case j @ Join(_, _, _, condition, _) if condition.nonEmpty => - val newCondition = condition.map(_.transformWithPruning(_.containsPattern(OR), ruleId) { + val newCondition = condition.map(optimizeCondition) + j.copy(condition = newCondition) + } + + // Rewriting the pattern to EqualNullSafe maps NULL to FALSE, so only recurse through And/Or. + private def optimizeCondition(condition: Expression): Expression = { + if (!condition.containsPattern(OR)) { + condition + } else { + condition match { case Or(EqualTo(l, r), And(IsNull(c1), IsNull(c2))) if (l.semanticEquals(c1) && r.semanticEquals(c2)) || (l.semanticEquals(c2) && r.semanticEquals(c1)) => @@ -39,7 +48,18 @@ object OptimizeJoinCondition extends Rule[LogicalPlan] with PredicateHelper { if (l.semanticEquals(c1) && r.semanticEquals(c2)) || (l.semanticEquals(c2) && r.semanticEquals(c1)) => EqualNullSafe(l, r) - }) - j.copy(condition = newCondition) + case and @ And(left, right) => + val newLeft = optimizeCondition(left) + val newRight = optimizeCondition(right) + if (newLeft.fastEquals(left) && newRight.fastEquals(right)) and + else And(newLeft, newRight) + case or @ Or(left, right) => + val newLeft = optimizeCondition(left) + val newRight = optimizeCondition(right) + if (newLeft.fastEquals(left) && newRight.fastEquals(right)) or + else Or(newLeft, newRight) + case other => other + } + } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala index 53ddc53091fbb..58b6479ddee30 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala @@ -125,6 +125,7 @@ abstract class Optimizer(catalogManager: CatalogManager) OptimizeRepartition, EliminateWindowPartitions, TransposeWindow, + PullUpProjectAliasThroughWindow, NullPropagation, // NullPropagation may introduce Exists subqueries, so RewriteNonCorrelatedExists must run // after. @@ -174,6 +175,13 @@ abstract class Optimizer(catalogManager: CatalogManager) PushDownPredicates)) val batches: Seq[Batch] = flattenBatches(Seq( + // UDF substitution rules should be executed before any other optimization rules + // so that the substituted Catalyst alternatives go through the same finalization + // (FinishAnalysis, RewriteWithExpression, etc.) and downstream optimization as + // any other expression. Anything ConvertToCatalyst leaves behind -- including + // RuntimeReplaceable nodes inside transpiled options -- is still rewritten by + // the FinishAnalysis batch that runs immediately after. + Batch("Convert python UDFs to Catalyst", Once, ConvertToCatalyst), Batch("Finish Analysis", FixedPoint(1), FinishAnalysis), // We must run this batch after `ReplaceExpressions`, as `RuntimeReplaceable` expression // may produce `With` expressions that need to be rewritten. @@ -255,6 +263,10 @@ abstract class Optimizer(catalogManager: CatalogManager) Batch("Eliminate Sorts", Once, EliminateSorts, RemoveRedundantSorts), + // Run after operator optimization normally folds accuracy expressions and before + // RewriteDistinctAggregates so fused distinct percentiles are rewritten correctly. + Batch("Combine Approximate Percentiles", Once, + CombineApproximatePercentiles), Batch("Decimal Optimizations", fixedPoint, DecimalAggregates), // This batch must run after "Decimal Optimizations", as that one may change the @@ -304,6 +316,10 @@ abstract class Optimizer(catalogManager: CatalogManager) */ def nonExcludableRules: Seq[String] = Seq( + // ConvertToCatalyst is the only rule that strips the Unevaluable + // TranspiledPythonUDF node; excluding it would leak that node into + // execution, so it must never be excludable. + ConvertToCatalyst.ruleName, FinishAnalysis.ruleName, RewriteDistinctAggregates.ruleName, ReplaceDeduplicateWithAggregate.ruleName, @@ -339,10 +355,10 @@ abstract class Optimizer(catalogManager: CatalogManager) NormalizeFloatingNumbers, RewriteNonCorrelatedExists, PullOutGroupingExpressions, - // Put `InsertMapSortInGroupingExpressions` after `PullOutGroupingExpressions`, - // so the grouping keys can only be attribute and literal which makes - // `InsertMapSortInGroupingExpressions` easy to insert `MapSort`. - InsertMapSortInGroupingExpressions, + // Put `InsertMapSortInAggregate` after `PullOutGroupingExpressions`, + // so grouping keys are attributes or literals. The rule also projects complex distinct + // aggregate arguments before inserting `MapSort`. + InsertMapSortInAggregate, InsertMapSortInRepartitionExpressions, ComputeCurrentTime, ReplaceCurrentLike(catalogManager), @@ -1006,6 +1022,80 @@ object LimitPushDown extends Rule[LogicalPlan] { } } +/** + * Attempt to convert UDFS to Catalyst expressions. + */ +object ConvertToCatalyst extends Rule[LogicalPlan] { + def apply(plan: LogicalPlan): LogicalPlan = { + // Short circuit if there are no Transpiled Python UDFs in the plan. + if (!plan.containsPattern(TRANSPILED_PYTHON_UDF)) { + return plan + } + // Traverse subquery plans too: this batch runs Once, and later rules (e.g. + // PullupCorrelatedPredicates) can move expressions from a subquery into the + // outer plan, so an Unevaluable TranspiledPythonUDF left inside a subquery + // here could otherwise escape and reach execution un-stripped. + plan.transformDownWithSubqueriesAndPruning( + _.containsPattern(TRANSPILED_PYTHON_UDF), ruleId) { + case p => p.transformExpressionsWithPruning(_.containsPattern(TRANSPILED_PYTHON_UDF)) { + case s: TranspiledPythonUDF => applyExpr(s, parentIsUdf = false) + } + } + } + + def applyExpr(expression: Expression, parentIsUdf: Boolean = false): Expression = { + expression match { + case s: TranspiledPythonUDF => + // We _shouldn't_ have these nodes if ANSI is not enabled or transpilation is disabled + // but if someone changed it while running we'll want to strip the nodes out. + if (!conf.getConf(SQLConf.ANSI_ENABLED)) { + logWarning(log"Skipping Python UDF transpilation: " + + log"${MDC(LogKeys.CONFIG, SQLConf.ANSI_ENABLED.key)} is disabled. The transpiler " + + log"targets ANSI semantics and refuses to rewrite plans under non-ANSI mode. " + + log"Enable ANSI or disable transpilation to silence this warning.") + s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true)) + } else if (!conf.getConf(SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS)) { + logWarning(log"Skipping Python UDF transpilation: " + + log"${MDC(LogKeys.CONFIG, SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key)} " + + log"is disabled but we still got TranspiledPythonUDFs in our plan.") + s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true)) + } else if (!parentIsUdf || !s.hasOnlyPythonUDFInputs) { + // Walk the full list of transpiled options and pick the first one, + // falling back to the original Python UDF if none are available. + // Options whose declared input-type categories don't match the bound + // column types are already pruned during analysis by + // ResolveTranspiledPythonUDFOptions, so any option that reaches here is + // safe to use. If you're plugging in your own transpilation, please add + // a separate ConvertToX so you can choose your desired transpiled nodes. + // NOTE: the substituted option is used as-is, with no cast back to the + // UDF's declared return type. The built-in transpiler guarantees each + // option's dataType already matches; a custom transpiler MUST do the + // same (or insert its own Cast), or it will silently change the output + // schema. + val firstEvaluable = s.transpiledOptions.headOption + firstEvaluable match { + case None => + s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true)) + case Some(catalystExpr) => + // Recursively apply to the children first because we may use them as inputs in parent + catalystExpr.mapChildren(applyExpr(_, parentIsUdf = false)) + } + } else { + // We should avoid converting a UDF node where that could break pipelining. + // For example: (UDF -> UDF -> UDF) is often cheaper than UDF -> Catalyst -> UDF. + s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true)) + } + case _ => + // Not a TranspiledPythonUDF: recurse down, telling the children whether + // this node is itself a scalar Python UDF so a transpiled child can + // preserve the UDF batch pipeline (e.g. an outer UDF that could not be + // transpiled wrapping one that could). + expression.mapChildren( + applyExpr(_, parentIsUdf = isScalarPythonUDF(expression))) + } + } +} + /** * Pushes Project operator to both sides of a Union operator. * Operations that are safe to pushdown are listed as follows. @@ -1664,6 +1754,10 @@ object OptimizeWindowFunctions extends Rule[LogicalPlan] { * Collapse Adjacent Window Expression. * - If the partition specs and order specs are the same and the window expression are * independent and are of the same window function type, collapse into the parent. + * - If the partition specs are the same and one of the order specs is empty, collapse into the + * parent when the window expressions of the empty-order window can be evaluated under any row + * order. The merged window keeps the non-empty order spec. Merging an empty-order child is + * gated by `spark.sql.optimizer.collapseWindowWithEmptyOrderSpecInChild`. */ object CollapseWindow extends Rule[LogicalPlan] { private def specCompatible(s1: Seq[Expression], s2: Seq[Expression]): Boolean = { @@ -1671,9 +1765,46 @@ object CollapseWindow extends Rule[LogicalPlan] { s1.zip(s2).forall(e => e._1.semanticEquals(e._2)) } + /** + * Returns true if the given window expression can still be evaluated correctly when the rows + * of the partition are reordered, so that it can be merged into another window with a different + * (non-empty) order spec. + * + * The frame determines whether reordering is safe. When the frame is the whole partition + * (`UNBOUNDED PRECEDING` to `UNBOUNDED FOLLOWING`), it always covers all the rows of the + * partition regardless of the ordering, so reordering changes only the order in which the rows + * are seen, never which rows are in the frame. Since the order spec of the window is empty, + * the query does not fix the row order, so evaluating its expressions under any ordering + * yields a valid result, even though the value may differ for order-dependent expressions + * such as `first`, `collect_list`, or floating-point `sum`/`avg`. On the other hand, a bounded + * frame (e.g. `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`) is order-sensitive: which + * rows are in the frame depends on the ordering, so even `count` or `sum` would change value, + * and such a window must not be merged. + */ + private def canEvaluateUnderAnyOrder(windowExpression: NamedExpression): Boolean = + windowExpression match { + case Alias(WindowExpression(_, WindowSpecDefinition(_, _, + SpecifiedWindowFrame(_, UnboundedPreceding, UnboundedFollowing))), _) => true + case _ => false + } + private def windowsCompatible(w1: Window, w2: Window): Boolean = { specCompatible(w1.partitionSpec, w2.partitionSpec) && - specCompatible(w1.orderSpec, w2.orderSpec) && + // The order specs can differ when one of them is empty, as long as the window expressions + // of the window with the empty order spec are safe to evaluate under any row order. In that + // case, they can be evaluated under the non-empty order spec of the other window. The + // operator then keeps the non-empty order spec while the merged-in expressions keep their + // own empty order spec; this divergence is safe because after analysis only + // OptimizeWindowFunctions reads an expression's own order spec, and it is a no-op without + // one. Merging an empty-order child into an ordered parent can disable + // InferWindowGroupLimit and LimitPushDownThroughWindow, so that direction is gated by + // conf.collapseWindowWithEmptyOrderSpecInChild. + (specCompatible(w1.orderSpec, w2.orderSpec) || + (w1.orderSpec.isEmpty && w2.orderSpec.nonEmpty && + w1.windowExpressions.forall(canEvaluateUnderAnyOrder)) || + (conf.collapseWindowWithEmptyOrderSpecInChild && w2.orderSpec.isEmpty && + w1.orderSpec.nonEmpty && + w2.windowExpressions.forall(canEvaluateUnderAnyOrder))) && w1.references.intersect(w2.windowOutputSet).isEmpty && w1.windowExpressions.nonEmpty && w2.windowExpressions.nonEmpty && // This assumes Window contains the same type of window expressions. This is ensured @@ -1686,13 +1817,19 @@ object CollapseWindow extends Rule[LogicalPlan] { _.containsPattern(WINDOW), ruleId) { case w1 @ Window(we1, _, _, w2 @ Window(we2, _, _, grandChild, _), _) if windowsCompatible(w1, w2) => - w1.copy(windowExpressions = we2 ++ we1, child = grandChild) + w1.copy( + orderSpec = if (w1.orderSpec.nonEmpty) w1.orderSpec else w2.orderSpec, + windowExpressions = we2 ++ we1, + child = grandChild) case w1 @ Window(we1, _, _, Project(pl, w2 @ Window(we2, _, _, grandChild, _)), _) if windowsCompatible(w1, w2) && w1.references.subsetOf(grandChild.outputSet) => Project( pl ++ w1.windowOutputSet, - w1.copy(windowExpressions = we2 ++ we1, child = grandChild)) + w1.copy( + orderSpec = if (w1.orderSpec.nonEmpty) w1.orderSpec else w2.orderSpec, + windowExpressions = we2 ++ we1, + child = grandChild)) } } @@ -2104,6 +2241,10 @@ object PushDownPredicates extends Rule[LogicalPlan] { * 2) the predicate is deterministic and the operator will not change any of rows. * 3) We don't add double evaluation OR double evaluation would be cheap OR we're configured to. * + * Note: if a new push-through case is added here, or the translation applied to pushed + * conditions changes (e.g. how aliases are substituted), also update + * `removePushedDownFilter` in [[PushdownPredicatesAndPruneColumnsForCTEDef]], which mirrors + * this rule's cases to locate and remove filters previously pushed into CTE definitions. */ object PushPredicateThroughNonJoin extends Rule[LogicalPlan] with PredicateHelper { def apply(plan: LogicalPlan): LogicalPlan = plan transform applyLocally @@ -2695,9 +2836,11 @@ object ConvertToLocalRelation extends Rule[LogicalPlan] { _.containsPattern(LOCAL_RELATION), ruleId) { case Project(projectList, LocalRelation(output, data, isStreaming, stream)) if !projectList.exists(hasUnevaluableExpr) => - val projection = new InterpretedMutableProjection(projectList, output) + val freshProjectList = projectList.map( + _.freshCopyIfContainsStatefulExpression().asInstanceOf[NamedExpression]) + val projection = new InterpretedMutableProjection(freshProjectList, output) projection.initialize(0) - LocalRelation(projectList.map(_.toAttribute), data.map(projection(_).copy()), + LocalRelation(freshProjectList.map(_.toAttribute), data.map(projection(_).copy()), isStreaming, stream) case Limit(IntegerLiteral(limit), LocalRelation(output, data, isStreaming, stream)) => @@ -2708,7 +2851,8 @@ object ConvertToLocalRelation extends Rule[LogicalPlan] { case Filter(condition, LocalRelation(output, data, isStreaming, stream)) if !hasUnevaluableExpr(condition) => - val predicate = Predicate.create(condition, output) + val freshCondition = condition.freshCopyIfContainsStatefulExpression() + val predicate = Predicate.create(freshCondition, output) predicate.initialize(0) LocalRelation(output, data.filter(row => predicate.eval(row)), isStreaming, stream) } @@ -2738,7 +2882,7 @@ object ReplaceDeduplicateWithAggregate extends Rule[LogicalPlan] { def apply(plan: LogicalPlan): LogicalPlan = plan transformUpWithNewOutput { case d @ Deduplicate(keys, child, _) if !child.isStreaming => val keyExprIds = keys.map(_.exprId) - val generatedAliasesMap = new mutable.HashMap[Attribute, Alias](); + val generatedAliasesMap = new mutable.HashMap[Attribute, Alias]() val aggCols = child.output.map { attr => if (keyExprIds.contains(attr.exprId)) { attr diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala deleted file mode 100644 index 1c43f91cee9dc..0000000000000 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PlanMerger.scala +++ /dev/null @@ -1,647 +0,0 @@ -/* - * 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.catalyst.optimizer - -import scala.collection.mutable - -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, Expression, If, Literal, NamedExpression, Or} -import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression -import org.apache.spark.sql.catalyst.plans.{Cross, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project} -import org.apache.spark.sql.catalyst.trees.TreeNodeTag -import org.apache.spark.sql.internal.SQLConf - -/** - * Result of attempting to merge a plan via [[PlanMerger.merge]]. - * - * @param mergedPlan The resulting plan, either: - * - An existing cached plan (if identical match found) - * - A newly merged plan combining the input with a cached plan - * - The original input plan (if no merge was possible) - * @param mergedPlanIndex The index of this plan in the PlanMerger's cache. - * @param outputMap Maps attributes of the input plan to their positional index in - * `mergedPlan.plan.output`. The index remains stable across subsequent - * [[PlanMerger.merge]] calls because outputs are only ever appended. - */ -case class MergeResult( - mergedPlan: MergedPlan, - mergedPlanIndex: Int, - outputMap: AttributeMap[Int]) - -/** - * Represents a plan in the PlanMerger's cache. - * - * @param plan The logical plan, which may have been merged from multiple original plans. - * @param merged Whether this plan is the result of merging two or more plans (true), or - * is an original unmerged plan (false). Merged plans typically require special - * handling such as wrapping in CTEs. - */ -case class MergedPlan(plan: LogicalPlan, merged: Boolean) - -object PlanMerger { - // Marker tag placed on Filter nodes that were produced by filter propagation. Its presence - // signals that the Filter's condition is already an OR of propagated filter attributes and - // its child Project already contains the corresponding aliases, so a subsequent merge only - // needs to add one new alias for the incoming plan rather than wrapping both sides again. - val MERGED_FILTER_TAG: TreeNodeTag[Unit] = TreeNodeTag("mergedFilter") - - // Global counter for generating unique names for propagated filter attributes across all - // PlanMerger instances. - private[optimizer] val curId = new java.util.concurrent.atomic.AtomicLong() - private[optimizer] def newId: Long = curId.getAndIncrement() -} - -/** - * A stateful utility for merging identical or similar logical plans to enable query plan reuse. - * - * `PlanMerger` maintains a cache of previously seen plans and attempts to either: - * 1. Reuse an identical plan already in the cache - * 2. Merge a new plan with a cached plan by combining their outputs - * - * The merging process preserves semantic equivalence while combining outputs from multiple - * plans into a single plan. This is primarily used by [[MergeSubplans]] to deduplicate subplan - * execution. - * - * Supported plan types for merging: - * - [[Project]]: Merges project lists - * - [[Aggregate]]: Merges aggregate expressions with identical grouping - * - [[Filter]]: Requires identical filter conditions - * - [[Join]]: Requires identical join type, hints, and conditions - * - * When `filterPropagationEnabled` is true, non-grouping [[Aggregate]]s over the same base plan - * with different [[Filter]] conditions can also be merged. The filter conditions are exposed as - * boolean [[Project]] attributes and consumed at the [[Aggregate]] as FILTER clauses. - * When both sides carry a [[Filter]] (the symmetric case), merging broadens the scan to OR(f1, f2), - * which may reduce IO pruning. This path is separately gated by - * `symmetricFilterPropagationEnabled`. - * When plans also differ in intermediate [[Project]] expressions, those are wrapped with - * `If(filterAttr, expr, null)` to avoid computing the expression for rows that do not match that - * side's filter condition. - * Filter propagation also works through [[Join]] nodes: a filter on one child of the join produces - * a boolean attribute that flows through the join output to the enclosing [[Aggregate]]. - * Propagation is only safe when the filter originates from the non-nullable side of the join, as - * enforced by `filterSafeForJoin`. When the filter is on the nullable side, the merged base plan - * restores rows that were filtered out of the nullable child, turning what were unmatched - * NULL-padded rows in the original plan into matched rows with real column values. This changes the - * result of expressions like `coalesce(col, default)` in the aggregate: an originally unmatched row - * would have contributed `default` via `coalesce(NULL, default)`, but in the merged plan it is - * matched, its real column value fails the filter, and `FILTER (WHERE false)` discards it entirely. - * Propagation is also skipped when both the left and right children simultaneously produce filter - * attributes, as combining them would require an additional AND alias above the join (not yet - * supported). - * - * {{{ - * // Input plans - * Aggregate [sum(a) AS sum_a] Aggregate [max(d) AS max_d] - * +- Filter (a < 1) +- Project [udf(a) AS d] - * +- Scan t +- Filter (a > 1) - * +- Scan t - * - * // Merged plan - * Aggregate [sum(a) FILTER f0 AS sum_a, max(d0) FILTER f1 AS max_d] - * +- Project [a, If(f1, udf(a), null) AS d0, f0, f1] - * +- Filter (f0 OR f1) [MERGED_FILTER_TAG] - * +- Project [a, (a < 1) AS f0, (a > 1) AS f1] - * +- Scan t - * }}} - * - * @example - * {{{ - * val merger = PlanMerger() - * val result1 = merger.merge(plan1) // Adds plan1 to cache - * val result2 = merger.merge(plan2) // Merges with plan1 if compatible - * // result2.mergedPlan.merged == true if plans were merged - * // result2.outputMap maps plan2's attributes to the merged plan's attributes - * }}} - */ -class PlanMerger( - filterPropagationEnabled: Boolean = - SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED), - symmetricFilterPropagationEnabled: Boolean = - SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED), - filterPropagationThroughJoinEnabled: Boolean = - SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED)) { - val cache = mutable.ArrayBuffer.empty[MergedPlan] - - /** - * Attempts to merge the given plan with cached plans, or adds it to the cache. - * - * The method tries the following in order: - * 1. Check if an identical plan exists in cache (using canonicalized comparison) - * 2. Try to merge with each cached plan using [[tryMergePlans]] - * 3. If no merge is possible, add as a new cache entry - * - * @param plan The logical plan to merge or cache. - * @param subqueryPlan If the logical plan is a subquery plan. - * @return A [[MergeResult]] containing: - * - The merged/cached plan to use - * - Its index in the cache - * - An attribute mapping for rewriting expressions - */ - def merge(plan: LogicalPlan, subqueryPlan: Boolean): MergeResult = { - cache.zipWithIndex.collectFirst(Function.unlift { - case (mp, i) => - checkIdenticalPlans(plan, mp.plan).map { _ => - // Identical subquery expression plans are not marked as `merged` as the - // `ReusedSubqueryExec` rule can handle them without extracting the plans to CTEs. - // But, when a non-subquery subplan is identical to a cached plan we need to mark the plan - // `merged` and so extract it to a CTE later. - val newMergedPlan = MergedPlan(mp.plan, mp.merged || !subqueryPlan) - cache(i) = newMergedPlan - val outputMap = AttributeMap(plan.output.zipWithIndex) - MergeResult(newMergedPlan, i, outputMap) - }.orElse { - tryMergePlans(plan, mp.plan, false).collect { - case TryMergeResult(mergedPlan, npMapping, None, None) => - val newMergedPlan = MergedPlan(mergedPlan, true) - cache(i) = newMergedPlan - val outputMap = AttributeMap(npMapping.iterator.map { case (origAttr, mergedAttr) => - origAttr -> mergedPlan.output.indexWhere(_.exprId == mergedAttr.exprId) - }.toSeq) - MergeResult(newMergedPlan, i, outputMap) - } - } - case _ => None - }).getOrElse { - val newMergedPlan = MergedPlan(plan, false) - cache += newMergedPlan - val outputMap = AttributeMap(plan.output.zipWithIndex) - MergeResult(newMergedPlan, cache.length - 1, outputMap) - } - } - - /** - * Returns all plans currently in the cache as an immutable indexed sequence. - * - * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan - * corresponds to the `mergedPlanIndex` returned by [[merge]]. - */ - def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq - - // If 2 plans are identical return the attribute mapping from the new to the cached version. - private def checkIdenticalPlans( - newPlan: LogicalPlan, - cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = { - if (newPlan.canonicalized == cachedPlan.canonicalized) { - Some(AttributeMap(newPlan.output.zip(cachedPlan.output))) - } else { - None - } - } - - /** - * Result of a successful [[tryMergePlans]] call. - * - * @param mergedPlan The combined logical plan. - * @param newPlanMapping Mapping from attributes in the new plan to the corresponding - * attributes in the merged plan. Used by parent nodes to remap - * new-plan-side expressions. - * @param newPlanFilter A boolean [[Attribute]] in the merged plan that encodes the filter - * condition from the new plan's side, to be applied as an aggregate - * `FILTER (WHERE ...)` clause when the propagation reaches an enclosing - * [[Aggregate]] node. The boolean component is `true` if the attribute was - * freshly aliased and must be appended to enclosing [[Project]] nodes, or - * `false` if it was reused from an existing alias already present in the - * merged plan. `None` when no differing filter was propagated. - * @param cachedPlanFilter Like `newPlanFilter` but for the cached plan's side. Always a freshly - * created alias when present, so no `isNew` flag is needed. - */ - case class TryMergeResult( - mergedPlan: LogicalPlan, - newPlanMapping: AttributeMap[Attribute], - newPlanFilter: Option[(Attribute, Boolean)] = None, - cachedPlanFilter: Option[Attribute] = None) - - /** - * Recursively attempts to merge two plans by traversing their tree structures. - * - * Two plans can be merged if: - * - They are identical (canonicalized forms match), OR - * - They have compatible root nodes with mergeable children - * - * Supported merge patterns: - * - Project nodes: Combines project lists from both plans - * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both - * support the same aggregate implementation (hash/object-hash/sort-based) - * - Filter nodes: Only if filter conditions are identical - * - Join nodes: Requires identical join type, hints, and conditions; filter propagation is - * forwarded into the join's children so a filter difference on one child can still be merged - * - * @param newPlan The plan to merge into the cached plan. - * @param cachedPlan The cached plan to merge with. - * @return Some([[TryMergeResult]]) if merge succeeds, None if plans cannot be merged. - */ - private def tryMergePlans( - newPlan: LogicalPlan, - cachedPlan: LogicalPlan, - filterPropagationSupported: Boolean): Option[TryMergeResult] = { - checkIdenticalPlans(newPlan, cachedPlan).map(TryMergeResult(cachedPlan, _)).orElse( - (newPlan, cachedPlan) match { - case (np: Project, cp: Project) => - tryMergePlans(np.child, cp.child, filterPropagationSupported).map { - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) => - val (mergedProjectList, newNPMapping) = - mergeNamedExpressions(np.projectList, cp.projectList, npMapping, npFilter, cpFilter) - TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, - cpFilter) - } - case (np, cp: Project) => - tryMergePlans(np, cp.child, filterPropagationSupported).map { - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) => - val (mergedProjectList, newNPMapping) = - mergeNamedExpressions(np.output, cp.projectList, npMapping, npFilter, cpFilter) - TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, - cpFilter) - } - case (np: Project, cp) => - tryMergePlans(np.child, cp, filterPropagationSupported).map { - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) => - val (mergedProjectList, newNPMapping) = - mergeNamedExpressions(np.projectList, cp.output, npMapping, npFilter, cpFilter) - TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, - cpFilter) - } - - case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) => - // Filter propagation into the aggregate is only safe when there is no grouping. - val childFilterPropagationSupported = filterPropagationEnabled && - np.groupingExpressions.isEmpty && cp.groupingExpressions.isEmpty - tryMergePlans(np.child, cp.child, childFilterPropagationSupported).flatMap { - case TryMergeResult(mergedChild, npMapping, None, None) => - val mappedNPGroupingExpression = - np.groupingExpressions.map(mapAttributes(_, npMapping)) - // Order of grouping expression does matter as merging different grouping orders can - // introduce "extra" shuffles/sorts that might not present in all of the original - // subqueries. - if (mappedNPGroupingExpression.map(_.canonicalized) == - cp.groupingExpressions.map(_.canonicalized)) { - val (mergedAggregateExpressions, newNPMapping) = - mergeNamedExpressions(np.aggregateExpressions, cp.aggregateExpressions, npMapping) - val mergedPlan = - Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild) - Some(TryMergeResult(mergedPlan, newNPMapping)) - } else { - None - } - case TryMergeResult(mergedChild, npMapping, npFilterOpt, cpFilterOpt) => - // childFilterPropagationSupported guarantees both aggregates have no grouping, so - // the grouping-match check is skipped. - assert(childFilterPropagationSupported) - - // Apply each propagated boolean attribute as a FILTER (WHERE ...) clause on the - // corresponding side's aggregate expressions. - // A None filter means the side's aggregate expressions already carry their individual - // FILTER attributes from a previous merge round and should be left unchanged. - // Filter propagation is consumed here and not passed further up. - val filteredNPAggregateExpressions = npFilterOpt.fold(np.aggregateExpressions) { - case (f, _) => applyFilterToAggregateExpressions(np.aggregateExpressions, f) - } - val filteredCPAggregateExpressions = cpFilterOpt.fold(cp.aggregateExpressions)( - applyFilterToAggregateExpressions(cp.aggregateExpressions, _)) - val (mergedAggregateExpressions, newNPMapping) = - mergeNamedExpressions(filteredNPAggregateExpressions, - filteredCPAggregateExpressions, npMapping) - val mergedPlan = Aggregate(Seq.empty, mergedAggregateExpressions, mergedChild) - Some(TryMergeResult(mergedPlan, newNPMapping)) - } - - case (np: Filter, cp: Filter) => - tryMergePlans(np.child, cp.child, filterPropagationSupported).flatMap { - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) => - val mappedNPCondition = mapAttributes(np.condition, npMapping) - // Comparing the canonicalized form is required to ignore different forms of the same - // expression. - if (mappedNPCondition.canonicalized == cp.condition.canonicalized) { - // Identical conditions: the filter node itself adds no new discrimination between - // the two sides, so we keep it unchanged and pass the child's mappings up. - val mergedPlan = Filter(cp.condition, mergedChild) - Some(TryMergeResult(mergedPlan, npMapping, npFilter, cpFilter)) - } else if (filterPropagationSupported && symmetricFilterPropagationEnabled) { - if (cp.getTagValue(PlanMerger.MERGED_FILTER_TAG).isDefined) { - // cp Filter is already a merged filter from a previous round: its condition - // is OR(f0, f1, ...) and its child Project already contains aliases for those - // attributes. Only create a new alias for the np side, and extend the OR - // condition. - val newNPCondition = npFilter.fold(mappedNPCondition) { - case (f, _) => And(f, mappedNPCondition) - } - val childProject = mergedChild match { - case p: Project => p - case other => throw new IllegalStateException( - "Expected Project child under MERGED_FILTER_TAG filter, got " + - s"${other.getClass.getSimpleName}") - } - // If newNPCondition is already aliased in the child Project (e.g. a third - // subplan whose filter matches one from a previous merge round), reuse the - // existing attribute instead of creating a redundant alias. - val existingNPFilter = childProject.projectList.collectFirst { - case a: Alias if a.child.canonicalized == newNPCondition.canonicalized => - a.toAttribute - } - existingNPFilter match { - case Some(reusedFilter) => - val newFilter = cp.withNewChildren(Seq(mergedChild)) - Some(TryMergeResult(newFilter, npMapping, Some((reusedFilter, false)), None)) - case None => - val newNPFilterAlias = - Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newNPFilter = newNPFilterAlias.toAttribute - val newProject = childProject.copy( - projectList = childProject.projectList ++ Seq(newNPFilterAlias)) - val newFilter = Filter(Or(cp.condition, newNPFilter), newProject) - newFilter.copyTagsFrom(cp) - Some(TryMergeResult(newFilter, npMapping, Some((newNPFilter, true)), None)) - } - } else { - // First-time filter propagation: alias both sides' conditions as boolean - // attributes in a new Project below the Filter, and set the Filter condition - // to OR(newNPFilter, newCPFilter). - // Note: the new Project always uses mergedChild as its child (rather than - // flattening into an existing Project below) because mergedChild.output may - // contain previously-propagated filter attributes that cp.condition references. - val newNPCondition = - npFilter.fold(mappedNPCondition) { case (f, _) => And(f, mappedNPCondition) } - val newCPCondition = cpFilter.fold(cp.condition)(And(_, cp.condition)) - val newNPFilterAlias = - Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newCPFilterAlias = - Alias(newCPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newNPFilter = newNPFilterAlias.toAttribute - val newCPFilter = newCPFilterAlias.toAttribute - val project = Project( - mergedChild.output.toList ++ Seq(newNPFilterAlias, newCPFilterAlias), - mergedChild) - val newFilter = Filter(Or(newNPFilter, newCPFilter), project) - newFilter.copyTagsFrom(cp) - newFilter.setTagValue(PlanMerger.MERGED_FILTER_TAG, ()) - Some(TryMergeResult(newFilter, npMapping, Some((newNPFilter, true)), - Some(newCPFilter))) - } - } else { - None - } - } - case (np: Filter, cp) if filterPropagationSupported => - tryMergePlans(np.child, cp, filterPropagationSupported).collect { - // If the cp side already propagated a filter from deeper recursion, the merge is - // effectively symmetric (both sides have a filter condition). Abort unless - // symmetricFilterPropagationEnabled. - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) - if cpFilter.isEmpty || symmetricFilterPropagationEnabled => - val mappedNPCondition = mapAttributes(np.condition, npMapping) - val newNPCondition = npFilter.fold(mappedNPCondition) { - case (f, _) => And(f, mappedNPCondition) - } - val newNPFilterAlias = - Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newNPFilter = newNPFilterAlias.toAttribute - val project = Project( - mergedChild.output.toList :+ newNPFilterAlias, - mergedChild) - TryMergeResult(project, npMapping, Some((newNPFilter, true)), cpFilter) - } - case (np, cp: Filter) if filterPropagationSupported => - tryMergePlans(np, cp.child, filterPropagationSupported).collect { - // If the np side already propagated a filter from deeper recursion, the merge is - // effectively symmetric (both sides have a filter condition). Abort unless - // symmetricFilterPropagationEnabled. - case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter) - if npFilter.isEmpty || symmetricFilterPropagationEnabled => - if (cp.getTagValue(PlanMerger.MERGED_FILTER_TAG).isDefined) { - // cp is a previously-merged Filter: its condition is `OR(pf_0, pf_1, ...)` and cp's - // aggregate expressions already carry individual `FILTER (WHERE pf_i)` clauses that - // restrict each aggregation to its originating side. Synthesising a new cpFilter - // alias for cp.condition would just produce `FILTER AND(OR(pf_0, pf_1, ...), pf_i)` - // upstream, which simplifies to `FILTER pf_i` -- wasted work and plan bloat. - // Drop cp's Filter and let the recursion's result flow up with cpFilter = None so - // cp's aggregates are left untouched. - TryMergeResult(mergedChild, npMapping, npFilter, None) - } else { - val newCPCondition = cpFilter.fold(cp.condition)(And(_, cp.condition)) - val newCPFilterAlias = - Alias(newCPCondition, s"propagatedFilter_${PlanMerger.newId}")() - val newCPFilter = newCPFilterAlias.toAttribute - val project = Project( - mergedChild.output.toList :+ newCPFilterAlias, - mergedChild) - TryMergeResult(project, npMapping, npFilter, Some(newCPFilter)) - } - } - - case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint => - tryMergePlans(np.left, cp.left, filterPropagationSupported).flatMap { - case TryMergeResult(mergedLeft, leftNPMapping, leftNPFilter, leftCPFilter) => - tryMergePlans(np.right, cp.right, filterPropagationSupported).flatMap { - case TryMergeResult(mergedRight, rightNPMapping, rightNPFilter, rightCPFilter) - // If both children independently propagate filter attributes we would need to - // AND them into a new alias above the join, which is not yet supported. - if !(leftNPFilter.isDefined && rightNPFilter.isDefined) && - !(leftCPFilter.isDefined && rightCPFilter.isDefined) && - // Gate join-crossing filter propagation behind its own config flag. - // When no filter attributes are in play the merge is unconditionally safe. - (leftNPFilter.isEmpty && leftCPFilter.isEmpty && - rightNPFilter.isEmpty && rightCPFilter.isEmpty || - filterPropagationThroughJoinEnabled) && - // A filter attribute is only safe to propagate through a join if it comes - // from the "preserved" (non-nullable) side. On the nullable side, unmatched - // rows are NULL-padded so f=NULL, causing FILTER (WHERE f) to incorrectly - // exclude rows that should contribute to the aggregate. Right-side - // attributes are also absent from semi/anti join output. - (leftNPFilter.isEmpty && leftCPFilter.isEmpty || - filterSafeForJoin(fromLeft = true, cp.joinType)) && - (rightNPFilter.isEmpty && rightCPFilter.isEmpty || - filterSafeForJoin(fromLeft = false, cp.joinType)) => - val npMapping = leftNPMapping ++ rightNPMapping - val mappedNPCondition = np.condition.map(mapAttributes(_, npMapping)) - // Comparing the canonicalized form is required to ignore different forms of the - // same expression and `AttributeReference.qualifier`s in `cp.condition`. - if (mappedNPCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) { - val npFilter = leftNPFilter.orElse(rightNPFilter) - val cpFilter = leftCPFilter.orElse(rightCPFilter) - Some(TryMergeResult(cp.withNewChildren(Seq(mergedLeft, mergedRight)), npMapping, - npFilter, cpFilter)) - } else { - None - } - case _ => None - } - case _ => None - } - - // Otherwise merging is not possible. - case _ => None - }) - } - - // Returns true when a filter attribute originating from `fromLeft` child of a join with - // `joinType` can be safely propagated through that join to a parent Aggregate. - // - // Two conditions must both hold: - // 1. The attribute is in the join's output (rules out the right side of LeftSemi/LeftAnti). - // 2. The filter must originate from the non-nullable ("preserved") side of the join. - // When a filter is on the nullable side, the merged base plan no longer applies it to the - // nullable child's scan, so rows that were previously absent from that child reappear as - // matched join rows instead of unmatched NULL-padded rows. This changes aggregate - // expressions that use the NULL-padded column: e.g. for `sum(coalesce(col, default))`, an - // originally unmatched row would have contributed `default` via `coalesce(NULL, default)`, - // but in the merged plan the row is now matched with its real column value, fails the - // filter, and FILTER (WHERE false) discards it -- losing the `default` contribution - // entirely. - private def filterSafeForJoin(fromLeft: Boolean, joinType: JoinType): Boolean = - if (fromLeft) { - // Left side is never NULL-padded in: Inner, LeftOuter, LeftSemi, LeftAnti, Cross. - joinType match { - case Inner | LeftOuter | LeftSemi | LeftAnti | Cross => true - case _ => false // RightOuter and FullOuter can NULL-pad the left side - } - } else { - // Right side is never NULL-padded AND is in the join output in: Inner, RightOuter, Cross. - joinType match { - case Inner | RightOuter | Cross => true - case _ => false // LeftOuter/FullOuter can NULL-pad right; LeftSemi/LeftAnti drop right - } - } - - private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[Attribute]) = { - expr.transform { - case a: Attribute => outputMap.getOrElse(a, a) - }.asInstanceOf[T] - } - - // Remaps attributes of `newPlanExpressions` through `newPlanMapping`, then merges them with - // `cachedPlanExpressions` into a single expression list. - // Returns a pair of: - // 1. The merged expression list - // 2. New plan output map: ne.toAttribute -> merged plan attr (for parent nodes to remap - // new-plan-side expressions) - // - // When `newPlanFilter`/`cachedPlanFilter` are provided (filter propagation active), non-matching - // expressions from each side are wrapped with `If(filterAttr, expr, null)`. This ensures that a - // non-matching expression from one side evaluates to null for rows that belong to the other side, - // which is safe for aggregate FILTER (WHERE ...) semantics and avoids computing values for - // irrelevant rows. The filter attributes themselves are appended to the merged expression list so - // they remain visible to the enclosing Aggregate that will consume them. A newPlanFilter with - // isNew=false was reused from a previous merge round and is already present in the merged child - // output, so it is not appended again. - private def mergeNamedExpressions( - newPlanExpressions: Seq[NamedExpression], - cachedPlanExpressions: Seq[NamedExpression], - newPlanMapping: AttributeMap[Attribute], - newPlanFilter: Option[(Attribute, Boolean)] = None, - cachedPlanFilter: Option[Attribute] = None) = { - val mergedExpressions = mutable.ArrayBuffer[NamedExpression](cachedPlanExpressions: _*) - val matchedCachedIndices = mutable.HashSet.empty[Int] - val newNPMapping = AttributeMap(newPlanExpressions.map { ne => - val mapped = mapAttributes(ne, newPlanMapping) - val withoutAlias = mapped match { - case Alias(child, _) => child - case e => e - } - val foundIdx = mergedExpressions.indexWhere { - case Alias(child, _) => child semanticEquals withoutAlias - case e => e semanticEquals withoutAlias - } - val resultAttr = if (foundIdx >= 0) { - // Matching expression: both sides compute the same value, no wrapping needed. - matchedCachedIndices += foundIdx - mergedExpressions(foundIdx).toAttribute - } else { - // Non-matching expression from the new plan side: wrap with the new plan filter so it - // is only computed for rows that belong to the new plan side. Plain attribute references - // are not wrapped since reading a column value is free. - val wrappedExpr: NamedExpression = newPlanFilter match { - case Some((f, _)) if !withoutAlias.isInstanceOf[Attribute] => - Alias(If(f, withoutAlias, Literal(null, withoutAlias.dataType)), mapped.name)() - case _ => mapped - } - mergedExpressions += wrappedExpr - wrappedExpr.toAttribute - } - ne.toAttribute -> resultAttr - }) - - // Wrap unmatched cached expressions with the cached plan's filter so they are only computed for - // rows that belong to the cached plan side. Plain attribute references are not wrapped. - cachedPlanFilter.foreach { f => - for (i <- 0 until cachedPlanExpressions.size if !matchedCachedIndices.contains(i)) { - mergedExpressions(i) match { - case ce @ Alias(child, _) if !child.isInstanceOf[Attribute] => - // Preserve the original ExprId so parent references to this cached attribute stay valid - // without a cp-side remapping. (The new-plan wrapping above uses a fresh ExprId because - // those aliases are appended rather than replacing an existing entry.) - mergedExpressions(i) = - Alias(If(f, child, Literal(null, child.dataType)), ce.name)( - exprId = ce.toAttribute.exprId) - case _ => // attribute or alias-of-attribute, no wrapping needed - } - } - } - - newPlanFilter.foreach { - case (f, true) => mergedExpressions += f - case _ => - } - cachedPlanFilter.foreach(mergedExpressions += _) - - (mergedExpressions.toSeq, newNPMapping) - } - - // Applies filter as a FILTER (WHERE ...) clause to every AggregateExpression in exprs, - // combining with any pre-existing filter on the aggregate via AND. - private def applyFilterToAggregateExpressions( - exprs: Seq[NamedExpression], - filter: Attribute): Seq[NamedExpression] = { - exprs.map(_.transform { - case ae: AggregateExpression => - val combinedFilter = ae.filter.fold[Expression](filter)(And(filter, _)) - val newAE = ae.copy(filter = Some(combinedFilter)) - newAE.copyTagsFrom(ae) - newAE - }.asInstanceOf[NamedExpression]) - } - - // Only allow aggregates of the same implementation because merging different implementations - // could cause performance regression. - private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = { - val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan => - plan.aggregateExpressions.flatMap(_.collect { - case a: AggregateExpression => a - }) - } - val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions) - - val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) = - aggregateExpressionsSeq.zip(groupByExpressionSeq).map { - case (aggregateExpressions, groupByExpressions) => - Aggregate.supportsHashAggregate( - aggregateExpressions.flatMap( - _.aggregateFunction.aggBufferAttributes), groupByExpressions) - } - - newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate || - newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && { - val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) = - aggregateExpressionsSeq.zip(groupByExpressionSeq).map { - case (aggregateExpressions, groupByExpressions) => - Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions) - } - newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate || - newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate - } - } -} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala new file mode 100644 index 0000000000000..d58404cc2a1a6 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindow.scala @@ -0,0 +1,140 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, AttributeSet, NamedExpression} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, Window} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.WINDOW + +/** + * A [[Window]] operator does not project its output partitioning or ordering through aliases: its + * physical `outputPartitioning`/`outputOrdering` are pure pass-throughs of the child's. So when a + * window partitioned/ordered by `k` is followed by a consumer (aggregate, repartition, a chained + * window, ...) that requires distribution/ordering on a rename `k AS a`, a redundant shuffle and/or + * sort is inserted: the window already shuffled/sorted on `k`, but `HashPartitioning(k)` does not + * satisfy `ClusteredDistribution(a)` because the rename lives in a [[Project]] *below* the window + * while the analyzer-inserted [[Project]] *above* references `a` as a bare [[Attribute]] (empty + * alias map), and `k` and `a`, though value-identical, have distinct expr ids. + * + * This rule pulls such aliases up from the bottom [[Project]] into the top [[Project]], across the + * chain of one or more [[Window]] operators in between (`Project - Window+ - Project`; adjacent + * windows that cannot be collapsed, e.g. with different order specs, leave no [[Project]] between + * them). The top project's alias map then maps `k -> a`, and the + * `PartitioningPreservingUnaryExecNode` / `OrderPreservingUnaryExecNode` machinery that + * `ProjectExec` mixes in projects the windows' `HashPartitioning(k)`/`SortOrder(k)` up through the + * alias, satisfying the consumer. The [[Window]] operators themselves are left untouched. + * + * Besides removing the redundant shuffle/sort, this also narrows the data crossing the window's + * shuffle: the alias no longer flows through the window as a separate column (only its input `k`, + * already needed for partitioning, does) and is recomputed cheaply in the top project above the + * exchange. + * + * An entry of the bottom project is pulled up only when: + * - it is an [[Alias]] (bare pass-through attributes stay below so the windows keep producing + * them for the top project to reference); + * - it is deterministic: a nondeterministic alias (e.g. `rand()`, `spark_partition_id()`) + * evaluated above the window's exchange/sort instead of below it could produce different + * values, so it must stay below; + * - no window in the chain references it, so window semantics are unaffected; + * - all of its input attributes remain produced by the pruned bottom project (i.e. they are + * referenced by some window and thus retained), so a computed alias whose inputs would be + * dropped is not lifted above windows that no longer output them; + * - the top project consumes its output solely as a bare pass-through attribute (never inside a + * larger expression), so replacing that attribute with the alias fully preserves it. + * + * The rewrite is a no-op on its own output (a moved alias is no longer a bare attribute above nor + * present below), so it converges immediately at the fixed point. + */ +object PullUpProjectAliasThroughWindow extends Rule[LogicalPlan] { + + /** + * Matches a non-empty chain of [[Window]] operators bottomed out by a [[Project]], returning the + * windows top-to-bottom and that bottom project. + */ + private object WindowChain { + def unapply(plan: LogicalPlan): Option[(Seq[Window], Project)] = plan match { + case w @ Window(_, _, _, lower: Project, _) => Some((Seq(w), lower)) + case w @ Window(_, _, _, WindowChain(windows, lower), _) => Some((w +: windows, lower)) + case _ => None + } + } + + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( + _.containsPattern(WINDOW), ruleId) { + // Match the `Project - Window+ - Project` shape: the top project is the windows' parent + // scaffolding, and the bottom project is where the rename (`key AS userid`) is defined. + case p @ Project(projectList, WindowChain(windows, lower)) => + // Entries any window references must stay below: they define what the bottom project must + // keep producing, and in particular carry the partition/order key attributes the pulled-up + // aliases reuse. + val windowRefs = AttributeSet(windows.flatMap(_.references)) + val (retained, candidates) = + lower.projectList.partition(e => windowRefs.contains(e.toAttribute)) + val retainedAttrs = AttributeSet(retained.map(_.toAttribute)) + // Attributes the top project consumes as bare pass-through entries, and those it consumes + // inside a larger expression (an alias child, a window-output reference, ...). + val bareAttrs = AttributeSet(projectList.collect { case a: Attribute => a }) + val referencedByExpr = AttributeSet(projectList.flatMap { + case _: Attribute => Nil + case other => other.references + }) + // Pull up an alias only when: it is deterministic (a nondeterministic alias must not move + // across the window's exchange/sort, which would change its per-partition/per-row values); + // its inputs survive in the pruned bottom project; the top project passes its output through + // as a bare attribute; and the top project does not also consume that output inside an + // expression (which would require the windows to keep it). + val pullUp = candidates.collect { + case a: Alias + if a.deterministic && + a.references.subsetOf(retainedAttrs) && + bareAttrs.contains(a.toAttribute) && + !referencedByExpr.contains(a.toAttribute) => a + } + if (pullUp.isEmpty) { + p + } else { + val pullUpSet: Set[NamedExpression] = pullUp.toSet + // Key the lookup by expr id (not by `Attribute.equals`, which also compares qualifier): + // the top project may reference these attributes with a different qualifier than the alias + // carries below (e.g. across a subquery alias). + val pullUpMap = AttributeMap(pullUp.map(a => a.toAttribute -> a)) + val newProjectList = projectList.map { + // Rebuild the alias so it keeps the lower alias's child and expr id but adopts the top + // attribute's full identity (name, qualifier, metadata). The lookup is keyed by expr id + // only, so a resolved top attribute may carry a different name/qualifier/metadata than + // the lower alias (e.g. a different qualifier across a subquery alias); taking them from + // the top attribute keeps the output schema byte-for-byte identical. + case attr: Attribute => + pullUpMap.get(attr) match { + case Some(a) => Alias(a.child, attr.name)( + exprId = a.exprId, + qualifier = attr.qualifier, + explicitMetadata = Some(attr.metadata)) + case None => attr + } + case other => other + } + val newLower = lower.copy(projectList = lower.projectList.filterNot(pullUpSet.contains)) + val newChild = windows.foldRight(newLower: LogicalPlan) { (w, child) => + w.withNewChildren(child :: Nil) + } + Project(newProjectList, newChild) + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala index 7c3331dbadc97..5454c6d88b001 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnion.scala @@ -27,8 +27,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.{JOIN, UNION} import org.apache.spark.sql.internal.SQLConf /** - * Pushes down `Join` through `Union` when the right side of the join is small enough - * to broadcast. + * Pushes down `Join` through `Union` when every resulting join would broadcast its right side. * * This rule transforms the pattern: * {{{ @@ -42,9 +41,10 @@ import org.apache.spark.sql.internal.SQLConf * where each `condK` has the Union output attributes rewritten to the corresponding child's * output attributes. * - * This is beneficial when the right side is small enough to broadcast, because each Union - * branch can directly perform a broadcast join with the right side, avoiding the need to - * materialize the entire (potentially very large) Union result before the Join. + * This is beneficial when each branch broadcasts the right side, because the branch can then join + * it directly instead of materializing the whole Union first. The right side has to be the + * broadcast side: the rewrite duplicates it once per branch, and a copy that is probed rather than + * broadcast is scanned on its own. * * Applicable join types: Inner, LeftOuter. */ @@ -57,9 +57,9 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf) plan.transformUpWithPruning( _.containsAllPatterns(JOIN, UNION), ruleId) { - case join @ Join(u: Union, right, joinType, joinCond, hint) + case join @ Join(u: Union, right, joinType, _, _) if (joinType == Inner || joinType == LeftOuter) && - canPlanAsBroadcastHashJoin(join, conf) && + broadcastsRightForEveryBranch(u, join) && // Exclude right subtrees containing subqueries, as DeduplicateRelations // may not correctly handle correlated references when cloning. !right.exists(_.expressions.exists(SubqueryExpression.hasSubquery)) && @@ -75,17 +75,71 @@ case class PushDownJoinThroughUnion(override val conf: SQLConf) val deduped = dedupRight(right) (deduped, AttributeMap(right.output.zip(deduped.output))) } - val leftRewrites = AttributeMap(unionHeadOutput.zip(child.output)) - val newCond = joinCond.map(_.transform { - case a: Attribute if leftRewrites.contains(a) => leftRewrites(a) - case a: Attribute if rightRewrites.contains(a) => rightRewrites(a) - }) - Join(child, newRight, joinType, newCond, hint) + branchJoin(join, unionHeadOutput, child, newRight, rightRewrites) } u.withNewChildren(newChildren) } } + /** + * The join for one `Union` branch: the branch on the left, `newRight` on the right, and the + * condition rewritten from the `Union` output to the outputs of both new children. + */ + private def branchJoin( + join: Join, + unionHeadOutput: Seq[Attribute], + child: LogicalPlan, + newRight: LogicalPlan, + rightRewrites: AttributeMap[Attribute]): Join = { + val leftRewrites = AttributeMap(unionHeadOutput.zip(child.output)) + val newCond = join.condition.map(_.transform { + case a: Attribute if leftRewrites.contains(a) => leftRewrites(a) + case a: Attribute if rightRewrites.contains(a) => rightRewrites(a) + }) + Join(child, newRight, join.joinType, newCond, join.hint) + } + + /** + * Whether every join produced by the rewrite is expected to broadcast its right side. + * + * Asking whether a broadcast hash join is possible is not enough: for an inner join the planner + * may build from either side, choosing the smaller one when both qualify. The rewrite replaces + * the `Union` on the left with one of its children, so the build side is decided per branch + * against a smaller left. Any branch that ends up building from the left leaves its copy of the + * right side as a plain probe input, which is not reused, so the right side is read once per such + * branch instead of once in total. + * + * `getBroadcastHashJoinBuildSide` returns `None` when the join has no equi-join keys, so this one + * check also carries the requirement the removed `canPlanAsBroadcastHashJoin` conjunct used to. + * It is not a statement about what the planner ends up choosing: with keys no hash join supports + * it still answers from the sizes, while `JoinSelection` falls through to a sort merge join. A + * `SHUFFLE_MERGE` or `SHUFFLE_REPLICATE_NL` hint is handled here rather than there: the planner + * tries those between a hinted broadcast and a size-based one, and falls back to the sizes only + * when the hinted strategy does not apply, so declining to rewrite errs on the safe side. + * + * The result predicts rather than guarantees the final build side, because later rules and AQE + * re-estimate the sizes it reads. It is also all or nothing: one branch small enough to be the + * build side blocks the rewrite for the others, trading a missed optimization for never + * duplicating a probe side. Only an inner join can build from the left, so for a left outer join + * this reduces to asking whether the right side is broadcastable at all. + */ + private def broadcastsRightForEveryBranch(u: Union, join: Join): Boolean = { + val hintPicksOtherStrategy = + hintToSortMergeJoin(join.hint) || hintToShuffleReplicateNL(join.hint) + val unionHeadOutput = u.children.head.output + u.children.forall { child => + // The condition has to be rewritten to the branch output: `getBroadcastHashJoinBuildSide` + // extracts the equi-join keys, which do not resolve against a branch other than the first. + val probe = + branchJoin(join, unionHeadOutput, child, join.right, AttributeMap.empty[Attribute]) + // A hinted broadcast outranks the hints above, so only a size-based answer is vetoed. + val hinted = getBroadcastBuildSide(probe, hintOnly = true, conf).isDefined + val planned = + if (hinted || !hintPicksOtherStrategy) getBroadcastHashJoinBuildSide(probe, conf) else None + planned.contains(BuildRight) + } + } + /** * Creates a copy of `plan` with fresh ExprIds on all output attributes, * using the same "fake self-join + DeduplicateRelations" pattern as InlineCTE. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala index 2339596920b01..7a6827cefcb3e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala @@ -19,7 +19,8 @@ package org.apache.spark.sql.catalyst.optimizer import scala.collection.mutable -import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet, Expression, Literal, Or, SubqueryExpression} +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet} +import org.apache.spark.sql.catalyst.expressions.{Expression, Literal, Or, PredicateHelper, SubqueryExpression} import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule @@ -30,7 +31,7 @@ import org.apache.spark.util.collection.Utils * Infer predicates and column pruning for [[CTERelationDef]] from its reference points, and push * the disjunctive predicates as well as the union of attributes down the CTE plan. */ -object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] { +object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] with PredicateHelper { // CTE_id - (CTE_definition, precedence, predicates_to_push_down, attributes_to_prune) private type CTEMap = mutable.HashMap[Long, (CTERelationDef, Int, Seq[Expression], AttributeSet)] @@ -113,32 +114,52 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] { * In order to guarantee idempotency, we keep the predicates (if any) being pushed down by the * last iteration of this rule in a temporary field of `CTERelationDef`, so that on the current * iteration, we only push down predicates for a CTE def if there exists any new predicate that - * has not been pushed before. Also, since part of a new predicate might overlap with some - * existing predicate and it can be hard to extract only the non-overlapping part, we also keep - * the original CTE definition plan without any predicate push-down in that temporary field so - * that when we do a new predicate push-down, we can construct a new plan with all latest - * predicates over the original plan without having to figure out the exact predicate difference. + * has not been pushed before. When such a new predicate push-down happens, the CTE definition + * is rebuilt from its CURRENT child: the push-down filter this rule placed in the previous + * iteration is removed (wherever it sits) and the result is wrapped with the latest combined + * predicate. This preserves any change other rules made to the CTE definition's child in + * between (e.g. filters injected by `InferFiltersFromConstraints`, which runs in the `Once` + * batch sandwiched between the two fixedPoint batches containing this rule). If the previous + * push-down can no longer be found (another rule rewrote or merged it with other filters), + * the current child is used as-is: re-pushing the combined predicate is redundant but always + * semantics-preserving, since the disjunction of the reference predicates is valid for every + * row of the CTE definition. */ private def pushdownPredicatesAndAttributes( plan: LogicalPlan, cteMap: CTEMap): LogicalPlan = plan.transformWithSubqueries { case cteDef @ CTERelationDef(child, id, originalPlanWithPredicates, _, _, _) => val (_, _, newPreds, newAttrSet) = cteMap(id) - val originalPlan = originalPlanWithPredicates.map(_._1).getOrElse(child) val preds = originalPlanWithPredicates.map(_._2).getOrElse(Seq.empty) if (!isTruePredicate(newPreds) && newPreds.exists(newPred => !preds.exists(_.semanticEquals(newPred)))) { + val basePlan = originalPlanWithPredicates match { + case Some((_, prevPreds)) if prevPreds.nonEmpty => + // Remove the push-down filter this rule placed in the previous iteration. It is + // usually the top-level node of the child, but rules sharing the fixedPoint + // batches with this rule (e.g. `PushDownPredicates`) may have moved it deeper, + // possibly across attribute-renaming projections - hence the comparison is done + // on canonicalized conditions. If the previous push-down can no longer be found + // (another rule rewrote or merged it with other filters), the current child is + // used as-is: re-pushing the combined predicate is redundant but always + // semantics-preserving, since the disjunction of the reference predicates is + // valid for every row of the CTE definition. + removePushedDownFilter(child, prevPreds.reduce(Or)) + case _ => child + } val newCombinedPred = newPreds.reduce(Or) - val newChild = if (needsPruning(originalPlan, newAttrSet)) { - Project(newAttrSet.toSeq, originalPlan) + val newChild = if (needsPruning(basePlan, newAttrSet)) { + Project(newAttrSet.toSeq, basePlan) } else { - originalPlan + basePlan } cteDef.copy(child = Filter(newCombinedPred, newChild), - originalPlanWithPredicates = Some((originalPlan, newPreds))) + // The plan component of `originalPlanWithPredicates` is recorded but never read + // back: on the next iteration only the pushed predicates are consulted (see the + // `basePlan` computation above). + originalPlanWithPredicates = Some((basePlan, newPreds))) } else if (needsPruning(cteDef.child, newAttrSet)) { - cteDef.copy(child = Project(newAttrSet.toSeq, cteDef.child), - originalPlanWithPredicates = Some((originalPlan, preds))) + cteDef.copy(child = Project(newAttrSet.toSeq, cteDef.child)) } else { cteDef } @@ -156,6 +177,84 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] { } } + /** + * Removes the previous push-down filter (identified by its condition, `predicate`) from + * `plan`, wherever predicate push-down rules sharing the fixedPoint batches with this rule + * (e.g. `PushDownPredicates`) may have moved it. The descent mirrors the cases of + * `PushPredicateThroughNonJoin` and `PushPredicateThroughJoin`, translating `predicate` back + * the same way they translate the pushed condition (for `Project` and `Aggregate` it uses + * the very same `AliasHelper` utilities, so the two cannot drift apart): through projection + * and grouping-key aliases, positionally into each branch (`Union`), and unchanged + * through operators that pass the referenced attributes verbatim (`Join`, `Window`, and + * output-preserving unary nodes like `Filter`, `Sort`, `Repartition`). Descent stops at + * operators that remap attributes in other ways (e.g. `Generate`, `Expand`): if the filter + * cannot be located, the input plan is returned unchanged, and the caller re-pushes on top, + * which is redundant but always semantics-preserving. + * + * Ancestors of the removed filter are rebuilt with `withNewChildren` rather than direct + * case-class copies so that `TreeNode` tags (e.g. `Project.hiddenOutputTag`, which the + * analyzer sets on the projection above a natural/USING join) survive the rebuild. + */ + private def removePushedDownFilter(plan: LogicalPlan, predicate: Expression): LogicalPlan = { + def remove(current: LogicalPlan, target: Expression): (LogicalPlan, Boolean) = current match { + case Filter(cond, inner) if cond.canonicalized == target.canonicalized => + (inner, true) + case p: Project => + // Mirror PushPredicateThroughNonJoin: translate the target through the projection's + // aliases with the same helper it uses to move the filter below the projection. + val translated = replaceAlias(target, getAliasMap(p)) + val (newChild, removed) = remove(p.child, translated) + if (removed) (p.withNewChildren(Seq(newChild)), true) else (p, false) + case j: Join => + val (newLeft, removedFromLeft) = remove(j.left, target) + if (removedFromLeft) { + (j.withNewChildren(Seq(newLeft, j.right)), true) + } else { + val (newRight, removedFromRight) = remove(j.right, target) + if (removedFromRight) (j.withNewChildren(Seq(j.left, newRight)), true) else (j, false) + } + case u: Union => + // PushDownPredicates copies the filter into every branch, mapping the union output + // attributes to each branch's output positionally; remove it from every branch where + // it is found. The ExprId-to-output-index map is built once for all branches. + val outputIndexByExprId = u.output.map(_.exprId).zipWithIndex.toMap + var removedAny = false + val newChildren = u.children.map { branch => + val branchTarget = target.transform { + case a: Attribute if outputIndexByExprId.contains(a.exprId) => + branch.output(outputIndexByExprId(a.exprId)) + } + val (newBranch, removed) = remove(branch, branchTarget) + if (removed) { + removedAny = true + newBranch + } else { + branch + } + } + if (removedAny) (u.withNewChildren(newChildren), true) else (u, false) + case agg: Aggregate => + // Mirror PushPredicateThroughNonJoin: translate the target through the grouping + // aliases with the same helper it uses to move grouping-key filters below the + // aggregate (filters referencing aggregate expressions stay up, so they are always + // found above this node). + val translated = replaceAlias(target, getAliasMap(agg)) + val (newChild, removed) = remove(agg.child, translated) + if (removed) (agg.withNewChildren(Seq(newChild)), true) else (agg, false) + case w: Window => + // PushDownPredicates pushes filters referencing only partition columns below the + // window unchanged (partition columns are input attributes, so no translation). + val (newChild, removed) = remove(w.child, target) + if (removed) (w.withNewChildren(Seq(newChild)), true) else (w, false) + case other if other.children.length == 1 && + other.outputSet == other.children.head.outputSet => + val (newChild, removed) = remove(other.children.head, target) + if (removed) (other.withNewChildren(Seq(newChild)), true) else (other, false) + case _ => (current, false) + } + remove(plan, predicate)._1 + } + private def isTruePredicate(predicates: Seq[Expression]): Boolean = { predicates.length == 1 && predicates.head == Literal.TrueLiteral } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReusableBroadcastValueProjection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReusableBroadcastValueProjection.scala new file mode 100644 index 0000000000000..8841396adbbe4 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ReusableBroadcastValueProjection.scala @@ -0,0 +1,111 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.trees.TreePattern.{CTE, DYNAMIC_PRUNING_SUBQUERY} +import org.apache.spark.sql.types.{DateType, StringType, TimestampNTZType, TimestampType} + +/** + * Finds an existing broadcast whose stored rows provide a safe superset of a pruning domain. + */ +private[sql] object ReusableBroadcastValueProjection extends PredicateHelper { + + // Keep this an allowlist: projected values are evaluated directly from broadcast rows, so + // supported expressions must be deterministic and row-local, without subqueries or outer + // references, and preserve their time-zone, collation, and equality semantics. Extend it only + // when those properties hold; runtime limits and recoverable failures disable projection. + private def isSafeValueExpression(expression: Expression): Boolean = { + expression.deterministic && !expression.exists { + case _: OuterReference | _: SubqueryExpression => true + case _ => false + } && (expression match { + case attribute: Attribute => + UnsafeRow.isFixedLength(attribute.dataType) || attribute.dataType == StringType + case _: Literal => true + case DateAdd(startDate, days: Literal) => + isSafeValueExpression(startDate) && isSafeValueExpression(days) + case cast: Cast if cast.child.dataType == DateType && + (cast.dataType == TimestampType || cast.dataType == TimestampNTZType) => + isSafeValueExpression(cast.child) + case DateFormatClass(timestamp, format: Literal, Some(_)) + if format.dataType == StringType && format.value != null => + isSafeValueExpression(timestamp) + case _ => false + }) + } + + private def isSafeSourceHashKey(expression: Expression): Boolean = + expression.isInstanceOf[Attribute] && expression.deterministic + + /** + * Follows deterministic projections to the first inner equality join. Unmatched broadcast rows + * can add partition values, but cannot remove a partition needed by the actual join. + */ + def find( + valueExpression: Expression, + valuePlan: LogicalPlan, + excludedPlan: LogicalPlan): Option[BroadcastValueProjection] = { + + def descend( + value: Expression, + plan: LogicalPlan): Option[BroadcastValueProjection] = { + plan match { + case project: Project if project.projectList.forall(_.deterministic) => + val rewritten = replaceAlias(value, getAliasMap(project)) + if (rewritten.references.subsetOf(project.child.outputSet) && + isSafeValueExpression(rewritten)) { + descend(rewritten, project.child) + } else { + None + } + + case ExtractEquiJoinKeys(Inner, leftKeys, rightKeys, _, _, left, right, _) => + val valueFromLeft = value.references.nonEmpty && + value.references.subsetOf(left.outputSet) + val valueFromRight = value.references.nonEmpty && + value.references.subsetOf(right.outputSet) + val candidate = (valueFromLeft, valueFromRight) match { + case (true, false) => Some((left, leftKeys)) + case (false, true) => Some((right, rightKeys)) + case _ => None + } + + candidate.filter { case (source, hashKeys) => + !source.isStreaming && source.deterministic && + hashKeys.nonEmpty && hashKeys.forall(isSafeSourceHashKey) && + !source.containsAnyPattern(CTE, DYNAMIC_PRUNING_SUBQUERY) && + !source.exists(_.sameResult(excludedPlan)) + }.map { case (source, hashKeys) => + BroadcastValueProjection(source, hashKeys, value) + } + + case _ => None + } + } + + if (valuePlan.deterministic && isSafeValueExpression(valueExpression)) { + descend(valueExpression, valuePlan) + } else { + None + } + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala index 5aef82b64ed32..b330988927460 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala @@ -21,7 +21,8 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LogicalPlan} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE +import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CASE_WHEN, IF} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.IntegerType import org.apache.spark.util.collection.Utils @@ -223,8 +224,9 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] { case a: Aggregate if mayNeedtoRewrite(a) => rewrite(a) } - def rewrite(a: Aggregate): Aggregate = { + def rewrite(origAgg: Aggregate): Aggregate = { + val a = normalizeCountDistinctConditional(origAgg) val aggExpressions = collectAggregateExprs(a) val distinctAggs = aggExpressions.filter(_.isDistinct) @@ -419,6 +421,63 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] { } } + /** + * Canonicalizes COUNT(DISTINCT IF(cond, base, NULL)) and + * COUNT(DISTINCT CASE WHEN cond THEN base END) to COUNT(DISTINCT base) FILTER (WHERE cond). + * This reduces the number of distinct groups: multiple conditional counts on the same base + * column collapse into one group, shrinking the Expand fan-out from Nx to 1x. + * + * Note that the rewrite moves `base` out of the protective conditional branch: after the + * rewrite the Expand operator evaluates the distinct child for every input row, while + * originally it was only evaluated on rows where `cond` holds. To preserve the + * short-circuit semantics of IF/CASE WHEN, the rewrite is restricted to base + * expressions that can be evaluated unconditionally, i.e. cannot raise errors or change + * results when evaluated on extra rows (see [[ExprUtils.canEvaluateUnconditionally]]). + */ + private def normalizeCountDistinctConditional(a: Aggregate): Aggregate = { + if (!SQLConf.get.rewriteCountDistinctConditionalEnabled) return a + a.transformExpressionsUpWithPruning( + _.containsAnyPattern(IF, CASE_WHEN)) { + case ae @ AggregateExpression(count: Count, _, true, None, _) + if count.children.size == 1 => + extractCondAndBase(count.children.head) match { + case Some((cond, base)) => + ae.copy( + aggregateFunction = Count(base), + filter = Some(cond)) + case None => ae + } + } + } + + /** + * Matches IF(cond, base, null), CASE WHEN cond THEN base END, and + * CASE WHEN cond THEN base ELSE NULL END (including null wrapped in Cast). + * Multi-branch CaseWhen is intentionally not rewritten -- Or-flattening is out of scope. + * The base must be safe to evaluate unconditionally (the rewrite evaluates it on rows + * where the original branch would not have been taken, see + * [[ExprUtils.canEvaluateUnconditionally]]); the condition needs no check because it is + * evaluated unconditionally as the IF/CASE WHEN predicate anyway. + * Returns None for anything else. + */ + private def extractCondAndBase(expr: Expression): Option[(Expression, Expression)] = + expr match { + case If(cond, base, e) if isNullExpr(e) && ExprUtils.canEvaluateUnconditionally(base) => + Some((cond, base)) + case CaseWhen(Seq((cond, base)), None) if ExprUtils.canEvaluateUnconditionally(base) => + Some((cond, base)) + case CaseWhen(Seq((cond, base)), Some(e)) + if isNullExpr(e) && ExprUtils.canEvaluateUnconditionally(base) => + Some((cond, base)) + case _ => None + } + + private def isNullExpr(e: Expression): Boolean = e match { + case Literal(null, _) => true + case Cast(child, _, _, _) => isNullExpr(child) + case _ => false + } + private def collectAggregateExprs(a: Aggregate): Seq[AggregateExpression] = { // Collect all aggregate expressions. a.aggregateExpressions.flatMap { _.collect { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala index 13fa3b09c23ca..7e8e818e37b8a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala @@ -84,6 +84,9 @@ object ConstantFolding extends Rule[LogicalPlan] { // object and running eval unnecessarily. case l: Literal => l + // This foldable expression carries planning identity that must survive later optimizer batches. + case p: PercentileFusionArray => p + case Size(c: CreateArray, _) if c.children.forall(hasNoSideEffect) => Literal(c.children.length) case Size(c: CreateMap, _) if c.children.forall(hasNoSideEffect) => @@ -1155,6 +1158,8 @@ object FoldablePropagation extends Rule[LogicalPlan] { object SimplifyCasts extends Rule[LogicalPlan] { def apply(plan: LogicalPlan): LogicalPlan = plan.transformAllExpressionsWithPruning( _.containsPattern(CAST), ruleId) { + // Annotated STRING (flag off) is not unconstrained STRING. Dropping CAST(... AS STRING) + // would hide the type change. First-class CHAR/VARCHAR already fail e.dataType == StringType. case c @ Cast(e: NamedExpression, StringType, _, _) if e.dataType == StringType && e.metadata.contains(CHAR_VARCHAR_TYPE_STRING_METADATA_KEY) => c case Cast(e, dataType, _, _) if e.dataType == dataType => e diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala index c8c00a3fa13ae..f1a9ffd54e9a1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern._ import org.apache.spark.sql.catalyst.trees.TreePatternBits import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.catalyst.util.DateTimeUtils.{convertSpecialDate, convertSpecialTimestamp, convertSpecialTimestampNTZ, instantToMicros, localDateTimeToMicros} -import org.apache.spark.sql.catalyst.util.SparkDateTimeUtils.{instantToNanosOfDay, truncateTimeToPrecision} +import org.apache.spark.sql.catalyst.util.SparkDateTimeUtils.{instantToNanosOfDay, instantToTimestampNanos, localDateTimeToTimestampNanos, truncateTimeToPrecision} import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLExpr import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.types._ @@ -119,6 +119,11 @@ object ComputeCurrentTime extends Rule[LogicalPlan] { val timezone = Literal.create(conf.sessionLocalTimeZone, StringType) val currentDates = collection.mutable.HashMap.empty[ZoneId, Literal] val localTimestamps = collection.mutable.HashMap.empty[ZoneId, Literal] + // Nanosecond current-timestamp literals depend on the requested precision (sub-precision + // digits are floored), so they are cached separately from the microsecond ones. LTZ is keyed + // by precision only; NTZ (localtimestamp) is keyed by (zone, precision) like its micro sibling. + val currentTimestampNanos = collection.mutable.HashMap.empty[Int, Literal] + val localTimestampNanos = collection.mutable.HashMap.empty[(ZoneId, Int), Literal] // CAST_TO_TIMESTAMP is a dedicated tree-pattern bit set on Cast nodes whose target type is // any timestamp type (NTZ or LTZ family). This lets the rule reach both TIME -> TIMESTAMP_NTZ @@ -186,12 +191,25 @@ object ComputeCurrentTime extends Rule[LogicalPlan] { currentTimeType.precision) Literal.create(truncatedTime, TimeType(currentTimeType.precision)) case CurrentTimestamp() | Now() => currentTime + case ct: CurrentTimestampNanos => + currentTimestampNanos.getOrElseUpdate(ct.precision, { + Literal.create( + instantToTimestampNanos(instant, ct.precision), + TimestampLTZNanosType(ct.precision)) + }) case CurrentTimeZone() => timezone case localTimestamp: LocalTimestamp => localTimestamps.getOrElseUpdate(localTimestamp.zoneId, { val asDateTime = LocalDateTime.ofInstant(instant, localTimestamp.zoneId) Literal.create(localDateTimeToMicros(asDateTime), TimestampNTZType) }) + case lt: LocalTimestampNanos => + localTimestampNanos.getOrElseUpdate((lt.zoneId, lt.precision), { + val asDateTime = LocalDateTime.ofInstant(instant, lt.zoneId) + Literal.create( + localDateTimeToTimestampNanos(asDateTime, lt.precision), + TimestampNTZNanosType(lt.precision)) + }) } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala index edd63829a7113..8988183fa3c94 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/joins.scala @@ -415,18 +415,38 @@ trait JoinSelectionHelper extends Logging { result } - def canPlanAsBroadcastHashJoin(join: Join, conf: SQLConf): Boolean = join match { + /** + * The build side a broadcast hash join would use, or `None` when one is ruled out by the join + * shape or by a hint. + * + * `Some` does not promise the planner picks a broadcast hash join: a `SHUFFLE_MERGE` or + * `SHUFFLE_REPLICATE_NL` hint is tried before the sizes are consulted, join keys no hash join + * supports send it to a sort merge join, and AQE re-estimates the sizes at runtime. Within the + * broadcast decision itself this does follow the planner's precedence: a hinted broadcast first, + * a hinted shuffle hash join as a veto, then the sizes. Callers that only need to know whether a + * broadcast hash join is possible should use `canPlanAsBroadcastHashJoin`. + */ + def getBroadcastHashJoinBuildSide(join: Join, conf: SQLConf): Option[BuildSide] = join match { case ExtractEquiJoinKeys(_, leftKeys, rightKeys, _, _, _, _, _) => - val hashJoinSupport = hashJoinSupported(leftKeys, rightKeys) - val noShufflePlannedBefore = - !hashJoinSupport || getShuffleHashJoinBuildSide(join, hintOnly = true, conf).isEmpty - getBroadcastBuildSide(join, hintOnly = true, conf).isDefined || - (noShufflePlannedBefore && - getBroadcastBuildSide(join, hintOnly = false, conf).isDefined) - case j @ ExtractSingleColumnNullAwareAntiJoin(_, _) => canBroadcastBySize(j.right, conf) - case _ => false + // A shuffle hash hint outranks a size-based broadcast, so it vetoes one. Keys no hash join + // supports cannot honor that hint either, so it does not veto here; the sizes then still + // produce an answer, which over-approximates `JoinSelection` (it falls through to a sort + // merge join). That over-approximation predates this method and is kept deliberately, so + // that `canPlanAsBroadcastHashJoin` keeps its truth table. + val noShufflePlannedBefore = !hashJoinSupported(leftKeys, rightKeys) || + getShuffleHashJoinBuildSide(join, hintOnly = true, conf).isEmpty + getBroadcastBuildSide(join, hintOnly = true, conf).orElse { + if (noShufflePlannedBefore) getBroadcastBuildSide(join, hintOnly = false, conf) else None + } + // `JoinSelection` always builds from the right for this shape. + case j @ ExtractSingleColumnNullAwareAntiJoin(_, _) => + if (canBroadcastBySize(j.right, conf)) Some(BuildRight) else None + case _ => None } + def canPlanAsBroadcastHashJoin(join: Join, conf: SQLConf): Boolean = + getBroadcastHashJoinBuildSide(join, conf).isDefined + def canPruneLeft(joinType: JoinType): Boolean = joinType match { case Inner | LeftSemi | RightOuter => true case _ => false diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala index bf3e63571dbbc..099c64ac3b76e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/subquery.scala @@ -136,7 +136,7 @@ object RewritePredicateSubquery extends Rule[LogicalPlan] with PredicateHelper { case _ => false } } - case _ => false; + case _ => false } } @@ -418,7 +418,7 @@ object RewritePredicateSubquery extends Rule[LogicalPlan] with PredicateHelper { val exists = AttributeReference("exists", BooleanType, nullable = false)() // Deduplicate conflicting attributes if any. val newSub = dedupSubqueryOnSelfJoin(newPlan, sub, Some(values)) - val inConditions = values.zip(sub.output).map(EqualTo.tupled) + val inConditions = values.zip(newSub.output).map(EqualTo.tupled) // To handle a null-aware predicate not-in-subquery in nested conditions // (e.g., `v > 0 OR t1.id NOT IN (SELECT id FROM t2)`), we transform // `inCondition` (t1.id=t2.id) into `(inCondition) OR ISNULL(inCondition)`. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala index e3a034c1cb5c5..d65a9c7a36442 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala @@ -2068,6 +2068,8 @@ class AstBuilder extends DataTypeAstBuilder relationPrimary match { case _: AliasedQueryContext => case _: TableValuedFunctionContext => + case _: UnnestTableContext => + case _: JsonTableRelationContext => case other => throw QueryParsingErrors.invalidLateralJoinRelationError(other) } @@ -2566,6 +2568,8 @@ class AstBuilder extends DataTypeAstBuilder ctx.right match { case _: AliasedQueryContext => case _: TableValuedFunctionContext => + case _: UnnestTableContext => + case _: JsonTableRelationContext => case other => throw QueryParsingErrors.invalidLateralJoinRelationError(other) } @@ -3231,6 +3235,117 @@ class AstBuilder extends DataTypeAstBuilder buildTvfFromTableFunctionCall(ctx.tableFunctionCall, ctx.tableAlias, ctx.watermarkClause) } + /** + * Create a plan for the ANSI SQL `UNNEST(array [, array ...]) [WITH ORDINALITY]` relation used + * in the FROM clause. It is desugared into a [[Generate]] over a [[OneRowRelation]] backed by the + * [[Unnest]] generator, then wrapped with the optional table/column aliases via the shared + * FROM-clause aliasing helper. Correlated references (e.g. `FROM t, LATERAL UNNEST(t.arr)`) are + * handled by the surrounding `LATERAL` machinery, exactly like generator table functions such as + * `explode`. + */ + override def visitUnnestTable(ctx: UnnestTableContext): LogicalPlan = withOrigin(ctx) { + val unnest = ctx.unnest + val expressions = expressionList(unnest.expression) + val withOrdinality = unnest.ORDINALITY != null + val generate = Generate( + Unnest(expressions, withOrdinality), + unrequiredChildIndex = Nil, + outer = false, + qualifier = None, + generatorOutput = Nil, + child = OneRowRelation()) + mayApplyAliasPlan(unnest.tableAlias, generate) + } + + /** + * Create a plan for the SQL:2016 `JSON_TABLE` table-valued function. This builds a + * [[Generate]] over the [[JsonTable]] generator (reusing the existing Generate operator), so a + * downstream `SELECT` sees one output column per COLUMNS entry. + */ + override def visitJsonTableRelation( + ctx: JsonTableRelationContext): LogicalPlan = withOrigin(ctx) { + val jt = ctx.jsonTable + val jsonExpr = expression(jt.jsonExpr) + val rowPath = string(visitStringLit(jt.rowPath)) + + val columns = jt.jsonTableColumn.asScala.map(buildJsonTableColumn).toSeq + // Column names must be unique within a single JSON_TABLE. Whether two names that differ only + // in case collide follows the configured resolver, so `a` and `A` stay distinct under + // `spark.sql.caseSensitive`. + val normalize: String => String = + if (conf.caseSensitiveAnalysis) identity else _.toLowerCase(Locale.ROOT) + val duplicate = columns.groupBy(c => normalize(c.name)).collectFirst { + case (_, cols) if cols.length > 1 => cols.head.name + } + duplicate.foreach { name => + throw QueryParsingErrors.duplicateJsonTableColumnError(name, jt) + } + + val errorMode = if (jt.jsonTableOnErrorClause != null && jt.jsonTableOnErrorClause.ERROR != null + && jt.jsonTableOnErrorClause.NULL == null) { + JsonTableErrorMode.ErrorOnError + } else { + JsonTableErrorMode.NullOnError + } + + val generator = JsonTable(jsonExpr, rowPath, columns, errorMode) + val generate = Generate( + generator, + unrequiredChildIndex = Nil, + outer = false, + qualifier = None, + generatorOutput = columns.map(c => UnresolvedAttribute.quoted(c.name)), + child = OneRowRelation()) + mayApplyAliasPlan(jt.tableAlias, generate) + } + + /** + * The implicit JSON path for a column with no explicit PATH: the column name as a single JSON + * object key. Bracket syntax (`$['name']`) is used rather than `$.name` so a column name that + * contains a dot (e.g. `a.b`) reads the literal key `"a.b"` instead of the nested path `a.b`. A + * name containing a single quote cannot be represented and yields an unparseable path, which + * `JsonTable.checkInputDataTypes` rejects (such a column must use an explicit PATH). + */ + private def implicitJsonTablePath(name: String): String = s"$$['$name']" + + /** + * Build a single [[JsonTableColumn]] from a `jsonTableColumn` grammar context. A value column + * with no explicit PATH gets an implicit path derived from its name (see + * [[implicitJsonTablePath]]), matching the SQL standard / Oracle behavior. + */ + private def buildJsonTableColumn(ctx: JsonTableColumnContext): JsonTableColumn = withOrigin(ctx) { + ctx match { + case ord: JsonTableOrdinalityColumnContext => + JsonTableColumn( + name = getIdentifierText(ord.colName), + dataType = LongType, + path = None, + kind = JsonTableColumnKind.Ordinality) + case ex: JsonTableExistsColumnContext => + val name = getIdentifierText(ex.colName) + val path = Option(ex.path).map(p => string(visitStringLit(p))) + .getOrElse(implicitJsonTablePath(name)) + JsonTableColumn( + name = name, + // A column value is produced by a `Cast` to the declared type, so normalize CHAR/VARCHAR + // to STRING exactly as `visitCast` does; a raw CHAR/VARCHAR target has no encoder. + dataType = CharVarcharUtils.replaceCharVarcharWithStringForCast( + typedVisit[DataType](ex.dataType)), + path = Some(path), + kind = JsonTableColumnKind.Exists) + case v: JsonTableValueColumnContext => + val name = getIdentifierText(v.colName) + val path = Option(v.path).map(p => string(visitStringLit(p))) + .getOrElse(implicitJsonTablePath(name)) + JsonTableColumn( + name = name, + dataType = CharVarcharUtils.replaceCharVarcharWithStringForCast( + typedVisit[DataType](v.dataType)), + path = Some(path), + kind = JsonTableColumnKind.Value) + } + } + /** * Extract the source name from an identifiedByClause context. */ @@ -4054,6 +4169,105 @@ class AstBuilder extends DataTypeAstBuilder } } + /** + * Resolve a `jsonValueBehavior` clause (`NULL` / `ERROR` / `DEFAULT <expr>`) into a + * [[JsonValueBehavior]] and, for the `DEFAULT` case, its expression. + */ + private def buildJsonValueBehavior( + ctx: JsonValueBehaviorContext): (JsonValueBehavior, Option[Expression]) = ctx match { + case _: JsonValueBehaviorNullContext => (JsonValueBehavior.Null, None) + case _: JsonValueBehaviorErrorContext => (JsonValueBehavior.Error, None) + case d: JsonValueBehaviorDefaultContext => + (JsonValueBehavior.Default, Some(expression(d.defaultExpr))) + } + + /** + * Create a [[JsonValue]] expression for the SQL:2016 `JSON_VALUE` scalar function. The `ON EMPTY` + * / `ON ERROR` clauses default to `NULL` when absent, per the standard. + */ + override def visitJsonValue(ctx: JsonValueContext): Expression = withOrigin(ctx) { + val jsonExpr = expression(ctx.jsonExpr) + val path = string(visitStringLit(ctx.path)) + // Default RETURNING type is STRING. Normalize CHAR/VARCHAR to STRING for the cast, as the value + // is produced by a `Cast` to the declared type (a raw CHAR/VARCHAR target has no encoder). + val returning = Option(ctx.returning) + .map(dt => CharVarcharUtils.replaceCharVarcharWithStringForCast(typedVisit[DataType](dt))) + .getOrElse(StringType) + val (onEmpty, emptyDefault) = Option(ctx.emptyBehavior) + .map(buildJsonValueBehavior).getOrElse((JsonValueBehavior.Null, None)) + val (onError, errorDefault) = Option(ctx.errorBehavior) + .map(buildJsonValueBehavior).getOrElse((JsonValueBehavior.Null, None)) + JsonValue(jsonExpr, path, returning, onEmpty, onError, emptyDefault, errorDefault) + } + + /** + * Create a [[JsonExists]] expression for the SQL:2016 `JSON_EXISTS` predicate. The `ON ERROR` + * clause defaults to `FALSE` when absent, per the standard. + */ + override def visitJsonExists(ctx: JsonExistsContext): Expression = withOrigin(ctx) { + val jsonExpr = expression(ctx.jsonExpr) + val path = string(visitStringLit(ctx.path)) + val onError = Option(ctx.errorBehavior).map { b => + if (b.TRUE != null) JsonExistsBehavior.True + else if (b.FALSE != null) JsonExistsBehavior.False + else if (b.UNKNOWN != null) JsonExistsBehavior.Unknown + else JsonExistsBehavior.Error + }.getOrElse(JsonExistsBehavior.False) + JsonExists(jsonExpr, path, onError) + } + + /** + * Resolve a `jsonQueryBehavior` clause (`NULL` / `ERROR` / `EMPTY ARRAY` / `EMPTY OBJECT`) into a + * [[JsonQueryBehavior]]. + */ + private def buildJsonQueryBehavior(ctx: JsonQueryBehaviorContext): JsonQueryBehavior = ctx match { + case _: JsonQueryBehaviorNullContext => JsonQueryBehavior.Null + case _: JsonQueryBehaviorErrorContext => JsonQueryBehavior.Error + case _: JsonQueryBehaviorEmptyArrayContext => JsonQueryBehavior.EmptyArray + case _: JsonQueryBehaviorEmptyObjectContext => JsonQueryBehavior.EmptyObject + } + + /** + * Create a [[JsonQuery]] expression for the SQL:2016 `JSON_QUERY` function. The array wrapper + * defaults to `WITHOUT ARRAY WRAPPER`, quotes to `KEEP QUOTES`, and both `ON EMPTY` / `ON ERROR` + * to `NULL`, per the standard. `OMIT QUOTES` cannot be combined with an array wrapper. + */ + override def visitJsonQuery(ctx: JsonQueryContext): Expression = withOrigin(ctx) { + val jsonExpr = expression(ctx.jsonExpr) + val path = string(visitStringLit(ctx.path)) + // Default RETURNING type is STRING; the result is JSON text. A CHAR/VARCHAR RETURNING is + // normalized to STRING truly unconditionally: JSON_QUERY returns the fragment verbatim without + // a length-enforcing cast, so the result type must never advertise a CHAR/VARCHAR length it + // cannot enforce. The CharVarcharUtils helpers cannot be used here: they honor + // spark.sql.preserveCharVarcharTypeInfo and would leave a VARCHAR(n) length in the output type + // when that flag is set. A non-string RETURNING is left intact for checkInputDataTypes to fail. + val returning = Option(ctx.returning).map(typedVisit[DataType]).map { + case c: CharType => c.toStringType + case v: VarcharType => v.toStringType + case other => other + }.getOrElse(StringType) + val wrapper = Option(ctx.wrapper).map { + case _: JsonQueryWrapperWithoutContext => JsonQueryWrapper.Without + case w: JsonQueryWrapperWithContext => + if (w.wrapperType != null && w.wrapperType.getType == SqlBaseParser.CONDITIONAL) { + JsonQueryWrapper.Conditional + } else { + JsonQueryWrapper.Unconditional + } + }.getOrElse(JsonQueryWrapper.Without) + val quotes = Option(ctx.quotes).map { + case _: JsonQueryQuotesKeepContext => JsonQueryQuotes.Keep + case _: JsonQueryQuotesOmitContext => JsonQueryQuotes.Omit + }.getOrElse(JsonQueryQuotes.Keep) + // The OMIT QUOTES + array-wrapper invariant is enforced in JsonQuery.checkInputDataTypes so it + // holds for directly-constructed expressions too, not only this parser path. + val onEmpty = + Option(ctx.emptyBehavior).map(buildJsonQueryBehavior).getOrElse(JsonQueryBehavior.Null) + val onError = + Option(ctx.errorBehavior).map(buildJsonQueryBehavior).getOrElse(JsonQueryBehavior.Null) + JsonQuery(jsonExpr, path, returning, wrapper, quotes, onEmpty, onError) + } + /** * Create a (windowed) Function expression. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala index ff471cd6f00f8..7242c39f07ddf 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala @@ -111,6 +111,8 @@ object NormalizePlan extends PredicateHelper { udf.copy(resultId = ExprId(0)) case udaf: PythonUDAF => udaf.copy(resultId = ExprId(0)) + case agg: PythonAggregate => + agg.copy(resultId = ExprId(0)) case a: FunctionTableSubqueryArgumentExpression => a.copy(plan = normalizeExprIds(a.plan), exprId = ExprId(0)) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala index 785131875114d..ed4e889aa5d98 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala @@ -258,13 +258,29 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] * query operator based on the mapped expressions. */ def mapExpressions(f: Expression => Expression): this.type = { + mapExpressions(f, useFastEquals = true) + } + + /** + * A variant of [[mapExpressions]] that retains structurally equal replacement expressions. + */ + private[sql] def mapExpressionsWithReferenceEquality( + f: Expression => Expression): this.type = { + mapExpressions(f, useFastEquals = false) + } + + private def mapExpressions( + f: Expression => Expression, + useFastEquals: Boolean): this.type = { var changed = false @inline def transformExpression(e: Expression): Expression = { val newE = CurrentOrigin.withOrigin(e.origin) { f(e) } - if (newE.fastEquals(e)) { + // Reference equality preserves fresh stateful copies that fastEquals sees as unchanged. + val unchanged = if (useFastEquals) newE.fastEquals(e) else newE.eq(e) + if (unchanged) { e } else { changed = true @@ -577,6 +593,30 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]] transformDownWithSubqueriesAndPruning(AlwaysProcess.fn, UnknownRuleId)(f) } + /** + * A variant of [[transformDownWithSubqueries]] that retains structurally equal replacement + * plans and expressions. + */ + private[sql] def transformDownWithSubqueriesAndReferenceEquality( + f: PartialFunction[PlanType, PlanType]): PlanType = { + val g: PartialFunction[PlanType, PlanType] = new PartialFunction[PlanType, PlanType] { + override def isDefinedAt(x: PlanType): Boolean = true + + override def apply(plan: PlanType): PlanType = { + val transformed = f.applyOrElse[PlanType, PlanType](plan, identity) + transformed.mapExpressionsWithReferenceEquality( + _.transformDownWithReferenceEquality { + case planExpression: PlanExpression[PlanType @unchecked] => + val newPlan = planExpression.plan + .transformDownWithSubqueriesAndReferenceEquality(f) + planExpression.withNewPlan(newPlan) + }) + } + } + + transformDownWithReferenceEquality(g) + } + /** * Same as `transformUpWithSubqueries` except allows for pruning opportunities. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala index bd277e92d11d2..e054c955cc774 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/Command.scala @@ -35,7 +35,12 @@ trait Command extends LogicalPlan { // is created. That said, the statistics of a command is useless. Here we just return a dummy // statistics to avoid unnecessary statistics calculation of command's children. override def stats: Statistics = Statistics.DUMMY - final override val nodePatterns: Seq[TreePattern] = Seq(COMMAND) + // Every command carries the shared COMMAND pattern. Keeping this `final` guarantees no subclass + // can drop COMMAND; subclasses add their own identity pattern(s) via `nodePatternsInternal()`. + final override val nodePatterns: Seq[TreePattern] = Seq(COMMAND) ++ nodePatternsInternal() + + // Subclasses can override this to contribute additional identity tree patterns. + protected def nodePatternsInternal(): Seq[TreePattern] = Seq() } trait LeafCommand extends Command with LeafLike[LogicalPlan] diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala index d573f48862541..757da1dccdd9e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala @@ -242,6 +242,44 @@ trait LeafNode extends LogicalPlan with LeafLike[LogicalPlan] { throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3114") } +/** + * A single observation of a fully materialized leaf's current cache generation. + */ +private[sql] case class MaterializedLeafMetadata( + rowCount: BigInt, + sizeInBytes: BigInt, + isOutputRepeatable: Boolean, + isDurable: Boolean) { + def statsAvailable: Boolean = isOutputRepeatable && isDurable +} + +/** + * A leaf node that exposes materialization metadata used by + * [[org.apache.spark.sql.catalyst.optimizer.InjectRuntimeFilter]] to determine whether its output + * can be scanned again safely and profitably to build a runtime filter. + */ +private[sql] trait MaterializedLeafNode extends LeafNode { + /** A complete, generation-consistent snapshot, if the leaf is fully materialized. */ + def materializedMetadata: Option[MaterializedLeafMetadata] + + /** Cheap prerequisite for reading potentially usable materialization metadata. */ + def mayHaveUsableMaterializedStats: Boolean + + /** + * Whether the current materialized output has complete, accurate statistics and durable storage + * for another scan. This excludes memory-only storage levels, whose blocks may be discarded + * under memory pressure and recomputed. + */ + def statsAvailable: Boolean = + mayHaveUsableMaterializedStats && materializedMetadata.exists(_.statsAvailable) + + /** Whether scanning the materialized output again returns the same rows. */ + def isOutputRepeatable: Boolean = materializedMetadata.exists(_.isOutputRepeatable) + + /** Whether the original plan contains a predicate that is likely to be selective. */ + def hasSelectivePredicate: Boolean +} + /** * A abstract class for LogicalQueryStage that is visible in logical rewrites. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala index 26ac617366655..fb2dffed92b78 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala @@ -184,13 +184,44 @@ object Project { if (other == target) { col } else if (Cast.canANSIStoreAssign(other, target)) { - Cast(col, target, Option(conf.sessionLocalTimeZone), ansiEnabled = true) + storeAssignCast(col, other, target, conf) } else { throw QueryCompilationErrors.invalidColumnOrFieldDataTypeError(columnPath, other, target) } } } + /** + * Cast `col` for ANSI store assignment without using character-to-character CAST + * truncation (ISO 6.13). + * + * For CHAR/VARCHAR targets the plan is Cast to unconstrained STRING, then + * `stringLengthCheck` (write-side overflow). Avoid replaceCharVarcharWithString + * so first-class types stay CHAR/VARCHAR. + */ + private def storeAssignCast( + col: Expression, + other: DataType, + target: DataType, + conf: SQLConf): Expression = { + val (castTarget, lengthCheckType) = target match { + case c: CharType => (c.toStringType, Some(c: DataType)) + case v: VarcharType => (v.toStringType, Some(v: DataType)) + case otherType => (otherType, None) + } + val casted = if (other == castTarget) { + col + } else { + Cast(col, castTarget, Option(conf.sessionLocalTimeZone), ansiEnabled = true) + } + lengthCheckType match { + case Some(dt) if CharVarcharUtils.shouldApplyWriteSideLengthCheck(conf) => + CharVarcharUtils.stringLengthCheck(casted, dt) + case _ => + casted + } + } + private def reorderFields( fields: Seq[(String, Expression)], expected: Seq[StructField], @@ -2069,6 +2100,18 @@ object SampleMethod { } object Sample { + /** + * Resolves the seed of a sample, generating a random one when the user did not specify one. + * + * Generated seeds are non-negative. A pushed-down sample renders its seed into SQL as + * `REPEATABLE (<seed>)`, and the seed in that grammar does not accept a sign. A + * user-specified seed is returned unchanged, negative values included. + */ + def resolveSeed(seed: Option[Long]): Long = { + // `Utils` in this file is o.a.s.util.collection.Utils, so qualify the one we want here. + seed.getOrElse(org.apache.spark.util.Utils.random.nextLong() & Long.MaxValue) + } + /** * Convenience constructor that wraps a concrete seed in [[Some]]. * Use the case-class constructor directly with [[None]] when no seed @@ -2546,14 +2589,8 @@ case class AsOfJoin( override protected def stringArgs: Iterator[Any] = super.stringArgs.take(5) - override def output: Seq[Attribute] = { - joinType match { - case LeftOuter => - left.output ++ right.output.map(_.withNullability(true)) - case _ => - left.output ++ right.output - } - } + override def output: Seq[Attribute] = + AsOfJoin.computeOutput(joinType, left.output, right.output) def duplicateResolved: Boolean = left.outputSet.intersect(right.outputSet).isEmpty @@ -2582,6 +2619,19 @@ case class AsOfJoin( object AsOfJoin { + /** + * Computes the output attributes of an [[AsOfJoin]] given its join type and child outputs. + */ + def computeOutput( + joinType: JoinType, + leftOutput: Seq[Attribute], + rightOutput: Seq[Attribute]): Seq[Attribute] = joinType match { + case LeftOuter => + leftOutput ++ rightOutput.map(_.withNullability(true)) + case _ => + leftOutput ++ rightOutput + } + def apply( left: LogicalPlan, right: LogicalPlan, @@ -2750,13 +2800,11 @@ object AsOfJoin { operand.isInstanceOf[CreateNamedStruct] private[catalyst] def normalizeMatchOperands( - left: LogicalPlan, - right: LogicalPlan, + leftSet: AttributeSet, + rightSet: AttributeSet, expr1: Expression, operator: MatchComparisonOperator, expr2: Expression): (Expression, Expression, MatchComparisonOperator) = { - val leftSet = left.outputSet - val rightSet = right.outputSet val expr1Side = operandJoinSide(expr1, leftSet, rightSet, syntacticIsLeft = true) val expr2Side = operandJoinSide(expr2, leftSet, rightSet, syntacticIsLeft = false) (expr1Side, expr2Side) match { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala index d7a2f0e76e371..dd9a14c4ce336 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/cteOperators.scala @@ -99,10 +99,16 @@ case class UnionLoopRef( * A wrapper for CTE definition plan with a unique ID. * @param child The CTE definition query plan. * @param id The unique ID for this CTE definition. - * @param originalPlanWithPredicates The original query plan before predicate pushdown and the - * predicates that have been pushed down into `child`. This is - * a temporary field used by optimization rules for CTE predicate - * pushdown to help ensure rule idempotency. + * @param originalPlanWithPredicates The base plan of the last predicate pushdown and the + * predicates that have been pushed down into `child`. The + * base plan (the definition's child at the time of the last + * pushdown, with the filter that pushdown inserted removed) + * is recorded but not read back: the rule rebuilds from the + * definition's current child and only consults the + * predicates, both to detect newly appeared predicates and + * to locate the previous pushdown for removal. This is a + * temporary field used by optimization rules for CTE + * predicate pushdown to help ensure rule idempotency. * @param underSubquery If true, it means we don't need to add a shuffle for this CTE relation as * subquery reuse will be applied to reuse CTE relation output. * @param maxDepth The maximal depth of a recursion in a recursive CTE. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/hints.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/hints.scala index bb316f5683d8d..dc8c4cf3dfa85 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/hints.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/hints.scala @@ -198,11 +198,11 @@ case object NO_BROADCAST_AND_REPLICATION extends JoinStrategyHint { override def hintAliases: Set[String] = Set.empty } -abstract class AggregateHint; +abstract class AggregateHint -abstract class WindowHint; +abstract class WindowHint -abstract class SortHint; +abstract class SortHint /** * The callback for implementing customized strategies of handling hint errors. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala index 85f3c726c7ec7..b816016a3ec84 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.DescribeCommandSchema import org.apache.spark.sql.catalyst.trees.BinaryLike +import org.apache.spark.sql.catalyst.trees.TreePattern.{DELETE_FROM_TABLE, MERGE_INTO_TABLE, REPLACE_DATA, TreePattern, UPDATE_TABLE, WRITE_DELTA} import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.catalyst.util.TypeUtils.{ordinalNumber, toSQLExpr} @@ -448,6 +449,8 @@ case class ReplaceData( override protected def withNewChildInternal(newChild: LogicalPlan): ReplaceData = { copy(query = newChild) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(REPLACE_DATA) } /** @@ -558,6 +561,8 @@ case class WriteDelta( override protected def withNewChildInternal(newChild: LogicalPlan): WriteDelta = { copy(query = newChild) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(WRITE_DELTA) } trait V2CreateTableAsSelectPlan @@ -1098,6 +1103,8 @@ case class DeleteFromTable( override def child: LogicalPlan = table override protected def withNewChildInternal(newChild: LogicalPlan): DeleteFromTable = copy(table = newChild) + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(DELETE_FROM_TABLE) } /** @@ -1138,6 +1145,8 @@ case class UpdateTable( case r: NamedRelation => r.skipSchemaResolution case _ => false } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(UPDATE_TABLE) } /** @@ -1229,6 +1238,8 @@ case class MergeIntoTable( newRight: LogicalPlan): MergeIntoTable = { copy(targetTable = newLeft, sourceTable = newRight) } + + override protected def nodePatternsInternal(): Seq[TreePattern] = Seq(MERGE_INTO_TABLE) } object MergeIntoTable { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index 76e667590a33c..12511c1c4224b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -446,19 +446,26 @@ case class CoalescedNullAwareHashPartitioning( * because data sources produce them that way and `GroupPartitionsExec` sorts while grouping. * Sorted order is not a hard requirement, but it is a useful property: when both sides of a * storage-partitioned join report sorted keys, `EnsureRequirements` can often match them - * without inserting an additional `GroupPartitionsExec`. After a narrowing projection through - * `PartitioningPreservingUnaryExecNode`, the projected keys may no longer be sorted; this is - * acceptable because `EnsureRequirements` can always reconcile both sides via - * `GroupPartitionsExec` with `expectedPartitionKeys`. + * without inserting an additional `GroupPartitionsExec`. The keys may no longer be sorted after + * a narrowing projection through `PartitioningPreservingUnaryExecNode` or after `UnionExec` + * concatenates its children's keys; `EnsureRequirements` reconciles both sides either via + * `GroupPartitionsExec` with `expectedPartitionKeys`, or -- when the chosen shuffle spec is a + * `KeyedShuffleSpec` and the other child's spec is not compatible with it -- by shuffling that + * other child onto these keys in the order given here. * * 2. '''In KeyedShuffleSpec''': When used within `KeyedShuffleSpec`, the `partitionKeys` may not - * be in sorted order. `EnsureRequirements` handles this by building a common ordered set of - * keys and pushing them down to `GroupPartitionsExec` on both sides. + * be in sorted order, and consumers must not assume otherwise. * * == Partition Keys == * - `partitionKeys`: The partition keys, one per partition. May contain duplicates initially * (ungrouped state), but becomes unique after `GroupPartitionsExec` applies grouping. * + * `partitionKeys` is a physical layout indexed by partition id, not a set: partition `i` holds key + * `partitionKeys(i)`. A consumer must therefore treat the given order as authoritative rather than + * re-derive one. In particular `ShuffleExchangeExec` builds its `KeyGroupedPartitioner` from this + * order, so that a side shuffled onto a `KeyedPartitioning` lands in the same partitions as the + * side that declared it. + * * == Grouping State == * A KeyedPartitioning can be in two states: * @@ -538,7 +545,7 @@ case class KeyedPartitioning( @transient lazy val expressionDataTypes: Seq[DataType] = expressions.map(_.dataType) @transient lazy val keyRowOrdering = - RowOrdering.createNaturalAscendingOrdering(expressionDataTypes) + KeyedPartitioning.groupedKeyRowOrdering(expressionDataTypes) @transient lazy val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row) @@ -613,9 +620,12 @@ case class KeyedPartitioning( val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2) val projectedExpressions = joinKeyPositions.map(expressions) val projectedKeys = projectKeys(joinKeyPositions)._2 - val distinctProjectedKeys = projectedKeys.distinct + // Sort the distinct projected keys the same way `GroupPartitionsExec` does (both sort with + // `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the keyed side is grouped + // and the other side is re-shuffled using this spec, the two `KeyedPartitioning`s carry the + // same keys in a different order and `PartitioningCollection.fromPartitionings` rejects them. val projectedPartitioning = - KeyedPartitioning(projectedExpressions, distinctProjectedKeys, isGrouped = true) + new KeyedPartitioning(projectedExpressions, projectedKeys, isGrouped = false).toGrouped result.copy(partitioning = projectedPartitioning, joinKeyPositions = Some(joinKeyPositions)) } else { result @@ -658,6 +668,24 @@ object KeyedPartitioning { } } + /** + * The ascending ordering in which grouped partition keys are laid out, for keys of the given + * data types. + * + * This is a shared contract, not a convenience: with `allowKeysSubsetOfPartitionKeys`, + * `createShuffleSpec` declares the keyed side's projected keys in this order (via `toGrouped`), + * and the other side of the join may be shuffled onto exactly those keys, while the + * `GroupPartitionsExec` inserted on the keyed side independently re-groups its partitions with + * the same key positions and sorts them with this same ordering + * (`GroupPartitionsExec.groupAndSortByKeys`). If the two sorts diverged, inner joins would fail + * loudly at planning time -- `ShuffledJoin` wraps both sides' partitionings into a + * `PartitioningCollection`, whose invariant requires equal partition keys -- but join types that + * expose only one side's partitioning (e.g. LEFT OUTER) run nothing that compares the two + * orders, and silently return wrong results. + */ + def groupedKeyRowOrdering(dataTypes: Seq[DataType]): BaseOrdering = + RowOrdering.createNaturalAscendingOrdering(dataTypes) + /** * Projects a sequence of partition keys by selecting only the specified positions. */ @@ -1374,6 +1402,13 @@ case class KeyedShuffleSpec( override def canCreatePartitioning: Boolean = SQLConf.get.v2BucketingShuffleEnabled && !SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled && + // Shuffling another child onto these partition keys assigns each key a single partition, so + // an ungrouped partitioning cannot be reproduced: its duplicate keys live in more than one + // partition. This is the local gate. Such a spec is not reachable today for a non-local + // reason -- `EnsureRequirements` wraps a child whose `KeyedPartitioning` does not satisfy the + // distribution in a `GroupPartitionsExec` before it builds any spec -- so do not read this + // clause as redundant. + partitioning.isGrouped && partitioning.expressions.forall { e => e.isInstanceOf[AttributeReference] || e.isInstanceOf[TransformExpression] } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala index 2953c09e183dc..9221637d4db59 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/rules/RuleIdCollection.scala @@ -127,12 +127,14 @@ object RuleIdCollection { "org.apache.spark.sql.catalyst.optimizer.CollapseRepartition" :: "org.apache.spark.sql.catalyst.optimizer.CollapseWindow" :: "org.apache.spark.sql.catalyst.optimizer.ColumnPruning" :: + "org.apache.spark.sql.catalyst.optimizer.CombineApproximatePercentiles" :: "org.apache.spark.sql.catalyst.optimizer.CombineConcats" :: "org.apache.spark.sql.catalyst.optimizer.CombineFilters" :: "org.apache.spark.sql.catalyst.optimizer.CombineTypedFilters" :: "org.apache.spark.sql.catalyst.optimizer.CombineUnions" :: "org.apache.spark.sql.catalyst.optimizer.ConstantFolding" :: "org.apache.spark.sql.catalyst.optimizer.ConstantPropagation" :: + "org.apache.spark.sql.catalyst.optimizer.ConvertToCatalyst" :: "org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation" :: "org.apache.spark.sql.catalyst.optimizer.CostBasedJoinReorder" :: "org.apache.spark.sql.catalyst.optimizer.DecimalAggregates" :: @@ -162,6 +164,7 @@ object RuleIdCollection { "org.apache.spark.sql.catalyst.optimizer.Optimizer$OptimizeSubqueries" :: "org.apache.spark.sql.catalyst.optimizer.PropagateEmptyRelation" :: "org.apache.spark.sql.catalyst.optimizer.PruneFilters" :: + "org.apache.spark.sql.catalyst.optimizer.PullUpProjectAliasThroughWindow" :: "org.apache.spark.sql.catalyst.optimizer.PushDownJoinThroughUnion" :: "org.apache.spark.sql.catalyst.optimizer.PushDownLeftSemiAntiJoin" :: "org.apache.spark.sql.catalyst.optimizer.PushExtraPredicateThroughJoin" :: diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala index e82e6a30b9bba..e88b0c214c76e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala @@ -471,6 +471,22 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] transformDownWithPruning(AlwaysProcess.fn, UnknownRuleId)(rule) } + /** + * A variant of [[transformDown]] that retains structurally equal replacement nodes. + */ + private[sql] def transformDownWithReferenceEquality( + rule: PartialFunction[BaseType, BaseType]): BaseType = { + val afterRule = CurrentOrigin.withOrigin(origin) { + rule.applyOrElse(this, identity[BaseType]) + } + if (this eq afterRule) { + mapChildrenWithReferenceEquality(_.transformDownWithReferenceEquality(rule)) + } else { + afterRule.copyTagsFrom(this) + afterRule.mapChildrenWithReferenceEquality(_.transformDownWithReferenceEquality(rule)) + } + } + /** * Returns a copy of this node where `rule` has been recursively applied to it and all of its * children (pre-order). When `rule` does not apply to a given node it is left unchanged. @@ -736,6 +752,22 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]] } } + private[sql] final def mapChildrenWithReferenceEquality( + f: BaseType => BaseType): BaseType = { + val newChildren = children.map(f) + if (children.iterator.zip(newChildren.iterator).forall { case (oldChild, newChild) => + oldChild eq newChild + }) { + this + } else { + CurrentOrigin.withOrigin(origin) { + val res = withNewChildrenInternal(asIndexedSeq(newChildren)) + res.copyTagsFrom(this) + res + } + } + } + /** * Args to the constructor that should be copied, but not transformed. * These are appended to the transformed args automatically by makeCopy diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala index dfb815414dd35..d63c57beb11e7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala @@ -106,6 +106,7 @@ object TreePattern extends Enumeration { val TEMP_RESOLVED_COLUMN: Value = Value val TIME_WINDOW: Value = Value val TIME_ZONE_AWARE_EXPRESSION: Value = Value + val TRANSPILED_PYTHON_UDF: Value = Value val TRUE_OR_FALSE_LITERAL: Value = Value val USER_DEFINED_AGGREGATION: Value = Value val VARIANT_GET: Value = Value @@ -145,6 +146,9 @@ object TreePattern extends Enumeration { val COLLECT_METRICS: Value = Value val COMMAND: Value = Value val CTE: Value = Value + val DATA_SOURCE_V2_RELATION: Value = Value + val DATA_SOURCE_V2_SCAN_RELATION: Value = Value + val DELETE_FROM_TABLE: Value = Value val DESERIALIZE_TO_OBJECT: Value = Value val DF_DROP_COLUMNS: Value = Value val DISTINCT_LIKE: Value = Value @@ -164,6 +168,7 @@ object TreePattern extends Enumeration { val LIMIT: Value = Value val LOCAL_RELATION: Value = Value val LOGICAL_QUERY_STAGE: Value = Value + val MERGE_INTO_TABLE: Value = Value val METRIC_VIEW_PLACEHOLDER: Value = Value val NATURAL_LIKE_JOIN: Value = Value val NEAREST_BY_JOIN: Value = Value @@ -180,6 +185,7 @@ object TreePattern extends Enumeration { val RELATION_TIME_TRAVEL: Value = Value val REPARTITION_OPERATION: Value = Value val REBALANCE_PARTITIONS: Value = Value + val REPLACE_DATA: Value = Value val RESOLVED_METRIC_VIEW: Value = Value val SEQUENTIAL_STREAMING_UNION: Value = Value val SERIALIZE_FROM_OBJECT: Value = Value @@ -190,10 +196,12 @@ object TreePattern extends Enumeration { val UNION: Value = Value val UNPIVOT: Value = Value val UPDATE_EVENT_TIME_WATERMARK_COLUMN: Value = Value + val UPDATE_TABLE: Value = Value val TYPED_FILTER: Value = Value val WINDOW: Value = Value val WINDOW_GROUP_LIMIT: Value = Value val WITH_WINDOW_DEFINITION: Value = Value + val WRITE_DELTA: Value = Value val ZIP: Value = Value // Unresolved Plan patterns (Alphabetically ordered) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/PhysicalDataType.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/PhysicalDataType.scala index 9acbf4be66f9e..d77da0b159af8 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/PhysicalDataType.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/types/PhysicalDataType.scala @@ -17,8 +17,8 @@ package org.apache.spark.sql.catalyst.types -import scala.reflect.runtime.universe.TypeTag -import scala.reflect.runtime.universe.typeTag +import scala.reflect.ClassTag +import scala.reflect.classTag import org.apache.spark.sql.catalyst.expressions.{Ascending, BoundReference, InterpretedOrdering, SortOrder} import org.apache.spark.sql.catalyst.types.ops.TypeOps @@ -31,7 +31,7 @@ import org.apache.spark.util.ArrayImplicits._ sealed abstract class PhysicalDataType { private[sql] type InternalType private[sql] def ordering: Ordering[InternalType] - private[sql] val tag: TypeTag[InternalType] + private[sql] val tag: ClassTag[InternalType] } object PhysicalDataType { @@ -133,7 +133,7 @@ class PhysicalBinaryType() extends PhysicalDataType { (x: Array[Byte], y: Array[Byte]) => ByteArray.compareBinary(x, y) private[sql] type InternalType = Array[Byte] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalBinaryType extends PhysicalBinaryType @@ -143,14 +143,14 @@ class PhysicalBooleanType extends PhysicalDataType with PhysicalPrimitiveType { // Defined with a private constructor so the companion object is the only possible instantiation. private[sql] type InternalType = Boolean private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalBooleanType extends PhysicalBooleanType class PhysicalByteType() extends PhysicalIntegralType with PhysicalPrimitiveType { private[sql] type InternalType = Byte private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Byte]] override private[sql] val exactNumeric = ByteExactNumeric private[sql] val integral = implicitly[Integral[Byte]] @@ -162,7 +162,7 @@ class PhysicalCalendarIntervalType() extends PhysicalDataType { throw QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( "PhysicalCalendarIntervalType") override private[sql] type InternalType = Any - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalCalendarIntervalType extends PhysicalCalendarIntervalType @@ -180,7 +180,7 @@ case object PhysicalCalendarIntervalType extends PhysicalCalendarIntervalType class PhysicalTimestampNTZNanosType() extends PhysicalDataType { override private[sql] type InternalType = TimestampNanosVal override private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalTimestampNTZNanosType extends PhysicalTimestampNTZNanosType @@ -198,14 +198,14 @@ case object PhysicalTimestampNTZNanosType extends PhysicalTimestampNTZNanosType class PhysicalTimestampLTZNanosType() extends PhysicalDataType { override private[sql] type InternalType = TimestampNanosVal override private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalTimestampLTZNanosType extends PhysicalTimestampLTZNanosType case class PhysicalDecimalType(precision: Int, scale: Int) extends PhysicalFractionalType { private[sql] type InternalType = Decimal private[sql] val ordering = Decimal.DecimalIsFractional - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = Decimal.DecimalIsFractional override private[sql] def exactNumeric = DecimalExactNumeric private[sql] val fractional = Decimal.DecimalIsFractional @@ -225,7 +225,7 @@ class PhysicalDoubleType() extends PhysicalFractionalType with PhysicalPrimitive private[sql] type InternalType = Double private[sql] val ordering = (x: Double, y: Double) => SQLOrderingUtil.compareDoubles(x, y) - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Double]] override private[sql] def exactNumeric = DoubleExactNumeric private[sql] val fractional = implicitly[Fractional[Double]] @@ -240,7 +240,7 @@ class PhysicalFloatType() extends PhysicalFractionalType with PhysicalPrimitiveT private[sql] type InternalType = Float private[sql] val ordering = (x: Float, y: Float) => SQLOrderingUtil.compareFloats(x, y) - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Float]] override private[sql] def exactNumeric = FloatExactNumeric private[sql] val fractional = implicitly[Fractional[Float]] @@ -254,7 +254,7 @@ class PhysicalIntegerType() extends PhysicalIntegralType with PhysicalPrimitiveT // Defined with a private constructor so the companion object is the only possible instantiation. private[sql] type InternalType = Int private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Int]] override private[sql] val exactNumeric = IntegerExactNumeric private[sql] val integral = implicitly[Integral[Int]] @@ -267,7 +267,7 @@ class PhysicalLongType() extends PhysicalIntegralType with PhysicalPrimitiveType // Defined with a private constructor so the companion object is the only possible instantiation. private[sql] type InternalType = Long private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Long]] override private[sql] val exactNumeric = LongExactNumeric private[sql] val integral = implicitly[Integral[Long]] @@ -279,7 +279,7 @@ case class PhysicalMapType(keyType: DataType, valueType: DataType, valueContains // maps are not orderable, we use `ordering` just to support group by queries override private[sql] def ordering = interpretedOrdering override private[sql] type InternalType = MapData - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] @transient private[sql] lazy val interpretedOrdering: Ordering[MapData] = new Ordering[MapData] { @@ -348,14 +348,14 @@ class PhysicalNullType() extends PhysicalDataType with PhysicalPrimitiveType { override private[sql] def ordering = implicitly[Ordering[Unit]].asInstanceOf[Ordering[Any]] override private[sql] type InternalType = Any - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalNullType extends PhysicalNullType class PhysicalShortType() extends PhysicalIntegralType with PhysicalPrimitiveType { private[sql] type InternalType = Short private[sql] val ordering = implicitly[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] val numeric = implicitly[Numeric[Short]] override private[sql] val exactNumeric = ShortExactNumeric private[sql] val integral = implicitly[Integral[Short]] @@ -368,7 +368,7 @@ case class PhysicalStringType(collationId: Int) extends PhysicalDataType { // Defined with a private constructor so the companion object is the only possible instantiation. private[sql] type InternalType = UTF8String private[sql] val ordering = CollationFactory.fetchCollation(collationId).comparator.compare(_, _) - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } object PhysicalStringType { def apply(collationId: Int): PhysicalStringType = new PhysicalStringType(collationId) @@ -378,7 +378,7 @@ case class PhysicalArrayType( elementType: DataType, containsNull: Boolean) extends PhysicalDataType { override private[sql] type InternalType = ArrayData override private[sql] def ordering = interpretedOrdering - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] @transient private[sql] lazy val interpretedOrdering: Ordering[ArrayData] = new Ordering[ArrayData] { @@ -425,7 +425,7 @@ case class PhysicalStructType(fields: Array[StructField]) extends PhysicalDataTy override private[sql] type InternalType = Any override private[sql] def ordering = forSchema(this.fields.map(_.dataType).toImmutableArraySeq).asInstanceOf[Ordering[InternalType]] - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] private[sql] def forSchema(dataTypes: Seq[DataType]): InterpretedOrdering = { new InterpretedOrdering(dataTypes.zipWithIndex.map { @@ -436,7 +436,7 @@ case class PhysicalStructType(fields: Array[StructField]) extends PhysicalDataTy class PhysicalVariantType extends PhysicalDataType { private[sql] type InternalType = VariantVal - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] // TODO(SPARK-45891): Support comparison for the Variant type. override private[sql] def ordering = @@ -451,7 +451,7 @@ object UninitializedPhysicalType extends PhysicalDataType { throw QueryExecutionErrors.orderedOperationUnsupportedByDataTypeError( "UninitializedPhysicalType") override private[sql] type InternalType = Any - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } // Physical type for opaque, variable-length byte payloads that are addressed as a zero-copy @@ -466,7 +466,7 @@ object UninitializedPhysicalType extends PhysicalDataType { class PhysicalBinaryViewType extends PhysicalDataType { private[sql] val ordering = (x: BinaryView, y: BinaryView) => x.compareTo(y) private[sql] type InternalType = BinaryView - @transient private[sql] lazy val tag = typeTag[InternalType] + @transient private[sql] lazy val tag = classTag[InternalType] } case object PhysicalBinaryViewType extends PhysicalBinaryViewType diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharVarcharUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharVarcharUtils.scala index 9501986bb0c5d..00adcfe69bc56 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharVarcharUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/CharVarcharUtils.scala @@ -82,9 +82,10 @@ object CharVarcharUtils extends Logging with SparkCharVarcharUtils { * warning message if it has char or varchar types */ def replaceCharVarcharWithStringForCast(dt: DataType): DataType = { - if (SQLConf.get.charVarcharAsString) { + // standardSemantics takes precedence over legacy charVarcharAsString. + if (SQLConf.get.charVarcharAsString && !SQLConf.get.charVarcharStandardSemantics) { replaceCharVarcharWithString(dt) - } else if (hasCharVarchar(dt) && !SQLConf.get.preserveCharVarcharTypeInfo) { + } else if (hasCharVarchar(dt) && !SQLConf.get.charVarcharFirstClassTypes) { logWarning(log"The Spark cast operator does not support char/varchar type and simply treats" + log" them as string type. Please use string type directly to avoid confusion. Otherwise," + log" you can set ${MDC(CONFIG, SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key)} " + @@ -166,6 +167,15 @@ object CharVarcharUtils extends Logging with SparkCharVarcharUtils { }.getOrElse(expr) } + /** + * Write-side CHAR/VARCHAR length checks apply unless the session is on the legacy + * `charVarcharAsString` path with no first-class types. `standardSemantics` and + * `preserveCharVarcharTypeInfo` keep first-class types even if the legacy flag is also on. + */ + def shouldApplyWriteSideLengthCheck(conf: SQLConf): Boolean = { + !conf.charVarcharAsString || conf.charVarcharFirstClassTypes + } + def stringLengthCheck(expr: Expression, dt: DataType): Expression = { processStringForCharVarchar( expr, @@ -183,7 +193,7 @@ object CharVarcharUtils extends Logging with SparkCharVarcharUtils { case c: CharType if charFuncName.isDefined => StaticInvoke( classOf[CharVarcharCodegenUtils], - if (SQLConf.get.preserveCharVarcharTypeInfo) { + if (SQLConf.get.charVarcharFirstClassTypes) { c } else { c.toStringType @@ -195,7 +205,7 @@ object CharVarcharUtils extends Logging with SparkCharVarcharUtils { case v: VarcharType if varcharFuncName.isDefined => StaticInvoke( classOf[CharVarcharCodegenUtils], - if (SQLConf.get.preserveCharVarcharTypeInfo) { + if (SQLConf.get.charVarcharFirstClassTypes) { v } else { v.toStringType @@ -256,9 +266,23 @@ object CharVarcharUtils extends Logging with SparkCharVarcharUtils { } def addPaddingForScan(attr: Attribute): Expression = { - getRawType(attr.metadata).map { rawType => - processStringForCharVarchar( - attr, rawType, charFuncName = Some("readSidePadding"), varcharFuncName = None) + // Driven by metadata rather than attr.dataType even when Char/Varchar are first-class types. + // The metadata is the "not yet padded" marker: ApplyCharTypePadding rebuilds the relation via + // cleanAttrMetadata, so a second application of the rule finds no raw type and leaves the plan + // alone. Keying off attr.dataType instead would re-pad an already-padded scan on every pass and + // break the Once strategy's idempotence check. + getRawType(attr.metadata).map { dt => + if (SQLConf.get.charVarcharStandardSemantics) { + // Pad CHAR and enforce length limits for CHAR/VARCHAR (trim trailing blanks first). + processStringForCharVarchar( + attr, + dt, + charFuncName = Some("charTypeReadSideCheck"), + varcharFuncName = Some("varcharTypeReadSideCheck")) + } else { + processStringForCharVarchar( + attr, dt, charFuncName = Some("readSidePadding"), varcharFuncName = None) + } }.getOrElse(attr) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala index 24aabdc988fb5..d605cbae7bb44 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/DateTimeUtils.scala @@ -283,6 +283,18 @@ object DateTimeUtils extends SparkDateTimeUtils { instantToMicros(microsToInstant(micros).atZone(zoneId).plusMonths(months).toInstant) } + /** + * Adds a year-month interval expressed in months to a nanosecond-precision timestamp value while + * preserving the `nanosWithinMicro` remainder. + */ + def timestampNanosAddMonths( + start: TimestampNanosVal, + months: Int, + zoneId: ZoneId): TimestampNanosVal = { + val epochMicros = timestampAddMonths(start.epochMicros, months, zoneId) + TimestampNanosVal.fromParts(epochMicros, start.nanosWithinMicro) + } + /** * Adds a day-time interval expressed in microseconds to a timestamp at the given time zone. * It converts the input timestamp to a local timestamp, and adds the interval by: diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala index f8ca60c3a40af..574d4edb21bb7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/GeneratedColumn.scala @@ -17,11 +17,12 @@ package org.apache.spark.sql.catalyst.util -import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} import org.apache.spark.sql.catalyst.plans.logical.ColumnDefinition import org.apache.spark.sql.connector.catalog.{Column, Identifier, Table, TableCapability, TableCatalog, TableCatalogCapability} import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation import org.apache.spark.sql.types.{Metadata, MetadataBuilder, StructField, StructType} /** @@ -35,6 +36,13 @@ object GeneratedColumn { */ val GENERATION_EXPRESSION_METADATA_KEY = "GENERATION_EXPRESSION" + /** + * The metadata key marking a generated column in a write target's output as one whose value + * Spark computed from the generation expression, so that it is not validated against the + * expression that produced it. Only set while resolving a write. + */ + val AUTO_FILLED_GENERATED_COLUMN_METADATA_KEY = "__auto_filled_generated_column" + /** * Whether the given `field` is a generated column */ @@ -114,18 +122,74 @@ object GeneratedColumn { } /** - * Returns an attribute with the generation expression metadata removed. - * Used when the catalog does not support auto-filling generated columns on write. + * Returns `relation`'s output with every generated column's generation expression recorded in + * the attribute's metadata, which is where [[TableOutputResolver]] looks when it auto-fills the + * generated columns a write does not provide a value for. + * + * Generation expressions are internal metadata: they should neither surface in a DataFrame's + * schema nor propagate into tables created from it. A table's V2 columns are the persisted form + * and the source of truth, so the write path copies the expression into plan attribute metadata + * only for as long as resolving the write takes. + * + * The output is returned unchanged if the table does not ask Spark to handle generated columns. */ - def removeGenerationExpressionMetadata(attr: Attribute): Attribute = { - if (isGeneratedColumn(attr.metadata)) { - val cleaned = new MetadataBuilder() - .withMetadata(attr.metadata) - .remove(GENERATION_EXPRESSION_METADATA_KEY) - .build() - attr.withMetadata(cleaned) - } else { - attr + def attachGenerationExpressions(relation: DataSourceV2Relation): Seq[AttributeReference] = { + if (!supportsGeneratedColumnsOnWrite(relation.table)) { + return relation.output + } + val genExprs = relation.table.columns() + .flatMap(col => Option(col.generationExpression()).map(col.name -> _)) + .toMap + relation.output.map { attr => + genExprs.get(attr.name) match { + case Some(genExpr) => + withMetadataEntry(attr, GENERATION_EXPRESSION_METADATA_KEY, genExpr) + case None => attr + } + } + } + + /** + * When a write supplies its own value for a generated column, that value must agree with the + * generation expression, so ResolveTableConstraints validates it with a CheckInvariant, the same + * way it enforces a table's CHECK constraints. In contrast, values computed by Spark from the + * generation expression pass that check by construction, so this mark is what tells the rule to + * skip them. + * + * Returns a write target's `output` with the generated columns named in `autoFilled` marked as + * computed by Spark, so that it can skip CheckInvariant validation. + * + * Columns are left unmarked by default, so a value that did not come from the generation + * expression is always validated. + */ + def markAutoFilledGeneratedColumns( + output: Seq[AttributeReference], + autoFilled: Set[String]): Seq[AttributeReference] = { + output.map { attr => + if (autoFilled.contains(attr.name)) { + withMetadataEntry(attr, AUTO_FILLED_GENERATED_COLUMN_METADATA_KEY, "true") + } else { + attr + } } } + + /** + * Whether `attr` is a generated column whose value Spark computed from the generation + * expression (see [[markAutoFilledGeneratedColumns]]). + */ + def isAutoFilledGeneratedColumn(attr: Attribute): Boolean = { + attr.metadata.contains(AUTO_FILLED_GENERATED_COLUMN_METADATA_KEY) + } + + private def withMetadataEntry( + attr: AttributeReference, + key: String, + value: String): AttributeReference = { + val metadata = new MetadataBuilder() + .withMetadata(attr.metadata) + .putString(key, value) + .build() + attr.withMetadata(metadata).asInstanceOf[AttributeReference] + } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala index 68529e41937e1..b9c0a0255f310 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ResolveDefaultColumnsUtil.scala @@ -439,7 +439,8 @@ object ResolveDefaultColumns extends QueryErrorsBase throw QueryCompilationErrors.defaultValuesDataTypeError( statementType, colName, defaultSQL, dataType, other.dataType)) } - if (!conf.charVarcharAsString && CharVarcharUtils.hasCharVarchar(dataType) && ret.foldable) { + if (CharVarcharUtils.hasCharVarchar(dataType) && + CharVarcharUtils.shouldApplyWriteSideLengthCheck(conf) && ret.foldable) { CharVarcharUtils.stringLengthCheck(ret, dataType).eval(EmptyRow) } ret diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala index 3fb20dcd6420d..d3cd4c3956bf9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala @@ -98,8 +98,14 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) extends L && SQLConf.get.getConfByKeyStrict[Boolean]("spark.sql.optimizer.datasourceV2ExprFolding") => // If the expression is context independent foldable, we can convert it to a literal. // This is useful for increasing the coverage of V2 expressions. + // Folding returns the expression unchanged when it failed to evaluate inside a conditional + // branch, and recursing on an unchanged expression would loop forever. val constantExpr = ConstantFolding.constantFolding(expr) - generateExpression(constantExpr, isPredicate) + if (constantExpr.fastEquals(expr)) { + None + } else { + generateExpression(constantExpr, isPredicate) + } case col @ ColumnOrField(nameParts) => val ref = FieldReference(nameParts) if (isPredicate && col.dataType.isInstanceOf[BooleanType]) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala index 4fd85ef4923d2..2a412c538b0c1 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala @@ -130,6 +130,11 @@ package object util extends Logging { case c: Cast if !c.containsTag(Cast.USER_SPECIFIED_CAST) => PrettyAttribute(usePrettyExpression(c.child, shouldTrimTempResolvedColumn).sql, c.dataType) case p: PythonFuncExpression => PrettyPythonUDF(p.name, p.dataType, p.children) + // Present a transpiled UDF exactly like the UDF it wraps, so auto-generated + // column names stay `f(a)` whether or not transpilation engages (the node + // carries the rewrite options as extra children, which must not leak into + // user-visible names). + case t: TranspiledPythonUDF => PrettyPythonUDF(t.name, t.dataType, t.pythonUDFExpr.children) } def quoteIdentifier(name: String): String = { @@ -248,7 +253,9 @@ package object util extends Logging { FileSourceGeneratedMetadataStructField.FILE_SOURCE_GENERATED_METADATA_COL_ATTR_KEY, MetadataColumn.PRESERVE_ON_DELETE, MetadataColumn.PRESERVE_ON_UPDATE, - MetadataColumn.PRESERVE_ON_REINSERT + MetadataColumn.PRESERVE_ON_REINSERT, + GeneratedColumn.GENERATION_EXPRESSION_METADATA_KEY, + GeneratedColumn.AUTO_FILLED_GENERATED_COLUMN_METADATA_KEY ) def removeInternalMetadata(schema: StructType, keepFieldIds: Boolean = false): StructType = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/StaxXmlParser.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/StaxXmlParser.scala index 82e127382b742..340dc61fb5112 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/StaxXmlParser.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/StaxXmlParser.scala @@ -560,7 +560,7 @@ class StaxXmlParser( } else { newRow(i) = row(i) } - i += 1; + i += 1 } if (badRecordException.isEmpty) { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/XmlOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/XmlOptions.scala index bd4da3b27c86e..49c82b2c0e56f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/XmlOptions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/xml/XmlOptions.scala @@ -78,15 +78,15 @@ class XmlOptions( } private def getBool(paramName: String, default: Boolean = false): Boolean = { - val param = parameters.getOrElse(paramName, default.toString) - if (param == null) { - default - } else if (param.toLowerCase(Locale.ROOT) == "true") { - true - } else if (param.toLowerCase(Locale.ROOT) == "false") { - false - } else { - throw QueryExecutionErrors.paramIsNotBooleanValueError(paramName) + val paramValue = parameters.get(paramName) + paramValue match { + case None => default + case Some(null) => default + case Some(value) => value.toLowerCase(Locale.ROOT) match { + case "true" => true + case "false" => false + case _ => throw QueryExecutionErrors.paramIsNotBooleanValueError(paramName) + } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala index a5f1ca7f1d289..37734a460986f 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Implicits.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.catalog.{BucketSpec, ClusterBySpec} import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, quoteNameParts, QuotingUtils} +import org.apache.spark.sql.catalyst.util.{quoteIfNeeded, quoteNameParts, removeInternalMetadata, QuotingUtils} import org.apache.spark.sql.connector.expressions.{BucketTransform, ClusterByTransform, FieldReference, IdentityTransform, LogicalExpressions, Transform} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} import org.apache.spark.sql.types.StructType @@ -252,6 +252,15 @@ private[sql] object CatalogV2Implicits { implicit class ColumnsHelper(columns: Array[Column]) { def asSchema: StructType = CatalogV2Util.v2ColumnsToStructType(columns) def toAttributes: Seq[AttributeReference] = DataTypeUtils.toAttributes(asSchema) + + /** + * Same as [[toAttributes]], but strips the internal metadata that must not surface in a + * relation's output, such as generation expressions. Column IDs are the exception: although + * the key is listed in INTERNAL_METADATA_KEYS so that other paths drop it, the column-ID + * feature deliberately surfaces field IDs on a relation's output. + */ + def toOutputAttributes: Seq[AttributeReference] = + DataTypeUtils.toAttributes(removeInternalMetadata(asSchema, keepFieldIds = true)) } def parseColumnPath(name: String): Seq[String] = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala index c0905ede4f0df..9660ede921214 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala @@ -36,6 +36,8 @@ import org.apache.spark.sql.connector.catalog.TableChange._ import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.catalog.functions.UnboundFunction import org.apache.spark.sql.connector.expressions.{ClusterByTransform, LiteralValue, Transform} +import org.apache.spark.sql.errors.DataTypeErrors.toSQLId +import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, MapType, Metadata, MetadataBuilder, StructField, StructType} @@ -472,35 +474,110 @@ private[sql] object CatalogV2Util { catalog: CatalogPlugin, ident: Identifier, timeTravelSpec: Option[TimeTravelSpec] = None, - writePrivilegesString: Option[String] = None): Option[Table] = + writePrivilegesString: Option[String] = None, + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): Option[Table] = try { - Option(getTable(catalog, ident, timeTravelSpec, writePrivilegesString)) + Option(getTable(catalog, ident, timeTravelSpec, writePrivilegesString, options)) } catch { case _: NoSuchTableException => None case _: NoSuchDatabaseException => None } + /** + * Extracts the options that may select table state from a complete option map. These are the + * only options passed to `loadTable` and used to identify a pinned table state. + */ + def extractTableStateOptions( + catalog: CatalogPlugin, + options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = { + val stateKeys = catalog.asTableCatalog.tableStateOptionKeys.asScala + .map(_.toLowerCase(Locale.ROOT)) + .toSet + val projected = options.asCaseSensitiveMap().asScala.collect { + case (key, value) if stateKeys.contains(key.toLowerCase(Locale.ROOT)) => key -> value + }.toMap + new CaseInsensitiveStringMap(projected.asJava) + } + + /** + * Loads a table from the catalog. Callers may pass the complete option map, but only the keys the + * catalog declares via `tableStateOptionKeys()` are forwarded to `loadTable`, so the loaded table + * state stays independent of non-state options and of how many times the table is referenced. + */ def getTable( catalog: CatalogPlugin, ident: Identifier, timeTravelSpec: Option[TimeTravelSpec] = None, - writePrivilegesString: Option[String] = None): Table = { - if (timeTravelSpec.nonEmpty) { - assert(writePrivilegesString.isEmpty, "Should not write to a table with time travel") - timeTravelSpec.get match { - case v: AsOfVersion => - catalog.asTableCatalog.loadTable(ident, v.version) - case ts: AsOfTimestamp => - catalog.asTableCatalog.loadTable(ident, ts.timestamp) - } - } else { - if (writePrivilegesString.isDefined) { - val writePrivileges = writePrivilegesString.get.split(",").map(_.trim) - .map(TableWritePrivilege.valueOf).toSet.asJava - catalog.asTableCatalog.loadTable(ident, writePrivileges) - } else { - catalog.asTableCatalog.loadTable(ident) - } + writePrivilegesString: Option[String] = None, + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): Table = { + val timeTravel: TimeTravel = timeTravelSpec match { + case Some(v: AsOfVersion) => new TimeTravel.AsOfVersion(v.version) + case Some(ts: AsOfTimestamp) => new TimeTravel.AsOfTimestamp(ts.timestamp) + case None => null + } + val context = new TableContext(timeTravel, parseWritePrivileges(writePrivilegesString)) + val stateOptions = extractTableStateOptions(catalog, options) + catalog.asTableCatalog.loadTable(ident, context, stateOptions) + } + + /** + * Loads a table for a write, forwarding the required privileges and only the write options that + * the catalog declares may affect table state. The complete option map remains on the write + * relation for write planning. + */ + def loadTableForV2Write( + catalog: CatalogPlugin, + ident: Identifier, + writePrivileges: Set[TableWritePrivilege], + options: CaseInsensitiveStringMap): Table = { + rejectTimeTravelOptionsForWrite(catalog, ident, options) + loadTableForWrite(catalog, ident, writePrivileges, options) + } + + /** + * Loads a table for a write without validating the complete write option map. This is used by + * callers that must inspect whether the loaded table falls back to V1 before applying V2-only + * option validation. + */ + def loadTableForWrite( + catalog: CatalogPlugin, + ident: Identifier, + writePrivileges: Set[TableWritePrivilege], + options: CaseInsensitiveStringMap): Table = { + val context = new TableContext(null, writePrivileges.asJava) + val stateOptions = extractTableStateOptions(catalog, options) + catalog.asTableCatalog.loadTable(ident, context, stateOptions) + } + + def rejectTimeTravelOptionsForWrite( + catalog: CatalogPlugin, + ident: Identifier, + options: CaseInsensitiveStringMap): Unit = { + if (containsTimeTravelOptions(options)) { + throw QueryCompilationErrors.timeTravelUnsupportedError( + toSQLId(ident.toQualifiedNameParts(catalog))) + } + } + + def containsTimeTravelOptions(options: CaseInsensitiveStringMap): Boolean = { + val conf = SQLConf.get + Seq( + conf.getConf(SQLConf.TIME_TRAVEL_TIMESTAMP_KEY), + conf.getConf(SQLConf.TIME_TRAVEL_VERSION_KEY)).exists(options.containsKey) + } + + /** + * Parses the comma-separated write-privileges string (as carried in the internal + * [[org.apache.spark.sql.catalyst.analysis.UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES]] + * option) into a set of [[TableWritePrivilege]]. Returns an empty set when absent (a read). + */ + private def parseWritePrivileges( + writePrivilegesString: Option[String]): util.Set[TableWritePrivilege] = { + writePrivilegesString match { + case Some(str) => + str.split(",").map(_.trim).map(TableWritePrivilege.valueOf).toSet.asJava + case None => + util.Set.of() } } @@ -527,25 +604,19 @@ private[sql] object CatalogV2Util { loadTable(catalog, ident).map(DataSourceV2Relation.create(_, Some(catalog), Some(ident))) } - def isSameTable( - rel: DataSourceV2Relation, - catalog: CatalogPlugin, - ident: Identifier, - table: Table): Boolean = { - rel.catalog.contains(catalog) && rel.identifier.contains(ident) && rel.table.id == table.id - } - def lookupCachedRelation( cache: RelationCache, catalog: CatalogPlugin, ident: Identifier, table: Table, + options: CaseInsensitiveStringMap, conf: SQLConf): Option[DataSourceV2Relation] = { - val nameParts = ident.toQualifiedNameParts(catalog) - val cached = cache.lookup(nameParts, conf.resolver) - cached.collect { - case r: DataSourceV2Relation if isSameTable(r, catalog, ident, table) => r - } + cache.lookup( + catalog, + ident, + Some(table.id), + extractTableStateOptions(catalog, options), + conf.resolver).collect { case r: DataSourceV2Relation => r } } def isSessionCatalog(catalog: CatalogPlugin): Boolean = { @@ -570,6 +641,7 @@ private[sql] object CatalogV2Util { .withQueryColumnNames(existing.queryColumnNames) Option(existing.currentCatalog).foreach(builder.withCurrentCatalog) Option(existing.schemaMode).foreach(builder.withSchemaMode) + Option(existing.viewDependencies).foreach(builder.withViewDependencies) builder } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala index 18d94969aa27e..0028f78d22f9a 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala @@ -388,13 +388,19 @@ private[sql] object HoursTransform { } private[sql] final case class LiteralValue[T](value: T, dataType: DataType) extends Literal[T] { - override def toString: String = dataType match { - case StringType => s"'${s"$value".replace("'", "''")}'" - case BinaryType => - assert(value.isInstanceOf[Array[Byte]]) - val bytes = value.asInstanceOf[Array[Byte]] - "0x" + HexFormat.of().withUpperCase().formatHex(bytes) - case _ => s"$value" + override def toString: String = { + if (value == null) { + "NULL" + } else { + dataType match { + case StringType => s"'${s"$value".replace("'", "''")}'" + case BinaryType => + assert(value.isInstanceOf[Array[Byte]]) + val bytes = value.asInstanceOf[Array[Byte]] + "0x" + HexFormat.of().withUpperCase().formatHex(bytes) + case _ => s"$value" + } + } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index bf11971e99fee..ada71af43d3b5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -67,6 +67,22 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat ) } + def invalidUDFParameterPlaceholder(placeholder: String): Throwable = { + new AnalysisException( + errorClass = "INVALID_UDF_PARAMETER_PLACEHOLDER", + messageParameters = Map("placeholder" -> placeholder) + ) + } + + def invalidUDFParameterPlaceholderIndex(index: Int, numParams: Int): Throwable = { + new AnalysisException( + errorClass = "INVALID_UDF_PARAMETER_PLACEHOLDER_INDEX", + messageParameters = Map( + "index" -> index.toString, + "numParams" -> numParams.toString) + ) + } + def positionalAndNamedArgumentDoubleReference( routineName: String, parameterName: String): Throwable = { val errorClass = @@ -1224,6 +1240,10 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat unsupportedTableOperationError(table.name(), "batch scan") } + def unsupportedBatchWriteError(table: Table): Throwable = { + unsupportedTableOperationError(table.name(), "batch write") + } + def unsupportedStreamingScanError(table: Table): Throwable = { unsupportedTableOperationError(table.name(), "either micro-batch or continuous scan") } @@ -1304,15 +1324,6 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat "namespaceB" -> toSQLId(namespaceB))) } - def cannotCreateTableWithBothProviderAndSerdeError( - provider: Option[String], maybeSerdeInfo: Option[SerdeInfo]): Throwable = { - new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1058", - messageParameters = Map( - "provider" -> provider.toString, - "serDeInfo" -> maybeSerdeInfo.get.describe)) - } - def invalidFileFormatForStoredAsError(serdeInfo: SerdeInfo): Throwable = { new AnalysisException( errorClass = "_LEGACY_ERROR_TEMP_1059", @@ -1439,7 +1450,7 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat def tableNotSpecifyLocationUriError(identifier: TableIdentifier): Throwable = { new AnalysisException( - errorClass = "_LEGACY_ERROR_TEMP_1081", + errorClass = "TABLE_LOCATION_URI_NOT_SPECIFIED", messageParameters = Map("identifier" -> identifier.toString)) } @@ -2425,6 +2436,23 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat "functionList" -> groupAggPandasUDFNames.map(toSQLId).mkString(", "))) } + def invalidPythonAggregatePlacementError( + pythonAggregateNames: Seq[String]): Throwable = { + new AnalysisException( + errorClass = "INVALID_PYTHON_UDF_PLACEMENT", + messageParameters = Map( + "functionList" -> pythonAggregateNames.map(toSQLId).mkString(", "))) + } + + def invalidIncrementalPythonAggregatorBufferError( + name: String, bufferType: DataType): Throwable = { + new AnalysisException( + errorClass = "INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA", + messageParameters = Map( + "functionName" -> toSQLId(name), + "bufferType" -> Option(bufferType).map(toSQLType).getOrElse("NULL"))) + } + def ambiguousAttributesInSelfJoinError( ambiguousAttrs: Seq[AttributeReference]): Throwable = { new AnalysisException( @@ -3001,6 +3029,12 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat messageParameters = Map("windowExpressions" -> windowExpressions.toString())) } + def multiplePythonUDFTypesInWindowError(functionNames: Seq[String]): Throwable = { + new AnalysisException( + errorClass = "UNSUPPORTED_FEATURE.MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW", + messageParameters = Map("functionList" -> functionNames.map(toSQLId).mkString(", "))) + } + def escapeCharacterInTheMiddleError(pattern: String, char: String): Throwable = { new AnalysisException( errorClass = "INVALID_FORMAT.ESC_IN_THE_MIDDLE", @@ -4819,6 +4853,38 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat ) } + def missingAttributesError( + operator: LogicalPlan, + missingInput: Iterable[Attribute], + input: Iterable[Attribute], + attributesWithSameName: Iterable[Attribute]): Throwable = { + val missingAttributes = missingInput.map(toSQLExpr).mkString(", ") + val inputAttributes = input.map(toSQLExpr).mkString(", ") + val operatorString = operator.simpleString(SQLConf.get.maxToStringFields) + if (attributesWithSameName.nonEmpty) { + new AnalysisException( + errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION", + messageParameters = Map( + "missingAttributes" -> missingAttributes, + "input" -> inputAttributes, + "operator" -> operatorString, + "operation" -> attributesWithSameName.map(toSQLExpr).mkString(", ") + ), + origin = operator.origin + ) + } else { + new AnalysisException( + errorClass = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", + messageParameters = Map( + "missingAttributes" -> missingAttributes, + "input" -> inputAttributes, + "operator" -> operatorString + ), + origin = operator.origin + ) + } + } + def resolutionValidationError(cause: Throwable, plan: LogicalPlan): Throwable = { new ExtendedAnalysisException( new AnalysisException( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala index a9ea3f9f1c26c..63ab10b694222 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala @@ -269,6 +269,15 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE summary = "") } + def bitmapInputTooLargeError(inputNumBytes: Int, maxNumBytes: Int): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "BITMAP_INPUT_TOO_LARGE", + messageParameters = Map( + "inputNumBytes" -> inputNumBytes.toString, + "maxNumBytes" -> maxNumBytes.toString), + cause = null) + } + def invalidFractionOfSecondError(secAndMicros: Double): DateTimeException = { new SparkDateTimeException( errorClass = "INVALID_FRACTION_OF_SECOND", @@ -936,6 +945,12 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE cause = e) } + def cannotReadZipEntry(entry: String, path: String): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "CANNOT_READ_ZIP_ENTRY", + messageParameters = Map("entry" -> entry, "path" -> path)) + } + def cannotCreateColumnarReaderError(): Throwable = { new SparkException( errorClass = "_LEGACY_ERROR_TEMP_2065", @@ -968,7 +983,7 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE def writingJobFailedError(cause: Throwable): Throwable = { new SparkException( - errorClass = "_LEGACY_ERROR_TEMP_2070", + errorClass = "WRITING_JOB_FAILED", messageParameters = Map.empty, cause = cause) } @@ -1217,12 +1232,14 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE def cannotAcquireMemoryForWindowAggregateError( requestedBytes: Long, - receivedBytes: Long): SparkOutOfMemoryError = { + receivedBytes: Long, + consumerBreakdown: String): SparkOutOfMemoryError = { new SparkOutOfMemoryError( "UNABLE_TO_ACQUIRE_MEMORY", java.util.Map.of( "requestedBytes", requestedBytes.toString, - "receivedBytes", receivedBytes.toString)) + "receivedBytes", receivedBytes.toString, + "consumerBreakdown", consumerBreakdown)) } def rowLargerThan256MUnsupportedError(): SparkUnsupportedOperationException = { @@ -1467,6 +1484,17 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE "functionName" -> toSQLId(prettyName))) } + def invalidElementCountForTrimArrayError( + prettyName: String, numElements: Int, length: Int): SparkRuntimeException = { + new SparkRuntimeException( + errorClass = "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + messageParameters = Map( + "parameter" -> toSQLId("n"), + "functionName" -> toSQLId(prettyName), + "numElements" -> numElements.toString, + "length" -> length.toString)) + } + def invalidIndexOfZeroError(context: QueryContext): RuntimeException = { new SparkRuntimeException( errorClass = "INVALID_INDEX_OF_ZERO", @@ -1617,6 +1645,40 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE cause = e) } + def jsonQueryOnEmptyError(functionName: String, path: String, cause: Throwable): Throwable = { + new SparkRuntimeException( + errorClass = "JSON_QUERY_ON_ERROR.EMPTY", + messageParameters = Map("functionName" -> toSQLId(functionName), "path" -> toSQLValue(path)), + cause = cause) + } + + def jsonQueryOnErrorError(functionName: String, path: String, cause: Throwable): Throwable = { + new SparkRuntimeException( + errorClass = "JSON_QUERY_ON_ERROR.ERROR", + messageParameters = Map("functionName" -> toSQLId(functionName), "path" -> toSQLValue(path)), + cause = cause) + } + + def jsonValueOnEmptyError(functionName: String, path: String, cause: Throwable): Throwable = { + new SparkRuntimeException( + errorClass = "JSON_VALUE_ON_ERROR.EMPTY", + messageParameters = Map("functionName" -> toSQLId(functionName), "path" -> toSQLValue(path)), + cause = cause) + } + + def jsonValueOnErrorError(functionName: String, path: String, cause: Throwable): Throwable = { + new SparkRuntimeException( + errorClass = "JSON_VALUE_ON_ERROR.ERROR", + messageParameters = Map("functionName" -> toSQLId(functionName), "path" -> toSQLValue(path)), + cause = cause) + } + + def jsonExistsOnError(functionName: String, path: String): Throwable = { + new SparkRuntimeException( + errorClass = "JSON_EXISTS_ON_ERROR", + messageParameters = Map("functionName" -> toSQLId(functionName), "path" -> toSQLValue(path))) + } + def invalidKerberosConfigForHiveServer2Error(): Throwable = { new SparkException( errorClass = "_LEGACY_ERROR_TEMP_2179", @@ -2544,6 +2606,15 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE "detailMessage" -> detailMessage)) } + def invalidNormalizeFormError(form: String): RuntimeException = { + new SparkRuntimeException( + errorClass = "INVALID_PARAMETER_VALUE.NORMALIZE_FORM", + messageParameters = Map( + "parameter" -> toSQLId("form"), + "functionName" -> toSQLId("normalize"), + "form" -> toSQLValue(form, StringType))) + } + def hiveTableWithAnsiIntervalsError( table: TableIdentifier): SparkUnsupportedOperationException = { new SparkUnsupportedOperationException( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala index f4c3bc08f403f..aaee630bfe517 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowWriter.scala @@ -688,7 +688,7 @@ private[arrow] class IntervalYearWriter(val valueVector: IntervalYearVector) } override def setValue(input: SpecializedGetters, ordinal: Int): Unit = { - valueVector.setSafe(count, input.getInt(ordinal)); + valueVector.setSafe(count, input.getInt(ordinal)) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala index 0b7c829939ff7..7b2cfacb652fe 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.datasources.v2 -import java.util.{Optional, OptionalLong} +import java.util.{Collections, Optional, OptionalLong} import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.analysis.{MultiInstanceRelation, NamedRelation, TimeTravelSpec} @@ -25,7 +25,9 @@ import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatisti import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned} +import org.apache.spark.sql.catalyst.trees.TreePattern.{DATA_SOURCE_V2_RELATION, DATA_SOURCE_V2_SCAN_RELATION, TreePattern} import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.{removeInternalMetadata, truncatedString, CharVarcharUtils} import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, SupportsMetadataColumns, Table, TableCapability, TableCatalog, V2TableUtil} @@ -34,6 +36,7 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics, Histogram => V2Histogram, HistogramBin => V2HistogramBin} import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream} +import org.apache.spark.sql.internal.connector.{SupportsRuntimeCatalystFiltering, V2StatisticsUtils} import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -95,7 +98,7 @@ abstract class DataSourceV2RelationBase( table.asReadable.newScanBuilder(options).build() match { case r: SupportsReportStatistics => val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) + DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output) case _ => Statistics(sizeInBytes = conf.defaultSizeInBytes) } @@ -143,6 +146,8 @@ case class DataSourceV2Relation( table.capabilities.contains(TableCapability.AUTOMATIC_SCHEMA_EVOLUTION) def isVersioned: Boolean = table.version != null + + override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_RELATION) } /** @@ -158,8 +163,22 @@ case class DataSourceV2Relation( * @param keyGroupedPartitioning if set, the partitioning expressions that are used to split the * rows in the scan across different partitions * @param ordering if set, the ordering provided by the scan - * @param pushedFilters Catalyst expressions for filters that were fully pushed to the data - * source and do not appear as post-scan filters + * @param pushedFilters Catalyst expressions for filters that were fully pushed to the data source + * and do not appear as post-scan filters. These reference the relation's + * (pre-pruning) output, so they may reference columns pruned out of `output` + * (e.g. an unselected partition column the source enforces internally). This + * complete set is what lets `PlanMerger` soundly compare and re-enforce a + * scan's filters when fusing two scans via a Spark-side scan merge + * (`TableCapability.SCAN_MERGING`). + * @param mergeableScan whether this scan may be fused with an equivalent scan by a Spark-side scan + * merge (see `TableCapability.SCAN_MERGING`). + * Default false (not mergeable): only the plain column-pruning + filter + * pushdown path in `V2ScanRelationPushDown` sets this true, and only when the + * scan carries nothing a rebuilt scan cannot reproduce. A scan with a + * non-reproducible pushdown (aggregate, join, variant extraction, limit, + * offset, top-N, sample) or by any other rule stays not-mergeable by default, + * so merging is safe by construction -- a new scan-relation build site need + * not opt out. */ case class DataSourceV2ScanRelation( relation: DataSourceV2Relation, @@ -167,28 +186,74 @@ case class DataSourceV2ScanRelation( output: Seq[AttributeReference], keyGroupedPartitioning: Option[Seq[Expression]] = None, ordering: Option[Seq[SortOrder]] = None, - pushedFilters: Seq[Expression] = Seq.empty) extends LeafNode with NamedRelation { + pushedFilters: Seq[Expression] = Seq.empty, + mergeableScan: Boolean = false) extends LeafNode with NamedRelation { // TODO: Override validConstraints to return ExpressionSet(pushedFilters) so that pushed // filters participate in constraint propagation (InferFiltersFromConstraints, PruneFilters). + // Note: pushedFilters may reference columns pruned out of `output`, so constraint use must first + // intersect with `outputSet` (a constraint has to reference the node's output). // This changes which filters InferFiltersFromConstraints adds or removes (e.g., it may // skip adding IsNotNull when the scan already implies it, or infer new filters across // joins), so plan stability testing is needed first. /** * Resolved attributes that the scan declares for runtime filtering via - * [[SupportsRuntimeV2Filtering.filterAttributes]]. Empty when the scan - * does not implement [[SupportsRuntimeV2Filtering]] or exposes no attributes. + * [[SupportsRuntimeV2Filtering.filterAttributes]] or + * [[SupportsRuntimeCatalystFiltering.filterAttributes]]. Empty when the scan + * implements neither interface or exposes no attributes. */ - lazy val runtimeFilterAttrs: AttributeSet = scan match { - case s: SupportsRuntimeV2Filtering => - AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( - s.filterAttributes.toImmutableArraySeq, this)) - case _ => AttributeSet.empty + lazy val runtimeFilterAttrs: AttributeSet = { + checkRuntimeFilteringInterfaces() + val filterAttrs = scan match { + case s: SupportsRuntimeV2Filtering => s.filterAttributes + case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() + case _ => Array.empty[NamedReference] + } + resolveTopLevelFilterAttrs(filterAttrs) } + /** + * Resolved attributes for which a Catalyst runtime-filtering scan fully evaluates predicates. + * Empty for a [[SupportsRuntimeV2Filtering]] scan, which keeps its post-scan filters. + */ + lazy val fullyPushedRuntimeFilterAttrs: AttributeSet = { + checkRuntimeFilteringInterfaces() + val filterAttrs = scan match { + case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes() + case _ => Array.empty[NamedReference] + } + resolveTopLevelFilterAttrs(filterAttrs) + } + + /** + * Resolves the given runtime-filter references against this relation's output. Both runtime + * filtering interfaces require each reference to be a top-level attribute of the read schema, + * so a nested reference is rejected. + */ + private def resolveTopLevelFilterAttrs(filterAttrs: Array[NamedReference]): AttributeSet = { + filterAttrs.find(_.fieldNames.length > 1).foreach { ref => + throw SparkException.internalError( + s"Runtime filter attribute '${ref.fieldNames.mkString(".")}' declared by " + + s"${scan.getClass.getName} must be a top-level attribute of the scan read schema, " + + "but it is a nested reference.") + } + AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( + filterAttrs.toImmutableArraySeq, this)) + } + + override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_SCAN_RELATION) + override def name: String = relation.name + // A leaf relation references no upstream attributes. `pushedFilters` (and, for that matter, + // partitioning/ordering) are scan metadata, not references to resolve, and `pushedFilters` may + // reference columns pruned out of `output` (e.g. an unselected partition column). Without this + // override those would surface as `missingInput`, which the optimizer's plan-change validation + // flags as dangling references. `mapExpressions`/`transformExpressions` still rewrite the + // metadata expressions -- they iterate the product directly, independent of `references`. + override def references: AttributeSet = AttributeSet.empty + override def simpleString(maxFields: Int): String = { val outputString = truncatedString(output, "[", ", ", "]", maxFields) val nameWithTimeTravelSpec = relation.timeTravelSpec match { @@ -199,15 +264,40 @@ case class DataSourceV2ScanRelation( } override def computeStats(): Statistics = { - scan match { - case r: SupportsReportStatistics => - val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) - case _ => - Statistics(sizeInBytes = conf.defaultSizeInBytes) + if (conf.cboEnabled || conf.planStatsEnabled) { + computeFullStats() + } else { + computeSizeOnlyStats() } } + private def computeFullStats(): Statistics = { + V2StatisticsUtils.computeStats(scan) match { + case Some(v2Stats) => + DataSourceV2Relation.transformV2Stats(v2Stats, conf.defaultSizeInBytes, output) + case _ => defaultSizeOnlyStats + } + } + + private def computeSizeOnlyStats(): Statistics = { + V2StatisticsUtils.computeSizeInBytes(scan, EstimationUtils.getSizePerRow(output)) match { + case Some(sizeInBytes) => Statistics(sizeInBytes = sizeInBytes) + case _ => defaultSizeOnlyStats + } + } + + private def defaultSizeOnlyStats: Statistics = { + Statistics(sizeInBytes = conf.defaultSizeInBytes) + } + + private def checkRuntimeFilteringInterfaces(): Unit = scan match { + case _: SupportsRuntimeV2Filtering with SupportsRuntimeCatalystFiltering => + throw SparkException.internalError( + "A scan must not implement both SupportsRuntimeV2Filtering and " + + s"SupportsRuntimeCatalystFiltering, but ${scan.getClass.getName} implements both.") + case _ => + } + override def doCanonicalize(): DataSourceV2ScanRelation = { this.copy( relation = this.relation.copy( @@ -220,7 +310,9 @@ case class DataSourceV2ScanRelation( ordering = ordering.map( _.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output))) ), - pushedFilters = pushedFilters.map(QueryPlan.normalizeExpressions(_, output)) + // pushedFilters may reference columns pruned out of `output` (see the field doc), so they are + // normalized against the relation's full output rather than `output`. + pushedFilters = pushedFilters.map(QueryPlan.normalizeExpressions(_, relation.output)) ) } } @@ -276,7 +368,7 @@ case class StreamingDataSourceV2ScanRelation( override def computeStats(): Statistics = scan match { case r: SupportsReportStatistics => val statistics = r.estimateStatistics() - DataSourceV2Relation.transformV2Stats(statistics, None, conf.defaultSizeInBytes, output) + DataSourceV2Relation.transformV2Stats(statistics, conf.defaultSizeInBytes, output) case _ => Statistics(sizeInBytes = conf.defaultSizeInBytes) } @@ -320,6 +412,10 @@ object ExtractV2ScanInfo { } object DataSourceV2Relation { + + private val EMPTY_V2_COLUMN_STATS = + Collections.emptyMap[NamedReference, ColumnStatistics]() + def create( table: Table, catalog: Option[CatalogPlugin], @@ -419,22 +515,23 @@ object DataSourceV2Relation { */ def transformV2Stats( v2Statistics: V2Statistics, - defaultRowCount: Option[BigInt], defaultSizeInBytes: Long, output: Seq[Attribute] = Seq.empty): Statistics = { val numRows: Option[BigInt] = if (v2Statistics.numRows().isPresent) { Some(v2Statistics.numRows().getAsLong) } else { - defaultRowCount + None } var colStats: Seq[(Attribute, ColumnStat)] = Seq.empty[(Attribute, ColumnStat)] - if (!v2Statistics.columnStats().isEmpty) { - val v2ColumnStat = v2Statistics.columnStats() - val keys = v2ColumnStat.keySet() + // columnStats() may be null even when numRows/sizeInBytes are present, so normalize it to an + // empty map before conversion to avoid an NPE. + val v2ColumnStats = Option(v2Statistics.columnStats()).getOrElse(EMPTY_V2_COLUMN_STATS) + if (!v2ColumnStats.isEmpty) { + val keys = v2ColumnStats.keySet() keys.forEach(key => { - val colStat = v2ColumnStat.get(key) + val colStat = v2ColumnStats.get(key) val distinct: Option[BigInt] = if (colStat.distinctCount().isPresent) Some(colStat.distinctCount().getAsLong) else None val min: Option[Any] = if (colStat.min().isPresent) Some(colStat.min().get) else None @@ -463,9 +560,20 @@ object DataSourceV2Relation { }) }) } + val attributeStats = AttributeMap(colStats) + // Prefer the source-reported size. Otherwise infer a projection-aware size from the row count + // (numRows * outputRowSize via getOutputSize). Fall back to the default size when neither is + // available. + val sizeInBytes = if (v2Statistics.sizeInBytes().isPresent) { + BigInt(v2Statistics.sizeInBytes().getAsLong) + } else if (numRows.isDefined) { + EstimationUtils.getOutputSize(output, numRows.get, attributeStats) + } else { + BigInt(defaultSizeInBytes) + } Statistics( - sizeInBytes = v2Statistics.sizeInBytes().orElse(defaultSizeInBytes), + sizeInBytes = sizeInBytes, rowCount = numRows, - attributeStats = AttributeMap(colStats)) + attributeStats = attributeStats) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 97594802e258d..1dad20fdb5130 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -574,6 +574,16 @@ object SQLConf { .booleanConf .createWithDefault(true) + val ANALYZER_SINGLE_PASS_RESOLVER_ENABLE_ASOF_JOIN_RESOLUTION = + buildConf("spark.sql.analyzer.singlePassResolver.enableAsOfJoinResolution") + .internal() + .doc("When true, enables ASOF JOIN resolution in single-pass analyzer. " + + "Otherwise, resolution falls back to fixed-point analyzer.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + val MULTI_COMMUTATIVE_OP_OPT_THRESHOLD = buildConf("spark.sql.analyzer.canonicalization.multiCommutativeOpMemoryOptThreshold") .internal() @@ -600,6 +610,34 @@ object SQLConf { .booleanConf .createWithDefault(true) + val ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS = + buildConf("spark.sql.experimental.optimizer.transpilePyUDFs") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .doc("When true, attempt to transpile Python UDFs to Catalyst expressions. " + + "Transpilation also requires ANSI mode (spark.sql.ansi.enabled=true) -- " + + "the rewritten expressions target ANSI semantics, so with ANSI off the " + + "transpiler falls back to interpreted Python and a warning is logged at " + + "UDF construction. Transpiled UDFS attempt to match the Python functionality but " + + "may not be 100% equivalent. Some known differences include: overflows from input types " + + "(you can precast to decimal to avoid), type coercion on comparison, and implicit " + + "returns. This initial version only works with non-Connect Spark; Spark Connect " + + "support is to follow. This is an *experimental* feature." + ) + .version("4.3.0") + .booleanConf + .createWithDefault(false) + + + val PYTHON_UDF_TRANSPILERS = + buildConf("spark.sql.experimental.optimizer.pyTranspilers") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .doc("Comma-separated list of Python transpilers to attempt, in order. " + + "The first transpiler that successfully produces a Catalyst expression " + + "is used. Default: catalyst.") + .version("4.3.0") + .stringConf + .createWithDefault("catalyst") + val OPTIMIZER_EXCLUDED_RULES = buildConf("spark.sql.optimizer.excludedRules") .doc("Configures a list of rules to be disabled in the optimizer, in which the rules are " + "specified by their rule names and separated by comma. It is not guaranteed that all the " + @@ -790,6 +828,53 @@ object SQLConf { .booleanConf .createWithDefault(true) + val DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED = + buildConf("spark.sql.optimizer.dynamicPartitionPruning.broadcastProjection.enabled") + .internal() + .doc("When true, dynamic partition pruning may project a bounded value domain from the " + + "full rows of an existing ancestor broadcast. When false, Spark uses its existing " + + "broadcast-key reuse and subquery fallback behavior.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + + val DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_ROWS = + buildConf("spark.sql.optimizer.dynamicPartitionPruning.broadcastProjection.maxRows") + .internal() + .doc("Maximum number of broadcast value rows visited when projecting a dynamic partition " + + "pruning domain. If the limit is exceeded, the projection fails open rather than " + + "using a partial domain.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .intConf + .checkValue(_ >= 0, "The maximum number of DPP broadcast rows must be nonnegative.") + .createWithDefault(10000) + + val DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_BYTES = + buildConf("spark.sql.optimizer.dynamicPartitionPruning.broadcastProjection.maxBytes") + .internal() + .doc("Maximum cumulative size of distinct values projected from a broadcast for dynamic " + + "partition pruning. If the limit is exceeded, the projection fails open rather than " + + "using a partial domain.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .bytesConf(ByteUnit.BYTE) + .checkValue(_ >= 0, "The maximum size of DPP broadcast values must be nonnegative.") + .createWithDefaultString("8MB") + + val DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_SOURCE_BYTES = + buildConf("spark.sql.optimizer.dynamicPartitionPruning.broadcastProjection.maxSourceBytes") + .internal() + .doc("Maximum materialized broadcast size that may be rehydrated on the driver when " + + "projecting dynamic partition pruning values. If runtime statistics are unavailable " + + "or exceed the limit, the projection fails open.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .bytesConf(ByteUnit.BYTE) + .checkValue(_ >= 0, "The maximum DPP source broadcast size must be nonnegative.") + .createWithDefaultString("128MB") + val RUNTIME_FILTER_NUMBER_THRESHOLD = buildConf("spark.sql.optimizer.runtimeFilter.number.threshold") .doc("The total number of injected runtime filters (non-DPP) for a single " + @@ -801,8 +886,9 @@ object SQLConf { val RUNTIME_BLOOM_FILTER_ENABLED = buildConf("spark.sql.optimizer.runtime.bloomFilter.enabled") - .doc("When true and if one side of a shuffle join has a selective predicate, we attempt " + - "to insert a bloom filter in the other side to reduce the amount of shuffle data.") + .doc("When true and if one side of a shuffle join has a selective predicate, or is a " + + "fully materialized, repeatable cache with evidence that pruning is beneficial, we " + + "attempt to insert a bloom filter in the other side to reduce shuffle data.") .version("3.3.0") .booleanConf .createWithDefault(true) @@ -810,11 +896,22 @@ object SQLConf { val RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD = buildConf("spark.sql.optimizer.runtime.bloomFilter.creationSideThreshold") .doc("Size threshold of the bloom filter creation side plan. Estimated size needs to be " + - "under this value to try to inject bloom filter.") + "under this value to try to inject a bloom filter, unless the creation side is fully " + + "materialized and has accurate statistics.") .version("3.3.0") .bytesConf(ByteUnit.BYTE) .createWithDefaultString("10MB") + val RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD = + buildConf("spark.sql.optimizer.runtime.bloomFilter.materializedCreationSideThreshold") + .doc("Size threshold of a fully materialized, repeatable bloom filter creation side with " + + "accurate statistics. This replaces the general creation-side threshold because scanning " + + "materialized output does not recompute its original plan.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .bytesConf(ByteUnit.BYTE) + .createWithDefaultString("100MB") + val RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD = buildConf("spark.sql.optimizer.runtime.bloomFilter.applicationSideScanSizeThreshold") .doc("Byte size threshold of the Bloom filter application side plan's aggregated scan " + @@ -878,14 +975,23 @@ object SQLConf { val PUSH_DOWN_JOIN_THROUGH_UNION_ENABLED = buildConf("spark.sql.optimizer.pushDownJoinThroughUnion.enabled") - .doc("When true, pushes down Join through Union when the right side is small enough " + - "to broadcast. This can improve performance by allowing each Union branch to " + - "directly perform a broadcast join, avoiding materializing the entire Union result.") + .doc("When true, pushes down Join through Union when every Union branch would broadcast " + + "the right side of the join. This can improve performance by allowing each Union branch " + + "to directly perform a broadcast join, avoiding materializing the entire Union result.") .version("4.2.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf .createWithDefault(false) + val COMBINE_APPROXIMATE_PERCENTILES_ENABLED = + buildConf("spark.sql.optimizer.combineApproximatePercentiles.enabled") + .doc("When true, combines compatible scalar approximate percentile aggregates into a " + + "single array-valued aggregate so they share one percentile digest.") + .version("5.0.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + val EXPRESSION_PROJECTION_CANDIDATE_LIMIT = buildConf("spark.sql.optimizer.expressionProjectionCandidateLimit") .doc("The maximum number of the candidate of output expressions whose alias are replaced." + @@ -970,6 +1076,18 @@ object SQLConf { .booleanConf .createWithDefault(true) + val SPLIT_STREAMED_SIDE_JOIN_CONDITION = + buildConf("spark.sql.join.splitStreamedSideJoinCondition") + .internal() + .doc("When true, split join conditions for LeftAnti, LeftOuter, RightOuter, and " + + "ExistenceJoin by referenced side, evaluating streamed-side-only conjuncts before " + + "the hash-bucket or merge walk. This avoids evaluating expensive streamed-side-only " + + "predicates once per matched buffered row.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + val SORT_MERGE_AS_OF_JOIN_ENABLED = buildConf("spark.sql.join.sortMergeAsOfJoin.enabled") .doc("When true, use a dedicated sort-merge physical operator for AS-OF joins " + @@ -996,7 +1114,11 @@ object SQLConf { .doc("When true, the planner requires all the clustering keys as the hash partition keys " + "of the children, to eliminate the shuffles for the operator that needs its children to " + "be co-partitioned, such as JOIN node. This is to avoid data skews which can lead to " + - "significant performance regression if shuffles are eliminated.") + "significant performance regression if shuffles are eliminated. For storage-partitioned " + + "join, every clustering key must be covered by some partition key, rather than matching " + + "the partition keys positionally, so a column partitioned by more than one transform " + + "does not prevent shuffle elimination; hash partitioning deliberately keeps the " + + "positional match.") .version("3.3.0") .booleanConf .createWithDefault(true) @@ -1579,6 +1701,21 @@ object SQLConf { .booleanConf .createWithDefault(true) + val REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED = + buildConf("spark.sql.optimizer.rewriteCountDistinctConditional.enabled") + .internal() + .doc("When true, rewrites COUNT(DISTINCT IF(cond, base, NULL)) and " + + "COUNT(DISTINCT CASE WHEN cond THEN base END) into " + + "COUNT(DISTINCT base) FILTER (WHERE cond). This reduces the Expand factor " + + "in RewriteDistinctAggregates from Nx to 1x when multiple conditional distinct " + + "counts share the same base column. The rewrite is only applied to base " + + "expressions that are safe to evaluate unconditionally (e.g. plain columns), " + + "so the short-circuit semantics of IF/CASE WHEN are preserved.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + val ESCAPED_STRING_LITERALS = buildConf("spark.sql.parser.escapedStringLiterals") .internal() .doc("When true, string literals (including regex patterns) remain escaped in our SQL " + @@ -2823,6 +2960,33 @@ object SQLConf { "disable logging or -1 to apply no limit.") .createWithDefault(1000) + val CODEGEN_COMPILER = buildConf("spark.sql.codegen.compiler") + .internal() + .doc("The compiler used to turn generated Java source into bytecode. " + + "Supported values are 'janino' (default) and 'jdk'. " + + "'janino' uses the Janino library; it is several times faster but the project " + + "is unmaintained upstream (last release 3.1.12, Feb 2024). " + + "'jdk' uses javax.tools.JavaCompiler from the JDK; it is maintained on the " + + "JDK release cadence but adds ~5x cold-start latency for large generated units " + + "and 30-300x for small ones. Switch to 'jdk' only if Janino's maintenance status " + + "is a concern; for most workloads the default remains the better trade-off. " + + "When 'jdk' is requested but javax.tools.JavaCompiler is unavailable " + + "(e.g. JRE-only image) Spark falls back to 'janino' with a warning. " + + "Regardless of this setting, codegen in REPL / interactive sessions (spark-shell, " + + "Spark Connect session artifacts), generated code referencing a class nested in " + + "a Scala package object, and generated code referencing an anonymous or local class " + + "that Spark determines cannot be soundly rewritten to a nameable supertype (that " + + "supertype, or a class enclosing it, is not public, or it does not offer a member the " + + "class exposes) always compile with 'janino'. A one-time INFO log records each such " + + "routing.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .stringConf + .transform(_.toLowerCase(Locale.ROOT)) + .checkValues(Set("janino", "jdk")) + .createWithDefault( + sys.env.get("SPARK_CODEGEN_COMPILER").filter(_.nonEmpty).getOrElse("janino")) + val WHOLESTAGE_HUGE_METHOD_LIMIT = buildConf("spark.sql.codegen.hugeMethodLimit") .internal() .doc("The maximum bytecode size of a single compiled Java function generated by whole-stage " + @@ -3000,6 +3164,17 @@ object SQLConf { .booleanConf .createWithDefault(true) + val INSERT_MAP_SORT_IN_DISTINCT_AGGREGATES_ENABLED = + buildConf("spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled") + .internal() + .doc("When true, map-typed arguments of distinct aggregates are normalized with MapSort. " + + "When false, MapSort is not added solely for distinct aggregate arguments; arguments " + + "that are also grouping expressions remain normalized.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val OPTIMIZE_EXPAND_RATIO = buildConf("spark.sql.optimizer.optimizeExpandRatio") .internal() @@ -3100,13 +3275,28 @@ object SQLConf { val NUM_STATE_STORE_MAINTENANCE_THREADS = buildConf("spark.sql.streaming.stateStore.numStateStoreMaintenanceThreads") .internal() - .doc("Number of threads in the thread pool that perform clean up and snapshotting tasks " + - "for stateful streaming queries. The default value is the number of cores * 0.25 " + - "so that this thread pool doesn't take too many resources " + - "away from the query and affect performance.") + .doc("Total number of threads split between the snapshot and cleanup " + + "maintenance pools for stateful streaming queries. Each pool needs at least " + + "1 thread, so the minimum is 2. The default value is the number of " + + "cores * 0.25 so that the pools don't take too many resources away from the " + + "query and affect performance. Use snapshotToCleanupThreadRatio to " + + "configure the split between snapshot and cleanup pools.") .intConf - .checkValue(_ > 0, "Must be greater than 0") - .createWithDefault(Math.max(Runtime.getRuntime.availableProcessors() / 4, 1)) + .checkValue(_ > 1, "Must be greater than 1") + .createWithDefault(Math.max(Runtime.getRuntime.availableProcessors() / 4, 2)) + + val STATE_STORE_MAINTENANCE_SNAPSHOT_THREAD_RATIO = + buildConf("spark.sql.streaming.stateStore.snapshotToCleanupThreadRatio") + .internal() + .version("5.0.0") + .doc("Ratio of total maintenance threads allocated to the snapshot " + + "pool. The remainder goes to the cleanup pool. The snapshot " + + "count is rounded to the nearest integer and clamped so each " + + "pool gets at least 1 thread and the total is never exceeded.") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .doubleConf + .checkValue(v => v > 0 && v < 1, "Must be between 0 and 1 (exclusive)") + .createWithDefault(0.5) val STATE_STORE_MAINTENANCE_SHUTDOWN_TIMEOUT = buildConf("spark.sql.streaming.stateStore.maintenanceShutdownTimeout") @@ -3435,6 +3625,20 @@ object SQLConf { .checkValue(v => Set(1, 2).contains(v), "Valid versions are 1 and 2") .createWithDefault(1) + val STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1 = + buildConf("spark.sql.streaming.realTimeMode.dangerouslyAllowCheckpointV1.enabled") + .internal() + .doc("Whether to allow a Real-Time Mode query to start with state store checkpoint format " + + "version 1. Real-Time Mode re-executes a failed batch, and with checkpoint format " + + "version 1 the re-execution can reuse the state file names of the partially-written " + + "failed batch, so starting on a version 1 checkpoint exposes the query to data loss on " + + "failure. Format version 2 avoids this with per-batch state store checkpoint ids. " + + "Escape hatch only; prefer a fresh checkpoint location.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + val STREAMING_MAX_NUM_STATE_SCHEMA_FILES = buildConf("spark.sql.streaming.stateStore.maxNumStateSchemaFiles") .internal() @@ -3597,6 +3801,37 @@ object SQLConf { .checkValue(v => Set(1, 2).contains(v), "Valid versions are 1 and 2") .createWithDefault(2) + val STREAMING_USE_STREAMLINE_AGGREGATOR = + buildConf("spark.sql.streaming.useStreamlineAggregator") + .internal() + .doc("Test/development only, not intended for production use. When true, plan a streaming " + + "aggregation with the streamline aggregation operator, which merges each input row " + + "against state and emits immediately, instead of the microbatch operators that only emit " + + "once the batch ends. Real-Time Mode queries use the streamline operator regardless of " + + "this config; this flag exists only so the operator can be exercised under an ordinary " + + "microbatch trigger in tests, and changes an aggregation's output timing when set.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + + val STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR = + buildConf("spark.sql.streaming.statefulOperator.incrementalCleanupFactor") + .internal() + .doc("For a stateful operator that evicts by watermark, the number of eviction-eligible " + + "records to remove per input record processed, spreading eviction cost across the batch " + + "instead of paying it all at batch end. When 0, incremental cleanup is disabled and all " + + "eviction happens at batch end. When k, up to k eligible records are removed per input; " + + "any still-eligible records left over are removed at batch end. Only applies to modes " + + "that evict (e.g. Append/Update); has no effect in Complete mode, which never evicts. " + + "Read by the streamline aggregation and streaming deduplication (dropDuplicates) " + + "operators; dropDuplicatesWithinWatermark always evicts at batch end.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .longConf + .checkValue(_ >= 0, "The incremental cleanup factor must not be negative.") + .createWithDefault(0L) + val STREAMING_STOP_ACTIVE_RUN_ON_RESTART = buildConf("spark.sql.streaming.stopActiveRunOnRestart") .doc("Running multiple runs of the same streaming query concurrently is not supported. " + @@ -3822,6 +4057,16 @@ object SQLConf { .booleanConf .createWithDefault(true) + val STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS = buildConf( + "spark.sql.streaming.realTimeMode.transformWithState.ttlEvictionIntervalMs") + .internal() + .doc("The threshold in milliseconds to perform eviction of TTL when using the JVM " + + "transformWithState operator with real-time mode.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .longConf + .createWithDefault(1 * 1000) + val STREAMING_ASYNC_PROGRESS_TRACKING_REAL_TIME_MODE_ENABLED_BY_DEFAULT = buildConf( "spark.sql.streaming.realTimeMode.asyncProgressTrackingByDefault.enabled") .internal() @@ -4104,6 +4349,47 @@ object SQLConf { .booleanConf .createWithDefault(false) + val ADAPTIVE_PARTIAL_AGGREGATION_ENABLED = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.enabled") + .doc("When true, hash aggregation adaptively bypasses the pre-shuffle partial aggregation " + + "at runtime when it observes that the partial aggregation is not reducing the number of " + + "rows enough to be worthwhile. Once bypassed, the remaining input rows are passed " + + "through as single-row partial aggregation buffers for the final aggregation to merge, " + + "which avoids the cost of maintaining and spilling a large aggregation map with little " + + "reduction benefit. Disabled by default. This applies only to hash aggregation with " + + "grouping keys.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + + val ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minRows") + .doc("The number of rows between periodic compaction-ratio evaluations by adaptive partial " + + s"aggregation (see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). Setting this to 0 " + + "disables the periodic evaluation. The ratio may still be evaluated when the aggregation " + + "map is about to spill.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .longConf + .checkValue(_ >= 0, "The minimum row count must not be negative.") + .createWithDefault(100000) + + val ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION = + buildConf("spark.sql.execution.aggregate.adaptivePartialAggregation.minCompaction") + .doc("The minimum compaction ratio required to keep the pre-shuffle partial aggregation " + + s"(see '${ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key}'). The compaction ratio is the " + + "number of processed rows divided by the number of keys held in the aggregation maps, " + + "so a ratio of 10 means the partial aggregation collapses ten rows into one. When an " + + "evaluation finds the ratio below this value, the partial aggregation is bypassed for " + + "the rest of the input. A larger value bypasses more aggressively.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .doubleConf + .checkValue(v => v >= 1.0 && v.isFinite, + "The minimum compaction ratio must be a finite value of at least 1.0.") + .createWithDefault(1.05) + val JSON_GENERATOR_IGNORE_NULL_FIELDS = buildConf("spark.sql.jsonGenerator.ignoreNullFields") .doc("Whether to ignore null fields when generating JSON objects in JSON data source and " + @@ -4527,6 +4813,17 @@ object SQLConf { "The threshold of window group limit must be -1, 0 or positive integer.") .createWithDefault(1000) + val COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD = + buildConf("spark.sql.optimizer.collapseWindowWithEmptyOrderSpecInChild") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .doc("When true, the optimizer collapses two adjacent windows with the same partition " + + "spec into one when the window with the empty order spec is the child (inner) window. " + + "This saves a WindowExec pass but can disable the WindowGroupLimit and the LocalLimit " + + "push-down optimizations for top-k queries.") + .version("4.4.0") + .booleanConf + .createWithDefault(false) + val WINDOW_SEGMENT_TREE_ENABLED = buildConf("spark.sql.window.segmentTree.enabled") .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) @@ -5171,6 +5468,23 @@ object SQLConf { " must be more than one.") .createOptional + val PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED = + buildConf("spark.sql.execution.pythonUDF.inHigherOrderFunction.enabled") + .doc("When true, a scalar Python UDF may be used inside the lambda of a higher-order " + + "function such as `transform` or `filter`. The UDF is not evaluated inside the lambda; " + + "it is applied to the whole array outside the lambda and the lambda reads the " + + "precomputed result. This uses an Arrow-based execution path, so PyArrow is required " + + "even for a UDF created with useArrow=False. Because the UDF is precomputed over the " + + "whole array, it runs once per element regardless of any short-circuiting (`exists`, " + + "`when`) in the lambda. In particular, an `array_sort` comparator that calls the UDF on " + + "both elements, `(a, b) -> udf(a, b)`, precomputes it over every pair, costing O(n^2) " + + "calls and memory for an array of length n; avoid it for large arrays. When false, such " + + "queries fail with UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF as before.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(true) + val PYTHON_UDF_ARROW_FALLBACK_ON_UDT = buildConf("spark.sql.execution.pythonUDF.arrow.legacy.fallbackOnUDT") .internal() @@ -5228,6 +5542,17 @@ object SQLConf { .booleanConf .createWithDefault(false) + val PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED = + buildConf("spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled") + .internal() + .doc("When true, mapInPandas and mapInArrow UDFs may return any iterable (e.g. a list) " + + "rather than a strict iterator, matching the behavior before 4.3.0. When false, the " + + "returned value must be an iterator, matching the declared Iterator[...] signatures.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val PYTHON_PLANNER_EXEC_MEMORY = buildConf("spark.sql.planner.pythonExecution.memory") .doc("Specifies the memory allocation for executing Python code in Spark driver, in MiB. " + @@ -5359,6 +5684,16 @@ object SQLConf { .booleanConf .createWithDefault(false) + val PARSE_SQL_ENABLED = + buildConf("spark.sql.function.parseSql.enabled") + .doc("When true, enables the parse_sql SQL function. This feature is under active " + + "development; the JSON contract may change across releases while the flag remains " + + "off by default.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val ELT_OUTPUT_AS_STRING = buildConf("spark.sql.function.eltOutputAsString") .doc("When this option is set to false and all inputs are binary, `elt` returns " + "an output as binary. Otherwise, it returns as a string.") @@ -5844,6 +6179,19 @@ object SQLConf { .booleanConf .createWithDefault(true) + val DECORRELATE_LIMIT_OFFSET_LEGACY_INCORRECT_ORDER_HANDLING_ENABLED = + buildConf("spark.sql.optimizer.decorrelateLimitOffsetLegacyIncorrectOrderHandling.enabled") + .internal() + .doc("If enabled, revert to the legacy incorrect behavior where the ORDER BY of a " + + "correlated subquery with LIMIT (and optional OFFSET) whose predicates only reference " + + "the outer table is dropped during decorrelation, turning ORDER BY ... LIMIT ... OFFSET " + + "into an arbitrary LIMIT/OFFSET. When disabled (default), the ORDER BY is preserved so " + + "the result is deterministic.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val DECORRELATE_UNION_OR_SET_OP_UNDER_LIMIT_ENABLED = buildConf("spark.sql.optimizer.decorrelateUnionOrSetOpUnderLimit.enabled") .internal() @@ -6612,6 +6960,18 @@ object SQLConf { .booleanConf .createWithDefault(false) + val PYTHON_LIMIT_PUSHDOWN_ENABLED = buildConf("spark.sql.python.limitPushdown.enabled") + .internal() + .doc("When true, enable limit pushdown to Python datasource. Pushing a limit runs a Python " + + "worker during planning; for a limit-only scan this replaces the worker that plans a " + + "plain read, while for a scan that also pushes down filters it runs in addition to " + + "filter pushdown. Spark always applies the limit again after the scan, so a pushed limit " + + "only lets the data source read less data.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + val CSV_FILTER_PUSHDOWN_ENABLED = buildConf("spark.sql.csv.filterPushdown.enabled") .doc("When true, enable filter pushdown to CSV datasource.") .version("3.0.0") @@ -6717,7 +7077,8 @@ object SQLConf { "per-row cast-error companion column (nullable string) so that the cast error is only " + "raised when the row is consumed by the user expression. Without this flag, the cast is " + "always evaluated and any failure raises immediately, even when the surrounding " + - "expression would not have consumed the failing row.") + "expression would not have consumed the failing row. This also allows throwable variant " + + "extractions to be hoisted across joins while preserving their original error timing.") .version("4.3.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) .booleanConf @@ -6764,6 +7125,33 @@ object SQLConf { .booleanConf .createWithDefault(true) + val VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED = + buildConf("spark.sql.variant.shreddedPredicatePushdown.enabled") + .internal() + .doc("When true, comparison predicates on shredded Variant fields produced by " + + "PushVariantIntoScan (e.g. variant_get(v, '$.a', 'bigint') > 999) are pushed to Parquet " + + "as a predicate on the physical shredded typed_value leaf column, guarded so that a row " + + "group is skipped only when the leaf min/max cannot match AND every value for the path " + + "is provably in the typed leaf (either every untyped residual value column along the " + + "path is entirely null, or the leaf column itself has no nulls). This enables row-group " + + "skipping for shredded Variant columns while never dropping rows that fall back to an " + + "untyped residual. The benefit depends on the data layout, like any Parquet min/max " + + "skipping: it helps most when the data is sorted on the filtered field (so each row " + + "group covers a narrow value range) and a file holds many row groups; unsorted data or " + + "a single row group per file gains little. Has no effect unless the Parquet column is " + + "shredded and spark.sql.variant.pushVariantIntoScan is also true. It also does not fire " + + "for a strict cast to a non-string type when " + + "spark.sql.variant.pushVariantIntoScan.deferCastError is true (the extraction is wrapped " + + "in UnwrapVariantCastError and is not translated to a pushable filter); try_variant_get " + + "and string targets are unaffected. Results are unaffected either way; this only " + + "controls whether row groups can be skipped.") + .version("4.4.0") + // Physical scan optimization only: it changes which Parquet row groups are read, not the + // resolved plan of a view/UDF/procedure body, so it does not participate in binding. + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(true) + val LEGACY_CSV_ENABLE_DATE_TIME_PARSING_FALLBACK = buildConf("spark.sql.legacy.csv.enableDateTimeParsingFallback") .internal() @@ -7032,13 +7420,29 @@ object SQLConf { .createWithDefault(false) val PRESERVE_CHAR_VARCHAR_TYPE_INFO = buildConf("spark.sql.preserveCharVarcharTypeInfo") - .doc("When true, Spark does not replace CHAR/VARCHAR types the STRING type, which is the " + - "default behavior of Spark 3.0 and earlier versions. This means the length checks for " + - "CHAR/VARCHAR types is enforced and CHAR type is also properly padded.") + .doc("When true, Spark does not replace CHAR/VARCHAR with STRING in schemas and plans. " + + "This is the Spark 4.0 experimental path: types can leak through transforming string " + + "functions via child.dataType. Prefer spark.sql.charVarchar.standardSemantics.enabled " + + "for SQL standard CHAR/VARCHAR behavior (CAST/LCT/STRING-returning transforms).") .version("4.0.0") .booleanConf .createWithDefault(false) + val CHAR_VARCHAR_STANDARD_SEMANTICS = + buildConf("spark.sql.charVarchar.standardSemantics.enabled") + .doc("When true, enable SQL standard CHAR/VARCHAR semantics: first-class types in " + + "schemas and CAST targets; least-common-type for COALESCE/CASE/UNION may return " + + "CHAR/VARCHAR; transforming string functions and operators return plain STRING. " + + "This is a breaking change from the annotated-STRING default and from " + + "preserveCharVarcharTypeInfo (which keeps Char/Varchar through transforms).") + .version("4.4.0") + // PERSISTED, like ANSI mode: the flag decides the types a view body resolves to, so a view + // created under standard semantics must keep computing CHAR/VARCHAR regardless of the + // caller's session setting. + .withBindingPolicy(ConfigBindingPolicy.PERSISTED) + .booleanConf + .createWithDefault(false) + val READ_FILE_SOURCE_TABLE_CACHE_IGNORE_OPTIONS = buildConf("spark.sql.legacy.readFileSourceTableCacheIgnoreOptions") .internal() @@ -7052,7 +7456,9 @@ object SQLConf { val READ_SIDE_CHAR_PADDING = buildConf("spark.sql.readSideCharPadding") .doc("When true, Spark applies string padding when reading CHAR type columns/fields, " + "in addition to the write-side padding. This config is true by default to better enforce " + - "CHAR type semantic in cases such as external tables.") + "CHAR type semantic in cases such as external tables. When " + + s"'${CHAR_VARCHAR_STANDARD_SEMANTICS.key}' is true, this config is ignored: read-side " + + "CHAR/VARCHAR checks are always applied, and setting it to false logs a warning.") .version("3.4.0") .booleanConf .createWithDefault(true) @@ -7263,6 +7669,51 @@ object SQLConf { .booleanConf .createWithDefault(false) + val MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED = buildConf( + "spark.sql.optimizer.mergeSubplans.filterPropagation.dsv2SymmetricFilterPropagation.enabled") + .doc("When true, two DataSource V2 scan subplans that pushed the same strict filters but " + + "carry different best-effort (post-scan) filters can merge into one scan even when " + + s"${MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key} is false. The strict filters " + + "are equal (re-enforced on the rebuilt scan) and the differing filters are pushed only as " + + "best-effort row-group/partition pruning (OR-widened), with the enclosing Filter " + + "re-checking exactness above the scan. Unlike the general symmetric case, OR-widening " + + "cannot change the strict (enforced) row set. Has no effect when " + + s"${MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key} is false.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + + val MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION = buildConf( + "spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.keyGroupedPartitioningDegradation.enabled") + .doc("When false, a DataSource V2 scan merge is declined if the rebuilt merged scan would " + + "report weaker key-grouped partitioning than an input reported (no longer clustering by " + + "the same expressions), which can force a shuffle the original plan avoided. The merge is " + + "also declined, before any rebuild, when the two input scans report different clustering " + + "expressions, since no single merged scan could preserve both. When true, the merge " + + "proceeds anyway in both cases, trading the partitioning for a single scan. Reported " + + "partitioning is re-derived on the merged scan, so a merge that does not weaken it is " + + "always allowed. Only affects sources that declare the SCAN_MERGING table capability.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + + val MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION = buildConf( + "spark.sql.optimizer.mergeSubplans.dsv2ScanMerge.orderingDegradation.enabled") + .doc("When false, a DataSource V2 scan merge is declined if the rebuilt merged scan would " + + "report a weaker output ordering than an input reported (an input ordering is no longer a " + + "prefix of the merged scan's), which can force a sort the original plan avoided. The merge " + + "is also declined, before any rebuild, when neither input's reported ordering satisfies " + + "the other, since no single merged scan could preserve both. When true, the merge proceeds " + + "anyway in both cases, trading the ordering for a single scan. Reported ordering is " + + "re-derived on the merged scan, so a merge that does not weaken it is always allowed. Only " + + "affects sources that declare the SCAN_MERGING table capability.") + .version("4.4.0") + .withBindingPolicy(ConfigBindingPolicy.SESSION) + .booleanConf + .createWithDefault(false) + val MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED = buildConf("spark.sql.optimizer.mergeSubplans.filterPropagation.throughJoin.enabled") .doc("When set to true, filter attributes can propagate through Join nodes during subplan " + @@ -7538,23 +7989,6 @@ object SQLConf { .booleanConf .createWithDefault(true) - // Deprecate "spark.connect.copyFromLocalToFs.allowDestLocal" in favor of this config. This is - // currently optional because we don't want to break existing users who are using the old config. - // If this config is set, then we override the deprecated config. - val ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL = - buildConf("spark.sql.artifact.copyFromLocalToFs.allowDestLocal") - .internal() - .doc(""" - |Allow `spark.copyFromLocalToFs` destination to be local file system - | path on spark driver node when - |`spark.sql.artifact.copyFromLocalToFs.allowDestLocal` is true. - |This will allow user to overwrite arbitrary file on spark - |driver node we should only enable it for testing purpose. - |""".stripMargin) - .version("4.0.0") - .booleanConf - .createOptional - val LEGACY_RETAIN_FRACTION_DIGITS_FIRST = buildConf("spark.sql.legacy.decimal.retainFractionDigitsOnTruncate") .internal() @@ -8036,6 +8470,21 @@ object SQLConf { .booleanConf .createWithDefault(false) + val PROTOBUF_DESCRIPTOR_CACHE_SIZE = + buildConf("spark.sql.protobuf.descriptorCacheSize") + .internal() + .doc("Maximum number of entries in the per-JVM cache of parsed Protobuf FileDescriptor " + + "graphs built from a binary FileDescriptorSet, used by from_protobuf/to_protobuf. This " + + "knob only bounds memory (each graph can be hundreds of KB). Setting it to 0 disables " + + "the cache entirely and acts as an immediate kill-switch; a non-zero size bound is fixed " + + "when the cache is first built for a JVM, so changing it has no effect on an " + + "already-built cache.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .intConf + .checkValue(_ >= 0, "descriptorCacheSize must be non-negative; 0 disables the cache") + .createWithDefault(8) + val LISTAGG_ALLOW_DISTINCT_CAST_WITH_ORDER = buildConf("spark.sql.listagg.allowDistinctCastWithOrder.enabled") .internal() @@ -8094,7 +8543,7 @@ object SQLConf { DeprecatedConfig(ESCAPED_STRING_LITERALS.key, "4.0", "Use raw string literals with the `r` prefix instead. "), DeprecatedConfig("spark.connect.copyFromLocalToFs.allowDestLocal", "4.0", - s"Use '${ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL.key}' instead."), + "Use 'spark.sql.artifact.copyFromLocalToFs.allowDestLocal' instead."), DeprecatedConfig(ALLOW_ZERO_INDEX_IN_FORMAT_STRING.key, "4.0", "Increase indexes by 1 " + "in `strfmt` of the `format_string` function. Refer to the first argument by \"1$\"."), DeprecatedConfig(SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key, "4.1", @@ -8235,12 +8684,27 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def dynamicPartitionPruningReuseBroadcastOnly: Boolean = getConf(DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY) + def dynamicPartitionPruningBroadcastProjectionEnabled: Boolean = + getConf(DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED) + + def dynamicPartitionPruningBroadcastProjectionMaxRows: Int = + getConf(DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_ROWS) + + def dynamicPartitionPruningBroadcastProjectionMaxBytes: Long = + getConf(DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_BYTES) + + def dynamicPartitionPruningBroadcastProjectionMaxSourceBytes: Long = + getConf(DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_SOURCE_BYTES) + def runtimeFilterBloomFilterEnabled: Boolean = getConf(RUNTIME_BLOOM_FILTER_ENABLED) def runtimeFilterCreationSideThreshold: Long = getConf(RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD) + def runtimeFilterMaterializedCreationSideThreshold: Long = + getConf(RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD) + def runtimeRowLevelOperationGroupFilterEnabled: Boolean = getConf(RUNTIME_ROW_LEVEL_OPERATION_GROUP_FILTER_ENABLED) @@ -8250,6 +8714,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def numStateStoreMaintenanceThreads: Int = getConf(NUM_STATE_STORE_MAINTENANCE_THREADS) + def snapshotToCleanupThreadRatio: Double = + getConf(STATE_STORE_MAINTENANCE_SNAPSHOT_THREAD_RATIO) + def numStateStoreInstanceMetricsToReport: Int = getConf(STATE_STORE_INSTANCE_METRICS_REPORT_LIMIT) @@ -8457,6 +8924,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def coalesceShufflePartitionsEnabled: Boolean = getConf(COALESCE_PARTITIONS_ENABLED) + def collapseWindowWithEmptyOrderSpecInChild: Boolean = + getConf(COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD) + def minBatchesToRetain: Int = getConf(MIN_BATCHES_TO_RETAIN) def maxVersionsToDeletePerMaintenance: Int = getConf(MAX_VERSIONS_TO_DELETE_PER_MAINTENANCE) @@ -8477,6 +8947,14 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def streamingOffsetLogFormatVersion: Int = getConf(STREAMING_OFFSET_LOG_FORMAT_VERSION) + // The commit log format version implied by the session config for a fresh checkpoint. A state + // store checkpoint format of v2 makes each batch write stateUniqueIds, which only a commit log at + // VERSION_2 or above can persist, so the commit log version tracks the state store format. It is + // capped at VERSION_2 here: VERSION_3 exists only to carry sink-evolution metadata and is written + // exclusively by the sink-evolution path in MicroBatchExecution, never derived from a config. + def streamingCommitLogFormatVersion: Int = + if (getConf(STATE_STORE_CHECKPOINT_FORMAT_VERSION) >= 2) 2 else 1 + def stateStoreEncodingFormat: String = getConf(STREAMING_STATE_STORE_ENCODING_FORMAT) def streamingValueStateSchemaEvolutionThreshold: Int = @@ -8557,6 +9035,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def codegenLogLevel: Level = getConf(CODEGEN_LOG_LEVEL) + def codegenCompiler: String = getConf(CODEGEN_COMPILER) + def loggingMaxLinesForCodegen: Int = getConf(CODEGEN_LOGGING_MAX_LINES) def hugeMethodLimit: Int = getConf(WHOLESTAGE_HUGE_METHOD_LIMIT) @@ -8699,6 +9179,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def preferSortMergeJoin: Boolean = getConf(PREFER_SORTMERGEJOIN) + def splitStreamedSideJoinCondition: Boolean = getConf(SPLIT_STREAMED_SIDE_JOIN_CONDITION) + def sortMergeAsOfJoinEnabled: Boolean = getConf(SORT_MERGE_AS_OF_JOIN_ENABLED) def sqlAsOfJoinEnabled: Boolean = getConf(SQL_ASOF_JOIN_ENABLED) @@ -8816,6 +9298,15 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def bypassPartialAggregation: Boolean = getConf(BYPASS_PARTIAL_AGGREGATION) + def adaptivePartialAggregationEnabled: Boolean = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_ENABLED) + + def adaptivePartialAggregationMinRows: Long = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS) + + def adaptivePartialAggregationMinCompaction: Double = + getConf(ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION) + def objectAggSortBasedFallbackThreshold: Int = getConf(OBJECT_AGG_SORT_BASED_FALLBACK_THRESHOLD) def variableSubstituteEnabled: Boolean = getConf(VARIABLE_SUBSTITUTE_ENABLED) @@ -8992,6 +9483,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pythonUDFArrowFallbackOnUDT: Boolean = getConf(PYTHON_UDF_ARROW_FALLBACK_ON_UDT) + def pythonUDFInHigherOrderFunctionEnabled: Boolean = + getConf(PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED) + def pysparkPlotMaxRows: Int = getConf(PYSPARK_PLOT_MAX_ROWS) def arrowSparkREnabled: Boolean = getConf(ARROW_SPARKR_EXECUTION_ENABLED) @@ -9046,6 +9540,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def legacyPandasConversionUDF: Boolean = getConf(PYTHON_UDF_LEGACY_PANDAS_CONVERSION_ENABLED) + def legacyMapInBatchAcceptAnyIterable: Boolean = + getConf(PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED) + def pythonPlannerExecMemory: Option[Long] = getConf(PYTHON_PLANNER_EXEC_MEMORY) def replaceExceptWithFilter: Boolean = getConf(REPLACE_EXCEPT_WITH_FILTER) @@ -9071,6 +9568,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def concatBinaryAsString: Boolean = getConf(CONCAT_BINARY_AS_STRING) + def parseSqlEnabled: Boolean = getConf(PARSE_SQL_ENABLED) + def eltOutputAsString: Boolean = getConf(ELT_OUTPUT_AS_STRING) def validatePartitionColumns: Boolean = getConf(VALIDATE_PARTITION_COLUMNS) @@ -9182,6 +9681,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def pythonFilterPushDown: Boolean = getConf(PYTHON_FILTER_PUSHDOWN_ENABLED) + def pythonLimitPushDown: Boolean = getConf(PYTHON_LIMIT_PUSHDOWN_ENABLED) + def csvFilterPushDown: Boolean = getConf(CSV_FILTER_PUSHDOWN_ENABLED) def jsonFilterPushDown: Boolean = getConf(JSON_FILTER_PUSHDOWN_ENABLED) @@ -9225,6 +9726,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def preserveCharVarcharTypeInfo: Boolean = getConf(SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO) + def charVarcharStandardSemantics: Boolean = getConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS) + def avoidDoubleFilterEval: Boolean = getConf(AVOID_DOUBLE_FILTER_EVAL) def structPredicateDecomposeEnabled: Boolean = getConf(STRUCT_PREDICATE_DECOMPOSE_ENABLED) @@ -9242,6 +9745,9 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def decorrelateInnerQueryEnabledForExistsIn: Boolean = !getConf(SQLConf.DECORRELATE_EXISTS_IN_SUBQUERY_LEGACY_INCORRECT_COUNT_HANDLING_ENABLED) + def rewriteCountDistinctConditionalEnabled: Boolean = + getConf(SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED) + def maxConcurrentOutputFileWriters: Int = getConf(SQLConf.MAX_CONCURRENT_OUTPUT_FILE_WRITERS) def plannedWriteEnabled: Boolean = getConf(SQLConf.PLANNED_WRITE_ENABLED) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala index 72e36250b068a..8b2cfc8a8d029 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/StaticSQLConf.scala @@ -338,6 +338,18 @@ object StaticSQLConf { .booleanConf .createWithDefault(true) + val ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL = + buildStaticConf("spark.sql.artifact.copyFromLocalToFs.allowDestLocal") + .internal() + .doc("Allow the `copyFromLocalToFs` destination to be a local file system path on the " + + "driver node. This lets the caller overwrite arbitrary files on the driver node, so it " + + "should only be enabled for testing purposes. This is a static conf: it can only be set " + + "when starting the driver, and not from a session.") + .version("4.3.0") + .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE) + .booleanConf + .createWithDefault(false) + val REFLECT_ALLOW_LIST = buildStaticConf("spark.sql.reflect.allowList") .doc("A comma-separated allow list of regular expressions matched against the canonical " + "static method name (in the form `class.method`, e.g. `java.util.UUID.randomUUID`) " + diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ExpressionWithToString.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ExpressionWithToString.scala index 8dd7662ce0240..18c934797bf5c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ExpressionWithToString.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ExpressionWithToString.scala @@ -21,5 +21,5 @@ import org.apache.spark.sql.connector.expressions.Expression abstract class ExpressionWithToString extends Expression with Serializable { private val builder = new ToStringSQLBuilder() - override def toString(): String = builder.build(this); + override def toString(): String = builder.build(this) } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionOffsetWithIndex.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionOffsetWithIndex.scala index 9db9bd2ac1243..36377553bbe01 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionOffsetWithIndex.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionOffsetWithIndex.scala @@ -14,11 +14,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package org.apache.spark.sql.internal.connector; +package org.apache.spark.sql.internal.connector -import org.apache.spark.sql.connector.read.streaming.PartitionOffset; +import org.apache.spark.sql.connector.read.streaming.PartitionOffset /** * Internal class for real time mode to pass partition offset from executors to the driver. */ -private[sql] case class PartitionOffsetWithIndex(index: Long, partitionOffset: PartitionOffset); +private[sql] case class PartitionOffsetWithIndex(index: Long, partitionOffset: PartitionOffset) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala new file mode 100644 index 0000000000000..eaa4857f1abe3 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala @@ -0,0 +1,103 @@ +/* + * 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.internal.connector + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.read.Scan + +/** + * A mix-in interface for [[Scan]]. Data sources can implement this interface if they can + * filter initially planned [[org.apache.spark.sql.connector.read.InputPartition]]s using + * Catalyst [[Expression]]s Spark infers at runtime. + * A scan must not implement this interface together with + * [[org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering]] or its subinterface + * [[org.apache.spark.sql.connector.read.SupportsRuntimeFiltering]]; Spark rejects such a scan. + * + * Spark considers a runtime predicate fully pushed when all attributes referenced by the + * predicate are returned by [[fullyPushedFilterAttributes]]. Fully pushed predicates are not + * evaluated again after the scan. + * + * Note that Spark will push runtime filters only if they are beneficial. + */ +trait SupportsRuntimeCatalystFiltering extends Scan { + + /** + * Returns attributes this scan can be filtered by at runtime. + * + * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes. + * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested + * references are rejected, and attributes pruned out of the read schema fail to resolve, when + * Spark builds the scan relation. + */ + def filterAttributes(): Array[NamedReference] + + /** + * Returns attributes for which this scan fully evaluates runtime predicates. + * + * Any runtime predicate that references only attributes in this set is considered fully pushed + * and will not be evaluated again after the scan. These attributes must also be returned by + * [[filterAttributes]]. Each attribute's value must therefore be fixed within every + * [[org.apache.spark.sql.connector.read.InputPartition]] the scan returns, since pruning + * partitions cannot fully evaluate a predicate on a column that varies within a partition. + * + * Spark relies on the scan alone here, so the scan must return only partitions it has proven + * satisfy such a predicate. Spark passes these expressions through as they are, leaving both + * translation and capability checking to the scan, so evaluating one can fail, for example on + * an ANSI cast or overflow error, or on a nested access that the scan matches differently + * against its partition layout. Declare an attribute only when the scan can evaluate every + * predicate over it. + * + * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested + * references are rejected, and attributes pruned out of the read schema fail to resolve, when + * Spark builds the scan relation. + */ + def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty + + /** + * Filters this scan using runtime Catalyst expressions. + * + * The provided expressions must be interpreted as a set of predicates that are ANDed together. + * Implementations may use the expressions to prune initially planned + * [[org.apache.spark.sql.connector.read.InputPartition]]s. + * + * An expression may access nested fields of an attribute returned by [[filterAttributes]], as + * that attribute is required to be top-level. The scan is responsible for matching such + * accesses against its own partition layout. + * + * Spark may call this method more than once for the same scan instance: a plan can hold several + * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes + * its own copy of the runtime filters. Implementations must treat successive calls as additive, + * ANDing the new expressions with those already pushed rather than replacing them. + * + * If the scan also implements + * [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve + * the originally reported partitioning during runtime filtering. While applying runtime + * predicates, the scan may detect that some + * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which + * case it can either replace the initially planned + * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with + * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of + * the original partition values (omitting those with no data) via + * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not + * report new partition values that were not present in the original partitioning. + * + * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime. + */ + def filter(expressions: Array[Expression]): Unit +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala new file mode 100644 index 0000000000000..805e87facff08 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/V2StatisticsUtils.scala @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import java.util.OptionalLong + +import org.apache.spark.sql.connector.read.{Scan, Statistics, SupportsReportStatistics} + +object V2StatisticsUtils { + + def isNotEmpty(stats: Statistics): Boolean = { + stats != null && hasAnyValue(stats) + } + + private def hasAnyValue(stats: Statistics): Boolean = { + stats.sizeInBytes.isPresent || + stats.numRows.isPresent || + (stats.columnStats != null && !stats.columnStats.isEmpty) + } + + def computeStats(scan: Scan): Option[Statistics] = scan match { + case s: SupportsReportStatistics => Some(s.estimateStatistics()).filter(isNotEmpty) + case _ => None + } + + def computeSizeInBytes( + scan: Scan, + avgRowSize: => BigInt): Option[BigInt] = scan match { + case s: SupportsReportStatistics => + // Prefer a cheap size estimate; otherwise read the full statistics once and derive the size + // from sizeInBytes, falling back to numRows * avgRowSize. + toBigInt(s.estimateSizeInBytes()).orElse { + Option(s.estimateStatistics()).flatMap { stats => + toBigInt(stats.sizeInBytes).orElse(toBigInt(stats.numRows).map(_ * avgRowSize)) + } + } + case _ => None + } + + private def toBigInt(value: OptionalLong): Option[BigInt] = { + if (value.isPresent) Some(BigInt(value.getAsLong)) else None + } +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataTypeExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataTypeExpression.scala index 6cb6b7d429d85..329af882391d3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataTypeExpression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/types/DataTypeExpression.scala @@ -87,6 +87,11 @@ private[sql] object AnyTimestampNanoTypeExpression { e.dataType.isInstanceOf[AnyTimestampNanoType] } +private[sql] object TimestampLTZNanosTypeExpression { + def unapply(e: Expression): Boolean = + e.dataType.isInstanceOf[TimestampLTZNanosType] +} + private[sql] object DecimalExpression { def unapply(e: Expression): Option[(Int, Int)] = e.dataType match { case t: DecimalType => Some((t.precision, t.scale)) diff --git a/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/codegen/StaticClashBase.java b/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/codegen/StaticClashBase.java new file mode 100644 index 0000000000000..661bc374d834e --- /dev/null +++ b/sql/catalyst/src/test/java/org/apache/spark/sql/catalyst/expressions/codegen/StaticClashBase.java @@ -0,0 +1,31 @@ +/* + * 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.catalyst.expressions.codegen; + +/** + * A narrowing-soundness fixture for {@code CodeCompilerSuite}: a public class declaring a + * public STATIC method. A Scala anonymous subclass can declare an INSTANCE method of the same + * erased signature, because scalac does not treat a Java static as an inherited member, while + * javac rejects the pair outright. That combination has to be written in Java to exist at all. + * + * Narrowing a reference to such a subclass is unsound: a static call is bound statically, so + * {@code ((StaticClashBase) ref).value()} reaches this method, not the subclass's. + */ +public class StaticClashBase { + public static int value() { return 1; } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala index cb5d77d445121..35f06a722f19f 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala @@ -424,6 +424,26 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { assert(!RangeShuffleSpec(10, distribution).canCreatePartitioning) } + test("SPARK-59022: canCreatePartitioning: KeyedShuffleSpec requires grouped partition keys") { + val a = $"a".int + val distribution = ClusteredDistribution(Seq(a)) + def keyedSpec(keys: Seq[Int]): KeyedShuffleSpec = KeyedShuffleSpec( + KeyedPartitioning(Seq(a), keys.map(k => InternalRow(k))), distribution) + + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + val grouped = keyedSpec(Seq(2, 1)) + assert(grouped.partitioning.isGrouped) + assert(grouped.canCreatePartitioning, + "unsorted keys are fine, the shuffle follows the declared order") + + // Duplicate keys mean one key spans several partitions, which a KeyGroupedPartitioner cannot + // reproduce, so this spec must not be the target for shuffling the other child. + val ungrouped = keyedSpec(Seq(1, 1, 2)) + assert(!ungrouped.partitioning.isGrouped) + assert(!ungrouped.canCreatePartitioning) + } + } + test("createPartitioning: HashShuffleSpec") { checkCreatePartitioning( HashShuffleSpec(HashPartitioning(Seq($"a"), 10), ClusteredDistribution(Seq($"a", $"b"))), diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala index 80d94cc3760df..c222655525eca 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnsiTypeCoercionSuite.scala @@ -1150,4 +1150,37 @@ class AnsiTypeCoercionSuite extends TypeCoercionSuiteBase { shouldNotCast(ArrayType(ArrayType(IntegerType)), AbstractArrayType(StringTypeWithCollation(supportsTrimCollation = true))) } + + test("SPARK-57811: ANSI string to nanosecond timestamp coercion in comparisons and predicates") { + // PromoteStrings has no standalone rule id (it only runs inside AnsiCombinedTypeCoercionRule in + // production), so wrap it the same way to get a runnable, id-registered rule for ruleTest. + val rule = + new AnsiTypeCoercion.AnsiCombinedTypeCoercionRule(Seq(AnsiTypeCoercion.PromoteStrings)) + // In ANSI mode the string operand is coerced to the nanosecond timestamp type for both the LTZ + // and NTZ families, via AnsiStringPromotionTypeCoercion.findWiderTypeForString (the atomic-type + // fall-through). This is config-blind: ANSI coercion never reads castDatetimeToString, so + // unlike the non-ANSI range path there is no legacy string-promotion branch. Both families + // behave exactly like their micros counterparts (TimestampType / TimestampNTZType) in ANSI + // mode, and the concrete operand type (family + precision) is preserved. Assert under both + // flag values to lock in that ANSI ignores it. + Seq("false", "true").foreach { legacy => + withSQLConf(SQLConf.LEGACY_CAST_DATETIME_TO_STRING.key -> legacy) { + Seq(7, 8, 9).foreach { p => + Seq(TimestampLTZNanosType(p), TimestampNTZNanosType(p)).foreach { nt => + val tsn = AttributeReference("tsn", nt)() + val strLit = Literal("2020-01-02 03:04:05.123456789") + // Equality (covers both 3VL EqualTo and null-safe EqualNullSafe). + ruleTest(rule, EqualTo(tsn, strLit), EqualTo(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualTo(strLit, tsn), EqualTo(Cast(strLit, nt), tsn)) + ruleTest(rule, EqualNullSafe(tsn, strLit), EqualNullSafe(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualNullSafe(strLit, tsn), EqualNullSafe(Cast(strLit, nt), tsn)) + // Range comparisons. + ruleTest(rule, LessThan(tsn, strLit), LessThan(tsn, Cast(strLit, nt))) + ruleTest(rule, GreaterThanOrEqual(strLit, tsn), + GreaterThanOrEqual(Cast(strLit, nt), tsn)) + } + } + } + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala new file mode 100644 index 0000000000000..f9295eb17b00c --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala @@ -0,0 +1,95 @@ +/* + * 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.catalyst.analysis + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat, Expression, Literal, PythonUDF, TranspiledPythonUDF} +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} +import org.apache.spark.sql.types.LongType + +/** + * Unit tests for [[ResolveTranspiledPythonUDFOptions]], which prunes a + * TranspiledPythonUDF's per-input-type options to those whose declared categories match the + * resolved argument types. func=null in the leaf PythonUDF is intentional: these structural + * tests don't execute Python. + */ +class ResolveTranspiledPythonUDFOptionsSuite extends PlanTest { + + private def pyUDF(children: Seq[Expression]): PythonUDF = + PythonUDF("udf", null, LongType, children, + PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true) + + // Runs the rule on a Project that wraps the node, and returns the (possibly pruned) node. + private def prune(node: TranspiledPythonUDF, rel: LocalRelation): TranspiledPythonUDF = { + val rewritten = ResolveTranspiledPythonUDFOptions(Project(Seq(Alias(node, "r")()), rel)) + rewritten.expressions.flatMap(_.collect { case t: TranspiledPythonUDF => t }).head + } + + test("keeps the numeric option for numeric columns and drops the string one") { + val a = $"a".long + val b = $"b".long + val numericOpt = Add(a, b) + val stringOpt = Concat(Seq(a, b)) + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a, b)), List(numericOpt, stringOpt), + List(List("numeric", "numeric"), List("string", "string"))) + val pruned = prune(node, LocalRelation(a, b)) + assert(pruned.transpiledOptions == List(numericOpt)) + assert(pruned.optionInputCategories.isEmpty) + } + + test("keeps the string option for string columns and drops the numeric one") { + val a = $"a".string + val b = $"b".string + val numericOpt = Add(a, b) + val stringOpt = Concat(Seq(a, b)) + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a, b)), List(numericOpt, stringOpt), + List(List("numeric", "numeric"), List("string", "string"))) + val pruned = prune(node, LocalRelation(a, b)) + assert(pruned.transpiledOptions == List(stringOpt)) + assert(pruned.optionInputCategories.isEmpty) + } + + test("empties the options when no category set matches (falls back to Python UDF)") { + val a = $"a".string + val b = $"b".long + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a, b)), + List(Add(a, b), Concat(Seq(a, b))), + List(List("numeric", "numeric"), List("string", "string"))) + val pruned = prune(node, LocalRelation(a, b)) + assert(pruned.transpiledOptions.isEmpty) + assert(pruned.optionInputCategories.isEmpty) + } + + test("matches binary columns against neither category (string is StringType only)") { + val a = $"a".binary + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)), + List(Concat(Seq(a, a))), List(List("string"))) + val pruned = prune(node, LocalRelation(a)) + assert(pruned.transpiledOptions.isEmpty) + } + + test("leaves options untouched when categories are empty (no restriction)") { + val a = $"a".long + val onlyOpt = Add(a, Literal(1L)) + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)), List(onlyOpt), Nil) + val pruned = prune(node, LocalRelation(a)) + assert(pruned.transpiledOptions == List(onlyOpt)) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala index 75846aa49616c..c078504ea33a6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala @@ -119,4 +119,21 @@ class TableLookupCacheSuite extends AnalysisTest with Matchers { verify(catalog, times(1)).getTable("default", "t1") } } + + test("nested view analysis shares both query-scoped caches") { + AnalysisContext.withNewAnalysisContext { + val outer = AnalysisContext.get + val viewDesc = CatalogTable( + TableIdentifier("view", Some("default")), + CatalogTableType.VIEW, + CatalogStorageFormat.empty, + StructType(Seq(StructField("a", IntegerType))), + viewText = Some("select * from t1")) + + AnalysisContext.withAnalysisContext(viewDesc) { + assert(AnalysisContext.get.relationCache eq outer.relationCache) + assert(AnalysisContext.get.tableCache eq outer.tableCache) + } + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionSuite.scala index d990a393ff653..577ff6de4f1cb 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TypeCoercionSuite.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.plans.ReferenceAllColumns import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.{Rule, RuleExecutor} import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.CalendarInterval @@ -353,6 +354,22 @@ abstract class TypeCoercionSuiteBase extends AnalysisTest { Concat(Seq(Literal("123".getBytes), Literal("456".getBytes))), Concat(Seq(Literal("123".getBytes), Literal("456".getBytes)))) } + + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charLit = Literal.create("ab", CharType(2)) + val collatedChar = Literal.create("ab", CharType(2, "UTF8_LCASE")) + val collatedString = StringType("UTF8_LCASE") + Seq(TypeCoercion.ConcatCoercion, AnsiTypeCoercion.ConcatCoercion).foreach { r => + ruleTest(r, + Concat(Seq(charLit, charLit)), + Concat(Seq(Cast(charLit, StringType), Cast(charLit, StringType)))) + ruleTest(r, + Concat(Seq(collatedChar, collatedChar)), + Concat(Seq( + Cast(collatedChar, collatedString), + Cast(collatedChar, collatedString)))) + } + } } test("type coercion for Elt") { @@ -407,6 +424,23 @@ abstract class TypeCoercionSuiteBase extends AnalysisTest { Elt(Seq(Literal(1), Literal("123".getBytes), Literal("456".getBytes))), Elt(Seq(Literal(1), Literal("123".getBytes), Literal("456".getBytes)))) } + + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charLit = Literal.create("ab", CharType(5)) + val collatedChar = Literal.create("ab", CharType(5, "UTF8_LCASE")) + val collatedString = StringType("UTF8_LCASE") + Seq(TypeCoercion.EltCoercion, AnsiTypeCoercion.EltCoercion).foreach { r => + ruleTest(r, + Elt(Seq(Literal(1), charLit, charLit)), + Elt(Seq(Literal(1), Cast(charLit, StringType), Cast(charLit, StringType)))) + ruleTest(r, + Elt(Seq(Literal(1), collatedChar, collatedChar)), + Elt(Seq( + Literal(1), + Cast(collatedChar, collatedString), + Cast(collatedChar, collatedString)))) + } + } } test("Datetime operations") { @@ -436,6 +470,67 @@ abstract class TypeCoercionSuiteBase extends AnalysisTest { ruleTest(rule, SubtractTimestamps(timestampNTZLiteral, timestampLiteral), SubtractTimestamps(timestampNTZLiteral, Cast(timestampLiteral, TimestampNTZType))) + + // SPARK-57832: subtraction accepts nanosecond-precision timestamps. A DATE operand takes the + // nanos type of the other side; a timestamp pair that differs only in precision or family is + // widened to a common type before being handed to SubtractTimestamps. + val ntzNanos9 = Literal.create( + DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2021-01-01T00:00:00"), precision = 9), + TimestampNTZNanosType(9)) + val ltzNanos9 = Literal.create( + DateTimeUtils.instantToTimestampNanos( + java.time.Instant.parse("2021-01-01T00:00:00Z"), precision = 9), + TimestampLTZNanosType(9)) + val ntzNanos7 = Literal.create( + DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2021-01-01T00:00:00"), precision = 7), + TimestampNTZNanosType(7)) + + // DATE - nanos and nanos - DATE cast the DATE side to the nanos type. Both the NTZ-nanos and + // the LTZ-nanos side are covered so the DATE-adopts-the-other-family arm is exercised in both + // time-zone families and both operand orders. + ruleTest(rule, + SubtractTimestamps(dateLiteral, ntzNanos9), + SubtractTimestamps(Cast(dateLiteral, TimestampNTZNanosType(9)), ntzNanos9)) + ruleTest(rule, + SubtractTimestamps(ntzNanos9, dateLiteral), + SubtractTimestamps(ntzNanos9, Cast(dateLiteral, TimestampNTZNanosType(9)))) + ruleTest(rule, + SubtractTimestamps(dateLiteral, ltzNanos9), + SubtractTimestamps(Cast(dateLiteral, TimestampLTZNanosType(9)), ltzNanos9)) + ruleTest(rule, + SubtractTimestamps(ltzNanos9, dateLiteral), + SubtractTimestamps(ltzNanos9, Cast(dateLiteral, TimestampLTZNanosType(9)))) + // Same-precision same-family pair is already the same type -> left untouched. + ruleTest(rule, + SubtractTimestamps(ntzNanos9, ntzNanos9), + SubtractTimestamps(ntzNanos9, ntzNanos9)) + // Cross-family nanos pair (LTZ vs NTZ) unifies in the NTZ family at the max precision. + ruleTest(rule, + SubtractTimestamps(ltzNanos9, ntzNanos7), + SubtractTimestamps( + Cast(ltzNanos9, TimestampNTZNanosType(9)), Cast(ntzNanos7, TimestampNTZNanosType(9)))) + // Micro NTZ vs nanos NTZ: same family, widen precision to the nanos type. + ruleTest(rule, + SubtractTimestamps(timestampNTZLiteral, ntzNanos9), + SubtractTimestamps(Cast(timestampNTZLiteral, TimestampNTZNanosType(9)), ntzNanos9)) + // Micro LTZ (TIMESTAMP) vs nanos LTZ: same LTZ family, widen precision to the nanos LTZ type + // so the subtraction still runs in the session time zone. + ruleTest(rule, + SubtractTimestamps(timestampLiteral, ltzNanos9), + SubtractTimestamps(Cast(timestampLiteral, TimestampLTZNanosType(9)), ltzNanos9)) + // Cross-family micro/nanos pairs: a micro operand on one side and a nanos operand of the other + // family on the other. Both unify in the NTZ family at the nanos precision (the cross-family + // rule prefers NTZ), so the micro operand widens across both axes at once. + ruleTest(rule, + SubtractTimestamps(timestampLiteral, ntzNanos9), + SubtractTimestamps(Cast(timestampLiteral, TimestampNTZNanosType(9)), ntzNanos9)) + ruleTest(rule, + SubtractTimestamps(timestampNTZLiteral, ltzNanos9), + SubtractTimestamps( + Cast(timestampNTZLiteral, TimestampNTZNanosType(9)), + Cast(ltzNanos9, TimestampNTZNanosType(9)))) } test("datetime comparison") { @@ -1021,6 +1116,34 @@ class TypeCoercionSuite extends TypeCoercionSuiteBase { ruleTest(TypeCoercion.ImplicitTypeCasts, NumericTypeUnaryExpression(Literal.create(null, NullType)), NumericTypeUnaryExpression(Literal.create(null, DoubleType))) + + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charLit = Literal.create("ab", CharType(2)) + ruleTest(TypeCoercion.ImplicitTypeCasts, + Upper(charLit), + Upper(Cast(charLit, StringType))) + } + } + + test("coerce JsonTuple children without the NullType rewrite") { + val json = Literal("""{"a":1}""") + val nullField = Literal.create(null, NullType) + val intField = Literal(1) + + // JsonTuple keeps its own NON_STRING_TYPE check, so these stay for checkInputDataTypes. + ruleTest(TypeCoercion.ImplicitTypeCasts, + JsonTuple(Seq(json, nullField)), + JsonTuple(Seq(json, nullField))) + ruleTest(TypeCoercion.ImplicitTypeCasts, + JsonTuple(Seq(json, intField)), + JsonTuple(Seq(json, intField))) + + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charLit = Literal.create("ab", CharType(2)) + ruleTest(TypeCoercion.ImplicitTypeCasts, + JsonTuple(Seq(charLit, charLit)), + JsonTuple(Seq(Cast(charLit, StringType), Cast(charLit, StringType)))) + } } test("cast NullType for binary operators") { @@ -1701,6 +1824,66 @@ class TypeCoercionSuite extends TypeCoercionSuiteBase { EqualTo(Cast(date0301, TimestampType), timestamp0301000000)) ruleTest(rule, LessThan(date0301, timestamp0301000001), LessThan(Cast(date0301, TimestampType), timestamp0301000001)) + + // SPARK-57811: a string operand is coerced to the nanosecond timestamp type in comparisons + // and predicates, mirroring the microsecond timestamp handling above. The concrete operand + // type (family + precision) is preserved, so the string is cast to that exact nanos type. + Seq(7, 8, 9).foreach { p => + Seq(TimestampLTZNanosType(p), TimestampNTZNanosType(p)).foreach { nt => + val tsn = AttributeReference("tsn", nt)() + val strLit = Literal("2020-01-02 03:04:05.123456789") + // Equality path: the string is cast to the nanos operand's own type so subsecond rounding + // does not affect the comparison. LTZ takes the explicit StringPromotionTypeCoercion + // Equality arm; NTZ has no arm and reaches the same cast via the general BinaryComparison + // fall-through (findCommonTypeForBinaryComparison returns the config-blind nanos type), + // mirroring how micros TimestampType (arm) vs TimestampNTZType (fall-through) are handled. + // The `Equality` extractor matches both 3VL `EqualTo` and null-safe `EqualNullSafe`, so + // both are covered. + ruleTest(rule, EqualTo(tsn, strLit), EqualTo(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualTo(strLit, tsn), EqualTo(Cast(strLit, nt), tsn)) + ruleTest(rule, EqualNullSafe(tsn, strLit), EqualNullSafe(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualNullSafe(strLit, tsn), EqualNullSafe(Cast(strLit, nt), tsn)) + // Range path (findCommonTypeForBinaryComparison). + ruleTest(rule, LessThan(tsn, strLit), LessThan(tsn, Cast(strLit, nt))) + ruleTest(rule, GreaterThanOrEqual(strLit, tsn), + GreaterThanOrEqual(Cast(strLit, nt), tsn)) + } + } + + // SPARK-57811: under legacy `castDatetimeToString`, the two nanos families mirror their micros + // counterparts. LTZ has a range arm in findCommonTypeForBinaryComparison, so its range + // comparisons promote both operands to string (like micros TimestampType); NTZ has no arm and + // stays config-blind (like micros TimestampNTZType), casting the string to the nanos type. In + // both families equality still casts the string to nanos, because the Equality arm fires before + // the range arm and reads no config. This block is the assertion that distinguishes the new + // production arms: in the default config the generic string-promotion fall-through already + // yields the nanos common type, so only the legacy branch separates the LTZ arm's effect. + withSQLConf(SQLConf.LEGACY_CAST_DATETIME_TO_STRING.key -> "true") { + val ltzType = TimestampLTZNanosType(9) + val ltz = AttributeReference("ltz", ltzType)() + val ntzType = TimestampNTZNanosType(9) + val ntz = AttributeReference("ntz", ntzType)() + val strLit = Literal("2020-01-02 03:04:05.123456789") + // LTZ range: both operands become strings (matches micros TimestampType). + ruleTest(rule, LessThan(ltz, strLit), LessThan(Cast(ltz, StringType), strLit)) + ruleTest(rule, GreaterThanOrEqual(strLit, ltz), + GreaterThanOrEqual(strLit, Cast(ltz, StringType))) + // NTZ range: config-blind, the string is cast to the nanos type (matches micros + // TimestampNTZType, which has no arm and falls through to canPromoteAsInBinaryComparison). + ruleTest(rule, LessThan(ntz, strLit), LessThan(ntz, Cast(strLit, ntzType))) + ruleTest(rule, GreaterThanOrEqual(strLit, ntz), + GreaterThanOrEqual(Cast(strLit, ntzType), ntz)) + // Equality: for both families the string is still cast to nanos so subseconds are compared + // exactly -- LTZ via the explicit Equality arm (which fires before the range arm), NTZ via + // the config-blind BinaryComparison fall-through. Holds for both 3VL `EqualTo` and null-safe + // `EqualNullSafe`. + Seq(ltzType -> ltz, ntzType -> ntz).foreach { case (nt, tsn) => + ruleTest(rule, EqualTo(tsn, strLit), EqualTo(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualTo(strLit, tsn), EqualTo(Cast(strLit, nt), tsn)) + ruleTest(rule, EqualNullSafe(tsn, strLit), EqualNullSafe(tsn, Cast(strLit, nt))) + ruleTest(rule, EqualNullSafe(strLit, tsn), EqualNullSafe(Cast(strLit, nt), tsn)) + } + } } test("cast WindowFrame boundaries to the type they operate upon") { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala index 293523b86f998..6b0b8abdcf15b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala @@ -30,8 +30,8 @@ import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._ import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.streaming.{GroupStateTimeout, OutputMode} -import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder} +import org.apache.spark.sql.streaming.{GroupStateTimeout, OutputMode, StatefulProcessor, TimeMode, TimerValues} +import org.apache.spark.sql.types.{IntegerType, LongType, MetadataBuilder, StructType} /** A dummy command for testing unsupported operations. */ case class DummyCommand() extends LeafCommand @@ -967,6 +967,58 @@ class UnsupportedOperationsSuite extends SparkFunSuite with SQLHelper { ) } + assertSupportedForRealTime( + "real-time with Scala transformWithState - update mode", + scalaTransformWithState(streamRelation), + Update + ) + + assertNotSupportedForRealTime( + "real-time with Scala transformWithState on both sides of union - update mode", + scalaTransformWithState(streamRelation) + .union(scalaTransformWithState(new TestStreamingRelation(attribute.newInstance()))), + Update, + "STREAMING_REAL_TIME_MODE.STATEFUL_OPERATORS_BEFORE_UNION_NOT_SUPPORTED" + ) + + assertSupportedForRealTime( + "real-time with batch aggregate before union - update mode", + streamRelation + .join(Aggregate(Nil, aggExprs("c"), batchRelation), joinType = Inner) + .select(attribute) + .union(new TestStreamingRelation(attribute.newInstance())), + Update + ) + + private def scalaTransformWithState(child: LogicalPlan): TransformWithState = { + val statefulProcessor = new StatefulProcessor[Any, Any, Any] { + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {} + + override def handleInputRows( + key: Any, + inputRows: Iterator[Any], + timerValues: TimerValues): Iterator[Any] = Iterator.empty + } + val keyEncoder = ExpressionEncoder(new StructType().add("a", IntegerType)) + .asInstanceOf[ExpressionEncoder[Any]] + new TransformWithState( + keyDeserializer = attribute, + valueDeserializer = attribute, + groupingAttributes = Seq(attribute), + dataAttributes = Seq(attribute), + statefulProcessor = statefulProcessor, + timeMode = NoTime, + outputMode = Update, + keyEncoder = keyEncoder, + outputObjAttr = attribute, + child = child, + hasInitialState = false, + initialStateGroupingAttrs = Seq(attribute), + initialStateDataAttrs = Seq(attribute), + initialStateDeserializer = attribute, + initialState = LocalRelation(Seq.empty[Attribute])) + } + /* ======================================================================================= TESTING FUNCTIONS diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/resolver/LogicalPlanDifferenceSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/resolver/LogicalPlanDifferenceSuite.scala index 36e739969b40f..6a6cc1c64e41f 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/resolver/LogicalPlanDifferenceSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/resolver/LogicalPlanDifferenceSuite.scala @@ -39,6 +39,10 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { private val FILTER = "Filter" private val PROJECT = "Project" + /** Splits `text` into lines, dropping any empty ones. */ + private def nonEmptyLines(text: String): Array[String] = + text.split("\n").filter(_.nonEmpty) + test("identical plans should return empty strings") { val plan1 = LocalRelation(idAttr, nameAttr) val plan2 = LocalRelation(idAttr, nameAttr) @@ -88,8 +92,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { assert(result2.contains("age") || result2.contains("50")) // Should NOT contain operations too far away (limit is at the top) - val lines1 = result1.split("\n").filter(_.nonEmpty).length - val lines2 = result2.split("\n").filter(_.nonEmpty).length + val lines1 = nonEmptyLines(result1).length + val lines2 = nonEmptyLines(result2).length assert( lines1 <= 7, s"Expected at most 7 lines (2 before + mismatch + 2 after + margins), got $lines1" @@ -159,8 +163,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { assert(result2.contains(FILTER)) // Both should have truncated output - val lines1 = result1.split("\n").filter(_.nonEmpty).length - val lines2 = result2.split("\n").filter(_.nonEmpty).length + val lines1 = nonEmptyLines(result1).length + val lines2 = nonEmptyLines(result2).length assert(lines1 >= 1 && lines1 <= 7, s"Expected 1-7 lines, got $lines1") assert(lines2 >= 1 && lines2 <= 7, s"Expected 1-7 lines, got $lines2") @@ -190,8 +194,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 0) // With 0 context lines, should only show the mismatched line - val lines1 = result1.split("\n").filter(_.nonEmpty) - val lines2 = result2.split("\n").filter(_.nonEmpty) + val lines1 = nonEmptyLines(result1) + val lines2 = nonEmptyLines(result2) // Should have exactly 1 line (the mismatch) assert(lines1.length == 1, s"Expected 1 line, got ${lines1.length}: ${lines1.mkString("; ")}") @@ -237,8 +241,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 5) // With 5 context lines, should show 5 before, mismatch, and 5 after (11 total) - val lines1 = result1.split("\n").filter(_.nonEmpty) - val lines2 = result2.split("\n").filter(_.nonEmpty) + val lines1 = nonEmptyLines(result1) + val lines2 = nonEmptyLines(result2) assert(lines1.length >= 10, s"Expected at least 10 lines, got ${lines1.length}") assert(lines2.length >= 10, s"Expected at least 10 lines, got ${lines2.length}") @@ -281,8 +285,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 2) // The mismatch is at the top (Project vs GlobalLimit), so both should show their top portions - val lines1 = result1.split("\n").filter(_.nonEmpty) - val lines2 = result2.split("\n").filter(_.nonEmpty) + val lines1 = nonEmptyLines(result1) + val lines2 = nonEmptyLines(result2) assert(lines1.length >= 1, s"Plan1 should have at least 1 line, got ${lines1.length}") assert(lines2.length >= 1, s"Plan2 should have at least 1 line, got ${lines2.length}") @@ -313,8 +317,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 2) - val lines1 = result1.split("\n").filter(_.nonEmpty) - val lines2 = result2.split("\n").filter(_.nonEmpty) + val lines1 = nonEmptyLines(result1) + val lines2 = nonEmptyLines(result2) // Both should have content assert(lines1.length >= 1, s"Plan1 should have at least 1 line") @@ -371,8 +375,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { assert(result1 != result2, "Plans should produce different output strings") // Should be truncated (not showing all operations) - val lines1 = result1.split("\n").filter(_.nonEmpty).length - val lines2 = result2.split("\n").filter(_.nonEmpty).length + val lines1 = nonEmptyLines(result1).length + val lines2 = nonEmptyLines(result2).length assert(lines1 <= 10, s"Expected truncated output with at most 10 lines, got $lines1") assert(lines2 <= 10, s"Expected truncated output with at most 10 lines, got $lines2") assert(lines1 >= 1, "Should have at least the mismatch line") @@ -415,10 +419,10 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 1000) // With very large context, should show entire plans - val lines1 = result1.split("\n").filter(_.nonEmpty).length - val lines2 = result2.split("\n").filter(_.nonEmpty).length - val planLines1 = plan1.toString.split("\n").filter(_.nonEmpty).length - val planLines2 = plan2.toString.split("\n").filter(_.nonEmpty).length + val lines1 = nonEmptyLines(result1).length + val lines2 = nonEmptyLines(result2).length + val planLines1 = nonEmptyLines(plan1.toString).length + val planLines2 = nonEmptyLines(plan2.toString).length assert(lines1 <= planLines1) assert(lines2 <= planLines2) @@ -445,8 +449,8 @@ class LogicalPlanDifferenceSuite extends SparkFunSuite with SQLConfHelper { val (result1, result2) = LogicalPlanDifference(plan1, plan2, 1) // With 1 context line, should show 1 line before, the mismatch, and 1 line after (3 total) - val lines1 = result1.split("\n").filter(_.nonEmpty) - val lines2 = result2.split("\n").filter(_.nonEmpty) + val lines1 = nonEmptyLines(result1) + val lines2 = nonEmptyLines(result2) assert( lines1.length <= 3, diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala index 69f4995220fe2..2a2889e441c77 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/encoders/RowEncoderSuite.scala @@ -21,7 +21,7 @@ import scala.collection.mutable import scala.util.Random import org.apache.spark.{SparkException, SparkRuntimeException} -import org.apache.spark.sql.{RandomDataGenerator, Row} +import org.apache.spark.sql.{AnalysisException, RandomDataGenerator, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.plans.CodegenInterpretedPlanTest import org.apache.spark.sql.catalyst.util.{ArrayData, DateTimeUtils, GenericArrayData, IntervalUtils} @@ -590,40 +590,77 @@ class RowEncoderSuite extends CodegenInterpretedPlanTest { val row = encoder.createSerializer()(data) } - test("do not allow serializing too long strings into char/varchar") { - Seq(CharType(5), VarcharType(5)).foreach { typ => - withSQLConf(SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true") { - val schema = new StructType().add("c", typ) - val encoder = ExpressionEncoder(schema).resolveAndBind() - val value = "abcdef" + test("SPARK-58794: encoderFor rejects CHAR/VARCHAR when first-class types are off") { + Seq(CharType(4), VarcharType(6)).foreach { dt => + val schema = new StructType().add("c", dt) + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false", + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "false") { checkError( - exception = intercept[SparkRuntimeException]({ - val row = toRow(encoder, Row(value)) - }), - condition = "EXCEED_LIMIT_LENGTH", - parameters = Map("limit" -> "5") + exception = intercept[AnalysisException] { + RowEncoder.encoderFor(schema) + }, + condition = "UNSUPPORTED_DATA_TYPE_FOR_ENCODER", + sqlState = "0A000", + parameters = Map("dataType" -> s"\"${dt.sql}\"") ) + // Engine-produced result schemas still decode on the Connect client. + RowEncoder.encoderForResultSchema(schema) + } + } + } + + test("do not allow serializing too long strings into char/varchar") { + Seq(CharType(5), VarcharType(5)).foreach { typ => + Seq( + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true", + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true").foreach { conf => + withSQLConf(conf) { + val schema = new StructType().add("c", typ) + val encoder = ExpressionEncoder(schema).resolveAndBind() + val value = "abcdef" + checkError( + exception = intercept[SparkRuntimeException]({ + val row = toRow(encoder, Row(value)) + }), + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "5") + ) + } } } } test("do not allow deserializing too long strings into char/varchar") { Seq(CharType(5), VarcharType(5)).foreach { typ => - withSQLConf(SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true") { - val fromSchema = new StructType().add("c", StringType) - val fromEncoder = ExpressionEncoder(fromSchema).resolveAndBind() - val toSchema = new StructType().add("c", typ) - val toEncoder = ExpressionEncoder(toSchema).resolveAndBind() - val value = "abcdef" - val row = toRow(fromEncoder, Row(value)) - checkError( - exception = intercept[SparkRuntimeException]({ - val value = fromRow(toEncoder, row) - }), - condition = "EXCEED_LIMIT_LENGTH", - parameters = Map("limit" -> "5") - ) + Seq( + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true", + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true").foreach { conf => + withSQLConf(conf) { + val fromSchema = new StructType().add("c", StringType) + val fromEncoder = ExpressionEncoder(fromSchema).resolveAndBind() + val toSchema = new StructType().add("c", typ) + val toEncoder = ExpressionEncoder(toSchema).resolveAndBind() + val value = "abcdef" + val row = toRow(fromEncoder, Row(value)) + checkError( + exception = intercept[SparkRuntimeException]({ + val value = fromRow(toEncoder, row) + }), + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "5") + ) + } } } } + + test("SPARK-58803: RowEncoder pads CHAR under standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val schema = new StructType().add("c", CharType(5)).add("v", VarcharType(5)) + val encoder = ExpressionEncoder(schema).resolveAndBind() + val row = toRow(encoder, Row("ab", "cd")) + assert(fromRow(encoder, row) === Row("ab ", "cd")) + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpressionSuite.scala new file mode 100644 index 0000000000000..9300f107a6aab --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ApplyFunctionExpressionSuite.scala @@ -0,0 +1,87 @@ +/* + * 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.catalyst.expressions + +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import scala.concurrent.{ExecutionContext, Future} +import scala.concurrent.duration._ + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.functions.ScalarFunction +import org.apache.spark.sql.types.{DataType, IntegerType} +import org.apache.spark.util.ThreadUtils + +class ApplyFunctionExpressionSuite extends SparkFunSuite { + + private val intIdentity = new ScalarFunction[Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "int_identity" + override def produceResult(input: InternalRow): Int = input.getInt(0) + } + + test("SPARK-58578: ApplyFunctionExpression is stateful and produces fresh copies") { + val expr = ApplyFunctionExpression( + intIdentity, Seq(BoundReference(0, IntegerType, nullable = false))) + assert(expr.stateful, "ApplyFunctionExpression.stateful should be true") + val copy = expr.freshCopyIfContainsStatefulExpression() + assert(copy ne expr, + "freshCopyIfContainsStatefulExpression should return a new instance " + + "for ApplyFunctionExpression") + assert(copy.eval(InternalRow(7)) === 7) + } + + test("SPARK-58578: fresh copies do not share the reused input row") { + val firstEvaluationStarted = new CountDownLatch(1) + val secondEvaluationStarted = new CountDownLatch(1) + val blockingIdentity = new ScalarFunction[Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "blocking_identity" + override def produceResult(input: InternalRow): Int = { + if (input.getInt(0) == 1) { + firstEvaluationStarted.countDown() + assert(secondEvaluationStarted.await(10, TimeUnit.SECONDS)) + } else { + secondEvaluationStarted.countDown() + } + input.getInt(0) + } + } + + val expr = ApplyFunctionExpression( + blockingIdentity, Seq(BoundReference(0, IntegerType, nullable = false))) + val firstEvaluator = expr.freshCopyIfContainsStatefulExpression() + val secondEvaluator = expr.freshCopyIfContainsStatefulExpression() + + val executor = ThreadUtils.newDaemonFixedThreadPool(2, "apply-function-expression-test") + val executionContext = ExecutionContext.fromExecutorService(executor) + try { + val firstResult = Future(firstEvaluator.eval(InternalRow(1)))(executionContext) + assert(firstEvaluationStarted.await(10, TimeUnit.SECONDS)) + val secondResult = Future(secondEvaluator.eval(InternalRow(2)))(executionContext) + + assert(ThreadUtils.awaitResult(firstResult, 10.seconds) === 1) + assert(ThreadUtils.awaitResult(secondResult, 10.seconds) === 2) + } finally { + executor.shutdownNow() + } + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtilsSuite.scala index 53935c66c6136..3cd4eea2e0ef5 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtilsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/BitmapExpressionUtilsSuite.scala @@ -92,4 +92,148 @@ class BitmapExpressionUtilsSuite extends SparkFunSuite { setBitmapBits(bitmap, bitmap.length - 1, 0x67) assert(BitmapExpressionUtils.bitmapCount(bitmap) == 15L) } + + test("scalar bitmap binary operations") { + val bitmap1 = Array(0xf0.toByte, 0x0f.toByte) + val bitmap2 = Array(0x70.toByte) + val originalBitmap1 = bitmap1.clone() + val originalBitmap2 = bitmap2.clone() + + val results = Seq( + BitmapExpressionUtils.bitmapAnd(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapOr(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapAndNot(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapXor(bitmap1, bitmap2)) + val expected = Seq( + Seq(0x70.toByte, 0x00.toByte), + Seq(0xf0.toByte, 0x0f.toByte), + Seq(0x80.toByte, 0x0f.toByte), + Seq(0x80.toByte, 0x0f.toByte)) + + results.zip(expected).foreach { case (result, expectedBytes) => + assert(result.length == BitmapExpressionUtils.NUM_BYTES) + assert(result.take(expectedBytes.length).toSeq == expectedBytes) + assert(result.drop(expectedBytes.length).forall(_ == 0)) + } + assert(bitmap1.sameElements(originalBitmap1)) + assert(bitmap2.sameElements(originalBitmap2)) + } + + test("scalar bitmap binary operations with a longer right input") { + val bitmap1 = Array(0xf0.toByte) + val bitmap2 = Array(0x70.toByte, 0x0f.toByte) + val originalBitmap1 = bitmap1.clone() + val originalBitmap2 = bitmap2.clone() + + val results = Seq( + BitmapExpressionUtils.bitmapAnd(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapOr(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapAndNot(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapXor(bitmap1, bitmap2)) + val expected = Seq( + Seq(0x70.toByte, 0x00.toByte), + Seq(0xf0.toByte, 0x0f.toByte), + Seq(0x80.toByte, 0x00.toByte), + Seq(0x80.toByte, 0x0f.toByte)) + + results.zip(expected).foreach { case (result, expectedBytes) => + assert(result.length == BitmapExpressionUtils.NUM_BYTES) + assert(result.take(expectedBytes.length).toSeq == expectedBytes) + assert(result.drop(expectedBytes.length).forall(_ == 0)) + } + assert(bitmap1.sameElements(originalBitmap1)) + assert(bitmap2.sameElements(originalBitmap2)) + } + + test("scalar bitmap binary operations with boundary lengths") { + val expectedBytes = Seq(0x80.toByte, 0xff.toByte, 0x00.toByte, 0x7f.toByte) + Seq(0, 1, BitmapExpressionUtils.NUM_BYTES - 1, BitmapExpressionUtils.NUM_BYTES).foreach { + size => + val bitmap1 = Array.fill[Byte](size)(0x80.toByte) + val bitmap2 = Array.fill[Byte](size)(0xff.toByte) + val results = Seq( + BitmapExpressionUtils.bitmapAnd(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapOr(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapAndNot(bitmap1, bitmap2), + BitmapExpressionUtils.bitmapXor(bitmap1, bitmap2)) + + results.zip(expectedBytes).foreach { case (result, expectedByte) => + assert(result.length == BitmapExpressionUtils.NUM_BYTES) + assert(result.take(size).forall(_ == expectedByte)) + assert(result.drop(size).forall(_ == 0)) + } + } + } + + test("scalar bitmap binary operations satisfy set operation properties") { + val bitmap = Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(0x5a.toByte) + val other = Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(0xa5.toByte) + val empty = Array.fill[Byte](BitmapExpressionUtils.NUM_BYTES)(0) + + assert(BitmapExpressionUtils.bitmapAnd(bitmap, bitmap).sameElements(bitmap)) + assert(BitmapExpressionUtils.bitmapOr(bitmap, bitmap).sameElements(bitmap)) + assert(BitmapExpressionUtils.bitmapXor(bitmap, bitmap).sameElements(empty)) + assert(BitmapExpressionUtils.bitmapAndNot(bitmap, bitmap).sameElements(empty)) + assert(BitmapExpressionUtils.bitmapAndNot(bitmap, empty).sameElements(bitmap)) + assert(BitmapExpressionUtils.bitmapAndNot(empty, bitmap).sameElements(empty)) + assert(BitmapExpressionUtils.bitmapOr(bitmap, empty).sameElements(bitmap)) + assert(BitmapExpressionUtils.bitmapXor(bitmap, empty).sameElements(bitmap)) + assert(BitmapExpressionUtils.bitmapOr(bitmap, other).sameElements( + BitmapExpressionUtils.bitmapOr(other, bitmap))) + assert(BitmapExpressionUtils.bitmapXor(bitmap, other).sameElements( + BitmapExpressionUtils.bitmapXor(other, bitmap))) + } + + test("bitmap_xor_merge equal length") { + val bitmap1 = Array[Byte](0x10, 0x30, 0x40) + val bitmap2 = Array[Byte](0x10, 0x20, 0x40) + // 0x10 ^ 0x10 = 0x00, 0x30 ^ 0x20 = 0x10, 0x40 ^ 0x40 = 0x00 + val expected = Array[Byte](0x00, 0x10, 0x00) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + for (i <- expected.indices) { + assert(bitmap1(i) == expected(i), s"bitmap1($i) should be ${expected(i)}") + } + } + + test("bitmap_xor_merge different lengths") { + val bitmap1 = Array[Byte](0x0A, 0x0B, 0x0C) + val bitmap2 = Array[Byte](0x0A) + // 0x0A ^ 0x0A = 0x00, remaining bytes unchanged because XOR 0 = X + val expected = Array[Byte](0x00, 0x0B, 0x0C) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + for (i <- expected.indices) { + assert(bitmap1(i) == expected(i), s"bitmap1($i) should be ${expected(i)}") + } + } + + test("bitmap_xor_merge all zeros") { + val bitmap1 = Array[Byte](0x10, 0x20) + val bitmap2 = Array[Byte](0x00, 0x00) + val expected = Array[Byte](0x10, 0x20) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + for (i <- expected.indices) { + assert(bitmap1(i) == expected(i), s"bitmap1($i) should be ${expected(i)}") + } + } + + test("bitmap_xor_merge self xor equals zero") { + val bitmap1 = Array[Byte](0x10, 0x30, 0x40) + val bitmap2 = Array[Byte](0x10, 0x30, 0x40) + val expected = Array[Byte](0x00, 0x00, 0x00) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + for (i <- expected.indices) { + assert(bitmap1(i) == expected(i), s"bitmap1($i) should be ${expected(i)}") + } + } + + test("bitmap_xor_merge with bytes containing sign bits") { + val bitmap1 = Array[Byte](0xFF.toByte, 0x80.toByte) + val bitmap2 = Array[Byte](0xF0.toByte, 0x0F.toByte) + // 0xFF ^ 0xF0 = 0x0F, 0x80 ^ 0x0F = 0x8F + val expected = Array[Byte](0x0F.toByte, 0x8F.toByte) + BitmapExpressionUtils.bitmapXorMerge(bitmap1, bitmap2) + for (i <- expected.indices) { + assert(bitmap1(i) == expected(i), s"bitmap1($i) should be ${expected(i)}") + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala index 7ce14bcedf4ba..d713ce4e82173 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala @@ -442,6 +442,48 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper { assert(ctx2.mutableStateInitCode.size == CodeGenerator.MUTABLESTATEARRAY_SIZE_LIMIT + 10) } + test("SPARK-58437: declare generic mutable state arrays with raw array creation") { + val ctx = new CodegenContext + val samplerType = + "org.apache.spark.util.random.BernoulliCellSampler" + + "<org.apache.spark.sql.catalyst.expressions.UnsafeRow>" + for (_ <- 1 to CodeGenerator.OUTER_CLASS_VARIABLES_THRESHOLD + 1) { + ctx.addMutableState(samplerType, "sampler") + } + + val states = ctx.declareMutableStates() + assert(states.contains(s"private $samplerType[]")) + assert(!states.contains(s"new $samplerType[")) + assert(states.contains("new org.apache.spark.util.random.BernoulliCellSampler[")) + } + + test("SPARK-58437: generate javac-compatible source for collection expressions") { + val intArray = + BoundReference(0, ArrayType(IntegerType, containsNull = false), nullable = false) + val intArray2 = + BoundReference(1, ArrayType(IntegerType, containsNull = false), nullable = false) + val collectionExprs = Seq( + ArrayDistinct(intArray), + ArrayUnion(intArray, intArray2), + ArrayIntersect(intArray, intArray2), + ArrayExcept(intArray, intArray2)) + + collectionExprs.foreach { expr => + val ctx = new CodegenContext + val exprCode = expr.genCode(ctx).code.toString + val code = exprCode + "\n" + ctx.declareAddedFunctions() + assert(!code.contains("ArrayBuilder$ofInt")) + assert(code.contains("ArrayBuilder.ofInt")) + } + + val sequence = new Sequence( + BoundReference(0, IntegerType, nullable = false), + BoundReference(1, IntegerType, nullable = false), + BoundReference(2, IntegerType, nullable = false)) + val sequenceCode = sequence.genCode(new CodegenContext).code.toString + assert(!sequenceCode.contains("final int[]")) + } + test("SPARK-22750: addImmutableStateIfNotExists") { val ctx = new CodegenContext val mutableState1 = "field1" @@ -626,7 +668,11 @@ class CodeGenerationSuite extends SparkFunSuite with ExpressionEvalHelper { |} |""".stripMargin - CodeGenerator.compile(new CodeAndComment(code, Map.empty)) + // Pin the Janino backend: this guards a Janino-specific compiler bug and must keep + // exercising Janino whatever spark.sql.codegen.compiler is set to. + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "janino") { + CodeGenerator.compile(new CodeAndComment(code, Map.empty)) + } } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGeneratorWithInterpretedFallbackSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGeneratorWithInterpretedFallbackSuite.scala index a843d43ae83b6..bf6351bba4b9f 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGeneratorWithInterpretedFallbackSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGeneratorWithInterpretedFallbackSuite.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.catalyst.expressions import java.util.concurrent.ExecutionException +import org.codehaus.commons.compiler.CompileException + import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.aggregate.NoOp @@ -83,13 +85,21 @@ class CodeGeneratorWithInterpretedFallbackSuite extends SparkFunSuite with PlanT } test("codegen failures in the CODEGEN_ONLY mode") { - val errMsg = intercept[ExecutionException] { + val e = intercept[ExecutionException] { val input = Seq(BoundReference(0, IntegerType, nullable = true)) withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenOnly) { FailedCodegenProjection.createObject(input) } - }.getMessage - assert(errMsg.contains("Failed to compile: org.codehaus.commons.compiler.CompileException:")) + } + // SPARK-23711/SPARK-25140 made this path catch the exception the compile cache raises rather + // than the compiler's own: a source-level failure goes through `compilerError`, whose checked + // CompileException the cache wraps in an ExecutionException. (The other branch, + // `internalCompilerError`, builds an unchecked InternalCompilerException, which + // `CodeGenerator.compile` unwraps and rethrows bare.) + assert(e.getCause.isInstanceOf[CompileException]) + // "Failed to compile: " comes from `QueryExecutionErrors.failedToCompileMsg`; the compiler's + // own diagnostic follows it, and its wording is not Spark's to assert on. + assert(e.getMessage.contains("Failed to compile:")) } test("SPARK-25358 Correctly handles NoOp in MutableProjection") { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala index 37e0b53dd46d0..f26285fb2d165 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CollectionExpressionsSuite.scala @@ -24,13 +24,13 @@ import java.util.TimeZone import scala.language.implicitConversions import scala.util.Random -import org.apache.spark.{SparkArrayIndexOutOfBoundsException, SparkFunSuite, SparkRuntimeException} +import org.apache.spark.{SparkArrayIndexOutOfBoundsException, SparkFunSuite, SparkIllegalArgumentException, SparkRuntimeException} import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TypeCheckResult import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch -import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, DateTimeTestUtils, DateTimeUtils, GenericArrayData} -import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.{outstandingZoneIds, LA, UTC} +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, DateTimeTestUtils, DateTimeUtils, GenericArrayData, TimestampNanosTestUtils} +import org.apache.spark.sql.catalyst.util.DateTimeTestUtils.{outstandingZoneIds, LA, UTC, UTC_OPT} import org.apache.spark.sql.catalyst.util.IntervalUtils._ import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber import org.apache.spark.sql.errors.DataTypeErrorsBase @@ -923,10 +923,68 @@ class CollectionExpressionsSuite checkEvaluation(Slice(a1, Literal(1), Literal(2)), Seq("a", "b")) checkEvaluation(Slice(a2, Literal(1), Literal(2)), Seq("", null)) checkEvaluation(Slice(a0, Literal(10), Literal(1)), Seq.empty[Int]) + // SPARK-57665: a large length must not overflow startIndex + length; both the interpreted and + // codegen paths must return the tail of the array, not an empty array. + checkEvaluation(Slice(a0, Literal(2), Literal(Int.MaxValue)), Seq(2, 3, 4, 5, 6)) + checkEvaluation(Slice(a0, Literal(1), Literal(Int.MaxValue)), Seq(1, 2, 3, 4, 5, 6)) + // Negative start with a large length exercises the start < 0 branch together with clamping. + checkEvaluation(Slice(a0, Literal(-2), Literal(Int.MaxValue)), Seq(5, 6)) + // Both extremes: an Int.MinValue start resolves far below 0, so the guard returns empty and + // the (harmlessly overflowing) resolved length is never used. + checkEvaluation(Slice(a0, Literal(Int.MinValue), Literal(Int.MaxValue)), Seq.empty[Int]) + // A negative length is still rejected when the start is out of range, because the length is + // resolved before the start guard. + checkErrorInExpression[SparkRuntimeException]( + expression = Slice(a0, Literal(-20), Literal(-1)), + condition = "INVALID_PARAMETER_VALUE.LENGTH", + parameters = Map( + "parameter" -> toSQLId("length"), + "length" -> (-1).toString, + "functionName" -> toSQLId("slice") + )) checkEvaluation(Slice(a1, Literal(10), Literal(1)), Seq.empty[String]) checkEvaluation(Slice(a3, Literal(2), Literal(3)), Seq(2, null, 4)) } + test("TrimArray") { + val a0 = Literal.create(Seq(1, 2, 3, 4, 5), ArrayType(IntegerType)) + val a1 = Literal.create(Seq[String]("a", "b", "c"), ArrayType(StringType)) + val a2 = Literal.create(Seq[String]("a", null, "b"), ArrayType(StringType, containsNull = true)) + val a3 = Literal.create(Seq.empty[Int], ArrayType(IntegerType)) + + // n between 0 and cardinality removes the last n elements. + checkEvaluation(TrimArray(a0, Literal(0)), Seq(1, 2, 3, 4, 5)) + checkEvaluation(TrimArray(a0, Literal(2)), Seq(1, 2, 3)) + checkEvaluation(TrimArray(a0, Literal(5)), Seq.empty[Int]) + checkEvaluation(TrimArray(a1, Literal(1)), Seq("a", "b")) + checkEvaluation(TrimArray(a2, Literal(1)), Seq("a", null)) + checkEvaluation(TrimArray(a3, Literal(0)), Seq.empty[Int]) + + // NULL array or NULL n yields NULL. + checkEvaluation(TrimArray(Literal.create(null, ArrayType(IntegerType)), Literal(1)), null) + checkEvaluation(TrimArray(a0, Literal.create(null, IntegerType)), null) + + // n < 0 and n > cardinality are rejected. + checkErrorInExpression[SparkRuntimeException]( + expression = TrimArray(a0, Literal(-1)), + condition = "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + parameters = Map( + "parameter" -> toSQLId("n"), + "functionName" -> toSQLId("trim_array"), + "numElements" -> "5", + "length" -> (-1).toString + )) + checkErrorInExpression[SparkRuntimeException]( + expression = TrimArray(a0, Literal(6)), + condition = "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + parameters = Map( + "parameter" -> toSQLId("n"), + "functionName" -> toSQLId("trim_array"), + "numElements" -> "5", + "length" -> 6.toString + )) + } + test("ArrayJoin") { def testArrays( arrays: Seq[Expression], @@ -1184,6 +1242,39 @@ class CollectionExpressionsSuite Seq(-1.toByte, -2.toByte, -3.toByte)) } + test("SPARK-58440: illegal sequence boundaries carry String message parameters") { + // Codegen-only: the interpreted path reports this through `require`, which throws a plain + // IllegalArgumentException with no error class or parameters, so the two-mode + // `checkErrorInExpression` cannot be used here. + // The parameter map is built in generated Java source. Janino erases the + // `Map<String, String>` type arguments, so non-String values used to slip in and reach + // `SparkThrowable.getMessageParameters`, whose declared value type is String. + withSQLConf( + SQLConf.CODEGEN_FACTORY_MODE.key -> CodegenObjectFactoryMode.CODEGEN_ONLY.toString) { + // Numeric start/stop/step (IntegralSequenceImpl). + checkError( + exception = intercept[SparkIllegalArgumentException] { + evaluateWithMutableProjection(new Sequence(Literal(1), Literal(2), Literal(0))) + }, + condition = "_LEGACY_ERROR_TEMP_3243", + parameters = Map("start" -> "1", "stop" -> "2", "step" -> "0")) + + // Interval step (InternalSequenceBase): `step` is a CalendarInterval, not a number. A + // month-granularity step is required to reach this path; a day-granularity one is + // delegated to the integral implementation with a plain `int` step. + checkError( + exception = intercept[SparkIllegalArgumentException] { + evaluateWithMutableProjection(Sequence( + Literal(Date.valueOf("1970-01-01")), + Literal(Date.valueOf("1970-02-01")), + Some(Literal(negateExact(stringToInterval("interval 1 month")))), + UTC_OPT)) + }, + condition = "_LEGACY_ERROR_TEMP_3243", + parameters = Map("start" -> "0", "stop" -> "2678400000000", "step" -> "-1 months")) + } + } + test("Sequence of timestamps") { checkEvaluation(new Sequence( Literal(Timestamp.valueOf("2018-01-01 00:00:00")), @@ -2131,6 +2222,47 @@ class CollectionExpressionsSuite checkEvaluation(ElementAt(dupTimeMap, Literal(t1, timeType)), 10) checkEvaluation(ElementAt(dupTimeMap, Literal(t2, timeType)), 20) + // Nanosecond timestamp keys (SPARK-57841). element_at routes through the same + // GetMapValueUtil map-lookup path as GetMapValue, so it must also distinguish keys that + // share epochMicros and differ only in nanosWithinMicro, on both the hash (threshold 0) and + // linear (threshold Int.MaxValue) paths. Physical TimestampNanosVal objects use the + // hashCode()/equals() fall-through arm (not the primitive-long arm). Built from internal + // values via ArrayBasedMapData (Literal.create of a Scala nanos-keyed Map would route + // through the schema-aware converter, which only accepts LocalDateTime / Instant keys). + val ntz9 = TimestampNTZNanosType(9) + val micros = 1577836800000000L // 2020-01-01T00:00:00Z + val n1 = TimestampNanosTestUtils.nanosVal(micros, 1) + val n2 = TimestampNanosTestUtils.nanosVal(micros, 999) // same micro as n1, must not alias + val n3 = TimestampNanosTestUtils.nanosVal(micros + 1, 0) + val ntzNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](n1, n2, n3)), + new GenericArrayData(Array[Any](10, 20, 30))), + MapType(ntz9, IntegerType)) + checkEvaluation(ElementAt(ntzNanosMap, Literal(n1, ntz9)), 10) + checkEvaluation(ElementAt(ntzNanosMap, Literal(n2, ntz9)), 20) + checkEvaluation(ElementAt(ntzNanosMap, Literal(n3, ntz9)), 30) + checkEvaluation(ElementAt(ntzNanosMap, + Literal(TimestampNanosTestUtils.nanosVal(micros, 500), ntz9)), null) + + // LTZ family + null value + duplicate-key first-wins. + val ltz9 = TimestampLTZNanosType(9) + val l1 = TimestampNanosTestUtils.nanosVal(micros, 100) + val l2 = TimestampNanosTestUtils.nanosVal(micros, 900) + val ltzNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](l1, l2)), + new GenericArrayData(Array[Any](10, null))), + MapType(ltz9, IntegerType)) + checkEvaluation(ElementAt(ltzNanosMap, Literal(l1, ltz9)), 10) + checkEvaluation(ElementAt(ltzNanosMap, Literal(l2, ltz9)), null) // present, null value + val dupNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](n1, n2, n1)), + new GenericArrayData(Array[Any](10, 20, 30))), + MapType(ntz9, IntegerType)) + checkEvaluation(ElementAt(dupNanosMap, Literal(n1, ntz9)), 10) + // Array Keys val arrayType = ArrayType(IntegerType) val arrayMap = Literal.create( diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala index f2a8e5b64d211..abc28ff203f8d 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala @@ -294,6 +294,53 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper { checkEvaluation(GetMapValue(dupTimeMap, Literal(t1, timeType)), 10) checkEvaluation(GetMapValue(dupTimeMap, Literal(t2, timeType)), 20) + // Nanosecond timestamp keys (SPARK-57841). Physically TimestampNanosVal objects, so they + // use the hashCode()/equals() fall-through arm of genHash / hashKeyOnDriver / genEqual -- + // NOT the primitive-long arm that lists only microsecond timestamps. Keys share epochMicros + // and differ only in nanosWithinMicro, so a correct lookup must consult the full + // sub-microsecond value on both the hash (threshold 0) and linear (threshold Int.MaxValue) + // paths. Built from internal TimestampNanosVal via ArrayBasedMapData because Literal.create + // of a Scala Map with a nanos type routes through the schema-aware converter, which accepts + // only LocalDateTime / Instant keys (see the Decimal / Time duplicate-key cases above). + val ntz9 = TimestampNTZNanosType(9) + val micros = 1577836800000000L // 2020-01-01T00:00:00Z + val n1 = TimestampNanosTestUtils.nanosVal(micros, 1) + val n2 = TimestampNanosTestUtils.nanosVal(micros, 999) // same micro as n1, must not alias + val n3 = TimestampNanosTestUtils.nanosVal(micros + 1, 0) + val ntzNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](n1, n2, n3)), + new GenericArrayData(Array[Any](10, 20, 30))), + MapType(ntz9, IntegerType)) + checkEvaluation(GetMapValue(ntzNanosMap, Literal(n1, ntz9)), 10) + checkEvaluation(GetMapValue(ntzNanosMap, Literal(n2, ntz9)), 20) + checkEvaluation(GetMapValue(ntzNanosMap, Literal(n3, ntz9)), 30) + checkEvaluation(GetMapValue(ntzNanosMap, + Literal(TimestampNanosTestUtils.nanosVal(micros, 500), ntz9)), null) + + // LTZ family behaves the same; also cover a null map value. + val ltz9 = TimestampLTZNanosType(9) + val l1 = TimestampNanosTestUtils.nanosVal(micros, 100) + val l2 = TimestampNanosTestUtils.nanosVal(micros, 900) + val ltzNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](l1, l2)), + new GenericArrayData(Array[Any](10, null))), + MapType(ltz9, IntegerType)) + checkEvaluation(GetMapValue(ltzNanosMap, Literal(l1, ltz9)), 10) + checkEvaluation(GetMapValue(ltzNanosMap, Literal(l2, ltz9)), null) // present, null value + checkEvaluation(GetMapValue(ltzNanosMap, + Literal(TimestampNanosTestUtils.nanosVal(micros, 500), ltz9)), null) // absent + + // Nanosecond duplicate keys: first match wins (ArrayBasedMapData first-wins semantics). + val dupNanosMap = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(Array[Any](n1, n2, n1)), + new GenericArrayData(Array[Any](10, 20, 30))), + MapType(ntz9, IntegerType)) + checkEvaluation(GetMapValue(dupNanosMap, Literal(n1, ntz9)), 10) + checkEvaluation(GetMapValue(dupNanosMap, Literal(n2, ntz9)), 20) + // 6. Binary Keys val binaryMap = Literal.create(Map(Array(1.toByte) -> 10, Array(2.toByte) -> 20), MapType(BinaryType, IntegerType)) @@ -368,6 +415,23 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper { assert(ElementAt(foldableLit, Literal(1)).usesFoldableHashLookup) } + // A nanosecond-timestamp key type is a hashable AtomicType, so a foldable nanos-keyed map above + // the threshold takes the hash path too (SPARK-57841). Keys are internal TimestampNanosVal. + val ntz9 = TimestampNTZNanosType(9) + val nanosEntries = (0 until 2000).map { i => + TimestampNanosTestUtils.nanosVal(1577836800000000L + i, i % 1000) -> i + } + val nanosFoldableLit = Literal.create( + new ArrayBasedMapData( + new GenericArrayData(nanosEntries.map(_._1.asInstanceOf[Any]).toArray), + new GenericArrayData(nanosEntries.map(_._2.asInstanceOf[Any]).toArray)), + MapType(ntz9, IntegerType)) + withSQLConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD.key -> "1000") { + assert(GetMapValue(nanosFoldableLit, + Literal(TimestampNanosTestUtils.nanosVal(1577836800000000L, 0), ntz9)) + .usesFoldableHashLookup) + } + // Foldable but below threshold --> LinearExecutor. withSQLConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD.key -> "10000") { assert(!GetMapValue(foldableLit, Literal(1)).usesFoldableHashLookup) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CsvExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CsvExpressionsSuite.scala index 631b08b5395f9..0cda94d7b188d 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CsvExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CsvExpressionsSuite.scala @@ -314,4 +314,18 @@ class CsvExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { CsvToStructs(schema, Map.empty, Literal.create(null, StringType), UTC_OPT), null) } + + test("CsvToStructs and StructsToCsv are stateful and produce fresh copies") { + val schema = StructType(StructField("a", IntegerType) :: Nil) + + val csvToStructs = CsvToStructs(schema, Map.empty, Literal("1"), UTC_OPT) + assert(csvToStructs.stateful) + assert(csvToStructs.freshCopyIfContainsStatefulExpression() ne csvToStructs) + + val struct = Literal.create(create_row(1), schema) + val structsToCsv = StructsToCsv(Map.empty, struct, UTC_OPT) + assert(structsToCsv.stateful) + assert(structsToCsv.freshCopyIfContainsStatefulExpression() ne structsToCsv) + } + } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala index 226165b7b36f5..ef90fa5d68a5f 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DateExpressionsSuite.scala @@ -28,7 +28,7 @@ import scala.language.postfixOps import scala.reflect.ClassTag import scala.util.Random -import org.apache.spark.{SparkArithmeticException, SparkDateTimeException, SparkFunSuite, SparkIllegalArgumentException, SparkRuntimeException, SparkUpgradeException} +import org.apache.spark.{SPARK_DOC_ROOT, SparkArithmeticException, SparkDateTimeException, SparkException, SparkFunSuite, SparkIllegalArgumentException, SparkRuntimeException, SparkUpgradeException} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch @@ -2341,6 +2341,158 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { assert(ltzMismatch.errorSubClass == "UNEXPECTED_INPUT_TYPE") } + test("SPARK-57825: add/subtract ANSI year-month interval on nanos timestamps") { + val interval = Period.ofYears(1).plusMonths(2) + val minusInterval = Period.ofMonths(-1) + + // A year-month shift moves only the month field: the whole fraction (including the + // sub-microsecond `789`) and the time of day are carried through unchanged. + val ntzType = TimestampNTZNanosType(9) + val ntzStart = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2020-01-02T03:04:05.123456789"), precision = 9) + val ntzExpectedAdd = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2021-03-02T03:04:05.123456789"), precision = 9) + val ntzExpectedSub = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2019-12-02T03:04:05.123456789"), precision = 9) + + checkEvaluation( + TimestampAddYMInterval(Literal.create(ntzStart, ntzType), Literal(interval), Some("UTC")), + ntzExpectedAdd) + checkEvaluation( + TimestampAddYMInterval( + Literal.create(ntzStart, ntzType), + UnaryMinus(Literal(interval)), + Some("UTC")), + DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2018-11-02T03:04:05.123456789"), precision = 9)) + checkEvaluation( + TimestampAddYMInterval( + Literal.create(ntzStart, ntzType), Literal(minusInterval), Some("UTC")), + ntzExpectedSub) + assert(ntzExpectedAdd.nanosWithinMicro == ntzStart.nanosWithinMicro) + assert(ntzExpectedSub.nanosWithinMicro == ntzStart.nanosWithinMicro) + + val ltzType = TimestampLTZNanosType(9) + val ltzStart = DateTimeUtils.instantToTimestampNanos( + Instant.parse("2020-01-02T03:04:05.123456789Z"), precision = 9) + val ltzExpectedAdd = DateTimeUtils.instantToTimestampNanos( + Instant.parse("2021-03-02T03:04:05.123456789Z"), precision = 9) + val ltzExpectedSub = DateTimeUtils.instantToTimestampNanos( + Instant.parse("2019-12-02T03:04:05.123456789Z"), precision = 9) + + checkEvaluation( + TimestampAddYMInterval(Literal.create(ltzStart, ltzType), Literal(interval), Some("UTC")), + ltzExpectedAdd) + checkEvaluation( + TimestampAddYMInterval( + Literal.create(ltzStart, ltzType), + UnaryMinus(Literal(interval)), + Some("UTC")), + DateTimeUtils.instantToTimestampNanos( + Instant.parse("2018-11-02T03:04:05.123456789Z"), precision = 9)) + checkEvaluation( + TimestampAddYMInterval( + Literal.create(ltzStart, ltzType), Literal(minusInterval), Some("UTC")), + ltzExpectedSub) + assert(ltzExpectedAdd.nanosWithinMicro == ltzStart.nanosWithinMicro) + assert(ltzExpectedSub.nanosWithinMicro == ltzStart.nanosWithinMicro) + + yearMonthIntervalTypes.foreach { it => + checkConsistencyBetweenInterpretedAndCodegen( + (ts: Expression, ym: Expression) => TimestampAddYMInterval(ts, ym, Some("UTC")), + ntzType, it) + checkConsistencyBetweenInterpretedAndCodegen( + (ts: Expression, ym: Expression) => TimestampAddYMInterval(ts, ym, Some("UTC")), + ltzType, it) + } + } + + test("SPARK-57832: subtract nanosecond-precision timestamps") { + // The difference between two nanosecond timestamps is reported on the microsecond grid: only + // each operand's epochMicros participates, so the sub-microsecond remainder is truncated. The + // first pair below differs by a whole day plus 0.123456 s at the microsecond level, and the + // 789/111 sub-microsecond digits drop out entirely. The zero-result case (two values inside the + // same microsecond) is exercised by the second pair further down. + val ntzType = TimestampNTZNanosType(9) + val ntzLeft = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2020-01-02T03:04:05.123456789"), precision = 9) + val ntzRight = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2020-01-01T03:04:05.000000111"), precision = 9) + // 1 day + 0.123456 s; the 789/111 sub-microsecond digits drop out. + checkEvaluation( + SubtractTimestamps( + Literal.create(ntzLeft, ntzType), + Literal.create(ntzRight, ntzType), + legacyInterval = false, + timeZoneId = Some("UTC")), + Duration.ofDays(1).plus(123456, ChronoUnit.MICROS)) + // Two values inside the same microsecond subtract to exactly zero (remainder truncated). + checkEvaluation( + SubtractTimestamps( + Literal.create(ntzLeft, ntzType), + Literal.create( + DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2020-01-02T03:04:05.123456001"), precision = 9), + ntzType), + legacyInterval = false, + timeZoneId = Some("UTC")), + Duration.ZERO) + // Legacy calendar-interval result carries the same microsecond difference. + checkEvaluation( + SubtractTimestamps( + Literal.create(ntzLeft, ntzType), + Literal.create(ntzRight, ntzType), + legacyInterval = true, + timeZoneId = Some("UTC")), + new CalendarInterval(0, 0, MICROS_PER_DAY + 123456L)) + + // LTZ nanos: subtraction reads the local wall clock at the session zone. Evaluated at UTC the + // instants and their local date-times coincide, so the difference matches the NTZ case above. + val ltzType = TimestampLTZNanosType(9) + val ltzLeft = DateTimeUtils.instantToTimestampNanos( + Instant.parse("2020-01-02T03:04:05.123456789Z"), precision = 9) + val ltzRight = DateTimeUtils.instantToTimestampNanos( + Instant.parse("2020-01-01T03:04:05.000000111Z"), precision = 9) + checkEvaluation( + SubtractTimestamps( + Literal.create(ltzLeft, ltzType), + Literal.create(ltzRight, ltzType), + legacyInterval = false, + timeZoneId = Some("UTC")), + Duration.ofDays(1).plus(123456, ChronoUnit.MICROS)) + + // Pre-epoch operand exercises the negative-epoch path; the result stays on the micros grid. + val ntzPreEpoch = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("1960-01-01T00:00:00.000000999"), precision = 9) + checkEvaluation( + SubtractTimestamps( + Literal.create(ntzLeft, ntzType), + Literal.create(ntzPreEpoch, ntzType), + legacyInterval = false, + timeZoneId = Some("UTC")), + Duration.between( + LocalDateTime.parse("1960-01-01T00:00:00"), + LocalDateTime.parse("2020-01-02T03:04:05.123456"))) + + // NULL operands propagate. + checkEvaluation( + SubtractTimestamps( + Literal.create(null, ntzType), + Literal.create(ntzRight, ntzType), + legacyInterval = false, + timeZoneId = Some("UTC")), + null) + + Seq(false, true).foreach { legacy => + checkConsistencyBetweenInterpretedAndCodegen( + (l: Expression, r: Expression) => SubtractTimestamps(l, r, legacy, Some("UTC")), + ntzType, ntzType) + checkConsistencyBetweenInterpretedAndCodegen( + (l: Expression, r: Expression) => SubtractTimestamps(l, r, legacy, Some("UTC")), + ltzType, ltzType) + } + } + test("SPARK-37552: convert a timestamp_ntz to another time zone") { checkEvaluation( ConvertTimezone( @@ -2381,6 +2533,104 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { } } + test("SPARK-57818: convert_timezone over nanosecond-precision timestamps") { + val ntzType9 = TimestampNTZNanosType(9) + + // The nanosWithinMicro remainder is carried through unchanged by a zone conversion; only the + // whole-microsecond part shifts with the zone offset. + val srcNanos = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2022-03-27T03:00:00.123456789"), 9) + val expectedNanos = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2022-03-27T04:00:00.123456789"), 9) + assert(srcNanos.nanosWithinMicro == expectedNanos.nanosWithinMicro) + checkEvaluation( + ConvertTimezone( + Literal("Europe/Brussels"), + Literal("Europe/Moscow"), + Literal.create(srcNanos, ntzType9)), + expectedNanos) + + // Pre-epoch values exercise the negative-epoch path. The expected epochMicros is derived from + // the already-verified micros-only conversion (SPARK-37552 tests that path); this only checks + // that the nanos wiring delegates to it correctly and carries the remainder through unchanged. + val preEpochSrc = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("1960-01-01T00:00:00.000000001"), 9) + val preEpochExpectedMicros = DateTimeUtils.convertTimestampNtzToAnotherTz( + "Europe/Moscow", "Europe/Brussels", preEpochSrc.epochMicros) + checkEvaluation( + ConvertTimezone( + Literal("Europe/Moscow"), + Literal("Europe/Brussels"), + Literal.create(preEpochSrc, ntzType9)), + TimestampNanosVal.fromParts(preEpochExpectedMicros, preEpochSrc.nanosWithinMicro)) + + // Precision (7/8/9) is preserved on the result; AnyTimestampNanoType.defaultConcreteType + // would incorrectly always widen it to 9. + Seq(7, 8, 9).foreach { precision => + val ntzType = TimestampNTZNanosType(precision) + val src = DateTimeUtils.localDateTimeToTimestampNanos( + LocalDateTime.parse("2022-03-27T03:00:00.123456789"), precision) + val convertExpr = ConvertTimezone( + Literal("Europe/Brussels"), Literal("Europe/Moscow"), Literal.create(src, ntzType)) + assert(convertExpr.dataType === ntzType) + } + + // LTZ(p) nanos values are rejected: this function is NTZ-only, matching the existing + // TimestampNTZType-only micro path. Unlike that micro path -- which implicitly casts a + // plain LTZ TimestampType argument down to TimestampNTZType -- a nanos source must not be + // silently reinterpreted from LTZ to NTZ, since that would drop the source time zone + // information without the user asking for it. + val ltzNanos = DateTimeUtils.instantToTimestampNanos(Instant.parse("2022-03-27T03:00:00Z"), 9) + val ltzMismatch = ConvertTimezone( + Literal("Europe/Brussels"), Literal("Europe/Moscow"), + Literal.create(ltzNanos, TimestampLTZNanosType(9))) + .checkInputDataTypes().asInstanceOf[DataTypeMismatch] + assert(ltzMismatch.errorSubClass == "UNEXPECTED_INPUT_TYPE") + + // A wholly invalid source type (not any kind of timestamp) hits the generic type check + // instead of the explicit LTZ guard above; both paths must report the same requiredType, + // since neither actually accepts an LTZ(p) source. + val wrongTypeMismatch = ConvertTimezone( + Literal("Europe/Brussels"), Literal("Europe/Moscow"), Literal(1)) + .checkInputDataTypes().asInstanceOf[DataTypeMismatch] + assert(wrongTypeMismatch.errorSubClass == "UNEXPECTED_INPUT_TYPE") + assert(wrongTypeMismatch.messageParameters("requiredType") === + ltzMismatch.messageParameters("requiredType")) + + // NULL handling: a NULL nanosecond timestamp, and NULL zone arguments with a non-NULL + // nanosecond timestamp. + checkEvaluation( + ConvertTimezone( + Literal("America/Los_Angeles"), + Literal("UTC"), + Literal.create(null, ntzType9)), + null) + checkEvaluation( + ConvertTimezone( + Literal.create(null, StringType), + Literal("UTC"), + Literal.create(srcNanos, ntzType9)), + null) + checkEvaluation( + ConvertTimezone( + Literal("America/Los_Angeles"), + Literal.create(null, StringType), + Literal.create(srcNanos, ntzType9)), + null) + + outstandingTimezonesIds.foreach { sourceTz => + outstandingTimezonesIds.foreach { targetTz => + checkConsistencyBetweenInterpretedAndCodegen( + (_: Expression, _: Expression, sourceTs: Expression) => + ConvertTimezone( + Literal(sourceTz), + Literal(targetTz), + sourceTs), + StringType, StringType, ntzType9) + } + } + } + test("SPARK-38195: add a quantity of interval units to a timestamp") { // Check case-insensitivity checkEvaluation( @@ -3236,4 +3486,104 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { val r5 = expr5.checkInputDataTypes().asInstanceOf[DataTypeMismatch] assert(r5.errorSubClass == "UNEXPECTED_INPUT_TYPE") } + + test("SPARK-57837: CurrentTimestampExpressionBuilder") { + // No argument keeps the historical micro TIMESTAMP expressions. + assert(CurrentTimestampExpressionBuilder.build("current_timestamp", Seq.empty) === + CurrentTimestamp()) + assert(CurrentTimestampExpressionBuilder.build("now", Seq.empty) === Now()) + + // Precision 6 stays on the micro type, per function name. + assert(CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Literal(6))) === + CurrentTimestamp()) + assert(CurrentTimestampExpressionBuilder.build("now", Seq(Literal(6))) === Now()) + + // Precisions 7-9 build the nanosecond TIMESTAMP_LTZ variant, including a foldable arg. + Seq(7, 8, 9).foreach { p => + val built = CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Literal(p))) + assert(built === CurrentTimestampNanos(p)) + assert(built.dataType === TimestampLTZNanosType(p)) + } + assert(CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Add(Literal(4), + Literal(5)))) === CurrentTimestampNanos(9)) + + // Out-of-range precision (other than 6) is rejected with INVALID_TIMESTAMP_PRECISION. + Seq(0, 3, 5, 10).foreach { p => + checkError( + exception = intercept[SparkException] { + CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Literal(p))) + }, + condition = "INVALID_TIMESTAMP_PRECISION", + parameters = Map("precision" -> p.toString, "type" -> "TIMESTAMP_LTZ")) + } + + // Non-foldable precision. + checkError( + exception = intercept[AnalysisException] { + CurrentTimestampExpressionBuilder.build( + "current_timestamp", Seq(AttributeReference("a", IntegerType)())) + }, + condition = "NON_FOLDABLE_ARGUMENT", + parameters = Map( + "funcName" -> "`current_timestamp`", + "paramName" -> "`precision`", + "paramType" -> "\"INT\"")) + + // Non-integral precision. + checkError( + exception = intercept[AnalysisException] { + CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Literal("9"))) + }, + condition = "UNEXPECTED_INPUT_TYPE", + parameters = Map( + "paramIndex" -> "first", + "functionName" -> "`current_timestamp`", + "requiredType" -> "\"INT\"", + "inputSql" -> "\"9\"", + "inputType" -> "\"STRING\"")) + + // Too many arguments. + checkError( + exception = intercept[AnalysisException] { + CurrentTimestampExpressionBuilder.build("current_timestamp", Seq(Literal(9), Literal(9))) + }, + condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + parameters = Map( + "functionName" -> "`current_timestamp`", + "expectedNum" -> "[0, 1]", + "actualNum" -> "2", + "docroot" -> SPARK_DOC_ROOT)) + } + + test("SPARK-57837: LocalTimestampExpressionBuilder") { + assert(LocalTimestampExpressionBuilder.build("localtimestamp", Seq.empty) === LocalTimestamp()) + assert(LocalTimestampExpressionBuilder.build("localtimestamp", Seq(Literal(6))) === + LocalTimestamp()) + + Seq(7, 8, 9).foreach { p => + val built = LocalTimestampExpressionBuilder.build("localtimestamp", Seq(Literal(p))) + assert(built === LocalTimestampNanos(p)) + assert(built.dataType === TimestampNTZNanosType(p)) + } + + Seq(0, 3, 5, 10).foreach { p => + checkError( + exception = intercept[SparkException] { + LocalTimestampExpressionBuilder.build("localtimestamp", Seq(Literal(p))) + }, + condition = "INVALID_TIMESTAMP_PRECISION", + parameters = Map("precision" -> p.toString, "type" -> "TIMESTAMP_NTZ")) + } + + checkError( + exception = intercept[AnalysisException] { + LocalTimestampExpressionBuilder.build("localtimestamp", Seq(Literal(9), Literal(9))) + }, + condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + parameters = Map( + "functionName" -> "`localtimestamp`", + "expectedNum" -> "[0, 1]", + "actualNum" -> "2", + "docroot" -> SPARK_DOC_ROOT)) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruningSubquerySuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruningSubquerySuite.scala index 614a29c5ac4a0..61700aa0d8604 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruningSubquerySuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/DynamicPruningSubquerySuite.scala @@ -17,8 +17,12 @@ package org.apache.spark.sql.catalyst.expressions +import java.lang.reflect.Modifier + import org.apache.spark.SparkFunSuite -import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} +import org.apache.spark.sql.catalyst.optimizer.ReusableBroadcastValueProjection +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, LocalRelation, Project} import org.apache.spark.sql.types.IntegerType class DynamicPruningSubquerySuite extends SparkFunSuite { @@ -31,59 +35,59 @@ class DynamicPruningSubquerySuite extends SparkFunSuite { buildKeys = Seq(pruningKeyExpression), broadcastKeyIndices = Seq(0), onlyInBroadcast = false - ) + )() test("pruningKey data type matches single buildKey") { val dynamicPruningSubquery = validDynamicPruningSubquery - .copy(buildKeys = Seq(Literal(2023))) + .copy(buildKeys = Seq(Literal(2023)))(None) assert(dynamicPruningSubquery.resolved == true) } test("pruningKey data type is a Struct and matches with Struct buildKey") { val dynamicPruningSubquery = validDynamicPruningSubquery .copy(pruningKey = CreateStruct(Seq(Literal(1), Literal.FalseLiteral)), - buildKeys = Seq(CreateStruct(Seq(Literal(2), Literal.TrueLiteral)))) + buildKeys = Seq(CreateStruct(Seq(Literal(2), Literal.TrueLiteral))))(None) assert(dynamicPruningSubquery.resolved == true) } test("multiple buildKeys but only one broadcastKeyIndex") { val dynamicPruningSubquery = validDynamicPruningSubquery .copy(buildKeys = Seq(Literal(0), Literal(2), Literal(0), Literal(9)), - broadcastKeyIndices = Seq(1)) + broadcastKeyIndices = Seq(1))(None) assert(dynamicPruningSubquery.resolved == true) } test("pruningKey data type does not match the single buildKey") { val dynamicPruningSubquery = validDynamicPruningSubquery.copy( pruningKey = Literal.TrueLiteral, - buildKeys = Seq(Literal(2013))) + buildKeys = Seq(Literal(2013)))(None) assert(dynamicPruningSubquery.resolved == false) } test("pruningKey data type is a Struct but mismatch with Struct buildKey") { val dynamicPruningSubquery = validDynamicPruningSubquery .copy(pruningKey = CreateStruct(Seq(Literal(1), Literal.FalseLiteral)), - buildKeys = Seq(CreateStruct(Seq(Literal.TrueLiteral, Literal(2))))) + buildKeys = Seq(CreateStruct(Seq(Literal.TrueLiteral, Literal(2)))))(None) assert(dynamicPruningSubquery.resolved == false) } test("DynamicPruningSubquery should only have a single broadcasting key") { val dynamicPruningSubquery = validDynamicPruningSubquery .copy(buildKeys = Seq(Literal(2025), Literal(2), Literal(1809)), - broadcastKeyIndices = Seq(0, 2)) + broadcastKeyIndices = Seq(0, 2))(None) assert(dynamicPruningSubquery.resolved == false) } test("duplicates in broadcastKeyIndices, and also should not be allowed") { val dynamicPruningSubquery = validDynamicPruningSubquery .copy(buildKeys = Seq(Literal(2)), - broadcastKeyIndices = Seq(0, 0)) + broadcastKeyIndices = Seq(0, 0))(None) assert(dynamicPruningSubquery.resolved == false) } test("broadcastKeyIndex out of bounds") { val dynamicPruningSubquery = validDynamicPruningSubquery - .copy(broadcastKeyIndices = Seq(1)) + .copy(broadcastKeyIndices = Seq(1))(None) assert(dynamicPruningSubquery.resolved == false) } @@ -98,17 +102,88 @@ class DynamicPruningSubquerySuite extends SparkFunSuite { buildQuery = LocalRelation(attr1), buildKeys = Seq(attr1), broadcastKeyIndices = Seq(0), - onlyInBroadcast = false) + onlyInBroadcast = false)() val dpq2 = DynamicPruningSubquery( pruningKey = Literal(1), buildQuery = LocalRelation(attr2), buildKeys = Seq(attr2), broadcastKeyIndices = Seq(0), - onlyInBroadcast = false) + onlyInBroadcast = false)() assert(dpq1.canonicalized == dpq2.canonicalized, "DynamicPruningSubquery with identical build queries but different ExprIds " + "must produce identical canonicalized forms so PlanMerger can deduplicate them") } + + test("transient broadcast value projection survives Catalyst copies without changing identity") { + val key = AttributeReference("key", IntegerType)() + val source = LocalRelation(key) + val projection = BroadcastValueProjection(source, Seq(key), key) + val pruning = DynamicPruningSubquery( + key, source, Seq(key), Seq(0), onlyInBroadcast = true)(Some(projection)) + + assert(pruning.resolved) + assert(pruning.usableBroadcastValueProjection.contains(projection)) + assert(pruning.productArity === 7) + assert(DynamicPruningSubquery.unapply(pruning).exists(_.productArity == 7)) + assert(Modifier.isTransient( + classOf[DynamicPruningSubquery].getDeclaredField("broadcastValueProjection").getModifiers)) + + Seq( + pruning.copy()(pruning.broadcastValueProjection), + pruning.withNewPlan(source), + pruning.withNewOuterAttrs(Seq(key)), + pruning.withNewHint(None).asInstanceOf[DynamicPruningSubquery], + pruning.withNewChildren(Seq(key)).asInstanceOf[DynamicPruningSubquery], + pruning.makeCopy(pruning.productIterator.map(_.asInstanceOf[AnyRef]).toArray) + .asInstanceOf[DynamicPruningSubquery] + ).foreach { rewritten => + assert(rewritten.broadcastValueProjection.contains(projection)) + } + + val unprojected = pruning.copy()(None) + assert(pruning === unprojected) + assert(pruning.canonicalized.asInstanceOf[DynamicPruningSubquery] + .broadcastValueProjection.isEmpty) + assert(pruning.canonicalized === unprojected.canonicalized) + + val missing = AttributeReference("missing", IntegerType)() + Seq( + projection.copy(sourceHashKeys = Seq(missing)), + projection.copy(valueExpression = missing), + projection.copy(sourceHashKeys = Seq.empty), + projection.copy(valueExpression = Literal(1L)) + ).foreach { invalidProjection => + val rewritten = pruning.copy()(Some(invalidProjection)) + assert(rewritten.resolved) + assert(rewritten.broadcastValueProjection.contains(invalidProjection)) + assert(rewritten.usableBroadcastValueProjection.isEmpty) + } + } + + test("extract broadcast hash keys without rejecting residual join predicates") { + val leftKey = AttributeReference("left_key", IntegerType)() + val leftRegion = AttributeReference("left_region", IntegerType)() + val leftValue = AttributeReference("left_value", IntegerType)() + val rightKey = AttributeReference("right_key", IntegerType)() + val rightRegion = AttributeReference("right_region", IntegerType)() + val rightValue = AttributeReference("right_value", IntegerType)() + val left = LocalRelation(leftKey, leftRegion, leftValue) + val right = LocalRelation(rightKey, rightRegion, rightValue) + val excluded = LocalRelation(AttributeReference("excluded", IntegerType)()) + val residual = LessThanOrEqual(leftValue, rightValue) + val condition = And( + And(EqualTo(leftKey, rightKey), EqualTo(rightRegion, leftRegion)), residual) + val join = Join(left, right, Inner, Some(condition), JoinHint.NONE) + + assert(ReusableBroadcastValueProjection.find(leftValue, join, excluded).contains( + BroadcastValueProjection(left, Seq(leftKey, leftRegion), leftValue))) + assert(ReusableBroadcastValueProjection.find(rightValue, join, excluded).contains( + BroadcastValueProjection(right, Seq(rightKey, rightRegion), rightValue))) + + val nullSafeJoin = Join( + left, right, Inner, Some(And(EqualNullSafe(leftKey, rightKey), residual)), JoinHint.NONE) + assert(ReusableBroadcastValueProjection.find(leftValue, nullSafeJoin, excluded).isEmpty) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExprUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExprUtilsSuite.scala new file mode 100644 index 0000000000000..c0ac407b04322 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExprUtilsSuite.scala @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.expressions + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.types.{ArrayType, IntegerType, LongType, StringType, StructField, StructType} + +class ExprUtilsSuite extends SparkFunSuite { + + private val a = AttributeReference("a", LongType)() + + test("canEvaluateUnconditionally: whitelisted total expressions") { + assert(ExprUtils.canEvaluateUnconditionally(a)) + assert(ExprUtils.canEvaluateUnconditionally(Literal(1L))) + assert(ExprUtils.canEvaluateUnconditionally( + And(LessThan(a, Literal(5L)), IsNotNull(a)))) + assert(ExprUtils.canEvaluateUnconditionally( + GetStructField( + AttributeReference("s", StructType(StructField("f1", IntegerType) :: Nil))(), 0))) + assert(ExprUtils.canEvaluateUnconditionally(Coalesce(Seq(a, Literal(0L))))) + assert(ExprUtils.canEvaluateUnconditionally(In(a, Seq(Literal(1L), Literal(2L))))) + } + + test("canEvaluateUnconditionally: expressions that can throw are excluded") { + // Arithmetic can throw (overflow, div-by-zero in ANSI mode) even though it does not + // override `throwable` and inherits false from its non-throwing children, which is + // why the whitelist is used instead of the throwable flag. + val arithmetic = LessThan(Add(a, Literal(1L)), Literal(5L)) + assert(!arithmetic.throwable) + assert(!ExprUtils.canEvaluateUnconditionally(arithmetic)) + assert(!ExprUtils.canEvaluateUnconditionally(EqualTo(Remainder(a, Literal(3L)), Literal(0L)))) + assert(!ExprUtils.canEvaluateUnconditionally(Cast(a, StringType))) + // GetArrayItem/ElementAt can throw on invalid ordinals in ANSI mode. + val arr = AttributeReference("arr", ArrayType(IntegerType))() + assert(!ExprUtils.canEvaluateUnconditionally(GetArrayItem(arr, Literal(0)))) + // A whitelisted predicate over a non-whitelisted child is still excluded. + assert(!ExprUtils.canEvaluateUnconditionally(IsNotNull(Add(a, Literal(1L))))) + } + + test("canEvaluateUnconditionally: non-deterministic expressions are excluded") { + assert(!ExprUtils.canEvaluateUnconditionally(LessThan(Rand(Literal(0L)), Literal(0.5)))) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelperSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelperSuite.scala index 3cc50da38906e..abee42e260e0b 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelperSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionEvalHelperSuite.scala @@ -45,9 +45,15 @@ class ExpressionEvalHelperSuite extends SparkFunSuite with ExpressionEvalHelper } test("SPARK-33619: make sure checkExceptionInExpression work as expected") { + // Assert on the unresolved type name rather than a compiler-specific phrasing: the + // codegen path's message depends on the configured compiler backend (Janino reports + // `Cannot determine simple type name "NoSuchElementException"`, the JDK compiler + // reports `cannot find symbol ... class NoSuchElementException`), while the + // non-codegen path throws the message defined in BadCodegenAndEvalExpression.eval. + // The type name is the common, meaningful token across all paths. checkExceptionInExpression[Exception]( BadCodegenAndEvalExpression(), - "Cannot determine simple type name \"NoSuchElementException\"") + "NoSuchElementException") } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtilsSuite.scala index e3a9177bfdbcb..608a13d47c594 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtilsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtilsSuite.scala @@ -480,4 +480,46 @@ class ExpressionImplUtilsSuite extends SparkFunSuite { tryValidateUTF8(UTF8String.fromBytes(Array[Byte](0xFF.toByte)), null) } + // These vectors follow Unicode Standard Annex #15's decomposition/composition algorithm and + // pin normalize() to Spark's bundled ICU4J/Unicode data (see icu4j.version in pom.xml): a + // future ICU4J upgrade that changed these mappings would fail this test, rather than silently + // changing query results. + test("Normalize with supported forms") { + def normalize(input: String, form: String): String = + ExpressionImplUtils.normalize( + UTF8String.fromString(input), UTF8String.fromString(form)).toString + + // scalastyle:off nonascii + assert(normalize("\uFB01", "NFKC") == "fi") + assert(normalize("A\u030A", "NFC") == "\u00C5") + assert(normalize("\u00C5", "NFD") == "A\u030A") + assert(normalize("\uFB01", "nfkc") == "fi") + assert(normalize("\u00BD", "NFKD") == "1" + "\u2044" + "2") + assert(normalize("\u00BD", "NFD") == "\u00BD") + assert(normalize("\uD835\uDC00", "NFKC") == "A") + assert(normalize("", "NFC") == "") + + // Combining marks with different combining classes are canonically reordered before + // composition (UAX #15 canonical ordering algorithm), so applying COMBINING DOT BELOW + // (U+0323) and COMBINING ACUTE ACCENT (U+0301) to the same base letter in either input + // order must normalize (NFC) to the same result. + val dotBelow = "\u0323" + val acute = "\u0301" + assert(normalize("a" + dotBelow + acute, "NFC") == normalize("a" + acute + dotBelow, "NFC")) + // scalastyle:on nonascii + } + + test("Normalize invalid form error") { + checkError( + exception = intercept[SparkRuntimeException] { + ExpressionImplUtils.normalize( + UTF8String.fromString("abc"), UTF8String.fromString("NFE")) + }, + condition = "INVALID_PARAMETER_VALUE.NORMALIZE_FORM", + parameters = Map( + "parameter" -> "`form`", + "functionName" -> "`normalize`", + "form" -> "'NFE'")) + } + } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/GeneratorExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/GeneratorExpressionSuite.scala index b6a3d61cb13a6..b9a245eb3ec35 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/GeneratorExpressionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/GeneratorExpressionSuite.scala @@ -21,6 +21,7 @@ import org.apache.spark.{SPARK_DOC_ROOT, SparkFunSuite} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.types._ class GeneratorExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { @@ -114,4 +115,106 @@ class GeneratorExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { ) ) } + + test("unnest - eval is lazy and only reads elements as rows are pulled") { + // Backing array that records which ordinals were read, to prove that pulling the first N rows + // touches only the first N elements rather than materializing the whole expansion up front. + val readOrdinals = scala.collection.mutable.ArrayBuffer.empty[Int] + val tracking = new GenericArrayData(Array[Any](10, 20, 30, 40, 50)) { + override def get(ordinal: Int, elementType: DataType): AnyRef = { + readOrdinals += ordinal + super.get(ordinal, elementType) + } + } + val result = Unnest(Seq(Literal(tracking, ArrayType(IntegerType))), withOrdinality = true) + .eval(null) + // The returned value is a lazy Iterator, not an eagerly materialized collection, and building + // it must not read any element. + assert(result.isInstanceOf[Iterator[_]]) + // Literal's constructor validates its value by reading element 0; ignore reads made before the + // iterator is created and observe only what pulling rows drives. + readOrdinals.clear() + + val it = result.iterator + assert(readOrdinals.isEmpty, "no element should be read before the iterator is advanced") + assert(it.next() === create_row(10, 1L)) + assert(it.next() === create_row(20, 2L)) + // Only the two consumed rows' elements were read; rows 2..4 remain untouched. + assert(readOrdinals.toSeq === Seq(0, 1)) + } + + test("unnest - single array") { + checkTuple(Unnest(Seq(empty_array), withOrdinality = false), Seq.empty) + checkTuple( + Unnest(Seq(int_array), withOrdinality = false), + Seq(create_row(1), create_row(2), create_row(3))) + // A null array is treated as empty and contributes no rows. + checkTuple( + Unnest(Seq(Literal.create(null, ArrayType(IntegerType))), withOrdinality = false), + Seq.empty) + } + + test("unnest - single column naming and ordinality") { + // With a single array the output column keeps the default name `col`. + assert(Unnest(Seq(int_array), withOrdinality = false).elementSchema === + new StructType().add("col", IntegerType, nullable = false)) + // WITH ORDINALITY appends a 1-based, non-nullable bigint column. + assert(Unnest(Seq(int_array), withOrdinality = true).elementSchema === + new StructType() + .add("col", IntegerType, nullable = false) + .add("ordinality", LongType, nullable = false)) + checkTuple( + Unnest(Seq(str_array), withOrdinality = true), + Seq(create_row("a", 1L), create_row("b", 2L), create_row("c", 3L))) + } + + test("unnest - multiple arrays are zipped and padded with nulls") { + val short_array = CreateArray(Seq(10, 20).map(Literal(_))) + // With several arrays the columns are named positionally and padded columns are nullable. + assert(Unnest(Seq(int_array, short_array), withOrdinality = false).elementSchema === + new StructType() + .add("col0", IntegerType, nullable = true) + .add("col1", IntegerType, nullable = true)) + checkTuple( + Unnest(Seq(int_array, short_array), withOrdinality = false), + Seq(create_row(1, 10), create_row(2, 20), create_row(3, null))) + // WITH ORDINALITY spans the full (longest) length. + checkTuple( + Unnest(Seq(int_array, short_array), withOrdinality = true), + Seq(create_row(1, 10, 1L), create_row(2, 20, 2L), create_row(3, null, 3L))) + } + + test("unnest - type checks") { + assert(Unnest(Seq(int_array), withOrdinality = false).checkInputDataTypes().isSuccess) + + // Providing no arguments is rejected. + checkError( + exception = intercept[AnalysisException] { + Unnest(Seq.empty, withOrdinality = false).checkInputDataTypes() + }, + condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + parameters = Map( + "functionName" -> "`unnest`", + "expectedNum" -> "> 0", + "actualNum" -> "0", + "docroot" -> SPARK_DOC_ROOT)) + + // A non-array argument is rejected, reporting the offending 1-based parameter index. + assert(Unnest(Seq(int_array, Literal(3)), withOrdinality = false).checkInputDataTypes() == + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> "second", + "requiredType" -> "\"ARRAY\"", + "inputSql" -> "\"3\"", + "inputType" -> "\"INT\""))) + } + + test("unnest - string representation hides the ordinality flag") { + // The `withOrdinality` boolean must not leak into plan output as a bare `true` argument; it is + // rendered as a readable `WITH ORDINALITY` suffix instead (see the EXPLAIN golden results). + assert(Unnest(Seq(int_array), withOrdinality = false).toString === "unnest(array(1, 2, 3))") + assert(Unnest(Seq(int_array), withOrdinality = true).toString === + "unnest(array(1, 2, 3), WITH ORDINALITY)") + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HashExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HashExpressionsSuite.scala index 02846b560ad00..d03db29257580 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HashExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HashExpressionsSuite.scala @@ -92,6 +92,17 @@ class HashExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { checkConsistencyBetweenInterpretedAndCodegen(Crc32, BinaryType) } + test("xxh3_64 and xxh3_128") { + // Concrete values match the reference implementation; XXH3Suite covers the full vector set. + checkEvaluation(Xxh364(Literal("Spark".getBytes(StandardCharsets.UTF_8))), 80997306238743657L) + checkEvaluation(Xxh3128(Literal("Spark".getBytes(StandardCharsets.UTF_8))), + "7d57dd84c60c86ca1f4e82ab91a12b5e") + checkEvaluation(Xxh364(Literal.create(null, BinaryType)), null) + checkEvaluation(Xxh3128(Literal.create(null, BinaryType)), null) + checkConsistencyBetweenInterpretedAndCodegen(Xxh364, BinaryType) + checkConsistencyBetweenInterpretedAndCodegen(Xxh3128, BinaryType) + } + def checkHiveHash(input: Any, dataType: DataType, expected: Long): Unit = { // Note : All expected hashes need to be computed using Hive 1.2.1 val actual = HiveHashFunction.hash( diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala index c33d258ac4de6..b9c870f310906 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/HigherOrderFunctionsSuite.scala @@ -940,6 +940,16 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper "actualType" -> toSQLType(StringType) ))) } + + test("NamedLambdaVariable is stateful and produces a fresh copy") { + val lv = NamedLambdaVariable("x", IntegerType, nullable = false) + assert(lv.stateful, "NamedLambdaVariable.stateful should be true") + val copy = lv.freshCopyIfContainsStatefulExpression() + assert(copy ne lv, + "freshCopyIfContainsStatefulExpression should return a new instance for NamedLambdaVariable") + assert(copy.asInstanceOf[NamedLambdaVariable].value ne lv.value, + "fresh copy should have an independent AtomicReference value") + } } case class CodegenFallbackExpr(child: Expression) extends UnaryExpression with CodegenFallback { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala index 37916f5a93be0..d8d68ef75c10a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/JsonExpressionsSuite.scala @@ -909,6 +909,33 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { } } + test("json_typeof") { + Seq( + // Invalid or empty inputs return null. + ("", null), + ("bad", null), + ("""{"key": 45, "random_string"}""", null), + // Trailing content after a valid value is not a single well-formed JSON document. + ("123 true", null), + // Valid JSON values return the type of the outermost value. + ("{}", "object"), + ("""{"key": 1, "arr": [1, 2]}""", "object"), + ("[]", "array"), + ("[1, 2, 3]", "array"), + ("\"hello\"", "string"), + ("123", "number"), + ("1.5", "number"), + ("-123", "number"), + ("-1.5", "number"), + ("true", "boolean"), + ("false", "boolean"), + ("null", "null") + ).foreach { + case (input, expected) => + checkEvaluation(JsonTypeof(Literal(input)), expected) + } + } + test("SPARK-35320: from_json should fail with a key type different of StringType") { Seq( (MapType(IntegerType, StringType), """{"1": "test"}"""), @@ -1041,4 +1068,31 @@ class JsonExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { input) } + test("JsonToStructs, GetJsonObject, JsonTuple, MultiGetJsonObject, JsonValue are stateful " + + "and produce fresh copies") { + val schema = StructType(StructField("a", IntegerType) :: Nil) + val jsonToStructs = JsonToStructs(schema, Map.empty, Literal("{}"), UTC_OPT) + assert(jsonToStructs.stateful) + assert(jsonToStructs.freshCopyIfContainsStatefulExpression() ne jsonToStructs) + + val getJsonObject = GetJsonObject(Literal("{}"), Literal("$.a")) + assert(getJsonObject.stateful) + assert(getJsonObject.freshCopyIfContainsStatefulExpression() ne getJsonObject) + + val jsonTuple = JsonTuple(Literal("{}") :: Literal("a") :: Nil) + assert(jsonTuple.stateful) + assert(jsonTuple.freshCopyIfContainsStatefulExpression() ne jsonTuple) + + val multiGetJsonObject = MultiGetJsonObject(Literal("{}"), Seq("$.a", "$.b")) + assert(multiGetJsonObject.stateful) + assert(multiGetJsonObject.freshCopyIfContainsStatefulExpression() ne multiGetJsonObject) + + // JsonValue reuses a mutable row to cast the extracted scalar, so it must be stateful. + val jsonValue = JsonValue( + Literal("{}"), "$.a", StringType, + JsonValueBehavior.Null, JsonValueBehavior.Null, None, None) + assert(jsonValue.stateful) + assert(jsonValue.freshCopyIfContainsStatefulExpression() ne jsonValue) + } + } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala index 0ea3bc77ec1cb..90148932237c5 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MathExpressionsSuite.scala @@ -727,6 +727,27 @@ class MathExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { checkConsistencyBetweenInterpretedAndCodegen(Logarithm, DoubleType, DoubleType) } + test("truncate") { + // Truncation toward zero, distinct from floor/ceil (toward -inf/+inf) and round (to nearest). + checkEvaluation(Truncate(Literal(Decimal(BigDecimal("1234.5678"))), Literal(2)), + Decimal(BigDecimal("1234.56"))) + checkEvaluation(Truncate(Literal(Decimal(BigDecimal("-1234.5678"))), Literal(2)), + Decimal(BigDecimal("-1234.56"))) + checkEvaluation(Truncate(Literal(Decimal(BigDecimal("1234.5678"))), Literal(-2)), + Decimal(BigDecimal("1200"))) + checkEvaluation(Truncate(Literal(Decimal(BigDecimal("-3.99"))), Literal(0)), + Decimal(BigDecimal("-3"))) + // Default scale is 0. + checkEvaluation(new Truncate(Literal(Decimal(BigDecimal("3.99")))), Decimal(BigDecimal("3"))) + // Integral input with negative scale. + checkEvaluation(Truncate(Literal(125), Literal(-1)), 120) + // Double input. + checkEvaluation(Truncate(Literal(3.1415926), Literal(3)), 3.141) + // Null propagation. + checkEvaluation(Truncate(Literal.create(null, DoubleType), Literal(2)), null) + checkEvaluation(Truncate(Literal(1.23), Literal.create(null, IntegerType)), null) + } + test("round/bround/floor/ceil") { val scales = -6 to 6 val doublePi: Double = math.Pi diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala index 0327624709001..a586e33afdd26 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/MiscExpressionsSuite.scala @@ -55,6 +55,10 @@ class MiscExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { ) } + test("SPARK-58627: RaiseError is throwable") { + assert(RaiseError(Literal("error!")).throwable) + } + test("SPARK-55109: RaiseError.sql uses single-argument form only for known error classes") { assert(RaiseError(Literal("error!")).sql === "raise_error('error!')") diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala index 215362c47b940..f41f3e7df1b56 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ObjectExpressionsSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.catalyst.expressions import java.sql.{Date, Timestamp} +import java.util.concurrent.ExecutionException import scala.collection.immutable import scala.collection.mutable @@ -37,6 +38,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, Genera import org.apache.spark.sql.catalyst.expressions.objects._ import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, DateTimeUtils, GenericArrayData, IntervalUtils} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -245,18 +247,47 @@ class ObjectExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { val initializeWithNonexistingMethod = InitializeJavaBean( Literal.fromObject(new java.util.LinkedList[Int]), Map("nonexistent" -> Literal(1))) - checkExceptionInExpression[Exception](initializeWithNonexistingMethod, - """A method named "nonexistent" is not declared in any enclosing class """ + - "nor any supertype") + // The two evaluation paths report this differently: interpreted execution raises Spark's own + // INTERNAL_ERROR from `methodNotDeclaredError`, while codegen fails when the Java compiler + // resolves the setter call in the generated source. Assert each on its own rather than + // through one shared substring. + checkError( + exception = intercept[SparkException] { + evaluateWithoutCodegen(initializeWithNonexistingMethod, InternalRow.fromSeq(Seq())) + }, + condition = "INTERNAL_ERROR", + parameters = Map("message" -> + ("""A method named "nonexistent" is not declared in any enclosing class """ + + "nor any supertype")), + sqlState = "XX000") + withSQLConf( + SQLConf.CODEGEN_FACTORY_MODE.key -> CodegenObjectFactoryMode.CODEGEN_ONLY.toString) { + // What Spark owns on this path is the wrapping - an ExecutionException around the + // CompileException that `compilerError` builds - and the "Failed to compile: " prefix from + // `failedToCompileMsg`. The diagnostic after the prefix is the compiler's own wording, so + // the setter name below is only a sanity check that it names the offending member: the + // name is also the input, and Spark echoes whole expressions into other errors, so on its + // own it would match unrelated failures too. + val errMsg = intercept[ExecutionException] { + evaluateWithMutableProjection(initializeWithNonexistingMethod, InternalRow.fromSeq(Seq())) + }.getMessage + assert(errMsg.contains("Failed to compile:")) + assert(errMsg.contains("nonexistent")) + } val initializeWithWrongParamType = InitializeJavaBean( Literal.fromObject(new TestBean), Map("setX" -> Literal("1"))) - intercept[Exception] { - evaluateWithoutCodegen(initializeWithWrongParamType, InternalRow.fromSeq(Seq())) - }.getMessage.contains( - """A method named "setX" is not declared in any enclosing class """ + - "nor any supertype") + // Same error as above, reached only through interpreted execution. + checkError( + exception = intercept[SparkException] { + evaluateWithoutCodegen(initializeWithWrongParamType, InternalRow.fromSeq(Seq())) + }, + condition = "INTERNAL_ERROR", + parameters = Map("message" -> + ("""A method named "setX" is not declared in any enclosing class """ + + "nor any supertype")), + sqlState = "XX000") } test("InitializeJavaBean doesn't call setters if input in null") { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala index 9e6b59b51138d..858c6e3b5028a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala @@ -64,4 +64,22 @@ class RandomSuite extends SparkFunSuite with ExpressionEvalHelper { testUniform(10.0F, 20.0F, 17.604954F) testUniform(10L, 20.0F, 17.604954F) } + + test("SPARK-58208: Uniform preserves its time zone when copied") { + val uniform = Uniform( + Literal(10), Literal(20), Literal(0), hideSeed = false, timeZoneId = Some("UTC")) + assert(uniform.resolved) + + val copied = uniform.freshCopyIfContainsStatefulExpression().asInstanceOf[Uniform] + assert(copied ne uniform) + assert(copied.timeZoneId == uniform.timeZoneId) + assert(copied.resolved) + + Seq(uniform.withNewSeed(1), uniform.withShiftedSeed(1)).foreach { + case copied: Uniform => + assert(copied.timeZoneId == uniform.timeZoneId) + assert(copied.resolved) + case other => fail(s"Expected Uniform but got ${other.getClass.getName}") + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RegexpExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RegexpExpressionsSuite.scala index 0bf29553ea33d..5377b9c7fc0b9 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RegexpExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RegexpExpressionsSuite.scala @@ -711,4 +711,28 @@ class RegexpExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { ) ) } + + test("RegExpReplace and RegExpExtractBase are stateful and produce fresh copies") { + val s = Literal("hello world") + val p = Literal("(\\w+)") + val r = Literal("X") + + val replace = RegExpReplace(s, p, r) + assert(replace.stateful, "RegExpReplace.stateful should be true") + val replaceCopy = replace.freshCopyIfContainsStatefulExpression() + assert(replaceCopy ne replace, + "freshCopyIfContainsStatefulExpression should return a new instance for RegExpReplace") + + val extract = RegExpExtract(s, p, Literal(1)) + assert(extract.stateful, "RegExpExtract.stateful should be true") + val extractCopy = extract.freshCopyIfContainsStatefulExpression() + assert(extractCopy ne extract, + "freshCopyIfContainsStatefulExpression should return a new instance for RegExpExtract") + + val extractAll = RegExpExtractAll(s, p, Literal(1)) + assert(extractAll.stateful, "RegExpExtractAll.stateful should be true") + val extractAllCopy = extractAll.freshCopyIfContainsStatefulExpression() + assert(extractAllCopy ne extractAll, + "freshCopyIfContainsStatefulExpression should return a new instance for RegExpExtractAll") + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SelectedFieldSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SelectedFieldSuite.scala index ddeac6cbb933d..8db237e1a68cb 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SelectedFieldSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SelectedFieldSuite.scala @@ -524,6 +524,43 @@ class SelectedFieldSuite extends AnalysisTest { StructField("subfield4", IntegerType) :: Nil)) :: Nil), containsNull = false))) } + // |-- col1: string (nullable = false) + // |-- col2: struct (nullable = true) (metadata = {"outer":"meta"}) + // | |-- field1: long (nullable = true) (metadata = {"inner":"meta"}) + // | |-- field2: long (nullable = true) + private val outerMetadata = new MetadataBuilder().putString("outer", "meta").build() + private val innerMetadata = new MetadataBuilder().putString("inner", "meta").build() + private val structWithMetadata = StructType(ignoredField :: + StructField("col2", StructType( + StructField("field1", LongType, nullable = true, innerMetadata) :: + StructField("field2", LongType) :: Nil), nullable = true, outerMetadata) :: Nil) + + testSelect(structWithMetadata, "col2.field1") { + StructField("col2", StructType( + StructField("field1", LongType, nullable = true, innerMetadata) :: Nil), + nullable = true, outerMetadata) + } + + // |-- col1: string (nullable = false) + // |-- col2: struct (nullable = true) + // | |-- field3: array (nullable = false) + // | | |-- element: struct (containsNull = true) + // | | | |-- subfield1: long (nullable = true) (metadata = {"elem":"meta"}) + // | | | |-- subfield2: long (nullable = true) + private val elementMetadata = new MetadataBuilder().putString("elem", "meta").build() + private val arrayOfStructWithMetadata = StructType(ignoredField :: + StructField("col2", StructType( + StructField("field3", ArrayType(StructType( + StructField("subfield1", LongType, nullable = true, elementMetadata) :: + StructField("subfield2", LongType) :: Nil)), nullable = false) :: Nil)) :: Nil) + + testSelect(arrayOfStructWithMetadata, "col2.field3.subfield1") { + StructField("col2", StructType( + StructField("field3", ArrayType(StructType( + StructField("subfield1", LongType, nullable = true, elementMetadata) :: Nil)), + nullable = false) :: Nil)) + } + def assertResult(expected: StructField)(actual: StructField)(selectExpr: String): Unit = { try { super.assertResult(expected)(actual) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala index 711b5edd72ad1..0cb944192cb8c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.expressions import java.math.{BigDecimal => JavaBigDecimal} -import org.apache.spark.{SPARK_DOC_ROOT, SparkFunSuite, SparkIllegalArgumentException} +import org.apache.spark.{SPARK_DOC_ROOT, SparkFunSuite, SparkIllegalArgumentException, SparkRuntimeException} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.analysis.TypeCheckResult import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, InvalidFormat} @@ -451,6 +451,33 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { } } + test("SPARK-48973: Mask with supplementary characters") { + def cp(codePoint: Int): String = new String(Character.toChars(codePoint)) + val smile = cp(0x1F642) + val boldA = cp(0x1D400) + val boldSmallA = cp(0x1D41A) + val boldZero = cp(0x1D7CE) + + checkEvaluation( + new Mask(Literal(smile), Literal('Y'), Literal('y'), Literal('n'), Literal('*')), "*") + checkEvaluation(new Mask(Literal("ABC"), Literal(smile)), smile * 3) + checkEvaluation(new Mask(Literal(s"A$boldA 1$boldZero")), "XX nn") + // Supplementary upper-case, lower-case and digit characters are each classified and + // replaced like their BMP counterparts. + checkEvaluation(new Mask(Literal(s"$boldA$boldSmallA$boldZero")), "Xxn") + // A supplementary replacement applied to a supplementary input. + checkEvaluation(new Mask(Literal(boldSmallA), Literal('Y'), Literal(smile)), smile) + + // A supplementary character must round-trip intact through the retain path, both when it + // falls into the otherChar category and when its own category is set to retain. + checkEvaluation(new Mask(Literal(smile)), smile) + checkEvaluation(new Mask(Literal(s"a${smile}1")), s"x${smile}n") + checkEvaluation(new Mask(Literal(boldA), Literal(null, StringType)), boldA) + checkEvaluation( + new Mask(Literal(boldZero), Literal('Y'), Literal('y'), Literal(null, StringType)), + boldZero) + } + test("SPARK-42384: Mask with null input") { val NULL_LITERAL = Literal(null, StringType) checkEvaluation( @@ -499,6 +526,30 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { checkEvaluation(UnBase64(Literal.create(null, StringType)), null, create_row("abdef")) } + test("to_base32/from_base32 for string") { + // RFC 4648 (section 10) test vectors. + checkEvaluation(Base32(Literal("".getBytes("UTF-8"))), "") + checkEvaluation(Base32(Literal("f".getBytes("UTF-8"))), "MY======") + checkEvaluation(Base32(Literal("fo".getBytes("UTF-8"))), "MZXQ====") + checkEvaluation(Base32(Literal("foo".getBytes("UTF-8"))), "MZXW6===") + checkEvaluation(Base32(Literal("foob".getBytes("UTF-8"))), "MZXW6YQ=") + checkEvaluation(Base32(Literal("fooba".getBytes("UTF-8"))), "MZXW6YTB") + checkEvaluation(Base32(Literal("foobar".getBytes("UTF-8"))), "MZXW6YTBOI======") + + assert(!Base32(Literal("foo".getBytes("UTF-8"))).nullable) + assert(Base32(Literal.create(null, BinaryType)).nullable) + assert(!UnBase32(Literal("MZXW6YTBOI======")).nullable) + assert(UnBase32(Literal.create(null, StringType)).nullable) + + checkEvaluation(UnBase32(Literal("MZXW6YTBOI======")), "foobar".getBytes("UTF-8")) + checkEvaluation(UnBase32(Literal("MY======")), "f".getBytes("UTF-8")) + + // Round trip. + checkEvaluation(Base32(UnBase32(Literal("MZXW6YTBOI======"))), "MZXW6YTBOI======") + checkEvaluation(Base32(UnBase32(Literal(""))), "") + checkEvaluation(Base32(UnBase32(Literal.create(null, StringType))), null) + } + test("encode/decode for string") { val a = $"a".string.at(0) val b = $"b".binary.at(0) @@ -2268,4 +2319,44 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { } } } + + test("StringTranslate and FormatNumber are stateful and produce fresh copies") { + val src = Literal("aeiou") + val matching = Literal("aeiou") + val replace = Literal("12345") + val translate = StringTranslate(src, matching, replace) + assert(translate.stateful, "StringTranslate.stateful should be true") + val translateCopy = translate.freshCopyIfContainsStatefulExpression() + assert(translateCopy ne translate, + "freshCopyIfContainsStatefulExpression should return a new instance for StringTranslate") + + val num = Literal(1234567.89) + val fmt = Literal(2) + val formatNumber = FormatNumber(num, fmt) + assert(formatNumber.stateful, "FormatNumber.stateful should be true") + val formatNumberCopy = formatNumber.freshCopyIfContainsStatefulExpression() + assert(formatNumberCopy ne formatNumber, + "freshCopyIfContainsStatefulExpression should return a new instance for FormatNumber") + } + + test("Normalize") { + // scalastyle:off nonascii + checkEvaluation(new Normalize(Literal("A\u030A")), "\u00C5") + checkEvaluation(Normalize(Literal("A\u030A"), Literal("NFC")), "\u00C5") + checkEvaluation(Normalize(Literal("\u00C5"), Literal("NFD")), "A\u030A") + checkEvaluation(Normalize(Literal("\uFB01"), Literal("NFKC")), "fi") + // scalastyle:on nonascii + checkEvaluation(Normalize(Literal.create(null, StringType), Literal("NFC")), null) + checkEvaluation(Normalize(Literal("abc"), Literal.create(null, StringType)), null) + } + + test("Normalize invalid form") { + checkErrorInExpression[SparkRuntimeException]( + Normalize(Literal("abc"), Literal("NFE")), + "INVALID_PARAMETER_VALUE.NORMALIZE_FORM", + Map( + "parameter" -> "`form`", + "functionName" -> "`normalize`", + "form" -> "'NFE'")) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TimestampNanosRowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TimestampNanosRowSuite.scala index d85350504f7f5..cc29ae3269338 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TimestampNanosRowSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TimestampNanosRowSuite.scala @@ -20,9 +20,9 @@ package org.apache.spark.sql.catalyst.expressions import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection -import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, GenericArrayData} import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.TimestampNanosVal +import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String} import org.apache.spark.util.ArrayImplicits._ class TimestampNanosRowSuite extends SparkFunSuite with ExpressionEvalHelper { @@ -156,6 +156,96 @@ class TimestampNanosRowSuite extends SparkFunSuite with ExpressionEvalHelper { assert(arr.getTimestampNTZNanos(1) === ntzValue) } + // Struct fields are stored in the UnsafeRow variable-length region. Exercises + // GenerateUnsafeProjection.writeStructToBuffer/writeExpressionsToBuffer (codegen) and the + // interpreted equivalent, then reads the nanos fields back through UnsafeRow.getStruct. + testBothCodegenAndInterpreted("UnsafeRow nested struct with nanos timestamp fields") { + val innerStruct = StructType(Seq( + StructField("inner_ntz", TimestampNTZNanosType(9), nullable = true), + StructField("inner_ltz", TimestampLTZNanosType(7), nullable = true))) + val fieldTypes = Array[DataType](innerStruct) + val converter = UnsafeProjection.create(fieldTypes) + + val innerRow = new GenericInternalRow(Array[Any](ntzValue, ltzValue)) + val input = new GenericInternalRow(Array[Any](innerRow)) + + val unsafeRow = converter.apply(input) + val nestedStruct = unsafeRow.getStruct(0, 2) + assert(nestedStruct.getTimestampNTZNanos(0) === ntzValue) + assert(nestedStruct.getTimestampLTZNanos(1) === ltzValue) + + // Null nanos fields inside a non-null struct must be marked null after serialization. + val innerRowNulls = new GenericInternalRow(Array[Any](null, null)) + val inputNulls = new GenericInternalRow(Array[Any](innerRowNulls)) + val unsafeRowNulls = converter.apply(inputNulls) + val nestedStructNulls = unsafeRowNulls.getStruct(0, 2) + assert(nestedStructNulls.isNullAt(0)) + assert(nestedStructNulls.isNullAt(1)) + } + + // A nanos field alongside a variable-length StringType field checks that the nanos composite + // (epochMicros Long + nanosWithinMicro Short) is written at the correct offset when other + // variable-length fields share the nested struct. + testBothCodegenAndInterpreted("UnsafeRow nested struct with nanos and non-nanos fields") { + val innerStruct = StructType(Seq( + StructField("ntz", TimestampNTZNanosType(8), nullable = true), + StructField("name", StringType, nullable = true))) + val fieldTypes = Array[DataType](innerStruct) + val converter = UnsafeProjection.create(fieldTypes) + + val innerRow = + new GenericInternalRow(Array[Any](ntzValue, UTF8String.fromString("test_string"))) + val input = new GenericInternalRow(Array[Any](innerRow)) + val unsafeRow = converter.apply(input) + + val nested = unsafeRow.getStruct(0, 2) + assert(nested.getTimestampNTZNanos(0) === ntzValue) + assert(nested.getUTF8String(1) === UTF8String.fromString("test_string")) + } + + // Map values are serialized as an UnsafeArrayData via GenerateUnsafeProjection.writeMapToBuffer + // delegating to writeArrayToBuffer. Reads back through UnsafeRow.getMap().valueArray(). + testBothCodegenAndInterpreted("UnsafeRow with map of nanos timestamp values") { + val mapType = MapType(StringType, TimestampNTZNanosType(9), valueContainsNull = true) + val fieldTypes = Array[DataType](mapType) + val converter = UnsafeProjection.create(fieldTypes) + + val keys = new GenericArrayData(Array[Any]( + UTF8String.fromString("a"), UTF8String.fromString("b"), UTF8String.fromString("c"))) + val values = new GenericArrayData(Array[Any](ntzValue, null, ntzValue)) + val mapData = new ArrayBasedMapData(keys, values) + val input = new GenericInternalRow(Array[Any](mapData)) + + val unsafeRow = converter.apply(input) + val unsafeMap = unsafeRow.getMap(0) + assert(unsafeMap.numElements() == 3) + val valueArray = unsafeMap.valueArray() + assert(valueArray.getTimestampNTZNanos(0) === ntzValue) + assert(valueArray.isNullAt(1)) + assert(valueArray.getTimestampNTZNanos(2) === ntzValue) + } + + // valueContainsNull = false exercises the codegen branch in writeArrayToBuffer that elides + // the per-element isNullAt check for map values. + testBothCodegenAndInterpreted("UnsafeRow with non-nullable map of nanos values") { + val mapType = MapType(StringType, TimestampLTZNanosType(7), valueContainsNull = false) + val fieldTypes = Array[DataType](mapType) + val converter = UnsafeProjection.create(fieldTypes) + + val keys = new GenericArrayData(Array[Any]( + UTF8String.fromString("x"), UTF8String.fromString("y"))) + val values = new GenericArrayData(Array[Any](ltzValue, ltzValue)) + val mapData = new ArrayBasedMapData(keys, values) + val input = new GenericInternalRow(Array[Any](mapData)) + + val unsafeRow = converter.apply(input) + val unsafeMap = unsafeRow.getMap(0) + assert(unsafeMap.numElements() == 2) + val valueArray = unsafeMap.valueArray() + assert(valueArray.getTimestampLTZNanos(0) === ltzValue) + assert(valueArray.getTimestampLTZNanos(1) === ltzValue) + } + testBothCodegenAndInterpreted("codegen projection reads nanos timestamp column") { val boundRef = BoundReference(0, TimestampNTZNanosType(9), nullable = false) val projection = GenerateUnsafeProjection.generate(Seq(boundRef)) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala new file mode 100644 index 0000000000000..1c41331eb9009 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala @@ -0,0 +1,93 @@ +/* + * 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.catalyst.expressions + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, ScalarFunction} +import org.apache.spark.sql.types.{DataType, IntegerType} + +class TransformExpressionSuite extends SparkFunSuite { + + /** + * A bound function with a stable canonical name and no `equals` of its own. A plain class, NOT a + * case class: a case class would bring structural `equals`, which is the connector-provided + * comparison these tests are trying to do without. + */ + private class NamedFunction(canonical: String) extends ScalarFunction[Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = canonical + override def canonicalName(): String = canonical + } + + /** Honours the contract in `BoundFunction#equals`. */ + private class ComparableFunction extends NamedFunction("test.comparable") { + override def equals(other: Any): Boolean = other.isInstanceOf[ComparableFunction] + override def hashCode(): Int = canonicalName().hashCode + } + + private val a = AttributeReference("a", IntegerType)() + private val b = AttributeReference("b", IntegerType)() + + private def bucket(function: BoundFunction, child: Expression, numBuckets: Int = 4) = + TransformExpression(function, Seq(child), Some(numBuckets)) + + test("SPARK-58769: expression equality follows the function's own equals") { + // Spark does not derive transform identity itself -- it defers to the connector, because only + // the connector knows which of its state matters. A function that does not implement `equals` + // therefore yields expressions that do not compare equal across separate binds, which costs + // deduplication and reuse but never correctness. See BoundFunction#equals. + assert( + bucket(new NamedFunction("test.bucket"), a) != bucket(new NamedFunction("test.bucket"), a), + "no equals on the function means no equality across binds") + + val shared = new NamedFunction("test.bucket") + assert(bucket(shared, a) == bucket(shared, a), "a shared instance is equal either way") + + // A function that does implement it gets the deduplication. + val left = bucket(new ComparableFunction, a) + val right = bucket(new ComparableFunction, a) + assert(left.function ne right.function, "the fixture must bind a fresh instance per call") + assert(left == right) + assert(left.semanticEquals(right)) + assert(ExpressionSet(Seq(left, right)).size == 1) + } + + test("SPARK-58769: the two comparisons agree for a function that follows the contract") { + // `equals` is the finer comparison and the canonical name the coarser one, so a function that + // overrides both answers both consistently. They still differ in what they take into account: + // `isSameFunction` ignores the arguments, because a join compares bucket(4, left.id) against + // bucket(4, right.id) and recovers the positions separately, while equality does not. + val left = bucket(new ComparableFunction, a) + val right = bucket(new ComparableFunction, b) + assert(left.function ne right.function, "the fixture must bind a fresh instance per call") + assert(left.function == right.function) + assert(left.function.canonicalName() == right.function.canonicalName()) + assert(left.isSameFunction(right), "the same partition function, arguments aside") + assert(left != right, "but not the same expression, since the arguments differ") + } + + test("SPARK-58769: the function's equals does not override the arguments") { + // A coarse comparison on the connector's side does not have to carry the whole identity: the + // transform's arguments and bucket count are compared separately, by Spark. + assert(bucket(new ComparableFunction, a) != bucket(new ComparableFunction, b), + "different argument") + assert(bucket(new ComparableFunction, a, 4) != bucket(new ComparableFunction, a, 8), + "different bucket count") + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala new file mode 100644 index 0000000000000..afd14131e5ec7 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WrapUDTExpressionSuite.scala @@ -0,0 +1,115 @@ +/* + * 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.catalyst.expressions + +import org.apache.spark.{SPARK_DOC_ROOT, SparkFunSuite} +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.analysis.FunctionRegistry +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.expressions.Cast.{toSQLExpr, toSQLType} +import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber +import org.apache.spark.sql.types.{BooleanType, IntegerType, StringType, TestUDT} + +class WrapUDTExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { + + test("WrapUDT should use target UDT with matching SQL type") { + val udt = new TestUDT.MyDenseVectorUDT() + val data = new GenericArrayData(Array[Any](1.0, 2.0)) + val wrapUDTExpression = WrapUDT(Literal.create(data, udt.sqlType), udt) + + assert(wrapUDTExpression.checkInputDataTypes().isSuccess) + assert(wrapUDTExpression.dataType == udt) + checkEvaluation(wrapUDTExpression, data) + } + + test("WrapUDT should parse target UDT from foldable expression") { + val udt = new TestUDT.MyDenseVectorUDT() + val data = new GenericArrayData(Array[Any](1.0, 2.0)) + val json = udt.json + val target = Concat(Seq( + Literal.create(json.substring(0, 8), StringType), + Literal.create(json.substring(8), StringType))) + val wrapUDTExpression = new WrapUDT(Literal.create(data, udt.sqlType), target) + + assert(wrapUDTExpression.checkInputDataTypes().isSuccess) + assert(wrapUDTExpression.dataType == udt) + checkEvaluation(wrapUDTExpression, data) + } + + test("WrapUDT target expression should be a UDT") { + val target = Literal.create("int", StringType) + checkError( + exception = intercept[AnalysisException] { + new WrapUDT(Literal.create(1, IntegerType), target) + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> toSQLExpr(target), + "paramIndex" -> ordinalNumber(1), + "requiredType" -> toSQLType("UserDefinedType"), + "inputSql" -> toSQLExpr(target), + "inputType" -> toSQLType(IntegerType))) + } + + test("WrapUDT target expression should be foldable") { + val target = AttributeReference("udt", StringType)() + checkError( + exception = intercept[AnalysisException] { + new WrapUDT(Literal.create(1, IntegerType), target) + }, + condition = "INVALID_SCHEMA.NON_STRING_LITERAL", + parameters = Map("inputSchema" -> toSQLExpr(target))) + } + + test("WrapUDT should reject wrong number of arguments through FunctionRegistry") { + val expression = Literal.create(1, IntegerType) + val target = Literal.create(new TestUDT.MyDenseVectorUDT().json, StringType) + + Seq( + Seq(expression) -> "1", + Seq(expression, target, Literal.create("extra", StringType)) -> "3").foreach { + case (arguments, actualNum) => + checkError( + exception = intercept[AnalysisException] { + FunctionRegistry.internal.lookupFunction(FunctionIdentifier("wrap_udt"), arguments) + }, + condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + parameters = Map( + "functionName" -> "`wrap_udt`", + "expectedNum" -> "2", + "actualNum" -> actualNum, + "docroot" -> SPARK_DOC_ROOT)) + } + } + + test("WrapUDT input type should match target UDT SQL type") { + val b1 = Literal.create(false, BooleanType) + val udt = new TestUDT.MyDenseVectorUDT() + val wrapUDTExpression = WrapUDT(b1, udt) + assert(wrapUDTExpression.checkInputDataTypes() == + DataTypeMismatch( + errorSubClass = "UNEXPECTED_INPUT_TYPE", + messageParameters = Map( + "paramIndex" -> ordinalNumber(0), + "requiredType" -> toSQLType(udt.sqlType), + "inputSql" -> "\"false\"", + "inputType" -> "\"BOOLEAN\""))) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XXH3Suite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XXH3Suite.scala new file mode 100644 index 0000000000000..0e5288e76d96e --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XXH3Suite.scala @@ -0,0 +1,143 @@ +/* + * 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.catalyst.expressions + +import org.apache.spark.SparkFunSuite + +/** + * Validates the [[XXH3]] port against the reference implementation's known-answer vectors + * (github.com/Cyan4973/xxHash). Inputs are prefixes of the same pseudo-random buffer the reference + * test harness uses (tests/sanity_test.c); expected values were produced with the reference + * `xxhash` library and cover every length branch (0, 1-3, 4-8, 9-16, 17-128, 129-240, and the + * long path) with both a zero and a non-zero seed. + */ +class XXH3Suite extends SparkFunSuite { + + private val buffer: Array[Byte] = { + val buf = new Array[Byte](4200) + var byteGen = 0x9E3779B1L + for (i <- buf.indices) { + buf(i) = (byteGen >>> 56).toByte + byteGen *= 0x9E3779B185EBCA8DL + } + buf + } + + private def input(len: Int): Array[Byte] = java.util.Arrays.copyOf(buffer, len) + + // 64-bit: (length, seed, expected) + private val vectors64: Seq[(Int, Long, Long)] = Seq( + (0, 0L, 0x2d06800538d394c2L), + (1, 0L, 0xc44bdff4074eecdbL), + (2, 0L, 0x7a9978044cb8a8bbL), + (3, 0L, 0x54247382a8d6b94dL), + (4, 0L, 0xe5dc74bc51848a51L), + (5, 0L, 0xe4243f00720306bbL), + (7, 0L, 0x9941e0007f555e50L), + (8, 0L, 0x24ccc9acaa9f65e4L), + (9, 0L, 0x14d5001c15dd3f2bL), + (12, 0L, 0xa713daf0dfbb77e7L), + (16, 0L, 0x981b17d36c7498c9L), + (17, 0L, 0x796f5acd3a60f862L), + (32, 0L, 0x9feaddbdbf57eed3L), + (64, 0L, 0x9cb48487720ec49dL), + (100, 0L, 0x93cd95432b7d483fL), + (128, 0L, 0xfcff24126754d861L), + (129, 0L, 0x98f1b0a679a2ca29L), + (160, 0L, 0x9d03a319ed4cbd2bL), + (200, 0L, 0xbddca58935d7c038L), + (240, 0L, 0x81c3c2b67f568ccfL), + (241, 0L, 0xc5a639ecd2030e5eL), + (256, 0L, 0x55de574ad89d0ac5L), + (512, 0L, 0x617e49599013cb6bL), + (1024, 0L, 0xdd85c9b5c1109c5cL), + (2048, 0L, 0xdd59e2c3a5f038e0L), + (4096, 0L, 0xe91206429d1f48f9L), + (0, 0x9e3779b185ebca8dL, 0xa8a6b918b2f0364aL), + (2, 0x9e3779b185ebca8dL, 0x764b35c90519ad88L), + (8, 0x9e3779b185ebca8dL, 0x8f973410999b8f6bL), + (16, 0x9e3779b185ebca8dL, 0x663f29333b4db6b1L), + (64, 0x9e3779b185ebca8dL, 0x4fe8895db9b8c077L), + (240, 0x9e3779b185ebca8dL, 0xcc0f58c27ef3d8eeL), + (256, 0x9e3779b185ebca8dL, 0x4d30234b7a3aa61cL), + (1024, 0x9e3779b185ebca8dL, 0xef368a8a2ebabaefL) + ) + + // 128-bit: (length, seed, expected canonical hex) + private val vectors128: Seq[(Int, Long, String)] = Seq( + (0, 0L, "99aa06d3014798d86001c324468d497f"), + (1, 0L, "a6cd5e9392000f6ac44bdff4074eecdb"), + (2, 0L, "76750c3c7bf956687a9978044cb8a8bb"), + (3, 0L, "20efc49ff02422ea54247382a8d6b94d"), + (4, 0L, "970d585ac632bf8e2e7d8d6876a39fe9"), + (5, 0L, "62ed587687606b4e057c7ed2c01fa1d1"), + (7, 0L, "dd9b6039f79ec416081c22dd284a2f0a"), + (8, 0L, "47a7f080d82bb45664c69cab4bb21dc5"), + (9, 0L, "564ef6078950d457ed7ccbc501eb7501"), + (12, 0L, "6e3efd8fc7802b18061a192713f69ad9"), + (16, 0L, "c68c368ecf8a9c05562980258a998629"), + (17, 0L, "955fa78643ed3669abbc12d11973d7db"), + (32, 0L, "98fc6458710dc2e8278410a17595e3f9"), + (64, 0L, "6d90e81a9b0fd622efdb6a44690721a9"), + (100, 0L, "9b50b05817ab158e5fcbc2e3295f2476"), + (128, 0L, "39992220e045260aebb15e34a7fb5ab1"), + (129, 0L, "03815fc91f1b30b686c9e3bc8f0a3b5c"), + (160, 0L, "ba5d218964b622ad737126c8d7c09cee"), + (200, 0L, "e76ff4780fe18439eb060f1bb3126f5a"), + (240, 0L, "aa4202daa2769dc85c9aae94c8ebe5a0"), + (241, 0L, "99a80ecf0ecfc647c5a639ecd2030e5e"), + (256, 0L, "8b1c66091423d28855de574ad89d0ac5"), + (512, 0L, "18d2d110dcc9bca1617e49599013cb6b"), + (1024, 0L, "0d30d24071c64c57dd85c9b5c1109c5c"), + (2048, 0L, "f736557fd47073a5dd59e2c3a5f038e0"), + (4096, 0L, "b9cfaea2ca5626a4e91206429d1f48f9"), + (0, 0x9e3779b185ebca8dL, "00feaa732a3ce25ea986dfc5d7605bfe"), + (2, 0x9e3779b185ebca8dL, "7b96e6a600dae67d764b35c90519ad88"), + (8, 0x9e3779b185ebca8dL, "f50cec145bcd5c5a7b29471dc729b5ff"), + (16, 0x9e3779b185ebca8dL, "6ffcb80cd33085c80346d13a7a5498c7"), + (64, 0x9e3779b185ebca8dL, "37b738968d40bda59405ba2affa95ceb"), + (240, 0x9e3779b185ebca8dL, "29d2133d6ea58c5b604e98db085c1864"), + (256, 0x9e3779b185ebca8dL, "aaa57235b92d5e7c4d30234b7a3aa61c"), + (1024, 0x9e3779b185ebca8dL, "17600efe2b493a18ef368a8a2ebabaef") + ) + + test("reference test buffer generation") { + val expectedFirst24 = Seq(0x00, 0x52, 0x92, 0x9B, 0xB7, 0x32, 0xA3, 0x24, 0x2D, 0x00, 0xAF, + 0x95, 0x0E, 0xEC, 0xB8, 0x93, 0xE3, 0xDF, 0xEF, 0x93, 0xAA, 0xD6, 0xCD, 0x2A) + expectedFirst24.zipWithIndex.foreach { case (b, i) => + assert((buffer(i) & 0xFF) == b, + s"buffer($i): expected 0x${b.toHexString} got 0x${(buffer(i) & 0xFF).toHexString}") + } + } + + test("XXH3 64-bit against reference vectors") { + vectors64.foreach { case (len, seed, expected) => + val actual = XXH3.hash64(input(len), seed) + assert(actual == expected, s"len=$len seed=0x${seed.toHexString}: " + + s"expected 0x${expected.toHexString} but got 0x${actual.toHexString}") + } + } + + test("XXH3 128-bit against reference vectors") { + vectors128.foreach { case (len, seed, expected) => + val actual = XXH3.hash128Hex(input(len), seed).toString + assert(actual == expected, s"len=$len seed=0x${seed.toHexString}: " + + s"expected $expected but got $actual") + } + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XmlExpressionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XmlExpressionsSuite.scala index 54bef646739fd..a0011e1dbbb94 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XmlExpressionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/XmlExpressionsSuite.scala @@ -522,4 +522,17 @@ class XmlExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper { ) } + test("XmlToStructs and StructsToXml are stateful and produce fresh copies") { + val schema = StructType(StructField("a", IntegerType) :: Nil) + + val xmlToStructs = XmlToStructs(schema, Map.empty, Literal("<a>1</a>"), UTC_OPT) + assert(xmlToStructs.stateful) + assert(xmlToStructs.freshCopyIfContainsStatefulExpression() ne xmlToStructs) + + val struct = Literal.create(InternalRow(1), schema) + val structsToXml = StructsToXml(Map.empty, struct, UTC_OPT) + assert(structsToXml.stateful) + assert(structsToXml.freshCopyIfContainsStatefulExpression() ne structsToXml) + } + } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala index 0f7f5ca54be01..4dcc3bf28d24e 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/DatasketchesHllSketchSuite.scala @@ -23,9 +23,9 @@ import scala.util.Random import org.apache.datasketches.hll.HllSketch import org.apache.datasketches.memory.Memory -import org.apache.spark.SparkFunSuite +import org.apache.spark.{SparkFunSuite, SparkRuntimeException} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{BoundReference, HllSketchEstimate} +import org.apache.spark.sql.catalyst.expressions.{BoundReference, HllSketchEstimate, HllUnion, Literal} import org.apache.spark.sql.types.{BinaryType, DataType, IntegerType, LongType, StringType} import org.apache.spark.unsafe.types.UTF8String @@ -153,4 +153,130 @@ class DatasketchesHllSketchSuite extends SparkFunSuite { s"but got: ${exception.getClass.getName}: ${exception.getMessage}" ) } + /** Runs HllUnionAgg over `inputs` (a NULL entry stands for a NULL sketch) for a single group. */ + private def unionAgg(inputs: Seq[Any], allowDifferentLgConfigK: Boolean): Array[Byte] = { + val aggFunc = new HllUnionAgg( + BoundReference(0, BinaryType, nullable = true), allowDifferentLgConfigK) + val buffer = inputs.foldLeft(aggFunc.createAggregationBuffer()) { (buf, input) => + aggFunc.update(buf, InternalRow(input)) + } + aggFunc.eval(buffer).asInstanceOf[Array[Byte]] + } + + /** Runs HllSketchAgg at `lgConfigK` over `values` (a NULL entry stands for a NULL value). */ + private def sketchAgg(values: Seq[Any], lgConfigK: Int): Array[Byte] = { + val aggFunc = new HllSketchAgg(BoundReference(0, StringType, nullable = true), lgConfigK) + val buffer = values.foldLeft(aggFunc.createAggregationBuffer()) { (buf, value) => + aggFunc.update(buf, InternalRow(value)) + } + aggFunc.eval(buffer).asInstanceOf[Array[Byte]] + } + + /** Evaluates the scalar hll_union over two serialized sketches. */ + private def scalarUnion( + left: Array[Byte], right: Array[Byte], allowDifferentLgConfigK: Boolean): Array[Byte] = + HllUnion( + Literal(left, BinaryType), + Literal(right, BinaryType), + Literal(allowDifferentLgConfigK)).eval(InternalRow.empty).asInstanceOf[Array[Byte]] + + private def lgConfigKOf(sketch: Array[Byte]): Int = + HllSketch.heapify(Memory.wrap(sketch)).getLgConfigK + + private def estimateOf(sketch: Array[Byte]): Long = + HllSketchEstimate(BoundReference(0, BinaryType, nullable = true)) + .eval(InternalRow(sketch)).asInstanceOf[Long] + + private def stringValues(n: Int): Seq[Any] = + Seq.tabulate(n)(i => UTF8String.fromString(i.toString)) + + test("hll_union_agg on a group with no non-NULL sketch yields an empty default-lgConfigK " + + "sketch") { + // The aggregate has no lgConfigK parameter and never saw a sketch, so it has no precision to + // report and falls back to the Datasketches default. Documented here because the resulting + // sketch is observable, and must stay harmless to later unions (see the tests below). + val allNull = unionAgg(Seq(null, null), allowDifferentLgConfigK = false) + assert(estimateOf(allNull) == 0L) + assert(lgConfigKOf(allNull) == HllSketch.DEFAULT_LG_K) + + // hll_sketch_agg does not share the problem only because it has the parameter: it builds its + // buffer eagerly at the requested lgConfigK. Without the argument it defaults to 12 as well. + val emptyAt15 = sketchAgg(Seq(null, null), 15) + assert(estimateOf(emptyAt15) == 0L) + assert(lgConfigKOf(emptyAt15) == 15) + } + + test("hll_union_agg merges an empty sketch of a different lgConfigK without an error") { + // An empty sketch holds no coupons, so unioning it cannot lose information at any lgConfigK. + // This is the shape produced by the test above, i.e. what a persisted table ends up holding + // for a group whose sketches were all NULL. + val emptyAtDefaultLgK = unionAgg(Seq(null), allowDifferentLgConfigK = false) + val sketchAt15 = sketchAgg(stringValues(1000), 15) + + Seq(true, false).foreach { allowDifferentLgConfigK => + Seq( + ("empty sketch first", Seq[Any](emptyAtDefaultLgK, sketchAt15)), + ("empty sketch last", Seq[Any](sketchAt15, emptyAtDefaultLgK)) + ).foreach { case (order, inputs) => + val merged = unionAgg(inputs, allowDifferentLgConfigK) + // The non-empty sketch decides the precision, whichever order the rows arrive in: the + // result must not depend on which row the aggregate happens to see first. + assert(lgConfigKOf(merged) == 15, + s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed the lgConfigK") + assert(estimateOf(merged) == estimateOf(sketchAt15), + s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed the estimate") + } + } + } + + test("hll_union_agg still rejects non-empty sketches with different lgConfigK") { + val sketchAt12 = sketchAgg(stringValues(1000), 12) + val sketchAt15 = sketchAgg(stringValues(1000), 15) + + Seq( + Seq[Any](sketchAt12, sketchAt15), + Seq[Any](sketchAt15, sketchAt12) + ).foreach { inputs => + val exception = intercept[SparkRuntimeException] { + unionAgg(inputs, allowDifferentLgConfigK = false) + } + assert(exception.getCondition == "HLL_UNION_DIFFERENT_LG_K") + } + + // And still downsamples rather than erroring when the caller opts in. + assert(lgConfigKOf(unionAgg(Seq(sketchAt15, sketchAt12), allowDifferentLgConfigK = true)) == 12) + } + + test("hll_union merges an empty sketch of a different lgConfigK without an error") { + val emptyAtDefaultLgK = unionAgg(Seq(null), allowDifferentLgConfigK = false) + val sketchAt15 = sketchAgg(stringValues(1000), 15) + + Seq(true, false).foreach { allowDifferentLgConfigK => + Seq( + ("empty sketch first", emptyAtDefaultLgK, sketchAt15), + ("empty sketch last", sketchAt15, emptyAtDefaultLgK) + ).foreach { case (order, left, right) => + val merged = scalarUnion(left, right, allowDifferentLgConfigK) + assert(lgConfigKOf(merged) == 15, + s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed the lgConfigK") + assert(estimateOf(merged) == estimateOf(sketchAt15), + s"$order (allowDifferentLgConfigK=$allowDifferentLgConfigK) changed the estimate") + } + } + } + + test("hll_union still rejects non-empty sketches with different lgConfigK") { + val sketchAt12 = sketchAgg(stringValues(1000), 12) + val sketchAt15 = sketchAgg(stringValues(1000), 15) + + Seq((sketchAt12, sketchAt15), (sketchAt15, sketchAt12)).foreach { case (left, right) => + val exception = intercept[SparkRuntimeException] { + scalarUnion(left, right, allowDifferentLgConfigK = false) + } + assert(exception.getCondition == "HLL_UNION_DIFFERENT_LG_K") + } + + // Two populated sketches still downsample to the smaller lgConfigK when opted in. + assert(lgConfigKOf(scalarUnion(sketchAt15, sketchAt12, allowDifferentLgConfigK = true)) == 12) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompilerSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompilerSuite.scala new file mode 100644 index 0000000000000..95692428349af --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeCompilerSuite.scala @@ -0,0 +1,1628 @@ +/* + * 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.catalyst.expressions.codegen + +import java.io.File +import java.lang.reflect.Modifier +import java.net.{URI, URL, URLClassLoader} +import java.util.Collections +import javax.tools.{JavaFileObject, SimpleJavaFileObject, StandardLocation, ToolProvider} + +import scala.jdk.CollectionConverters._ + +import org.codehaus.commons.compiler.CompileException +import org.mockito.Mockito.{mock, when} + +import org.apache.spark.{JobArtifactSet, JobArtifactState, SparkConf, SparkEnv, SparkFunSuite} +import org.apache.spark.executor.ExecutorClassLoader +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BoundReference, GreaterThan, LessThan, Literal} +import org.apache.spark.sql.catalyst.plans.SQLHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.util.Utils + +/** + * Tests for the [[CodeCompiler]] trait, its backend selection, and behavioural + * parity between the [[JaninoCodeCompiler]] and [[JdkCodeCompiler]] backends. + */ +class CodeCompilerSuite extends SparkFunSuite with SQLHelper { + + // A self-contained class body that exercises common shapes Spark's code generators + // produce: an override of generate(), a nested concrete class, fields, and a + // straightforward arithmetic method. Designed to compile under both backends with + // no dependencies on Spark types beyond GeneratedClass. + private val sampleClassBody: String = + s""" + |public java.lang.Object generate(Object[] references) { + | return new SpecificEvaluator(references); + |} + | + |static class SpecificEvaluator { + | private Object[] references; + | private long counter = 0L; + | + | public SpecificEvaluator(Object[] refs) { + | this.references = refs; + | } + | + | public long evaluate(long input) { + | counter += 1L; + | return (input * 31L) + counter; + | } + |} + |""".stripMargin + + private def newCodeAndComment(body: String): CodeAndComment = + new CodeAndComment(body, scala.collection.Map.empty[String, String]) + + // ---------------- backend selection ---------------- + + test("forBackend: 'janino' returns JaninoCodeCompiler") { + assert(CodeCompiler.forBackend("janino") eq JaninoCodeCompiler) + } + + test("forBackend: 'jdk' returns JdkCodeCompiler when available, else falls back") { + val backend = CodeCompiler.forBackend("jdk") + if (JdkCodeCompiler.isAvailable) { + assert(backend eq JdkCodeCompiler) + } else { + assert(backend eq JaninoCodeCompiler) + } + } + + test("forBackend: name is case-insensitive") { + assert(CodeCompiler.forBackend("JANINO") eq JaninoCodeCompiler) + assert(CodeCompiler.forBackend("Janino") eq JaninoCodeCompiler) + if (JdkCodeCompiler.isAvailable) { + assert(CodeCompiler.forBackend("JDK") eq JdkCodeCompiler) + assert(CodeCompiler.forBackend("Jdk") eq JdkCodeCompiler) + } + } + + test("forBackend: unknown name throws IllegalArgumentException") { + val ex = intercept[IllegalArgumentException] { + CodeCompiler.forBackend("acme-compiler") + } + assert(ex.getMessage.contains("acme-compiler")) + assert(ex.getMessage.contains("janino")) + assert(ex.getMessage.contains("jdk")) + } + + test("active() honors SQLConf.CODEGEN_COMPILER") { + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "janino") { + assert(CodeCompiler.active().name == CodeCompiler.JANINO) + } + if (JdkCodeCompiler.isAvailable) { + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active().name == CodeCompiler.JDK) + } + } + } + + test("active() routes REPL-context codegen to Janino regardless of config") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // An ExecutorClassLoader in the active class loader chain marks REPL / interactive + // codegen. Classes defined there carry self-inconsistent reflection metadata that the + // JDK compiler cannot resolve, so the backend must deterministically route to Janino + // even when `jdk` is configured. A `spark://` class URI keeps construction cheap: the + // RPC fetch function is only referenced, never invoked, at construction time. + val replLoader = new ExecutorClassLoader( + new SparkConf(), null, "spark://localhost:0", getClass.getClassLoader, false) + val childOfRepl = new URLClassLoader(Array.empty[URL], replLoader) + val prev = Thread.currentThread().getContextClassLoader + try { + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + // Baseline: an ordinary (non-REPL) loader honors the configured `jdk` backend. + Thread.currentThread().setContextClassLoader(prev) + assert(CodeCompiler.active() eq JdkCodeCompiler) + // An ExecutorClassLoader at the head of the chain forces Janino. + Thread.currentThread().setContextClassLoader(replLoader) + assert(CodeCompiler.active() eq JaninoCodeCompiler) + // An ExecutorClassLoader anywhere in the parent chain forces Janino too. + Thread.currentThread().setContextClassLoader(childOfRepl) + assert(CodeCompiler.active() eq JaninoCodeCompiler) + } + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("active() routes artifact/REPL-session codegen (replClassDirUri) to Janino") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + // Without an artifact class-dir URI, the configured `jdk` backend is honored. + assert(CodeCompiler.active() eq JdkCodeCompiler) + // A session/job carrying a `replClassDirUri` (Spark Connect per-session artifacts, + // spark-shell) must route to Janino, including driver-side codegen where no + // ExecutorClassLoader is in the loader chain (e.g. a Connect UDF over a local + // relation referencing an Ammonite `$sess` class). + JobArtifactSet.withActiveJobArtifactState( + JobArtifactState("test-uuid", Some("spark://localhost:0/classes"))) { + assert(CodeCompiler.active() eq JaninoCodeCompiler) + } + } + } + + test("active() routes spark-shell codegen (spark.repl.class.uri in SparkEnv conf) to Janino") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // spark-shell publishes the REPL class URI in the SparkEnv conf. That signal alone - + // no ExecutorClassLoader in the chain, no artifact state - must route to Janino. + val env = mock(classOf[SparkEnv]) + when(env.conf).thenReturn( + new SparkConf().set("spark.repl.class.uri", "spark://localhost:0/classes")) + val prevEnv = SparkEnv.get + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + // Baseline: without the REPL conf signal, the configured backend is honored. + assert(CodeCompiler.active() eq JdkCodeCompiler) + SparkEnv.set(env) + try { + assert(CodeCompiler.active() eq JaninoCodeCompiler) + } finally { + SparkEnv.set(prevEnv) + } + } + } + + test("active(code) routes code referencing a package-object class to Janino") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + // Ordinary code honors the configured jdk backend. + assert(CodeCompiler.active(newCodeAndComment("int x = 1;")) eq JdkCodeCompiler) + // A reference to a Scala `package object`'s nested class (`...package$Inner`) - whose + // `package` segment the JDK compiler can name in no form - must route to Janino, the + // same always-Janino bucket as REPL classes javac cannot name. + val pkgObjBody = "org.apache.spark.sql.foo.package$Inner v = " + + "(org.apache.spark.sql.foo.package$Inner) references[0];" + assert(CodeCompiler.active(newCodeAndComment(pkgObjBody)) eq JaninoCodeCompiler) + // A legal identifier that merely contains the text "package" is unaffected. + assert(CodeCompiler.active(newCodeAndComment("com.mypackage.Inner v;")) eq JdkCodeCompiler) + } + } + + test("SQLConf rejects invalid backend names at set time") { + val ex = intercept[IllegalArgumentException] { + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "acme") {} + } + // SQLConf checkValues surfaces the allowed set in the message + val msg = ex.getMessage.toLowerCase(java.util.Locale.ROOT) + assert(msg.contains("janino") || msg.contains("jdk")) + } + + // ---------------- compilation parity ---------------- + + /** + * Invoke `evaluate(long)` on the object returned by `generate()`. The nested + * `SpecificEvaluator` class lives inside the generated `GeneratedClass`, and the + * generated `GeneratedClass` is loaded by a different classloader than this test + * suite. Even though the method is `public`, reflection requires explicit access + * because the test cannot statically reach the enclosing class. + */ + private def invokeEvaluate(result: Any, input: Long): Any = { + val m = result.getClass.getMethod("evaluate", classOf[Long]) + m.setAccessible(true) + m.invoke(result, Long.box(input)) + } + + test("Janino backend compiles a simple class body and produces working bytecode") { + val (generated, stats) = JaninoCodeCompiler.compile(newCodeAndComment(sampleClassBody)) + assert(generated != null) + val result = generated.generate(Array.empty[Any]) + assert(invokeEvaluate(result, 10L) == 311L) // 10*31 + 1 + assert(invokeEvaluate(result, 20L) == 622L) // 20*31 + 2 + assert(stats.maxMethodCodeSize > 0) + assert(stats.maxConstPoolSize > 0) + } + + test("JDK backend compiles the same class body and produces equivalent results") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val (generated, stats) = JdkCodeCompiler.compile(newCodeAndComment(sampleClassBody)) + assert(generated != null) + val result = generated.generate(Array.empty[Any]) + assert(invokeEvaluate(result, 10L) == 311L) + assert(invokeEvaluate(result, 20L) == 622L) + assert(stats.maxMethodCodeSize > 0) + assert(stats.maxConstPoolSize > 0) + } + + test("Both backends produce class bytecodes the ByteCodeStats parser accepts") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val code = newCodeAndComment(sampleClassBody) + val (_, janinoStats) = JaninoCodeCompiler.compile(code) + val (_, jdkStats) = JdkCodeCompiler.compile(code) + // Bytecode size will not be identical - the two compilers emit different instruction + // sequences for the same source - but both must be non-trivial and non-error. + assert(janinoStats.maxMethodCodeSize > 0) + assert(jdkStats.maxMethodCodeSize > 0) + assert(janinoStats.maxConstPoolSize > 0) + assert(jdkStats.maxConstPoolSize > 0) + // numInnerClasses must agree between backends. Both wrap the same body in a single + // outer class declaration, so for K nested classes the formula `size - 2` yields the + // same K-1 value under either backend. + assert(janinoStats.numInnerClasses == jdkStats.numInnerClasses, + s"numInnerClasses divergence: Janino=${janinoStats.numInnerClasses}, " + + s"JDK=${jdkStats.numInnerClasses}") + } + + test("Both backends compile a reference to a class nested in a Scala object") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // End-to-end check of the mllib legacy save/load shape: `NewInstance` emits the + // binary class name (CodeGenerator.javaSourceName == Class#getName). For a class + // nested in a Scala object that name carries module `$`s (Outer$SaveLoadV1$Leaf), + // and only that binary form resolves under javac - the dotted canonical form makes + // javac reconstruct a non-existent Outer$SaveLoadV1$$Leaf. Both backends must + // accept the same generated source. + val binary = CodeGenerator.javaSourceName(classOf[CodeCompilerSuite.SaveLoadV1.Leaf]) + assert(binary.contains("$SaveLoadV1$"), s"fixture is not object-nested: $binary") + val body = + s""" + |public java.lang.Object generate(Object[] references) { + | $binary leaf = new $binary(7); + | return Integer.valueOf(leaf.x()); + |} + |""".stripMargin + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) === Integer.valueOf(7)) + } + + test("JDK backend resolves a class available only via the context classloader") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // dyn.DynHelper is compiled into a temp dir that is NOT on java.class.path and is + // served by a custom (non-URLClassLoader) loader - the shape of REPL-generated and + // Spark Connect session-artifact classes, which live only on a runtime loader. The + // JDK backend must resolve it the way Janino does (via the classloader), not via a + // file-based -classpath. + val dir = compileDynHelper() + // DirClassLoader is deliberately a plain ClassLoader, not a URLClassLoader, so the + // retired -classpath harvesting would never have found its classes. + val loader = new DirClassLoader(dir, getClass.getClassLoader) + assert(loader.loadClass("dyn.DynHelper") != null) + + val body = + """ + |public java.lang.Object generate(Object[] references) { + | return Integer.valueOf(dyn.DynHelper.magic()); + |} + |""".stripMargin + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + // Both backends must resolve dyn.DynHelper through the context classloader. + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) === Integer.valueOf(4242)) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("JDK backend resolves a class from a non-enumerable (REPL-style) classloader") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A loader that serves classes by name (getResource / getResourceAsStream) but does + // NOT support package enumeration (getResources) - the shape of the Scala REPL and + // Spark Connect session loaders, which hold generated classes only in memory. + // Resolution must fall back to the source's referenced names, the way Janino does. + val dir = compileDynHelper() + val loader = new NonEnumerableDirClassLoader(dir, getClass.getClassLoader) + assert(!loader.getResources("dyn").hasMoreElements, + "fixture loader must not support package enumeration") + assert(loader.loadClass("dyn.DynHelper") != null) + + val body = + """ + |public java.lang.Object generate(Object[] references) { + | return Integer.valueOf(dyn.DynHelper.magic()); + |} + |""".stripMargin + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) === Integer.valueOf(4242)) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("JDK backend resolves a class from a getResourceAsStream-only classloader") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A loader that serves .class bytes only via getResourceAsStream and exposes no + // resource URL (getResource is null) and no enumeration - the shape of the Scala + // REPL / Ammonite (`ammonite.$sess`) and similar in-memory loaders. Resolution must + // probe with getResourceAsStream, the way Janino does, not getResource. + val dir = compileDynHelper() + val loader = new StreamOnlyDirClassLoader(dir, getClass.getClassLoader) + assert(loader.getResource("dyn/DynHelper.class") == null, + "fixture loader must expose no resource URL") + assert(loader.getResourceAsStream("dyn/DynHelper.class") != null) + assert(!loader.getResources("dyn").hasMoreElements) + + val body = + """ + |public java.lang.Object generate(Object[] references) { + | return Integer.valueOf(dyn.DynHelper.magic()); + |} + |""".stripMargin + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) === Integer.valueOf(4242)) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("JDK backend ignores a phantom class served for a package path") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // Some artifact / session loaders return a non-null, non-class stream for a + // package-shaped resource path. The generated unit's own package is + // org.apache.spark.sql.catalyst.expressions; if such a phantom were treated as a + // class, javac would fail with "package ... clashes with class of same name". + // Resolution must validate the class-file magic and ignore the phantom. + val loader = new PhantomPackageClassLoader(getClass.getClassLoader) + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(sampleClassBody)) + assert(generated != null) + assert(invokeEvaluate(generated.generate(Array.empty[Any]), 10L) == 311L) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("JDK backend resolves an object-nested class + enclosing chain (non-enumerable loader)") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // SaveLoadV1.Leaf is object-nested (binary ...CodeCompilerSuite$SaveLoadV1$Leaf, + // canonical carries `$`), the same shape as a Scala REPL class + // ($line.$read$$iw$X). Served by a loader that resolves classes by name but does + // NOT enumerate packages, javac needs the leaf AND its enclosing classes; the + // by-name fallback must add the whole `$`-prefix chain. + val binary = classOf[CodeCompilerSuite.SaveLoadV1.Leaf].getName + val loader = new NonEnumerableWrapper(getClass.getClassLoader) + val pkgPath = binary.substring(0, binary.lastIndexOf('.')).replace('.', '/') + assert(!loader.getResources(pkgPath).hasMoreElements, "loader must not enumerate") + assert(loader.getResourceAsStream(binary.replace('.', '/') + ".class") != null) + + val body = + s""" + |public java.lang.Object generate(Object[] references) { + | return new $binary(7); + |} + |""".stripMargin + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) != null) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + // A wrapper that delegates class/resource loading to its parent but refuses to + // enumerate packages (getResources returns empty), forcing the by-name fallback - + // the shape of an in-memory REPL / session loader that serves getResourceAsStream + // but not getResources. + private class NonEnumerableWrapper(parent: ClassLoader) extends ClassLoader(parent) { + override def getResources(name: String): java.util.Enumeration[URL] = + Collections.emptyEnumeration() + } + + // Compile `dyn.DynHelper` into a fresh temp dir using the system Java compiler, so + // the class exists only under that dir (never on java.class.path). + test("reflection that raises a LinkageError does not route the unit or escape") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A partial or shaded jar can leave a class loadable while its enclosing class is not. + // `getCanonicalName` and `getEnclosingClass` both throw NoClassDefFoundError then, and + // `NonFatal` does not cover a LinkageError, so an escaping Error would bypass the + // codegen fallbacks. The token cannot be evaluated, which is no evidence that narrowing + // it is unsafe, so the unit must stay on the configured backend rather than gain a + // permanent Janino arm. + val dir = compileLinkageFixture() + // Drop the enclosing class, keeping the anonymous one that references it. + assert(new File(dir, "lnk/Holder$Mid.class").delete(), "fixture setup: enclosing class") + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val anon = loader.loadClass("lnk.Holder$Mid$1") + // Precondition: the reflective calls the predicate makes really do throw here. + intercept[LinkageError](anon.getCanonicalName) + intercept[LinkageError](anon.getEnclosingClass) + + val body = s"${anon.getName} v = (${anon.getName}) references[0];" + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(!JdkCodeCompiler.referencesUnnarrowableClass(body), + "an unevaluable token must not route the unit to Janino") + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(newCodeAndComment(body)) eq JdkCodeCompiler) + } + // The rewrite degrades to the binary name instead of propagating the Error. + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === body) + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("reflection that fails only on member enumeration does not narrow the reference") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The dangerous asymmetry: a class can be missing a type named in one of its method + // signatures, so the supertype climb succeeds while `getMethods` throws. Narrowing then + // looks safe, since the emitted supertype name compiles, even though no member was ever + // checked, and an overload collision would bind to the supertype's method. The verdict + // has to keep the binary name instead, which javac rejects and Spark falls back from. + val dir = compileMemberLinkageFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val anon = loader.loadClass("lnk2.Holder$1") + val body = s"${anon.getName} v = (${anon.getName}) references[0];" + + // Positive control: while the classpath is complete the reference IS narrowed, which + // proves the token reaches the verdict rather than being inert in this body. + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === + "lnk2.Foo v = (lnk2.Foo) references[0];", + "with a complete classpath the reference must narrow to the interface") + + // Now break only the method signature's parameter type. A fresh loader is needed because + // the one above has already resolved members for this class. + assert(new File(dir, "lnk2/Missing.class").delete(), "fixture setup: signature type") + val brokenLoader = new DirClassLoader(dir, getClass.getClassLoader) + val brokenAnon = brokenLoader.loadClass("lnk2.Holder$1") + // Preconditions: the climb reads fine, only member enumeration throws. + assert(brokenAnon.getCanonicalName == null, "fixture must be unnameable") + assert(brokenAnon.getInterfaces.map(_.getName).contains("lnk2.Foo"), "climb input must read") + intercept[LinkageError](brokenAnon.getMethods) + + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(brokenLoader) + // Not routed: an unevaluable class is no evidence that narrowing is unsafe. + assert(!JdkCodeCompiler.referencesUnnarrowableClass(body), + "an unevaluable token must not route the unit to Janino") + // And not narrowed either: narrowing to lnk2.Foo would compile and hide the risk. + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, brokenLoader) === body, + "an unevaluable class must keep its binary name rather than be narrowed") + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(newCodeAndComment(body)) eq JdkCodeCompiler) + } + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("reflection that fails with a non-LinkageError does not narrow the reference") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The other half of the same failure: a classloader is user code, so it may reject a + // class with an ordinary RuntimeException instead of ClassNotFoundException, as a + // relocating or artifact loader can, and the JVM propagates that out of `getMethods` + // unwrapped rather than as a LinkageError. The verdict has to treat it the same way, + // which is why its catch does not stop at LinkageError. + val dir = compileMemberLinkageFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) { + override def findClass(name: String): Class[_] = + if (name == "lnk2.Missing") throw new IllegalStateException("relocated away") + else super.findClass(name) + } + val anon = loader.loadClass("lnk2.Holder$1") + // Preconditions. What matters is not the exact type but that it is NOT a LinkageError, + // since that is the property the verdict's `NonFatal` arm exists for; the enumeration + // that throws is the one over `lnk2.Foo`, which declares the missing parameter type. + assert(anon.getCanonicalName == null, "fixture must be unnameable") + assert(anon.getInterfaces.map(_.getName).contains("lnk2.Foo"), "climb input must read") + intercept[IllegalStateException](anon.getMethods) + + val body = s"${anon.getName} v = (${anon.getName}) references[0];" + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(!JdkCodeCompiler.referencesUnnarrowableClass(body), + "an unevaluable token must not route the unit to Janino") + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === body, + "an unevaluable class must keep its binary name rather than be narrowed") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + /** + * Compile `nrw.Holder`, whose local class holds a named inner class extending `ArrayList`. + * `Holder$<n>Local$Inner` is a member class of a local class: neither anonymous nor local + * itself, with a null canonical name. javac's outer reference `this$0` is package-private + * and comes with no accessor method, so `getFields`/`getMethods` report nothing beyond + * `ArrayList`'s and the class is narrowable. + */ + private def compileNestedLocalFixture(): File = { + val dir = Utils.createTempDir() + val src = + """package nrw; + |public class Holder { + | public static Object make() { + | class Local { + | public class Inner extends java.util.ArrayList<String> { } + | } + | return new Local().new Inner(); + | } + |} + |""".stripMargin + compileJavaFixtures(dir, Seq("nrw/Holder.java" -> src)) + dir + } + + /** + * Compile `shd.Holder`, whose anonymous subclass redeclares its supertype's public field. + * `getFields` reports both, so a field check that compares names rather than declaring + * classes accepts the pair, and narrowing then reads the supertype's value, since field + * access is resolved statically. Java, not Scala: a Scala `val` compiles to a private + * field plus an accessor, so it cannot shadow a public field. + * + * `makePublic` carries the same shadowing pair on a member class of a local class. javac + * strips `ACC_PUBLIC` from an anonymous class, which Janino then refuses to reference from + * the generated unit's package, so an end-to-end test needs this shape instead, for the same + * reason [[compileStaticHiderFixture]] uses one. + */ + private def compileShadowedFieldFixture(): File = { + val dir = Utils.createTempDir() + val src = + """package shd; + |public class Holder { + | public static class Base { public int shadowed = 1; } + | public static Object make() { return new Base() { public int shadowed = 99; }; } + | public static Object makePublic() { + | class Local { + | public class Inner extends Base { public int shadowed = 99; } + | } + | return new Local().new Inner(); + | } + |} + |""".stripMargin + compileJavaFixtures(dir, Seq("shd/Holder.java" -> src)) + dir + } + + /** + * Compile `sth.Holder`, whose member-of-local class hides its supertype's public static + * method. `getMethods` reports both, with identical erased signatures, so a signature-based + * check accepts the pair, and a narrowed call binds the supertype's method, since a static + * call is bound statically. A member class of a local class is used rather than an anonymous + * one because only the former keeps `ACC_PUBLIC`, which the Janino path needs. + */ + private def compileStaticHiderFixture(): File = { + val dir = Utils.createTempDir() + val src = + """package sth; + |public class Holder { + | public static class Base { public static int hidden() { return 1; } } + | public static Object make() { + | class Local { + | public class Inner extends Base { public static int hidden() { return 99; } } + | } + | return new Local().new Inner(); + | } + | public static Object makePlain() { + | class Local2 { public class Inner extends Base { } } + | return new Local2().new Inner(); + | } + |} + |""".stripMargin + compileJavaFixtures(dir, Seq("sth/Holder.java" -> src)) + dir + } + + /** + * Compile `lnk2.Holder`, whose anonymous `Foo` declares a method taking `lnk2.Missing`. + * Deleting `Missing.class` afterwards leaves the anonymous class loadable with a readable + * supertype while `getMethods` throws: the asymmetric partial-jar shape. + */ + private def compileMemberLinkageFixture(): File = { + val dir = Utils.createTempDir() + compileJavaFixtures(dir, Seq( + "lnk2/Missing.java" -> "package lnk2; public class Missing {}", + "lnk2/Foo.java" -> "package lnk2; public interface Foo { String extra(Missing m); }", + "lnk2/Holder.java" -> + """package lnk2; + |public class Holder { + | public static Object make() { + | return new Foo() { public String extra(Missing m) { return "e"; } }; + | } + |} + |""".stripMargin)) + dir + } + + /** Compile Java source strings, given as (path, source) pairs, into `dir`. */ + private def compileJavaFixtures(dir: File, sources: Seq[(String, String)]): Unit = { + val compiler = ToolProvider.getSystemJavaCompiler + val fm = compiler.getStandardFileManager(null, null, null) + try { + fm.setLocation(StandardLocation.CLASS_OUTPUT, Collections.singletonList(dir)) + val files = sources.map { case (path, src) => + new SimpleJavaFileObject(URI.create(s"string:///$path"), JavaFileObject.Kind.SOURCE) { + override def getCharContent(ignoreEncodingErrors: Boolean): CharSequence = src + }.asInstanceOf[JavaFileObject] + } + assert(compiler.getTask(null, fm, null, null, null, files.asJava).call(), + s"failed to compile fixture: ${sources.map(_._1).mkString(", ")}") + } finally { + fm.close() + } + } + + /** + * Compile `lnk.Holder`, whose nested `Mid` holds an anonymous `Runnable`. Deleting + * `Holder$Mid.class` afterwards leaves `Holder$Mid$1` loadable but makes reflection over + * its enclosing class throw, which is the partial-jar shape the guards have to survive. + */ + private def compileLinkageFixture(): File = { + val dir = Utils.createTempDir() + val src = + """package lnk; + |public class Holder { + | public static class Mid { + | public static Object make() { return new Runnable() { public void run() {} }; } + | } + |} + |""".stripMargin + compileJavaFixtures(dir, Seq("lnk/Holder.java" -> src)) + dir + } + + private def compileDynHelper(): File = { + val dir = Utils.createTempDir() + val src = + "package dyn; public class DynHelper { public static int magic() { return 4242; } }" + compileJavaFixtures(dir, Seq("dyn/DynHelper.java" -> src)) + dir + } + + // A loader that is NOT a URLClassLoader (so the retired -classpath harvesting would + // miss it) but loads classes and serves their resources from `dir` - the shape of the + // loaders Spark uses for REPL-generated / Connect session-artifact classes. + private class DirClassLoader(dir: File, parent: ClassLoader) extends ClassLoader(parent) { + override def findClass(name: String): Class[_] = { + val f = new File(dir, name.replace('.', '/') + ".class") + if (!f.isFile) throw new ClassNotFoundException(name) + val bytes = java.nio.file.Files.readAllBytes(f.toPath) + defineClass(name, bytes, 0, bytes.length) + } + override def findResource(name: String): URL = { + val f = new File(dir, name) + if (f.exists) f.toURI.toURL else null + } + override def findResources(name: String): java.util.Enumeration[URL] = { + val f = new File(dir, name) + if (f.exists) Collections.enumeration(Collections.singletonList(f.toURI.toURL)) + else Collections.emptyEnumeration() + } + } + + // Like DirClassLoader but serves only individual .class resources by name and does + // NOT implement findResources, so getResources(package) yields nothing - the shape of + // an in-memory REPL / Connect session loader that cannot enumerate its packages. + private class NonEnumerableDirClassLoader(dir: File, parent: ClassLoader) + extends ClassLoader(parent) { + override def findClass(name: String): Class[_] = { + val f = new File(dir, name.replace('.', '/') + ".class") + if (!f.isFile) throw new ClassNotFoundException(name) + val bytes = java.nio.file.Files.readAllBytes(f.toPath) + defineClass(name, bytes, 0, bytes.length) + } + override def findResource(name: String): URL = { + val f = new File(dir, name) + if (name.endsWith(".class") && f.isFile) f.toURI.toURL else null + } + } + + // Serves .class bytes ONLY via getResourceAsStream, with no resource URL + // (getResource is null) and no enumeration - the shape of an in-memory REPL / + // Ammonite loader. The class is still loadable (findClass) and readable as a stream. + private class StreamOnlyDirClassLoader(dir: File, parent: ClassLoader) + extends ClassLoader(parent) { + override def findClass(name: String): Class[_] = { + val f = new File(dir, name.replace('.', '/') + ".class") + if (!f.isFile) throw new ClassNotFoundException(name) + val bytes = java.nio.file.Files.readAllBytes(f.toPath) + defineClass(name, bytes, 0, bytes.length) + } + override def getResourceAsStream(name: String): java.io.InputStream = { + val f = new File(dir, name) + if (name.endsWith(".class") && f.isFile) new java.io.FileInputStream(f) + else super.getResourceAsStream(name) + } + } + + // Returns a non-null, non-class stream for the package-shaped path + // `org/apache/spark/sql/catalyst/expressions.class` (the generated unit's own + // package), reproducing the artifact-loader behaviour that surfaced a phantom class + // clashing with that package. All other resources delegate to the parent. + private class PhantomPackageClassLoader(parent: ClassLoader) extends ClassLoader(parent) { + override def getResourceAsStream(name: String): java.io.InputStream = { + if (name == "org/apache/spark/sql/catalyst/expressions.class") { + new java.io.ByteArrayInputStream(Array[Byte]('n', 'o', 't')) + } else { + super.getResourceAsStream(name) + } + } + } + + // ---------------- error paths ---------------- + + test("Janino backend surfaces compile errors as QueryExecutionErrors.compilerError") { + val malformedBody = + "public boolean evaluate() { return missing_identifier; }" + val ex = intercept[CompileException] { + JaninoCodeCompiler.compile(newCodeAndComment(malformedBody)) + } + // The "Failed to compile:" prefix is what QueryExecutionErrors.compilerError adds; + // a raw Janino CompileException would not carry it. + assert(ex.getMessage.contains("Failed to compile:")) + } + + test("JDK backend surfaces compile errors with the same exception type as Janino") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val malformedBody = + "public boolean evaluate() { return missing_identifier; }" + val janinoEx = intercept[Exception] { + JaninoCodeCompiler.compile(newCodeAndComment(malformedBody)) + } + val jdkEx = intercept[Exception] { + JdkCodeCompiler.compile(newCodeAndComment(malformedBody)) + } + // Both backends surface source errors through QueryExecutionErrors.compilerError, so + // callers matching on the exception type behave identically whichever is active. + assert(jdkEx.getClass === janinoEx.getClass, + s"backend exception types diverge: janino=${janinoEx.getClass}, jdk=${jdkEx.getClass}") + assert(jdkEx.isInstanceOf[CompileException]) + // The diagnostic names the offending identifier so users can locate the failure. + assert(jdkEx.getMessage.contains("missing_identifier"), + s"diagnostic lost the offending symbol:\n${jdkEx.getMessage}") + } + + // ---------------- end-to-end through GeneratePredicate ---------------- + // + // These exercise a real Spark code generator under each backend. They smoke-test + // that each backend is compatible with code shapes Spark generators actually + // produce (nested classes, references arrays, multiple methods, the GeneratedClass + // contract). The backend cache key includes the backend, so withSQLConf + // guarantees a fresh compilation regardless of prior in-JVM state. + + test("end-to-end: GeneratePredicate with Janino backend evaluates correctly") { + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> CodeCompiler.JANINO) { + val predicate = GeneratePredicate.generate( + GreaterThan(BoundReference(0, IntegerType, nullable = false), Literal(5))) + assert(predicate.eval(InternalRow(10)) === true) + assert(predicate.eval(InternalRow(3)) === false) + } + } + + test("end-to-end: GeneratePredicate with JDK backend evaluates correctly") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> CodeCompiler.JDK) { + val predicate = GeneratePredicate.generate( + LessThan(BoundReference(0, IntegerType, nullable = false), Literal(5))) + assert(predicate.eval(InternalRow(3)) === true) + assert(predicate.eval(InternalRow(10)) === false) + } + } + + test("CodeGenerator.compile caches the same source separately per backend") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The backend is part of the cache key; if it were dropped, flipping the config + // would silently serve the other backend's cached class for identical source. + val code = newCodeAndComment(sampleClassBody) + val (janinoGenerated, _) = withSQLConf(SQLConf.CODEGEN_COMPILER.key -> CodeCompiler.JANINO) { + CodeGenerator.compile(code) + } + val (jdkGenerated, _) = withSQLConf(SQLConf.CODEGEN_COMPILER.key -> CodeCompiler.JDK) { + CodeGenerator.compile(code) + } + assert(janinoGenerated.getClass.getClassLoader ne jdkGenerated.getClass.getClassLoader, + "each backend must produce (and cache) its own compilation of the same source") + } + + // ---------------- wrapAsCompilationUnit shape ---------------- + + test("JDK backend's wrapAsCompilationUnit produces a well-formed source unit") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val wrapped = JdkCodeCompiler.wrapAsCompilationUnit(sampleClassBody, getClass.getClassLoader) + assert(wrapped.startsWith("package org.apache.spark.sql.catalyst.expressions;"), + s"missing package declaration:\n$wrapped") + assert(wrapped.contains("import org.apache.spark.unsafe.Platform;"), + s"missing default imports:\n$wrapped") + assert(wrapped.contains( + "public class GeneratedClass extends " + + "org.apache.spark.sql.catalyst.expressions.codegen.GeneratedClass"), + s"missing class declaration:\n$wrapped") + } + + // ---------------- rewriteInnerClassRefs ---------------- + + // Classloader used to resolve candidate type references in the tests below. + private val rewriteLoader: ClassLoader = getClass.getClassLoader + private def rewrite(body: String): String = + JdkCodeCompiler.rewriteInnerClassRefs(body, rewriteLoader) + + test("rewriteInnerClassRefs: converts binary inner-class refs to dotted form") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // Unresolvable synthetic names fall back to the conservative regex. + assert(rewrite("a.b.Outer$Inner x;") === "a.b.Outer.Inner x;") + // Doubly-nested names rewrite every separator. + assert(rewrite("Outer$Mid$Inner") === "Outer.Mid.Inner") + } + + test("rewriteInnerClassRefs: restores the trailing dot of a line-wrapped member access") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A member access wrapped onto the next line scans as a token ending in '.', which + // `split('.')` would otherwise drop along with the trailing empty segment. + assert(rewrite("java.util.Map$Entry.\ncomparingByKey()") === + "java.util.Map.Entry.\ncomparingByKey()") + } + + test("rewriteInnerClassRefs: dots regular nested classes resolved via reflection") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // java.util.Map$Entry is a real class; its canonical name java.util.Map.Entry + // is a plain dotted name, so the reflection path emits the dotted form. A + // trailing member access is preserved. + assert(rewrite("java.util.Map$Entry e;") === "java.util.Map.Entry e;") + assert(rewrite("java.util.Map$Entry.class") === "java.util.Map.Entry.class") + } + + test("rewriteInnerClassRefs: preserves binary names of object-nested classes") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A case class nested inside a Scala object has a canonical name carrying a + // module `$` (Outer.SaveLoadV1$.Leaf) that the JDK compiler cannot resolve; + // the binary name must be kept verbatim. + val binary = classOf[CodeCompilerSuite.SaveLoadV1.Leaf].getName + assert(binary.count(_ == '$') === 2, s"unexpected binary form: $binary") + assert(rewrite(s"$binary x = null;") === s"$binary x = null;") + // Companion-object access on such a class is also preserved (longest loadable + // prefix is the companion class, trailing MODULE$.apply is kept). + assert(rewrite(s"$binary$$.MODULE$$.apply(1)") === s"$binary$$.MODULE$$.apply(1)") + } + + test("rewriteInnerClassRefs: preserves Scala companion and mangled names") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // `$` followed by `.` (companion access) on an unresolvable name must not be touched. + assert(rewrite("Foo$.MODULE$.apply()") === "Foo$.MODULE$.apply()") + // A real top-level operator-named class: canonical is `scala.collection.immutable.::` + // which is not a valid Java identifier, so the binary name is kept. + assert(rewrite("scala.collection.immutable.$colon$colon") === + "scala.collection.immutable.$colon$colon") + } + + test("rewriteInnerClassRefs: does not corrupt string or char literals or comments") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A `$Upper` sequence inside a string literal (e.g. a column name) is preserved. + assert(rewrite("""String s = "col$Name";""") === """String s = "col$Name";""") + // Escaped quote inside the string does not end it early. + assert(rewrite("""x = "a\"b$Cd"; Outer$Inner y;""") === """x = "a\"b$Cd"; Outer.Inner y;""") + // Char literal preserved. + assert(rewrite("""char c = '$'; Outer$Inner z;""") === """char c = '$'; Outer.Inner z;""") + // Line and block comments preserved. + assert(rewrite("// see Foo$Bar\nOuter$Inner w;") === "// see Foo$Bar\nOuter.Inner w;") + assert(rewrite("/* Foo$Bar */ Outer$Inner w;") === "/* Foo$Bar */ Outer.Inner w;") + } + + test("rewriteInnerClassRefs: rewrites anonymous-class refs to a nameable supertype") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A Scala `new HashMap() {...}` compiles to an anonymous class (`...$$anon$N`) that + // cannot be named in Java source; the JDK compiler rejects a qualified reference to + // it even with the bytecode present. It must be rewritten to its nearest nameable + // supertype (java.util.HashMap), which is a sound cast target for the generated code. + val anon = new java.util.HashMap[String, String]() {} + val anonName = anon.getClass.getName + assert(anon.getClass.isAnonymousClass, s"expected an anonymous class, got: $anonName") + assert(rewrite(s"$anonName m = ($anonName) references[0];") === + "java.util.HashMap m = (java.util.HashMap) references[0];") + } + + test("rewriteInnerClassRefs: anonymous interface impl rewrites to the interface") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // An anonymous class implementing only an interface has Object as its superclass; + // the nameable-supertype climb must pick the interface, not Object. + val anon = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = a.compareTo(b) + } + val anonName = anon.getClass.getName + assert(anon.getClass.isAnonymousClass || anon.getClass.isLocalClass, + s"expected an anonymous/local class, got: $anonName") + assert(rewrite(s"$anonName c = ($anonName) references[0];") === + "java.util.Comparator c = (java.util.Comparator) references[0];") + } + + test("JDK backend resolves an anonymous class reference via its nameable supertype") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // Mirrors SPARK-23589 (ExternalMapToCatalyst over an anonymous java.util.HashMap): + // codegen casts the literal object to its anonymous binary name. Janino loads that + // name directly; the JDK backend must rewrite it to a nameable supertype and still + // produce working bytecode. + val anon = new java.util.HashMap[String, String]() { + put("k", "v") + } + val anonName = anon.getClass.getName + val body = + s""" + |public java.lang.Object generate(Object[] references) { + | $anonName m = ($anonName) references[0]; + | return m.get("k"); + |} + """.stripMargin + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array[Any](anon)) === "v") + } + + // ---------------- unnarrowable anonymous/local classes ---------------- + + test("containsDollarDigit gates the scan on names Java cannot spell") { + // The `$`-digit gate must admit every unnameable shape and reject the nameable `$` + // forms the rewrite handles, or the scan either misses a class or runs needlessly. + for (name <- Seq("Outer$1", "Outer$1Local", "Outer$$anon$1", "Outer$1$Inner", + "a.b.Outer$$anon$12")) { + assert(JdkCodeCompiler.containsDollarDigit(name), s"expected $name to be gated in") + } + for (name <- Seq("", "$", "a$", "java.util.Map$Entry", "Foo$", "Model$SaveLoad$Leaf", + "scala.Function1$mcII$sp", "scala.collection.immutable.$colon$colon", + "pkg.package$Inner", "$line21.$read$$iw$T")) { + assert(!JdkCodeCompiler.containsDollarDigit(name), s"expected $name to be gated out") + } + } + + test("referencesUnnarrowableClass: false when the supertype offers every member") { + // The shapes codegen actually produces: the anonymous or local class only overrides + // methods the supertype already declares, so narrowing the reference loses nothing. + // `specializedPrimitiveOverride` is the boxing case - scalac specializes the override to + // `apply(int)` while the bridge that reaches it takes `Object`. `anonWithLambdaBody` is + // the synthetic-static case: a closure in the body adds a `public static final + // $anonfun$...` helper javac cannot name, which must not make the class unnarrowable. + val anonSubclass = new java.util.HashMap[String, String]() { put("k", "v") } + val anonInterface = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = a.compareTo(b) + } + val narrowable = Seq[Any](anonSubclass, anonInterface, CodeCompilerSuite.plainLocal, + CodeCompilerSuite.specializedPrimitiveOverride, CodeCompilerSuite.anonWithLambdaBody) + for (o <- narrowable) { + val cls = o.getClass + // Guard the fixture's own precondition: a nameable class would pass the assertion + // below for the wrong reason. + assert(cls.getCanonicalName == null, s"expected an unnameable class, got: ${cls.getName}") + val name = cls.getName + assert(!JdkCodeCompiler.referencesUnnarrowableClass(s"$name v = ($name) references[0];"), + s"expected $name to be narrowable") + } + } + + test("referencesUnnarrowableClass: a synthetic static does not make a class unnarrowable") { + // Pins the fixture's own shape, so the case above cannot pass by accident if a future + // scalac stops emitting the helper as a public static on the anonymous class. + val cls = CodeCompilerSuite.anonWithLambdaBody.getClass + val statics = cls.getMethods.filter(m => Modifier.isStatic(m.getModifiers)) + assert(statics.nonEmpty, s"fixture must carry a public static, got none on ${cls.getName}") + assert(statics.forall(_.isSynthetic), + s"fixture's statics must all be synthetic, got: ${statics.mkString(", ")}") + assert(statics.exists(_.getDeclaringClass eq cls), + s"at least one static must be declared by the anonymous class itself, got: " + + statics.map(_.getDeclaringClass.getName).mkString(", ")) + } + + test("referencesUnnarrowableClass: true when narrowing would lose access") { + // An extra public method, public fields inherited from a second interface, a member on + // a second interface, an overload that shadows nothing, a local class with an extra + // method, a bridged override sharing its name and arity with an unrelated overload, and + // an instance method whose only counterpart on the supertype is static: each puts + // something out of reach of the nearest nameable supertype. + val unnarrowable = Seq[Any]( + CodeCompilerSuite.anonWithExtraMethod, + CodeCompilerSuite.anonWithPublicFields, + CodeCompilerSuite.anonWithSecondInterface, + CodeCompilerSuite.anonWithOverload, + CodeCompilerSuite.localWithExtraMethod, + CodeCompilerSuite.anonWithBridgeAndPrimitiveOverload, + CodeCompilerSuite.anonWithBridgeAndReferenceOverload, + CodeCompilerSuite.anonOverStaticClash) + for (o <- unnarrowable) { + val cls = o.getClass + assert(cls.getCanonicalName == null, s"expected an unnameable class, got: ${cls.getName}") + val name = cls.getName + assert(JdkCodeCompiler.referencesUnnarrowableClass(s"$name v = ($name) references[0];"), + s"expected $name to be rejected") + } + } + + test("referencesUnnarrowableClass: a bridge covers only the method it forwards to") { + // Guards the fixture shapes behind the two bridge tests: both classes carry a genuine + // bridge for the generic override, so a check that excused the whole name/arity group + // would let the unrelated overload ride along on it. + for (o <- Seq[Any](CodeCompilerSuite.anonWithBridgeAndPrimitiveOverload, + CodeCompilerSuite.anonWithBridgeAndReferenceOverload)) { + val compares = o.getClass.getMethods.filter(_.getName == "compare") + assert(compares.exists(_.isBridge), "fixture precondition: expected a bridge method") + assert(compares.count(m => !m.isBridge && m.getParameterCount == 2) === 2, + "fixture precondition: expected the override and the overload to share name and arity") + } + } + + test("referencesUnnarrowableClass: true when the supertype itself cannot be named") { + // Members line up here, but the nearest nameable supertype is a private nested class + // (scala.collection.mutable.HashSet$HashSetIterator), so javac could not write the + // narrowed cast at all. + val cls = scala.collection.mutable.HashSet("a").iterator.getClass + assert(cls.getCanonicalName == null, s"expected an unnameable class, got: ${cls.getName}") + val name = cls.getName + assert(JdkCodeCompiler.referencesUnnarrowableClass(s"$name v = ($name) references[0];"), + s"expected $name to be rejected for an unnameable supertype") + } + + test("referencesUnnarrowableClass: a field shadowing the supertype's cannot be narrowed") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The one silent-wrong-answer shape that methods do not have. `getFields` reports both + // the shadowing and the shadowed field, so comparing names would accept the pair, but + // field access is resolved statically: the narrowed reference would read the supertype's + // 1 instead of the anonymous class's 99, compiling cleanly the whole way. + val dir = compileShadowedFieldFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val anon = loader.loadClass("shd.Holder").getMethod("make").invoke(null).getClass + assert(anon.getCanonicalName == null, s"expected an unnameable class, got: ${anon.getName}") + // Fixture preconditions: the shadowing pair is exactly what a name check would miss. + val vs = anon.getFields.filter(_.getName == "shadowed") + assert(vs.length === 2, s"expected two public `shadowed` fields: ${vs.mkString(", ")}") + assert(vs.exists(_.getDeclaringClass eq anon), + "one `shadowed` must be declared by the anonymous class") + + val body = s"${anon.getName} v = (${anon.getName}) references[0];" + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JdkCodeCompiler.referencesUnnarrowableClass(body), + "a shadowed field must make the class unnarrowable") + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === body, + "an unnarrowable class must keep its binary name") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("referencesUnnarrowableClass: a static method hiding the supertype's cannot be narrowed") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The method counterpart of the field case: `hidden()` erases identically on both + // classes, so a signature check accepts the pair, but a static call is bound statically: + // the narrowed reference would call the supertype's `hidden()` for 1 instead of 99, and + // Java permits the instance-qualified form so even a plain cast rebinds. + val dir = compileStaticHiderFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val inner = loader.loadClass("sth.Holder").getMethod("make").invoke(null).getClass + assert(inner.getCanonicalName == null, s"expected an unnameable class, got: ${inner.getName}") + // Fixture preconditions: `f` is static and declared here, while the target declares an + // identically-erased static of its own: the pair a declaring-class check must reject and + // a signature check would accept. (`getMethods` reports only the hiding one; a static is + // hidden, not inherited.) + val fs = inner.getMethods.filter(_.getName == "hidden") + assert(fs.length === 1, s"expected one visible `f`, got: ${fs.mkString(", ")}") + assert((fs.head.getDeclaringClass eq inner) && Modifier.isStatic(fs.head.getModifiers), + s"`f` must be a static declared by the member class, got: ${fs.head}") + val base = loader.loadClass("sth.Holder$Base") + assert(base.getMethods.exists(m => + m.getName == "hidden" && m.getParameterCount == 0 && Modifier.isStatic(m.getModifiers)), + "the target must declare the same static signature, or nothing would be hidden") + assert(inner.getSuperclass eq base, s"expected Base as superclass, got ${inner.getSuperclass}") + // And the complement, so both branches of the declaring-class check are pinned: a static + // the class merely INHERITS from the target stays narrowable. + val plain = loader.loadClass("sth.Holder").getMethod("makePlain").invoke(null).getClass + assert(plain.getCanonicalName == null, s"expected an unnameable class, got: ${plain.getName}") + assert(plain.getMethods.exists(m => + m.getName == "hidden" && Modifier.isStatic(m.getModifiers) && + (m.getDeclaringClass eq base)), + "the complement fixture must inherit the target's static rather than hide it") + + val body = s"${inner.getName} v = (${inner.getName}) references[0];" + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JdkCodeCompiler.referencesUnnarrowableClass(body), + "a hidden static method must make the class unnarrowable") + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === body, + "an unnarrowable class must keep its binary name") + val plainBody = s"${plain.getName} v = (${plain.getName}) references[0];" + assert(!JdkCodeCompiler.referencesUnnarrowableClass(plainBody), + "merely inheriting the target's static must stay narrowable") + assert(JdkCodeCompiler.rewriteInnerClassRefs(plainBody, loader) === + s"sth.Holder.Base v = (sth.Holder.Base) references[0];", + "the complement fixture must narrow to the target") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("rewriteInnerClassRefs: narrows a class nested inside a local class") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A named class declared inside a local class is neither anonymous nor local, so the + // climb has to key off `getCanonicalName == null` to catch it. Its outer reference is + // reachable from neither side: javac's `this$0` is package-private with no accessor, and + // scalac's `$outer` pair is synthetic, so the member check sees nothing beyond + // ArrayList's either way (the two tests below cover the Scala shape). + val dir = compileNestedLocalFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + // Derive the binary name rather than hardcoding it: the JLS fixes the shape + // (`Holder$<digits>Local$Inner`) but leaves the digit sequence to the compiler. + val cls = loader.loadClass("nrw.Holder") + .getMethod("make").invoke(null).getClass + assert(!cls.isAnonymousClass && !cls.isLocalClass && cls.isMemberClass, + s"expected a member class of a local class, got: ${cls.getName}") + assert(cls.getCanonicalName == null, s"expected no canonical name for ${cls.getName}") + val name = cls.getName + val body = s"$name v = ($name) references[0];" + // Fixture guard: this shape must stay on the JDK backend, or the rewrite is moot. The + // context loader has to be swapped because the fixture lives only in a temp dir. + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(!JdkCodeCompiler.referencesUnnarrowableClass(body), + "fixture must be narrowable, otherwise it would route to Janino instead") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + assert(JdkCodeCompiler.rewriteInnerClassRefs(body, loader) === + "java.util.ArrayList v = (java.util.ArrayList) references[0];") + } + + test("rewriteInnerClassRefs: keeps the binary name of an unnarrowable nested local class") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The complement of the test above: the same shape plus a non-synthetic public method + // `ArrayList` does not have, so narrowing would put it out of reach. The rewrite refuses + // rather than emitting a name that compiles but drops the member. The `$outer` field and + // accessor are NOT what makes this unnarrowable: being synthetic, they are unreachable + // from generated source either way. No context-loader swap is needed here: this fixture + // is on the suite's own classpath, unlike the temp-dir one above. + val cls = CodeCompilerSuite.memberOfLocalClassWithExtra.getClass + assert(cls.getCanonicalName == null, s"expected no canonical name for ${cls.getName}") + assert(cls.getMethods.exists(m => m.getName == "extra" && !m.isSynthetic), + "fixture must carry a non-synthetic public method absent from ArrayList") + val name = cls.getName + val body = s"$name v = ($name) references[0];" + assert(JdkCodeCompiler.referencesUnnarrowableClass(body), + "a member class with an extra public method must be reported unnarrowable") + assert(rewrite(body) === body, "an unnarrowable class must keep its binary name") + } + + test("referencesUnnarrowableClass: synthetic outer accessors do not block narrowing") { + // `memberOfLocalClass` carries scalac's public `$outer` field and accessor, both synthetic + // and therefore invisible to javac's source lookup ("cannot find symbol"), so nothing a + // generator emits can reach them and narrowing to `ArrayList` loses nothing. + val cls = CodeCompilerSuite.memberOfLocalClass.getClass + assert(cls.getCanonicalName == null, s"expected no canonical name for ${cls.getName}") + val outerFields = cls.getFields.filter(_.getName.contains("outer")) + assert(outerFields.nonEmpty && outerFields.forall(_.isSynthetic), + s"fixture must carry a synthetic public outer field, got: ${outerFields.mkString(", ")}") + val extraMethods = cls.getMethods.filterNot(m => + classOf[java.util.ArrayList[_]].getMethods.exists(_.getName == m.getName)) + assert(extraMethods.nonEmpty && extraMethods.forall(_.isSynthetic), + s"fixture's extra methods must all be synthetic, got: ${extraMethods.mkString(", ")}") + val name = cls.getName + assert(!JdkCodeCompiler.referencesUnnarrowableClass(s"$name v = ($name) references[0];"), + "synthetic-only extras must not make a class unnarrowable") + } + + test("active(code) routes an unnarrowable class reference to Janino") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val narrowableName = (new java.util.HashMap[String, String]() { put("k", "v") }) + .getClass.getName + val unnarrowableName = CodeCompilerSuite.anonWithExtraMethod.getClass.getName + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(newCodeAndComment(s"$narrowableName v;")) eq JdkCodeCompiler) + assert(CodeCompiler.active(newCodeAndComment(s"$unnarrowableName v;")) eq JaninoCodeCompiler) + } + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "janino") { + assert(CodeCompiler.active(newCodeAndComment(s"$unnarrowableName v;")) eq JaninoCodeCompiler) + } + } + + test("JDK backend cannot compile a member access that narrowing drops") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // Both halves of the routing decision: javac rejects the rewritten unit (the reference + // narrows to java.util.HashMap, which has no `extra()`), and `active` therefore hands + // the unit to Janino, which compiles it and produces the right answer. + val anon = CodeCompilerSuite.anonWithExtraMethod + val anonName = anon.getClass.getName + val code = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $anonName m = ($anonName) references[0]; + | return m.extra(); + |} + """.stripMargin) + intercept[CompileException] { + JdkCodeCompiler.compile(code) + } + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(code) eq JaninoCodeCompiler) + val (generated, _) = CodeGenerator.compile(code) + assert(generated.generate(Array[Any](anon)) === "extra") + } + } + + test("JDK backend compiles a call that narrowing preserves through a bridge") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The counterpart to the test above: an override of a generic method erases narrower + // than the interface declares, but the compiler-emitted bridge keeps dispatch correct, + // so this must stay on the JDK backend rather than being routed away. + val anon = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = a.compareTo(b) + } + val anonName = anon.getClass.getName + val code = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $anonName c = ($anonName) references[0]; + | return Integer.valueOf(c.compare("a", "b")); + |} + """.stripMargin) + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(code) eq JdkCodeCompiler) + } + val (generated, _) = JdkCodeCompiler.compile(code) + assert(generated.generate(Array[Any](anon)) === -1) + } + + test("end-to-end: a shadowing field reads the class's own value, not the target's") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The verdict tests pin the decision; this pins the outcome. `Inner` redeclares `Base`'s + // public `shadowed`, and a field read is bound statically, so the two names answer + // differently at run time: 99 through the class, 1 through the supertype. Routing keeps 99. + val dir = compileShadowedFieldFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val value = loader.loadClass("shd.Holder").getMethod("makePublic").invoke(null) + val name = value.getClass.getName + assert(value.getClass.getCanonicalName == null, s"expected an unnameable class: $name") + val hazard = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $name v = ($name) references[0]; + | return Integer.valueOf(v.shadowed); + |} + """.stripMargin) + // The counterfactual: the same read spelled with the name narrowing would emit. Janino + // compiles it too, and it answers 1, which is exactly why it may not be emitted. + val narrowed = newCodeAndComment( + """ + |public java.lang.Object generate(Object[] references) { + | shd.Holder.Base v = (shd.Holder.Base) references[0]; + | return Integer.valueOf(v.shadowed); + |} + """.stripMargin) + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + // Janino names the class itself, so it binds the class's own field. + assert(JaninoCodeCompiler.compile(hazard)._1.generate(Array[Any](value)) === 99) + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(hazard) eq JaninoCodeCompiler) + val (generated, _) = CodeGenerator.compile(hazard) + assert(generated.generate(Array[Any](value)) === 99, + "routing must preserve the shadowing field's value") + } + assert(JaninoCodeCompiler.compile(narrowed)._1.generate(Array[Any](value)) === 1, + "counterfactual: the narrowed name reads the supertype's field") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("end-to-end: a hidden static call reaches the class's own method, not the target's") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The method counterpart. `Inner.hidden()` hides `Base.hidden()` with an identical + // erasure, and a static call is bound statically even in the instance-qualified form, so + // the narrowed reference would answer 1 where the class answers 99. + val dir = compileStaticHiderFixture() + val loader = new DirClassLoader(dir, getClass.getClassLoader) + val value = loader.loadClass("sth.Holder").getMethod("make").invoke(null) + val name = value.getClass.getName + assert(value.getClass.getCanonicalName == null, s"expected an unnameable class: $name") + val hazard = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $name v = ($name) references[0]; + | return Integer.valueOf(v.hidden()); + |} + """.stripMargin) + val narrowed = newCodeAndComment( + """ + |public java.lang.Object generate(Object[] references) { + | sth.Holder.Base v = (sth.Holder.Base) references[0]; + | return Integer.valueOf(v.hidden()); + |} + """.stripMargin) + val prev = Thread.currentThread().getContextClassLoader + try { + Thread.currentThread().setContextClassLoader(loader) + assert(JaninoCodeCompiler.compile(hazard)._1.generate(Array[Any](value)) === 99) + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(hazard) eq JaninoCodeCompiler) + val (generated, _) = CodeGenerator.compile(hazard) + assert(generated.generate(Array[Any](value)) === 99, + "routing must preserve the hiding static's value") + } + assert(JaninoCodeCompiler.compile(narrowed)._1.generate(Array[Any](value)) === 1, + "counterfactual: the narrowed name calls the supertype's static") + } finally { + Thread.currentThread().setContextClassLoader(prev) + } + } + + test("end-to-end: an instance method clashing with a static keeps the class's value") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // The third shape: `value()` is an INSTANCE method on the anonymous class whose only + // counterpart on `StaticClashBase` is STATIC. Java permits the instance-qualified form + // for a static, so the narrowed call compiles and binds the target's `value()`, giving 1 + // instead of 7. Both fixtures are on the suite's own classpath, so no loader swap here. + val anon = CodeCompilerSuite.anonOverStaticClash + val name = anon.getClass.getName + assert(anon.getClass.getCanonicalName == null, s"expected an unnameable class: $name") + val target = classOf[StaticClashBase].getName + val hazard = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $name v = ($name) references[0]; + | return Integer.valueOf(v.value()); + |} + """.stripMargin) + val narrowed = newCodeAndComment( + s""" + |public java.lang.Object generate(Object[] references) { + | $target v = ($target) references[0]; + | return Integer.valueOf(v.value()); + |} + """.stripMargin) + assert(JaninoCodeCompiler.compile(hazard)._1.generate(Array[Any](anon)) === 7) + withSQLConf(SQLConf.CODEGEN_COMPILER.key -> "jdk") { + assert(CodeCompiler.active(hazard) eq JaninoCodeCompiler) + val (generated, _) = CodeGenerator.compile(hazard) + assert(generated.generate(Array[Any](anon)) === 7, + "routing must preserve the instance method's value") + } + assert(JdkCodeCompiler.compile(narrowed)._1.generate(Array[Any](anon)) === 1, + "counterfactual: the narrowed name calls the target's static") + } + + // ---------------- Function1 apply(Object) bridge ---------------- + + test("stripFunction1ApplyBridges removes the bridge but keeps apply(InternalRow)") { + val src = + s"""${CodeGenerator.function1ApplyBridge("i")} + |public UnsafeRow apply(InternalRow i) { return null; }""".stripMargin + assert(src.contains("apply(java.lang.Object")) + val stripped = JdkCodeCompiler.stripFunction1ApplyBridges(src) + assert(!stripped.contains("apply(java.lang.Object"), + s"bridge not stripped:\n$stripped") + assert(stripped.contains("apply(InternalRow i)"), s"typed apply lost:\n$stripped") + } + + test("both backends compile a projection with the Function1 apply(Object) bridge") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // A class extending a Scala `InternalRow => *` must carry an explicit apply(Object) + // bridge for Janino (which does not synthesize it) but not for javac (which does, and + // rejects an explicit duplicate). Generators emit the bridge; the JDK backend strips + // it. The same generated source must compile under both backends. + val body = + s""" + |public java.lang.Object generate(Object[] references) { + | return new SpecificProj(); + |} + |static class SpecificProj + | extends org.apache.spark.sql.catalyst.expressions.UnsafeProjection { + | ${CodeGenerator.function1ApplyBridge("i")} + | public UnsafeRow apply(InternalRow i) { return null; } + | public void initialize(int partitionIndex) {} + |} + """.stripMargin + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + assert(JdkCodeCompiler.compile(newCodeAndComment(body))._1 != null) + } + + // ---------------- extractLeadingImports ---------------- + + test("extractLeadingImports: hoists leading imports and leaves the rest") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val body = + """import a.B; + | + |import c.D; + |public void foo() {} + |import e.F;""".stripMargin + val (imports, rest) = JdkCodeCompiler.extractLeadingImports(body) + // The blank line between the leading imports is consumed and not preserved - + // intentional: blank lines are cosmetic in the import block. + assert(imports === "import a.B;\nimport c.D;\n") + // An import after the first non-import line stays in the body. + assert(rest === "public void foo() {}\nimport e.F;") + } + + test("extractLeadingImports: no leading imports leaves body unchanged") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + val body = "public void foo() {}\n" + val (imports, rest) = JdkCodeCompiler.extractLeadingImports(body) + assert(imports === "") + assert(rest === body) + } + + test("JDK backend hoists leading imports and compiles (GenerateColumnAccessor shape)") { + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + // GenerateColumnAccessor emits `import` lines at the top of the class body. Janino + // accepts them there; javac only accepts them at the compilation-unit level, so + // wrapAsCompilationUnit must hoist them - and the result must actually compile. + val body = + """import java.util.ArrayList; + | + |public java.lang.Object generate(Object[] references) { + | ArrayList list = new ArrayList(); + | list.add("ok"); + | return list.get(0); + |}""".stripMargin + assert(JaninoCodeCompiler.compile(newCodeAndComment(body))._1 != null) + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(body)) + assert(generated.generate(Array.empty[Any]) === "ok") + } + + // ---------------- interrupt isolation ---------------- + + test("JDK backend compiles successfully even if the calling thread is interrupted") { + // javac reads classpath jars via interruptible NIO channels; running on a Spark + // task thread whose interrupt flag is set must NOT break compilation. The compile + // runs on a dedicated worker thread, so the caller's interrupt does not reach the + // jar reads. The caller's interrupt status must be preserved on return. + assume(JdkCodeCompiler.isAvailable, "javax.tools.JavaCompiler not available") + Thread.currentThread().interrupt() + try { + val (generated, _) = JdkCodeCompiler.compile(newCodeAndComment(sampleClassBody)) + assert(generated != null) + assert(Thread.currentThread().isInterrupted, + "caller's interrupt status should be preserved") + } finally { + // Clear the interrupt flag so it does not leak into subsequent tests. + Thread.interrupted() + } + } +} + +object CodeCompilerSuite { + // Mirrors the mllib legacy save/load shape (a case class nested inside a Scala + // `object`), whose binary name uses `$` as both the module suffix and the + // nesting separator (CodeCompilerSuite$SaveLoadV1$Leaf) - the case the + // reflection-based rewrite must preserve verbatim. + object SaveLoadV1 { + case class Leaf(x: Int) + } + + private[codegen] trait Greeter { + def hello(): String + } + + private[codegen] abstract class Converter { + def convert(o: Any): String = "base" + } + + // Anonymous and local classes for the narrowing-soundness tests. They are held in vals + // rather than built inline in the tests because scalac keeps an anonymous class's extra + // members `public` only when the binding's type is inferred as the refined type; giving + // the val an explicit type, or passing the expression as `Any`, makes them private. + val anonWithExtraMethod = new java.util.HashMap[String, String]() { + def extra(): String = "extra" + } + + // Mixing in a Java constants interface is what actually yields public FIELDS: a Scala + // `val` compiles to a private field plus an accessor, which only exercises the method + // check. ObjectStreamConstants contributes 30 public static final fields and no method + // beyond Comparator's, so this isolates the field clause of `narrowingVerdict`. + val anonWithPublicFields = new java.util.Comparator[String] with java.io.ObjectStreamConstants { + override def compare(a: String, b: String): Int = a.compareTo(b) + } + + val anonWithSecondInterface = new java.util.Comparator[String] with Greeter { + override def compare(a: String, b: String): Int = a.compareTo(b) + override def hello(): String = "hi" + } + + // An INSTANCE method whose only counterpart on the supertype is STATIC. scalac allows it, + // since it does not treat a Java static as an inherited member, so `value()` is not an + // override; javac rejects the pair outright, which is why `StaticClashBase` is written in + // Java. A narrowed call binds the supertype's static `value()`, returning 1 rather than 7. + val anonOverStaticClash = new StaticClashBase { + def value(): Int = 7 + } + + // An OVERLOAD, not an override: `convert(String)` does not implement `convert(Any)`, so + // no bridge is emitted. Narrowing would silently bind the call to the supertype's method. + val anonWithOverload = new Converter { + def convert(s: String): String = "anon" + } + + // A bridged generic override alongside an unrelated overload of the same name and arity. + // `compare(String, String)` is reached through a `compare(Object, Object)` bridge, but + // `compare(Int, Int)` has no bridge of its own: after narrowing to `Comparator`, an + // integer call would bind to `compare(Object, Object)` and land in the String-casting + // bridge instead of the overload. + val anonWithBridgeAndPrimitiveOverload = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = a.compareTo(b) + def compare(a: Int, b: Int): Int = a - b + } + + // The same collision with a reference-typed overload. Boxing-blind parameter matching + // would accept this one, since the bridge's `Object` parameters do accept `Integer`. + val anonWithBridgeAndReferenceOverload = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = a.compareTo(b) + def compare(a: Integer, b: Integer): Int = a - b + } + + // Sound counterpart of the two above, and the reason bridge matching treats boxing as + // equivalence: scalac specializes this override to `apply(int)`, which is reached through + // an `apply(Object)` bridge. Requiring the bridge parameter to accept the override's + // parameter without boxing would reject it, since `Object` is not assignable from `int`. + val specializedPrimitiveOverride = new scala.runtime.AbstractFunction1[Int, Boolean] { + def apply(i: Int): Boolean = i > 0 + } + + // A lambda in the body makes scalac emit a `public static final $anonfun$...` helper on the + // anonymous class. It is synthetic, so javac cannot name it ("cannot find symbol") and no + // generated reference can reach it, so narrowing loses nothing and the static check has to + // excuse it or every closure-carrying anonymous class would route to Janino. + val anonWithLambdaBody = new java.util.Comparator[String] { + override def compare(a: String, b: String): Int = { + val len: String => Int = _.length + len(a) - len(b) + } + } + + val plainLocal: java.util.ArrayList[String] = { + class PlainLocal extends java.util.ArrayList[String] + new PlainLocal + } + + val localWithExtraMethod = { + class LocalWithExtra extends java.util.ArrayList[String] { + def extra(): String = "extra" + } + new LocalWithExtra + } + + // A named class declared inside a local class: neither anonymous nor local itself + // (`isMemberClass` is true), yet Java cannot name it either. Its supertype offers every + // member, so only the canonical-name test keeps it from being narrowed to a binary name + // javac would reject. The `$outer` field and accessor scalac adds are synthetic, hence + // unreachable from generated source and no obstacle to narrowing. + val memberOfLocalClass: Any = { + class Holder { + class Inner extends java.util.ArrayList[String] + def make(): Any = new Inner + } + new Holder().make() + } + + // The same shape plus one NON-synthetic public method, which is what actually puts a member + // out of reach of `ArrayList`. Used where the test needs an unnarrowable member-of-local. + val memberOfLocalClassWithExtra: Any = { + class Holder { + class Inner extends java.util.ArrayList[String] { + def extra(): String = "e" + } + def make(): Any = new Inner + } + new Holder().make() + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala index eaf9ac267ff0b..ed790e643adc0 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.catalyst.util.DateTimeConstants._ import org.apache.spark.sql.catalyst.util.DateTimeTestUtils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ -import org.apache.spark.types.variant.VariantBuilder +import org.apache.spark.types.variant.{Variant, VariantBuilder} import org.apache.spark.types.variant.VariantUtil._ import org.apache.spark.unsafe.types.{UTF8String, VariantVal} import org.apache.spark.util.collection.Utils.createArray @@ -548,6 +548,57 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { testVariantGet(json, "$." + numKeys, IntegerType, null) } + test("SPARK-58949: object keys use unsigned UTF-8 order") { + val bmpKey = new String(Character.toChars(65535)) + val supplementaryKey = new String(Character.toChars(0x10000)) + val quote = 34.toChar.toString + val asciiFields = (0 until 32).map(i => quote + i + quote + ":" + i) + val objectJson = (asciiFields ++ Seq( + quote + supplementaryKey + quote + ":99", + quote + bmpKey + quote + ":98")).mkString("{", ",", "}") + + val variant = VariantBuilder.parseJson(objectJson, false) + assert(variant.getFieldAtIndex(32).key === bmpKey) + assert(variant.getFieldAtIndex(33).key === supplementaryKey) + assert(variant.getFieldByKey(bmpKey).getLong === 98L) + assert(variant.getFieldByKey(supplementaryKey).getLong === 99L) + assert(variant.getFieldByKey("missing") === null) + + val nestedJson = "{" + quote + "nested" + quote + ":" + objectJson + "}" + val nested = VariantBuilder.parseJson(nestedJson, false) + .getFieldByKey("nested") + assert(nested.getFieldAtIndex(32).key === bmpKey) + assert(nested.getFieldByKey(supplementaryKey).getLong === 99L) + + // Reorder the last two field entries to reproduce the UTF-16 order written by older Spark. + val legacyValue = variant.getValue.clone() + handleObject[Unit](legacyValue, 0, + (size, idSize, offsetSize, idStart, offsetStart, _dataStart) => { + def swap(start: Int, width: Int): Unit = { + val left = start + (size - 2) * width + val right = left + width + val leftValue = readUnsigned(legacyValue, left, width) + val rightValue = readUnsigned(legacyValue, right, width) + writeLong(legacyValue, left, rightValue, width) + writeLong(legacyValue, right, leftValue, width) + } + swap(idStart, idSize) + swap(offsetStart, offsetSize) + }) + val legacy = new Variant(legacyValue, variant.getMetadata) + assert(legacy.getFieldAtIndex(32).key === supplementaryKey) + assert(legacy.getFieldByKey("31").getLong === 31L) + assert(legacy.getFieldByKey(bmpKey).getLong === 98L) + assert(legacy.getFieldByKey("missing") === null) + + val expectedSchemaNames = ((0 until 32).map(_.toString).sorted ++ + Seq(supplementaryKey, bmpKey)).toArray + Seq(variant, legacy).foreach { v => + val schema = SchemaOfVariant.schemaOf(v).asInstanceOf[StructType] + assert(schema.fieldNames === expectedSchemaNames) + } + } + test("variant_get timestamp") { DateTimeTestUtils.outstandingZoneIds.foreach { zid => withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> zid.getId) { @@ -686,6 +737,16 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { // scalastyle:on nonascii testVariantGet("[1, 2, 3]", "$[2147483647]", IntegerType, null) + Seq("variant_get" -> true, "try_variant_get" -> false).foreach { + case (name, failOnError) => + checkErrorInExpression[SparkRuntimeException]( + VariantGet(BoundReference(0, VariantType, nullable = true), Literal(".a"), + IntegerType, failOnError), + InternalRow(null), + "INVALID_VARIANT_GET_PATH", + Map("path" -> ".a", "functionName" -> s"`$name`")) + } + checkInvalidPath("") checkInvalidPath(".a") checkInvalidPath("$1") @@ -698,6 +759,31 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { checkInvalidPath("$[\"\\\"\"]") } + test("SPARK-58672: validate char/varchar target types in variant_get") { + def check(dataType: DataType, expected: Boolean): Unit = { + assert( + variantGet("""{"a": 1}""", "$", dataType) + .checkInputDataTypes().isSuccess == expected) + } + + def targetTypes(stringType: StringType): Seq[DataType] = Seq( + stringType, + ArrayType(stringType), + MapType(stringType, IntegerType), + MapType(StringType, stringType), + StructType(Seq(StructField("v", stringType)))) + + targetTypes(StringType).foreach { dataType => + check(dataType, expected = true) + } + + Seq(CharType(10), VarcharType(10)).foreach { stringType => + targetTypes(stringType).foreach { dataType => + check(dataType, expected = false) + } + } + } + test("cast from variant") { // We do not test too many type combinations, as the cast implementation is mostly the same as // variant_get. @@ -976,6 +1062,94 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { checkFailure(Map(1 -> 1), toVariantObject = true) } + test("SPARK-58672: validate char/varchar input types in to_variant_object") { + def check(dataType: DataType, expected: Boolean): Unit = { + assert( + ToVariantObject(Literal.create(null, dataType)) + .checkInputDataTypes().isSuccess == expected) + } + + def nestedTypes(stringType: StringType): Seq[DataType] = Seq( + ArrayType(stringType), + MapType(stringType, IntegerType), + MapType(StringType, stringType), + StructType(Seq(StructField("v", stringType)))) + + nestedTypes(StringType).foreach { dataType => + check(dataType, expected = true) + } + + Seq(CharType(10), VarcharType(10)).foreach { stringType => + nestedTypes(stringType).foreach { dataType => + check(dataType, expected = false) + } + } + } + + test("variant_from_arrays and variant_from_entries") { + def keysValues(keys: Any, values: Any, valueType: DataType): VariantFromArrays = + VariantFromArrays( + Literal.create(keys, ArrayType(StringType)), + Literal.create(values, ArrayType(valueType))) + + def entriesOf(entries: Any, valueType: DataType, + containsNull: Boolean = false): VariantFromEntries = + VariantFromEntries(Literal.create(entries, ArrayType( + StructType(Seq(StructField("k", StringType), StructField("v", valueType))), containsNull))) + + // Basic object construction; keys are sorted in the resulting variant object. + checkEvaluation(StructsToJson(Map.empty, + keysValues(Array("z", "a"), Array(1, 2), IntegerType)), """{"a":2,"z":1}""") + checkEvaluation(StructsToJson(Map.empty, + entriesOf(Array(Row("a", 1), Row("b", 2)), IntegerType)), """{"a":1,"b":2}""") + + // Empty input produces an empty object. + checkEvaluation(StructsToJson(Map.empty, + keysValues(Array.empty[String], Array.empty[Int], IntegerType)), "{}") + + // Null values are kept as variant null; nested values are converted recursively. + checkEvaluation(StructsToJson(Map.empty, + entriesOf(Array(Row("a", 1), Row("b", null)), IntegerType)), """{"a":1,"b":null}""") + checkEvaluation(StructsToJson(Map.empty, + keysValues(Array("a"), Array(Array(1, 2, 3)), ArrayType(IntegerType))), """{"a":[1,2,3]}""") + checkEvaluation(StructsToJson(Map.empty, keysValues(Array("a"), Array(Row(1)), + StructType(Seq(StructField("i", IntegerType))))), """{"a":{"i":1}}""") + + // A null entry makes the whole result null. + checkEvaluation(StructsToJson(Map.empty, + entriesOf(Array(Row("a", 1), null), IntegerType, containsNull = true)), null) + + // A null entry dominates a value-conversion failure in an earlier entry (matches + // map_from_entries: the null check runs for every entry before any value is converted). + checkEvaluation(StructsToJson(Map.empty, + entriesOf(Array(Row("a", Row(1, 2)), null), + StructType(Seq(StructField("x", IntegerType), StructField("x", IntegerType))), + containsNull = true)), null) + + // A null array input produces null. + checkEvaluation(StructsToJson(Map.empty, VariantFromArrays( + Literal.create(null, ArrayType(StringType)), + Literal.create(Array(1), ArrayType(IntegerType)))), null) + + // A null key is rejected. + checkErrorInExpression[SparkRuntimeException]( + keysValues(Array("a", null), Array(1, 2), IntegerType), + "NULL_MAP_KEY", Map.empty[String, String]) + + // Duplicate keys are rejected for both forms. + checkErrorInExpression[SparkRuntimeException]( + keysValues(Array("a", "a"), Array(1, 2), IntegerType), + "VARIANT_DUPLICATE_KEY", Map("key" -> "a")) + checkErrorInExpression[SparkRuntimeException]( + entriesOf(Array(Row("a", 1), Row("a", 2)), IntegerType), + "VARIANT_DUPLICATE_KEY", Map("key" -> "a")) + + // Mismatched array lengths are rejected. + checkErrorInExpression[SparkRuntimeException]( + keysValues(Array("a", "b"), Array(1), IntegerType), + "_LEGACY_ERROR_TEMP_2128", Map.empty[String, String]) + } + test("schema_of_variant - unknown type") { val emptyMetadata = Array[Byte](VERSION, 0, 0) @@ -1487,6 +1661,16 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { "INVALID_VARIANT_PATH", name => Map("path" -> "abc", "functionName" -> s"`$name`")) + Seq("variant_insert" -> true, "try_variant_insert" -> false).foreach { + case (name, failOnError) => + checkErrorInExpression[SparkRuntimeException]( + VariantInsert(BoundReference(0, VariantType, nullable = true), + Literal.create("abc", StringType), Literal(1), failOnError), + InternalRow(null), + "INVALID_VARIANT_PATH", + Map("path" -> "abc", "functionName" -> s"`$name`")) + } + val tooBig = "x".repeat(16 * 1024 * 1024) checkInsertUnrecoverableError("{}", "$.a[2000000000]", Literal(1), "VARIANT_SIZE_LIMIT", @@ -1617,20 +1801,34 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { checkSet("""{"a": 1}""", "$.a", Literal.create(null, NullType), null) Seq(true, false).foreach { failOnError => - val dynamic = VariantSet( + val dynamicCreate = VariantSet( Literal(parseJson("""{"a": 1}""")), BoundReference(0, StringType, nullable = true), Literal(2), - BoundReference(1, BooleanType, nullable = true), + Literal(true), failOnError) checkEvaluation( - ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)), + ResolveTimeZone.resolveTimeZones(Cast(dynamicCreate, StringType)), """{"a":1,"b":2}""", - InternalRow(UTF8String.fromString("$.b"), true)) + InternalRow(UTF8String.fromString("$.b"))) + val dynamicNoCreate = VariantSet( + Literal(parseJson("""{"a": 1}""")), + BoundReference(0, StringType, nullable = true), + Literal(2), + Literal(false), + failOnError) checkEvaluation( - ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)), + ResolveTimeZone.resolveTimeZones(Cast(dynamicNoCreate, StringType)), """{"a":1}""", - InternalRow(UTF8String.fromString("$.b"), false)) + InternalRow(UTF8String.fromString("$.b"))) + } + + // create_if_missing must be a constant, for both variant_set and try_variant_set. + Seq(true, false).foreach { failOnError => + assert(VariantSet( + Literal(parseJson("""{"a": 1}""")), Literal("$.a"), Literal(2), + BoundReference(0, BooleanType, nullable = true), failOnError) + .checkInputDataTypes().isFailure) } // Recoverable errors: `variant_set` throws; `try_variant_set` returns NULL. Every shape of @@ -1840,6 +2038,16 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { checkAppendUnrecoverableError("[]", "abc", Literal(1), "INVALID_VARIANT_PATH", name => Map("path" -> "abc", "functionName" -> s"`$name`")) + + Seq("variant_array_append" -> true, "try_variant_array_append" -> false).foreach { + case (name, failOnError) => + checkErrorInExpression[SparkRuntimeException]( + VariantArrayAppend(BoundReference(0, VariantType, nullable = true), + Literal.create("abc", StringType), Literal(1), failOnError), + InternalRow(null), + "INVALID_VARIANT_PATH", + Map("path" -> "abc", "functionName" -> s"`$name`")) + } checkErrorInExpression[SparkRuntimeException]( VariantArrayAppend(Literal(parseJson("[]")), Literal(""), Literal(1)), "INVALID_VARIANT_PATH", @@ -1849,4 +2057,74 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper { "VARIANT_SIZE_LIMIT", name => Map("sizeLimit" -> "16.0 MiB", "functionName" -> s"`$name`")) } + + test("variant_strip_nulls") { + // Strip `input`, render the result back to JSON, and compare. `includeArrays` defaults to true. + def check(input: String, expected: String, includeArrays: Boolean = true): Unit = { + val expr = VariantStripNulls(Literal(parseJson(input)), Literal(includeArrays)) + val result = replace(expr).eval().asInstanceOf[VariantVal] + val json = if (result == null) null + else new Variant(result.getValue, result.getMetadata).toJson(ZoneOffset.UTC) + assert(json == expected) + } + + // The optional `include_arrays` argument defaults to true in the function signature. + assert(VariantStripNullsExpressionBuilder.functionSignature.get.parameters.last.default + .contains(Literal.create(true, BooleanType))) + + // include_arrays must be a constant; a non-foldable expression is rejected. + assert(VariantStripNulls( + Literal(parseJson("[1, null]")), BoundReference(0, BooleanType, nullable = true)) + .checkInputDataTypes().isFailure) + + check("""{"a": 1, "b": null, "c": 3}""", """{"a":1,"c":3}""") + check("[1, null, 3]", "[1,3]") + check("""{"user": {"name": "Alice", "age": null}}""", """{"user":{"name":"Alice"}}""") + check("""{"a": [1, null, {"b": null, "c": 2}]}""", """{"a":[1,{"c":2}]}""") + check("[[1, null], [null]]", "[[1],[]]") + + // Empty containers are preserved; the parent is never collapsed. + check("""{"a": null}""", "{}") + check("[null]", "[]") + check("""{"a": {"b": null}}""", """{"a":{}}""") + check("""{"a": [null]}""", """{"a":[]}""") + check("{}", "{}") + check("[]", "[]") + + // Top-level variant null and scalars are returned unchanged. + check("null", "null") + check("42", "42") + check("\"hi\"", "\"hi\"") + + check("""{"a": {"b": {"c": null, "d": 4}}}""", """{"a":{"b":{"d":4}}}""") + + check("""{"a": 300, "b": null, "c": 100000, "d": 10000000000}""", + """{"a":300,"c":100000,"d":10000000000}""") + check("""[1000000, null, "hello world", null, 10000000000]""", + """[1000000,"hello world",10000000000]""") + val bigStr = "x".repeat(300) + check(s"""{"k": "$bigStr", "n": null, "m": 3}""", s"""{"k":"$bigStr","m":3}""") + check(s"""[null, "$bigStr", null, 100000]""", s"""["$bigStr",100000]""") + + // `includeArrays = false`. + check("""{"a": [1, null, 3], "b": null}""", """{"a":[1,null,3]}""", includeArrays = false) + check( + """[{"a": 1, "b": null}, null, {"c": null, "d": 4}]""", + """[{"a":1},null,{"d":4}]""", + includeArrays = false) + // `includeArrays = true` (explicit) strips array nulls. + check("""{"a": [1, null, 3]}""", """{"a":[1,3]}""", includeArrays = true) + + // SQL NULL variant input yields SQL NULL. + checkEvaluation( + Cast(VariantStripNulls(Literal.create(null, VariantType), Literal(true)), StringType), + null) + // NULL `includeArrays` yields SQL NULL (the expression is null intolerant). + checkEvaluation( + Cast( + VariantStripNulls( + Literal(parseJson("""{"a": null}""")), Literal.create(null, BooleanType)), + StringType), + null) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala index 515203da7caf6..fb2f0957d74d2 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CollapseWindowSuite.scala @@ -19,9 +19,15 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ +import org.apache.spark.sql.catalyst.expressions.{ + CurrentRow, RangeFrame, RowFrame, RowNumber, SpecifiedWindowFrame, + UnboundedFollowing, UnboundedPreceding} +import org.apache.spark.sql.catalyst.expressions.aggregate.{ + AggregateExpression, Complete, Count, First, Sum} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.internal.SQLConf class CollapseWindowSuite extends PlanTest { object Optimize extends RuleExecutor[LogicalPlan] { @@ -168,4 +174,248 @@ class CollapseWindowSuite extends PlanTest { comparePlans(optimized, correctAnswer) } + + test("collapse windows when one has an empty order spec " + + "(row_number + count over the whole partition)") { + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window has multiple window expressions") { + // Every window expression of the empty-order window must be order-insensitive for the merge. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + val sm = windowExpr( + AggregateExpression(Sum(b), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("sm") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt, sm), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt, sm), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window has first() over the whole partition") { + // `first` is non-deterministic when the order is not determined by the query, so evaluating it + // under the other window's order spec yields a valid result. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val fr = windowExpr( + First(a, ignoreNulls = true).toAggregateExpression(), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("fr") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(fr), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, fr), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("don't collapse windows when the empty-order window has a bounded frame") { + // The frame `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` is order-sensitive: which rows + // fall in the frame depends on the ordering, so the window cannot be evaluated under the other + // window's order spec. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt), partitionSpec1, Nil) + + val optimized = Optimize.execute(query.analyze) + val correctAnswer = query.analyze + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window has a RANGE whole-partition frame") { + // `RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING` covers the whole partition just + // like `ROWS`, so it also collapses. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RangeFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .window(Seq(cnt), partitionSpec1, Nil) + + val analyzed = query.analyze + val optimized = Optimize.execute(analyzed) + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows when the empty-order window is the inner window") { + // The empty-order window can also be the child of the ordered window. In that case its + // expressions are evaluated under the ordered window's order spec, which is valid because all + // of them are order-insensitive. This direction can disable InferWindowGroupLimit, so it is + // gated by `spark.sql.optimizer.collapseWindowWithEmptyOrderSpecInChild`. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(cnt), partitionSpec1, Nil) + .window(Seq(rk), partitionSpec1, orderSpec1) + + val analyzed = query.analyze + val optimized = withSQLConf( + SQLConf.COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD.key -> "true") { + Optimize.execute(analyzed) + } + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(cnt, rk), partitionSpec1, orderSpec1) + + comparePlans(optimized, correctAnswer) + } + + test("don't collapse the inner empty-order window by default") { + // Merging an empty-order child into an ordered parent can disable InferWindowGroupLimit for + // top-k queries, so it is off by default. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(cnt), partitionSpec1, Nil) + .window(Seq(rk), partitionSpec1, orderSpec1) + + val optimized = Optimize.execute(query.analyze) + val correctAnswer = query.analyze + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows with a Project between them when one has an empty order spec") { + // The same merge applies when a Project sits between the two windows and only passes through + // columns that are available below the inner window (SPARK-34565 shape). The empty-order + // window is the inner window here, so the config must be enabled. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(cnt), partitionSpec1, Nil) + .select($"a", $"b", $"c", $"cnt") + .window(Seq(rk), partitionSpec1, orderSpec1) + .select($"a", $"b", $"c", $"cnt", $"rk") + + val analyzed = query.analyze + val optimized = withSQLConf( + SQLConf.COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD.key -> "true") { + Optimize.execute(analyzed) + } + assert(analyzed.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(cnt, rk), partitionSpec1, orderSpec1) + .select(a, b, c, $"cnt", $"rk") + .analyze + + comparePlans(optimized, correctAnswer) + } + + test("collapse windows with a Project between them, empty order spec as parent") { + // The empty-order window is the parent here, so this merges by default without the config. + val rk = windowExpr( + RowNumber(), + windowSpec(partitionSpec1, orderSpec1, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rk") + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(partitionSpec1, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + + val query = testRelation + .window(Seq(rk), partitionSpec1, orderSpec1) + .select($"a", $"b", $"c", $"rk") + .window(Seq(cnt), partitionSpec1, Nil) + .select($"a", $"b", $"c", $"rk", $"cnt") + + val optimized = Optimize.execute(query.analyze) + assert(query.analyze.output === optimized.output) + + val correctAnswer = testRelation + .window(Seq(rk, cnt), partitionSpec1, orderSpec1) + .select(a, b, c, $"rk", $"cnt") + .analyze + + comparePlans(optimized, correctAnswer) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala new file mode 100644 index 0000000000000..96bf9c7db256b --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/CombineApproximatePercentilesSuite.scala @@ -0,0 +1,304 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.analysis.FunctionRegistry +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.dsl.plans._ +import org.apache.spark.sql.catalyst.expressions.{Alias, CreateArray, Expression, GetArrayItem, Literal} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile} +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.catalyst.util.ArrayData +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.IntegerType + +class CombineApproximatePercentilesSuite extends PlanTest { + + private object Optimize extends RuleExecutor[LogicalPlan] { + override val batches: Seq[Batch] = + Batch("Combine Approximate Percentiles", Once, CombineApproximatePercentiles) :: Nil + } + + private val relation = LocalRelation($"value".int, $"other".int, $"group".int) + private val value = relation.output(0) + private val other = relation.output(1) + private val group = relation.output(2) + + private def percentile( + child: Expression, + percentage: Double, + accuracy: Int = 10000, + isDistinct: Boolean = false, + filter: Option[Expression] = None): AggregateExpression = { + new ApproximatePercentile(child, Literal(percentage), Literal(accuracy)) + .toAggregateExpression(isDistinct = isDistinct, filter = filter) + } + + private def optimizedAggregate(expressions: Alias*): Aggregate = { + val plan = Aggregate(Seq.empty, expressions, relation) + withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true") { + Optimize.execute(plan.analyze).asInstanceOf[Aggregate] + } + } + + private def percentileAggregates(aggregate: Aggregate): Seq[AggregateExpression] = { + aggregate.aggregateExpressions.flatMap(_.collect { + case expression @ AggregateExpression(_: ApproximatePercentile, _, _, _, _) => expression + }) + } + + private def percentageValues(percentile: ApproximatePercentile): Seq[Double] = { + percentile.percentageExpression.eval().asInstanceOf[ArrayData].toDoubleArray().toSeq + } + + private def ordinals(aggregate: Aggregate): Seq[Option[Int]] = { + aggregate.aggregateExpressions.map(_.collectFirst { + case GetArrayItem(_, Literal(index: Int, _), false) => index + }) + } + + private def assertNotCombined(aggregate: Aggregate): Unit = { + val expressions = percentileAggregates(aggregate) + assert(expressions.map(_.resultId).distinct.size == expressions.size) + assert(!aggregate.aggregateExpressions.exists(_.exists(_.isInstanceOf[GetArrayItem]))) + } + + test("do not combine approximate percentiles when disabled") { + assert(!new SQLConf().getConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED)) + val original = Aggregate( + Seq.empty, + Seq( + Alias(percentile(value, 0.5), "p50")(), + Alias(percentile(value, 0.9), "p90")()), + relation).analyze.asInstanceOf[Aggregate] + + withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") { + assertNotCombined(Optimize.execute(original).asInstanceOf[Aggregate]) + } + } + + test("combine compatible percentiles and preserve output shape") { + val firstPercentile = percentile(value, 0.9) + firstPercentile.aggregateFunction.setTagValue(FunctionRegistry.FUNC_ALIAS, "approx_percentile") + val aliases = Seq( + Alias(firstPercentile, "first")(), + Alias(percentile(value, 0.5), "second")(), + Alias(percentile(value, 0.9), "third")()) + val original = Aggregate(Seq(group), group +: aliases, relation).analyze + .asInstanceOf[Aggregate] + val result = withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true") { + Optimize.execute(original).asInstanceOf[Aggregate] + } + val aggregates = percentileAggregates(result) + val combined = aggregates.head.aggregateFunction.asInstanceOf[ApproximatePercentile] + + assert(aggregates.map(_.resultId).distinct.size == 1) + assert(aggregates.head.resultId != firstPercentile.resultId) + assert(percentageValues(combined) == Seq(0.9, 0.5, 0.9)) + assert(combined.percentageExpression.toString == "[0.9,0.5,0.9]") + assert(combined.percentageExpression.sql == "ARRAY(0.9D, 0.5D, 0.9D)") + assert(combined.prettyName == "approx_percentile") + assert(ordinals(result) == Seq(None, Some(0), Some(1), Some(2))) + assert(result.groupingExpressions == original.groupingExpressions) + assert(result.output.map(_.exprId) == original.output.map(_.exprId)) + } + + test("preserve fusion identity through later constant folding") { + val result = optimizedAggregate( + Alias(percentile(value, 0.5), "p50")(), + Alias(percentile(value, 0.9), "p90")()) + val folded = ConstantFolding(result).asInstanceOf[Aggregate] + val combined = percentileAggregates(folded).head.aggregateFunction + .asInstanceOf[ApproximatePercentile] + + assert(combined.percentageExpression.isInstanceOf[PercentileFusionArray]) + assert(!combined.percentageExpression.contextIndependentFoldable) + } + + test("avoid redundant entries for duplicate physical aggregates") { + val duplicateOnly = optimizedAggregate( + Alias(percentile(value, 0.5), "first")(), + Alias(percentile(value, 0.5), "second")()) + assertNotCombined(duplicateOnly) + + val p50 = percentile(value, 0.5) + val mixed = optimizedAggregate( + Alias(p50, "first")(), + Alias(p50, "second")(), + Alias(percentile(value, 0.9), "third")()) + val combined = percentileAggregates(mixed).head.aggregateFunction + .asInstanceOf[ApproximatePercentile] + assert(percentageValues(combined) == Seq(0.5, 0.9)) + assert(ordinals(mixed) == Seq(Some(0), Some(0), Some(1))) + } + + test("respect basic compatibility boundaries") { + val incompatible = Seq( + "input" -> ( + percentile(value, 0.5), + percentile(other, 0.9)), + "accuracy" -> ( + percentile(value, 0.5, accuracy = 1000), + percentile(value, 0.9, accuracy = 10000)), + "filter" -> ( + percentile(value, 0.5), + percentile(value, 0.9, filter = Some(value > Literal(0)))), + "distinct" -> ( + percentile(value, 0.5), + percentile(value, 0.9, isDistinct = true))) + + incompatible.foreach { case (name, (first, second)) => + withClue(name) { + assertNotCombined(optimizedAggregate( + Alias(first, "first")(), + Alias(second, "second")())) + } + } + + val filter = value > Literal(0) + val compatible = optimizedAggregate( + Alias(percentile(value, 0.5, filter = Some(filter)), "filtered_p50")(), + Alias(percentile(value, 0.9, filter = Some(filter)), "filtered_p90")(), + Alias(percentile(value, 0.5, isDistinct = true, filter = Some(filter)), "distinct_p50")(), + Alias(percentile(value, 0.9, isDistinct = true, filter = Some(filter)), "distinct_p90")()) + assert(percentileAggregates(compatible).map(_.resultId).distinct.size == 2) + } + + test("require structural equality for inputs and filters") { + val firstInput = value + (other + group) + val secondInput = (value + other) + group + val inputResult = optimizedAggregate( + Alias(percentile(firstInput, 0.5), "first_p50")(), + Alias(percentile(secondInput, 0.5), "second_p50")(), + Alias(percentile(secondInput, 0.9), "second_p90")(), + Alias(percentile(firstInput, 0.9), "first_p90")()) + + assert(firstInput != secondInput) + assert(firstInput.canonicalized == secondInput.canonicalized) + assertNotCombined(inputResult) + + val disjointInputPercentages = optimizedAggregate( + Alias(percentile(firstInput, 0.5), "first_p50")(), + Alias(percentile(firstInput, 0.9), "first_p90")(), + Alias(percentile(secondInput, 0.25), "second_p25")(), + Alias(percentile(secondInput, 0.75), "second_p75")()) + assert(percentileAggregates(disjointInputPercentages).map(_.resultId).distinct.size == 2) + assert(ordinals(disjointInputPercentages) == + Seq(Some(0), Some(1), Some(0), Some(1))) + + val crossDistinctInput = optimizedAggregate( + Alias(percentile(firstInput, 0.5), "first_p50")(), + Alias(percentile(firstInput, 0.9), "first_p90")(), + Alias(percentile(secondInput, 0.5, isDistinct = true), "distinct_p50")(), + Alias(percentile(secondInput, 0.9, isDistinct = true), "distinct_p90")()) + assertNotCombined(crossDistinctInput) + + val firstFilter = firstInput > Literal(0) + val secondFilter = secondInput > Literal(0) + val filterResult = optimizedAggregate( + Alias(percentile(value, 0.5, filter = Some(firstFilter)), "first_p50")(), + Alias(percentile(value, 0.5, filter = Some(secondFilter)), "second_p50")(), + Alias(percentile(value, 0.9, filter = Some(secondFilter)), "second_p90")(), + Alias(percentile(value, 0.9, filter = Some(firstFilter)), "first_p90")()) + assertNotCombined(filterResult) + + val crossDistinctFilter = optimizedAggregate( + Alias(percentile(value, 0.5, filter = Some(firstFilter)), "first_p50")(), + Alias(percentile(value, 0.9, filter = Some(firstFilter)), "first_p90")(), + Alias(percentile( + value, 0.5, isDistinct = true, filter = Some(secondFilter)), "distinct_p50")(), + Alias(percentile( + value, 0.9, isDistinct = true, filter = Some(secondFilter)), "distinct_p90")()) + assertNotCombined(crossDistinctFilter) + } + + test("do not fuse canonically equivalent but differently evaluated accuracies") { + val firstAccuracy = + ((Literal(1.0e16) + Literal(-1.0e16)) + Literal(3.0)).cast(IntegerType) + val secondAccuracy = + (Literal(1.0e16) + (Literal(-1.0e16) + Literal(3.0))).cast(IntegerType) + + def percentileWithAccuracy( + percentage: Double, + accuracy: Expression): AggregateExpression = { + new ApproximatePercentile(value, Literal(percentage), accuracy).toAggregateExpression() + } + + val result = optimizedAggregate( + Alias(percentileWithAccuracy(0.5, firstAccuracy), "first_p50")(), + Alias(percentileWithAccuracy(0.5, secondAccuracy), "second_p50")(), + Alias(percentileWithAccuracy(0.9, secondAccuracy), "second_p90")(), + Alias(percentileWithAccuracy(0.9, firstAccuracy), "first_p90")()) + + assert(firstAccuracy.eval() == 3) + assert(secondAccuracy.eval() == 4) + assert(firstAccuracy.canonicalized == secondAccuracy.canonicalized) + assertNotCombined(result) + } + + test("do not fuse canonically equal percentages that evaluate differently") { + val firstPercentage = + (Literal(1.0e16) + Literal(-1.0e16)) + Literal(0.5) + val secondPercentage = + Literal(1.0e16) + (Literal(-1.0e16) + Literal(0.5)) + def percentileWithPercentage(percentage: Expression): AggregateExpression = { + new ApproximatePercentile(value, percentage, Literal(10000)).toAggregateExpression() + } + + val scalarResult = optimizedAggregate( + Alias(percentileWithPercentage(firstPercentage), "p50")(), + Alias(percentileWithPercentage(secondPercentage), "p0")(), + Alias(percentile(value, 0.9), "p90")()) + assert(firstPercentage.eval() == 0.5d) + assert(secondPercentage.eval() == 0.0d) + assert(firstPercentage.canonicalized == secondPercentage.canonicalized) + assertNotCombined(scalarResult) + + val arrayPercentile = new ApproximatePercentile( + value, + CreateArray(Seq(firstPercentage, Literal(0.9))), + Literal(10000)).toAggregateExpression() + val arrayResult = optimizedAggregate( + Alias(arrayPercentile, "percentiles")(), + Alias(percentileWithPercentage(secondPercentage), "p0")(), + Alias(percentile(value, 0.9), "p90")()) + val arrayAggregates = percentileAggregates(arrayResult) + assert(arrayAggregates.map(_.resultId).distinct.size == 2) + assert(percentageValues(arrayAggregates.head.aggregateFunction + .asInstanceOf[ApproximatePercentile]) == Seq(0.5, 0.9)) + assert(percentageValues(arrayAggregates(1).aggregateFunction + .asInstanceOf[ApproximatePercentile]) == Seq(0.0, 0.9)) + assert(ordinals(arrayResult) == Seq(None, Some(0), Some(1))) + } + + test("preserve existing arrays while combining compatible scalars") { + val arrayPercentile = new ApproximatePercentile( + value, + CreateArray(Seq(Literal(0.25), Literal(0.75))), + Literal(10000)).toAggregateExpression() + val result = optimizedAggregate( + Alias(arrayPercentile, "percentiles")(), + Alias(percentile(value, 0.5), "p50")(), + Alias(percentile(value, 0.9), "p90")()) + + assert(percentileAggregates(result).map(_.resultId).distinct.size == 2) + assert(result.aggregateExpressions.drop(1).forall(_.exists(_.isInstanceOf[GetArrayItem]))) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala index be24f9c9f01f9..dc8cfd74c7851 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ComputeCurrentTimeSuite.scala @@ -18,13 +18,13 @@ package org.apache.spark.sql.catalyst.optimizer import java.lang.Thread.sleep -import java.time.{LocalDateTime, ZoneId} +import java.time.{Instant, LocalDateTime, ZoneId} import scala.concurrent.duration._ import scala.jdk.CollectionConverters.MapHasAsScala import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Cast, CurrentDate, CurrentTime, CurrentTimestamp, CurrentTimeZone, Expression, InSubquery, ListQuery, Literal, LocalTimestamp, Now} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Cast, CurrentDate, CurrentTime, CurrentTimestamp, CurrentTimestampNanos, CurrentTimeZone, Expression, InSubquery, ListQuery, Literal, LocalTimestamp, LocalTimestampNanos, Now} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, LogicalPlan, Project} import org.apache.spark.sql.catalyst.rules.RuleExecutor @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.trees.TreePattern import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{DateType, IntegerType, StringType, TimestampLTZNanosType, TimestampNTZNanosType, TimestampNTZType, TimestampType, TimeType} -import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String} class ComputeCurrentTimeSuite extends PlanTest { object Optimize extends RuleExecutor[LogicalPlan] { @@ -197,6 +197,70 @@ class ComputeCurrentTimeSuite extends PlanTest { assert(lits(0) == lits(1)) } + test("SPARK-57837: analyzer should replace current_timestamp(p) with nanos literals") { + Seq(7, 8, 9).foreach { p => + val in = Project( + Seq(Alias(CurrentTimestampNanos(p), "a")(), Alias(CurrentTimestampNanos(p), "b")()), + LocalRelation()) + + val min = DateTimeUtils.instantToMicros(Instant.now()) + val plan = Optimize.execute(in.analyze).asInstanceOf[Project] + val max = DateTimeUtils.instantToMicros(Instant.now()) + + // The literals carry the nanosecond TIMESTAMP_LTZ(p) type. + val typedLits = literalsWithType(plan, TimestampLTZNanosType(p)) + assert(typedLits.size == 2, s"precision $p should yield two TIMESTAMP_LTZ($p) literals") + val vals = typedLits.map(_.asInstanceOf[TimestampNanosVal]) + // All calls in one query return the same value. + assert(vals(0) == vals(1)) + assert(vals(0).epochMicros >= min && vals(0).epochMicros <= max) + // Sub-precision digits are floored: at p == 7 the last two nanos-within-micro digits are 0, + // at p == 8 the last digit is 0, at p == 9 all three are kept. + val step = math.pow(10, 9 - p).toInt % 1000 + if (step != 0) { + assert(vals(0).nanosWithinMicro % step == 0, + s"nanosWithinMicro ${vals(0).nanosWithinMicro} should be floored to precision $p") + } + } + } + + test("SPARK-57837: analyzer should replace localtimestamp(p) with nanos literals") { + Seq(7, 8, 9).foreach { p => + val in = Project( + Seq(Alias(LocalTimestampNanos(p), "a")(), Alias(LocalTimestampNanos(p), "b")()), + LocalRelation()) + + val plan = Optimize.execute(in.analyze).asInstanceOf[Project] + + val typedLits = literalsWithType(plan, TimestampNTZNanosType(p)) + assert(typedLits.size == 2, s"precision $p should yield two TIMESTAMP_NTZ($p) literals") + val vals = typedLits.map(_.asInstanceOf[TimestampNanosVal]) + assert(vals(0) == vals(1)) + val step = math.pow(10, 9 - p).toInt % 1000 + if (step != 0) { + assert(vals(0).nanosWithinMicro % step == 0, + s"nanosWithinMicro ${vals(0).nanosWithinMicro} should be floored to precision $p") + } + } + } + + test("SPARK-57837: nanos current-timestamp respects time flow across analyses") { + val in = Project(Alias(CurrentTimestampNanos(9), "t1")() :: Nil, LocalRelation()) + + val planT1 = Optimize.execute(in.analyze).asInstanceOf[Project] + sleep(5) + val planT2 = Optimize.execute(in.analyze).asInstanceOf[Project] + + val t1 = literalsWithType(planT1, TimestampLTZNanosType(9)) + .head.asInstanceOf[TimestampNanosVal] + val t2 = literalsWithType(planT2, TimestampLTZNanosType(9)) + .head.asInstanceOf[TimestampNanosVal] + + // A later analysis observes a strictly newer instant (each re-analysis re-reads the clock). + assert(t2.epochMicros > t1.epochMicros, + s"Expected a newer time in the second analysis, but got t1=$t1, t2=$t2") + } + test("analyzer should use equal timestamps across subqueries") { val timestampInSubQuery = Project(Seq(Alias(LocalTimestamp(), "timestamp1")()), LocalRelation()) val listSubQuery = ListQuery(timestampInSubQuery) @@ -344,6 +408,20 @@ class ComputeCurrentTimeSuite extends PlanTest { literals } + private def literalsWithType( + plan: LogicalPlan, + dataType: org.apache.spark.sql.types.DataType) + : scala.collection.mutable.ArrayBuffer[Any] = { + val buf = new scala.collection.mutable.ArrayBuffer[Any] + plan.transformWithSubqueries { case subQuery => + subQuery.transformAllExpressions { case lit: Literal if lit.dataType == dataType => + buf += lit.value + lit + } + } + buf + } + test("SPARK-57748: TIME->TIMESTAMP cast is rewritten even with no CURRENT_LIKE node") { val timeLit = Literal(0L, TimeType(6)) Seq(TimestampNTZType, TimestampType).foreach { target => diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala new file mode 100644 index 0000000000000..23baf25a5897a --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala @@ -0,0 +1,237 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count} +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, Project} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{BooleanType, IntegerType, LongType} + +/** + * Unit tests for the ConvertToCatalyst optimizer rule, which rewrites + * TranspiledPythonUDF nodes to their Catalyst equivalents. + * + * These tests exercise the rule directly via applyExpr rather than running the + * full optimizer pipeline, which means no JVM/Python bridge is required. + */ +class ConvertToCatalystSuite extends PlanTest { + + private val attrA = $"a".long + + // A leaf PythonUDF that takes one column argument. func=null is intentional: + // structural tests don't need an executable PythonFunction. + private def makePyUDF(input: Expression = attrA): PythonUDF = + PythonUDF("udf", null, LongType, Seq(input), + PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true) + + // A leaf PythonUDAF (grouped-agg pandas eval type, return type Long for parity + // with Count's output). func=null is intentional, as with makePyUDF. + private def makePyUDAF(input: Expression = attrA): PythonUDAF = + PythonUDAF("agg", null, LongType, Seq(input), + udfDeterministic = true, + evalType = PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF) + + // A TranspiledPythonUDF wrapping pyUDF with a single Catalyst option. + private def makeTPUDF(pyUDF: PythonUDF, catalystOpt: Expression): TranspiledPythonUDF = + TranspiledPythonUDF("udf", pyUDF, List(catalystOpt)) + + private val catalystExpr: Expression = Add(attrA, Literal(4L)) + + // ---- helpers ---- + + // Both ANSI and ATTEMPT_TRANSPILATION must be true for the transpile path to fire. + private def transpileOn[T](block: => T): T = + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key -> "true") { block } + + private def ansiOff[T](block: => T): T = + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "false", + SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key -> "true") { block } + + private def transpileOff[T](block: => T): T = + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key -> "false") { block } + + // ---- tests ---- + + test("transpiles when not nested (parentIsUdf = false)") { + transpileOn { + val tpudf = makeTPUDF(makePyUDF(), catalystExpr) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + assert(!result.isInstanceOf[PythonUDF]) + } + } + + test("prevents transpilation when parentIsUdf=true and inputs are plain PythonUDFs") { + // PythonUDF -> TranspiledPythonUDF -> PythonUDF: the middle node should NOT be + // transpiled when called from an outer UDF context, to preserve the batch pipeline. + transpileOn { + val innerPyUDF = makePyUDF(attrA) + val outerPyUDF = makePyUDF(innerPyUDF) + val outerTPUDF = makeTPUDF(outerPyUDF, Add(innerPyUDF, Literal(4L))) + val result = ConvertToCatalyst.applyExpr(outerTPUDF, parentIsUdf = true) + assert(result.isInstanceOf[PythonUDF]) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + } + } + + test("does not prevent transpilation when input to pythonUDFExpr is a TranspiledPythonUDF") { + // When the input to a TPUDF is itself a TranspiledPythonUDF (has a Catalyst alternative), + // hasOnlyPythonUDFInputs returns false so the outer TPUDF still transpiles. + transpileOn { + val innerPyUDF = makePyUDF(attrA) + val innerTPUDF = makeTPUDF(innerPyUDF, catalystExpr) + val outerPyUDF = makePyUDF(innerTPUDF) + val outerTPUDF = makeTPUDF(outerPyUDF, Add(innerTPUDF, Literal(4L))) + val result = ConvertToCatalyst.applyExpr(outerTPUDF, parentIsUdf = true) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + assert(!result.isInstanceOf[PythonUDF]) + } + } + + test("hasOnlyPythonUDFInputs unit test") { + val innerPyUDF = makePyUDF(attrA) + val innerTPUDF = makeTPUDF(innerPyUDF, catalystExpr) + + // pythonUDFExpr's child is a plain PythonUDF -> true + assert(makeTPUDF(makePyUDF(innerPyUDF), catalystExpr).hasOnlyPythonUDFInputs) + // pythonUDFExpr's child is a TranspiledPythonUDF -> false + assert(!makeTPUDF(makePyUDF(innerTPUDF), catalystExpr).hasOnlyPythonUDFInputs) + // pythonUDFExpr's child is a plain column (leaf) -> false + assert(!makeTPUDF(makePyUDF(attrA), catalystExpr).hasOnlyPythonUDFInputs) + // zero-arg pythonUDFExpr -> false (nonEmpty guard) + val zeroPyUDF = PythonUDF("udf", null, LongType, Seq.empty, + PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true) + assert(!TranspiledPythonUDF("udf", zeroPyUDF, List(Literal(42L))).hasOnlyPythonUDFInputs) + } + + test("falls back to PythonUDF when ANSI is disabled") { + ansiOff { + val tpudf = makeTPUDF(makePyUDF(), catalystExpr) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(result.isInstanceOf[PythonUDF]) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + } + } + + test("falls back to PythonUDF when transpilation is disabled") { + transpileOff { + val tpudf = makeTPUDF(makePyUDF(), catalystExpr) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(result.isInstanceOf[PythonUDF]) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + } + } + + test("falls back to PythonUDF when transpiledOptions is empty") { + transpileOn { + val pyUDF = makePyUDF() + val tpudf = TranspiledPythonUDF("udf", pyUDF, List()) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(result.isInstanceOf[PythonUDF]) + assert(!result.isInstanceOf[TranspiledPythonUDF]) + } + } + + test("apply(plan) reaches TranspiledPythonUDF nodes below the root") { + // Regression test for the traversal bug where ``plan.mapExpressions`` only + // walks expressions on the root plan node. With that bug, a TPUDF inside a + // Filter (or any non-root node) would survive the optimizer rule as an + // ``Unevaluable`` expression and crash at execution. The fix uses + // ``transformAllExpressionsWithPruning`` which descends through child + // plans; this test pins that contract. + transpileOn { + val attrB = $"b".long + val relation = LocalRelation(attrA, attrB) + // The TPUDF lives in the Filter's condition (boolean), not at the root. + val booleanTPUDF = TranspiledPythonUDF( + "udf", + PythonUDF("udf", null, BooleanType, Seq(attrA), + PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true), + List(GreaterThan(attrA, Literal(0L)))) + val plan = Project(Seq(attrB), Filter(booleanTPUDF, relation)) + val rewritten = ConvertToCatalyst.apply(plan) + // No TranspiledPythonUDF should remain anywhere in the rewritten plan. + val leftover = rewritten.collect { + case p if p.expressions.exists(_.find(_.isInstanceOf[TranspiledPythonUDF]).isDefined) => + p + } + assert(leftover.isEmpty, + s"TranspiledPythonUDF survived ConvertToCatalyst.apply: $rewritten") + // The Filter's condition must be the resolved Catalyst expression, not a fallback PythonUDF. + val filterCond = rewritten.asInstanceOf[Project].child.asInstanceOf[Filter].condition + assert(filterCond == GreaterThan(attrA, Literal(0L)), + s"Filter condition was not rewritten to GreaterThan: $filterCond") + } + } + + test("uses pre-coerced transpiledOptions as-is (analysis is responsible for coercion)") { + // The Analyzer coerces transpiledOptions before the optimizer runs, because + // TranspiledPythonUDF.children exposes them to the resolver's generic coercion pass. + // ConvertToCatalyst must not re-run coercion; it simply selects the first non-null option. + // This test simulates what analysis would produce for `def f(x: Long): return x + 4` + // where the integer literal has already been cast to LongType. + transpileOn { + val preCoerced = Add(attrA, Cast(Literal(4, IntegerType), LongType)) + val tpudf = makeTPUDF(makePyUDF(), preCoerced) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(result == preCoerced, + s"Expected pre-coerced expression unchanged, got: $result") + } + } + + // ---- UDAF cases (post-fromUDFExpr shape) ---- + // + // After UserDefinedPythonFunction.fromUDFExpr lifts a PythonUDAF inside a + // TranspiledPythonUDF, the wrapper holds an AggregateExpression instead of a + // bare PythonUDAF. These tests pin the optimizer rule's behavior on that shape. + + test("transpiles TranspiledPythonUDF wrapping AggregateExpression(PythonUDAF)") { + transpileOn { + val pyAgg = makePyUDAF().toAggregateExpression() + val catalystAgg = Count(Seq(attrA)).toAggregateExpression() + val tpudf = TranspiledPythonUDF("agg", pyAgg, List(catalystAgg)) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + assert(result == catalystAgg, + s"Expected catalyst aggregate alternative, got: $result") + } + } + + test("falls back to AggregateExpression(PythonUDAF) when ANSI is off (UDAF)") { + ansiOff { + val pyAgg = makePyUDAF().toAggregateExpression() + val catalystAgg = Count(Seq(attrA)).toAggregateExpression() + val tpudf = TranspiledPythonUDF("agg", pyAgg, List(catalystAgg)) + val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false) + result match { + case ae: AggregateExpression => + assert(ae.aggregateFunction.isInstanceOf[PythonUDAF], + s"Expected aggregateFunction to be PythonUDAF, got: ${ae.aggregateFunction}") + case other => fail(s"Expected AggregateExpression(PythonUDAF, ...), got: $other") + } + } + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala index 622af60d85d93..f4d412153b0ed 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala @@ -21,12 +21,12 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow, LessThan, Literal, UnaryExpression} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ArrayTransform, Expression, GenericInternalRow, LambdaFunction, LessThan, Literal, NamedLambdaVariable, UnaryExpression} import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} import org.apache.spark.sql.catalyst.plans.PlanTest -import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Project} import org.apache.spark.sql.catalyst.rules.RuleExecutor -import org.apache.spark.sql.types.{DataType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, StructType} class ConvertToLocalRelationSuite extends PlanTest { @@ -87,6 +87,18 @@ class ConvertToLocalRelationSuite extends PlanTest { comparePlans(optimized, correctAnswer) } + + test("SPARK-58208: ConvertToLocalRelation uses fresh stateful project expressions") { + val element = NamedLambdaVariable("x", IntegerType, nullable = false) + val transform = ArrayTransform( + Literal.create(Seq(1, 2), ArrayType(IntegerType, containsNull = false)), + LambdaFunction(Add(element, Literal(1)), Seq(element))) + val project = Project(Seq(Alias(transform, "v")()), LocalRelation(Nil, Seq(InternalRow.empty))) + + Optimize.execute(project) + + assert(element.value.get() == null) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuerySuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuerySuite.scala index 1923c7836b236..a451c6961ba84 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuerySuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/DecorrelateInnerQuerySuite.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.dsl.plans._ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.IntegerType class DecorrelateInnerQuerySuite extends PlanTest { @@ -736,4 +737,81 @@ class DecorrelateInnerQuerySuite extends PlanTest { DomainJoin(Seq(a), testRelation2)))))) check(outputPlan, joinCond, correctAnswer, Seq(a <=> a)) } + + test("SPARK-58411: order by is preserved for limit with correlation only on outer table") { + // The correlated predicate references only outer columns (a, b), so no domain join is + // needed and the limit is computed directly over the inner query (partitionFields.isEmpty + // branch). The ORDER BY must be preserved so ORDER BY ... LIMIT stays deterministic. + val outerPlan = testRelation + val innerPlan = + Project(Seq(x), + Limit(1, Sort(Seq(SortOrder(x, Ascending)), true, + Filter(OuterReference(a) < OuterReference(b), + testRelation2)))) + val (outputPlan, joinCond) = DecorrelateInnerQuery(innerPlan, outerPlan.select()) + + val correctAnswer = + Project(Seq(x), + Limit(1, Sort(Seq(SortOrder(x, Ascending)), global = true, + testRelation2))) + check(outputPlan, joinCond, correctAnswer, Seq(a < b)) + } + + test("SPARK-58411: explicit ORDER BY ASC NULLS LAST is preserved for limit with correlation " + + "only on outer table") { + // An explicitly specified null ordering (NULLS LAST, the opposite of ASC's NULLS FIRST + // default) must be carried through decorrelation unchanged. + val outerPlan = testRelation + val nullsLast = SortOrder(x, Ascending, NullsLast, Seq.empty) + val innerPlan = + Project(Seq(x), + Limit(1, Sort(Seq(nullsLast), true, + Filter(OuterReference(a) < OuterReference(b), + testRelation2)))) + val (outputPlan, joinCond) = DecorrelateInnerQuery(innerPlan, outerPlan.select()) + + val correctAnswer = + Project(Seq(x), + Limit(1, Sort(Seq(nullsLast), global = true, + testRelation2))) + check(outputPlan, joinCond, correctAnswer, Seq(a < b)) + } + + test("SPARK-58411: order by is preserved for limit with offset and correlation only on " + + "outer table") { + // Same as above but with an OFFSET between the LIMIT and the ORDER BY. The ordering must + // be preserved for the LIMIT ... OFFSET case as well. + val outerPlan = testRelation + val innerPlan = + Project(Seq(x), + Limit(1, Offset(2, Sort(Seq(SortOrder(x, Ascending)), true, + Filter(OuterReference(a) < OuterReference(b), + testRelation2))))) + val (outputPlan, joinCond) = DecorrelateInnerQuery(innerPlan, outerPlan.select()) + + val correctAnswer = + Project(Seq(x), + Limit(1, Offset(2, Sort(Seq(SortOrder(x, Ascending)), global = true, + testRelation2)))) + check(outputPlan, joinCond, correctAnswer, Seq(a < b)) + } + + test("SPARK-58411: legacy flag restores the incorrect dropped-order behavior") { + withSQLConf( + SQLConf.DECORRELATE_LIMIT_OFFSET_LEGACY_INCORRECT_ORDER_HANDLING_ENABLED.key -> "true") { + val outerPlan = testRelation + val innerPlan = + Project(Seq(x), + Limit(1, Sort(Seq(SortOrder(x, Ascending)), true, + Filter(OuterReference(a) < OuterReference(b), + testRelation2)))) + val (outputPlan, joinCond) = DecorrelateInnerQuery(innerPlan, outerPlan.select()) + + // Legacy behavior: the Sort is dropped, leaving an arbitrary (non-deterministic) limit. + val correctAnswer = + Project(Seq(x), + Limit(1, testRelation2)) + check(outputPlan, joinCond, correctAnswer, Seq(a < b)) + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala index a43be9a1c0a66..517d310216327 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala @@ -1645,6 +1645,55 @@ class FilterPushdownSuite extends PlanTest { comparePlans(optimizedQueryWithoutStep, correctAnswer) } + test("SPARK-58627: do not push down predicate with raise_error through joins") { + val x = testStringRelation.subquery("x") + val y = testRelation1.subquery("y") + + // raise_error over literals references no columns, so it looks evaluable on either side and + // gets pushed into the left relation, where it fires on rows the join would have dropped. + val queryWithRaiseError = x.join(y, joinType = Inner, condition = Some($"x.a" === $"y.d")) + .where(IsNull(RaiseError(Literal("boom")))) + .analyze + comparePlans(Optimize.execute(queryWithRaiseError), queryWithRaiseError) + } + + test("SPARK-58627: do not push down predicate with a throwing child of sequence through joins") { + val x = testStringRelation.subquery("x") + val y = testRelation1.subquery("y") + + // Sequence overrides `throwable` for its step check, so it also has to fall back to its + // children. Without that fallback a RaiseError under a stepless sequence reports + // non-throwable and the predicate gets pushed below the join. + val raiseErrorInt = RaiseError( + Literal("USER_RAISED_EXCEPTION"), + CreateMap(Seq(Literal("errorMessage"), $"x.e")), + IntegerType) + val queryWithRaiseError = x.join(y, joinType = Inner, condition = Some($"x.a" === $"y.d")) + .where(IsNotNull(Sequence($"x.a", raiseErrorInt, None))) + .analyze + comparePlans(Optimize.execute(queryWithRaiseError), queryWithRaiseError) + } + + test("SPARK-58627: do not combine predicate with raise_error with other filters") { + val x = testStringRelation.subquery("x") + + // Do not combine. Two stacked Filters pin raise_error above the inner predicate, while a + // single merged And does not: execution does not guarantee the conjuncts are evaluated in + // order, and later rules are free to re-split and relocate them independently. Either way + // raise_error can end up evaluated on rows the inner filter would have removed. + val queryWithRaiseError = x.where($"x.a" > 1) + .where(IsNull(RaiseError($"x.e"))) + .analyze + comparePlans(Optimize.execute(queryWithRaiseError), queryWithRaiseError) + + // The same shape without raise_error is combined into a single filter. + val queryWithoutRaiseError = x.where($"x.a" > 1) + .where(IsNotNull($"x.e")) + .analyze + val correctAnswer = x.where(IsNotNull($"x.e") && $"x.a" > 1).analyze + comparePlans(Optimize.execute(queryWithoutRaiseError), correctAnswer) + } + test("push down deterministic predicate through BinBy") { // Relation: ts_start, ts_end, value (DISTRIBUTE), label (pass-through). val tsStart = AttributeReference("ts_start", TimestampType, nullable = false)() diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferWindowGroupLimitSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferWindowGroupLimitSuite.scala index 5aa7a27f65fba..5e0f978300ccb 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferWindowGroupLimitSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InferWindowGroupLimitSuite.scala @@ -20,9 +20,10 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{CurrentRow, DenseRank, Literal, NthValue, NTile, PercentRank, Rank, RowFrame, RowNumber, SpecifiedWindowFrame, UnboundedPreceding} +import org.apache.spark.sql.catalyst.expressions.{CurrentRow, DenseRank, Literal, NthValue, NTile, PercentRank, Rank, RowFrame, RowNumber, SpecifiedWindowFrame, UnboundedFollowing, UnboundedPreceding} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, Count} import org.apache.spark.sql.catalyst.plans.PlanTest -import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, WindowGroupLimit} import org.apache.spark.sql.catalyst.rules.RuleExecutor import org.apache.spark.sql.internal.SQLConf @@ -46,6 +47,17 @@ class InferWindowGroupLimitSuite extends PlanTest { LimitPushDownThroughWindow) :: Nil } + private object WithCollapseWindow extends RuleExecutor[LogicalPlan] { + val batches = + Batch("Insert WindowGroupLimit with CollapseWindow", FixedPoint(10), + CollapseWindow, + CollapseProject, + RemoveNoopOperators, + PushDownPredicates, + InferWindowGroupLimit, + LimitPushDownThroughWindow) :: Nil + } + private val testRelation = LocalRelation.fromExternalRows( Seq("a".attr.int, "b".attr.int, "c".attr.int), 1.to(6).map(_ => Row(1, 2, 3))) @@ -354,4 +366,32 @@ class InferWindowGroupLimitSuite extends PlanTest { Optimize.execute(originalQuery.analyze), WithoutOptimize.execute(originalQuery.analyze)) } + + test("SPARK-58757: collapseWindowWithEmptyOrderSpecInChild keeps WindowGroupLimit by default") { + val cnt = windowExpr( + AggregateExpression(Count(c), Complete, isDistinct = false, None), + windowSpec(a :: Nil, Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, UnboundedFollowing))).as("cnt") + val rn = windowExpr( + RowNumber(), + windowSpec(a :: Nil, c.desc :: Nil, + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))).as("rn") + + def analyzed: LogicalPlan = testRelation + .window(Seq(cnt), a :: Nil, Nil) + .window(Seq(rn), a :: Nil, c.desc :: Nil) + .where($"rn" <= 2) + .analyze + + // By default the config is off, so the empty-order child is not collapsed into the ordered + // parent and the WindowGroupLimit is preserved. + val defaultPlan = WithCollapseWindow.execute(analyzed) + assert(defaultPlan.collect { case _: WindowGroupLimit => 1 }.size == 1) + + // With the config on, the empty-order child is collapsed, disabling the WindowGroupLimit. + withSQLConf(SQLConf.COLLAPSE_WINDOW_WITH_EMPTY_ORDER_SPEC_IN_CHILD.key -> "true") { + val enabledPlan = WithCollapseWindow.execute(analyzed) + assert(enabledPlan.collect { case _: WindowGroupLimit => 1 }.isEmpty) + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala index 0d515c4824017..2daad7dfb64bb 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InlineCTESuite.scala @@ -108,4 +108,22 @@ class InlineCTESuite extends PlanTest { assert(e.getMessage.contains( "found a subquery with outer-scope reference")) } + + test("SPARK-58779: optimizer InlineCTE (isAnalysis = false) fails on a ref with no definition") { + // During analysis a CTERelationRef whose definition is not in the plan is tolerated -- it is + // owned by a surrounding scope (e.g. when `ResolveSQLTableFunctions` runs `checkAnalysis` on a + // table-function subplan whose argument references an outer CTE). In the optimizer the plan is + // complete, so a missing definition indicates corruption and must fail loudly rather than be + // dropped. + val defX = CTERelationDef(TestRelation(Seq($"a".int)).select($"a")) + val refX = CTERelationRef(defX.id, defX.resolved, defX.output, defX.isStreaming) + val danglingRef = CTERelationRef(defX.id + 1000, true, Seq($"a".int), false) + val plan = WithCTE(refX.union(danglingRef), Seq(defX)) + val e = intercept[SparkException] { + InlineCTE(isAnalysis = false).apply(plan) + } + assert(e.getCondition == "INTERNAL_ERROR") + // The analysis path tolerates the out-of-scope reference (no throw). + InlineCTE(isAnalysis = true).apply(plan) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInAggregateSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInAggregateSuite.scala new file mode 100644 index 0000000000000..fa2340318cbca --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/InsertMapSortInAggregateSuite.scala @@ -0,0 +1,117 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions.{Alias, MapSort} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.PlanTest +import org.apache.spark.sql.catalyst.plans.logical.{ + Aggregate, LocalRelation, LogicalPlan, Project, Union} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} + +class InsertMapSortInAggregateSuite extends PlanTest { + private val input = LocalRelation(Symbol("m").map(StringType, IntegerType)) + private val mapAttribute = input.output.head + + private def aliasesNamed(plan: LogicalPlan, name: String): Seq[Alias] = { + plan.flatMap { node => + node.expressions.flatMap(_.collect { + case alias @ Alias(_, aliasName) if aliasName == name => alias + }) + } + } + + test("reuse map sort when a grouping key is also a distinct argument") { + val plan = Aggregate( + Seq(mapAttribute), + Seq(mapAttribute, countDistinct(mapAttribute).as("count")), + input) + val rewritten = InsertMapSortInAggregate(plan) + val groupingAliases = aliasesNamed(rewritten, "_groupingmapsort") + + assert(groupingAliases.size == 1) + assert(groupingAliases.head.child.isInstanceOf[MapSort]) + rewritten match { + case Aggregate(Seq(groupingExpression), aggregateExpressions, _: Project, _) => + assert(groupingExpression.semanticEquals(groupingAliases.head.toAttribute)) + val distinctChildren = aggregateExpressions.flatMap(_.collect { + case expression: AggregateExpression if expression.isDistinct => + expression.aggregateFunction.children + }).flatten + assert(distinctChildren.size == 1) + assert(distinctChildren.head.semanticEquals(groupingAliases.head.toAttribute)) + case other => + fail(s"Unexpected plan:\n$other") + } + } + + test("project complex distinct arguments only when needed") { + val attributePlan = Aggregate( + Nil, + Seq(countDistinct(mapAttribute).as("count")), + input) + val complexPlan = Aggregate( + Nil, + Seq(countDistinct(namedStruct("m", mapAttribute)).as("count")), + input) + + val rewrittenAttributePlan = InsertMapSortInAggregate(attributePlan) + assert(rewrittenAttributePlan.collect { case _: Project => 1 }.size == 1) + assert(aliasesNamed(rewrittenAttributePlan, "_distinctaggregateexpression").isEmpty) + assert(aliasesNamed(rewrittenAttributePlan, "_distinctmapsort").size == 1) + + val rewrittenComplexPlan = InsertMapSortInAggregate(complexPlan) + assert(rewrittenComplexPlan.collect { case _: Project => 1 }.size == 2) + assert(aliasesNamed(rewrittenComplexPlan, "_distinctaggregateexpression").size == 1) + assert(aliasesNamed(rewrittenComplexPlan, "_distinctmapsort").size == 1) + } + + test("skip distinct argument normalization when disabled") { + val plan = Aggregate( + Nil, + Seq(countDistinct(mapAttribute).as("count")), + input) + + withSQLConf(SQLConf.INSERT_MAP_SORT_IN_DISTINCT_AGGREGATES_ENABLED.key -> "false") { + comparePlans(InsertMapSortInAggregate(plan), plan) + } + } + + test("leave map-free aggregates untouched when another aggregate needs rewriting") { + val scalarInput = LocalRelation(Symbol("i").int) + val scalarAggregate = Aggregate( + Nil, + Seq(count(scalarInput.output.head).as("count")), + scalarInput) + val mapAggregate = Aggregate( + Nil, + Seq(countDistinct(mapAttribute).as("count")), + input) + + InsertMapSortInAggregate(Union(Seq(scalarAggregate, mapAggregate))) match { + case Union(Seq(rewrittenScalarAggregate, rewrittenMapAggregate), _, _) => + comparePlans(rewrittenScalarAggregate, scalarAggregate) + assert(rewrittenScalarAggregate.collect { case _: Project => 1 }.isEmpty) + assert(rewrittenMapAggregate.collect { case _: Project => 1 }.size == 1) + case other => + fail(s"Unexpected plan:\n$other") + } + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala index bac20c6ed3533..61ef8be59a883 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/JoinSelectionHelperSuite.scala @@ -171,4 +171,56 @@ class JoinSelectionHelperSuite extends PlanTest with JoinSelectionHelper { } } + test("getBroadcastHashJoinBuildSide returns the hinted side") { + val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head))) + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + assert(getBroadcastHashJoinBuildSide( + equiJoin.copy(hint = JoinHint(hintBroadcast, None)), SQLConf.get) === Some(BuildLeft)) + assert(getBroadcastHashJoinBuildSide( + equiJoin.copy(hint = JoinHint(None, hintBroadcast)), SQLConf.get) === Some(BuildRight)) + // Both sides hinted: the smaller one wins, as in `getBroadcastBuildSide`. + assert(getBroadcastHashJoinBuildSide( + equiJoin.copy(hint = JoinHint(hintBroadcast, hintBroadcast)), SQLConf.get) === + Some(BuildRight)) + } + } + + test("getBroadcastHashJoinBuildSide falls back to the smaller side") { + val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head))) + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + // Only the right side is under the threshold, so it is the only candidate. + assert(getBroadcastHashJoinBuildSide(equiJoin, SQLConf.get) === Some(BuildRight)) + } + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + assert(getBroadcastHashJoinBuildSide(equiJoin, SQLConf.get).isEmpty) + } + } + + test("getBroadcastHashJoinBuildSide returns None when a shuffle hash hint applies") { + val equiJoin = join.copy(condition = Some(EqualTo(left.output.head, right.output.head))) + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + assert(getBroadcastHashJoinBuildSide( + equiJoin.copy(hint = JoinHint(None, hintShuffleHash)), SQLConf.get).isEmpty) + } + } + + test("getBroadcastHashJoinBuildSide returns None without equi-join keys") { + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + assert(getBroadcastHashJoinBuildSide(join, SQLConf.get).isEmpty) + } + } + + test("getBroadcastHashJoinBuildSide builds from the right for a null-aware anti join") { + val leftKey = left.output.head + val rightKey = right.output.head + val condition = Or(EqualTo(leftKey, rightKey), IsNull(EqualTo(leftKey, rightKey))) + val nullAwareAntiJoin = Join(left, right, LeftAnti, Some(condition), JoinHint.NONE) + + withSQLConf( + SQLConf.OPTIMIZE_NULL_AWARE_ANTI_JOIN.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10MB") { + assert(getBroadcastHashJoinBuildSide(nullAwareAntiJoin, SQLConf.get) === Some(BuildRight)) + } + } + } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinConditionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinConditionSuite.scala index e7f090ec4d0dc..c686eb6e680de 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinConditionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJoinConditionSuite.scala @@ -46,4 +46,41 @@ class OptimizeJoinConditionSuite extends PlanTest { comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) }) } + + test("SPARK-58384: do not replace null-safe equality pattern under NOT") { + val x = testRelation.subquery("x") + val y = testRelation1.subquery("y") + val originalQuery = + x.join(y, Inner, Option(!($"a" === $"c" || ($"a".isNull && $"c".isNull)))) + + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-58384: replace null-safe equality pattern under AND and OR") { + val x = testRelation.subquery("x") + val y = testRelation1.subquery("y") + val pattern = $"a" === $"c" || ($"a".isNull && $"c".isNull) + val optimizedPattern = $"a" <=> $"c" + val otherCondition = $"b" === $"d" + val conditions = Seq( + (pattern && otherCondition) -> (optimizedPattern && otherCondition), + (pattern || otherCondition) -> (optimizedPattern || otherCondition)) + + conditions.foreach { case (condition, optimizedCondition) => + val originalQuery = x.join(y, Inner, Option(condition)) + val correctAnswer = x.join(y, Inner, Option(optimizedCondition)) + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + } + + test("SPARK-58384: preserve unrelated AND and OR nodes") { + val x = testRelation.subquery("x") + val y = testRelation1.subquery("y") + val condition = ($"a" === $"c") && (($"b" === $"d") || ($"a" === 1)) + val originalQuery = x.join(y, Inner, Option(condition)).analyze.asInstanceOf[Join] + + val optimized = OptimizeJoinCondition(originalQuery).asInstanceOf[Join] + + assert(optimized.condition.get eq originalQuery.condition.get) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala index dd0f3dd013277..801404ba55a06 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeJsonExprsSuite.scala @@ -690,6 +690,67 @@ class OptimizeJsonExprsSuite extends PlanTest with ExpressionEvalHelper { comparePlans(optimized2, query2.analyze) } + test("SPARK-58373: do not simplify named_struct + from_json if options is not empty") { + val schema = StructType.fromDDL("a int, b int, c long") + def query(options: Map[String, String]): LogicalPlan = testRelation2.select(namedStruct( + "a", GetStructField(JsonToStructs(schema, options, $"json"), 0), + "b", GetStructField(JsonToStructs(schema, options, $"json"), 1)).as("struct")).analyze + + // Any option disables the rewrite, not just a parse mode: pruning `c` out of the schema + // stops the parser from converting it, which a caller may be relying on. + Seq(Map("mode" -> "failfast"), Map("timestampFormat" -> "yyyy")).foreach { options => + comparePlans(Optimizer.execute(query(options)), query(options)) + } + + // Control: the same shape with no options is still rewritten, so the guard above is what + // blocks it rather than some other precondition of the rule. + val optimized = Optimizer.execute(query(Map.empty)) + assert(optimized != query(Map.empty), "expected the empty-options plan to be rewritten") + } + + test("SPARK-58707: do not prune a from_json schema down to the corrupt record column") { + val schema = StructType.fromDDL("a int, b int, _corrupt_record string") + + def getField(ordinal: Int): LogicalPlan = testRelation2 + .select(GetStructField(JsonToStructs(schema, Map.empty, $"json"), ordinal)).analyze + + // Pruning to the corrupt record column alone leaves the parser nothing to convert, so a + // malformed value in a dropped field would no longer populate the column. + comparePlans(Optimizer.execute(getField(2)), getField(2)) + + // Control: a non-corrupt field is still pruned, so the guard above is what blocks it rather + // than some other precondition of the rule. + val prunedSchema = StructType.fromDDL("a int") + comparePlans( + Optimizer.execute(getField(0)), + testRelation2 + .select(GetStructField(JsonToStructs(prunedSchema, Map.empty, $"json"), 0)).analyze) + } + + test("SPARK-58707: simplify named_struct + from_json when no field is dropped") { + val schema = StructType.fromDDL("a int, _corrupt_record string") + + def query(fields: (String, Int)*): LogicalPlan = testRelation2.select( + namedStruct(fields.flatMap { case (name, ordinal) => + Seq(Literal(name), GetStructField(JsonToStructs(schema, Map.empty, $"json"), ordinal)) + }: _*).as("struct")).analyze + + // Selecting the corrupt record column together with every other field drops nothing, so the + // rewrite is safe and still collapses the repeated parses into one. + val nullStruct = namedStruct( + "a", Literal(null, IntegerType), "_corrupt_record", Literal(null, StringType)) + comparePlans( + Optimizer.execute(query("a" -> 0, "_corrupt_record" -> 1)), + testRelation2.select( + If(IsNull($"json"), + nullStruct, + KnownNotNull(JsonToStructs(schema, Map.empty, $"json"))).as("struct")).analyze) + + // Dropping `a` while selecting the corrupt record column is what makes the column unreliable. + val pruning = query("_corrupt_record" -> 1) + comparePlans(Optimizer.execute(pruning), pruning) + } + test("SPARK-33007: simplify named_struct + from_json") { val options = Map.empty[String, String] val schema = StructType.fromDDL("a int, b int, c long, d string") diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala new file mode 100644 index 0000000000000..25d62cb016b2d --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullUpProjectAliasThroughWindowSuite.scala @@ -0,0 +1,315 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.rules._ +import org.apache.spark.sql.types.MetadataBuilder + +class PullUpProjectAliasThroughWindowSuite extends PlanTest { + + private object Optimize extends RuleExecutor[LogicalPlan] { + val batches = + Batch("Pull up project alias through window", FixedPoint(20), + PullUpProjectAliasThroughWindow) :: Nil + } + + private val testRelation = LocalRelation($"key".int, $"key2".int, $"value".string) + private val key = testRelation.output(0) + private val key2 = testRelation.output(1) + private val value = testRelation.output(2) + private val windowFrame = SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow) + + // Builds `Window [row_number() ... AS <name>], <partitionSpec>, <orderSpec>` over `child`. + private def windowOver( + child: LogicalPlan, + name: String, + partitionSpec: Seq[Expression], + orderSpec: Seq[SortOrder] = Nil): Window = { + val spec = windowSpec(partitionSpec, orderSpec, windowFrame) + val winExpr = windowExpr(RowNumber(), spec).as(name) + Window(Seq(winExpr), partitionSpec, orderSpec, child) + } + + test("pull up a rename of the window partition key") { + // Project [userid, w] <- both bare attributes + // +- Window [... AS w], [key] + // +- Project [key AS userid, value, key] + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, w), window) + + // `userid` is pulled up into the parent project; the lower project keeps the rest. + val prunedBottom = Project(Seq(value, key), testRelation) + val expected = Project( + Seq(Alias(key, "userid")(exprId = userid.exprId), w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("pull up every applicable key in a multi-key window") { + val u1 = Alias(key, "u1")() + val u2 = Alias(key2, "u2")() + val bottom = Project(Seq(u1, u2, value, key, key2), testRelation) + val window = windowOver(bottom, "w", key :: key2 :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(u1.toAttribute, u2.toAttribute, w), window) + + val prunedBottom = Project(Seq(value, key, key2), testRelation) + val expected = Project( + Seq( + Alias(key, "u1")(exprId = u1.exprId), + Alias(key2, "u2")(exprId = u2.exprId), + w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("pull up a rename of the window order key") { + // Window ordered by `value`; the top Project passes `tstamp` (a rename of `value`) through as + // a bare attribute. It must be pulled up so the window's output ordering projects through it. + val tstamp = Alias(value, "tstamp")() + val bottom = Project(Seq(key, tstamp, value), testRelation) + val window = windowOver(bottom, "w", key :: Nil, SortOrder(value, Ascending) :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(key, tstamp.toAttribute, w), window) + + val prunedBottom = Project(Seq(key, value), testRelation) + val expected = Project( + Seq(key, Alias(value, "tstamp")(exprId = tstamp.exprId), w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("pull up renames of both partition and order keys") { + val userid = Alias(key, "userid")() + val tstamp = Alias(value, "tstamp")() + val bottom = Project(Seq(userid, tstamp, key, value), testRelation) + val window = windowOver(bottom, "w", key :: Nil, SortOrder(value, Ascending) :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, tstamp.toAttribute, w), window) + + val prunedBottom = Project(Seq(key, value), testRelation) + val expected = Project( + Seq( + Alias(key, "userid")(exprId = userid.exprId), + Alias(value, "tstamp")(exprId = tstamp.exprId), + w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("keep a bare pass-through column the window does not reference below the window") { + // `key2` is a bare pass-through that the window neither partitions nor orders by. It must stay + // below the window so the window keeps producing it for the parent project to reference; only + // the rename `key AS userid` is pulled up. + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, key2, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, key2, w), window) + + val prunedBottom = Project(Seq(key2, value, key), testRelation) + val expected = Project( + Seq(Alias(key, "userid")(exprId = userid.exprId), key2, w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("no rewrite when the renamed column is not a window key") { + // `userid` renames `key`, but the window is partitioned by `value`. + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val window = windowOver(bottom, "w", value :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, w), window) + + comparePlans(Optimize.execute(originalQuery), originalQuery) + } + + test("no rewrite for a computed alias whose input is not a window key") { + // `(key2 + 1) AS z` depends on `key2`, which the window does not produce once pruned, so it + // cannot be pulled above the window and must stay below. + val z = Alias(key2 + 1, "z")() + val bottom = Project(Seq(z, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(z.toAttribute, w), window) + + comparePlans(Optimize.execute(originalQuery), originalQuery) + } + + test("pull up a computed alias whose inputs are all window keys") { + // `(key + 1) AS z` is computed rather than a pure rename, but its only input `key` is the + // window partition key and is retained below. It is pulled up so the expression is evaluated + // in the top project, above the window's shuffle rather than being carried across it. This + // narrows the shuffled data even though it yields no partitioning benefit (`HashPartitioning` + // on `key` does not project through `key + 1`). + val z = Alias(key + 1, "z")() + val bottom = Project(Seq(z, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(z.toAttribute, w), window) + + val prunedBottom = Project(Seq(value, key), testRelation) + val expected = Project( + Seq(Alias(key + 1, "z")(exprId = z.exprId), w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("pull up a computed alias combining a partition key and an order key") { + // `(key + key2) AS z` combines the partition key `key` and the order key `key2`; both are + // retained below, so the whole expression lifts above the window. + val z = Alias(key + key2, "z")() + val bottom = Project(Seq(z, key, key2), testRelation) + val window = windowOver(bottom, "w", key :: Nil, SortOrder(key2, Ascending) :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(z.toAttribute, w), window) + + val prunedBottom = Project(Seq(key, key2), testRelation) + val expected = Project( + Seq(Alias(key + key2, "z")(exprId = z.exprId), w), + window.copy(child = prunedBottom)) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("no rewrite when the window child is not a Project") { + val window = windowOver(testRelation, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(key, w), window) + + comparePlans(Optimize.execute(originalQuery), originalQuery) + } + + test("pull up an alias across a chain of windows") { + // Project [userid, w1, w2] + // +- Window [... AS w2], [key], [value DESC] + // +- Window [... AS w1], [key], [value] <- adjacent, no Project between + // +- Project [key AS userid, value, key] + // `userid` renames `key` and is referenced by neither window, so it is pulled all the way up + // to the top project, through both windows. + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val w1 = windowOver(bottom, "w1", key :: Nil, SortOrder(value, Ascending) :: Nil) + val w1a = w1.windowExpressions.head.toAttribute + val w2 = windowOver(w1, "w2", key :: Nil, SortOrder(value, Descending) :: Nil) + val w2a = w2.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, w1a, w2a), w2) + + val prunedBottom = Project(Seq(value, key), testRelation) + val newW1 = w1.copy(child = prunedBottom) + val newW2 = w2.copy(child = newW1) + val expected = Project( + Seq(Alias(key, "userid")(exprId = userid.exprId), w1a, w2a), newW2) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("keep an alias referenced by an inner window below the chain") { + // The chain's inner window orders by `tstamp` (a rename of `value`), so `tstamp` is referenced + // by a window and must stay below; only `userid` (referenced by no window) is pulled up. + val userid = Alias(key, "userid")() + val tstamp = Alias(value, "tstamp")() + val bottom = Project(Seq(userid, tstamp, key, value), testRelation) + val w1 = windowOver(bottom, "w1", key :: Nil, SortOrder(tstamp.toAttribute, Ascending) :: Nil) + val w1a = w1.windowExpressions.head.toAttribute + val w2 = windowOver(w1, "w2", key :: Nil) + val w2a = w2.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, tstamp.toAttribute, w1a, w2a), w2) + + // `tstamp` stays below (referenced by w1's order spec); `userid` is pulled up. + val prunedBottom = Project(Seq(tstamp, key, value), testRelation) + val newW1 = w1.copy(child = prunedBottom) + val newW2 = w2.copy(child = newW1) + val expected = Project( + Seq(Alias(key, "userid")(exprId = userid.exprId), tstamp.toAttribute, w1a, w2a), newW2) + comparePlans(Optimize.execute(originalQuery), expected) + } + + test("no rewrite for a nondeterministic alias") { + // `spark_partition_id() AS pid` is a leaf with no references, so it would pass the input- + // survival check vacuously, but moving it above the window's exchange/sort would change its + // per-partition value. It must stay below. + val pid = Alias(SparkPartitionID(), "pid")() + val bottom = Project(Seq(pid, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(pid.toAttribute, w), window) + + comparePlans(Optimize.execute(originalQuery), originalQuery) + } + + test("rewrite preserves the top attribute's name, qualifier, and metadata") { + // The lookup is keyed by expr id only, so a resolved top attribute can carry a different name, + // qualifier, and metadata than the lower alias (same expr id). The rebuilt alias must adopt the + // top attribute's full identity so the output schema is byte-for-byte unchanged. + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + // The top project references `userid` under a different name, qualifier, and metadata. + val metadata = new MetadataBuilder().putString("comment", "external").build() + val topAttr = AttributeReference("external", key.dataType, key.nullable, metadata)( + exprId = userid.exprId, qualifier = Seq("sub")) + val originalQuery = Project(Seq(topAttr, w), window) + + val prunedBottom = Project(Seq(value, key), testRelation) + val expected = Project( + Seq( + Alias(key, "external")( + exprId = userid.exprId, qualifier = Seq("sub"), explicitMetadata = Some(metadata)), + w), + window.copy(child = prunedBottom)) + val optimized = Optimize.execute(originalQuery) + comparePlans(optimized, expected) + // The output attribute must match the top attribute exactly, not the lower alias's identity. + val out = optimized.output.head + assert(out.name == "external") + assert(out.qualifier == Seq("sub")) + assert(out.metadata == metadata) + } + + test("rewrite is idempotent") { + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, w), window) + + val once = Optimize.execute(originalQuery) + val twice = Optimize.execute(once) + comparePlans(once, twice) + } + + test("rewrite preserves the output schema (exprId, name, type, nullability)") { + val userid = Alias(key, "userid")() + val bottom = Project(Seq(userid, value, key), testRelation) + val window = windowOver(bottom, "w", key :: Nil) + val w = window.windowExpressions.head.toAttribute + val originalQuery = Project(Seq(userid.toAttribute, w), window) + + val optimized = Optimize.execute(originalQuery) + // The output must be byte-for-byte identical so that downstream references still resolve. + assert(optimized.output === originalQuery.output) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnionSuite.scala index 65d4621063f84..37ba77312bbe1 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushDownJoinThroughUnionSuite.scala @@ -19,10 +19,12 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Explode, Rand} +import org.apache.spark.sql.catalyst.expressions.{AttributeMap, Explode, Rand} import org.apache.spark.sql.catalyst.plans._ -import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Union} +import org.apache.spark.sql.catalyst.plans.logical.{BROADCAST, HintInfo, Join, JoinHint, LocalRelation, + LogicalPlan, SHUFFLE_HASH, SHUFFLE_MERGE, SHUFFLE_REPLICATE_NL, Union} import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.catalyst.statsEstimation.StatsTestPlan import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.IntegerType @@ -42,6 +44,16 @@ class PushDownJoinThroughUnionSuite extends PlanTest { val testRelation3 = LocalRelation($"e".int, $"f".int) val testRelation4 = LocalRelation($"g".int, $"h".int) + // A Union whose branches are far too large to broadcast. The 0-byte empty LocalRelations above + // cannot express that: any non-leaf right side is estimated at 1 byte at least, so + // `getSmallerSide` picks the 0-byte left one and the rule does not fire. A `def` so that every + // test gets fresh ExprIds. + private def largeStatsUnion = Union( + StatsTestPlan(Seq($"a".int, $"b".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024)), + StatsTestPlan(Seq($"c".int, $"d".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024))) + test("Push down Inner Join through Union when right side is small") { val union = Union(testRelation1, testRelation2) val query = union.join(testRelation3, Inner, Some($"a" === $"e")) @@ -139,7 +151,7 @@ class PushDownJoinThroughUnionSuite extends PlanTest { val complexRight = testRelation3 .where($"f" > 0) .select($"e", ($"f" + 1).as("f_plus_1")) - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(complexRight, Inner, Some($"a" === $"e")) val optimized = Optimize.execute(query.analyze) @@ -161,7 +173,7 @@ class PushDownJoinThroughUnionSuite extends PlanTest { val rightWithGenerate = arrayRelation .generate(Explode($"arr"), outputNames = Seq("exploded_val")) .select($"k", $"exploded_val") - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(rightWithGenerate, Inner, Some($"a" === $"k")) val optimized = Optimize.execute(query.analyze) @@ -179,7 +191,7 @@ class PushDownJoinThroughUnionSuite extends PlanTest { test("Push down when right side contains SubqueryAlias") { val rightWithAlias = testRelation3.subquery("dim") - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(rightWithAlias, Inner, Some($"a" === $"e")) val optimized = Optimize.execute(query.analyze) @@ -198,7 +210,7 @@ class PushDownJoinThroughUnionSuite extends PlanTest { test("Push down when right side contains Project with Alias") { val rightWithAlias = testRelation3 .select($"e", ($"f" + 1).as("f_plus_1")) - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(rightWithAlias, Inner, Some($"a" === $"e")) val optimized = Optimize.execute(query.analyze) @@ -215,7 +227,7 @@ class PushDownJoinThroughUnionSuite extends PlanTest { test("Push down when right side contains Aggregate") { val rightWithAgg = testRelation3 .groupBy($"e")(count($"f").as("cnt"), $"e") - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(rightWithAgg, Inner, Some($"a" === $"e")) val optimized = Optimize.execute(query.analyze) @@ -232,10 +244,153 @@ class PushDownJoinThroughUnionSuite extends PlanTest { test("Do not push down when right side contains non-deterministic expressions") { val rightWithRand = testRelation3 .select($"e", Rand(10).as("rand_val")) - val union = Union(testRelation1, testRelation2) + val union = largeStatsUnion val query = union.join(rightWithRand, Inner, Some($"a" === $"e")) val optimized = Optimize.execute(query.analyze) comparePlans(optimized, query.analyze) } + + test("SPARK-58449: do not push down Inner Join when only the Union side is broadcastable") { + // For an inner join the planner may broadcast either side, so a broadcastable Union on the + // left is enough to make the join a broadcast hash join. Pushing down then clones the large + // right side once per branch, and nothing reuses those scans, so the right side is read N + // times instead of once. + val smallUnion = Union( + StatsTestPlan(Seq($"a".int, $"b".int), 10, AttributeMap.empty, Some(100)), + StatsTestPlan(Seq($"c".int, $"d".int), 10, AttributeMap.empty, Some(100))) + val largeRight = StatsTestPlan(Seq($"e".int, $"f".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024)) + + val query = smallUnion.join(largeRight, Inner, Some($"a" === $"e")) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: push down Inner Join when the right side is broadcastable") { + // Control for the case above: the same shape with a small right side is still pushed down. + val smallRight = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(100)) + + val query = largeStatsUnion.join(smallRight, Inner, Some($"a" === $"e")) + val optimized = Optimize.execute(query.analyze) + + val union = optimized.asInstanceOf[Union] + assert(union.children.size == 2) + assert(union.children.forall(_.isInstanceOf[Join])) + } + + test("SPARK-58449: do not push down when the right side is the build side only for the Union") { + // The right side is smaller than the whole Union but larger than either branch, so the join + // broadcasts the right side before the rewrite and would build from the left after it. A guard + // that checks the whole join instead of each branch misses this. + val union = Union( + StatsTestPlan(Seq($"a".int, $"b".int), 10, AttributeMap.empty, Some(400)), + StatsTestPlan(Seq($"c".int, $"d".int), 10, AttributeMap.empty, Some(400))) + val right = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(500)) + + val query = union.join(right, Inner, Some($"a" === $"e")) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: push down Left Outer Join when the right side is the build side only for " + + "the Union") { + // Only an inner join can build from the left, so a left outer join broadcasts the right side + // whatever the sizes are. The shape that blocks the inner join above is still pushed down. + val union = Union( + StatsTestPlan(Seq($"a".int, $"b".int), 10, AttributeMap.empty, Some(400)), + StatsTestPlan(Seq($"c".int, $"d".int), 10, AttributeMap.empty, Some(400))) + val right = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(500)) + + val query = union.join(right, LeftOuter, Some($"a" === $"e")) + val optimized = Optimize.execute(query.analyze) + + val optimizedUnion = optimized.asInstanceOf[Union] + assert(optimizedUnion.children.size == 2) + assert(optimizedUnion.children.forall(_.isInstanceOf[Join])) + } + + test("SPARK-58449: push down when a broadcast hint names the right side") { + // A hint is honored ahead of the sizes, so the right side is the build side even though it is + // the larger one. + val largeRight = StatsTestPlan(Seq($"e".int, $"f".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024)) + val hint = JoinHint(None, Some(HintInfo(Some(BROADCAST)))) + + val query = Join(largeStatsUnion, largeRight, Inner, Some($"a" === $"e"), hint) + val optimized = Optimize.execute(query.analyze) + + val union = optimized.asInstanceOf[Union] + assert(union.children.size == 2) + assert(union.children.forall(_.isInstanceOf[Join])) + } + + test("SPARK-58449: do not push down when a hint picks a non-broadcast strategy") { + // The planner honors the merge hint over a size-based broadcast, so no branch would broadcast + // the right side and the rewrite would only multiply the joins. + val smallRight = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(100)) + val hint = JoinHint(None, Some(HintInfo(Some(SHUFFLE_MERGE)))) + + val query = Join(largeStatsUnion, smallRight, Inner, Some($"a" === $"e"), hint) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: do not push down when a hint asks to replicate the right side") { + // A cartesian product requires no distribution on either side, so nothing would broadcast and + // nothing would put the duplicated right side behind a reusable exchange. + val smallRight = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(100)) + val hint = JoinHint(None, Some(HintInfo(Some(SHUFFLE_REPLICATE_NL)))) + + val query = Join(largeStatsUnion, smallRight, Inner, Some($"a" === $"e"), hint) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: do not push down when a shuffle hash hint decided the strategy") { + // The co-guard rejects the join once a shuffle hash hint applies, so the rewrite never gets to + // ask about the build side. Nothing else covers that delegation. + val smallRight = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(100)) + val hint = JoinHint(None, Some(HintInfo(Some(SHUFFLE_HASH)))) + + val query = Join(largeStatsUnion, smallRight, Inner, Some($"a" === $"e"), hint) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: do not push down when only some branches would broadcast the right side") { + // The first branch is far too large to build from, the second is smaller than the right side. + // The rewrite is all or nothing: the second branch would probe its own copy of the right side, + // so neither branch is pushed down. + val union = Union( + StatsTestPlan(Seq($"a".int, $"b".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024)), + StatsTestPlan(Seq($"c".int, $"d".int), 10, AttributeMap.empty, Some(100))) + val right = StatsTestPlan(Seq($"e".int, $"f".int), 10, AttributeMap.empty, Some(500)) + + val query = union.join(right, Inner, Some($"a" === $"e")) + val optimized = Optimize.execute(query.analyze) + + comparePlans(optimized, query.analyze) + } + + test("SPARK-58449: a broadcast hint outranks a non-broadcast hint on the other side") { + // The planner tries a hinted broadcast before the merge hint, so the right side is still the + // build side and the rewrite applies. + val largeRight = StatsTestPlan(Seq($"e".int, $"f".int), 1000000, AttributeMap.empty, + Some(100 * 1024 * 1024)) + val hint = JoinHint(Some(HintInfo(Some(SHUFFLE_MERGE))), Some(HintInfo(Some(BROADCAST)))) + + val query = Join(largeStatsUnion, largeRight, Inner, Some($"a" === $"e"), hint) + val optimized = Optimize.execute(query.analyze) + + val union = optimized.asInstanceOf[Union] + assert(union.children.size == 2) + assert(union.children.forall(_.isInstanceOf[Join])) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala new file mode 100644 index 0000000000000..b696ec307deeb --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesForCTEDefStalenessSuite.scala @@ -0,0 +1,366 @@ +/* + * 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.catalyst.optimizer + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.{Ascending, CurrentRow, RowFrame, RowNumber} +import org.apache.spark.sql.catalyst.expressions.{EqualTo, Expression, GreaterThan} +import org.apache.spark.sql.catalyst.expressions.{IsNotNull, Literal} +import org.apache.spark.sql.catalyst.expressions.{SortOrder, SpecifiedWindowFrame, UnboundedPreceding} +import org.apache.spark.sql.catalyst.expressions.{WindowExpression, WindowSpecDefinition} +import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, Join, JoinHint, LocalRelation} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project, Union, Window, WithCTE} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.types.IntegerType + +/** + * Regression test for a staleness defect in + * [[PushdownPredicatesAndPruneColumnsForCTEDef]]: on its second pass the rule rebuilt a + * CTE definition from the plan snapshot stored in `originalPlanWithPredicates`, which was + * taken during its first pass. Any change made to the CTE definition's child by other rules + * in between (e.g. filters injected by `InferFiltersFromConstraints`, which runs in the + * "Infer Filters" batch sandwiched between the two fixed-point batches that both contain + * this rule) was silently discarded by the rebuild. + * + * The inter-pass mutations are produced by applying the real rules + * ([[InferFiltersFromConstraints]] and [[PushPredicateThroughNonJoin]]) between two real + * applications of the rule under test, so the suite keeps exercising the actual cross-rule + * interaction as those rules evolve. + */ +class PushdownPredicatesForCTEDefStalenessSuite extends PlanTest { + + test("CTE def rebuild must not discard filters injected between rule passes") { + val t1a = AttributeReference("a", IntegerType, nullable = true)() + val t2b = AttributeReference("b", IntegerType, nullable = true)() + val t1 = LocalRelation(t1a) + val t2 = LocalRelation(t2b) + val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE) + + val cteId = 0L + val cteDef = CTERelationDef(join, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + // Pass 1 (batch "Operator Optimization before Inferring Filters"): the rule pushes the + // combined reference predicates into the CTE definition and records them in + // originalPlanWithPredicates. + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + assert(theOnlyDef(afterPass1).originalPlanWithPredicates.isDefined) + + // Run the real "Infer Filters" batch rule between the two passes (it is sandwiched + // between the two fixed-point batches that both contain this rule). It strengthens the + // pushed filter itself (propagating a = 5 | a = 7 through the join condition onto b), + // injects isnotnull filters below the join, and enriches the reference sites with + // isnotnull, which re-arms the rule's guard on the next pass. + val afterInfer = InferFiltersFromConstraints.apply(afterPass1) + val injectedConds = theOnlyDef(afterInfer).child.collect { case f: Filter => f.condition } + val pass1FilterCount = theOnlyDef(afterPass1).child.collect { case f: Filter => f }.length + assert(injectedConds.length > pass1FilterCount, + "test setup failed: InferFiltersFromConstraints did not inject filters into the " + + "CTE definition") + + // Pass 2 (batch "Operator Optimization after Inferring Filters"): the rule sees the + // enriched reference predicates and rebuilds the definition. Every filter present + // after the Infer Filters batch must survive the rebuild; before the fix the rebuild + // used the stale first-pass snapshot and discarded all of them. + val defAfterPass2 = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(afterInfer)) + injectedConds.foreach { cond => + assert(hasFilterOn(defAfterPass2.child, cond), + s"PushdownPredicatesAndPruneColumnsForCTEDef discarded a filter that " + + s"InferFiltersFromConstraints injected between its two passes: $cond") + } + } + + test("rule is idempotent when no new predicates appear between passes") { + val t1a = AttributeReference("a", IntegerType, nullable = true)() + val t2b = AttributeReference("b", IntegerType, nullable = true)() + val t1 = LocalRelation(t1a) + val t2 = LocalRelation(t2b) + val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE) + + val cteId = 0L + val cteDef = CTERelationDef(join, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + // First application pushes the combined reference predicates and records them. + val once = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + // A second application with no new reference predicates must be a no-op: + // no re-push, no stacked filters. + val twice = PushdownPredicatesAndPruneColumnsForCTEDef.apply(once) + comparePlans(once, twice) + + // Even after a foreign mutation inside the CTE definition (e.g. a filter injected + // by InferFiltersFromConstraints), as long as no NEW reference-site predicate + // appears, the rule must leave the plan untouched. (The mutation is hand-written + // because no real rule mutates only the definition without also enriching the + // reference sites, which is exactly what would re-arm the guard.) + val mutated = once.transform { + case d @ CTERelationDef(Filter(cond, j: Join), `cteId`, Some(_), _, _, _) => + d.copy(child = Filter(cond, j.copy(right = Filter(GreaterThan(t2b, Literal(0)), j.right)))) + } + val afterMutation = PushdownPredicatesAndPruneColumnsForCTEDef.apply(mutated) + comparePlans(mutated, afterMutation) + } + + test("rebuild removes the previous push-down even after it was pushed deeper") { + val t1a = AttributeReference("a", IntegerType, nullable = true)() + val t2b = AttributeReference("b", IntegerType, nullable = true)() + val t1 = LocalRelation(t1a) + val t2 = LocalRelation(t2b) + val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE) + + // The CTE definition ends in a renaming projection, like a view selecting aliased columns. + val pa = Alias(t1a, "pa")() + val pb = Alias(t2b, "pb")() + val project = Project(Seq(pa, pb), join) + + val cteId = 0L + val cteDef = CTERelationDef(project, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + // Pass 1: pushes the combined predicate on top of the definition's projection. + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + + // Consume the pushed filter with the real push-down rule that shares the fixedPoint + // batches with the rule under test: it moves below the renaming projection with the + // attributes rewritten to the projection's input. + val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1), + cteId, Literal(7)) + assert(theOnlyDef(consumed).child match { + case Project(_, Filter(_, _: Join)) => true + case _ => false + }, "test setup failed: push-down did not move the filter below the projection") + + // Pass 2: the rule must remove its previous push-down from wherever push-down left + // it; otherwise the rebuilt definition carries the old and new combined predicates + // as redundant stacked filters. + assertSingleEnrichedTopFilter(consumed, + "previous push-down was not removed before re-pushing") + } + + test("rebuild removes the previous push-down pushed into union branches") { + val l1 = AttributeReference("a", IntegerType, nullable = true)() + val l2 = AttributeReference("b", IntegerType, nullable = true)() + val r1 = AttributeReference("a", IntegerType, nullable = true)() + val r2 = AttributeReference("b", IntegerType, nullable = true)() + // The union output shares exprIds with the first branch; the def output is the union output. + val union = Union(Seq(LocalRelation(l1, l2), LocalRelation(r1, r2))) + + val cteId = 0L + val cteDef = CTERelationDef(union, cteId) + val plan = withTwoRefs(cteDef)( + out => And(EqualTo(out(0), Literal(5)), GreaterThan(out(1), Literal(0))), + out => And(EqualTo(out(0), Literal(7)), GreaterThan(out(1), Literal(0)))) + + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + + // The real push-down rule copies the pushed filter into every union branch, + // translating the union output attributes to each branch's output positionally. + val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1), + cteId, Literal(7)) + assert(theOnlyDef(consumed).child match { + case u: Union => u.children.forall(_.isInstanceOf[Filter]) + case _ => false + }, "test setup failed: push-down did not copy the filter into every union branch") + + assertSingleEnrichedTopFilter(consumed, + "previous push-down was not removed from the union branches") + } + + test("rebuild removes the previous push-down pushed below an aggregate") { + val a = AttributeReference("a", IntegerType, nullable = true)() + val b = AttributeReference("b", IntegerType, nullable = true)() + val pa = Alias(a, "pa")() + // Grouping-only aggregate (distinct), so the definition has a single output column. + val aggregate = Aggregate(Seq(pa), Seq(pa), LocalRelation(a, b)) + + val cteId = 0L + val cteDef = CTERelationDef(aggregate, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + + // The real push-down rule moves the filter below the aggregate, translating the + // grouping alias back to its child attribute. + val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1), + cteId, Literal(7)) + assert(theOnlyDef(consumed).child match { + case Aggregate(_, _, _: Filter, _) => true + case _ => false + }, "test setup failed: push-down did not move the filter below the aggregate") + + assertSingleEnrichedTopFilter(consumed, + "previous push-down was not removed below the aggregate") + } + + test("rebuild removes the previous push-down pushed below a window") { + val a = AttributeReference("a", IntegerType, nullable = true)() + val b = AttributeReference("b", IntegerType, nullable = true)() + val rn = Alias( + WindowExpression(RowNumber(), WindowSpecDefinition(Seq(a), Seq(SortOrder(a, Ascending)), + SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))), + "rn")() + val window = Window(Seq(rn), Seq(a), Seq(SortOrder(a, Ascending)), LocalRelation(a, b)) + + val cteId = 0L + val cteDef = CTERelationDef(window, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + + // The real push-down rule moves the filter below the window unchanged (the predicate + // references only the partition column, which is an input attribute). + val consumed = addIsNotNullToRef(PushPredicateThroughNonJoin.apply(afterPass1), + cteId, Literal(7)) + assert(theOnlyDef(consumed).child match { + case w: Window => w.child.isInstanceOf[Filter] + case _ => false + }, "test setup failed: push-down did not move the filter below the window") + + assertSingleEnrichedTopFilter(consumed, + "previous push-down was not removed below the window") + } + + test("rebuild preserves TreeNode tags on the nodes it rebuilds") { + val t1a = AttributeReference("a", IntegerType, nullable = true)() + val t2b = AttributeReference("b", IntegerType, nullable = true)() + val t1 = LocalRelation(t1a) + val t2 = LocalRelation(t2b) + val join = Join(t1, t2, Inner, Some(EqualTo(t1a, t2b)), JoinHint.NONE) + // Identity projection mimicking the analyzer-inserted projection above a natural or + // USING join, which carries the hidden join-key columns in Project.hiddenOutputTag. + val project = Project(Seq(t1a, t2b), join) + + val cteId = 0L + val cteDef = CTERelationDef(project, cteId) + val plan = withTwoRefs(cteDef)( + out => EqualTo(out(0), Literal(5)), + out => EqualTo(out(0), Literal(7))) + + val afterPass1 = PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan) + + // The real push-down rules move the pushed filter below the projection and then into + // the join's left branch (the predicate references only the left side of the inner + // join), so the next rebuild has to rebuild both the projection and the join. + val consumed = addIsNotNullToRef( + PushPredicateThroughJoin.apply(PushPredicateThroughNonJoin.apply(afterPass1)), + cteId, Literal(7)) + assert(theOnlyDef(consumed).child match { + case Project(_, Join(_: Filter, _, _, _, _)) => true + case _ => false + }, "test setup failed: push-down did not move the filter into the join branch") + + // Tag the nodes on the removal path in place, so the tags are present when the rule + // under test rebuilds them. (The intermediate push-down rules rebuild nodes with + // plain copies of their own; their tag handling is out of scope here.) + val taggedProject = theOnlyDef(consumed).child.asInstanceOf[Project] + val taggedJoin = taggedProject.child.asInstanceOf[Join] + taggedProject.setTagValue(Project.hiddenOutputTag, Seq(t2b)) + taggedJoin.setTagValue(testTag, "preserved") + + // Pass 2: removing the previous push-down from the join branch rebuilds both the + // join and the projection. The rebuild must carry their tags over. + val rebuilt = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(consumed)).child + val rebuiltProject = rebuilt.collect { case p: Project => p }.head + val rebuiltJoin = rebuilt.collect { case j: Join => j }.head + assert(rebuiltProject.getTagValue(Project.hiddenOutputTag).contains(Seq(t2b)), + s"rebuild dropped Project.hiddenOutputTag: $rebuilt") + assert(rebuiltJoin.getTagValue(testTag).contains("preserved"), + s"rebuild dropped tags on the join: $rebuilt") + } + + /** A tag with no consumer, used to verify that rebuilds preserve arbitrary tags. */ + private val testTag = TreeNodeTag[String]("cte_pushdown_test_tag") + + /** + * Applies the rule under test to `plan` (pass 2) and asserts that the rebuilt CTE + * definition carries exactly one filter: the freshly re-pushed combined predicate + * (recognizable by the IsNotNull enrichment). Any leftover copy of the previous + * push-down shows up as an additional filter and fails the test. + */ + private def assertSingleEnrichedTopFilter(plan: LogicalPlan, message: String): Unit = { + val cteDef = theOnlyDef(PushdownPredicatesAndPruneColumnsForCTEDef.apply(plan)) + val filters = cteDef.child.collect { case f: Filter => f } + assert(filters.length == 1, s"$message: ${cteDef.child}") + assert(filters.head.condition.find(_.isInstanceOf[IsNotNull]).isDefined, + s"expected the enriched combined predicate on top of the def: ${cteDef.child}") + } + + /** + * Builds one reference to the given CTE definition: `Project -> Filter -> CTERelationRef` + * with fresh output attribute instances (`CTERelationRef` is a `MultiInstanceRelation`, + * so every reference must own fresh exprIds). The project list covers all definition + * output columns so that column pruning never kicks in. `pred` receives the fresh + * output attributes and returns the reference-site predicate. + */ + private def mkRef(cteId: Long, defOutput: Seq[Attribute])( + pred: Seq[Attribute] => Expression): Project = { + val out = defOutput.map(_.newInstance()) + Project(out, Filter(pred(out), CTERelationRef(cteId, true, out, false))) + } + + /** + * Builds `WithCTE(Union(refs), Seq(cteDef))` with two reference sites carrying the given + * predicates, like a view referenced from multiple call sites. + */ + private def withTwoRefs(cteDef: CTERelationDef)( + pred1: Seq[Attribute] => Expression, + pred2: Seq[Attribute] => Expression): LogicalPlan = { + val refs = Seq(pred1, pred2).map(p => mkRef(cteDef.id, cteDef.output)(p)) + WithCTE(Union(refs), Seq(cteDef)) + } + + /** The single CTE definition in the given plan. */ + private def theOnlyDef(plan: LogicalPlan): CTERelationDef = { + plan.collect { case d: CTERelationDef => d }.head + } + + /** + * Simulates the reference-site enrichment `InferFiltersFromConstraints` performs between + * the rule's two passes: conjoins `IsNotNull` on the column the site predicate compares + * to `marker` (e.g. the second reference's `Literal(7)`), which re-arms the rule's + * rebuild on the next pass. The removal tests use this targeted mutation instead of the + * real rule because the real rule would also rewrite the pushed filter inside the + * definition, defeating the exact-match removal those tests isolate. + */ + private def addIsNotNullToRef( + plan: LogicalPlan, cteId: Long, marker: Literal): LogicalPlan = { + plan.transform { + case f @ Filter(cond, ref: CTERelationRef) if ref.cteId == cteId => + cond.collectFirst { case EqualTo(a: Attribute, m: Literal) if m == marker => a } match { + case Some(a) => Filter(And(cond, IsNotNull(a)), ref) + case _ => f + } + } + } + + private def hasFilterOn(plan: LogicalPlan, condition: Expression): Boolean = { + plan.collect { case f: Filter => f.condition }.exists(_.semanticEquals(condition)) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala index 08dd4011f04d6..b67eb88f0d68a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala @@ -18,11 +18,12 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Literal, Round} -import org.apache.spark.sql.catalyst.expressions.aggregate.CollectSet +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CaseWhen, Cast, Coalesce, Expression, GetMapValue, GetStructField, If, Literal, Lower, Round} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, CollectSet, Count, Sum} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LocalRelation, LogicalPlan} -import org.apache.spark.sql.types.{IntegerType, StringType} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, MapType, StringType, StructField, StructType} class RewriteDistinctAggregatesSuite extends PlanTest { val nullInt = Literal(null, IntegerType) @@ -125,4 +126,305 @@ class RewriteDistinctAggregatesSuite extends PlanTest { fail(s"Plan is not as expected:\n$rewrite") } } + + // --------------------------------------------------------------------------- + // COUNT(DISTINCT IF/CASE) canonicalization (SPARK-56898) + // --------------------------------------------------------------------------- + + val conditionalTestRelation = LocalRelation( + Symbol("a").int, Symbol("b").int, Symbol("c").int, Symbol("d").string) + + private def countDistinctIf(cond: Expression, base: Expression): Expression = { + Count(If(cond, base, Literal(null))).toAggregateExpression(isDistinct = true) + } + + private def countDistinctCaseWhen(cond: Expression, base: Expression): Expression = { + val caseWhen = CaseWhen( + Seq((cond, base)), + None) + Count(caseWhen).toAggregateExpression(isDistinct = true) + } + + private def countDistinctCaseWhenElseNull(cond: Expression, base: Expression): Expression = { + val caseWhen = CaseWhen( + Seq((cond, base)), + Some(Literal(null))) + Count(caseWhen).toAggregateExpression(isDistinct = true) + } + + private def collectAggregateExpressions(plan: LogicalPlan): Seq[AggregateExpression] = { + plan.collect { case a: Aggregate => a.aggregateExpressions } + .flatten + .flatMap(_.collect { case ae: AggregateExpression => ae }) + } + + /** + * Asserts that the optimized plan has exactly one Expand node with one projection, + * that the projection contains `baseColName` as a plain attribute, that it + * contains no expression of `removedWrapperType` (the IF/CaseWhen that was stripped), + * and that the outer aggregate has moved the condition into a FILTER clause. + */ + private def assertSingleDistinctGroupExpand( + optimized: LogicalPlan, + baseColName: String, + removedWrapperType: Class[_]): Unit = { + val expand = optimized.collectFirst { case e: Expand => e }.get + assert(expand.projections.size == 1, + s"expected 1 distinct group but got ${expand.projections.size}") + val baseAttr = conditionalTestRelation.output.find(_.name == baseColName).get + assert(expand.projections.head.exists(_.semanticEquals(baseAttr)), + s"expected base column $baseColName in Expand projection") + assert(!expand.projections.head.exists(e => removedWrapperType.isInstance(e)), + s"${removedWrapperType.getSimpleName} wrapper should have been removed " + + "from Expand projection") + assert(collectAggregateExpressions(optimized).exists(_.filter.isDefined), + "expected at least one AggregateExpression with a FILTER clause") + } + + test("conditional: disabled when config is false") { + withSQLConf(SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"), + countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + // The general RewriteDistinctAggregates still fires, but conditional + // canonicalization is disabled so the two conditional counts stay as 2 + // distinct groups instead of collapsing to 1. + val expands = optimized.collect { case e: Expand => e } + assert(expands.head.projections.size == 2, + "expected 2 distinct groups when conditional canonicalization is disabled") + } + } + + test("conditional: rewrite COUNT(DISTINCT IF(cond, col, NULL)) to COUNT(DISTINCT col) FILTER") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"), + countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + assertSingleDistinctGroupExpand(optimized, "c", classOf[If]) + } + + test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col END) to " + + "COUNT(DISTINCT col) FILTER") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctCaseWhen(Symbol("b") > 1, Symbol("c")).as("cnt1"), + countDistinctCaseWhen(Symbol("b") > 2, Symbol("c")).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + assertSingleDistinctGroupExpand(optimized, "c", classOf[CaseWhen]) + } + + test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col ELSE NULL END)") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctCaseWhenElseNull(Symbol("b") > 1, Symbol("c")).as("cnt1"), + countDistinctCaseWhenElseNull(Symbol("b") > 2, Symbol("c")).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + assertSingleDistinctGroupExpand(optimized, "c", classOf[CaseWhen]) + } + + test("conditional: multiple conditional distinct counts collapse to single distinct group") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"), + countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"), + countDistinctIf(Symbol("b") > 3, Symbol("c")).as("cnt3")) + .analyze + val optimized = RewriteDistinctAggregates(input) + // All three counts share the same base column c, collapsed to 1 distinct group. + assertSingleDistinctGroupExpand(optimized, "c", classOf[If]) + } + + test("conditional: single conditional distinct count is gated out by mayNeedtoRewrite") { + // A lone COUNT(DISTINCT IF(...)) must NOT be pushed onto the Expand path. + // The canonicalization runs inside rewrite(), which is only called when + // mayNeedtoRewrite returns true. A single conditional distinct count has no filter + // and forms only one distinct group, so mayNeedtoRewrite returns false, rewrite() + // is never called, and no Expand is produced. + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1")) + .analyze + val optimized = RewriteDistinctAggregates(input) + val expands = optimized.collect { case e: Expand => e } + assert(expands.isEmpty, "single conditional distinct count should not produce an Expand") + } + + test("conditional: do not rewrite IF with non-null else branch") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + Count(If(Symbol("b") > 1, Symbol("c"), Literal(0, IntegerType))) + .toAggregateExpression(isDistinct = true) + .as("cnt1"), + Count(If(Symbol("b") > 2, Symbol("c"), Literal(0, IntegerType))) + .toAggregateExpression(isDistinct = true) + .as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + // Plan is rewritten (2 distinct groups) but not canonicalized + checkRewrite(optimized) + val expands = optimized.collect { case e: Expand => e } + assert(expands.head.projections.size == 2, + "non-null else branch should not be collapsed to 1 distinct group") + } + + test("conditional: do not rewrite non-distinct COUNT") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + Count(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) + .toAggregateExpression(isDistinct = false) + .as("cnt1"), + countDistinct(Symbol("c")).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + // Still a single distinct group (cnt2), no canonicalization of the non-distinct agg + comparePlans(optimized, input) + } + + test("conditional: do not rewrite when FILTER already exists") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + Count(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) + .toAggregateExpression(isDistinct = true, filter = Some(Symbol("d") === "x")) + .as("cnt1"), + Count(If(Symbol("b") > 2, Symbol("c"), Literal(null, IntegerType))) + .toAggregateExpression(isDistinct = true, filter = Some(Symbol("d") === "y")) + .as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + checkRewrite(optimized) + val expands = optimized.collect { case e: Expand => e } + // 2 groups because existing FILTERs prevent canonicalization + assert(expands.head.projections.size == 2) + } + + test("conditional: do not rewrite multi-branch CASE WHEN") { + val caseWhen1 = new CaseWhen( + Seq( + (Symbol("b") > Literal(1), Symbol("c")), + (Symbol("b") > Literal(2), Symbol("a"))), + Some(Literal(null))) + val caseWhen2 = new CaseWhen( + Seq( + (Symbol("b") > Literal(3), Symbol("c")), + (Symbol("b") > Literal(4), Symbol("a"))), + Some(Literal(null))) + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + Count(caseWhen1).toAggregateExpression(isDistinct = true).as("cnt1"), + Count(caseWhen2).toAggregateExpression(isDistinct = true).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + checkRewrite(optimized) + val expands = optimized.collect { case e: Expand => e } + // 2 groups - multi-branch CASE was not canonicalized + assert(expands.head.projections.size == 2) + } + + test("conditional: do not rewrite SUM(DISTINCT IF(...))") { + val input = conditionalTestRelation + .groupBy(Symbol("a"))( + Sum(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) + .toAggregateExpression(isDistinct = true) + .as("sum1"), + Sum(If(Symbol("b") > 2, Symbol("c"), Literal(null, IntegerType))) + .toAggregateExpression(isDistinct = true) + .as("sum2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + checkRewrite(optimized) + val expands = optimized.collect { case e: Expand => e } + // 2 groups - SUM(DISTINCT IF) is not canonicalized + assert(expands.head.projections.size == 2) + } + + /** + * Asserts that two conditional distinct counts over `base` collapse into a single + * distinct group with a FILTER clause, i.e. the conditional canonicalization fired. + */ + private def assertCanonicalized(relation: LocalRelation, base: Expression): Unit = { + val input = relation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, base).as("cnt1"), + countDistinctIf(Symbol("b") > 2, base).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + val expand = optimized.collectFirst { case e: Expand => e }.get + assert(expand.projections.size == 1, + s"branch-safe base should collapse to 1 distinct group: $base") + // The Expand rewrite adds gid-routing filters to every distinct aggregate; the + // canonicalization additionally introduces filters referencing the user condition. + val hasUserConditionFilter = collectAggregateExpressions(optimized) + .flatMap(_.filter) + .exists(_.references.exists(_.name != "gid")) + assert(hasUserConditionFilter, + s"expected a user-condition FILTER clause after canonicalizing base: $base") + } + + /** + * Asserts that two conditional distinct counts over `base` are rewritten by the + * general Expand path but are NOT canonicalized, i.e. they stay 2 distinct groups + * and no FILTER clause is introduced. + */ + private def assertNotCanonicalized(relation: LocalRelation, base: Expression): Unit = { + val input = relation + .groupBy(Symbol("a"))( + countDistinctIf(Symbol("b") > 1, base).as("cnt1"), + countDistinctIf(Symbol("b") > 2, base).as("cnt2")) + .analyze + val optimized = RewriteDistinctAggregates(input) + checkRewrite(optimized) + val expand = optimized.collectFirst { case e: Expand => e }.get + assert(expand.projections.size == 2, + s"non-branch-safe base should stay 2 distinct groups: $base") + // The Expand rewrite adds gid-routing filters to every distinct aggregate; a + // user-condition filter introduced by the canonicalization would reference the + // condition's columns instead of only gid. + val hasUserConditionFilter = collectAggregateExpressions(optimized) + .flatMap(_.filter) + .exists(_.references.exists(_.name != "gid")) + assert(!hasUserConditionFilter, + s"non-branch-safe base should not produce a user-condition FILTER clause: $base") + } + + test("conditional: do not rewrite when base is not branch-safe") { + val unsafeBases = Seq( + // Division may throw (e.g. divide-by-zero under ANSI mode) on rows where the + // original IF branch is not taken. + Symbol("c") / Symbol("b"), + // Cast may throw under ANSI mode. + Cast(Symbol("c"), StringType), + // String functions are not in the branch-safe whitelist. + Lower(Symbol("d"))) + unsafeBases.foreach { base => + assertNotCanonicalized(conditionalTestRelation, base) + } + } + + test("conditional: rewrite branch-safe expression bases") { + // Predicate and null-handling bases on the flat test relation. + Seq( + Symbol("b") > Symbol("c"), + Coalesce(Seq(Symbol("b"), Literal(0, IntegerType)))) + .foreach { base => + assertCanonicalized(conditionalTestRelation, base) + } + + // Total accessor bases need struct/map attributes. + val structAttr = AttributeReference("s", StructType(Seq(StructField("f", IntegerType))))() + val mapAttr = AttributeReference("m", MapType(IntegerType, StringType))() + val richRelation = LocalRelation(Symbol("a").int, Symbol("b").int, structAttr, mapAttr) + Seq( + GetStructField(structAttr, 0), + GetMapValue(mapAttr, Literal(1, IntegerType))) + .foreach { base => + assertCanonicalized(richRelation, base) + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSubquerySuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSubquerySuite.scala index dca1d503e3fdd..d98c043960823 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSubquerySuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteSubquerySuite.scala @@ -20,9 +20,9 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Cast, Exists, IsNull, ListQuery, Literal, Not} +import org.apache.spark.sql.catalyst.expressions.{Cast, EqualTo, Exists, InSubquery, IsNull, ListQuery, Literal, Not, Or} import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, LeftSemi, PlanTest} -import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, Join, LocalRelation, LogicalPlan} import org.apache.spark.sql.catalyst.rules.RuleExecutor import org.apache.spark.sql.types.LongType @@ -72,6 +72,33 @@ class RewriteSubquerySuite extends PlanTest { comparePlans(optimized, correctAnswer) } + test("SPARK-58365: NOT-IN nested in OR deduplicates conflicting attrs on the join right side") { + // Nesting the NOT-IN under an OR routes it through rewriteExistentialExprWithAttrs, + // whose join condition must be built from the deduplicated subquery output. Build the + // colliding-attribute plan directly: the analyzer's DeduplicateRelations would renew + // the shared exprId before the optimizer runs, hiding the rule's behavior. + val a = $"a".int + val b = $"b".int + val relation = LocalRelation(a, b) + // The subquery reuses `a`, so its output conflicts with the outer plan by exprId. + val subquery = relation.select(a) + val query = Filter( + Or(EqualTo(b, Literal(1)), Not(InSubquery(Seq(a), ListQuery(subquery, numCols = 1)))), + relation) + + val optimized = Optimize.execute(query) + + val join = optimized.collectFirst { case j: Join => j }.get + // dedupSubqueryOnSelfJoin aliases the conflicting attr to a fresh exprId, so the join + // is duplicate-resolved regardless. The real check: the condition must reference that + // deduplicated right-side attr. The pre-fix code zips against the stale pre-dedup + // output, so the condition only touches the outer side and the right side dangles. + assert(join.duplicateResolved) + assert(join.condition.get.references.intersect(join.right.outputSet).nonEmpty, + s"join condition ${join.condition.get} must reference the right child output " + + s"${join.right.outputSet}") + } + test("SPARK-34598: Filters without subquery must not be modified by RewritePredicateSubquery") { val relation = LocalRelation($"a".int, $"b".int, $"c".int, $"d".int) val query = relation.where(($"a" === 1 || $"b" === 2) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala index 71acbdfdd2fca..9d4cf9dff953c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/complexTypesSuite.scala @@ -409,6 +409,24 @@ class ComplexTypesSuite extends PlanTest with ExpressionEvalHelper { checkRule(mapRel, mapExpected) } + test("SPARK-58431: don't simplify ANSI out-of-bounds array access to null") { + val ansiQuery = relation + .select( + GetArrayItem( + CreateArray(Seq($"nullable_id", $"nullable_id" + 1L)), + 5, + failOnError = true) as "a1") + val nonAnsiQuery = relation + .select( + GetArrayItem( + CreateArray(Seq($"nullable_id", $"nullable_id" + 1L)), + 5, + failOnError = false) as "a1") + + checkRule(ansiQuery, ansiQuery) + checkRule(nonAnsiQuery, relation.select(Literal.create(null, LongType) as "a1")) + } + test("SPARK-23500: Ensure that aggregation expressions are not simplified") { // Make sure that aggregation exprs are correctly ignored. Maps can't be used in // grouping exprs so aren't tested here. diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/ExpressionParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/ExpressionParserSuite.scala index 3812dd2878378..d46c85d8ed093 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/ExpressionParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/ExpressionParserSuite.scala @@ -307,6 +307,90 @@ class ExpressionParserSuite extends AnalysisTest { assertEqual("-+~~a", -( +(~(~$"a")))) } + test("JSON_VALUE expressions") { + import org.apache.spark.sql.catalyst.expressions.JsonValueBehavior + // Bare form: default STRING RETURNING, NULL ON EMPTY / NULL ON ERROR. + assertEqual( + "json_value(a, '$.b')", + JsonValue($"a", "$.b", StringType, JsonValueBehavior.Null, JsonValueBehavior.Null, + None, None)) + // RETURNING. + assertEqual( + "json_value(a, '$.b' RETURNING INT)", + JsonValue($"a", "$.b", IntegerType, JsonValueBehavior.Null, JsonValueBehavior.Null, + None, None)) + // ERROR ON EMPTY / ERROR ON ERROR. + assertEqual( + "json_value(a, '$.b' ERROR ON EMPTY ERROR ON ERROR)", + JsonValue($"a", "$.b", StringType, JsonValueBehavior.Error, JsonValueBehavior.Error, + None, None)) + // DEFAULT ON EMPTY / DEFAULT ON ERROR carry expressions. + assertEqual( + "json_value(a, '$.b' DEFAULT 'x' ON EMPTY DEFAULT 'y' ON ERROR)", + JsonValue($"a", "$.b", StringType, JsonValueBehavior.Default, JsonValueBehavior.Default, + Some(Literal("x")), Some(Literal("y")))) + } + + test("JSON_EXISTS expressions") { + import org.apache.spark.sql.catalyst.expressions.JsonExistsBehavior + // Bare form defaults to FALSE ON ERROR. + assertEqual("json_exists(a, '$.b')", JsonExists($"a", "$.b", JsonExistsBehavior.False)) + assertEqual( + "json_exists(a, '$.b' TRUE ON ERROR)", JsonExists($"a", "$.b", JsonExistsBehavior.True)) + assertEqual( + "json_exists(a, '$.b' FALSE ON ERROR)", JsonExists($"a", "$.b", JsonExistsBehavior.False)) + assertEqual( + "json_exists(a, '$.b' UNKNOWN ON ERROR)", + JsonExists($"a", "$.b", JsonExistsBehavior.Unknown)) + assertEqual( + "json_exists(a, '$.b' ERROR ON ERROR)", JsonExists($"a", "$.b", JsonExistsBehavior.Error)) + } + + test("JSON_QUERY expressions") { + import org.apache.spark.sql.catalyst.expressions.{JsonQueryBehavior, JsonQueryQuotes, + JsonQueryWrapper} + // Bare form: default STRING RETURNING, WITHOUT wrapper, KEEP quotes, NULL ON EMPTY / ON ERROR. + assertEqual( + "json_query(a, '$.b')", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Without, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + // WITH ARRAY WRAPPER defaults to UNCONDITIONAL; the ARRAY word is optional. + assertEqual( + "json_query(a, '$.b' WITH ARRAY WRAPPER)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Unconditional, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + assertEqual( + "json_query(a, '$.b' WITH CONDITIONAL WRAPPER)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Conditional, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + // WITHOUT ARRAY WRAPPER with OMIT QUOTES. + assertEqual( + "json_query(a, '$.b' WITHOUT ARRAY WRAPPER OMIT QUOTES)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Without, JsonQueryQuotes.Omit, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + // EMPTY ARRAY ON EMPTY / EMPTY OBJECT ON ERROR, plus RETURNING STRING. + assertEqual( + "json_query(a, '$.b' RETURNING STRING EMPTY ARRAY ON EMPTY EMPTY OBJECT ON ERROR)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Without, JsonQueryQuotes.Keep, + JsonQueryBehavior.EmptyArray, JsonQueryBehavior.EmptyObject)) + // The ARRAY word is optional in every wrapper spelling, and WITH alone means UNCONDITIONAL. + Seq("WITH WRAPPER", "WITH UNCONDITIONAL WRAPPER").foreach { spelling => + assertEqual( + s"json_query(a, '$$.b' $spelling)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Unconditional, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + } + assertEqual( + "json_query(a, '$.b' WITHOUT WRAPPER)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Without, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + // Explicit NULL ON EMPTY / NULL ON ERROR (same as the omitted default) and KEEP QUOTES. + assertEqual( + "json_query(a, '$.b' KEEP QUOTES NULL ON EMPTY NULL ON ERROR)", + JsonQuery($"a", "$.b", StringType, JsonQueryWrapper.Without, JsonQueryQuotes.Keep, + JsonQueryBehavior.Null, JsonQueryBehavior.Null)) + } + test("cast expressions") { // Note that DataType parsing is tested elsewhere. assertEqual("cast(a as int)", $"a".cast(IntegerType)) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala index 2d8e60aee22ce..a0fd03ff00c0a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/PlanParserSuite.scala @@ -673,6 +673,85 @@ class PlanParserSuite extends AnalysisTest { stop = 115)) } + test("unnest in FROM clause") { + def unnest( + exprs: Seq[Expression], + withOrdinality: Boolean = false): LogicalPlan = + Generate( + Unnest(exprs, withOrdinality), + unrequiredChildIndex = Nil, + outer = false, + qualifier = None, + generatorOutput = Nil, + child = OneRowRelation()) + + // Single array. + assertEqual( + "select * from unnest(array(1, 2, 3))", + unnest(Seq(UnresolvedFunction("array", Seq(Literal(1), Literal(2), Literal(3)), + isDistinct = false))).select(star())) + + // Multiple arrays. + assertEqual( + "select * from unnest(a, b)", + unnest(Seq(UnresolvedAttribute("a"), UnresolvedAttribute("b"))).select(star())) + + // WITH ORDINALITY. + assertEqual( + "select * from unnest(a) with ordinality", + unnest(Seq(UnresolvedAttribute("a")), withOrdinality = true).select(star())) + + // Table alias only. + assertEqual( + "select * from unnest(a) t", + unnest(Seq(UnresolvedAttribute("a"))).as("t").select(star())) + + // Table alias with column aliases. + assertEqual( + "select * from unnest(a, b) t(x, y)", + SubqueryAlias( + "t", + UnresolvedSubqueryColumnAliases( + Seq("x", "y"), + unnest(Seq(UnresolvedAttribute("a"), UnresolvedAttribute("b"))))).select(star())) + + // Correlated via LATERAL, with WITH ORDINALITY and column aliases. + assertEqual( + "select * from t, lateral unnest(t.arr) with ordinality u(v, o)", + table("t").lateralJoin( + SubqueryAlias( + "u", + UnresolvedSubqueryColumnAliases( + Seq("v", "o"), + unnest(Seq(UnresolvedAttribute(Seq("t", "arr"))), withOrdinality = true)))) + .select(star())) + + // With no arguments the dedicated UNNEST relation rule does not match; the statement instead + // parses as a generic table-valued function call named `unnest`, which is not registered and + // therefore fails later during analysis rather than at parse time. + assertEqual( + "select * from unnest()", + UnresolvedTableValuedFunction("unnest", Nil).select(star())) + + // `UNNEST` and `ORDINALITY` are non-reserved keywords, so they remain usable as regular + // table and column identifiers for backwards compatibility. + assertEqual( + "select ordinality from unnest", + table("unnest").select($"ordinality")) + assertEqual( + "select unnest.ordinality from unnest", + table("unnest").select($"unnest.ordinality")) + + // Quoting the name bypasses the UNNEST relation syntax, so a table-valued function named + // `unnest` can still be invoked. This is the escape hatch for the non-reserved keyword. + assertEqual( + "select * from `unnest`(array(1, 2))", + UnresolvedTableValuedFunction( + "unnest", + Seq(UnresolvedFunction("array", Seq(Literal(1), Literal(2)), isDistinct = false))) + .select(star())) + } + test("joins") { // Test single joins. val testUnconditionalJoin = (sql: String, jt: JoinType) => { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/SampleSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/SampleSuite.scala new file mode 100644 index 0000000000000..0f3efb0c76983 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/SampleSuite.scala @@ -0,0 +1,50 @@ +/* + * 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.catalyst.plans.logical + +import org.apache.spark.SparkFunSuite + +class SampleSuite extends SparkFunSuite { + + test("resolveSeed returns a user-specified seed unchanged") { + assert(Sample.resolveSeed(Some(42L)) === 42L) + assert(Sample.resolveSeed(Some(0L)) === 0L) + assert(Sample.resolveSeed(Some(Long.MaxValue)) === Long.MaxValue) + // Only generated seeds are constrained to be non-negative. The Dataset API accepts a + // negative seed even though the SQL REPEATABLE grammar does not, so it must pass through. + assert(Sample.resolveSeed(Some(-5L)) === -5L) + assert(Sample.resolveSeed(Some(Long.MinValue)) === Long.MinValue) + } + + test("resolveSeed generates non-negative seeds") { + // A pushed-down sample renders its seed into SQL as `REPEATABLE (<seed>)`, and the seed + // in that grammar does not accept a sign. + for (_ <- 0 until 10000) { + assert(Sample.resolveSeed(None) >= 0L) + } + } + + test("resolveSeed draws from a wide range of values") { + // Guards against SPARK-56573, where the generated seed was limited to 1000 distinct + // values. Drawing from 2^63 makes 1000 collisions in 10000 draws effectively impossible. + val seeds = Seq.fill(10000)(Sample.resolveSeed(None)).toSet + assert(seeds.size > 9000, s"expected nearly all seeds to be distinct, got ${seeds.size}") + // The old implementation could never exceed 999. + assert(seeds.exists(_ > 1000L)) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala new file mode 100644 index 0000000000000..3b20eae1a75ed --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/logical/V2CommandTreePatternSuite.scala @@ -0,0 +1,97 @@ +/* + * 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.catalyst.plans.logical + +import java.util + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.ProjectingInternalRow +import org.apache.spark.sql.catalyst.dsl.expressions._ +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.catalyst.trees.TreePattern +import org.apache.spark.sql.catalyst.util.{ReplaceDataProjections, WriteDeltaProjections} +import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Pins the tree-pattern identity contract for the DSv2 row-level command nodes: each node carries + * both the shared `COMMAND` bit and its own identity bit, so rules can prune on either. + */ +class V2CommandTreePatternSuite extends SparkFunSuite { + + private val target = LocalRelation($"a".int, $"b".int) + private val source = LocalRelation($"c".int, $"d".int) + + // A minimal `NamedRelation` for the row-level write nodes, whose `nodePatternsInternal()` + // returns a compile-time constant, so no resolution is needed to pin the identity bit. + private val v2Relation: DataSourceV2Relation = { + val table = new Table { + override def name(): String = "t" + override def columns(): Array[Column] = Array(Column.create("a", IntegerType)) + override def capabilities(): util.Set[TableCapability] = util.Set.of[TableCapability]() + } + DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty()) + } + + private val emptyRow = ProjectingInternalRow(new StructType(), IndexedSeq.empty) + + test("DeleteFromTable declares COMMAND and DELETE_FROM_TABLE") { + val plan = DeleteFromTable(target, Literal.TrueLiteral) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.DELETE_FROM_TABLE)) + } + + test("UpdateTable declares COMMAND and UPDATE_TABLE") { + val plan = UpdateTable(target, Seq.empty, Some(Literal.TrueLiteral)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.UPDATE_TABLE)) + } + + test("MergeIntoTable declares COMMAND and MERGE_INTO_TABLE") { + val plan = MergeIntoTable( + target, + source, + mergeCondition = $"a" === $"c", + matchedActions = Seq(DeleteAction(None)), + notMatchedActions = Seq(InsertAction(None, + Seq(Assignment($"a", $"c"), Assignment($"b", $"d")))), + notMatchedBySourceActions = Seq(DeleteAction(None)), + withSchemaEvolution = false) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.MERGE_INTO_TABLE)) + } + + test("ReplaceData declares COMMAND and REPLACE_DATA") { + val plan = ReplaceData( + v2Relation, + Literal.TrueLiteral, + source, + v2Relation, + ReplaceDataProjections(emptyRow, None)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.REPLACE_DATA)) + } + + test("WriteDelta declares COMMAND and WRITE_DELTA") { + val plan = WriteDelta( + v2Relation, + Literal.TrueLiteral, + source, + v2Relation, + WriteDeltaProjections(None, emptyRow, None)) + assert(plan.containsAllPatterns(TreePattern.COMMAND, TreePattern.WRITE_DELTA)) + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/BasicStatsEstimationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/BasicStatsEstimationSuite.scala index f07c19120438d..afb4618fea79c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/BasicStatsEstimationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/statsEstimation/BasicStatsEstimationSuite.scala @@ -249,14 +249,14 @@ test("range with invalid long value") { } test("sample estimation") { - val sample = Sample(0.0, 0.5, withReplacement = false, (math.random() * 1000).toLong, plan) + val sample = Sample(0.0, 0.5, withReplacement = false, Sample.resolveSeed(None), plan) checkStats(sample, Statistics(sizeInBytes = 60, rowCount = Some(5))) // Child doesn't have rowCount in stats val childStats = Statistics(sizeInBytes = 120) val childPlan = DummyLogicalPlan(childStats, childStats) val sample2 = - Sample(0.0, 0.11, withReplacement = false, (math.random() * 1000).toLong, childPlan) + Sample(0.0, 0.11, withReplacement = false, Sample.resolveSeed(None), childPlan) checkStats(sample2, Statistics(sizeInBytes = 14)) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala index 196d32e1a4f9d..c46daf8035989 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/DateTimeUtilsSuite.scala @@ -1199,6 +1199,49 @@ class DateTimeUtilsSuite extends SparkFunSuite with Matchers with SQLHelper { } } + test("SPARK-57825: timestamp nanos add year-month interval preserves nanosWithinMicro") { + def nanos(epochMicros: Long, nanosWithinMicro: Int): TimestampNanosVal = + TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort) + + // The epoch-micros part follows the micro `timestampAddMonths` (including the Jan-31 -> Feb-29 + // day clamp in a leap year) while the sub-microsecond remainder is carried through unchanged. + assert(timestampNanosAddMonths( + nanos(date(2020, 1, 31, 12, 0, 0, 123000, LA), 789), 1, LA) === + nanos(date(2020, 2, 29, 12, 0, 0, 123000, LA), 789)) + + outstandingZoneIds.foreach { zid => + // The sub-microsecond remainder is preserved for the boundary values 0, 1 and 999. + Seq(0, 1, 999).foreach { rem => + // Zero interval is a no-op on both the epoch-micros and the remainder. + assert(timestampNanosAddMonths( + nanos(date(2021, 3, 18, 19, 44, 1, 123456, zid), rem), 0, zid) === + nanos(date(2021, 3, 18, 19, 44, 1, 123456, zid), rem)) + // Adding whole years/months shifts only the month field, never the fraction. + assert(timestampNanosAddMonths( + nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), rem), 14, zid) === + nanos(date(2021, 3, 2, 3, 4, 5, 123456, zid), rem)) + // Subtracting months is symmetric. + assert(timestampNanosAddMonths( + nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), rem), -1, zid) === + nanos(date(2019, 12, 2, 3, 4, 5, 123456, zid), rem)) + // Pre-epoch (negative epochMicros) value. + assert(timestampNanosAddMonths( + nanos(date(1960, 1, 2, 3, 4, 5, 123456, zid), rem), 1, zid) === + nanos(date(1960, 2, 2, 3, 4, 5, 123456, zid), rem)) + } + } + + // Consistency with the micro helper: epochMicros matches `timestampAddMonths` exactly and the + // remainder is independent of the interval amount. + outstandingZoneIds.foreach { zid => + val start = nanos(date(2020, 1, 2, 3, 4, 5, 123456, zid), 789) + val months = 15 + val result = timestampNanosAddMonths(start, months, zid) + assert(result.epochMicros === timestampAddMonths(start.epochMicros, months, zid)) + assert(result.nanosWithinMicro === start.nanosWithinMicro) + } + } + test("SPARK-57159: timestampNanosToEpochNanos packs into int64 epoch-nanoseconds") { def nanos(epochMicros: Long, nanosWithinMicro: Int): TimestampNanosVal = TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/QuotingUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/QuotingUtilsSuite.scala index 259284d7c4055..c7d1280220a3c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/QuotingUtilsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/QuotingUtilsSuite.scala @@ -40,4 +40,17 @@ class QuotingUtilsSuite extends SparkFunSuite { assert(quoteIfNeeded("0d") == "`0d`") assert(quoteIfNeeded("") === "``") } + + test("quoteIdentifier escapes back-ticks and always wraps") { + assert(QuotingUtils.quoteIdentifier("a") === "`a`") + assert(QuotingUtils.quoteIdentifier("") === "``") + assert(QuotingUtils.quoteIdentifier("a`b") === "`a``b`") + assert(QuotingUtils.quoteIdentifier("`") === "````") + } + + test("escapeSingleQuotedString escapes single quotes only") { + assert(QuotingUtils.escapeSingleQuotedString("abc") === "abc") + assert(QuotingUtils.escapeSingleQuotedString("a'b") === "a\\'b") + assert(QuotingUtils.escapeSingleQuotedString("''") === "\\'\\'") + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala index eda401ceb6bdf..e128adebdff81 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala @@ -17,15 +17,41 @@ package org.apache.spark.sql.connector.catalog -import org.mockito.Mockito.{mock, when} +import org.mockito.ArgumentCaptor +import org.mockito.ArgumentMatchers.{any, eq => mockEq} +import org.mockito.Mockito.{mock, verify, when} -import org.apache.spark.SparkFunSuite +import org.apache.spark.{SparkFunSuite, SparkIllegalArgumentException} +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.catalyst.analysis.{ + AsOfTimestamp, AsOfVersion, TimeTravelSpec, UnresolvedRelation} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation -import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap class CatalogV2UtilSuite extends SparkFunSuite { - test("Load relation should encode the identifiers for V2Relations") { + + private def catalogWithStateOptions(keys: java.util.Set[String]): TableCatalog = { + val catalog = mock(classOf[TableCatalog]) + when(catalog.tableStateOptionKeys()).thenReturn(keys) + catalog + } + + // CatalogV2Util.getTable routes through the options-aware TableCatalog.loadTable, whose default + // implementation dispatches to the existing overloads. Stub only that method to run the real + // default so the dispatch is exercised; the leaf overloads stay as plain mock methods (returning + // null) that we then `verify`. + private def mockCatalogWithRealDispatch(): TableCatalog = { val testCatalog = mock(classOf[TableCatalog]) + when(testCatalog.tableStateOptionKeys()).thenCallRealMethod() + when(testCatalog.loadTable( + any[Identifier], any[TableContext], any[CaseInsensitiveStringMap])).thenCallRealMethod() + testCatalog + } + + test("Load relation should encode the identifiers for V2Relations") { + val testCatalog = mockCatalogWithRealDispatch() val ident = mock(classOf[Identifier]) val table = mock(classOf[Table]) when(table.columns()).thenReturn(Array(Column.create("i", IntegerType))) @@ -37,4 +63,223 @@ class CatalogV2UtilSuite extends SparkFunSuite { assert(v2Relation.catalog.exists(_ == testCatalog)) assert(v2Relation.identifier.exists(_ == ident)) } + + private def getTableAndVerifyDispatch( + timeTravelSpec: Option[TimeTravelSpec], + writePrivilegesString: Option[String])( + verifyOverload: TableCatalog => Unit): Unit = { + val testCatalog = mockCatalogWithRealDispatch() + val ident = mock(classOf[Identifier]) + CatalogV2Util.getTable(testCatalog, ident, timeTravelSpec, writePrivilegesString) + verifyOverload(testCatalog) + } + + test("getTable dispatches to loadTable(ident) with no time travel and no write privileges") { + getTableAndVerifyDispatch(None, None) { c => verify(c).loadTable(any[Identifier]) } + } + + test("getTable dispatches to loadTable(ident, writePrivileges) with write privileges") { + getTableAndVerifyDispatch(None, Some("INSERT,DELETE")) { c => + verify(c).loadTable( + any[Identifier], + mockEq(java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE))) + } + } + + test("getTable dispatches to loadTable(ident, version) for version time travel") { + getTableAndVerifyDispatch(Some(AsOfVersion("v1")), None) { c => + verify(c).loadTable(any[Identifier], mockEq("v1")) + } + } + + test("getTable dispatches to loadTable(ident, timestamp) for timestamp time travel") { + getTableAndVerifyDispatch(Some(AsOfTimestamp(123L)), None) { c => + verify(c).loadTable(any[Identifier], mockEq(123L)) + } + } + + test("getTable rejects combining time travel and write privileges") { + val testCatalog = mockCatalogWithRealDispatch() + val ident = mock(classOf[Identifier]) + val e = intercept[SparkIllegalArgumentException] { + CatalogV2Util.getTable(testCatalog, ident, Some(AsOfVersion("v1")), Some("INSERT")) + } + assert(e.getMessage.contains("Cannot set both time travel and write privileges")) + } + + test("loadTableForV2Write forwards write privileges and only table-state options") { + val testCatalog = mock(classOf[TableCatalog]) + when(testCatalog.tableStateOptionKeys()).thenReturn(java.util.Set.of("state")) + val ident = mock(classOf[Identifier]) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("state", "branch", "custom", "value")) + val contextCaptor = ArgumentCaptor.forClass(classOf[TableContext]) + + CatalogV2Util.loadTableForV2Write( + testCatalog, ident, Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE), options) + + val expectedStateOptions = + new CaseInsensitiveStringMap(java.util.Map.of("state", "branch")) + verify(testCatalog).loadTable( + mockEq(ident), contextCaptor.capture(), mockEq(expectedStateOptions)) + assert(contextCaptor.getValue.timeTravel().isEmpty) + assert(contextCaptor.getValue.writePrivileges() === + java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + } + + test("loadTableForV2Write rejects configured time travel options") { + val testCatalog = mock(classOf[TableCatalog]) + when(testCatalog.name()).thenReturn("testcat") + val ident = Identifier.of(Array("ns"), "table") + val conf = new SQLConf + conf.setConf(SQLConf.TIME_TRAVEL_VERSION_KEY, "customVersion") + conf.setConf(SQLConf.TIME_TRAVEL_TIMESTAMP_KEY, "customTimestamp") + + SQLConf.withExistingConf(conf) { + Seq("customVersion", "customTimestamp").foreach { key => + val options = new CaseInsensitiveStringMap(java.util.Map.of(key, "value")) + val e = intercept[AnalysisException] { + CatalogV2Util.loadTableForV2Write( + testCatalog, ident, Set(TableWritePrivilege.INSERT), options) + } + assert(e.getCondition === "UNSUPPORTED_FEATURE.TIME_TRAVEL") + } + } + } + + test("UnresolvedRelation preserves option key case while updating write privileges") { + val options = new CaseInsensitiveStringMap(java.util.Map.of( + "targetLoadOption", "loadValue", + "targetWriteOption", "writeValue")) + val relation = UnresolvedRelation(Seq("catalog", "table"), options) + + val withPrivileges = relation.requireWritePrivileges(Set(TableWritePrivilege.INSERT)) + assert(withPrivileges.options.asCaseSensitiveMap().containsKey("targetLoadOption")) + assert(withPrivileges.options.asCaseSensitiveMap().containsKey("targetWriteOption")) + assert(!withPrivileges.options.asCaseSensitiveMap().containsKey("targetloadoption")) + assert(withPrivileges.options.get(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) === "INSERT") + + val cleared = withPrivileges.clearWritePrivileges + assert(cleared.options.asCaseSensitiveMap() === options.asCaseSensitiveMap()) + } + + test("TableContext normalizes null time travel and null write privileges to empty") { + val context = new TableContext(null, null) + assert(context.timeTravel().isEmpty) + assert(context.writePrivileges().isEmpty) + } + + test("TableContext equals / hashCode / toString") { + val emptyPrivileges = java.util.Set.of[TableWritePrivilege]() + val a = new TableContext(new TimeTravel.AsOfVersion("v1"), emptyPrivileges) + val b = new TableContext(new TimeTravel.AsOfVersion("v1"), emptyPrivileges) + val c = new TableContext(new TimeTravel.AsOfTimestamp(1L), emptyPrivileges) + assert(a == b) + assert(a.hashCode() == b.hashCode()) + assert(a != c) + assert(a.toString.contains("timeTravel")) + assert(a.toString.contains("writePrivileges")) + } + + test("getTable forwards only declared table-state options") { + val catalog = mock(classOf[TableCatalog]) + when(catalog.tableStateOptionKeys()).thenReturn(java.util.Set.of("snapshot")) + val ident = mock(classOf[Identifier]) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "5")) + + CatalogV2Util.getTable(catalog, ident, options = options) + + val expected = new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "s1")) + verify(catalog).loadTable( + mockEq(ident), + any[TableContext], + mockEq(expected)) + } + + test("getTable forwards no options when a catalog declares no table-state options") { + val catalog = mock(classOf[TableCatalog]) + when(catalog.tableStateOptionKeys()).thenCallRealMethod() + val ident = mock(classOf[Identifier]) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "5")) + + CatalogV2Util.getTable(catalog, ident, options = options) + + verify(catalog).loadTable( + mockEq(ident), + any[TableContext], + mockEq(CaseInsensitiveStringMap.empty())) + } + + test("extractTableStateOptions projects declared keys case-insensitively") { + val catalog = catalogWithStateOptions(java.util.Set.of("BrAnCh", "tag")) + val options = new CaseInsensitiveStringMap(java.util.Map.of( + "branch", "Main", + "TAG", "Release", + "split-size", "5")) + + val stateOptions = CatalogV2Util.extractTableStateOptions(catalog, options) + + assert(stateOptions.size() == 2) + assert(stateOptions.get("BRANCH") == "Main") + assert(stateOptions.get("tag") == "Release") + assert(!stateOptions.containsKey("split-size")) + } + + test("extractTableStateOptions compares option keys case-insensitively") { + val catalog = catalogWithStateOptions(java.util.Set.of("SnApShOt")) + val lowerCaseKey = CatalogV2Util.extractTableStateOptions( + catalog, + new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "main"))) + val upperCaseKey = CatalogV2Util.extractTableStateOptions( + catalog, + new CaseInsensitiveStringMap(java.util.Map.of("SNAPSHOT", "main"))) + + assert(lowerCaseKey == upperCaseKey) + } + + test("extractTableStateOptions compares option values case-sensitively") { + val catalog = catalogWithStateOptions(java.util.Set.of("snapshot")) + val lowerCaseValue = CatalogV2Util.extractTableStateOptions( + catalog, + new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "main"))) + val upperCaseValue = CatalogV2Util.extractTableStateOptions( + catalog, + new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "MAIN"))) + + assert(lowerCaseValue != upperCaseValue) + } + + test("extractTableStateOptions returns no options by default") { + val catalog = mock(classOf[TableCatalog]) + when(catalog.tableStateOptionKeys()).thenCallRealMethod() + val options = new CaseInsensitiveStringMap( + java.util.Map.of("branch", "Main", "split-size", "5")) + + val stateOptions = CatalogV2Util.extractTableStateOptions(catalog, options) + + assert(stateOptions.isEmpty) + } + + test("viewInfoBuilderFrom preserves the dependency list") { + val dependencies = DependencyList.of(Array(Dependency.table(Array("cat", "ns", "events")))) + val existing = viewWithDependencies(Some(dependencies)) + val rebuilt = CatalogV2Util.viewInfoBuilderFrom(existing).build() + assert(rebuilt.viewDependencies() === dependencies) + } + + test("viewInfoBuilderFrom leaves an absent dependency list absent") { + val existing = viewWithDependencies(None) + val rebuilt = CatalogV2Util.viewInfoBuilderFrom(existing).build() + assert(rebuilt.viewDependencies() === null) + } + + private def viewWithDependencies(dependencies: Option[DependencyList]): View = { + val builder = new View.Builder() + .withSchema(new StructType().add("i", IntegerType)) + .withQueryText("SELECT i FROM cat.ns.events") + dependencies.foreach(builder.withViewDependencies) + builder.build() + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala index fba80eb3d4cbe..d9540b8fff854 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala @@ -25,11 +25,11 @@ import java.util.OptionalLong import java.util.concurrent.atomic.AtomicLong import scala.collection.mutable -import scala.collection.mutable.ListBuffer +import scala.collection.mutable.{ArrayBuffer, ListBuffer} import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Cast, EvalMode, GenericInternalRow, JoinedRow, Literal, MetadataStructFieldWithLogicalName} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, CaseInsensitiveMap, CharVarcharUtils, DateTimeUtils, GenericArrayData, MapData, ResolveDefaultColumns} import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} @@ -43,7 +43,7 @@ import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset} import org.apache.spark.sql.connector.write._ import org.apache.spark.sql.connector.write.streaming.{StreamingDataWriterFactory, StreamingWrite} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.{ColumnImpl, SupportsStreamingUpdateAsAppend} +import org.apache.spark.sql.internal.connector.{ColumnImpl, SupportsRuntimeCatalystFiltering, SupportsStreamingUpdateAsAppend} import org.apache.spark.sql.sources._ import org.apache.spark.sql.types._ import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -521,23 +521,35 @@ abstract class InMemoryBaseTable( private var _pushedFilters: Array[Filter] = Array.empty override def build: Scan = { - val scan = if (InMemoryBaseTable.this.ordering.nonEmpty) { - new InMemoryBatchScanWithOrdering( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } else { - InMemoryBatchScan( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } - if (evaluableFilters.nonEmpty) { - scan.filter(evaluableFilters) + val scan = createScan( + data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, options) + scan match { + case s: InMemoryBatchScan => + if (evaluableFilters.nonEmpty) { + s.filter(evaluableFilters) + } + s.pushedFilters = _pushedFilters + case _ => } - scan.pushedFilters = _pushedFilters recordScanEvent(_pushedFilters) scan } + /** + * Creates the batch scan for [[build]]. + */ + protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (InMemoryBaseTable.this.ordering.nonEmpty) { + new InMemoryBatchScanWithOrdering(partitions, readSchema, tableSchema, options) + } else { + InMemoryBatchScan(partitions, readSchema, tableSchema, options) + } + } + override def pruneColumns(requiredSchema: StructType): Unit = { // The required schema could contain conflict-renamed metadata columns, so we need to match // them by their logical (original) names, not their current names. @@ -691,6 +703,116 @@ abstract class InMemoryBaseTable( new InMemoryMicroBatchStream(readSchema, tableSchema) } + /** + * Reference implementation of [[SupportsRuntimeCatalystFiltering.filter]] for the in-memory + * fixtures: records what was pushed, and for expressions referencing only partition columns + * binds them against the partition key and drops partitions that do not match. Binding and + * interpreting rather than pattern matching a fixed set of operators is what lets the fixture + * honor an arbitrary pushed expression, the same way `PartitionPredicateImpl` does. Mixing + * classes supply their own `filterAttributes()`. + */ + trait CatalystRuntimeFilteringScan extends SupportsRuntimeCatalystFiltering { + self: BatchScanBaseClass => + + /** The full table schema, used to locate partition columns pruned out of `readSchema`. */ + protected def tableSchema: StructType + + private val catalystPredicates = ArrayBuffer.empty[CatalystExpression] + private var filterCalls = 0 + + override def filter(expressions: Array[CatalystExpression]): Unit = { + catalystPredicates ++= expressions + filterCalls += 1 + val partAttrs = partitionAttributes + if (partAttrs.isEmpty) return + val partAttrRefs = partAttrs.map(_._2) + + expressions.foreach { expr => + // Top down, so `s.part` is rewritten before its `s` child is considered. + val remapped = expr.transformDown { + case e => partitionAttrFor(e, partAttrs).getOrElse(e) + } + // Only evaluate expressions whose refs are all partition columns, so we can bind + // against the partition key InternalRow (same approach as PartitionPredicateImpl). + if (remapped.references.forall(r => partAttrRefs.exists(_.exprId == r.exprId))) { + val bound = BindReferences.bindReference(remapped, partAttrRefs) + val pred = CatalystPredicate.createInterpreted(bound) + self.data = self.data.filter { p => + try { + pred.eval(p.asInstanceOf[BufferedRows].partitionKey()) + } catch { + // Keep the partition on eval failure, which is safe here because every predicate + // the fixture pushes evaluates cleanly. `PartitionPredicateImpl` fails open for a + // reason of its own: Spark keeps the post-scan `FilterExec` on that path, so + // failing open costs just a pruning opportunity. A scan declaring an attribute in + // `fullyPushedFilterAttributes()` stands alone as the evaluator, so keeping an + // unevaluated partition would return nonmatching rows. + case _: Exception => true + } + } + } + } + } + + /** Predicates recorded by [[filter]], for test assertions only. */ + def pushedCatalystPredicates: Seq[CatalystExpression] = catalystPredicates.toSeq + + def filterCallCount: Int = filterCalls + + /** + * The `AttributeReference`s standing for the partition key InternalRow fields, in its field + * order, each paired with the name-part sequence of its partition column. The parts are kept + * unflattened so a quoted top-level column `a.b` (parts `Seq("a.b")`) stays distinct from a + * nested column `a`.`b` (parts `Seq("a", "b")`). Example: + * - `PARTITIONED BY (part, s.nested)` -> `(Seq("part"), AttributeReference(part))`, then + * `(Seq("s", "nested"), AttributeReference(s.nested))` + */ + private def partitionAttributes: Seq[(Seq[String], AttributeReference)] = { + partitioning.flatMap(_.references()).flatMap { ref => + val path = ref.fieldNames.toImmutableArraySeq + readSchema.findNestedField(path).orElse(tableSchema.findNestedField(path)).map { + case (_, f) => + path -> AttributeReference(ref.fieldNames.mkString("."), f.dataType, f.nullable)() + } + }.toSeq + } + + /** + * The partition key `AttributeReference` that `e` reads, or None if `e` reads no partition + * column. The path `e` reads is compared to each partition column's name parts component-wise + * with the resolver, so a quoted top-level column `a.b` cannot collide with a nested column + * `a`.`b`. Examples, under `PARTITIONED BY (part, s.nested)` where `nested` is field 0 of `s`: + * - `AttributeReference(part)` -> `AttributeReference(part)` + * - `GetStructField(AttributeReference(s), 0)` -> `AttributeReference(s.nested)` + * - `AttributeReference(s)` -> None if `s` itself is not a partition column, only `s.nested` + */ + private def partitionAttrFor( + e: CatalystExpression, + partAttrs: Seq[(Seq[String], AttributeReference)]): Option[AttributeReference] = { + val resolver = SQLConf.get.resolver + partitionKeyPath(e).flatMap { path => + partAttrs.collectFirst { + case (parts, attr) if parts.length == path.length && + parts.lazyZip(path).forall((part, name) => resolver(part, name)) => attr + } + } + } + + /** + * The name parts `e` reads, or None if it reads neither a column nor a struct field. Each + * `GetStructField` ordinal is the field's position in its parent struct. Examples: + * - `AttributeReference(a)` -> `Seq("a")`, the top level column a + * - `GetStructField(AttributeReference(a), 0)` -> `Seq("a", "b")`, the nested column a.b + * - `GetStructField(GetStructField(AttributeReference(a), 0), 0)` -> `Seq("a", "b", "c")` + */ + private def partitionKeyPath(e: CatalystExpression): Option[Seq[String]] = e match { + case a: AttributeReference => Some(Seq(a.name)) + case g: GetStructField => + partitionKeyPath(g.child).map(parent => parent :+ g.childSchema(g.ordinal).name) + case _ => None + } + } + case class InMemoryBatchScan( var _data: Seq[InputPartition], readSchema: StructType, diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala new file mode 100644 index 0000000000000..369e8123f1d7e --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala @@ -0,0 +1,132 @@ +/* + * 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.connector.catalog + +import java.util + +import InMemoryCatalystRuntimeFilterTable._ + +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, SortOrder, Transform} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.util.ArrayImplicits._ + +/** + * In-memory table whose batch scan mixes in [[CatalystRuntimeFilteringScan]], so runtime filters + * arrive as Catalyst expressions rather than connector predicates. + * + * Table properties: + * - `filter-attributes` (default: all partition cols): comma-separated list of + * column names to expose from `filterAttributes`. + * - `fully-pushed-filter-attributes` (default: none): comma-separated list of + * column names to expose from `fullyPushedFilterAttributes`. + */ +class InMemoryCatalystRuntimeFilterTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint] = Array.empty, + distribution: Distribution = Distributions.unspecified(), + ordering: Array[SortOrder] = Array.empty, + numPartitions: Option[Int] = None, + advisoryPartitionSize: Option[Long] = None, + isDistributionStrictlyRequired: Boolean = true, + numRowsPerSplit: Int = Int.MaxValue) + extends InMemoryTableWithV2Filter(name, columns, partitioning, properties, constraints, + distribution, ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired, + numRowsPerSplit) { + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new InMemoryCatalystRuntimeFilterScanBuilder(schema, options) + } + + class InMemoryCatalystRuntimeFilterScanBuilder( + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends InMemoryScanBuilder(tableSchema, options) { + override def build: Scan = InMemoryCatalystRuntimeFilterBatchScan( + data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, + schema, tableSchema, options) + } + + /** + * Scan that receives runtime filters as Catalyst expressions. Pruning comes from + * [[CatalystRuntimeFilteringScan]], so fully-pushed predicates are enforced when Spark drops + * the post-scan [[org.apache.spark.sql.execution.FilterExec]]. + */ + case class InMemoryCatalystRuntimeFilterBatchScan( + var _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + with CatalystRuntimeFilteringScan { + + private val restrictedFilterAttrs: Option[Set[String]] = + Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey)) + .map(_.split(",").map(_.trim).toSet) + + private val fullyPushedFilterAttrs: Set[String] = Option( + InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey)) + .map(_.split(",").map(_.trim).toSet) + .getOrElse(Set.empty) + + /** + * The partition columns, each named by the top level read schema column it lives under, the + * form both interface methods require. Columns pruned out of the read schema are dropped, + * since neither method may name one. Examples: + * - `PARTITIONED BY (part)` -> `"part"` + * - `PARTITIONED BY (s.nested)` -> `"s"`, the struct column holding the partition field + */ + private def partitionAttrNames: Array[String] = { + val scanFields = readSchema.fields.map(_.name).toSet + partitioning.flatMap(_.references()).map(_.fieldNames.head).distinct + .filter(scanFields.contains) + } + + override def filterAttributes(): Array[NamedReference] = { + partitionAttrNames + .filter(name => restrictedFilterAttrs.forall(_.contains(name))) + .map(FieldReference.column) + } + + // Not intersected with `filterAttributes()`, so a table can declare a fully pushed attribute + // that is not a filter attribute, a combination the interface forbids. + override def fullyPushedFilterAttributes(): Array[NamedReference] = { + partitionAttrNames.filter(fullyPushedFilterAttrs.contains).map(FieldReference.column) + } + } +} + +object InMemoryCatalystRuntimeFilterTable { + /** + * Table property: comma-separated column names to expose from + * filterAttributes. Default: all partition columns. + */ + private[catalog] val FilterAttributesKey = "filter-attributes" + + /** + * Table property: comma-separated column names to expose from + * fullyPushedFilterAttributes. Default: none. + */ + private[catalog] val FullyPushedFilterAttributesKey = "fully-pushed-filter-attributes" +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala index fd02f926e8141..6cb23d505f784 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} import org.apache.spark.sql.connector.expressions.{FieldReference, LogicalExpressions, NamedReference, SortDirection, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} import org.apache.spark.sql.connector.write.{BatchWrite, DeltaBatchWrite, DeltaWrite, DeltaWriteBuilder, DeltaWriter, DeltaWriterFactory, LogicalWriteInfo, PhysicalWriteInfo, RequiresDistributionAndOrdering, RowLevelOperation, RowLevelOperationBuilder, RowLevelOperationInfo, SupportsDelta, Write, WriteBuilder, WriterCommitMessage} import org.apache.spark.sql.connector.write.RowLevelOperation.Command import org.apache.spark.sql.types.StructType @@ -78,7 +78,10 @@ class InMemoryRowLevelOperationTable private ( private final val SUPPORTS_DELTAS = "supports-deltas" private final val SPLIT_UPDATES = "split-updates" private final val NO_METADATA = "no-metadata" + private final val USE_CATALYST_RUNTIME_FILTERING = "use-catalyst-runtime-filtering" private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" + private final val useCatalystRuntimeFiltering = + properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty @@ -133,7 +136,7 @@ class InMemoryRowLevelOperationTable private ( case class PartitionBasedOperation(command: Command, options: CaseInsensitiveStringMap) extends RowLevelOperation with RowLevelOperationWithOptions { - var configuredScan: InMemoryBatchScan = _ + var configuredScan: BatchScanBaseClass = _ override def requiredMetadataAttributes(): Array[NamedReference] = { if (noMetadata) { @@ -144,12 +147,8 @@ class InMemoryRowLevelOperationTable private ( } override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) { - override def build: Scan = { - val scan = super.build() - configuredScan = scan.asInstanceOf[InMemoryBatchScan] - scan - } + newRowLevelScanBuilder(options) { scan => + configuredScan = scan } } @@ -186,7 +185,7 @@ class InMemoryRowLevelOperationTable private ( override def description(): String = "InMemoryPartitionReplaceOperation" } - private case class PartitionBasedReplaceData(scan: InMemoryBatchScan) + private case class PartitionBasedReplaceData(scan: BatchScanBaseClass) extends TestBatchWrite { override protected def doCommit( @@ -216,7 +215,7 @@ class InMemoryRowLevelOperationTable private ( override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) + newRowLevelScanBuilder(options)(_ => ()) } override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = { @@ -267,6 +266,55 @@ class InMemoryRowLevelOperationTable private ( override def abort(messages: Array[WriterCommitMessage]): Unit = {} } + + /** + * Builds a scan for row-level operations. When `use-catalyst-runtime-filtering` is set, the + * scan mixes in [[CatalystRuntimeFilteringScan]] so group filtering goes through the Catalyst + * path. + */ + private def newRowLevelScanBuilder( + options: CaseInsensitiveStringMap)( + onBuild: BatchScanBaseClass => Unit): ScanBuilder = { + new InMemoryScanBuilder(schema, options) { + override protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (useCatalystRuntimeFiltering) { + InMemoryCatalystRowLevelBatchScan(partitions, readSchema, tableSchema, options) + } else { + super.createScan(partitions, readSchema, tableSchema, options) + } + } + + override def build: Scan = { + val scan = super.build().asInstanceOf[BatchScanBaseClass] + onBuild(scan) + scan + } + } + } + + /** + * Row-level batch scan that receives runtime filters as Catalyst expressions. Pruning comes + * from [[CatalystRuntimeFilteringScan]], so group filtering actually drops partitions (needed + * for `replacedPartitions` assertions). + */ + case class InMemoryCatalystRowLevelBatchScan( + var _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + with CatalystRuntimeFilteringScan { + + override def filterAttributes(): Array[NamedReference] = { + val scanFields = readSchema.fields.map(_.name).toSet + partitioning.flatMap(_.references()) + .filter(ref => scanFields.contains(ref.fieldNames.mkString("."))) + } + } } private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala new file mode 100644 index 0000000000000..6fd043c364640 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryScanMergingPartitionFilterTable.scala @@ -0,0 +1,131 @@ +/* + * 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.connector.catalog + +import java.util + +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.read.{Batch, Scan, ScanBuilder} +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Catalog that hands out [[InMemoryScanMergingPartitionFilterTable]]s, an iterative-pushdown source + * that additionally opts in to Spark-side scan merging. Used to exercise merging two DSv2 scans + * whose (equal) filter is strict only via the iterative PartitionPredicate second pass. + */ +class InMemoryScanMergingPartitionFilterCatalog + extends InMemoryTableEnhancedPartitionFilterCatalog { + import CatalogV2Implicits._ + + /** + * The table this catalog hands out. Overridden by [[InMemoryScanMergingReportingCatalog]], which + * shares the `createTable` body below and differs only in the table it creates. + */ + protected def newScanMergingTable( + tableName: String, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = + new InMemoryScanMergingPartitionFilterTable(tableName, columns, partitions, properties) + + override def createTable( + ident: Identifier, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = { + if (tables.containsKey(ident)) { + throw new TableAlreadyExistsException(ident.asMultipartIdentifier) + } + InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) + val tableName = s"$name.${ident.quoted}" + val table = newScanMergingTable(tableName, columns, partitions, properties) + tables.put(ident, table) + namespaces.putIfAbsent(ident.namespace.toList, Map()) + table + } +} + +/** + * Like [[InMemoryScanMergingPartitionFilterCatalog]] but hands out tables that KEEP their reported + * partitioning/ordering (no [[NonReportingScan]] wrapper), so a scan merge that must preserve the + * reported key-grouped partitioning across the merge can be exercised. + */ +class InMemoryScanMergingReportingCatalog extends InMemoryScanMergingPartitionFilterCatalog { + override protected def newScanMergingTable( + tableName: String, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = + new InMemoryScanMergingReportingTable(tableName, columns, partitions, properties) +} + +/** + * An [[InMemoryEnhancedPartitionFilterTable]] that opts into `TableCapability.SCAN_MERGING`, so + * [[org.apache.spark.sql.execution.planmerging.PlanMerger]] may fuse two scans of it. It does NOT + * strip the reported partitioning: a partitioned table's scan reports `KeyGroupedPartitioning` as + * usual, so this is the fixture for checking that a merge preserves that report on the rebuilt + * merged scan (re-derived by V2ScanPartitioningAndOrdering). It is also the base of + * [[InMemoryScanMergingPartitionFilterTable]], which drops the report again. + */ +class InMemoryScanMergingReportingTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String]) + extends InMemoryEnhancedPartitionFilterTable(name, columns, partitioning, properties) { + + override def capabilities(): util.Set[TableCapability] = { + val caps = new util.HashSet[TableCapability](super.capabilities()) + caps.add(TableCapability.SCAN_MERGING) + caps + } +} + +/** + * An [[InMemoryScanMergingReportingTable]] whose scan is wrapped in a thin [[NonReportingScan]], so + * a partitioned table does not set the scan relation's `keyGroupedPartitioning`, keeping this + * fixture focused on the iterative-pushdown behavior under test. Preserving a reported partitioning + * across a merge is exercised by the base [[InMemoryScanMergingReportingTable]] instead. + */ +class InMemoryScanMergingPartitionFilterTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String]) + extends InMemoryScanMergingReportingTable(name, columns, partitioning, properties) { + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new InMemoryEnhancedPartitionFilterScanBuilder(schema()) { + override def build(): Scan = NonReportingScan(super.build()) + } +} + +/** + * Thin scan decorator that exposes only `readSchema`, `toBatch` and `description`, dropping the + * base scan's `SupportsReportPartitioning`/`SupportsReportStatistics`. So the scan relation carries + * no reported partitioning/ordering/statistics -- for a partitioned table this keeps + * `keyGroupedPartitioning` unset, so the fixture stays focused on pushdown; preserving reported + * partitioning across a merge is exercised by [[InMemoryScanMergingReportingTable]]. + */ +case class NonReportingScan(inner: Scan) extends Scan { + override def readSchema(): StructType = inner.readSchema() + override def toBatch: Batch = inner.toBatch + override def description(): String = inner.description() +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala index bb137ba4830df..cc5c75dd779be 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala @@ -22,6 +22,7 @@ import java.util.Collections import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger +import scala.collection.mutable import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.InternalRow @@ -47,12 +48,31 @@ class BasicInMemoryTableCatalog extends TableCatalog { private var _name: Option[String] = None private var copyOnLoad: Boolean = false + private var stateOptionKeys: util.Set[String] = util.Set.of() + + // Records every (TableContext, table-state options) pair passed to the options-aware + // loadTable(), in call order, so tests can verify that the analyzer / DataFrame API correctly + // constructed and forwarded them -- including how many times loadTable was called when the same + // table is referenced more than once in a statement with different options. + // "loadTable" is in the name because the subclass InMemoryChangelogCatalog has an analogous + // `lastOptions` recording the options passed to loadChangelog(); the two must not collide. + private val _loadTableCalls = mutable.ArrayBuffer.empty[(TableContext, CaseInsensitiveStringMap)] + def loadTableCalls: Seq[(TableContext, CaseInsensitiveStringMap)] = _loadTableCalls.toSeq + def resetLoadTableCalls(): Unit = _loadTableCalls.clear() + + def lastTableContext: Option[TableContext] = _loadTableCalls.lastOption.map(_._1) + def lastLoadTableOptions: Option[CaseInsensitiveStringMap] = _loadTableCalls.lastOption.map(_._2) override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = { _name = Some(name) copyOnLoad = options.getBoolean("copyOnLoad", false) + stateOptionKeys = Option(options.get("tableStateOptionKeys")) + .map(_.split(",").iterator.map(_.trim).filter(_.nonEmpty).toSet.asJava) + .getOrElse(util.Set.of()) } + override def tableStateOptionKeys(): util.Set[String] = stateOptionKeys + override def name: String = _name.get override def listTables(namespace: Array[String]): Array[Identifier] = { @@ -124,6 +144,16 @@ class BasicInMemoryTableCatalog extends TableCatalog { } } + // Records the forwarded context/state options so tests can verify they reached the catalog, then + // defers to the default dispatch in TableCatalog (rather than reimplementing it here). + override def loadTable( + ident: Identifier, + context: TableContext, + stateOptions: CaseInsensitiveStringMap): Table = { + _loadTableCalls += ((context, stateOptions)) + super.loadTable(ident, context, stateOptions) + } + override def invalidateTable(ident: Identifier): Unit = { invalidatedTables.add(ident) } @@ -163,14 +193,39 @@ class BasicInMemoryTableCatalog extends TableCatalog { InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) val tableName = s"$name.${ident.quoted}" - val table = new InMemoryTable(tableName, columns, partitions, properties, constraints, - distribution, ordering, requiredNumPartitions, advisoryPartitionSize, - distributionStrictlyRequired, numRowsPerSplit) + val table = newInMemoryTable( + tableName, columns, partitions, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit, + util.UUID.randomUUID().toString) tables.put(ident, table) namespaces.putIfAbsent(ident.namespace.toList, Map()) table } + /** + * Builds the in-memory table this catalog serves. Subclasses that expose a specialized table + * type must override this so both CREATE and ALTER reconstruct the same class. + */ + // scalastyle:off argcount + protected def newInMemoryTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryTable(name, columns, partitioning, properties, constraints, distribution, + ordering, requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, + numRowsPerSplit, id) + } + override def alterTable(ident: Identifier, changes: TableChange*): Table = { val table = loadTable(ident).asInstanceOf[InMemoryBaseTable] val properties = CatalogV2Util.applyPropertiesChanges(table.properties, changes) @@ -208,22 +263,13 @@ class BasicInMemoryTableCatalog extends TableCatalog { val currentVersion = table.version() val columnsWithIds = InMemoryBaseTable.assignMissingIds( CatalogV2Util.structTypeToV2Columns(schema)) + val reconstructedId = Option(table.id()).getOrElse(util.UUID.randomUUID().toString) val newTable = table match { - case _: InMemoryTable => - new InMemoryTable( - name = table.name, - columns = columnsWithIds, - partitioning = finalPartitioning, - properties = properties, - constraints = constraints, - id = table.id) - .alterTableWithData(table.data, schemaAfterDrops) - case _: InMemoryTableWithV2Filter => - new InMemoryTableWithV2Filter( - name = table.name, - columns = columnsWithIds, - partitioning = finalPartitioning, - properties = properties) + case _: InMemoryTable | _: InMemoryTableWithV2Filter => + newInMemoryTable( + table.name, columnsWithIds, finalPartitioning, properties, constraints, + table.distribution, table.ordering, table.numPartitions, table.advisoryPartitionSize, + table.isDistributionStrictlyRequired, table.numRowsPerSplit, reconstructedId) .alterTableWithData(table.data, schemaAfterDrops) case other => throw new UnsupportedOperationException( diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala new file mode 100644 index 0000000000000..e699f439223f9 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala @@ -0,0 +1,61 @@ +/* + * 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.connector.catalog + +import java.util + +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.Distribution +import org.apache.spark.sql.connector.expressions.{SortOrder, Transform} + +/** + * Mix-in that constructs [[InMemoryCatalystRuntimeFilterTable]] from the shared in-memory + * catalog factory used by both CREATE TABLE and ALTER TABLE. + */ +trait InMemoryCatalystRuntimeFilterTableFactory { self: BasicInMemoryTableCatalog => + // scalastyle:off argcount + override protected def newInMemoryTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryCatalystRuntimeFilterTable( + name, columns, partitioning, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit) + } +} + +class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog + with InMemoryCatalystRuntimeFilterTableFactory + +/** + * The [[InMemoryCatalog]] counterpart of [[InMemoryTableCatalystRuntimeFilterCatalog]]: it hands + * out tables whose scans take runtime filters as Catalyst expressions, and honors + * `numRowsPerSplit` so that a partition key can have several splits. + */ +class InMemoryCatalystRuntimeFilterCatalog extends InMemoryCatalog + with InMemoryCatalystRuntimeFilterTableFactory diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala index e9d73d0f9fe1e..b59a81cc29092 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala @@ -21,7 +21,9 @@ import java.util import org.scalatest.Assertions.assert -import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, NamedReference, Transform} +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} +import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, NamedReference, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.write.{LogicalWriteInfo, SupportsOverwriteV2, WriteBuilder, WriterCommitMessage} @@ -33,8 +35,17 @@ class InMemoryTableWithV2Filter( name: String, columns: Array[Column], partitioning: Array[Transform], - properties: util.Map[String, String]) - extends InMemoryBaseTable(name, columns, partitioning, properties) with SupportsDeleteV2 { + properties: util.Map[String, String], + constraints: Array[Constraint] = Array.empty, + distribution: Distribution = Distributions.unspecified(), + ordering: Array[SortOrder] = Array.empty, + numPartitions: Option[Int] = None, + advisoryPartitionSize: Option[Long] = None, + isDistributionStrictlyRequired: Boolean = true, + numRowsPerSplit: Int = Int.MaxValue) + extends InMemoryBaseTable(name, columns, partitioning, properties, constraints, distribution, + ordering, numPartitions, advisoryPartitionSize, isDistributionStrictlyRequired, + numRowsPerSplit) with SupportsDeleteV2 { override def canDeleteWhere(predicates: Array[Predicate]): Boolean = { InMemoryTableWithV2Filter.supportsPredicates(predicates) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala index ef2f5e26f0029..be5a907f52125 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2FilterCatalog.scala @@ -19,31 +19,28 @@ package org.apache.spark.sql.connector.catalog import java.util -import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException -import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.distributions.Distribution +import org.apache.spark.sql.connector.expressions.{SortOrder, Transform} class InMemoryTableWithV2FilterCatalog extends InMemoryTableCatalog { - import CatalogV2Implicits._ - - override def createTable( - ident: Identifier, + // scalastyle:off argcount + override protected def newInMemoryTable( + name: String, columns: Array[Column], - partitions: Array[Transform], - properties: util.Map[String, String]): Table = { - if (tables.containsKey(ident)) { - throw new TableAlreadyExistsException(ident.asMultipartIdentifier) - } - - InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) - - val tableName = s"$name.${ident.quoted}" - val table = new InMemoryTableWithV2Filter(tableName, columns, partitions, properties) - tables.put(ident, table) - namespaces.putIfAbsent(ident.namespace.toList, Map()) - table - } - - override def createTable(ident: Identifier, tableInfo: TableInfo): Table = { - createTable(ident, tableInfo.columns(), tableInfo.partitions(), tableInfo.properties) + partitioning: Array[Transform], + properties: util.Map[String, String], + constraints: Array[Constraint], + distribution: Distribution, + ordering: Array[SortOrder], + requiredNumPartitions: Option[Int], + advisoryPartitionSize: Option[Long], + distributionStrictlyRequired: Boolean, + numRowsPerSplit: Int, + id: String): InMemoryBaseTable = { + // scalastyle:on argcount + new InMemoryTableWithV2Filter( + name, columns, partitioning, properties, constraints, distribution, ordering, + requiredNumPartitions, advisoryPartitionSize, distributionStrictlyRequired, numRowsPerSplit) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala index f2b16b731f41a..f48f99994a8cc 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala @@ -136,6 +136,7 @@ class TxnTable( override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = { catalog.writeTarget = this + lastWriteInfo = info super.newWriteBuilder(info) } @@ -171,6 +172,7 @@ class TxnTable( class TxnTableCatalog(delegate: InMemoryRowLevelOperationTableCatalog) extends TableCatalog { private val tables: util.Map[Identifier, TxnTable] = new ConcurrentHashMap[Identifier, TxnTable]() + val loadTableCalls: ArrayBuffer[(TableContext, CaseInsensitiveStringMap)] = ArrayBuffer.empty var writeTarget: TxnTable = _ @@ -178,6 +180,8 @@ class TxnTableCatalog(delegate: InMemoryRowLevelOperationTableCatalog) extends T override def capabilities: java.util.Set[TableCatalogCapability] = delegate.capabilities + override def tableStateOptionKeys(): util.Set[String] = delegate.tableStateOptionKeys() + override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {} override def listTables(namespace: Array[String]): Array[Identifier] = { @@ -196,6 +200,22 @@ class TxnTableCatalog(delegate: InMemoryRowLevelOperationTableCatalog) extends T }) } + override def loadTable( + ident: Identifier, + context: TableContext, + options: CaseInsensitiveStringMap): Table = { + loadTableCalls += ((context, options)) + super.loadTable(ident, context, options) + } + + override def loadTable(ident: Identifier, version: String): Table = { + delegate.loadTable(ident, version) + } + + override def loadTable(ident: Identifier, timestamp: Long): Table = { + delegate.loadTable(ident, timestamp) + } + override def alterTable(ident: Identifier, changes: TableChange*): Table = { // AlterTable may be called by ResolveSchemaEvolution when schema evolution is enabled. Thus, // it needs to be transactional. The schema changes are only propagated to the delegate at diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/expressions/LiteralValueSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/expressions/LiteralValueSuite.scala new file mode 100644 index 0000000000000..9a21f14e022a0 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/expressions/LiteralValueSuite.scala @@ -0,0 +1,52 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.expressions + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.types.{BinaryType, BooleanType, IntegerType, StringType} + +class LiteralValueSuite extends SparkFunSuite { + + test("SPARK-58782: null literals should render as NULL") { + assert(LiteralValue(null, StringType).toString === "NULL") + assert(LiteralValue(null, IntegerType).toString === "NULL") + assert(LiteralValue(null, BooleanType).toString === "NULL") + assert(LiteralValue(null, BinaryType).toString === "NULL") + } + + test("non-null string literals should be quoted") { + assert(LiteralValue("test", StringType).toString === "'test'") + assert(LiteralValue("", StringType).toString === "''") + assert(LiteralValue("it's", StringType).toString === "'it''s'") + } + + test("non-null numeric literals should not be quoted") { + assert(LiteralValue(42, IntegerType).toString === "42") + assert(LiteralValue(0, IntegerType).toString === "0") + } + + test("non-null boolean literals should not be quoted") { + assert(LiteralValue(true, BooleanType).toString === "true") + assert(LiteralValue(false, BooleanType).toString === "false") + } + + test("non-null binary literals should render as hex") { + val bytes = Array[Byte](0x12, 0x34, 0xAB.toByte, 0xCD.toByte) + assert(LiteralValue(bytes, BinaryType).toString === "0x1234ABCD") + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala index 46ed870cb3221..79fd057f037e4 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala @@ -18,18 +18,26 @@ package org.apache.spark.sql.execution.datasources.v2 import java.util +import java.util.OptionalLong import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics} +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.plans.logical.{Histogram, HistogramBin} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils +import org.apache.spark.sql.catalyst.trees.TreePattern import org.apache.spark.sql.catalyst.util.FieldMetadataUtils.FIELD_ID_METADATA_KEY import org.apache.spark.sql.catalyst.util.INTERNAL_METADATA_KEYS import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability} -import org.apache.spark.sql.connector.expressions.FieldReference +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} +import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics} +import org.apache.spark.sql.connector.read.colstats.ColumnStatistics +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap -class DataSourceV2RelationSuite extends SparkFunSuite { +class DataSourceV2RelationSuite extends SparkFunSuite with SQLHelper { test("DataSourceV2Relation.v1StatsToV2Stats") { val schema = StructType(Seq( @@ -114,6 +122,360 @@ class DataSourceV2RelationSuite extends SparkFunSuite { assert(!idV2NoHist.histogram().isPresent) } + private def scanRel( + output: Seq[AttributeReference], + scan: Scan, + table: Table = new FakeTableWithSchema()): DataSourceV2ScanRelation = { + DataSourceV2ScanRelation( + DataSourceV2Relation(table, output, None, None, CaseInsensitiveStringMap.empty()), + scan, + output) + } + + test("DataSourceV2ScanRelation.computeStats uses non-empty scan stats with CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(42L) + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("id"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(40L) + override def avgLen(): OptionalLong = OptionalLong.of(4L) + }) + stats + } + } + } + + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.rowCount.contains(BigInt(42))) + assert(stats.attributeStats.size === 1) + assert(stats.attributeStats(idAttr).distinctCount.contains(BigInt(40))) + assert(stats.attributeStats(idAttr).avgLen.contains(4L)) + assert(stats.sizeInBytes === + EstimationUtils.getOutputSize(output, BigInt(42), stats.attributeStats)) + } + } + + test("DataSourceV2ScanRelation.computeStats derives size 1 for a zero-row scan with CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(0L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.rowCount.contains(BigInt(0))) + assert(stats.sizeInBytes === BigInt(1)) + } + } + + test("DataSourceV2ScanRelation.computeStats uses full stats with plan stats enabled") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(50L) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(1000L) + override def numRows(): OptionalLong = OptionalLong.of(7L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "true") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(1000)) + assert(stats.rowCount.contains(BigInt(7))) + } + } + + test("DataSourceV2ScanRelation.computeStats treats column-only scan stats as non-empty") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("id"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(7L) + }) + stats + } + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.size === 1) + assert(stats.attributeStats(idAttr).distinctCount.contains(BigInt(7))) + } + } + + test("DataSourceV2ScanRelation.computeStats uses size-only estimates without CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(50L) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(1000L) + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(50)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default estimateSizeInBytes without CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(64L) + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(64)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats infers size-only estimates from row count") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === EstimationUtils.getSizePerRow(output) * BigInt(5)) + assert(stats.rowCount.isEmpty) + assert(stats.attributeStats.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats invokes estimateStatistics once without CBO") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + // Scan with only a row count and the default estimateSizeInBytes(): the size-only path must + // consult the full statistics exactly once, not once for the (empty) size and again for the row + // count. + var estimateStatisticsCalls = 0 + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = { + estimateStatisticsCalls += 1 + new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(5L) + } + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === EstimationUtils.getSizePerRow(output) * BigInt(5)) + assert(estimateStatisticsCalls === 1, + s"estimateStatistics should be called at most once, was $estimateStatisticsCalls") + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size without CBO for empty stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size without reported stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size for empty scan stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val scan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + test("DataSourceV2ScanRelation.computeStats uses default size for null scan stats") { + val output = Seq(AttributeReference("id", IntegerType)()) + val nullStatsScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = null + } + val nullColumnStatsScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = null + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "true", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + Seq(nullStatsScan, nullColumnStatsScan).foreach { scan => + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false", + SQLConf.DEFAULT_SIZE_IN_BYTES.key -> "12345") { + Seq(nullStatsScan, nullColumnStatsScan).foreach { scan => + val stats = scanRel(output, scan).computeStats() + + assert(stats.sizeInBytes === BigInt(12345)) + assert(stats.rowCount.isEmpty) + } + } + } + + test("DataSourceV2ScanRelation.computeStats tolerates null column stats with row count") { + val idAttr = AttributeReference("id", IntegerType)() + val output = Seq(idAttr) + // numRows present, columnStats null: isNotEmpty is true (via numRows), so the conversion runs + // and must not NPE on the null column-stats map. + val rowCountScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(42L) + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = null + } + } + // sizeInBytes present, columnStats null: same null-tolerance requirement, reached via + // sizeInBytes instead of numRows. + val sizeScan = new Scan with SupportsReportStatistics { + override def readSchema(): StructType = StructType(Seq(StructField("id", IntegerType))) + override def estimateStatistics(): V2Statistics = new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(1000L) + override def numRows(): OptionalLong = OptionalLong.empty() + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = null + } + } + + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val rowCountStats = scanRel(output, rowCountScan).computeStats() + assert(rowCountStats.rowCount.contains(BigInt(42))) + assert(rowCountStats.attributeStats.isEmpty) + assert(rowCountStats.sizeInBytes === + EstimationUtils.getOutputSize(output, BigInt(42), rowCountStats.attributeStats)) + + val sizeStats = scanRel(output, sizeScan).computeStats() + assert(sizeStats.sizeInBytes === BigInt(1000)) + assert(sizeStats.rowCount.isEmpty) + assert(sizeStats.attributeStats.isEmpty) + } + } + test("create strips leaked internal metadata but preserves column IDs") { // A column carrying both a column ID (surfaced on purpose) and every internal metadata key // (listed in INTERNAL_METADATA_KEYS), simulating a v2 source that leaks internal metadata. @@ -146,4 +508,37 @@ class DataSourceV2RelationSuite extends SparkFunSuite { assert(field.id.contains("1")) assert(field.metadata.contains(FIELD_ID_METADATA_KEY)) } + + test("nodePatterns declare the DSv2 relation identity tree patterns") { + val table = new Table { + override def name(): String = "t" + override def columns(): Array[Column] = + Array(Column.create("id", IntegerType)) + override def capabilities(): util.Set[TableCapability] = + util.Set.of[TableCapability]() + } + + val relation = + DataSourceV2Relation.create(table, None, None, CaseInsensitiveStringMap.empty()) + assert(relation.containsPattern(TreePattern.DATA_SOURCE_V2_RELATION)) + + val scan = new Scan { + override def readSchema(): StructType = relation.schema + } + val scanRelation = DataSourceV2ScanRelation(relation, scan, relation.output) + assert(scanRelation.containsPattern(TreePattern.DATA_SOURCE_V2_SCAN_RELATION)) + // The scan leaf must not inherit the pre-pushdown relation's identity pattern from its + // `relation` field (which is a case-class arg, not a child), so pruning on the two bits + // stays distinct. + assert(!scanRelation.containsPattern(TreePattern.DATA_SOURCE_V2_RELATION)) + } +} + +private class FakeTableWithSchema( + tableSchema: StructType = StructType(Seq(StructField("id", IntegerType)))) + extends Table { + + override def name(): String = "fake" + override def schema(): StructType = tableSchema + override def capabilities(): java.util.Set[TableCapability] = java.util.Set.of() } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DataTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DataTypeSuite.scala index 285d840eed6c8..6aae5047395da 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DataTypeSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DataTypeSuite.scala @@ -24,6 +24,7 @@ import org.json4s.jackson.JsonMethods import org.apache.spark.{SparkClassNotFoundException, SparkException, SparkFunSuite} import org.apache.spark.SparkIllegalArgumentException +import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.analysis.{caseInsensitiveResolution, caseSensitiveResolution} import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParseException} import org.apache.spark.sql.catalyst.plans.SQLHelper @@ -1656,22 +1657,27 @@ class DataTypeSuite extends SparkFunSuite with SQLHelper { } } - test("SPARK-56965: JSON parser rejects nanos timestamp types when preview flag is off") { + test("SPARK-57835: JSON parser reconstructs nanos timestamp types when preview flag is off") { + // Read-through policy: unlike the SQL parser (DataTypeAstBuilder), the JSON path is how a + // persisted schema is restored from the catalog, so it must reconstruct nanos types even + // when the preview flag is off. Otherwise a table written with the flag on would become + // completely inaccessible (DESCRIBE / SHOW CREATE TABLE / DROP) once it is off. The flag is + // instead enforced at analysis/execution time (TypeUtils.failUnsupportedDataType). This + // mirrors how TIME types round-trip through fromJson regardless of their own flag. withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "false") { Seq( - "\"timestamp_ltz(7)\"" -> "Nanosecond-precision timestamp types", - "\"timestamp_ntz(9)\"" -> "Nanosecond-precision timestamp types").foreach { - case (json, featureName) => - checkError( - exception = intercept[SparkException] { - DataType.fromJson(json) - }, - condition = "FEATURE_NOT_ENABLED", - parameters = Map( - "featureName" -> featureName, - "configKey" -> "spark.sql.timestampNanosTypes.enabled", - "configValue" -> "true")) + "\"timestamp_ltz(7)\"" -> TimestampLTZNanosType(7), + "\"timestamp_ntz(9)\"" -> TimestampNTZNanosType(9)).foreach { + case (json, expected) => + assert(DataType.fromJson(json) === expected) } + // Nested nanos types (inside struct/array/map) also reconstruct with the flag off, since + // that is exactly how catalog schemas are shaped. + val nested = StructType(Seq( + StructField("s", StructType(Seq(StructField("ntz", TimestampNTZNanosType(7))))), + StructField("a", ArrayType(TimestampLTZNanosType(8), containsNull = false)), + StructField("m", MapType(StringType, TimestampNTZNanosType(9))))) + assert(DataType.fromJson(nested.json) === nested) // Precision 6 maps to the GA types and stays accepted with the gate off. assert(DataType.fromJson("\"timestamp_ltz(6)\"") === TimestampType) assert(DataType.fromJson("\"timestamp_ntz(6)\"") === TimestampNTZType) @@ -1729,4 +1735,14 @@ class DataTypeSuite extends SparkFunSuite with SQLHelper { s"${clazz.getSimpleName}: PhysicalDataType should recognize non-singleton instance") } } + + test("SPARK-58350: DecimalType with scale greater than precision throws " + + "DECIMAL_SCALE_EXCEEDS_PRECISION") { + checkError( + exception = intercept[AnalysisException] { + DecimalType(2, 3) + }, + condition = "DECIMAL_SCALE_EXCEEDS_PRECISION", + parameters = Map("scale" -> "3", "precision" -> "2")) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DecimalSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DecimalSuite.scala index 080c6693280a0..a987aa52e1010 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DecimalSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/DecimalSuite.scala @@ -28,7 +28,8 @@ import org.apache.spark.unsafe.types.UTF8String class DecimalSuite extends SparkFunSuite with PrivateMethodTester with SQLHelper { - val allSupportedRoundModes = Seq(ROUND_HALF_UP, ROUND_HALF_EVEN, ROUND_CEILING, ROUND_FLOOR) + val allSupportedRoundModes = + Seq(ROUND_HALF_UP, ROUND_HALF_EVEN, ROUND_CEILING, ROUND_FLOOR, ROUND_DOWN) /** Check that a Decimal has the given string representation, precision and scale */ private def checkDecimal(d: Decimal, string: String, precision: Int, scale: Int): Unit = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala index 42579f6cc6ee3..3e72cdf0f392e 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/types/StructTypeSuite.scala @@ -506,6 +506,54 @@ class StructTypeSuite extends SparkFunSuite with SQLHelper { assert(struct.toString() === "StructType(StructField(a,IntegerType,true))") } + test("SPARK-58525: case-insensitive merge folds case-only field names, including nested") { + // Top-level: a case-only difference merges into a single field keeping the left name. + assert( + new StructType().add("value", StringType) + .merge(new StructType().add("Value", StringType), caseSensitive = false) === + new StructType().add("value", StringType)) + + // The case-sensitive default keeps both as distinct fields. + assert( + new StructType().add("value", StringType) + .merge(new StructType().add("Value", StringType)) === + new StructType().add("value", StringType).add("Value", StringType)) + + // Nested struct: the flag must propagate through the recursive merge (SPARK-58525), so a + // case-only difference nested under a matched parent also folds rather than producing both. + assert( + new StructType().add("s", new StructType().add("value", StringType)) + .merge( + new StructType().add("s", new StructType().add("Value", StringType)), + caseSensitive = false) === + new StructType().add("s", new StructType().add("value", StringType))) + + // array<struct> element recursion. + assert( + new StructType().add("a", ArrayType(new StructType().add("value", StringType))) + .merge( + new StructType().add("a", ArrayType(new StructType().add("Value", StringType))), + caseSensitive = false) === + new StructType().add("a", ArrayType(new StructType().add("value", StringType)))) + + // map<_, struct> value recursion. + assert( + new StructType().add("m", MapType(StringType, new StructType().add("value", StringType))) + .merge( + new StructType().add("m", MapType(StringType, new StructType().add("Value", StringType))), + caseSensitive = false) === + new StructType().add("m", MapType(StringType, new StructType().add("value", StringType)))) + + // map<struct, _> key recursion. mergeInternal recurses through map keys and values + // independently, so the key path needs its own coverage. + assert( + new StructType().add("m", MapType(new StructType().add("value", StringType), StringType)) + .merge( + new StructType().add("m", MapType(new StructType().add("Value", StringType), StringType)), + caseSensitive = false) === + new StructType().add("m", MapType(new StructType().add("value", StringType), StringType))) + } + test("SPARK-37191: Merge DecimalType") { val source1 = StructType.fromDDL("c1 DECIMAL(12, 2)") .merge(StructType.fromDDL("c1 DECIMAL(12, 2)")) diff --git a/sql/connect/client/jdbc/pom.xml b/sql/connect/client/jdbc/pom.xml index 461a01dcf2eee..ada9bba9e4762 100644 --- a/sql/connect/client/jdbc/pom.xml +++ b/sql/connect/client/jdbc/pom.xml @@ -47,6 +47,14 @@ <classifier>tests</classifier> <scope>test</scope> </dependency> + <!-- Compile-time only: the shade plugin relocates io.grpc references to the + copies embedded in the spark-connect-client-jvm jar, so consumers must + not pull the unshaded artifact transitively. --> + <dependency> + <groupId>io.grpc</groupId> + <artifactId>grpc-api</artifactId> + <scope>provided</scope> + </dependency> </dependencies> <build> <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory> diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaData.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaData.scala index 823984f666bcd..b49894e92809b 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaData.scala +++ b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaData.scala @@ -540,7 +540,7 @@ class SparkConnectDatabaseMetaData(conn: SparkConnectConnection) extends Databas if (field.nullable) columnNullable else columnNoNulls, // NULLABLE field.getComment().orNull, // REMARKS field.getCurrentDefaultValue().orNull, // COLUMN_DEF - 0, // CHAR_OCTET_LENGTH + JdbcTypeUtils.getCharOctetLength(field), // CHAR_OCTET_LENGTH i + 1, // ORDINAL_POSITION if (field.nullable) "YES" else "NO", // IS_NULLABLE "", // IS_AUTOINCREMENT @@ -683,7 +683,7 @@ class SparkConnectDatabaseMetaData(conn: SparkConnectConnection) extends Databas "SQL_DATA_TYPE", "SQL_DATETIME_SUB", "NUM_PREC_RADIX") - .orderBy("DATA_TYPE") + .orderBy("DATA_TYPE", "TYPE_NAME") new SparkConnectResultSet(df.collectResult()) } @@ -944,7 +944,8 @@ object SparkConnectDatabaseMetaData { numPrecRadix) // Static JDBC type metadata for the Spark SQL atomic types, mirroring the - // JdbcTypeUtils type-code/precision mapping. Only STRING is case-sensitive. + // JdbcTypeUtils type-code/precision mapping. CHAR, VARCHAR, and STRING are + // case-sensitive. // TIMESTAMP_NTZ is omitted because it maps to the same JDBC type code // (Types.TIMESTAMP) as TIMESTAMP, so the TIMESTAMP row already covers it. // TIME is omitted for now because its maximum PRECISION/scale representation @@ -958,6 +959,8 @@ object SparkConnectDatabaseMetaData { typeRow("FLOAT", Types.FLOAT, 7, null, null, false, 0, 0, 10), typeRow("DOUBLE", Types.DOUBLE, 15, null, null, false, 0, 0, 10), typeRow("DECIMAL", Types.DECIMAL, 38, null, "precision,scale", false, 0, 38, 10), + typeRow("CHAR", Types.CHAR, Int.MaxValue, "'", "length", true, 0, 0, null), + typeRow("VARCHAR", Types.VARCHAR, Int.MaxValue, "'", "length", true, 0, 0, null), typeRow("STRING", Types.VARCHAR, Int.MaxValue, "'", null, true, 0, 0, null), typeRow("BINARY", Types.VARBINARY, Int.MaxValue, "X'", null, false, 0, 0, null, literalSuffix = "'"), diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSet.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSet.scala index 6b02f655f04c7..12bd6cd133aba 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSet.scala +++ b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSet.scala @@ -54,17 +54,21 @@ class SparkConnectResultSet( override def next(): Boolean = { checkOpen() - val hasNext = iterator.hasNext - if (hasNext) { - currentRow = iterator.next() - cursor += 1 - } else { - currentRow = null - if (cursor > 0 && cursor == sparkResult.length) { + // rows are fetched lazily, so an error raised while the query is still + // streaming results surfaces here rather than in execute() + JdbcErrorUtils.mapToSQLException { + val hasNext = iterator.hasNext + if (hasNext) { + currentRow = iterator.next() cursor += 1 + } else { + currentRow = null + if (cursor > 0 && cursor == sparkResult.length) { + cursor += 1 + } } + hasNext } - hasNext } @volatile private var closed: Boolean = false diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatement.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatement.scala index 245832087268e..66f9cc8fbec4d 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatement.scala +++ b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatement.scala @@ -99,18 +99,20 @@ class SparkConnectStatement(conn: SparkConnectConnection) extends Statement { resultSet = null resultsExhausted = false - var df = conn.spark.sql(sql) - if (maxRows > 0) { - df = df.limit(maxRows) - } - val sparkResult = df.collectResult() - operationId = sparkResult.operationId - if (hasResultSet(sparkResult)) { - resultSet = new SparkConnectResultSet(sparkResult, this) - true - } else { - sparkResult.close() - false + JdbcErrorUtils.mapToSQLException { + var df = conn.spark.sql(sql) + if (maxRows > 0) { + df = df.limit(maxRows) + } + val sparkResult = df.collectResult() + operationId = sparkResult.operationId + if (hasResultSet(sparkResult)) { + resultSet = new SparkConnectResultSet(sparkResult, this) + true + } else { + sparkResult.close() + false + } } } diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtils.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtils.scala index 6480c5d768f3f..35236b32e2208 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtils.scala +++ b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtils.scala @@ -19,6 +19,12 @@ package org.apache.spark.sql.connect.client.jdbc.util import java.sql.{Array => _, _} +import scala.util.control.NonFatal + +import io.grpc.{Status, StatusRuntimeException} + +import org.apache.spark.SparkThrowable + private[jdbc] object JdbcErrorUtils { def stringifyTransactionIsolationLevel(level: Int): String = level match { @@ -53,4 +59,62 @@ private[jdbc] object JdbcErrorUtils { case _ => throw new IllegalArgumentException(s"Invalid fetch direction: $direction") } + + // SQLState class 08 is "connection exception"; HYT00 is the conventional + // (ODBC-derived) state for an elapsed timeout. + private val CONNECTION_EXCEPTION_ERROR_CLASS = "08" + private val CONNECTION_FAILURE = "08006" + private val TIMEOUT_EXPIRED = "HYT00" + + /** + * Maps the unchecked exceptions raised by the Spark Connect client (a + * [[SparkThrowable]] once GrpcExceptionConverter has converted the gRPC error) to + * the [[SQLException]] a JDBC method is required to throw: + * + * - a server error in SQLState class 08 ("connection exception", e.g. the + * `INVALID_HANDLE.SESSION_*` conditions with 08003) means the session backing + * the connection is gone; the connection is unusable and retrying on it is + * pointless, so it maps to a [[SQLNonTransientConnectionException]] keeping + * the server-provided SQLState. + * - a gRPC UNAVAILABLE (e.g. a server restart or a network blip) maps to a + * [[SQLTransientConnectionException]] with SQLState 08006 ("connection + * failure"), since a fresh connection can succeed. + * - a gRPC DEADLINE_EXCEEDED means the RPC deadline elapsed, which a slow query + * fires on a perfectly healthy connection, so it maps to a + * [[SQLTimeoutException]] rather than a connection error. + * - any other error keeps the server-provided SQLState when one is available. + * + * The gRPC status is read from the [[StatusRuntimeException]] that + * GrpcExceptionConverter preserves in the cause chain, never from message text, + * so a server-side error merely quoting a gRPC exception cannot be mistaken for + * a transport failure. SQLState class 08 is how connection pools and BI tools + * detect a dead connection and reconnect. + */ + def toSQLException(t: Throwable): SQLException = t match { + case e: SQLException => e + case e => + val chain = causeChain(e) + val sparkThrowableOpt = chain.collectFirst { case st: SparkThrowable => st } + val sqlState = sparkThrowableOpt.flatMap(st => Option(st.getSqlState)) + val grpcCode = chain.collectFirst { case sre: StatusRuntimeException => + sre.getStatus.getCode + } + if (sqlState.exists(_.startsWith(CONNECTION_EXCEPTION_ERROR_CLASS))) { + new SQLNonTransientConnectionException(e.getMessage, sqlState.get, e) + } else if (grpcCode.contains(Status.Code.UNAVAILABLE)) { + new SQLTransientConnectionException(e.getMessage, CONNECTION_FAILURE, e) + } else if (grpcCode.contains(Status.Code.DEADLINE_EXCEEDED)) { + new SQLTimeoutException(e.getMessage, TIMEOUT_EXPIRED, e) + } else { + new SQLException(e.getMessage, sqlState.orNull, e) + } + } + + /** Runs `body`, rethrowing any non-fatal failure as the mapped [[SQLException]]. */ + def mapToSQLException[T](body: => T): T = + try body catch { case NonFatal(e) => throw toSQLException(e) } + + // The take(20) caps the walk in case of a cause cycle. + private def causeChain(t: Throwable): Seq[Throwable] = + Iterator.iterate(t)(_.getCause).takeWhile(_ != null).take(20).toSeq } diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcTypeUtils.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcTypeUtils.scala index 48a42cd9ec9cb..45bb8fcb43d82 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcTypeUtils.scala +++ b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcTypeUtils.scala @@ -42,6 +42,8 @@ private[jdbc] object JdbcTypeUtils { case LongType => Types.BIGINT case FloatType => Types.FLOAT case DoubleType => Types.DOUBLE + case _: CharType => Types.CHAR + case _: VarcharType => Types.VARCHAR case StringType => Types.VARCHAR case _: DecimalType => Types.DECIMAL case DateType => Types.DATE @@ -65,7 +67,7 @@ private[jdbc] object JdbcTypeUtils { case LongType => classOf[JLong].getName case FloatType => classOf[JFloat].getName case DoubleType => classOf[JDouble].getName - case StringType => classOf[String].getName + case _: StringType => classOf[String].getName case _: DecimalType => classOf[JBigDecimal].getName case DateType => classOf[Date].getName case TimestampType => classOf[Timestamp].getName @@ -82,7 +84,7 @@ private[jdbc] object JdbcTypeUtils { def isSigned(field: StructField): Boolean = field.dataType match { case ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType | _: DecimalType => true - case NullType | BooleanType | StringType | DateType | BinaryType | _: TimeType | + case NullType | BooleanType | _: StringType | DateType | BinaryType | _: TimeType | TimestampType | TimestampNTZType | _: ArrayType | _: MapType | _: StructType => false case other => throw new SQLFeatureNotSupportedException(s"DataType $other is not supported yet.") @@ -97,6 +99,8 @@ private[jdbc] object JdbcTypeUtils { case LongType => 19 case FloatType => 7 case DoubleType => 15 + case c: CharType => c.length + case v: VarcharType => v.length case StringType => Int.MaxValue case DecimalType.Fixed(p, _) => p case DateType => 10 @@ -120,7 +124,7 @@ private[jdbc] object JdbcTypeUtils { case DoubleType => 15 case TimestampType => 6 case TimestampNTZType => 6 - case NullType | BooleanType | ByteType | ShortType | IntegerType | LongType | StringType | + case NullType | BooleanType | ByteType | ShortType | IntegerType | LongType | _: StringType | DateType | BinaryType | _: TimeType | _: ArrayType | _: MapType | _: StructType => 0 case DecimalType.Fixed(_, s) => s case other => @@ -134,7 +138,7 @@ private[jdbc] object JdbcTypeUtils { getPrecision(field) + 1 // may have leading negative sign case FloatType => 14 case DoubleType => 24 - case StringType => + case _: StringType => getPrecision(field) case DateType => 10 // length of `YYYY-MM-DD` case TimestampType => 29 // length of `YYYY-MM-DD HH:MM:SS.SSSSSS` @@ -168,6 +172,22 @@ private[jdbc] object JdbcTypeUtils { case _ => null } + /** + * JDBC `CHAR_OCTET_LENGTH` is a byte capacity. Spark CHAR/VARCHAR lengths are in + * characters, so report `4 * n` (UTF-8 maximum bytes per character), saturating at + * `Int.MaxValue`. Unbounded STRING and non-character types keep 0 (not applicable). + */ + def getCharOctetLength(field: StructField): Int = field.dataType match { + case c: CharType => maxUtf8OctetLength(c.length) + case v: VarcharType => maxUtf8OctetLength(v.length) + case _ => 0 + } + + private def maxUtf8OctetLength(numChars: Int): Int = { + val maxChars = Int.MaxValue / 4 + if (numChars > maxChars) Int.MaxValue else numChars * 4 + } + /** * Converts a value materialized by the Spark Connect client (Scala Seq / Map / Row for * complex types) into the corresponding standard JDBC object, recursively: diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaDataSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaDataSuite.scala index 3d3260c750f6f..f9276f3cd7152 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaDataSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectDatabaseMetaDataSuite.scala @@ -210,7 +210,7 @@ class SparkConnectDatabaseMetaDataSuite extends ConnectFunSuite with RemoteSpark val metadata = conn.getMetaData // scalastyle:off line.size.limit // CURRENT_PATH and SYSTEM are excluded: getSQLKeywords drops SQL:2003 reserved words (see companion). - assert(metadata.getSQLKeywords === "ADD,AFTER,AGGREGATE,ALIGN,ALWAYS,ANALYZE,ANTI,ANY_VALUE,APPLY,APPROX,ARCHIVE,ASC,ASOF,AUTO,BERNOULLI,BIN,BINDING,BIN_DISTRIBUTE_RATIO,BIN_END,BIN_START,BUCKET,BUCKETS,BYTE,CACHE,CASCADE,CATALOG,CATALOGS,CDC,CHANGE,CHANGES,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATION,COLLATIONS,COLLECTION,COLUMNS,COMMENT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONTAINS,CONTINUE,COST,CURRENT_DATABASE,CURRENT_SCHEMA,DATA,DATABASE,DATABASES,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAYOFYEAR,DAYS,DBPROPERTIES,DEFAULT_PATH,DEFINED,DEFINER,DELAY,DELIMITED,DESC,DFS,DIRECTORIES,DIRECTORY,DISTANCE,DISTRIBUTE,DIV,DO,ELSEIF,ENFORCED,ESCAPED,EVOLUTION,EXACT,EXCHANGE,EXCLUDE,EXCLUSIVE,EXIT,EXPLAIN,EXPORT,EXTEND,EXTENDED,FIELDS,FILEFORMAT,FIRST,FLOW,FOLLOWING,FORMAT,FORMATTED,FOUND,FUNCTIONS,GENERATED,GEOGRAPHY,GEOMETRY,HANDLER,HISTORY,HOURS,IDENTIFIED,IDENTIFIER,IF,IGNORE,ILIKE,IMMEDIATE,INCLUDE,INCLUSIVE,INCREMENT,INDEX,INDEXES,INPATH,INPUT,INPUTFORMAT,INVOKER,ITEMS,ITERATE,JSON,KEY,KEYS,LAST,LAZY,LEAVE,LEVEL,LIMIT,LINES,LIST,LOAD,LOCATION,LOCK,LOCKS,LOGICAL,LONG,LOOP,MACRO,MAP,MATCHED,MATCH_CONDITION,MATERIALIZED,MEASURE,METRICS,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTES,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NEAREST,NORELY,NULLS,OFFSET,OPTION,OPTIONS,OUTPUTFORMAT,OVERWRITE,PARTITIONED,PARTITIONS,PATH,PERCENT,PIVOT,PLACING,PRECEDING,PRINCIPALS,PROCEDURES,PROPERTIES,PURGE,QUALIFY,QUARTER,QUERY,RECORDREADER,RECORDWRITER,RECOVER,RECURSION,REDUCE,REFRESH,RELY,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,ROLE,ROLES,SCD,SCHEMA,SCHEMAS,SECONDS,SECURITY,SEMI,SEPARATED,SEQUENCE,SERDE,SERDEPROPERTIES,SETS,SHORT,SHOW,SIMILARITY,SINGLE,SKEWED,SORT,SORTED,SOURCE,STATISTICS,STORED,STRATIFY,STREAM,STREAMING,STRING,STRUCT,SUBSTR,SYNC,SYSTEM_PATH,SYSTEM_TIME,SYSTEM_VERSION,TABLES,TARGET,TBLPROPERTIES,TERMINATED,TIMEDIFF,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TOUCH,TRACK,TRANSACTION,TRANSACTIONS,TRANSFORM,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNIFORM,UNLOCK,UNPIVOT,UNSET,UNTIL,USE,VAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WATERMARK,WEEK,WEEKS,WHILE,WIDTH,X,YEARS,ZONE") + assert(metadata.getSQLKeywords === "ADD,AFTER,AGGREGATE,ALIGN,ALWAYS,ANALYZE,ANTI,ANY_VALUE,APPLY,APPROX,ARCHIVE,ASC,ASOF,AUTO,BERNOULLI,BIN,BINDING,BIN_DISTRIBUTE_RATIO,BIN_END,BIN_START,BUCKET,BUCKETS,BYTE,CACHE,CASCADE,CATALOG,CATALOGS,CDC,CHANGE,CHANGES,CLEAR,CLUSTER,CLUSTERED,CODEGEN,COLLATION,COLLATIONS,COLLECTION,COLUMNS,COMMENT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONDITIONAL,CONTAINS,CONTINUE,COST,CURRENT_DATABASE,CURRENT_SCHEMA,DATA,DATABASE,DATABASES,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAYOFYEAR,DAYS,DBPROPERTIES,DEFAULT_PATH,DEFINED,DEFINER,DELAY,DELIMITED,DESC,DFS,DIRECTORIES,DIRECTORY,DISTANCE,DISTRIBUTE,DIV,DO,ELSEIF,EMPTY,ENFORCED,ERROR,ESCAPED,EVOLUTION,EXACT,EXCHANGE,EXCLUDE,EXCLUSIVE,EXIT,EXPLAIN,EXPORT,EXTEND,EXTENDED,FIELDS,FILEFORMAT,FIRST,FLOW,FOLLOWING,FORMAT,FORMATTED,FOUND,FUNCTIONS,GENERATED,GEOGRAPHY,GEOMETRY,HANDLER,HISTORY,HOURS,IDENTIFIED,IDENTIFIER,IF,IGNORE,ILIKE,IMMEDIATE,INCLUDE,INCLUSIVE,INCREMENT,INDEX,INDEXES,INPATH,INPUT,INPUTFORMAT,INVOKER,ITEMS,ITERATE,JSON,JSON_EXISTS,JSON_QUERY,JSON_TABLE,JSON_VALUE,KEEP,KEY,KEYS,LAST,LAZY,LEAVE,LEVEL,LIMIT,LINES,LIST,LOAD,LOCATION,LOCK,LOCKS,LOGICAL,LONG,LOOP,MACRO,MAP,MATCHED,MATCH_CONDITION,MATERIALIZED,MEASURE,METRICS,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTES,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NEAREST,NORELY,NULLS,OBJECT,OFFSET,OMIT,OPTION,OPTIONS,ORDINALITY,OUTPUTFORMAT,OVERWRITE,PARTITIONED,PARTITIONS,PATH,PERCENT,PIVOT,PLACING,PRECEDING,PRINCIPALS,PROCEDURES,PROPERTIES,PURGE,QUALIFY,QUARTER,QUERY,QUOTES,RECORDREADER,RECORDWRITER,RECOVER,RECURSION,REDUCE,REFRESH,RELY,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURNING,ROLE,ROLES,SCD,SCHEMA,SCHEMAS,SECONDS,SECURITY,SEMI,SEPARATED,SEQUENCE,SERDE,SERDEPROPERTIES,SETS,SHORT,SHOW,SIMILARITY,SINGLE,SKEWED,SORT,SORTED,SOURCE,STATISTICS,STORED,STRATIFY,STREAM,STREAMING,STRING,STRUCT,SUBSTR,SYNC,SYSTEM_PATH,SYSTEM_TIME,SYSTEM_VERSION,TABLES,TARGET,TBLPROPERTIES,TERMINATED,TIMEDIFF,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TOUCH,TRACK,TRANSACTION,TRANSACTIONS,TRANSFORM,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNCONDITIONAL,UNIFORM,UNLOCK,UNPIVOT,UNSET,UNTIL,USE,VAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WATERMARK,WEEK,WEEKS,WHILE,WIDTH,WRAPPER,X,YEARS,ZONE") // scalastyle:on line.size.limit } } @@ -886,17 +886,19 @@ class SparkConnectDatabaseMetaDataSuite extends ConnectFunSuite with RemoteSpark Option(rs.getObject("NUM_PREC_RADIX")).map(_.asInstanceOf[Integer].toInt)) }.toSeq - // results are ordered by DATA_TYPE + // results are ordered by DATA_TYPE, then TYPE_NAME assert(types.map(t => (t.name, t.dataType)) === Seq( ("TINYINT", Types.TINYINT), ("BIGINT", Types.BIGINT), ("BINARY", Types.VARBINARY), + ("CHAR", Types.CHAR), ("DECIMAL", Types.DECIMAL), ("INT", Types.INTEGER), ("SMALLINT", Types.SMALLINT), ("FLOAT", Types.FLOAT), ("DOUBLE", Types.DOUBLE), ("STRING", Types.VARCHAR), + ("VARCHAR", Types.VARCHAR), ("BOOLEAN", Types.BOOLEAN), ("DATE", Types.DATE), ("TIMESTAMP", Types.TIMESTAMP))) @@ -904,12 +906,15 @@ class SparkConnectDatabaseMetaDataSuite extends ConnectFunSuite with RemoteSpark // every type is nullable and searchable assert(types.forall(_.nullable == DatabaseMetaData.typeNullable)) assert(types.forall(_.searchable == DatabaseMetaData.typeSearchable)) - // only STRING is case-sensitive - assert(types.filter(_.caseSensitive).map(_.name) === Seq("STRING")) + // CHAR, VARCHAR and STRING are case-sensitive + assert(types.filter(_.caseSensitive).map(_.name) === + Seq("CHAR", "STRING", "VARCHAR")) // string-like types are quoted with a single quote on both sides, except BINARY, // whose literals use the hex syntax X'...'. Numeric types carry no literal quote. val quoted = Map( + "CHAR" -> ("'", "'"), + "VARCHAR" -> ("'", "'"), "STRING" -> ("'", "'"), "DATE" -> ("'", "'"), "TIMESTAMP" -> ("'", "'"), @@ -923,14 +928,18 @@ class SparkConnectDatabaseMetaDataSuite extends ConnectFunSuite with RemoteSpark // PRECISION mirrors JdbcTypeUtils.getPrecision for every type val precisions = Map( "BOOLEAN" -> 1, "TINYINT" -> 3, "SMALLINT" -> 5, "INT" -> 10, "BIGINT" -> 19, - "FLOAT" -> 7, "DOUBLE" -> 15, "DECIMAL" -> 38, "STRING" -> Int.MaxValue, + "FLOAT" -> 7, "DOUBLE" -> 15, "DECIMAL" -> 38, + "CHAR" -> Int.MaxValue, "VARCHAR" -> Int.MaxValue, "STRING" -> Int.MaxValue, "BINARY" -> Int.MaxValue, "DATE" -> 10, "TIMESTAMP" -> 29) types.foreach { t => assert(t.precision === precisions(t.name), s"unexpected PRECISION for ${t.name}") } // CREATE_PARAMS is set only for the parameterized types - val createParams = Map("DECIMAL" -> "precision,scale") + val createParams = Map( + "DECIMAL" -> "precision,scale", + "CHAR" -> "length", + "VARCHAR" -> "length") types.foreach { t => assert(t.createParams === createParams.getOrElse(t.name, null), s"unexpected CREATE_PARAMS for ${t.name}") diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectJdbcDataTypeSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectJdbcDataTypeSuite.scala index 4ad8ab91d4b2f..59507c816b306 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectJdbcDataTypeSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectJdbcDataTypeSuite.scala @@ -223,6 +223,66 @@ class SparkConnectJdbcDataTypeSuite extends ConnectFunSuite with RemoteSparkSess } } + test("SPARK-58794: get char and varchar types") { + withStatement { statement => + statement.execute("SET spark.sql.charVarchar.standardSemantics.enabled=true") + withExecuteQuery(statement, + "SELECT CAST('ab' AS CHAR(4)) AS c, CAST('cd' AS VARCHAR(6)) AS v") { rs => + assert(rs.next()) + assert(rs.getString(1) === "ab ") + assert(rs.getString(2) === "cd") + assert(!rs.next()) + + val metaData = rs.getMetaData + assert(metaData.getColumnType(1) === Types.CHAR) + assert(metaData.getColumnTypeName(1) === "CHAR(4)") + assert(metaData.getColumnClassName(1) === "java.lang.String") + assert(metaData.isSigned(1) === false) + assert(metaData.getPrecision(1) === 4) + assert(metaData.getScale(1) === 0) + assert(metaData.getColumnDisplaySize(1) === 4) + + assert(metaData.getColumnType(2) === Types.VARCHAR) + assert(metaData.getColumnTypeName(2) === "VARCHAR(6)") + assert(metaData.getColumnClassName(2) === "java.lang.String") + assert(metaData.isSigned(2) === false) + assert(metaData.getPrecision(2) === 6) + assert(metaData.getScale(2) === 0) + assert(metaData.getColumnDisplaySize(2) === 6) + } + } + } + + test("SPARK-58794: getColumns for CHAR/VARCHAR table") { + val table = "char_varchar_jdbc_cols" + withStatement { statement => + statement.execute("SET spark.sql.charVarchar.standardSemantics.enabled=true") + statement.execute(s"DROP TABLE IF EXISTS $table") + statement.execute( + s"CREATE TABLE $table (c CHAR(4), v VARCHAR(6)) USING parquet") + try { + Using.resource( + statement.getConnection.getMetaData.getColumns(null, null, table, null)) { rs => + assert(rs.next()) + assert(rs.getString("COLUMN_NAME") === "c") + assert(rs.getInt("DATA_TYPE") === Types.CHAR) + assert(rs.getString("TYPE_NAME") === "CHAR(4)") + assert(rs.getInt("COLUMN_SIZE") === 4) + assert(rs.getInt("CHAR_OCTET_LENGTH") === 16) + assert(rs.next()) + assert(rs.getString("COLUMN_NAME") === "v") + assert(rs.getInt("DATA_TYPE") === Types.VARCHAR) + assert(rs.getString("TYPE_NAME") === "VARCHAR(6)") + assert(rs.getInt("COLUMN_SIZE") === 6) + assert(rs.getInt("CHAR_OCTET_LENGTH") === 24) + assert(!rs.next()) + } + } finally { + statement.execute(s"DROP TABLE IF EXISTS $table") + } + } + } + test("get decimal type") { withStatement { stmt => Seq( diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSetSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSetSuite.scala index 21b8e261aef45..8470a91515331 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSetSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectResultSetSuite.scala @@ -186,4 +186,18 @@ class SparkConnectResultSetSuite extends ConnectFunSuite with RemoteSparkSession assert(!rs.next()) } } + + test("next throws SQLException when the query fails during result streaming") { + // The server sends the schema response before any arrow batch, so executeQuery + // returns successfully and the runtime error surfaces while iterating rows. + withStatement { stmt => + val rs = stmt.executeQuery("SELECT 10 / (5 - id) FROM range(10)") + val e = intercept[SQLException] { + while (rs.next()) {} + } + // DIVIDE_BY_ZERO under ANSI mode + assert(e.getSQLState === "22012") + assert(e.getCause != null) + } + } } diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatementSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatementSuite.scala index 6ace1fa9cf785..1163feaf5e90b 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatementSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectStatementSuite.scala @@ -173,4 +173,17 @@ class SparkConnectStatementSuite extends ConnectFunSuite with RemoteSparkSession } } } + + test("execute throws SQLException with server-provided SQLState on query failure") { + withStatement { stmt => + val e = intercept[SQLException] { + stmt.execute("SELECT * FROM this_table_does_not_exist") + } + // TABLE_OR_VIEW_NOT_FOUND + assert(e.getSQLState === "42P01") + assert(!e.isInstanceOf[SQLNonTransientConnectionException]) + assert(!e.isInstanceOf[SQLTransientConnectionException]) + assert(e.getCause != null) + } + } } diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtilsSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtilsSuite.scala new file mode 100644 index 0000000000000..e64dd1f8681e6 --- /dev/null +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/util/JdbcErrorUtilsSuite.scala @@ -0,0 +1,166 @@ +/* + * 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.connect.client.jdbc.util + +import java.sql.{SQLException, SQLNonTransientConnectionException, SQLTimeoutException, SQLTransientConnectionException} + +import io.grpc.{Status, StatusRuntimeException} + +import org.apache.spark.SparkThrowable +import org.apache.spark.sql.connect.test.ConnectFunSuite + +/** + * Tests for [[JdbcErrorUtils.toSQLException]]. A RuntimeException mixing in + * [[SparkThrowable]] stands in for converted server errors, carrying the sqlState + * GrpcExceptionConverter would supply. A real [[StatusRuntimeException]] cause + * models transport errors the way GrpcExceptionConverter preserves them. + */ +class JdbcErrorUtilsSuite extends ConnectFunSuite { + + private def sparkError( + condition: String, + sqlState: String, + msg: String, + cause: Throwable = null): RuntimeException = + new RuntimeException(msg, cause) with SparkThrowable { + override def getCondition: String = condition + override def getSqlState: String = sqlState + } + + private def grpcError(status: Status): StatusRuntimeException = + new StatusRuntimeException(status) + + test("a class-08 SQLSTATE maps to a non-transient connection exception") { + val e = JdbcErrorUtils.toSQLException( + sparkError("INVALID_HANDLE.SESSION_CLOSED", "08003", "Session was closed")) + assert(e.isInstanceOf[SQLNonTransientConnectionException]) + assert(e.getSQLState === "08003") + assert(e.getMessage === "Session was closed") + } + + test("the class-08 mapping does not depend on the condition name") { + val e = JdbcErrorUtils.toSQLException(sparkError("SOME_FUTURE_CONDITION", "08004", "gone")) + assert(e.isInstanceOf[SQLNonTransientConnectionException]) + assert(e.getSQLState === "08004") + } + + test("operation-level INVALID_HANDLE subconditions do not map to a connection exception") { + // the session, and thus the connection, is still healthy: no SQLState class 08 + Seq( + "INVALID_HANDLE.OPERATION_NOT_FOUND", + "INVALID_HANDLE.OPERATION_ABANDONED", + "INVALID_HANDLE.OPERATION_ALREADY_EXISTS", + "INVALID_HANDLE.FORMAT").foreach { condition => + val e = JdbcErrorUtils.toSQLException(sparkError(condition, "HY000", "operation gone")) + assert(!e.isInstanceOf[SQLNonTransientConnectionException], condition) + assert(!e.isInstanceOf[SQLTransientConnectionException], condition) + } + } + + test("a non-connection Spark error keeps the server-provided SQLState") { + val e = JdbcErrorUtils.toSQLException(sparkError("DIVIDE_BY_ZERO", "22012", "boom")) + assert(!e.isInstanceOf[SQLNonTransientConnectionException]) + assert(!e.isInstanceOf[SQLTransientConnectionException]) + assert(e.getSQLState === "22012") + assert(e.getMessage === "boom") + } + + test("an UNAVAILABLE StatusRuntimeException maps to a transient connection exception") { + val e = JdbcErrorUtils.toSQLException( + grpcError(Status.UNAVAILABLE.withDescription("Channel shutdown invoked"))) + assert(e.isInstanceOf[SQLTransientConnectionException]) + assert(e.getSQLState === "08006") + } + + test("an UNAVAILABLE cause preserved by GrpcExceptionConverter maps to 08006") { + // the shape GrpcExceptionConverter produces for transport errors + val wrapped = sparkError( + "CONNECT_CLIENT_UNEXPECTED_MISSING_SQL_STATE", + "XXKCM", + "io.grpc.StatusRuntimeException: UNAVAILABLE: io exception", + cause = grpcError(Status.UNAVAILABLE.withDescription("io exception"))) + val e = JdbcErrorUtils.toSQLException(wrapped) + assert(e.isInstanceOf[SQLTransientConnectionException]) + assert(e.getSQLState === "08006") + } + + test("a server error merely quoting a gRPC exception is not a connection failure") { + // the gRPC status comes from exception instances, never from message text, so a + // server error embedding gRPC text (e.g. from a UDF) is not a transport failure + val e = JdbcErrorUtils.toSQLException(sparkError( + "FAILED_EXECUTE_UDF", + "39000", + "Job aborted: io.grpc.StatusRuntimeException: UNAVAILABLE: backend down")) + assert(!e.isInstanceOf[SQLTransientConnectionException]) + assert(!e.isInstanceOf[SQLNonTransientConnectionException]) + } + + test("a DEADLINE_EXCEEDED cause maps to a timeout, not a connection failure") { + // a slow query fires the RPC deadline on a healthy connection: timeout, not class 08 + val wrapped = sparkError( + "CONNECT_CLIENT_UNEXPECTED_MISSING_SQL_STATE", + "XXKCM", + "io.grpc.StatusRuntimeException: DEADLINE_EXCEEDED: deadline exceeded", + cause = grpcError(Status.DEADLINE_EXCEEDED.withDescription("deadline exceeded"))) + val e = JdbcErrorUtils.toSQLException(wrapped) + assert(e.isInstanceOf[SQLTimeoutException]) + assert(e.getSQLState === "HYT00") + assert(!e.isInstanceOf[SQLTransientConnectionException]) + } + + test("a non-retryable gRPC status does not map to a connection exception") { + val e = JdbcErrorUtils.toSQLException( + grpcError(Status.INVALID_ARGUMENT.withDescription("bad plan"))) + assert(!e.isInstanceOf[SQLTransientConnectionException]) + assert(!e.isInstanceOf[SQLNonTransientConnectionException]) + assert(!e.isInstanceOf[SQLTimeoutException]) + } + + test("a class-08 error takes precedence over a transport code") { + // both signals present: the gone session wins + val e = JdbcErrorUtils.toSQLException(sparkError( + "INVALID_HANDLE.SESSION_CLOSED", + "08003", + "closed", + cause = grpcError(Status.UNAVAILABLE.withDescription("x")))) + assert(e.isInstanceOf[SQLNonTransientConnectionException]) + assert(e.getSQLState === "08003") + } + + test("the connection error is found through the cause chain") { + val root = sparkError("INVALID_HANDLE.SESSION_CLOSED", "08003", "closed") + val e = JdbcErrorUtils.toSQLException(new RuntimeException("wrapper", root)) + assert(e.isInstanceOf[SQLNonTransientConnectionException]) + assert(e.getSQLState === "08003") + } + + test("a non-Spark exception becomes a plain SQLException carrying the cause") { + val cause = new RuntimeException("raw failure") + val e = JdbcErrorUtils.toSQLException(cause) + assert(!e.isInstanceOf[SQLNonTransientConnectionException]) + assert(!e.isInstanceOf[SQLTransientConnectionException]) + assert(e.getSQLState === null) + assert(e.getMessage === "raw failure") + assert(e.getCause eq cause) + } + + test("an existing SQLException passes through unchanged") { + val original = new SQLException("already mapped") + assert(JdbcErrorUtils.toSQLException(original) eq original) + } +} diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala index 266de35706da3..b117f1a0b5cd0 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala @@ -106,6 +106,8 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { private val printer = JsonFormat.printer().usingTypeRegistry(registry) + private val testUDT = new TestUDT.NewArrayUDT() + private var session: SparkSession = _ override protected def beforeAll(): Unit = { @@ -1623,6 +1625,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.round(fn.col("b"), 2) } + functionTest("truncate") { + fn.truncate(fn.col("b"), 2) + } + functionTest("sec") { fn.sec(fn.col("b")) } @@ -1711,6 +1717,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.crc32(fn.col("g").cast("binary")) } + functionTest("xxh3_64") { + fn.xxh3_64(fn.col("g").cast("binary")) + } + + functionTest("xxh3_128") { + fn.xxh3_128(fn.col("g").cast("binary")) + } + functionTest("hash") { fn.hash(fn.col("b"), fn.col("id")) } @@ -1855,6 +1869,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.unbase64(fn.col("g")) } + functionTest("to_base32") { + fn.to_base32(fn.col("g").cast("binary")) + } + + functionTest("from_base32") { + fn.from_base32(fn.col("g")) + } + functionTest("rpad") { fn.rpad(fn.col("g"), 10, "-") } @@ -2246,6 +2268,26 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { binary.select(fn.bitmap_or_agg(fn.col("bytes"))) } + test("function bitmap_and") { + binary.select(fn.bitmap_and(fn.col("bytes"), fn.col("bytes"))) + } + + test("function bitmap_or") { + binary.select(fn.bitmap_or(fn.col("bytes"), fn.col("bytes"))) + } + + test("function bitmap_andnot") { + binary.select(fn.bitmap_andnot(fn.col("bytes"), fn.col("bytes"))) + } + + test("function bitmap_xor") { + binary.select(fn.bitmap_xor(fn.col("bytes"), fn.col("bytes"))) + } + + test("function bitmap_xor_agg") { + binary.select(fn.bitmap_xor_agg(fn.col("bytes"))) + } + private def temporalFunctionTest(name: String)(f: => Column): Unit = { test("function " + name) { temporals.select(f) @@ -2569,6 +2611,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.slice(fn.col("e"), 0, 5) } + functionTest("trim_array") { + fn.trim_array(fn.col("e"), 2) + } + functionTest("array_join") { fn.array_join(fn.col("e"), ";") } @@ -2795,6 +2841,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.try_variant_array_append(fn.parse_json(fn.col("g")), "$.a", fn.lit(1)) } + functionTest("variant_strip_nulls") { + fn.variant_strip_nulls(fn.parse_json(fn.col("g")), false) + } + functionTest("variant_get") { fn.variant_get(fn.parse_json(fn.col("g")), "$", "int") } @@ -2807,6 +2857,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.schema_of_variant(fn.parse_json(fn.col("g"))) } + functionTest("variant_from_arrays") { + fn.variant_from_arrays(fn.array(lit("a"), lit("b")), fn.array(lit(1), lit(2))) + } + + functionTest("variant_from_entries") { + fn.variant_from_entries(fn.array(fn.struct(lit("a"), lit(1)), fn.struct(lit("b"), lit(2)))) + } + functionTest("schema_of_variant_agg") { fn.schema_of_variant_agg(fn.parse_json(fn.col("g"))) } @@ -3129,6 +3187,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.json_object_keys(fn.col("g")) } + functionTest("json_typeof") { + fn.json_typeof(fn.col("g")) + } + functionTest("mask with specific upperChar lowerChar digitChar otherChar") { fn.mask(fn.col("g"), fn.lit('X'), fn.lit('x'), fn.lit('n'), fn.lit('*')) } @@ -3240,6 +3302,14 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging { fn.typeof(fn.col("g")) } + functionTest("wrap_udt") { + fn.wrap_udt(fn.array(fn.col("b")), testUDT) + } + + functionTest("unwrap_udt") { + fn.unwrap_udt(fn.wrap_udt(fn.array(fn.col("b")), testUDT)) + } + functionTest("stack") { fn.stack(lit(2), fn.col("g"), fn.col("g"), fn.col("g")) } diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala new file mode 100644 index 0000000000000..77704b8a57c30 --- /dev/null +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/TestUDT.scala @@ -0,0 +1,42 @@ +/* + * 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 + +import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType, SQLUserDefinedType, UserDefinedType} + +object TestUDT { + + @SQLUserDefinedType(udt = classOf[NewArrayUDT]) + private[sql] class NewArray(val values: Array[Double]) extends Serializable + + private[sql] class NewArrayUDT extends UserDefinedType[NewArray] { + + override def sqlType: DataType = ArrayType(DoubleType, containsNull = false) + + override def serialize(obj: NewArray): Any = obj.values + + override def deserialize(datum: Any): NewArray = { + datum match { + case values: Array[_] => + new NewArray(values.map(_.asInstanceOf[Double])) + } + } + + override def userClass: Class[NewArray] = classOf[NewArray] + } +} diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/FunctionTestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/FunctionTestSuite.scala index 8c5fc2b2b8ec9..ec191e4707f24 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/FunctionTestSuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/FunctionTestSuite.scala @@ -180,6 +180,7 @@ class FunctionTestSuite extends ConnectFunSuite { window(a, "10 seconds")) testEquals("session_window", session_window(a, "1 second"), session_window(a, lit("1 second"))) testEquals("slice", slice(a, 1, 2), slice(a, lit(1), lit(2))) + testEquals("trim_array", trim_array(a, 2), trim_array(a, lit(2))) testEquals("bucket", bucket(lit(3), a), bucket(3, a)) testEquals( "lag", diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionE2ESuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionE2ESuite.scala index 9c149a858018a..d8061386527bf 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionE2ESuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/SparkSessionE2ESuite.scala @@ -433,6 +433,7 @@ class SparkSessionE2ESuite extends ConnectFunSuite with RemoteSparkSession { } assert(e.getMessage.contains("[INVALID_HANDLE.SESSION_CHANGED]")) + assert(e.getSqlState == "08003") assert(!session1.client.isSessionValid) assert(SparkSession.getActiveSession.isEmpty) assert(SparkSession.getDefaultSession.isEmpty) diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala index 24e637a1f2153..6a456a554c5c5 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala @@ -45,8 +45,11 @@ class SparkConnectClientSuite extends ConnectFunSuite { private var service: DummySparkConnectService = _ private var server: Server = _ - private def startDummyServer(port: Int, interceptors: Seq[ServerInterceptor] = Seq()): Unit = { - service = new DummySparkConnectService + private def startDummyServer( + port: Int, + interceptors: Seq[ServerInterceptor] = Seq(), + dummyService: DummySparkConnectService = new DummySparkConnectService): Unit = { + service = dummyService val serverBuilder = NettyServerBuilder .forPort(port) .addService(service) @@ -78,6 +81,111 @@ class SparkConnectClientSuite extends ConnectFunSuite { assert(client.userId == System.getProperty("user.name")) } + Seq(false, true).foreach { reattachable => + test(s"client generates an operation ID for ExecutePlan requests ($reattachable)") { + val operationIdHeaders = mutable.Map.empty[String, Option[String]] + val interceptor = new ServerInterceptor { + override def interceptCall[ReqT, RespT]( + call: ServerCall[ReqT, RespT], + headers: Metadata, + next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = { + val key = Metadata.Key.of( + SparkConnectClient.OPERATION_ID_HEADER, + Metadata.ASCII_STRING_MARSHALLER) + operationIdHeaders.synchronized { + operationIdHeaders(call.getMethodDescriptor.getBareMethodName) = + Option(headers.get(key)) + } + next.startCall(call, headers) + } + } + val dummyService = if (reattachable) { + new DummySparkConnectService { + override def executePlan( + request: ExecutePlanRequest, + responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { + responseObserver.onNext( + ExecutePlanResponse + .newBuilder() + .setSessionId(request.getSessionId) + .setOperationId(request.getOperationId) + .setResponseId("initial-response") + .build()) + responseObserver.onCompleted() + } + + override def reattachExecute( + request: proto.ReattachExecuteRequest, + responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { + responseObserver.onNext( + ExecutePlanResponse + .newBuilder() + .setSessionId(request.getSessionId) + .setOperationId(request.getOperationId) + .setResponseId("result-complete") + .setResultComplete(proto.ExecutePlanResponse.ResultComplete.newBuilder().build()) + .build()) + responseObserver.onCompleted() + } + } + } else { + new DummySparkConnectService + } + startDummyServer(0, Seq(interceptor), dummyService) + val builder = SparkConnectClient + .builder() + .connectionString(s"sc://localhost:${server.getPort}") + .option(SparkConnectClient.OPERATION_ID_HEADER, "ignored") + if (reattachable) builder.enableReattachableExecute() + else builder.disableReattachableExecute() + client = builder.build() + + val responses = client.execute(buildPlan("select 1")).toSeq + val operationId = responses.head.getOperationId + + UUID.fromString(operationId) + assert(responses.forall(_.getOperationId == operationId)) + assert(operationIdHeaders.synchronized { + operationIdHeaders("ExecutePlan").contains(operationId) + }) + if (reattachable) { + assert(operationIdHeaders.synchronized { + operationIdHeaders("ReattachExecute").contains(operationId) + }) + Eventually.eventually(timeout(5.seconds)) { + assert(operationIdHeaders.synchronized { + operationIdHeaders("ReleaseExecute").contains(operationId) + }) + } + } + } + } + + test("ExecutePlan exceptions expose the client-generated operation ID") { + val failingService = new DummySparkConnectService { + override def executePlan( + request: ExecutePlanRequest, + responseObserver: StreamObserver[ExecutePlanResponse]): Unit = { + responseObserver.onError(Status.INTERNAL.withDescription("expected").asRuntimeException()) + } + } + server = NettyServerBuilder.forPort(0).addService(failingService).build().start() + service = failingService + client = SparkConnectClient + .builder() + .connectionString(s"sc://localhost:${server.getPort}") + .disableReattachableExecute() + .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false, name = "NoRetry")) + .build() + + val error = intercept[SparkException] { + client.execute(buildPlan("select 1")).foreach(_ => ()) + } + val operationId = SparkConnectClient.getOperationId(error) + assert(operationId.isDefined) + UUID.fromString(operationId.get) + } + test("Placeholder test: Create SparkConnectClient") { client = SparkConnectClient.builder().userId("abc123").build() assert(client.userId == "abc123") @@ -898,6 +1006,34 @@ class SparkConnectClientSuite extends ConnectFunSuite { } } + test("transport errors preserve the original gRPC status exception as the cause") { + val failingService = new DummySparkConnectService { + override def analyzePlan( + request: AnalyzePlanRequest, + responseObserver: StreamObserver[AnalyzePlanResponse]): Unit = { + responseObserver.onError( + Status.UNAVAILABLE.withDescription("injected failure").asRuntimeException()) + } + } + server = NettyServerBuilder.forPort(0).addService(failingService).build().start() + service = failingService + client = SparkConnectClient + .builder() + .connectionString(s"sc://localhost:${server.getPort}") + .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false, name = "NoRetry")) + .build() + + val ex = intercept[SparkException] { + client.analyze(proto.AnalyzePlanRequest.newBuilder().setSessionId("abc123").build()) + } + // The original StatusRuntimeException must survive as the cause so that callers + // can programmatically inspect the gRPC status code instead of parsing the message. + assert(ex.getCause.isInstanceOf[StatusRuntimeException]) + assert( + ex.getCause.asInstanceOf[StatusRuntimeException].getStatus.getCode == + Status.Code.UNAVAILABLE) + } + test( "SPARK-58094: gRPC keepalive surfaces a bounded failure on a silently dropped " + "connection instead of hanging forever") { @@ -963,9 +1099,8 @@ class SparkConnectClientSuite extends ConnectFunSuite { scala.concurrent.Await.result(resultPromise.future, FiniteDuration(15, TimeUnit.SECONDS)) } // scalastyle:on awaitresult - // A keepalive-triggered UNAVAILABLE carries no wrapped cause (same as DEADLINE_EXCEEDED, - // see GrpcExceptionConverter.toThrowable), so the status code/description is only in the - // message. + // The status code/description of a keepalive-triggered UNAVAILABLE is part of the + // message (see GrpcExceptionConverter.toThrowable). assert(ex.getMessage.contains("UNAVAILABLE")) assert(ex.getMessage.contains("Keepalive failed")) } finally { diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala index efb6a7a2304f0..9963ec5620ff2 100644 --- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala +++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/arrow/ArrowEncoderSuite.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.{Encoders, Row} import org.apache.spark.sql.catalyst.{DefinedByConstructorParams, JavaTypeInference, ScalaReflection} import org.apache.spark.sql.catalyst.encoders.{AgnosticEncoder, Codec, OuterScopes} import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{agnosticEncoderFor, BinaryEncoder, BoxedBooleanEncoder, BoxedByteEncoder, BoxedDoubleEncoder, BoxedFloatEncoder, BoxedIntEncoder, BoxedLongEncoder, BoxedShortEncoder, CalendarIntervalEncoder, DateEncoder, DayTimeIntervalEncoder, EncoderField, InstantEncoder, IterableEncoder, JavaDecimalEncoder, LocalDateEncoder, LocalDateTimeEncoder, NullEncoder, PrimitiveBooleanEncoder, PrimitiveByteEncoder, PrimitiveDoubleEncoder, PrimitiveFloatEncoder, PrimitiveIntEncoder, PrimitiveLongEncoder, PrimitiveShortEncoder, RowEncoder, ScalaDecimalEncoder, StringEncoder, TimestampEncoder, TransformingEncoder, UDTEncoder, YearMonthIntervalEncoder} -import org.apache.spark.sql.catalyst.encoders.RowEncoder.{encoderFor => toRowEncoder} +import org.apache.spark.sql.catalyst.encoders.RowEncoder.{encoderFor => toRowEncoder, encoderForResultSchema => toResultRowEncoder} import org.apache.spark.sql.catalyst.util.{DateFormatter, TimestampFormatter} import org.apache.spark.sql.catalyst.util.DateTimeConstants.MICROS_PER_SECOND import org.apache.spark.sql.catalyst.util.IntervalStringStyles.ANSI_STYLE @@ -48,7 +48,7 @@ import org.apache.spark.sql.catalyst.util.SparkDateTimeUtils._ import org.apache.spark.sql.catalyst.util.SparkIntervalUtils._ import org.apache.spark.sql.connect.client.arrow.FooEnum.FooEnum import org.apache.spark.sql.connect.test.ConnectFunSuite -import org.apache.spark.sql.types.{ArrayType, DataType, DayTimeIntervalType, Decimal, DecimalType, Geography, Geometry, IntegerType, Metadata, SQLUserDefinedType, StringType, StructType, UserDefinedType, YearMonthIntervalType} +import org.apache.spark.sql.types.{ArrayType, CharType, DataType, DayTimeIntervalType, Decimal, DecimalType, Geography, Geometry, IntegerType, Metadata, SQLUserDefinedType, StringType, StructType, UserDefinedType, VarcharType, YearMonthIntervalType} import org.apache.spark.sql.util.CloseableIterator import org.apache.spark.unsafe.types.VariantVal import org.apache.spark.util.{MaybeNull, SparkStringUtils} @@ -309,6 +309,28 @@ class ArrowEncoderSuite extends ConnectFunSuite { } } + test("SPARK-58794: char/varchar round trip") { + // The client cannot see the server's charVarchar configuration, so a result schema carrying + // CHAR/VARCHAR must be decodable regardless of the local one. Values are padded and length + // checked by the server, so the client passes them through unchanged. + val encoder = toResultRowEncoder( + new StructType() + .add("c", CharType(4)) + .add("v", VarcharType(6)) + .add("s", new StructType().add("c", CharType(4))) + .add("a", ArrayType(VarcharType(6)))) + roundTripAndCheckIdentical(encoder) { () => + val maybeNull = MaybeNull(7) + Iterator.tabulate(101) { i => + Row( + maybeNull("ab "), + maybeNull("cd"), + maybeNull(Row("ef ")), + maybeNull(mutable.ArraySeq.make[String](Array("gh")))) + } + } + } + test("single batch") { val inspector = new CountingBatchInspector roundTripAndCheckIdentical(singleIntEncoder, inspectBatch = inspector) { () => diff --git a/sql/connect/common/src/main/protobuf/spark/connect/base.proto b/sql/connect/common/src/main/protobuf/spark/connect/base.proto index c7247129f1907..be4ab641c7b5b 100644 --- a/sql/connect/common/src/main/protobuf/spark/connect/base.proto +++ b/sql/connect/common/src/main/protobuf/spark/connect/base.proto @@ -253,6 +253,9 @@ message AnalyzePlanResponse { JsonToDDL json_to_ddl = 16; } + // Support arbitrary result objects. + repeated google.protobuf.Any extensions = 999; + message Schema { DataType schema = 1; } diff --git a/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto b/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto index 18f8f294f0c02..432bc13e918dc 100644 --- a/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto +++ b/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto @@ -475,6 +475,9 @@ message PythonUDF { string python_ver = 4; // (Optional) Additional includes for the Python UDF. repeated string additional_includes = 5; + // (Optional) Intermediate buffer schema for an incremental Python aggregator + // (see PythonAggregate). Set only for the incremental aggregator eval types. + optional DataType buffer_type = 6; } message ScalarScalaUDF { diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala index a4406c2a68fda..688029f4883ad 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala @@ -45,14 +45,27 @@ private[connect] class CustomSparkConnectBlockingStub( // GrpcExceptionConverter with a GRPC stub for fetching error details from server. private val grpcExceptionConverter = stubState.exceptionConverter + private def executePlanStub(operationId: String) = { + Option(operationId) + .filter(_.nonEmpty) + .map { id => + stub.withInterceptors( + new SparkConnectClient.MetadataHeaderClientInterceptor( + Map(SparkConnectClient.OPERATION_ID_HEADER -> id))) + } + .getOrElse(stub) + } + // Non-reattachable executePlan intentionally has no deadline: a timeout here would kill the // server-side execution with no way to recover (there is no ReattachExecute for this path). // Use reattachable execution for long-running queries that need deadline protection. def executePlan(request: ExecutePlanRequest): CloseableIterator[ExecutePlanResponse] = { + val stubWithOperationId = executePlanStub(request.getOperationId) grpcExceptionConverter.convert( request.getSessionId, request.getUserContext, - request.getClientType) { + request.getClientType, + Option(request.getOperationId).filter(_.nonEmpty)) { grpcExceptionConverter.convertIterator[ExecutePlanResponse]( request.getSessionId, request.getUserContext, @@ -61,8 +74,9 @@ private[connect] class CustomSparkConnectBlockingStub( request, r => { stubState.responseValidator.wrapIterator( - CloseableIterator(stub.executePlan(r).asScala)) - })) + CloseableIterator(stubWithOperationId.executePlan(r).asScala)) + }), + Option(request.getOperationId).filter(_.nonEmpty)) } } @@ -71,7 +85,8 @@ private[connect] class CustomSparkConnectBlockingStub( grpcExceptionConverter.convert( request.getSessionId, request.getUserContext, - request.getClientType) { + request.getClientType, + Option(request.getOperationId).filter(_.nonEmpty)) { grpcExceptionConverter.convertIterator[ExecutePlanResponse]( request.getSessionId, request.getUserContext, @@ -83,7 +98,8 @@ private[connect] class CustomSparkConnectBlockingStub( channel, stubState.retryHandler, stubState.rpcDeadlines.reattachableExecutePlan, - stubState.rpcDeadlines.reattachExecute))) + stubState.rpcDeadlines.reattachExecute)), + Option(request.getOperationId).filter(_.nonEmpty)) } } diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/ExecutePlanResponseReattachableIterator.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/ExecutePlanResponseReattachableIterator.scala index e6428270c20ba..c7073c96aba64 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/ExecutePlanResponseReattachableIterator.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/ExecutePlanResponseReattachableIterator.scala @@ -81,8 +81,14 @@ class ExecutePlanResponseReattachableIterator( // - this does it's own custom retry handling // - error conversion is wrapped around this in CustomSparkConnectBlockingStub, // this needs raw GRPC errors for retries. - private val rawBlockingStub = proto.SparkConnectServiceGrpc.newBlockingStub(channel) - private val rawAsyncStub = proto.SparkConnectServiceGrpc.newStub(channel) + private val operationIdInterceptor = new SparkConnectClient.MetadataHeaderClientInterceptor( + Map(SparkConnectClient.OPERATION_ID_HEADER -> operationId)) + private val rawBlockingStub = + proto.SparkConnectServiceGrpc + .newBlockingStub(channel) + .withInterceptors(operationIdInterceptor) + private val rawAsyncStub = + proto.SparkConnectServiceGrpc.newStub(channel).withInterceptors(operationIdInterceptor) private def stubWithDeadline(deadline: Option[FiniteDuration]) : proto.SparkConnectServiceGrpc.SparkConnectServiceBlockingStub = diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala index b52261e5efde5..3b6c16f06e4da 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala @@ -64,12 +64,18 @@ private[client] class GrpcExceptionConverter( .map(d => grpcStub.withDeadline(Deadline.after(d.toMillis, TimeUnit.MILLISECONDS))) .getOrElse(grpcStub) - def convert[T](sessionId: String, userContext: UserContext, clientType: String)(f: => T): T = { + def convert[T]( + sessionId: String, + userContext: UserContext, + clientType: String, + operationId: Option[String] = None)(f: => T): T = { try { f } catch { case e: StatusRuntimeException => - throw toThrowable(e, sessionId, userContext, clientType) + val converted = toThrowable(e, sessionId, userContext, clientType) + operationId.foreach(SparkConnectClient.attachOperationId(converted, _)) + throw converted } } @@ -77,25 +83,26 @@ private[client] class GrpcExceptionConverter( sessionId: String, userContext: UserContext, clientType: String, - iter: CloseableIterator[T]): CloseableIterator[T] = { + iter: CloseableIterator[T], + operationId: Option[String] = None): CloseableIterator[T] = { new WrappedCloseableIterator[T] { override def innerIterator: Iterator[T] = iter override def hasNext: Boolean = { - convert(sessionId, userContext, clientType) { + convert(sessionId, userContext, clientType, operationId) { iter.hasNext } } override def next(): T = { - convert(sessionId, userContext, clientType) { + convert(sessionId, userContext, clientType, operationId) { iter.next() } } override def close(): Unit = { - convert(sessionId, userContext, clientType) { + convert(sessionId, userContext, clientType, operationId) { iter.close() } } @@ -170,23 +177,22 @@ private[client] class GrpcExceptionConverter( } // If no ErrorInfo is found, create a SparkException based on the StatusRuntimeException. - val (message, cause) = if (ex.getStatus.getCode == Status.Code.DEADLINE_EXCEEDED) { - val msg = s"${ex.toString}: RPC deadline exceeded. Deadlines can be configured via " + + val message = if (ex.getStatus.getCode == Status.Code.DEADLINE_EXCEEDED) { + s"${ex.toString}: RPC deadline exceeded. Deadlines can be configured via " + "SparkConnectClient.Builder.rpcDeadlines(). To disable all deadlines: " + "SparkConnectClient.builder().rpcDeadlines(RpcDeadlines.disabled).build()" - // For DEADLINE_EXCEEDED, we pass `ex` itself as the cause rather than `ex.getCause`. - // StatusRuntimeException.getCause() returns status.getCause(), which is always null for - // client-side deadline fires (gRPC constructs the status without a wrapped cause). Using - // ex.getCause would produce a SparkException with cause = null, losing the gRPC status - // code and description from the exception chain. Passing ex preserves full context and - // allows callers to programmatically inspect the status code via getCause().getStatus(). - (msg, ex) } else { - (ex.toString, ex.getCause) + ex.toString } + // Pass `ex` itself as the cause rather than `ex.getCause`. StatusRuntimeException.getCause() + // returns status.getCause(), which is often null (e.g. always for client-side deadline + // fires, since gRPC constructs the status without a wrapped cause). Using ex.getCause + // would produce a SparkException with cause = null, losing the gRPC status code and + // description from the exception chain. Passing ex preserves full context and allows + // callers to programmatically inspect the status code via getCause().getStatus(). new SparkException( message = message, - cause = cause, + cause = ex, errorClass = Some("CONNECT_CLIENT_UNEXPECTED_MISSING_SQL_STATE"), messageParameters = Map("message" -> message), context = Array.empty) diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala index eeb14fcfc2c72..5a3bd0b892c59 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala @@ -56,6 +56,14 @@ private[sql] class SparkConnectClient( private val userContext: UserContext = configuration.userContext private[this] val stubState = new SparkConnectStubState(channel, configuration) + + if (configuration.metadata.keys.exists( + _.equalsIgnoreCase(SparkConnectClient.OPERATION_ID_HEADER))) { + logWarning( + s"Connection option ${SparkConnectClient.OPERATION_ID_HEADER} is ignored because " + + "Spark Connect sets it for each ExecutePlan request.") + } + private[this] val bstub = new CustomSparkConnectBlockingStub(channel, stubState) private[this] val stub = @@ -318,13 +326,12 @@ private[sql] class SparkConnectClient( serverSideSessionId.foreach(session => request.setClientObservedServerSideSessionId(session)) - operationId.foreach { opId => - require( - isValidUUID(opId), - s"Invalid operationId: $opId. The id must be an UUID string of " + - "the format `00112233-4455-6677-8899-aabbccddeeff`") - request.setOperationId(opId) - } + val resolvedOperationId = operationId.getOrElse(UUID.randomUUID.toString) + require( + isValidUUID(resolvedOperationId), + s"Invalid operationId: $resolvedOperationId. The id must be an UUID string of " + + "the format `00112233-4455-6677-8899-aabbccddeeff`") + request.setOperationId(resolvedOperationId) if (configuration.useReattachableExecute) { bstub.executePlanReattachable(request.build()) } else { @@ -699,8 +706,31 @@ private[sql] class SparkConnectClient( // Options for plan compression case class PlanCompressionOptions(thresholdBytes: Int, algorithm: String) +private final class SparkConnectOperationIdException(val operationId: String) + extends RuntimeException(s"Spark Connect operation ID: $operationId", null, false, false) + object SparkConnectClient { + private[connect] val OPERATION_ID_HEADER = "spark-connect-operation-id" + + /** + * Returns the ExecutePlan operation ID attached to a Spark Connect failure, when available. + * + * @since 4.3.0 + */ + @DeveloperApi + def getOperationId(error: Throwable): Option[String] = { + error.getSuppressed.collectFirst { case marker: SparkConnectOperationIdException => + marker.operationId + } + } + + private[client] def attachOperationId(error: Throwable, operationId: String): Unit = { + if (getOperationId(error).isEmpty) { + error.addSuppressed(new SparkConnectOperationIdException(operationId)) + } + } + private[sql] val SPARK_REMOTE: String = "SPARK_REMOTE" private val DEFAULT_USER_AGENT: String = "_SPARK_CONNECT_SCALA" @@ -1158,9 +1188,11 @@ object SparkConnectClient { // Workaround LocalChannelCredentials are added in // https://github.com/grpc/grpc-java/issues/9900 - var metadataWithOptionalToken = metadata + var metadataWithOptionalToken = metadata.filterNot { case (key, _) => + key.equalsIgnoreCase(OPERATION_ID_HEADER) + } if (!isSslEnabled.contains(true) && isLocal && token.isDefined) { - metadataWithOptionalToken = metadata + (("Authorization", s"Bearer ${token.get}")) + metadataWithOptionalToken += (("Authorization", s"Bearer ${token.get}")) } if (metadataWithOptionalToken.nonEmpty) { @@ -1207,7 +1239,7 @@ object SparkConnectClient { applier.apply(headers) } catch { case e: Throwable => - applier.fail(Status.UNAUTHENTICATED.withCause(e)); + applier.fail(Status.UNAUTHENTICATED.withCause(e)) } }) } diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala index 375d18514fbf3..dd6e0ec474697 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkResult.scala @@ -101,7 +101,7 @@ private[sql] class SparkResult[T]( case UnboundRowEncoder => // Replace the row encoder with the encoder inferred from the schema. RowEncoder - .encoderFor(dataType.asInstanceOf[StructType]) + .encoderForResultSchema(dataType.asInstanceOf[StructType]) .asInstanceOf[AgnosticEncoder[E]] case ProductEncoder(clsTag, fields, outer) if ProductEncoder.isTuple(clsTag) => // Recursively continue updating the tuple product encoder diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowDeserializer.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowDeserializer.scala index fe0ecace05cc0..4441d85b76292 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowDeserializer.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowDeserializer.scala @@ -40,7 +40,7 @@ import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders._ import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema import org.apache.spark.sql.connect.common.types.ops.ConnectTypeOps import org.apache.spark.sql.errors.{CompilationErrors, ExecutionErrors} -import org.apache.spark.sql.types.Decimal +import org.apache.spark.sql.types.{Decimal, StringHelper} import org.apache.spark.sql.util.{CloseableIterator, ConcatenatingArrowStreamReader, MessageIterator} import org.apache.spark.unsafe.types.VariantVal @@ -135,8 +135,14 @@ object ArrowDeserializers { new Deserializer[Any] { def get(i: Int): Any = null } - case (StringEncoder, v: FieldVector) => - new LeafFieldDeserializer[String](encoder, v, timeZoneId) { + // CHAR/VARCHAR travel as plain Arrow string vectors; the length is part of the type, not of + // the encoding, and the values arrive already padded and length checked by the server. Read + // them against the unconstrained string type, since narrowing STRING to CHAR(n)/VARCHAR(n) + // is not an up-cast and the reader would reject the vector. + case (StringEncoder | _: CharEncoder | _: VarcharEncoder, v: FieldVector) => + val stringReader = + ArrowVectorReader(StringHelper.plainStringType(encoder.dataType), v, timeZoneId) + new LeafFieldDeserializer[String](stringReader) { override def value(i: Int): String = reader.getString(i) } case (JavaEnumEncoder(tag), v: FieldVector) => diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowSerializer.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowSerializer.scala index 02918e4eb1403..10375fa32c875 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowSerializer.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/ArrowSerializer.scala @@ -281,11 +281,11 @@ object ArrowSerializer { new FieldSerializer[Unit, NullVector](v) { override def set(index: Int, value: Unit): Unit = vector.setNull(index) } - case (StringEncoder, v: VarCharVector) => + case (StringEncoder | _: CharEncoder | _: VarcharEncoder, v: VarCharVector) => new FieldSerializer[String, VarCharVector](v) { override def set(index: Int, value: String): Unit = setString(v, index, value) } - case (StringEncoder, v: LargeVarCharVector) => + case (StringEncoder | _: CharEncoder | _: VarcharEncoder, v: LargeVarCharVector) => new FieldSerializer[String, LargeVarCharVector](v) { override def set(index: Int, value: String): Unit = setString(v, index, value) } diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/types/ops/TimestampNanosTypeConnectOps.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/types/ops/TimestampNanosTypeConnectOps.scala new file mode 100644 index 0000000000000..b29e5da958fb3 --- /dev/null +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/arrow/types/ops/TimestampNanosTypeConnectOps.scala @@ -0,0 +1,226 @@ +/* + * 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.connect.client.arrow.types.ops + +import java.time.{Instant, LocalDateTime} + +import org.apache.arrow.vector.FieldVector + +import org.apache.spark.connect.proto +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder +import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.{InstantNanosEncoder, LocalDateTimeNanosEncoder} +import org.apache.spark.sql.catalyst.util.SparkDateTimeUtils +import org.apache.spark.sql.connect.client.arrow.{ArrowDeserializers, ArrowSerializer, ArrowVectorReader} +import org.apache.spark.sql.connect.common.InvalidPlanInput +import org.apache.spark.sql.connect.common.types.ops.ConnectTypeOps +import org.apache.spark.sql.types.{DataType, TimestampLTZNanosType, TimestampNTZNanosType} +import org.apache.spark.unsafe.types.TimestampNanosVal + +/** + * Combined Connect operations shared by the nanosecond-capable timestamp types + * ([[TimestampNTZNanosType]] and [[TimestampLTZNanosType]], precision 7..9). + * + * Implements the proto DataType/Literal side of [[ConnectTypeOps]]. The physical value is + * [[org.apache.spark.unsafe.types.TimestampNanosVal]] (epoch micros + nanos within the micro), + * which the proto carries as `epoch_micros` + `nanos_within_micro` rather than a single int64 of + * nanoseconds, because nanoseconds-since-epoch cannot span the supported 0001..9999 year range. + * The two concrete subclasses differ only in their proto message arm, external java.time value + * ([[LocalDateTime]] for NTZ, [[Instant]] for LTZ) and the conversion helpers used. + * + * Arrow IPC serialization is out of scope for these types (SPARK-57161), so the ops is not + * registered in the Arrow dispatch of [[ConnectTypeOps]] and the Arrow methods below are never + * reached; they throw to make an accidental wiring obvious. + * + * Lives under the arrow.types.ops sub-package to co-locate with [[TimeTypeConnectOps]], the + * reference implementation this mirrors. + * + * @since 4.3.0 + */ +private[connect] abstract class TimestampNanosTypeConnectOps extends ConnectTypeOps { + + /** + * Rebuilds the physical value from the two proto components. `nanosWithinMicro` is an int32 on + * the wire, so its range is checked here before narrowing to `Short`: without the check + * `.toShort` would truncate an out-of-range value modulo 2^16 (e.g. 65536 -> 0) and slip it + * past the `[0, 999]` guard in `fromParts`, yielding a silently wrong value instead of a clear + * error. + */ + protected def toTimestampNanosVal( + epochMicros: Long, + nanosWithinMicro: Int): TimestampNanosVal = { + if (nanosWithinMicro < 0 || nanosWithinMicro > TimestampNanosVal.MAX_NANOS_WITHIN_MICRO) { + throw InvalidPlanInput( + s"nanos_within_micro must be in [0, ${TimestampNanosVal.MAX_NANOS_WITHIN_MICRO}], got: " + + nanosWithinMicro) + } + TimestampNanosVal.fromParts(epochMicros, nanosWithinMicro.toShort) + } + + // ==================== Arrow Serialization (unsupported) ==================== + + private def arrowUnsupported: Nothing = + throw new UnsupportedOperationException( + s"Arrow serialization is not supported for ${dataType.sql} over Spark Connect.") + + override def createArrowSerializer(vector: AnyRef): ArrowSerializer.Serializer = + arrowUnsupported + + override def createArrowDeserializer( + enc: AgnosticEncoder[_], + data: AnyRef, + timeZoneId: String): ArrowDeserializers.Deserializer[Any] = arrowUnsupported + + override def createArrowVectorReader(vector: FieldVector): ArrowVectorReader = arrowUnsupported +} + +/** + * Connect operations for [[TimestampNTZNanosType]]. The external java.time value is + * [[LocalDateTime]] (interpreted at UTC), matching the server-side TypeOps and RowEncoder. + * + * @param t + * The TimestampNTZNanosType with precision information + * @since 4.3.0 + */ +private[connect] class TimestampNTZNanosTypeConnectOps(val t: TimestampNTZNanosType) + extends TimestampNanosTypeConnectOps { + + override def dataType: DataType = t + + override def encoder: AgnosticEncoder[_] = LocalDateTimeNanosEncoder(t.precision) + + // ==================== Proto Conversions ==================== + + override def toCatalystTypeFromProto(t: proto.DataType): DataType = { + val nanos = t.getTimestampNtzNanos + if (nanos.hasPrecision) TimestampNTZNanosType(nanos.getPrecision) else TimestampNTZNanosType() + } + + override def toConnectProtoType: proto.DataType = { + proto.DataType + .newBuilder() + .setTimestampNtzNanos( + proto.DataType.TimestampNTZNanos.newBuilder().setPrecision(t.precision).build()) + .build() + } + + override def toLiteralProto( + value: Any, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = + setLiteral(value, TimestampNTZNanosType.DEFAULT_PRECISION, builder) + + override def toLiteralProtoWithType( + value: Any, + dt: DataType, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = + setLiteral(value, dt.asInstanceOf[TimestampNTZNanosType].precision, builder) + + private def setLiteral( + value: Any, + precision: Int, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = { + val v = SparkDateTimeUtils + .localDateTimeToTimestampNanos(value.asInstanceOf[LocalDateTime], precision) + builder.setTimestampNtzNanos( + builder.getTimestampNtzNanosBuilder + .setEpochMicros(v.epochMicros) + .setNanosWithinMicro(v.nanosWithinMicro.toInt) + .setPrecision(precision)) + } + + override def getScalaConverter: proto.Expression.Literal => Any = { v => + val nanos = v.getTimestampNtzNanos + SparkDateTimeUtils.timestampNanosToLocalDateTime( + toTimestampNanosVal(nanos.getEpochMicros, nanos.getNanosWithinMicro)) + } + + override def getProtoDataTypeFromLiteral(literal: proto.Expression.Literal): proto.DataType = { + val typeBuilder = proto.DataType.TimestampNTZNanos.newBuilder() + if (literal.getTimestampNtzNanos.hasPrecision) { + typeBuilder.setPrecision(literal.getTimestampNtzNanos.getPrecision) + } + proto.DataType.newBuilder().setTimestampNtzNanos(typeBuilder.build()).build() + } +} + +/** + * Connect operations for [[TimestampLTZNanosType]]. The external java.time value is [[Instant]], + * matching the server-side TypeOps and RowEncoder. + * + * @param t + * The TimestampLTZNanosType with precision information + * @since 4.3.0 + */ +private[connect] class TimestampLTZNanosTypeConnectOps(val t: TimestampLTZNanosType) + extends TimestampNanosTypeConnectOps { + + override def dataType: DataType = t + + override def encoder: AgnosticEncoder[_] = InstantNanosEncoder(t.precision) + + // ==================== Proto Conversions ==================== + + override def toCatalystTypeFromProto(t: proto.DataType): DataType = { + val nanos = t.getTimestampLtzNanos + if (nanos.hasPrecision) TimestampLTZNanosType(nanos.getPrecision) else TimestampLTZNanosType() + } + + override def toConnectProtoType: proto.DataType = { + proto.DataType + .newBuilder() + .setTimestampLtzNanos( + proto.DataType.TimestampLTZNanos.newBuilder().setPrecision(t.precision).build()) + .build() + } + + override def toLiteralProto( + value: Any, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = + setLiteral(value, TimestampLTZNanosType.DEFAULT_PRECISION, builder) + + override def toLiteralProtoWithType( + value: Any, + dt: DataType, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = + setLiteral(value, dt.asInstanceOf[TimestampLTZNanosType].precision, builder) + + private def setLiteral( + value: Any, + precision: Int, + builder: proto.Expression.Literal.Builder): proto.Expression.Literal.Builder = { + val v = SparkDateTimeUtils.instantToTimestampNanos(value.asInstanceOf[Instant], precision) + builder.setTimestampLtzNanos( + builder.getTimestampLtzNanosBuilder + .setEpochMicros(v.epochMicros) + .setNanosWithinMicro(v.nanosWithinMicro.toInt) + .setPrecision(precision)) + } + + override def getScalaConverter: proto.Expression.Literal => Any = { v => + val nanos = v.getTimestampLtzNanos + SparkDateTimeUtils.timestampNanosToInstant( + toTimestampNanosVal(nanos.getEpochMicros, nanos.getNanosWithinMicro)) + } + + override def getProtoDataTypeFromLiteral(literal: proto.Expression.Literal): proto.DataType = { + val typeBuilder = proto.DataType.TimestampLTZNanos.newBuilder() + if (literal.getTimestampLtzNanos.hasPrecision) { + typeBuilder.setPrecision(literal.getTimestampLtzNanos.getPrecision) + } + proto.DataType.newBuilder().setTimestampLtzNanos(typeBuilder.build()).build() + } +} diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala index 33da07cc1b5bf..43995df77d72c 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/LiteralValueProtoConverter.scala @@ -497,6 +497,14 @@ object LiteralValueProtoConverter { true case (proto.Expression.Literal.LiteralTypeCase.TIME, proto.DataType.KindCase.TIME) => true + case ( + proto.Expression.Literal.LiteralTypeCase.TIMESTAMP_NTZ_NANOS, + proto.DataType.KindCase.TIMESTAMP_NTZ_NANOS) => + true + case ( + proto.Expression.Literal.LiteralTypeCase.TIMESTAMP_LTZ_NANOS, + proto.DataType.KindCase.TIMESTAMP_LTZ_NANOS) => + true case (proto.Expression.Literal.LiteralTypeCase.ARRAY, proto.DataType.KindCase.ARRAY) => true case (proto.Expression.Literal.LiteralTypeCase.MAP, proto.DataType.KindCase.MAP) => diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/types/ops/ConnectTypeOps.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/types/ops/ConnectTypeOps.scala index 7225bc05feefe..84f0c5012b59f 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/types/ops/ConnectTypeOps.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/common/types/ops/ConnectTypeOps.scala @@ -23,8 +23,8 @@ import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.encoders.AgnosticEncoder import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.LocalTimeEncoder import org.apache.spark.sql.connect.client.arrow.{ArrowDeserializers, ArrowSerializer, ArrowVectorReader} -import org.apache.spark.sql.connect.client.arrow.types.ops.TimeTypeConnectOps -import org.apache.spark.sql.types.{DataType, TimeType} +import org.apache.spark.sql.connect.client.arrow.types.ops.{TimestampLTZNanosTypeConnectOps, TimestampNTZNanosTypeConnectOps, TimeTypeConnectOps} +import org.apache.spark.sql.types.{DataType, TimestampLTZNanosType, TimestampNTZNanosType, TimeType} /** * Optional type operations for Spark Connect infrastructure. @@ -97,6 +97,8 @@ object ConnectTypeOps { /** DataType-keyed dispatch for proto conversions. */ def apply(dt: DataType): Option[ConnectTypeOps] = dt match { case tt: TimeType => Some(new TimeTypeConnectOps(tt)) + case t: TimestampNTZNanosType => Some(new TimestampNTZNanosTypeConnectOps(t)) + case t: TimestampLTZNanosType => Some(new TimestampLTZNanosTypeConnectOps(t)) // Add new framework types here case _ => None } @@ -119,6 +121,10 @@ object ConnectTypeOps { private def opsForKindCase(kindCase: proto.DataType.KindCase): Option[ConnectTypeOps] = kindCase match { case proto.DataType.KindCase.TIME => Some(new TimeTypeConnectOps(TimeType())) + case proto.DataType.KindCase.TIMESTAMP_NTZ_NANOS => + Some(new TimestampNTZNanosTypeConnectOps(TimestampNTZNanosType())) + case proto.DataType.KindCase.TIMESTAMP_LTZ_NANOS => + Some(new TimestampLTZNanosTypeConnectOps(TimestampLTZNanosType())) // Add new framework proto kinds here - single registration for all KindCase lookups case _ => None } @@ -142,6 +148,10 @@ object ConnectTypeOps { litCase: proto.Expression.Literal.LiteralTypeCase): proto.DataType.KindCase = litCase match { case proto.Expression.Literal.LiteralTypeCase.TIME => proto.DataType.KindCase.TIME + case proto.Expression.Literal.LiteralTypeCase.TIMESTAMP_NTZ_NANOS => + proto.DataType.KindCase.TIMESTAMP_NTZ_NANOS + case proto.Expression.Literal.LiteralTypeCase.TIMESTAMP_LTZ_NANOS => + proto.DataType.KindCase.TIMESTAMP_LTZ_NANOS // Add new framework literal-to-kind mappings here case _ => proto.DataType.KindCase.KIND_NOT_SET } diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_and.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_and.explain new file mode 100644 index 0000000000000..d8892a7cfa822 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_and.explain @@ -0,0 +1,2 @@ +Project [bitmap_and(bytes#0, bytes#0) AS bitmap_and(bytes, bytes)#0] ++- LocalRelation <empty>, [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_andnot.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_andnot.explain new file mode 100644 index 0000000000000..7a081fa449751 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_andnot.explain @@ -0,0 +1,2 @@ +Project [bitmap_andnot(bytes#0, bytes#0) AS bitmap_andnot(bytes, bytes)#0] ++- LocalRelation <empty>, [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_or.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_or.explain new file mode 100644 index 0000000000000..f58db6eb69084 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_or.explain @@ -0,0 +1,2 @@ +Project [bitmap_or(bytes#0, bytes#0) AS bitmap_or(bytes, bytes)#0] ++- LocalRelation <empty>, [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor.explain new file mode 100644 index 0000000000000..a6f81bbc9a2b4 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor.explain @@ -0,0 +1,2 @@ +Project [bitmap_xor(bytes#0, bytes#0) AS bitmap_xor(bytes, bytes)#0] ++- LocalRelation <empty>, [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor_agg.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor_agg.explain new file mode 100644 index 0000000000000..ee4bb977934ce --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_bitmap_xor_agg.explain @@ -0,0 +1,2 @@ +Aggregate [bitmap_xor_agg(bytes#0, 0, 0) AS bitmap_xor_agg(bytes)#0] ++- LocalRelation <empty>, [id#0L, bytes#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_base32.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_base32.explain new file mode 100644 index 0000000000000..482fb5b0f232f --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_from_base32.explain @@ -0,0 +1,2 @@ +Project [static_invoke(UnBase32.decode(g#0)) AS from_base32(g)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_json_typeof.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_json_typeof.explain new file mode 100644 index 0000000000000..c8d3652edcb13 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_json_typeof.explain @@ -0,0 +1,2 @@ +Project [static_invoke(JsonExpressionUtils.jsonTypeof(g#0)) AS json_typeof(g)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_base32.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_base32.explain new file mode 100644 index 0000000000000..340a45adc3619 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_to_base32.explain @@ -0,0 +1,2 @@ +Project [static_invoke(Base32.encode(cast(g#0 as binary))) AS to_base32(CAST(g AS BINARY))#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_trim_array.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_trim_array.explain new file mode 100644 index 0000000000000..e4f13ce71a193 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_trim_array.explain @@ -0,0 +1,2 @@ +Project [trim_array(e#0, 2) AS trim_array(e, 2)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_truncate.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_truncate.explain new file mode 100644 index 0000000000000..6f07efb1efd8f --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_truncate.explain @@ -0,0 +1,2 @@ +Project [truncate(b#0, 2) AS truncate(b, 2)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain new file mode 100644 index 0000000000000..08b1a3eb011d9 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_unwrap_udt.explain @@ -0,0 +1,2 @@ +Project [unwrap_udt(wrap_udt(array(b#0), org.apache.spark.sql.TestUDT$NewArrayUDT)) AS unwrap_udt(wrap_udt(array(b), org.apache.spark.sql.TestUDT$NewArrayUDT))#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_arrays.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_arrays.explain new file mode 100644 index 0000000000000..522406c52b069 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_arrays.explain @@ -0,0 +1,2 @@ +Project [variant_from_arrays(array(a, b), array(1, 2)) AS variant_from_arrays(array(a, b), array(1, 2))#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_entries.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_entries.explain new file mode 100644 index 0000000000000..550854623fa9e --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_from_entries.explain @@ -0,0 +1,2 @@ +Project [variant_from_entries(array(struct(col1, a, col2, 1), struct(col1, b, col2, 2))) AS variant_from_entries(array(struct(a, 1), struct(b, 2)))#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_strip_nulls.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_strip_nulls.explain new file mode 100644 index 0000000000000..f6d8330de8280 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_strip_nulls.explain @@ -0,0 +1,2 @@ +Project [static_invoke(VariantExpressionEvalUtils.stripNulls(static_invoke(VariantExpressionEvalUtils.parseJson(g#0, false, true, true)), false)) AS variant_strip_nulls(parse_json(g), false)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_wrap_udt.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_wrap_udt.explain new file mode 100644 index 0000000000000..a1fb72ba2b775 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_wrap_udt.explain @@ -0,0 +1,2 @@ +Project [wrap_udt(array(b#0), org.apache.spark.sql.TestUDT$NewArrayUDT) AS wrap_udt(array(b), org.apache.spark.sql.TestUDT$NewArrayUDT)#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_128.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_128.explain new file mode 100644 index 0000000000000..189ed0c540e9c --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_128.explain @@ -0,0 +1,2 @@ +Project [xxh3_128(cast(g#0 as binary)) AS xxh3_128(CAST(g AS BINARY))#0] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_64.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_64.explain new file mode 100644 index 0000000000000..38bb03544c003 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_xxh3_64.explain @@ -0,0 +1,2 @@ +Project [xxh3_64(cast(g#0 as binary)) AS xxh3_64(CAST(g AS BINARY))#0L] ++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0] diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.json b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.json new file mode 100644 index 0000000000000..1abf52ef603ef --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,bytes:binary\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "bitmap_and", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "bitmap_and", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.proto.bin new file mode 100644 index 0000000000000..4ed1b96834a21 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_and.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.json b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.json new file mode 100644 index 0000000000000..a3363b09d4301 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,bytes:binary\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "bitmap_andnot", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "bitmap_andnot", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.proto.bin new file mode 100644 index 0000000000000..b0cdbd7ea9bd4 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_andnot.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.json b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.json new file mode 100644 index 0000000000000..7c0acff13b79a --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,bytes:binary\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "bitmap_or", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "bitmap_or", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.proto.bin new file mode 100644 index 0000000000000..b8dbed026128b Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_or.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.json b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.json new file mode 100644 index 0000000000000..59c78fa16801c --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,bytes:binary\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "bitmap_xor", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "bitmap_xor", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.proto.bin new file mode 100644 index 0000000000000..dfc994306a53a Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.json b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.json new file mode 100644 index 0000000000000..12f2878b6b97d --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.json @@ -0,0 +1,60 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,bytes:binary\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "bitmap_xor_agg", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "bytes" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "bitmap_xor_agg", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.proto.bin new file mode 100644 index 0000000000000..27ca4a6c38805 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_bitmap_xor_agg.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.json b/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.json new file mode 100644 index 0000000000000..d53a88995bfda --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.json @@ -0,0 +1,60 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "from_base32", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "from_base32", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.proto.bin new file mode 100644 index 0000000000000..551dda1b878c8 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_from_base32.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.json b/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.json new file mode 100644 index 0000000000000..a2f300b98ee81 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.json @@ -0,0 +1,60 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "json_typeof", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "json_typeof", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.proto.bin new file mode 100644 index 0000000000000..7177b03938e2a Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_json_typeof.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.json b/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.json new file mode 100644 index 0000000000000..4abce6e102574 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.json @@ -0,0 +1,85 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "to_base32", + "arguments": [{ + "cast": { + "expr": { + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, + "type": { + "binary": { + } + } + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.Column", + "methodName": "cast", + "fileName": "Column.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "to_base32", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.proto.bin new file mode 100644 index 0000000000000..93ef08635f3de Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_to_base32.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.json b/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.json new file mode 100644 index 0000000000000..1cd8f3d1b8761 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "trim_array", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "e" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "integer": 2 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "trim_array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "trim_array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.proto.bin new file mode 100644 index 0000000000000..d32a2256dad8e Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_trim_array.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.json b/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.json new file mode 100644 index 0000000000000..48f46452c670c --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.json @@ -0,0 +1,81 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "truncate", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "b" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "integer": 2 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "truncate", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "truncate", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.proto.bin new file mode 100644 index 0000000000000..f9f6342267523 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_truncate.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.json b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.json new file mode 100644 index 0000000000000..9a8af401c78d3 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.json @@ -0,0 +1,127 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "unwrap_udt", + "arguments": [{ + "unresolvedFunction": { + "functionName": "wrap_udt", + "arguments": [{ + "unresolvedFunction": { + "functionName": "array", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "b" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "string": "{\"type\":\"udt\",\"class\":\"org.apache.spark.sql.TestUDT$NewArrayUDT\",\"pyClass\":null,\"sqlType\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false}}" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "wrap_udt", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": true + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "wrap_udt", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": true + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "unwrap_udt", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.proto.bin new file mode 100644 index 0000000000000..1ec310e22f772 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_unwrap_udt.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.json b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.json new file mode 100644 index 0000000000000..1a9d2c5ae1483 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.json @@ -0,0 +1,169 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "variant_from_arrays", + "arguments": [{ + "unresolvedFunction": { + "functionName": "array", + "arguments": [{ + "literal": { + "string": "a" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "string": "b" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedFunction": { + "functionName": "array", + "arguments": [{ + "literal": { + "integer": 1 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "integer": 2 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "variant_from_arrays", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.proto.bin new file mode 100644 index 0000000000000..a005fd07d20e8 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_arrays.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.json b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.json new file mode 100644 index 0000000000000..d41b6b387bc30 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.json @@ -0,0 +1,192 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "variant_from_entries", + "arguments": [{ + "unresolvedFunction": { + "functionName": "array", + "arguments": [{ + "unresolvedFunction": { + "functionName": "struct", + "arguments": [{ + "literal": { + "string": "a" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "integer": 1 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "struct", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "unresolvedFunction": { + "functionName": "struct", + "arguments": [{ + "literal": { + "string": "b" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "integer": 2 + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "lit", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "struct", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "variant_from_entries", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.proto.bin new file mode 100644 index 0000000000000..09d0cc19b9164 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_from_entries.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.json b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.json new file mode 100644 index 0000000000000..19a249ce81888 --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.json @@ -0,0 +1,104 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "variant_strip_nulls", + "arguments": [{ + "unresolvedFunction": { + "functionName": "parse_json", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "parse_json", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "boolean": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "variant_strip_nulls", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "variant_strip_nulls", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.proto.bin new file mode 100644 index 0000000000000..e1a82aceffcb3 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_strip_nulls.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json new file mode 100644 index 0000000000000..7376698e8516b --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.json @@ -0,0 +1,104 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "wrap_udt", + "arguments": [{ + "unresolvedFunction": { + "functionName": "array", + "arguments": [{ + "unresolvedAttribute": { + "unparsedIdentifier": "b" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "array", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, { + "literal": { + "string": "{\"type\":\"udt\",\"class\":\"org.apache.spark.sql.TestUDT$NewArrayUDT\",\"pyClass\":null,\"sqlType\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false}}" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "wrap_udt", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": true + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "wrap_udt", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.proto.bin new file mode 100644 index 0000000000000..dd9898055ff57 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_wrap_udt.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.json b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.json new file mode 100644 index 0000000000000..27c099a4146fe --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.json @@ -0,0 +1,85 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "xxh3_128", + "arguments": [{ + "cast": { + "expr": { + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, + "type": { + "binary": { + } + } + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.Column", + "methodName": "cast", + "fileName": "Column.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "xxh3_128", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.proto.bin new file mode 100644 index 0000000000000..34feea650f4e4 Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_128.proto.bin differ diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.json b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.json new file mode 100644 index 0000000000000..1b802376f84bb --- /dev/null +++ b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.json @@ -0,0 +1,85 @@ +{ + "common": { + "planId": "1" + }, + "project": { + "input": { + "common": { + "planId": "0" + }, + "localRelation": { + "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e" + } + }, + "expressions": [{ + "unresolvedFunction": { + "functionName": "xxh3_64", + "arguments": [{ + "cast": { + "expr": { + "unresolvedAttribute": { + "unparsedIdentifier": "g" + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "col", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }, + "type": { + "binary": { + } + } + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.Column", + "methodName": "cast", + "fileName": "Column.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }], + "isInternal": false + }, + "common": { + "origin": { + "jvmOrigin": { + "stackTrace": [{ + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.functions$", + "methodName": "xxh3_64", + "fileName": "functions.scala" + }, { + "classLoaderName": "app", + "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite", + "methodName": "~~trimmed~anonfun~~", + "fileName": "PlanGenerationTestSuite.scala" + }] + } + } + } + }] + } +} \ No newline at end of file diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.proto.bin new file mode 100644 index 0000000000000..3cd206c11e4fb Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_xxh3_64.proto.bin differ diff --git a/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala b/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala new file mode 100644 index 0000000000000..77704b8a57c30 --- /dev/null +++ b/sql/connect/common/src/test/scala/org/apache/spark/sql/TestUDT.scala @@ -0,0 +1,42 @@ +/* + * 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 + +import org.apache.spark.sql.types.{ArrayType, DataType, DoubleType, SQLUserDefinedType, UserDefinedType} + +object TestUDT { + + @SQLUserDefinedType(udt = classOf[NewArrayUDT]) + private[sql] class NewArray(val values: Array[Double]) extends Serializable + + private[sql] class NewArrayUDT extends UserDefinedType[NewArray] { + + override def sqlType: DataType = ArrayType(DoubleType, containsNull = false) + + override def serialize(obj: NewArray): Any = obj.values + + override def deserialize(datum: Any): NewArray = { + datum match { + case values: Array[_] => + new NewArray(values.map(_.asInstanceOf[Double])) + } + } + + override def userClass: Class[NewArray] = classOf[NewArray] + } +} diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala index f5a71a6674657..10531e063224d 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala @@ -22,7 +22,6 @@ import java.util.concurrent.TimeUnit import org.apache.spark.SparkEnv import org.apache.spark.network.util.ByteUnit import org.apache.spark.sql.connect.common.config.ConnectCommon -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.buildConf object Connect { @@ -302,17 +301,6 @@ object Connect { .intConf .createWithDefault(200) - val CONNECT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL = - buildStaticConf("spark.connect.copyFromLocalToFs.allowDestLocal") - .internal() - .doc(s""" - |(Deprecated since Spark 4.0, please set - |'${SQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL.key}' instead. - |""".stripMargin) - .version("3.5.0") - .booleanConf - .createWithDefault(false) - val CONNECT_UI_SESSION_LIMIT = buildStaticConf("spark.sql.connect.ui.retainedSessions") .doc("The number of client sessions kept in the Spark Connect UI history.") .version("3.5.0") diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala index 9f606b698d30c..ff447ea170162 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/ExecuteThreadRunner.scala @@ -25,7 +25,7 @@ import scala.util.control.NonFatal import com.google.protobuf.Message -import org.apache.spark.SparkSQLException +import org.apache.spark.{SparkContext, SparkSQLException} import org.apache.spark.connect.proto import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.sql.connect.common.ProtoUtils @@ -224,6 +224,9 @@ private[connect] class ExecuteThreadRunner(executeHolder: ExecuteHolder) extends "callSite.short", s"Spark Connect - ${Utils.abbreviate(debugString, 128)}") session.sparkContext.setLocalProperty("callSite.long", Utils.abbreviate(debugString, 2048)) + session.sparkContext.setLocalProperty( + SparkContext.SPARK_CONNECT_OPERATION_ID_PROPERTY, + executeHolder.operationId) executeHolder.request.getPlan.getOpTypeCase match { case proto.Plan.OpTypeCase.ROOT | proto.Plan.OpTypeCase.COMMAND => diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala index 5deb9ce1c3f19..91bc09afacf70 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala @@ -23,6 +23,7 @@ import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, TimeUnit} import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong} import scala.collection.mutable +import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal import com.google.common.cache.{CacheBuilder, RemovalNotification} @@ -74,14 +75,28 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_OFFLOADING_TIMEOUT) } + private case class ModelMetadata( + className: String, + modelString: String, + estimatedSizeBytes: Option[Long]) + + // Keep lightweight metadata after a model is evicted from memory so the UI can report + // offloaded models without loading them back into memory. + private val cachedModelMetadata = new ConcurrentHashMap[String, ModelMetadata]() + private val inMemoryModelIds = ConcurrentHashMap.newKeySet[String]() + private[ml] case class CacheItem(obj: Object, sizeBytes: Long) private[ml] val cachedModel: ConcurrentMap[String, CacheItem] = { if (getMemoryControlEnabled) { CacheBuilder .newBuilder() .softValues() - .removalListener((removed: RemovalNotification[String, CacheItem]) => - totalMLCacheInMemorySizeBytes.addAndGet(-removed.getValue.sizeBytes)) + .removalListener((removed: RemovalNotification[String, CacheItem]) => { + Option(removed.getValue).foreach { value => + totalMLCacheInMemorySizeBytes.addAndGet(-value.sizeBytes) + } + inMemoryModelIds.remove(removed.getKey) + }) .maximumWeight(getMaxInMemoryCacheSizeKB) .weigher((key: String, value: CacheItem) => { Math.ceil(value.sizeBytes.toDouble / 1024).toInt @@ -149,6 +164,7 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { } else { 0L // Don't need to calculate size if disables memory-control. } + inMemoryModelIds.add(objectId) cachedModel.put(objectId, CacheItem(obj, sizeBytes)) if (getMemoryControlEnabled) { val savePath = getModelOffloadingPath(objectId) @@ -163,6 +179,12 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { totalMLCacheInMemorySizeBytes.addAndGet(sizeBytes) totalMLCacheSizeBytes.addAndGet(sizeBytes) } + cachedModelMetadata.put( + objectId, + ModelMetadata( + obj.getClass.getName, + obj.toString, + if (getMemoryControlEnabled) Some(sizeBytes) else None)) } else { throw new RuntimeException("'MLCache.register' only accepts model or summary objects.") } @@ -219,6 +241,7 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { loadPath.toString, loadFromLocal = true) val sizeBytes = estimateObjectSize(obj) + inMemoryModelIds.add(refId) cachedModel.put(refId, CacheItem(obj, sizeBytes)) totalMLCacheInMemorySizeBytes.addAndGet(sizeBytes) } @@ -231,6 +254,7 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { verifyObjectId(refId) val removedModel = cachedModel.remove(refId) val removedFromMem = removedModel != null + inMemoryModelIds.remove(refId) val removedFromDisk = if (!evictOnly && removedModel != null && getMemoryControlEnabled) { totalMLCacheSizeBytes.addAndGet(-removedModel.sizeBytes) val removePath = getModelOffloadingPath(refId) @@ -244,6 +268,9 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { } else { false } + if (removedFromMem && (!evictOnly || !getMemoryControlEnabled)) { + cachedModelMetadata.remove(refId) + } removedFromMem || removedFromDisk } @@ -264,6 +291,9 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { def clear(): Int = this.synchronized { val size = cachedModel.size() cachedModel.clear() + cachedModelMetadata.clear() + inMemoryModelIds.clear() + totalMLCacheInMemorySizeBytes.set(0) totalMLCacheSizeBytes.set(0) if (getMemoryControlEnabled) { SparkFileUtils.cleanDirectory(new File(offloadedModelsDir.toString)) @@ -280,4 +310,39 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { } info.result() } + + /** Returns a cache snapshot without loading or touching any cached model. */ + def getStatus: MLCacheStatus = this.synchronized { + val models = cachedModelMetadata.asScala.iterator.map { case (id, metadata) => + MLCacheModelInfo( + id = id, + className = metadata.className, + modelString = metadata.modelString, + estimatedSizeBytes = metadata.estimatedSizeBytes, + inMemory = inMemoryModelIds.contains(id)) + }.toSeq + MLCacheStatus( + memoryControlEnabled = getMemoryControlEnabled, + inMemorySizeBytes = totalMLCacheInMemorySizeBytes.get(), + maxInMemorySizeBytes = sessionHolder.session.conf.get( + Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_IN_MEMORY_SIZE), + totalSizeBytes = totalMLCacheSizeBytes.get(), + maxTotalSizeBytes = getMLCacheMaxSize, + models = models) + } } + +private[connect] case class MLCacheModelInfo( + id: String, + className: String, + modelString: String, + estimatedSizeBytes: Option[Long], + inMemory: Boolean) + +private[connect] case class MLCacheStatus( + memoryControlEnabled: Boolean, + inMemorySizeBytes: Long, + maxInMemorySizeBytes: Long, + totalSizeBytes: Long, + maxTotalSizeBytes: Long, + models: Seq[MLCacheModelInfo]) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala index 841ce26402cf0..78562e8ebf92d 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala @@ -2201,6 +2201,15 @@ class SparkConnectPlanner( createUserDefinedPythonFunction(fun) .builder(fun.getArgumentsList.asScala.map(transformExpression).toSeq) match { case udaf: PythonUDAF => udaf.toAggregateExpression() + case agg: PythonAggregate => + // The two-stage incremental aggregation operators do not implement DISTINCT. The SQL path + // rejects it in FunctionResolution, but a Connect aggregate is already resolved and skips + // that guard, so reject `is_distinct` here rather than silently dropping it and returning a + // non-distinct result. + if (fun.getIsDistinct) { + throw QueryCompilationErrors.functionWithUnsupportedSyntaxError(agg.name, "DISTINCT") + } + agg.toAggregateExpression() case other => other } } @@ -2214,7 +2223,9 @@ class SparkConnectPlanner( func = function, dataType = transformDataType(udf.getOutputType), pythonEvalType = udf.getEvalType, - udfDeterministic = fun.getDeterministic) + udfDeterministic = fun.getDeterministic, + // Set only for incremental Python aggregators (see PythonAggregate). + bufferType = if (udf.hasBufferType) transformDataType(udf.getBufferType) else null) } private def transformPythonFunction(fun: proto.PythonUDF): SimplePythonFunction = { diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala index 2276230545e67..e97bcb0c786e3 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connect.IllegalStateErrors import org.apache.spark.sql.connect.common.InvalidPlanInput import org.apache.spark.sql.connect.config.Connect -import org.apache.spark.sql.connect.ml.MLCache +import org.apache.spark.sql.connect.ml.{MLCache, MLCacheStatus} import org.apache.spark.sql.connect.pipelines.DataflowGraphRegistry import org.apache.spark.sql.connect.planner.PythonStreamingQueryListener import org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper @@ -132,7 +132,16 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio new ConcurrentHashMap() // ML model cache - private[connect] lazy val mlCache = new MLCache(this) + @volatile private var mlCacheInitialized = false + private[connect] lazy val mlCache = { + val cache = new MLCache(this) + mlCacheInitialized = true + cache + } + + private[connect] def getMLCacheStatus: Option[MLCacheStatus] = { + if (mlCacheInitialized) Some(mlCache.getStatus) else None + } // Mapping from id to StreamingQueryListener. Used for methods like removeListener() in // StreamingQueryManager. diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala index 06bc24b6ccae6..73eba3a6f476e 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala @@ -18,12 +18,15 @@ package org.apache.spark.sql.connect.service import scala.jdk.CollectionConverters._ +import scala.util.matching.Regex import io.grpc.stub.StreamObserver import org.apache.spark.connect.proto import org.apache.spark.internal.Logging +import org.apache.spark.internal.config.SECRET_REDACTION_PATTERN import org.apache.spark.sql.RuntimeConfig +import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.internal.SQLConf class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigResponse]) @@ -44,18 +47,22 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes } private def doHandle(r: proto.ConfigRequest, h: SessionHolder): Unit = h.withSession { s => + // Read the pattern from the SparkConf rather than from the session config. The session config + // is writable by the client, so taking the pattern from there would let a client widen its own + // view of the configuration before reading it back. + val redactionPattern = s.sparkContext.conf.get(SECRET_REDACTION_PATTERN) // Make sure we're using the current running session. val builder = r.getOperation.getOpTypeCase match { case proto.ConfigRequest.Operation.OpTypeCase.SET => handleSet(r.getOperation.getSet, s.conf) case proto.ConfigRequest.Operation.OpTypeCase.GET => - handleGet(r.getOperation.getGet, s.conf) + handleGet(r.getOperation.getGet, s.conf, redactionPattern) case proto.ConfigRequest.Operation.OpTypeCase.GET_WITH_DEFAULT => - handleGetWithDefault(r.getOperation.getGetWithDefault, s.conf) + handleGetWithDefault(r.getOperation.getGetWithDefault, s.conf, redactionPattern) case proto.ConfigRequest.Operation.OpTypeCase.GET_OPTION => - handleGetOption(r.getOperation.getGetOption, s.conf) + handleGetOption(r.getOperation.getGetOption, s.conf, redactionPattern) case proto.ConfigRequest.Operation.OpTypeCase.GET_ALL => - handleGetAll(r.getOperation.getGetAll, s.conf) + handleGetAll(r.getOperation.getGetAll, s.conf, redactionPattern) case proto.ConfigRequest.Operation.OpTypeCase.UNSET => handleUnset(r.getOperation.getUnset, s.conf) case proto.ConfigRequest.Operation.OpTypeCase.IS_MODIFIABLE => @@ -94,10 +101,15 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes private def handleGet( operation: proto.ConfigRequest.Get, - conf: RuntimeConfig): proto.ConfigResponse.Builder = { + conf: RuntimeConfig, + redactionPattern: Regex): proto.ConfigResponse.Builder = { val builder = proto.ConfigResponse.newBuilder() operation.getKeysList.asScala.iterator.foreach { key => val value = conf.get(key) + if (SparkConnectConfigHandler.isRedacted(key, Option(value), redactionPattern)) { + // This operation reports an unset key by failing, so a redacted key fails the same way. + throw QueryExecutionErrors.sqlConfigNotFoundError(key) + } builder.addPairs(SparkConnectConfigHandler.toProtoKeyValue(key, Option(value))) getWarning(key).foreach(builder.addWarnings) } @@ -106,12 +118,19 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes private def handleGetWithDefault( operation: proto.ConfigRequest.GetWithDefault, - conf: RuntimeConfig): proto.ConfigResponse.Builder = { + conf: RuntimeConfig, + redactionPattern: Regex): proto.ConfigResponse.Builder = { val builder = proto.ConfigResponse.newBuilder() operation.getPairsList.asScala.iterator.foreach { pair => val (key, default) = SparkConnectConfigHandler.toKeyValue(pair) - val value = conf.get(key, default.orNull) - builder.addPairs(SparkConnectConfigHandler.toProtoKeyValue(key, Option(value))) + val stored = Option(conf.get(key, default.orNull)) + // A redacted entry falls back to the caller's default, as an unset key does. + val value = if (SparkConnectConfigHandler.isRedacted(key, stored, redactionPattern)) { + default + } else { + stored + } + builder.addPairs(SparkConnectConfigHandler.toProtoKeyValue(key, value)) getWarning(key).foreach(builder.addWarnings) } builder @@ -119,10 +138,16 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes private def handleGetOption( operation: proto.ConfigRequest.GetOption, - conf: RuntimeConfig): proto.ConfigResponse.Builder = { + conf: RuntimeConfig, + redactionPattern: Regex): proto.ConfigResponse.Builder = { val builder = proto.ConfigResponse.newBuilder() operation.getKeysList.asScala.iterator.foreach { key => - val value = conf.getOption(key) + val stored = conf.getOption(key) + val value = if (SparkConnectConfigHandler.isRedacted(key, stored, redactionPattern)) { + None + } else { + stored + } builder.addPairs(SparkConnectConfigHandler.toProtoKeyValue(key, value)) getWarning(key).foreach(builder.addWarnings) } @@ -131,15 +156,22 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes private def handleGetAll( operation: proto.ConfigRequest.GetAll, - conf: RuntimeConfig): proto.ConfigResponse.Builder = { + conf: RuntimeConfig, + redactionPattern: Regex): proto.ConfigResponse.Builder = { val builder = proto.ConfigResponse.newBuilder() + // Drop redacted entries before the prefix is stripped below. Matching afterwards would let a + // GetAll with prefix `spark.my.secret.` return `spark.my.secret.value` as `value`, which no + // longer matches the pattern. + val visible = conf.getAll.iterator.filterNot { case (key, value) => + SparkConnectConfigHandler.isRedacted(key, Option(value), redactionPattern) + } val results = if (operation.hasPrefix) { val prefix = operation.getPrefix - conf.getAll.iterator + visible .filter { case (key, _) => key.startsWith(prefix) } .map { case (key, value) => (key.substring(prefix.length), value) } } else { - conf.getAll.iterator + visible } results.foreach { case (key, value) => builder.addPairs(SparkConnectConfigHandler.toProtoKeyValue(key, Option(value))) @@ -185,6 +217,21 @@ object SparkConnectConfigHandler { private[connect] val unsupportedConfigurations = Set("spark.sql.execution.arrow.enabled", "spark.sql.execution.arrow.pyspark.fallback.enabled") + /** + * Whether a configuration entry is considered sensitive and must not be disclosed by the Config + * RPC. Follows `spark.redaction.regex` and matches the key or the value, the way `Utils.redact` + * does for the environment UI and event logs and `SetCommand` does for `SET`. Matching the + * value is what covers a secret carried by an innocuous key, such as a password inside a JDBC + * URL. + */ + private[connect] def isRedacted( + key: String, + value: Option[String], + redactionPattern: Regex): Boolean = { + redactionPattern.findFirstIn(key).isDefined || + value.exists(redactionPattern.findFirstIn(_).isDefined) + } + def toKeyValue(pair: proto.KeyValue): (String, Option[String]) = { val key = pair.getKey val value = if (pair.hasValue) { diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectGetStatusHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectGetStatusHandler.scala index c37061ff2ebae..fcc7089f27aaf 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectGetStatusHandler.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectGetStatusHandler.scala @@ -179,7 +179,8 @@ class SparkConnectGetStatusHandler(responseObserver: StreamObserver[proto.GetSta SparkConnectPluginRegistry.getStatusRegistry.flatMap { plugin => try { plugin.processRequestExtensions(sessionHolder, requestExtensions).toScala match { - case Some(extensions) => extensions.asScala.toSeq + // Filter nulls inside the isolation boundary so a null element can't NPE addExtensions. + case Some(extensions) => extensions.asScala.iterator.filter(_ != null).toSeq case None => Seq.empty } } catch { @@ -202,7 +203,8 @@ class SparkConnectGetStatusHandler(responseObserver: StreamObserver[proto.GetSta plugin .processOperationExtensions(operationId, sessionHolder, operationExtensions) .toScala match { - case Some(extensions) => extensions.asScala.toSeq + // Filter nulls inside the isolation boundary so a null element can't NPE addExtensions. + case Some(extensions) => extensions.asScala.iterator.filter(_ != null).toSeq case None => Seq.empty } } catch { diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala index c76794e3b6ec1..1fa9fbeda1e53 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala @@ -386,7 +386,8 @@ object SparkConnectService extends Logging { Some( new SparkConnectServerTab( new SparkConnectServerAppStatusStore(kvStore), - SparkConnectServerTab.getSparkUI(sc))) + SparkConnectServerTab.getSparkUI(sc), + Some(sessionManager))) } else { None } @@ -558,19 +559,6 @@ object SparkConnectService extends Logging { listenerBus.post(eventBuilder(bindingAddress)) } - - def extractErrorMessage(st: Throwable): String = { - val message = Utils.abbreviate(st.getMessage, 2048) - convertNullString(message) - } - - def convertNullString(str: String): String = { - if (str != null) { - str - } else { - "" - } - } } /** diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala index d3ddf592e9e7d..c7558465a0b9a 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala @@ -32,6 +32,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{INTERVAL, SESSION_HOLD_INFO} import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connect.config.Connect.{CONNECT_SESSION_MANAGER_CLOSED_SESSIONS_TOMBSTONES_SIZE, CONNECT_SESSION_MANAGER_DEFAULT_SESSION_TIMEOUT, CONNECT_SESSION_MANAGER_MAINTENANCE_INTERVAL} +import org.apache.spark.sql.connect.ml.MLCacheStatus import org.apache.spark.util.ThreadUtils /** @@ -284,6 +285,16 @@ class SparkConnectSessionManager extends Logging { closedSessionsCache.asMap.asScala.values.toSeq } + // Read live cache state directly without updating the sessions' last-access times. + private[connect] def getMLCacheStatuses: Map[SessionKey, Option[MLCacheStatus]] = { + sessionStore + .entrySet() + .asScala + .iterator + .map(entry => entry.getKey -> entry.getValue.getMLCacheStatus) + .toMap + } + /** * Schedules periodic maintenance checks if it is not already scheduled. * diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala index ea78b2dc59f6a..0d55c2e7866ff 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala @@ -26,6 +26,8 @@ import scala.xml.Node import jakarta.servlet.http.HttpServletRequest import org.apache.spark.internal.Logging +import org.apache.spark.sql.connect.ml.{MLCacheModelInfo, MLCacheStatus} +import org.apache.spark.sql.connect.service.SessionKey import org.apache.spark.sql.connect.ui.ToolTips._ import org.apache.spark.ui._ import org.apache.spark.ui.UIUtils._ @@ -52,6 +54,12 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) /** Render the page */ def render(request: HttpServletRequest): Seq[Node] = { + // Do not hold the status store lock while waiting for live ML cache snapshots. ML cache + // operations can perform disk I/O while holding their own locks. + val mlCacheStatuses = parent.getMLCacheStatuses + val initializedMLCacheStatuses = mlCacheStatuses.toSeq.flatMap(_.iterator.collect { + case (key, Some(status)) => key -> status + }) val content = store.synchronized { // make sure all parts in this page are consistent generateBasicStats() ++ <br/> ++ @@ -62,8 +70,9 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) {store.getTotalRunning} Request(s) </h4> ++ - generateSessionStatsTable(request) ++ - generateSQLStatsTable(request) + generateSessionStatsTable(request, mlCacheStatuses) ++ + generateSQLStatsTable(request) ++ + generateMLCacheStatsTable(request, initializedMLCacheStatuses) } UIUtils.headerSparkPage(request, "Spark Connect", content, parent) } @@ -132,7 +141,9 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) } /** Generate stats of batch sessions of the Spark Connect server */ - private def generateSessionStatsTable(request: HttpServletRequest): Seq[Node] = { + private def generateSessionStatsTable( + request: HttpServletRequest, + mlCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]]): Seq[Node] = { val numSessions = store.getSessionList.size val table = if (numSessions > 0) { @@ -149,7 +160,8 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) store.getSessionList, "connect", UIUtils.prependBaseUri(request, parent.basePath), - sessionTableTag).table(sessionTablePage)) + sessionTableTag, + mlCacheStatuses).table(sessionTablePage)) } catch { case e @ (_: IllegalArgumentException | _: IndexOutOfBoundsException) => Some(<div class="alert alert-danger"> @@ -179,6 +191,83 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) content } + + /** Generate live ML cache statistics for active Spark Connect sessions. */ + private def generateMLCacheStatsTable( + request: HttpServletRequest, + mlCacheStatuses: Seq[(SessionKey, MLCacheStatus)]): Seq[Node] = { + val populatedCacheStatuses = mlCacheStatuses.filter(_._2.models.nonEmpty) + val models = populatedCacheStatuses.flatMap { case (key, status) => + status.models.map(MLCacheModelTableRow(key.userId, key.sessionId, _)) + } + if (models.isEmpty) { + return Seq.empty + } + + val tableTag = "mlcachemodels" + val tablePage = Option(request.getParameter(s"$tableTag.page")).map(_.toInt).getOrElse(1) + val table = + try { + new MLCacheModelStatsPagedTable( + request, + parent, + models, + "connect", + UIUtils.prependBaseUri(request, parent.basePath), + tableTag).table(tablePage) + } catch { + case e @ (_: IllegalArgumentException | _: IndexOutOfBoundsException) => + <div class="alert alert-danger"> + <p>Error while rendering ML cache table:</p> + <pre> + {Utils.exceptionString(e)} + </pre> + </div> + } + + val inMemoryModels = models.count(_.model.inMemory) + val memoryControlledStatuses = + populatedCacheStatuses.map(_._2).filter(_.memoryControlEnabled) + val inMemorySize = memoryControlledStatuses.map(s => BigInt(s.inMemorySizeBytes)).sum + val maxInMemorySize = memoryControlledStatuses.map(s => BigInt(s.maxInMemorySizeBytes)).sum + val totalSize = memoryControlledStatuses.map(s => BigInt(s.totalSizeBytes)).sum + val maxTotalSize = memoryControlledStatuses.map(s => BigInt(s.maxTotalSizeBytes)).sum + val sizeStats = if (memoryControlledStatuses.nonEmpty) { + Seq( + <li> + <strong>Estimated size (In-memory): </strong> + {Utils.bytesToString(inMemorySize)} / {Utils.bytesToString(maxInMemorySize)} + </li>, + <li> + <strong>Estimated size (In-memory and Offloaded data): </strong> + {Utils.bytesToString(totalSize)} / {Utils.bytesToString(maxTotalSize)} + </li>) + } else { + Seq.empty + } + + <span id="mlcachestat" class="collapse-table" data-bs-toggle="collapse" + data-bs-target="#aggregated-mlcachestat" + aria-expanded="true" aria-controls="aggregated-mlcachestat" + data-collapse-name="collapse-aggregated-mlcachestat"> + <h4> + <span class="collapse-table-arrow arrow-open"></span> + <a>ML Cache Statistics ({models.size})</a> + </h4> + </span> ++ + <div class="collapsible-table collapse show" id="aggregated-mlcachestat"> + <ul class="list-unstyled"> + <li><strong>Sessions with cached models: </strong>{populatedCacheStatuses.size}</li> + <li> + <strong>Cached models: </strong> + {models.size} ({inMemoryModels} in memory, {models.size - inMemoryModels} offloaded) + </li> + {sizeStats} + </ul> + <h5>Cached Models</h5> + {table} + </div> + } } private[ui] class SqlStatsPagedTable( @@ -361,7 +450,8 @@ private[ui] class SessionStatsPagedTable( data: Seq[SessionInfo], subPath: String, basePath: String, - sessionStatsTableTag: String) + sessionStatsTableTag: String, + mlCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]]) extends PagedTable[SessionInfo] { private val (sortColumn, desc, pageSize) = @@ -404,7 +494,8 @@ private[ui] class SessionStatsPagedTable( ("Start Time", true, None), ("Finish Time", true, None), ("Duration", true, Some(SPARK_CONNECT_SESSION_DURATION)), - ("Total Execute", true, Some(SPARK_CONNECT_SESSION_TOTAL_EXECUTE))) + ("Total Execute", true, Some(SPARK_CONNECT_SESSION_TOTAL_EXECUTE)), + ("ML Cache", false, Some(SPARK_CONNECT_SESSION_ML_CACHE))) isSortColumnValid(sessionTableHeadersAndTooltips, sortColumn) @@ -431,8 +522,153 @@ private[ui] class SessionStatsPagedTable( <td> {if (session.finishTimestamp > 0) formatDate(session.finishTimestamp)} </td> <td> {formatDurationVerbose(session.totalTime)} </td> <td> {session.totalExecution.toString} </td> + <td> {renderMLCacheStatus(session)} </td> </tr> } + + private def renderMLCacheStatus(session: SessionInfo): Seq[Node] = { + if (session.finishTimestamp > 0) { + <span>N/A</span> + } else { + mlCacheStatuses + .flatMap(_.get(SessionKey(session.userId, session.sessionId))) + .map { + case Some(status) if status.models.nonEmpty => + if (status.memoryControlEnabled) { + val inMemoryModels = status.models.count(_.inMemory) + val modelLabel = if (inMemoryModels == 1) "model" else "models" + <span> + {s"$inMemoryModels $modelLabel in memory"}<br/> + { + s"${Utils.bytesToString(status.inMemorySizeBytes)} / " + + s"${Utils.bytesToString(status.maxInMemorySizeBytes)} memory" + }<br/> + { + s"${Utils.bytesToString(status.totalSizeBytes)} / " + + s"${Utils.bytesToString(status.maxTotalSizeBytes)} total" + } + </span> + } else { + val cachedModels = status.models.size + val modelLabel = if (cachedModels == 1) "model" else "models" + <span> + Memory control disabled<br/> + {s"$cachedModels cached $modelLabel"} + </span> + } + case Some(_) | None => + <span>Not used</span> + } + .getOrElse(<span>N/A</span>) + } + } +} + +private[ui] case class MLCacheModelTableRow( + userId: String, + sessionId: String, + model: MLCacheModelInfo) + +private[ui] class MLCacheModelStatsPagedTable( + request: HttpServletRequest, + parent: SparkConnectServerTab, + data: Seq[MLCacheModelTableRow], + subPath: String, + basePath: String, + tableTag: String) + extends PagedTable[MLCacheModelTableRow] { + + private val (sortColumn, desc, pageSize) = + getTableParameters(request, tableTag, "Estimated Size") + + private val encodedSortColumn = URLEncoder.encode(sortColumn, UTF_8.name()) + private val parameterPath = s"$basePath/$subPath/?${getParameterOtherTable(request, tableTag)}" + + override val dataSource = + new MLCacheModelTableDataSource(data, pageSize, sortColumn, desc) + + override def tableId: String = tableTag + + override def tableCssClass: String = + "table table-bordered table-sm table-striped table-head-clickable table-cell-width-limited" + + override def pageLink(page: Int): String = { + parameterPath + + s"&$pageNumberFormField=$page" + + s"&$tableTag.sort=$encodedSortColumn" + + s"&$tableTag.desc=$desc" + + s"&$pageSizeFormField=$pageSize" + + s"#$tableTag" + } + + override def pageSizeFormField: String = s"$tableTag.pageSize" + + override def pageNumberFormField: String = s"$tableTag.page" + + override def goButtonFormPath: String = + s"$parameterPath&$tableTag.sort=$encodedSortColumn" + + s"&$tableTag.desc=$desc#$tableTag" + + override def headers: Seq[Node] = { + val headersAndTooltips: Seq[(String, Boolean, Option[String])] = Seq( + ("User", true, None), + ("Session ID", true, None), + ("Model ID", true, None), + ("Model Class", true, None), + ("Model Details", true, Some(SPARK_CONNECT_ML_CACHE_MODEL_DETAILS)), + ("Estimated Size", true, Some(SPARK_CONNECT_ML_CACHE_ESTIMATED_SIZE)), + ("Storage", true, Some(SPARK_CONNECT_ML_CACHE_STORAGE))) + + isSortColumnValid(headersAndTooltips, sortColumn) + headerRow(headersAndTooltips, desc, pageSize, sortColumn, parameterPath, tableTag, tableTag) + } + + override def row(row: MLCacheModelTableRow): Seq[Node] = { + val model = row.model + val sessionLink = "%s/%s/session/?id=%s&userId=%s".format( + UIUtils.prependBaseUri(request, parent.basePath), + parent.prefix, + URLEncoder.encode(row.sessionId, UTF_8.name()), + ConnectUiUtils.encodeUserId(row.userId)) + <tr> + <td>{row.userId}</td> + <td><a href={sessionLink}>{row.sessionId}</a></td> + <td>{model.id}</td> + <td>{model.className}</td> + <td>{model.modelString}</td> + <td>{model.estimatedSizeBytes.map(Utils.bytesToString).getOrElse("N/A")}</td> + <td>{if (model.inMemory) "In memory" else "Offloaded"}</td> + </tr> + } +} + +private[ui] class MLCacheModelTableDataSource( + info: Seq[MLCacheModelTableRow], + pageSize: Int, + sortColumn: String, + desc: Boolean) + extends PagedDataSource[MLCacheModelTableRow](pageSize) { + + private val data = info.sorted(ordering(sortColumn, desc)) + + override def dataSize: Int = data.size + + override def sliceData(from: Int, to: Int): Seq[MLCacheModelTableRow] = data.slice(from, to) + + private def ordering(sortColumn: String, desc: Boolean): Ordering[MLCacheModelTableRow] = { + val ordering: Ordering[MLCacheModelTableRow] = sortColumn match { + case "User" => Ordering.by(_.userId) + case "Session ID" => Ordering.by(_.sessionId) + case "Model ID" => Ordering.by(_.model.id) + case "Model Class" => Ordering.by(_.model.className) + case "Model Details" => Ordering.by(_.model.modelString) + case "Estimated Size" => + Ordering.by((row: MLCacheModelTableRow) => row.model.estimatedSizeBytes) + case "Storage" => Ordering.by(_.model.inMemory) + case unknownColumn => throw new IllegalArgumentException(s"Unknown column: $unknownColumn") + } + if (desc) ordering.reverse else ordering + } } private[ui] class SqlStatsTableRow( diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala index c5ea0bf618b52..0de43122a194f 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala @@ -21,12 +21,15 @@ import java.util.Date import org.apache.spark.SparkContext import org.apache.spark.internal.Logging +import org.apache.spark.sql.connect.ml.MLCacheStatus +import org.apache.spark.sql.connect.service.{SessionKey, SparkConnectSessionManager} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.ui.{SparkUI, SparkUITab} private[connect] class SparkConnectServerTab( val store: SparkConnectServerAppStatusStore, - sparkUI: SparkUI) + sparkUI: SparkUI, + sessionManager: Option[SparkConnectSessionManager] = None) extends SparkUITab(sparkUI, "connect") with Logging { @@ -47,6 +50,9 @@ private[connect] class SparkConnectServerTab( parent.detachTab(this) } + def getMLCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]] = + sessionManager.map(_.getMLCacheStatuses) + override def displayOrder: Int = 3 } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala index 9b51ace83c6c1..30762b9f0f5d1 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala @@ -36,4 +36,16 @@ private[ui] object ToolTips { val SPARK_CONNECT_SESSION_DURATION = "Elapsed time since session start, or until closed if the session was closed" + val SPARK_CONNECT_SESSION_ML_CACHE = + "Current cached ML model usage for an active Spark Connect session" + + val SPARK_CONNECT_ML_CACHE_ESTIMATED_SIZE = + "Approximate model size recorded when it was added to the Spark Connect ML cache" + + val SPARK_CONNECT_ML_CACHE_MODEL_DETAILS = + "Output of model.toString recorded when the model was added to the cache" + + val SPARK_CONNECT_ML_CACHE_STORAGE = + "Whether the model is currently in driver memory or offloaded to driver-local disk" + } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala index 588493aeaa929..0a5a411093b79 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala @@ -276,7 +276,7 @@ private[connect] object ErrorUtils extends Logging { .newBuilder() .setCode(RPCCode.INTERNAL_VALUE) .addDetails(ProtoAny.pack(withStackTrace.build())) - .setMessage(SparkConnectService.extractErrorMessage(st)) + .setMessage(errorDescription(st)) .build() } @@ -287,6 +287,21 @@ private[connect] object ErrorUtils extends Logging { .exists(_.toString.contains("org.apache.spark.sql.execution.python")) } + /** + * Returns a non-empty description for the given throwable: its abbreviated message when one is + * present, and its fully qualified class name otherwise. Both the INTERNAL status built for a + * non-fatal throwable and the UNKNOWN fallback status use this, so that a client never receives + * an error whose description is empty. + */ + private def errorDescription(e: Throwable): String = { + val message = e.getMessage + if (message != null && message.nonEmpty) { + Utils.abbreviate(message, 2048) + } else { + e.getClass.getName + } + } + /** * Process an error by retrieving session context, converting to gRPC status, logging, posting * events, and executing callbacks. This is the core error handling logic shared by both @@ -343,7 +358,7 @@ private[connect] object ErrorUtils extends Logging { case e: Throwable => Status.UNKNOWN .withCause(e) - .withDescription(Utils.abbreviate(e.getMessage, 2048)) + .withDescription(errorDescription(e)) .asRuntimeException() } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTest.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTest.scala index 89e7ea90f8f49..13d9cf48638b0 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTest.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTest.scala @@ -40,7 +40,7 @@ trait SessionQueryTest extends sql.SessionQueryTest with SparkSessionBinder { /** * Approximates [[sql.SessionQueryTest.isDfSorted]] by inspecting the explain string. */ - override def isDfSorted(df: sql.DataFrame): Boolean = df match { + override def isDfSorted(df: org.apache.spark.sql.DataFrame): Boolean = df match { case df: DataFrame => sortOperator.unanchored.matches(df.explainString(extended = false)) case df => super.isDfSorted(df) } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTestBeforeAfterHooksSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTestBeforeAfterHooksSuite.scala new file mode 100644 index 0000000000000..9cc30139859bd --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTestBeforeAfterHooksSuite.scala @@ -0,0 +1,24 @@ +/* + * 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.connect + +import org.apache.spark.sql + +class SessionQueryTestBeforeAfterHooksSuite + extends sql.SessionQueryTestBeforeAfterHooksSuite + with SessionQueryTest diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionBinder.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionBinder.scala index b30bc55f7df14..fea98122ed8ae 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionBinder.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionBinder.scala @@ -38,7 +38,8 @@ trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite => protected override def spark: SparkSession = _connectSpark override protected def beforeAll(): Unit = { - super.beforeAll() + initializeSession() + // Other suites using mocks leave a mess in the global executionManager, // shut it down so that it's cleared before starting server. SparkConnectService.executionManager.shutdown() @@ -60,14 +61,18 @@ trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite => .builder() .client(client) .create() + super.beforeAll() } override def afterAll(): Unit = { - if (_connectSpark != null) { - _connectSpark.close() - _connectSpark = null + try { + super.afterAll() + } finally { + if (_connectSpark != null) { + _connectSpark.close() + _connectSpark = null + } + SparkConnectService.stop() } - SparkConnectService.stop() - super.afterAll() } } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionProvider.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionProvider.scala index 959e7fd899397..d139f2076df7b 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionProvider.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SparkSessionProvider.scala @@ -25,4 +25,6 @@ import org.apache.spark.sql */ trait SparkSessionProvider extends sql.SparkSessionProvider { protected override def spark: SparkSession + + override protected def sql(query: String): DataFrame = spark.sql(query) } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala index 9ba5a499ba8fe..76dca9c938a7f 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala @@ -387,6 +387,34 @@ class MLSuite extends MLHelper { } } + test("MLCache status") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_ENABLED.key, "true") + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_IN_MEMORY_SIZE.key, 16384) + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_STORAGE_SIZE.key, 65536) + + // Reading UI status should not initialize an unused cache. + assert(sessionHolder.getMLCacheStatus.isEmpty) + + val modelId = trainLogisticRegressionModel(sessionHolder) + val status = sessionHolder.getMLCacheStatus.get + assert(status.memoryControlEnabled) + assert(status.inMemorySizeBytes > 0) + assert(status.maxInMemorySizeBytes === 16384) + assert(status.totalSizeBytes === status.inMemorySizeBytes) + assert(status.maxTotalSizeBytes === 65536) + assert(status.models.size === 1) + val modelInfo = status.models.head + assert(modelInfo.id === modelId) + assert(modelInfo.className === classOf[LogisticRegressionModel].getName) + assert(modelInfo.modelString.startsWith("LogisticRegressionModel: uid=")) + assert(modelInfo.estimatedSizeBytes.contains(status.totalSizeBytes)) + assert(modelInfo.inMemory) + } + test("MLCache offloading works") { val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) sessionHolder.session.conf @@ -420,6 +448,11 @@ class MLSuite extends MLHelper { assert(sessionHolder.mlCache.totalMLCacheInMemorySizeBytes.get() <= memorySizeBytes) } + val status = sessionHolder.getMLCacheStatus.get + assert(status.models.size === modelIdList.size) + assert(status.models.count(_.inMemory) === maxNumModels) + assert(status.models.map(_.id).toSet === modelIdList.toSet) + // Assert all models can be loaded back from disk after they are offloaded. for (modelId <- modelIdList) { assert(sessionHolder.mlCache.get(modelId) != null) @@ -461,6 +494,7 @@ class MLSuite extends MLHelper { assert(mlCache2.get(modelId) != null) mlCache2.close() assert(mlCache2.cachedModel.isEmpty) + assert(mlCache2.getStatus.models.isEmpty) // Test 3: Edge case - register then remove model, close should still run cleanup val edgeCaseSessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/PythonPipelineSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/PythonPipelineSuite.scala index c21a4a9ac25bf..ba44051ee6926 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/PythonPipelineSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/PythonPipelineSuite.scala @@ -137,14 +137,16 @@ class PythonPipelineSuite TableIdentifier(catalog = Option("spark_catalog"), database = Option("default"), table = name) } + private def sessionCaseSensitive: Boolean = spark.sessionState.conf.caseSensitiveAnalysis + test("basic") { val graph = buildGraph(""" |@dp.table |def table1(): | return spark.readStream.format("rate").load() |""".stripMargin) - .resolve() - .validate() + .resolve(sessionCaseSensitive) + .validate(sessionCaseSensitive) assert(graph.flows.size == 1) assert(graph.tables.size == 1) } @@ -302,7 +304,7 @@ class PythonPipelineSuite |def a(): | return spark.range(5) |""".stripMargin) - val resolvedGraph = graph.resolve().validate() + val resolvedGraph = graph.resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert(resolvedGraph.tables.size == 4) assert(resolvedGraph.resolvedFlows.size == 4) } @@ -316,7 +318,7 @@ class PythonPipelineSuite |@dp.append_flow(target = "a") |def supplement(): | return spark.readStream.format("rate").load() - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert(graph.tables.map(_.identifier.table).toSet == Set("a")) assert(graph.resolvedFlows.size == 2) @@ -377,7 +379,7 @@ class PythonPipelineSuite |@dp.table |def d(): | return spark.sql("SELECT * FROM STREAM src") - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert( graph.table.keySet == Set( @@ -424,7 +426,7 @@ class PythonPipelineSuite |@dp.table |def e(): | return spark.sql("SELECT * FROM STREAM spark_catalog.default.src") - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert( graph.tables.map(_.identifier).toSet == Set( @@ -466,7 +468,7 @@ class PythonPipelineSuite |@dp.table |def e(): | return spark.sql("SELECT * FROM STREAM src") - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) assert(graph.resolutionFailedFlows.size == 5) graph.resolutionFailedFlows.foreach { flow => @@ -496,7 +498,7 @@ class PythonPipelineSuite |@dp.table |def e(): | return spark.sql("SELECT * FROM STREAM spark_catalog.default.src") - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) assert(graph.resolutionFailedFlows.size == 5) graph.resolutionFailedFlows.foreach { flow => assert(flow.failure.head.getMessage.contains("[TABLE_OR_VIEW_NOT_FOUND]")) @@ -517,7 +519,7 @@ class PythonPipelineSuite |@dp.materialized_view |def mv_from_read_table_df(): | return read_table_df - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert( graph.resolvedFlows.map(_.identifier).toSet == Set( @@ -541,7 +543,7 @@ class PythonPipelineSuite |def mv_from_read_table_df(): | return read_table_df | - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert( graph.resolvedFlows.map(_.identifier).toSet == Set( graphIdentifier("mv_from_read_table_df"), @@ -608,7 +610,7 @@ class PythonPipelineSuite |@dp.table(name = "schema_b.st_2") |def irrelevant_3(): | return spark.readStream.format("rate").load() - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) // validate these dataset are properly fully qualified assert( @@ -659,7 +661,7 @@ class PythonPipelineSuite |@dp.table(name = "some_catalog.some_schema.st") |def irrelevant_2(): | return spark.readStream.format("rate").load() - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) } assert(graphTry.isSuccess) assert( @@ -689,7 +691,7 @@ class PythonPipelineSuite |@dp.temporary_view(name= "view_3") |def irrelevant_2(): | return spark.read.table("view_1") - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) // views are temporary views, so they're not fully qualified. assert( Set("view_1", "view_2", "view_3").subsetOf( @@ -725,7 +727,7 @@ class PythonPipelineSuite |@dp.append_flow(target = "default.a") |def supplement(): | return spark.readStream.format("rate").load() - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) assert(graph.tables.map(_.identifier) == Seq(graphIdentifier("a"))) assert( @@ -894,8 +896,8 @@ class PythonPipelineSuite |def table_with_string_schema(): | return spark.range(5).withColumn("name", lit("test")) |""".stripMargin) - .resolve() - .validate() + .resolve(sessionCaseSensitive) + .validate(sessionCaseSensitive) assert(graph.flows.size == 1) assert(graph.tables.size == 1) @@ -917,8 +919,8 @@ class PythonPipelineSuite |def table_with_struct_schema(): | return spark.range(5).withColumn("name", lit("test")) |""".stripMargin) - .resolve() - .validate() + .resolve(sessionCaseSensitive) + .validate(sessionCaseSensitive) assert(graph.flows.size == 1) assert(graph.tables.size == 1) @@ -936,9 +938,9 @@ class PythonPipelineSuite |def table_with_wrong_schema(): | return spark.range(5).withColumn("wrong_column", lit("test")) |""".stripMargin) - .resolve() + .resolve(sessionCaseSensitive) - val ex = intercept[AnalysisException] { graph.validate() } + val ex = intercept[AnalysisException] { graph.validate(sessionCaseSensitive) } assert(ex.getMessage.contains("has a user-specified schema that is incompatible")) assert(ex.getMessage.contains("table_with_wrong_schema")) } @@ -955,9 +957,9 @@ class PythonPipelineSuite |def table_with_wrong_struct_schema(): | return spark.range(5).withColumn("different_column", lit("test")) |""".stripMargin) - .resolve() + .resolve(sessionCaseSensitive) - val ex = intercept[AnalysisException] { graph.validate() } + val ex = intercept[AnalysisException] { graph.validate(sessionCaseSensitive) } assert(ex.getMessage.contains("has a user-specified schema that is incompatible")) assert(ex.getMessage.contains("table_with_wrong_struct_schema")) } @@ -1103,6 +1105,26 @@ class PythonPipelineSuite assert(flow.destinationIdentifier == graphIdentifier("target")) } + test("AutoCDC API: spark_conf is forwarded to the flow's sqlConf") { + val flow = buildAutoCdcFlow(""" + |@dp.table + |def src(): + | return spark.readStream.format("rate").load() + | + |dp.create_streaming_table("target") + | + |dp.create_auto_cdc_flow( + | target = "target", + | source = "src", + | keys = ["value"], + | sequence_by = "timestamp", + | spark_conf = {"spark.sql.shuffle.partitions": "8"}, + |) + |""".stripMargin) + + assert(flow.sqlConf == Map("spark.sql.shuffle.partitions" -> "8")) + } + test("AutoCDC API: multi-part `keys` column is rejected at flow registration") { val ex = intercept[RuntimeException] { buildAutoCdcFlow(""" @@ -1185,7 +1207,7 @@ class PythonPipelineSuite | keys = ["value"], | sequence_by = "timestamp", |) - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) val resolvedFlow = graph.resolvedFlow(graphIdentifier("target")) assert(resolvedFlow.inputs == Set(graphIdentifier("src"))) @@ -1214,7 +1236,8 @@ class PythonPipelineSuite |""".stripMargin, defaultCatalog = Some("my_catalog"), defaultDatabase = Some("my_db"), - setupSql = Some("CREATE NAMESPACE IF NOT EXISTS my_catalog.my_db")).resolve() + setupSql = Some("CREATE NAMESPACE IF NOT EXISTS my_catalog.my_db")) + .resolve(sessionCaseSensitive) val resolvedFlow = graph.resolvedFlow(TableIdentifier("target", Some("my_db"), Some("my_catalog"))) @@ -1237,7 +1260,7 @@ class PythonPipelineSuite | keys = ["value"], | sequence_by = "timestamp", |) - |""".stripMargin).resolve() + |""".stripMargin).resolve(sessionCaseSensitive) val targetIdent = TableIdentifier("target", Some("some_schema"), Some("some_catalog")) val srcIdent = TableIdentifier("src", Some("some_schema"), Some("some_catalog")) @@ -1309,7 +1332,7 @@ class PythonPipelineSuite | apply_as_deletes = "value % 2 = 0", | column_list = ["value", "timestamp"], |) - |""".stripMargin).resolve().validate() + |""".stripMargin).resolve(sessionCaseSensitive).validate(sessionCaseSensitive) val resolvedFlow = graph.resolvedFlow(graphIdentifier("target")) assert(resolvedFlow.isInstanceOf[AutoCdcMergeFlow]) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala index 7ff8504a3e6ae..ec06d15ad8db5 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala @@ -30,6 +30,9 @@ import org.apache.spark.sql.connect.service.{SessionKey, SparkConnectService} class SparkDeclarativePipelinesServerSuite extends SparkDeclarativePipelinesServerTest with Logging { + + private def sessionCaseSensitive: Boolean = spark.sessionState.conf.caseSensitiveAnalysis + test("CreateDataflowGraph request creates a new graph") { withRawBlockingStub { implicit stub => assert(Option(createDataflowGraph(stub)).isDefined) @@ -265,7 +268,7 @@ class SparkDeclarativePipelinesServerSuite val definition = getDefaultSessionHolder.dataflowGraphRegistry.getDataflowGraphOrThrow(graphId) - val graph = definition.toDataflowGraph.resolve() + val graph = definition.toDataflowGraph.resolve(sessionCaseSensitive) assert(graph.flows.size == 3) assert(graph.tables.size == 2) @@ -312,7 +315,7 @@ class SparkDeclarativePipelinesServerSuite registerPipelineOutputs(pipeline) val graph = definition.toDataflowGraph - .resolve() + .resolve(sessionCaseSensitive) assert(graph.flows.size == 3) assert(graph.tables.size == 2) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/LiteralExpressionProtoConverterSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/LiteralExpressionProtoConverterSuite.scala index 9c176a7a54f2b..242c241d2b7e4 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/LiteralExpressionProtoConverterSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/LiteralExpressionProtoConverterSuite.scala @@ -17,17 +17,20 @@ package org.apache.spark.sql.connect.planner -import java.time.LocalTime +import java.time.{Instant, LocalDateTime, LocalTime} import org.scalatest.funsuite.AnyFunSuite // scalastyle:ignore funsuite +import org.apache.spark.SparkException import org.apache.spark.connect.proto import org.apache.spark.sql.catalyst.{expressions, CatalystTypeConverters} import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema +import org.apache.spark.sql.connect.common.DataTypeProtoConverter import org.apache.spark.sql.connect.common.InvalidPlanInput import org.apache.spark.sql.connect.common.LiteralValueProtoConverter import org.apache.spark.sql.connect.common.LiteralValueProtoConverter.ToLiteralProtoOptions import org.apache.spark.sql.connect.planner.LiteralExpressionProtoConverter +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ class LiteralExpressionProtoConverterSuite extends AnyFunSuite { // scalastyle:ignore funsuite @@ -77,6 +80,127 @@ class LiteralExpressionProtoConverterSuite extends AnyFunSuite { // scalastyle:i assertResult(LocalTime.of(1, 2, 3))(LiteralValueProtoConverter.toScalaValue(literalProto)) } + test("SPARK-57161: nanosecond timestamp DataType proto round-trip across precisions") { + for (precision <- + TimestampNTZNanosType.MIN_PRECISION to TimestampNTZNanosType.MAX_PRECISION) { + for (dt <- Seq(TimestampNTZNanosType(precision), TimestampLTZNanosType(precision))) { + val protoType = DataTypeProtoConverter.toConnectProtoType(dt) + assertResult(dt)(DataTypeProtoConverter.toCatalystType(protoType)) + } + } + } + + test("SPARK-57161: TIMESTAMP_NTZ nanosecond literal proto and catalyst value round-trip") { + // Boundary and pre-epoch values, plus a sub-microsecond value that exercises the extra nanos. + val values = Seq( + LocalDateTime.of(1, 1, 1, 0, 0, 0, 0), + LocalDateTime.of(1969, 12, 31, 23, 59, 59, 999999999), + LocalDateTime.of(1970, 1, 1, 0, 0, 0, 123456789), + LocalDateTime.of(2023, 6, 15, 12, 34, 56, 987654321), + LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999999)) + for (precision <- TimestampNTZNanosType.MIN_PRECISION to TimestampNTZNanosType.MAX_PRECISION; + v <- values) { + val t = TimestampNTZNanosType(precision) + val literalProto = toLiteralProto(v, t) + // The literal carries the nanos proto arm with the expected precision. + assert(literalProto.getTimestampNtzNanos.getPrecision == precision) + // Scala value -> proto -> Catalyst value equals converting the Scala value directly, so the + // same sub-microsecond truncation to `precision` is applied on both paths. + val convert = CatalystTypeConverters.createToCatalystConverter(t) + val expected = expressions.Literal(convert(v), t) + assertResult(expected)(LiteralExpressionProtoConverter.toCatalystExpression(literalProto)) + } + } + + test("SPARK-57161: TIMESTAMP_LTZ nanosecond literal proto and catalyst value round-trip") { + val values = Seq( + Instant.parse("0001-01-01T00:00:00Z"), + Instant.parse("1969-12-31T23:59:59.999999999Z"), + Instant.parse("1970-01-01T00:00:00.123456789Z"), + Instant.parse("2023-06-15T12:34:56.987654321Z"), + Instant.parse("9999-12-31T23:59:59.999999999Z")) + for (precision <- TimestampLTZNanosType.MIN_PRECISION to TimestampLTZNanosType.MAX_PRECISION; + v <- values) { + val t = TimestampLTZNanosType(precision) + val literalProto = toLiteralProto(v, t) + assert(literalProto.getTimestampLtzNanos.getPrecision == precision) + val convert = CatalystTypeConverters.createToCatalystConverter(t) + val expected = expressions.Literal(convert(v), t) + assertResult(expected)(LiteralExpressionProtoConverter.toCatalystExpression(literalProto)) + } + } + + test("SPARK-57161: nanosecond timestamp literal proto carries epoch micros + extra nanos") { + // 1970-01-01T00:00:00.000001500Z is 1 microsecond and 500 extra nanoseconds past the epoch. + val ntzProto = toLiteralProto( + LocalDateTime.of(1970, 1, 1, 0, 0, 0, 1500), + TimestampNTZNanosType(TimestampNTZNanosType.NANOS_PRECISION)) + assert(ntzProto.getTimestampNtzNanos.getEpochMicros == 1L) + assert(ntzProto.getTimestampNtzNanos.getNanosWithinMicro == 500) + + val ltzProto = toLiteralProto( + Instant.parse("1970-01-01T00:00:00.000001500Z"), + TimestampLTZNanosType(TimestampLTZNanosType.NANOS_PRECISION)) + assert(ltzProto.getTimestampLtzNanos.getEpochMicros == 1L) + assert(ltzProto.getTimestampLtzNanos.getNanosWithinMicro == 500) + } + + test( + "SPARK-57161: nanosecond timestamp literal with out-of-range nanos_within_micro is " + + "rejected") { + // nanos_within_micro is an int32 on the wire but must be in [0, 999]. Build literals directly + // with out-of-range values, including 65536 which would truncate to 0 if narrowed to Short + // before validation, and confirm the read path rejects them instead of wrapping. + for (nanos <- Seq(1000, 65536, 65636, -1)) { + val ntzProto = proto.Expression.Literal + .newBuilder() + .setTimestampNtzNanos( + proto.Expression.Literal.TimestampNTZNanos + .newBuilder() + .setEpochMicros(0L) + .setNanosWithinMicro(nanos) + .setPrecision(TimestampNTZNanosType.NANOS_PRECISION)) + .build() + val ltzProto = proto.Expression.Literal + .newBuilder() + .setTimestampLtzNanos( + proto.Expression.Literal.TimestampLTZNanos + .newBuilder() + .setEpochMicros(0L) + .setNanosWithinMicro(nanos) + .setPrecision(TimestampLTZNanosType.NANOS_PRECISION)) + .build() + for (literalProto <- Seq(ntzProto, ltzProto)) { + val e = intercept[InvalidPlanInput] { + LiteralValueProtoConverter.toScalaValue(literalProto) + } + assert(e.getMessage.contains("nanos_within_micro")) + } + } + } + + test("SPARK-57161: nanosecond timestamp literals are rejected when the feature is disabled") { + // Build the proto with the feature enabled (default in tests), mirroring a message arriving + // over the wire, then convert it on the server with the feature turned off. + val ntzProto = toLiteralProto( + LocalDateTime.of(2023, 1, 1, 0, 0, 0, 0), + TimestampNTZNanosType(TimestampNTZNanosType.NANOS_PRECISION)) + val ltzProto = toLiteralProto( + Instant.parse("2023-01-01T00:00:00Z"), + TimestampLTZNanosType(TimestampLTZNanosType.NANOS_PRECISION)) + + val disabledConf = new SQLConf() + disabledConf.setConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED, false) + SQLConf.withExistingConf(disabledConf) { + for (literalProto <- Seq(ntzProto, ltzProto)) { + val e = intercept[SparkException] { + LiteralExpressionProtoConverter.toCatalystExpression(literalProto) + } + assert(e.getCondition == "FEATURE_NOT_ENABLED") + } + } + } + // The goal of this test is to check that converting a Scala value -> Proto -> Catalyst value // is equivalent to converting a Scala value directly to a Catalyst value. Seq[(Any, DataType)]( diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/GetStatusHandlerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/GetStatusHandlerSuite.scala index 21d96fd8a87d8..e8e960fdb578b 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/GetStatusHandlerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/GetStatusHandlerSuite.scala @@ -88,6 +88,28 @@ class NoOpGetStatusPlugin extends GetStatusPlugin { Optional.empty() } +/** + * A plugin that returns a list containing a null element, to test null filtering. + */ +class NullElementGetStatusPlugin extends GetStatusPlugin { + override def processRequestExtensions( + sessionHolder: SessionHolder, + requestExtensions: util.List[protobuf.Any]): Optional[util.List[protobuf.Any]] = + listWithNull() + + override def processOperationExtensions( + operationId: String, + sessionHolder: SessionHolder, + operationExtensions: util.List[protobuf.Any]): Optional[util.List[protobuf.Any]] = + listWithNull() + + private def listWithNull(): Optional[util.List[protobuf.Any]] = { + val result = new util.ArrayList[protobuf.Any]() + result.add(null) + Optional.of(result) + } +} + /** * A plugin that always throws a RuntimeException. */ @@ -344,6 +366,32 @@ class GetStatusHandlerSuite extends SharedSparkSession { assert(opExtValues.contains(s"op-echo:${executeHolder.operationId}:safe")) assert(opExtValues.contains(s"second-op:${executeHolder.operationId}:safe")) } + + test("GetStatus filters null extension elements returned by a plugin") { + SparkConnectPluginRegistry.setGetStatusPluginsForTesting( + Seq(new NullElementGetStatusPlugin(), new EchoGetStatusPlugin())) + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + val command = proto.Command.newBuilder().build() + val executeHolder = SparkConnectTestUtils.createDummyExecuteHolder(sessionHolder, command) + + val reqExt = protobuf.Any.pack(StringValue.of("data")) + val opExt = protobuf.Any.pack(StringValue.of("data")) + val response = sendGetOperationStatusRequest( + sessionHolder.sessionId, + operationIds = Seq(executeHolder.operationId), + userId = sessionHolder.userId, + requestExtensions = Seq(reqExt), + operationExtensions = Seq(opExt)) + + // The null element is dropped; only the healthy plugin's extension survives, at both levels. + val responseExtValues = response.getExtensionsList.asScala + .map(_.unpack(classOf[StringValue]).getValue) + assert(responseExtValues == Seq("request-echo:data")) + + val opExtValues = response.getOperationStatusesList.asScala.head.getExtensionsList.asScala + .map(_.unpack(classOf[StringValue]).getValue) + assert(opExtValues == Seq(s"op-echo:${executeHolder.operationId}:data")) + } } private class GetStatusResponseObserver extends StreamObserver[proto.GetStatusResponse] { diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectAuthSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectAuthSuite.scala index 30f186ab7c2b1..93c910c19aca2 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectAuthSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectAuthSuite.scala @@ -23,8 +23,13 @@ import org.apache.spark.sql.connect.{SparkConnectServerTest, SparkSession} import org.apache.spark.sql.connect.service.SparkConnectService class SparkConnectAuthSuite extends SparkConnectServerTest { + private val tokenKey = "spark.connect.authenticate.token" + private val markerKey = "spark.connect.test.marker" + override protected def sparkConf = { - super.sparkConf.set("spark.connect.authenticate.token", "deadbeef") + super.sparkConf + .set(tokenKey, "deadbeef") + .set(markerKey, "visible") } test("Test local authentication") { @@ -43,4 +48,19 @@ class SparkConnectAuthSuite extends SparkConnectServerTest { } assert(exception.getMessage.contains("Invalid authentication token")) } + + test("Test the authentication token is not readable through the Config RPC") { + val session = SparkSession + .builder() + .remote(s"sc://localhost:${SparkConnectService.localPort}/;token=deadbeef") + .create() + + assert(session.conf.getOption(tokenKey).isEmpty) + assert(session.conf.get(tokenKey, "absent") === "absent") + intercept[NoSuchElementException](session.conf.get(tokenKey)) + assert(!session.conf.getAll.contains(tokenKey)) + + // Server-side configurations that are not sensitive stay readable. + assert(session.conf.get(markerKey) === "visible") + } } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandlerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandlerSuite.scala new file mode 100644 index 0000000000000..269f7836506bb --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandlerSuite.scala @@ -0,0 +1,172 @@ +/* + * 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.connect.service + +import scala.concurrent.Promise +import scala.concurrent.duration._ +import scala.jdk.CollectionConverters._ + +import io.grpc.stub.StreamObserver + +import org.apache.spark.SparkNoSuchElementException +import org.apache.spark.connect.proto +import org.apache.spark.internal.config.SECRET_REDACTION_PATTERN +import org.apache.spark.sql.connect.SparkConnectTestUtils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.util.ThreadUtils + +class SparkConnectConfigHandlerSuite extends SharedSparkSession { + + // Matches the default spark.redaction.regex on "password". + private val secretKey = "spark.test.connect.password" + private val secretValue = "hunter2" + private val plainKey = "spark.test.connect.endpoint" + private val plainValue = "localhost:15002" + + protected override def afterEach(): Unit = { + super.afterEach() + SparkConnectService.sessionManager.invalidateAllSessions() + } + + private def sendConfigRequest( + sessionHolder: SessionHolder, + customize: proto.ConfigRequest.Operation.Builder => Unit): proto.ConfigResponse = { + val operation = proto.ConfigRequest.Operation.newBuilder() + customize(operation) + val request = proto.ConfigRequest + .newBuilder() + .setUserContext(proto.UserContext.newBuilder().setUserId(sessionHolder.userId).build()) + .setSessionId(sessionHolder.sessionId) + .setOperation(operation) + .build() + val responseObserver = new ConfigResponseObserver() + new SparkConnectConfigHandler(responseObserver).handle(request) + ThreadUtils.awaitResult(responseObserver.promise.future, 10.seconds) + } + + /** The returned pairs, with an absent value mapped to None. */ + private def pairs(response: proto.ConfigResponse): Map[String, Option[String]] = { + response.getPairsList.asScala + .map(pair => pair.getKey -> (if (pair.hasValue) Some(pair.getValue) else None)) + .toMap + } + + test("GetAll does not return keys matching the redaction pattern") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + withSQLConf(secretKey -> secretValue, plainKey -> plainValue) { + val returned = pairs(sendConfigRequest(sessionHolder, _.getGetAllBuilder)) + assert(!returned.contains(secretKey)) + assert(returned(plainKey) === Some(plainValue)) + } + } + + test("GetAll matches the redaction pattern on the full key, not the prefixed one") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + // Stripping the prefix leaves "value", which no longer matches the pattern. The filter has to + // run before the prefix comes off. + val prefix = "spark.test.connect.secret." + withSQLConf(prefix + "value" -> secretValue) { + val returned = pairs(sendConfigRequest(sessionHolder, _.getGetAllBuilder.setPrefix(prefix))) + assert(returned.isEmpty) + } + } + + test("a secret carried by an innocuous key is withheld too") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + // `SET spark.test.connect.jdbc.url` already masks this value, because SetCommand redacts + // through Utils.redact, which matches the value as well. The Config RPC has to agree. + val urlKey = "spark.test.connect.jdbc.url" + val urlValue = "jdbc:postgresql://db:5432/app?user=app&password=hunter2" + withSQLConf(urlKey -> urlValue) { + assert(!pairs(sendConfigRequest(sessionHolder, _.getGetAllBuilder)).contains(urlKey)) + val returned = + pairs(sendConfigRequest(sessionHolder, _.getGetOptionBuilder.addKeys(urlKey))) + assert(returned(urlKey) === None) + } + } + + test("Get reports a redacted key the way it reports an unset one") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + withSQLConf(secretKey -> secretValue) { + val redacted = intercept[SparkNoSuchElementException] { + sendConfigRequest(sessionHolder, _.getGetBuilder.addKeys(secretKey)) + } + val unset = intercept[SparkNoSuchElementException] { + sendConfigRequest(sessionHolder, _.getGetBuilder.addKeys("spark.test.connect.absent")) + } + assert(redacted.getCondition === unset.getCondition) + } + } + + test("GetOption returns no value for a redacted key") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + withSQLConf(secretKey -> secretValue, plainKey -> plainValue) { + val returned = pairs( + sendConfigRequest( + sessionHolder, + _.getGetOptionBuilder.addKeys(secretKey).addKeys(plainKey))) + assert(returned(secretKey) === None) + assert(returned(plainKey) === Some(plainValue)) + } + } + + test("GetWithDefault returns the caller's default for a redacted key") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + withSQLConf(secretKey -> secretValue) { + val returned = pairs( + sendConfigRequest( + sessionHolder, + _.getGetWithDefaultBuilder.addPairsBuilder().setKey(secretKey).setValue("fallback"))) + assert(returned(secretKey) === Some("fallback")) + } + } + + test("the redaction pattern is not read from the session config") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + // Both of these are reachable through the Set operation: the legacy flag is what otherwise + // stops a client from writing spark.redaction.regex into its own session config. + withSQLConf( + SQLConf.SET_COMMAND_REJECTS_SPARK_CORE_CONFS.key -> "false", + secretKey -> secretValue) { + spark.conf.set(SECRET_REDACTION_PATTERN.key, "matches-no-key") + try { + val returned = pairs(sendConfigRequest(sessionHolder, _.getGetAllBuilder)) + assert(!returned.contains(secretKey)) + } finally { + spark.conf.unset(SECRET_REDACTION_PATTERN.key) + } + } + } + + test("IsModifiable still answers for a redacted key") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + withSQLConf(secretKey -> secretValue) { + // Whether a key is modifiable is a property of the key, so it discloses no value. + val returned = + pairs(sendConfigRequest(sessionHolder, _.getIsModifiableBuilder.addKeys(secretKey))) + assert(returned(secretKey) === Some("false")) + } + } +} + +private class ConfigResponseObserver extends StreamObserver[proto.ConfigResponse] { + val promise: Promise[proto.ConfigResponse] = Promise() + override def onNext(value: proto.ConfigResponse): Unit = promise.success(value) + override def onError(t: Throwable): Unit = promise.failure(t) + override def onCompleted(): Unit = {} +} diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala index a433534b7511a..56e69486e1fa4 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceE2ESuite.scala @@ -18,14 +18,17 @@ package org.apache.spark.sql.connect.service import java.io.ByteArrayOutputStream import java.util.UUID +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicReference import com.github.luben.zstd.{Zstd, ZstdOutputStreamNoFinalizer} import com.google.protobuf.ByteString import org.scalatest.concurrent.Eventually import org.scalatest.time.SpanSugar._ -import org.apache.spark.SparkException +import org.apache.spark.{SparkContext, SparkException} import org.apache.spark.connect.proto +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart} import org.apache.spark.sql.connect.SparkConnectServerTest import org.apache.spark.sql.connect.config.Connect @@ -37,6 +40,32 @@ class SparkConnectServiceE2ESuite extends SparkConnectServerTest { // were all already in the buffer. val BIG_ENOUGH_QUERY = "select * from range(1000000)" + test("ExecutePlan operation ID is available as a Spark local property") { + val operationIdFromJob = new AtomicReference[String]() + val jobStarted = new CountDownLatch(1) + val listener = new SparkListener { + override def onJobStart(jobStart: SparkListenerJobStart): Unit = { + val operationId = + jobStart.properties.getProperty(SparkContext.SPARK_CONNECT_OPERATION_ID_PROPERTY) + if (operationId != null) { + operationIdFromJob.set(operationId) + jobStarted.countDown() + } + } + } + spark.sparkContext.addSparkListener(listener) + try { + withClient { client => + val responses = client.execute(buildPlan("select count(*) from range(10)")).toSeq + val operationId = responses.head.getOperationId + assert(jobStarted.await(10, TimeUnit.SECONDS)) + assert(operationIdFromJob.get() == operationId) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + test("Execute is sent eagerly to the server upon iterator creation") { // This behavior changed with grpc upgrade from 1.56.0 to 1.59.0. // Testing to be aware of future changes. diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceInternalServerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceInternalServerSuite.scala index 173dc5c672bc3..eb52fdb8da64a 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceInternalServerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceInternalServerSuite.scala @@ -302,6 +302,7 @@ private class SparkConnectServiceLifeCycleListener extends SparkListener { SparkConnectServiceLifeCycleListener.checksOnServiceEndEvent.foreach { checks => checks.foreach(_(serviceEnd)) } + case _ => } } } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala index a0b249a002aba..ccf7a00ae1f20 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectServiceKeepAliveSuite.scala @@ -108,9 +108,8 @@ class SparkConnectServiceKeepAliveSuite extends SparkConnectServerTest { FiniteDuration(15, TimeUnit.SECONDS)) } // scalastyle:on awaitresult - // A keepalive-triggered UNAVAILABLE carries no wrapped cause (same as - // DEADLINE_EXCEEDED, see GrpcExceptionConverter.toThrowable), so the status - // code/description is only in the message. + // The status code/description of a keepalive-triggered UNAVAILABLE is part of the + // message (see GrpcExceptionConverter.toThrowable). assert(ex.getMessage.contains("UNAVAILABLE")) assert(ex.getMessage.contains("Keepalive failed")) } finally { diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala index 4b846631d7b73..5fae135bcf18a 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManagerSuite.scala @@ -76,6 +76,7 @@ class SparkConnectSessionManagerSuite extends SharedSparkSession { Some(sessionHolder.session.sessionUUID + "invalid")) } assert(exGet.getCondition == "INVALID_HANDLE.SESSION_CHANGED") + assert(exGet.getSqlState == "08003") } test( @@ -89,11 +90,13 @@ class SparkConnectSessionManagerSuite extends SharedSparkSession { SparkConnectService.sessionManager.getOrCreateIsolatedSession(key, None) } assert(exGetOrCreate.getCondition == "INVALID_HANDLE.SESSION_CLOSED") + assert(exGetOrCreate.getSqlState == "08003") val exGet = intercept[SparkSQLException] { SparkConnectService.sessionManager.getIsolatedSession(key, None) } assert(exGet.getCondition == "INVALID_HANDLE.SESSION_CLOSED") + assert(exGet.getSqlState == "08003") val sessionGetIfPresent = SparkConnectService.sessionManager.getIsolatedSessionIfPresent(key) assert(sessionGetIfPresent.isEmpty) @@ -106,6 +109,7 @@ class SparkConnectSessionManagerSuite extends SharedSparkSession { SparkConnectService.sessionManager.getIsolatedSession(key, None) } assert(exGet.getCondition == "INVALID_HANDLE.SESSION_NOT_FOUND") + assert(exGet.getSqlState == "08003") val sessionGetIfPresent = SparkConnectService.sessionManager.getIsolatedSessionIfPresent(key) assert(sessionGetIfPresent.isEmpty) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala index 7f6af17bc41b3..8adaa729d2bc6 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.connect.ui -import java.util.{Calendar, Locale} +import java.util.{Calendar, Locale, UUID} import jakarta.servlet.http.HttpServletRequest import org.apache.commons.text.StringEscapeUtils @@ -26,8 +26,11 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.{SharedSparkContext, SparkConf, SparkFunSuite} import org.apache.spark.scheduler.SparkListenerJobStart +import org.apache.spark.sql.classic.SparkSession +import org.apache.spark.sql.connect.ml.{MLCacheModelInfo, MLCacheStatus} import org.apache.spark.sql.connect.service._ import org.apache.spark.status.ElementTrackingStore +import org.apache.spark.ui.SparkUI import org.apache.spark.util.kvstore.InMemoryStore class SparkConnectServerPageSuite @@ -47,7 +50,10 @@ class SparkConnectServerPageSuite /** * Run a dummy session and return the store */ - private def getStatusStore: SparkConnectServerAppStatusStore = { + private def getStatusStore( + closeSession: Boolean = true, + sessionId: String = "sessionId", + userId: String = "userId"): SparkConnectServerAppStatusStore = { kvstore = new ElementTrackingStore(new InMemoryStore, new SparkConf()) // val server = mock(classOf[SparkConnectServer], RETURNS_SMART_NULLS) val sparkConf = new SparkConf @@ -56,14 +62,14 @@ class SparkConnectServerPageSuite val statusStore = new SparkConnectServerAppStatusStore(kvstore) listener.onOtherEvent( - SparkListenerConnectSessionStarted("sessionId", "userId", System.currentTimeMillis())) + SparkListenerConnectSessionStarted(sessionId, userId, System.currentTimeMillis())) listener.onOtherEvent( SparkListenerConnectOperationStarted( "jobTag", "operationId", System.currentTimeMillis(), - "sessionId", - "userId", + sessionId, + userId, "userName", "dummy query", Set())) @@ -74,14 +80,16 @@ class SparkConnectServerPageSuite SparkListenerConnectOperationFinished("jobTag", "operationId", System.currentTimeMillis())) listener.onOtherEvent( SparkListenerConnectOperationClosed("jobTag", "operationId", System.currentTimeMillis())) - listener.onOtherEvent( - SparkListenerConnectSessionClosed("sessionId", "userId", System.currentTimeMillis())) + if (closeSession) { + listener.onOtherEvent( + SparkListenerConnectSessionClosed(sessionId, userId, System.currentTimeMillis())) + } statusStore } test("Spark Connect Server page should load successfully") { - val store = getStatusStore + val store = getStatusStore() val request = mock(classOf[HttpServletRequest]) val tab = mock(classOf[SparkConnectServerTab], RETURNS_SMART_NULLS) @@ -89,6 +97,7 @@ class SparkConnectServerPageSuite when(tab.store).thenReturn(store) when(tab.appName).thenReturn("testing") when(tab.headerTabs).thenReturn(Seq.empty) + when(tab.getMLCacheStatuses).thenReturn(None) val page = new SparkConnectServerPage(tab) val html = page.render(request).toString().toLowerCase(Locale.ROOT) @@ -96,6 +105,9 @@ class SparkConnectServerPageSuite assert(html.contains("session statistics (1)")) assert(html.contains("request statistics (1)")) assert(html.contains("dummy query")) + assert(html.contains("ml cache")) + assert(html.contains("<span>n/a</span>")) + assert(!html.contains("ml cache statistics")) // Pagination support assert(html.contains("<label class=\"text-nowrap\">1 pages. jump to</label>")) @@ -107,7 +119,7 @@ class SparkConnectServerPageSuite } test("Spark Connect Server session page should load successfully") { - val store = getStatusStore + val store = getStatusStore() val request = mock(classOf[HttpServletRequest]) when(request.getParameter("id")).thenReturn("sessionId") @@ -134,6 +146,162 @@ class SparkConnectServerPageSuite " data-bs-target=\"#aggregated-sqlsessionstat\"")) } + test("Spark Connect Server page should show live ML cache statistics and model details") { + val store = getStatusStore(closeSession = false) + + val request = mock(classOf[HttpServletRequest]) + val tab = mock(classOf[SparkConnectServerTab], RETURNS_SMART_NULLS) + when(tab.startTime).thenReturn(Calendar.getInstance().getTime) + when(tab.store).thenReturn(store) + when(tab.appName).thenReturn("testing") + when(tab.headerTabs).thenReturn(Seq.empty) + val status = MLCacheStatus( + memoryControlEnabled = true, + inMemorySizeBytes = 1024, + maxInMemorySizeBytes = 4096, + totalSizeBytes = 2048, + maxTotalSizeBytes = 8192, + models = Seq( + MLCacheModelInfo( + id = "model-id-1", + className = "org.apache.spark.ml.classification.LogisticRegressionModel", + modelString = "LogisticRegressionModel: uid=logreg-1", + estimatedSizeBytes = Some(1024), + inMemory = true), + MLCacheModelInfo( + id = "model-id-2", + className = "org.apache.spark.ml.classification.LogisticRegressionModel", + modelString = "LogisticRegressionModel: uid=logreg-2", + estimatedSizeBytes = Some(1024), + inMemory = false))) + when(tab.getMLCacheStatuses).thenReturn( + Some(Map(SessionKey("userId", "sessionId") -> Some(status)))) + + val page = new SparkConnectServerPage(tab) + val html = page.render(request).toString().toLowerCase(Locale.ROOT) + + val sessionStatsIndex = html.indexOf("session statistics") + val mlCacheStatsIndex = html.indexOf("ml cache statistics (2)") + val requestStatsIndex = html.indexOf("request statistics") + assert(sessionStatsIndex < requestStatsIndex && requestStatsIndex < mlCacheStatsIndex) + assert(html.contains("1 model in memory")) + assert(html.contains("1024.0 b / 4.0 kib memory")) + assert(html.contains("2.0 kib / 8.0 kib total")) + assert(html.contains("2 (1 in memory, 1 offloaded)")) + assert(html.contains("estimated size (in-memory)")) + assert(html.contains("1024.0 b / 4.0 kib")) + assert(html.contains("estimated size (in-memory and offloaded data)")) + assert(html.contains("2.0 kib / 8.0 kib")) + assert(html.contains("model-id-1")) + assert(html.contains("logisticregressionmodel: uid=logreg-1")) + assert(html.contains("in memory")) + assert(html.contains("offloaded")) + } + + test("Spark Connect Server page should show unused and unavailable ML cache states") { + val store = getStatusStore(closeSession = false) + + val request = mock(classOf[HttpServletRequest]) + val tab = mock(classOf[SparkConnectServerTab], RETURNS_SMART_NULLS) + when(tab.startTime).thenReturn(Calendar.getInstance().getTime) + when(tab.store).thenReturn(store) + when(tab.appName).thenReturn("testing") + when(tab.headerTabs).thenReturn(Seq.empty) + + when(tab.getMLCacheStatuses).thenReturn(None) + val unavailableHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(unavailableHtml.contains("<span>n/a</span>")) + + when(tab.getMLCacheStatuses).thenReturn(Some(Map(SessionKey("userId", "sessionId") -> None))) + val uninitializedHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(uninitializedHtml.contains("<span>not used</span>")) + + val clearedStatus = MLCacheStatus( + memoryControlEnabled = true, + inMemorySizeBytes = 0, + maxInMemorySizeBytes = 4096, + totalSizeBytes = 0, + maxTotalSizeBytes = 8192, + models = Seq.empty) + when(tab.getMLCacheStatuses).thenReturn( + Some(Map(SessionKey("userId", "sessionId") -> Some(clearedStatus)))) + val clearedHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(clearedHtml.contains("<span>not used</span>")) + assert(!clearedHtml.contains("ml cache statistics")) + + when(tab.getMLCacheStatuses).thenReturn(Some(Map.empty)) + val removedHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(removedHtml.contains("<span>n/a</span>")) + + val disabledStatus = MLCacheStatus( + memoryControlEnabled = false, + inMemorySizeBytes = 0, + maxInMemorySizeBytes = 0, + totalSizeBytes = 0, + maxTotalSizeBytes = 0, + models = Seq( + MLCacheModelInfo( + id = "model-id", + className = "model-class", + modelString = "model-details", + estimatedSizeBytes = None, + inMemory = true))) + when(tab.getMLCacheStatuses).thenReturn( + Some(Map(SessionKey("userId", "sessionId") -> Some(disabledStatus)))) + val disabledHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(disabledHtml.contains("memory control disabled")) + assert(disabledHtml.contains("1 cached model")) + } + + test("Spark Connect Server page should read ML cache status without touching the session") { + val key = SessionKey("userId", UUID.randomUUID().toString) + val store = + getStatusStore(closeSession = false, sessionId = key.sessionId, userId = key.userId) + val sessionManager = new SparkConnectSessionManager() + sessionManager.initializeBaseSession(() => new SparkSession(sc)) + val sessionHolder = sessionManager.getOrCreateIsolatedSession(key, None) + + val request = mock(classOf[HttpServletRequest]) + val sparkUI = SparkUI.create( + Some(sc), + sc.statusStore, + sc.conf, + sc.env.securityManager, + sc.appName, + "", + sc.startTime) + val tab = + new SparkConnectServerTab(store, sparkUI, Some(sessionManager)) + + try { + val lastAccessTime = sessionHolder.getSessionHolderInfo.lastAccessTimeMs + Thread.sleep(10) + val unusedHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(unusedHtml.contains("<span>not used</span>")) + assert(sessionHolder.getMLCacheStatus.isEmpty) + assert(sessionHolder.getSessionHolderInfo.lastAccessTimeMs === lastAccessTime) + + sessionHolder.mlCache + val initializedLastAccessTime = sessionHolder.getSessionHolderInfo.lastAccessTimeMs + Thread.sleep(10) + val initializedHtml = + new SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT) + assert(initializedHtml.contains("<span>not used</span>")) + assert(sessionHolder.getMLCacheStatus.exists(_.models.isEmpty)) + assert(sessionHolder.getSessionHolderInfo.lastAccessTimeMs === initializedLastAccessTime) + } finally { + tab.detach() + sessionManager.closeSession(key) + sessionManager.shutdown() + } + } + test("SPARK-58097: session page only shows the requested user's operations") { // Two users share the same session UUID, each running a distinct query. kvstore = new ElementTrackingStore(new InMemoryStore, new SparkConf()) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/utils/ErrorUtilsSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/utils/ErrorUtilsSuite.scala new file mode 100644 index 0000000000000..97b9c8db673b2 --- /dev/null +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/utils/ErrorUtilsSuite.scala @@ -0,0 +1,73 @@ +/* + * 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.connect.utils + +import java.util.UUID + +import com.google.rpc.Code +import io.grpc.{Status, StatusRuntimeException} +import io.grpc.stub.StreamObserver + +import org.apache.spark.sql.test.SharedSparkSession + +class ErrorUtilsSuite extends SharedSparkSession { + + test("handleError fallback uses throwable class when fatal throwable has no message") { + val observer = new StreamObserver[Unit] { + override def onNext(value: Unit): Unit = { + fail(s"Unexpected response: $value") + } + + override def onError(t: Throwable): Unit = { + throw t + } + + override def onCompleted(): Unit = { + fail("Unexpected completion") + } + } + + val error = intercept[StatusRuntimeException] { + ErrorUtils.handleError("execute", observer, "user1", UUID.randomUUID().toString)( + new InterruptedException()) + } + + assert(error.getStatus.getCode == Status.Code.UNKNOWN) + assert(error.getStatus.getDescription == classOf[InterruptedException].getName) + } + + test("buildStatusFromThrowable uses throwable class when non-fatal throwable has no message") { + val status = ErrorUtils.buildStatusFromThrowable(new RuntimeException(), None) + + assert(status.getCode == Code.INTERNAL_VALUE) + assert(status.getMessage == classOf[RuntimeException].getName) + } + + test("buildStatusFromThrowable uses throwable class when non-fatal message is empty") { + val status = ErrorUtils.buildStatusFromThrowable(new RuntimeException(""), None) + + assert(status.getCode == Code.INTERNAL_VALUE) + assert(status.getMessage == classOf[RuntimeException].getName) + } + + test("buildStatusFromThrowable keeps the message when the non-fatal throwable has one") { + val status = ErrorUtils.buildStatusFromThrowable(new RuntimeException("boom"), None) + + assert(status.getCode == Code.INTERNAL_VALUE) + assert(status.getMessage == "boom") + } +} diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt new file mode 100644 index 0000000000000..dd08f21eca704 --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk21-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 4069 4145 108 2.1 485.0 1.0X +codegen = true, adaptive = T 2585 2632 67 3.2 308.1 1.6X +codegen = false, adaptive = F 4818 4849 44 1.7 574.3 0.8X +codegen = false, adaptive = T 2983 3150 236 2.8 355.6 1.4X + + +================================================================================================ +low-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 281 292 7 59.6 16.8 1.0X +codegen = true, adaptive = T 291 323 52 57.7 17.3 1.0X +codegen = false, adaptive = F 1289 1303 20 13.0 76.8 0.2X +codegen = false, adaptive = T 1266 1270 7 13.3 75.4 0.2X + + +================================================================================================ +high-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 8120 8126 9 1.0 968.0 1.0X +codegen = true, adaptive = T 4460 4476 22 1.9 531.7 1.8X +codegen = false, adaptive = F 9331 9389 82 0.9 1112.3 0.9X +codegen = false, adaptive = T 5369 5382 18 1.6 640.1 1.5X + + +================================================================================================ +low-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 768 782 12 21.8 45.8 1.0X +codegen = true, adaptive = T 769 780 13 21.8 45.8 1.0X +codegen = false, adaptive = F 1366 1369 4 12.3 81.4 0.6X +codegen = false, adaptive = T 1358 1358 0 12.4 80.9 0.6X + + diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt new file mode 100644 index 0000000000000..f54c40877aafb --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-jdk25-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1021-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 4338 4408 100 1.9 517.1 1.0X +codegen = true, adaptive = T 2586 2608 30 3.2 308.3 1.7X +codegen = false, adaptive = F 5146 5199 76 1.6 613.4 0.8X +codegen = false, adaptive = T 3271 3287 22 2.6 390.0 1.3X + + +================================================================================================ +low-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1021-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 269 275 4 62.5 16.0 1.0X +codegen = true, adaptive = T 284 294 6 59.1 16.9 0.9X +codegen = false, adaptive = F 1317 1321 6 12.7 78.5 0.2X +codegen = false, adaptive = T 1309 1382 104 12.8 78.0 0.2X + + +================================================================================================ +high-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1021-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 8332 8462 184 1.0 993.2 1.0X +codegen = true, adaptive = T 4677 4683 9 1.8 557.5 1.8X +codegen = false, adaptive = F 9932 9966 47 0.8 1184.0 0.8X +codegen = false, adaptive = T 5412 5618 291 1.6 645.1 1.5X + + +================================================================================================ +low-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1021-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 773 786 17 21.7 46.1 1.0X +codegen = true, adaptive = T 790 803 20 21.2 47.1 1.0X +codegen = false, adaptive = F 1370 1377 10 12.2 81.6 0.6X +codegen = false, adaptive = T 1366 1373 11 12.3 81.4 0.6X + + diff --git a/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt new file mode 100644 index 0000000000000..5523166ca01fb --- /dev/null +++ b/sql/core/benchmarks/AdaptivePartialAggregationBenchmark-results.txt @@ -0,0 +1,56 @@ +================================================================================================ +high-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +codegen = true, adaptive = F 4118 4125 11 2.0 490.9 1.0X +codegen = true, adaptive = T 2411 2427 23 3.5 287.4 1.7X +codegen = false, adaptive = F 4846 4861 21 1.7 577.7 0.8X +codegen = false, adaptive = T 2991 3009 25 2.8 356.6 1.4X + + +================================================================================================ +low-cardinality input, pass-through at the periodic check +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, no spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 270 311 37 62.2 16.1 1.0X +codegen = true, adaptive = T 293 318 22 57.3 17.5 0.9X +codegen = false, adaptive = F 1305 1312 10 12.9 77.8 0.2X +codegen = false, adaptive = T 1329 1435 150 12.6 79.2 0.2X + + +================================================================================================ +high-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, high card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 8196 8247 73 1.0 977.0 1.0X +codegen = true, adaptive = T 4433 4530 137 1.9 528.5 1.8X +codegen = false, adaptive = F 9544 9555 15 0.9 1137.8 0.9X +codegen = false, adaptive = T 5243 5395 215 1.6 625.0 1.6X + + +================================================================================================ +low-cardinality input, pass-through at the spill check +================================================================================================ + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1020-azure +AMD EPYC 7763 64-Core Processor +adaptive partial agg, low card, spill: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +codegen = true, adaptive = F 803 813 11 20.9 47.8 1.0X +codegen = true, adaptive = T 829 839 9 20.2 49.4 1.0X +codegen = false, adaptive = F 1362 1387 35 12.3 81.2 0.6X +codegen = false, adaptive = T 1365 1376 15 12.3 81.4 0.6X + + diff --git a/sql/core/benchmarks/ExpandBenchmark-jdk21-results.txt b/sql/core/benchmarks/ExpandBenchmark-jdk21-results.txt index a4c9d9a062e27..9cc85852fafce 100644 --- a/sql/core/benchmarks/ExpandBenchmark-jdk21-results.txt +++ b/sql/core/benchmarks/ExpandBenchmark-jdk21-results.txt @@ -2,58 +2,69 @@ Expand: varying number of COUNT(DISTINCT) ================================================================================================ -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 7763 64-Core Processor +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro 2 distinct aggregates: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -with sum - baseline (ratio 3) 3150 3221 71 3.3 300.4 1.0X -with sum - optimized (ratio 3) 3095 3136 33 3.4 295.2 1.0X -pure distinct - baseline (ratio 2) 2038 2064 26 5.1 194.4 1.5X -pure distinct - optimized (ratio 2) 675 711 31 15.5 64.3 4.7X +with sum - baseline (ratio 3) 2518 2821 425 4.2 240.1 1.0X +with sum - optimized (ratio 3) 2363 2429 52 4.4 225.4 1.1X +pure distinct - baseline (ratio 2) 1382 1471 149 7.6 131.8 1.8X +pure distinct - optimized (ratio 2) 400 409 11 26.2 38.2 6.3X -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 7763 64-Core Processor +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro 4 distinct aggregates: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -with sum - baseline (ratio 5) 6006 6042 30 1.7 572.8 1.0X -with sum - optimized (ratio 5) 5981 6030 33 1.8 570.4 1.0X -pure distinct - baseline (ratio 4) 4540 4602 49 2.3 432.9 1.3X -pure distinct - optimized (ratio 4) 767 800 22 13.7 73.2 7.8X +with sum - baseline (ratio 5) 4569 4650 60 2.3 435.8 1.0X +with sum - optimized (ratio 5) 4584 4644 41 2.3 437.2 1.0X +pure distinct - baseline (ratio 4) 3393 3471 64 3.1 323.6 1.3X +pure distinct - optimized (ratio 4) 471 510 44 22.3 44.9 9.7X -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 7763 64-Core Processor +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro 6 distinct aggregates: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -with sum - baseline (ratio 7) 9223 9268 50 1.1 879.6 1.0X -with sum - optimized (ratio 7) 9253 9346 122 1.1 882.4 1.0X -pure distinct - baseline (ratio 6) 7420 7484 45 1.4 707.6 1.2X -pure distinct - optimized (ratio 6) 918 969 38 11.4 87.6 10.0X +with sum - baseline (ratio 7) 6311 6527 288 1.7 601.8 1.0X +with sum - optimized (ratio 7) 5659 5853 206 1.9 539.7 1.1X +pure distinct - baseline (ratio 6) 5145 5184 28 2.0 490.6 1.2X +pure distinct - optimized (ratio 6) 506 518 10 20.7 48.2 12.5X -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 7763 64-Core Processor +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro 8 distinct aggregates: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -with sum - baseline (ratio 9) 13382 13516 91 0.8 1276.2 1.0X -with sum - optimized (ratio 9) 13465 13538 108 0.8 1284.1 1.0X -pure distinct - baseline (ratio 8) 11329 11434 110 0.9 1080.4 1.2X -pure distinct - optimized (ratio 8) 1869 1920 37 5.6 178.3 7.2X +with sum - baseline (ratio 9) 8226 8286 36 1.3 784.5 1.0X +with sum - optimized (ratio 9) 8218 8272 46 1.3 783.8 1.0X +pure distinct - baseline (ratio 8) 7086 7134 40 1.5 675.8 1.2X +pure distinct - optimized (ratio 8) 955 976 12 11.0 91.1 8.6X ================================================================================================ Expand: varying data characteristics (pure distinct) ================================================================================================ -OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1020-azure -AMD EPYC 7763 64-Core Processor +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro 6 pure distinct aggs with varying data: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative ------------------------------------------------------------------------------------------------------------------------ -1K groups, moderate card - baseline 7437 7471 24 1.4 709.3 1.0X -1K groups, moderate card - optimized 898 922 18 11.7 85.7 8.3X -100K groups, moderate card - baseline 15063 15226 135 0.7 1436.5 0.5X -100K groups, moderate card - optimized 4454 4561 95 2.4 424.8 1.7X -1K groups, low card (5 vals) - baseline 7293 7342 40 1.4 695.5 1.0X -1K groups, low card (5 vals) - optimized 762 783 21 13.8 72.6 9.8X -no grouping key - baseline 4888 4948 39 2.1 466.1 1.5X -no grouping key - optimized 377 379 2 27.8 35.9 19.7X +1K groups, moderate card - baseline 5175 5218 37 2.0 493.6 1.0X +1K groups, moderate card - optimized 497 507 10 21.1 47.4 10.4X +100K groups, moderate card - baseline 9323 9420 130 1.1 889.1 0.6X +100K groups, moderate card - optimized 1993 2069 139 5.3 190.1 2.6X +1K groups, low card (5 vals) - baseline 5277 5315 23 2.0 503.2 1.0X +1K groups, low card (5 vals) - optimized 466 474 8 22.5 44.5 11.1X +no grouping key - baseline 3577 3610 23 2.9 341.1 1.4X +no grouping key - optimized 218 221 5 48.2 20.8 23.8X +================================================================================================ +Expand: subexpression elimination across branches +================================================================================================ +OpenJDK 64-Bit Server VM 21.0.3+9-LTS on Mac OS X 26.5.2 +Apple M3 Pro +9 conditional COUNT(DISTINCT) + 9 conditional SUM sharing one subexpression: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +----------------------------------------------------------------------------------------------------------------------------------------------------------- +rewrite off, CSE off (10x amplify, 18 evals/row) 84450 84850 382 0.1 16107.6 1.0X +rewrite on, CSE off (2x amplify, 18 evals/row) 41354 41544 312 0.1 7887.7 2.0X +rewrite on, CSE on (2x amplify, 1 eval/row) 10712 10843 120 0.5 2043.2 7.9X +rewrite off, CSE on (10x amplify, 1 eval/row) 53281 54021 722 0.1 10162.6 1.6X diff --git a/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk21-results.txt b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk21-results.txt new file mode 100644 index 0000000000000..f91bddaedecb9 --- /dev/null +++ b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk21-results.txt @@ -0,0 +1,35 @@ +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip all row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 1948 2008 87 10.8 92.9 1.0X +With shredded predicate pushdown 92 107 10 227.1 4.4 21.1X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 2030 2048 10 10.3 96.8 1.0X +With shredded predicate pushdown 108 115 5 194.1 5.2 18.8X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 2440 2468 28 8.6 116.3 1.0X +With shredded predicate pushdown 2500 2530 21 8.4 119.2 1.0X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups (default block size): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 2166 2196 23 9.7 103.3 1.0X +With shredded predicate pushdown 2175 2211 27 9.6 103.7 1.0X + +OpenJDK 64-Bit Server VM 21.0.12+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups (partial object): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 2347 2355 10 8.9 111.9 1.0X +With shredded predicate pushdown 164 174 9 127.6 7.8 14.3X + diff --git a/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk25-results.txt b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk25-results.txt new file mode 100644 index 0000000000000..71a0f414335e5 --- /dev/null +++ b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-jdk25-results.txt @@ -0,0 +1,35 @@ +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip all row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 1515 1574 79 13.8 72.3 1.0X +With shredded predicate pushdown 67 76 8 312.8 3.2 22.6X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 1507 1526 18 13.9 71.9 1.0X +With shredded predicate pushdown 84 91 8 250.7 4.0 18.0X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 1876 1893 17 11.2 89.5 1.0X +With shredded predicate pushdown 1940 1958 19 10.8 92.5 1.0X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups (default block size): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 1658 1695 38 12.6 79.1 1.0X +With shredded predicate pushdown 1659 1666 10 12.6 79.1 1.0X + +OpenJDK 64-Bit Server VM 25.0.4+7-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups (partial object): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 1772 1789 20 11.8 84.5 1.0X +With shredded predicate pushdown 116 123 7 180.2 5.5 15.2X + diff --git a/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-results.txt b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-results.txt new file mode 100644 index 0000000000000..78372cd131f1a --- /dev/null +++ b/sql/core/benchmarks/VariantShreddedPredicatePushdownBenchmark-results.txt @@ -0,0 +1,35 @@ +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip all row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 2118 2175 132 9.9 101.0 1.0X +With shredded predicate pushdown 91 103 10 229.9 4.3 23.2X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 2133 2162 29 9.8 101.7 1.0X +With shredded predicate pushdown 112 124 9 186.8 5.4 19.0X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------ +Without shredded predicate pushdown 2562 2583 21 8.2 122.2 1.0X +With shredded predicate pushdown 2637 2663 26 8.0 125.7 1.0X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip no row groups (default block size): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +--------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 2261 2296 40 9.3 107.8 1.0X +With shredded predicate pushdown 2260 2284 21 9.3 107.8 1.0X + +OpenJDK 64-Bit Server VM 17.0.20+8-LTS on Linux 6.17.0-1022-azure +AMD EPYC 9V74 80-Core Processor +Can skip some row groups (partial object): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative +------------------------------------------------------------------------------------------------------------------------- +Without shredded predicate pushdown 2461 2475 9 8.5 117.3 1.0X +With shredded predicate pushdown 156 161 5 134.3 7.4 15.8X + diff --git a/sql/core/pom.xml b/sql/core/pom.xml index d96b71a9d8058..281546b901a7a 100644 --- a/sql/core/pom.xml +++ b/sql/core/pom.xml @@ -273,8 +273,8 @@ <artifactId>htmlunit3-driver</artifactId> <scope>test</scope> </dependency> - <!-- Explicit declaration of bouncy-castle dependencies are - needed for maven test builds on later hadoop releases.--> + <!-- Explicit declaration of bouncy-castle dependencies is + needed for Maven test builds on later Hadoop releases. --> <dependency> <groupId>org.bouncycastle</groupId> <artifactId>bcprov-jdk18on</artifactId> @@ -348,7 +348,7 @@ so that the tests classes of external modules can use them. The two execution profiles are necessary - first one for 'mvn package', second one for 'mvn test-compile'. Ideally, 'mvn compile' should not compile test classes and therefore should not need this. - However, a closed due to "Cannot Reproduce" Maven bug (https://issues.apache.org/jira/browse/MNG-3559) + However, a Maven bug closed as "Cannot Reproduce" (https://issues.apache.org/jira/browse/MNG-3559) causes the compilation to fail if catalyst test-jar is not generated. Hence, the second execution profile for 'mvn test-compile'. --> diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java index af8d5a4610f64..d850d0d18befe 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/UnsafeFixedWidthAggregationMap.java @@ -227,6 +227,14 @@ public double getAvgHashProbesPerKey() { return map.getAvgHashProbesPerKey(); } + /** + * Returns the number of distinct keys currently stored in the underlying `BytesToBytesMap`. + * Used by adaptive partial aggregation to estimate the pre-shuffle reduction ratio. + */ + public int getNumKeys() { + return map.numKeys(); + } + /** * Sorts the map's records in place, spill them to disk, and returns an [[UnsafeKVExternalSorter]] * diff --git a/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java index 3f67d050a8c8e..859a2fdf60d91 100644 --- a/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java +++ b/sql/core/src/main/java/org/apache/spark/sql/execution/aggregate/RowBasedAggregateHashMap.java @@ -110,6 +110,14 @@ public final KVIterator<UnsafeRow, UnsafeRow> rowIterator() { return batch.rowIterator(); } + /** + * Returns the number of distinct keys currently held by this map. Used by adaptive partial + * aggregation to measure the reduction ratio across both aggregation maps. + */ + public final int getNumKeys() { + return numRows; + } + @Override public final void close() { batch.close(); diff --git a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/allexecutionspage.js b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/allexecutionspage.js index 5973ab896de1b..e06b8bd804628 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/allexecutionspage.js +++ b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/allexecutionspage.js @@ -17,7 +17,7 @@ /* global $, uiRoot, appBasePath, createSqlApiBase, getSqlTableColumns, withResolvedAppId, statusBadge, jobIdLinks, formatDurationSql, - descriptionHtml */ + formatTotalTaskTime, descriptionHtml */ $(document).ready(function () { // Read the cluster-level grouping toggle rendered into the page by Scala @@ -132,7 +132,8 @@ $(document).ready(function () { var html = '<table id="' + childId + '" class="table table-sm table-bordered mb-0 sub-exec-table">'; html += '<thead><tr><th>ID</th><th>Status</th><th>Description</th>' + - '<th>Duration</th><th>Succeeded Jobs</th></tr></thead><tbody>'; + '<th>Duration</th><th>Total Task Time</th>' + + '<th>Succeeded Jobs</th></tr></thead><tbody>'; subs.forEach(function (child) { html += '<tr><td><a href="' + basePath + '/SQL/execution/?id=' + child.id + '">' + child.id + '</a></td>'; @@ -141,6 +142,7 @@ $(document).ready(function () { id: child.id, description: child.description || "" }) + '</td>'; html += '<td>' + formatDurationSql(child.duration) + '</td>'; + html += '<td>' + formatTotalTaskTime(child.totalTaskTime) + '</td>'; html += '<td>' + jobIdLinks(child.jobIds || []) + '</td></tr>'; }); html += '</tbody></table>'; diff --git a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/executionpage.js b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/executionpage.js index 4cb9a65c05100..1f50c202219ec 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/executionpage.js +++ b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/executionpage.js @@ -43,6 +43,7 @@ $(document).ready(function () { description: data.description || "", submissionTime: data.submissionTime, duration: data.duration, + totalTaskTime: data.totalTaskTime, jobIds: data.successJobIds || [], errorMessage: data.errorMessage || "" }; diff --git a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/sql-table-utils.js b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/sql-table-utils.js index 7b0f4ffccde58..ecefa11139a42 100644 --- a/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/sql-table-utils.js +++ b/sql/core/src/main/resources/org/apache/spark/sql/execution/ui/static/sql-table-utils.js @@ -30,6 +30,14 @@ function formatDurationSql(milliseconds) { return hours.toFixed(1) + " h"; } +// Format a total task time value. A negative or absent value means "unknown" +// (e.g. the execution has no stages to aggregate), which is shown as "N/A" +// rather than a misleading "0 ms". +function formatTotalTaskTime(value) { + if (value === null || value === undefined || value < 0) return "N/A"; + return formatDurationSql(value); +} + function formatDateSql(dateStr) { if (!dateStr) return ""; try { @@ -232,6 +240,14 @@ function getSqlTableColumns(opts) { } }; + var totalTaskTimeColumn = { + data: "totalTaskTime", name: "totalTaskTime", title: "Total Task Time", + render: function (data, type) { + if (type !== "display") return data; + return formatTotalTaskTime(data); + } + }; + var jobsColumn = { data: "jobIds", name: "jobIds", title: "Succeeded Jobs", orderable: false, @@ -258,5 +274,6 @@ function getSqlTableColumns(opts) { }; return [idColumn, queryIdColumn, statusColumn, descriptionColumn, - submissionColumn, durationColumn, jobsColumn, errorColumn]; + submissionColumn, durationColumn, totalTaskTimeColumn, jobsColumn, + errorColumn]; } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala index 804b5269c929c..9b22cf4d9dd6a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/artifact/ArtifactManager.scala @@ -36,7 +36,7 @@ import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.internal.config.{CONNECT_SCALA_UDF_STUB_PREFIXES, EXECUTOR_USER_CLASS_PATH_FIRST} import org.apache.spark.sql.Artifact import org.apache.spark.sql.classic.SparkSession -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.util.ArtifactUtils import org.apache.spark.storage.{BlockManager, CacheId, StorageLevel} import org.apache.spark.util.{ChildFirstURLClassLoader, StubClassLoader, Utils} @@ -504,10 +504,8 @@ class ArtifactManager(session: SparkSession) extends AutoCloseable with Logging val localPath = serverLocalStagingPath val fs = destFSPath.getFileSystem(hadoopConf) if (fs.isInstanceOf[LocalFileSystem]) { - val allowDestLocalConf = - session.sessionState.conf.getConf(SQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL) - .getOrElse( - session.conf.get("spark.connect.copyFromLocalToFs.allowDestLocal").contains("true")) + val allowDestLocalConf = session.sessionState.conf.getConf( + StaticSQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL) if (!allowDestLocalConf) { // To avoid security issue, by default, @@ -517,7 +515,7 @@ class ArtifactManager(session: SparkSession) extends AutoCloseable with Logging // We can temporarily allow the behavior by setting spark config // `spark.sql.artifact.copyFromLocalToFs.allowDestLocal` // to `true` when starting spark driver, we should only enable it for testing - // purpose. + // purpose. It is a static conf, so it cannot be set from a session. throw new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_3161") } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala index ce16c4a2cc3ae..25d0c75849081 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroDeserializer.scala @@ -30,6 +30,7 @@ import org.apache.avro.Schema.Type._ import org.apache.avro.generic._ import org.apache.avro.util.Utf8 +import org.apache.spark.SparkRuntimeException import org.apache.spark.sql.avro.AvroUtils.{nonNullUnionBranches, toFieldStr, AvroMatchedField} import org.apache.spark.sql.catalyst.{InternalRow, NoopFilters, StructFilters} import org.apache.spark.sql.catalyst.expressions.{SpecificInternalRow, UnsafeArrayData} @@ -249,7 +250,7 @@ private[sql] class AvroDeserializer( case (DOUBLE, DoubleType) => (updater, ordinal, value) => updater.setDouble(ordinal, value.asInstanceOf[Double]) - case (STRING, StringType) => (updater, ordinal, value) => + case (STRING, _: StringType) => (updater, ordinal, value) => val str = value match { case s: String => UTF8String.fromString(s) case s: Utf8 => @@ -259,7 +260,7 @@ private[sql] class AvroDeserializer( } updater.set(ordinal, str) - case (ENUM, StringType) => (updater, ordinal, value) => + case (ENUM, _: StringType) => (updater, ordinal, value) => updater.set(ordinal, UTF8String.fromString(value.toString)) case (FIXED, BinaryType) => (updater, ordinal, value) => @@ -275,7 +276,8 @@ private[sql] class AvroDeserializer( bytes case b: Array[Byte] => b case other => - throw new RuntimeException(errorPrefix + s"$other is not a valid avro binary.") + throw new IncompatibleSchemaException( + errorPrefix + s"$other is not a valid Avro binary.") } updater.set(ordinal, bytes) @@ -329,8 +331,9 @@ private[sql] class AvroDeserializer( val element = iter.next() if (element == null) { if (!containsNull) { - throw new RuntimeException( - s"Array value at path ${toFieldStr(avroElementPath)} is not allowed to be null") + throw new SparkRuntimeException( + errorClass = "AVRO_CANNOT_READ_NULL_FIELD", + messageParameters = Map("name" -> toFieldStr(avroElementPath))) } else { elementUpdater.setNullAt(i) } @@ -342,8 +345,8 @@ private[sql] class AvroDeserializer( updater.set(ordinal, result) - case (MAP, MapType(keyType, valueType, valueContainsNull)) if keyType == StringType => - val keyWriter = newWriter(SchemaBuilder.builder().stringType(), StringType, + case (MAP, MapType(keyType: StringType, valueType, valueContainsNull)) => + val keyWriter = newWriter(SchemaBuilder.builder().stringType(), keyType, avroPath :+ "key", catalystPath :+ "key") val valueWriter = newWriter(avroType.getValueType, valueType, avroPath :+ "value", catalystPath :+ "value") @@ -361,8 +364,9 @@ private[sql] class AvroDeserializer( keyWriter(keyUpdater, i, entry.getKey) if (entry.getValue == null) { if (!valueContainsNull) { - throw new RuntimeException( - s"Map value at path ${toFieldStr(avroPath :+ "value")} is not allowed to be null") + throw new SparkRuntimeException( + errorClass = "AVRO_CANNOT_READ_NULL_FIELD", + messageParameters = Map("name" -> toFieldStr(avroPath :+ "value"))) } else { valueUpdater.setNullAt(i) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala index c2a05a8c9d8b9..5988724070961 100755 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala @@ -199,7 +199,9 @@ private[sql] class AvroFileFormat extends FileFormat } else { new NoopFilters } - SupportsArchiveFormat.readArchiveEntries(file.toPath, conf) { (_, in) => + val entryGlob = parsedOptions.archivePathFilterPattern + SupportsArchiveFormat.readArchiveEntries( + file.toPath, conf, archivePathFilter = entryGlob) { (_, in) => val datumReader = userProvidedSchema match { case Some(schema) => new GenericDatumReader[GenericRecord](schema) case None => new GenericDatumReader[GenericRecord]() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala index d5405f69ff051..353ac2eecbf11 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroSerializer.scala @@ -137,7 +137,7 @@ private[sql] class AvroSerializer( decimalConversions.toBytes(decimal.toJavaBigDecimal, avroType, LogicalTypes.decimal(d.precision, d.scale)) - case (StringType, ENUM) => + case (_: StringType, ENUM) => val enumSymbols: Set[String] = avroType.getEnumSymbols.asScala.toSet (getter, ordinal) => val data = getter.getUTF8String(ordinal).toString @@ -148,7 +148,7 @@ private[sql] class AvroSerializer( } new EnumSymbol(avroType, data) - case (StringType, STRING) => + case (_: StringType, STRING) => (getter, ordinal) => new Utf8(getter.getUTF8String(ordinal).getBytes) case (BinaryType, FIXED) => @@ -259,7 +259,7 @@ private[sql] class AvroSerializer( case (LongType, UNION) if nonNullUnionTypes(avroType) == Set(INT, LONG) => (getter, ordinal) => getter.getLong(ordinal) - case (MapType(kt, vt, valueContainsNull), MAP) if kt == StringType => + case (MapType(_: StringType, vt, valueContainsNull), MAP) => val valueConverter = newConverter( vt, resolveNullableType(avroType.getValueType, valueContainsNull), catalystPath :+ "value", avroPath :+ "value") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala index a1bf478f14f10..266e6ee835ced 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala @@ -27,7 +27,7 @@ import org.apache.avro.generic.{GenericDatumReader, GenericRecord} import org.apache.avro.mapred.{AvroOutputFormat, FsInput} import org.apache.avro.mapreduce.AvroJob import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.mapreduce.Job import org.apache.spark.{SparkException, SparkIllegalArgumentException} @@ -97,7 +97,8 @@ private[sql] object AvroUtils extends Logging { } if (archives.nonEmpty) { inferAvroSchemaFromArchives(archives, nonArchives, conf, parsedOptions.ignoreExtension, - fileSourceOptions.ignoreCorruptFiles, fileSourceOptions.ignoreMissingFiles) + fileSourceOptions.ignoreCorruptFiles, fileSourceOptions.ignoreMissingFiles, + fileSourceOptions.archivePathFilterPattern) } else { inferAvroSchemaFromFiles(files, conf, parsedOptions.ignoreExtension, fileSourceOptions.ignoreCorruptFiles) @@ -250,10 +251,12 @@ private[sql] object AvroUtils extends Logging { conf: Configuration, ignoreExtension: Boolean, ignoreCorruptFiles: Boolean, - ignoreMissingFiles: Boolean): Schema = { + ignoreMissingFiles: Boolean, + archivePathFilter: Option[GlobPattern]): Schema = { archives.iterator .flatMap { f => - firstArchiveEntrySchema(f.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles) + firstArchiveEntrySchema( + f.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles, archivePathFilter) } .nextOption() .getOrElse { @@ -274,11 +277,13 @@ private[sql] object AvroUtils extends Logging { path: Path, conf: Configuration, ignoreCorruptFiles: Boolean, - ignoreMissingFiles: Boolean): Option[Schema] = { + ignoreMissingFiles: Boolean, + archivePathFilter: Option[GlobPattern]): Option[Schema] = { try { // `readArchiveEntries` returns a Closeable iterator; take the first entry's schema and close // it so the archive stream is released without draining the remaining entries. - val entries = SupportsArchiveFormat.readArchiveEntries(path, conf) { (_, in) => + val entries = SupportsArchiveFormat.readArchiveEntries( + path, conf, archivePathFilter = archivePathFilter) { (_, in) => val stream = new DataFileStream[GenericRecord](in, new GenericDatumReader[GenericRecord]()) try { Iterator.single(stream.getSchema) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala index a5700a0481531..aeeb608c83ffc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/CustomDecimal.scala @@ -33,7 +33,7 @@ private[spark] class CustomDecimal(schema: Schema) extends LogicalType(CustomDec val obj = schema.getObjectProp("scale") obj match { case null => - throw new IllegalArgumentException(s"Invalid ${CustomDecimal.TYPE_NAME}: missing scale"); + throw new IllegalArgumentException(s"Invalid ${CustomDecimal.TYPE_NAME}: missing scale") case i : Integer => i case other => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala b/sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala index c6fba163309f1..9ecbe1b6a0f59 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/avro/SchemaConverters.scala @@ -32,6 +32,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{FIELD_NAME, FIELD_TYPE, RECURSIVE_DEPTH} import org.apache.spark.sql.avro.AvroOptions.RECURSIVE_FIELD_MAX_DEPTH_LIMIT import org.apache.spark.sql.catalyst.parser.CatalystSqlParser +import org.apache.spark.sql.catalyst.util.CharVarcharUtils import org.apache.spark.sql.types._ import org.apache.spark.sql.types.Decimal.minBytesForPrecision @@ -95,6 +96,33 @@ object SchemaConverters extends Logging { // The property specifies Catalyst type of the given field private val CATALYST_TYPE_PROP_NAME = "spark.sql.catalyst.type" + // Avro map keys are always STRING; stamp CHAR/VARCHAR key types on the map schema. + private val CATALYST_MAP_KEY_TYPE_PROP_NAME = "spark.sql.catalyst.mapKey.type" + + private def avroStringSchema(catalystType: StringType): Schema = { + val stringSchema = SchemaBuilder.builder().stringType() + CharVarcharUtils.charVarcharTypeName(catalystType).foreach { name => + stringSchema.addProp(CATALYST_TYPE_PROP_NAME, name) + } + stringSchema + } + + private def avroMapSchema(keyType: StringType, valueSchema: Schema): Schema = { + val mapSchema = SchemaBuilder.builder().map().values(valueSchema) + CharVarcharUtils.charVarcharTypeName(keyType).foreach { name => + mapSchema.addProp(CATALYST_MAP_KEY_TYPE_PROP_NAME, name) + } + mapSchema + } + + private def parseStampedStringType(catalystTypeAttrValue: String): StringType = { + CatalystSqlParser.parseDataType(catalystTypeAttrValue) match { + case s: StringType => s + case other => + throw new IncompatibleSchemaException( + s"Avro $CATALYST_TYPE_PROP_NAME for STRING must be a STRING subtype, got $other") + } + } private def toSqlTypeHelper( avroSchema: Schema, @@ -114,7 +142,14 @@ object SchemaConverters extends Logging { } SchemaType(catalystType, nullable = false) } - case STRING => SchemaType(StringType, nullable = false) + case STRING => + val catalystTypeAttrValue = avroSchema.getProp(CATALYST_TYPE_PROP_NAME) + val catalystType = if (catalystTypeAttrValue == null) { + StringType + } else { + parseStampedStringType(catalystTypeAttrValue) + } + SchemaType(catalystType, nullable = false) case BOOLEAN => SchemaType(BooleanType, nullable = false) case BYTES | FIXED => avroSchema.getLogicalType match { // For FIXED type, if the precision requires more bytes than fixed size, the logical @@ -244,8 +279,14 @@ object SchemaConverters extends Logging { ) null } else { + val keyAttr = avroSchema.getProp(CATALYST_MAP_KEY_TYPE_PROP_NAME) + val keyType = if (keyAttr == null) { + StringType + } else { + parseStampedStringType(keyAttr) + } SchemaType( - MapType(StringType, schemaType.dataType, valueContainsNull = schemaType.nullable), + MapType(keyType, schemaType.dataType, valueContainsNull = schemaType.nullable), nullable = false) } @@ -369,7 +410,9 @@ object SchemaConverters extends Logging { case FloatType => builder.floatType() case DoubleType => builder.doubleType() - case StringType => builder.stringType() + // CharType/VarcharType are StringType subclasses, not the StringType singleton. + // Stamp spark.sql.catalyst.type so inference restores the length constraint. + case s: StringType => avroStringSchema(s) case NullType => builder.nullType() case d: DecimalType => val avroType = LogicalTypes.decimal(d.precision, d.scale) @@ -385,9 +428,8 @@ object SchemaConverters extends Logging { case ArrayType(et, containsNull) => builder.array() .items(toAvroType(et, containsNull, recordName, nameSpace)) - case MapType(StringType, vt, valueContainsNull) => - builder.map() - .values(toAvroType(vt, valueContainsNull, recordName, nameSpace)) + case MapType(kt: StringType, vt, valueContainsNull) => + avroMapSchema(kt, toAvroType(vt, valueContainsNull, recordName, nameSpace)) case st: StructType => val childNameSpace = if (nameSpace != "") s"$nameSpace.$recordName" else recordName val fieldsAssembler = builder.record(recordName).namespace(nameSpace).fields() @@ -489,7 +531,7 @@ object SchemaConverters extends Logging { case LongType => builder.longType() case FloatType => builder.floatType() case DoubleType => builder.doubleType() - case StringType => builder.stringType() + case s: StringType => avroStringSchema(s) case NullType => builder.nullType() case DateType => LogicalTypes.date().addToSchema(builder.intType()) case TimestampType => LogicalTypes.timestampMicros().addToSchema(builder.longType()) @@ -522,9 +564,10 @@ object SchemaConverters extends Logging { // Make array types nullable Schema.createUnion(nullSchema, arraySchema) - case MapType(StringType, valueType, _) => - val mapSchema = builder.map() - .values(toAvroTypeWithDefaults(valueType, recordName = recordName, + case MapType(kt: StringType, valueType, _) => + val mapSchema = avroMapSchema( + kt, + toAvroTypeWithDefaults(valueType, recordName = recordName, namespace = namespace, nestingLevel = nestingLevel + 1)) // Make map types nullable Schema.createUnion(nullSchema, mapSchema) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala index 2f139393ade38..8015d3d0f9cff 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveDataSource.scala @@ -25,7 +25,6 @@ import org.apache.spark.sql.catalyst.analysis.NamedStreamingRelation import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnresolvedDataSource} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 -import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connector.catalog.{SupportsRead, TableProvider} @@ -112,7 +111,7 @@ class ResolveDataSource(sparkSession: SparkSession) extends Rule[LogicalPlan] { import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ StreamingRelationV2( Some(provider), source, table, dsOptions, - toAttributes(table.columns.asSchema), None, None, v1Relation, + table.columns.toOutputAttributes, None, None, v1Relation, v1DataSource.streamingSourceIdentifyingName) // fallback to v1 @@ -173,7 +172,7 @@ class ResolveDataSource(sparkSession: SparkSession) extends Rule[LogicalPlan] { import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ StreamingRelationV2( Some(provider), source, table, dsOptions, - toAttributes(table.columns.asSchema), None, None, v1Relation, + table.columns.toOutputAttributes, None, None, v1Relation, v1DataSource.streamingSourceIdentifyingName) // fallback to v1 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala index 01f39e60afd3c..e6fc6d8d862ce 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala @@ -762,8 +762,9 @@ class ResolveSessionCatalog(val catalogManager: CatalogManager) if (provider.isDefined) { // The parser guarantees that USING and STORED AS/ROW FORMAT won't co-exist. if (maybeSerdeInfo.isDefined) { - throw QueryCompilationErrors.cannotCreateTableWithBothProviderAndSerdeError( - provider, maybeSerdeInfo) + throw SparkException.internalError( + s"Cannot create table with both USING ${provider.get} and " + + s"${maybeSerdeInfo.get.describe}") } (nonHiveStorageFormat, provider.get) } else if (maybeSerdeInfo.isDefined) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/expressions/ParseSql.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/expressions/ParseSql.scala new file mode 100644 index 0000000000000..649feeee9adf7 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/expressions/ParseSql.scala @@ -0,0 +1,113 @@ +/* + * 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.catalyst.expressions + +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.catalyst.analysis.{FunctionRegistry, FunctionRegistryBase, TypeCheckResult} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.parser.ParseSqlResult +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.types.StringTypeWithCollation +import org.apache.spark.sql.types.{AbstractDataType, DataType, StringType} +import org.apache.spark.unsafe.types.UTF8String + +/** + * Parses a SQL statement string and returns a compact JSON description of the + * unresolved statement (identifier/code, lineage references, select-list names, + * parameters), or a STANDARD-format error object when the statement does not + * parse. + * + * Behind [[SQLConf.PARSE_SQL_ENABLED]] while the JSON contract is still + * evolving. Designed for batch evaluation over DataFrames of SQL text. + * User-facing parse errors become JSON; unexpected internal failures propagate. + */ +// scalastyle:off line.size.limit +@ExpressionDescription( + usage = """_FUNC_(sqlStmt) - Parses `sqlStmt` with the stock Spark SQL parser and + returns a JSON string describing the statement (parse success, Table 39 statement + identifier/code, target and source table references for lineage, select-list column + names, and parameter markers). Session parser extensions are not applied. + Requires spark.sql.function.parseSql.enabled=true. On syntax / parse error returns JSON + with `parse_success` false, source location, and a nested STANDARD error object + instead of throwing.""", + arguments = """ + Arguments: + * sqlStmt - A SQL statement string to parse. + An expression that evaluates to a string. + """, + examples = """ + Examples: + > SELECT _FUNC_('SELECT a, b FROM t'); + {"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["t"]],"select_list":[{"name":["a"]},{"name":["b"]}]} + > SELECT get_json_object(_FUNC_('SELEC'), '$.error.errorClass'); + PARSE_SYNTAX_ERROR + """, + group = "misc_funcs", + since = "4.4.0") +// scalastyle:on line.size.limit +case class ParseSql(child: Expression) + extends UnaryExpression + with ImplicitCastInputTypes + with CodegenFallback { + + override def prettyName: String = "parse_sql" + + override def nullable: Boolean = true + + override def nullIntolerant: Boolean = true + + override def dataType: DataType = StringType + + override def inputTypes: Seq[AbstractDataType] = + Seq(StringTypeWithCollation(supportsTrimCollation = true)) + + override def checkInputDataTypes(): TypeCheckResult = { + if (!SQLConf.get.parseSqlEnabled) { + throw new AnalysisException( + errorClass = "FEATURE_NOT_ENABLED", + messageParameters = Map( + "featureName" -> "parse_sql", + "configKey" -> SQLConf.PARSE_SQL_ENABLED.key, + "configValue" -> "true")) + } + super.checkInputDataTypes() + } + + override def nullSafeEval(input: Any): Any = { + val sql = input.asInstanceOf[UTF8String].toString + UTF8String.fromString(ParseSqlResult.fromSql(sql)) + } + + override protected def withNewChildInternal(newChild: Expression): ParseSql = + copy(child = newChild) +} + +object ParseSql { + /** Register the builtin with a session function registry and the global builtin set. */ + def register(registry: FunctionRegistry): Unit = { + val (info, builder) = FunctionRegistryBase.build[ParseSql]("parse_sql", Some("4.4.0")) + // Keep the session registry in sync for the first session (cloned before this runs). + registry.registerFunction( + FunctionRegistry.builtinFunctionIdentifier("parse_sql"), + info, + builder) + // Also publish into FunctionRegistry.builtin / functionSet so SHOW USER/SYSTEM + // FUNCTIONS classify parse_sql as SYSTEM rather than a user/temp function. + FunctionRegistry.registerExtraBuiltin("parse_sql", info, builder) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResult.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResult.scala new file mode 100644 index 0000000000000..228da54c3c7d1 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResult.scala @@ -0,0 +1,426 @@ +/* + * 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.catalyst.parser + +import scala.collection.mutable + +import org.json4s._ +import org.json4s.jackson.JsonMethods.{compact, parse => parseJson, render} + +import org.apache.spark.{ErrorMessageFormat, SparkThrowable, SparkThrowableHelper} +import org.apache.spark.sql.catalyst.analysis._ +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin, SQLQueryContext} +import org.apache.spark.sql.exceptions.SqlScriptingException +import org.apache.spark.sql.execution.SparkSqlParser +import org.apache.spark.sql.execution.command.{CreateViewCommand, DescribeQueryCommand, ExplainCommand} +import org.apache.spark.sql.execution.datasources.CreateTempViewUsing + +/** + * Parses a SQL statement string and returns a compact JSON description of the + * unresolved plan (parse-only; no catalog resolution). + * + * Uses a stock [[SparkSqlParser]] (ThreadLocal) so statement coverage matches + * the default production parser (EXPLAIN / SET / ADD JAR / temp views / etc.). + * Session-specific [[org.apache.spark.sql.SparkSessionExtensions]] parser + * wrappers are intentionally not applied: `parse_sql` must evaluate on + * executors without a session, so only the stock parser is available under + * distributed eval. + * + * On success the JSON always includes `parse_success`, the statement + * identifier/code (ISO/IEC 9075-2:2023 Table 39), and omits unused optional + * fields (`target_table_references`, `source_table_references`, + * `function_references`, `select_list`, `parameter_markers`) when empty. On parse + * failure it returns `parse_success: false` with source location and a nested + * STANDARD-format error object, and does not throw. Only [[ParseException]] / + * [[SqlScriptingException]] are converted to JSON; unexpected / internal + * failures propagate so the function fails. + */ +object ParseSqlResult { + + private val parser: ThreadLocal[SparkSqlParser] = + ThreadLocal.withInitial(() => new SparkSqlParser()) + + /** Parse `sql` and render the JSON result string. */ + def fromSql(sql: String): String = { + try { + // Do not inherit the outer query's origin from the parse_sql expression. + // Errors and parsed nodes must refer to the SQL string passed to this function. + val origin = if (sql.nonEmpty) { + Origin(startIndex = Some(0), stopIndex = Some(sql.length - 1), sqlText = Some(sql)) + } else { + Origin(sqlText = Some(sql)) + } + CurrentOrigin.withOrigin(origin) { + val plan = parser.get().parsePlan(sql) + fromPlan(plan) + } + } catch { + // User-facing parse / scripting failures become JSON; everything else fails. + case e: ParseException => + errorJson(e) + case e: SqlScriptingException => + errorJson(e) + } + } + + /** Build success JSON from an already-parsed unresolved plan. */ + def fromPlan(plan: LogicalPlan): String = { + val classification = SqlStatementCodes.classify(plan) + val fields = mutable.ListBuffer.empty[JField] + fields += "parse_success" -> JBool(true) + fields += "statement_identifier" -> JString(classification.statementIdentifier) + fields += "statement_code" -> JInt(classification.statementCode) + // Omit unused collections / markers so consumers can treat absence as empty. + val refs = collectPlanReferences(plan) + if (refs.targetTables.nonEmpty) { + fields += "target_table_references" -> + JArray(refs.targetTables.map(partsToJArray).toList) + } + if (refs.sourceTables.nonEmpty) { + fields += "source_table_references" -> + JArray(refs.sourceTables.map(partsToJArray).toList) + } + if (refs.functions.nonEmpty) { + fields += "function_references" -> JArray(refs.functions.map(partsToJArray).toList) + } + val selectList = collectSelectList(plan) + if (selectList.nonEmpty) { + fields += "select_list" -> JArray(selectList.toList) + } + refs.parameterMarkers.foreach(markers => fields += "parameter_markers" -> markers) + compact(render(JObject(fields.toList))) + } + + private def errorJson(e: SparkThrowable with Throwable): String = { + val errorObj = parseJson( + SparkThrowableHelper.getMessage(e, ErrorMessageFormat.STANDARD)).asInstanceOf[JObject] + val origin = e match { + case p: ParseException => Some(p.start) + case s: SqlScriptingException => Some(s.origin) + case _ => None + } + val locationFields = origin.toSeq.flatMap(originFields) + val contextFields = if (errorObj.obj.exists(_._1 == "queryContext")) { + Nil + } else { + origin.toSeq.flatMap(queryContextField) + } + compact(render(JObject( + "parse_success" -> JBool(false), + "error" -> JObject(errorObj.obj ++ contextFields ++ locationFields) + ))) + } + + private def queryContextField(origin: Origin): Option[JField] = origin.context match { + case context: SQLQueryContext if context.isValid => + Some("queryContext" -> JArray(List(JObject( + "objectType" -> JString(context.objectType), + "objectName" -> JString(context.objectName), + "startIndex" -> JInt(context.startIndex + 1), + "stopIndex" -> JInt(context.stopIndex + 1), + "fragment" -> JString(context.fragment) + )))) + case _ => None + } + + private def originFields(origin: Origin): Seq[JField] = Seq( + origin.line.map(line => "line" -> JInt(line)), + origin.startPosition.map(position => "position" -> JInt(position))).flatten + + private def partsToJArray(parts: Seq[String]): JArray = + JArray(parts.map(JString).toList) + + /** + * Walk expressions in all product fields, including wrappers such as column + * definitions that [[LogicalPlan.expressions]] does not descend into. + */ + private def foreachExpressionDeep(plan: LogicalPlan)(f: Expression => Unit): Unit = { + def visit(value: Any): Unit = value match { + case e: Expression => f(e) + case _: LogicalPlan => + case values: Iterable[_] => values.foreach(visit) + case value: Product => value.productIterator.foreach(visit) + case _ => + } + plan.productIterator.foreach(visit) + } + + /** Whether a command's single `child` names the statement target table/view. */ + private def isTargetTableChild(parent: LogicalPlan): Boolean = parent match { + case _: DeleteFromTable | _: DeleteFromTableWithFilters | _: UpdateTable | + _: CreateTable | _: ReplaceTable | _: DropTable | _: DropView | + _: TruncateTable | _: TruncatePartition | _: AlterTableCommand | + _: RenameTable | _: SetViewProperties | _: RefreshTable | + _: UncacheTable | _: CommentOnTable | _: CreateIndex => + true + case _ => false + } + + /** CTE aliases visible at a given point of the walk, normalized for lookup. */ + private type CteScope = Set[String] + + /** Whether a table/view reference names the statement target or a read source. */ + private sealed trait TableRefRole + private object TableRefRole { + case object Target extends TableRefRole + case object Source extends TableRefRole + } + + private def normalizeCteName(name: String): String = + name.toLowerCase(java.util.Locale.ROOT) + + private def childScope(parent: LogicalPlan, scope: CteScope): CteScope = parent match { + case w: UnresolvedWith => scope ++ w.cteRelations.map(r => normalizeCteName(r._1)) + case _ => scope + } + + /** + * Deep plan walk that hands each node the CTE aliases in scope there and the + * table-reference role for relation nodes at that position, and covers tree + * slots that standard `collect` / `collectWithSubqueries` miss. + */ + private def foreachPlanDeep( + plan: LogicalPlan)(f: (LogicalPlan, CteScope, TableRefRole) => Unit): Unit = + visitPlan(plan, Set.empty, TableRefRole.Source)(f) + + private def visitPlan( + plan: LogicalPlan, + scope: CteScope, + role: TableRefRole)(f: (LogicalPlan, CteScope, TableRefRole) => Unit): Unit = { + f(plan, scope, role) + plan match { + case s: SingleStatement => + visitPlan(s.parsedPlan, scope, role)(f) + case _ => + visitNonChildSlots(plan, scope, role)(f) + visitChildPlans(plan, scope, role)(f) + } + } + + private def visitChildPlans( + parent: LogicalPlan, + scope: CteScope, + role: TableRefRole)(f: (LogicalPlan, CteScope, TableRefRole) => Unit): Unit = { + val nextScope = childScope(parent, scope) + parent match { + case m: MergeIntoTable => + visitPlan(m.targetTable, scope, TableRefRole.Target)(f) + visitPlan(m.sourceTable, scope, TableRefRole.Source)(f) + case i: InsertIntoStatement => + visitPlan(i.query, nextScope, TableRefRole.Source)(f) + case c: CreateTableAsSelect => + visitPlan(c.name, scope, TableRefRole.Target)(f) + visitPlan(c.query, nextScope, TableRefRole.Source)(f) + case r: ReplaceTableAsSelect => + visitPlan(r.name, scope, TableRefRole.Target)(f) + visitPlan(r.query, nextScope, TableRefRole.Source)(f) + case cv: CreateView => + visitPlan(cv.child, scope, TableRefRole.Target)(f) + visitPlan(cv.query, nextScope, TableRefRole.Source)(f) + case cts: CacheTableAsSelect => + visitPlan(cts.plan, nextScope, TableRefRole.Source)(f) + case _: CacheTable => + // Name is taken from multipartIdentifier or the non-child `table` slot. + case SubqueryAlias(_, child) => + visitPlan(child, nextScope, role)(f) + case _ => + parent.subqueries.foreach(sq => visitPlan(sq, nextScope, TableRefRole.Source)(f)) + parent.children.foreach { child => + val childRole = + if (isTargetTableChild(parent)) TableRefRole.Target else role + visitPlan(child, nextScope, childRole)(f) + } + } + } + + /** Visit plan slots that are not exposed via `children` / subqueries. */ + private def visitNonChildSlots( + plan: LogicalPlan, + scope: CteScope, + role: TableRefRole)(f: (LogicalPlan, CteScope, TableRefRole) => Unit): Unit = { + plan match { + case w: UnresolvedWith => + // A CTE definition sees the aliases defined before it, plus its own + // name when the clause is RECURSIVE. Later aliases are not in scope, + // so a definition naming one refers to the real table. + var definitionScope = scope + w.cteRelations.foreach { case (name, ctePlan, _) => + val normalized = normalizeCteName(name) + val bodyScope = + if (w.allowRecursion) definitionScope + normalized else definitionScope + visitPlan(ctePlan, bodyScope, TableRefRole.Source)(f) + definitionScope += normalized + } + case i: InsertIntoStatement => + visitPlan(i.table, scope, TableRefRole.Target)(f) + case c: CacheTable if c.multipartIdentifier.isEmpty => + visitPlan(c.table, scope, TableRefRole.Target)(f) + case c: CompoundBody => + c.handlers.foreach(h => visitPlan(h, scope, TableRefRole.Source)(f)) + case s: SimpleCaseStatement => + s.elseBody.foreach(b => visitPlan(b, scope, TableRefRole.Source)(f)) + case ExplainCommand(logicalPlan, _) => + visitPlan(logicalPlan, scope, TableRefRole.Source)(f) + case DescribeQueryCommand(_, queryPlan) => + visitPlan(queryPlan, scope, TableRefRole.Source)(f) + case _ => + } + } + + private def tableIdentifierParts(id: org.apache.spark.sql.catalyst.TableIdentifier): Seq[String] = + id.catalog.toSeq ++ id.database.toSeq :+ id.table + + private final case class PlanReferences( + targetTables: Seq[Seq[String]], + sourceTables: Seq[Seq[String]], + functions: Seq[Seq[String]], + parameterMarkers: Option[JObject]) + + /** + * Collect multipart table/view identifiers, function names, and parameter + * markers for lineage in a single deep walk. Target references name the + * table/view a DML or DDL statement writes to or alters; source references + * name tables read in FROM clauses and query bodies. A single-part name is + * dropped only when a CTE alias in scope at that node shadows it. Function / + * variable identifiers are not collected as tables. Deduplicates while + * preserving first-seen order within each category. + */ + private def collectPlanReferences(plan: LogicalPlan): PlanReferences = { + val targetTables = mutable.LinkedHashSet.empty[Seq[String]] + val sourceTables = mutable.LinkedHashSet.empty[Seq[String]] + val functions = mutable.LinkedHashSet.empty[Seq[String]] + val namedParams = mutable.LinkedHashSet.empty[String] + var unnamedCount = 0 + + def isCteName(parts: Seq[String], scope: CteScope): Boolean = parts match { + case Seq(name) => scope.contains(normalizeCteName(name)) + case _ => false + } + + def addTable(parts: Seq[String], scope: CteScope, role: TableRefRole): Unit = { + if (parts.nonEmpty && !isCteName(parts, scope)) { + role match { + case TableRefRole.Target => targetTables += parts + case TableRefRole.Source => sourceTables += parts + } + } + } + + def addFunction(parts: Seq[String]): Unit = { + if (parts.nonEmpty) functions += parts + } + + def visitExpr(e: Expression): Unit = e.foreach { + case f: UnresolvedFunction => addFunction(f.nameParts) + case n: NamedParameter => namedParams += n.name + case _: PosParameter => unnamedCount += 1 + case _ => + } + + def addTarget(parts: Seq[String]): Unit = { + if (parts.nonEmpty) targetTables += parts + } + + foreachPlanDeep(plan) { (p, scope, role) => + foreachExpressionDeep(p)(visitExpr) + def add(parts: Seq[String]): Unit = addTable(parts, scope, role) + p match { + case u: UnresolvedRelation => add(u.multipartIdentifier) + case u: UnresolvedTable => add(u.multipartIdentifier) + case u: UnresolvedView => add(u.multipartIdentifier) + case u: UnresolvedTableOrView => add(u.multipartIdentifier) + case u: UnresolvedIdentifier if role == TableRefRole.Target => add(u.nameParts) + case c: CreateViewCommand => addTarget(tableIdentifierParts(c.name)) + case c: CreateTempViewUsing => addTarget(tableIdentifierParts(c.tableIdent)) + case c: CacheTable if c.multipartIdentifier.nonEmpty => + addTarget(c.multipartIdentifier) + case u: UnresolvedTableValuedFunction => addFunction(u.name) + case _ => + } + } + + val markers = + if (namedParams.isEmpty && unnamedCount == 0) { + None + } else { + val markerFields = mutable.ListBuffer.empty[JField] + if (namedParams.nonEmpty) { + markerFields += "named" -> JArray(namedParams.toList.map(JString)) + } + if (unnamedCount > 0) markerFields += "unnamed_count" -> JInt(unnamedCount) + Some(JObject(markerFields.toList)) + } + PlanReferences(targetTables.toSeq, sourceTables.toSeq, functions.toSeq, markers) + } + + /** + * Collect the primary select list as `{name}` objects (multipart name + * parts only). Empty for non-query statements without a projected query body. + */ + private def collectSelectList(plan: LogicalPlan): Seq[JObject] = { + val query = primaryQueryPlan(plan) + val named: Seq[NamedExpression] = query match { + case p: Project => p.projectList + case a: Aggregate => a.aggregateExpressions + case _ => Nil + } + named.map(selectListItem) + } + + private def primaryQueryPlan(plan: LogicalPlan): LogicalPlan = plan match { + case UnresolvedWith(child, _, _) => primaryQueryPlan(child) + case InsertIntoStatement(_, _, _, query, _, _, _, _, _) => + primaryQueryPlan(query) + case c: CreateTableAsSelect => primaryQueryPlan(c.query) + case r: ReplaceTableAsSelect => primaryQueryPlan(r.query) + case c: CreateView => primaryQueryPlan(c.query) + case c: CreateViewCommand => primaryQueryPlan(c.plan) + case c: CacheTableAsSelect => primaryQueryPlan(c.plan) + case ExplainCommand(logicalPlan, _) => primaryQueryPlan(logicalPlan) + case DescribeQueryCommand(_, queryPlan) => primaryQueryPlan(queryPlan) + case SubqueryAlias(_, child) => primaryQueryPlan(child) + case Sort(_, _, child, _) => primaryQueryPlan(child) + case Filter(_, child) => primaryQueryPlan(child) + case UnresolvedHaving(_, child) => primaryQueryPlan(child) + case UnresolvedQualify(_, child) => primaryQueryPlan(child) + case Distinct(child) => primaryQueryPlan(child) + case GlobalLimit(_, child) => primaryQueryPlan(child) + case LocalLimit(_, child) => primaryQueryPlan(child) + case Offset(_, child) => primaryQueryPlan(child) + case Repartition(_, _, child) => primaryQueryPlan(child) + case RepartitionByExpression(_, child, _, _) => primaryQueryPlan(child) + case Sample(_, _, _, _, child, _) => primaryQueryPlan(child) + case other => other + } + + private def selectListItem(ne: NamedExpression): JObject = ne match { + case Alias(_, name) => + JObject("name" -> partsToJArray(Seq(name))) + case _: UnresolvedAlias => + JObject("name" -> partsToJArray(Nil)) + case s: UnresolvedStar => + val name = s.target.map(_ :+ "*").getOrElse(Seq("*")) + JObject("name" -> partsToJArray(name)) + case a: UnresolvedAttribute => + JObject("name" -> partsToJArray(a.nameParts)) + case other => + JObject("name" -> partsToJArray(Seq(other.name))) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlStatementCodes.scala b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlStatementCodes.scala new file mode 100644 index 0000000000000..cb2fae1d9af9c --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlStatementCodes.scala @@ -0,0 +1,214 @@ +/* + * 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.catalyst.parser + +import org.apache.spark.sql.catalyst.analysis.{ + RelationTimeTravel, + ResolvedInlineTable, + UnresolvedExecuteImmediate, + UnresolvedHaving, + UnresolvedInlineTable, + UnresolvedRelation, + UnresolvedTableValuedFunction +} +import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.execution.command._ +import org.apache.spark.sql.execution.datasources.{CreateTempViewUsing, RefreshResource} +import org.apache.spark.sql.metricview.logical.CreateMetricView + +/** + * Classification of a parsed SQL statement using ISO/IEC 9075-2:2023 Table 39, + * "SQL-statement codes" (clause 23.1 <get diagnostics statement>). + * + * @param statementIdentifier Table 39 Identifier column (or Spark product name) + * @param statementCode Table 39 Code column; Spark-only statements use negative + * implementation-defined codes (Table 39 IE005 / IV190) + */ +case class SqlStatementClassification( + statementIdentifier: String, + statementCode: Int) + +/** + * Maps unresolved [[LogicalPlan]]s to Table 39 statement codes. + * + * Spark-only statements use the standard's implementation-defined escape hatch: + * a product-specific identifier and a distinct negative code. Codes are + * append-only and must never be renumbered. + * + * Unknown plans map to [[Unrecognized]] (empty identifier, code 0). Query + * shapes are allowlisted; unknown non-commands are not assumed to be SELECT. + */ +object SqlStatementCodes { + + // Standard Table 39 entries used by Spark SQL (ISO/IEC 9075-2:2023). + val Select: SqlStatementClassification = SqlStatementClassification("SELECT", 21) + val Insert: SqlStatementClassification = SqlStatementClassification("INSERT", 50) + val DeleteWhere: SqlStatementClassification = SqlStatementClassification("DELETE WHERE", 19) + val UpdateWhere: SqlStatementClassification = SqlStatementClassification("UPDATE WHERE", 82) + val Merge: SqlStatementClassification = SqlStatementClassification("MERGE", 128) + val CreateTable: SqlStatementClassification = SqlStatementClassification("CREATE TABLE", 77) + val CreateView: SqlStatementClassification = SqlStatementClassification("CREATE VIEW", 84) + val DropTable: SqlStatementClassification = SqlStatementClassification("DROP TABLE", 32) + val DropView: SqlStatementClassification = SqlStatementClassification("DROP VIEW", 36) + val AlterTable: SqlStatementClassification = SqlStatementClassification("ALTER TABLE", 4) + val CreateSchema: SqlStatementClassification = SqlStatementClassification("CREATE SCHEMA", 64) + val DropSchema: SqlStatementClassification = SqlStatementClassification("DROP SCHEMA", 31) + val SetSchema: SqlStatementClassification = SqlStatementClassification("SET SCHEMA", 74) + val TruncateTable: SqlStatementClassification = + SqlStatementClassification("TRUNCATE TABLE", 139) + val CreateRoutine: SqlStatementClassification = SqlStatementClassification("CREATE ROUTINE", 14) + val DropRoutine: SqlStatementClassification = SqlStatementClassification("DROP ROUTINE", 30) + val ExecuteImmediate: SqlStatementClassification = + SqlStatementClassification("EXECUTE IMMEDIATE", 43) + val Call: SqlStatementClassification = SqlStatementClassification("CALL", 7) + + // Table 39 "Unrecognized statements": empty identifier, code 0. + val Unrecognized: SqlStatementClassification = SqlStatementClassification("", 0) + + // Spark product-specific identifiers with append-only negative codes + // (Table 39 implementation-defined / IE005 row: negative Code values). + val CacheTable: SqlStatementClassification = spark("CACHE TABLE", -1) + val CacheTableAsSelect: SqlStatementClassification = spark("CACHE TABLE AS SELECT", -2) + val UncacheTable: SqlStatementClassification = spark("UNCACHE TABLE", -3) + val RefreshTable: SqlStatementClassification = spark("REFRESH TABLE", -4) + val ShowTables: SqlStatementClassification = spark("SHOW TABLES", -5) + val DescribeTable: SqlStatementClassification = spark("DESCRIBE TABLE", -6) + val AnalyzeTable: SqlStatementClassification = spark("ANALYZE TABLE", -7) + val DeclareVariable: SqlStatementClassification = spark("DECLARE VARIABLE", -8) + val SetVariable: SqlStatementClassification = spark("SET VARIABLE", -9) + val DropVariable: SqlStatementClassification = spark("DROP VARIABLE", -10) + val ShowTableProperties: SqlStatementClassification = spark("SHOW TBLPROPERTIES", -11) + val DescribeNamespace: SqlStatementClassification = spark("DESCRIBE NAMESPACE", -12) + val ShowFunctions: SqlStatementClassification = spark("SHOW FUNCTIONS", -13) + val DescribeFunction: SqlStatementClassification = spark("DESCRIBE FUNCTION", -14) + val ShowCreateTable: SqlStatementClassification = spark("SHOW CREATE TABLE", -15) + val ShowColumns: SqlStatementClassification = spark("SHOW COLUMNS", -16) + val ShowPartitions: SqlStatementClassification = spark("SHOW PARTITIONS", -17) + val ShowViews: SqlStatementClassification = spark("SHOW VIEWS", -18) + val RefreshFunction: SqlStatementClassification = spark("REFRESH FUNCTION", -19) + val CommentOnNamespace: SqlStatementClassification = spark("COMMENT ON NAMESPACE", -20) + val CommentOnTable: SqlStatementClassification = spark("COMMENT ON TABLE", -21) + // SQL/PSM-style scripting (9075-4); not in Foundation Table 39. + val BeginEnd: SqlStatementClassification = spark("BEGIN END", -22) + // SparkSqlParser-only session / resource commands (append-only). + val Explain: SqlStatementClassification = spark("EXPLAIN", -23) + val Set: SqlStatementClassification = spark("SET", -24) + val Reset: SqlStatementClassification = spark("RESET", -25) + val AddJar: SqlStatementClassification = spark("ADD JAR", -26) + val AddFile: SqlStatementClassification = spark("ADD FILE", -27) + val AddArchive: SqlStatementClassification = spark("ADD ARCHIVE", -28) + val ListJar: SqlStatementClassification = spark("LIST JAR", -29) + val ListFile: SqlStatementClassification = spark("LIST FILE", -30) + val ClearCache: SqlStatementClassification = spark("CLEAR CACHE", -31) + val RefreshResourceCmd: SqlStatementClassification = spark("REFRESH RESOURCE", -32) + val DescribeQuery: SqlStatementClassification = spark("DESCRIBE QUERY", -33) + val ShowCatalogs: SqlStatementClassification = spark("SHOW CATALOGS", -34) + val ShowCurrentNamespace: SqlStatementClassification = + spark("SHOW CURRENT NAMESPACE", -35) + val SetCatalog: SqlStatementClassification = spark("SET CATALOG", -36) + val CreateMetricViewStmt: SqlStatementClassification = spark("CREATE METRIC VIEW", -37) + + private def spark(identifier: String, code: Int): SqlStatementClassification = { + assert(code < 0, s"Spark statement codes must be negative, got $code") + SqlStatementClassification(statementIdentifier = identifier, statementCode = code) + } + + /** Classify an unresolved logical plan. */ + def classify(plan: LogicalPlan): SqlStatementClassification = plan match { + case UnresolvedWith(child, _, _) => classify(child) + case _: CompoundBody => BeginEnd + case _: InsertIntoStatement => Insert + case _: DeleteFromTable | _: DeleteFromTableWithFilters => DeleteWhere + case _: UpdateTable => UpdateWhere + case _: MergeIntoTable => Merge + case _: CreateTableAsSelect | _: ReplaceTableAsSelect => CreateTable + case _: CreateTable | _: CreateTableLike | _: ReplaceTable => CreateTable + case _: CreateView | _: CreateViewCommand | _: CreateTempViewUsing => CreateView + case _: DropTable => DropTable + case _: DropView => DropView + case _: CreateNamespace => CreateSchema + case _: DropNamespace => DropSchema + case _: SetCatalogAndNamespace | _: SetNamespaceCommand => SetSchema + case _: SetCatalogCommand => SetCatalog + case _: TruncateTable => TruncateTable + case _: CreateFunction | _: CreateFunctionCommand | + _: CreateUserDefinedFunction | _: CreateUserDefinedFunctionCommand => + CreateRoutine + case _: DropFunction | _: DropFunctionCommand => DropRoutine + case _: UnresolvedExecuteImmediate => ExecuteImmediate + case _: Call => Call + case _: CommentOnTable => CommentOnTable + case _: AlterTableCommand | _: RenameTable => AlterTable + case _: CacheTable => CacheTable + case _: CacheTableAsSelect => CacheTableAsSelect + case _: UncacheTable => UncacheTable + case _: RefreshTable => RefreshTable + case _: ShowTables | _: ShowTablesExtended => ShowTables + case _: DescribeRelation | _: DescribeTablePartition | _: DescribeColumn => + DescribeTable + case _: DescribeQueryCommand => DescribeQuery + case _: AnalyzeTable | _: AnalyzeTables | _: AnalyzeColumn => AnalyzeTable + case _: CreateVariable => DeclareVariable + case _: SetVariable => SetVariable + case _: DropVariable => DropVariable + case _: ShowTableProperties => ShowTableProperties + case _: DescribeNamespace => DescribeNamespace + case _: ShowFunctions => ShowFunctions + case _: DescribeFunction => DescribeFunction + case _: ShowCreateTable => ShowCreateTable + case _: ShowColumns => ShowColumns + case _: ShowPartitions | _: ShowTablePartition => ShowPartitions + case _: ShowViews => ShowViews + case _: RefreshFunction => RefreshFunction + case _: CommentOnNamespace => CommentOnNamespace + case _: ExplainCommand => Explain + case _: SetCommand => Set + case _: ResetCommand => Reset + case _: AddJarsCommand => AddJar + case _: AddFilesCommand => AddFile + case _: AddArchivesCommand => AddArchive + case _: ListJarsCommand => ListJar + case _: ListFilesCommand => ListFile + case ClearCacheCommand => ClearCache + case _: RefreshResource => RefreshResourceCmd + case _: ShowCatalogsCommand => ShowCatalogs + case _: ShowCurrentNamespaceCommand => ShowCurrentNamespace + case _: CreateMetricView | _: CreateMetricViewCommand => CreateMetricViewStmt + case _: Command => Unrecognized + case p if isQueryPlan(p) => Select + case _ => Unrecognized + } + + /** + * Allowlisted query-shaped plans. Unknown non-command plans are not assumed + * to be SELECT. + */ + private def isQueryPlan(plan: LogicalPlan): Boolean = plan match { + case _: Project | _: Aggregate | _: Distinct | _: Filter | _: Sort | + _: GlobalLimit | _: LocalLimit | _: Join | _: Union | _: Except | + _: Intersect | _: SubqueryAlias | _: Repartition | + _: RepartitionByExpression | _: Sample | _: Range | + _: OneRowRelation | _: LocalRelation | _: Deduplicate | + _: Expand | _: Generate | _: Window | _: Tail | _: Offset | + _: LateralJoin | _: UnresolvedHaving | _: CollectMetrics | + _: WithCTE | _: UnresolvedRelation | _: UnresolvedInlineTable | + _: ResolvedInlineTable | _: RelationTimeTravel | + _: UnresolvedTableValuedFunction => true + case _ => false + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala index a9f16ffa87be1..3d0036f865225 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriter.scala @@ -41,6 +41,7 @@ import org.apache.spark.sql.execution.datasources.{DataSource, DataSourceUtils} import org.apache.spark.sql.execution.datasources.v2._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode +import org.apache.spark.sql.sources.CreatableRelationProvider import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -176,16 +177,22 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram val catalog = CatalogV2Util.getTableProviderCatalog( supportsExtract, catalogManager, dsOptions) - (catalog.loadTable(ident), Some(catalog), Some(ident)) + val table = CatalogV2Util.loadTableForV2Write( + catalog, ident, getWritePrivileges, dsOptions) + (table, Some(catalog), Some(ident)) case _: TableProvider => val t = getTable if (t.supports(BATCH_WRITE)) { (t, None, None) - } else { - // Streaming also uses the data source V2 API. So it may be that the data source - // implements v2, but has no v2 implementation for batch writes. In that case, we - // fall back to saving as though it's a V1 source. + } else if (t.supports(V1_BATCH_WRITE) || + provider.isInstanceOf[CreatableRelationProvider]) { + // Fall back to V1 write path. V1_BATCH_WRITE is the explicit capability + // for DSv2 tables that opt into the V1 write path. We also check + // CreatableRelationProvider for backward compatibility with sources that + // predate the V1_BATCH_WRITE capability (SPARK-28334). return saveToV1SourceCommand(path) + } else { + throw QueryCompilationErrors.unsupportedBatchWriteError(t) } } @@ -229,15 +236,17 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram finalOptions, ignoreIfExists = createMode == SaveMode.Ignore) case _: TableProvider => - if (getTable.supports(BATCH_WRITE)) { + val t = getTable + if (t.supports(BATCH_WRITE)) { throw QueryCompilationErrors.writeWithSaveModeUnsupportedBySourceError( source, createMode.name()) - } else { - // Streaming also uses the data source V2 API. So it may be that the data source - // implements v2, but has no v2 implementation for batch writes. In that case, we - // fallback to saving as though it's a V1 source. + } else if (t.supports(V1_BATCH_WRITE) || + provider.isInstanceOf[CreatableRelationProvider]) { + // Fall back to V1 write path (see comment in Append/Overwrite branch above). assertSchemaEvolutionNotEnabledForV1Write() saveToV1SourceCommand(path) + } else { + throw QueryCompilationErrors.unsupportedBatchWriteError(t) } } } @@ -308,7 +317,9 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram * +---+---+ * }}} * - * Because it inserts data to an existing table, format or options will be ignored. + * Because it inserts data to an existing table, the format is ignored. For data source V2 + * tables, catalog-declared table-state options are forwarded to the table load and all options + * are forwarded to the write; for V1 tables the options are ignored. * * @since 1.4.0 */ @@ -347,13 +358,14 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram } private def insertIntoCommand(catalog: CatalogPlugin, ident: Identifier): LogicalPlan = { - import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ - - val table = catalog.asTableCatalog.loadTable(ident, getWritePrivileges.toSet.asJava) match { + val tableOptions = new CaseInsensitiveStringMap(extraOptions.toMap.asJava) + val table = CatalogV2Util.loadTableForWrite( + catalog, ident, getWritePrivileges, tableOptions) match { case _: V1Table => return insertIntoCommand(TableIdentifier(ident.name(), ident.namespace().headOption)) case t => - DataSourceV2Relation.create(t, Some(catalog), Some(ident)) + CatalogV2Util.rejectTimeTravelOptionsForWrite(catalog, ident, tableOptions) + DataSourceV2Relation.create(t, Some(catalog), Some(ident), tableOptions) } curmode match { @@ -479,18 +491,27 @@ final class DataFrameWriter[T] private[sql](ds: Dataset[T]) extends sql.DataFram v2ProviderOpt: Option[TableProvider], ident: Identifier, nameParts: Seq[String]): LogicalPlan = { - val tableOpt = try Option(catalog.loadTable(ident, getWritePrivileges.toSet.asJava)) catch { + val tableOptions = new CaseInsensitiveStringMap(extraOptions.toMap.asJava) + val tableOpt = try { + Option(CatalogV2Util.loadTableForWrite(catalog, ident, getWritePrivileges, tableOptions)) + } catch { case _: NoSuchTableException => None } - (curmode, tableOpt) match { - case (_, Some(_: V1Table)) => + tableOpt match { + case Some(_: V1Table) => assertSchemaEvolutionNotEnabledForV1Write() - saveAsV1TableCommand(TableIdentifier(ident.name(), ident.namespace().headOption)) + return saveAsV1TableCommand(TableIdentifier(ident.name(), ident.namespace().headOption)) + case _ => () + } + + CatalogV2Util.rejectTimeTravelOptionsForWrite(catalog, ident, tableOptions) + (curmode, tableOpt) match { case (SaveMode.Append, Some(table)) => checkPartitioningMatchesV2Table(table) - val v2Relation = DataSourceV2Relation.create(table, Some(catalog), Some(ident)) + val v2Relation = + DataSourceV2Relation.create(table, Some(catalog), Some(ident), tableOptions) AppendData.byName(v2Relation, df.logicalPlan, extraOptions.toMap, _withSchemaEvolution) case (SaveMode.Overwrite, _) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriterV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriterV2.scala index 0da60e98cdbe7..a8cc7c79f7fc0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriterV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/DataFrameWriterV2.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.classic import java.util import scala.collection.mutable -import scala.jdk.CollectionConverters.MapHasAsScala +import scala.jdk.CollectionConverters._ import org.apache.spark.annotation.Experimental import org.apache.spark.sql @@ -29,12 +29,14 @@ import org.apache.spark.sql.catalyst.analysis.{NoSuchTableException, UnresolvedF import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap +import org.apache.spark.sql.connector.catalog.TableWritePrivilege import org.apache.spark.sql.connector.catalog.TableWritePrivilege._ import org.apache.spark.sql.connector.expressions._ import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.util.CaseInsensitiveStringMap /** * Interface used to write a [[org.apache.spark.sql.classic.Dataset]] to external storage using @@ -197,7 +199,7 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T]) private[sql] def appendCommand(): LogicalPlan = { AppendData.byName( - UnresolvedRelation(tableName).requireWritePrivileges(Set(INSERT)), + createUnresolvedWriteTarget(Set(INSERT)), logicalPlan, options.toMap, withSchemaEvolution = _withSchemaEvolution) } @@ -209,7 +211,7 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T]) private[sql] def overwriteCommand(condition: Column): LogicalPlan = { OverwriteByExpression.byName( - UnresolvedRelation(tableName).requireWritePrivileges(Set(INSERT, DELETE)), + createUnresolvedWriteTarget(Set(INSERT, DELETE)), logicalPlan, expression(condition), options.toMap, _withSchemaEvolution) } @@ -221,10 +223,16 @@ final class DataFrameWriterV2[T] private[sql](table: String, ds: Dataset[T]) private[sql] def overwritePartitionsCommand(): LogicalPlan = { OverwritePartitionsDynamic.byName( - UnresolvedRelation(tableName).requireWritePrivileges(Set(INSERT, DELETE)), + createUnresolvedWriteTarget(Set(INSERT, DELETE)), logicalPlan, options.toMap, _withSchemaEvolution) } + private def createUnresolvedWriteTarget( + privileges: Set[TableWritePrivilege]): UnresolvedRelation = { + val tableOptions = new CaseInsensitiveStringMap(options.toMap.asJava) + UnresolvedRelation(tableName, tableOptions).requireWritePrivileges(privileges) + } + /** * Wrap an action to track the QueryExecution and time cost, then report to the user-registered * callback functions. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala index c5f01de37000c..ffbfed21bb29d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala @@ -763,7 +763,10 @@ class SparkSession private( } /** @inheritdoc */ - override def removeTag(tag: String): Unit = managedJobTags.get().remove(tag) + override def removeTag(tag: String): Unit = { + SparkContext.throwIfInvalidTag(tag) + managedJobTags.get().remove(tag) + } /** @inheritdoc */ override def getTags(): Set[String] = managedJobTags.get().keySet.toSet diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/StreamingQueryManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/StreamingQueryManager.scala index c934be69bc699..4755cfa9960a3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/StreamingQueryManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/StreamingQueryManager.scala @@ -35,7 +35,7 @@ import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.streaming._ import org.apache.spark.sql.execution.streaming.continuous.ContinuousExecution import org.apache.spark.sql.execution.streaming.runtime.{AsyncProgressTrackingMicroBatchExecution, MicroBatchExecution, StreamingQueryListenerBus, StreamingQueryWrapper} -import org.apache.spark.sql.execution.streaming.state.StateStoreCoordinatorRef +import org.apache.spark.sql.execution.streaming.state.{RocksDBStateStoreProvider, StateStoreCoordinatorRef} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.StaticSQLConf.STREAMING_QUERY_LISTENERS import org.apache.spark.sql.streaming @@ -186,6 +186,54 @@ class StreamingQueryManager private[sql] ( SQLConf.STREAMING_ASYNC_PROGRESS_TRACKING_REAL_TIME_MODE_ENABLED_BY_DEFAULT)) } + /** + * Rejects, before the query is built, session configurations that are incompatible with + * Real-Time Mode. Real-Time Mode defaults these when the user has not set them + * (see StreamExecution.setSparkSessionConfigsForRealTimeMode), but an explicit incompatible + * value is a mistake worth surfacing up front rather than silently overriding or failing later + * mid-query. Mirrors the Databricks runtime's throwIfConfsAreRTMIncompatible, limited to the + * checks that apply in OSS. + * + * - state store checkpoint format below v2 (unless the v1 escape hatch is set): Real-Time Mode + * requires v2 (see the fail-fast in MicroBatchExecution.initializeExecution); + * - a non-RocksDB state store provider: v2 requires RocksDB; + * - sortBeforeRepartition left true: the blocking sort never completes on an unbounded stream. + */ + private def throwIfConfsAreRealTimeModeIncompatible(sparkSession: SparkSession): Unit = { + val conf = sparkSession.sessionState.conf + // Preserve declaration order for a stable message. + val invalidReasons = new mutable.LinkedHashMap[String, String] + + val allowCheckpointV1 = + conf.getConf(SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1) + if (!allowCheckpointV1 && + conf.contains(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key) && + conf.getConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION) < 2) { + invalidReasons += (SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2 or above") + } + + if (conf.contains(SQLConf.STATE_STORE_PROVIDER_CLASS.key) && + conf.getConf(SQLConf.STATE_STORE_PROVIDER_CLASS) != + classOf[RocksDBStateStoreProvider].getName) { + invalidReasons += + (SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName) + } + + if (conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key) && + conf.getConf(SQLConf.SORT_BEFORE_REPARTITION)) { + invalidReasons += (SQLConf.SORT_BEFORE_REPARTITION.key -> "false") + } + + if (invalidReasons.nonEmpty) { + throw new SparkIllegalArgumentException( + errorClass = "STREAMING_REAL_TIME_MODE.SQL_CONFIGURATION_NOT_SUPPORTED", + messageParameters = Map( + "invalidReasons" -> invalidReasons.zipWithIndex.map { + case ((confName, req), index) => s"${index + 1}. $confName must be $req" + }.mkString("; "))) + } + } + // scalastyle:off argcount private def createQuery( userSpecifiedName: Option[String], @@ -217,6 +265,7 @@ class StreamingQueryManager private[sql] ( ) ) } + throwIfConfsAreRealTimeModeIncompatible(sparkSession) } val dataStreamWritePlan = WriteToStreamStatement( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/conversions.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/conversions.scala index 3cfdace73c458..fd5771778b277 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/classic/conversions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/conversions.scala @@ -59,7 +59,21 @@ trait ClassicConversions { } @DeveloperApi -object ClassicConversions extends ClassicConversions +object ClassicConversions extends ClassicConversions { + /** + * Convert an [[Expression]] into a [[Column]]. This is the counterpart of + * [[ColumnConversions.expression]], for callers that would rather name the conversion than rely + * on the implicit [[ClassicConversions.ColumnConstructorExt]]. + * + * This is intentionally defined on the object rather than on the [[ClassicConversions]] trait: + * on the trait it would shadow `functions.column(colName: String)` for everyone who mixes the + * trait in, which is the trait's documented use case. + * + * @since 4.4.0 + */ + @DeveloperApi + def column(e: Expression): Column = ExpressionUtils.column(e) +} /** * Conversions from a [[Column]] to an [[Expression]]. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala index 2d8fec806f7f3..a313e2c671bec 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala @@ -85,7 +85,7 @@ trait BaseScriptTransformationExec extends UnaryExecNode { SparkFiles.getRootDirectory() builder.environment().put("PATH", path) // if OMP_NUM_THREADS is not explicitly set, override it with the value of "spark.task.cpus" - // which may be fractional, so round up to an integer (at least 1) of threads. + // which may be fractional, so round up to an integer number of threads (at least 1). if (System.getenv("OMP_NUM_THREADS") == null) { val taskCpus = conf.getConfString("spark.task.cpus", "1.0").toDouble builder.environment().put("OMP_NUM_THREADS", math.max(1, taskCpus.ceil.toInt).toString) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/BroadcastValueProjector.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/BroadcastValueProjector.scala new file mode 100644 index 0000000000000..bbff16a10f193 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/BroadcastValueProjector.scala @@ -0,0 +1,158 @@ +/* + * 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.execution + +import java.io.InterruptedIOException +import java.nio.channels.ClosedByInterruptException +import java.util.IdentityHashMap +import java.util.concurrent.CancellationException + +import scala.collection.mutable +import scala.util.control.{ControlThrowable, NonFatal} + +import org.apache.spark.TaskKilledException +import org.apache.spark.broadcast.Broadcast +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{BindReferences, Expression, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.plans.logical.Statistics +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, BroadcastQueryStageExec, ResultQueryStageExec} +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} +import org.apache.spark.sql.execution.joins.{EmptyHashedRelation, HashedRelation, HashedRelationWithAllNullKeys, LongHashedRelation} + +private[sql] case class BroadcastValueProjectionLimits( + maxInputRows: Long, + maxOutputBytes: Long, + maxSourceBytes: Long) + +private[sql] sealed trait BroadcastValueResult[+T] + +private[sql] object BroadcastValueResult { + case class Available[T](value: T) extends BroadcastValueResult[T] + case object Unavailable extends BroadcastValueResult[Nothing] +} + +/** Projects a complete, bounded value domain from rows already stored in a hash broadcast. */ +private[sql] object BroadcastValueProjector { + import BroadcastValueResult._ + + def collectExactValueDomain( + broadcast: Broadcast[HashedRelation], + child: SparkPlan, + valueExpression: Expression, + limits: BroadcastValueProjectionLimits, + onInputRow: () => Unit, + onError: () => Unit): BroadcastValueResult[Array[InternalRow]] = { + try { + if (sourceBroadcastCannotBeSafelyRehydrated(child, limits)) { + return Unavailable + } + + val broadcastRelation = broadcast.value + if (broadcastRelation == HashedRelationWithAllNullKeys) { + return Unavailable + } + if (broadcastRelation == EmptyHashedRelation) { + return Available(Array.empty[InternalRow]) + } + + val relation = broadcastRelation.asReadOnlyCopy() + val projection = UnsafeProjection.create( + BindReferences.bindReference(valueExpression, child.output)) + val projectedRows = mutable.LinkedHashSet.empty[UnsafeRow] + val valueRows = relation match { + case longRelation: LongHashedRelation => + longRelation.keys().flatMap { key => + val rows = longRelation.get(key) + if (rows == null) Iterator.empty else rows + } + case otherRelation => + otherRelation.valuesWithKeyIndex().map(_.getValue) + } + var visitedRows = 0L + var projectedBytes = 0L + + while (valueRows.hasNext) { + if (visitedRows >= limits.maxInputRows) { + return Unavailable + } + val valueRow = valueRows.next() + visitedRows += 1 + onInputRow() + val projected = projection(valueRow) + if (!projected.isNullAt(0) && !projectedRows.contains(projected)) { + val rowSize = projected.getSizeInBytes.toLong + if (rowSize > limits.maxOutputBytes - projectedBytes) { + return Unavailable + } + projectedBytes += rowSize + projectedRows += projected.copy() + } + } + + Available(projectedRows.toArray[InternalRow]) + } catch { + case NonFatal(error) if !mustPropagateFailure(error) => + onError() + Unavailable + } + } + + private def sourceBroadcastRuntimeStatistics(plan: SparkPlan): Option[Statistics] = plan match { + case exchange: BroadcastExchangeLike => Some(exchange.runtimeStatistics) + case ReusedExchangeExec(_, exchange: BroadcastExchangeLike) => + Some(exchange.runtimeStatistics) + case stage: BroadcastQueryStageExec => Some(stage.getRuntimeStatistics) + case stage: ResultQueryStageExec => + sourceBroadcastRuntimeStatistics(stage.plan) + case adaptive: AdaptiveSparkPlanExec => + sourceBroadcastRuntimeStatistics(adaptive.executedPlan) + case _ => None + } + + private def sourceBroadcastCannotBeSafelyRehydrated( + child: SparkPlan, + limits: BroadcastValueProjectionLimits): Boolean = { + val maxRows = BigInt(limits.maxInputRows) + val maxSourceBytes = BigInt(limits.maxSourceBytes) + sourceBroadcastRuntimeStatistics(child) match { + case Some(statistics) => + statistics.rowCount.forall(_ > maxRows) || statistics.sizeInBytes > maxSourceBytes + case None => true + } + } + + private[sql] def mustPropagateFailure(error: Throwable): Boolean = { + if (Thread.currentThread().isInterrupted) { + return true + } + + val seen = new IdentityHashMap[Throwable, java.lang.Boolean]() + var current = error + while (current != null && !seen.containsKey(current)) { + seen.put(current, java.lang.Boolean.TRUE) + current match { + case _: InterruptedException | _: InterruptedIOException | + _: ClosedByInterruptException | _: CancellationException | _: TaskKilledException | + _: VirtualMachineError | _: ThreadDeath | _: LinkageError | _: ControlThrowable => + return true + case _ => current = current.getCause + } + } + false + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala index 345c1d5d635f2..a7adfdf4049e1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{Command, LogicalPlan, Resolv import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION import org.apache.spark.sql.catalyst.util.sideBySide import org.apache.spark.sql.classic.{Dataset, SparkSession} -import org.apache.spark.sql.connector.catalog.CatalogPlugin +import org.apache.spark.sql.connector.catalog.{CatalogPlugin, CatalogV2Util} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.{IdentifierHelper, MultipartIdentifierHelper} import org.apache.spark.sql.connector.catalog.Identifier import org.apache.spark.sql.connector.catalog.transactions.Transaction @@ -42,6 +42,7 @@ import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.{FileIndex, HadoopFsRelation, LogicalRelation, LogicalRelationWithTable} import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, ExtractV2CatalogAndIdentifier, ExtractV2Table, FileTable, V2TableRefreshUtil} import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.storage.StorageLevel import org.apache.spark.storage.StorageLevel.MEMORY_AND_DISK @@ -419,9 +420,9 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { try { EliminateSubqueryAliases(plan) match { case r @ ExtractV2CatalogAndIdentifier(catalog, ident) if r.timeTravelSpec.isEmpty => - val table = catalog.loadTable(ident) + val table = CatalogV2Util.getTable(catalog, ident, options = r.options) if (r.table.id == table.id) { - Some(DataSourceV2Relation.create(table, Some(catalog), Some(ident))) + Some(DataSourceV2Relation.create(table, Some(catalog), Some(ident), r.options)) } else { None } @@ -436,17 +437,25 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper { } private[sql] def lookupCachedTable( - name: Seq[String], + catalog: CatalogPlugin, + ident: Identifier, + tableId: Option[String], + stateOptions: CaseInsensitiveStringMap, resolver: Resolver): Option[LogicalPlan] = { + val name = ident.toQualifiedNameParts(catalog) val cachedRelations = findCachedRelations(name, resolver) - cachedRelations match { - case cachedRelation +: _ => - CacheManager.logCacheOperation( - log"Relation cache hit for table ${MDC(TABLE_NAME, name.quoted)}") - Some(cachedRelation) - case _ => - None + val cachedRelation = cachedRelations.collectFirst { + case r: DataSourceV2Relation + if r.catalog.contains(catalog) && r.identifier.contains(ident) && + tableId.forall(_ == r.table.id) && + CatalogV2Util.extractTableStateOptions(catalog, r.options) == stateOptions => + r + } + cachedRelation.foreach { _ => + CacheManager.logCacheOperation( + log"Relation cache hit for table ${MDC(TABLE_NAME, name.quoted)}") } + cachedRelation } private def findCachedRelations( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala index 00643e1638fff..f9d841615d427 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CombineAdjacentAggregation.scala @@ -17,13 +17,15 @@ package org.apache.spark.sql.execution -import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final, Partial} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Complete, Final, Partial, PartialMerge} import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf /** - * This rule combines adjacent aggregation with `Partial` and `Final` to `Complete` mode. + * This rule combines adjacent aggregation with `Partial` and `Final` to `Complete` mode. For + * [[HashAggregateExec]], it also combines `PartialMerge` and `Final` to `Final` mode. The latter + * can be produced by physical plan extensions that add an extra aggregation stage. * Example for hash aggregate: * HashAggregate (Final) HashAggregate (Complete) * | | @@ -43,19 +45,20 @@ import org.apache.spark.sql.internal.SQLConf * It supports [[HashAggregateExec]], [[SortAggregateExec]] and [[ObjectHashAggregateExec]]. */ object CombineAdjacentAggregation extends Rule[SparkPlan] { + private case class CombinedAggregate( + aggregateExpressions: Seq[AggregateExpression], + initialInputBufferOffset: Int) + override def apply(plan: SparkPlan): SparkPlan = { if (!conf.getConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED)) { return plan } plan.transformDown { - case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg: HashAggregateExec) - if isPartialAgg(partialAgg, finalAgg) => - finalAgg.copy( - groupingExpressions = partialAgg.groupingExpressions, - aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), - initialInputBufferOffset = 0, - child = partialAgg.child) + case finalAgg @ HashAggregateExec(_, _, _, _, _, _, _, _, partialAgg: HashAggregateExec) => + combinedAggregate(partialAgg, finalAgg) + .map(combineHashAggregates(partialAgg, finalAgg, _)) + .getOrElse(finalAgg) case finalAgg @ SortAggregateExec(_, _, _, _, _, _, _, _, partialAgg: SortAggregateExec) if isPartialAgg(partialAgg, finalAgg) => @@ -76,6 +79,41 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] { } } + private def combineHashAggregates( + partialAgg: HashAggregateExec, + finalAgg: HashAggregateExec, + combined: CombinedAggregate): HashAggregateExec = { + // Keep the final aggregate's distribution requirement because the rule runs after + // EnsureRequirements. The other child-facing metadata comes from the removed aggregate. + finalAgg.copy( + isStreaming = partialAgg.isStreaming, + numShufflePartitions = partialAgg.numShufflePartitions, + groupingExpressions = partialAgg.groupingExpressions, + aggregateExpressions = combined.aggregateExpressions, + initialInputBufferOffset = combined.initialInputBufferOffset, + child = partialAgg.child) + } + + private def combinedAggregate( + partialAgg: HashAggregateExec, + finalAgg: HashAggregateExec): Option[CombinedAggregate] = { + if (!isCompatibleAggregates(partialAgg, finalAgg)) { + None + } else if (partialAgg.aggregateExpressions.forall(_.mode == Partial)) { + Some(CombinedAggregate( + partialAgg.aggregateExpressions.map(_.copy(mode = Complete)), + initialInputBufferOffset = 0)) + } else if (partialAgg.aggregateExpressions.forall(_.mode == PartialMerge) && + partialAgg.aggregateExpressions.forall(_.filter.isEmpty) && + finalAgg.aggregateExpressions.forall(_.filter.isEmpty)) { + Some(CombinedAggregate( + finalAgg.aggregateExpressions, + partialAgg.initialInputBufferOffset)) + } else { + None + } + } + /** * Check if `partialAgg` is the partial aggregate of `finalAgg`. */ @@ -83,7 +121,13 @@ object CombineAdjacentAggregation extends Rule[SparkPlan] { partialAgg: BaseAggregateExec, finalAgg: BaseAggregateExec): Boolean = { partialAgg.aggregateExpressions.forall(_.mode == Partial) && - finalAgg.aggregateExpressions.forall(_.mode == Final) && + isCompatibleAggregates(partialAgg, finalAgg) + } + + private def isCompatibleAggregates( + partialAgg: BaseAggregateExec, + finalAgg: BaseAggregateExec): Boolean = { + finalAgg.aggregateExpressions.forall(_.mode == Final) && partialAgg.groupingExpressions.map(_.canonicalized) == finalAgg.groupingExpressions.map(_.canonicalized) && finalAgg.logicalLink.isDefined && diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala index ba1238564348e..c40e79e61d784 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExpandExec.scala @@ -137,6 +137,11 @@ case class ExpandExec( * * We use a for loop here so we only includes one copy of the consume code and avoid code * size explosion. + * + * In addition, when subexpression elimination is enabled, common subexpressions shared by + * the branch expressions (e.g. an expensive condition repeated in many branches) are + * evaluated only once per input row, before the loop, since all the branches consume the + * same input row. */ // Tracks whether a column has the same output for all rows. @@ -147,42 +152,62 @@ case class ExpandExec( projections.map(p => p(colIndex)).toSet.size == 1 }.toArray + // Bind all the branch expressions once up front, so that identical expressions appearing in + // different branches bind to identical trees and can be deduplicated by subexpression + // elimination below. + val boundProjections: Seq[Seq[Expression]] = projections.map { exprs => + BindReferences.bindReferences(exprs, child.output) + } + + // Set up subexpression elimination over all the branch expressions. This deduplicates + // repeated subexpressions both within a branch and across branches. The code evaluating the + // common subexpressions is emitted once before the branch loop (see the end of this method). + val subExprs: SubExprCodes = if (conf.subexpressionEliminationEnabled) { + ctx.subexpressionEliminationForWholeStageCodegen(boundProjections.flatten) + } else { + SubExprCodes(Map.empty, Seq.empty) + } + // Part 1: declare variables for each column // If a column has the same value for all output rows, then we also generate its computation // right after declaration. Otherwise its value is computed in the part 2. - lazy val attributeSeq: AttributeSeq = child.output - val outputColumns = output.indices.map { col => - val firstExpr = projections.head(col) - if (sameOutput(col)) { - // This column is the same across all output rows. Just generate code for it here. - BindReferences.bindReference(firstExpr, attributeSeq).genCode(ctx) - } else { - val isNull = ctx.addMutableState( - CodeGenerator.JAVA_BOOLEAN, - "resultIsNull", - v => s"$v = true;") - val value = ctx.addMutableState( - CodeGenerator.javaType(firstExpr.dataType), - "resultValue", - v => s"$v = ${CodeGenerator.defaultValue(firstExpr.dataType)};") + val outputColumns = ctx.withSubExprEliminationExprs(subExprs.states) { + output.indices.map { col => + val firstExpr = boundProjections.head(col) + if (sameOutput(col)) { + // This column is the same across all output rows. Just generate code for it here. + firstExpr.genCode(ctx) + } else { + val isNull = ctx.addMutableState( + CodeGenerator.JAVA_BOOLEAN, + "resultIsNull", + v => s"$v = true;") + val value = ctx.addMutableState( + CodeGenerator.javaType(firstExpr.dataType), + "resultValue", + v => s"$v = ${CodeGenerator.defaultValue(firstExpr.dataType)};") - ExprCode( - JavaCode.isNullVariable(isNull), - JavaCode.variable(value, firstExpr.dataType)) + ExprCode( + JavaCode.isNullVariable(isNull), + JavaCode.variable(value, firstExpr.dataType)) + } } } // Part 2: switch/case statements - val switchCaseExprs = projections.zipWithIndex.map { case (exprs, row) => - val (exprCodesWithIndices, inputVarSets) = exprs.indices.flatMap { col => - if (!sameOutput(col)) { - val boundExpr = BindReferences.bindReference(exprs(col), attributeSeq) - val exprCode = boundExpr.genCode(ctx) - val inputVars = CodeGenerator.getLocalInputVariableValues(ctx, boundExpr)._1 - Some(((col, exprCode), inputVars)) - } else { - None - } + val switchCaseExprs = projections.indices.map { row => + val colsToGenerate = projections(row).indices.filter(col => !sameOutput(col)) + val exprCodes = ctx.withSubExprEliminationExprs(subExprs.states) { + colsToGenerate.map(col => boundProjections(row)(col).genCode(ctx)) + } + val (exprCodesWithIndices, inputVarSets) = colsToGenerate.zip(exprCodes).map { + case (col, exprCode) => + // Pass `subExprs.states` so that the input variables of the split switch/case + // functions below include the variables holding the common subexpression values, + // which are evaluated outside of the split functions. + val inputVars = CodeGenerator.getLocalInputVariableValues( + ctx, boundProjections(row)(col), subExprs.states)._1 + ((col, exprCode), inputVars) }.unzip val inputVars = inputVarSets.foldLeft(Set.empty[VariableValue])(_ ++ _) @@ -239,7 +264,14 @@ case class ExpandExec( val i = ctx.freshName("i") // these column have to declared before the loop. val evaluate = evaluateVariables(outputColumns) + // The input variables used by the common subexpressions have to be evaluated first, then + // the common subexpressions themselves, both before the loop since every branch consumes + // the same input row. + val evaluateSubExprInputs = evaluateVariables(subExprs.exprCodesNeedEvaluate) + val evaluateSubExprs = ctx.evaluateSubExprEliminationState(subExprs.states.values) s""" + |$evaluateSubExprInputs + |$evaluateSubExprs |$evaluate |for (int $i = 0; $i < ${projections.length}; $i ++) { | switch ($i) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExplainUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExplainUtils.scala index 2e878c21dc7a6..a81ddb1b1da1b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExplainUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExplainUtils.scala @@ -87,32 +87,7 @@ object ExplainUtils extends AdaptiveSparkPlanHelper { // overwrites and to allow intentional overwriting of IDs generated in previous AQE iteration val idMap = new IdentityHashMap[QueryPlan[_], Int]() localIdMap.set(idMap) - // Initialize an array of ReusedExchanges to help find Adaptively Optimized Out - // Exchanges as part of SPARK-42753 - val reusedExchanges = ArrayBuffer.empty[ReusedExchangeExec] - - var currentOperatorID = 0 - currentOperatorID = generateOperatorIDs(plan, currentOperatorID, idMap, reusedExchanges, - true) - - val subqueries = ArrayBuffer.empty[(SparkPlan, Expression, BaseSubqueryExec)] - getSubqueries(plan, subqueries) - - currentOperatorID = subqueries.foldLeft(currentOperatorID) { - (curId, plan) => generateOperatorIDs(plan._3.child, curId, idMap, reusedExchanges, - true) - } - - // SPARK-42753: Process subtree for a ReusedExchange with unknown child - val optimizedOutExchanges = ArrayBuffer.empty[Exchange] - reusedExchanges.foreach{ reused => - val child = reused.child - if (!idMap.containsKey(child)) { - optimizedOutExchanges.append(child) - currentOperatorID = generateOperatorIDs(child, currentOperatorID, idMap, - reusedExchanges, false) - } - } + val (subqueries, optimizedOutExchanges) = assignOperatorIds(plan, idMap) val collectedOperators = BitSet.empty processPlanSkippingSubqueries(plan, append, collectedOperators) @@ -150,6 +125,41 @@ object ExplainUtils extends AdaptiveSparkPlanHelper { } } + /** + * Assigns operator IDs to all operators across the full plan tree -- the main plan, + * any subqueries, and any adaptively-optimized-out exchanges (SPARK-42753) -- by populating + * the supplied idMap. Returns the discovered subqueries and optimized-out exchanges so + * the caller ([[processPlan]]) can perform post-assignment work (text output) without + * rediscovering them. + */ + private def assignOperatorIds( + plan: QueryPlan[_], + idMap: java.util.Map[QueryPlan[_], Int]) + : (ArrayBuffer[(SparkPlan, Expression, BaseSubqueryExec)], ArrayBuffer[Exchange]) = { + // Initialize an array of ReusedExchanges to help find Adaptively Optimized Out + // Exchanges as part of SPARK-42753 + val reusedExchanges = ArrayBuffer.empty[ReusedExchangeExec] + var currentOperatorID = generateOperatorIDs(plan, 0, idMap, reusedExchanges, true) + + val subqueries = ArrayBuffer.empty[(SparkPlan, Expression, BaseSubqueryExec)] + getSubqueries(plan, subqueries) + currentOperatorID = subqueries.foldLeft(currentOperatorID) { + (curId, sub) => generateOperatorIDs(sub._3.child, curId, idMap, reusedExchanges, true) + } + + // SPARK-42753: Process subtree for a ReusedExchange with unknown child + val optimizedOutExchanges = ArrayBuffer.empty[Exchange] + reusedExchanges.foreach { reused => + val child = reused.child + if (!idMap.containsKey(child)) { + optimizedOutExchanges.append(child) + currentOperatorID = generateOperatorIDs(child, currentOperatorID, idMap, + reusedExchanges, false) + } + } + (subqueries, optimizedOutExchanges) + } + /** * Traverses the supplied input plan in a bottom-up fashion and records the operator id via * setting a tag in the operator. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExternalAppendOnlyUnsafeRowArray.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExternalAppendOnlyUnsafeRowArray.scala index 7ca9d9f85f484..5b9944772f16c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExternalAppendOnlyUnsafeRowArray.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExternalAppendOnlyUnsafeRowArray.scala @@ -122,7 +122,7 @@ class ExternalAppendOnlyUnsafeRowArray( spillableArray = null } else if (inMemoryBuffer != null) { inMemoryBuffer.clear() - inMemoryBufferSizeInBytes = 0; + inMemoryBufferSizeInBytes = 0 } numFieldsPerRow = 0 numRows = 0 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala index aa29128cda7e0..88f0de2555f70 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala @@ -44,7 +44,7 @@ object InsertSortForLimitAndOffset extends Rule[SparkPlan] { _, // Should not match AQE shuffle stage because we only target un-submitted stages which // we can still rewrite the query plan. - s @ ShuffleExchangeExec(SinglePartition, child, _, _), + s @ ShuffleExchangeExec(SinglePartition, child, _, _, _), _) if child.logicalLink.isDefined => extractOrderingAndPropagateOrderingColumns(child) match { case Some((ordering, newChild)) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ProjectedBroadcastValueSubqueryExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ProjectedBroadcastValueSubqueryExec.scala new file mode 100644 index 0000000000000..25132c38f9c29 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/ProjectedBroadcastValueSubqueryExec.scala @@ -0,0 +1,136 @@ +/* + * 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.execution + +import java.util.concurrent.{Future => JFuture} + +import scala.concurrent.duration.Duration + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Cast, Expression, NamedExpression, UnsafeRow} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.BroadcastValueResult.{Available, Unavailable} +import org.apache.spark.sql.execution.joins.HashedRelation +import org.apache.spark.sql.execution.metric.SQLMetrics +import org.apache.spark.util.ThreadUtils + +/** Collects pruning values from the full rows of an already required hash broadcast. */ +case class ProjectedBroadcastValueSubqueryExec( + name: String, + valueExpression: Expression, + child: SparkPlan) extends BaseSubqueryExec with UnaryExecNode { + + override def output: Seq[Attribute] = { + val outputName = valueExpression match { + case named: NamedExpression => named.name + case Cast(named: NamedExpression, _, _, _) => named.name + case _ => "key" + } + Seq(AttributeReference(outputName, valueExpression.dataType, valueExpression.nullable)()) + } + + override lazy val metrics = Map( + "numInputRows" -> SQLMetrics.createMetric(sparkContext, "number of input rows"), + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "dataSize" -> SQLMetrics.createMetric(sparkContext, "data size (bytes)"), + "projectionDisabled" -> SQLMetrics.createMetric(sparkContext, "projection disabled"), + "projectionErrors" -> SQLMetrics.createMetric(sparkContext, "projection errors"), + "collectTime" -> SQLMetrics.createMetric(sparkContext, "time to collect (ms)")) + + override def doCanonicalize(): SparkPlan = { + ProjectedBroadcastValueSubqueryExec( + "dpp", + QueryPlan.normalizeExpressions(valueExpression, child.output), + child.canonicalized) + } + + @transient + private lazy val relationFuture: JFuture[BroadcastValueResult[Array[InternalRow]]] = { + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + SQLExecution.withThreadLocalCaptured[BroadcastValueResult[Array[InternalRow]]]( + session, SubqueryBroadcastExec.executionContext) { + SQLExecution.withExecutionId(session, executionId) { + val beforeCollect = System.nanoTime() + val broadcast = child.executeBroadcast[HashedRelation]() + val limits = BroadcastValueProjectionLimits( + maxInputRows = conf.dynamicPartitionPruningBroadcastProjectionMaxRows.toLong, + maxOutputBytes = conf.dynamicPartitionPruningBroadcastProjectionMaxBytes, + maxSourceBytes = conf.dynamicPartitionPruningBroadcastProjectionMaxSourceBytes) + val result = BroadcastValueProjector.collectExactValueDomain( + broadcast, + child, + valueExpression, + limits, + onInputRow = () => longMetric("numInputRows") += 1, + onError = () => longMetric("projectionErrors") += 1) + + longMetric("collectTime") += (System.nanoTime() - beforeCollect) / 1000000 + result match { + case Available(rows) => + longMetric("numOutputRows") += rows.length + longMetric("dataSize") += + rows.iterator.map(_.asInstanceOf[UnsafeRow].getSizeInBytes.toLong).sum + case Unavailable => + longMetric("projectionDisabled") += 1 + } + SQLMetrics.postDriverMetricUpdates(sparkContext, executionId, metrics.values.toSeq) + result + } + } + } + + override protected def doPrepare(): Unit = { + relationFuture + } + + override protected def doExecute(): RDD[InternalRow] = { + throw QueryExecutionErrors.executeCodePathUnsupportedError( + "ProjectedBroadcastValueSubqueryExec") + } + + override def executeCollect(): Array[InternalRow] = executeCollectResult() match { + case Available(rows) => rows + case Unavailable => + throw SparkException.internalError( + "An unavailable projected broadcast value domain must be consumed through " + + "executeCollectResult.") + } + + private[execution] def executeCollectResult(): BroadcastValueResult[Array[InternalRow]] = + ThreadUtils.awaitResult(relationFuture, Duration.Inf) + + override def stringArgs: Iterator[Any] = super.stringArgs ++ Iterator(s"[id=#$id]") + + override protected def withNewChildInternal( + newChild: SparkPlan): ProjectedBroadcastValueSubqueryExec = copy(child = newChild) +} + +object ProjectedBroadcastValueSubqueryExec { + private[execution] def resultOf( + plan: BaseSubqueryExec): Option[BroadcastValueResult[Array[InternalRow]]] = { + plan match { + case projected: ProjectedBroadcastValueSubqueryExec => + Some(projected.executeCollectResult()) + case ReusedSubqueryExec(child) => resultOf(child) + case _ => None + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 0a1185aa4ea57..51e54ef5f7536 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -308,6 +308,13 @@ class QueryExecution( def assertCommandExecuted(): Unit = commandExecuted + private def cloneWithFreshStatefulExpressions(plan: LogicalPlan): LogicalPlan = { + plan.clone().transformDownWithSubqueriesAndReferenceEquality { + case node => + node.mapExpressionsWithReferenceEquality(_.freshCopyIfContainsStatefulExpression()) + } + } + private val lazyOptimizedPlan = LazyTry { // We need to materialize the commandExecuted here because optimizedPlan is also tracked under // the optimizing phase @@ -315,8 +322,8 @@ class QueryExecution( executePhase(QueryPlanningTracker.OPTIMIZATION) { // clone the plan to avoid sharing the plan instance between different stages like analyzing, // optimizing and planning. - val plan = - sparkSession.sessionState.optimizer.executeAndTrack(withCachedData.clone(), tracker) + val plan = sparkSession.sessionState.optimizer.executeAndTrack( + cloneWithFreshStatefulExpressions(withCachedData), tracker) // We do not want optimized plans to be re-analyzed as literals that have been constant // folded and such can cause issues during analysis. While `clone` should maintain the // `analyzed` state of the LogicalPlan, we set the plan as analyzed here as well out of @@ -376,7 +383,10 @@ class QueryExecution( def assertExecutedPlanPrepared(): Unit = executedPlan val lazyToRdd = LazyTry { - new SQLExecutionRDD(executedPlan.execute(), sparkSession.sessionState.conf) + new SQLExecutionRDD( + executedPlan.execute(), + sparkSession.sessionState.conf, + SparkPlanInfo.fromSparkPlan(executedPlan)) } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecutionRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecutionRDD.scala index 45b9cadc4aeda..7c38ccee49831 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecutionRDD.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecutionRDD.scala @@ -29,9 +29,18 @@ import org.apache.spark.sql.internal.SQLConf * * @param sqlRDD the `RDD` generated by the SQL plan * @param conf the `SQLConf` to apply to the execution of the SQL plan + * @param sparkPlanInfo the physical plan information for `sqlRDD` */ class SQLExecutionRDD( - var sqlRDD: RDD[InternalRow], @transient conf: SQLConf) extends RDD[InternalRow](sqlRDD) { + var sqlRDD: RDD[InternalRow], + @transient conf: SQLConf, + @transient val sparkPlanInfo: SparkPlanInfo = SparkPlanInfo.EMPTY) + extends RDD[InternalRow](sqlRDD) { + + def this(sqlRDD: RDD[InternalRow], conf: SQLConf) = { + this(sqlRDD, conf, SparkPlanInfo.EMPTY) + } + private val sqlConfigs = conf.getAllConfs private lazy val sqlConfExecutorSide = { val newConf = new SQLConf() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala index 74afa396264e4..d0f7a6bc32910 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala @@ -146,7 +146,7 @@ case class SortExec( v => s"$v = $thisPlan.createSorter();", forceInline = true) val metrics = ctx.addMutableState(classOf[TaskMetrics].getName, "metrics", v => s"$v = org.apache.spark.TaskContext.get().taskMetrics();", forceInline = true) - val sortedIterator = ctx.addMutableState("scala.collection.Iterator<UnsafeRow>", "sortedIter", + val sortedIterator = ctx.addMutableState("scala.collection.Iterator<InternalRow>", "sortedIter", forceInline = true) val addToSorter = ctx.freshName("addToSorter") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala index 0994220afe0c1..2e9520742b0ab 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkOptimizer.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.execution.datasources.{MarkSingleTaskExecution, PruneFileSourcePartitions, PullOutVariantExtractions, PushVariantIntoScan, SchemaPruning, V1Writes} import org.apache.spark.sql.execution.datasources.v2.{GroupBasedRowLevelOperationScanPlanning, OptimizeMetadataOnlyDeleteFromTable, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown, V2Writes} import org.apache.spark.sql.execution.dynamicpruning.{CleanupDynamicPruningFilters, PartitionPruning, RowLevelOperationRuntimeGroupFiltering} +import org.apache.spark.sql.execution.planmerging.MergeSubplans import org.apache.spark.sql.execution.python.{ExtractGroupingPythonUDFFromAggregate, ExtractPythonUDFFromAggregate, ExtractPythonUDFs, ExtractPythonUDTFs} class SparkOptimizer( @@ -71,6 +72,7 @@ class SparkOptimizer( InjectRuntimeFilter), Batch("MergeSubplans", Once, MergeSubplans, + CombineApproximatePercentiles, RewriteDistinctAggregates), Batch("Pushdown Filters from PartitionPruning", fixedPoint, PushDownPredicates), @@ -88,6 +90,8 @@ class SparkOptimizer( ExtractPythonUDFFromAggregate, // This must be executed after `ExtractPythonUDFFromAggregate` and before `ExtractPythonUDFs`. ExtractGroupingPythonUDFFromAggregate, + // `ExtractPythonUDFs` first lifts Python UDFs out of higher-order function lambdas + // (via `ExtractPythonUDFFromLambda`) and then extracts them as ordinary top-level UDFs. ExtractPythonUDFs, ExtractPythonUDTFs, // The eval-python node may be between Project/Filter and the scan node, which breaks @@ -114,6 +118,8 @@ class SparkOptimizer( ExtractPythonUDFFromJoinCondition.ruleName, ExtractPythonUDFFromAggregate.ruleName, ExtractGroupingPythonUDFFromAggregate.ruleName, + // Non-excludable: a plan with a Python UDF in a higher-order function lambda only works + // because `ExtractPythonUDFs` lifts it out (via `ExtractPythonUDFFromLambda`). ExtractPythonUDFs.ruleName, GroupBasedRowLevelOperationScanPlanning.ruleName, V2ScanRelationPushDown.ruleName, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala index 4410fe50912f7..32bff876d454d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala @@ -17,7 +17,10 @@ package org.apache.spark.sql.execution +import scala.util.control.NonFatal + import org.apache.spark.annotation.DeveloperApi +import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.plans.logical.{EmptyRelation, LogicalPlan} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} import org.apache.spark.sql.execution.adaptive.LogicalQueryStage @@ -74,6 +77,7 @@ private[execution] object SparkPlanInfo { case a: AdaptiveSparkPlanExec => a.executedPlan :: Nil case stage: QueryStageExec => stage.plan :: Nil case inMemTab: InMemoryTableScanExec => inMemTab.relation.cachedPlan :: Nil + case rddScan: RDDScanExec => sparkPlanInfosFromRDD(rddScan.rdd) case EmptyRelationExec(logical) => (logical :: Nil) case _ => plan.children ++ plan.subqueries } @@ -91,6 +95,8 @@ private[execution] object SparkPlanInfo { Some(fromSparkPlan(child)) case child: LogicalPlan => Some(fromLogicalPlan(child)) + case child: SparkPlanInfo => + Some(child) case _ => None } new SparkPlanInfo( @@ -102,4 +108,31 @@ private[execution] object SparkPlanInfo { } final lazy val EMPTY: SparkPlanInfo = new SparkPlanInfo("", "", Nil, Map.empty, Nil) + + private def sparkPlanInfosFromRDD(rdd: RDD[_]): Seq[SparkPlanInfo] = { + // Walk only driver-side RDD dependency metadata. Dedupe by RDD id so shared lineage does not + // duplicate the same internal SQL plan under a single RDDScanExec. + val visitedRDDs = scala.collection.mutable.HashSet.empty[Int] + val rddsToVisit = scala.collection.mutable.Queue.empty[RDD[_]] + val planInfos = scala.collection.mutable.ArrayBuffer.empty[SparkPlanInfo] + + rddsToVisit.enqueue(rdd) + while (rddsToVisit.nonEmpty) { + val current = rddsToVisit.dequeue() + if (visitedRDDs.add(current.id)) { + current match { + case sqlRDD: SQLExecutionRDD if sqlRDD.sparkPlanInfo != EMPTY => + planInfos += sqlRDD.sparkPlanInfo + case _ => + try { + current.dependencies.foreach(dep => rddsToVisit.enqueue(dep.rdd)) + } catch { + case NonFatal(_) => + } + } + } + } + + planInfos.toSeq + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala index 99bb63041136a..5070a40259e37 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala @@ -686,7 +686,7 @@ class SparkSqlAstBuilder extends AstBuilder { } override def visitFailSetRole(ctx: FailSetRoleContext): LogicalPlan = withOrigin(ctx) { - invalidStatement("SET ROLE", ctx); + invalidStatement("SET ROLE", ctx) } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala index 996ecac61b728..d395b3f986e83 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala @@ -39,7 +39,7 @@ import org.apache.spark.sql.execution.aggregate.AggUtils import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.command._ import org.apache.spark.sql.execution.datasources.{LogicalRelation, WriteFiles, WriteFilesExec} -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, StreamingDataSourceV2ScanRelation} import org.apache.spark.sql.execution.exchange.{REBALANCE_PARTITIONS_BY_COL, REBALANCE_PARTITIONS_BY_NONE, REPARTITION_BY_COL, REPARTITION_BY_NUM, ShuffleExchangeExec} import org.apache.spark.sql.execution.python._ import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec} @@ -86,6 +86,19 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { } } + /** + * Whether this plan reads a streaming source in Real-Time Mode, which is the case when the + * relation carries a real-time mode duration -- the same signal that decides whether to plan + * a [[org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec]] for it. + */ + private def isRealTimeMode(plan: LogicalPlan): Boolean = { + plan.collectLeaves().exists { + case s: StreamingDataSourceV2ScanRelation => + s.relation.realTimeModeDuration.isDefined + case _ => false + } + } + /** * Plans special cases of limit operators. */ @@ -566,7 +579,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { case PhysicalAggregation( namedGroupingExpressions, aggregateExpressions, rewrittenResultExpressions, child) => - if (aggregateExpressions.exists(_.aggregateFunction.isInstanceOf[PythonUDAF])) { + if (aggregateExpressions.exists(ae => + ae.aggregateFunction.isInstanceOf[PythonUDAF] || + ae.aggregateFunction.isInstanceOf[PythonAggregate])) { throw new AnalysisException( errorClass = "_LEGACY_ERROR_TEMP_3067", messageParameters = Map.empty) @@ -601,12 +616,26 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { case None => val stateVersion = conf.getConf(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION) - AggUtils.planStreamingAggregation( - normalizedGroupingExpressions, - aggregateExpressions, - rewrittenResultExpressions, - stateVersion, - planLater(child)) + // A Real-Time Mode batch runs until its duration elapses rather than until its input + // is exhausted, so an aggregation that only emits once the batch ends would hold every + // result back for the whole batch. Plan the streamline operator instead, which merges + // each input row against state and emits immediately. + if (isRealTimeMode(child) || + conf.getConf(SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR)) { + AggUtils.planStreamlineStreamingAggregation( + normalizedGroupingExpressions, + aggregateExpressions, + rewrittenResultExpressions, + stateVersion, + planLater(child)) + } else { + AggUtils.planStreamingAggregation( + normalizedGroupingExpressions, + aggregateExpressions, + rewrittenResultExpressions, + stateVersion, + planLater(child)) + } } case _ => Nil @@ -692,7 +721,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { object Aggregation extends Strategy { def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { case PhysicalAggregation(groupingExpressions, aggExpressions, resultExpressions, child) - if !aggExpressions.exists(_.aggregateFunction.isInstanceOf[PythonUDAF]) => + if !aggExpressions.exists(ae => ae.aggregateFunction.isInstanceOf[PythonUDAF] || + ae.aggregateFunction.isInstanceOf[PythonAggregate]) => val (functionsWithDistinct, functionsWithoutDistinct) = aggExpressions.partition(_.isDistinct) val distinctAggChildSets = functionsWithDistinct.map { ae => @@ -771,13 +801,40 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { resultExpressions, planLater(child))) + case PhysicalAggregation(groupingExpressions, aggExpressions, resultExpressions, child) + if aggExpressions.forall(_.aggregateFunction.isInstanceOf[PythonAggregate]) => + // Ideally this should be done in `NormalizeFloatingNumbers`, but we do it here because + // `groupingExpressions` is not extracted during logical phase. Without this, 0.0/-0.0 (and + // distinct NaN bit patterns) would split one logical group across output rows. + val normalizedGroupingExpressions = groupingExpressions.map { e => + NormalizeFloatingNumbers.normalize(e) match { + case n: NamedExpression => n + // Keep the name of the original expression. + case other => Alias(other, e.name)(exprId = e.exprId) + } + } + Seq(execution.python.PythonIncrementalAggregateExec.plan( + normalizedGroupingExpressions, + aggExpressions, + resultExpressions, + planLater(child))) + case PhysicalAggregation(_, aggExpressions, _, _) => - val groupAggPandasUDFNames = aggExpressions - .map(_.aggregateFunction) - .filter(_.isInstanceOf[PythonUDAF]) - .map(_.asInstanceOf[PythonUDAF].name) - // If cannot match the two cases above, then it's an error - throw QueryCompilationErrors.invalidPandasUDFPlacementError(groupAggPandasUDFNames.distinct) + // Reached when Python aggregate UDFs cannot be planned by the two cases above -- e.g. a + // grouped-agg pandas/arrow UDF or an incremental Python aggregator is mixed with other + // (SQL or differently-typed Python) aggregate functions in the same Aggregate. + val aggFunctions = aggExpressions.map(_.aggregateFunction) + val incrementalNames = aggFunctions.collect { case p: PythonAggregate => p.name } + val pandasUDFNames = aggFunctions.collect { case p: PythonUDAF => p.name } + if (incrementalNames.nonEmpty) { + // The message for pandas UDFs is misleading for an Arrow-based incremental aggregator, so + // report the dedicated, aggregator-neutral error. Name every offending Python aggregate + // function -- both the incremental aggregators and any grouped-agg pandas/arrow UDAFs + // mixed in -- so the diagnostic is complete rather than dropping the co-offenders. + throw QueryCompilationErrors.invalidPythonAggregatePlacementError( + (incrementalNames ++ pandasUDFNames).distinct) + } + throw QueryCompilationErrors.invalidPandasUDFPlacementError(pandasUDFNames.distinct) case _ => Nil } @@ -900,6 +957,7 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { eventTimeWatermarkForEviction = None, planLater(child), isStreaming = true, + isRealTimeMode = isRealTimeMode(plan), hasInitialState, initialStateGroupingAttrs, initialStateDataAttrs, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala index 555f4f41d3cd2..c435f06ffddc4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala @@ -38,7 +38,11 @@ case class SubqueryAdaptiveBroadcastExec( onlyInBroadcast: Boolean, @transient buildPlan: LogicalPlan, buildKeys: Seq[Expression], - child: SparkPlan) extends BaseSubqueryExec with UnaryExecNode { + child: SparkPlan)( + @transient private[sql] val broadcastValueProjection: Option[BroadcastValueProjection] = None) + extends BaseSubqueryExec with UnaryExecNode { + + override protected def otherCopyArgs: Seq[AnyRef] = Seq(broadcastValueProjection) protected override def doExecute(): RDD[InternalRow] = { throw QueryExecutionErrors.executeCodePathUnsupportedError("SubqueryAdaptiveBroadcastExec") @@ -46,9 +50,9 @@ case class SubqueryAdaptiveBroadcastExec( protected override def doCanonicalize(): SparkPlan = { val keys = buildKeys.map(k => QueryPlan.normalizeExpressions(k, child.output)) - copy(name = "dpp", buildKeys = keys, child = child.canonicalized) + copy(name = "dpp", buildKeys = keys, child = child.canonicalized)(None) } override protected def withNewChildInternal(newChild: SparkPlan): SubqueryAdaptiveBroadcastExec = - copy(child = newChild) + copy(child = newChild)(broadcastValueProjection) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala index e2a013b9e814c..d9107c6a128e7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEPropagateEmptyRelation.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.planning.ExtractSingleColumnNullAwareAntiJo import org.apache.spark.sql.catalyst.plans.logical.EmptyRelation import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.trees.TreePattern.{LOCAL_RELATION, LOGICAL_QUERY_STAGE, TRUE_OR_FALSE_LITERAL} +import org.apache.spark.sql.execution.{BaseLimitExec, SortExec, SparkPlan} import org.apache.spark.sql.execution.aggregate.BaseAggregateExec import org.apache.spark.sql.execution.exchange.{REPARTITION_BY_COL, REPARTITION_BY_NUM, ShuffleExchangeLike} import org.apache.spark.sql.execution.joins.HashedRelationWithAllNullKeys @@ -52,19 +53,39 @@ object AQEPropagateEmptyRelation extends PropagateEmptyRelationBase { // - positive value means an estimated row count which can be over-estimated // - none means the plan has not materialized or the plan can not be estimated private def getEstimatedRowCount(plan: LogicalPlan): Option[BigInt] = plan match { - case LogicalQueryStage(_, stage: QueryStageExec) if stage.isMaterialized => + case LogicalQueryStage(_, physicalPlan) => + getEstimatedRowCount(physicalPlan) + + case _: EmptyRelation => Some(0) + + case _ => None + } + + private def getEstimatedRowCount(plan: SparkPlan): Option[BigInt] = plan match { + case stage: QueryStageExec if stage.isMaterialized => stage.getRuntimeStatistics.rowCount - case LogicalQueryStage(_, agg: BaseAggregateExec) if agg.groupingExpressions.nonEmpty && - agg.child.isInstanceOf[QueryStageExec] => - val stage = agg.child.asInstanceOf[QueryStageExec] - if (stage.isMaterialized) { - stage.getRuntimeStatistics.rowCount + case sort: SortExec => + getEstimatedRowCount(sort.child) + + // A global limit can also drop rows through an offset, so only propagate + // proven emptiness; a zero limit is unconditionally empty. + case limit: BaseLimitExec => + if (limit.limit == 0) { + Some(BigInt(0)) } else { - None + getEstimatedRowCount(limit.child).filter(_ == 0) } - case _: EmptyRelation => Some(0) + // Match the global-aggregate invariant in LogicalQueryStage.computeStats without adopting + // its logical-statistics fallback or its unrestricted physical-plan traversal. + case aggregate: BaseAggregateExec if aggregate.groupingExpressions.isEmpty => + getEstimatedRowCount(aggregate.child).map { rowCount => + if (rowCount == 0) BigInt(1) else rowCount + } + + case aggregate: BaseAggregateExec => + getEstimatedRowCount(aggregate.child) case _ => None } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala index 578e0acd80525..cc3ca90739090 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala @@ -30,7 +30,7 @@ object AQEUtils { // Project/Filter/LocalSort/CollectMetrics. // Note: we only care about `HashPartitioning` as `EnsureRequirements` can only optimize out // user-specified repartition with `HashPartitioning`. - case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _) + case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _, _) if shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM => val numPartitions = if (shuffleOrigin == REPARTITION_BY_NUM) { Some(h.numPartitions) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala index bfe6a9a3f6332..a12b202b5b009 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala @@ -18,7 +18,8 @@ package org.apache.spark.sql.execution.adaptive import java.util -import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue} +import java.util.concurrent.{CompletableFuture, ConcurrentHashMap, LinkedBlockingQueue} +import java.util.concurrent.atomic.AtomicReference import scala.collection.concurrent.TrieMap import scala.collection.mutable @@ -77,6 +78,17 @@ case class AdaptiveSparkPlanExec( @transient private val lock = new Object() + // Access is serialized by the execution context's stage lifecycle lock. + @transient private val uncancelledObsoleteStageIds = mutable.HashSet.empty[Int] + + // Remember every local alias so cancellation and failure handling cover every stage ID. + @transient private val stageIdsByResult = + mutable.HashMap.empty[AtomicReference[Option[Any]], mutable.LinkedHashSet[Int]] + + // A failed native cancellation can leave its exchange marked cancelled while its job still runs. + @transient private val failedStageCancellationResults = + mutable.HashSet.empty[AtomicReference[Option[Any]]] + @transient private val logOnLevel: ( => MessageWithContext) => Unit = logBasedOnLevel(conf.adaptiveExecutionLogLevel) @@ -301,6 +313,7 @@ case class AdaptiveSparkPlanExec( var result = createQueryStages(fun, currentPhysicalPlan, firstRun = true) val events = new LinkedBlockingQueue[StageMaterializationEvent]() val errors = new mutable.ArrayBuffer[Throwable]() + val obsoleteCancelledStageIds = new mutable.HashSet[Int] var stagesToReplace = Seq.empty[QueryStageExec] while (!result.allChildStagesMaterialized) { currentPhysicalPlan = result.newPlan @@ -356,7 +369,9 @@ case class AdaptiveSparkPlanExec( stage.resultOption.set(Some(res)) case StageFailure(stage, ex) => stage.error.set(Some(ex)) - errors.append(ex) + if (!shouldIgnoreObsoleteStageFailure(stage, ex, obsoleteCancelledStageIds)) { + errors.append(ex) + } } // In case of errors, we cancel all running stages and throw exception. @@ -397,6 +412,9 @@ case class AdaptiveSparkPlanExec( currentPhysicalPlan.treeString, newPhysicalPlan.treeString).mkString("\n") logOnLevel(log"Plan changed:\n${MDC(QUERY_PLAN, plans)}") cleanUpTempTags(newPhysicalPlan) + obsoleteCancelledStageIds ++= + cancelObsoleteStages( + newPhysicalPlan, obsoleteStageCandidates(currentPhysicalPlan, stagesToReplace)) currentPhysicalPlan = newPhysicalPlan currentLogicalPlan = newLogicalPlan stagesToReplace = Seq.empty[QueryStageExec] @@ -420,6 +438,204 @@ case class AdaptiveSparkPlanExec( .get.asInstanceOf[T] } + /** Include stages retained by earlier adopted plans as well as newly created query stages. */ + private def obsoleteStageCandidates( + oldPhysicalPlan: SparkPlan, + stagesToReplace: Seq[QueryStageExec]): Seq[QueryStageExec] = { + val visitedStageIds = mutable.HashSet.empty[Int] + (oldPhysicalPlan.collect { case stage: QueryStageExec => stage } ++ stagesToReplace) + .filter(stage => visitedStageIds.add(stage.id)) + } + + /** + * Cancel unfinished exchange stages that are no longer referenced by the adopted plan. + * + * An unfinished stage cannot be hidden inside another stage's plan because parent stages are + * created only after all child stages have materialized. Results reused by another adaptive + * plan, such as a subquery, are separately protected by the context's shared-result tracking. + * A failed or unsafe cancellation leaves the stage cached until its eventual materialization + * event proves whether the stage is still required. + */ + private def cancelObsoleteStages( + newPhysicalPlan: SparkPlan, + candidateStages: Seq[QueryStageExec]): Seq[Int] = { + val retainedStageIds = mutable.HashSet.empty[Int] + val retainedStageResults = mutable.HashSet.empty[AtomicReference[Option[Any]]] + newPhysicalPlan.foreach { + case stage: QueryStageExec => + retainedStageIds += stage.id + retainedStageResults += stage.resultOption + case _ => + } + val obsoleteStages = mutable.LinkedHashMap.empty[ + AtomicReference[Option[Any]], mutable.ArrayBuffer[ExchangeQueryStageExec]] + candidateStages.foreach { + case stage: ExchangeQueryStageExec + if !retainedStageIds.contains(stage.id) && + !retainedStageResults.contains(stage.resultOption) => + obsoleteStages.getOrElseUpdate(stage.resultOption, mutable.ArrayBuffer.empty) += stage + case _ => + } + obsoleteStages.values.toSeq.flatMap { stages => + val stage = stages.head + val reservation = context.withStageLifecycleLock { + if (!stage.isMaterialized && !context.isSharedStageResult(stage.resultOption)) { + val recordedStageIds = stageIdsByResult.get(stage.resultOption).toSeq.flatten + val obsoleteStageIds = (recordedStageIds ++ stages.map(_.id)).distinct + if (failedStageCancellationResults.contains(stage.resultOption)) { + uncancelledObsoleteStageIds ++= obsoleteStageIds + None + } else { + context.reserveStageCancellation(stage.resultOption).map { cancellation => + (cancellation, obsoleteStageIds) + } + } + } else { + None + } + } + reservation.toSeq.flatMap { case (cancellation, obsoleteStageIds) => + try { + withShuffleCancellationLock(stage) { + val shouldCancel = context.withStageLifecycleLock { + if (stage.isMaterialized || context.isSharedStageResult(stage.resultOption)) { + false + } else if (canCancelStage(stage)) { + true + } else { + if (canTrackObsoleteStageFailure(stage)) { + uncancelledObsoleteStageIds ++= obsoleteStageIds + } + false + } + } + if (shouldCancel) { + try { + recordSubmittedShuffleIds(stage) + stage.cancel( + "The query stage is no longer referenced by the current adaptive plan.") + context.withStageLifecycleLock { + removeStageFromCache(stage) + } + obsoleteStageIds + } catch { + case NonFatal(t) => + context.withStageLifecycleLock { + failedStageCancellationResults += stage.resultOption + uncancelledObsoleteStageIds ++= obsoleteStageIds + } + logError(s"Exception in cancelling obsolete query stage: ${stage.treeString}", t) + Seq.empty + } + } else { + Seq.empty + } + } + } finally { + context.withStageLifecycleLock { + context.finishStageCancellation(stage.resultOption, cancellation) + } + } + } + } + } + + /** Serialize shuffle eligibility and cancellation with the exchange's job submission. */ + private def withShuffleCancellationLock[T]( + stage: ExchangeQueryStageExec)(body: => T): T = stage match { + case shuffleStage: ShuffleQueryStageExec => shuffleStage.shuffle.synchronized(body) + case _ => body + } + + /** + * Broadcast cancellation cannot prevent its supplier from submitting another job afterward. + * Delegating shuffles may conceal an already submitted job behind an empty futureAction, and + * submitted shuffle cancellation is asynchronous, so removing its files can race map writers. + */ + private def canCancelStage(stage: ExchangeQueryStageExec): Boolean = stage match { + case _: BroadcastQueryStageExec => false + case shuffleStage: ShuffleQueryStageExec => + val shuffle = shuffleStage.shuffle + val submitted = shuffle.futureAction.get().isDefined + (shuffle.isInstanceOf[ShuffleExchangeExec] || submitted) && + (!submitted || context.qe.shuffleCleanupMode != RemoveShuffleFiles) + case _ => true + } + + /** Opaque delegated shuffles can conceal submitted jobs and their required cleanup IDs. */ + private def canTrackObsoleteStageFailure(stage: ExchangeQueryStageExec): Boolean = stage match { + case shuffleStage: ShuffleQueryStageExec => + shuffleStage.shuffle.isInstanceOf[ShuffleExchangeExec] || + shuffleStage.shuffle.futureAction.get().isDefined + case _ => true + } + + /** Never hide fatal stage failures, even when their stage was cancelled as obsolete. */ + private def shouldIgnoreObsoleteStageFailure( + stage: QueryStageExec, + error: Throwable, + obsoleteCancelledStageIds: scala.collection.Set[Int]): Boolean = { + !error.isInstanceOf[SparkFatalException] && NonFatal(error) && + (obsoleteCancelledStageIds.contains(stage.id) || ignoreFailedObsoleteStageFailure(stage)) + } + + /** + * Suppress an uncancelled obsolete stage's later failure only while its result remains private + * and absent from the current physical plan. This covers failed cancellation and stages whose + * cancellation would be unsafe. Reused stages must still fail their consumers; otherwise, the + * failed cache entry can now be safely discarded. + */ + private def ignoreFailedObsoleteStageFailure(stage: QueryStageExec): Boolean = stage match { + case exchangeStage: ExchangeQueryStageExec => + context.withStageLifecycleLock { + if (!uncancelledObsoleteStageIds.remove(stage.id)) { + false + } else { + val stillReferenced = currentPhysicalPlan.exists { + case retainedStage: QueryStageExec => + retainedStage.id == stage.id || retainedStage.resultOption.eq(stage.resultOption) + case _ => false + } + if (!stillReferenced && !context.isSharedStageResult(stage.resultOption)) { + recordSubmittedShuffleIds(exchangeStage) + removeStageFromCache(exchangeStage) + true + } else { + false + } + } + } + case _ => false + } + + /** Record only submitted shuffles, avoiding initialization of an unsubmitted lazy dependency. */ + private def recordSubmittedShuffleIds(stage: ExchangeQueryStageExec): Unit = { + def record(shuffle: ShuffleExchangeLike): Unit = { + if (shuffle.futureAction.get().isDefined) { + context.shuffleIds.put(shuffle.shuffleId, true) + } + } + stage match { + // A reuse instance wraps its exchange in ReusedExchangeExec, which is a leaf node. + case shuffleStage: ShuffleQueryStageExec => record(shuffleStage.shuffle) + case _ => + stage.plan.foreach { + case shuffle: ShuffleExchangeLike => record(shuffle) + case _ => + } + } + } + + /** Remove the exact canonical exchange-cache entry only when it still owns this result. */ + private def removeStageFromCache(stage: ExchangeQueryStageExec): Unit = { + val cacheKey = stage.plan.canonicalized + context.stageCache.get(cacheKey).foreach { cachedStage => + if (cachedStage.resultOption.eq(stage.resultOption)) { + context.stageCache.remove(cacheKey) + } + } + } + // Use a lazy val to avoid this being called more than once. @transient private lazy val finalPlanUpdate: Unit = { // Do final plan update after result stage has materialized. @@ -623,16 +839,25 @@ case class AdaptiveSparkPlanExec( private def createNonResultQueryStages(plan: SparkPlan): CreateStageResult = plan match { case e: Exchange => // First have a quick check in the `stageCache` without having to traverse down the node. - context.stageCache.get(e.canonicalized) match { - case Some(existingStage) if conf.exchangeReuseEnabled => - val stage = reuseQueryStage(existingStage, e) + val reusedStage = if (conf.exchangeReuseEnabled) { + withStageCacheEntry(e.canonicalized) { + _.map { existingStage => + reuseQueryStage(existingStage, e) + } + } + } else { + None + } + + reusedStage match { + case Some(stage) => val isMaterialized = stage.isMaterialized CreateStageResult( newPlan = stage, allChildStagesMaterialized = isMaterialized, newStages = if (isMaterialized) Seq.empty else Seq(stage)) - case _ => + case None => val result = createNonResultQueryStages(e.child) val newPlan = e.withNewChildren(Seq(result.newPlan)).asInstanceOf[Exchange] // Create a query stage only when all the child query stages are ready. @@ -642,10 +867,14 @@ case class AdaptiveSparkPlanExec( // Check the `stageCache` again for reuse. If a match is found, ditch the new stage // and reuse the existing stage found in the `stageCache`, otherwise update the // `stageCache` with the new stage. - val queryStage = context.stageCache.getOrElseUpdate( - newStage.plan.canonicalized, newStage) - if (queryStage.ne(newStage)) { - newStage = reuseQueryStage(queryStage, e) + val cacheKey = newStage.plan.canonicalized + withStageCacheEntry(cacheKey) { + case Some(queryStage) => + newStage = reuseQueryStage(queryStage, e) + case None => + context.registerStageOwner(newStage.resultOption, this) + recordStageId(newStage) + context.stageCache.put(cacheKey, newStage) } } val isMaterialized = newStage.isMaterialized @@ -685,6 +914,26 @@ case class AdaptiveSparkPlanExec( } } + /** Wait for cancellation of this result without blocking unrelated exchange-cache access. */ + @scala.annotation.tailrec + private def withStageCacheEntry[T]( + cacheKey: SparkPlan)( + useStage: Option[ExchangeQueryStageExec] => T): T = { + val result: Either[CompletableFuture[Unit], T] = context.withStageLifecycleLock { + val cachedStage = context.stageCache.get(cacheKey) + cachedStage.flatMap(stage => context.pendingStageCancellation(stage.resultOption)) match { + case Some(cancellation) => Left(cancellation) + case None => Right(useStage(cachedStage)) + } + } + result match { + case Left(cancellation) => + cancellation.join() + withStageCacheEntry(cacheKey)(useStage) + case Right(value) => value + } + } + private def newResultQueryStage( resultHandler: SparkPlan => Any, plan: SparkPlan): ResultQueryStageExec = { @@ -740,12 +989,19 @@ case class AdaptiveSparkPlanExec( private def reuseQueryStage( existing: ExchangeQueryStageExec, exchange: Exchange): ExchangeQueryStageExec = { + context.markSharedStageResult(existing.resultOption, this) val queryStage = existing.newReuseInstance(currentStageId, exchange.output) currentStageId += 1 setLogicalLinkForNewQueryStage(queryStage, exchange) + recordStageId(queryStage) queryStage } + /** Track local reuse aliases whose materialization may report the same eventual failure. */ + private def recordStageId(stage: ExchangeQueryStageExec): Unit = { + stageIdsByResult.getOrElseUpdate(stage.resultOption, mutable.LinkedHashSet.empty) += stage.id + } + /** * Set the logical node link of the `stage` as the corresponding logical node of the `plan` it * encloses. If an `plan` has been transformed from a `Repartition`, it should have `logicalLink` @@ -973,7 +1229,6 @@ object AdaptiveSparkPlanExec { * The execution context shared between the main query and all sub-queries. */ case class AdaptiveExecutionContext(session: SparkSession, qe: QueryExecution) { - /** * The subquery-reuse map shared across the entire query. */ @@ -986,6 +1241,84 @@ case class AdaptiveExecutionContext(session: SparkSession, qe: QueryExecution) { val stageCache: TrieMap[SparkPlan, ExchangeQueryStageExec] = new TrieMap[SparkPlan, ExchangeQueryStageExec]() + private val stageLifecycleLock = new Object + + /** + * Serialize exchange-cache lookup, reuse, and cancellation reservations across subqueries. + * Blocking cancellation and waits for a particular result happen outside this query-wide lock. + */ + private[adaptive] def withStageLifecycleLock[T](body: => T): T = { + stageLifecycleLock.synchronized(body) + } + + /** Stage-scoped reservations keep a cancelling result cached without blocking other stages. */ + private val stageCancellationReservations = + new ConcurrentHashMap[AtomicReference[Option[Any]], CompletableFuture[Unit]]() + + /** Claim a result before waiting for its shuffle monitor or cancelling its exchange. */ + private[adaptive] def reserveStageCancellation( + resultOption: AtomicReference[Option[Any]]): Option[CompletableFuture[Unit]] = { + val cancellation = new CompletableFuture[Unit]() + if (stageCancellationReservations.putIfAbsent(resultOption, cancellation) == null) { + Some(cancellation) + } else { + None + } + } + + /** Return the cancellation that a potential consumer of this particular result must await. */ + private[adaptive] def pendingStageCancellation( + resultOption: AtomicReference[Option[Any]]): Option[CompletableFuture[Unit]] = { + Option(stageCancellationReservations.get(resultOption)) + } + + /** Publish the final cache state before allowing consumers of this result to retry lookup. */ + private[adaptive] def finishStageCancellation( + resultOption: AtomicReference[Option[Any]], + cancellation: CompletableFuture[Unit]): Unit = { + stageCancellationReservations.remove(resultOption, cancellation) + cancellation.complete(()) + } + + /** + * The adaptive plan that first cached each result. Identity, rather than case-class equality, + * distinguishes independently planned subqueries with equivalent physical input plans. + */ + private val stageResultOwners = + new ConcurrentHashMap[AtomicReference[Option[Any]], AdaptiveSparkPlanExec]() + + /** Record the original owner atomically with exchange-cache insertion. */ + private[adaptive] def registerStageOwner( + resultOption: AtomicReference[Option[Any]], owner: AdaptiveSparkPlanExec): Unit = { + stageResultOwners.putIfAbsent(resultOption, owner) + } + + /** + * Results reused by another adaptive plan in this execution context. Cross-plan protection + * remains for the context's lifetime; aliases within one plan are not considered shared. + */ + private val sharedStageResults = + new ConcurrentHashMap[AtomicReference[Option[Any]], Boolean]() + + /** Conservatively protect a result whose owner is unknown or whose protection is test-forced. */ + private[adaptive] def markSharedStageResult(resultOption: AtomicReference[Option[Any]]): Unit = { + sharedStageResults.put(resultOption, true) + } + + /** Conservatively protect a result with unknown ownership or reuse by another adaptive plan. */ + private[adaptive] def markSharedStageResult( + resultOption: AtomicReference[Option[Any]], owner: AdaptiveSparkPlanExec): Unit = { + val originalOwner = stageResultOwners.get(resultOption) + if (originalOwner == null || (originalOwner ne owner)) { + markSharedStageResult(resultOption) + } + } + + /** Return whether this result is conservatively protected from obsolete-stage cancellation. */ + private[adaptive] def isSharedStageResult(resultOption: AtomicReference[Option[Any]]): Boolean = { + sharedStageResults.containsKey(resultOption) + } + val shuffleIds: ConcurrentHashMap[Int, Boolean] = new ConcurrentHashMap[Int, Boolean]() } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala index 62e00d1ea6eda..09fc97bb46fa5 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala @@ -71,7 +71,21 @@ case class LogicalQueryStage( physicalStats.getOrElse(logicalPlan.stats) } - override def maxRows: Option[Long] = stats.rowCount.map(_.min(Long.MaxValue).toLong) + override def maxRows: Option[Long] = { + // A query stage's `rowCount` is an exact, valid upper bound only when it comes from the + // runtime statistics of a materialized stage. Checking `isMaterialized` alone is not enough: + // `computeStats()` can still fall back to `logicalPlan.stats` (a cost estimate, with + // `isRuntime = false`) when the physical-stage lookup yields no statistics, and that estimate + // can under-count (e.g., returning 0). Treating such an estimate as a hard upper bound lets + // rules such as EliminateLimits wrongly drop a LIMIT and change the result cardinality + // (SPARK-57956). Only trust `rowCount` when the stats are materialized runtime stats; + // otherwise fall back to `logicalPlan.maxRows`, which is always a sound upper bound. + if (isMaterialized && stats.isRuntime) { + stats.rowCount.map(_.min(Long.MaxValue).toLong) + } else { + logicalPlan.maxRows + } + } override def isMaterialized: Boolean = physicalPlan.exists { case s: QueryStageExec => s.isMaterialized diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala index 6b721d5442e28..9374758d47311 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.adaptive -import org.apache.spark.sql.catalyst.expressions.{Alias, BindReferences, DynamicPruningExpression, Literal} +import org.apache.spark.sql.catalyst.expressions.{Alias, BindReferences, BroadcastValueProjection, DynamicPruningExpression, Expression, Literal} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} import org.apache.spark.sql.catalyst.plans.logical.Aggregate import org.apache.spark.sql.catalyst.rules.Rule @@ -35,6 +35,61 @@ case class PlanAdaptiveDynamicPruningFilters( override def conf: SQLConf = rootPlan.context.session.sessionState.conf + private def reusableBroadcast( + name: String, + indices: Seq[Int], + buildKeys: Seq[Expression], + valueExpression: Option[Expression], + adaptivePlan: AdaptiveSparkPlanExec): Option[BaseSubqueryExec] = { + if (!conf.exchangeReuseEnabled || buildKeys.isEmpty) { + return None + } + + val packedKeys = BindReferences.bindReferences( + HashJoin.rewriteKeyExpr(buildKeys), adaptivePlan.executedPlan.output) + val exchange = BroadcastExchangeExec( + HashedRelationBroadcastMode(packedKeys, isNullAware = false), + adaptivePlan.executedPlan) + val canReuseExchange = find(rootPlan) { + case join: BroadcastHashJoinExec if !join.isNullAwareAntiJoin => + val candidatePlan = join.buildSide match { + case BuildLeft => join.left + case BuildRight => join.right + } + candidatePlan.sameResult(exchange) + case _ => false + }.isDefined + + if (canReuseExchange) { + adaptivePlan.executedPlan.logicalLink.foreach(exchange.setLogicalLink) + val newAdaptivePlan = adaptivePlan.copy(inputPlan = exchange) + Some(valueExpression match { + case Some(value) => + ProjectedBroadcastValueSubqueryExec(name, value, newAdaptivePlan) + case None => + SubqueryBroadcastExec(name, indices, buildKeys, newAdaptivePlan) + }) + } else { + None + } + } + + private def projectedBroadcast( + name: String, + projection: BroadcastValueProjection, + context: AdaptiveExecutionContext): Option[BaseSubqueryExec] = { + QueryExecution.prepareExecutedPlan(projection.sourcePlan, context) match { + case adaptive: AdaptiveSparkPlanExec => + reusableBroadcast( + name, + Seq(0), + projection.sourceHashKeys, + Some(projection.valueExpression), + adaptive) + case _ => None + } + } + def apply(plan: SparkPlan): SparkPlan = { if (!conf.dynamicPartitionPruningEnabled) { return plan @@ -43,42 +98,34 @@ case class PlanAdaptiveDynamicPruningFilters( plan.transformAllExpressionsWithPruning( _.containsAllPatterns(DYNAMIC_PRUNING_EXPRESSION, IN_SUBQUERY_EXEC)) { case DynamicPruningExpression(InSubqueryExec( - value, SubqueryAdaptiveBroadcastExec(name, indices, onlyInBroadcast, buildPlan, buildKeys, - adaptivePlan: AdaptiveSparkPlanExec), exprId, _, _, _)) => - val packedKeys = BindReferences.bindReferences( - HashJoin.rewriteKeyExpr(buildKeys), adaptivePlan.executedPlan.output) - val mode = HashedRelationBroadcastMode(packedKeys) - // plan a broadcast exchange of the build side of the join - val exchange = BroadcastExchangeExec(mode, adaptivePlan.executedPlan) - - val canReuseExchange = conf.exchangeReuseEnabled && buildKeys.nonEmpty && - find(rootPlan) { - case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _, _) => - left.sameResult(exchange) - case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _, _) => - right.sameResult(exchange) - case _ => false - }.isDefined - - if (canReuseExchange) { - exchange.setLogicalLink(adaptivePlan.executedPlan.logicalLink.get) - val newAdaptivePlan = adaptivePlan.copy(inputPlan = exchange) - - val broadcastValues = SubqueryBroadcastExec( - name, indices, buildKeys, newAdaptivePlan) - DynamicPruningExpression(InSubqueryExec(value, broadcastValues, exprId)) - } else if (onlyInBroadcast) { - DynamicPruningExpression(Literal.TrueLiteral) - } else { - // we need to apply an aggregate on the buildPlan in order to be column pruned - val aliases = indices.map(idx => Alias(buildKeys(idx), buildKeys(idx).toString)()) - val aggregate = Aggregate(aliases, aliases, buildPlan) + value, subquery @ SubqueryAdaptiveBroadcastExec( + name, indices, onlyInBroadcast, buildPlan, buildKeys, + adaptivePlan: AdaptiveSparkPlanExec), exprId, _, _, _)) => + val directBroadcast = reusableBroadcast( + name, indices, buildKeys, None, adaptivePlan) + val reusedBroadcast = directBroadcast.orElse { + if (onlyInBroadcast) { + subquery.broadcastValueProjection + .flatMap(projectedBroadcast(name, _, adaptivePlan.context)) + } else { + None + } + } - val sparkPlan = QueryExecution.prepareExecutedPlan(aggregate, adaptivePlan.context) - assert(sparkPlan.isInstanceOf[AdaptiveSparkPlanExec]) - val newAdaptivePlan = sparkPlan.asInstanceOf[AdaptiveSparkPlanExec] - val values = SubqueryExec(name, newAdaptivePlan) - DynamicPruningExpression(InSubqueryExec(value, values, exprId)) + reusedBroadcast match { + case Some(broadcastValues) => + DynamicPruningExpression(InSubqueryExec(value, broadcastValues, exprId)) + case None if onlyInBroadcast => + DynamicPruningExpression(Literal.TrueLiteral) + case None => + val aliases = indices.map(idx => + Alias(buildKeys(idx), buildKeys(idx).toString)()) + val aggregate = Aggregate(aliases, aliases, buildPlan) + val sparkPlan = QueryExecution.prepareExecutedPlan(aggregate, adaptivePlan.context) + assert(sparkPlan.isInstanceOf[AdaptiveSparkPlanExec]) + val newAdaptivePlan = sparkPlan.asInstanceOf[AdaptiveSparkPlanExec] + val values = SubqueryExec(name, newAdaptivePlan) + DynamicPruningExpression(InSubqueryExec(value, values, exprId)) } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala index 5f2638655c37c..f97339cde51f3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveSubqueries.scala @@ -46,11 +46,16 @@ case class PlanAdaptiveSubqueries( } val subquery = SubqueryExec(s"subquery#${exprId.id}", subqueryMap(exprId.id)) InSubqueryExec(expr, subquery, exprId, isDynamicPruning = false) - case expressions.DynamicPruningSubquery(value, buildPlan, + case pruning @ expressions.DynamicPruningSubquery(value, buildPlan, buildKeys, broadcastKeyIndices, onlyInBroadcast, exprId, _) => val name = s"dynamicpruning#${exprId.id}" - val subquery = SubqueryAdaptiveBroadcastExec(name, broadcastKeyIndices, onlyInBroadcast, - buildPlan, buildKeys, subqueryMap(exprId.id)) + val subquery = SubqueryAdaptiveBroadcastExec( + name, + broadcastKeyIndices, + onlyInBroadcast, + buildPlan, + buildKeys, + subqueryMap(exprId.id))(pruning.usableBroadcastValueProjection) DynamicPruningExpression(InSubqueryExec(value, subquery, exprId)) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala index 7fda560863809..16c34b404de64 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/AggUtils.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.plans.logical.Aggregate import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.streaming.{ProjectAggregationBufferExec, StatefulStreamlineAggregateExec} import org.apache.spark.sql.execution.streaming.operators.stateful._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.util.Utils @@ -425,6 +426,95 @@ object AggUtils { finalAndCompleteAggregate :: Nil } + /** + * Plans a "streamlined" streaming aggregation using the following progression: + * (Here the term "streamline" represents the loop of "read-process-output" for each input.) + * + * - Initialize Aggregation Buffer (Passthrough aggregation) + * - Shuffle + * - Streamlined Stateful Aggregation + * - For each input, do the following + * - Read the previous value for grouping key in state store + * - Merge the input and previous value (if any) + * - Store the new value to the state store + * - Complete (output the current result of the aggregation) + * + * The concept is to aggregate only between input and the state store, enabling an output to + * be available just from processing a single input. In update mode, each input will produce + * an output, which won't have any blocking operation in real time mode, leading to the + * lowest output latency. + */ + def planStreamlineStreamingAggregation( + groupingExpressions: Seq[NamedExpression], + functionsWithoutDistinct: Seq[AggregateExpression], + resultExpressions: Seq[NamedExpression], + stateFormatVersion: Int, + child: SparkPlan): Seq[SparkPlan] = { + val groupingAttributes = groupingExpressions.map(_.toAttribute) + + val initAggBuffer: SparkPlan = { + val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Partial)) + val aggregateAttributes = aggregateExpressions.map(_.resultAttribute) + + ProjectAggregationBufferExec( + numShufflePartitions = None, + groupingExpressions = groupingExpressions, + aggregateExpressions = aggregateExpressions, + aggregateAttributes = aggregateAttributes, + resultExpressions = groupingAttributes ++ + aggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes), + isFinalAggregate = false, + child = child) + } + + val aggregate: SparkPlan = { + val aggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = PartialMerge)) + val aggregateAttributes = aggregateExpressions.map(_.resultAttribute) + + StatefulStreamlineAggregateExec( + requiredChildDistributionExpressions = + Some(groupingAttributes), + groupingExpressions = groupingAttributes, + aggregateExpressions = mayRemoveAggFilters(aggregateExpressions), + aggregateAttributes = aggregateAttributes, + initialInputBufferOffset = groupingAttributes.length, + resultExpressions = groupingAttributes ++ + aggregateExpressions.flatMap(_.aggregateFunction.inputAggBufferAttributes), + isFinalAggregate = false, + numShufflePartitions = None, + outputMode = None, + stateFormatVersion = stateFormatVersion, + child = initAggBuffer, + stateInfo = None, + eventTimeWatermarkForLateEvents = None, + eventTimeWatermarkForEviction = None) + } + + val finalAndCompleteAggregate: SparkPlan = { + val finalAggregateExpressions = functionsWithoutDistinct.map(_.copy(mode = Final)) + // The attributes of the final aggregation buffer, which is presented as input to the result + // projection: + val finalAggregateAttributes = finalAggregateExpressions.map(_.resultAttribute) + + ProjectAggregationBufferExec( + requiredChildDistributionExpressions = Some(groupingAttributes), + numShufflePartitions = None, + // The child here is the post-shuffle output of the stateful aggregate, whose grouping + // columns are already resolved attributes, so group by those rather than by the original + // expressions. This matches planStreamingAggregation's final stage and the stateful + // aggregate above, both of which group by the attributes. + groupingExpressions = groupingAttributes, + aggregateExpressions = finalAggregateExpressions, + aggregateAttributes = finalAggregateAttributes, + initialInputBufferOffset = groupingAttributes.length, + resultExpressions = resultExpressions, + isFinalAggregate = true, + child = aggregate) + } + + finalAndCompleteAggregate :: Nil + } + /** * Plans a streaming session aggregation using the following progression: * diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala index 62c4f896f2ee4..f3561fb91394c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExec.scala @@ -72,7 +72,16 @@ case class HashAggregateExec( "aggTime" -> SQLMetrics.createTimingMetric(sparkContext, "time in aggregation build"), "avgHashProbe" -> SQLMetrics.createAverageMetric(sparkContext, "avg hash probes per key"), - "numTasksFallBacked" -> SQLMetrics.createMetric(sparkContext, "number of sort fallback tasks")) + "numTasksFallBacked" -> SQLMetrics.createMetric(sparkContext, "number of sort fallback tasks") + ) ++ { + // Only the aggregates that can actually bypass report this, so the rest do not show a + // constant 0 in the SQL UI (see `UnionExec.metrics` for the same approach). + if (adaptivePartialAggEnabled) { + Map("numBypassingRows" -> SQLMetrics.createMetric(sparkContext, "number of bypassing rows")) + } else { + Map.empty[String, SQLMetric] + } + } // This is for testing. We force TungstenAggregationIterator to fall back to the unsafe row hash // map and/or the sort-based aggregation once it has processed a given number of input rows. @@ -94,6 +103,8 @@ case class HashAggregateExec( val avgHashProbe = longMetric("avgHashProbe") val aggTime = longMetric("aggTime") val numTasksFallBacked = longMetric("numTasksFallBacked") + // Registered only when the feature applies, and only read from the pass-through path. + val numBypassingRows = if (adaptivePartialAggEnabled) longMetric("numBypassingRows") else null child.execute().mapPartitionsWithIndex { (partIndex, iter) => @@ -121,7 +132,12 @@ case class HashAggregateExec( peakMemory, spillSize, avgHashProbe, - numTasksFallBacked) + numTasksFallBacked, + numBypassingRows, + aggTime, + adaptivePartialAggEnabled, + adaptiveMinRows, + adaptiveMinCompaction) if (!hasInput && groupingExpressions.isEmpty) { numOutputRows += 1 Iterator.single[UnsafeRow](aggregationIterator.outputForEmptyGroupingKeyWithoutInput()) @@ -141,6 +157,59 @@ case class HashAggregateExec( .map(_.asInstanceOf[DeclarativeAggregate]) private val bufferSchema = DataTypeUtils.fromAttributes(aggregateBufferAttributes) + /** + * Whether adaptive partial aggregation applies to this operator. When it does, the aggregation + * may bypass partial aggregation at runtime and pass the remaining input rows through as + * single-row partial buffers (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). It only + * applies to a pre-shuffle partial aggregation with grouping keys: + * - `Partial` and `PartialMerge` modes only: the downstream `Final` aggregation merges the + * passed-through single-row buffers. `Final`/`Complete` produce the result themselves and + * have no such downstream. A `PartialMerge` member is the non-distinct aggregate of the + * DISTINCT intermediate phase (`AggUtils.planAggregateWithOneDistinct`): its input row is + * already a partial buffer, so the pass-through applies the merge to an empty buffer, which + * leaves the incoming buffer unchanged, and the downstream `Final` re-merges it. A pure + * `PartialMerge` phase (the de-duplication on keys ++ distinct columns) must not bypass, or + * duplicate (key, distinct column) rows would over-count DISTINCT. The built-in planner + * never emits such a phase without a required distribution, so the + * `requiredChildDistributionExpressions` check below already keeps it out. The + * `exists(_.mode == Partial)` check is a defensive guard on top of it, and only for a + * de-duplication phase that carries non-distinct aggregates: one with none at all has an + * empty `aggregateExpressions`, is admitted by the `isEmpty` disjunct, and still relies on + * the distribution check alone. + * - grouping keys present: a global aggregation produces a single output row, so partial + * aggregation achieves the maximum reduction and must never be bypassed. + * - no required distribution: with no aggregate functions (a group-by-only aggregate) the + * mode check is vacuously true for both phases, so `requiredChildDistributionExpressions` + * tells them apart, the `Final` phase requiring a distribution and the pre-shuffle phase + * not. It is not a general pre-shuffle test, because `AggUtils.planAggregateWithOneDistinct` + * leaves it `None` on a post-shuffle aggregate. DISTINCT aggregate functions are allowed: + * the phase that de-duplicates on (keys ++ distinct columns) requires a distribution, so + * this check keeps it out, while the distinct partial phase that groups on the keys alone + * (carrying the non-distinct aggregates as `PartialMerge`) is eligible even though it sits + * after a shuffle: another `Exchange` and a `Final` follow it, so its passed-through + * buffers are still merged. + * - batch only: a streaming partial aggregate keeps state across batches, and it is built + * with all-`Partial` modes and no required distribution, so it would otherwise qualify. + * A batch `session_window` grouping likewise qualifies, but its partial aggregate feeds a + * `MergingSessionsExec` that merges overlapping sessions, so it is kept out for now. The + * static sibling `spark.sql.execution.bypassPartialAggregation` likewise stays away from + * streaming and `session_window` groupings. + */ + private val adaptivePartialAggEnabled: Boolean = { + conf.adaptivePartialAggregationEnabled && + groupingExpressions.nonEmpty && + !isStreaming && + !groupingExpressions.exists(_.metadata.contains(SessionWindow.marker)) && + aggregateExpressions.forall(a => a.mode == Partial || a.mode == PartialMerge) && + (aggregateExpressions.exists(_.mode == Partial) || aggregateExpressions.isEmpty) && + requiredChildDistributionExpressions.isEmpty + } + + // The number of rows between two compaction-ratio evaluations, and the ratio below which the + // partial aggregation is considered ineffective. Only read when the feature applies. + private val adaptiveMinRows: Long = conf.adaptivePartialAggregationMinRows + private val adaptiveMinCompaction: Double = conf.adaptivePartialAggregationMinCompaction + // The name for Fast HashMap private var fastHashMapTerm: String = _ private var isFastHashMapEnabled: Boolean = false @@ -154,6 +223,47 @@ case class HashAggregateExec( private var hashMapTerm: String = _ private var sorterTerm: String = _ + // Codegen state for adaptive partial aggregation. When the aggregation maps stop collapsing + // enough rows, the operator stops populating them and instead streams each remaining row through + // as a single-row partial buffer for the Final aggregate to merge. The compaction ratio is + // measured at the operator level: all processed rows against the keys held by both the fast and + // the regular map, so two-level-map routing does not change the decision. + private var adaptivePassThroughTerm: String = _ + private var processedRowsTerm: String = _ + private var adaptiveChildrenConsumedTerm: String = _ + // Whether the map output has already been emitted (and the maps freed). Once pass-through is + // active the maps are frozen; draining them -- which also frees them -- releases their memory + // before the rest of the input is streamed. The drain starts as soon as the first passed-through + // row queues a copy behind the maps, advancing one map row per queued row (see + // `handlePassThroughRow`); this flag lets the later output skip the maps. + private var adaptiveMapOutputDoneTerm: String = _ + // Whether the map iterators have been set up (`finishHashMap`). The map-output function may be + // re-entered when its loops return via `shouldStop()` to drain the buffer, and `finishAggregate` + // destructs the map, so the setup must run only once. + private var adaptiveMapSetupDoneTerm: String = _ + // The processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and the count is reset after a spill so the new in-memory map + // epoch is judged on its own rows. + private var adaptiveNextCheckRowTerm: String = _ + // The fast map's key count as of the last spill. The fast map never spills and is never cleared, + // so its keys outlive the epoch the processed-row count is reset for; subtracting this baseline + // keeps both sides of the ratio on the same epoch. Once the fast map fills it stops accepting + // keys, so the difference settles at zero and the ratio is the regular map's alone. + private var adaptiveFastKeysAtSpillTerm: String = _ + // The name of the generated output function, promoted to a field so `doConsumeWithKeys` can emit + // pass-through rows directly from within the build loop. + private var outputFunc: String = _ + // The name of the generated function that drains the frozen maps and, once they are fully + // drained, flushes the queue of passed-through rows held behind them. Promoted to a field so + // `doConsumeWithKeys` can advance the maps one row per passed-through row (see + // `handlePassThroughRow`). + private var adaptiveOutputMapAndFlushFuncName: String = _ + // The queue of passed-through rows, held behind the frozen maps so the maps are emitted first. + // Each queued row advances the map output by one row, so the queue stays bounded by how many + // passed-through rows a child emits before it honours `shouldStop()` (see + // `handlePassThroughRow`). + private var adaptivePendingRowsTerm: String = _ + /** * This is called by generated Java class, should be public. */ @@ -436,6 +546,25 @@ case class HashAggregateExec( protected override def doProduceWithKeys(ctx: CodegenContext): String = { val initAgg = ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "initAgg") + if (adaptivePartialAggEnabled) { + adaptivePassThroughTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptivePassThrough") + processedRowsTerm = ctx.addMutableState(CodeGenerator.JAVA_LONG, "processedRows") + adaptiveNextCheckRowTerm = + ctx.addMutableState(CodeGenerator.JAVA_LONG, "adaptiveNextCheckRow", + v => s"$v = ${adaptiveMinRows}L;") + adaptiveChildrenConsumedTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveChildrenConsumed") + adaptiveMapOutputDoneTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapOutputDone") + adaptiveMapSetupDoneTerm = + ctx.addMutableState(CodeGenerator.JAVA_BOOLEAN, "adaptiveMapSetupDone") + adaptiveFastKeysAtSpillTerm = + ctx.addMutableState(CodeGenerator.JAVA_INT, "adaptiveFastKeysAtSpill") + adaptivePendingRowsTerm = ctx.addMutableState( + "java.util.LinkedList<UnsafeRow[]>", "adaptivePendingRows", + v => s"$v = new java.util.LinkedList<UnsafeRow[]>();", forceInline = true) + } if (conf.enableTwoLevelAggMap) { enableTwoLevelHashMap() } else if (conf.enableVectorizedHashMap) { @@ -535,19 +664,31 @@ case class HashAggregateExec( // `addNewFunction` spills this helper into a nested class (as can happen // once the outer class passes the code-size threshold), the bare field // reference fails with `IllegalAccessError`. - val doAggFuncName = ctx.addNewFunction(doAgg, - s""" - |private void $doAgg(int partitionIndex) throws java.io.IOException { - | ${child.asInstanceOf[CodegenSupport].produce(ctx, this)} - | $finishHashMap - |} - """.stripMargin) - // generate code for output + // Generate code for output. With adaptive partial aggregation enabled this must happen before + // the `doAgg` helper below, because `doConsumeWithKeys` (invoked from the child's produce + // inside `doAgg`) emits pass-through rows by calling this output function directly. Otherwise + // the output function is generated after `doAgg`, so the early `consume(...)` inside it cannot + // change the codegen layout for plans the feature never touches. val keyTerm = ctx.freshName("aggKey") val bufferTerm = ctx.freshName("aggBuffer") - val outputFunc = generateResultFunction(ctx) + if (adaptivePartialAggEnabled) { + outputFunc = generateResultFunction(ctx) + } + // After the child input is consumed, finish the build: with adaptive partial aggregation mark + // that the child is fully consumed (to support re-entry; the map iterators are set up inside + // the map-output function), otherwise set up the map iterators for the output below. + // A child may leave its produce loop early rather than exhausting its input: `UnionExec` runs + // each child inside its own helper, so a streamed row that fills the output buffer returns + // only as far as here. Reaching the end of `produce` therefore does not by itself mean the + // input is consumed -- pending output says the child parked instead, and the build resumes on + // re-entry. Without this the rest of that partition is silently dropped. + val postChildProduce = if (adaptivePartialAggEnabled) { + s"if (!shouldStop()) { $adaptiveChildrenConsumedTerm = true; }" + } else { + finishHashMap + } val limitNotReachedCondition = limitNotReachedCond def outputFromFastHashMap: String = { @@ -615,8 +756,109 @@ case class HashAggregateExec( """.stripMargin } + // With adaptive partial aggregation the maps are frozen once pass-through is active, so their + // output (which also frees them) can happen as soon as pass-through fires, releasing the memory + // before the remaining input is streamed. The output loops are wrapped in a function so the + // same code runs either early (once pass-through freezes the maps) or at the end of the build. + // The done flag is set inside, after the loops, so a mid-output drain (the loops return via + // `shouldStop()`) leaves it unset and the caller resumes the map iterator on re-entry; once it + // is set the maps have been fully output and freed and will not be touched again. The iterator + // setup (`finishHashMap`, which destructs the map) is guarded to run only once. The queue of + // passed-through rows held behind the maps is flushed here, right after the done flag is set: + // the maps must precede the held rows so the downstream Final merges a group's map buffer + // before its pass-through buffers (a group can straddle the freeze, since the maps only freeze + // once pass-through fires), and the copies were made when the row was queued, so emitting them + // here cannot corrupt the build's reusable rows. + adaptiveOutputMapAndFlushFuncName = if (adaptivePartialAggEnabled) { + val name = ctx.freshName("outputMapAndFlush") + val pair = ctx.freshName("pendingRow") + ctx.addNewFunction(name, + s""" + |private void $name() throws java.io.IOException { + | if (!$adaptiveMapSetupDoneTerm) { + | $finishHashMap + | $adaptiveMapSetupDoneTerm = true; + | } + | $outputFromFastHashMap + | $outputFromRegularHashMap + | $adaptiveMapOutputDoneTerm = true; + | while (!$adaptivePendingRowsTerm.isEmpty()) { + | UnsafeRow[] $pair = (UnsafeRow[]) $adaptivePendingRowsTerm.poll(); + | $outputFunc($pair[0], $pair[1]); + | } + |} + """.stripMargin) + } else { + "" + } + + val doAggFuncName = ctx.addNewFunction(doAgg, + s""" + |private void $doAgg(int partitionIndex) throws java.io.IOException { + | ${child.asInstanceOf[CodegenSupport].produce(ctx, this)} + | $postChildProduce + |} + """.stripMargin) + // For a non-adaptive plan, generate the output function only after `doAgg` so its + // `consume(...)` cannot reorder the codegen layout (see above). It must still land before + // `adaptiveFinalOutput` reads it below. + if (!adaptivePartialAggEnabled) { + outputFunc = generateResultFunction(ctx) + } + val aggTime = metricTerm(ctx, "aggTime") val beforeAgg = ctx.freshName("beforeAgg") + // Split by an exchange, `doAgg` may start appending pass-through rows to the output buffer + // mid-build. In that case `shouldStop()` becomes true and we must return so the buffered rows + // are drained; on re-entry the frozen maps are output first (`adaptiveOutputMapAndFlush`) and + // then the build resumes (guarded by `childrenConsumed`) until the child input is exhausted. + // Fused with the Final, nothing is buffered, so `shouldStop()` never fires and the maps are + // drained inside the build (see `handlePassThroughRow`). + val adaptiveStopCheck = if (adaptivePartialAggEnabled) { + "if (shouldStop()) return;" + } else { + "" + } + // Once pass-through is active the maps are frozen, so output them (releasing their memory) + // ahead of the passed-through rows. The output starts the moment the first passed-through row + // queues behind the maps (`handlePassThroughRow` advances the maps by one row per queued row), + // and continues on re-entry here; `adaptiveFinalOutput` covers the case where pass-through + // never fired during the build, when the full maps are the result. The output loops return via + // `shouldStop()` when the buffer fills, so the done flag is set inside the output function and + // re-entry resumes the map iterator. Only once the maps are fully drained are the queued + // pass-through rows flushed, so every map buffer precedes the rows that collided with it. + val adaptiveOutputMapAndFlush = if (adaptivePartialAggEnabled) { + s""" + |if (!$adaptiveMapOutputDoneTerm) { + | $adaptiveOutputMapAndFlushFuncName(); + | if (shouldStop()) return; + |} + """.stripMargin + } else { + "" + } + val adaptiveResumeBuild = if (adaptivePartialAggEnabled) { + val beforeResumedAgg = ctx.freshName("beforeResumedAgg") + s""" + |if (!$adaptiveChildrenConsumedTerm) { + | $adaptiveOutputMapAndFlush + | long $beforeResumedAgg = System.nanoTime(); + | $doAggFuncName(partitionIndex); + | $aggTime.add((System.nanoTime() - $beforeResumedAgg) / $NANOS_PER_MILLIS); + | if (shouldStop()) return; + |} + """.stripMargin + } else { + "" + } + val adaptiveFinalOutput = if (adaptivePartialAggEnabled) { + adaptiveOutputMapAndFlush + } else { + s""" + |$outputFromFastHashMap + |$outputFromRegularHashMap + """.stripMargin + } s""" |if (!$initAgg) { | $initAgg = true; @@ -626,13 +868,39 @@ case class HashAggregateExec( | long $beforeAgg = System.nanoTime(); | $doAggFuncName(partitionIndex); | $aggTime.add((System.nanoTime() - $beforeAgg) / $NANOS_PER_MILLIS); + | $adaptiveStopCheck |} - |// output the result - |$outputFromFastHashMap - |$outputFromRegularHashMap + |$adaptiveResumeBuild + |$adaptiveFinalOutput """.stripMargin } + // Blocking operators normally suppress the child's `shouldStop()` check because they buffer all + // output. With adaptive partial aggregation, pass-through rows are appended to the output buffer + // while consuming child input, so the stop check is re-enabled to let the child yield between + // rows. + // + // This bounds the buffer only as far as the child honours it. Each passed-through row queues a + // copy and advances the frozen-map output by one row, and that one appended map row makes + // `shouldStop()` true, so a child that checks between rows never queues more than one row. A + // one-to-many child that does not check `shouldStop()` inside its fan-out (`GenerateExec` emits + // `for (index ...) { consume }` and `while (iterator.hasNext()) { consume }` with no check) + // appends every row produced from one input row before it can yield: the fan-out batch lands in + // `BufferedRowIterator.currentRows` (which is not spillable), and the map advances one row per + // queued row, so the buffer grows to roughly the batch width rather than `minRows`. + override def needStopCheck: Boolean = adaptivePartialAggEnabled + + // Blocking operators normally do not copy their result because every output row is drained (via + // `shouldStop()`) before the next one is produced. Adaptive pass-through breaks that assumption: + // `outputFunc` writes every output row into the same reusable result `UnsafeRow`, and under a + // fan-out child several such rows are appended without an intervening drain - the frozen-map + // rows, emitted one per queued row inside the child's fan-out loop, and the held pass-through + // batch, flushed from `outputMapAndFlush` after the maps drain. Without a copy they would all + // alias the single result row. Such children report `needCopyResult` themselves, so propagate + // their requirement rather than copying for every adaptive aggregate. + override def needCopyResult: Boolean = adaptivePartialAggEnabled && + child.asInstanceOf[CodegenSupport].needCopyResult + protected override def doConsumeWithKeys(ctx: CodegenContext, input: Seq[ExprCode]): String = { // create grouping key val unsafeRowKeyCode = GenerateUnsafeProjection.createCode( @@ -644,6 +912,18 @@ case class HashAggregateExec( val unsafeRowBuffer = ctx.freshName("unsafeRowAggBuffer") val fastRowBuffer = ctx.freshName("fastAggBuffer") + // For adaptive partial aggregation pass-through, each bypassed row is emitted as a single-row + // partial buffer: start from the initial aggregation buffer, apply the update expressions once, + // and output `key ++ buffer` for the Final aggregate to merge. This projects the initial + // buffer. + val emptyAggBufferCode = if (adaptivePartialAggEnabled) { + GenerateUnsafeProjection.createCode(ctx, declFunctions.flatMap(f => f.initialValues)) + } else { + null + } + // Per-row local flag marking that the current row is being streamed through (held by no map). + val adaptiveRowBypassedTerm = ctx.freshName("adaptiveRowBypassed") + // To individually generate code for each aggregate function, an element in `updateExprs` holds // all the expressions for the buffer of an aggregation function. val updateExprs = aggregateExpressions.map { e => @@ -663,46 +943,126 @@ case class HashAggregateExec( case _ => ("true", "", "") } - val findOrInsertRegularHashMap: String = - s""" - |// generate grouping key - |${unsafeRowKeyCode.code} - |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode(); - |if ($checkFallbackForBytesToBytesMap) { - | // try to get the buffer from hash map - | $unsafeRowBuffer = - | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys, $unsafeRowKeyHash); - |} - |// Can't allocate buffer from the hash map. Spill the map and fallback to sort-based - |// aggregation after processing all input rows. - |if ($unsafeRowBuffer == null) { - | if ($sorterTerm == null) { - | $sorterTerm = $hashMapTerm.destructAndCreateExternalSorter(); - | } else { - | $sorterTerm.merge($hashMapTerm.destructAndCreateExternalSorter()); - | } - | $resetCounter - | // the hash map had be spilled, it should have enough memory now, - | // try to allocate buffer again. - | $unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( - | $unsafeRowKeys, $unsafeRowKeyHash); - | if ($unsafeRowBuffer == null) { - | // failed to allocate the first page - | throw QueryExecutionErrors.aggregateOutOfMemoryError(); - | } - |} - """.stripMargin + // The compaction ratio is measured at the operator level: all processed rows against the keys + // held by both maps, so two-level-map routing does not change the decision. The same predicate + // decides both check points -- periodically every `minRows` rows, and right before the map + // would spill (in which case the spill is skipped entirely). `minRows = 0` disables the + // periodic check: the row count is only ever compared after being incremented past 0, so it + // never matches and only the spill check remains. + val adaptiveIneffective = if (adaptivePartialAggEnabled) { + val totalKeys = if (isFastHashMapEnabled) { + s"($fastHashMapTerm.getNumKeys() - $adaptiveFastKeysAtSpillTerm + " + + s"$hashMapTerm.getNumKeys())" + } else { + s"$hashMapTerm.getNumKeys()" + } + s"$processedRowsTerm < (double) $totalKeys * ${adaptiveMinCompaction}D" + } else { + "" + } + + val findOrInsertRegularHashMap: String = { + // Assumes the grouping key projection (`unsafeRowKeyCode.code`) has already run for this row, + // so `unsafeRowKeyCode.value` holds the current key. The projection is emitted exactly once + // per regular-map row (see below); emitting it in more than one runtime branch is unsafe + // because the projection's subexpression/writer state assigned in one branch would be read + // stale from another (e.g. the adaptive pass-through path would reuse the last probed key). + val probeRegularMap = + s""" + |int $unsafeRowKeyHash = ${unsafeRowKeyCode.value}.hashCode(); + |if ($checkFallbackForBytesToBytesMap) { + | // try to get the buffer from hash map + | $unsafeRowBuffer = + | $hashMapTerm.getAggregationBufferFromUnsafeRow($unsafeRowKeys, $unsafeRowKeyHash); + |} + """.stripMargin + + val spillMap = + s""" + |$sorterTerm = org.apache.spark.sql.execution.aggregate.HashAggregateExec + | .spillHashMapToSorter($hashMapTerm, $sorterTerm); + |$resetCounter + |// the hash map had been spilled, so it should have enough memory now, + |// try to allocate buffer again. + |$unsafeRowBuffer = $hashMapTerm.getAggregationBufferFromUnsafeRow( + | $unsafeRowKeys, $unsafeRowKeyHash); + |if ($unsafeRowBuffer == null) { + | // failed to allocate the first page + | throw QueryExecutionErrors.aggregateOutOfMemoryError(); + |} + """.stripMargin + + if (adaptivePartialAggEnabled) { + // Spilling starts a new in-memory epoch, so the counters restart and the ratio of that + // epoch alone decides the remaining rows. The fast map neither spills nor clears, so its + // keys outlive the epoch; snapshotting them here keeps both sides of the ratio on the + // same rows. Once the fast map fills it stops accepting keys and the difference settles + // at zero, leaving the regular map's keys alone. + val snapshotFastKeys = if (isFastHashMapEnabled) { + s"$adaptiveFastKeysAtSpillTerm = $fastHashMapTerm.getNumKeys();" + } else { + "" + } + val spillAndRestartEpoch = + s""" + |$spillMap + |$processedRowsTerm = 0L; + |$adaptiveNextCheckRowTerm = ${adaptiveMinRows}L; + |$snapshotFastKeys + """.stripMargin + s""" + |// generate grouping key + |${unsafeRowKeyCode.code} + |if (!$adaptivePassThroughTerm) { + | $probeRegularMap + | if ($unsafeRowBuffer == null) { + | if ($processedRowsTerm > 0 && $adaptiveIneffective) { + | $adaptivePassThroughTerm = true; + | } else { + | $spillAndRestartEpoch + | } + | } + |} + """.stripMargin + } else { + s""" + |// generate grouping key + |${unsafeRowKeyCode.code} + |$probeRegularMap + |// Can't allocate buffer from the hash map. Spill the map and fallback to sort-based + |// aggregation after processing all input rows. + |if ($unsafeRowBuffer == null) { + | $spillMap + |} + """.stripMargin + } + } val findOrInsertHashMap: String = { - if (isFastHashMapEnabled) { + val findCode = if (isFastHashMapEnabled) { // If fast hash map is on, we first generate code to probe and update the fast hash map. // If the probe is successful the corresponding fast row buffer will hold the mutable row. + // Once adaptive pass-through is active, skip the fast map entirely so the row is streamed + // through instead of being inserted anywhere. + val fastMapProbe = + s""" + |${fastRowKeys.map(_.code).mkString("\n")} + |if (${fastRowKeys.map("!" + _.isNull).mkString(" && ")}) { + | $fastRowBuffer = $fastHashMapTerm.findOrInsert( + | ${fastRowKeys.map(_.value).mkString(", ")}); + |} + """.stripMargin + val guardedFastMapProbe = if (adaptivePartialAggEnabled) { + s""" + |if (!$adaptivePassThroughTerm) { + | $fastMapProbe + |} + """.stripMargin + } else { + fastMapProbe + } s""" - |${fastRowKeys.map(_.code).mkString("\n")} - |if (${fastRowKeys.map("!" + _.isNull).mkString(" && ")}) { - | $fastRowBuffer = $fastHashMapTerm.findOrInsert( - | ${fastRowKeys.map(_.value).mkString(", ")}); - |} + |$guardedFastMapProbe |// Cannot find the key in fast hash map, try regular hash map. |if ($fastRowBuffer == null) { | $findOrInsertRegularHashMap @@ -711,6 +1071,54 @@ case class HashAggregateExec( } else { findOrInsertRegularHashMap } + + // Every row is either accepted by an aggregation map or streamed through -- the fast map + // serves a row without it ever reaching the regular map, so both buffers are consulted to + // tell the two apart. + // + // An accepted row counts toward the compaction ratio, so the numerator matches the + // operator-level denominator. Counting inside the regular-map branch alone would drop the + // rows the fast map absorbed from the ratio and bypass an aggregation that is in fact + // reducing. A row no map holds is streamed once pass-through is active: both probes are + // skipped (guarded above), so neither buffer is set, and `rowBypassed` marks exactly those + // rows. The row that fails to insert at the spill boundary lands here too, while the row + // that merely flipped pass-through at the check point is already aggregated in the map that + // took it and must not be re-emitted. + val countOrPassThroughRow = if (adaptivePartialAggEnabled) { + val heldByAMap = if (isFastHashMapEnabled) { + s"($fastRowBuffer != null || $unsafeRowBuffer != null)" + } else { + s"($unsafeRowBuffer != null)" + } + // The grouping key was already projected in `findOrInsertRegularHashMap` + // (`unsafeRowKeyCode.code`), so `unsafeRowKeyCode.value` holds this row's key. Only build + // the single-row partial buffer here. + s""" + |if ($heldByAMap) { + | if (!$adaptivePassThroughTerm) { + | $processedRowsTerm += 1; + | if ($processedRowsTerm == $adaptiveNextCheckRowTerm) { + | if ($adaptiveIneffective) { + | $adaptivePassThroughTerm = true; + | } else { + | $adaptiveNextCheckRowTerm += ${adaptiveMinRows}L; + | } + | } + | } + |} else if ($adaptivePassThroughTerm) { + | $adaptiveRowBypassedTerm = true; + | ${emptyAggBufferCode.code} + | $unsafeRowBuffer = ${emptyAggBufferCode.value}; + |} + """.stripMargin + } else { + "" + } + + s""" + |$findCode + |$countOrPassThroughRow + """.stripMargin } val inputAttrs = aggregateBufferAttributes ++ inputAttributes @@ -845,29 +1253,75 @@ case class HashAggregateExec( } } - val declareRowBuffer: String = if (isFastHashMapEnabled) { - val fastRowType = if (isVectorizedHashMapEnabled) { - classOf[MutableColumnarRow].getName + val declareRowBuffer: String = { + val declareBuffers = if (isFastHashMapEnabled) { + val fastRowType = if (isVectorizedHashMapEnabled) { + classOf[MutableColumnarRow].getName + } else { + "UnsafeRow" + } + s""" + |UnsafeRow $unsafeRowBuffer = null; + |$fastRowType $fastRowBuffer = null; + """.stripMargin + } else { + s"UnsafeRow $unsafeRowBuffer = null;" + } + val declareBypassed = if (adaptivePartialAggEnabled) { + s"boolean $adaptiveRowBypassedTerm = false;" } else { - "UnsafeRow" + "" } s""" - |UnsafeRow $unsafeRowBuffer = null; - |$fastRowType $fastRowBuffer = null; + |$declareBuffers + |$declareBypassed """.stripMargin - } else { - s"UnsafeRow $unsafeRowBuffer = null;" } // We try to do hash map based in-memory aggregation first. If there is not enough memory (the // hash map will return null for new key), we spill the hash map to disk to free memory, then // continue to do in-memory aggregation and spilling until all the rows had been processed. // Finally, sort the spilled aggregate buffers by key, and merge them together for same key. + // + // With adaptive partial aggregation, once pass-through is active `updateRowInHashMap` fills the + // single-row buffer built above; we then emit `key ++ buffer` straight to the parent so the row + // skips both the fast map and the regular map. + // + // The maps were frozen when pass-through fired, so a group can straddle the freeze and hold + // both a map buffer and pass-through buffers. The downstream Final merges in emit order, so + // every map buffer must precede the pass-through buffers of the same group. We hold each + // passed-through row in a queue behind the maps: the row is queued as a copy, and queuing it + // advances the map output by one row. That one appended map row is what makes `shouldStop()` + // true, so a child that honours the check yields at the end of its batch and never queues more + // than one row; a one-to-many child that cannot yield mid-fan-out queues its whole batch at + // once, bounding the queue by the batch width. The map drains one row per queued row (the same + // cadence the 1:1 shape already uses), and once the maps are fully drained the queue is + // flushed, so the maps always precede the held rows. Fused with the Final, nothing is + // buffered, `shouldStop()` never fires, and the single output call drains the whole map before + // the held row follows. + val handlePassThroughRow = if (adaptivePartialAggEnabled) { + val numBypassingRows = metricTerm(ctx, "numBypassingRows") + s""" + |if ($adaptiveRowBypassedTerm) { + | $numBypassingRows.add(1); + | if ($adaptiveMapOutputDoneTerm) { + | $outputFunc(${unsafeRowKeyCode.value}, $unsafeRowBuffer); + | } else { + | $adaptivePendingRowsTerm.add(new UnsafeRow[] { + | ${unsafeRowKeyCode.value}.copy(), $unsafeRowBuffer.copy() }); + | $adaptiveOutputMapAndFlushFuncName(); + | } + |} + """.stripMargin + } else { + "" + } s""" |$declareRowBuffer |$findOrInsertHashMap |$incCounter |$updateRowInHashMap + |$handlePassThroughRow """.stripMargin } @@ -897,3 +1351,26 @@ case class HashAggregateExec( override protected def withNewChildInternal(newChild: SparkPlan): HashAggregateExec = copy(child = newChild) } + +object HashAggregateExec { + /** + * Spills the in-memory hash map to disk and returns the sorter holding the spilled data. Called + * by the generated code of [[HashAggregateExec]] (through the static forwarder) when a buffer + * cannot be allocated from the hash map: the first spill destructs the map into a new sorter, and + * later spills merge into the existing one. Returning the sorter lets the generated code keep it + * in a single mutable field. + * + * Extracting this type-independent spill machinery into a shared helper keeps it compiled once + * per JVM instead of being re-emitted into every HashAggregateExec stage's generated code. + */ + def spillHashMapToSorter( + hashMap: UnsafeFixedWidthAggregationMap, + sorter: UnsafeKVExternalSorter): UnsafeKVExternalSorter = { + if (sorter == null) { + hashMap.destructAndCreateExternalSorter() + } else { + sorter.merge(hashMap.destructAndCreateExternalSorter()) + sorter + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala index 00d18a2f79a81..9521a68ae4b2e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TungstenAggregationIterator.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.execution.aggregate +import java.util.concurrent.TimeUnit.NANOSECONDS + import org.apache.spark.{SparkException, TaskContext} import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow @@ -57,9 +59,10 @@ import org.apache.spark.util.ArrayImplicits._ * - Part 3: Methods and fields used by hash-based aggregation. * - Part 4: Methods and fields used when we switch to sort-based aggregation. * - Part 5: Methods and fields used by sort-based aggregation. - * - Part 6: Loads input and process input rows. - * - Part 7: Public methods of this iterator. - * - Part 8: A utility function used to generate a result when there is no + * - Part 6: Methods and fields used by adaptive partial aggregation pass-through. + * - Part 7: Loads input and processes input rows. + * - Part 8: Public methods of this iterator. + * - Part 9: A utility function used to generate a result when there is no * input and there is no grouping expression. * * @param partIndex @@ -95,7 +98,12 @@ class TungstenAggregationIterator( peakMemory: SQLMetric, spillSize: SQLMetric, avgHashProbe: SQLMetric, - numTasksFallBacked: SQLMetric) + numTasksFallBacked: SQLMetric, + numBypassingRows: SQLMetric, + aggTime: SQLMetric, + adaptivePartialAggEnabled: Boolean, + adaptiveMinRows: Long, + adaptiveMinCompaction: Double) extends AggregationIterator( partIndex, groupingExpressions, @@ -121,8 +129,9 @@ class TungstenAggregationIterator( /////////////////////////////////////////////////////////////////////////// // Creates a new aggregation buffer and initializes buffer values. - // This function should be only called at most two times (when we create the hash map, - // and when we create the re-used buffer for sort-based aggregation). + // This function should be called at most three times: when we create the hash map, when we + // create the re-used buffer for sort-based aggregation, and when we create the buffer for + // pass-through rows. private def createNewAggregationBuffer(): UnsafeRow = { val bufferSchema = aggregateFunctions.flatMap(_.aggBufferAttributes) val buffer: UnsafeRow = UnsafeProjection.create(bufferSchema.map(_.dataType)) @@ -179,6 +188,21 @@ class TungstenAggregationIterator( // hashMap. If there is not enough memory, it will multiple hash-maps, spilling // after each becomes full then using sort to merge these spills, finally do sort // based aggregation. + // + // When adaptive partial aggregation is enabled (`adaptivePartialAggEnabled`), the processing may + // stop early and switch to pass-through mode: the remaining input rows are not added to the map + // but are instead emitted as single-row partial buffers by the output stage (see + // `nextPassThroughOutput`). One predicate decides both check points -- the aggregation is + // ineffective when `processedRows < distinctKeys * minCompaction`, i.e. it does not collapse + // `minCompaction` rows into one key: + // - periodically, every `minRows` processed rows; + // - when the map is about to spill, in which case the spill is skipped entirely and the row + // that could not be inserted becomes the first pass-through row. + // A spill starts a new in-memory map epoch: the row counters restart so that epoch is judged on + // its own rows, which lets an input whose cardinality only turns unfavorable later still be + // passed through. Pass-through may therefore coexist with earlier spills. The output order + // matches the generated path: the frozen map (or sort-based) output comes first, and the + // pass-through rows stream afterwards. private def processInputs(fallbackStartsAt: (Int, Int)): Unit = { if (groupingExpressions.isEmpty) { // If there is no grouping expressions, we can just reuse the same buffer over and over again. @@ -191,7 +215,19 @@ class TungstenAggregationIterator( } } else { var i = 0 - while (inputIter.hasNext) { + var processedRows = 0L + val minRows = adaptiveMinRows + // The processed-row count at which the compaction ratio is evaluated next. It advances by + // `minRows` after every check, and restarts after a spill so the new in-memory map epoch is + // judged on its own rows. `minRows = 0` disables the periodic check: the count is only ever + // compared after being incremented past 0, so it never matches and only the spill check + // below remains. + var nextCheckRow = minRows + // The partial aggregation is ineffective when it does not collapse `minCompaction` rows into + // one key. There is no fast map on this path, so the map's keys are all the operator holds. + def ineffective(): Boolean = + processedRows < hashMap.getNumKeys().toDouble * adaptiveMinCompaction + while (inputIter.hasNext && !passThrough) { val newInput = inputIter.next() val groupingKey = groupingProjection.apply(newInput) var buffer: UnsafeRow = null @@ -199,21 +235,47 @@ class TungstenAggregationIterator( buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) } if (buffer == null) { - val sorter = hashMap.destructAndCreateExternalSorter() - if (externalSorter == null) { - externalSorter = sorter + // The map is full and would normally spill. Adaptive partial aggregation may instead + // bypass: keep the in-memory map as-is, pass this row and all remaining rows through, + // and skip the spill entirely. + if (adaptivePartialAggEnabled && processedRows > 0 && ineffective()) { + passThrough = true + // `newInput` could not be inserted; stash a copy as the first pass-through row so it + // is not lost when we drain the rest of `inputIter`. + pendingPassThroughRow = newInput.copy() } else { - externalSorter.merge(sorter) + val sorter = hashMap.destructAndCreateExternalSorter() + if (externalSorter == null) { + externalSorter = sorter + } else { + externalSorter.merge(sorter) + } + i = 0 + // The map starts a new in-memory epoch, so the counters restart and the ratio of that + // epoch alone decides whether the remaining rows are passed through. + processedRows = 0L + nextCheckRow = minRows + buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) + if (buffer == null) { + // failed to allocate the first page + throw QueryExecutionErrors.aggregateOutOfMemoryError() + } } - i = 0 - buffer = hashMap.getAggregationBufferFromUnsafeRow(groupingKey) - if (buffer == null) { - // failed to allocate the first page - throw QueryExecutionErrors.aggregateOutOfMemoryError() + } + if (!passThrough) { + processRow(buffer, newInput) + i += 1 + processedRows += 1 + // Periodic check: if the aggregation is not collapsing enough rows, bypass it for the + // rest of the input. + if (adaptivePartialAggEnabled && processedRows == nextCheckRow) { + if (ineffective()) { + passThrough = true + } else { + nextCheckRow += minRows + } } } - processRow(buffer, newInput) - i += 1 } if (externalSorter != null) { @@ -355,7 +417,55 @@ class TungstenAggregationIterator( } /////////////////////////////////////////////////////////////////////////// - // Part 6: Loads input rows and setup aggregationBufferMapIterator if we + // Part 6: Methods and fields used by adaptive partial aggregation pass-through. + /////////////////////////////////////////////////////////////////////////// + + // Indicates that partial aggregation has been bypassed and the remaining input rows should be + // passed through as single-row partial buffers. Set in `processInputs` by either check point. + // It may coexist with earlier spills. The output order matches the generated path: the frozen + // map (or sort-based) output comes first, and the pass-through rows stream afterwards. + private[this] var passThrough: Boolean = false + + // The row that could not be inserted at the spill check. It is stashed here (as + // a copy) so it becomes the first pass-through row rather than being lost. + private[this] var pendingPassThroughRow: InternalRow = null + + // A reused aggregation buffer for building single-row partial buffers during pass-through. It is + // re-initialized from `initialAggregationBuffer` for every passed-through row. + private[this] lazy val passThroughAggregationBuffer: UnsafeRow = createNewAggregationBuffer() + + // Whether there are remaining pass-through rows to emit. + private def passThroughHasNext: Boolean = + passThrough && (pendingPassThroughRow != null || inputIter.hasNext) + + // `aggTime` only covers the build, which runs to completion in the constructor before + // pass-through is known; the interpreted path consumes the rest of the input lazily through + // `next()`. The drain is therefore timed there, bracketed once per drain (not per row) so the + // timing call does not distort the per-row work, mirroring the generated path which times the + // drain around each resume of the build. + private var passThroughDrainStart: Long = -1L + + // Emits the next input row as a single-row partial aggregation buffer, i.e. a group of size one. + // The output (grouping key ++ buffer) is a valid partial buffer that the downstream Final + // aggregation merges, so the result is identical to running partial aggregation on this row. + private def nextPassThroughOutput(): UnsafeRow = { + val row = if (pendingPassThroughRow != null) { + val stashed = pendingPassThroughRow + pendingPassThroughRow = null + stashed + } else { + inputIter.next() + } + val groupingKey = groupingProjection.apply(row) + // Reset the buffer to initial values, then update it with this single row. + passThroughAggregationBuffer.copyFrom(initialAggregationBuffer) + processRow(passThroughAggregationBuffer, row) + numBypassingRows += 1 + generateOutput(groupingKey, passThroughAggregationBuffer) + } + + /////////////////////////////////////////////////////////////////////////// + // Part 7: Loads input rows and sets up aggregationBufferMapIterator if we // have not switched to sort-based aggregation. /////////////////////////////////////////////////////////////////////////// @@ -394,16 +504,17 @@ class TungstenAggregationIterator( }) /////////////////////////////////////////////////////////////////////////// - // Part 7: Iterator's public methods. + // Part 8: Iterator's public methods. /////////////////////////////////////////////////////////////////////////// override final def hasNext: Boolean = { - (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) + (sortBased && sortedInputHasNewGroup) || (!sortBased && mapIteratorHasNext) || + passThroughHasNext } override final def next(): UnsafeRow = { if (hasNext) { - val res = if (sortBased) { + val res = if (sortBased && sortedInputHasNewGroup) { // Process the current group. processCurrentSortedGroup() // Generate output row for the current group. @@ -412,7 +523,7 @@ class TungstenAggregationIterator( sortBasedAggregationBuffer.copyFrom(initialAggregationBuffer) outputRow - } else { + } else if (mapIteratorHasNext) { // We did not fall back to sort-based aggregation. val result = generateOutput( @@ -426,13 +537,25 @@ class TungstenAggregationIterator( if (!mapIteratorHasNext) { // If there is no input from aggregationBufferMapIterator, we copy current result. val resultCopy = result.copy() - // Then, we free the map. + // Then, we free the map. Pass-through (if any) does not use the map. hashMap.free() resultCopy } else { result } + } else { + // Adaptive partial aggregation bypassed partial aggregation: emit the remaining input + // rows as single-row partial buffers, after the frozen map output above. + if (passThroughDrainStart == -1L) { + passThroughDrainStart = System.nanoTime() + } + val out = nextPassThroughOutput() + if (!passThroughHasNext) { + aggTime += NANOSECONDS.toMillis(System.nanoTime() - passThroughDrainStart) + passThroughDrainStart = -1L + } + out } numOutputRows += 1 @@ -444,7 +567,7 @@ class TungstenAggregationIterator( } /////////////////////////////////////////////////////////////////////////// - // Part 8: Utility functions + // Part 9: Utility functions /////////////////////////////////////////////////////////////////////////// /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala index df0addad7861a..2a5ab8dd04284 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/TypedAggregateExpression.scala @@ -107,7 +107,7 @@ trait TypedAggregateExpression extends AggregateFunction { // aggregator.getClass.getSimpleName can cause Malformed class name error, // call safer `Utils.getSimpleName` instead - override def nodeName: String = Utils.getSimpleName(aggregator.getClass).stripSuffix("$"); + override def nodeName: String = Utils.getSimpleName(aggregator.getClass).stripSuffix("$") } // TODO: merge these 2 implementations once we refactor the `AggregateFunction` interface. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala index f9c4ecc14e6c7..e1f9682439a31 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/VectorizedHashMapGenerator.scala @@ -83,6 +83,10 @@ class VectorizedHashMapGenerator( | buckets = new int[numBuckets]; | java.util.Arrays.fill(buckets, -1); | } + | + | public int getNumKeys() { + | return numRows; + | } """.stripMargin } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index e94b774a37d19..bea86501e6f3a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.optimizer.CollapseProject +import org.apache.spark.sql.catalyst.plans.logical.Sample import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -207,10 +208,19 @@ trait GeneratePredicateHelper extends PredicateHelper { } }.mkString("\n").trim + val nestedNullChecks = notNullPreds.zipWithIndex.collect { + case (p @ IsNotNull(n), idx) + if !generatedIsNotNullChecks(idx) && !n.isInstanceOf[Attribute] && + c.exists(_.semanticEquals(n)) => + generatedIsNotNullChecks(idx) = true + genPredicate(p, inputExprCode, inputAttrs) + }.mkString("\n").trim + // Here we use *this* operator's output with this output's nullability since we already // enforced them with the IsNotNull checks above. s""" |$nullChecks + |$nestedNullChecks |${genPredicate(c, inputExprCode, outputAttrs)} """.stripMargin.trim }.mkString("\n") @@ -400,6 +410,20 @@ case class FilterExec(condition: Expression, child: SparkPlan) parts.append('\n') } } + notNullPreds.zipWithIndex.foreach { + case (p @ IsNotNull(n), ni) + if !generatedIsNotNullChecks(ni) && !n.isInstanceOf[Attribute] && + orig.exists(_.semanticEquals(n)) => + generatedIsNotNullChecks(ni) = true + var checkCode: String = null + ctx.withSubExprEliminationExprs(Map.empty) { + checkCode = genNotNull(p) + Seq.empty + } + parts.append(checkCode) + parts.append('\n') + case _ => + } statesByFirstUse.get(idx).foreach { states => parts.append(ctx.evaluateSubExprEliminationState(states)) parts.append('\n') @@ -497,7 +521,7 @@ case class SampleExec( seed: Option[Long], child: SparkPlan) extends UnaryExecNode with CodegenSupport { - val resolvedSeed: Long = seed.getOrElse((math.random() * 1000).toLong) + val resolvedSeed: Long = Sample.resolveSeed(seed) override def output: Seq[Attribute] = child.output @@ -891,22 +915,24 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } /** - * Returns the output partitionings of the children, with the attributes converted to - * the first child's attributes at the same position. + * Returns the output partitionings of the children, with the attributes converted to this + * union's output attributes at the same position. */ private def prepareOutputPartitioning(): Seq[Partitioning] = { - // Create a map of attributes from the other children to the first child. - val firstAttrs = children.head.output - val attributesMap = children.tail.map(_.output).map { otherAttrs => - AttributeMap(otherAttrs.zip(firstAttrs)) + // Map every child's partitioning attributes to this union's output attributes, so all + // partitionings are expressed in the same attribute space before comparison. A child's + // `outputPartitioning` may reference attributes that differ from its own `output` in any + // field `AttributeReference.equals` compares other than `dataType`, which two attributes + // sharing an `ExprId` agree on (so name, nullability, metadata, qualifier): a Filter + // narrows nullability via `IsNotNull` while passing its child's partitioning through, and + // a partitioning built inside a view or subquery carries that relation's qualifier. + // `AttributeMap` is keyed by `ExprId`, so remapping every child (including the first) + // normalizes all of those. + val unionOutput = output + val attributesMap = children.map(_.output).map { childAttrs => + AttributeMap(childAttrs.zip(unionOutput)) } - - val partitionings = children.map(_.outputPartitioning) - val firstPartitioning = partitionings.head - val otherPartitionings = partitionings.tail - - val convertedOtherPartitionings = otherPartitionings.zipWithIndex.map { case (p, idx) => - val attributeMap = attributesMap(idx) + children.map(_.outputPartitioning).zip(attributesMap).map { case (p, attributeMap) => p match { case e: Expression => e.transform { @@ -916,7 +942,6 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup case _ => p } } - Seq(firstPartitioning) ++ convertedOtherPartitionings } // Compares two leaf partitionings for union pass-through equivalence. Callers pass leaf @@ -926,8 +951,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup (left, right) match { case (SinglePartition, SinglePartition) => true case (l: HashPartitioningLike, r: HashPartitioningLike) => l == r - // For `KeyedPartitioning`, only the partition expressions must match (the other child's - // expressions have already been remapped to the first child's attributes by + // For `KeyedPartitioning`, only the partition expressions must match (both sides' + // expressions have already been remapped to this union's output attributes by // `prepareOutputPartitioning`). The partition keys are intentionally not compared here: // children typically carry different key sets, and `outputPartitioning` merges them. case (l: KeyedPartitioning, r: KeyedPartitioning) => @@ -944,17 +969,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup return super.outputPartitioning } - // Children's partitionings with attributes remapped to the first child's attributes. + // Children's partitionings with attributes remapped to this union's output attributes. val partitionings = prepareOutputPartitioning() - // Map from the first child's attributes to this union's own output attributes. - val attributeMap = children.head.output.zip(output).toMap - def toUnionOutput(p: Partitioning): Partitioning = p match { - case e: Expression => - e.transform { - case a: Attribute if attributeMap.contains(a) => attributeMap(a) - }.asInstanceOf[Partitioning] - case _ => p - } // Case A: every child is a single `KeyedPartitioning`. A `UnionExec` concatenates its // children's partitions in order (one child's partitions after another's), so the merged @@ -971,9 +987,7 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup val compatible = kps.forall(comparePartitioning(_, headKp)) if (compatible) { val mergedKeys = kps.flatMap(_.partitionKeys) - val mergedExpressions = headKp.expressions.map(_.transform { - case a: Attribute if attributeMap.contains(a) => attributeMap(a) - }) + val mergedExpressions = headKp.expressions val isGrouped = mergedKeys.distinct.size == mergedKeys.size val isNarrowed = kps.exists(_.isNarrowed) return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed) @@ -1004,8 +1018,8 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup } intersection match { case Seq() => super.outputPartitioning - case Seq(p) => toUnionOutput(p) - case ps => PartitioningCollection.fromPartitionings(ps.map(toUnionOutput)) + case Seq(p) => p + case ps => PartitioningCollection.fromPartitionings(ps) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala index f90f1a9dcd226..baba2c1365c82 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala @@ -23,10 +23,14 @@ import java.nio.channels.Channels import scala.jdk.CollectionConverters._ import org.apache.arrow.compression.{Lz4CompressionCodec, ZstdCompressionCodec} -import org.apache.arrow.vector.{VectorLoader, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.flatbuf.{RecordBatch => FlatBufRecordBatch} +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{TypeLayout, VectorLoader, VectorSchemaRoot, VectorUnloader} import org.apache.arrow.vector.compression.{CompressionCodec, NoCompressionCodec} import org.apache.arrow.vector.ipc.{ReadChannel, WriteChannel} +import org.apache.arrow.vector.ipc.message.{ArrowBodyCompression, ArrowFieldNode} import org.apache.arrow.vector.ipc.message.{ArrowRecordBatch, MessageSerializer} +import org.apache.arrow.vector.types.pojo.Field import org.apache.spark.{SparkException, TaskContext} import org.apache.spark.rdd.RDD @@ -144,8 +148,7 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { conf: SQLConf): RDD[ColumnarBatch] = { val cacheSchema = DataTypeUtils.fromAttributes(cacheAttributes) val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes) - val columnIndices = - selectedAttributes.map(a => cacheAttributes.map(o => o.exprId).indexOf(a.exprId)).toArray + val columnIndices = CachedColumnIndices(cacheAttributes, selectedAttributes) // Capture config on driver val timeZoneId = conf.sessionLocalTimeZone val prefetchEnabled = conf.arrowCachePrefetchEnabled @@ -183,11 +186,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { val selectedSchema = DataTypeUtils.fromAttributes(selectedAttributes) val timeZoneId = conf.sessionLocalTimeZone - // Calculate column indices for projection - val selectedIndices = selectedAttributes.map { attr => - cacheAttributes.indexWhere(_.exprId == attr.exprId) - }.toArray - // Check if all selected types can use the fast path. // Types not handled by ArrowColumnReader must use the fallback path. val needsFallback = selectedSchema.fields.exists { f => @@ -227,6 +225,9 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { } } } else { + // Only the fast path consumes the column indices; the fallback branch above delegates to + // convertCachedBatchToColumnarBatch, which resolves them itself. + val selectedIndices = CachedColumnIndices(cacheAttributes, selectedAttributes) val prefetchEnabled = conf.arrowCachePrefetchEnabled input.mapPartitionsInternal { batchIterator => new ArrowCachedBatchToInternalRowIterator( @@ -288,6 +289,135 @@ private object ArrowCachedBatchSerializer { out.toByteArray } + /** + * Number of Arrow buffers a field occupies in a RecordBatch body, including all of its + * descendants, in the depth-first order `VectorLoader` consumes them. The type's own buffer + * count comes from `TypeLayout` (validity + offset/data buffers), then each child contributes + * its whole subtree recursively. Used to map each top-level column to its run of buffers. + */ + private def fieldBufferCount(field: Field): Int = + TypeLayout.getTypeBufferCount(field.getType) + + field.getChildren.asScala.map(fieldBufferCount).sum + + /** Number of Arrow field nodes a field occupies (itself plus every descendant). */ + private def fieldNodeCount(field: Field): Int = + 1 + field.getChildren.asScala.map(fieldNodeCount).sum + + /** Number of variadic buffer counts a field contributes (one per view-type buffer, recursive). */ + private def fieldVariadicCount(field: Field): Int = { + val own = field.getType match { + // View types (Utf8View/BinaryView) carry a variadic-buffer count in the RecordBatch; + // no other type does. The cache never writes view vectors today, but account for them so + // the span arithmetic stays correct if that changes. + case _: org.apache.arrow.vector.types.pojo.ArrowType.Utf8View | + _: org.apache.arrow.vector.types.pojo.ArrowType.BinaryView => 1 + case _ => 0 + } + own + field.getChildren.asScala.map(fieldVariadicCount).sum + } + + /** + * Read an encapsulated IPC RecordBatch message from `data`, materializing off-heap only the + * buffers of the requested top-level columns. This is the projection-pushdown read path: the + * message metadata (a small flatbuffer) lists every buffer's (offset, length) within the body, + * so we copy just the byte ranges belonging to the selected columns straight out of the + * in-memory `data` array, never touching (or allocating off-heap for) the unselected columns. + * + * The body is a flat, depth-first sequence of buffers in schema order, so each top-level column + * owns a contiguous run of buffers whose span is `fieldBufferCount`; field nodes and variadic + * counts run in the same order. The selected columns' bytes are copied into one off-heap buffer + * (each buffer 8-byte aligned, matching Arrow's IPC body layout) and the returned batch's + * buffers are windows into it, exactly like the standard reader slices one body buffer -- so the + * batch has a single underlying allocation and no per-buffer bookkeeping. The returned batch + * owns its buffers (the constructor retains each), so the caller closes it as usual. + * + * Compression is preserved unchanged: buffer (offset, length) spans cover the on-body bytes + * including any per-buffer uncompressed-length prefix, so the copied windows are still compressed + * as written; `VectorLoader.load` decompresses only the selected ones later. + */ + def readProjectedRecordBatch( + data: Array[Byte], + schemaFields: Seq[Field], + selectedIndices: Array[Int], + allocator: BufferAllocator): ArrowRecordBatch = { + val in = new ByteArrayInputStream(data) + val readChannel = new ReadChannel(Channels.newChannel(in)) + // Read only the message metadata; the body bytes stay in `data` and are copied selectively. + val metadata = MessageSerializer.readMessage(readChannel) + require(metadata != null, "Unexpected end of input reading cached batch message") + val batch = + metadata.getMessage.header(new FlatBufRecordBatch()).asInstanceOf[FlatBufRecordBatch] + // serializeBatch writes exactly [encapsulated message][body] with no end-of-stream marker, so + // the body is the tail of `data`: it starts at data.length minus the declared body length. + val bodyStart = data.length - metadata.getMessageBodyLength().toInt + + val compression: ArrowBodyCompression = + if (batch.compression() == null) NoCompressionCodec.DEFAULT_BODY_COMPRESSION + else new ArrowBodyCompression(batch.compression().codec(), batch.compression().method()) + + val nodeStarts = schemaFields.scanLeft(0)(_ + fieldNodeCount(_)).toArray + val bufferStarts = schemaFields.scanLeft(0)(_ + fieldBufferCount(_)).toArray + val variadicStarts = schemaFields.scanLeft(0)(_ + fieldVariadicCount(_)).toArray + val hasVariadic = batch.variadicBufferCountsLength() > 0 + + // Enumerate the selected columns' nodes, buffer indices and variadic counts, in output order. + val selectedNodes = new java.util.ArrayList[ArrowFieldNode]() + val selectedBufferIdx = new scala.collection.mutable.ArrayBuffer[Int]() + val selectedVariadic = new java.util.ArrayList[java.lang.Long]() + selectedIndices.foreach { i => + val field = schemaFields(i) + val nStart = nodeStarts(i) + (nStart until nStart + fieldNodeCount(field)).foreach { j => + val node = batch.nodes(j) + selectedNodes.add(new ArrowFieldNode(node.length(), node.nullCount())) + } + val bStart = bufferStarts(i) + (bStart until bStart + fieldBufferCount(field)).foreach(selectedBufferIdx += _) + if (hasVariadic) { + val vStart = variadicStarts(i) + (vStart until vStart + fieldVariadicCount(field)).foreach(j => + selectedVariadic.add(batch.variadicBufferCounts(j))) + } + } + + val layout = selectedBufferIdx.map { j => + val buf = batch.buffers(j) + (buf.offset(), buf.length()) + } + val alignedSizes = layout.map { case (_, len) => ((len + 7) / 8) * 8 } + val body = allocator.buffer(math.max(alignedSizes.sum, 1)) + try { + val selectedBuffers = new java.util.ArrayList[org.apache.arrow.memory.ArrowBuf]() + var pos = 0L + layout.indices.foreach { k => + val (srcOffset, len) = layout(k) + if (len > 0) { + body.setBytes(pos, data, bodyStart + srcOffset.toInt, len.toInt) + } + val window = body.slice(pos, len) + window.writerIndex(len) + selectedBuffers.add(window) + pos += alignedSizes(k) + } + val recordBatch = new ArrowRecordBatch( + batch.length().toInt, + selectedNodes, + selectedBuffers, + compression, + selectedVariadic, + false) + // The constructor retained each window (slice() itself does not), so the batch now holds one + // reference per window into `body`. Drop `body`'s original allocation reference; the batch is + // then the sole owner and the caller's recordBatch.close() frees the single allocation. + body.close() + recordBatch + } catch { + case t: Throwable => + body.close() + throw t + } + } + /** * Byte offset of the unscaled low-order word within a 16-byte Arrow Decimal128 slot, for the * given native byte order. Arrow Java writes decimal values in native byte order @@ -1146,6 +1276,21 @@ private class ArrowCachedBatchToColumnarBatchIterator( private val arrowSchema = ArrowUtils.toArrowSchema( cacheSchema, timeZoneId, false, false, losslessInternalTypes = true) + // Projection pushdown: the cached batch stores all cache columns, but only the selected ones + // are needed. When every selected column maps to a distinct cached column, read a batch holding + // only the selected columns' buffers (in columnIndices order) so unselected columns are never + // copied off-heap, loaded, or decompressed. The projected schema's field order matches + // columnIndices, so the loaded root's vectors are already in output order. If any selected + // attribute is absent from the cache schema (index -1), fall back to reading the full batch. + private val cacheFields = arrowSchema.getFields.asScala.toSeq + private val canProjectOnLoad = columnIndices.forall(_ >= 0) + private val projectedSchema = + if (canProjectOnLoad) { + new org.apache.arrow.vector.types.pojo.Schema(columnIndices.map(cacheFields).toList.asJava) + } else { + arrowSchema + } + // Track only the previous root to close it when next batch is produced private var previousRoot: VectorSchemaRoot = null @@ -1205,11 +1350,16 @@ private class ArrowCachedBatchToColumnarBatchIterator( previousRoot = root - // Wrap vectors in ArrowColumnVector and project to selected columns. - val allColumns = root.getFieldVectors.asScala.map { vector => - new ArrowColumnVector(vector) - }.toArray[ColumnVector] - val selectedColumns = columnIndices.map(allColumns(_)) + // When projected on load, the root already holds only the selected columns in output order, + // so wrap its vectors directly. Otherwise it holds all cache columns and must be selected. + val selectedColumns = if (canProjectOnLoad) { + root.getFieldVectors.asScala.map(v => new ArrowColumnVector(v)).toArray[ColumnVector] + } else { + val allColumns = root.getFieldVectors.asScala.map { vector => + new ArrowColumnVector(vector) + }.toArray[ColumnVector] + columnIndices.map(allColumns(_)) + } val batch = new ColumnarBatch(selectedColumns, root.getRowCount) // Start prefetching the next batch while this one is being consumed. @@ -1220,11 +1370,18 @@ private class ArrowCachedBatchToColumnarBatchIterator( /** Deserialize a cached batch into its own freshly-created root. Does not touch other roots. */ private def deserializeToRoot(cachedBatch: ArrowCachedBatch): VectorSchemaRoot = { - val in = new ByteArrayInputStream(cachedBatch.arrowData) - val readChannel = new ReadChannel(Channels.newChannel(in)) - val recordBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator) + // Projection pushdown: read only the selected columns' buffers out of the cached bytes, so + // unselected columns are never copied off-heap, loaded, or decompressed. + val recordBatch = if (canProjectOnLoad) { + ArrowCachedBatchSerializer.readProjectedRecordBatch( + cachedBatch.arrowData, cacheFields, columnIndices, allocator) + } else { + val in = new ByteArrayInputStream(cachedBatch.arrowData) + val readChannel = new ReadChannel(Channels.newChannel(in)) + MessageSerializer.deserializeRecordBatch(readChannel, allocator) + } Utils.tryWithSafeFinally { - val root = VectorSchemaRoot.create(arrowSchema, allocator) + val root = VectorSchemaRoot.create(projectedSchema, allocator) // VectorLoader.load fills vectors incrementally, so a failure (malformed data, decompression // error, OOM) can occur after earlier vectors have allocated buffers. Close the partially // loaded root on failure, otherwise it becomes unreachable and the later allocator.close() @@ -1427,6 +1584,20 @@ private class ArrowCachedBatchToInternalRowIterator( private val arrowSchema = ArrowUtils.toArrowSchema( cacheSchema, timeZoneId, false, false, losslessInternalTypes = true) + // Projection pushdown: see ArrowCachedBatchToColumnarBatchIterator. When every selected column + // maps to a distinct cached column, read a batch holding only the selected columns' buffers so + // unselected columns are never copied off-heap, loaded, or decompressed, and readers bind + // positionally. If any selected attribute is absent from the cache (index -1), fall back to the + // full batch and bind readers via columnIndices. + private val cacheFields = arrowSchema.getFields.asScala.toSeq + private val canProjectOnLoad = columnIndices.forall(_ >= 0) + private val projectedSchema = + if (canProjectOnLoad) { + new org.apache.arrow.vector.types.pojo.Schema(columnIndices.map(cacheFields).toList.asJava) + } else { + arrowSchema + } + // Pre-build typed readers per column at init time -- no per-row pattern match private val columnReaders: Array[ArrowColumnReader] = selectedSchema.fields.map(f => ArrowColumnReader.create(f.dataType)) @@ -1507,11 +1678,17 @@ private class ArrowCachedBatchToInternalRowIterator( /** Deserialize a cached batch into a VectorSchemaRoot. */ private def deserializeBatch(cachedBatch: ArrowCachedBatch): VectorSchemaRoot = { - val in = new ByteArrayInputStream(cachedBatch.arrowData) - val readChannel = new ReadChannel(Channels.newChannel(in)) - val recordBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator) + // Projection pushdown: read only the selected columns' buffers out of the cached bytes. + val recordBatch = if (canProjectOnLoad) { + ArrowCachedBatchSerializer.readProjectedRecordBatch( + cachedBatch.arrowData, cacheFields, columnIndices, allocator) + } else { + val in = new ByteArrayInputStream(cachedBatch.arrowData) + val readChannel = new ReadChannel(Channels.newChannel(in)) + MessageSerializer.deserializeRecordBatch(readChannel, allocator) + } try { - val root = VectorSchemaRoot.create(arrowSchema, allocator) + val root = VectorSchemaRoot.create(projectedSchema, allocator) // VectorLoader.load fills vectors incrementally, so a failure (malformed data, decompression // error, OOM) can occur after earlier vectors have allocated buffers. Close the partially // loaded root on failure, otherwise it becomes unreachable and the later allocator.close() @@ -1563,10 +1740,13 @@ private class ArrowCachedBatchToInternalRowIterator( currentRoot = root - // Update pre-built readers with new vectors + // Update pre-built readers with new vectors. When projected on load, the root holds the + // selected columns positionally; otherwise it holds all cache columns, selected via + // columnIndices. var i = 0 while (i < numFields) { - columnReaders(i).setVector(root.getVector(columnIndices(i))) + val vectorIndex = if (canProjectOnLoad) i else columnIndices(i) + columnReaders(i).setVector(root.getVector(vectorIndex)) i += 1 } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala new file mode 100644 index 0000000000000..14fd150dfa7f1 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/CachedColumnIndices.scala @@ -0,0 +1,40 @@ +/* + * 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.execution.columnar + +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSeq} + +/** + * Resolves where the columns a scan requests sit within a cached batch's schema. Every + * `CachedBatchSerializer` needs this mapping to project the requested columns out of a cached + * batch, so it lives here rather than in each serializer. + */ +private[columnar] object CachedColumnIndices { + + /** + * Returns, for each attribute in `selectedAttributes`, its ordinal in `cacheAttributes`, or -1 + * if the attribute is not part of the cached schema. + */ + def apply(cacheAttributes: Seq[Attribute], selectedAttributes: Seq[Attribute]): Array[Int] = { + // The explicit AttributeSeq wrapper is required. On a bare Seq[Attribute] the inherited + // Seq.indexOf[B >: Attribute](elem: B) wins overload resolution with B inferred as Any, so + // the implicit conversion never fires and every column silently resolves to -1. + val cacheAttributeSeq = AttributeSeq(cacheAttributes) + selectedAttributes.map(a => cacheAttributeSeq.indexOf(a.exprId)).toArray + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala index 14ab652b4f077..b597f5e238862 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/GenerateColumnAccessor.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.execution.columnar +import scala.collection.AbstractIterator + import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow @@ -27,8 +29,17 @@ import org.apache.spark.unsafe.types.CalendarInterval /** * An Iterator to walk through the InternalRows from a CachedBatch + * + * Extends `AbstractIterator` rather than mixing in `Iterator` directly. A subclass of this type + * is produced by runtime codegen, and javac (the alternative codegen backend) cannot subclass a + * scalac-compiled class that fixes a trait's type parameter while mixing the trait in directly: + * the mixin forwarders scalac emits on such a class are raw, carrying no `Signature` attribute, + * so javac sees e.g. `Object minBy(Function1, Ordering)` clash with `IterableOnceOps`' + * `<B> A minBy(Function1<A, B>, Ordering<B>)`. `AbstractIterator` stays generic in its element + * type, so its forwarders keep their signatures and no clash arises. Janino does not perform + * this check, so the direct form worked there. */ -abstract class ColumnarIterator extends Iterator[InternalRow] { +abstract class ColumnarIterator extends AbstractIterator[InternalRow] { def initialize(input: Iterator[DefaultCachedBatch], columnTypes: Array[DataType], columnIndexes: Array[Int]): Unit } @@ -165,7 +176,6 @@ object GenerateColumnAccessor extends CodeGenerator[Seq[DataType], ColumnarItera import java.nio.ByteOrder; import scala.collection.Iterator; import org.apache.spark.sql.types.DataType; - import org.apache.spark.sql.catalyst.expressions.codegen.BufferHolder; import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter; import org.apache.spark.sql.execution.columnar.MutableUnsafeRow; diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index e19d562c19b5b..e6bf65ec89e6b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -22,17 +22,25 @@ import com.esotericsoftware.kryo.io.{Input => KryoInput, Output => KryoOutput} import org.apache.spark.{SparkException, TaskContext} import org.apache.spark.network.util.JavaUtils -import org.apache.spark.rdd.RDD -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.rdd.{DeterministicLevel, RDD} +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.analysis.MultiInstanceRelation import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.plans.{logical, QueryPlan} import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, LogicalPlan, Statistics} +import org.apache.spark.sql.catalyst.trees.TreePattern.CURRENT_LIKE import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.catalyst.util.truncatedString +import org.apache.spark.sql.catalyst.util.{truncatedString, CaseInsensitiveMap} import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer, SimpleMetricsCachedBatch, SimpleMetricsCachedBatchSerializer} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.datasources.{FileFormat, FileScanRDD, HadoopFsRelation, LogicalRelation} +import org.apache.spark.sql.execution.datasources.binaryfile.BinaryFileFormat +import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat +import org.apache.spark.sql.execution.datasources.json.JsonFileFormat +import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat +import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat +import org.apache.spark.sql.execution.datasources.text.TextFileFormat import org.apache.spark.sql.execution.vectorized.{OffHeapColumnVector, OnHeapColumnVector, WritableColumnVector} import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.types._ @@ -197,8 +205,7 @@ class DefaultCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { conf: SQLConf): RDD[ColumnarBatch] = { val offHeapColumnVectorEnabled = conf.offHeapColumnVectorEnabled val outputSchema = DataTypeUtils.fromAttributes(selectedAttributes) - val columnIndices = - selectedAttributes.map(a => cacheAttributes.map(o => o.exprId).indexOf(a.exprId)).toArray + val columnIndices = CachedColumnIndices(cacheAttributes, selectedAttributes) def createAndDecompressColumn(cb: CachedBatch): ColumnarBatch = { val cachedColumnarBatch = cb.asInstanceOf[DefaultCachedBatch] @@ -231,21 +238,18 @@ class DefaultCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer { selectedAttributes: Seq[Attribute], conf: SQLConf): RDD[InternalRow] = { // Find the ordinals and data types of the requested columns. - val (requestedColumnIndices, requestedColumnDataTypes) = - selectedAttributes.map { a => - cacheAttributes.map(_.exprId).indexOf(a.exprId) -> a.dataType - }.unzip + val requestedColumnIndices = CachedColumnIndices(cacheAttributes, selectedAttributes) - val columnTypes = requestedColumnDataTypes.map { + val columnTypes = selectedAttributes.map(_.dataType match { case udt: UserDefinedType[_] => udt.sqlType case other => other - }.toArray + }).toArray input.mapPartitionsInternal { cachedBatchIterator => val columnarIterator = GenerateColumnAccessor.generate(columnTypes.toImmutableArraySeq) columnarIterator.initialize(cachedBatchIterator.asInstanceOf[Iterator[DefaultCachedBatch]], columnTypes, - requestedColumnIndices.toArray) + requestedColumnIndices) columnarIterator } } @@ -257,10 +261,14 @@ case class CachedRDDBuilder( storageLevel: StorageLevel, @transient cachedPlan: SparkPlan, tableName: Option[String], - @transient logicalPlan: LogicalPlan) { + @transient logicalPlan: LogicalPlan, + isCachedLogicalPlanRepeatable: Boolean = false, + hasSelectivePredicate: Boolean = false, + fileSourceOptions: Seq[Map[String, String]] = Seq.empty) { @transient @volatile private var _cachedColumnBuffers: RDD[CachedBatch] = null - @transient @volatile private var _cachedColumnBuffersAreLoaded: Boolean = false + @volatile private var isCachedRDDRepeatable = false + private var hasStrictFileSourceReads = true // The cache's materialization bookkeeping: a partition-keyed accumulator storing // (rowCount, sizeInBytes) per partition. AQE creates a separate cache scan stage per reference to @@ -300,32 +308,42 @@ case class CachedRDDBuilder( if (_cachedColumnBuffers != null) { _cachedColumnBuffers.unpersist(blocking) _cachedColumnBuffers = null - // The buffers no longer back a live RDD. Reset the one-way "loaded" latch and install new - // bookkeeping so a rebuild on this builder does not inherit stale state or late updates from - // tasks that captured the previous generation's accumulator. - _cachedColumnBuffersAreLoaded = false partitionStats = newPartitionStats() } + isCachedRDDRepeatable = false + // Read strictness is derived independently for each cache generation. + hasStrictFileSourceReads = true } def isCachedColumnBuffersLoaded: Boolean = synchronized { - _cachedColumnBuffers != null && isCachedRDDLoaded - } - - private def isCachedRDDLoaded: Boolean = { - _cachedColumnBuffersAreLoaded || { - // We must make sure the statistics of `sizeInBytes` and `rowCount` are accurate if - // `isCachedRDDLoaded` return true. Otherwise, AQE would do a wrong optimization, - // e.g., convert a non-empty plan to empty local relation if `rowCount` is 0. - // Count DISTINCT materialized partitions (the keyed accumulator's key set), so the cache is - // only reported loaded once every partition has been computed -- sound even if a partition is - // computed more than once by concurrent or speculative tasks. - val numMaterialized = partitionStats.accumulatedNumPartitions - val rddLoaded = _cachedColumnBuffers.partitions.length.toLong == numMaterialized - if (rddLoaded) { - _cachedColumnBuffersAreLoaded = rddLoaded + _cachedColumnBuffers != null && + partitionStats.accumulatedNumPartitions == _cachedColumnBuffers.partitions.length + } + + private[sql] def isCachedPlanRepeatable: Boolean = + isCachedLogicalPlanRepeatable && isCachedRDDRepeatable + + /** Reads completeness and exact statistics from one cache generation atomically. */ + private[sql] def loadedMaterializedStats: Option[(Long, Long)] = synchronized { + if (_cachedColumnBuffers == null) { + None + } else { + partitionStats.foldValuesIfComplete( + _cachedColumnBuffers.partitions.length, + (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) } - rddLoaded + } + } + + private[sql] def materializedMetadata: Option[logical.MaterializedLeafMetadata] = synchronized { + loadedMaterializedStats.map { case (rowCount, sizeInBytes) => + logical.MaterializedLeafMetadata( + rowCount = rowCount, + sizeInBytes = sizeInBytes, + isOutputRepeatable = isCachedPlanRepeatable, + isDurable = storageLevel.useDisk) } } @@ -347,16 +365,51 @@ case class CachedRDDBuilder( } private def buildBuffers(): RDD[CachedBatch] = { + def buildInputRDD[T](input: => RDD[T]): RDD[T] = { + if (fileSourceOptions.isEmpty) { + input + } else { + // File scans initialize their input RDD lazily. Check both configuration domains while + // constructing that RDD so a temporary best-effort setting cannot be mistaken for a + // repeatable strict read. + val materializationConf = SQLConf.get.clone() + val cachedPlanConf = cachedPlan.conf.clone() + + def hasStrictReads(conf: SQLConf): Boolean = SQLConf.withExistingConf(conf) { + fileSourceOptions.forall { options => + val effectiveOptions = new FileSourceOptions(options) + !effectiveOptions.ignoreMissingFiles && !effectiveOptions.ignoreCorruptFiles + } + } + + val (inputRDD, strictPhysicalReads) = SQLConf.withExistingConf(materializationConf) { + val result = input + val fileScans = cachedPlan.collect { case scan: FileSourceScanExec => scan } + val scansAreStrict = fileScans.size == fileSourceOptions.size && fileScans.forall { + scan => scan.inputRDD match { + case fileRDD: FileScanRDD => fileRDD.hasStrictFileReads + case _ => false + } + } + (result, scansAreStrict) + } + hasStrictFileSourceReads = hasStrictFileSourceReads && + hasStrictReads(materializationConf) && hasStrictReads(cachedPlanConf) && + strictPhysicalReads + inputRDD + } + } + val cb = try { if (supportsColumnarInput) { serializer.convertColumnarBatchToCachedBatch( - cachedPlan.executeColumnar(), + buildInputRDD(cachedPlan.executeColumnar()), cachedPlan.output, storageLevel, cachedPlan.conf) } else { serializer.convertInternalRowToCachedBatch( - cachedPlan.execute(), + buildInputRDD(cachedPlan.execute()), cachedPlan.output, storageLevel, cachedPlan.conf) @@ -374,9 +427,8 @@ case class CachedRDDBuilder( // id. Bound to a local so the task closure below captures only the accumulator, not the // enclosing CachedRDDBuilder (whose cachedPlan is not serializable). val accumulator = partitionStats - val cached = cb.mapPartitionsInternal { it => + val cached = cb.mapPartitionsWithIndexInternal { (partitionId, it) => val taskContext = TaskContext.get() - val partitionId = taskContext.partitionId() // This task computes exactly one partition. Tally its totals so the completion listener // records them once, keyed by partition id (covering empty-output partitions, which produce // no batches). @@ -403,11 +455,127 @@ case class CachedRDDBuilder( } }.persist(storageLevel) cached.setName(cachedName) + isCachedRDDRepeatable = hasStrictFileSourceReads && + cached.outputDeterministicLevel != DeterministicLevel.INDETERMINATE && + InMemoryRelation.hasRepeatablePhysicalPlan(cachedPlan) cached } } -object InMemoryRelation { +object InMemoryRelation extends PredicateHelper { + + private val trustedFileFormatClasses: Set[Class[_ <: FileFormat]] = Set( + classOf[BinaryFileFormat], + classOf[CSVFileFormat], + classOf[JsonFileFormat], + classOf[OrcFileFormat], + classOf[ParquetFileFormat], + classOf[TextFileFormat]) + + private val trustedExternalFileFormatNames = Set( + "org.apache.spark.sql.avro.AvroFileFormat", + "org.apache.spark.sql.hive.orc.OrcFileFormat") + + private val catalystExpressionPackage = "org.apache.spark.sql.catalyst.expressions." + + private def hasSafeExpressions(plan: QueryPlan[_]): Boolean = { + // Treat the Catalyst namespace as the trust boundary for Expression.deterministic's + // repeatability contract. Expressions outside it fail closed; reject AesEncrypt and + // opaque/user-defined expressions explicitly. + plan.expressions.forall { expression => + !expression.exists { + case _: AesEncrypt | _: NonSQLExpression | _: UserDefinedExpression => true + case value => !value.deterministic || value.containsPattern(CURRENT_LIKE) || + !value.getClass.getName.startsWith(catalystExpressionPackage) + } + } + } + + private def hasRepeatableLogicalPlan( + analyzedPlan: LogicalPlan, + plan: LogicalPlan, + onOptimizedNode: LogicalPlan => Unit): Boolean = { + // Runtime-replaceable expressions such as AES encryption can become deterministic-looking + // StaticInvoke nodes during optimization despite using a fresh random initialization vector. + // Inspect the original analyzed expressions before trusting the optimized execution shape. + var repeatable = analyzedPlan.deterministic + if (repeatable) { + analyzedPlan.foreachWithSubqueries { node => + if (repeatable && !hasSafeExpressions(node)) { + repeatable = false + } + } + } + if (repeatable) { + repeatable = plan.deterministic + } + plan.foreachWithSubqueries { node => + onOptimizedNode(node) + if (repeatable) { + repeatable = hasSafeExpressions(node) && (node match { + case _: logical.Project | _: logical.Filter | _: logical.SubqueryAlias | + _: logical.Range | _: logical.LocalRelation => true + case relation: LogicalRelation => relation.relation match { + case fileRelation: HadoopFsRelation => + val fileFormatClass = fileRelation.fileFormat.getClass + trustedFileFormatClasses.contains(fileFormatClass) || + trustedExternalFileFormatNames.contains(fileFormatClass.getName) + case _ => false + } + case _ => false + }) + } + } + repeatable + } + + private[columnar] def hasRepeatablePhysicalPlan(plan: SparkPlan): Boolean = { + !plan.exists { node => + val supported = node match { + case _: ColumnarToRowExec | _: FileSourceScanExec | _: FilterExec | + _: InputAdapter | _: LocalTableScanExec | _: ProjectExec | + _: RangeExec | _: WholeStageCodegenExec => true + case _ => false + } + !supported || node.subqueries.nonEmpty || !hasSafeExpressions(node) + } + } + + private def newCacheBuilder( + serializer: CachedBatchSerializer, + storageLevel: StorageLevel, + cachedPlan: SparkPlan, + tableName: Option[String], + logicalPlan: LogicalPlan, + analyzedPlan: LogicalPlan, + optimizedPlan: LogicalPlan): CachedRDDBuilder = { + val relevantOptions = Seq( + FileSourceOptions.IGNORE_MISSING_FILES, + FileSourceOptions.IGNORE_CORRUPT_FILES) + val fileSourceOptions = Seq.newBuilder[Map[String, String]] + var hasSelectivePredicate = false + val repeatable = hasRepeatableLogicalPlan(analyzedPlan, optimizedPlan, { + case logical.Filter(condition, _) => + if (!hasSelectivePredicate && condition.deterministic && isLikelySelective(condition)) { + hasSelectivePredicate = true + } + case relation: LogicalRelation if relation.relation.isInstanceOf[HadoopFsRelation] => + val options = CaseInsensitiveMap(relation.relation.asInstanceOf[HadoopFsRelation].options) + fileSourceOptions += relevantOptions.flatMap { key => + options.get(key).map(key -> _) + }.toMap + case _ => + }) + CachedRDDBuilder( + serializer, + storageLevel, + cachedPlan, + tableName, + logicalPlan, + isCachedLogicalPlanRepeatable = repeatable, + hasSelectivePredicate = hasSelectivePredicate, + fileSourceOptions = fileSourceOptions.result()) + } private[this] var ser: Option[CachedBatchSerializer] = None private[this] def getSerializer(sqlConf: SQLConf): CachedBatchSerializer = synchronized { @@ -434,8 +602,8 @@ object InMemoryRelation { } else { qe.executedPlan } - val cacheBuilder = - CachedRDDBuilder(serializer, storageLevel, child, tableName, qe.logical) + val cacheBuilder = newCacheBuilder( + serializer, storageLevel, child, tableName, qe.logical, qe.analyzed, optimizedPlan) val relation = new InMemoryRelation(child.output, cacheBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats relation @@ -450,8 +618,8 @@ object InMemoryRelation { child: SparkPlan, tableName: Option[String], optimizedPlan: LogicalPlan): InMemoryRelation = { - val cacheBuilder = - CachedRDDBuilder(serializer, storageLevel, child, tableName, optimizedPlan) + val cacheBuilder = newCacheBuilder( + serializer, storageLevel, child, tableName, optimizedPlan, optimizedPlan, optimizedPlan) val relation = new InMemoryRelation(child.output, cacheBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats relation @@ -465,7 +633,14 @@ object InMemoryRelation { } else { qe.executedPlan } - val newBuilder = cacheBuilder.copy(cachedPlan = newCachedPlan, logicalPlan = qe.logical) + val newBuilder = newCacheBuilder( + serializer, + cacheBuilder.storageLevel, + newCachedPlan, + cacheBuilder.tableName, + qe.logical, + qe.analyzed, + optimizedPlan) val relation = new InMemoryRelation( newBuilder.cachedPlan.output, newBuilder, optimizedPlan.outputOrdering) relation.statsOfPlanToCache = optimizedPlan.stats @@ -487,7 +662,7 @@ case class InMemoryRelation( output: Seq[Attribute], @transient cacheBuilder: CachedRDDBuilder, override val outputOrdering: Seq[SortOrder]) - extends logical.LeafNode with MultiInstanceRelation { + extends logical.MaterializedLeafNode with MultiInstanceRelation { @volatile var statsOfPlanToCache: Statistics = null @@ -500,6 +675,17 @@ case class InMemoryRelation( def cachedPlan: SparkPlan = cacheBuilder.cachedPlan + override def mayHaveUsableMaterializedStats: Boolean = + cacheBuilder.storageLevel.useDisk && cacheBuilder.isCachedPlanRepeatable + + override def materializedMetadata: Option[logical.MaterializedLeafMetadata] = + cacheBuilder.materializedMetadata + + override def isOutputRepeatable: Boolean = + cacheBuilder.isCachedPlanRepeatable && materializedMetadata.exists(_.isOutputRepeatable) + + override def hasSelectivePredicate: Boolean = cacheBuilder.hasSelectivePredicate + private[sql] def updateStats( rowCount: Long, newColStats: Map[Attribute, ColumnStat]): Unit = this.synchronized { @@ -511,14 +697,11 @@ case class InMemoryRelation( } override def computeStats(): Statistics = { - if (!cacheBuilder.isCachedColumnBuffersLoaded) { + cacheBuilder.loadedMaterializedStats.map { case (rowCount, sizeInBytes) => + statsOfPlanToCache.copy(sizeInBytes = sizeInBytes, rowCount = Some(rowCount)) + }.getOrElse { // Underlying columnar RDD hasn't been materialized, use the stats from the plan to cache. statsOfPlanToCache - } else { - statsOfPlanToCache.copy( - sizeInBytes = cacheBuilder.materializedSizeInBytes, - rowCount = Some(cacheBuilder.materializedRowCount) - ) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala index 0fe1fbcf94501..e0f68a70daae6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/CompressionScheme.scala @@ -84,4 +84,10 @@ private[columnar] object CompressionScheme { // null count + null positions 4 + 4 * nullCount } + + def createNullsBuffer(buffer: ByteBuffer): ByteBuffer = { + val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) + nullsBuffer.rewind() + nullsBuffer + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala index 86d76856e12bb..72bebb96ed3ab 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/compression/compressionSchemes.scala @@ -18,7 +18,6 @@ package org.apache.spark.sql.execution.columnar.compression import java.nio.ByteBuffer -import java.nio.ByteOrder import scala.collection.mutable @@ -108,8 +107,7 @@ private[columnar] case object PassThrough extends CompressionScheme { capacity: Int, unitSize: Int, putFunction: (WritableColumnVector, Int, Int, Int) => Unit): Unit = { - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind() + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else capacity var pos = 0 @@ -307,8 +305,7 @@ private[columnar] case object RunLengthEncoding extends CompressionScheme { capacity: Int, getFunction: (ByteBuffer) => Long, putFunction: (WritableColumnVector, Int, Long) => Unit): Unit = { - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind() + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1 var pos = 0 @@ -486,8 +483,7 @@ private[columnar] case object DictionaryEncoding extends CompressionScheme { override def hasNext: Boolean = buffer.hasRemaining override def decompress(columnVector: WritableColumnVector, capacity: Int): Unit = { - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind() + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1 var pos = 0 @@ -619,8 +615,7 @@ private[columnar] case object BooleanBitSet extends CompressionScheme { override def decompress(columnVector: WritableColumnVector, capacity: Int): Unit = { var currentWordLocal: Long = 0 var visitedLocal: Int = 0 - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind() + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1 var pos = 0 @@ -730,8 +725,7 @@ private[columnar] case object IntDelta extends CompressionScheme { override def decompress(columnVector: WritableColumnVector, capacity: Int): Unit = { var prevLocal: Int = 0 - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind() + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1 var pos = 0 @@ -837,8 +831,7 @@ private[columnar] case object LongDelta extends CompressionScheme { override def decompress(columnVector: WritableColumnVector, capacity: Int): Unit = { var prevLocal: Long = 0 - val nullsBuffer = buffer.duplicate().order(ByteOrder.nativeOrder()) - nullsBuffer.rewind + val nullsBuffer = CompressionScheme.createNullsBuffer(buffer) val nullCount = ByteBufferHelper.getInt(nullsBuffer) var nextNullIndex = if (nullCount > 0) ByteBufferHelper.getInt(nullsBuffer) else -1 var pos = 0 diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/DataWritingCommand.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/DataWritingCommand.scala index 90050b25e9543..e5ff8cd5a564b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/DataWritingCommand.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/DataWritingCommand.scala @@ -103,14 +103,14 @@ object DataWritingCommand { } /** * When execute CTAS operators, and the location is not empty, throw [[AnalysisException]]. - * For CTAS, the SaveMode is always [[ErrorIfExists]] * * @param tablePath Table location. * @param saveMode Save mode of the table. * @param hadoopConf Configuration. */ def assertEmptyRootPath(tablePath: URI, saveMode: SaveMode, hadoopConf: Configuration): Unit = { - if (saveMode == SaveMode.ErrorIfExists && !SQLConf.get.allowNonEmptyLocationInCTAS) { + if ((saveMode == SaveMode.ErrorIfExists || saveMode == SaveMode.Ignore) && + !SQLConf.get.allowNonEmptyLocationInCTAS) { val filePath = new org.apache.hadoop.fs.Path(tablePath) val fs = filePath.getFileSystem(hadoopConf) if (fs.exists(filePath) && diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/tables.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/tables.scala index df050742ec666..4d513b4d3b06f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/tables.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/tables.scala @@ -654,15 +654,16 @@ case class DescribeTableCommand( } describeSchema(metadata.schema, result, header = false) } else { - if (metadata.schema.isEmpty) { + val schema = if (metadata.schema.isEmpty) { // In older version(prior to 2.1) of Spark, the table schema can be empty and should be // inferred at runtime. We should still support it. - describeSchema(sparkSession.table(metadata.identifier).schema, result, header = false) + sparkSession.table(metadata.identifier).schema } else { - describeSchema(metadata.schema, result, header = false) + metadata.schema } + describeSchema(schema, result, header = false) - describePartitionInfo(metadata, result) + describePartitionInfo(metadata, schema, result) describeClusteringInfo(metadata, result) if (partitionSpec.nonEmpty) { @@ -682,10 +683,29 @@ case class DescribeTableCommand( result.toSeq } - private def describePartitionInfo(table: CatalogTable, buffer: ArrayBuffer[Row]): Unit = { - if (table.partitionColumnNames.nonEmpty) { - append(buffer, "# Partition Information", "", "") - describeSchema(table.partitionSchema, buffer, header = true) + private def describePartitionInfo( + table: CatalogTable, + schema: StructType, + buffer: ArrayBuffer[Row]): Unit = { + val partitionColumnNames = table.partitionColumnNames + if (partitionColumnNames.nonEmpty) { + // Same positional convention as `CatalogTable.partitionSchema`, but reported instead of + // asserted so that a table with inconsistent metadata stays describable. + val partitionFields = schema.takeRight(partitionColumnNames.length) + val consistent = partitionFields.length == partitionColumnNames.length && + partitionFields.map(_.name).zip(partitionColumnNames).forall { + case (schemaColumn, partitionColumn) => conf.resolver(schemaColumn, partitionColumn) + } + if (consistent) { + append(buffer, "# Partition Information", "", "") + describeSchema(StructType(partitionFields), buffer, header = true) + } else { + append(buffer, "# Invalid Partition Information", "", "") + append(buffer, "Declared Partition Columns", + partitionColumnNames.mkString("[", ", ", "]"), "") + append(buffer, "Last Columns in Table Schema", + partitionFields.map(_.name).mkString("[", ", ", "]"), "") + } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala index 411682f35f6df..4338bd367fe08 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala @@ -599,9 +599,9 @@ object ViewHelper extends SQLConfHelper with Logging with CapturesConfig { plan.children.foreach(child => checkCyclicViewReference(child, path, viewIdent)) } - // Detect cyclic references from subqueries. + // Detect cyclic references from subqueries nested in expressions. plan.expressions.foreach { expr => - expr match { + expr.foreach { case s: SubqueryExpression => checkCyclicViewReference(s.plan, path, viewIdent) case _ => // Do nothing. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ApplyCharTypePadding.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ApplyCharTypePadding.scala index d952927f9d30a..a1cc576b2b236 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ApplyCharTypePadding.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/ApplyCharTypePadding.scala @@ -17,6 +17,9 @@ package org.apache.spark.sql.execution.datasources +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.spark.internal.LogKeys import org.apache.spark.sql.catalyst.analysis.ApplyCharTypePaddingHelper import org.apache.spark.sql.catalyst.catalog.HiveTableRelation import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan @@ -36,12 +39,28 @@ import org.apache.spark.sql.internal.SQLConf */ object ApplyCharTypePadding extends Rule[LogicalPlan] { + private val readSidePaddingOverrideWarned = new AtomicBoolean(false) + + private def warnReadSidePaddingOverride(): Unit = { + if (readSidePaddingOverrideWarned.compareAndSet(false, true)) { + logWarning(log"${MDC(LogKeys.CONFIG, SQLConf.READ_SIDE_CHAR_PADDING.key)} is disabled but " + + log"${MDC(LogKeys.CONFIG2, SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key)} is enabled; " + + log"read-side CHAR/VARCHAR checks are still applied because standard semantics require " + + log"a read to observe the value a write would have produced.") + } + } + override def apply(plan: LogicalPlan): LogicalPlan = { - if (conf.charVarcharAsString) { + // standardSemantics takes precedence over legacy charVarcharAsString. + if (conf.charVarcharAsString && !conf.charVarcharStandardSemantics) { return plan } - if (conf.readSideCharPadding) { + if (conf.charVarcharStandardSemantics && !conf.readSideCharPadding) { + warnReadSidePaddingOverride() + } + + if (conf.readSideCharPadding || conf.charVarcharStandardSemantics) { val newPlan = plan.resolveOperatorsUpWithNewOutput { case r: LogicalRelation => ApplyCharTypePaddingHelper.readSidePadding(r, () => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala index de642a3e850af..50d5dd28ab74e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala @@ -189,9 +189,13 @@ object FileFormatWriter extends Logging { if (writeFilesOpt.isDefined) { // build `WriteFilesSpec` for `WriteFiles` - val concurrentOutputWriterSpecFunc = (plan: SparkPlan) => { - val sortPlan = createSortPlan(plan, requiredOrdering, outputSpec) - createConcurrentOutputWriterSpec(sparkSession, sortPlan, sortColumns) + val concurrentOutputWriterSpecFunc = if (orderingMatched) { + (_: SparkPlan) => None + } else { + (plan: SparkPlan) => { + val sortPlan = createSortPlan(plan, requiredOrdering, outputSpec) + createConcurrentOutputWriterSpec(sparkSession, sortPlan, sortColumns) + } } val writeSpec = WriteFilesSpec( description = description, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala index b591573c00afe..ee6bf9c0abe90 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileScanRDD.scala @@ -89,6 +89,10 @@ class FileScanRDD( private val ignoreCorruptFiles = options.ignoreCorruptFiles private val ignoreMissingFiles = options.ignoreMissingFiles + + /** Whether this reader fails instead of silently skipping missing or corrupt input files. */ + private[sql] def hasStrictFileReads: Boolean = !ignoreCorruptFiles && !ignoreMissingFiles + // Evaluated on the driver (sparkSession is @transient) and serialized to executors so the // `compute` iterator below can pass it through to ColumnVectorUtils.populate. private val memoryMode: MemoryMode = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala index 6695e24ad6186..0ed27b12edb2f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PullOutVariantExtractions.scala @@ -107,18 +107,13 @@ import org.apache.spark.sql.types.VariantType * * Cast-error surface: relocating a strict extraction (`variant_get(..., failOnError = true)` or a * strict `Cast`) below a `Join` means it is evaluated at the scan on rows the join later eliminates - * -- so a cast failure can surface for a row the un-hoisted plan would never have cast (here the - * eliminating rows come from the *other* joined table). This is the same pre-existing trade-off as - * [[PushVariantIntoScan]] pushing casts below a `Filter`, not a new error class: in both, the - * strict cast runs before the operator that would have discarded the failing row. This rule only - * relocates the extraction into a `Project`; [[PushVariantIntoScan]] still does the scan-level - * materialization and, when [[SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR]] is set (default - * false), wraps the cast with a per-row cast-error companion slot so the error is only raised when - * the original expression consumes the failing row. That deferral is provenance-agnostic -- it acts - * on the relocated extraction regardless of how it reached the `Project` -- so enabling the flag - * suppresses the join-eliminated-row error exactly as it does the filter-eliminated-row case. With - * the flag off (the default), the strict cast raises immediately on any failing scanned row, as - * documented for below-`Filter` pushdown. + * -- so a cast failure can surface for a row the un-hoisted plan would never have cast. This is the + * same pre-existing trade-off as [[PushVariantIntoScan]] pushing casts below a `Filter`. When + * [[SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR]] is set, the cast error is deferred until the + * original expression consumes the failing row, suppressing errors from rows eliminated by either + * operator. With the flag off, strict `VariantGet` and `Cast` retain their existing behavior + * because they are not currently classified as [[Expression.throwable]]. If an extraction is + * classified as throwable, it crosses a join only when cast-error deferral is enabled. */ object PullOutVariantExtractions extends Rule[LogicalPlan] { @@ -177,11 +172,18 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { case _ => false } + private def isJoinHoistable(e: Expression): Boolean = { + // Unlike optimizer rules that must stop at throwable expressions, scan pushdown can preserve + // the original error timing with its per-row cast-error companion column when enabled. + isHoistable(e) && (!e.throwable || + SQLConf.get.getConf(SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR)) + } + /** * Collects hoisted extractions as aliases, de-duplicated by canonical form so a repeated * extraction maps to a single output slot. */ - private class ExtractionHoister { + private class ExtractionHoister(hoistable: Expression => Boolean = isHoistable) { private val extracted = mutable.LinkedHashMap.empty[Expression, Alias] def aliases: Seq[NamedExpression] = extracted.values.toSeq @@ -194,10 +196,25 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { /** Replaces every hoistable extraction in `e` with a reference to its (new) alias. */ def hoist(e: Expression): Expression = e.transformDown { - case ex if isHoistable(ex) => aliasFor(ex) + case ex if hoistable(ex) => aliasFor(ex) } } + // Mirrors the traversal in pushSideAliases for one extraction. A Project is pass-through only + // when the extraction resolves against its child; all other operators stop alias placement. + private def extractionCrossesJoin(child: LogicalPlan, e: Expression): Boolean = child match { + case _: Join => true + case Project(_, grandChild) if e.references.subsetOf(grandChild.outputSet) => + extractionCrossesJoin(grandChild, e) + case _ => false + } + + private def extractionHoister(child: LogicalPlan): ExtractionHoister = { + new ExtractionHoister(e => + if (extractionCrossesJoin(child, e)) isJoinHoistable(e) else isHoistable(e) + ) + } + // Recursively pushes hoisted extraction aliases down through a join tree until each lands in a // `Project` directly above a non-join child (the scan side, or the `Filter`/`Project` chain above // it that `PhysicalOperation` collapses). Hoisting an extraction into a `Project` above a `Join` @@ -299,7 +316,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { rightHoister: ExtractionHoister): Seq[NamedExpression] = { projectList.map { e => e.transformDown { - case ex if isHoistable(ex) => + case ex if isJoinHoistable(ex) => if (ex.references.subsetOf(leftOutput)) { leftHoister.aliasFor(ex) } else if (ex.references.subsetOf(rightOutput)) { @@ -312,7 +329,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { } private def rewriteAggregate(agg: Aggregate): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(agg.child) // Only hoist extractions that sit inside an aggregate function's arguments (or filter). A // top-level extraction in `aggregateExpressions` is a grouping-key reference (grouping keys are // already pulled out by `PullOutGroupingExpressions`); hoisting it would leave the `Aggregate` @@ -388,7 +405,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { private def rewriteSortUnderProject( project: Project, projectList: Seq[NamedExpression], sort: Sort): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(sort.child) val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder]) if (hoister.isEmpty) { project @@ -419,7 +436,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { // yielding a multi-slot shredded struct with a full-variant slot -- the same shape as a // Sort-under-Project whose `v` is also selected. See the class doc. private def rewriteBareSort(sort: Sort): LogicalPlan = { - val hoister = new ExtractionHoister + val hoister = extractionHoister(sort.child) val newOrder = sort.order.map(hoister.hoist(_).asInstanceOf[SortOrder]) if (hoister.isEmpty) { sort @@ -442,6 +459,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { project: Project, projectList: Seq[NamedExpression], join: Join): LogicalPlan = { val leftOutput = join.left.outputSet val rightOutput = join.right.outputSet + // Join-crossing eligibility is checked at the routing sites before aliasFor is called. val leftHoister = new ExtractionHoister val rightHoister = new ExtractionHoister @@ -451,7 +469,7 @@ object PullOutVariantExtractions extends Rule[LogicalPlan] { // join (its extractions -- e.g. aggregate arguments hoisted here by `rewriteAggregate`, or a // user's `SELECT variant_get(...)`), so the pushdown sees them below the join. val newCondition = join.condition.map(_.transformDown { - case ex if isHoistable(ex) => + case ex if isJoinHoistable(ex) => if (ex.references.subsetOf(leftOutput)) { leftHoister.aliasFor(ex) } else if (ex.references.subsetOf(rightOutput)) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormat.scala index 1c6c635820393..dbfcff3f1fa5d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormat.scala @@ -24,16 +24,17 @@ import java.util.Locale import java.util.regex.Pattern import java.util.zip.GZIPInputStream +import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal -import org.apache.commons.compress.archivers.{ArchiveEntry, ArchiveInputStream} -import org.apache.commons.compress.archivers.sevenz.{SevenZArchiveEntry, SevenZFile} +import org.apache.commons.compress.archivers.ArchiveEntry +import org.apache.commons.compress.archivers.sevenz.SevenZFile import org.apache.commons.compress.archivers.tar.TarArchiveInputStream -import org.apache.commons.compress.archivers.zip.ZipArchiveInputStream +import org.apache.commons.compress.archivers.zip.ZipFile import org.apache.commons.io.ByteOrderMark import org.apache.commons.io.input.{BOMInputStream, CloseShieldInputStream} import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FSDataInputStream, Path} +import org.apache.hadoop.fs.{FSDataInputStream, GlobPattern, Path} import org.apache.hadoop.io.Text import org.apache.hadoop.util.LineReader @@ -41,6 +42,7 @@ import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.internal.Logging import org.apache.spark.paths.SparkPath import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.util.{HadoopFSUtils, Utils} /** @@ -102,16 +104,18 @@ trait SupportsArchiveFormat extends Logging { * Reads an archive by unpacking each entry to a temp file and applying `readEntry`, for a * format that needs a complete file on disk (random access). * - * @param file the archive as a [[PartitionedFile]] - * @param conf Hadoop configuration used to open the archive - * @param tempPrefix prefix for the per-task temp dir the entries are unpacked into - * @param readEntry reads one unpacked entry file into rows + * @param file the archive as a [[PartitionedFile]] + * @param conf Hadoop configuration used to open the archive + * @param tempPrefix prefix for the per-task temp dir the entries are unpacked into + * @param archivePathFilter optional glob matched against the entry's full path + * @param readEntry reads one unpacked entry file into rows * @return iterator of rows across all entries */ protected def readLocalizedEntries( file: PartitionedFile, conf: Configuration, - tempPrefix: String)( + tempPrefix: String, + archivePathFilter: Option[GlobPattern])( readEntry: PartitionedFile => Iterator[InternalRow]): Iterator[InternalRow] = { val tempDir = Utils.createTempDir(Utils.getLocalDir(SparkEnv.get.conf), tempPrefix) // Register cleanup before constructing `entries`, which can throw before returning an iterator @@ -120,7 +124,8 @@ trait SupportsArchiveFormat extends Logging { Utils.deleteRecursively(tempDir) }) val entries = - try SupportsArchiveFormat.localizeEntries(file.toPath, conf, tempDir, archiveEntryFilter) + try SupportsArchiveFormat.localizeEntries( + file.toPath, conf, tempDir, archiveEntryFilter, archivePathFilter) catch { case NonFatal(e) => Utils.deleteRecursively(tempDir) @@ -203,28 +208,20 @@ object SupportsArchiveFormat { name.endsWith(".zip") || name.endsWith(".7z") } + /** An archive's entries as lazy `(entry, stream)` pairs; closing releases the container. */ + private type ArchiveEntries = Iterator[(ArchiveEntry, InputStream)] with Closeable + /** - * Opens the archive at `path` as a commons-compress stream, selecting the container by extension. + * Opens the archive at `path`, selecting the container by extension and exposing its entries as + * `(entry, stream)` pairs. */ - private def openArchiveStream( - path: Path, - conf: Configuration): ArchiveInputStream[_ <: ArchiveEntry] = { + private def openArchiveStream(path: Path, conf: Configuration): ArchiveEntries = { val name = path.getName.toLowerCase(Locale.ROOT) name match { case n if n.endsWith(".tar") || n.endsWith(".tar.gz") || n.endsWith(".tgz") => - val base = CodecStreams.createInputStreamWithCloseResource(conf, path) - try { - // GZIPInputStream reads the gzip header in its constructor, so a corrupt archive can - // throw here -- after `base` is already open -- and `base` must not leak. - val tarBytes = if (n.endsWith(".tgz")) new GZIPInputStream(base) else base - new TarArchiveInputStream(tarBytes) - } catch { - case NonFatal(e) => - try base.close() catch { case NonFatal(_) => } - throw e - } + openTarStream(path, conf) case n if n.endsWith(".zip") => - new ZipArchiveInputStream(CodecStreams.createInputStreamWithCloseResource(conf, path)) + openZipStream(path, conf) case n if n.endsWith(".7z") => openSevenZStream(path, conf) case _ => @@ -233,25 +230,111 @@ object SupportsArchiveFormat { } } - /** - * Opens a `.7z` archive by seeking. - * - * @param path the archive path - * @param conf Hadoop configuration used to open the archive - * @return the archive's entries as an [[ArchiveInputStream]] cursor - */ - private def openSevenZStream( - path: Path, - conf: Configuration): ArchiveInputStream[_ <: ArchiveEntry] = { + /** Opens a `.tar`/`.tar.gz`/`.tgz` archive, streaming its entries through one forward cursor. */ + private def openTarStream(path: Path, conf: Configuration): ArchiveEntries = { + val gzipped = path.getName.toLowerCase(Locale.ROOT).endsWith(".tgz") + val base = CodecStreams.createInputStreamWithCloseResource(conf, path) + try { + // GZIPInputStream reads the gzip header in its constructor, so a corrupt archive can throw + // here -- after `base` is already open -- and `base` must not leak. + val tar = new TarArchiveInputStream(if (gzipped) new GZIPInputStream(base) else base) + val entries = Iterator.continually(tar.getNextEntry).takeWhile(_ != null) + .map((_, tar: InputStream)) + closeable(entries, () => tar.close()) + } catch { + case NonFatal(e) => + try base.close() catch { case NonFatal(_) => } + throw e + } + } + + /** Pairs an entry iterator with the resource it reads from, closed when the caller is done. */ + private def closeable( + entries: Iterator[(ArchiveEntry, InputStream)], + closeFn: () => Unit): ArchiveEntries = + new Iterator[(ArchiveEntry, InputStream)] with Closeable { + override def hasNext: Boolean = entries.hasNext + override def next(): (ArchiveEntry, InputStream) = entries.next() + override def close(): Unit = closeFn() + } + + /** Opens a `.7z` archive by seeking. */ + private def openSevenZStream(path: Path, conf: Configuration): ArchiveEntries = { + val fs = path.getFileSystem(conf) + val length = fs.getFileStatus(path).getLen + var channel: SeekableByteChannel = null + var sevenZ: SevenZFile = null + try { + channel = new HadoopSeekableByteChannel(fs.open(path), length) + sevenZ = SevenZFile.builder().setSeekableByteChannel(channel).get() + // SevenZFile is a forward cursor: one stream serves whichever entry getNextEntry selects. + val entryStream = new SevenZEntryInputStream(sevenZ) + val entries = Iterator.continually(sevenZ.getNextEntry).takeWhile(_ != null) + .map((_, entryStream)) + closeable(entries, () => { + try { + sevenZ.close() + } finally { + channel.close() + } + }) + } catch { + case NonFatal(e) => + Utils.closeQuietly(sevenZ) + Utils.closeQuietly(channel) + throw e + } + } + + /** Opens a `.zip` archive by seeking, reading the central directory first. */ + private def openZipStream(path: Path, conf: Configuration): ArchiveEntries = { val fs = path.getFileSystem(conf) - val in = fs.open(path) + val length = fs.getFileStatus(path).getLen + var channel: SeekableByteChannel = null + var zipFile: ZipFile = null try { - val channel = new HadoopSeekableByteChannel(in, fs.getFileStatus(path).getLen) - new SevenZArchiveInputStream( - SevenZFile.builder().setSeekableByteChannel(channel).get(), channel) + channel = new HadoopSeekableByteChannel(fs.open(path), length) + zipFile = ZipFile.builder().setSeekableByteChannel(channel).get() + var current: InputStream = null + def closeCurrentStream(): Unit = if (current != null) { + try current.close() catch { case NonFatal(_) => } + current = null + } + // Open each entry's stream lazily, on first read. An entry the engine skips (directory or + // dotfile) is never read, so it is never opened and its readability never checked -- else a + // skipped-but-encrypted entry would throw CANNOT_READ_ZIP_ENTRY, unlike the streaming reader + // ZipFile replaced. Opening an entry releases the previous one's inflater. + val entries = zipFile.getEntries.asScala.map { entry => + val stream = new InputStream { + private var delegate: InputStream = _ + private def open(): InputStream = { + if (delegate == null) { + closeCurrentStream() + if (!zipFile.canReadEntryData(entry)) { + throw QueryExecutionErrors.cannotReadZipEntry(entry.getName, path.toString) + } + delegate = zipFile.getInputStream(entry) + current = delegate + } + delegate + } + override def read(): Int = open().read() + override def read(b: Array[Byte], off: Int, len: Int): Int = open().read(b, off, len) + } + (entry: ArchiveEntry, stream) + } + closeable(entries, () => { + try { + closeCurrentStream() + zipFile.close() + } finally { + channel.close() + } + }) } catch { case NonFatal(e) => - try in.close() catch { case NonFatal(_) => } + Utils.closeQuietly(zipFile) + Utils.closeQuietly(channel) throw e } } @@ -262,10 +345,16 @@ object SupportsArchiveFormat { * * @param entry the archive entry to test * @param ignoredPathSegmentRegex per-segment filter matched against each `/`-separated component - * @return true if the entry is a directory or any path component is filtered out + * @param archivePathFilter optional glob matched against the entry's full path + * @return true if the entry is a directory, any path component is filtered out, or the entry's + * path does not match `archivePathFilter` */ - private def shouldSkipEntry(entry: ArchiveEntry, ignoredPathSegmentRegex: Pattern): Boolean = { + private def shouldSkipEntry( + entry: ArchiveEntry, + ignoredPathSegmentRegex: Pattern, + archivePathFilter: Option[GlobPattern]): Boolean = { if (entry.isDirectory) return true + if (archivePathFilter.exists(!_.matches(entry.getName))) return true entry.getName.split("/").exists(c => c.nonEmpty && HadoopFSUtils.shouldFilterOutPathName(c, ignoredPathSegmentRegex)) } @@ -279,13 +368,15 @@ object SupportsArchiveFormat { * @param ignoredPathSegmentRegex per-segment filter for entries to skip (defaults to the * `InMemoryFileIndex` filter); pass a custom one to match a * loose-file scan + * @param archivePathFilter optional glob matched against the entry's full path * @param parseEntry turns one entry's `(entry, stream)` into an iterator of results * @return the concatenated results across kept entries, lazily one entry at a time */ def readArchiveEntries[T]( path: Path, conf: Configuration, - ignoredPathSegmentRegex: Pattern = HadoopFSUtils.defaultIgnoredPathSegmentRegexPattern)( + ignoredPathSegmentRegex: Pattern = HadoopFSUtils.defaultIgnoredPathSegmentRegexPattern, + archivePathFilter: Option[GlobPattern])( parseEntry: (ArchiveEntry, InputStream) => Iterator[T]): Iterator[T] = { val archive = openArchiveStream(path, conf) var closed = false @@ -312,17 +403,19 @@ object SupportsArchiveFormat { case c: Closeable => try c.close() catch { case NonFatal(_) => } case _ => } - var entry = archive.getNextEntry - while (entry != null && shouldSkipEntry(entry, ignoredPathSegmentRegex)) { - entry = archive.getNextEntry + var next: (ArchiveEntry, InputStream) = null + while (next == null && archive.hasNext) { + val entry = archive.next() + if (!shouldSkipEntry(entry._1, ignoredPathSegmentRegex, archivePathFilter)) { + next = entry + } } - if (entry == null) { + if (next == null) { done = true cleanup() } else { - // CloseShieldInputStream ignores close(), so a parser closing its input does not close - // the archive; any unread remainder is skipped by getNextEntry() when advancing. - currentIter = parseEntry(entry, CloseShieldInputStream.wrap(archive)) + // Parse the entry stream; any unread remainder is skipped when the archive advances. + currentIter = parseEntry(next._1, CloseShieldInputStream.wrap(next._2)) } } } @@ -380,17 +473,19 @@ object SupportsArchiveFormat { * companion so executor-side callers (a format's distributed archive inference) can use it * without a trait instance. * - * @param path the archive path - * @param conf Hadoop configuration used to open the archive - * @param localDir directory the per-entry temp files are created under - * @param entryFilter which entry names to keep + * @param path the archive path + * @param conf Hadoop configuration used to open the archive + * @param localDir directory the per-entry temp files are created under + * @param entryFilter which entry names to keep + * @param archivePathFilter optional glob matched against the entry's full path */ def localizeEntries( path: Path, conf: Configuration, localDir: File, - entryFilter: String => Boolean): Iterator[(String, File)] = - readArchiveEntries(path, conf) { (entry, in) => + entryFilter: String => Boolean, + archivePathFilter: Option[GlobPattern]): Iterator[(String, File)] = + readArchiveEntries(path, conf, archivePathFilter = archivePathFilter) { (entry, in) => val name = entry.getName if (entryFilter(name)) { Iterator.single((name, copyEntryToLocalFile(in, localDir, name))) @@ -447,18 +542,3 @@ private class SevenZEntryInputStream(sevenZ: SevenZFile) extends InputStream { override def read(): Int = sevenZ.read() override def read(b: Array[Byte], off: Int, len: Int): Int = sevenZ.read(b, off, len) } - -/** - * Adapts a [[SevenZFile]] to the [[ArchiveInputStream]] cursor the engine consumes. Closing it - * closes both the `SevenZFile` and the channel it reads from, which `SevenZFile` does not own. - * - * @param sevenZ the 7z file to adapt - * @param channel the channel `sevenZ` reads from, closed alongside it - */ -private class SevenZArchiveInputStream(sevenZ: SevenZFile, channel: SeekableByteChannel) - extends ArchiveInputStream[SevenZArchiveEntry](new SevenZEntryInputStream(sevenZ), "UTF-8") { - - override def getNextEntry(): SevenZArchiveEntry = sevenZ.getNextEntry - - override def close(): Unit = try sevenZ.close() finally channel.close() -} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala index 4493d1a6e6895..abd30e9c65a4d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala @@ -190,7 +190,7 @@ object V1WritesUtils { val partitionSet = AttributeSet(partitionColumns) var needConvert = false val projectList: Seq[NamedExpression] = output.map { - case p if partitionSet.contains(p) && p.dataType == StringType && p.nullable => + case p if partitionSet.contains(p) && p.dataType.isInstanceOf[StringType] && p.nullable => needConvert = true Alias(Empty2Null(p), p.name)() case attr => attr diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala index 33de63072ef39..c0e50ee02e522 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala @@ -113,16 +113,17 @@ case class BinaryFileFormat() extends FileFormat val caseInsensitiveOptions = CaseInsensitiveMap(options) val archiveReadEnabled = !caseInsensitiveOptions.get(WHOLE_FILE).forall(_.toBoolean) && getSqlConf(sparkSession).getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED) + val fileSourceOptions = new FileSourceOptions(caseInsensitiveOptions) file: PartitionedFile => { val path = file.toPath val fs = path.getFileSystem(broadcastedHadoopConf.value.value) val status = fs.getFileStatus(path) if (archiveReadEnabled && SupportsArchiveFormat.isArchivePath(path)) { - val ignoredPathSegmentRegex = - new FileSourceOptions(caseInsensitiveOptions).ignoredPathSegmentRegexPattern SupportsArchiveFormat.readArchiveEntries( - path, broadcastedHadoopConf.value.value, ignoredPathSegmentRegex) { (entry, in) => + path, broadcastedHadoopConf.value.value, + fileSourceOptions.ignoredPathSegmentRegexPattern, + fileSourceOptions.archivePathFilterPattern) { (entry, in) => val entryStatus = new FileStatus( entry.getSize, false, 0, 0, status.getModificationTime, new Path(s"${status.getPath}!/${entry.getName}")) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVDataSource.scala index 2eb0ad44f4723..d2075a3304dfb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVDataSource.scala @@ -25,7 +25,7 @@ import scala.util.control.NonFatal import com.univocity.parsers.csv.CsvParser import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.mapreduce.Job import org.apache.hadoop.mapreduce.lib.input.FileInputFormat @@ -36,7 +36,7 @@ import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.paths.SparkPath import org.apache.spark.rdd.{BinaryFileRDD, RDD} import org.apache.spark.sql.{Dataset, Encoders, SparkSession} -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.csv.{CSVHeaderChecker, CSVInferSchema, CSVOptions, UnivocityParser} import org.apache.spark.sql.classic.ClassicConversions.castToImpl import org.apache.spark.sql.errors.QueryExecutionErrors @@ -76,14 +76,10 @@ abstract class CSVDataSource extends Serializable with Logging with SupportsArch val hasArchive = parsedOptions.archiveFormatEnabled && inputPaths.exists(f => SupportsArchiveFormat.isArchivePath(f.getPath)) if (hasArchive && supportsArchiveScan) { - // Archives (and any loose files alongside them) are inferred in a single CSVInferSchema - // pass over all inputs -- archive entries are streamed, never unpacked -- so the result - // matches what the scan returns for the same files. Some(inferWithArchives(sparkSession, inputPaths, parsedOptions)) } else if (hasArchive) { - // The caller's scan path cannot read archives (e.g. the DSv2 reader), so refuse to infer - // a schema when any input is an archive: returning None raises UNABLE_TO_INFER_SCHEMA, - // which fails loudly instead of letting the scan parse raw archive bytes as CSV. + // The caller's scan cannot read archives (e.g. DSv2), so refuse rather than let it + // mis-read raw archive bytes as CSV. None } else if (inputPaths.nonEmpty) { Some(infer(sparkSession, inputPaths, parsedOptions)) @@ -122,6 +118,7 @@ abstract class CSVDataSource extends Serializable with Logging with SupportsArch * @param getHeaderChecker builds a fresh [[CSVHeaderChecker]] for `(isStartOfFile, source)`. * @param ignoredPathSegmentRegex the compiled effective `ignoredPathSegmentRegex` option, so * hidden entries are skipped exactly like Spark's file listing would. + * @param archivePathFilter optional glob matched against the entry's full path */ def readArchive( conf: Configuration, @@ -129,7 +126,8 @@ abstract class CSVDataSource extends Serializable with Logging with SupportsArch getParser: () => UnivocityParser, getHeaderChecker: (Boolean, String) => CSVHeaderChecker, requiredSchema: StructType, - ignoredPathSegmentRegex: Pattern): Iterator[InternalRow] + ignoredPathSegmentRegex: Pattern, + archivePathFilter: Option[GlobPattern]): Iterator[InternalRow] /** * Shared driver used by the [[readArchive]] implementations: streams each non-skipped entry's @@ -142,11 +140,12 @@ abstract class CSVDataSource extends Serializable with Logging with SupportsArch file: PartitionedFile, getParser: () => UnivocityParser, getHeaderChecker: (Boolean, String) => CSVHeaderChecker, - ignoredPathSegmentRegex: Pattern)( + ignoredPathSegmentRegex: Pattern, + archivePathFilter: Option[GlobPattern])( parseEntry: (UnivocityParser, CSVHeaderChecker, InputStream) => Iterator[InternalRow]) : Iterator[InternalRow] = { SupportsArchiveFormat.readArchiveEntries( - file.toPath, conf, ignoredPathSegmentRegex) { (entry, in) => + file.toPath, conf, ignoredPathSegmentRegex, archivePathFilter) { (entry, in) => val headerChecker = getHeaderChecker(true, s"CSV archive entry: ${file.urlEncodedPath}!/${entry.getName}") val parser = getParser() @@ -172,29 +171,41 @@ abstract class CSVDataSource extends Serializable with Logging with SupportsArch inputPaths: Seq[FileStatus], parsedOptions: CSVOptions): StructType = { val baseRdd = CSVDataSource.createBaseRdd(sparkSession, inputPaths, parsedOptions) - def tokens(dropHeader: Boolean): RDD[Array[String]] = baseRdd.flatMap { stream => - val path = new Path(stream.getPath()) - try { - if (SupportsArchiveFormat.isArchivePath(path)) { - SupportsArchiveFormat.readArchiveEntries(path, stream.getConfiguration) { (_, in) => - tokenizeForInference(in, dropHeader, parsedOptions) + // Inference must see the same entries the scan reads, so it honors archivePathFilter too. + // Capture the glob string: the compiled GlobPattern is not serializable, so each task + // compiles it once when the archive branch is taken. + val archivePathFilterGlob = parsedOptions.archivePathFilter + def tokens(dropHeader: Boolean): RDD[Array[String]] = baseRdd.mapPartitions { streams => + // Compile at most once per partition: lazy so a partition of only loose files never + // compiles, while a partition with archives reuses one matcher across all of them. + lazy val archivePathFilter = + archivePathFilterGlob.map(FileSourceOptions.compileArchivePathFilter) + streams.flatMap { stream => + val path = new Path(stream.getPath()) + try { + if (SupportsArchiveFormat.isArchivePath(path)) { + SupportsArchiveFormat.readArchiveEntries( + path, stream.getConfiguration, archivePathFilter = archivePathFilter) { + (_, in) => + tokenizeForInference(in, dropHeader, parsedOptions) + } + } else { + tokenizeForInference( + CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path), + dropHeader, parsedOptions) } - } else { - tokenizeForInference( - CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path), - dropHeader, parsedOptions) + } catch { + case e: FileNotFoundException if parsedOptions.ignoreMissingFiles => + logWarning(log"Skipped missing input: ${MDC(PATH, stream.getPath())}", e) + Iterator.empty + case e: FileNotFoundException => throw e + case e @ (_: RuntimeException | _: IOException) if parsedOptions.ignoreCorruptFiles => + logWarning(log"Skipped the corrupted input: ${MDC(PATH, stream.getPath())}", e) + Iterator.empty + case NonFatal(e) => + throw QueryExecutionErrors.cannotReadFilesError( + e, SparkPath.fromPathString(stream.getPath()).urlEncoded) } - } catch { - case e: FileNotFoundException if parsedOptions.ignoreMissingFiles => - logWarning(log"Skipped missing input: ${MDC(PATH, stream.getPath())}", e) - Iterator.empty - case e: FileNotFoundException => throw e - case e @ (_: RuntimeException | _: IOException) if parsedOptions.ignoreCorruptFiles => - logWarning(log"Skipped the corrupted input: ${MDC(PATH, stream.getPath())}", e) - Iterator.empty - case NonFatal(e) => - throw QueryExecutionErrors.cannotReadFilesError( - e, SparkPath.fromPathString(stream.getPath()).urlEncoded) } } tokens(dropHeader = false).take(1).headOption match { @@ -303,10 +314,12 @@ object TextInputCSVDataSource extends CSVDataSource { getParser: () => UnivocityParser, getHeaderChecker: (Boolean, String) => CSVHeaderChecker, requiredSchema: StructType, - ignoredPathSegmentRegex: Pattern): Iterator[InternalRow] = + ignoredPathSegmentRegex: Pattern, + archivePathFilter: Option[GlobPattern]): Iterator[InternalRow] = // Stream each tar entry through the line-based parser, treating the entry exactly like a // standalone CSV file (a fresh parser/header checker is built per entry). - streamArchiveEntries(conf, file, getParser, getHeaderChecker, ignoredPathSegmentRegex) { + streamArchiveEntries( + conf, file, getParser, getHeaderChecker, ignoredPathSegmentRegex, archivePathFilter) { (parser, headerChecker, in) => UnivocityParser.parseIterator( entryLines(in, parser.options), parser, headerChecker, requiredSchema) @@ -432,10 +445,12 @@ object MultiLineCSVDataSource extends CSVDataSource { getParser: () => UnivocityParser, getHeaderChecker: (Boolean, String) => CSVHeaderChecker, requiredSchema: StructType, - ignoredPathSegmentRegex: Pattern): Iterator[InternalRow] = + ignoredPathSegmentRegex: Pattern, + archivePathFilter: Option[GlobPattern]): Iterator[InternalRow] = // Stream each tar entry whole through the multi-line parser (a fresh parser/header checker is // built per entry). - streamArchiveEntries(conf, file, getParser, getHeaderChecker, ignoredPathSegmentRegex) { + streamArchiveEntries( + conf, file, getParser, getHeaderChecker, ignoredPathSegmentRegex, archivePathFilter) { (parser, headerChecker, in) => UnivocityParser.parseStream(in, parser, headerChecker, requiredSchema) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVFileFormat.scala index a406fbe4844de..388030dcf42df 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/CSVFileFormat.scala @@ -57,8 +57,6 @@ case class CSVFileFormat() extends TextBasedFileFormat with DataSourceRegister { options: Map[String, String], files: Seq[FileStatus]): Option[StructType] = { val parsedOptions = getCsvOptions(sparkSession, options) - // The v1 file format routes archives to `readArchive` (see `buildReader`), so archive schema - // inference is supported here. CSVDataSource(parsedOptions) .inferSchema(sparkSession, files, parsedOptions, supportsArchiveScan = true) } @@ -144,7 +142,8 @@ case class CSVFileFormat() extends TextBasedFileFormat with DataSourceRegister { // archive reads are enabled; otherwise the file is parsed directly. if (parsedOptions.archiveFormatEnabled && SupportsArchiveFormat.isArchivePath(file.toPath)) { CSVDataSource(parsedOptions).readArchive( - conf, file, () => newParser(), getHeaderChecker, requiredSchema, ignoredPathSegmentRegex) + conf, file, () => newParser(), getHeaderChecker, requiredSchema, ignoredPathSegmentRegex, + parsedOptions.archivePathFilterPattern) } else { val parser = newParser() val headerChecker = getHeaderChecker(file.start == 0, s"CSV file: ${file.urlEncodedPath}") diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCOptions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCOptions.scala index 7188c2b8b2e8e..e4bd4df8cf2f8 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCOptions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCOptions.scala @@ -255,10 +255,21 @@ class JDBCOptions( .map(_.toBoolean) .getOrElse(SQLConf.get.timestampType == TimestampNTZType) + // Infers driver TIMESTAMP columns that report a sub-microsecond fractional-second scale (7-9) + // as the nanosecond-capable timestamp types TIMESTAMP_NTZ(p) / TIMESTAMP_LTZ(p). When disabled + // (the default) such columns keep the historical microsecond mapping. Only takes effect when the + // `spark.sql.timestampNanosTypes.enabled` preview flag is on. + val preferTimestampNanos = + parameters + .get(JDBC_PREFER_TIMESTAMP_NANOS) + .map(_.toBoolean) + .getOrElse(false) + val hint = parameters.get(JDBC_HINT_STRING).map(value => { require(value.matches("(?s)^/\\*\\+ .* \\*/$"), s"Invalid value `$value` for option `$JDBC_HINT_STRING`." + - s" It should start with `/*+ ` and end with ` */`.") + s" It should start with `/*+ ` and end with ` */`," + + s" for example `/*+ INDEX(t1 id_idx) */`.") s"$value " }).getOrElse("") @@ -366,5 +377,6 @@ object JDBCOptions { val JDBC_CONNECTION_PROVIDER = newOption("connectionProvider") val JDBC_PREPARE_QUERY = newOption("prepareQuery") val JDBC_PREFER_TIMESTAMP_NTZ = newOption("preferTimestampNTZ") + val JDBC_PREFER_TIMESTAMP_NANOS = newOption("preferTimestampNanos") val JDBC_HINT_STRING = newOption("hint") } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala index 2989c0975143f..6636db6e70bb0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRDD.scala @@ -106,7 +106,8 @@ object JDBCRDD extends Logging { statement.setQueryTimeout(options.queryTimeout) Using.resource(statement.executeQuery()) { rs => JdbcUtils.getSchema(conn, rs, dialect, alwaysNullable = true, - isTimestampNTZ = options.preferTimestampNTZ) + isTimestampNTZ = options.preferTimestampNTZ, + preferTimestampNanos = options.preferTimestampNanos) } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala index 972bb3e35ee6f..1c529279744e0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCRelation.scala @@ -27,14 +27,14 @@ import org.apache.spark.rdd.RDD import org.apache.spark.sql.{DataFrame, Row, SaveMode, SparkSession, SQLContext} import org.apache.spark.sql.catalyst.analysis._ import org.apache.spark.sql.catalyst.util.{DateFormatter, DateTimeUtils, TimestampFormatter} -import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp} +import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp, stringToTimestampWithoutTimeZone} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.jdbc.JdbcDialects import org.apache.spark.sql.sources._ -import org.apache.spark.sql.types.{DataType, DateType, NumericType, StructType, TimestampType} +import org.apache.spark.sql.types.{DataType, DateType, NumericType, StructType, TimestampNTZType, TimestampType} import org.apache.spark.unsafe.types.UTF8String /** @@ -188,7 +188,7 @@ private[sql] object JDBCRelation extends Logging { columnName, schema.simpleString(maxNumToStringFields)) } column.dataType match { - case _: NumericType | DateType | TimestampType => + case _: NumericType | DateType | TimestampType | TimestampNTZType => case _ => throw QueryCompilationErrors.invalidPartitionColumnTypeError(column) } @@ -209,6 +209,7 @@ private[sql] object JDBCRelation extends Logging { case _: NumericType => value.toLong case DateType => parse(stringToDate).toLong case TimestampType => parse(stringToTimestamp(_, getZoneId(timeZoneId))) + case TimestampNTZType => parse(stringToTimestampWithoutTimeZone(_, allowTimeZone = false)) } } @@ -224,12 +225,17 @@ private[sql] object JDBCRelation extends Logging { val timestampFormatter = TimestampFormatter.getFractionFormatter( DateTimeUtils.getZoneId(timeZoneId)) timestampFormatter.format(value) + case TimestampNTZType => + // NTZ micros are zoneless wall-clock values; format in UTC so no zone shift is applied. + val timestampFormatter = TimestampFormatter.getFractionFormatter( + DateTimeUtils.getZoneId("UTC")) + timestampFormatter.format(value) } s"'$dateTimeStr'" } columnType match { case _: NumericType => value.toString - case DateType | TimestampType => dateTimeToString() + case DateType | TimestampType | TimestampNTZType => dateTimeToString() } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCValueGetter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCValueGetter.scala index bd8d04760da65..3461bcf493e4f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCValueGetter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JDBCValueGetter.scala @@ -22,13 +22,13 @@ import java.nio.charset.StandardCharsets import java.sql.{Date, ResultSet, Time, Timestamp} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.util.DateTimeConstants.MICROS_PER_MILLIS +import org.apache.spark.sql.catalyst.util.DateTimeConstants.{MICROS_PER_MILLIS, NANOS_PER_MICROS} import org.apache.spark.sql.catalyst.util.DateTimeUtils._ import org.apache.spark.sql.catalyst.util.GenericArrayData import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.jdbc.JdbcDialect import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.UTF8String +import org.apache.spark.unsafe.types.{TimestampNanosVal, UTF8String} /** * A `JDBCValueGetter` is responsible for getting a value from a `ResultSet` into a field of an @@ -189,6 +189,44 @@ private[jdbc] object JDBCValueGetter { } } + // Reads a driver TIMESTAMP into a nanosecond-precision local date-time. Uses + // `getObject(LocalDateTime)` to fetch the stored wall-clock directly (mirroring `TimeGetter`), + // which preserves the full sub-microsecond fraction and, being time-zone independent, avoids the + // JVM-default-zone shift of the microsecond `getTimestamp` path. The value is then floored to the + // column precision. + final case class TimestampNTZNanosGetter(precision: Int) extends JDBCValueGetter { + def apply(rs: ResultSet, row: InternalRow, pos: Int): Unit = { + val localDateTime = rs.getObject(pos + 1, classOf[java.time.LocalDateTime]) + if (localDateTime != null) { + val fullNanos = + localDateTimeToTimestampNanos(localDateTime, TimestampNTZNanosType.NANOS_PRECISION) + row.update(pos, truncateTimestampNanosToPrecision(fullNanos, precision)) + } else { + row.update(pos, null) + } + } + } + + // Reads a driver TIMESTAMP into a nanosecond-precision instant. Mirrors the microsecond + // `TimestampGetter` (including the Julian->Gregorian rebase in `fromJavaTimestamp`), then + // re-attaches the sub-microsecond digits from `java.sql.Timestamp.getNanos` that the micro path + // drops, before truncating to the column precision. + final case class TimestampLTZNanosGetter(dialect: JdbcDialect, precision: Int) + extends JDBCValueGetter { + def apply(rs: ResultSet, row: InternalRow, pos: Int): Unit = { + val t = rs.getTimestamp(pos + 1) + if (t != null) { + val converted = dialect.convertJavaTimestampToTimestamp(t) + val epochMicros = fromJavaTimestamp(converted) + val subMicroNanos = (converted.getNanos % NANOS_PER_MICROS).toShort + val fullNanos = TimestampNanosVal.fromParts(epochMicros, subMicroNanos) + row.update(pos, truncateTimestampNanosToPrecision(fullNanos, precision)) + } else { + row.update(pos, null) + } + } + } + case object BinaryBitGetter extends JDBCValueGetter { def apply(rs: ResultSet, row: InternalRow, pos: Int): Unit = { val bytes = rs.getBytes(pos + 1) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala index ca44e8b710b13..ccec597e6e0b1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala @@ -163,6 +163,12 @@ object JdbcUtils extends Logging with SQLConfHelper { case TimestampNTZType => Option(JdbcType("TIMESTAMP", java.sql.Types.TIMESTAMP)) case DateType => Option(JdbcType("DATE", java.sql.Types.DATE)) case t: TimeType => Option(JdbcType(s"TIME(${t.precision})", java.sql.Types.TIME)) + // Nanosecond-capable timestamps (precision 7-9) map to SQL TIMESTAMP(p). Dialects may + // override this, e.g. to emit TIMESTAMP(p) WITH TIME ZONE for the LTZ variant. + case t: TimestampNTZNanosType => + Option(JdbcType(s"TIMESTAMP(${t.precision})", java.sql.Types.TIMESTAMP)) + case t: TimestampLTZNanosType => + Option(JdbcType(s"TIMESTAMP(${t.precision})", java.sql.Types.TIMESTAMP)) case t: DecimalType => Option( JdbcType(s"DECIMAL(${t.precision},${t.scale})", java.sql.Types.DECIMAL)) case _ => None @@ -178,6 +184,10 @@ object JdbcUtils extends Logging with SQLConfHelper { if (isTimestampNTZ) TimestampNTZType else TimestampType } + private def getTimestampNanosType(isTimestampNTZ: Boolean, precision: Int): DataType = { + if (isTimestampNTZ) TimestampNTZNanosType(precision) else TimestampLTZNanosType(precision) + } + /** * Maps a JDBC type to a Catalyst type. This function is called only when * the JdbcDialect class corresponding to your database driver returns null. @@ -191,7 +201,8 @@ object JdbcUtils extends Logging with SQLConfHelper { precision: Int, scale: Int, signed: Boolean, - isTimestampNTZ: Boolean): DataType = sqlType match { + isTimestampNTZ: Boolean, + preferTimestampNanos: Boolean = false): DataType = sqlType match { case java.sql.Types.BIGINT => if (signed) LongType else DecimalType(20, 0) case java.sql.Types.BINARY => BinaryType case java.sql.Types.BIT => BooleanType // @see JdbcDialect for quirks @@ -233,7 +244,16 @@ object JdbcUtils extends Logging with SQLConfHelper { else TimeType.DEFAULT_PRECISION TimeType(timePrecision) } else getTimestampType(isTimestampNTZ) - case java.sql.Types.TIMESTAMP => getTimestampType(isTimestampNTZ) + case java.sql.Types.TIMESTAMP => + // When nanosecond timestamps are requested (and the preview feature is enabled), a driver + // TIMESTAMP that reports a sub-microsecond fractional-second scale (7-9) is mapped to the + // nanosecond-capable type. Otherwise the historical microsecond mapping is preserved. + if (preferTimestampNanos && + scale >= TimestampNTZNanosType.MIN_PRECISION && + scale <= TimestampNTZNanosType.MAX_PRECISION && + conf.timestampNanosTypesEnabled) { + getTimestampNanosType(isTimestampNTZ, scale) + } else getTimestampType(isTimestampNTZ) case java.sql.Types.TINYINT => IntegerType case java.sql.Types.VARBINARY => BinaryType case java.sql.Types.VARCHAR if conf.charVarcharAsString => StringType @@ -262,7 +282,8 @@ object JdbcUtils extends Logging with SQLConfHelper { try { statement.setQueryTimeout(options.queryTimeout) Some(getSchema(conn, statement.executeQuery(), dialect, - isTimestampNTZ = options.preferTimestampNTZ)) + isTimestampNTZ = options.preferTimestampNTZ, + preferTimestampNanos = options.preferTimestampNanos)) } catch { case _: SQLException => None } finally { @@ -286,7 +307,8 @@ object JdbcUtils extends Logging with SQLConfHelper { resultSet: ResultSet, dialect: JdbcDialect, alwaysNullable: Boolean = false, - isTimestampNTZ: Boolean = false): StructType = { + isTimestampNTZ: Boolean = false, + preferTimestampNanos: Boolean = false): StructType = { val rsmd = resultSet.getMetaData val ncols = rsmd.getColumnCount val fields = new Array[StructField](ncols) @@ -326,13 +348,15 @@ object JdbcUtils extends Logging with SQLConfHelper { } metadata.putBoolean("isSigned", isSigned) metadata.putBoolean("isTimestampNTZ", isTimestampNTZ) + metadata.putBoolean("preferTimestampNanos", preferTimestampNanos) metadata.putLong("scale", fieldScale) metadata.putString("jdbcClientType", typeName) dialect.updateExtraColumnMeta(conn, rsmd, i + 1, metadata) val columnType = dialect.getCatalystType(dataType, typeName, fieldSize, metadata).getOrElse( - getCatalystType(dataType, typeName, fieldSize, fieldScale, isSigned, isTimestampNTZ)) + getCatalystType(dataType, typeName, fieldSize, fieldScale, isSigned, isTimestampNTZ, + preferTimestampNanos)) fields(i) = StructField(columnName, columnType, nullable, metadata.build()) i = i + 1 } @@ -448,6 +472,8 @@ object JdbcUtils extends Logging with SQLConfHelper { case TimestampNTZType if metadata.contains("logical_time_type") => JDBCValueGetter.LogicalTimeNTZGetter(dialect) case TimestampNTZType => JDBCValueGetter.TimestampNTZGetter(dialect) + case t: TimestampNTZNanosType => JDBCValueGetter.TimestampNTZNanosGetter(t.precision) + case t: TimestampLTZNanosType => JDBCValueGetter.TimestampLTZNanosGetter(dialect, t.precision) case BinaryType if metadata.contains("binarylong") => JDBCValueGetter.BinaryBitGetter case BinaryType => JDBCValueGetter.BytesGetter case _: YearMonthIntervalType => JDBCValueGetter.YearMonthIntervalGetter(dialect) @@ -518,6 +544,23 @@ object JdbcUtils extends Logging with SQLConfHelper { stmt.setTimestamp(pos + 1, dialect.convertTimestampNTZToJavaTimestamp(row.getAs[java.time.LocalDateTime](pos))) + // Nanosecond-precision timestamps are always materialized as java.time values by the Nanos + // encoders (independent of the datetimeJava8Api flag). + // NTZ is time-zone independent: write the LocalDateTime wall-clock directly (mirroring the + // TimeType setter) so the full sub-microsecond fraction survives without a zone shift. + case _: TimestampNTZNanosType => + (stmt: PreparedStatement, row: Row, pos: Int) => + stmt.setObject(pos + 1, row.getAs[java.time.LocalDateTime](pos)) + + // LTZ is an absolute instant: go through the micro java.sql.Timestamp path (matching the + // TimestampType setter), then restore the full nanosecond-of-second the micro path drops. + case _: TimestampLTZNanosType => + (stmt: PreparedStatement, row: Row, pos: Int) => + val instant = row.getAs[Instant](pos) + val ts = toJavaTimestamp(instantToMicros(instant)) + ts.setNanos(instant.getNano) + stmt.setTimestamp(pos + 1, ts) + case DateType => if (conf.datetimeJava8ApiEnabled) { (stmt: PreparedStatement, row: Row, pos: Int) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonDataSource.scala index a5f19ae353d42..14e23a5fbdeec 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonDataSource.scala @@ -17,14 +17,13 @@ package org.apache.spark.sql.execution.datasources.json -import java.io.{ByteArrayInputStream, FileNotFoundException, InputStream, IOException} +import java.io.{ByteArrayInputStream, Closeable, FileNotFoundException, InputStream, IOException} -import scala.reflect.ClassTag import scala.util.control.NonFatal import com.fasterxml.jackson.core.{JsonFactory, JsonParser} import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.io.Text import org.apache.hadoop.mapreduce.Job import org.apache.hadoop.mapreduce.lib.input.FileInputFormat @@ -36,7 +35,7 @@ import org.apache.spark.internal.LogKeys.PATH import org.apache.spark.paths.SparkPath import org.apache.spark.rdd.{BinaryFileRDD, RDD} import org.apache.spark.sql.{Dataset, Encoders, SparkSession} -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.json.{CreateJacksonParser, JacksonParser, JsonInferSchema, JSONOptions} import org.apache.spark.sql.catalyst.util.FailureSafeParser import org.apache.spark.sql.classic.ClassicConversions.castToImpl @@ -85,13 +84,16 @@ abstract class JsonDataSource extends Serializable with Logging with SupportsArc * and is intentionally left untouched. * * @param parser builds a fresh JSON parser for each entry. + * @param archivePathFilter optional glob matched against the entry's full path */ def readArchive( conf: Configuration, file: PartitionedFile, parser: () => JacksonParser, - schema: StructType): Iterator[InternalRow] = - SupportsArchiveFormat.readArchiveEntries(file.toPath, conf) { (_, in) => + schema: StructType, + archivePathFilter: Option[GlobPattern]): Iterator[InternalRow] = + SupportsArchiveFormat.readArchiveEntries( + file.toPath, conf, archivePathFilter = archivePathFilter) { (_, in) => readStream(in, parser(), schema) } @@ -105,15 +107,9 @@ abstract class JsonDataSource extends Serializable with Logging with SupportsArc case None => val hasArchive = parsedOptions.archiveFormatEnabled && inputPaths.exists(f => SupportsArchiveFormat.isArchivePath(f.getPath)) - if (hasArchive && supportsArchiveScan) { - // Archives (and any loose files alongside them) are inferred in a single JsonInferSchema - // pass over all inputs -- archive entries are streamed, never unpacked -- so the result - // matches what the scan returns for the same files. - Some(inferWithArchives(sparkSession, inputPaths, parsedOptions)) - } else if (hasArchive) { - // The caller's scan path cannot read archives (e.g. the DSv2 reader), so refuse to infer - // a schema when any input is an archive: returning None raises UNABLE_TO_INFER_SCHEMA, - // which fails loudly instead of letting the scan parse raw archive bytes as JSON. + if (hasArchive && !supportsArchiveScan) { + // The caller's scan cannot read archives (e.g. DSv2), so refuse rather than let it + // mis-read raw archive bytes as JSON. None } else if (inputPaths.nonEmpty) { Some(infer(sparkSession, inputPaths, parsedOptions)) @@ -127,92 +123,6 @@ abstract class JsonDataSource extends Serializable with Logging with SupportsArc sparkSession: SparkSession, inputPaths: Seq[FileStatus], parsedOptions: JSONOptions): StructType - - /** - * Infers a JSON schema when at least one input is a tar archive. Every archive entry (streamed - * via `SupportsArchiveFormat`, never unpacked to disk) and every loose file is read as JSON - * records -- each line is a record, or the whole input is one document in multi-line mode -- and - * all of them feed a single [[JsonInferSchema]] pass, exactly as a directory of the same files - * would infer. - * Because [[JsonInferSchema]] already merges every record's type by field name across all inputs, - * one pass is itself the union: a field empty in one input but typed in another widens to the - * real type, and a `NullType` field survives to the single final canonicalization rather than - * being collapsed per-input. A corrupt/missing input is skipped as a unit (a whole archive or a - * whole file) when `ignoreCorruptFiles`/`ignoreMissingFiles` are set. - */ - private def inferWithArchives( - sparkSession: SparkSession, - inputPaths: Seq[FileStatus], - parsedOptions: JSONOptions): StructType = { - val baseRdd = JsonDataSource.createBaseRdd(sparkSession, inputPaths, parsedOptions) - val multiLine = parsedOptions.multiLine - val lineSeparator = parsedOptions.lineSeparatorInRead - val encoding = parsedOptions.encoding - val ignoreCorruptFiles = parsedOptions.ignoreCorruptFiles - val ignoreMissingFiles = parsedOptions.ignoreMissingFiles - - // Applies `perEntry` to each input -- once per archive entry, once for a loose file -- skipping - // a whole input when it is corrupt/missing and the ignore flags are set. The entry/file stream - // is consumed lazily by `perEntry`, never buffered whole; mirrors CSV's `inferWithArchives`. - // - // An archive entry's stream is a `CloseShieldInputStream` view over the one shared - // `TarArchiveInputStream` cursor, valid only until `readEntries` advances to the next entry - // (`getNextEntry` skips the prior entry's unread bytes). A consumer that emits the raw stream - // (the multiLine path below) must therefore fully consume each element before the iterator - // advances; `JsonInferSchema.infer` does, parsing each record before pulling the next. - // Buffering, look-ahead, or parallelizing the per-partition consumption would read from an - // advanced cursor and infer a wrong schema. The line-delimited path sidesteps this by - // materializing each record (`copyBytes()`). - def perInput[T: ClassTag](perEntry: InputStream => Iterator[T]): RDD[T] = baseRdd.flatMap { - stream => - val path = new Path(stream.getPath()) - try { - if (SupportsArchiveFormat.isArchivePath(path)) { - SupportsArchiveFormat.readArchiveEntries(path, stream.getConfiguration) { (_, in) => - perEntry(in) - } - } else { - perEntry( - CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path)) - } - } catch { - case e: FileNotFoundException if ignoreMissingFiles => - logWarning(log"Skipped missing input: ${MDC(PATH, stream.getPath())}", e) - Iterator.empty - case e: FileNotFoundException => throw e - case e @ (_: RuntimeException | _: IOException) if ignoreCorruptFiles => - logWarning(log"Skipped the corrupted input: ${MDC(PATH, stream.getPath())}", e) - Iterator.empty - case NonFatal(e) => - throw QueryExecutionErrors.cannotReadFilesError( - e, SparkPath.fromPathString(stream.getPath()).urlEncoded) - } - } - - SQLExecution.withSQLConfPropagated(sparkSession) { - val inferSchema = new JsonInferSchema(parsedOptions) - if (multiLine) { - // Each input/entry is one JSON document: hand its stream straight to the parser - // (`CreateJacksonParser.inputStream`, matching MultiLineJsonDataSource and its charset - // auto-detect) so the document is parsed incrementally rather than buffered. - val docs = perInput(in => Iterator.single(in)) - val docParser: (JsonFactory, InputStream) => JsonParser = encoding - .map(enc => CreateJacksonParser.inputStream(enc, _: JsonFactory, _: InputStream)) - .getOrElse(CreateJacksonParser.inputStream(_: JsonFactory, _: InputStream)) - inferSchema.infer[InputStream]( - JsonUtils.sample(docs, parsedOptions), docParser, isReadFile = true) - } else { - // Line-delimited: each line is a record, copied off the reused line buffer and parsed from - // its bytes (`CreateJacksonParser.bytes`, matching TextInputJsonDataSource). - val lines = perInput(in => lineIterator(in, lineSeparator).map(_.copyBytes())) - val lineParser: (JsonFactory, Array[Byte]) => JsonParser = encoding - .map(enc => CreateJacksonParser.bytes(enc, _: JsonFactory, _: Array[Byte])) - .getOrElse(CreateJacksonParser.bytes(_: JsonFactory, _: Array[Byte])) - inferSchema.infer[Array[Byte]]( - JsonUtils.sample(lines, parsedOptions), lineParser, isReadFile = false) - } - } - } } object JsonDataSource { @@ -344,6 +254,12 @@ object MultiLineJsonDataSource extends JsonDataSource { sparkSession: SparkSession, inputPaths: Seq[FileStatus], parsedOptions: JSONOptions): StructType = { + val hasArchive = parsedOptions.archiveFormatEnabled && + inputPaths.exists(f => SupportsArchiveFormat.isArchivePath(f.getPath)) + if (hasArchive) { + return inferWithArchives(sparkSession, inputPaths, parsedOptions) + } + val json: RDD[PortableDataStream] = JsonDataSource.createBaseRdd(sparkSession, inputPaths, parsedOptions) val sampled: RDD[PortableDataStream] = JsonUtils.sample(json, parsedOptions) @@ -357,6 +273,109 @@ object MultiLineJsonDataSource extends JsonDataSource { } } + /** + * Infers a multi-line JSON schema when at least one input is an archive. Each archive entry + * (streamed via `SupportsArchiveFormat`, never unpacked to disk) and each loose file is one whole + * JSON document, and all feed a single [[JsonInferSchema]] pass -- exactly as a directory of the + * same files would infer. Single-line archive inference does not come here: the Text data source + * reads archives directly, so it flows through [[TextInputJsonDataSource.infer]] like any + * directory read. Corrupt/missing inputs are skipped when the ignore flags are set (see + * [[skipInputOnError]]). + */ + private def inferWithArchives( + sparkSession: SparkSession, + inputPaths: Seq[FileStatus], + parsedOptions: JSONOptions): StructType = { + val baseRdd = JsonDataSource.createBaseRdd(sparkSession, inputPaths, parsedOptions) + // Inference must see the same entries the scan reads, so it honors archivePathFilter too. + // Capture the glob string: the compiled GlobPattern is not serializable, so each task + // compiles it once when the archive branch is taken. + val archivePathFilterGlob = parsedOptions.archivePathFilter + val encoding = parsedOptions.encoding + val ignoreCorruptFiles = parsedOptions.ignoreCorruptFiles + val ignoreMissingFiles = parsedOptions.ignoreMissingFiles + + // An archive entry's stream is only valid until the shared cursor advances, so each document + // must be consumed before the next is pulled; `JsonInferSchema.infer` does. + val docs: RDD[InputStream] = baseRdd.mapPartitions { streams => + // Compile at most once per partition: lazy so a partition of only loose files never + // compiles, while a partition with archives reuses one matcher across all of them. + lazy val archivePathFilter = + archivePathFilterGlob.map(FileSourceOptions.compileArchivePathFilter) + streams.flatMap { stream => + val path = new Path(stream.getPath()) + skipInputOnError(stream.getPath(), ignoreMissingFiles, ignoreCorruptFiles) { + if (SupportsArchiveFormat.isArchivePath(path)) { + SupportsArchiveFormat.readArchiveEntries( + path, stream.getConfiguration, archivePathFilter = archivePathFilter) { + (_, in) => + Iterator.single(in) + } + } else { + Iterator.single( + CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path)) + } + } + } + } + + SQLExecution.withSQLConfPropagated(sparkSession) { + val docParser: (JsonFactory, InputStream) => JsonParser = encoding + .map(enc => CreateJacksonParser.inputStream(enc, _: JsonFactory, _: InputStream)) + .getOrElse(CreateJacksonParser.inputStream(_: JsonFactory, _: InputStream)) + new JsonInferSchema(parsedOptions).infer[InputStream]( + JsonUtils.sample(docs, parsedOptions), docParser, isReadFile = true) + } + } + + /** + * Builds one input's document iterator, catching a missing/corrupt error when the ignore flags + * are set. `readArchiveEntries` advances to later entries lazily, so a corrupt later entry throws + * on `hasNext`, not at construction; the returned iterator catches both. A construction failure + * skips the whole input; a mid-advance failure keeps the entries already yielded and skips only + * the remainder of the archive. + */ + private def skipInputOnError( + inputPath: String, + ignoreMissingFiles: Boolean, + ignoreCorruptFiles: Boolean)( + build: => Iterator[InputStream]): Iterator[InputStream] = { + def handle(e: Throwable): Iterator[InputStream] = e match { + case e: FileNotFoundException if ignoreMissingFiles => + logWarning(log"Skipped missing input: ${MDC(PATH, inputPath)}", e) + Iterator.empty + case e: FileNotFoundException => throw e + case e @ (_: RuntimeException | _: IOException) if ignoreCorruptFiles => + logWarning(log"Skipped the corrupted input: ${MDC(PATH, inputPath)}", e) + Iterator.empty + case NonFatal(e) => + throw QueryExecutionErrors.cannotReadFilesError( + e, SparkPath.fromPathString(inputPath).urlEncoded) + } + + val underlying = + try build + catch { case NonFatal(e) => return handle(e) } + + new Iterator[InputStream] { + private var delegate = underlying + override def hasNext: Boolean = + try delegate.hasNext + catch { + case NonFatal(e) => + // A mid-advance throw reaches neither exhaustion nor close, so close the failed + // iterator here to release its archive stream promptly rather than at task completion. + delegate match { + case c: Closeable => try c.close() catch { case NonFatal(_) => } + case _ => + } + delegate = handle(e) + delegate.hasNext + } + override def next(): InputStream = delegate.next() + } + } + private def dataToInputStream(dataStream: PortableDataStream): InputStream = { val path = new Path(dataStream.getPath()) CodecStreams.createInputStreamWithCloseResource(dataStream.getConfiguration, path) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala index 4ede461a4d513..eeeb96d0e6e51 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/json/JsonFileFormat.scala @@ -109,7 +109,8 @@ case class JsonFileFormat() extends TextBasedFileFormat with DataSourceRegister filters) if (parsedOptions.archiveFormatEnabled && SupportsArchiveFormat.isArchivePath(file.toPath)) { JsonDataSource(parsedOptions).readArchive( - broadcastedHadoopConf.value.value, file, () => parser(), requiredSchema) + broadcastedHadoopConf.value.value, file, () => parser(), requiredSchema, + parsedOptions.archivePathFilterPattern) } else { JsonDataSource(parsedOptions).readFile( broadcastedHadoopConf.value.value, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcDeserializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcDeserializer.scala index 04dd37dec50d0..42fcb0cea2eb4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcDeserializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcDeserializer.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{SpecificInternalRow, UnsafeArrayData} import org.apache.spark.sql.catalyst.util._ import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns._ +import org.apache.spark.sql.execution.datasources.orc.types.ops.OrcTypeOps import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -126,7 +127,7 @@ class OrcDeserializer( case IntegerType | _: YearMonthIntervalType => (ordinal, value) => updater.setInt(ordinal, value.asInstanceOf[IntWritable].get) - case LongType | _: DayTimeIntervalType | _: TimestampNTZType | _: TimeType => + case LongType | _: DayTimeIntervalType | _: TimestampNTZType => (ordinal, value) => updater.setLong(ordinal, value.asInstanceOf[LongWritable].get) case FloatType => (ordinal, value) => @@ -149,16 +150,11 @@ class OrcDeserializer( case TimestampType => (ordinal, value) => updater.setLong(ordinal, DateTimeUtils.fromJavaTimestamp(value.asInstanceOf[OrcTimestamp])) - case t: TimestampLTZNanosType => (ordinal, value) => - val ts = value.asInstanceOf[OrcTimestamp] - val instant = ts.toInstant - updater.set(ordinal, DateTimeUtils.instantToTimestampNanos(instant, t.precision)) - case t: TimestampNTZNanosType => (ordinal, value) => - val ts = value.asInstanceOf[OrcTimestamp] - val localDateTime = ts.toLocalDateTime - updater.set( - ordinal, - DateTimeUtils.localDateTimeToTimestampNanos(localDateTime, t.precision)) + + // Framework types (TimeType, nanosecond timestamps) provide their own ORC row reader. The + // setter callbacks decouple the ops sub-package from the sealed CatalystDataUpdater trait. + case OrcTypeOps(ops) => + ops.makeDeserializer(updater.setLong, updater.set) case DecimalType.Fixed(precision, scale) => (ordinal, value) => val v = OrcShimUtils.getDecimal(value) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala index bf0729bf4a0a9..c069abe3b4380 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFileFormat.scala @@ -33,12 +33,13 @@ import org.apache.orc.mapreduce._ import org.apache.spark.TaskContext import org.apache.spark.memory.MemoryMode import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.execution.datasources._ import org.apache.spark.sql.internal.SessionStateHelper +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.sources._ import org.apache.spark.sql.types._ import org.apache.spark.util.{SerializableConfiguration, Utils} @@ -50,8 +51,12 @@ class OrcFileFormat extends FileFormat with DataSourceRegister with SessionStateHelper + with SupportsArchiveFormat with Serializable { + // ORC part-files are often extensionless, so read every archive entry. + override protected def archiveEntryFilter(name: String): Boolean = true + override def shortName(): String = "orc" override def toString: String = "ORC" @@ -114,6 +119,10 @@ class OrcFileFormat sparkSession: SparkSession, options: Map[String, String], path: Path): Boolean = { + val orcOptions = new OrcOptions(options, getSqlConf(sparkSession)) + if (orcOptions.archiveFormatEnabled && SupportsArchiveFormat.isArchivePath(path)) { + return false + } true } @@ -168,8 +177,10 @@ class OrcFileFormat SerializableConfiguration.broadcast(sparkSession.sparkContext, hadoopConf) val isCaseSensitive = sqlConf.caseSensitiveAnalysis val orcFilterPushDown = sqlConf.orcFilterPushDown + val archiveFormatEnabled = sqlConf.getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED) + val fileSourceOptions = new FileSourceOptions(options) - (file: PartitionedFile) => { + def readSingleFile(file: PartitionedFile): Iterator[InternalRow] = { val conf = broadcastedConf.value.value val filePath = file.toPath @@ -244,6 +255,18 @@ class OrcFileFormat } } } + + (file: PartitionedFile) => { + if (archiveFormatEnabled && SupportsArchiveFormat.isArchivePath(file.toPath)) { + readLocalizedEntries( + file, broadcastedConf.value.value, "orc-archive", + fileSourceOptions.archivePathFilterPattern) { entryFile => + readSingleFile(entryFile) + } + } else { + readSingleFile(file) + } + } } override def supportDataType(dataType: DataType): Boolean = dataType match { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilters.scala index 66770870dabe9..6ca2511b33162 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilters.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilters.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.datasources.orc -import java.time.{Duration, Instant, LocalDate, LocalDateTime, LocalTime, Period} +import java.time.{Duration, Instant, LocalDate, LocalDateTime, Period} import org.apache.hadoop.hive.common.`type`.HiveDecimal import org.apache.hadoop.hive.ql.io.sarg.{PredicateLeaf, SearchArgument} @@ -25,9 +25,10 @@ import org.apache.hadoop.hive.ql.io.sarg.SearchArgument.Builder import org.apache.hadoop.hive.ql.io.sarg.SearchArgumentFactory.newBuilder import org.apache.hadoop.hive.serde2.io.HiveDecimalWritable -import org.apache.spark.sql.catalyst.util.DateTimeUtils.{instantToMicros, localDateTimeToMicros, localDateToDays, localTimeToNanos, toJavaDate, toJavaTimestamp} +import org.apache.spark.sql.catalyst.util.DateTimeUtils.{instantToMicros, localDateTimeToMicros, localDateToDays, toJavaDate, toJavaTimestamp} import org.apache.spark.sql.catalyst.util.IntervalUtils import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.datasources.orc.types.ops.OrcTypeOps import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types._ @@ -142,13 +143,17 @@ private[sql] object OrcFilters extends OrcFiltersBase { def getPredicateLeafType(dataType: DataType): PredicateLeaf.Type = dataType match { case BooleanType => PredicateLeaf.Type.BOOLEAN case ByteType | ShortType | IntegerType | LongType | - _: AnsiIntervalType | TimestampNTZType | _: TimeType => PredicateLeaf.Type.LONG + _: AnsiIntervalType | TimestampNTZType => PredicateLeaf.Type.LONG case FloatType | DoubleType => PredicateLeaf.Type.FLOAT case StringType => PredicateLeaf.Type.STRING case DateType => PredicateLeaf.Type.DATE case TimestampType => PredicateLeaf.Type.TIMESTAMP case _: DecimalType => PredicateLeaf.Type.DECIMAL - case _ => throw QueryExecutionErrors.unsupportedOperationForDataTypeError(dataType) + // Framework types (e.g. TimeType, the nanosecond-timestamp types) supply their own + // predicate-leaf type. A framework type whose predicateLeafType is None, or any other unmapped + // type, reaches the same unsupported-type error as before this change. + case dt => OrcTypeOps(dt).flatMap(_.predicateLeafType) + .getOrElse(throw QueryExecutionErrors.unsupportedOperationForDataTypeError(dt)) } /** @@ -170,12 +175,13 @@ private[sql] object OrcFilters extends OrcFiltersBase { toJavaTimestamp(instantToMicros(value.asInstanceOf[Instant])) case _: TimestampNTZType if value.isInstanceOf[LocalDateTime] => localDateTimeToMicros(value.asInstanceOf[LocalDateTime]) - case _: TimeType if value.isInstanceOf[LocalTime] => - localTimeToNanos(value.asInstanceOf[LocalTime]) case _: YearMonthIntervalType => IntervalUtils.periodToMonths(value.asInstanceOf[Period]).longValue() case _: DayTimeIntervalType => IntervalUtils.durationToMicros(value.asInstanceOf[Duration]) + // Framework types (e.g. TimeType) cast their own filter literals; the default is identity, so + // non-framework types and non-matching values fall through unchanged. + case OrcTypeOps(ops) => ops.castFilterLiteral(value) case _ => value } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcSerializer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcSerializer.scala index 64bc3e4292124..3587700645182 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcSerializer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcSerializer.scala @@ -27,8 +27,8 @@ import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.SpecializedGetters import org.apache.spark.sql.catalyst.util._ +import org.apache.spark.sql.execution.datasources.orc.types.ops.OrcTypeOps import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.TimestampNanosVal /** * A serializer to serialize Spark rows to ORC structs. @@ -107,7 +107,7 @@ class OrcSerializer(dataSchema: StructType) { } - case LongType | _: DayTimeIntervalType | _: TimestampNTZType | _: TimeType => + case LongType | _: DayTimeIntervalType | _: TimestampNTZType => if (reuseObj) { val result = new LongWritable() (getter, ordinal) => @@ -155,19 +155,9 @@ class OrcSerializer(dataSchema: StructType) { val result = new OrcTimestamp(ts.getTime) result.setNanos(ts.getNanos) result - case t: TimestampLTZNanosType => (getter, ordinal) => - val v = getter.get(ordinal, t).asInstanceOf[TimestampNanosVal] - val instant = DateTimeUtils.timestampNanosToInstant(v) - val result = new OrcTimestamp(instant.toEpochMilli) - result.setNanos(instant.getNano) - result - case t: TimestampNTZNanosType => (getter, ordinal) => - val v = getter.get(ordinal, t).asInstanceOf[TimestampNanosVal] - val localDateTime = DateTimeUtils.timestampNanosToLocalDateTime(v) - val ts = java.sql.Timestamp.valueOf(localDateTime) - val result = new OrcTimestamp(ts.getTime) - result.setNanos(ts.getNanos) - result + + // Framework types (TimeType, nanosecond timestamps) provide their own ORC value writer. + case OrcTypeOps(ops) => ops.makeSerializer(reuseObj) case DecimalType.Fixed(precision, scale) => OrcShimUtils.getHiveDecimalWritable(precision, scale) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcUtils.scala index 1073b53dcaf66..f533ded6243c2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/OrcUtils.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.datasources.orc -import java.io.FileNotFoundException +import java.io.{File, FileNotFoundException, IOException} import java.nio.charset.StandardCharsets.UTF_8 import java.util.Locale @@ -27,7 +27,7 @@ import scala.util.control.NonFatal import org.apache.commons.lang3.exception.ExceptionUtils import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.hive.serde2.io.DateWritable import org.apache.hadoop.io.{BooleanWritable, ByteWritable, DoubleWritable, FloatWritable, IntWritable, LongWritable, ShortWritable, WritableComparable} import org.apache.hadoop.mapreduce.lib.input.FileSplit @@ -35,7 +35,7 @@ import org.apache.orc.{BooleanColumnStatistics, ColumnStatistics, DateColumnStat import org.apache.orc.mapred.{OrcInputFormat => OrcMapredInputFormat, OrcStruct} import org.apache.orc.mapreduce.OrcMapreduceRecordReader -import org.apache.spark.{SPARK_VERSION_SHORT, SparkException} +import org.apache.spark.{SPARK_VERSION_SHORT, SparkEnv, SparkException} import org.apache.spark.deploy.SparkHadoopUtil import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.PATH @@ -47,11 +47,13 @@ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.catalyst.util.{quoteIdentifier, CaseInsensitiveMap, CharVarcharUtils} import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Count, CountStar, Max, Min} import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, SchemaMergeUtils} +import org.apache.spark.sql.execution.datasources.{AggregatePushDownUtils, SchemaMergeUtils, SupportsArchiveFormat} +import org.apache.spark.sql.execution.datasources.orc.types.ops.OrcTypeOps import org.apache.spark.sql.execution.datasources.v2.V2ColumnUtils +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.spark.util.{ThreadUtils, Utils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.ThreadUtils object OrcUtils extends Logging { @@ -174,8 +176,8 @@ object OrcUtils extends Logging { MapType(catalystKeyType, catalystValueType) } - // The Spark query engine has not completely supported CHAR/VARCHAR type yet, and here we - // replace the orc CHAR/VARCHAR with STRING type. + // Annotate metadata for the legacy path; under charVarcharFirstClassTypes the constrained + // types are kept so ORC catalyst attributes / native char/varchar round-trip as first-class. CharVarcharUtils.replaceCharVarcharWithStringInSchema(toStructType(schema)) } @@ -183,11 +185,23 @@ object OrcUtils extends Logging { : Option[StructType] = { val ignoreCorruptFiles = new FileSourceOptions(CaseInsensitiveMap(options)).ignoreCorruptFiles val conf = sparkSession.sessionState.newHadoopConfWithOptions(options) - files.iterator.map(file => readSchema(file.getPath, conf, ignoreCorruptFiles)).collectFirst { - case Some(schema) => - logDebug(s"Reading schema from file $files, got Hive schema string: $schema") - toCatalystSchema(schema) - } + val ignoreMissingFiles = + new FileSourceOptions(CaseInsensitiveMap(options)).ignoreMissingFiles + val archiveFormatEnabled = SQLConf.get.getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED) + val archivePathFilter = + new FileSourceOptions(CaseInsensitiveMap(options)).archivePathFilterPattern + files.iterator.flatMap { file => + if (archiveFormatEnabled && SupportsArchiveFormat.isArchivePath(file.getPath)) { + readArchiveSchemas(conf, file, ignoreCorruptFiles, ignoreMissingFiles, + stopAtFirst = true, archivePathFilter) + .headOption + } else { + readSchema(file.getPath, conf, ignoreCorruptFiles).map { schema => + logDebug(s"Reading schema from file $files, got Hive schema string: $schema") + toCatalystSchema(schema) + } + } + }.collectFirst { case schema => schema } } /** @@ -197,12 +211,85 @@ object OrcUtils extends Logging { def readOrcSchemasInParallel( files: Seq[FileStatus], conf: Configuration, ignoreCorruptFiles: Boolean, ignoreMissingFiles: Boolean): Seq[StructType] = { + // Read outside `parmap`: its worker threads do not inherit the caller's `SQLConf` thread-local, + // so `SQLConf.get` there would fall back to defaults and never take the archive branch. + val archiveEnabled = SQLConf.get.getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED) + // The signature is fixed by `SchemaMergeUtils.mergeSchemasInParallel`'s `schemaReader` type, + // so the glob comes from `conf`, which that caller builds with the user options via + // `newHadoopConfWithOptions`. + val archivePathFilter = Option(conf.get(FileSourceOptions.ARCHIVE_PATH_FILTER)) + .filter(_.nonEmpty).map(FileSourceOptions.compileArchivePathFilter) ThreadUtils.parmap(files, "readingOrcSchemas", 8) { currentFile => - OrcUtils.readSchema(currentFile.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles) - .map(toCatalystSchema) + if (archiveEnabled && SupportsArchiveFormat.isArchivePath(currentFile.getPath)) { + readArchiveSchemas(conf, currentFile, ignoreCorruptFiles, ignoreMissingFiles, + archivePathFilter = archivePathFilter, + stopAtFirst = false) + } else { + OrcUtils.readSchema(currentFile.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles) + .map(toCatalystSchema).toSeq + } }.flatten } + /** + * Reads ORC entry schemas from one archive. + * + * @param stopAtFirst when true, returns just the first entry's schema (sample-one); otherwise + * reads every entry so a corrupt entry fails the whole archive. + */ + private def readArchiveSchemas( + conf: Configuration, + archive: FileStatus, + ignoreCorruptFiles: Boolean, + ignoreMissingFiles: Boolean, + stopAtFirst: Boolean, + archivePathFilter: Option[GlobPattern]): Seq[StructType] = { + val tempDir = Utils.createTempDir(Utils.getLocalDir(SparkEnv.get.conf), "orc-archive-infer") + // localizeEntries eagerly opens the first entry, so build it inside the try; the finally must + // still delete tempDir when a corrupt archive throws there. + var entries: Iterator[(String, File)] = Iterator.empty + try { + entries = SupportsArchiveFormat.localizeEntries( + archive.getPath, conf, tempDir, _ => true, archivePathFilter) + // With ignore flags off, a corrupt entry throws to the per-archive catch below. `.toList` + // reads every entry (whole archive atomic); `stopAtFirst` stays lazy and stops at the first. + val schemas = entries.flatMap { case (_, entryFile) => + try { + readSchema(new Path(entryFile.toURI), conf, ignoreCorruptFiles = false, + ignoreMissingFiles = false).map(toCatalystSchema) + } finally entryFile.delete() + } + if (stopAtFirst) schemas.take(1).toList else schemas.toList + } catch { + // A corrupt container throws IOException at open; a corrupt entry footer throws + // cannotReadFooterForFileError, whose root cause is also an IOException. + case e: Exception if ignoreMissingFiles && + ExceptionUtils.getThrowables(e).exists(_.isInstanceOf[FileNotFoundException]) => + logWarning(log"Skipped missing archive during inference: ${MDC(PATH, archive.getPath)}", e) + Seq.empty + case e: Exception if { + val root = Utils.getRootCause(e) + root.isInstanceOf[IOException] && !root.isInstanceOf[FileNotFoundException] + } => + if (ignoreCorruptFiles) { + logWarning(log"Skipped the corrupt archive during inference: " + + log"${MDC(PATH, archive.getPath)}", e) + Seq.empty + } else if (e.isInstanceOf[SparkException]) { + throw e // a corrupt entry footer already is cannotReadFooterForFileError + } else { + // Match the loose-file footer error rather than leaking the raw container IOException. + throw QueryExecutionErrors.cannotReadFooterForFileError(archive.getPath, e) + } + } finally { + entries match { + case c: java.io.Closeable => c.close() + case _ => + } + Utils.deleteRecursively(tempDir) + } + } + def inferSchema(sparkSession: SparkSession, files: Seq[FileStatus], options: Map[String, String]) : Option[StructType] = { val orcOptions = new OrcOptions(options, sparkSession.sessionState.conf) @@ -350,11 +437,10 @@ object OrcUtils extends Logging { s"array<${getOrcSchemaString(a.elementType)}>" case m: MapType => s"map<${getOrcSchemaString(m.keyType)},${getOrcSchemaString(m.valueType)}>" - case _: DayTimeIntervalType | _: TimestampNTZType | _: TimeType => LongType.catalogString - case _: TimestampLTZNanosType => "timestamp with local time zone" - case _: TimestampNTZNanosType => "timestamp" + case _: DayTimeIntervalType | _: TimestampNTZType => LongType.catalogString case _: YearMonthIntervalType => IntegerType.catalogString - case _ => dt.catalogString + // Framework types (TimeType, nanosecond timestamps) supply their own ORC schema string. + case _ => OrcTypeOps(dt).map(_.orcSchemaString).getOrElse(dt.catalogString) } def orcTypeDescription(dt: DataType): TypeDescription = { @@ -372,25 +458,28 @@ object OrcUtils extends Logging { val typeDesc = new TypeDescription(TypeDescription.Category.LONG) typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, n.typeName) Some(typeDesc) - case tm: TimeType => - val typeDesc = new TypeDescription(TypeDescription.Category.LONG) - typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, tm.typeName) - Some(typeDesc) case t: TimestampType => val typeDesc = new TypeDescription(TypeDescription.Category.TIMESTAMP) typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, t.typeName) Some(typeDesc) - case t: TimestampLTZNanosType => - val typeDesc = new TypeDescription(TypeDescription.Category.TIMESTAMP_INSTANT) - typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, t.typeName) - Some(typeDesc) - case t: TimestampNTZNanosType => - val typeDesc = new TypeDescription(TypeDescription.Category.TIMESTAMP) - typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, t.typeName) + // Framework types (TimeType, nanosecond timestamps) supply their own ORC category; the + // CATALYST_TYPE_ATTRIBUTE_NAME is stamped uniformly here so the true Spark type + // round-trips on read. + case OrcTypeOps(ops) => + val typeDesc = new TypeDescription(ops.orcCategory) + typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, dt.typeName) Some(typeDesc) - case _: StringType => + // Write CHAR/VARCHAR as ORC STRING plus spark.sql.catalyst.type, not native + // ORC CHAR/VARCHAR. Native ORC maxLength would truncate/pad independently of + // Spark store assignment. Hive-written native CHAR still round-trips on read + // via toCatalystSchema. Unbounded STRING (including collated) stamps + // StringType.typeName ("string"), matching Avro: this PR does not round-trip + // collation on file-only reads. + case s: StringType => val typeDesc = new TypeDescription(TypeDescription.Category.STRING) - typeDesc.setAttribute(CATALYST_TYPE_ATTRIBUTE_NAME, StringType.typeName) + typeDesc.setAttribute( + CATALYST_TYPE_ATTRIBUTE_NAME, + CharVarcharUtils.charVarcharTypeName(s).getOrElse(StringType.typeName)) Some(typeDesc) case _ => None } @@ -508,7 +597,7 @@ object OrcUtils extends Logging { // Get column statistics with column name. def getColumnStatistics(columnName: String): ColumnStatistics = { - val columnIndex = dataSchema.fieldNames.indexOf(columnName) + val columnIndex = dataSchema.getFieldIndex(columnName).getOrElse(-1) columnsStatistics.get(columnIndex).getStatistics } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/OrcTypeOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/OrcTypeOps.scala new file mode 100644 index 0000000000000..b5a2a60103687 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/OrcTypeOps.scala @@ -0,0 +1,194 @@ +/* + * 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.execution.datasources.orc.types.ops + +import org.apache.hadoop.hive.ql.io.sarg.PredicateLeaf +import org.apache.hadoop.io.WritableComparable +import org.apache.orc.TypeDescription + +import org.apache.spark.sql.catalyst.expressions.SpecializedGetters +import org.apache.spark.sql.types.{DataType, TimestampLTZNanosType, TimestampNTZNanosType, TimeType} + +/** + * Optional trait for ORC storage-format integration in the Types Framework. + * + * Implement this trait to enable ORC read/write support for a framework type. Each framework + * type that supports ORC provides a concrete implementation and registers it in the companion + * object's apply() method. + * + * The trait covers the ORC concerns that are homogeneous across the row-based code paths: + * - Schema mapping: Spark DataType -> ORC schema string (OrcUtils.getOrcSchemaString) and the + * ORC TypeDescription category (OrcUtils.orcTypeDescription). ORC has no native TIME / + * nanosecond-timestamp category, so framework types map onto an existing physical ORC + * category and round-trip the true Spark type via the CATALYST_TYPE_ATTRIBUTE_NAME attribute + * (stamped uniformly by the caller, so the ops only chooses the category). + * - Value write (serialize): Catalyst value -> ORC WritableComparable + * (OrcSerializer.newConverter) + * - Row-based read (deserialize): ORC WritableComparable -> Catalyst value + * (OrcDeserializer.newWriter) + * - Predicate pushdown: PredicateLeaf.Type mapping + literal casting (OrcFilters) + * + * DELIBERATELY NOT ON THE TRAIT: + * - Vectorized read. ORC's vectorized path (OrcAtomicColumnVector, a Java class) dispatches via + * boolean `instanceof` flags set in the constructor plus typed accessor methods + * (getTimestampNTZNanos/getTimestampLTZNanos), NOT via an `Ops(dt).map(_.x)` closure. It has + * no per-type extension seam, so it stays inline. This mirrors Parquet, whose vectorized read + * is likewise not routed through ParquetTypeOps (it dispatches on the Spark type inline in + * ParquetVectorUpdaterFactory.getUpdater). + * + * CAVEAT for a new type author: the vectorized reader is a SEPARATE registration you must + * handle in addition to the ops class. OrcUtils.supportColumnarReads returns true for every + * `AtomicType` and OrcAtomicColumnVector dispatches by `instanceof`, so a new framework + * `AtomicType` wired only into this ops registry is still routed to the vectorized reader, + * where it would be read as raw physical values (silently wrong) rather than failing loudly. + * Add an OrcAtomicColumnVector arm (or exclude the type from supportColumnarReads) as well. + * The current types are wired correctly; this note is so the "one ops class + one registry + * arm" framing does not mislead. + * - supportDataType. OrcFileFormat.supportDataType / OrcTable.supportsDataType already admit + * every `AtomicType` via `case _: AtomicType => true`, so framework types are supported with + * no per-type arm; no gate method is needed (unlike Parquet, whose default differs). + * + * DISPATCH PATTERN: each ORC integration site keeps its existing built-in-type arms unchanged and + * routes only framework types through the ops. The built-in types are matched first; framework + * types are handled by an added arm reached after them (i.e. framework types are dispatched last, + * not first), so this change never alters how a built-in type is handled. Two shapes are used: + * - Inside a `dataType match` (OrcSerializer.newConverter, OrcDeserializer.newWriter, + * OrcUtils.orcTypeDescription, OrcFilters.castLiteralValue): an added + * `case OrcTypeOps(ops) => ops.method(...)` arm placed among the existing arms. The `unapply` + * extractor binds the ops in a single registry lookup. + * - In expression position (OrcUtils.getOrcSchemaString, OrcFilters.getPredicateLeafType): the + * original fallback stays inline and framework types are folded in with + * `OrcTypeOps(dt).map(_.method).getOrElse(<original fallback>)`. + * There are no extracted `*Default` methods; the original ORC code is left in place as the + * fallback arm/expression. + * + * DECOUPLING NOTE: makeDeserializer takes the Catalyst setter callbacks it needs + * ((Int, Long) => Unit and (Int, Any) => Unit) rather than OrcDeserializer's CatalystDataUpdater, + * because that updater is a sealed trait nested in OrcDeserializer and is not visible to this + * sub-package. The callbacks are the only two setter shapes the current framework types use. + * + * @see TimeTypeOrcOps for a reference implementation (primitive Long-backed type) + * @see TimestampLTZNanosOrcOps / TimestampNTZNanosOrcOps for OrcTimestamp-backed types + * @since 5.0.0 + */ +private[orc] trait OrcTypeOps extends Serializable { + + // ==================== Schema Mapping ==================== + + /** + * The ORC schema-string fragment for this type (OrcUtils.getOrcSchemaString). Examples: + * TimeType -> "bigint" (LongType.catalogString), TimestampLTZNanosType -> "timestamp with local + * time zone", TimestampNTZNanosType -> "timestamp". + */ + def orcSchemaString: String + + /** + * The ORC TypeDescription category for this type (OrcUtils.orcTypeDescription). The caller + * stamps CATALYST_TYPE_ATTRIBUTE_NAME = sparkType.typeName onto the returned descriptor, so the + * ops only chooses the physical category. + */ + def orcCategory: TypeDescription.Category + + // ==================== Value Write (serialize) ==================== + + /** + * Creates a converter that turns a Catalyst value at an ordinal into an ORC WritableComparable + * (OrcSerializer.newConverter). + * + * @param reuseObj whether the serializer may reuse a single mutable Writable across rows + * (OrcSerializer passes this through; the primitive Long-backed TimeType reuses, + * the OrcTimestamp-backed nanos types do not). + */ + def makeSerializer(reuseObj: Boolean): (SpecializedGetters, Int) => WritableComparable[_] + + // ==================== Row-Based Read (deserialize) ==================== + + /** + * Creates a writer that sets a decoded ORC value into the Catalyst row at an ordinal + * (OrcDeserializer.newWriter). The WritableComparable is the raw ORC value (LongWritable for + * LONG-category types, OrcTimestamp for the timestamp categories). + * + * @param setLong callback into the row's setLong (used by primitive Long-backed types) + * @param set callback into the row's generic set (used by object-backed types, e.g. nanos) + */ + def makeDeserializer( + setLong: (Int, Long) => Unit, + set: (Int, Any) => Unit): (Int, WritableComparable[_]) => Unit + + // ==================== Predicate Pushdown ==================== + + /** + * The ORC PredicateLeaf.Type for this type, consumed by OrcFilters.getPredicateLeafType. Return + * Some(type) to enable predicate pushdown for this type. + * + * A None keeps the pre-existing OrcFilters behavior for a type with no leaf mapping: it reaches + * getPredicateLeafType's final arm, which throws unsupportedOperationForDataTypeError. In + * practice that arm is only reached for a column already deemed pushdown-eligible upstream + * (OrcFiltersBase.getSearchableTypeMap admits any AtomicType), so a framework AtomicType that + * returns None and is used in a pushed filter would fail during planning. A future type that + * must be pushdown-eligible therefore has to return Some here (and, if its literal needs + * conversion, override castFilterLiteral); returning None is only safe for a type that can never + * reach getPredicateLeafType. + * + * TimeType and the nanosecond-timestamp types return Some (LONG and TIMESTAMP respectively) so + * their pushed filters convert to search arguments. A type that returns None (there are none + * today) would reach the throw arm if ever pushed. + */ + def predicateLeafType: Option[PredicateLeaf.Type] = None + + /** + * Casts a filter literal to the ORC search-argument representation (OrcFilters.castLiteralValue). + * Only meaningful for types that also return a predicateLeafType. Default is identity. + */ + def castFilterLiteral(value: Any): Any = value +} + +/** + * Factory object for creating OrcTypeOps instances. + * + * Provides forward lookup (DataType -> ops) for the framework-type dispatch arms at ORC + * integration sites. apply() returns Some only for framework-managed types, so callers fall back + * to the original inline path for everything else. + */ +private[orc] object OrcTypeOps { + + /** + * Returns an OrcTypeOps instance for the given DataType, if supported. + * + * Returns None if the type has no ORC ops. This is the single registration point for all ORC + * type operations. + */ + def apply(dt: DataType): Option[OrcTypeOps] = dt match { + case tt: TimeType => Some(TimeTypeOrcOps(tt)) + case t: TimestampLTZNanosType => Some(TimestampLTZNanosOrcOps(t)) + case t: TimestampNTZNanosType => Some(TimestampNTZNanosOrcOps(t)) + // Add new types here - single registration point + case _ => None + } + + /** + * Extractor so a `dataType match` arm can bind the ops in a single lookup: + * {{{ + * case OrcTypeOps(ops) => ops.method(...) + * }}} + * Delegates to apply(); returning the same Option keeps the registry the single source of truth + * and avoids the double lookup of a `case _ if OrcTypeOps(dt).isDefined => OrcTypeOps(dt).get` + * guard. + */ + def unapply(dt: DataType): Option[OrcTypeOps] = apply(dt) +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimeTypeOrcOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimeTypeOrcOps.scala new file mode 100644 index 0000000000000..1896a2712eec9 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimeTypeOrcOps.scala @@ -0,0 +1,87 @@ +/* + * 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.execution.datasources.orc.types.ops + +import java.time.LocalTime + +import org.apache.hadoop.hive.ql.io.sarg.PredicateLeaf +import org.apache.hadoop.io.{LongWritable, WritableComparable} +import org.apache.orc.TypeDescription + +import org.apache.spark.sql.catalyst.expressions.SpecializedGetters +import org.apache.spark.sql.catalyst.util.DateTimeUtils.localTimeToNanos +import org.apache.spark.sql.types.{LongType, TimeType} + +/** + * ORC operations for TimeType. + * + * TimeType is a primitive Long-backed type (nanoseconds since midnight). ORC has no TIME category, + * so it is stored in the LONG physical category with the true Spark type recovered on read from + * the CATALYST_TYPE_ATTRIBUTE_NAME attribute. Read and write are pure Long pass-through: the + * internal nanos-of-day value IS the stored LONG (no unit conversion, unlike Parquet which may + * store MICROS). Precision affects only display, so nothing is truncated here. + * + * @param t the TimeType (precision is unused on the ORC path; kept for symmetry with the other + * storage ops and possible future precision-aware behavior). + * @since 5.0.0 + */ +case class TimeTypeOrcOps(t: TimeType) extends OrcTypeOps { + + // ==================== Schema Mapping ==================== + + // Was: OrcUtils.getOrcSchemaString `case ... | _: TimeType => LongType.catalogString` + override def orcSchemaString: String = LongType.catalogString + + // Was: OrcUtils.orcTypeDescription `case tm: TimeType => new TypeDescription(LONG) ...` + override def orcCategory: TypeDescription.Category = TypeDescription.Category.LONG + + // ==================== Value Write (serialize) ==================== + + // Was: OrcSerializer.newConverter `case ... | _: TimeType => ... LongWritable ...` + override def makeSerializer( + reuseObj: Boolean): (SpecializedGetters, Int) => WritableComparable[_] = + if (reuseObj) { + val result = new LongWritable() + (getter: SpecializedGetters, ordinal: Int) => { + result.set(getter.getLong(ordinal)) + result + } + } else { + (getter: SpecializedGetters, ordinal: Int) => new LongWritable(getter.getLong(ordinal)) + } + + // ==================== Row-Based Read (deserialize) ==================== + + // Was: OrcDeserializer.newWriter `case ... | _: TimeType => ... setLong ...` + override def makeDeserializer( + setLong: (Int, Long) => Unit, + set: (Int, Any) => Unit): (Int, WritableComparable[_]) => Unit = + (ordinal: Int, value: WritableComparable[_]) => + setLong(ordinal, value.asInstanceOf[LongWritable].get) + + // ==================== Predicate Pushdown ==================== + + // Was: OrcFilters.getPredicateLeafType `case ... | _: TimeType => PredicateLeaf.Type.LONG` + override def predicateLeafType: Option[PredicateLeaf.Type] = Some(PredicateLeaf.Type.LONG) + + // Was: OrcFilters.castLiteralValue `case _: TimeType if value.isInstanceOf[LocalTime] => ...` + override def castFilterLiteral(value: Any): Any = value match { + case lt: LocalTime => localTimeToNanos(lt) + case other => other + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimestampNanosOrcOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimestampNanosOrcOps.scala new file mode 100644 index 0000000000000..5d5923b5b571e --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/orc/types/ops/TimestampNanosOrcOps.scala @@ -0,0 +1,157 @@ +/* + * 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.execution.datasources.orc.types.ops + +import java.sql.Timestamp +import java.time.{Instant, LocalDateTime, ZoneOffset} + +import org.apache.hadoop.hive.ql.io.sarg.PredicateLeaf +import org.apache.hadoop.io.WritableComparable +import org.apache.orc.TypeDescription +import org.apache.orc.mapred.OrcTimestamp + +import org.apache.spark.sql.catalyst.expressions.SpecializedGetters +import org.apache.spark.sql.catalyst.util.DateTimeUtils +import org.apache.spark.sql.types.{TimestampLTZNanosType, TimestampNTZNanosType} +import org.apache.spark.unsafe.types.TimestampNanosVal + +/** + * ORC operations for TimestampLTZNanosType (nanosecond precision, with time zone). + * + * Stored in ORC as the TIMESTAMP_INSTANT category (matching TIMESTAMP with local time zone). The + * internal TimestampNanosVal is converted through java.time.Instant in both directions; the true + * Spark type and precision are recovered on read from the CATALYST_TYPE_ATTRIBUTE_NAME attribute. + * ORC does not reuse the OrcTimestamp object (the conversion is already expensive), so the + * serializer runs with reuse disabled regardless of the reuseObj hint. + * + * Predicate pushdown uses PredicateLeaf.Type.TIMESTAMP (matching the physical category). ORC's + * SearchArgument builder requires the literal's class to be exactly java.sql.Timestamp (it rejects + * the OrcTimestamp subclass), and ORC compares against a TIMESTAMP_INSTANT column by first shifting + * the stored UTC value into the JVM default zone (the default useUTCTimestamp=false read mode Spark + * uses). So the filter literal (an external java.time.Instant, already truncated to the column + * precision upstream) is cast to a java.sql.Timestamp whose *local* wall clock equals the instant's + * UTC wall clock, via Timestamp.valueOf(LocalDateTime.ofInstant(instant, UTC)). This keeps the + * pushed predicate consistent with the ORC min/max statistics regardless of the JVM time zone. + * + * @since 5.0.0 + */ +case class TimestampLTZNanosOrcOps(t: TimestampLTZNanosType) extends OrcTypeOps { + + // Was: OrcUtils.getOrcSchemaString + // `case _: TimestampLTZNanosType => "timestamp with local time zone"` + override def orcSchemaString: String = "timestamp with local time zone" + + // Was: OrcUtils.orcTypeDescription `case t: TimestampLTZNanosType => ... TIMESTAMP_INSTANT ...` + override def orcCategory: TypeDescription.Category = TypeDescription.Category.TIMESTAMP_INSTANT + + // Was: OrcSerializer.newConverter `case t: TimestampLTZNanosType => ...` + override def makeSerializer( + reuseObj: Boolean): (SpecializedGetters, Int) => WritableComparable[_] = + (getter: SpecializedGetters, ordinal: Int) => { + val v = getter.get(ordinal, t).asInstanceOf[TimestampNanosVal] + val instant = DateTimeUtils.timestampNanosToInstant(v) + val result = new OrcTimestamp(instant.toEpochMilli) + result.setNanos(instant.getNano) + result + } + + // Was: OrcDeserializer.newWriter `case t: TimestampLTZNanosType => ...` + override def makeDeserializer( + setLong: (Int, Long) => Unit, + set: (Int, Any) => Unit): (Int, WritableComparable[_]) => Unit = + (ordinal: Int, value: WritableComparable[_]) => { + val ts = value.asInstanceOf[OrcTimestamp] + val instant = ts.toInstant + set(ordinal, DateTimeUtils.instantToTimestampNanos(instant, t.precision)) + } + + // The physical ORC category is a timestamp, so the search argument uses the TIMESTAMP leaf type. + override def predicateLeafType: Option[PredicateLeaf.Type] = Some(PredicateLeaf.Type.TIMESTAMP) + + // The filter literal is an external java.time.Instant (see CatalystTypeConverters). ORC evaluates + // a TIMESTAMP_INSTANT predicate against the stored UTC value shifted into the JVM default zone, + // so the literal must be a java.sql.Timestamp whose local wall clock equals the instant's UTC + // wall clock; Timestamp.valueOf(LocalDateTime at UTC) produces exactly that. Any non-Instant + // value is passed through unchanged. + override def castFilterLiteral(value: Any): Any = value match { + case i: Instant => Timestamp.valueOf(LocalDateTime.ofInstant(i, ZoneOffset.UTC)) + case other => other + } +} + +/** + * ORC operations for TimestampNTZNanosType (nanosecond precision, without time zone). + * + * Stored in ORC as the TIMESTAMP category. The internal TimestampNanosVal is converted through + * java.time.LocalDateTime in both directions; the true Spark type and precision are recovered on + * read from the CATALYST_TYPE_ATTRIBUTE_NAME attribute. As with the LTZ variant, the OrcTimestamp + * object is not reused. + * + * Predicate pushdown uses PredicateLeaf.Type.TIMESTAMP (matching the physical category). The filter + * literal is an external java.time.LocalDateTime (already truncated to the column precision + * upstream); it is cast to a plain java.sql.Timestamp via Timestamp.valueOf, the same conversion + * the write path applies before wrapping the result in an OrcTimestamp. ORC's SearchArgument + * builder requires the literal's class to be exactly java.sql.Timestamp (it rejects the + * OrcTimestamp subclass), while the stored OrcTimestamp compares as its java.sql.Timestamp + * super-value, so the two are value-equal and the pushed predicate is consistent with the ORC + * min/max statistics. + * + * @since 5.0.0 + */ +case class TimestampNTZNanosOrcOps(t: TimestampNTZNanosType) extends OrcTypeOps { + + // Was: OrcUtils.getOrcSchemaString `case _: TimestampNTZNanosType => "timestamp"` + override def orcSchemaString: String = "timestamp" + + // Was: OrcUtils.orcTypeDescription `case t: TimestampNTZNanosType => ... TIMESTAMP ...` + override def orcCategory: TypeDescription.Category = TypeDescription.Category.TIMESTAMP + + // Was: OrcSerializer.newConverter `case t: TimestampNTZNanosType => ...` + override def makeSerializer( + reuseObj: Boolean): (SpecializedGetters, Int) => WritableComparable[_] = + (getter: SpecializedGetters, ordinal: Int) => { + val v = getter.get(ordinal, t).asInstanceOf[TimestampNanosVal] + val localDateTime = DateTimeUtils.timestampNanosToLocalDateTime(v) + val ts = Timestamp.valueOf(localDateTime) + val result = new OrcTimestamp(ts.getTime) + result.setNanos(ts.getNanos) + result + } + + // Was: OrcDeserializer.newWriter `case t: TimestampNTZNanosType => ...` + override def makeDeserializer( + setLong: (Int, Long) => Unit, + set: (Int, Any) => Unit): (Int, WritableComparable[_]) => Unit = + (ordinal: Int, value: WritableComparable[_]) => { + val ts = value.asInstanceOf[OrcTimestamp] + val localDateTime = ts.toLocalDateTime + set(ordinal, DateTimeUtils.localDateTimeToTimestampNanos(localDateTime, t.precision)) + } + + // The physical ORC category is a timestamp, so the search argument uses the TIMESTAMP leaf type. + override def predicateLeafType: Option[PredicateLeaf.Type] = Some(PredicateLeaf.Type.TIMESTAMP) + + // The filter literal is an external java.time.LocalDateTime (see CatalystTypeConverters); convert + // it to a plain java.sql.Timestamp via the same Timestamp.valueOf the serializer uses, so it is + // value-equal to the written OrcTimestamp. Any non-LocalDateTime value is passed through + // unchanged. + override def castFilterLiteral(value: Any): Any = value match { + case ldt: LocalDateTime => Timestamp.valueOf(ldt) + case other => other + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala index bbfbbfde0ba4b..98385c796d850 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/InferVariantShreddingSchema.scala @@ -365,16 +365,23 @@ class InferVariantShreddingSchema(val schema: StructType) { v.getType match { case Type.OBJECT => val size = v.objectSize() - // Validate fields are sorted (per variant spec) + var utf8Sorted = true + var utf16Sorted = true + var previousKey = if (size > 0) v.getFieldAtIndex(0).key else null + var previousKeyBytes = if (size > 0) VariantUtil.encodeKey(previousKey) else null for (i <- 1 until size) { - val prevKey = v.getFieldAtIndex(i - 1).key - val currKey = v.getFieldAtIndex(i).key - if (prevKey >= currKey) { - throw new SparkRuntimeException( - errorClass = "MALFORMED_VARIANT", - messageParameters = Map.empty - ) - } + val currentKey = v.getFieldAtIndex(i).key + val currentKeyBytes = VariantUtil.encodeKey(currentKey) + utf8Sorted &&= VariantUtil.compareKeys(previousKeyBytes, currentKeyBytes) < 0 + utf16Sorted &&= previousKey.compareTo(currentKey) < 0 + previousKey = currentKey + previousKeyBytes = currentKeyBytes + } + if (!utf8Sorted && !utf16Sorted) { + throw new SparkRuntimeException( + errorClass = "MALFORMED_VARIANT", + messageParameters = Map.empty + ) } // Process each field diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala index 235aa428fbce1..c108e74eca0ff 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormat.scala @@ -27,7 +27,7 @@ import scala.util.{Failure, Try} import org.apache.commons.lang3.exception.ExceptionUtils import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.mapred.FileSplit import org.apache.hadoop.mapreduce._ import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl @@ -41,12 +41,12 @@ import org.apache.spark.{SparkEnv, TaskContext} import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{PATH, SCHEMA} import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.catalyst.parser.LegacyTypeStringParser import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes -import org.apache.spark.sql.catalyst.util.{DateTimeUtils, RebaseDateTime} +import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils, RebaseDateTime} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.datasources._ import org.apache.spark.sql.execution.datasources.parquet.types.ops.ParquetTypeOps @@ -210,6 +210,18 @@ class ParquetFileFormat val pushDownStringPredicate = sqlConf.parquetFilterPushDownStringPredicate val pushDownInFilterThreshold = sqlConf.parquetFilterPushDownInFilterThreshold val isCaseSensitive = sqlConf.caseSensitiveAnalysis + // When shredded-variant predicate pushdown is enabled and `requiredSchema` actually carries a + // variant-extraction struct produced by PushVariantIntoScan, passing it lets ParquetFilters map + // logical paths like "v.`0`" to the physical shredded columns for row-group skipping. The + // `isVariantStruct` check keeps the per-file ParquetFilters construction free of any shredded + // traversal for the common case of scans with no variant-extraction columns. + val variantExtractionSchema = + if (sqlConf.getConf(SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED) && + requiredSchema.existsRecursively(VariantMetadata.isVariantStruct)) { + Some(requiredSchema) + } else { + None + } val parquetOptions = new ParquetOptions(options, sqlConf) val datetimeRebaseModeInRead = parquetOptions.datetimeRebaseModeInRead val int96RebaseModeInRead = parquetOptions.int96RebaseModeInRead @@ -266,7 +278,8 @@ class ParquetFileFormat pushDownStringPredicate, pushDownInFilterThreshold, isCaseSensitive, - datetimeRebaseSpec) + datetimeRebaseSpec, + variantExtractionSchema = variantExtractionSchema) filters // Collects all converted Parquet filter predicates. Notice that not all predicates // can be converted (`ParquetFilters.createFilter` returns an `Option`). That's why @@ -322,7 +335,9 @@ class ParquetFileFormat // An archive is read by unpacking each Parquet entry to a local temp file and reading it with // the plain reader (readSingleFile); readLocalizedEntries owns the unpack/iterate/cleanup. def readArchiveFile(file: PartitionedFile): Iterator[InternalRow] = - readLocalizedEntries(file, broadcastedHadoopConf.value.value, "parquet-archive") { + readLocalizedEntries( + file, broadcastedHadoopConf.value.value, "parquet-archive", + parquetOptions.archivePathFilterPattern) { entryFile => readSingleFile(entryFile) } @@ -552,6 +567,11 @@ object ParquetFileFormat extends Logging { ignoreCorruptFiles: Boolean, ignoreMissingFiles: Boolean = false): Seq[Footer] = { val archiveEnabled = SQLConf.get.getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED) + // This signature is shared with `SchemaMergeUtils.mergeSchemasInParallel`'s `schemaReader`, so + // the glob comes from `conf`, which that caller builds with the user options via + // `newHadoopConfWithOptions`. + val archivePathFilter = Option(conf.get(FileSourceOptions.ARCHIVE_PATH_FILTER)) + .filter(_.nonEmpty).map(FileSourceOptions.compileArchivePathFilter) ThreadUtils.parmap(partFiles, "readingParquetFooters", 8, preserveSparkThrowable = true) { currentFile => try { @@ -561,7 +581,7 @@ object ParquetFileFormat extends Logging { if (archiveEnabled && SupportsArchiveFormat.isArchivePath(currentFile.getPath)) { // An archive is one file here; read each of its Parquet entries' footers (the archive is // atomic under ignoreCorruptFiles, see readArchiveFooters). - readArchiveFooters(conf, currentFile) + readArchiveFooters(conf, currentFile, archivePathFilter) } else { Seq(new Footer(currentFile.getPath, ParquetFooterReader.readFooter( @@ -591,13 +611,17 @@ object ParquetFileFormat extends Logging { } /** Reads every Parquet entry's footer in one archive. */ - private def readArchiveFooters(conf: Configuration, archive: FileStatus): Seq[Footer] = { + private def readArchiveFooters( + conf: Configuration, + archive: FileStatus, + archivePathFilter: Option[GlobPattern]): Seq[Footer] = { val tempDir = Utils.createTempDir(Utils.getLocalDir(SparkEnv.get.conf), "parquet-archive-infer") // localizeEntries eagerly opens/copies the first entry, so build it inside the try -- a corrupt // archive throws there and the finally must still delete tempDir. var entries: Iterator[(String, File)] = Iterator.empty try { - entries = SupportsArchiveFormat.localizeEntries(archive.getPath, conf, tempDir, _ => true) + entries = SupportsArchiveFormat.localizeEntries( + archive.getPath, conf, tempDir, _ => true, archivePathFilter) entries.map { case (_, entryFile) => try { val status = new FileStatus(entryFile.length(), false, 0, 0, entryFile.lastModified(), @@ -654,6 +678,11 @@ object ParquetFileFormat extends Logging { timestampNanosTypesEnabled = timestampNanosTypesEnabled, respectUnknownTypeAnnotation = respectUnknownTypeAnnotation) + // readParquetFootersInParallel reads archivePathFilter from the conf (its signature is fixed + // by SchemaMergeUtils' schemaReader type), so put the option there. + new FileSourceOptions(CaseInsensitiveMap(parameters)).archivePathFilter.foreach( + conf.set(FileSourceOptions.ARCHIVE_PATH_FILTER, _)) + readParquetFootersInParallel(conf, files, ignoreCorruptFiles, ignoreMissingFiles) .map(ParquetFileFormat.readSchemaFromFooter(_, converter)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala index f60ced3eb5973..0adece577c58b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilters.scala @@ -37,11 +37,14 @@ import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName._ import org.apache.parquet.schema.Type.Repetition +import org.apache.spark.sql.catalyst.expressions.variant.{ObjectExtraction, VariantPathParser} import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils, IntervalUtils} import org.apache.spark.sql.catalyst.util.RebaseDateTime.{rebaseGregorianToJulianDays, rebaseGregorianToJulianMicros, RebaseSpec} +import org.apache.spark.sql.execution.datasources.VariantMetadata import org.apache.spark.sql.execution.datasources.parquet.types.ops.{ParquetFilterOps, ParquetTypeOps} import org.apache.spark.sql.internal.LegacyBehaviorPolicy import org.apache.spark.sql.sources +import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.ArrayImplicits._ @@ -56,25 +59,18 @@ class ParquetFilters( pushDownStringPredicate: Boolean, pushDownInFilterThreshold: Int, caseSensitive: Boolean, - datetimeRebaseSpec: RebaseSpec) { + datetimeRebaseSpec: RebaseSpec, + variantExtractionSchema: Option[StructType] = None) { + // Shredded-variant physical field-name constants. + private val TYPED_VALUE = "typed_value" + private val VALUE = "value" + // A map which contains parquet field name and data type, if predicate push down applies. // // Each key in `nameToParquetField` represents a column; `dots` are used as separators for // nested columns. If any part of the names contains `dots`, it is quoted to avoid confusion. // See `org.apache.spark.sql.connector.catalog.quote` for implementation details. private val nameToParquetField : Map[String, ParquetPrimitiveField] = { - def getNormalizedLogicalType(p: PrimitiveType): LogicalTypeAnnotation = { - // SPARK-40280: Signed 64 bits on an INT64 and signed 32 bits on an INT32 are optional, but - // the rest of the code here assumes they are not set, so normalize them to not being set. - (p.getPrimitiveTypeName, p.getLogicalTypeAnnotation) match { - case (INT32, intType: IntLogicalTypeAnnotation) - if intType.getBitWidth() == 32 && intType.isSigned() => null - case (INT64, intType: IntLogicalTypeAnnotation) - if intType.getBitWidth() == 64 && intType.isSigned() => null - case (_, otherType) => otherType - } - } - // Recursively traverse the parquet schema to get primitive fields that can be pushed-down. // `parentFieldNames` is used to keep track of the current nested level when traversing. def getPrimitiveFields( @@ -128,6 +124,278 @@ class ParquetFilters( fieldNames: Array[String], fieldType: ParquetSchemaType) + /** + * Holds the mapping from a logical shredded-variant path (e.g. "v.`0`") to the physical + * shredded columns needed to push a sound row-group-skipping predicate. + * + * @param leaf the physical `typed_value` scalar leaf carrying min/max statistics + * @param residualFieldNames the untyped `value` residual columns along the path, from the + * top-level residual down to the leaf's own-level sibling. Each is a + * physical field-name array. Only residuals that exist in this file's + * schema are included; a value for the path can only be hiding in one + * of these residuals when the typed leaf is NULL, which is what the + * pushed predicate's guard checks (see `makeShreddedFilter`). + * Spark's writer (`VariantShreddingWriter.castShredded`) always shreds + * an object field that is in the shredding schema and routes only + * non-schema keys into a level's own `value`, so with Spark-written + * files a value can only fall back to the leaf's own-level residual; + * the ancestor-level residuals guard against writers that legitimately + * decline to shred an intermediate level. + */ + private case class ShreddedVariantField( + leaf: ParquetPrimitiveField, + residualFieldNames: Seq[Array[String]]) + + // Maps logical shredded-variant paths produced by PushVariantIntoScan (e.g. "v.`0`") to the + // physical shredded columns. Populated only when `variantExtractionSchema` is provided and the + // physical file schema actually shreds the requested path. + // + // Soundness: shredding is per-row and per-file best-effort. A row whose value does not fit the + // shredded type (type mismatch or overflow), or whose field is not shredded in this file, is + // stored in an untyped `value` residual with `typed_value` NULL. Parquet min/max excludes NULLs, + // so pushing the predicate on the typed leaf alone could skip a row group that still holds a + // matching row in a residual. To stay sound we push a guarded predicate that skips a row group + // only when the leaf cannot match AND every value for the path is provably in the typed leaf. + // See `makeShreddedFilter`. + // + // Lazy so it is computed after the `Parquet*Type` vals below are initialized (resolution reads + // them via `expectedLeafType`); a strict val here would see them as null under Scala's + // declaration-order initialization. + private lazy val nameToShreddedVariantField: Map[String, ShreddedVariantField] = { + variantExtractionSchema match { + case Some(variantSchema) => + val entries = shreddedVariantEntries( + variantSchema.fields.toSeq, schema.asGroupType(), Array.empty, Array.empty) + if (caseSensitive) { + entries.toMap + } else { + // Mirror `nameToParquetField`: drop names that are ambiguous under case-insensitive + // matching rather than risk pushing a filter on the wrong physical column. + val dedup = entries + .groupBy(_._1.toLowerCase(Locale.ROOT)) + .filter(_._2.size == 1) + .transform((_, v) => v.head._2) + CaseInsensitiveMap(dedup) + } + case None => Map.empty + } + } + + // Look up a child of `group` by name, always case-sensitively. Returns the child type together + // with its actual physical name so callers build paths from the on-disk names. + // + // The match is always exact-case because every caller resolves either a variant object key or a + // structural `typed_value`/`value` name. Variant object keys are data, not Spark identifiers, and + // the reader resolves them case-sensitively (VariantSchema.objectSchemaMap and + // Variant.getFieldByKey use exact equals). A file may legally shred sibling keys differing only + // in case (e.g. `A` and `a`), so a case-insensitive first-match could bind the predicate to the + // wrong physical subtree and skip a row group that holds matching rows -- silent data loss. The + // top-level variant column name is a Spark identifier and is matched by `caseSensitive` + // separately, in `shreddedVariantEntries`. + private def findChild(group: GroupType, name: String): Option[Type] = { + group.getFields.asScala.find(_.getName == name) + } + + // Look up the untyped `value` residual sibling in `group`, if it exists as a non-REPEATED + // primitive. Returns the physical field name. `value` is a fixed structural name; matched exact. + private def residualIn(group: GroupType): Option[String] = + findChild(group, VALUE).collect { + case p: PrimitiveType if p.getRepetition != Repetition.REPEATED => p.getName + } + + // Shared by both the `nameToParquetField` traversal and the shredded leaf resolution so the two + // pushdown paths normalize physical types identically. + private def getNormalizedLogicalType(p: PrimitiveType): LogicalTypeAnnotation = { + // SPARK-40280: Signed 64 bits on an INT64 and signed 32 bits on an INT32 are optional, but + // the rest of the code here assumes they are not set, so normalize them to not being set. + (p.getPrimitiveTypeName, p.getLogicalTypeAnnotation) match { + case (INT32, intType: IntLogicalTypeAnnotation) + if intType.getBitWidth() == 32 && intType.isSigned() => null + case (INT64, intType: IntLogicalTypeAnnotation) + if intType.getBitWidth() == 64 && intType.isSigned() => null + case (_, otherType) => otherType + } + } + + // Navigate the regular shredding layout from a variant column's physical group, resolving both + // the typed leaf and the residual `value` columns along the path. The layout is: + // <col> / typed_value / k0 / typed_value / ... / kN / typed_value (leaf) + // <col> / value (L0 residual) + // <col> / typed_value / k0 / value (L1 residual) + // ... + // <col> / typed_value / k0 / ... / kN / value (leaf-level residual) + // Paths are built from the on-disk field names (via `findChild`). Object keys and the structural + // typed_value/value names are matched case-sensitively (variant keys are data; see `findChild`). + // A value for the path can only be hiding in one of these residual `value` columns when the typed + // leaf is NULL, so IS NOT NULL on all of them is the soundness guard. + // Residuals absent in this file's schema are skipped (that level cannot hold a fallback here). + // Returns None if the file does not shred this path down to a non-REPEATED scalar leaf (nothing + // is pushed and the row group is simply read). + private def resolveShredded( + physCol: GroupType, + physColPath: Array[String], + keys: Array[String], + targetType: DataType): Option[ShreddedVariantField] = { + if (keys.isEmpty) return None + val residuals = scala.collection.mutable.ArrayBuffer.empty[Array[String]] + // L0: the variant column's own residual. + residualIn(physCol).foreach(r => residuals += (physColPath :+ r)) + // Descend key by key: <group>/typed_value/<key>. Collect each level's residual sibling. + var group = physCol + var namePath = physColPath + var idx = 0 + while (idx < keys.length) { + val typedChild = findChild(group, TYPED_VALUE) match { + case Some(g: GroupType) => g + case _ => return None + } + val typedName = typedChild.getName + // Variant object keys are data, matched case-sensitively (see `findChild`). + val keyChild = findChild(typedChild, keys(idx)) match { + case Some(g: GroupType) => g + case _ => return None + } + namePath = namePath ++ Array(typedName, keyChild.getName) + group = keyChild + residualIn(group).foreach(r => residuals += (namePath :+ r)) + idx += 1 + } + // The leaf is the typed_value of the last key group. + findChild(group, TYPED_VALUE) match { + case Some(p: PrimitiveType) if p.getRepetition != Repetition.REPEATED => + val leafType = + ParquetSchemaType(getNormalizedLogicalType(p), p.getPrimitiveTypeName, p.getTypeLength) + // Accept the leaf when the extraction's target type maps to it exactly, or when the leaf is + // a narrower signed integer than the target (safe widening). Narrowing must be rejected: + // for a narrower extraction such as smallint over an int leaf, the leaf min/max is over int + // values, so a row group holding only out-of-int16-range values (residuals null) would be + // skipped, changing an eager INVALID_VARIANT_CAST into an empty result. Widening is sound: + // every value in a narrower leaf casts to the wider target losslessly and ordering is + // preserved, so the leaf stats bound the target predicate; `valueMatchesParquetType` still + // rejects a literal outside the leaf's representable range at push time. + if (!expectedLeafType(targetType).contains(leafType) && + !isSafeIntegerWidening(leafType, targetType)) { + None + } else { + val leaf = ParquetPrimitiveField(namePath :+ p.getName, leafType) + Some(ShreddedVariantField(leaf, residuals.toSeq)) + } + case _ => None + } + } + + // Whether an extraction of `targetType` may be pushed against a physical `leafType` that is a + // narrower signed integer (e.g. bigint extraction over an int/smallint/tinyint leaf). Only the + // integer family widens soundly here. + private def isSafeIntegerWidening( + leafType: ParquetSchemaType, targetType: DataType): Boolean = { + def intRank(t: ParquetSchemaType): Option[Int] = t match { + case ParquetByteType => Some(0) + case ParquetShortType => Some(1) + case ParquetIntegerType => Some(2) + case ParquetLongType => Some(3) + case _ => None + } + (intRank(leafType), expectedLeafType(targetType).flatMap(intRank)) match { + case (Some(leafRank), Some(targetRank)) => leafRank < targetRank + case _ => false + } + } + + // The physical Parquet leaf type a shredded scalar of `targetType` is written as, matching + // `SparkShreddingUtils.variantShreddingSchema` (which writes the scalar's natural type) and the + // `Parquet*Type` normalization used for the leaf. Returns None for types that are not shredded as + // a comparable scalar leaf (or that this pushdown does not handle), so the path is not pushed. + private def expectedLeafType(targetType: DataType): Option[ParquetSchemaType] = targetType match { + case BooleanType => Some(ParquetBooleanType) + case ByteType => Some(ParquetByteType) + case ShortType => Some(ParquetShortType) + case IntegerType => Some(ParquetIntegerType) + case LongType => Some(ParquetLongType) + case FloatType => Some(ParquetFloatType) + case DoubleType => Some(ParquetDoubleType) + case _: StringType => Some(ParquetStringType) + case BinaryType => Some(ParquetBinaryType) + case DateType => Some(ParquetDateType) + case d: DecimalType if DecimalType.is32BitDecimalType(d) => + Some(ParquetSchemaType(LogicalTypeAnnotation.decimalType(d.scale, d.precision), INT32, 0)) + case d: DecimalType if DecimalType.is64BitDecimalType(d) => + Some(ParquetSchemaType(LogicalTypeAnnotation.decimalType(d.scale, d.precision), INT64, 0)) + case d: DecimalType => + Some(ParquetSchemaType(LogicalTypeAnnotation.decimalType(d.scale, d.precision), + FIXED_LEN_BYTE_ARRAY, Decimal.minBytesForPrecision(d.precision))) + case _ => None + } + + // Walk the variant-extraction schema alongside the physical Parquet group, collecting + // logicalName -> ShreddedVariantField entries for shredded scalar object paths that this file + // actually shreds. Only object-extraction, scalar-leaf paths are eligible; array-index paths and + // synthetic (empty / placeholder / companion / full-variant passthrough) paths resolve to None. + // + // `logicalParentNames` accumulates the logical field names (used to build the map key that the + // pushed filter references); `physParentNames` accumulates the on-disk field names (used to build + // the physical Parquet column paths). They differ only in case under case-insensitive matching. + private def shreddedVariantEntries( + variantFields: Seq[StructField], + physGroup: GroupType, + logicalParentNames: Array[String], + physParentNames: Array[String]): Seq[(String, ShreddedVariantField)] = { + import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper + variantFields.flatMap { field => + val physChildOpt = physGroup.getFields.asScala.collectFirst { + case g: GroupType if + (if (caseSensitive) g.getName == field.name + else g.getName.equalsIgnoreCase(field.name)) => g + } + physChildOpt match { + case None => Nil + case Some(physChild) => + val logicalColPath = logicalParentNames :+ field.name + val physColPath = physParentNames :+ physChild.getName + field.dataType match { + // Variant struct: each child is a requested extraction carrying VariantMetadata. + case s: StructType if VariantMetadata.isVariantStruct(s) => + s.fields.toSeq.flatMap { extraction => + if (!extraction.metadata.contains(VariantMetadata.METADATA_KEY)) { + Nil + } else { + val meta = VariantMetadata.fromMetadata(extraction.metadata) + // `VariantPathParser.parse` returns None for an unparseable path; use it directly + // rather than `parsedPath()` (which throws) so a bad path is a clean no-push and + // we don't swallow unrelated exceptions. + VariantPathParser.parse(meta.path) match { + case None => Nil + case Some(segments) => + // Only scalar object-extraction paths are eligible. Reject array-index paths + // (a mix of ObjectExtraction and ArrayExtraction fails this check) and empty + // paths ("$" / passthrough / companion), which yield no keys. + val keys = segments.collect { case o: ObjectExtraction => o.key } + if (keys.isEmpty || keys.length != segments.length) { + Nil + } else { + // `physChild` is this variant column's physical group; `physColPath` holds + // its on-disk name path. Navigate the shredding layout from there. The + // extraction's target type must match the physical leaf type exactly. + resolveShredded(physChild, physColPath, keys, extraction.dataType) match { + case None => Nil + case Some(shredded) => + val logicalName = + (logicalColPath :+ extraction.name).toImmutableArraySeq.quoted + Seq(logicalName -> shredded) + } + } + } + } + } + // Ordinary struct: recurse into nested fields. + case s: StructType if !VariantMetadata.isVariantStruct(s) => + shreddedVariantEntries(s.fields.toSeq, physChild, logicalColPath, physColPath) + case _ => Nil + } + } + } + } + private case class ParquetSchemaType( logicalTypeAnnotation: LogicalTypeAnnotation, primitiveTypeName: PrimitiveTypeName, @@ -624,6 +892,9 @@ class ParquetFilters( } else { Some(sources.Or(leftResultOptional.get, rightResultOptional.get)) } + // A negated shredded-variant predicate cannot be pushed soundly (see the matching guard in + // `createFilterHelper`), so it is not convertible either. + case sources.Not(pred) if referencesShreddedName(pred) => None case sources.Not(pred) => val resultOptional = convertibleFiltersHelper(pred, canPartialPushDown = false) resultOptional.map(sources.Not) @@ -647,7 +918,12 @@ class ParquetFilters( // Parquet's type in the given file should be matched to the value's type // in the pushed filter in order to push down the filter to Parquet. private def valueCanMakeFilterOn(name: String, value: Any): Boolean = { - value == null || (nameToParquetField(name).fieldType match { + valueMatchesParquetType(nameToParquetField(name).fieldType, value) + } + + // The value's type must match the field's on-disk Parquet type for a filter to be pushed. + private def valueMatchesParquetType(fieldType: ParquetSchemaType, value: Any): Boolean = { + value == null || (fieldType match { case ParquetBooleanType => value.isInstanceOf[JBoolean] case ParquetIntegerType if value.isInstanceOf[Period] => true case ParquetByteType | ParquetShortType | ParquetIntegerType => value match { @@ -692,6 +968,74 @@ class ParquetFilters( nameToParquetField.contains(name) && valueCanMakeFilterOn(name, value) } + // Whether `name` is a shredded-variant logical path whose typed leaf accepts `value`. `value` + // must be non-null: shredded pushdown only handles comparison predicates. + private def canMakeShreddedFilterOn(name: String, value: Any): Boolean = { + value != null && nameToShreddedVariantField.get(name).exists { f => + valueMatchesParquetType(f.leaf.fieldType, value) + } + } + + // Whether `predicate` references a shredded-variant logical path anywhere. Used to refuse + // conversion under negation: the shredded predicate is `or(leaf, and(anyResidualNotNull, + // isNull(leaf)))`, and parquet-mr's LogicalInverseRewriter pushes `not(...)` inside to + // `and(not(leaf), or(and(eq(residual, null)...), notEq(leaf, null)))`, which is row-group- + // droppable as soon as some residual has no nulls AND the leaf is entirely NULL -- exactly an + // all-fallback row group, whose matching values are all in a residual. Since a negated shredded + // predicate cannot be expressed soundly with row-group statistics, we do not push it at all. + // + // `sources.Filter.references` already recurses through And/Or/Not and every leaf filter, so this + // stays correct if new Filter subtypes are added. + private def referencesShreddedName(predicate: sources.Filter): Boolean = + predicate.references.exists(nameToShreddedVariantField.contains) + + // Build the sound shredded-variant predicate: + // or(leafPredicate, and(anyResidualNotNull, isNull(leaf))) + // where `anyResidualNotNull` is `or(notEq(residual_0, null), ..., notEq(residual_n, null))` and + // `isNull(leaf)` is `eq(leaf, null)`. + // + // Parquet's statistics drop logic: `or(a, b)` is row-group-droppable iff BOTH `a` and `b` are + // droppable; `and(a, b)` iff EITHER is; `notEq(col, null)` (IS NOT NULL) iff the column is + // entirely NULL (no non-nulls); `eq(col, null)` (IS NULL) iff the column has no nulls. So the + // whole `or` drops the row group iff the leaf min/max cannot match AND (every residual is + // entirely NULL OR the leaf column has no nulls). The second arm is what makes this sound and + // still effective: a value for the path can be outside the typed leaf only on a row where the + // leaf is NULL, so a leaf with zero nulls means every value is provably in the typed leaf and the + // leaf min/max is a complete summary -- regardless of what the residual columns hold (they may be + // non-null because a sibling key outside the shredding schema landed in the level's `value`, + // which is the normal layout for real Variant data). Per record it still keeps every row that + // could match: a row whose value fell back to a residual has a NULL leaf and a non-null residual, + // so `and(anyResidualNotNull, isNull(leaf))` holds for it. + // + // The naive `and(leafPredicate, isNull(residual))` is UNSOUND: `and` drops iff EITHER conjunct is + // droppable, so the leaf predicate alone would drop the row group regardless of the residual. The + // earlier flat `or(leaf, isNotNull(residual)...)` was sound but could never drop a row group once + // any residual was non-null (e.g. a partial object), i.e. it paid the pushdown cost without ever + // skipping on that common layout. + // + // `makeLeaf` produces the leaf predicate from the leaf's field-name array; it returns None if the + // leaf type has no comparison encoding. + private def makeShreddedFilter( + name: String, + makeLeaf: (ParquetSchemaType, Array[String]) => Option[FilterPredicate] + ): Option[FilterPredicate] = { + val field = nameToShreddedVariantField(name) + makeLeaf(field.leaf.fieldType, field.leaf.fieldNames).map { leafPredicate => + field.residualFieldNames + .map(n => FilterApi.notEq(binaryColumn(n), null.asInstanceOf[Binary])) + .reduceLeftOption[FilterPredicate](FilterApi.or) match { + case None => + // No residuals exist in this file for the path, so the leaf is a complete summary. + leafPredicate + case Some(anyResidualNotNull) => + // `makeLeaf` returned Some, so the leaf type is one `makeEq` also covers (its case list + // is a superset of the comparison ops), hence `makeEq.lift` is defined here. + val leafIsNull = makeEq.lift(field.leaf.fieldType).get(field.leaf.fieldNames, null) + FilterApi.or(leafPredicate, FilterApi.and(anyResidualNotNull, leafIsNull)) + } + } + } + /** * @param predicate the input filter predicates. Not all the predicates can be pushed down. * @param canPartialPushDownConjuncts whether a subset of conjuncts of predicates can be pushed @@ -718,6 +1062,36 @@ class ParquetFilters( // Probably I missed something and obviously this should be changed. predicate match { + // Shredded-variant paths (e.g. "v.`0`"). Only comparison predicates that use min/max + // statistics are eligible. Each pushes the guarded predicate built by `makeShreddedFilter`. + // IS NULL / IS NOT NULL on the logical variant field are intentionally out of scope: "the + // extracted field is null" is not the same as "typed_value is null", so we must not conflate + // them. + case sources.EqualTo(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeEq.lift(t).map(_(n, value))) + case sources.EqualNullSafe(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeEq.lift(t).map(_(n, value))) + case sources.LessThan(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeLt.lift(t).map(_(n, value))) + case sources.LessThanOrEqual(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeLtEq.lift(t).map(_(n, value))) + case sources.GreaterThan(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeGt.lift(t).map(_(n, value))) + case sources.GreaterThanOrEqual(name, value) if canMakeShreddedFilterOn(name, value) => + makeShreddedFilter(name, (t, n) => makeGtEq.lift(t).map(_(n, value))) + case sources.In(name, values) if pushDownInFilterThreshold > 0 && values.nonEmpty && + values.forall(v => canMakeShreddedFilterOn(name, v)) => + // Build the leaf predicate once (an OR of per-value equalities under the threshold, or a + // single FilterApi.in above it), then OR the residual isNotNull guards on once via + // `makeShreddedFilter`. Mirrors the regular `In` path below: threshold on `values.length`, + // and `makeInPredicate`/`FilterApi.in` for large lists so those still get skipping. + makeShreddedFilter(name, (t, n) => + if (values.length <= pushDownInFilterThreshold) { + values.distinct.flatMap(v => makeEq.lift(t).map(_(n, v))).reduceLeftOption(FilterApi.or) + } else { + makeInPredicate.lift(t).map(_(n, values)) + }) + case sources.IsNull(name) if canMakeFilterOn(name, null) => makeEq.lift(nameToParquetField(name).fieldType) .map(_(nameToParquetField(name).fieldNames, null)) @@ -796,12 +1170,23 @@ class ParquetFilters( rhsFilter <- createFilterHelper(rhs, canPartialPushDownConjuncts) } yield FilterApi.or(lhsFilter, rhsFilter) + // Refuse to push a negated predicate that references a shredded-variant path: not() of the + // guarded shredded predicate is rewritten by parquet-mr into a form that can drop an + // all-fallback row group (see `referencesShreddedName`). + case sources.Not(pred) if referencesShreddedName(pred) => None case sources.Not(pred) => createFilterHelper(pred, canPartialPushDownConjuncts = false) .map(FilterApi.not) + // Every value must be pushable, not just the head: both branches below convert *all* the + // values (per-element `makeEq` when under the threshold, `makeInPredicate` otherwise), so a + // head-only check lets a non-pushable tail element reach the converter. For a value-range- + // sensitive type (e.g. nanosecond timestamps, whose encoder throws outside the int64 + // epoch-nanos range) that would crash filter creation instead of falling back to a full + // scan. `forall` short-circuits and, for type-only `valueCanMakeFilterOn` checks, is + // equivalent to the previous head check (all `In` values share the coerced column type). case sources.In(name, values) if pushDownInFilterThreshold > 0 && values.nonEmpty && - canMakeFilterOn(name, values.head) => + values.forall(canMakeFilterOn(name, _)) => val fieldType = nameToParquetField(name).fieldType val fieldNames = nameToParquetField(name).fieldNames if (values.length <= pushDownInFilterThreshold) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala index 7e67f14e0ad97..af143955ed707 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetUtils.scala @@ -344,7 +344,7 @@ object ParquetUtils extends Logging { agg match { case max: Max if V2ColumnUtils.extractV2Column(max.column).isDefined => val colName = V2ColumnUtils.extractV2Column(max.column).get - index = dataSchema.fieldNames.toList.indexOf(colName) + index = dataSchema.getFieldIndex(colName).getOrElse(-1) schemaName = "max(" + colName + ")" val currentMax = getCurrentBlockMaxOrMin(filePath, blockMetaData, index, true) if (value == None || currentMax.asInstanceOf[Comparable[Any]].compareTo(value) > 0) { @@ -352,7 +352,7 @@ object ParquetUtils extends Logging { } case min: Min if V2ColumnUtils.extractV2Column(min.column).isDefined => val colName = V2ColumnUtils.extractV2Column(min.column).get - index = dataSchema.fieldNames.toList.indexOf(colName) + index = dataSchema.getFieldIndex(colName).getOrElse(-1) schemaName = "min(" + colName + ")" val currentMin = getCurrentBlockMaxOrMin(filePath, blockMetaData, index, false) if (value == None || currentMin.asInstanceOf[Comparable[Any]].compareTo(value) < 0) { @@ -363,12 +363,12 @@ object ParquetUtils extends Logging { schemaName = "count(" + colName + ")" rowCount += block.getRowCount var isPartitionCol = false - if (partitionSchema.fields.map(_.name).toSet.contains(colName)) { + if (partitionSchema.getFieldIndex(colName).isDefined) { isPartitionCol = true } isCount = true if (!isPartitionCol) { - index = dataSchema.fieldNames.toList.indexOf(colName) + index = dataSchema.getFieldIndex(colName).getOrElse(-1) // Count(*) includes the null values, but Count(colName) doesn't. rowCount -= getNumNulls(filePath, blockMetaData, index) } @@ -428,7 +428,7 @@ object ParquetUtils extends Logging { "filePath" -> filePath, "config" -> PARQUET_AGGREGATE_PUSHDOWN_ENABLED.key)) } - statistics.getNumNulls; + statistics.getNumNulls } // Replaces each VariantType in the schema with the corresponding type in the shredding schema. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/ParquetTypeOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/ParquetTypeOps.scala index 7c2091c8c0810..f839e6c90043a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/ParquetTypeOps.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/ParquetTypeOps.scala @@ -306,5 +306,8 @@ private[parquet] object ParquetTypeOps { * predicate pushdown lists its [[ParquetFilterOps]] here. This is what `filterOpsFor` * scans, so a new type participates in pushdown by adding its ops to this Seq. */ - private val filterOpsList: Seq[ParquetFilterOps] = Seq(TimeTypeParquetOps.filterOps) + private val filterOpsList: Seq[ParquetFilterOps] = Seq( + TimeTypeParquetOps.filterOps, + TimestampNanosParquetOps.ltzFilterOps, + TimestampNanosParquetOps.ntzFilterOps) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala index 6eb19f5f34b01..04f6633afd125 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOps.scala @@ -17,6 +17,9 @@ package org.apache.spark.sql.execution.datasources.parquet.types.ops +import java.lang.{Long => JLong} +import java.time.{Instant, LocalDateTime} + import org.apache.parquet.column.{ColumnDescriptor, Dictionary} import org.apache.parquet.io.api.{Converter, RecordConsumer} import org.apache.parquet.schema.{LogicalTypeAnnotation, Type, Types} @@ -185,6 +188,62 @@ private[ops] object TimestampNanosParquetOps { case ts: TimestampLogicalTypeAnnotation => ts.getUnit == TimeUnit.NANOS case _ => false }) + + // Repacks an externalized nanos filter value into the signed INT64 epoch-nanoseconds the write + // path produces. Conversion is at precision 9 (a lossless repack): the literal has already been + // floored to the column precision upstream, so no sub-microsecond digits are dropped here. The + // single-arg `timestampNanosToEpochNanos` throws `ArithmeticException` outside the int64 range; + // callers reach it only after `acceptsValue` has cleared the value (see [[epochNanosInRange]]). + private def instantToEpochNanos(v: Instant): JLong = + DateTimeUtils.timestampNanosToEpochNanos( + DateTimeUtils.instantToTimestampNanos(v, TimestampLTZNanosType.NANOS_PRECISION)) + + private def localDateTimeToEpochNanos(v: LocalDateTime): JLong = + DateTimeUtils.timestampNanosToEpochNanos( + DateTimeUtils.localDateTimeToTimestampNanos(v, TimestampNTZNanosType.NANOS_PRECISION)) + + // SPARK-46092-style guard: only push down when the value is representable as int64 + // epoch-nanoseconds. An out-of-range value would throw in the encoder, and -- worse -- a + // wrapped/truncated encoding could silently mis-skip row groups; rejecting it falls back to a + // full scan, which is always correct. + private def epochNanosInRange(encode: => JLong): Boolean = + try { encode; true } catch { case _: ArithmeticException => false } + + /** + * Parquet filter-pushdown ops for the nanosecond timestamp types, registered in + * [[ParquetTypeOps.filterOpsList]]. Filter dispatch is keyed on the file's on-disk encoding, so + * each type gets its own ops: both are stored as INT64 TIMESTAMP(NANOS) and differ only in the + * `isAdjustedToUTC` flag (LTZ = true, NTZ = false), which also fixes the externalized filter + * value (`java.time.Instant` for LTZ, `java.time.LocalDateTime` for NTZ). Values are encoded to + * the same signed INT64 epoch-nanoseconds `TimestampNanosParquetOps` writes, never truncated to + * micros. This replaces the inline nanos arms once carried in `ParquetFilters`, matching how + * TimeType routes its pushdown through [[TimeTypeParquetOps.filterOps]]. + */ + private[ops] val ltzFilterOps: ParquetFilterOps = new LongParquetFilterOps { + override val logicalTypeAnnotation: LogicalTypeAnnotation = + LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS) + + override def acceptsValue(value: Any): Boolean = value match { + case i: Instant => epochNanosInRange(instantToEpochNanos(i)) + case _ => false + } + + override protected def toLong(value: Any): JLong = + instantToEpochNanos(value.asInstanceOf[Instant]) + } + + private[ops] val ntzFilterOps: ParquetFilterOps = new LongParquetFilterOps { + override val logicalTypeAnnotation: LogicalTypeAnnotation = + LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS) + + override def acceptsValue(value: Any): Boolean = value match { + case ldt: LocalDateTime => epochNanosInRange(localDateTimeToEpochNanos(ldt)) + case _ => false + } + + override protected def toLong(value: Any): JLong = + localDateTimeToEpochNanos(value.asInstanceOf[LocalDateTime]) + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/TextFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/TextFileFormat.scala index c220dd5a957f9..62d2cd4a2eddc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/TextFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/text/TextFileFormat.scala @@ -142,7 +142,9 @@ case class TextFileFormat() textOptions: TextOptions): PartitionedFile => Iterator[UnsafeRow] = { (file: PartitionedFile) => { val confValue = conf.value.value - SupportsArchiveFormat.readArchiveEntries(file.toPath, confValue) { (_, in) => + val entryGlob = textOptions.archivePathFilterPattern + SupportsArchiveFormat.readArchiveEntries( + file.toPath, confValue, archivePathFilter = entryGlob) { (_, in) => // Each entry is read as a standalone text file, so it gets its own row writer, exactly as // `readToUnsafeMem` builds one per file. val emptyUnsafeRow = new UnsafeRow(0) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index bcdbf5fe6ee29..8805bfe75298f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -171,22 +171,39 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat // Extract scalar subquery filters on runtime-filterable columns for runtime pushdown. // These filters stay in postScanFilters for correctness (FilterExec above scan), // but are also routed into runtimeFilters so BatchScanExec can use them for - // partition pruning via SupportsRuntimeV2Filtering.filter(). + // partition pruning via SupportsRuntimeV2Filtering.filter(). The exceptions are filters + // that only reference attributes the scan fully evaluates, which are dropped from + // postScanFilters below. + // Non-deterministic filters are not routed: they would be pushed to the source for + // pruning while the FilterExec above the scan re-evaluates them, so the two evaluations + // may disagree and rows the source pruned away could not be recovered. This is the + // runtime counterpart of the pushFilters guard in PushDownUtils (SPARK-58207). val scalarSubqueryFilters = if (relation.runtimeFilterAttrs.nonEmpty) { postScanFilters.filter { f => - f.containsPattern(SCALAR_SUBQUERY) && + f.deterministic && + f.containsPattern(SCALAR_SUBQUERY) && f.references.nonEmpty && f.references.subsetOf(relation.runtimeFilterAttrs) } } else { Seq.empty } + // Screen with the same test pushdown applies, or a filter dropped here and rejected there + // would be evaluated nowhere. + val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f => + f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) && + PushDownUtils.isPushablePartitionFilter(f, includeSubquery = true) + } + // dynamicFilters need no such check: a DynamicPruningSubquery over a non-deterministic + // filtering plan is itself non-deterministic, so CleanupDynamicPruningFilters has already + // rewritten it to TrueLiteral by the time we get here. val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters, relation.ordering, relation.relation.table, relation.keyGroupedPartitioning) DataSourceV2Strategy.withProjectAndFilter( - project, postScanFilters, batchExec, !batchExec.supportsColumnar) :: Nil + project, postScanFilters.diff(fullyPushedRuntimeFilters), + batchExec, !batchExec.supportsColumnar) :: Nil case PhysicalOperation(p, f, r: StreamingDataSourceV2ScanRelation) if r.startOffset.isDefined && r.endOffset.isDefined => @@ -955,6 +972,9 @@ private[sql] object DataSourceV2Strategy extends Logging { * If the underlying subquery hasn't completed yet, this method will throw an exception. */ protected[sql] def translateRuntimeFilterV2(expr: Expression): Option[Predicate] = expr match { + case TrueLiteral => None + case in: InSubqueryExec if in.isResultUnavailable => + None case in @ InSubqueryExec(PushableColumnAndNestedColumn(name), _, _, _, _, _) => val values = in.values().getOrElse { throw SparkException.internalError( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Utils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Utils.scala index a3b5c5aeb7995..6f7447faa6214 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Utils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Utils.scala @@ -141,7 +141,7 @@ private[sql] object DataSourceV2Utils extends Logging { } val timeTravel = TimeTravelSpec.create( timeTravelTimestamp, timeTravelVersion, conf.sessionLocalTimeZone) - val tbl = CatalogV2Util.getTable(catalog, ident, timeTravel) + val tbl = CatalogV2Util.getTable(catalog, ident, timeTravel, options = dsOptions) (tbl, Some(catalog), Some(ident), timeTravel) case _ => // TODO: Non-catalog paths for DSV2 are currently not well defined. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala index 0cb7967cb7369..a41aad05d4351 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.planning.{GroupBasedRowLevelOperation, PhysicalOperation} import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan, ReplaceData} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.REPLACE_DATA import org.apache.spark.sql.connector.expressions.filter.{Predicate => V2Filter} import org.apache.spark.sql.connector.read.ScanBuilder import org.apache.spark.sql.connector.write.RowLevelOperation.Command.MERGE @@ -40,7 +41,8 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr import DataSourceV2Implicits._ - override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformDownWithPruning( + _.containsPattern(REPLACE_DATA)) { // push down the filter from the command condition instead of the filter in the rewrite plan, // which is negated for data sources that only support replacing groups of data (e.g. files) case GroupBasedRowLevelOperation(rd: ReplaceData, cond, _, relation: DataSourceV2Relation) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index ccff1f5b3311c..6f671bfcb80be 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -101,12 +101,14 @@ case class GroupPartitionsExec( } /** - * Groups and sorts partitions by their keys in ascending order. + * Groups and sorts partitions by their keys in ascending order. The sort must match + * `KeyedPartitioning.toGrouped`, which is why both use + * `KeyedPartitioning.groupedKeyRowOrdering` -- see its documentation for the contract. */ private def groupAndSortByKeys( keyMap: Map[InternalRowComparableWrapper, Seq[Int]], dataTypes: Seq[DataType]) = { - val keyOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes) + val keyOrdering = KeyedPartitioning.groupedKeyRowOrdering(dataTypes) keyMap.toSeq.sorted(keyOrdering.on((t: (InternalRowComparableWrapper, _)) => t._1.row)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala index c052edea53e08..0a0afb73dbb25 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeMetadataOnlyDeleteFromTable.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.expressions.{Expression, PredicateHelper, S import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{DeleteFromTable, DeleteFromTableWithFilters, LogicalPlan, ReplaceData, RowLevelWrite, WriteDelta} import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.{REPLACE_DATA, WRITE_DELTA} import org.apache.spark.sql.connector.catalog.{SupportsDeleteV2, TruncatableTable} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.write.RowLevelOperation @@ -37,7 +38,8 @@ import org.apache.spark.util.ArrayImplicits._ */ object OptimizeMetadataOnlyDeleteFromTable extends Rule[LogicalPlan] with PredicateHelper { - override def apply(plan: LogicalPlan): LogicalPlan = plan transform { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformWithPruning( + _.containsAnyPattern(REPLACE_DATA, WRITE_DELTA)) { case RewrittenRowLevelCommand(rowLevelPlan, DELETE, cond, relation: DataSourceV2Relation) => relation.table match { case table: SupportsDeleteV2 if !SubqueryExpression.hasSubquery(cond) => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala index b74a785aafc3a..fa330212370fb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala @@ -22,7 +22,7 @@ import scala.collection.mutable import org.apache.spark.SparkException import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.logical.SampleMethod import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning} import org.apache.spark.sql.catalyst.types.DataTypeUtils @@ -32,10 +32,10 @@ import org.apache.spark.sql.connector.catalog.Table import org.apache.spark.sql.connector.expressions.{IdentityTransform, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, SampleMethod => SampleMethodV2, Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownLimit, SupportsPushDownOffset, SupportsPushDownRequiredColumns, SupportsPushDownTableSample, SupportsPushDownTopN, SupportsPushDownV2Filters, SupportsRuntimeV2Filtering} -import org.apache.spark.sql.execution.{ScalarSubquery => ExecScalarSubquery} +import org.apache.spark.sql.execution.{InSubqueryExec, ScalarSubquery => ExecScalarSubquery} import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, DataSourceUtils} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters} +import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters, SupportsRuntimeCatalystFiltering} import org.apache.spark.sql.sources import org.apache.spark.sql.types.{StructField, StructType} import org.apache.spark.util.ArrayImplicits.SparkArrayOps @@ -167,8 +167,19 @@ object PushDownUtils extends Logging { * the first pass are used to derive PartitionPredicates in the second pass, avoiding duplicate * pushdown. * - * Note: Do not call multiple times for the same `scan` instance; - * [[SupportsRuntimeV2Filtering.filter]] is mutating. + * Note: `filter` is mutating, and Spark may call this more than once for the same `scan` + * instance: a plan can hold several scan nodes sharing one scan (e.g. the two branches of a + * group-based UPDATE), each pushing its own copy of the runtime filters. Successive calls are + * additive: a scan ANDs the newly pushed predicates with those it already holds. + * + * Note: `runtimeFilters` must not contain non-deterministic filters. A runtime filter is also + * evaluated by the `FilterExec` above the scan, so pushing a non-deterministic one would + * evaluate it twice with different results. `DataSourceV2Strategy` enforces this where + * `runtimeFilters` is built (SPARK-58207). + * + * A scan implementing [[SupportsRuntimeCatalystFiltering]] takes a separate path: all + * runtime filters are pushed as Catalyst expressions in a single call, with no translation to + * connector predicates and no `filterAttributes` gating. The two paths are mutually exclusive. * * @return true if any filters were pushed to the data source */ @@ -178,6 +189,11 @@ object PushDownUtils extends Logging { table: Table, output: Seq[AttributeReference]): Boolean = { scan match { + case _: SupportsRuntimeV2Filtering with SupportsRuntimeCatalystFiltering => + throw SparkException.internalError( + "A scan must not implement both SupportsRuntimeV2Filtering and " + + s"SupportsRuntimeCatalystFiltering, but ${scan.getClass.getName} implements both.") + case filterableScan: SupportsRuntimeV2Filtering if runtimeFilters.nonEmpty => // Push down translatable runtime filters. val filtersToTranslated = runtimeFilters.flatMap { f => @@ -213,6 +229,41 @@ object PushDownUtils extends Logging { } translatedFiltersPushed || partPredicatesPushed + + case catalystScan: SupportsRuntimeCatalystFiltering if runtimeFilters.nonEmpty => + // A runtime filter is normally evaluated twice: the source prunes with it, and the + // FilterExec above the scan applies it again. The two have to agree, so this screen + // pushes only predicates the source can be trusted to evaluate in Spark's place. + // + // But pushing a non-deterministic one would let the source prune on its own coin flip and + // Spark flip again for the rows that survive, so we handle that specially. + // + // Not pushing a filter is safe only while its FilterExec still evaluates it. A fully + // pushed filter has none, so the source is its only evaluator. That is why + // DataSourceV2Strategy deletes a FilterExec at planning time only for a filter we push + // below, and we use the same method (isPushablePartitionFilter) to determine if we + // should push it. + // + // Note: today every filter reaching either site passes this check + // (deleted by DataSourceV2Strategy, and pushed by this method), since runtimeFilters + // holds only deterministic filters (SPARK-58207) and ExtractPythonUDFs has already + // lifted any Python UDF out of the post-scan filters. But sharing the check keeps the two + // decisions consistent if non-deterministic filters reach the site. + // + // A DPP filter degrades to TrueLiteral once its subquery is pruned away, so it matches + // every row. The V2 path above drops these implicitly, since translateRuntimeFilterV2 + // returns None; here we push Catalyst expressions directly, so we remove them explicitly. + val catalystFilters = runtimeFilters + .filter(isPushablePartitionFilter(_, includeSubquery = true)) + .flatMap(unwrapRuntimeFilterExpression) + .filterNot(_ == Literal.TrueLiteral) + if (catalystFilters.nonEmpty) { + catalystScan.filter(catalystFilters.toArray) + true + } else { + false + } + case _ => false } @@ -225,8 +276,8 @@ object PushDownUtils extends Logging { * pre-filter partition set. * * Notes: - * - Do not call multiple times for the same `scan` instance; - * [[SupportsRuntimeV2Filtering.filter]] is mutating. + * - `filter` is mutating, and Spark may call this more than once for the same `scan` instance + * (see [[pushRuntimeFilters]]); successive calls are additive. * - When `outputPartitioning` is a [[KeyedPartitioning]], every split from * `planInputPartitions()` used on this path must implement [[HasPartitionKey]]. * @@ -398,7 +449,7 @@ object PushDownUtils extends Logging { val partitionAttributes = partitionFields.map(_.attrRef) val (partFilters, nonPartitionFilters) = DataSourceUtils.getPartitionFiltersAndDataFilters(partitionAttributes, flattenedFilters) - val (pushable, nonPushable) = partFilters.partition(isPushablePartitionFilter) + val (pushable, nonPushable) = partFilters.partition(isPushablePartitionFilter(_)) val (partitionPredicates, errorPartitionPredicates) = pushable.partitionMap { e => PartitionPredicateImpl(e, partitionFields).toLeft(e) } @@ -433,19 +484,40 @@ object PushDownUtils extends Logging { private[v2] def createRuntimePartitionPredicates( runtimeFilters: Seq[Expression], partitionFields: Seq[PartitionPredicateField]): Seq[PartitionPredicateImpl] = { - val catalystExprs = runtimeFilters.flatMap { + val catalystExprs = runtimeFilters.flatMap(unwrapRuntimeFilterExpression) + val flattened = flattenNestedPartitionFilters(catalystExprs, partitionFields).keys + createPartitionPredicates(flattened.toSeq, partitionFields)._1 + } + + /** Unwraps a runtime filter to the Catalyst predicate for pushdown. */ + private def unwrapRuntimeFilterExpression(rf: Expression): Option[Expression] = + rf match { + case DynamicPruningExpression(in: InSubqueryExec) if in.isResultUnavailable => + None case DynamicPruningExpression(e) => Some(e) case _: DynamicPruning => None case f => Some(f.transform { case s: ExecScalarSubquery => s.toLiteral }) } - val flattened = flattenNestedPartitionFilters(catalystExprs, partitionFields).keys - createPartitionPredicates(flattened.toSeq, partitionFields)._1 - } - private def isPushablePartitionFilter(f: Expression) = + /** + * Whether the data source can be trusted to evaluate `f` in place of Spark: a non-deterministic + * expression would not give the same answer twice, and a Python UDF or a subquery cannot be + * evaluated by the source at all. Spark will attempt to push only these expressions to the + * data source. + * + * Spark will also call this before dropping a fully pushed runtime filter from the + * post-scan filters, so that it never removes the only evaluator of a filter that Spark + * refuses to push. + * + * @param includeSubquery whether a subquery in `f` should be tolerated. Only safe for a runtime + * filter, which is pushed at execution time, after its subquery has been evaluated. + */ + private[v2] def isPushablePartitionFilter( + f: Expression, + includeSubquery: Boolean = false): Boolean = f.deterministic && - !SubqueryExpression.hasSubquery(f) && - !f.exists(_.isInstanceOf[PythonUDF]) + !f.exists(_.isInstanceOf[PythonUDF]) && + (includeSubquery || !SubqueryExpression.hasSubquery(f)) /** * Replaces all partition column references with canonical [[AttributeReference]]s diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala index 479154e7a2852..165425aea19dc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitioningAndOrdering.scala @@ -21,6 +21,7 @@ import org.apache.spark.internal.LogKeys.CLASS_NAME import org.apache.spark.sql.catalyst.expressions.V2ExpressionUtils import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.DATA_SOURCE_V2_SCAN_RELATION import org.apache.spark.sql.connector.read.{SupportsReportOrdering, SupportsReportPartitioning} import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, UnknownPartitioning} import org.apache.spark.util.ArrayImplicits._ @@ -40,7 +41,8 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { } } - private def partitioning(plan: LogicalPlan) = plan.transformDown { + private def partitioning(plan: LogicalPlan) = plan.transformDownWithPruning( + _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { case d @ ExtractV2ScanInfo(relation, scan: SupportsReportPartitioning, _) if d.keyGroupedPartitioning.isEmpty => val catalystPartitioning = scan.outputPartitioning() match { @@ -68,7 +70,8 @@ object V2ScanPartitioningAndOrdering extends Rule[LogicalPlan] with Logging { d.copy(keyGroupedPartitioning = catalystPartitioning) } - private def ordering(plan: LogicalPlan) = plan.transformDown { + private def ordering(plan: LogicalPlan) = plan.transformDownWithPruning( + _.containsPattern(DATA_SOURCE_V2_SCAN_RELATION)) { case d @ ExtractV2ScanInfo(relation, scan: SupportsReportOrdering, _) => val ordering = V2ExpressionUtils.toCatalystOrdering(scan.outputOrdering(), relation, relation.funCatalog) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index 14d9f754f95c5..d8817b9aa7582 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.datasources.v2 -import java.util.Locale +import java.util.{Locale, OptionalLong} import scala.collection.mutable @@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder} import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg, Count, CountStar, Max, Min, Sum} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, V1Scan, VariantExtraction} +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, Statistics => V2Statistics, SupportsPushDownAggregates, SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns, SupportsPushDownVariantExtractions, SupportsReportStatistics, V1Scan, VariantExtraction} import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, VariantInRelation, VariantMetadata} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.VariantExtractionImpl @@ -69,6 +69,24 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { } } + /** + * Rebuilds a single scan for a Spark-side scan merge (see + * `TableCapability.SCAN_MERGING`): runs this rule on a synthetic + * `Project(projectList, Filter(conditions, relation))` and returns the resulting + * [[DataSourceV2ScanRelation]]. This lets `PlanMerger` fuse two scans of the same table by + * expressing what the merged scan should project and filter, while the pushdown lifecycle (filter + * translation, column pruning, the iterative PartitionPredicate second pass) stays owned here. + * `conditions` are pushed as a single `Filter` directly above the relation; the caller decides + * which of them must come back fully enforced. + */ + def rebuildScan( + relation: DataSourceV2Relation, + projectList: Seq[NamedExpression], + conditions: Seq[Expression]): Option[DataSourceV2ScanRelation] = { + val child = conditions.reduceOption(And).map(Filter(_, relation)).getOrElse(relation) + apply(Project(projectList, child)).collectFirst { case s: DataSourceV2ScanRelation => s } + } + private def collapseGroupedSumOfCount(plan: LogicalPlan): LogicalPlan = { val excludedRules = SQLConf.get.optimizerExcludedRules.toSeq.flatMap(Utils.stringToSeq) if (excludedRules.contains(CollapseGroupedSumOfCount.ruleName)) { @@ -111,8 +129,9 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val postScanFilters = postScanFiltersWithoutSubquery ++ normalizedFiltersWithSubquery // Compute the pushed filter expressions: the normalized filters that were fully pushed - // down (i.e., not in postScanFilters). These are stored on the scan relation for - // potential future use in constraint propagation. + // down (i.e., not in postScanFilters). These are stored on the scan relation for potential + // future use in constraint propagation, and are read by a Spark-side scan merge (see + // TableCapability.SCAN_MERGING) to compare and re-enforce a scan's filters. val postScanFilterSet = ExpressionSet(postScanFiltersWithoutSubquery) sHolder.pushedFilterExpressions = normalizedFiltersWithoutSubquery .filterNot(postScanFilterSet.contains) @@ -826,6 +845,37 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { } } + /** + * Whether the plain-scan path ([[pruneColumns]]) carries a pushdown that a Spark-side scan merge + * (see `TableCapability.SCAN_MERGING`) cannot reproduce by + * rebuilding the scan from a fresh `ScanBuilder` and re-applying the pushed filters and pruned + * columns. Used only to decide the plain-scan `mergeableScan` flag (`!hasBlockingPushdown`). + * + * Scan merging is default-safe: [[DataSourceV2ScanRelation.mergeableScan]] defaults to false and + * this is the ONLY site that ever sets it true. Every other scan-relation build path -- the + * dedicated aggregate/join/variant rules ([[buildScanWithPushedAggregate]] / + * [[buildScanWithPushedJoin]] / [[buildScanWithPushedVariants]]), row-level-operation planning, + * and anything added later -- leaves the scan not-mergeable by default, so those sites need no + * change here. + * + * The plain-scan non-reproducible pushdowns are: + * - `pushedLimit` -- pushed LIMIT + * - `pushedOffset` -- pushed OFFSET + * - `pushedSample` -- pushed table sample + * - `sortOrders` -- pushed sort order (top-N / ordering) + * + * Reproducible (re-applied by the merge), so NOT blocking: `output` (column pruning, via + * `SupportsPushDownRequiredColumns`) and `pushedPredicates` / `pushedFilterExpressions` + * (deterministic filters, re-pushed via `SupportsPushDownV2Filters`). A pushed aggregate, join or + * variant never reaches here -- their build rules run first. When a new pushdown capability adds + * plain-scan state to [[ScanBuilderHolder]], list it above so it keeps the scan not-mergeable. + */ + private def hasBlockingPushdown(holder: ScanBuilderHolder): Boolean = + holder.pushedLimit.isDefined || + holder.pushedOffset.isDefined || + holder.pushedSample.isDefined || + holder.sortOrders.nonEmpty + def buildScanWithPushedAggregate(plan: LogicalPlan): LogicalPlan = plan.transform { case holder: ScanBuilderHolder if holder.pushedAggregate.isDefined => // No need to do column pruning because only the aggregate columns are used as @@ -838,6 +888,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes to aggregate columns. When validConstraints is wired up, this needs revisiting. + // A pushed aggregate cannot be reproduced by rebuilding the scan, so this scan stays + // not-mergeable (the default `mergeableScan = false`). val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) val projectList = realOutput.zip(holder.output).map { case (a1, a2) => // The data source may return columns with arbitrary data types and it's safer to cast them @@ -857,6 +909,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes with pushed join. When validConstraints is wired up, this needs revisiting. + // A pushed join cannot be reproduced by rebuilding the scan, so this scan stays not-mergeable + // (the default `mergeableScan = false`). val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) // When join is pushed down, the real output is going to be, for example, @@ -885,6 +939,8 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { val wrappedScan = getWrappedScan(scan, holder) // Note: holder.pushedFilterExpressions is not propagated here because the output schema // changes with variant extraction. When validConstraints is wired up, this needs revisiting. + // Pushed variant extraction cannot be reproduced by rebuilding the scan, so this scan stays + // not-mergeable (the default `mergeableScan = false`). val scanRelation = DataSourceV2ScanRelation(holder.relation, wrappedScan, realOutput) // Create projection to map real output to expected output (with transformed types) @@ -941,10 +997,11 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { case projectionOverSchema(newExpr) => newExpr } - // Remap pushed filter attributes to the pruned output schema and drop filters - // whose references are no longer in the pruned output. Catch FIELD_NOT_FOUND - // because ProjectionOverSchema throws when a pushed filter references a nested - // struct field that was pruned from the schema. + // Remap pushed filter attributes to the pruned output schema and drop filters whose + // references are no longer in the pruned output. Catch FIELD_NOT_FOUND because + // ProjectionOverSchema throws when a pushed filter references a nested struct field that was + // pruned from the schema. This feeds only the Spark post-pushdown adjustment below; the scan + // relation's own pushedFilters keep the complete set (see the next comment). val remappedPushedFilters = sHolder.pushedFilterExpressions.flatMap { filter => try Some(projectionFunc(filter)) catch { @@ -952,22 +1009,40 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { None } }.filter(_.references.subsetOf(AttributeSet(output))) + + // Record the fully-pushed filter expressions on the scan relation, keeping their references + // to the relation's (pre-pruning) output. These include filters on columns that were pruned + // out of the scan output -- e.g. an unselected partition column the source still enforces + // internally. PlanMerger needs this complete set to compare two scans' filters and to + // re-enforce them when it rebuilds a merged scan; remapping to the pruned output (as the + // post-scan filters below are) would silently drop pruned-out filter columns and make the + // merge unsound. See DataSourceV2ScanRelation.pushedFilters. val scanRelation = DataSourceV2ScanRelation(sHolder.relation, wrappedScan, output, - pushedFilters = remappedPushedFilters) + pushedFilters = sHolder.pushedFilterExpressions, + // The one site that grants mergeability: a plain scan carrying only reproducible pushdowns + // (column pruning + deterministic filters) may be fused. See hasBlockingPushdown. + mergeableScan = !hasBlockingPushdown(sHolder)) val finalFilters = normalizedFilters.map(projectionFunc) // bottom-most filters are put in the left of the list. val withFilter = finalFilters.foldLeft[LogicalPlan](scanRelation)((plan, cond) => { Filter(cond, plan) }) + // Best effort: column pruning can make fully-pushed filters unavailable in the scan output. + // `remappedPushedFilters` already drops those filters, so Spark post-pushdown adjustment can + // only re-add predicates that still reference the pruned scan output. + val withPostPushdownAdjustmentFilters = + withSparkPostPushdownAdjustmentFilters(withFilter, remappedPushedFilters) - if (withFilter.output != project) { + if (withPostPushdownAdjustmentFilters.output != project) { val newProjects = normalizedProjects .map(projectionFunc) .asInstanceOf[Seq[NamedExpression]] - Project(restoreOriginalOutputNames(newProjects, project.map(_.name)), withFilter) + Project( + restoreOriginalOutputNames(newProjects, project.map(_.name)), + withPostPushdownAdjustmentFilters) } else { - withFilter + withPostPushdownAdjustmentFilters } } @@ -978,10 +1053,7 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { sample.lowerBound, sample.upperBound, sample.withReplacement, - // TODO(SPARK-56573): The * 1000 limits the seed to only 1000 distinct values. - // Kept here for consistency with SampleExec.resolvedSeed; will be fixed - // across all call sites in SPARK-56573. - sample.seed.getOrElse((math.random() * 1000).toLong), + Sample.resolveSeed(sample.seed), sampleMethod = sample.sampleMethod) val pushed = PushDownUtils.pushTableSample(sHolder.builder, tableSample) if (pushed) { @@ -1161,6 +1233,34 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper { sHolder.joinedRelationsPushedDownOperators, optRelationName) } + private def withSparkPostPushdownAdjustmentFilters( + plan: LogicalPlan, + pushedFilters: Seq[Expression]): LogicalPlan = { + pushedFilters.reduceLeftOption(And) match { + case None => plan + case Some(pushedCondition) => + def shouldAddPushedFilter(scanRelation: DataSourceV2ScanRelation): Boolean = { + scanRelation.scan match { + case s: SupportsReportStatistics => !s.reflectsFullyPushedDownFilters() + case _ => false + } + } + + def addToScan(plan: LogicalPlan): LogicalPlan = plan match { + case Filter(condition, scanRelation: DataSourceV2ScanRelation) + if shouldAddPushedFilter(scanRelation) => + Filter(And(condition, pushedCondition), scanRelation) + case Filter(condition, child) => + Filter(condition, addToScan(child)) + case scanRelation: DataSourceV2ScanRelation if shouldAddPushedFilter(scanRelation) => + Filter(pushedCondition, scanRelation) + case other => other + } + + addToScan(plan) + } + } + } case class ScanBuilderHolder( @@ -1199,6 +1299,30 @@ case class ScanBuilderHolder( case class V1ScanWrapper( v1Scan: V1Scan, handledFilters: Seq[sources.Filter], - pushedDownOperators: PushedDownOperators) extends Scan { + pushedDownOperators: PushedDownOperators) extends Scan with SupportsReportStatistics { override def readSchema(): StructType = v1Scan.readSchema() + + override def estimateStatistics(): V2Statistics = { + v1Scan match { + case r: SupportsReportStatistics => r.estimateStatistics() + case _ => new V2Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.empty() + } + } + } + + override def estimateSizeInBytes(): OptionalLong = { + v1Scan match { + case r: SupportsReportStatistics => r.estimateSizeInBytes() + case _ => OptionalLong.empty() + } + } + + override def reflectsFullyPushedDownFilters(): Boolean = { + v1Scan match { + case r: SupportsReportStatistics => r.reflectsFullyPushedDownFilters() + case _ => true + } + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala index f1ff11b1a4a65..6653a10c87c1f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala @@ -26,11 +26,14 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connector.catalog.{Identifier, Table, TableCatalog, V2TableUtil} import org.apache.spark.sql.connector.catalog.CatalogV2Util import org.apache.spark.sql.errors.QueryCompilationErrors +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.sql.util.SchemaValidationMode import org.apache.spark.sql.util.SchemaValidationMode.ALLOW_NEW_FIELDS import org.apache.spark.sql.util.SchemaValidationMode.PROHIBIT_CHANGES private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { + private type CurrentTableKey = (TableCatalog, Identifier, CaseInsensitiveStringMap) + /** * Refreshes table metadata for tables in the plan. * @@ -81,19 +84,20 @@ private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { plan: LogicalPlan, versionedOnly: Boolean, schemaValidationMode: SchemaValidationMode): LogicalPlan = { - val currentTables = mutable.HashMap.empty[(TableCatalog, Identifier), Table] + val currentTables = mutable.HashMap.empty[CurrentTableKey, Table] plan transformWithSubqueries { case r @ ExtractV2CatalogAndIdentifier(catalog, ident) if (r.isVersioned || !versionedOnly) && r.timeTravelSpec.isEmpty => - val currentTable = currentTables.getOrElseUpdate((catalog, ident), { + val stateOptions = CatalogV2Util.extractTableStateOptions(catalog, r.options) + val currentTable = currentTables.getOrElseUpdate((catalog, ident, stateOptions), { val tableName = V2TableUtil.toQualifiedName(catalog, ident) - lookupCachedRelation(spark, catalog, ident, r.table) match { + lookupCachedRelation(spark, catalog, ident, r.table, r.options) match { case Some(cached) => logDebug(s"Refreshing table metadata for $tableName using shared relation cache") cached.table - case None => + case _ => logDebug(s"Refreshing table metadata for $tableName using catalog") - catalog.loadTable(ident) + CatalogV2Util.getTable(catalog, ident, options = r.options) } }) validateTableIdentity(currentTable, r) @@ -107,8 +111,10 @@ private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { spark: SparkSession, catalog: TableCatalog, ident: Identifier, - table: Table): Option[DataSourceV2Relation] = { - CatalogV2Util.lookupCachedRelation(spark.sharedState.relationCache, catalog, ident, table, conf) + table: Table, + options: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { + CatalogV2Util.lookupCachedRelation( + spark.sharedState.relationCache, catalog, ident, table, options, conf) } // it is not safe to allow any schema changes in commands (e.g. CTAS, RTAS, MERGE) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala index be8e96e8034d2..5340937115d96 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala @@ -129,7 +129,6 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper { commandOptions: Map[String, String], dsOptions: Map[String, String]): Map[String, String] = { // for DataFrame API cases, same options are carried by both Command and DataSourceV2Relation - // for DataFrameV2 API cases, options are only carried by Command // for SQL cases, options are only carried by DataSourceV2Relation assert(commandOptions == dsOptions || commandOptions.isEmpty || dsOptions.isEmpty) commandOptions ++ dsOptions diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala index d280076622b0b..b9d01153e5ee4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala @@ -40,6 +40,7 @@ import org.apache.spark.sql.execution.{EmptyRDDWithPartitions, QueryExecution, S import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLLastAttemptMetric, SQLLastAttemptMetrics, SQLMetric, SQLMetrics} import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.sql.util.SchemaValidationMode.PROHIBIT_CHANGES import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils @@ -98,7 +99,7 @@ case class CreateTableAsSelectExec( .withProperties(properties.asJava) .build() val table = Option(catalog.createTable(ident, tableInfo)) - .getOrElse(catalog.loadTable(ident, Set(TableWritePrivilege.INSERT).asJava)) + .getOrElse(loadTableForInsert(catalog, ident, writeOptions)) val result = writeToTable(catalog, table, writeOptions, ident, query, overwrite = false) transaction.foreach(TransactionUtils.commit) result @@ -142,7 +143,7 @@ case class AtomicCreateTableAsSelectExec( .withProperties(properties.asJava) .build() val stagedTable = Option(catalog.stageCreate(ident, tableInfo) - ).getOrElse(catalog.loadTable(ident, Set(TableWritePrivilege.INSERT).asJava)) + ).getOrElse(loadTableForInsert(catalog, ident, writeOptions)) writeToTable(catalog, stagedTable, writeOptions, ident, query, overwrite = false) } } @@ -205,7 +206,7 @@ case class ReplaceTableAsSelectExec( .withProperties(properties.asJava) .build() val table = Option(catalog.createTable(ident, tableInfo)) - .getOrElse(catalog.loadTable(ident, Set(TableWritePrivilege.INSERT).asJava)) + .getOrElse(loadTableForInsert(catalog, ident, writeOptions)) val result = writeToTable( catalog, table, writeOptions, ident, refreshedQuery, overwrite = true, refreshPhaseEnabled = false) @@ -274,8 +275,7 @@ case class AtomicReplaceTableAsSelectExec( ident, CatalogV2Util.searchPathForTableIdentifier(catalog, ident)) } - val table = Option(staged).getOrElse( - catalog.loadTable(ident, Set(TableWritePrivilege.INSERT).asJava)) + val table = Option(staged).getOrElse(loadTableForInsert(catalog, ident, writeOptions)) writeToTable(catalog, table, writeOptions, ident, query, overwrite = true) } } @@ -964,6 +964,15 @@ case class DeltaWithMetadataWritingSparkTask( private[v2] trait V2CreateTableAsSelectBaseExec extends LeafV2CommandExec { override def output: Seq[Attribute] = Nil + protected def loadTableForInsert( + catalog: TableCatalog, + ident: Identifier, + writeOptions: Map[String, String]): Table = { + val options = new CaseInsensitiveStringMap(writeOptions.asJava) + CatalogV2Util.loadTableForV2Write( + catalog, ident, Set(TableWritePrivilege.INSERT), options) + } + protected def getV2Columns(schema: StructType, forceNullable: Boolean): Array[Column] = { val rawSchema = CharVarcharUtils.getRawSchema(removeInternalMetadata(schema), conf) val tableSchema = if (forceNullable) rawSchema.asNullable else rawSchema @@ -979,7 +988,9 @@ private[v2] trait V2CreateTableAsSelectBaseExec extends LeafV2CommandExec { overwrite: Boolean, refreshPhaseEnabled: Boolean = true): Seq[InternalRow] = { Utils.tryWithSafeFinallyAndFailureCallbacks({ - val relation = DataSourceV2Relation.create(table, Some(catalog), Some(ident)) + val tableOptions = new CaseInsensitiveStringMap(writeOptions.asJava) + val relation = + DataSourceV2Relation.create(table, Some(catalog), Some(ident), tableOptions) val writeCommand = if (overwrite) { OverwriteByExpression.byPosition(relation, query, Literal.TrueLiteral, writeOptions) } else { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/csv/CSVTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/csv/CSVTable.scala index 184eb41b6c49f..e7b206bd9c58c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/csv/CSVTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/csv/CSVTable.scala @@ -46,9 +46,7 @@ case class CSVTable( columnPruning = sparkSession.sessionState.conf.csvColumnPruning, sparkSession.sessionState.conf.sessionLocalTimeZone) - // The DSv2 reader does not route archives to `readArchive` (it calls `readFile` directly), so - // archive scans aren't supported here; pass supportsArchiveScan = false so an archive input - // keeps failing with UNABLE_TO_INFER_SCHEMA rather than having its raw bytes parsed as CSV. + // The DSv2 scan cannot read archives, so refuse archive inference here. CSVDataSource(parsedOptions) .inferSchema(sparkSession, files, parsedOptions, supportsArchiveScan = false) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala index c6eec4487f5b7..bcc332f4bd3cc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCScanBuilder.scala @@ -148,7 +148,7 @@ case class JDBCScanBuilder( JDBCOptions.JDBC_QUERY_STRING filteredJDBCOptions == otherSideFilteredJDBCOptions - }; + } /** * Helper method to calculate StructType based on the SupportsPushDownJoin.ColumnWithAlias and diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/json/JsonTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/json/JsonTable.scala index 74095e85a0f6c..c635a7d06beda 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/json/JsonTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/json/JsonTable.scala @@ -45,8 +45,6 @@ case class JsonTable( options.asScala.toMap, sparkSession.sessionState.conf.sessionLocalTimeZone, sparkSession.sessionState.conf.columnNameOfCorruptRecord) - // The DSv2 reader calls `readFile` directly and cannot read archives, so refuse to infer a - // schema for archive inputs (supportsArchiveScan = false) rather than mis-reading raw bytes. JsonDataSource(parsedOptions).inferSchema( sparkSession, files, parsedOptions, supportsArchiveScan = false) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScanBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScanBuilder.scala index 149d7e6f0b720..e702cfc3ad256 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScanBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/parquet/ParquetScanBuilder.scala @@ -66,6 +66,15 @@ case class ParquetScanBuilder( val isCaseSensitive = sqlConf.caseSensitiveAnalysis val parquetSchema = new SparkToParquetSchemaConverter(sparkSession.sessionState.conf).convert(readDataSchema()) + // Shredded-variant predicate pushdown (SPARK-55817) is not wired here: it applies to the + // DSv1 path only. DSv2 does rewrite variant extractions into `v.`0`` struct accesses, but + // only in `V2ScanRelationPushDown.buildScanWithPushedVariants`, which runs *after* + // `pushDownFilters`. So the filters reaching this method are still `variant_get(v, ...)` + // predicates, which do not translate to a source `Filter` at all -- there is no + // shredded-variant logical name for ParquetFilters to resolve here, and nothing would be + // reported convertible even with a variantExtractionSchema. DSv2 reads remain correct (the + // variant filter is applied post-scan); they just do not get row-group skipping on shredded + // columns. val parquetFilters = new ParquetFilters( parquetSchema, pushDownDate, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala index 7c113c1cb03a9..9fcf81746fe30 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonDataSourceV2.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2.python import org.apache.spark.sql.SparkSession import org.apache.spark.sql.connector.catalog._ import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -53,24 +54,35 @@ class PythonDataSourceV2 extends TableProvider { } private var readInfo: PythonDataSourceReadInfo = _ + // The pushdown flags in effect when `readInfo` was cached, as (filter, limit). The cached read + // info is identical regardless of these flags, but planning it also validates -- via + // DATA_SOURCE_PUSHDOWN_DISABLED -- that the reader does not implement a pushdown method whose + // config is off. That validation runs only when the planning worker runs, so reusing the cache + // across a change in these flags (e.g. a reused relation scanned first with pushdown on, then + // off) would skip it. Track the flags and recompute when they differ so the check always runs + // for the flags currently in effect. + private var readInfoPushdownFlags: Option[(Boolean, Boolean)] = None + // Caches the pushdown-free read info for reads that push nothing down. This is safe to keep on + // the (provider-scoped) data source because every such scan produces the same full read info. + // Pushdown-specific read info is never stored here; it is carried by the `PythonScan` instead, + // so that it cannot leak across scans that share this data source (e.g. a base DataFrame and + // its `.limit(n)`). def getOrCreateReadInfo( shortName: String, options: CaseInsensitiveStringMap, outputSchema: StructType, isStreaming: Boolean ): PythonDataSourceReadInfo = { - if (readInfo == null) { + val flags = (SQLConf.get.pythonFilterPushDown, SQLConf.get.pythonLimitPushDown) + if (readInfo == null || !readInfoPushdownFlags.contains(flags)) { val creationResult = getOrCreateDataSourceInPython(shortName, options, Some(outputSchema)) readInfo = source.createReadInfoInPython(creationResult, outputSchema, isStreaming) + readInfoPushdownFlags = Some(flags) } readInfo } - def setReadInfo(readInfo: PythonDataSourceReadInfo): Unit = { - this.readInfo = readInfo - } - override def inferSchema(options: CaseInsensitiveStringMap): StructType = { getOrCreateDataSourceInPython(shortName, options, None).schema } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala index 9e3effe7d441d..b911659432272 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScan.scala @@ -31,9 +31,13 @@ class PythonScan( shortName: String, outputSchema: StructType, options: CaseInsensitiveStringMap, - supportedFilters: Array[Filter] + supportedFilters: Array[Filter], + pushedLimit: Option[Int] = None, + // Read info computed during filter/limit pushdown, if any. Scoped to this scan so it does not + // leak across scans that share the same `PythonDataSourceV2` (see `PythonScanBuilder`). + readInfo: Option[PythonDataSourceReadInfo] = None ) extends Scan with SupportsMetadata { - override def toBatch: Batch = new PythonBatch(ds, shortName, outputSchema, options) + override def toBatch: Batch = new PythonBatch(ds, shortName, outputSchema, options, readInfo) override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { val runner = PythonMicroBatchStream.createPythonStreamingSourceRunner( @@ -67,7 +71,7 @@ class PythonScan( Map( "PushedFilters" -> supportedFilters.mkString("[", ", ", "]"), "ReadSchema" -> outputSchema.simpleString - ) + ) ++ pushedLimit.map(limit => "PushedLimit" -> s"LIMIT $limit") } } @@ -75,7 +79,11 @@ class PythonBatch( ds: PythonDataSourceV2, shortName: String, outputSchema: StructType, - options: CaseInsensitiveStringMap) extends Batch { + options: CaseInsensitiveStringMap, + // Read info already computed during pushdown, if any. When empty (no pushdown), it is + // computed lazily below via the provider, which is safe because that path always produces + // the full, pushdown-free read info. + readInfo: Option[PythonDataSourceReadInfo] = None) extends Batch { private val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) private val sessionUUID = { SparkSession.getActiveSession.collect { @@ -85,7 +93,8 @@ class PythonBatch( } private lazy val infoInPython: PythonDataSourceReadInfo = { - ds.getOrCreateReadInfo(shortName, options, outputSchema, isStreaming = false) + readInfo.getOrElse( + ds.getOrCreateReadInfo(shortName, options, outputSchema, isStreaming = false)) } override def planInputPartitions(): Array[InputPartition] = diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala index 3dabbcb8af05b..2ea1bec70cd16 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/PythonScanBuilder.scala @@ -16,7 +16,8 @@ */ package org.apache.spark.sql.execution.datasources.v2.python -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters} +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownLimit} +import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.sources.Filter import org.apache.spark.sql.types.StructType @@ -29,25 +30,59 @@ class PythonScanBuilder( outputSchema: StructType, options: CaseInsensitiveStringMap) extends ScanBuilder - with SupportsPushDownFilters { + with SupportsPushDownFilters + with SupportsPushDownLimit { private var supportedFilters: Array[Filter] = Array.empty + // All filters handed to `pushFilters`, kept so that `pushLimit` can replay them and bring the + // Python reader back to the same state before calling `pushLimit` on it. + private var allFilters: Array[Filter] = Array.empty + private var pushedLimit: Option[Int] = None + // Read info (partitions + read function) produced as a side effect of filter/limit pushdown, + // carried into the `PythonScan` so it stays scoped to this scan. It must NOT be stored on the + // provider-scoped `PythonDataSourceV2`: a single data source instance is shared by every scan + // built from it -- e.g. a base DataFrame and its `.limit(n)` reuse the same relation -- so a + // pushdown-specific read function stored there would leak into an unrelated scan and make it + // read too few rows. + private var readInfo: Option[PythonDataSourceReadInfo] = None + // True when a pushdown pass ran but produced no read info, so build() must plan the read + // itself. Two paths set it: (1) the filter-pushdown pass defers planning while limit pushdown + // is enabled and a limit pass might follow; (2) a pushed limit was rejected, so planning falls + // back to a fresh reader. In both, build() plans the read with the filters only. This is + // distinct from `readInfo.isEmpty` when both pushdowns are disabled -- in that case, no + // pushdown pass ran and `PythonScan` plans the read later instead, so build() must not plan. + private var deferredPlanning: Boolean = false - override def build(): Scan = - new PythonScan(ds, shortName, outputSchema, options, supportedFilters) + override def build(): Scan = { + if (deferredPlanning && readInfo.isEmpty) { + // Reached when a pushdown pass ran but planned no read info: either the filter pass + // deferred planning and no limit followed, or a pushed limit was rejected (including a + // limit-only scan with no filters). Plan now, with the filters only and no limit, so that + // partitions()/read() run exactly once and only after all pushdowns are known -- never in + // the filter pass before a possible pushLimit. + val dataSource = ds.getOrCreateDataSourceInPython(shortName, options, Some(outputSchema)) + val result = ds.source.pushdownLimitInPython(dataSource, outputSchema, allFilters, None) + checkFiltersReplayedDeterministically(result) + readInfo = result.readInfo + } + new PythonScan(ds, shortName, outputSchema, options, supportedFilters, pushedLimit, readInfo) + } // Optionally called by DSv2 once to push down filters before the scan is built. override def pushFilters(filters: Array[Filter]): Array[Filter] = { if (!SQLConf.get.pythonFilterPushDown) { return filters } + allFilters = filters val dataSource = ds.getOrCreateDataSourceInPython(shortName, options, Some(outputSchema)) ds.source.pushdownFiltersInPython(dataSource, outputSchema, filters) match { case None => filters // No filters are supported. case Some(result) => - // Filter pushdown also returns partitions and the read function. - // This helps reduce the number of Python worker calls. - ds.setReadInfo(result.readInfo) + // Filter pushdown may also return the read function and partitions. It defers that + // planning when limit pushdown is enabled -- so partitions()/read() are not planned + // before a possible pushLimit -- in which case the limit pass or build() plans instead. + readInfo = result.readInfo + deferredPlanning = result.readInfo.isEmpty // Partition the filters into supported and unsupported ones. val isPushed = result.isFilterPushed.zip(filters) @@ -58,4 +93,53 @@ class PythonScanBuilder( } override def pushedFilters(): Array[Filter] = supportedFilters + + // Optionally called by DSv2 once to push down a LIMIT before the scan is built. DSv2 calls this + // after `pushFilters`, so the filters that were pushed there (none, for a query without + // pushable filters) are replayed here to rebuild the same reader state before `pushLimit` is + // invoked on it in Python. + override def pushLimit(limit: Int): Boolean = { + if (!SQLConf.get.pythonLimitPushDown) { + return false + } + + val dataSource = ds.getOrCreateDataSourceInPython(shortName, options, Some(outputSchema)) + val result = ds.source.pushdownLimitInPython(dataSource, outputSchema, allFilters, Some(limit)) + checkFiltersReplayedDeterministically(result) + + // Adopt the read info the worker planned. It is empty when the reader rejected the limit -- + // the worker plans nothing then, to keep the rejected (possibly mutated) reader state out of + // the scan -- in which case build() plans the filters-only read from a fresh reader and + // re-validates the replayed filters. Record the pushed limit only when the reader accepted it. + readInfo = result.readInfo + deferredPlanning = result.readInfo.isEmpty + if (result.isLimitPushed) { + pushedLimit = Some(limit) + } + result.isLimitPushed + } + + // The read is replayed on a fresh reader (to push a limit, or to plan after filter pushdown + // deferred). That replay runs `pushFilters` again, which must reach the same decision as the + // first pass. Spark has already committed to that first decision -- `pushedFilters()` was read + // by the optimizer and the filters it reported were dropped from the plan -- so a reader whose + // `pushFilters` is not deterministic would leave Spark applying the first pass's filters while + // reading with the replayed reader, silently returning wrong rows. Fail fast instead. + private def checkFiltersReplayedDeterministically(result: PythonFilterPushdownResult): Unit = { + val replayedFilters = result.isFilterPushed.zip(allFilters).collect { + case (true, filter) => filter + }.toArray + if (!replayedFilters.sameElements(supportedFilters)) { + throw QueryCompilationErrors.pythonDataSourceError( + action = "plan", + tpe = "read", + msg = "pushFilters() returned a different set of supported filters when it was replayed " + + s"during planning: [${supportedFilters.mkString(", ")}] then " + + s"[${replayedFilters.mkString(", ")}]. pushFilters() must be deterministic.") + } + } + + // Spark always applies the LIMIT again after the scan: a Python data source is not trusted to + // return at most `limit` rows, and it is free to over-deliver. + override def isPartiallyPushed(): Boolean = true } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala index 7611cf6764995..7a36118a1879f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala @@ -81,7 +81,11 @@ case class UserDefinedPythonDataSource(dataSourceCls: PythonFunction) { val runner = new UserDefinedPythonDataSourceFilterPushdownRunner( createPythonFunction(pythonResult.dataSource), outputSchema, - filters + filters, + limit = None, + // Defer planning while limit pushdown is enabled: a later limit pass, or build(), plans + // once it is known whether a limit is pushed. Plan here only when limit pushdown is off. + planReadInfo = !SQLConf.get.pythonLimitPushDown ) if (runner.isAnyFilterSupported) { Some(runner.runInPython()) @@ -90,6 +94,36 @@ case class UserDefinedPythonDataSource(dataSourceCls: PythonFunction) { } } + /** + * (Driver-side) Run Python process to plan the read while optionally pushing down a limit, and + * return a [[PythonFilterPushdownResult]] carrying which filters were pushed, whether the limit + * was pushed, and the planned read info. The worker plans the read function and partitions in + * this pass, except when a pushed `limit` is rejected: the reader may have mutated itself while + * rejecting it, so the result's `readInfo` is then `None` and the caller plans the filters-only + * read from a fresh reader instead. + * + * `limit` is the limit to push, or `None` to plan the read from the filters only -- the + * build-time path uses `None` when the filter-pushdown pass deferred planning and no limit was + * ultimately pushed. + * + * `filters` must be the filters that were previously pushed via [[pushdownFiltersInPython]] -- + * empty for a limit-only scan, where no filter pushdown ran -- so that replaying them brings + * the reader to the same state it had during filter pushdown before `pushLimit` is called. + */ + def pushdownLimitInPython( + pythonResult: PythonDataSourceCreationResult, + outputSchema: StructType, + filters: Array[Filter], + limit: Option[Int]): PythonFilterPushdownResult = { + new UserDefinedPythonDataSourceFilterPushdownRunner( + createPythonFunction(pythonResult.dataSource), + outputSchema, + filters, + limit = limit, + planReadInfo = true + ).runInPython() + } + /** * (Driver-side) Run Python process, and get the partition read functions, and * partition information. @@ -345,15 +379,22 @@ private class UserDefinedPythonDataSourceRunner( } /** + * @param readInfo The read function and partitions, when the worker planned them. Empty in two + * cases: the filter-pushdown pass deferred planning (limit pushdown enabled), or a + * pushed limit was rejected (the rejecting reader may be mutated). Planning then + * happens on a later limit pass or at build time, from a fresh reader. * @param isFilterPushed A sequence of bools indicating whether each filter is pushed down. + * @param isLimitPushed Whether the limit is pushed down. False when no limit was sent, and also + * when a limit was sent but the reader rejected it. */ case class PythonFilterPushdownResult( - readInfo: PythonDataSourceReadInfo, - isFilterPushed: collection.Seq[Boolean] + readInfo: Option[PythonDataSourceReadInfo], + isFilterPushed: collection.Seq[Boolean], + isLimitPushed: Boolean = false ) /** - * Push down filters to a Python data source. + * Push down filters and optionally a limit to a Python data source. * * @param dataSource * a Python data source instance @@ -361,11 +402,19 @@ case class PythonFilterPushdownResult( * output schema of the Python data source * @param filters * all filters to be pushed down + * @param limit + * the limit to be pushed down, if any */ private class UserDefinedPythonDataSourceFilterPushdownRunner( dataSource: PythonFunction, schema: StructType, - filters: collection.Seq[Filter]) + filters: collection.Seq[Filter], + limit: Option[Int], + // Whether the worker should plan (and send back) the read function and partitions in this + // pass. Kept separate from `limit` so a build-time, filter-only planning pass can request + // planning without a limit -- deriving it from `limit.isDefined` would force such a pass to + // smuggle a dummy limit value in just to trigger planning. + planReadInfo: Boolean) extends PythonPlannerRunner[PythonFilterPushdownResult](dataSource) { private case class SerializedFilter( @@ -476,14 +525,41 @@ private class UserDefinedPythonDataSourceFilterPushdownRunner( // Send the filters PythonWorkerUtils.writeUTF(mapper.writeValueAsString(serializedFilters), dataOut) + // Send the limit, if any. -1 means no limit is pushed down. + dataOut.writeInt(limit.getOrElse(-1)) + // Send configurations dataOut.writeInt(SQLConf.get.arrowMaxRecordsPerBatch) + dataOut.writeBoolean(SQLConf.get.pythonFilterPushDown) + dataOut.writeBoolean(SQLConf.get.pythonLimitPushDown) dataOut.writeBoolean(SQLConf.get.pysparkBinaryAsBytes) + // Whether the worker should plan the read function and partitions. The filter-pushdown pass + // sets this false while limit pushdown is enabled, because a later limit pass -- or + // build-time planning -- will plan once it is known whether a limit is pushed, so + // `partitions()`/`read()` never run on a reader whose plan would be discarded. + dataOut.writeBoolean(planReadInfo) } override protected def receiveFromPython(dataIn: DataInputStream): PythonFilterPushdownResult = { - // Receive the read function and the partitions. Also check for exceptions. - val readInfo = PythonDataSourceReadInfo.receive(dataIn) + // Whether a read function + partitions follow. Read first so that an exception raised in the + // worker (sent as PYTHON_EXCEPTION_THROWN) is surfaced here instead of being misread as data. + val hasReadInfo = dataIn.readInt() + if (hasReadInfo == SpecialLengths.PYTHON_EXCEPTION_THROWN) { + val msg = PythonWorkerUtils.readUTF(dataIn) + throw QueryCompilationErrors.pythonDataSourceError( + action = "plan", + tpe = "read", + msg = msg + ) + } + // Receive the read function and partitions, when the worker planned them. None follow when the + // filter-pushdown pass defers planning (limit pushdown enabled) or when a pushed limit was + // rejected; build() (or a later limit pass) then plans them from a fresh reader instead. + val readInfo = if (hasReadInfo != 0) { + Some(PythonDataSourceReadInfo.receive(dataIn)) + } else { + None + } // Receive the pushed filters as a list of indices. val numFiltersPushed = dataIn.readInt() @@ -493,9 +569,13 @@ private class UserDefinedPythonDataSourceFilterPushdownRunner( isFilterPushed(serializedFilters(i).index) = true } + // Receive whether the limit was pushed down, sent as 1 or 0. + val isLimitPushed = dataIn.readInt() != 0 + PythonFilterPushdownResult( readInfo = readInfo, - isFilterPushed = isFilterPushed + isFilterPushed = isFilterPushed, + isLimitPushed = isLimitPushed ) } } @@ -575,6 +655,7 @@ private class UserDefinedPythonDataSourceReadRunner( // Send configurations dataOut.writeInt(SQLConf.get.arrowMaxRecordsPerBatch) dataOut.writeBoolean(SQLConf.get.pythonFilterPushDown) + dataOut.writeBoolean(SQLConf.get.pythonLimitPushDown) dataOut.writeBoolean(isStreaming) dataOut.writeBoolean(SQLConf.get.pysparkBinaryAsBytes) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala index bcf9c5db367a8..5dfe2c717b8dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSource.scala @@ -320,7 +320,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging // Read the schema file path from operator metadata version v2 onwards // for the transformWithState operator - val oldSchemaFilePaths = if (storeMetadata.length > 0 && storeMetadata.head.version == 2) { + val oldSchemaFilePaths = if (storeMetadata.nonEmpty && storeMetadata.head.version == 2) { val opName = storeMetadata.head.operatorName if (StatefulOperatorsUtils.TRANSFORM_WITH_STATE_OP_NAMES.exists(opName.contains)) { val storeMetadataEntry = storeMetadata.head @@ -442,7 +442,7 @@ class StateDataSource extends TableProvider with DataSourceRegister with Logging storeMetadata: Array[StateMetadataTableEntry]): KeyStateEncoderSpec = { // If operator metadata is not found, then log a warning and continue with using the no-prefix // key state encoder - val keyStateEncoderSpec = if (storeMetadata.length == 0) { + val keyStateEncoderSpec = if (storeMetadata.isEmpty) { logWarning("Metadata for state store not found, possible cause is this checkpoint " + "is created by older version of spark. If the query has session window aggregation, " + "the state can't be read correctly and runtime exception will be thrown. " + diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/text/TextTable.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/text/TextTable.scala index d8880b84c6211..289c41ca1d0b2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/text/TextTable.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/text/TextTable.scala @@ -46,7 +46,8 @@ case class TextTable( } } - override def supportsDataType(dataType: DataType): Boolean = dataType == StringType + override def supportsDataType(dataType: DataType): Boolean = + dataType.isInstanceOf[StringType] override def formatName: String = "Text" } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlDataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlDataSource.scala index ef34711e041eb..4ec07656f3f4b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlDataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlDataSource.scala @@ -17,13 +17,13 @@ package org.apache.spark.sql.execution.datasources.xml -import java.io.{ByteArrayInputStream, FileNotFoundException, InputStream, IOException} +import java.io.{ByteArrayInputStream, Closeable, FileNotFoundException, InputStream, IOException} import java.nio.charset.{Charset, StandardCharsets} import scala.util.control.NonFatal import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.{FileStatus, Path} +import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path} import org.apache.hadoop.hdfs.BlockMissingException import org.apache.hadoop.mapreduce.Job import org.apache.hadoop.mapreduce.lib.input.FileInputFormat @@ -34,7 +34,7 @@ import org.apache.spark.input.{PortableDataStream, StreamInputFormat} import org.apache.spark.internal.Logging import org.apache.spark.rdd.{BinaryFileRDD, RDD} import org.apache.spark.sql.{Dataset, Encoders, SparkSession} -import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.{FileSourceOptions, InternalRow} import org.apache.spark.sql.catalyst.util.FailureSafeParser import org.apache.spark.sql.catalyst.xml.{StaxXmlParser, StaxXMLRecordReader, XmlInferSchema, XmlOptions} import org.apache.spark.sql.classic.ClassicConversions.castToImpl @@ -81,13 +81,16 @@ abstract class XmlDataSource extends Serializable with Logging with SupportsArch * `XmlFileFormat` read path supports archives; XML has no DSv2 reader. * * @param parser builds a fresh XML parser for each entry. + * @param archivePathFilter optional glob matched against the entry's full path */ def readArchive( conf: Configuration, file: PartitionedFile, parser: () => StaxXmlParser, - schema: StructType): Iterator[InternalRow] = - SupportsArchiveFormat.readArchiveEntries(file.toPath, conf) { (_, in) => + schema: StructType, + archivePathFilter: Option[GlobPattern]): Iterator[InternalRow] = + SupportsArchiveFormat.readArchiveEntries( + file.toPath, conf, archivePathFilter = archivePathFilter) { (_, in) => readStream(in, parser(), schema) } @@ -101,15 +104,7 @@ abstract class XmlDataSource extends Serializable with Logging with SupportsArch parsedOptions.singleVariantColumn match { case Some(columnName) => Some(StructType(Array(StructField(columnName, VariantType)))) case None => - // When any input is a tar archive, infer over all inputs in a single pass -- archive - // entries are streamed (never unpacked to disk) and tokenized as XML records alongside any - // loose files -- so the result matches a directory read of the same files. XML has no DSv2 - // reader, so this archive scan is always V1. - val hasArchive = parsedOptions.archiveFormatEnabled && - inputPaths.exists(f => SupportsArchiveFormat.isArchivePath(f.getPath)) - if (hasArchive) { - Some(inferWithArchives(sparkSession, inputPaths, parsedOptions)) - } else if (inputPaths.nonEmpty) { + if (inputPaths.nonEmpty) { Some(infer(sparkSession, inputPaths, parsedOptions)) } else { None @@ -122,69 +117,6 @@ abstract class XmlDataSource extends Serializable with Logging with SupportsArch inputPaths: Seq[FileStatus], parsedOptions: XmlOptions): StructType - /** - * Infers an XML schema when at least one input is a tar archive (`.tar`/`.tar.gz`/`.tgz`). Every - * archive entry (streamed through `SupportsArchiveFormat`, never unpacked to disk) and every - * loose file is tokenized into records and fed to a single [[XmlInferSchema]] pass, exactly as a - * directory of the same files would infer. Tokenization is per-mode so it matches the scan: - * multi-line splits the whole stream into `rowTag`-delimited records, single-line treats each - * line as a record (mirroring [[readFile]] and JSON's `inferWithArchives`). - */ - private def inferWithArchives( - sparkSession: SparkSession, - inputPaths: Seq[FileStatus], - parsedOptions: XmlOptions): StructType = { - val baseRdd = createBaseRdd(sparkSession, inputPaths, parsedOptions) - val ignoreCorruptFiles = parsedOptions.ignoreCorruptFiles - val ignoreMissingFiles = parsedOptions.ignoreMissingFiles - - // Applies `perEntry` to each input -- an archive entry by entry (streamed, so only one entry's - // bytes are in flight at a time), a loose file directly -- skipping a whole input when it is - // corrupt/missing and the ignore flags are set. - def perInput(perEntry: InputStream => Iterator[String]): RDD[String] = baseRdd.flatMap { - stream => - val path = new Path(stream.getPath()) - try { - if (SupportsArchiveFormat.isArchivePath(path)) { - SupportsArchiveFormat.readArchiveEntries(path, stream.getConfiguration) { (_, in) => - perEntry(in) - } - } else { - perEntry( - CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path)) - } - } catch { - case e: FileNotFoundException if ignoreMissingFiles => - logWarning("Skipped missing file", e) - Iterator.empty[String] - case NonFatal(e) => - Utils.getRootCause(e) match { - case root @ (_: AccessControlException | _: BlockMissingException) => throw root - case _: RuntimeException | _: IOException if ignoreCorruptFiles => - logWarning("Skipped the rest of the content in the corrupted file", e) - Iterator.empty[String] - case other => throw other - } - } - } - - // Tokenize each input the way this mode's scan reads records, so the inferred schema matches a - // directory read: multi-line splits the whole stream into rowTag-delimited records, single-line - // treats each line as a record (mirroring TextInputXmlDataSource.readFile). - val tokenRDD: RDD[String] = if (parsedOptions.multiLine) { - perInput(in => StaxXmlParser.tokenizeStream(in, parsedOptions)) - } else { - val charset = parsedOptions.charset - perInput(in => lineIterator(in, None).map { line => - new String(line.getBytes, 0, line.getLength, charset) - }) - } - SQLExecution.withSQLConfPropagated(sparkSession) { - new XmlInferSchema(parsedOptions, sparkSession.sessionState.conf.caseSensitiveAnalysis) - .infer(tokenRDD) - } - } - protected def createBaseRdd( sparkSession: SparkSession, inputPaths: Seq[FileStatus], @@ -358,6 +290,12 @@ object MultiLineXmlDataSource extends XmlDataSource { inputPaths: Seq[FileStatus], parsedOptions: XmlOptions): StructType = { + val hasArchive = parsedOptions.archiveFormatEnabled && + inputPaths.exists(f => SupportsArchiveFormat.isArchivePath(f.getPath)) + if (hasArchive) { + return inferWithArchives(sparkSession, inputPaths, parsedOptions) + } + if (!parsedOptions.useLegacyXMLParser) { return inferOptimized(sparkSession, inputPaths, parsedOptions) } @@ -418,4 +356,103 @@ object MultiLineXmlDataSource extends XmlDataSource { schema } } + + /** + * Infers a multi-line XML schema when at least one input is an archive. Every archive entry + * (streamed through `SupportsArchiveFormat`, never unpacked to disk) and every loose file is + * tokenized into `rowTag`-delimited records and fed to a single [[XmlInferSchema]] pass, exactly + * as a directory of the same files would infer. Single-line archive inference does not come here: + * the Text data source reads archives directly, so it flows through + * [[TextInputXmlDataSource.infer]] like any directory read. Corrupt/missing inputs are skipped + * when the ignore flags are set (see [[skipInputOnError]]). Uses the legacy tokenizer because the + * optimized parser re-opens its input, which a single-use archive entry stream does not support. + */ + private def inferWithArchives( + sparkSession: SparkSession, + inputPaths: Seq[FileStatus], + parsedOptions: XmlOptions): StructType = { + val baseRdd = createBaseRdd(sparkSession, inputPaths, parsedOptions) + // Inference must see the same entries the scan reads, so it honors archivePathFilter too. + // Capture the glob string: the compiled GlobPattern is not serializable, so each task + // compiles it once when the archive branch is taken. + val archivePathFilterGlob = parsedOptions.archivePathFilter + val ignoreCorruptFiles = parsedOptions.ignoreCorruptFiles + val ignoreMissingFiles = parsedOptions.ignoreMissingFiles + + val tokenRDD: RDD[String] = baseRdd.mapPartitions { streams => + // Compile at most once per partition: lazy so a partition of only loose files never + // compiles, while a partition with archives reuses one matcher across all of them. + lazy val archivePathFilter = + archivePathFilterGlob.map(FileSourceOptions.compileArchivePathFilter) + streams.flatMap { stream => + val path = new Path(stream.getPath()) + skipInputOnError(ignoreMissingFiles, ignoreCorruptFiles) { + if (SupportsArchiveFormat.isArchivePath(path)) { + SupportsArchiveFormat.readArchiveEntries( + path, stream.getConfiguration, archivePathFilter = archivePathFilter) { + (_, in) => + StaxXmlParser.tokenizeStream(in, parsedOptions) + } + } else { + StaxXmlParser.tokenizeStream( + CodecStreams.createInputStreamWithCloseResource(stream.getConfiguration, path), + parsedOptions) + } + } + } + } + SQLExecution.withSQLConfPropagated(sparkSession) { + new XmlInferSchema(parsedOptions, sparkSession.sessionState.conf.caseSensitiveAnalysis) + .infer(tokenRDD) + } + } + + /** + * Builds one input's token iterator, catching a missing/corrupt error when the ignore flags are + * set. `readArchiveEntries` advances to later entries lazily, so a corrupt later entry throws on + * `hasNext`, not at construction; the returned iterator catches both. A construction failure + * skips the whole input; a mid-advance failure keeps the records already yielded and skips only + * the remainder of the archive. Access/block errors are always rethrown (unwrapped), matching the + * non-archive read path. + */ + private def skipInputOnError( + ignoreMissingFiles: Boolean, + ignoreCorruptFiles: Boolean)( + build: => Iterator[String]): Iterator[String] = { + def handle(e: Throwable): Iterator[String] = e match { + case e: FileNotFoundException if ignoreMissingFiles => + logWarning("Skipped missing file", e) + Iterator.empty[String] + case NonFatal(e) => + Utils.getRootCause(e) match { + case root @ (_: AccessControlException | _: BlockMissingException) => throw root + case _: RuntimeException | _: IOException if ignoreCorruptFiles => + logWarning("Skipped the rest of the content in the corrupted file", e) + Iterator.empty[String] + case other => throw other + } + } + + val underlying = + try build + catch { case NonFatal(e) => return handle(e) } + + new Iterator[String] { + private var delegate = underlying + override def hasNext: Boolean = + try delegate.hasNext + catch { + case NonFatal(e) => + // A mid-advance throw reaches neither exhaustion nor close, so close the failed + // iterator here to release its archive stream promptly rather than at task completion. + delegate match { + case c: Closeable => try c.close() catch { case NonFatal(_) => } + case _ => + } + delegate = handle(e) + delegate.hasNext + } + override def next(): String = delegate.next() + } + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlFileFormat.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlFileFormat.scala index 26bf2cd3ccab6..7823efd6ae811 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlFileFormat.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/xml/XmlFileFormat.scala @@ -131,7 +131,8 @@ case class XmlFileFormat() extends TextBasedFileFormat with DataSourceRegister { broadcastedHadoopConf.value.value, file, () => parser(), - requiredSchema) + requiredSchema, + xmlOptions.archivePathFilterPattern) } else { XmlDataSource(xmlOptions).readFile( broadcastedHadoopConf.value.value, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala index 3e1b39b564c79..1bf12a695bc4c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.dynamicpruning import org.apache.spark.sql.catalyst.catalog.HiveTableRelation import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.optimizer.JoinSelectionHelper +import org.apache.spark.sql.catalyst.optimizer.{JoinSelectionHelper, ReusableBroadcastValueProjection} import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule @@ -28,6 +28,7 @@ import org.apache.spark.sql.execution.LogicalRDD import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} import org.apache.spark.sql.execution.datasources.v2.ExtractV2Scan +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering /** * Dynamic partition pruning optimization is performed based on the type and * selectivity of the join operation. During query optimization, we insert a @@ -86,6 +87,14 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } else { None } + case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) => + val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( + scan.filterAttributes(), r.output) + if (resExp.references.subsetOf(filterAttrs)) { + Some(r) + } else { + None + } case _ => None } } @@ -110,6 +119,11 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join require(filteringKeys.size == 1, "DPP Filters should only have a single broadcasting key " + "since there are no usage for multiple broadcasting keys at the moment.") val indices = Seq(joinKeys.indexOf(filteringKeys.head)) + val broadcastValueProjection = if (conf.dynamicPartitionPruningBroadcastProjectionEnabled) { + ReusableBroadcastValueProjection.find(filteringKeys.head, filteringPlan, partScan) + } else { + None + } lazy val hasBenefit = pruningHasBenefit( pruningKey, partScan, filteringKeys.head, filteringPlan, hasSelectivePredicate(filteringPlan)) if (reuseEnabled || hasBenefit) { @@ -120,7 +134,7 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join filteringPlan, joinKeys, indices, - conf.dynamicPartitionPruningReuseBroadcastOnly || !hasBenefit), + conf.dynamicPartitionPruningReuseBroadcastOnly || !hasBenefit)(broadcastValueProjection), pruningPlan) } else { // abort dynamic partition pruning diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala index fdcb78fdb55a1..ca7407eb2a88a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala @@ -17,14 +17,14 @@ package org.apache.spark.sql.execution.dynamicpruning -import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeSeq, BindReferences, DynamicPruningExpression, DynamicPruningSubquery, Expression, Literal} +import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeSeq, BindReferences, BroadcastValueProjection, DynamicPruningExpression, DynamicPruningSubquery, Expression, Literal} import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.catalyst.plans.logical.Aggregate +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.DYNAMIC_PRUNING_SUBQUERY import org.apache.spark.sql.classic.SparkSession -import org.apache.spark.sql.execution.{InSubqueryExec, QueryExecution, SparkPlan, SubqueryBroadcastExec, SubqueryExec} +import org.apache.spark.sql.execution.{BaseSubqueryExec, InSubqueryExec, ProjectedBroadcastValueSubqueryExec, QueryExecution, SparkPlan, SubqueryBroadcastExec, SubqueryExec} import org.apache.spark.sql.execution.exchange.BroadcastExchangeExec import org.apache.spark.sql.execution.joins._ import org.apache.spark.sql.internal.SQLConf @@ -46,47 +46,89 @@ case class PlanDynamicPruningFilters(sparkSession: SparkSession) extends Rule[Sp HashedRelationBroadcastMode(packedKeys) } + private def reusableBroadcast( + plan: SparkPlan, + name: String, + indices: Seq[Int], + buildPlan: LogicalPlan, + buildKeys: Seq[Expression], + projection: Option[BroadcastValueProjection]): Option[BaseSubqueryExec] = { + if (!conf.exchangeReuseEnabled || buildKeys.isEmpty) { + return None + } + + val sparkPlan = QueryExecution.createSparkPlan( + sparkSession.sessionState.planner, buildPlan) + val requiredMode = broadcastMode(buildKeys, sparkPlan.output) + val canReuseExchange = plan.exists { + case join: BroadcastHashJoinExec => + val (candidateKeys, candidatePlan) = join.buildSide match { + case BuildLeft => (join.leftKeys, join.left) + case BuildRight => (join.rightKeys, join.right) + } + // Preserve existing direct reuse; only a projected domain requires the exact hash mode. + candidatePlan.sameResult(sparkPlan) && + (projection.isEmpty || + (!join.isNullAwareAntiJoin && + broadcastMode(candidateKeys, candidatePlan.output) == requiredMode)) + case _ => false + } + + if (canReuseExchange) { + val executedPlan = QueryExecution.prepareExecutedPlan(sparkSession, sparkPlan) + val exchange = BroadcastExchangeExec( + broadcastMode(buildKeys, executedPlan.output), executedPlan) + Some(projection match { + case Some(valueProjection) => + ProjectedBroadcastValueSubqueryExec( + name, valueProjection.valueExpression, exchange) + case None => + SubqueryBroadcastExec(name, indices, buildKeys, exchange) + }) + } else { + None + } + } + override def apply(plan: SparkPlan): SparkPlan = { if (!conf.dynamicPartitionPruningEnabled) { return plan } plan.transformAllExpressionsWithPruning(_.containsPattern(DYNAMIC_PRUNING_SUBQUERY)) { - case DynamicPruningSubquery( + case pruning @ DynamicPruningSubquery( value, buildPlan, buildKeys, broadcastKeyIndices, onlyInBroadcast, exprId, _) => - val sparkPlan = QueryExecution.createSparkPlan(sparkSession.sessionState.planner, buildPlan) val name = s"dynamicpruning#${exprId.id}" - // Using `sparkPlan` is a little hacky as it is based on the assumption that this rule is - // the first to be applied (apart from `InsertAdaptiveSparkPlan`). - val canReuseExchange = conf.exchangeReuseEnabled && buildKeys.nonEmpty && - plan.exists { - case BroadcastHashJoinExec(_, _, _, BuildLeft, _, left, _, _, _) => - left.sameResult(sparkPlan) - case BroadcastHashJoinExec(_, _, _, BuildRight, _, _, right, _, _) => - right.sameResult(sparkPlan) - case _ => false + val directBroadcast = reusableBroadcast( + plan, name, broadcastKeyIndices, buildPlan, buildKeys, None) + val reusedBroadcast = directBroadcast.orElse { + if (onlyInBroadcast) { + pruning.usableBroadcastValueProjection.flatMap { projection => + reusableBroadcast( + plan, + name, + Seq(0), + projection.sourcePlan, + projection.sourceHashKeys, + Some(projection)) + } + } else { + None } + } - if (canReuseExchange) { - val executedPlan = QueryExecution.prepareExecutedPlan(sparkSession, sparkPlan) - val mode = broadcastMode(buildKeys, executedPlan.output) - // plan a broadcast exchange of the build side of the join - val exchange = BroadcastExchangeExec(mode, executedPlan) - // place the broadcast adaptor for reusing the broadcast results on the probe side - val broadcastValues = - SubqueryBroadcastExec(name, broadcastKeyIndices, buildKeys, exchange) - DynamicPruningExpression(InSubqueryExec(value, broadcastValues, exprId)) - } else if (onlyInBroadcast) { - // it is not worthwhile to execute the query, so we fall-back to a true literal - DynamicPruningExpression(Literal.TrueLiteral) - } else { - // we need to apply an aggregate on the buildPlan in order to be column pruned - val aliases = broadcastKeyIndices.map(idx => - Alias(buildKeys(idx), buildKeys(idx).toString)()) - val aggregate = Aggregate(aliases, aliases, buildPlan) - val sparkPlan = QueryExecution.prepareExecutedPlan(sparkSession, aggregate) - val values = SubqueryExec(name, sparkPlan) - DynamicPruningExpression(InSubqueryExec(value, values, exprId)) + reusedBroadcast match { + case Some(broadcastValues) => + DynamicPruningExpression(InSubqueryExec(value, broadcastValues, exprId)) + case None if onlyInBroadcast => + DynamicPruningExpression(Literal.TrueLiteral) + case None => + val aliases = broadcastKeyIndices.map(idx => + Alias(buildKeys(idx), buildKeys(idx).toString)()) + val aggregate = Aggregate(aliases, aliases, buildPlan) + val sparkPlan = QueryExecution.prepareExecutedPlan(sparkSession, aggregate) + val values = SubqueryExec(name, sparkPlan) + DynamicPruningExpression(InSubqueryExec(value, values, exprId)) } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala index 9f8409efa360e..ae7b432d72f5b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala @@ -24,9 +24,12 @@ import org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation, GroupBasedRowLevelOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, RowLevelWrite} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering +import org.apache.spark.sql.catalyst.trees.TreePattern.{REPLACE_DATA, WRITE_DELTA} +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan} +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.util.ArrayImplicits._ /** @@ -49,28 +52,42 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla import DataSourceV2Implicits._ - override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { + override def apply(plan: LogicalPlan): LogicalPlan = plan.transformDownWithPruning( + _.containsAnyPattern(REPLACE_DATA, WRITE_DELTA)) { case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(replaceData, cond, scan) + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) + if canInjectGroupFilters(cond, scan.filterAttributes) => + injectGroupFilters(replaceData, cond, scan, scan.filterAttributes) + + case GroupBasedRowLevelOperation(replaceData, _, Some(cond), + ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) + if canInjectGroupFilters(cond, scan.filterAttributes()) => + injectGroupFilters(replaceData, cond, scan, scan.filterAttributes()) + + case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) + if canInjectGroupFilters(cond, scan.filterAttributes) => + injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes) case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(writeDelta, cond, scan) + ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) + if canInjectGroupFilters(cond, scan.filterAttributes()) => + injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes()) } private def canInjectGroupFilters( cond: Expression, - scan: SupportsRuntimeV2Filtering): Boolean = { + filterAttrs: Array[NamedReference]): Boolean = { conf.runtimeRowLevelOperationGroupFilterEnabled && cond != TrueLiteral && - scan.filterAttributes.nonEmpty + filterAttrs.nonEmpty } private def injectGroupFilters( write: RowLevelWrite, cond: Expression, - scan: SupportsRuntimeV2Filtering): LogicalPlan = { + scan: Scan, + filterAttrs: Array[NamedReference]): LogicalPlan = { // use reference equality on scan to find required scan relations val newQuery = write.query transformUp { case r: DataSourceV2ScanRelation if r.scan eq scan => @@ -79,9 +96,9 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla val originalTable = r.relation.table.asRowLevelOperationTable.table val relation = r.relation.copy(table = originalTable) val matchingRowsPlan = buildMatchingRowsPlan(write, relation, cond) - val filterAttrs = scan.filterAttributes.toImmutableArraySeq - val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, matchingRowsPlan) - val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r) + val filterAttrsSeq = filterAttrs.toImmutableArraySeq + val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, matchingRowsPlan) + val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, r) Filter(buildDynamicPruningCond(matchingRowsPlan, buildKeys, pruningKeys), r) } // optimize subqueries to rewrite them as joins and trigger job planning diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index c632b3d841e61..c2051ac91b062 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.execution.exchange +import scala.annotation.tailrec import scala.collection.mutable import scala.collection.mutable.ArrayBuffer @@ -248,13 +249,16 @@ case class EnsureRequirements( child case ((child, dist), idx) => if (bestSpecOpt.isDefined && bestSpecOpt.get.isCompatibleWith(specs(idx))) { - bestSpecOpt match { + // If the child's partitioning is a `PartitioningCollection`, its spec is a + // `ShuffleSpecCollection` whose `createPartitioning` delegates to the head spec, + // so unwrap to the head spec to stay aligned with the re-shuffled side below. + unwrapSpecCollection(bestSpecOpt.get) match { // If `areChildrenCompatible` is false, we can still perform SPJ // by shuffling the other side based on join keys (see the else case below). // Hence we need to ensure that after this call, the outputPartitioning of the // partitioned side's BatchScanExec is grouped by join keys to match, // and we do that by pushing down the join keys - case Some(KeyedShuffleSpec(_, _, Some(joinKeyPositions))) => + case KeyedShuffleSpec(_, _, Some(joinKeyPositions)) => withJoinKeyPositions(child, joinKeyPositions) case _ => child } @@ -272,8 +276,8 @@ case class EnsureRequirements( } child match { - case ShuffleExchangeExec(_, c, so, ps) => - ShuffleExchangeExec(newPartitioning, c, so, ps) + case s: ShuffleExchangeExec => + s.copy(outputPartitioning = newPartitioning) case gpe: GroupPartitionsExec => ShuffleExchangeExec(newPartitioning, gpe.child) case _ => ShuffleExchangeExec(newPartitioning, child) } @@ -760,6 +764,14 @@ case class EnsureRequirements( } } + // Unwraps a `ShuffleSpecCollection` (possibly nested) to the spec that its + // `createPartitioning` delegates to, i.e. the head spec. + @tailrec + private def unwrapSpecCollection(spec: ShuffleSpec): ShuffleSpec = spec match { + case ShuffleSpecCollection(specs) => unwrapSpecCollection(specs.head) + case other => other + } + /** * Applies join key positions to a plan by wrapping or updating GroupPartitionsExec. */ @@ -782,20 +794,19 @@ case class EnsureRequirements( partitioning: Partitioning, distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = { def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = { - // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one - // attribute per partition expression. - val attributes = partitioning.expressions.flatMap(_.references) - val clustering = distribution.clustering - - val satisfies = if (SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION)) { - attributes.length == clustering.length && attributes.zip(clustering).forall { - case (l, r) => l.semanticEquals(r) - } - } else { - partitioning.satisfies(distribution) + // The config requires all the cluster keys to be covered by the partition keys, to avoid + // the skew of joining on keys that are coarser than the join keys. Key order and duplicated + // cluster keys don't matter. + def allClusterKeysCovered: Boolean = { + // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one + // attribute per partition expression. + val attributes = partitioning.expressions.flatMap(_.references) + distribution.clustering.forall(c => attributes.exists(_.semanticEquals(c))) } - if (satisfies) { + if (partitioning.satisfies(distribution) && + (!SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION) || + allClusterKeysCovered)) { Some(partitioning.createShuffleSpec(distribution).asInstanceOf[KeyedShuffleSpec]) } else { None @@ -896,7 +907,7 @@ case class EnsureRequirements( def apply(plan: SparkPlan): SparkPlan = { val newPlan = plan.transformUp { - case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _) + case operator @ ShuffleExchangeExec(upper: HashPartitioning, child, shuffleOrigin, _, _) if optimizeOutRepartition && (shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin == REPARTITION_BY_NUM) => def hasSemanticEqualPartitioning(partitioning: Partitioning): Boolean = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala index dd829e697df61..0f3bb973000c1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala @@ -191,7 +191,8 @@ case class ShuffleExchangeExec( override val outputPartitioning: Partitioning, child: SparkPlan, shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS, - advisoryPartitionSize: Option[Long] = None) + advisoryPartitionSize: Option[Long] = None, + pipelined: Boolean = false) extends ShuffleExchangeLike { private lazy val writeMetrics = @@ -205,6 +206,17 @@ case class ShuffleExchangeExec( override def nodeName: String = "Exchange" + // `pipelined` is only meaningful for a Real-Time Mode plan, and the default arg string is + // positional, so printing it unconditionally would add a bare `false` to every shuffle in every + // plan. Show it only when set, and name it when shown. + override def stringArgs: Iterator[Any] = { + // `pipelined` is the last field, so drop it positionally rather than by value; argString drops + // the child on its own. Exchange's `[plan_id=...]` suffix is re-appended here. + val argsWithoutPipelined = productIterator.toSeq.dropRight(1).iterator + val pipelinedArg = if (pipelined) Iterator("isPipelined=true") else Iterator.empty + argsWithoutPipelined ++ pipelinedArg ++ Iterator(s"[plan_id=$id]") + } + private lazy val serializer: Serializer = new UnsafeRowSerializer(child.output.size, longMetric("dataSize")) @@ -252,7 +264,8 @@ case class ShuffleExchangeExec( child.output, outputPartitioning, serializer, - writeMetrics) + writeMetrics, + pipelined) metrics("numPartitions").set(dep.partitioner.numPartitions) val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) SQLMetrics.postDriverMetricUpdates( @@ -346,7 +359,8 @@ object ShuffleExchangeExec { outputAttributes: Seq[Attribute], newPartitioning: Partitioning, serializer: Serializer, - writeMetrics: Map[String, SQLMetric]) + writeMetrics: Map[String, SQLMetric], + pipelined: Boolean = false) : ShuffleDependency[Int, InternalRow, InternalRow] = { val part: Partitioner = newPartitioning match { case RoundRobinPartitioning(numPartitions) => new HashPartitioner(numPartitions) @@ -385,11 +399,18 @@ object ShuffleExchangeExec { samplePointsPerPartitionHint = SQLConf.get.rangeExchangeSampleSizePerPartition) case SinglePartition => new ConstantPartitioner case k: KeyedPartitioning => - val keyGroupedPartitioning = k.toGrouped - val valueMap = keyGroupedPartitioning.partitionKeys.zipWithIndex.map { - case (key, index) => (key.row.toSeq(keyGroupedPartitioning.expressionDataTypes), index) + // `partitionKeys` is the physical layout its producer declared: partition `i` holds key + // `partitionKeys(i)`. Keep that order, whatever it is -- it need not be sorted, and + // re-deriving one here would disagree with the side this shuffle co-partitions with. + // Unique keys are a precondition, since each key gets exactly one partition; + // `KeyedShuffleSpec.canCreatePartitioning` is what refuses an ungrouped partitioning. + assert(k.isGrouped, + s"Expected a grouped KeyedPartitioning on ${k.expressions}, but got ${k.numPartitions} " + + "partition keys with duplicates among them") + val valueMap = k.partitionKeys.zipWithIndex.map { + case (key, index) => (key.row.toSeq(k.expressionDataTypes), index) }.toMap - new KeyGroupedPartitioner(mutable.Map.from(valueMap), keyGroupedPartitioning.numPartitions) + new KeyGroupedPartitioner(mutable.Map.from(valueMap), k.numPartitions) case _ => throw SparkException.internalError(s"Exchange not implemented for $newPartitioning") // TODO: Handle BroadcastPartitioning. } @@ -509,8 +530,16 @@ object ShuffleExchangeExec { // round-robin function is order sensitive if we don't sort the input. // Stateful partition assignment is order-sensitive when it depends on row visitation order. - val isOrderSensitive = - (isRoundRobin || isNullAwareHashPartitioning) && !SQLConf.get.sortBeforeRepartition + // + // A pipelined shuffle is exempt. Marking the map RDD order-sensitive only serves to make it + // INDETERMINATE when its own input is UNORDERED (see + // MapPartitionsRDD.getOutputDeterministicLevel), which tells the scheduler a retry cannot be + // trusted and the stage must be rolled back and recomputed. A pipelined stage is never + // retried (a pipelined task set gets a single attempt) and never recomputed, and the + // DAGScheduler rejects an indeterminate pipelined producer outright -- so keeping the flag + // would reject a chain of round-robin repartitions rather than protect anything. + val isOrderSensitive = (isRoundRobin || isNullAwareHashPartitioning) && + !SQLConf.get.sortBeforeRepartition && !pipelined if (needToCopyObjectsBeforeShuffle(part)) { newRdd.mapPartitionsWithIndexInternal((_, iter) => { val getPartitionKey = getPartitionKeyExtractor() @@ -537,15 +566,30 @@ object ShuffleExchangeExec { } } val dependency = - new ShuffleDependency[Int, InternalRow, InternalRow]( - rddWithPartitionIds, - new PartitionIdPassthrough(part.numPartitions), - serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), - rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), - _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, - checksumMismatchQueryLevelRollbackEnabled = - SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + if (pipelined) { + // A pipelined shuffle is transient and incrementally readable: the DAGScheduler + // co-schedules its producer and consumer stages instead of materializing the shuffle + // first. The PipelinedShuffleDependency type is the entire opt-in -- routing to the + // streaming shuffle manager and pipelined-group co-scheduling both follow from it. The + // checksum-mismatch retry knobs are intentionally not carried over: a transient shuffle is + // never recomputed, so PipelinedShuffleDependency does not expose them (they stay off). + new PipelinedShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize)) + } else { + new ShuffleDependency[Int, InternalRow, InternalRow]( + rddWithPartitionIds, + new PartitionIdPassthrough(part.numPartitions), + serializer, + shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + rowBasedChecksums = UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize), + _checksumMismatchFullRetryEnabled = SQLConf.get.shuffleChecksumMismatchFullRetryEnabled, + checksumMismatchQueryLevelRollbackEnabled = + SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled) + } dependency } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala index a7292ee1f8fa7..cdc27bf316cb9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastNestedLoopJoinExec.scala @@ -177,7 +177,7 @@ case class BroadcastNestedLoopJoinExec( nextIndex += 1 if (boundCondition(resultRow)) { if (foundMatch && singleJoin) { - throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows() } foundMatch = true return true diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala index 9df791aa8de0c..e86db7299ab65 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashJoin.scala @@ -152,6 +152,25 @@ trait HashJoin extends JoinCodegenSupport { (r: InternalRow) => true } + /** + * For join types that preserve all streamed rows, split the condition into + * streamed-only and rest. The streamed-only part can be evaluated once per streamed row + * before probing the hash table. + */ + protected lazy val (streamedOnlyCondition, restCondition): + (Option[Expression], Option[Expression]) = { + StreamedSideJoinCondition.split( + condition, joinType, streamedPlan, conf.splitStreamedSideJoinCondition) + } + + @transient protected[this] lazy val boundStreamedOnlyCondition = streamedOnlyCondition.map { + Predicate.create(_, streamedPlan.output).eval _ + }.getOrElse((_: InternalRow) => true) + + @transient protected[this] lazy val boundRestCondition = restCondition.map { + Predicate.create(_, streamedPlan.output ++ buildPlan.output).eval _ + }.getOrElse((_: InternalRow) => true) + protected def createResultProjection(): (InternalRow) => InternalRow = joinType match { case LeftExistence(_) => UnsafeProjection.create(output, output) @@ -205,8 +224,14 @@ trait HashJoin extends JoinCodegenSupport { streamedIter.map { currentRow => val rowKey = keyGenerator(currentRow) joinedRow.withLeft(currentRow) - val matched = hashedRelation.getValue(rowKey) - if (matched != null && boundCondition(joinedRow.withRight(matched))) { + // If streamed-only condition is false/null, the full condition can never be true, + // so the row is emitted with null build side (no probe needed). + val matched = if (boundStreamedOnlyCondition(currentRow)) { + hashedRelation.getValue(rowKey) + } else { + null + } + if (matched != null && boundRestCondition(joinedRow.withRight(matched))) { joinedRow } else { joinedRow.withRight(nullRow) @@ -216,15 +241,19 @@ trait HashJoin extends JoinCodegenSupport { streamedIter.flatMap { currentRow => val rowKey = keyGenerator(currentRow) joinedRow.withLeft(currentRow) - val buildIter = hashedRelation.get(rowKey) + val buildIter = if (boundStreamedOnlyCondition(currentRow)) { + hashedRelation.get(rowKey) + } else { + null + } new RowIterator { private var found = false override def advanceNext(): Boolean = { while (buildIter != null && buildIter.hasNext) { val nextBuildRow = buildIter.next() - if (boundCondition(joinedRow.withRight(nextBuildRow))) { + if (boundRestCondition(joinedRow.withRight(nextBuildRow))) { if (found && singleJoin) { - throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows(); + throw QueryExecutionErrors.scalarSubqueryReturnsMultipleRows() } found = true return true @@ -279,19 +308,23 @@ trait HashJoin extends JoinCodegenSupport { if (hashedRelation.keyIsUnique) { streamIter.map { current => val key = joinKeys(current) - lazy val matched = hashedRelation.getValue(key) - val exists = !key.anyNull && matched != null && - (condition.isEmpty || boundCondition(joinedRow(current, matched))) + val exists = !key.anyNull && boundStreamedOnlyCondition(current) && { + val matched = hashedRelation.getValue(key) + matched != null && (restCondition.isEmpty || + boundRestCondition(joinedRow(current, matched))) + } result.setBoolean(0, exists) joinedRow(current, result) } } else { streamIter.map { current => val key = joinKeys(current) - lazy val buildIter = hashedRelation.get(key) - val exists = !key.anyNull && buildIter != null && (condition.isEmpty || buildIter.exists { - (row: InternalRow) => boundCondition(joinedRow(current, row)) - }) + val exists = !key.anyNull && boundStreamedOnlyCondition(current) && { + val buildIter = hashedRelation.get(key) + buildIter != null && (restCondition.isEmpty || buildIter.exists { + (row: InternalRow) => boundRestCondition(joinedRow(current, row)) + }) + } result.setBoolean(0, exists) joinedRow(current, result) } @@ -313,16 +346,21 @@ trait HashJoin extends JoinCodegenSupport { streamIter.filter { current => val key = joinKeys(current) lazy val matched = hashedRelation.getValue(key) - key.anyNull || matched == null || - (condition.isDefined && !boundCondition(joinedRow(current, matched))) + // If streamed-only condition is false/null, the full condition can never be true, + // so the row is guaranteed emitted (no probe needed). + key.anyNull || !boundStreamedOnlyCondition(current) || matched == null || + (restCondition.isDefined && !boundRestCondition(joinedRow(current, matched))) } } else { streamIter.filter { current => val key = joinKeys(current) lazy val buildIter = hashedRelation.get(key) - key.anyNull || buildIter == null || (condition.isDefined && !buildIter.exists { - row => boundCondition(joinedRow(current, row)) - }) + // If streamed-only condition is false/null, the full condition can never be true, + // so the row is guaranteed emitted (no probe needed). + key.anyNull || !boundStreamedOnlyCondition(current) || buildIter == null || + (restCondition.isDefined && !buildIter.exists { + row => boundRestCondition(joinedRow(current, row)) + }) } } } @@ -357,6 +395,46 @@ trait HashJoin extends JoinCodegenSupport { } } + /** + * Generates the code for evaluating a streamed-side-only condition on the given stream vars. + */ + protected def genStreamedOnlyCondition( + ctx: CodegenContext, + expr: Expression, + streamVars: Seq[ExprCode]): ExprCode = { + ctx.currentVars = streamVars + val boundExpr = BindReferences.bindReference(expr, streamedPlan.output) + boundExpr.genCode(ctx) + } + + /** + * Generates code evaluating the hoisted streamed-only condition, if any. Returns the + * generated code and the name of the boolean variable it declares, which holds whether the + * streamed row may satisfy the full join condition. Returns ("", "") when no streamed-only + * condition was hoisted. + * + * Callers must fold the boolean into their probe-skip condition instead of emitting a row + * and returning early: HashJoin's doConsume is inlined into the streamed producer's row + * loop, where a bare `return` would skip the producer's loop-cursor write-back (batching + * producers such as ColumnarToRowExec would reprocess the same batch). Folding also keeps + * a single consume site per streamed row, so lazy streamed variables are materialized in + * the same scope as all their uses. + */ + protected def genStreamedOnlyCheck( + ctx: CodegenContext, + input: Seq[ExprCode]): (String, String) = { + streamedOnlyCondition match { + case Some(expr) => + val ev = genStreamedOnlyCondition(ctx, expr, input) + val passed = ctx.freshName("streamedOnlyPassed") + (s""" + |${ev.code} + |boolean $passed = !${ev.isNull} && ${ev.value}; + """.stripMargin, passed) + case None => ("", "") + } + } + override def doProduce(ctx: CodegenContext): String = { streamedPlan.asInstanceOf[CodegenSupport].produce(ctx, this) } @@ -457,14 +535,11 @@ trait HashJoin extends JoinCodegenSupport { val buildVars = genOneSideJoinVars(ctx, matched, buildPlan, setDefaultValue = true) val numOutput = metricTerm(ctx, "numOutputRows") - // filter the output via condition. When there is no condition, skip the `conditionPassed` - // variable and the wrapping `if (!conditionPassed)` / `if (conditionPassed)` branches that - // would always be dead / unconditional. - val hasCondition = condition.isDefined - val conditionPassed = if (hasCondition) ctx.freshName("conditionPassed") else "" - val checkCondition = if (hasCondition) { - val expr = condition.get - // evaluate the variables from build side that used by condition + // Evaluate the rest of the condition (cross-side conjuncts) inside the match loop. + val hasRestCondition = restCondition.isDefined + val conditionPassed = if (hasRestCondition) ctx.freshName("conditionPassed") else "" + val checkCondition = if (hasRestCondition) { + val expr = restCondition.get val eval = evaluateRequiredVariables(buildPlan.output, buildVars, expr.references) ctx.currentVars = input ++ buildVars val ev = @@ -486,8 +561,19 @@ trait HashJoin extends JoinCodegenSupport { case BuildRight => input ++ buildVars } + // Fold the hoisted streamed-only predicate (if any) into the probe below: a streamed row + // that fails it cannot match, so the probe is skipped and the regular path emits the row + // with a null build side. See genStreamedOnlyCheck for why this must not be an early + // emit + return. + val (streamedOnlyCheckCode, streamedOnlyPassed) = genStreamedOnlyCheck(ctx, input) + val skipProbe = if (streamedOnlyPassed.nonEmpty) { + s"$anyNull || !$streamedOnlyPassed" + } else { + anyNull + } + if (keyIsUnique) { - val resetWhenConditionFails = if (hasCondition) { + val resetWhenConditionFails = if (hasRestCondition) { s""" |if (!$conditionPassed) { | $matched = null; @@ -501,8 +587,9 @@ trait HashJoin extends JoinCodegenSupport { s""" |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// find matches from HashedRelation - |UnsafeRow $matched = $anyNull ? null: (UnsafeRow)$relationTerm.getValue(${keyEv.value}); + |UnsafeRow $matched = $skipProbe ? null: (UnsafeRow)$relationTerm.getValue(${keyEv.value}); |${checkCondition.trim} |$resetWhenConditionFails |$numOutput.add(1); @@ -525,13 +612,15 @@ trait HashJoin extends JoinCodegenSupport { } val (conditionGuardOpen, conditionGuardClose) = - if (hasCondition) (s"if ($conditionPassed) {", "}") else ("", "") + if (hasRestCondition) (s"if ($conditionPassed) {", "}") else ("", "") s""" |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// find matches from HashRelation - |$iteratorCls $matches = $anyNull ? null : ($iteratorCls)$relationTerm.get(${keyEv.value}); + |$iteratorCls $matches = $skipProbe ? + | null : ($iteratorCls)$relationTerm.get(${keyEv.value}); |boolean $found = false; |// the last iteration of this loop is to emit an empty row if there is no matched rows. |while ($matches != null && $matches.hasNext() || !$found) { @@ -617,7 +706,23 @@ trait HashJoin extends JoinCodegenSupport { } val (keyEv, anyNull) = genStreamSideJoinKey(ctx, input) - val (matched, checkCondition, _) = getJoinCondition(ctx, input, streamedPlan, buildPlan) + val (matched, checkCondition, _) = restCondition match { + case Some(expr) => getJoinCondition(ctx, expr, input, streamedPlan, buildPlan, None) + case None => + val dummy = ctx.freshName("matched") + (dummy, "", Nil) + } + + // Fold the hoisted streamed-only predicate (if any) into the probe below: a streamed row + // that fails it is guaranteed to be emitted, so the probe is skipped and the regular + // !found path emits the row. See genStreamedOnlyCheck for why this must not be an early + // emit + return. + val (streamedOnlyCheckCode, streamedOnlyPassed) = genStreamedOnlyCheck(ctx, input) + val probeAllowed = if (streamedOnlyPassed.nonEmpty) { + s"!($anyNull) && $streamedOnlyPassed" + } else { + s"!($anyNull)" + } if (keyIsUnique) { val found = ctx.freshName("found") @@ -625,8 +730,9 @@ trait HashJoin extends JoinCodegenSupport { |boolean $found = false; |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// Check if the key has nulls. - |if (!($anyNull)) { + |if ($probeAllowed) { | // Check if the HashedRelation exists. | UnsafeRow $matched = (UnsafeRow)$relationTerm.getValue(${keyEv.value}); | if ($matched != null) { @@ -649,8 +755,9 @@ trait HashJoin extends JoinCodegenSupport { |boolean $found = false; |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// Check if the key has nulls. - |if (!($anyNull)) { + |if ($probeAllowed) { | // Check if the HashedRelation exists. | $iteratorCls $matches = ($iteratorCls)$relationTerm.get(${keyEv.value}); | if ($matches != null) { @@ -682,21 +789,32 @@ trait HashJoin extends JoinCodegenSupport { val matched = ctx.freshName("matched") val buildVars = genOneSideJoinVars(ctx, matched, buildPlan, setDefaultValue = false) - val checkCondition = if (condition.isDefined) { - val expr = condition.get - // evaluate the variables from build side that used by condition - val eval = evaluateRequiredVariables(buildPlan.output, buildVars, expr.references) - // filter the output via condition - ctx.currentVars = input ++ buildVars - val ev = - BindReferences.bindReference(expr, streamedPlan.output ++ buildPlan.output).genCode(ctx) - s""" - |$eval - |${ev.code} - |$existsVar = !${ev.isNull} && ${ev.value}; - """.stripMargin + + // Evaluate the rest of the condition (cross-side conjuncts) inside the match loop. + val checkCondition = restCondition match { + case Some(expr) => + val eval = evaluateRequiredVariables(buildPlan.output, buildVars, expr.references) + ctx.currentVars = input ++ buildVars + val ev = BindReferences.bindReference( + expr, streamedPlan.output ++ buildPlan.output).genCode(ctx) + s""" + |$eval + |${ev.code} + |$existsVar = !${ev.isNull} && ${ev.value}; + """.stripMargin + case None => + s"$existsVar = true;" + } + + // Fold the hoisted streamed-only predicate (if any) into the probe below: a streamed row + // that fails it yields exists = false, so the probe is skipped and the single emit below + // produces the row. See genStreamedOnlyCheck for why this must not be an early emit + + // return. + val (streamedOnlyCheckCode, streamedOnlyPassed) = genStreamedOnlyCheck(ctx, input) + val probeAllowed = if (streamedOnlyPassed.nonEmpty) { + s"!($anyNull) && $streamedOnlyPassed" } else { - s"$existsVar = true;" + s"!($anyNull)" } val resultVar = input ++ Seq(ExprCode.forNonNullValue( @@ -704,13 +822,16 @@ trait HashJoin extends JoinCodegenSupport { if (keyIsUnique) { s""" + |boolean $existsVar = false; |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// find matches from HashedRelation - |UnsafeRow $matched = $anyNull ? null: (UnsafeRow)$relationTerm.getValue(${keyEv.value}); - |boolean $existsVar = false; - |if ($matched != null) { - | $checkCondition + |if ($probeAllowed) { + | UnsafeRow $matched = (UnsafeRow)$relationTerm.getValue(${keyEv.value}); + | if ($matched != null) { + | $checkCondition + | } |} |$numOutput.add(1); |${consume(ctx, resultVar)} @@ -719,15 +840,18 @@ trait HashJoin extends JoinCodegenSupport { val matches = ctx.freshName("matches") val iteratorCls = classOf[Iterator[UnsafeRow]].getName s""" + |boolean $existsVar = false; |// generate join key for stream side |${keyEv.code} + |$streamedOnlyCheckCode |// find matches from HashRelation - |$iteratorCls $matches = $anyNull ? null : ($iteratorCls)$relationTerm.get(${keyEv.value}); - |boolean $existsVar = false; - |if ($matches != null) { - | while (!$existsVar && $matches.hasNext()) { - | UnsafeRow $matched = (UnsafeRow) $matches.next(); - | $checkCondition + |if ($probeAllowed) { + | $iteratorCls $matches = ($iteratorCls)$relationTerm.get(${keyEv.value}); + | if ($matches != null) { + | while (!$existsVar && $matches.hasNext()) { + | UnsafeRow $matched = (UnsafeRow) $matches.next(); + | $checkCondition + | } | } |} |$numOutput.add(1); diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/JoinCodegenSupport.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/JoinCodegenSupport.scala index 6496f9a0006e2..056ddca2072d7 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/JoinCodegenSupport.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/JoinCodegenSupport.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.joins -import org.apache.spark.sql.catalyst.expressions.{BindReferences, BoundReference} +import org.apache.spark.sql.catalyst.expressions.{BindReferences, BoundReference, Expression} import org.apache.spark.sql.catalyst.expressions.codegen._ import org.apache.spark.sql.catalyst.expressions.codegen.Block._ import org.apache.spark.sql.execution.{CodegenSupport, SparkPlan} @@ -40,6 +40,25 @@ trait JoinCodegenSupport extends CodegenSupport with BaseJoinExec { streamPlan: SparkPlan, buildPlan: SparkPlan, buildRow: Option[String] = None): (String, String, Seq[ExprCode]) = { + if (condition.isDefined) { + getJoinCondition(ctx, condition.get, streamVars, streamPlan, buildPlan, buildRow) + } else { + val buildSideRow = buildRow.getOrElse(ctx.freshName("buildRow")) + val buildVars = genOneSideJoinVars(ctx, buildSideRow, buildPlan, setDefaultValue = false) + (buildSideRow, "", buildVars) + } + } + + /** + * Generate the (non-equi) condition used to filter joined rows from an explicit expression. + */ + protected def getJoinCondition( + ctx: CodegenContext, + conditionExpr: Expression, + streamVars: Seq[ExprCode], + streamPlan: SparkPlan, + buildPlan: SparkPlan, + buildRow: Option[String]): (String, String, Seq[ExprCode]) = { val buildSideRow = buildRow.getOrElse(ctx.freshName("buildRow")) val buildVars = genOneSideJoinVars(ctx, buildSideRow, buildPlan, setDefaultValue = false) // We want to evaluate the passed streamVars. However, evaluation modifies the contained @@ -47,25 +66,21 @@ trait JoinCodegenSupport extends CodegenSupport with BaseJoinExec { // full outer join will want to evaluate streamVars in a different scope than the // condition check). Because of this, we first make a copy. val streamVars2 = streamVars.map(_.copy()) - val checkCondition = if (condition.isDefined) { - val expr = condition.get - // evaluate the variables that are used by the condition - val eval = evaluateRequiredVariables(streamPlan.output ++ buildPlan.output, - streamVars2 ++ buildVars, expr.references) + // evaluate the variables that are used by the condition + val eval = evaluateRequiredVariables(streamPlan.output ++ buildPlan.output, + streamVars2 ++ buildVars, conditionExpr.references) - // filter the output via condition - ctx.currentVars = streamVars2 ++ buildVars - val ev = - BindReferences.bindReference(expr, streamPlan.output ++ buildPlan.output).genCode(ctx) - val skipRow = s"${ev.isNull} || !${ev.value}" + // filter the output via condition + ctx.currentVars = streamVars2 ++ buildVars + val ev = BindReferences.bindReference( + conditionExpr, streamPlan.output ++ buildPlan.output).genCode(ctx) + val skipRow = s"${ev.isNull} || !${ev.value}" + val checkCondition = s""" |$eval |${ev.code} |if (!($skipRow)) """.stripMargin - } else { - "" - } (buildSideRow, checkCondition, buildVars) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinEvaluatorFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinEvaluatorFactory.scala index 2b6a19dfa8a8d..cf358960c7105 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinEvaluatorFactory.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinEvaluatorFactory.scala @@ -37,7 +37,9 @@ class SortMergeJoinEvaluatorFactory( sizeInBytesSpillThreshold: Long, numOutputRows: SQLMetric, spillSize: SQLMetric, - onlyBufferFirstMatchedRow: Boolean) + onlyBufferFirstMatchedRow: Boolean, + streamedOnlyCondition: Option[Expression] = None, + restCondition: Option[Expression] = None) extends PartitionEvaluatorFactory[InternalRow, InternalRow] { override def createEvaluator(): PartitionEvaluator[InternalRow, InternalRow] = new SortMergeJoinEvaluator @@ -136,10 +138,17 @@ class SortMergeJoinEvaluatorFactory( spillSize, cleanupResources) val rightNullRow = new GenericInternalRow(right.output.length) + val boundStreamedOnly: InternalRow => Boolean = streamedOnlyCondition.map { + Predicate.create(_, left.output).eval _ + }.getOrElse((_: InternalRow) => true) + val boundRest: InternalRow => Boolean = restCondition.map { + Predicate.create(_, left.output ++ right.output).eval _ + }.getOrElse((_: InternalRow) => true) new LeftOuterIterator( smjScanner, rightNullRow, - boundCondition, + boundStreamedOnly, + boundRest, resultProj, numOutputRows).toScala @@ -156,10 +165,17 @@ class SortMergeJoinEvaluatorFactory( spillSize, cleanupResources) val leftNullRow = new GenericInternalRow(left.output.length) + val boundStreamedOnly: InternalRow => Boolean = streamedOnlyCondition.map { + Predicate.create(_, right.output).eval _ + }.getOrElse((_: InternalRow) => true) + val boundRest: InternalRow => Boolean = restCondition.map { + Predicate.create(_, left.output ++ right.output).eval _ + }.getOrElse((_: InternalRow) => true) new RightOuterIterator( smjScanner, leftNullRow, - boundCondition, + boundStreamedOnly, + boundRest, resultProj, numOutputRows).toScala @@ -217,6 +233,12 @@ class SortMergeJoinEvaluatorFactory( }.toScala case LeftAnti => + val boundStreamedOnly: InternalRow => Boolean = streamedOnlyCondition.map { + Predicate.create(_, left.output).eval _ + }.getOrElse((_: InternalRow) => true) + val boundRest: InternalRow => Boolean = restCondition.map { + Predicate.create(_, left.output ++ right.output).eval _ + }.getOrElse((_: InternalRow) => true) new RowIterator { private[this] var currentLeftRow: InternalRow = _ private[this] val smjScanner = new SortMergeJoinScanner( @@ -236,6 +258,11 @@ class SortMergeJoinEvaluatorFactory( override def advanceNext(): Boolean = { while (smjScanner.findNextOuterJoinRows()) { currentLeftRow = smjScanner.getStreamedRow + if (!boundStreamedOnly(currentLeftRow)) { + // streamed-only predicate is false/null -> full condition is false -> emit row + numOutputRows += 1 + return true + } val currentRightMatches = smjScanner.getBufferedMatches if (currentRightMatches == null || currentRightMatches.length == 0) { numOutputRows += 1 @@ -245,7 +272,7 @@ class SortMergeJoinEvaluatorFactory( val rightMatchesIterator = currentRightMatches.generateIterator() while (!found && rightMatchesIterator.hasNext) { joinRow(currentLeftRow, rightMatchesIterator.next()) - if (boundCondition(joinRow)) { + if (boundRest(joinRow)) { found = true } } @@ -261,6 +288,12 @@ class SortMergeJoinEvaluatorFactory( }.toScala case j: ExistenceJoin => + val boundStreamedOnly: InternalRow => Boolean = streamedOnlyCondition.map { + Predicate.create(_, left.output).eval _ + }.getOrElse((_: InternalRow) => true) + val boundRest: InternalRow => Boolean = restCondition.map { + Predicate.create(_, left.output ++ right.output).eval _ + }.getOrElse((_: InternalRow) => true) new RowIterator { private[this] var currentLeftRow: InternalRow = _ private[this] val result: InternalRow = new GenericInternalRow(Array[Any](null)) @@ -281,13 +314,20 @@ class SortMergeJoinEvaluatorFactory( override def advanceNext(): Boolean = { while (smjScanner.findNextOuterJoinRows()) { currentLeftRow = smjScanner.getStreamedRow + if (!boundStreamedOnly(currentLeftRow)) { + // streamed-only predicate is false/null -> full condition is false -> + // exists=false + result.setBoolean(0, false) + numOutputRows += 1 + return true + } val currentRightMatches = smjScanner.getBufferedMatches var found = false if (currentRightMatches != null && currentRightMatches.length > 0) { val rightMatchesIterator = currentRightMatches.generateIterator() while (!found && rightMatchesIterator.hasNext) { joinRow(currentLeftRow, rightMatchesIterator.next()) - if (boundCondition(joinRow)) { + if (boundRest(joinRow)) { found = true } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala index 51604cdfedf1c..92445fb3b554a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/SortMergeJoinExec.scala @@ -121,6 +121,17 @@ case class SortMergeJoinExec( } } + /** + * For join types that preserve all streamed rows, split the condition into + * streamed-only and rest. The streamed-only part can be evaluated once per streamed row + * before walking the buffered matches. + */ + private lazy val (streamedOnlyCondition, restCondition): + (Option[Expression], Option[Expression]) = { + StreamedSideJoinCondition.split( + condition, joinType, streamedPlan, conf.splitStreamedSideJoinCondition) + } + protected override def doExecute(): RDD[InternalRow] = { val numOutputRows = longMetric("numOutputRows") val spillSize = longMetric("spillSize") @@ -140,7 +151,9 @@ case class SortMergeJoinExec( sizeInBytesSpillThreshold, numOutputRows, spillSize, - onlyBufferFirstMatchedRow + onlyBufferFirstMatchedRow, + streamedOnlyCondition, + restCondition ) if (conf.usePartitionEvaluator) { left.execute().zipPartitionsWithEvaluator(right.execute(), evaluatorFactory) @@ -513,7 +526,28 @@ case class SortMergeJoinExec( s"SortMergeJoin.doProduce should not take $x as the JoinType") } - val (streamedBeforeLoop, condCheck, loadStreamed) = if (condition.isDefined) { + // Generate streamed-only condition check for join types that preserve streamed rows. + val (streamedOnlyPre, streamedOnlyGuard) = + if (streamedOnlyCondition.isDefined) { + ctx.currentVars = streamedVars + val ev = BindReferences.bindReference( + streamedOnlyCondition.get, streamedPlan.output).genCode(ctx) + val isNullVar = ctx.freshName("streamedOnlyIsNull") + val valueVar = ctx.freshName("streamedOnlyValue") + val pre = + s""" + |${ev.code} + |boolean $isNullVar = ${ev.isNull}; + |boolean $valueVar = ${ev.value}; + """.stripMargin + (pre, Some(s"!$isNullVar && $valueVar")) + } else { + ("", None) + } + + val conditionForCodegen = if (streamedOnlyGuard.isDefined) restCondition else condition + + val (streamedBeforeLoop, condCheck, loadStreamed) = if (conditionForCodegen.isDefined) { // Split the code of creating variables based on whether it's used by condition or not. val loaded = ctx.freshName("loaded") val (streamedBefore, streamedAfter) = splitVarsByCondition(streamedOutput, streamedVars) @@ -521,7 +555,7 @@ case class SortMergeJoinExec( // Generate code for condition ctx.currentVars = streamedVars ++ bufferedVars val cond = BindReferences.bindReference( - condition.get, streamedPlan.output ++ bufferedPlan.output).genCode(ctx) + conditionForCodegen.get, streamedPlan.output ++ bufferedPlan.output).genCode(ctx) // Evaluate the columns those used by condition before loop val before = joinType match { case LeftAnti => @@ -572,10 +606,14 @@ case class SortMergeJoinExec( (evaluateVariables(streamedVars), "", "") } - val beforeLoop = + val existsVarDecl = existsVar.map(v => s"boolean $v = false;").getOrElse("") + + val beforeLoopWithoutGuard = s""" |${streamedVarDecl.mkString("\n")} |${streamedBeforeLoop.trim} + |$streamedOnlyPre + |$existsVarDecl |scala.collection.Iterator<UnsafeRow> $iterator = $matches.generateIterator(); """.stripMargin val outputRow = @@ -583,10 +621,40 @@ case class SortMergeJoinExec( |$numOutput.add(1); |${consume(ctx, resultVars)} """.stripMargin + val guardOutputRow = joinType match { + case LeftOuter | RightOuter => + val defaultBufferedVars = + genOneSideJoinVars(ctx, bufferedRow, bufferedPlan, setDefaultValue = true) + val guardResultVars = joinType match { + case RightOuter => defaultBufferedVars ++ streamedVars + case _ => streamedVars ++ defaultBufferedVars + } + s""" + |$numOutput.add(1); + |${consume(ctx, guardResultVars)} + """.stripMargin + case _ => outputRow + } val findNextJoinRows = s"$findNextJoinRowsFuncName($streamedInput, $bufferedInput)" val thisPlan = ctx.addReferenceObj("plan", this) val eagerCleanup = s"$thisPlan.cleanupResources();" + // For join types with a streamed-only guard, prepend the guard to beforeLoop + // so the row is emitted before the inner loop. + val beforeLoop = streamedOnlyGuard match { + case Some(guard) => + s""" + |$beforeLoopWithoutGuard + |if (!($guard)) { + | InternalRow $bufferedRow = null; + | $loadStreamed + | $guardOutputRow + | continue; + |} + """.stripMargin + case None => beforeLoopWithoutGuard + } + val doJoin = joinType match { case _: InnerLike => val cleanedFlag = @@ -781,7 +849,6 @@ case class SortMergeJoinExec( |while ($streamedInput.hasNext()) { | $findNextJoinRows; | $beforeLoop - | boolean $exists = false; | | while (!$exists && $matchIterator.hasNext()) { | InternalRow $bufferedRow = (InternalRow) $matchIterator.next(); @@ -1299,11 +1366,13 @@ private[joins] class SortMergeJoinScanner( private class LeftOuterIterator( smjScanner: SortMergeJoinScanner, rightNullRow: InternalRow, - boundCondition: InternalRow => Boolean, + boundStreamedOnly: InternalRow => Boolean, + boundRest: InternalRow => Boolean, resultProj: InternalRow => InternalRow, numOutputRows: SQLMetric) extends OneSideOuterIterator( - smjScanner, rightNullRow, boundCondition, resultProj, numOutputRows) { + smjScanner, rightNullRow, boundStreamedOnly, boundRest, + resultProj, numOutputRows) { protected override def setStreamSideOutput(row: InternalRow): Unit = joinedRow.withLeft(row) protected override def setBufferedSideOutput(row: InternalRow): Unit = joinedRow.withRight(row) @@ -1315,10 +1384,13 @@ private class LeftOuterIterator( private class RightOuterIterator( smjScanner: SortMergeJoinScanner, leftNullRow: InternalRow, - boundCondition: InternalRow => Boolean, + boundStreamedOnly: InternalRow => Boolean, + boundRest: InternalRow => Boolean, resultProj: InternalRow => InternalRow, numOutputRows: SQLMetric) - extends OneSideOuterIterator(smjScanner, leftNullRow, boundCondition, resultProj, numOutputRows) { + extends OneSideOuterIterator( + smjScanner, leftNullRow, boundStreamedOnly, boundRest, + resultProj, numOutputRows) { protected override def setStreamSideOutput(row: InternalRow): Unit = joinedRow.withRight(row) protected override def setBufferedSideOutput(row: InternalRow): Unit = joinedRow.withLeft(row) @@ -1336,14 +1408,17 @@ private class RightOuterIterator( * * @param smjScanner a scanner that streams rows and buffers any matching rows * @param bufferedSideNullRow the default row to return when a streamed row has no matches - * @param boundCondition an additional filter condition for buffered rows + * @param boundStreamedOnly a predicate evaluated on the streamed row only + * @param boundRest a predicate evaluated on the joined (left ++ right) row for the residual + * condition, bound to the physical row order produced by this iterator * @param resultProj how the output should be projected * @param numOutputRows an accumulator metric for the number of rows output */ private abstract class OneSideOuterIterator( smjScanner: SortMergeJoinScanner, bufferedSideNullRow: InternalRow, - boundCondition: InternalRow => Boolean, + boundStreamedOnly: InternalRow => Boolean, + boundRest: InternalRow => Boolean, resultProj: InternalRow => InternalRow, numOutputRows: SQLMetric) extends RowIterator { @@ -1368,12 +1443,15 @@ private abstract class OneSideOuterIterator( rightMatchesIterator = null if (smjScanner.findNextOuterJoinRows()) { setStreamSideOutput(smjScanner.getStreamedRow) - if (smjScanner.getBufferedMatches.isEmpty) { + if (!boundStreamedOnly(smjScanner.getStreamedRow)) { + // Streamed-only predicate is false/null -> full condition is false -> emit null-padded row. + setBufferedSideOutput(bufferedSideNullRow) + } else if (smjScanner.getBufferedMatches.isEmpty) { // There are no matching rows in the buffer, so return the null row setBufferedSideOutput(bufferedSideNullRow) } else { - // Find the next row in the buffer that satisfied the bound condition - if (!advanceBufferUntilBoundConditionSatisfied()) { + // Find the next row in the buffer that satisfied the rest condition + if (!advanceBufferUntilRestConditionSatisfied()) { setBufferedSideOutput(bufferedSideNullRow) } } @@ -1385,10 +1463,10 @@ private abstract class OneSideOuterIterator( } /** - * Advance to the next row in the buffer that satisfies the bound condition. + * Advance to the next row in the buffer that satisfies the rest condition. * @return whether there is such a row in the current buffer. */ - private def advanceBufferUntilBoundConditionSatisfied(): Boolean = { + private def advanceBufferUntilRestConditionSatisfied(): Boolean = { var foundMatch: Boolean = false if (rightMatchesIterator == null) { rightMatchesIterator = smjScanner.getBufferedMatches.generateIterator() @@ -1396,13 +1474,18 @@ private abstract class OneSideOuterIterator( while (!foundMatch && rightMatchesIterator.hasNext) { setBufferedSideOutput(rightMatchesIterator.next()) - foundMatch = boundCondition(joinedRow) + foundMatch = boundRest(joinedRow) } foundMatch } override def advanceNext(): Boolean = { - val r = advanceBufferUntilBoundConditionSatisfied() || advanceStream() + // Only walk the buffered matches if we are in the middle of iterating them for the + // current streamed row. If the iterator is null, advanceStream() has just emitted a + // null-padded row (either no matches or the streamed-only predicate was false), so we + // must move to the next streamed row rather than re-create the match iterator. + val r = (rightMatchesIterator != null && advanceBufferUntilRestConditionSatisfied()) || + advanceStream() if (r) numOutputRows += 1 r } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/StreamedSideJoinCondition.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/StreamedSideJoinCondition.scala new file mode 100644 index 0000000000000..e5e5778e5d2f5 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/StreamedSideJoinCondition.scala @@ -0,0 +1,68 @@ +/* + * 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.execution.joins + +import org.apache.spark.sql.catalyst.expressions.{And, Expression, ExprUtils, PredicateHelper} +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, JoinType, LeftAnti, LeftOuter, RightOuter} +import org.apache.spark.sql.execution.SparkPlan + +/** + * Helper for splitting a join condition into a streamed-side-only part and the remaining part, + * for join types that preserve all streamed rows. + * + * The split is valid because for these join types, if a streamed-only predicate S is FALSE/NULL + * for a streamed row, the full join condition S AND other(S, B) is FALSE/NULL for ANY buffered + * row B. The streamed row outcome is therefore already determined before any probe: it is + * emitted for outer/existence joins, and emitted for left anti joins when no match exists. + */ +private[joins] object StreamedSideJoinCondition extends PredicateHelper { + + /** + * Splits `condition` into conjuncts that reference only the streamed side and the remaining + * conjuncts, when `splitEnabled` is true and the join type preserves all streamed rows. + * The streamed-only part can be evaluated once per streamed row before probing/walking the + * buffered matches. Returns `(None, condition)` unchanged otherwise. + * + * Only conjuncts that can be evaluated unconditionally are hoisted (see + * [[ExprUtils.canEvaluateUnconditionally]]): the hoisted part runs for every streamed + * row, including rows that have no buffered match and would never evaluate the conjunct + * otherwise, so hoisting a conjunct that can throw could turn a valid outer/anti result + * into an exception, and hoisting a non-deterministic conjunct would change how many + * times it is evaluated per streamed row. A whitelist of total expression families is + * used instead of [[Expression.throwable]], which is opt-in metadata that UDFs such as + * ScalaUDF do not override, so a throwing UDF would report non-throwable. + */ + def split( + condition: Option[Expression], + joinType: JoinType, + streamedPlan: SparkPlan, + splitEnabled: Boolean): (Option[Expression], Option[Expression]) = { + val supported = joinType match { + case LeftAnti | LeftOuter | RightOuter | _: ExistenceJoin => true + case _ => false + } + if (condition.isDefined && splitEnabled && supported) { + val conjuncts = splitConjunctivePredicates(condition.get) + val (streamedOnly, rest) = conjuncts.partition(p => + ExprUtils.canEvaluateUnconditionally(p) && p.references.subsetOf(streamedPlan.outputSet)) + (streamedOnly.reduceOption(And), rest.reduceOption(And)) + } else { + (None, condition) + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptAccumulator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptAccumulator.scala index 245e5204104c6..c2b11c9776765 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptAccumulator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptAccumulator.scala @@ -23,7 +23,7 @@ import scala.util.control.NonFatal import org.apache.spark.SparkContext import org.apache.spark.internal.{LogEntry, Logging} import org.apache.spark.sql.Dataset -import org.apache.spark.sql.execution.{BaseSubqueryExec, QueryExecution, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{BaseSubqueryExec, ProjectedBroadcastValueSubqueryExec, QueryExecution, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, BroadcastExchangeLike, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.util.{AccumulatorV2, LastAttemptAccumulator} @@ -368,7 +368,7 @@ object SQLLastAttemptAccumulator extends Logging { // ``` // will launch stages in scope of child. scopeIds(s.child) - case _: SubqueryBroadcastExec => + case _: SubqueryBroadcastExec | _: ProjectedBroadcastValueSubqueryExec => // Used by DPP filter only, not part of main flow of query execution. Nil case _: SubqueryAdaptiveBroadcastExec => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala similarity index 89% rename from sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala rename to sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala index 037abc207298a..45c06f86bd416 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplans.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/MergeSubplans.scala @@ -15,16 +15,16 @@ * limitations under the License. */ -package org.apache.spark.sql.catalyst.optimizer +package org.apache.spark.sql.execution.planmerging import scala.collection.mutable.ArrayBuffer import org.apache.spark.sql.catalyst.expressions._ -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LeafNode, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE} +import org.apache.spark.sql.catalyst.optimizer.{NonGroupingAggregateReference, ScalarSubqueryReference} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, CTERelationDef, CTERelationRef, LogicalPlan, OneRowRelation, Project, Subquery, WithCTE} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CTE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE, TreePattern} +import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CTE, NO_GROUPING_AGGREGATE_REFERENCE, SCALAR_SUBQUERY, SCALAR_SUBQUERY_REFERENCE} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.DataType /** * This rule tries to merge multiple subplans that have one row result. This can be either the plan @@ -340,40 +340,3 @@ object MergeSubplans extends Rule[LogicalPlan] { } } -/** - * Temporal reference to a subquery which is added to a `PlanMerger`. - * - * @param level The level of the replaced subquery. It defines the `PlanMerger` instance into which - * the subquery is merged. - * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`. - * @param outputIndex The index of the output attribute of the merged plan. - * @param dataType The dataType of original scalar subquery. - * @param exprId The expression id of the original scalar subquery. - */ -case class ScalarSubqueryReference( - level: Int, - mergedPlanIndex: Int, - outputIndex: Int, - override val dataType: DataType, - exprId: ExprId) extends LeafExpression with Unevaluable { - override def nullable: Boolean = true - - final override val nodePatterns: Seq[TreePattern] = Seq(SCALAR_SUBQUERY_REFERENCE) -} - -/** - * Temporal reference to a non-grouping aggregate which is added to a `PlanMerger`. - * - * @param level The level of the replaced aggregate. It defines the `PlanMerger` instance into which - * the aggregate is merged. - * @param mergedPlanIndex The index of the merged plan in the `PlanMerger`. - * @param outputIndices The indices of the output attributes of the merged plan. - * @param output The output of original aggregate. - */ -case class NonGroupingAggregateReference( - level: Int, - mergedPlanIndex: Int, - outputIndices: Seq[Int], - override val output: Seq[Attribute]) extends LeafNode { - final override val nodePatterns: Seq[TreePattern] = Seq(NO_GROUPING_AGGREGATE_REFERENCE) -} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala new file mode 100644 index 0000000000000..0c967421283f3 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala @@ -0,0 +1,1070 @@ +/* + * 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.execution.planmerging + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeMap, AttributeSet, Expression, ExpressionSet, If, Literal, NamedExpression, Or, SortOrder} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.{Cross, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project} +import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown} +import org.apache.spark.sql.internal.SQLConf + +/** + * Result of attempting to merge a plan via [[PlanMerger.merge]]. + * + * @param mergedPlan The resulting plan, either: + * - An existing cached plan (if identical match found) + * - A newly merged plan combining the input with a cached plan + * - The original input plan (if no merge was possible) + * @param mergedPlanIndex The index of this plan in the PlanMerger's cache. + * @param outputMap Maps attributes of the input plan to their positional index in + * `mergedPlan.plan.output`. The index remains stable across subsequent + * [[PlanMerger.merge]] calls because outputs are only ever appended. + */ +case class MergeResult( + mergedPlan: MergedPlan, + mergedPlanIndex: Int, + outputMap: AttributeMap[Int]) + +/** + * Represents a plan in the PlanMerger's cache. + * + * @param plan The logical plan, which may have been merged from multiple original plans. + * @param merged Whether this plan is the result of merging two or more plans (true), or + * is an original unmerged plan (false). Merged plans typically require special + * handling such as wrapping in CTEs. + */ +case class MergedPlan(plan: LogicalPlan, merged: Boolean) + +object PlanMerger { + // Marker tag placed on Filter nodes that were produced by filter propagation. Its presence + // signals that the Filter's condition is already an OR of propagated filter attributes and + // its child Project already contains the corresponding aliases, so a subsequent merge only + // needs to add one new alias for the incoming plan rather than wrapping both sides again. + val MERGED_FILTER_TAG: TreeNodeTag[Unit] = TreeNodeTag("mergedFilter") + + // Global counter for generating unique names for propagated filter attributes across all + // PlanMerger instances. + private[planmerging] val curId = new java.util.concurrent.atomic.AtomicLong() + private[planmerging] def newId: Long = curId.getAndIncrement() +} + +/** + * A stateful utility for merging identical or similar logical plans to enable query plan reuse. + * + * `PlanMerger` maintains a cache of previously seen plans and attempts to either: + * 1. Reuse an identical plan already in the cache + * 2. Merge a new plan with a cached plan by combining their outputs + * + * The merging process preserves semantic equivalence while combining outputs from multiple + * plans into a single plan. This is primarily used by [[MergeSubplans]] to deduplicate subplan + * execution. + * + * Supported plan types for merging: + * - [[Project]]: Merges project lists + * - [[Aggregate]]: Merges aggregate expressions with identical grouping + * - [[Filter]]: Requires identical filter conditions + * - [[Join]]: Requires identical join type, hints, and conditions + * + * When `filterPropagationEnabled` is true, non-grouping [[Aggregate]]s over the same base plan + * with different [[Filter]] conditions can also be merged. The filter conditions are exposed as + * boolean [[Project]] attributes and consumed at the [[Aggregate]] as FILTER clauses. + * When both sides carry a [[Filter]] (the symmetric case), merging broadens the scan to OR(f1, f2), + * which may reduce IO pruning. This path is separately gated by + * `symmetricFilterPropagationEnabled`. + * When plans also differ in intermediate [[Project]] expressions, those are wrapped with + * `If(filterAttr, expr, null)` to avoid computing the expression for rows that do not match that + * side's filter condition. + * Filter propagation also works through [[Join]] nodes: a filter on one child of the join produces + * a boolean attribute that flows through the join output to the enclosing [[Aggregate]]. + * Propagation is only safe when the filter originates from the non-nullable side of the join, as + * enforced by `filterSafeForJoin`. When the filter is on the nullable side, the merged base plan + * restores rows that were filtered out of the nullable child, turning what were unmatched + * NULL-padded rows in the original plan into matched rows with real column values. This changes the + * result of expressions like `coalesce(col, default)` in the aggregate: an originally unmatched row + * would have contributed `default` via `coalesce(NULL, default)`, but in the merged plan it is + * matched, its real column value fails the filter, and `FILTER (WHERE false)` discards it entirely. + * Propagation is also skipped when both the left and right children simultaneously produce filter + * attributes, as combining them would require an additional AND alias above the join (not yet + * supported). + * + * {{{ + * // Input plans + * Aggregate [sum(a) AS sum_a] Aggregate [max(d) AS max_d] + * +- Filter (a < 1) +- Project [udf(a) AS d] + * +- Scan t +- Filter (a > 1) + * +- Scan t + * + * // Merged plan + * Aggregate [sum(a) FILTER f0 AS sum_a, max(d0) FILTER f1 AS max_d] + * +- Project [a, If(f1, udf(a), null) AS d0, f0, f1] + * +- Filter (f0 OR f1) [MERGED_FILTER_TAG] + * +- Project [a, (a < 1) AS f0, (a > 1) AS f1] + * +- Scan t + * }}} + * + * @example + * {{{ + * val merger = PlanMerger() + * val result1 = merger.merge(plan1) // Adds plan1 to cache + * val result2 = merger.merge(plan2) // Merges with plan1 if compatible + * // result2.mergedPlan.merged == true if plans were merged + * // result2.outputMap maps plan2's attributes to the merged plan's attributes + * }}} + */ +class PlanMerger( + filterPropagationEnabled: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED), + symmetricFilterPropagationEnabled: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED), + filterPropagationThroughJoinEnabled: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_THROUGH_JOIN_ENABLED), + dsv2SymmetricFilterPropagationEnabled: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED), + dsv2AllowKeyGroupedPartitioningDegradation: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION), + dsv2AllowOrderingDegradation: Boolean = + SQLConf.get.getConf(SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION)) { + val cache = mutable.ArrayBuffer.empty[MergedPlan] + + /** + * Attempts to merge the given plan with cached plans, or adds it to the cache. + * + * The method tries the following in order: + * 1. Check if an identical plan exists in cache (using canonicalized comparison) + * 2. Try to merge with each cached plan using [[tryMergePlans]] + * 3. If no merge is possible, add as a new cache entry + * + * @param plan The logical plan to merge or cache. + * @param subqueryPlan If the logical plan is a subquery plan. + * @return A [[MergeResult]] containing: + * - The merged/cached plan to use + * - Its index in the cache + * - An attribute mapping for rewriting expressions + */ + def merge(plan: LogicalPlan, subqueryPlan: Boolean): MergeResult = { + cache.zipWithIndex.collectFirst(Function.unlift { + case (mp, i) => + checkIdenticalPlans(plan, mp.plan).map { _ => + // Identical subquery expression plans are not marked as `merged` as the + // `ReusedSubqueryExec` rule can handle them without extracting the plans to CTEs. + // But, when a non-subquery subplan is identical to a cached plan we need to mark the plan + // `merged` and so extract it to a CTE later. + val newMergedPlan = MergedPlan(mp.plan, mp.merged || !subqueryPlan) + cache(i) = newMergedPlan + val outputMap = AttributeMap(plan.output.zipWithIndex) + MergeResult(newMergedPlan, i, outputMap) + }.orElse { + tryMergePlans(plan, mp.plan, MergeContext(filterPropagationSupported = false)).collect { + case TryMergeResult(mergedPlan, npMapping, None, None, None, _) => + val newMergedPlan = MergedPlan(mergedPlan, true) + cache(i) = newMergedPlan + val outputMap = AttributeMap(npMapping.iterator.map { case (origAttr, mergedAttr) => + origAttr -> mergedPlan.output.indexWhere(_.exprId == mergedAttr.exprId) + }.toSeq) + MergeResult(newMergedPlan, i, outputMap) + } + } + case _ => None + }).getOrElse { + val newMergedPlan = MergedPlan(plan, false) + cache += newMergedPlan + val outputMap = AttributeMap(plan.output.zipWithIndex) + MergeResult(newMergedPlan, cache.length - 1, outputMap) + } + } + + /** + * Returns all plans currently in the cache as an immutable indexed sequence. + * + * @return An indexed sequence of [[MergedPlan]]s in cache order. The index of each plan + * corresponds to the `mergedPlanIndex` returned by [[merge]]. + */ + def mergedPlans(): IndexedSeq[MergedPlan] = cache.toIndexedSeq + + // If 2 plans are identical return the attribute mapping from the new to the cached version. + private def checkIdenticalPlans( + newPlan: LogicalPlan, + cachedPlan: LogicalPlan): Option[AttributeMap[Attribute]] = { + if (newPlan.canonicalized == cachedPlan.canonicalized) { + Some(AttributeMap(newPlan.output.zip(cachedPlan.output))) + } else { + None + } + } + + /** + * Result of a successful [[tryMergePlans]] call. + * + * @param mergedPlan The combined logical plan. + * @param newPlanMapping Mapping from attributes in the new plan to the corresponding + * attributes in the merged plan. Used by parent nodes to remap + * new-plan-side expressions. + * @param newPlanFilter A boolean [[Attribute]] in the merged plan that encodes the filter + * condition from the new plan's side, to be applied as an aggregate + * `FILTER (WHERE ...)` clause when the propagation reaches an enclosing + * [[Aggregate]] node. The boolean component is `true` if the attribute was + * freshly aliased and must be appended to enclosing [[Project]] nodes, or + * `false` if it was reused from an existing alias already present in the + * merged plan. `None` when no differing filter was propagated. + * @param cachedPlanFilter Like `newPlanFilter` but for the cached plan's side. Always a freshly + * created alias when present, so no `isNew` flag is needed. + * @param dsv2Merged Whether an (equal-strict) DSv2 scan merge occurred anywhere in this merged + * subtree. Unlike `dsv2DeferredScan` (consumed at the enclosing Filter that + * builds the scan), this fact is propagated all the way up: it lets a Filter + * pair that is NOT the innermost one still recognize the merge below it and + * apply the DSv2-symmetric exemption. It only ever gates behaviour when + * `dsv2SymmetricFilterPropagationEnabled` is on. + */ + case class TryMergeResult( + mergedPlan: LogicalPlan, + newPlanMapping: AttributeMap[Attribute], + newPlanFilter: Option[(Attribute, Boolean)] = None, + cachedPlanFilter: Option[Attribute] = None, + dsv2DeferredScan: Option[DSv2DeferredScan] = None, + dsv2Merged: Boolean = false) + + /** + * Carries a DSv2 scan whose build has been deferred from the leaf up to the enclosing [[Filter]], + * so the merged scan is built exactly once per merge round (strict + best-effort filters + * together) rather than once strict-only at the leaf and then rebuilt at the Filter. + * + * The relation to rebuild from is not carried here: the deferring leaf leaves it in the plan as + * the placeholder [[TryMergeResult.mergedPlan]] (the sole bare [[DataSourceV2Relation]] in the + * subtree), and `tryBuildFilterDSv2ScanChild` recovers it from there. Only what the tree does NOT + * hold is carried: the projected `unionAttrs` and the `strictFilters` to re-enforce. + * + * @param unionAttrs The union of both sides' projected columns the merged scan must produce. + * @param strictFilters The strict pushed filters that must be re-enforced by the rebuilt scan. + * @param requiredKeyGroupedPartitioning The key-grouped partitioning the merged scan must + * reproduce to keep both inputs not-worse (the inputs' combined report, in the merged + * relation's attribute space); empty means no requirement. Enforced unless + * `dsv2AllowKeyGroupedPartitioningDegradation` is set. + * @param requiredOrdering The output ordering the merged scan must satisfy likewise; empty means + * no requirement. Enforced unless `dsv2AllowOrderingDegradation` is set. + */ + case class DSv2DeferredScan( + unionAttrs: Seq[Attribute], + strictFilters: Seq[Expression], + requiredKeyGroupedPartitioning: Seq[Expression], + requiredOrdering: Seq[SortOrder]) + + /** + * Context threaded DOWN through [[tryMergePlans]] recursion. + * + * Invariants: + * - `filterAboveScan` is set true ONLY by the `(Filter, Filter)` arm; the leaf defers building + * the merged DSv2 scan ONLY when it is true. A deferred result flows leaf -> (pass-through + * Project arms, which propagate `dsv2DeferredScan` unchanged) -> the enclosing `(Filter, + * Filter)` arm, which builds it once (via `tryBuildFilterDSv2ScanChild`) and returns + * `dsv2DeferredScan = None`. + * - A `(Filter, Filter)`'s children CAN themselves be Filters: `PartitionPruning` and + * `InjectRuntimeFilter` insert a Filter above an existing one after `CombineFilters` has run, + * and the `PushDownPredicates` pass that would re-combine them runs only in a later batch. The + * deferral is consumed at the INNERMOST `(Filter, Filter)` pair (which builds the scan); an + * enclosing pair sees `dsv2DeferredScan = None`, but `dsv2Merged` still marks the merge below + * it, so it can apply the DSv2-symmetric exemption too. This is sound because + * `V2ScanRelationPushDown` only pushes the innermost (scan-adjacent) Filter chain to the scan: + * an outer stacked Filter was never scan pruning (a runtime filter inserted after pushdown, or + * one separated from the scan by an operator that blocks pushdown), so OR-widening it above the + * built scan drops no pruning. + * - The `merge()` pattern requiring `dsv2DeferredScan = None` is the fail-safe backstop: if a + * deferred scan somehow reached `merge()` unbuilt, it declines the merge rather than emitting a + * plan with a placeholder relation. + * - Eligibility gating stays at the leaf (read-only, inspects only the two input scans) with one + * exception: whether the rebuilt merged scan degrades a partitioning/ordering an input reported + * can only be decided once that scan exists, so on the deferred path that one check runs at the + * Filter, next to the build (see `tryBuildFilterDSv2ScanChild`). + */ + case class MergeContext(filterPropagationSupported: Boolean, filterAboveScan: Boolean = false) + + /** + * Recursively attempts to merge two plans by traversing their tree structures. + * + * Two plans can be merged if: + * - They are identical (canonicalized forms match), OR + * - They have compatible root nodes with mergeable children + * + * Supported merge patterns: + * - Project nodes: Combines project lists from both plans + * - Aggregate nodes: Combines aggregate expressions if grouping is identical and both + * support the same aggregate implementation (hash/object-hash/sort-based) + * - Filter nodes: Only if filter conditions are identical + * - Join nodes: Requires identical join type, hints, and conditions; filter propagation is + * forwarded into the join's children so a filter difference on one child can still be merged + * + * @param newPlan The plan to merge into the cached plan. + * @param cachedPlan The cached plan to merge with. + * @return Some([[TryMergeResult]]) if merge succeeds, None if plans cannot be merged. + */ + private def tryMergePlans( + newPlan: LogicalPlan, + cachedPlan: LogicalPlan, + context: MergeContext): Option[TryMergeResult] = { + // The plain "reuse the cached plan as-is" result, shared by every branch below. Lazy because + // the DSv2 merge path under a Filter does not need it when the merge itself succeeds. + lazy val identical = checkIdenticalPlans(newPlan, cachedPlan).map(TryMergeResult(cachedPlan, _)) + // A DSv2 scan pair is handled here in one place, rather than split between this leading check + // and the structural match below. Under a Filter, merging DEFERS the scan build so the + // enclosing Filter's row-group pruning is pushed in a single rebuild -- so try the merge first + // and fall back to plain reuse if the scans cannot merge. Trying the merge first even when the + // two scans are identical is deliberate: the deferred rebuild re-pushes the enclosing Filter's + // condition as a best-effort filter, which plain reuse would leave as a post-scan Filter, so + // reusing identical scans here would forfeit that pruning. Without a Filter there is no pruning + // to recover, so reuse an identical scan as-is (no rebuild) and only merge scans that differ + // (projected columns / strict filters). Any other plan pair uses the general reuse. + val earlyResult = (newPlan, cachedPlan) match { + case (np: DataSourceV2ScanRelation, cp: DataSourceV2ScanRelation) => + if (context.filterAboveScan) { + tryMergeScanRelations(np, cp, context).orElse(identical) + } else { + identical.orElse(tryMergeScanRelations(np, cp, context)) + } + case _ => identical + } + earlyResult.orElse( + (newPlan, cachedPlan) match { + case (np: Project, cp: Project) => + tryMergePlans(np.child, cp.child, context).map { + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, deferred, dsv2Merged) => + val (mergedProjectList, newNPMapping) = + mergeNamedExpressions(np.projectList, cp.projectList, npMapping, npFilter, cpFilter) + TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, + cpFilter, deferred, dsv2Merged) + } + case (np, cp: Project) => + tryMergePlans(np, cp.child, context).map { + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, deferred, dsv2Merged) => + val (mergedProjectList, newNPMapping) = + mergeNamedExpressions(np.output, cp.projectList, npMapping, npFilter, cpFilter) + TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, + cpFilter, deferred, dsv2Merged) + } + case (np: Project, cp) => + tryMergePlans(np.child, cp, context).map { + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, deferred, dsv2Merged) => + val (mergedProjectList, newNPMapping) = + mergeNamedExpressions(np.projectList, cp.output, npMapping, npFilter, cpFilter) + TryMergeResult(Project(mergedProjectList, mergedChild), newNPMapping, npFilter, + cpFilter, deferred, dsv2Merged) + } + + case (np: Aggregate, cp: Aggregate) if supportedAggregateMerge(np, cp) => + // Filter propagation into the aggregate is only safe when there is no grouping. + val childFilterPropagationSupported = filterPropagationEnabled && + np.groupingExpressions.isEmpty && cp.groupingExpressions.isEmpty + tryMergePlans(np.child, cp.child, + MergeContext(childFilterPropagationSupported, filterAboveScan = false)).flatMap { + case TryMergeResult(mergedChild, npMapping, None, None, _, dsv2Merged) => + val mappedNPGroupingExpression = + np.groupingExpressions.map(mapAttributes(_, npMapping)) + // Order of grouping expression does matter as merging different grouping orders can + // introduce "extra" shuffles/sorts that might not present in all of the original + // subqueries. + if (mappedNPGroupingExpression.map(_.canonicalized) == + cp.groupingExpressions.map(_.canonicalized)) { + val (mergedAggregateExpressions, newNPMapping) = + mergeNamedExpressions(np.aggregateExpressions, cp.aggregateExpressions, npMapping) + val mergedPlan = + Aggregate(cp.groupingExpressions, mergedAggregateExpressions, mergedChild) + Some(TryMergeResult(mergedPlan, newNPMapping, dsv2Merged = dsv2Merged)) + } else { + None + } + case TryMergeResult(mergedChild, npMapping, npFilterOpt, cpFilterOpt, _, dsv2Merged) => + // childFilterPropagationSupported guarantees both aggregates have no grouping, so + // the grouping-match check is skipped. + assert(childFilterPropagationSupported) + + // Apply each propagated boolean attribute as a FILTER (WHERE ...) clause on the + // corresponding side's aggregate expressions. + // A None filter means the side's aggregate expressions already carry their individual + // FILTER attributes from a previous merge round and should be left unchanged. + // Filter propagation is consumed here and not passed further up. + val filteredNPAggregateExpressions = npFilterOpt.fold(np.aggregateExpressions) { + case (f, _) => applyFilterToAggregateExpressions(np.aggregateExpressions, f) + } + val filteredCPAggregateExpressions = cpFilterOpt.fold(cp.aggregateExpressions)( + applyFilterToAggregateExpressions(cp.aggregateExpressions, _)) + val (mergedAggregateExpressions, newNPMapping) = + mergeNamedExpressions(filteredNPAggregateExpressions, + filteredCPAggregateExpressions, npMapping) + val mergedPlan = Aggregate(Seq.empty, mergedAggregateExpressions, mergedChild) + Some(TryMergeResult(mergedPlan, newNPMapping, dsv2Merged = dsv2Merged)) + } + + case (np: Filter, cp: Filter) => + tryMergePlans(np.child, cp.child, context.copy(filterAboveScan = true)).flatMap { + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, deferred, dsv2Merged) => + val mappedNPCondition = mapAttributes(np.condition, npMapping) + // Comparing the canonicalized form is required to ignore different forms of the same + // expression. + if (mappedNPCondition.canonicalized == cp.condition.canonicalized) { + // Identical conditions: the filter node adds no new discrimination between the two + // sides, so keep it unchanged. If it sits above a deferred merged DSv2 scan, build + // that scan once here with this condition as a best-effort filter. If it cannot be + // built to spec -- the strict filters do not come back re-enforced, or the rebuilt + // scan's re-derived partitioning/ordering would degrade what the inputs reported -- + // decline the whole merge (the leaf's strict-only build would have failed + // identically). + tryBuildFilterDSv2ScanChild(mergedChild, deferred, Some(cp.condition)) + .map { prunedChild => + val mergedPlan = Filter(cp.condition, prunedChild) + TryMergeResult( + mergedPlan, npMapping, npFilter, cpFilter, dsv2Merged = dsv2Merged) + } + // Symmetric propagation broadens the merged scan to OR(f1, f2); it is off by default + // because the two sides may read disjoint data. A DSv2 scan merge is exempt under its + // own config: `dsv2Merged` marks an (equal-strict) DSv2 merge below -- whose equal + // strict filters mean both sides read the same base set, so the OR only weakens the + // best-effort filter. Unlike `deferred`, `dsv2Merged` survives past the innermost + // Filter that built the scan, so a stacked outer Filter pair recognizes it too. + } else if (context.filterPropagationSupported && + (symmetricFilterPropagationEnabled || + (dsv2SymmetricFilterPropagationEnabled && dsv2Merged))) { + if (cp.getTagValue(PlanMerger.MERGED_FILTER_TAG).isDefined) { + // cp Filter is already a merged filter from a previous round: its condition + // is OR(f0, f1, ...) and its child Project already contains aliases for those + // attributes. Only create a new alias for the np side, and extend the OR + // condition. A tagged filter is always built with a Project child (see the + // first-time branch below), so a non-Project child should not happen; decline the + // merge rather than fail, keeping the merge best-effort. + mergedChild match { + case childProject: Project => + val newNPCondition = npFilter.fold(mappedNPCondition) { + case (f, _) => And(f, mappedNPCondition) + } + // If newNPCondition is already aliased in the child Project (e.g. a third + // subplan whose filter matches one from a previous merge round), reuse the + // existing attribute instead of creating a redundant alias. + val existingNPFilter = childProject.projectList.collectFirst { + case a: Alias if a.child.canonicalized == newNPCondition.canonicalized => + a.toAttribute + } + val (newProjectList, newCondition, newNPFilterOut) = + existingNPFilter match { + case Some(reusedFilter) => + // np matches an existing side: no new alias, OR condition unchanged. + (childProject.projectList, cp.condition, (reusedFilter, false)) + case None => + val newNPFilterAlias = + Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() + (childProject.projectList :+ newNPFilterAlias, + Or(cp.condition, newNPFilterAlias.toAttribute): Expression, + (newNPFilterAlias.toAttribute, true)) + } + // Phase 2: the leaf re-merge rebuilt the scan with strict filters only, + // dropping the OR best-effort filter established in earlier rounds, so + // re-establish it here from ALL propagated conditions, not just the new + // side's. Only the aliases the OR condition references are filter sides; + // other aliases in the Project are computed columns, not filters. + val conditions = newProjectList.collect { + case a: Alias if newCondition.references.contains(a.toAttribute) => a.child + } + tryBuildFilterDSv2ScanChild( + childProject.child, deferred, conditions.reduceOption(Or)) + .map { prunedChild => + val newProject = childProject.copy( + projectList = newProjectList, child = prunedChild) + val newFilter = Filter(newCondition, newProject) + newFilter.copyTagsFrom(cp) + TryMergeResult(newFilter, npMapping, Some(newNPFilterOut), None, + dsv2Merged = dsv2Merged) + } + case _ => + None + } + } else { + // First-time filter propagation: alias both sides' conditions as boolean + // attributes in a new Project below the Filter, and set the Filter condition + // to OR(newNPFilter, newCPFilter). + // Note: the new Project always uses mergedChild as its child (rather than + // flattening into an existing Project below) because mergedChild.output may + // contain previously-propagated filter attributes that cp.condition references. + val newNPCondition = + npFilter.fold(mappedNPCondition) { case (f, _) => And(f, mappedNPCondition) } + val newCPCondition = cpFilter.fold(cp.condition)(And(_, cp.condition)) + // The OR-widen moves both conditions into a boolean Project above the merged + // scan, so the scan itself would read the full table. Build the deferred scan + // here with OR(np condition, cp condition) as the best-effort filter (the Filter + // above still enforces exactness). The best-effort filter is derived from the + // scan-level conditions, not the propagated filter attributes in newNP/newCP. + tryBuildFilterDSv2ScanChild( + mergedChild, deferred, Some(Or(mappedNPCondition, cp.condition))) + .map { prunedChild => + val newNPFilterAlias = + Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() + val newCPFilterAlias = + Alias(newCPCondition, s"propagatedFilter_${PlanMerger.newId}")() + val newNPFilter = newNPFilterAlias.toAttribute + val newCPFilter = newCPFilterAlias.toAttribute + val project = Project( + prunedChild.output.toList ++ Seq(newNPFilterAlias, newCPFilterAlias), + prunedChild) + val newFilter = Filter(Or(newNPFilter, newCPFilter), project) + newFilter.copyTagsFrom(cp) + newFilter.setTagValue(PlanMerger.MERGED_FILTER_TAG, ()) + TryMergeResult(newFilter, npMapping, Some((newNPFilter, true)), + Some(newCPFilter), dsv2Merged = dsv2Merged) + } + } + } else { + None + } + } + case (np: Filter, cp) if context.filterPropagationSupported => + tryMergePlans(np.child, cp, context.copy(filterAboveScan = false)).collect { + // If the cp side already propagated a filter from deeper recursion, the merge is + // effectively symmetric (both sides have a filter condition). Abort unless + // symmetricFilterPropagationEnabled. + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, _, dsv2Merged) + if cpFilter.isEmpty || symmetricFilterPropagationEnabled => + val mappedNPCondition = mapAttributes(np.condition, npMapping) + val newNPCondition = npFilter.fold(mappedNPCondition) { + case (f, _) => And(f, mappedNPCondition) + } + val newNPFilterAlias = + Alias(newNPCondition, s"propagatedFilter_${PlanMerger.newId}")() + val newNPFilter = newNPFilterAlias.toAttribute + val project = Project( + mergedChild.output.toList :+ newNPFilterAlias, + mergedChild) + TryMergeResult(project, npMapping, Some((newNPFilter, true)), cpFilter, + dsv2Merged = dsv2Merged) + } + case (np, cp: Filter) if context.filterPropagationSupported => + tryMergePlans(np, cp.child, context.copy(filterAboveScan = false)).collect { + // If the np side already propagated a filter from deeper recursion, the merge is + // effectively symmetric (both sides have a filter condition). Abort unless + // symmetricFilterPropagationEnabled. + case TryMergeResult(mergedChild, npMapping, npFilter, cpFilter, _, dsv2Merged) + if npFilter.isEmpty || symmetricFilterPropagationEnabled => + if (cp.getTagValue(PlanMerger.MERGED_FILTER_TAG).isDefined) { + // cp is a previously-merged Filter: its condition is `OR(pf_0, pf_1, ...)` and cp's + // aggregate expressions already carry individual `FILTER (WHERE pf_i)` clauses that + // restrict each aggregation to its originating side. Synthesising a new cpFilter + // alias for cp.condition would just produce `FILTER AND(OR(pf_0, pf_1, ...), pf_i)` + // upstream, which simplifies to `FILTER pf_i` -- wasted work and plan bloat. + // Drop cp's Filter and let the recursion's result flow up with cpFilter = None so + // cp's aggregates are left untouched. + TryMergeResult(mergedChild, npMapping, npFilter, None, dsv2Merged = dsv2Merged) + } else { + val newCPCondition = cpFilter.fold(cp.condition)(And(_, cp.condition)) + val newCPFilterAlias = + Alias(newCPCondition, s"propagatedFilter_${PlanMerger.newId}")() + val newCPFilter = newCPFilterAlias.toAttribute + val project = Project( + mergedChild.output.toList :+ newCPFilterAlias, + mergedChild) + TryMergeResult(project, npMapping, npFilter, Some(newCPFilter), + dsv2Merged = dsv2Merged) + } + } + + case (np: Join, cp: Join) if np.joinType == cp.joinType && np.hint == cp.hint => + tryMergePlans(np.left, cp.left, context.copy(filterAboveScan = false)).flatMap { + case TryMergeResult(mergedLeft, leftNPMapping, leftNPFilter, leftCPFilter, _, + leftDsv2Merged) => + tryMergePlans(np.right, cp.right, context.copy(filterAboveScan = false)).flatMap { + case TryMergeResult(mergedRight, rightNPMapping, rightNPFilter, rightCPFilter, _, + rightDsv2Merged) + // If both children independently propagate filter attributes we would need to + // AND them into a new alias above the join, which is not yet supported. + if !(leftNPFilter.isDefined && rightNPFilter.isDefined) && + !(leftCPFilter.isDefined && rightCPFilter.isDefined) && + // Gate join-crossing filter propagation behind its own config flag. + // When no filter attributes are in play the merge is unconditionally safe. + (leftNPFilter.isEmpty && leftCPFilter.isEmpty && + rightNPFilter.isEmpty && rightCPFilter.isEmpty || + filterPropagationThroughJoinEnabled) && + // A filter attribute is only safe to propagate through a join if it comes + // from the "preserved" (non-nullable) side. On the nullable side, unmatched + // rows are NULL-padded so f=NULL, causing FILTER (WHERE f) to incorrectly + // exclude rows that should contribute to the aggregate. Right-side + // attributes are also absent from semi/anti join output. + (leftNPFilter.isEmpty && leftCPFilter.isEmpty || + filterSafeForJoin(fromLeft = true, cp.joinType)) && + (rightNPFilter.isEmpty && rightCPFilter.isEmpty || + filterSafeForJoin(fromLeft = false, cp.joinType)) => + val npMapping = leftNPMapping ++ rightNPMapping + val mappedNPCondition = np.condition.map(mapAttributes(_, npMapping)) + // Comparing the canonicalized form is required to ignore different forms of the + // same expression and `AttributeReference.qualifier`s in `cp.condition`. + if (mappedNPCondition.map(_.canonicalized) == cp.condition.map(_.canonicalized)) { + val npFilter = leftNPFilter.orElse(rightNPFilter) + val cpFilter = leftCPFilter.orElse(rightCPFilter) + Some(TryMergeResult(cp.withNewChildren(Seq(mergedLeft, mergedRight)), npMapping, + npFilter, cpFilter, dsv2Merged = leftDsv2Merged || rightDsv2Merged)) + } else { + None + } + case _ => None + } + case _ => None + } + + // Otherwise merging is not possible. + case _ => None + }) + } + + /** + * The DSv2 scan merge: fuse two scans of the same table that differ only in projected columns + * (and carry the same strict pushed filters) into a single scan reading the union of their + * columns. The connector opts in via + * `TableCapability.SCAN_MERGING`; Spark runs the real DSv2 pushdown + * ([[V2ScanRelationPushDown]]) on a synthetic `Filter` over the relation, extracts the merged + * scan, and verifies the (equal) strict filters remain fully enforced. The `mergeable` gate is + * read-only, inspecting only the two input scans; two further checks can decline the merge -- the + * inputs' reported partitioning/ordering must combine into a single report (checked here, before + * any rebuild), and the rebuilt scan must not degrade that report (checked once it is built). + * + * When this scan sits under a [[Filter]] (`context.filterAboveScan`), the build is DEFERRED: the + * merged scan is not built here but carried up as a [[DSv2DeferredScan]] and built exactly once + * at the enclosing Filter (via `tryBuildFilterDSv2ScanChild`), where the strict filters and the + * Filter's best-effort row-group pruning are known and can be pushed together in a single + * rebuild. The placeholder plan returned in that case is the bare relation (its output is a + * superset of the union columns); the Filter arm splices the built scan in by reference identity. + * Otherwise (no enclosing Filter) the scan is built here, strict-only. + * + * A differing post-scan filter is handled by filter propagation at the Filter, not here. There is + * no fallback, so any anomaly (a strict filter the rebuilt scan does not fully enforce, an + * unexpected output schema) results in `None` (no merge): it must be correct on its own. + */ + private def tryMergeScanRelations( + np: DataSourceV2ScanRelation, + cp: DataSourceV2ScanRelation, + context: MergeContext): Option[TryMergeResult] = { + // np's relation attributes paired with cp's by position, reused below. Safe because the + // `mergeable` gate requires the two relations to be canonically equal, so both list the table's + // columns in the same order; lazy so it is only built once that check has passed. A scan's + // (pruned) `output` is a subset of its `relation.output` -- even where nested field pruning + // narrows a column's type -- so this maps each scan's output, and its pushed filters (which + // reference the relation's full output), through to the other scan's relation. + lazy val npRelationMapping = AttributeMap[Attribute](np.relation.output.zip(cp.relation.output)) + + // Each side must read a subset of its relation's columns AT THE SAME type. Column pruning may + // drop columns, but a struct/array/map column read at a narrower type via nested schema pruning + // leaves its extractors (e.g. GetStructField ordinals) resolved against the narrow layout; the + // merge rebuilds the column at the relation's full type, so those ordinals would read the wrong + // field. This subset property also guarantees every output attribute maps through + // npRelationMapping (and, in tryBuildMergedDSv2Scan, back to the relation). Widening the + // merged scan to the union of nested fields and remapping the ordinals is a possible follow-up. + def readsSubsetOfRelation(scan: DataSourceV2ScanRelation): Boolean = { + val relTypes = AttributeMap(scan.relation.output.map(a => a -> a.dataType)) + scan.output.forall(a => relTypes.get(a).contains(a.dataType)) + } + + val mergeable = + // Same table, options, catalog and identifier: the relation's canonical form covers all of + // these (options compares by content via `CaseInsensitiveStringMap.equals`). + np.relation.canonicalized == cp.relation.canonicalized && + // Both scans are mergeable: each came out of the plain column-pruning + filter pushdown + // path carrying only pushdowns a rebuilt scan can reproduce. A scan with a non-reproducible + // pushdown (aggregate, join, variant, limit, offset, top-N, sample) or built by any other + // rule is not mergeable by default. + np.mergeableScan && cp.mergeableScan && + // Reported partitioning/ordering is not reconstructed by the rebuilt scan + // (V2ScanPartitioningAndOrdering is a separate early rule that rebuildScan does not run). + // Rather than decline here, the merged scan re-derives its own when built (see + // tryBuildMergedDSv2Scan), and mergeDegradesReporting declines the merge if that would + // degrade what an input reported (unless a dsv2ScanMerge config allows it). + // The table opts in to Spark-side merging (a table capability, so a V1-fallback source + // whose scan Spark wraps can still opt in). Both relations are the same table (canonically + // equal, checked above), but check each to be safe. + np.relation.table.capabilities().contains(TableCapability.SCAN_MERGING) && + cp.relation.table.capabilities().contains(TableCapability.SCAN_MERGING) && + // Each side reads a subset of its relation's columns at the same type (see above). + readsSubsetOfRelation(np) && readsSubsetOfRelation(cp) && + // Both pushed the same strict filters, so re-pushing reproduces both sides' row sets. Remap + // np's pushed filters onto cp's via npRelationMapping (the full relation-to-relation + // mapping, since a pushed filter may reference a column pruned out of np.output) before + // comparing as sets. + ExpressionSet(np.pushedFilters.map(mapAttributes(_, npRelationMapping))) == + ExpressionSet(cp.pushedFilters) + + // The read-only gate above is settled. The report combine below can still decline (and so can + // the degradation check once the scan is built); everything else below constructs the merge. + if (!mergeable) { + return None + } + + // np's columns expressed in cp's relation space; the subset check above guarantees each maps. + val npMapping = AttributeMap(np.output.map(a => a -> npRelationMapping(a))) + // cp's columns are already cp.relation attributes; append the np-only columns, in np.output + // order (npMapping.values would be exprId-hash-ordered). + val unionAttrs = cp.output ++ np.output.map(npMapping).filterNot(cp.outputSet.contains) + + // The reported key-grouped partitioning / ordering the merged scan must preserve so BOTH inputs + // stay not-worse. Each input reports its own, remapped into cp's relation space (cp's already + // is; np's via npRelationMapping). The two normally agree -- the clustering expressions come + // from the table, and for partitioning, pruning only ever drops a report wholesale (an empty + // side, which never constrains) -- but a source is free to report per scan, e.g. an ordering + // that holds only for the file set the filters pushed into that scan left behind. Combine them + // into the single report the merge must keep (kGP: equal; ordering: the stronger, which + // satisfies both). None from combine* means the inputs are INCOMPATIBLE -- no rebuilt scan + // could keep both not-worse -- so decline HERE, before rebuilding, unless the matching config + // accepts degrading that dimension. + val combinedKeyGroupedPartitioning = combineRequiredKeyGroupedPartitioning( + np.keyGroupedPartitioning.map(_.map(mapAttributes(_, npRelationMapping))).getOrElse(Nil), + cp.keyGroupedPartitioning.getOrElse(Nil)) + val combinedOrdering = combineRequiredOrdering( + np.ordering.map(_.map(mapAttributes(_, npRelationMapping))).getOrElse(Nil), + cp.ordering.getOrElse(Nil)) + if ((combinedKeyGroupedPartitioning.isEmpty && !dsv2AllowKeyGroupedPartitioningDegradation) || + (combinedOrdering.isEmpty && !dsv2AllowOrderingDegradation)) { + return None + } + // Empty = no requirement (both inputs reported none, or they were incompatible but the config + // accepts the degradation). Otherwise the single report the merged scan must reproduce/satisfy. + val expectedKeyGroupedPartitioning = combinedKeyGroupedPartitioning.getOrElse(Nil) + val expectedOrdering = combinedOrdering.getOrElse(Nil) + + if (context.filterAboveScan) { + // Defer the build to the enclosing Filter so the scan is built once with strict + + // best-effort filters. The placeholder mergedPlan is the bare relation (its output is a + // superset of unionAttrs). rebuildScan reuses the relation's attributes, so mapping + // np.output to cp's relation attributes is consistent with the eventual built scan. + Some(TryMergeResult(cp.relation, npMapping, + dsv2DeferredScan = Some(DSv2DeferredScan(unionAttrs, cp.pushedFilters, + expectedKeyGroupedPartitioning, expectedOrdering)), dsv2Merged = true)) + } else { + // No enclosing Filter: build the merged scan here enforcing the (equal) strict filters over + // the union of columns, with no best-effort filter (no post-scan Filter to prune on). + tryBuildMergedDSv2Scan(cp.relation, unionAttrs, cp.pushedFilters, bestEffortFilter = None) + .filterNot(mergeDegradesReporting(_, expectedKeyGroupedPartitioning, expectedOrdering)) + .map(TryMergeResult(_, npMapping, dsv2Merged = true)) + } + } + + /** + * Rebuilds the merged DSv2 scan via [[V2ScanRelationPushDown.rebuildScan]], projecting + * `unionAttrs` and filtering by `strictFilters` (plus the `bestEffortFilter`). This reuses + * the production pushdown end to end -- the same filter translation, column pruning, + * determinism/subquery handling and iterative PartitionPredicate second pass -- rather than + * reimplementing a slice of it here. + * + * `strictFilters` must come back fully enforced (present in the rebuilt scan's `pushedFilters`); + * otherwise `None`, because nothing above the leaf re-checks it. The `bestEffortFilter` is + * offered to the source only when sound: it is dropped unless it is deterministic (a + * non-deterministic predicate the source prunes on would drop rows the enclosing Filter cannot + * recover) and references only the relation's own columns (propagated boolean filter attributes + * are not columns of the relation). + * + * The rebuilt scan's reported partitioning/ordering is re-derived here, but NOT checked against + * what the inputs reported -- callers do that via [[mergeDegradesReporting]], so a degradation + * stays distinguishable from a filter-enforcement failure. + */ + private def tryBuildMergedDSv2Scan( + relation: DataSourceV2Relation, + unionAttrs: Seq[Attribute], + strictFilters: Seq[Expression], + bestEffortFilter: Option[Expression]): Option[DataSourceV2ScanRelation] = { + val relationOut = relation.outputSet + // Defensive: strict filters come from `pushedFilters`, which reference only relation columns, + // so this holds today. If a future caller offers a filter over non-relation attributes, decline + // the merge rather than build an unsound scan (there is no fallback above the leaf). + if (!strictFilters.forall(_.references.subsetOf(relationOut))) { + return None + } + // `unionAttrs` are the relation's own attributes (the caller builds the union in the relation's + // space), so they are the projection directly. + // strictFilters are enforced; the bestEffortFilter is offered too, but only when it is + // expressible over the relation -- a condition referencing propagated boolean filter aliases + // rather than relation columns is dropped. A non-deterministic predicate needs no handling + // here: SPARK-58207 keeps non-deterministic filters from being pushed to a V2 source, so an + // offered one is simply not pushed (and dropped when the scan is extracted). + val conds = strictFilters ++ bestEffortFilter.filter(_.references.subsetOf(relationOut)) + V2ScanRelationPushDown.rebuildScan(relation, unionAttrs, conds).filter { scan => + // The rebuilt scan must itself be mergeable. rebuildScan re-runs the full pushdown, so any + // non-reproducible pushdown it introduces would make the merged scan unsound. Today's + // Project-over-Filter input only triggers the plain path (always mergeable), so this is + // defensive: it re-validates the rebuild's output against the same gate applied to its + // inputs, rather than trusting the rebuild if its input plan ever broadens. + scan.mergeableScan && + // Every intended-strict filter must be fully enforced by the rebuilt scan (nothing above + // re-checks it), and the scan must produce exactly the requested union of columns. + strictFilters.forall(ExpressionSet(scan.pushedFilters).contains) && + scan.outputSet == AttributeSet(unionAttrs) + }.map { scan => + // rebuildScan returns the merged scan with reported partitioning/ordering unset + // (V2ScanPartitioningAndOrdering is a separate early rule the rebuild does not run), so + // re-derive them on this single node. Safe on one node: the partitioning pass is idempotent + // and the ordering pass is applied once to a fresh node. + V2ScanPartitioningAndOrdering(scan).asInstanceOf[DataSourceV2ScanRelation] + } + } + + /** + * Threads the deferred DSv2 scan build through a [[Filter]] arm. If the merged child carries a + * [[DSv2DeferredScan]], build the scan once here -- at the enclosing Filter, with the strict + * filters plus the Filter's `bestEffortFilter` -- and splice it in place of the placeholder + * relation. Tries strict + best-effort first, then strict-only (the best-effort filter is + * droppable); a build returns `None` only if the strict filters cannot be re-enforced at all (the + * leaf's strict-only build would have failed identically). Each built scan is also checked + * against the required report the leaf computed, and rejected if it would degrade that -- per + * attempt, so a source whose report depends on which filters were pushed can still satisfy it + * strict-only. Either way the caller must decline the merge. If there is no deferred scan (a + * non-DSv2 child, or a scan not under a Filter), there is nothing to build and the child is + * returned unchanged. + */ + private def tryBuildFilterDSv2ScanChild( + child: LogicalPlan, + dsv2DeferredScan: Option[DSv2DeferredScan], + bestEffortFilter: Option[Expression]): Option[LogicalPlan] = dsv2DeferredScan match { + case None => Some(child) + case Some(d) => + // The deferring leaf left the relation to rebuild from in the plan as the placeholder + // mergedPlan. It is the sole bare DataSourceV2Relation in the subtree (early pushdown has + // turned every other relation into a DataSourceV2ScanRelation), so recover it by type here + // rather than carrying it on DSv2DeferredScan. + child.collectFirst { case r: DataSourceV2Relation => r }.flatMap { relation => + // Check the report per attempt, and at the caller rather than inside the build: the strict + // filters and the report are independent reasons to reject a build, so a source whose + // report depends on what got pushed still gets its second chance from the strict-only + // attempt, and `tryBuildMergedDSv2Scan`'s `None` keeps its single meaning. The strict-only + // attempt is what the leaf builds when no Filter is above the scan, so checking only the + // first attempt would leave the deferred path weaker than the leaf path. + def build(offeredBestEffortFilter: Option[Expression]) = + tryBuildMergedDSv2Scan( + relation, d.unionAttrs, d.strictFilters, offeredBestEffortFilter) + .filterNot( + mergeDegradesReporting(_, d.requiredKeyGroupedPartitioning, d.requiredOrdering)) + + build(bestEffortFilter) + .orElse(build(None)) + .map { built => + child.transformUp { case r: DataSourceV2Relation if r eq relation => built } + } + } + } + + // The key-grouped partitioning the merged scan must reproduce to keep both inputs not-worse: the + // two must be equal, so a differing non-empty pair is INCOMPATIBLE (None); an empty side imposes + // no constraint. Compared in cp's relation space (np's report was remapped into it by the + // caller). A report carrying a `TransformExpression` compares equal across the two inputs only if + // the connector's `BoundFunction` implements `equals` -- Spark does not derive that identity + // itself, see `BoundFunction#equals` -- so a connector that does not gives up this merge. + private def combineRequiredKeyGroupedPartitioning( + a: Seq[Expression], b: Seq[Expression]): Option[Seq[Expression]] = { + if (a.isEmpty) Some(b) + else if (b.isEmpty) Some(a) + else if (a.map(_.canonicalized) == b.map(_.canonicalized)) Some(a) + else None + } + + // The ordering the merged scan must satisfy to keep both inputs not-worse: the stronger of the + // two (the one that satisfies the other -- satisfying it implies satisfying the weaker). If + // neither satisfies the other they are INCOMPATIBLE (None). An empty ordering never constrains. + private def combineRequiredOrdering( + a: Seq[SortOrder], b: Seq[SortOrder]): Option[Seq[SortOrder]] = { + if (SortOrder.orderingSatisfies(a, b)) Some(a) + else if (SortOrder.orderingSatisfies(b, a)) Some(b) + else None + } + + // True when the rebuilt merged scan does not reproduce the required key-grouped partitioning, or + // does not satisfy the required ordering (the combined report the merge must preserve, computed + // at the leaf) -- that can force a shuffle/sort the original plan avoided. Gated per dimension by + // the dsv2ScanMerge degradation configs; an empty required report imposes no constraint. Compared + // in cp's relation space. + // + // Only the reported EXPRESSIONS are compared, which is all a DataSourceV2ScanRelation carries; + // the merged scan's split count and partition values can still differ from an input's, since it + // may push a different best-effort filter and so prune differently. And a report the merged scan + // GAINS is not a degradation either. For partitioning that is because an input dropped its own + // only where a pruned column left the expressions inexpressible over that scan's output + // (V2ScanPartitioningAndOrdering's partitioning pass is reference-subset guarded), not because + // the source stopped reporting; the ordering pass has no such guard, so an ordering report is + // never dropped by pruning and a gained one can only come from the source. Either way, keeping it + // is exactly the win this merge is after. + private def mergeDegradesReporting( + merged: DataSourceV2ScanRelation, + requiredKeyGroupedPartitioning: Seq[Expression], + requiredOrdering: Seq[SortOrder]): Boolean = { + val kgpDegraded = !dsv2AllowKeyGroupedPartitioningDegradation && + requiredKeyGroupedPartitioning.nonEmpty && + !merged.keyGroupedPartitioning.exists( + _.map(_.canonicalized) == requiredKeyGroupedPartitioning.map(_.canonicalized)) + val orderingDegraded = !dsv2AllowOrderingDegradation && + requiredOrdering.nonEmpty && + !SortOrder.orderingSatisfies(merged.ordering.getOrElse(Nil), requiredOrdering) + kgpDegraded || orderingDegraded + } + + // Returns true when a filter attribute originating from `fromLeft` child of a join with + // `joinType` can be safely propagated through that join to a parent Aggregate. + // + // Two conditions must both hold: + // 1. The attribute is in the join's output (rules out the right side of LeftSemi/LeftAnti). + // 2. The filter must originate from the non-nullable ("preserved") side of the join. + // When a filter is on the nullable side, the merged base plan no longer applies it to the + // nullable child's scan, so rows that were previously absent from that child reappear as + // matched join rows instead of unmatched NULL-padded rows. This changes aggregate + // expressions that use the NULL-padded column: e.g. for `sum(coalesce(col, default))`, an + // originally unmatched row would have contributed `default` via `coalesce(NULL, default)`, + // but in the merged plan the row is now matched with its real column value, fails the + // filter, and FILTER (WHERE false) discards it -- losing the `default` contribution + // entirely. + private def filterSafeForJoin(fromLeft: Boolean, joinType: JoinType): Boolean = + if (fromLeft) { + // Left side is never NULL-padded in: Inner, LeftOuter, LeftSemi, LeftAnti, Cross. + joinType match { + case Inner | LeftOuter | LeftSemi | LeftAnti | Cross => true + case _ => false // RightOuter and FullOuter can NULL-pad the left side + } + } else { + // Right side is never NULL-padded AND is in the join output in: Inner, RightOuter, Cross. + joinType match { + case Inner | RightOuter | Cross => true + case _ => false // LeftOuter/FullOuter can NULL-pad right; LeftSemi/LeftAnti drop right + } + } + + private def mapAttributes[T <: Expression](expr: T, outputMap: AttributeMap[_ <: Attribute]) = { + expr.transform { + case a: Attribute => outputMap.getOrElse(a, a) + }.asInstanceOf[T] + } + + // Remaps attributes of `newPlanExpressions` through `newPlanMapping`, then merges them with + // `cachedPlanExpressions` into a single expression list. + // Returns a pair of: + // 1. The merged expression list + // 2. New plan output map: ne.toAttribute -> merged plan attr (for parent nodes to remap + // new-plan-side expressions) + // + // When `newPlanFilter`/`cachedPlanFilter` are provided (filter propagation active), non-matching + // expressions from each side are wrapped with `If(filterAttr, expr, null)`. This ensures that a + // non-matching expression from one side evaluates to null for rows that belong to the other side, + // which is safe for aggregate FILTER (WHERE ...) semantics and avoids computing values for + // irrelevant rows. The filter attributes themselves are appended to the merged expression list so + // they remain visible to the enclosing Aggregate that will consume them. A newPlanFilter with + // isNew=false was reused from a previous merge round and is already present in the merged child + // output, so it is not appended again. + private def mergeNamedExpressions( + newPlanExpressions: Seq[NamedExpression], + cachedPlanExpressions: Seq[NamedExpression], + newPlanMapping: AttributeMap[Attribute], + newPlanFilter: Option[(Attribute, Boolean)] = None, + cachedPlanFilter: Option[Attribute] = None) = { + val mergedExpressions = mutable.ArrayBuffer[NamedExpression](cachedPlanExpressions: _*) + val matchedCachedIndices = mutable.HashSet.empty[Int] + val newNPMapping = AttributeMap(newPlanExpressions.map { ne => + val mapped = mapAttributes(ne, newPlanMapping) + val withoutAlias = mapped match { + case Alias(child, _) => child + case e => e + } + val foundIdx = mergedExpressions.indexWhere { + case Alias(child, _) => child semanticEquals withoutAlias + case e => e semanticEquals withoutAlias + } + val resultAttr = if (foundIdx >= 0) { + // Matching expression: both sides compute the same value, no wrapping needed. + matchedCachedIndices += foundIdx + mergedExpressions(foundIdx).toAttribute + } else { + // Non-matching expression from the new plan side: wrap with the new plan filter so it + // is only computed for rows that belong to the new plan side. Plain attribute references + // are not wrapped since reading a column value is free. + val wrappedExpr: NamedExpression = newPlanFilter match { + case Some((f, _)) if !withoutAlias.isInstanceOf[Attribute] => + Alias(If(f, withoutAlias, Literal(null, withoutAlias.dataType)), mapped.name)() + case _ => mapped + } + mergedExpressions += wrappedExpr + wrappedExpr.toAttribute + } + ne.toAttribute -> resultAttr + }) + + // Wrap unmatched cached expressions with the cached plan's filter so they are only computed for + // rows that belong to the cached plan side. Plain attribute references are not wrapped. + cachedPlanFilter.foreach { f => + for (i <- 0 until cachedPlanExpressions.size if !matchedCachedIndices.contains(i)) { + mergedExpressions(i) match { + case ce @ Alias(child, _) if !child.isInstanceOf[Attribute] => + // Preserve the original ExprId so parent references to this cached attribute stay valid + // without a cp-side remapping. (The new-plan wrapping above uses a fresh ExprId because + // those aliases are appended rather than replacing an existing entry.) + mergedExpressions(i) = + Alias(If(f, child, Literal(null, child.dataType)), ce.name)( + exprId = ce.toAttribute.exprId) + case _ => // attribute or alias-of-attribute, no wrapping needed + } + } + } + + newPlanFilter.foreach { + case (f, true) => mergedExpressions += f + case _ => + } + cachedPlanFilter.foreach(mergedExpressions += _) + + (mergedExpressions.toSeq, newNPMapping) + } + + // Applies filter as a FILTER (WHERE ...) clause to every AggregateExpression in exprs, + // combining with any pre-existing filter on the aggregate via AND. + private def applyFilterToAggregateExpressions( + exprs: Seq[NamedExpression], + filter: Attribute): Seq[NamedExpression] = { + exprs.map(_.transform { + case ae: AggregateExpression => + val combinedFilter = ae.filter.fold[Expression](filter)(And(filter, _)) + val newAE = ae.copy(filter = Some(combinedFilter)) + newAE.copyTagsFrom(ae) + newAE + }.asInstanceOf[NamedExpression]) + } + + // Only allow aggregates of the same implementation because merging different implementations + // could cause performance regression. + private def supportedAggregateMerge(newPlan: Aggregate, cachedPlan: Aggregate) = { + val aggregateExpressionsSeq = Seq(newPlan, cachedPlan).map { plan => + plan.aggregateExpressions.flatMap(_.collect { + case a: AggregateExpression => a + }) + } + val groupByExpressionSeq = Seq(newPlan, cachedPlan).map(_.groupingExpressions) + + val Seq(newPlanSupportsHashAggregate, cachedPlanSupportsHashAggregate) = + aggregateExpressionsSeq.zip(groupByExpressionSeq).map { + case (aggregateExpressions, groupByExpressions) => + Aggregate.supportsHashAggregate( + aggregateExpressions.flatMap( + _.aggregateFunction.aggBufferAttributes), groupByExpressions) + } + + newPlanSupportsHashAggregate && cachedPlanSupportsHashAggregate || + newPlanSupportsHashAggregate == cachedPlanSupportsHashAggregate && { + val Seq(newPlanSupportsObjectHashAggregate, cachedPlanSupportsObjectHashAggregate) = + aggregateExpressionsSeq.zip(groupByExpressionSeq).map { + case (aggregateExpressions, groupByExpressions) => + Aggregate.supportsObjectHashAggregate(aggregateExpressions, groupByExpressions) + } + newPlanSupportsObjectHashAggregate && cachedPlanSupportsObjectHashAggregate || + newPlanSupportsObjectHashAggregate == cachedPlanSupportsObjectHashAggregate + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala index 4a09898a58c76..0170d4354a6a9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala @@ -71,6 +71,12 @@ private[spark] class BatchIterator[T](iter: Iterator[T], batchSize: Int) * <li> SQL_SCALAR_ARROW_ITER_UDF for Scalar Iterator Arrow UDF * <li> SQL_SCALAR_PANDAS_UDF for Scalar Pandas UDF * <li> SQL_SCALAR_PANDAS_ITER_UDF for Scalar Iterator Pandas UDF + * <li> SQL_ARROW_ELEMENTWISE_UDF for a row-at-a-time UDF lifted out of a higher-order function's + * lambda (see ExtractPythonUDFFromLambda) + * <li> SQL_SCALAR_PANDAS_ELEMENTWISE_UDF for a Scalar Pandas UDF lifted out of such a lambda + * <li> SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF for a Scalar Iterator Pandas UDF lifted out of one + * <li> SQL_SCALAR_ARROW_ELEMENTWISE_UDF for a Scalar Arrow UDF lifted out of such a lambda + * <li> SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF for a Scalar Iterator Arrow UDF lifted out of one * </ul> * */ @@ -162,6 +168,11 @@ case class ArrowEvalPythonExec( private def supportedPythonEvalTypes: Array[Int] = Array( PythonEvalType.SQL_ARROW_BATCHED_UDF, + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF, PythonEvalType.SQL_SCALAR_ARROW_UDF, PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, PythonEvalType.SQL_SCALAR_PANDAS_UDF, @@ -205,7 +216,10 @@ class ArrowEvalPythonEvaluatorFactory( pythonRunnerConf, pythonMetrics, jobArtifactUUID, - sessionUUID) with BatchedPythonArrowInput + sessionUUID, + // Parallel to `funcs` (both come from `udfs` in order); tells the worker how many `array` + // levels each element-wise UDF flattens/re-nests (see `PythonUDF.elementwiseNestingDepth`). + udfs.map(_.elementwiseNestingDepth)) with BatchedPythonArrowInput val columnarBatchIter = pyRunner.compute(batchIter, context.partitionId(), context) columnarBatchIter.flatMap { batch => diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala index 75b8465e2607a..4d380a9cf4311 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala @@ -130,22 +130,20 @@ class ArrowPythonWithNamedArgumentRunner( pythonRunnerConf: Map[String, String], pythonMetrics: Map[String, SQLMetric], jobArtifactUUID: Option[String], - sessionUUID: Option[String]) + sessionUUID: Option[String], + // Per-UDF element-wise nesting depth, parallel to `funcs` (see + // `PythonUDF.elementwiseNestingDepth`). Empty (the default) means depth 1 for every UDF, so the + // many non-element-wise construction sites need not pass it. + elementwiseNestingDepths: Seq[Int] = Nil) extends RowInputArrowPythonRunner( funcs, evalType, argMetas.map(_.map(_.offset)), schema, timeZoneId, largeVarTypes, pythonMetrics, jobArtifactUUID, sessionUUID) { override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf - override protected def evalConf: Map[String, String] = { - if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) { - super.evalConf ++ Map( - "input_type" -> schema.json - ) - } else { - super.evalConf - } - } + override protected def evalConf: Map[String, String] = + ArrowPythonRunner.elementwiseEvalConf( + super.evalConf, evalType, schema, elementwiseNestingDepths) override protected def writeUDF(dataOut: DataOutputStream): Unit = { PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas) @@ -153,6 +151,29 @@ class ArrowPythonWithNamedArgumentRunner( } object ArrowPythonRunner { + /** + * Adds the `input_type` (and, for element-wise UDFs, the per-UDF `elementwise_nesting`) entries + * an Arrow runner sends to the worker via eval-conf. An element-wise UDF receives each argument + * as an `array<T>` column and flattens it, so like the Arrow batched UDF it needs the input + * schema to convert the incoming Arrow types; a lifted UDF from *nested* lambdas re-nests more + * than one level, so it also needs its per-UDF nesting depth. Shared by row and columnar runners. + */ + def elementwiseEvalConf( + base: Map[String, String], + evalType: Int, + schema: StructType, + elementwiseNestingDepths: Seq[Int]): Map[String, String] = { + if (PythonEvalType.isElementwiseUDF(evalType)) { + base ++ Map( + "input_type" -> schema.json, + "elementwise_nesting" -> elementwiseNestingDepths.mkString(",")) + } else if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) { + base ++ Map("input_type" -> schema.json) + } else { + base + } + } + /** Return Map with conf settings to be used in ArrowPythonRunner */ def getPythonRunnerConfMap(conf: SQLConf): Map[String, String] = { val confMap = collection.mutable.Map.empty[String, String] @@ -163,6 +184,7 @@ object ArrowPythonRunner { SQLConf.ARROW_EXECUTION_USE_LARGE_VAR_TYPES, SQLConf.PYTHON_TABLE_UDF_LEGACY_PANDAS_CONVERSION_ENABLED, SQLConf.PYTHON_UDF_LEGACY_PANDAS_CONVERSION_ENABLED, + SQLConf.PYTHON_UDF_MAP_IN_BATCH_LEGACY_ACCEPT_ANY_ITERABLE_ENABLED, SQLConf.PYTHON_UDF_PANDAS_INT_TO_DECIMAL_COERCION_ENABLED, SQLConf.PYTHON_UDF_PANDAS_PREFER_INT_EXTENSION_DTYPE, SQLConf.PYSPARK_BINARY_AS_BYTES, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala index ab9671c022f97..7a6c90478a806 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala @@ -26,7 +26,7 @@ import org.apache.spark.{JobArtifactSet, PartitionEvaluator, PartitionEvaluatorF import org.apache.spark.api.python.ChainedPythonFunctions import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, BoundReference, EmptyRow, Expression, JoinedRow, NamedArgumentExpression, NamedExpression, PythonFuncExpression, PythonUDAF, SortOrder, SpecificInternalRow, UnsafeProjection, UnsafeRow, WindowExpression} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, BoundReference, EmptyRow, Expression, JoinedRow, NamedArgumentExpression, NamedExpression, PythonFuncExpression, SortOrder, SpecificInternalRow, UnsafeProjection, UnsafeRow, WindowExpression} import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray @@ -156,8 +156,12 @@ class ArrowWindowPythonEvaluatorFactory( // Extract window expressions and window functions private val windowExpressions = expressions.flatMap(_.collect { case e: WindowExpression => e }) + // The window aggregate function is either a `PythonUDAF` (grouped-agg pandas/arrow UDF) or the + // incremental `PythonAggregate`; both are `PythonFuncExpression`s and share the per-frame Arrow + // window path (only the worker-side per-frame computation and eval type differ). private val udfExpressions = windowExpressions.map { e => - e.windowFunction.asInstanceOf[AggregateExpression].aggregateFunction.asInstanceOf[PythonUDAF] + e.windowFunction.asInstanceOf[AggregateExpression].aggregateFunction + .asInstanceOf[PythonFuncExpression] } // We shouldn't be chaining anything here. diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala index dd360bbf230b6..92956e4ecb5d2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala @@ -22,6 +22,7 @@ import org.apache.spark.api.python.PythonEvalType import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} import org.apache.spark.sql.execution.window._ @@ -33,6 +34,7 @@ import org.apache.spark.sql.execution.window._ * <ul> * <li> SQL_WINDOW_AGG_ARROW_UDF for Arrow UDF * <li> SQL_WINDOW_AGG_PANDAS_UDF for Pandas UDF + * <li> SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF for the incremental Arrow aggregator (`udaf`) * </ul> * * This is similar to [[WindowExec]]. The main difference is that this node does not compute @@ -127,7 +129,8 @@ case class ArrowWindowPythonExec( private def supportedPythonEvalTypes: Array[Int] = Array( PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, - PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF) + PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, + PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF) } object ArrowWindowPythonExec { @@ -139,8 +142,16 @@ object ArrowWindowPythonExec { val evalTypes = windowExpression.flatMap(w => WindowFunctionType.pythonEvalType(w)) assert(evalTypes.nonEmpty, "Cannot extract eval type from PythonUDAFs in ArrowWindowPythonExec") - assert(evalTypes.distinct.size == 1, - "All window functions must have the same eval type in ArrowWindowPythonExec") + // Distinct eval types here means incompatible Python window UDFs were grouped into one window + // (e.g. an incremental aggregator mixed with a grouped-agg pandas/arrow UDAF). They share the + // `WindowFunctionType.Python` type, so they pass the check in `PhysicalWindow`, but cannot run + // in a single operator; surface a clear analysis error rather than an internal AssertionError. + if (evalTypes.distinct.size != 1) { + val functionNames = windowExpression.flatMap(_.collect { + case e: PythonFuncExpression => e.name + }).distinct + throw QueryCompilationErrors.multiplePythonUDFTypesInWindowError(functionNames) + } ArrowWindowPythonExec(windowExpression, partitionSpec, orderSpec, child, evalTypes.head) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala index bad609919a4b1..3981602875adb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala @@ -202,7 +202,7 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( pyFuncs, evalType, argMetas, udfInputSchema, sessionLocalTimeZone, largeVarTypes, pythonRunnerConf, pythonMetrics, jobArtifactUUID, sessionUUID, - columnIndices) + columnIndices, udfs.map(_.elementwiseNestingDepth)) val resultIter = pyRunner.compute( bufferedIter, context.partitionId(), context) @@ -262,7 +262,7 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( pyFuncs, evalType, argMetas, udfInputSchema, sessionLocalTimeZone, largeVarTypes, pythonRunnerConf, pythonMetrics, jobArtifactUUID, sessionUUID, - inputColumnIndices.get) + inputColumnIndices.get, udfs.map(_.elementwiseNestingDepth)) pyRunner.compute( bufferedIter, context.partitionId(), context) } else { @@ -279,7 +279,8 @@ private[python] class ColumnarArrowEvalPythonEvaluatorFactory( val pyRunner = new ArrowPythonWithNamedArgumentRunner( pyFuncs, evalType, argMetas, udfInputSchema, sessionLocalTimeZone, largeVarTypes, pythonRunnerConf, - pythonMetrics, jobArtifactUUID, sessionUUID + pythonMetrics, jobArtifactUUID, sessionUUID, + udfs.map(_.elementwiseNestingDepth) ) with BasicPythonArrowInput pyRunner.compute( batchIter, context.partitionId(), context) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala index 41a563a7302f3..de19b43d32c1d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala @@ -42,7 +42,10 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner( override val pythonMetrics: Map[String, SQLMetric], jobArtifactUUID: Option[String], sessionUUID: Option[String], - override protected val inputColumnIndices: Array[Int]) + override protected val inputColumnIndices: Array[Int], + // Per-UDF element-wise nesting depth, parallel to `funcs` (see + // `PythonUDF.elementwiseNestingDepth`); empty means depth 1 for every UDF. + elementwiseNestingDepths: Seq[Int] = Nil) extends BaseArrowPythonRunner[ColumnarBatch, ColumnarBatch]( funcs, evalType, argMetas.map(_.map(_.offset)), schema, timeZoneId, largeVarTypes, pythonMetrics, jobArtifactUUID, sessionUUID) @@ -51,20 +54,14 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner( override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf - // The input schema travels in evalConf (key "input_type"), exactly like - // ArrowPythonWithNamedArgumentRunner. It must NOT be written into the UDF command section: + // The input schema (and per-UDF nesting depth) travel in evalConf, exactly like + // ArrowPythonWithNamedArgumentRunner. They must NOT be written into the UDF command section: // the worker's WorkerInitInfo.from_stream reads that section as a UDF count followed by UDF // entries, so an extra UTF string there desyncs the whole init-message parse -- the worker // then blocks waiting for bytes past the end of the init message and the task hangs forever. - override protected def evalConf: Map[String, String] = { - if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) { - super.evalConf ++ Map( - "input_type" -> schema.json - ) - } else { - super.evalConf - } - } + override protected def evalConf: Map[String, String] = + ArrowPythonRunner.elementwiseEvalConf( + super.evalConf, evalType, schema, elementwiseNestingDepths) override protected def writeUDF(dataOut: DataOutputStream): Unit = { PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala new file mode 100644 index 0000000000000..dfaa4732ca9d7 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambda.scala @@ -0,0 +1,662 @@ +/* + * 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.execution.python + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern._ +import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType} + +/** + * Rewrites scalar Python UDFs inside a higher-order function's lambda so they can be evaluated. + * + * A `PythonUDF` runs in a separate operator that [[ExtractPythonUDFs]] pulls out, but a lambda's + * [[NamedLambdaVariable]]s only exist while the function iterates, so the UDF can neither stay in + * the lambda nor be lifted out normally. Instead this rule applies the UDF once to the *whole + * array*, outside every lambda, and has the lambda read the result positionally: + * + * {{{ + * -- before (rejected) + * transform(values, x -> plus_one(x)) + * + * -- after (the PythonUDF is outside every lambda) + * transform(arrays_zip(values AS c0, plus_one_over_array(values) AS u0), s -> s.u0) + * }}} + * + * `plus_one_over_array` is the same function re-typed as `array<T> => array<R>` and run with an + * element-wise eval type chosen from the UDF's own flavor (see + * [[PythonUDF.liftedElementwiseEvalType]]): the row-at-a-time UDFs share the pickle-based + * [[org.apache.spark.api.python.PythonEvalType]]'s SQL_ARROW_ELEMENTWISE_UDF, while a scalar pandas + * / Arrow UDF (and its iterator variant) lifts to its own element-wise type so the worker keeps + * that flavor's batching contract. + * The array-at-a-time behaviour lives in the Python worker: it flattens each list column once, + * calls the function over all elements of the batch, and re-nests by the input's offsets - one row + * in, one row out, one Python round trip per batch. + * + * Every lifted argument is a single-level `array<T>` aligned with the iterated array (an + * element-independent value is repeated into one with a native `transform`), so the worker flattens + * them uniformly with no per-argument metadata. With the result now an ordinary column, arithmetic, + * `when`, casts, the element index, multiple UDFs and nested calls `f(g(x))` all just work. + * + * Runs before [[ExtractPythonUDFs]]. Handles all ten single-lambda functions: `transform`, + * `filter`, `exists`, `forall`, `zip_with`, `array_sort`, and the four map functions (desugared to + * `map_keys`/`map_values` arrays and rebuilt with `map_from_arrays`). `array_sort` precomputes a + * per-element key, or, when one call takes both elements, the UDF over the cross product of pairs. + * A UDF in a *nested* lambda, `transform(arr, i -> transform(i, x -> f(x)))`, is handled too: the + * whole nest is rewritten root-first (see `apply` / `rewriteNest`), lifting the UDF out one lambda + * level at a time so it ends up applied to the fully flattened leaves. + * + * `CheckAnalysis` still rejects what this rule does not handle: + * - a UDF in `aggregate` / `reduce`: the fold is sequential (the UDF sees earlier steps' outputs, + * not array elements), so it cannot be applied once to the whole array. + */ +object ExtractPythonUDFFromLambda extends Rule[LogicalPlan] { + + def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.pythonUDFInHigherOrderFunctionEnabled) { + plan + } else { + // Rewrite each expression tree top-down. The first (outermost) higher-order function that is + // a liftable *nest root* - it iterates real columns and its whole, possibly nested, nest is + // rewritable - has its entire nest rewritten in one action by `rewriteNest`. Handling the + // nest atomically means a UDF in a nested lambda is never momentarily left as a free-variable + // `PythonUDF` that downstream could not evaluate (SPARK-48706): the rule either rewrites a + // nest completely or leaves it untouched for `CheckAnalysis` to reject. The gate mirrors + // `CheckAnalysis` exactly, so analysis accepts precisely what this rewrites. + plan.transformUpWithPruning( + _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) { + case p => + p.transformExpressionsDownWithPruning( + _.containsAllPatterns(PYTHON_UDF, HIGH_ORDER_FUNCTION)) { + case hof: HigherOrderFunction + if hof.functions.exists(_.exists(_.isInstanceOf[PythonUDF])) && + PythonUDF.canRewritePythonUDFInLambda(hof) => + rewriteNest(hof) + } + } + } + } + + /** + * Rewrites a whole nest of higher-order functions rooted at `root`, already validated liftable by + * [[PythonUDF.canRewritePythonUDFInLambda]]. `transformUp` visits the innermost function first, + * so each inner lambda's UDFs are lifted onto that lambda's (enclosing-variable) argument - + * becoming element-wise UDFs in an argument position - and then the enclosing function re-lifts + * them onto its own array one level deeper (see `buildCarrier` / `overArray`). A single bottom-up + * pass therefore lifts every UDF out of every lambda in the nest, innermost first. + */ + private def rewriteNest(root: HigherOrderFunction): Expression = root.transformUp(rewriteOne) + + /** + * Whether one UDF call in an `array_sort` comparator takes both elements, e.g. + * `(a, b) -> udf(a, b)`. Such a call has no per-element key, so it is precomputed over the cross + * product of pairs rather than per element. + */ + private def comparatorTakesBothElements(function: Expression): Boolean = function match { + case LambdaFunction(body, Seq(left: NamedLambdaVariable, right: NamedLambdaVariable), _) => + body.exists { + case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) => + def reads(id: ExprId) = udf.exists { + case v: NamedLambdaVariable => v.exprId == id + case _ => false + } + reads(left.exprId) && reads(right.exprId) + case _ => false + } + case _ => false + } + + /** + * Rewrites one higher-order function whose own lambda holds a rewritable Python UDF. The generic + * path reads arguments, lambdas and parameter roles off the [[HigherOrderFunction]] API and + * rebuilds with `withNewChildren` (children are `arguments` then `functions`), naming a concrete + * class only where a shape cannot be inferred otherwise (`ArraySort`'s comparator, and the + * result-type traits telling `ArrayFilter` from `ArrayTransform`). A pairwise `array_sort` + * comparator, whose single call takes both elements, needs its own path. + * + * Applied by `rewriteNest` to every function in a validated nest, innermost first. Unlike the + * nest-root gate in `apply`, `liftableHof` here does not re-check free lambda variables: within + * an already-validated nest an inner function iterating an enclosing variable is expected, and + * `rewriteNest` guarantees the enclosing function re-lifts whatever an inner step leaves in an + * argument position. + */ + private val rewriteOne: PartialFunction[Expression, Expression] = { + case sort @ ArraySort(_, function, _) + if liftableHof(sort) && comparatorTakesBothElements(function) => + rewritePairwiseComparator(sort) + + // Every result-typed higher-order function is handled; anything else is left alone. + case hof: HigherOrderFunction + if liftableHof(hof) && + (hof.isInstanceOf[ResultTypeFromArgument] || hof.isInstanceOf[ResultTypeFromFunction]) => + rewriteMapping(hof) + } + + /** + * Rewrites `array_sort(arr, (a, b) -> udf(a, b))`, where one call takes both elements so there is + * no per-element key. Precomputes the UDF over every ordered pair - an n x n matrix with + * `udf(arr[i], arr[j])` at (i, j) - and the comparator reads it by the two elements' positions, + * so no Python runs while sorting. Costs O(n^2) calls and memory vs. O(n) for a per-element key. + */ + private def rewritePairwiseComparator(sort: ArraySort): Expression = { + val ArraySort(argument, function, allowNull) = sort + val LambdaFunction(body, Seq(leftVar: NamedLambdaVariable, rightVar: NamedLambdaVariable), _) = + function + val arrayType = argument.dataType.asInstanceOf[ArrayType] + val elementType = arrayType.elementType + val containsNull = arrayType.containsNull + val n = Size(argument) + + // The two sides of the cross product. `array_repeat` avoids introducing a lambda that could + // capture the UDF; the one lambda here holds only the repeat, never the UDF. + val repeatVar = NamedLambdaVariable("a", elementType, containsNull) + val lefts = Flatten( + ArrayTransform(argument, LambdaFunction(ArrayRepeat(repeatVar, n), Seq(repeatVar)))) + val rights = Flatten(ArrayRepeat(argument, n)) + + // The UDF over all n*n pairs: this is just the element-wise rewrite with the pair arrays as the + // iterated arguments, so `buildCarrier` lifts the UDF and a `transform` runs the rest of the + // comparator body (cast, `when`, arithmetic) once per pair in the JVM. + val pairLambda = LambdaFunction(body, Seq(leftVar, rightVar)) + val pairCarrier = buildCarrier(Seq(lefts, rights), pairLambda, Seq(leftVar, rightVar), None) + val flatCells = ArrayTransform( + pairCarrier.carrier, LambdaFunction(pairCarrier.body, Seq(pairCarrier.boundVar))) + + // Carry each element's position and the shared flat result array so the comparator can read + // its pair's precomputed cell, sort, then drop them again. `flatCells` must be built here, in + // the sort's *argument*, not inside the comparator: `ArraySort` re-evaluates the whole + // comparator body on every comparison, and `ExtractPythonUDFs` hoists only the `PythonUDF` + // node - the surrounding `arrays_zip`/`transform`/`flatten`/`array_repeat` that build the cells + // would otherwise be rebuilt O(n^2) per comparison (O(n^3 log n) overall). In interpreted + // evaluation `array_repeat` stores n references to the one computed `flatCells` array, so the + // carry is O(n^2); a later copy of the carrier into the Unsafe format would materialize each + // reference into O(n^3) bytes. Either way the whole pairwise path is already O(n^2) in Python + // calls, so it is only intended for small arrays (see the config doc). + // Also carry the row width `n` so the comparator does not re-evaluate `Size(argument)` on every + // comparison (negligible for a column, but `argument` may be a computed expression). Like the + // cells, it is repeated into the carrier once and read as a struct field. + val posElem = NamedLambdaVariable("x", elementType, containsNull) + val posIdx = NamedLambdaVariable("i", IntegerType, nullable = false) + val cellsField = "cells" + val sizeField = "n" + val indexed = ArraysZip( + Seq( + argument, + ArrayTransform(argument, LambdaFunction(posIdx, Seq(posElem, posIdx))), + ArrayRepeat(flatCells, n), + ArrayRepeat(n, n)), + Seq( + Literal(s"${carrierElementPrefix}0"), + Literal(carrierIndexField), + Literal(cellsField), + Literal(sizeField))) + val indexedElement = indexed.dataType.asInstanceOf[ArrayType].elementType + + // Index the flat n*n results directly: cell (i, j) is at `i * n + j`. The cells and `n` live in + // struct fields carried by every element, so the comparator only does field reads plus an + // `element_at`, all O(1). `element_at` is 1-based. + val cmpLeft = NamedLambdaVariable("a", indexedElement, nullable = false) + val cmpRight = NamedLambdaVariable("b", indexedElement, nullable = false) + def idxOf(v: NamedLambdaVariable): Expression = GetStructField(v, 1, Some(carrierIndexField)) + val cells = GetStructField(cmpLeft, 2, Some(cellsField)) + val width = GetStructField(cmpLeft, 3, Some(sizeField)) + val comparison = ElementAt( + cells, + Add(Add(Multiply(idxOf(cmpLeft), width), idxOf(cmpRight)), Literal(1)), + None, + failOnError = false) + + unwrapCarrier( + ArraySort(indexed, LambdaFunction(comparison, Seq(cmpLeft, cmpRight)), allowNull), 0) + } + + /** + * The generic rewrite for a mapping higher-order function. + * + * A map-valued argument is first desugared to its key and value arrays, so everything below works + * in terms of arrays; the result is rebuilt as a map afterwards. The lambda's parameters are then + * matched to those arrays, the UDFs are lifted onto them, and the node is rebuilt around a + * carrier that the single new lambda parameter reads. + */ + private def rewriteMapping(hof: HigherOrderFunction): Expression = { + val lambda = hof.functions.head.asInstanceOf[LambdaFunction] + // The result is the input elements (so the carrier is unwrapped afterwards) rather than the + // lambda's value: `filter` / `array_sort` / `map_filter` keep the input's type. + val isFromElements = hof.isInstanceOf[ResultTypeFromArgument] + + // Desugar maps into arrays. `map_zip_with` visits the union of both key sets and looks each map + // up per key, which yields null for a key missing from one side - exactly its own semantics. + val mapValued = hof.arguments.exists(_.dataType.isInstanceOf[MapType]) + val (arrays, rebuildResult): (Seq[Expression], Expression => Expression) = + if (!mapValued) { + (hof.arguments, identity) + } else if (hof.arguments.length == 1) { + val map = hof.arguments.head + val keys = MapKeys(map) + val values = MapValues(map) + // Rebuild by the concrete function, not the result type: `transform_keys` replaces the + // keys, `transform_values` the values, `map_filter` keeps whichever pairs survive. Keying + // off the type would be wrong for e.g. `transform_values` on `map<string, string>`, whose + // lambda result type equals the key type. + val rebuild: Expression => Expression = hof match { + case _: TransformKeys => (newKeys: Expression) => MapFromArrays(newKeys, values) + case _: TransformValues => (newValues: Expression) => MapFromArrays(keys, newValues) + case _: MapFilter => (kept: Expression) => + MapFromArrays(unwrapCarrier(kept, 0), unwrapCarrier(kept, 1)) + } + (Seq(keys, values), rebuild) + } else { + val Seq(left, right) = hof.arguments + val keys = ArrayUnion(MapKeys(left), MapKeys(right)) + val keyType = keys.dataType.asInstanceOf[ArrayType] + def valuesFor(map: Expression): Expression = { + val k = NamedLambdaVariable("k", keyType.elementType, keyType.containsNull) + ArrayTransform(keys, LambdaFunction(ElementAt(map, k, None, failOnError = false), Seq(k))) + } + (Seq(keys, valuesFor(left), valuesFor(right)), + (newValues: Expression) => MapFromArrays(keys, newValues)) + } + + // Match lambda parameters to the arrays they iterate: leading ones map to the arrays, a + // trailing extra one is the element index. `array_sort` is the one exception - its lambda is a + // comparator whose two parameters are two elements of the *same* array, indistinguishable from + // an indexed lambda by types alone (both `(T, Int)`), so it is special-cased by class here. + val params = lambda.arguments.map(_.asInstanceOf[NamedLambdaVariable]) + val (elementVars, indexVar, alsoBind) = + if (hof.isInstanceOf[ArraySort]) { + (Seq(params.head), None, Seq(params.last)) + } else { + (params.take(arrays.length), params.drop(arrays.length).headOption, Nil) + } + + val built = buildCarrier(arrays, lambda, elementVars, indexVar, alsoBind) + val newLambda = LambdaFunction(built.body, built.boundVar +: built.extraBoundVars) + + // Rebuild the node over the single carrier. A single-array function keeps its own class (via + // `withNewChildren`, children being arguments then functions); a desugared map or a multi-array + // one becomes a `transform`, or an `ArrayFilter` when the carrier must survive the filtering so + // both key and value sides can be projected out. + val keepsOwnNode = hof.arguments.length == 1 && !mapValued + val iterated = + if (keepsOwnNode) { + hof.withNewChildren(IndexedSeq(built.carrier, newLambda)).asInstanceOf[Expression] + } else if (isFromElements) { + ArrayFilter(built.carrier, newLambda) + } else { + ArrayTransform(built.carrier, newLambda) + } + + // A from-elements result (e.g. `filter`) is the input elements, so project them back out of the + // carrier; for a map `rebuildResult` knows which of the key/value sides to keep. + if (!mapValued && isFromElements) rebuildResult(unwrapCarrier(iterated, 0)) + else rebuildResult(iterated) + } + + /** + * True if `hof`'s single lambda holds a UDF to lift at *this* level - either directly in the body + * or in a nested function's argument (which `hasDirectRewritableUDF` reaches, but a nested + * function's own lambda is not this level's concern; `rewriteNest` handles that level itself). + * + * Does not re-check free lambda variables: `apply` enters `rewriteNest` only on a validated nest + * root, so within the nest an inner function iterating an enclosing variable is expected, and the + * enclosing function is guaranteed to re-lift whatever this level leaves in an argument position. + */ + private def liftableHof(hof: HigherOrderFunction): Boolean = + hof.functions.length == 1 && (hof.functions.head match { + case LambdaFunction(body, args, _) => + hasDirectRewritableUDF(body) && args.forall(_.isInstanceOf[NamedLambdaVariable]) + case _ => false + }) + + /** + * Whether `body` holds a rewritable UDF belonging to *this* lambda. A nested function's lambda is + * skipped (its UDF reads that lambda's variable), but its *arguments* are not: in + * `transform(arr, x -> transform(udf(x), y -> y))`, `udf(x)` is in the inner argument and lifts + * onto `arr`. + */ + private def hasDirectRewritableUDF(body: Expression): Boolean = body match { + case e if PythonUDF.isElementwiseRewritableUDF(e) => true + case hof: HigherOrderFunction => hof.arguments.exists(hasDirectRewritableUDF) + case e => e.children.exists(hasDirectRewritableUDF) + } + + /** The pieces produced by [[buildCarrier]]. */ + private case class Carrier( + carrier: Expression, + body: Expression, + boundVar: NamedLambdaVariable, + extraBoundVars: Seq[NamedLambdaVariable]) + + /** + * Builds the carrier array and the rewritten lambda body. + * + * The carrier is `arrays_zip` of the original arrays, one array per lifted UDF, and - when the + * lambda declares an index parameter - an index array. The rewritten body reads each of those + * through a struct field of the lambda variable bound to the carrier. + * + * `alsoBind` names further lambda variables that should read the same carrier; it exists for + * `array_sort`'s comparator, whose two parameters are both elements of the same array. + */ + private def buildCarrier( + arguments: Seq[Expression], + function: Expression, + elementVars: Seq[NamedLambdaVariable], + indexVar: Option[NamedLambdaVariable], + alsoBind: Seq[NamedLambdaVariable] = Nil): Carrier = { + val LambdaFunction(body, _, _) = function + val lambdaExprIds = + (elementVars ++ indexVar.toSeq ++ alsoBind).map(_.exprId).toSet + + // Collect the UDF calls to lift. Innermost first, so that a nested call like `f(g(x))` has + // `g` lifted before `f`, letting `f`'s array UDF consume `g`'s array result. + val liftableUDFs = collectLiftableUDFs(body) + + // With more than one argument the arrays may be ragged (`zip_with` / `map_zip_with` pad with + // nulls), so flattening them independently would misalign the elements. Projecting each out of + // one common `arrays_zip` pads them to the same per-row length, which the positional rewrite + // requires. + val alignedArguments = + if (arguments.length > 1) { + val names = arguments.indices.map(i => s"$carrierElementPrefix$i") + val zipped = ArraysZip(arguments, names.map(Literal(_))) + arguments.indices.map(i => unwrapCarrier(zipped, i)) + } else { + arguments + } + + // An index array, when the lambda asked for the element index. + val indexArray = indexVar.map { _ => + val head = alignedArguments.head + val headType = head.dataType.asInstanceOf[ArrayType] + val v = NamedLambdaVariable("x", headType.elementType, headType.containsNull) + val i = NamedLambdaVariable("i", IntegerType, nullable = false) + ArrayTransform(head, LambdaFunction(i, Seq(v, i))) + } + + // Maps each element/index variable to the array it stands for, so a UDF argument written in + // terms of the variables can be rewritten as an expression over whole arrays. For a + // comparator, `alsoBind`'s variables denote the same array as the element variable. + val arrayOfVar: Map[ExprId, Expression] = + elementVars.map(_.exprId).zip(alignedArguments).toMap ++ + indexVar.map(_.exprId -> indexArray.get).toMap ++ + alsoBind.map(_.exprId -> alignedArguments.head).toMap + + // One lifted array UDF per distinct call. `arrayResults` maps each original call (by `liftKey`) + // to the array holding its per-element results, so a nested call `f(g(x))` and the carrier + // lookups can find it. Deterministic calls that lift to the *same* function over the *same* + // array arguments share one lifted UDF: a key-form comparator's `udf(a)` and `udf(b)` + // canonicalize differently (`a` != `b`) but both read the whole array, so without this the + // Python function would run 2n times instead of n. Nondeterministic calls stay distinct (their + // signature is the lifted node itself, carrying a distinct `resultId`), matching `liftKey`. + var arrayResults = Map.empty[Expression, Expression] + val distinctLifted = scala.collection.mutable.ArrayBuffer.empty[PythonUDF] + val ordinalBySignature = scala.collection.mutable.HashMap.empty[Expression, Int] + val udfFieldByKey = scala.collection.mutable.LinkedHashMap.empty[Expression, Int] + liftableUDFs.foreach { udf => + // `overArray` turns each argument into an `array<T>` aligned with the iterated array, so the + // worker flattens every one exactly once (no per-argument shape to track). A keyword argument + // keeps its `NamedArgumentExpression` wrapper as a direct child of the lifted UDF - only its + // value is lifted - so the runner still derives the kwargs mapping from the direct children. + val arrayArgs = udf.children.map { + case NamedArgumentExpression(key, value) => + NamedArgumentExpression( + key, overArray(value, alignedArguments.head, arrayOfVar, lambdaExprIds, arrayResults)) + case child => + overArray(child, alignedArguments.head, arrayOfVar, lambdaExprIds, arrayResults) + } + // Each lift wraps the arguments in exactly one more `array` level. Lifting a base UDF gives + // depth 1; re-lifting an already-lifted element-wise UDF (a UDF from a *nested* lambda, + // lifted once onto the inner variable and now again onto the enclosing array) adds one more + // level, so the worker flattens `depth` levels down to the scalar element and re-nests them. + val newDepth = + if (PythonEvalType.isElementwiseUDF(udf.evalType)) udf.elementwiseNestingDepth + 1 else 1 + val lifted = PythonUDF( + udf.name, + udf.func, + // The wrapper returns one element per input element, i.e. one array level on top of the + // UDF's previous return type. Elements may be null (the UDF can return null), hence + // containsNull = true. + ArrayType(udf.dataType, containsNull = true), + arrayArgs, + // Each rewritable flavor lifts to its own element-wise eval type so the worker keeps that + // flavor's batching contract (pickle row-at-a-time, pandas Series, Arrow Array, or an + // iterator of batches); an already-lifted type maps to itself. See + // `PythonUDF.liftedElementwiseEvalType`. + PythonUDF.liftedElementwiseEvalType(udf.evalType), + udf.udfDeterministic, + elementwiseNestingDepth = newDepth) + val signature: Expression = if (udf.udfDeterministic) lifted.canonicalized else lifted + val ordinal = ordinalBySignature.getOrElseUpdate(signature, { + val o = distinctLifted.length + distinctLifted += lifted + o + }) + arrayResults += (liftKey(udf) -> distinctLifted(ordinal)) + udfFieldByKey += (liftKey(udf) -> ordinal) + } + val liftedArrays = distinctLifted.toSeq + + // The carrier: the original arrays first, then one field per lifted UDF, then the index. + val carrierFields = alignedArguments ++ liftedArrays ++ indexArray.toSeq + val carrierNames = + arguments.indices.map(i => s"$carrierElementPrefix$i") ++ + liftedArrays.indices.map(i => s"$carrierUDFFieldPrefix$i") ++ + indexArray.map(_ => carrierIndexField).toSeq + val carrier = ArraysZip(carrierFields, carrierNames.map(Literal(_))) + + val structType = carrier.dataType.asInstanceOf[ArrayType].elementType + val boundVar = NamedLambdaVariable("s", structType, nullable = false) + val extraBoundVars = alsoBind.map(v => + NamedLambdaVariable(v.name, structType, nullable = false)) + + // Which struct field each lambda variable reads. For a comparator, `alsoBind`'s variable reads + // the same ordinals but through its own bound variable. + val fieldOfVar: Map[ExprId, Int] = + elementVars.map(_.exprId).zipWithIndex.toMap ++ + indexVar.map(_.exprId -> (carrierFields.length - 1)).toMap + val extraVarOf: Map[ExprId, NamedLambdaVariable] = + alsoBind.map(_.exprId).zip(extraBoundVars).toMap + + // Rewrite the body. This must be top-down: a UDF call is matched by its canonicalized form, + // and rewriting its arguments first (a variable becoming a struct field read) would change + // that form so the call no longer matches and would be left inside the lambda. Replacing the + // call outright also stops the traversal descending into arguments that no longer exist. + def readerFor(v: NamedLambdaVariable, udfOrdinal: Option[Int]): Expression = { + val base = extraVarOf.getOrElse(v.exprId, boundVar) + udfOrdinal match { + case Some(u) => + GetStructField(base, arguments.length + u, Some(s"$carrierUDFFieldPrefix$u")) + case None => + val ordinal = fieldOfVar(v.exprId) + GetStructField(base, ordinal, Some(carrierNames(ordinal))) + } + } + + val rewrittenBody = body.transformDown { + case udf: PythonUDF if udfFieldByKey.contains(liftKey(udf)) => + val ordinal = udfFieldByKey(liftKey(udf)) + // A UDF over a comparator's right-hand element must read that element's key, so the + // struct field is read through whichever bound variable the call's own arguments used. + val side = udf.collectFirst { + case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => v + } + side match { + case Some(v) => readerFor(v, Some(ordinal)) + case None => + GetStructField(boundVar, arguments.length + ordinal, + Some(s"$carrierUDFFieldPrefix$ordinal")) + } + case v: NamedLambdaVariable if fieldOfVar.contains(v.exprId) => readerFor(v, None) + case v: NamedLambdaVariable if extraVarOf.contains(v.exprId) => + // A comparator's right-hand element itself, read through its own bound variable. + GetStructField(extraVarOf(v.exprId), 0, Some(carrierNames.head)) + } + + Carrier(carrier, rewrittenBody, boundVar, extraBoundVars) + } + + /** + * Collects the Python UDF calls directly in `body` that must be lifted, innermost first. + * + * Every rewritable UDF directly in the lambda body is lifted, even one whose arguments do not + * read the lambda variable (`transform(arr, _ -> f(lit(10)))`). It cannot stay inside the lambda, + * and lifting it - `overArray` repeats a constant/outer-column argument into an aligned array - + * gives it the lambda's own call domain: once per element, and zero times for a null or empty + * array. Leaving it to [[ExtractPythonUDFs]] would instead call it once per row, including rows + * whose array is null or empty where the lambda never runs. + */ + private def collectLiftableUDFs(body: Expression): Seq[PythonUDF] = { + val collected = Seq.newBuilder[PythonUDF] + def visit(e: Expression): Unit = { + // A nested higher-order function's lambda is not ours to rewrite, but its arguments are + // evaluated outside that lambda and so belong to this body. See `hasDirectRewritableUDF`. + val children = e match { + case hof: HigherOrderFunction => hof.arguments + case other => other.children + } + // Children first, so nested calls come out innermost-first. + children.foreach(visit) + e match { + case udf: PythonUDF if PythonUDF.isElementwiseRewritableUDF(udf) => + collected += udf + case _ => + } + } + visit(body) + // Deduplicate identical calls so the same UDF is evaluated once per array. Nondeterministic + // calls are kept distinct (see `liftKey`). + val seen = scala.collection.mutable.LinkedHashMap.empty[Expression, PythonUDF] + collected.result().foreach(udf => seen.getOrElseUpdate(liftKey(udf), udf)) + seen.values.toSeq + } + + private def readsLambdaVariable(e: Expression, lambdaExprIds: Set[ExprId]): Boolean = + e.exists { + case v: NamedLambdaVariable => lambdaExprIds.contains(v.exprId) + case _ => false + } + + /** + * The key that decides whether two UDF calls are "the same call" for lifting. A deterministic + * call is deduplicated by canonical form, so an identical call is evaluated once. A + * nondeterministic call must stay distinct - `transform(arr, x -> f(x) + f(x))` calls `f` twice + * and each call may return a different value - so it is keyed by its own `resultId`-bearing node + * (`canonicalized` erases `resultId`, which would collapse the two calls into one). + */ + private def liftKey(udf: PythonUDF): Expression = + if (udf.udfDeterministic) udf.canonicalized else udf + + /** + * Turns a UDF argument expression, written in terms of single elements, into the equivalent + * `array<T>` aligned with the iterated array, so the worker can flatten every argument uniformly. + * + * - a lambda variable becomes the array it stands for; + * - an already-lifted UDF call (a nested call `f(g(x))`) becomes its array result; + * - an expression independent of the lambda is repeated into an aligned array with a native + * `transform`, rather than passed through as a scalar to broadcast; + * - anything else is an expression over the elements, computed for every element by a native + * `transform` that stays inside the JVM. + * + * An already-lifted UDF call nested inside a composite argument (`f(g(x) + 1)`, `f(-g(x))`) is + * handled by the last case too: each such call is replaced by a synthetic variable standing for + * that call's aligned array result, which is then zipped in like any element array. Leaving the + * raw `g` inside the generated `transform` lambda would put a `PythonUDF` back inside a lambda - + * `ExtractPythonUDFs` would then extract a `g` whose child is a `NamedLambdaVariable` (the + * SPARK-48706 failure mode), so this substitution is for correctness, not just efficiency. + */ + private def overArray( + child: Expression, + firstArgument: Expression, + arrayOfVar: Map[ExprId, Expression], + lambdaExprIds: Set[ExprId], + arrayResults: Map[Expression, Expression]): Expression = child match { + case v: NamedLambdaVariable if arrayOfVar.contains(v.exprId) => arrayOfVar(v.exprId) + case udf: PythonUDF if arrayResults.contains(liftKey(udf)) => + arrayResults(liftKey(udf)) + case e => + // Replace each already-lifted nested UDF call with a synthetic variable standing for that + // call's aligned array result, then fold those arrays into the variable-to-array map so the + // logic below treats them exactly like element variables. This keeps a lifted UDF buried in + // a composite argument (`f(g(x) + 1)`) from being left as a raw UDF inside a lambda. + val nestedVars = scala.collection.mutable.LinkedHashMap.empty[Expression, NamedLambdaVariable] + val expr = e.transformUp { + case u: PythonUDF if arrayResults.contains(liftKey(u)) => + nestedVars.getOrElseUpdate(liftKey(u), { + val arrType = arrayResults(liftKey(u)).dataType.asInstanceOf[ArrayType] + NamedLambdaVariable("g", arrType.elementType, arrType.containsNull) + }) + } + val fullArrayOf = arrayOfVar ++ + nestedVars.map { case (key, v) => v.exprId -> arrayResults(key) } + + if (!readsLambdaVariable(expr, fullArrayOf.keySet)) { + // Independent of the element (an outer column or constant): repeat it into an array aligned + // with the iterated array, so every UDF argument is a single-level array the worker + // flattens the same way. `transform(arr, _ -> e)` keeps the value constant, matching shape. + val arrType = firstArgument.dataType.asInstanceOf[ArrayType] + val v = NamedLambdaVariable("x", arrType.elementType, arrType.containsNull) + ArrayTransform(firstArgument, LambdaFunction(expr, Seq(v))) + } else { + // An expression over the element(s) and/or nested results, e.g. `udf(x * 2)` or + // `f(g(x) + 1)`. Compute it for every element with a native `transform`, which stays inside + // the JVM. Multi-argument shapes need the values side by side, so zip them first. + val arrays = fullArrayOf.values.toSeq.distinct + if (arrays.length == 1) { + val arr = arrays.head + val elemType = arr.dataType.asInstanceOf[ArrayType] + val v = NamedLambdaVariable("x", elemType.elementType, elemType.containsNull) + val substituted = expr.transformUp { + case old: NamedLambdaVariable if fullArrayOf.contains(old.exprId) => v + } + ArrayTransform(arr, LambdaFunction(substituted, Seq(v))) + } else { + // Zip every array the expression may read, then project the fields it needs. + val names = arrays.indices.map(i => s"$carrierElementPrefix$i") + val zipped = ArraysZip(arrays, names.map(Literal(_))) + val structType = zipped.dataType.asInstanceOf[ArrayType].elementType + val v = NamedLambdaVariable("z", structType, nullable = false) + val ordinalOf = fullArrayOf.map { case (id, arr) => id -> arrays.indexOf(arr) } + val substituted = expr.transformUp { + case old: NamedLambdaVariable if ordinalOf.contains(old.exprId) => + GetStructField(v, ordinalOf(old.exprId), Some(names(ordinalOf(old.exprId)))) + } + ArrayTransform(zipped, LambdaFunction(substituted, Seq(v))) + } + } + } + + /** + * Projects field `ordinal` back out of a carrier array. Used where the result is built from the + * input elements rather than from the lambda's return value: `filter`, `array_sort` and the map + * family. + */ + private def unwrapCarrier(carrierArray: Expression, ordinal: Int): Expression = { + val structType = carrierArray.dataType.asInstanceOf[ArrayType].elementType + val v = NamedLambdaVariable("s", structType, nullable = false) + ArrayTransform( + carrierArray, + LambdaFunction( + GetStructField(v, ordinal, Some(s"$carrierElementPrefix$ordinal")), Seq(v))) + } + + private val carrierElementPrefix = "c" + private val carrierUDFFieldPrefix = "u" + private val carrierIndexField = "idx" +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala index af9b41e93a029..c290ec04dedab 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala @@ -241,7 +241,12 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with Logging { def canChainWithParallelUDFs(evalType: Int): Boolean = { if (evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF || - evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF) { + evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF || + // The lifted iterator flavors carry the same one-UDF-per-operator constraint as their base + // iterator types: the worker feeds a single iterator UDF, so they must not be + // parallel-fused with siblings (e.g. `transform(arr, x -> f(x) + g(x))` for iterator UDFs). + evalType == PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF || + evalType == PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF) { false } else { evalType == firstVisitedScalarUDFEvalType.get @@ -264,7 +269,14 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with Logging { expressions.flatMap(collectEvaluableUDFs) } - def apply(plan: LogicalPlan): LogicalPlan = plan match { + def apply(plan: LogicalPlan): LogicalPlan = applyInternal( + // Lift Python UDFs out of higher-order function lambdas first, turning them into ordinary + // top-level UDFs this rule can then extract. Running it here rather than as a separate batch + // entry makes the ordering structural: `CheckAnalysis` accepts a lambda-UDF plan only because + // this rewrite will run, so the two must not be separable by rule reordering or exclusion. + ExtractPythonUDFFromLambda(plan)) + + private def applyInternal(plan: LogicalPlan): LogicalPlan = plan match { // SPARK-26293: A subquery will be rewritten into join later, and will go through this rule // eventually. Here we skip subquery, as Python UDF only needs to be extracted once. case s: Subquery if s.correlated => plan @@ -342,6 +354,11 @@ object ExtractPythonUDFs extends Rule[LogicalPlan] with Logging { case PythonEvalType.SQL_SCALAR_PANDAS_UDF | PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF | PythonEvalType.SQL_ARROW_BATCHED_UDF + | PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF + | PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF + | PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF + | PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF + | PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF | PythonEvalType.SQL_SCALAR_ARROW_UDF | PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF => ArrowEvalPython(validUdfs, resultAttrs, child, evalType) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala new file mode 100644 index 0000000000000..9d8931c7f9192 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala @@ -0,0 +1,393 @@ +/* + * 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.execution.python + +import java.io.File + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{JobArtifactSet, SparkEnv, TaskContext} +import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType} +import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.physical.{AllTuples, ClusteredDistribution, Distribution, Partitioning, UnspecifiedDistribution} +import org.apache.spark.sql.execution.{GroupedIterator, SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata +import org.apache.spark.sql.types.{DataType, StructField, StructType} +import org.apache.spark.util.Utils + +/** + * Execution logic for the post-shuffle FINAL stage of an incremental Python aggregation (see + * [[org.apache.spark.sql.catalyst.expressions.PythonAggregate]]). It groups the (shuffled) child + * rows by the grouping expressions via a local sort + [[GroupedIterator]], sends each group's + * intermediate-buffer columns to the Python worker as Arrow record batches, and joins the single + * result row the worker returns per group back with the grouping key. + * + * The map-side PARTIAL stage does not extend this base: it hash-combines many groups per batch + * inside the worker and needs no sort (see [[PythonIncrementalAggregatePartialExec]]). This base is + * kept as an abstract class so the per-stage inputs, eval type, output attributes and final + * projection stay explicit hooks. + */ +abstract class PythonIncrementalAggregateExecBase extends UnaryExecNode with PythonSQLMetrics { + + def groupingExpressions: Seq[NamedExpression] + def aggExpressions: Seq[AggregateExpression] + + protected val udfExpressions: Seq[PythonAggregate] = + aggExpressions.map(_.aggregateFunction.asInstanceOf[PythonAggregate]) + + /** The Python eval type for this stage. */ + protected def evalType: Int + + /** Per-UDF input expressions to project out of the child and send to the Python worker. */ + protected def udfInputs: Seq[Seq[Expression]] + + /** Attributes of the row the Python worker returns per group (right side of the join). */ + protected def pythonOutputAttributes: Seq[Attribute] + + /** Expressions producing this operator's output from (groupingKey ++ pythonOutput). */ + protected def outputExpressions: Seq[NamedExpression] + + /** The grouping attributes as seen in the child's output. */ + protected def groupingAttributes: Seq[Attribute] = groupingExpressions.map(_.toAttribute) + + /** + * Whether to still invoke Python on an empty partition. Only the FINAL stage of a *global* + * (no grouping) aggregation sets this: it must emit the identity row `finish(zero)` for empty + * input, matching SQL aggregate semantics. Everywhere else an empty partition yields no rows. + */ + protected def emitOnEmptyPartition: Boolean = false + + override def output: Seq[Attribute] = outputExpressions.map(_.toAttribute) + + override def producedAttributes: AttributeSet = AttributeSet(output) + + override def requiredChildOrdering: Seq[Seq[SortOrder]] = + Seq(groupingExpressions.map(SortOrder(_, Ascending))) + + override protected def doExecute(): RDD[InternalRow] = { + val inputRDD = child.execute() + + val sessionLocalTimeZone = conf.sessionLocalTimeZone + val largeVarTypes = conf.arrowUseLargeVarTypes + val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf) + + val pyFuncs = udfExpressions.map { u => + (ChainedPythonFunctions(Seq(u.func)), u.resultId.id) + } + + // Filter child output attributes down to only those that are UDF inputs, and eliminate + // duplicates, mirroring ArrowAggregatePythonExec. + val allInputs = new ArrayBuffer[Expression] + val dataTypes = new ArrayBuffer[DataType] + val argMetas = PythonIncrementalAggregateExec.buildArgMetas(udfInputs, allInputs, dataTypes) + + val aggInputSchema = StructType(dataTypes.zipWithIndex.map { case (dt, i) => + StructField(s"_$i", dt) + }.toArray) + + val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) + val sessionUUID = Option(session).collect { + case s if s.sessionState.conf.pythonWorkerLoggingEnabled => s.sessionUUID + } + + val groupingExprs = groupingExpressions + val childOutput = child.output + val joinedAttributes = groupingAttributes ++ pythonOutputAttributes + val resultExprs = outputExpressions + val localEvalType = evalType + + val emitIdentityOnEmpty = emitOnEmptyPartition + inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty && !emitIdentityOnEmpty) iter else { + val prunedProj = UnsafeProjection.create(allInputs.toSeq, childOutput) + + val groupedItr = if (groupingExprs.isEmpty) { + Iterator((new UnsafeRow(), iter)) + } else { + GroupedIterator(iter, groupingExprs, childOutput) + } + + // For a global aggregation with empty input, feed one all-null buffer row so the Python + // worker still emits `finish(zero)`. An empty group cannot be sent through + // GroupedPythonArrowInput (it asserts a non-empty batch per group), and the worker treats a + // null partial buffer as contributing nothing to `merge`. + lazy val nullInputRow: UnsafeRow = + UnsafeProjection.create(aggInputSchema.map(_.dataType).toArray) + .apply(new GenericInternalRow(aggInputSchema.length)).copy() + val grouped = groupedItr.map { case (key, rows) => + val projected = rows.map(prunedProj) + val toSend = if (emitIdentityOnEmpty && groupingExprs.isEmpty && !projected.hasNext) { + Iterator(nullInputRow) + } else { + projected + } + (key, toSend) + } + + val context = TaskContext.get() + + // In pipelined mode the queue's add() runs in the writer thread and remove() in the task + // thread; use lock-free mode to skip per-row synchronization (as ArrowAggregatePythonExec). + val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION) + val queue = HybridRowQueue(context.taskMemoryManager(), + new File(Utils.getLocalDir(SparkEnv.get.conf)), groupingExprs.length, lockFree = pipelined) + context.addTaskCompletionListener[Unit] { _ => queue.close() } + + val projectedRowIter = grouped.map { case (groupingKey, rows) => + queue.add(groupingKey.asInstanceOf[UnsafeRow]) + rows + } + + val runner = new ArrowPythonWithNamedArgumentRunner( + pyFuncs, + localEvalType, + argMetas, + aggInputSchema, + sessionLocalTimeZone, + largeVarTypes, + pythonRunnerConf, + pythonMetrics, + jobArtifactUUID, + sessionUUID) with GroupedPythonArrowInput + + val columnarBatchIter = runner.compute(projectedRowIter, context.partitionId(), context) + + val joined = new JoinedRow + val resultProj = UnsafeProjection.create(resultExprs, joinedAttributes) + + columnarBatchIter.map(_.rowIterator.next()).map { pythonOutputRow => + val leftRow = queue.remove() + resultProj(joined(leftRow, pythonOutputRow)) + } + }} + } +} + +/** + * Map-side PARTIAL stage: hash-combines the input rows of each group into a per-group intermediate + * buffer via the aggregator's `reduce`, inside the Python worker. Unlike the FINAL stage it needs + * neither a clustered distribution nor an ordering: the whole point of the map-side combine is to + * avoid a full pre-shuffle sort. It streams ordinary (multi-group) Arrow batches to the worker, + * which maintains one running buffer per distinct grouping key and emits, at end of partition, one + * row per key -- the grouping key columns followed by one intermediate-buffer struct column per + * aggregator. + * + * Because keys may be split across partitions (no shuffle here) and the worker's map-side grouping + * is only a combine optimization, correctness does not depend on it being exhaustive: + * [[PythonIncrementalAggregateFinalExec]] re-groups the emitted partial buffers authoritatively + * (by JVM `UnsafeRow` key, after the shuffle) and merges any that share a key. + */ +case class PythonIncrementalAggregatePartialExec( + groupingExpressions: Seq[NamedExpression], + aggExpressions: Seq[AggregateExpression], + bufferAttributes: Seq[Attribute], + child: SparkPlan) extends UnaryExecNode with PythonSQLMetrics { + + private val udfExpressions: Seq[PythonAggregate] = + aggExpressions.map(_.aggregateFunction.asInstanceOf[PythonAggregate]) + + private def groupingAttributes: Seq[Attribute] = groupingExpressions.map(_.toAttribute) + + override def output: Seq[Attribute] = groupingAttributes ++ bufferAttributes + + override def producedAttributes: AttributeSet = AttributeSet(output) + + override def requiredChildDistribution: Seq[Distribution] = Seq(UnspecifiedDistribution) + + // No ordering: the worker hash-combines rather than relying on grouped input. + override def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq(Nil) + + override def outputPartitioning: Partitioning = child.outputPartitioning + + override protected def doExecute(): RDD[InternalRow] = { + val inputRDD = child.execute() + + val sessionLocalTimeZone = conf.sessionLocalTimeZone + val largeVarTypes = conf.arrowUseLargeVarTypes + val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf) + + val pyFuncs = udfExpressions.map { u => + (ChainedPythonFunctions(Seq(u.func)), u.resultId.id) + } + + // The columns sent to Python are the grouping keys first (so the worker can hash-group by them + // and echo them back with each partial buffer), followed by the deduplicated aggregator input + // columns. A UDF input that coincides with a grouping key simply reuses that leading column. + val allInputs = new ArrayBuffer[Expression] + val dataTypes = new ArrayBuffer[DataType] + groupingExpressions.foreach { g => + allInputs += g + dataTypes += g.dataType + } + val numGroupingKeys = groupingExpressions.length + + val argMetas = PythonIncrementalAggregateExec.buildArgMetas( + udfExpressions.map(_.children), allInputs, dataTypes) + + val aggInputSchema = StructType(dataTypes.zipWithIndex.map { case (dt, i) => + StructField(s"_$i", dt) + }.toArray) + // The leading `numGroupingKeys` columns are the grouping keys; hand their schema to the worker. + val groupingKeySchemaJson = StructType(aggInputSchema.fields.take(numGroupingKeys)).json + + val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid) + val sessionUUID = Option(session).collect { + case s if s.sessionState.conf.pythonWorkerLoggingEnabled => s.sessionUUID + } + + val childOutput = child.output + val outputAttrs = output + + inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty) Iterator.empty else { + val prunedProj = UnsafeProjection.create(allInputs.toSeq, childOutput) + val projectedRowIter = iter.map(prunedProj) + + val context = TaskContext.get() + + val runner = new ArrowPythonWithNamedArgumentRunner( + pyFuncs, + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF, + argMetas, + aggInputSchema, + sessionLocalTimeZone, + largeVarTypes, + pythonRunnerConf, + pythonMetrics, + jobArtifactUUID, + sessionUUID) with BatchedPythonArrowInput { + // Tell the worker how many leading columns are grouping keys, so it can hash-group by them + // and re-emit them alongside each partial buffer. + override protected def evalConf: Map[String, String] = + super.evalConf + ("grouping_key_schema" -> groupingKeySchemaJson) + } + + val columnarBatchIter = runner.compute( + Iterator(projectedRowIter), context.partitionId(), context) + + // Each batch the worker returns holds (grouping key columns ++ one buffer struct column per + // aggregator), i.e. this operator's output columns; copy each row out as an UnsafeRow. + val resultProj = UnsafeProjection.create(outputAttrs, outputAttrs) + columnarBatchIter.flatMap(_.rowIterator().asScala).map(resultProj) + }} + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} + +/** + * Post-shuffle FINAL stage: clusters the partial buffers by the grouping key, merges the buffers + * of each group via the aggregator's `merge`, and produces the output via `finish`. Its input is + * the [[PythonIncrementalAggregatePartialExec]] output (grouping key columns followed by the + * intermediate-buffer columns); it sends the buffer columns to Python and outputs + * `resultExpressions`. + */ +case class PythonIncrementalAggregateFinalExec( + groupingExpressions: Seq[NamedExpression], + aggExpressions: Seq[AggregateExpression], + bufferAttributes: Seq[Attribute], + resultExpressions: Seq[NamedExpression], + child: SparkPlan) extends PythonIncrementalAggregateExecBase { + + override protected def evalType: Int = + PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF + + // Each aggregator reads its own intermediate-buffer column from the (shuffled) child. + override protected def udfInputs: Seq[Seq[Expression]] = bufferAttributes.map(Seq(_)) + + override protected def pythonOutputAttributes: Seq[Attribute] = + aggExpressions.map(_.resultAttribute) + + override protected def outputExpressions: Seq[NamedExpression] = resultExpressions + + // A global (no-grouping) aggregation must return the identity row even for empty input. This + // stage runs on a single partition (AllTuples), so exactly one identity row is produced. + override protected def emitOnEmptyPartition: Boolean = groupingExpressions.isEmpty + + override def requiredChildDistribution: Seq[Distribution] = { + if (groupingExpressions.isEmpty) { + AllTuples :: Nil + } else { + ClusteredDistribution(groupingExpressions) :: Nil + } + } + + override def outputPartitioning: Partitioning = child.outputPartitioning + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} + +object PythonIncrementalAggregateExec { + + /** + * Deduplicates the per-aggregator input expressions into a shared column list (matched by + * [[Expression.semanticEquals]]), returning one [[ArgumentMetadata]] array per aggregator that + * points into that list. `allInputs`/`dataTypes` may be pre-seeded -- the PARTIAL stage prepends + * its grouping-key columns so a UDF input equal to a grouping key reuses that leading column -- + * and any newly seen columns are appended to them in place. Shared by both stages (and mirrors + * the same dedup in [[ArrowAggregatePythonExec]]). + */ + private[python] def buildArgMetas( + udfInputs: Seq[Seq[Expression]], + allInputs: ArrayBuffer[Expression], + dataTypes: ArrayBuffer[DataType]): Array[Array[ArgumentMetadata]] = { + udfInputs.map { input => + input.map { e => + val (key, value) = e match { + case NamedArgumentExpression(key, value) => (Some(key), value) + case _ => (None, e) + } + if (allInputs.exists(_.semanticEquals(value))) { + ArgumentMetadata(allInputs.indexWhere(_.semanticEquals(value)), key) + } else { + allInputs += value + dataTypes += value.dataType + ArgumentMetadata(allInputs.length - 1, key) + } + }.toArray + }.toArray + } + + /** + * Builds the two-stage physical plan (PARTIAL -> [Exchange, inserted by EnsureRequirements] -> + * FINAL) for a logical aggregation whose aggregate functions are all [[PythonAggregate]]. + */ + def plan( + groupingExpressions: Seq[NamedExpression], + aggExpressions: Seq[AggregateExpression], + resultExpressions: Seq[NamedExpression], + child: SparkPlan): SparkPlan = { + // One intermediate-buffer attribute per aggregator, threaded from the PARTIAL output into the + // FINAL inputs (matched by expression id). + val bufferAttributes = aggExpressions.map { ae => + val agg = ae.aggregateFunction.asInstanceOf[PythonAggregate] + AttributeReference(s"buf_${agg.resultId.id}", agg.bufferSchema, nullable = true)() + } + val partial = PythonIncrementalAggregatePartialExec( + groupingExpressions, aggExpressions, bufferAttributes, child) + // After the PARTIAL stage the grouping expressions are materialized as plain attributes. + val groupingAttributes = groupingExpressions.map(_.toAttribute) + PythonIncrementalAggregateFinalExec( + groupingAttributes, aggExpressions, bufferAttributes, resultExpressions, partial) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala index 41bffaca65cfd..9f53f078b9176 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala @@ -18,18 +18,22 @@ package org.apache.spark.sql.execution.python import java.io.{DataInputStream, DataOutputStream} +import java.util.{List => JList} import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ import net.razorvine.pickle.Pickler import org.apache.spark.api.python.{PythonEvalType, PythonFunction, PythonWorkerUtils, SpecialLengths} import org.apache.spark.sql.{Column, TableArg} -import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, Expression, FunctionTableSubqueryArgumentExpression, NamedArgumentExpression, NullsFirst, NullsLast, PythonUDAF, PythonUDF, PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression, SortOrder, UnresolvedPolymorphicPythonUDTF, UnresolvedTableArgPlanId} +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, Expression, FunctionTableSubqueryArgumentExpression, NamedArgumentExpression, NullsFirst, NullsLast, PythonAggregate, PythonUDAF, PythonUDF, PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression, SortOrder, TranspiledPythonUDF, UnresolvedPolymorphicPythonUDTF, UnresolvedTableArgPlanId} import org.apache.spark.sql.catalyst.parser.ParserInterface import org.apache.spark.sql.catalyst.plans.logical.{Generate, LogicalPlan, NamedParametersSupport, OneRowRelation} import org.apache.spark.sql.classic.{DataFrame, Dataset, SparkSession} import org.apache.spark.sql.classic.ClassicConversions._ +import org.apache.spark.sql.classic.ColumnConversions import org.apache.spark.sql.classic.ExpressionUtils.expression import org.apache.spark.sql.errors.QueryCompilationErrors import org.apache.spark.sql.internal.{SQLConf, TableValuedFunctionArgument} @@ -43,7 +47,20 @@ case class UserDefinedPythonFunction( func: PythonFunction, dataType: DataType, pythonEvalType: Int, - udfDeterministic: Boolean) { + udfDeterministic: Boolean, + // TODO: Add support for transpilation with Spark Connect and remove the default value. + transpiled: JList[Column] = Nil.asJava, + // Per-option input-type categories ("numeric"/"string" per public param), + // parallel to `transpiled` (same length). The analyzer rule + // ResolveTranspiledPythonUDFOptions later keeps only the options whose + // categories match the bound argument types; when none match, the call + // falls back to the plain Python UDF. `builder` requires the two lists to + // be parallel and skips transpilation otherwise. + transpiledInputTypes: JList[JList[String]] = Nil.asJava, + // Schema of the intermediate aggregation buffer, set only for the incremental Python + // aggregator eval types (see [[PythonAggregate]]); `null` otherwise. Nullable rather than + // `Option` so it can be passed positionally from Python over Py4J. + bufferType: DataType = null) { def builder(e: Seq[Expression]): Expression = { if (pythonEvalType == PythonEvalType.SQL_BATCHED_UDF @@ -53,7 +70,8 @@ case class UserDefinedPythonFunction( || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF || pythonEvalType == PythonEvalType.SQL_SCALAR_ARROW_UDF || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF - || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF) { + || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF + || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF) { /* * Check if the named arguments: * - don't have duplicated names @@ -63,15 +81,100 @@ case class UserDefinedPythonFunction( } else if (e.exists(_.isInstanceOf[NamedArgumentExpression])) { throw QueryCompilationErrors.namedArgumentsNotSupported(name) } + val transpiledExprs: List[Expression] = transpiled.asScala.map( + column => ColumnConversions.expression(column)).toList + val optionInputTypes: List[List[String]] = + transpiledInputTypes.asScala.map(_.asScala.toList).toList + - if (pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF + val udfExpr = if (pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF) { PythonUDAF(name, func, dataType, e, udfDeterministic, pythonEvalType) + } else if (pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF) { + // The incremental Python aggregator. `bufferType` (the intermediate buffer schema) must have + // been supplied as a struct when the UDF was created. The single expression carries the + // aggregator for both the PARTIAL and FINAL stages; the physical operator picks the per-stage + // eval type. `udaf()` enforces this up front, but a malformed Connect proto (eval type set, + // `buffer_type` missing or non-struct) or a direct `UserDefinedFunction(f, evalType=...)` can + // reach here without it, so return a classed error rather than a bare require / + // ClassCastException. + val bufferStruct = bufferType match { + case s: StructType => s + case _ => + throw QueryCompilationErrors.invalidIncrementalPythonAggregatorBufferError( + name, bufferType) + } + PythonAggregate(name, func, dataType, e, udfDeterministic, bufferStruct) } else { PythonUDF(name, func, dataType, e, pythonEvalType, udfDeterministic) } + // The ``_udf_param_N`` substitution below is positional, so a UDF + // call site that supplied named arguments (e.g. SQL ``name => val`` + // or pyspark ``udf(b=col)``) would splice ``NamedArgumentExpression`` + // wrappers into the rewritten Catalyst tree and confuse downstream + // function resolution (``isnotnull`` etc. reject named parameters). + // The Python ``__call__`` shim resolves kwargs to positional before + // they reach this builder; SQL named arguments don't go through that + // shim, so we conservatively skip transpilation here when any child + // is a ``NamedArgumentExpression`` and let the regular Python UDF + // path execute. + val transpiledExprsForUse = + if (e.exists(_.isInstanceOf[NamedArgumentExpression])) Nil else transpiledExprs + val optionInputTypesForUse = + if (e.exists(_.isInstanceOf[NamedArgumentExpression])) Nil else optionInputTypes + // If we have possible transpiled expressions insert the node carrying every + // option plus its declared input-type categories. We can't pick here: this + // builder runs at call-construction time, before the argument columns are + // bound, so their types aren't known yet. ResolveTranspiledPythonUDFOptions + // prunes the options to those matching the resolved input types (once known, + // and before CheckAnalysis), and ConvertToCatalyst picks the survivor. + // Only build the node when every option carries its parallel input-type + // categories. ResolveTranspiledPythonUDFOptions prunes type-incompatible + // options using those categories, but only when they are present (its guard + // is `optionInputCategories.nonEmpty`); an empty or mismatched categories + // list would leave a type-invalid option to fail CheckAnalysis instead of + // falling back. If the two lists don't line up, skip transpilation. + // A call-site arity mismatch (user passed more or fewer args than the UDF's + // parameters) must fall back to the plain Python UDF so the standard runtime + // TypeError surfaces. Each option's category list has exactly one entry per + // public parameter, so a length mismatch against the bound children detects + // both directions: too few args would otherwise trip the placeholder bounds + // check below as a misleading "internal error", and too many args on a + // zero/fewer-param UDF would otherwise silently succeed where Python raises. + if (transpiledExprsForUse.nonEmpty && + optionInputTypesForUse.length == transpiledExprsForUse.length && + optionInputTypesForUse.forall(_.length == e.length)) { + val udfChildren = udfExpr.children.toArray + // Resolve the `_udf_param_N` placeholders the transpiler emits into the bound + // UDF arguments. Apply this ONLY to the transpiled options -- never to + // `udfExpr` itself, whose children are the user's argument expressions. A user + // column literally named `_udf_param_N` passed as an argument must not be + // rewritten, so we leave `udfExpr` untouched. + def resolveUDFParams(expression: Expression, children: Array[Expression]): Expression = { + expression match { + case UnresolvedAttribute(nameParts) + if nameParts.length == 1 && nameParts.head.startsWith("_udf_param_") => + val suffix = nameParts.head.stripPrefix("_udf_param_") + val index = suffix.toIntOption.getOrElse { + throw QueryCompilationErrors.invalidUDFParameterPlaceholder(nameParts.head) + } + if (index >= 0 && index < children.length) { + children(index) + } else { + throw QueryCompilationErrors.invalidUDFParameterPlaceholderIndex( + index, children.length) + } + case _ => + expression.mapChildren(resolveUDFParams(_, children)) + } + } + val resolvedOptions = transpiledExprsForUse.map(resolveUDFParams(_, udfChildren)) + TranspiledPythonUDF(name, udfExpr, resolvedOptions, optionInputTypesForUse) + } else { + udfExpr + } } def builderWithColumns(e: Seq[Column]): Expression = builder(e.map(expression)) @@ -86,7 +189,10 @@ case class UserDefinedPythonFunction( */ def fromUDFExpr(expr: Expression): Column = { Column(expr match { + case TranspiledPythonUDF(name, udaf: PythonUDAF, transpiled, inputCategories) => + TranspiledPythonUDF(name, udaf.toAggregateExpression(), transpiled, inputCategories) case udaf: PythonUDAF => udaf.toAggregateExpression() + case agg: PythonAggregate => agg.toAggregateExpression() case _ => expr }) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala index d3fd757784e0c..f3084be7722a2 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkExec.scala @@ -195,17 +195,7 @@ case class TransformWithStateInPySparkExec( groupingKeySchema, driverProcessorHandle ) - // runner initialization - runner.init() - try { - // execute UDF on the python runner - runner.process() - } catch { - case e: Throwable => - throw new SparkException("TransformWithStateInPySpark driver worker " + - "exited unexpectedly (crashed)", e) - } - runner.stop() + TransformWithStateInPySparkExec.runPreInitRunner(runner) val info = getStateInfo val stateSchemaDir = stateSchemaDirPath() @@ -448,6 +438,22 @@ case class TransformWithStateInPySparkExec( // scalastyle:off argcount object TransformWithStateInPySparkExec { + private[streaming] def runPreInitRunner( + runner: TransformWithStateInPySparkPythonPreInitRunner): Unit = { + Utils.tryWithSafeFinally { + runner.init() + try { + runner.process() + } catch { + case e: Throwable => + throw new SparkException("TransformWithStateInPySpark driver worker " + + "exited unexpectedly (crashed)", e) + } + } { + runner.stop() + } + } + // Plan logical transformWithStateInPySpark for batch queries def generateSparkPlanForBatchQueries( functionExpr: Expression, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala index 33fe0cdfee3f1..0b7e6b938689f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPythonRunner.scala @@ -333,8 +333,10 @@ class TransformWithStateInPySparkPythonPreInitRunner( override def stop(): Unit = { super.stop() + if (daemonThread != null) { + daemonThread.interrupt() + } closeServerSocketChannelSilently(stateServerSocket) - daemonThread.interrupt() } private def startStateServer(): Unit = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala index 4fee6a6e71d30..c2cea67061bcc 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServer.scala @@ -18,7 +18,13 @@ package org.apache.spark.sql.execution.python.streaming import java.io.{BufferedInputStream, BufferedOutputStream, DataInputStream, DataOutputStream, EOFException, InterruptedIOException} -import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel} +import java.nio.channels.{ + Channels, + ClosedByInterruptException, + ClosedChannelException, + ServerSocketChannel, + SocketChannel +} import java.time.Duration import scala.collection.mutable @@ -40,6 +46,7 @@ import org.apache.spark.sql.execution.streaming.state.StateMessage.KeyAndValuePa import org.apache.spark.sql.execution.streaming.state.StateMessage.StateResponseWithListGet import org.apache.spark.sql.streaming.{ListState, MapState, TTLConfig, ValueState} import org.apache.spark.sql.types.StructType +import org.apache.spark.util.Utils /** * This class is used to handle the state requests from the Python side. It runs on a separate @@ -138,8 +145,27 @@ class TransformWithStateInPySparkStateServer( } else new mutable.HashMap[String, Iterator[Long]]() def run(): Unit = { - val listeningSocket = stateServerSocket.accept() + val listeningSocket = try { + stateServerSocket.accept() + } catch { + case _: InterruptedException | _: InterruptedIOException | _: ClosedByInterruptException => + logInfo(log"State server listener interrupted before the Python worker connected") + Thread.currentThread().interrupt() + statefulProcessorHandle.setHandleState(StatefulProcessorHandleState.CLOSED) + return + case _: ClosedChannelException => + logInfo(log"State server socket closed before the Python worker connected") + statefulProcessorHandle.setHandleState(StatefulProcessorHandleState.CLOSED) + return + } + + // The task completion listener closes only the listening server socket, and the + // request loop has several early returns, so the accepted connection is closed + // through tryWithResource. + Utils.tryWithResource(listeningSocket)(serveRequests) + } + private def serveRequests(listeningSocket: SocketChannel): Unit = { // SPARK-51667: We have a pattern of sending messages continuously from one side // (Python -> JVM, and vice versa) before getting response from other side. Since most // messages we are sending are small, this triggers the bad combination of Nagle's algorithm diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/stat/StatFunctions.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/stat/StatFunctions.scala index e812af229524f..57b7bf7ac41be 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/stat/StatFunctions.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/stat/StatFunctions.scala @@ -81,7 +81,8 @@ object StatFunctions extends Logging { val accuracy = if (relativeError == 0.0) { Int.MaxValue } else { - math.min(Int.MaxValue, (1.0 / relativeError).ceil.toLong).toInt + val raw = (1.0 / relativeError).ceil.toLong + math.max(1, math.min(Int.MaxValue, raw)).toInt } val results = Array.fill(cols.size)(Seq.empty[Double]) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointVersionManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointVersionManager.scala index 4f41f7cc46cfc..7f7a614e9d33a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointVersionManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CheckpointVersionManager.scala @@ -35,6 +35,7 @@ case class StreamingCheckpointVersion(offsetLogVersion: Int) { sealed trait CheckpointLogType case object OffsetLogType extends CheckpointLogType +case object CommitLogType extends CheckpointLogType /** * The `CheckpointVersionManager` is responsible for managing the versioning of the streaming @@ -109,6 +110,35 @@ object CheckpointVersionManager extends Logging { logType: CheckpointLogType): Int = { logType match { case OffsetLogType => getOffsetLogVersion(sparkSessionForStream) + case CommitLogType => getCommitLogVersion(sparkSessionForStream) + } + } + + /** + * The commit log format version requested by the session config. + * `streamingCommitLogFormatVersion` tracks [[SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION]]: a + * state store checkpoint format of v2 makes each batch write `stateUniqueIds`, which only a + * commit log at [[CommitLog.VERSION_2]] or above can persist, so a v2 state store format raises + * the commit log version to v2. + */ + private def getCommitLogVersion(sparkSessionForStream: SparkSession): Int = { + val result = sparkSessionForStream.sessionState.conf.streamingCommitLogFormatVersion + logInfo(s"Retrieved commit log writer version=$result") + result + } + + /** + * Determines the commit log format version for ordinary writes in this query run. An existing + * checkpoint wins, so a session config change cannot start writing a format the checkpoint was + * not created with. Only a fresh checkpoint takes the version from the session config. Sink + * evolution may independently upgrade writes to [[CommitLog.VERSION_3]]. + */ + def resolveCommitLogVersion( + sparkSessionForStream: SparkSession, + latestCommittedBatch: Option[(Long, CommitMetadataBase)]): Int = { + latestCommittedBatch match { + case Some((_, commitMetadata)) => commitMetadata.version + case None => getFormatVersionFromSession(sparkSessionForStream, CommitLogType) } } @@ -142,10 +172,38 @@ object CheckpointVersionManager extends Logging { def setFormatVersion( sparkSessionForStream: SparkSession, logType: CheckpointLogType, - version: Int): Unit = { + version: Int, + commitMetadata: Option[CommitMetadataBase] = None): Unit = { logType match { case OffsetLogType => setSparkSessionConfigsForOffsetLog(sparkSessionForStream, version) + case CommitLogType => + setSparkSessionConfigsForCommitLog(sparkSessionForStream, version, commitMetadata) + } + } + + /** + * Records the state store checkpoint format used by the existing commit log. VERSION_1 uses + * state store format v1 and VERSION_2 uses v2. VERSION_3 can represent either format because it + * was introduced independently for sink metadata; the presence of state store checkpoint ids + * distinguishes v2 from v1. + * + * The state store format is clamped to 2 rather than set to the commit log version, because the + * two version spaces are separate: the commit log has a VERSION_3 (sink metadata) with no state + * store counterpart, while the state store config accepts only versions 1 and 2. Writing 3 here + * would report a state store format that does not exist. + */ + private def setSparkSessionConfigsForCommitLog( + sparkSessionForStream: SparkSession, + commitLogFormatVersion: Int, + commitMetadata: Option[CommitMetadataBase]): Unit = { + val stateStoreVersion = commitMetadata match { + case Some(metadata) if metadata.version == CommitLog.VERSION_3 => + if (metadata.stateUniqueIds.isDefined) 2 else 1 + case _ => + if (commitLogFormatVersion >= CommitLog.VERSION_2) 2 else 1 } + sparkSessionForStream.conf + .set(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key, stateStoreVersion.toString) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CommitLog.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CommitLog.scala index b5271f664cd76..59ad6b986b0ae 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CommitLog.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/CommitLog.scala @@ -28,7 +28,6 @@ import org.json4s.jackson.Serialization import org.apache.spark.sql.SparkSession import org.apache.spark.sql.connector.read.streaming.{Offset => OffsetV2} import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.internal.SQLConf /** * Used to write log files that represent batch commit points in structured streaming. @@ -56,11 +55,6 @@ class CommitLog( import CommitLog._ - // The configured commit log format version. Used as the default version when callers - // construct metadata through [[createMetadata]]. - private[sql] val defaultVersion: Int = sparkSession.conf.get( - SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key).toInt - override protected[sql] def deserialize(in: InputStream): CommitMetadataBase = { CommitLog.readCommitMetadata(in) } @@ -76,7 +70,10 @@ class CommitLog( /** * Factory for creating a [[CommitMetadataBase]] for the requested wire format version. - * Defaults to the version configured via [[SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION]]. + * + * [[commitLogFormatVersion]] is a required parameter rather than a field read from the session + * config, so that a caller always supplies the version it resolved for this query run. Reading + * the config here would ignore the format an existing checkpoint was created with. * * For [[VERSION_3]], [[sinkMetadataMap]] must be non-empty and contain exactly one active * sink; [[CommitMetadataV3]] enforces this invariant. @@ -85,7 +82,7 @@ class CommitLog( nextBatchWatermarkMs: Long = 0, stateUniqueIds: Option[Map[Long, Array[Array[String]]]] = None, sinkMetadataMap: Map[String, SinkMetadataInfo] = Map.empty, - commitLogFormatVersion: Int = defaultVersion): CommitMetadataBase = { + commitLogFormatVersion: Int): CommitMetadataBase = { commitLogFormatVersion match { case VERSION_3 => CommitMetadataV3(nextBatchWatermarkMs, stateUniqueIds, sinkMetadataMap) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala index 4c7a8437a46fd..aae0b7c7121c1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala @@ -26,7 +26,7 @@ import scala.collection.mutable.{Map => MutableMap} import org.apache.spark.SparkEnv import org.apache.spark.internal.LogKeys._ -import org.apache.spark.sql.catalyst.expressions.{CurrentDate, CurrentTimestampLike, LocalTimestamp} +import org.apache.spark.sql.catalyst.expressions.{CurrentDate, CurrentTimestampLike, LocalTimestamp, LocalTimestampNanos} import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.streaming.{StreamingRelationV2, WriteToStream} import org.apache.spark.sql.catalyst.trees.TreePattern.CURRENT_LIKE @@ -234,7 +234,8 @@ class ContinuousExecution( } withNewSources.transformAllExpressionsWithPruning(_.containsPattern(CURRENT_LIKE)) { - case (_: CurrentTimestampLike | _: CurrentDate | _: LocalTimestamp) => + case (_: CurrentTimestampLike | _: CurrentDate | _: LocalTimestamp | + _: LocalTimestampNanos) => throw new IllegalStateException("CurrentTimestamp, Now, CurrentDate and LocalTimestamp" + " not yet supported for continuous processing") } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala new file mode 100644 index 0000000000000..839e52bde8fbd --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/GenericBufferAggregationIterator.scala @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.streaming + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, MutableProjection, NamedExpression, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.execution.aggregate.AggregationIterator + +/** + * This base class extends [[AggregationIterator]] and produces a new aggregation buffer on demand. + * The instance of aggregation buffer is optimal for given aggregate functions. + * + * This base class is useful for cases where aggregation buffer needs to be initialized frequently, + * e.g. more than the cardinality of grouping keys. + */ +abstract class GenericBufferAggregationIterator( + partIndex: Int, + groupingExpressions: Seq[NamedExpression], + originalInputAttributes: Seq[Attribute], + aggregateExpressions: Seq[AggregateExpression], + aggregateAttributes: Seq[Attribute], + initialInputBufferOffset: Int, + resultExpressions: Seq[NamedExpression], + newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection) + extends AggregationIterator( + partIndex, + groupingExpressions, + originalInputAttributes, + aggregateExpressions, + aggregateAttributes, + initialInputBufferOffset, + resultExpressions, + newMutableProjection) { + + protected val useUnsafeBuffer = aggregateFunctions.flatMap(_.aggBufferAttributes) + .map(_.dataType).forall(UnsafeRow.isMutable) + + /** + * Returns an aggregation buffer containing initial buffer values. Each call will produce the + * different buffer instance. + */ + protected def newAggregationBuffer(): InternalRow = { + val buffer = initialAggregationBuffer.copy() + // if we are using a GenericInternalRow which + // is just a wrapper for an underlying data structured + // we need to re-initialize the buffer since + // copy does not actually create a new copy + // of the underlying data structure + if (!useUnsafeBuffer) { + initializeBuffer(buffer) + } + buffer + } + + // An aggregation buffer containing initial buffer values. It is used to + // initialize other aggregation buffers. + private val initialAggregationBuffer: InternalRow = createNewAggregationBuffer() + + private def createNewAggregationBuffer(): InternalRow = { + val bufferSchema = aggregateFunctions.flatMap(_.aggBufferAttributes) + val bufferRowSize: Int = bufferSchema.length + val genericMutableBuffer = new GenericInternalRow(bufferRowSize) + + val buffer = if (useUnsafeBuffer) { + val unsafeProjection = + UnsafeProjection.create(bufferSchema.map(_.dataType)) + val buf = unsafeProjection.apply(genericMutableBuffer) + initializeBuffer(buf) + buf + } else { + genericMutableBuffer + } + buffer + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala new file mode 100644 index 0000000000000..7046e4236aa0e --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/ProjectAggregationBufferExec.scala @@ -0,0 +1,121 @@ +/* + * 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.execution.streaming + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, MutableProjection, NamedExpression, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.metric.SQLMetrics + +/** + * This class handles the part of aggregation functions in the input rows, based on the function's + * mode. This class intends to either initialize the aggregation buffer or complete the aggregate + * buffer and produce the result, so it is expected to be used for two aggregate modes: + * 1) partial merge 2) final. This class is pass-through and does not perform the actual + * aggregation. + */ +case class ProjectAggregationBufferExec( + requiredChildDistributionExpressions: Option[Seq[Expression]] = None, + numShufflePartitions: Option[Int], + groupingExpressions: Seq[NamedExpression] = Nil, + aggregateExpressions: Seq[AggregateExpression] = Nil, + aggregateAttributes: Seq[Attribute] = Nil, + initialInputBufferOffset: Int = 0, + resultExpressions: Seq[NamedExpression] = Nil, + isFinalAggregate: Boolean, + child: SparkPlan) + extends BaseAggregateExec { + + override val isStreaming: Boolean = true + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + override protected def doExecute(): RDD[InternalRow] = { + metrics // force lazy initialization at driver + + val numOutputRows = longMetric("numOutputRows") + + child.execute().mapPartitionsWithIndex { case (partIdx, iter) => + val aggProcessor = new ProjectAggregationBufferProcessor( + partIdx, + groupingExpressions, + inputAttributes, + aggregateExpressions, + aggregateAttributes, + initialInputBufferOffset, + resultExpressions, + (expressions, inputSchema) => + MutableProjection.create(expressions, inputSchema)) + + iter.map { row => + numOutputRows += 1 + aggProcessor.process(row) + } + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) + + override def toSortAggregate: SortAggregateExec = { + throw new IllegalStateException("This class cannot be replaced with SortAggregate!") + } +} + +/** + * This class is an implementation of GenericBufferAggregationIterator which only handles the + * aggregation buffer of input, depending on the aggregate mode. This class is pass-through + * and does not perform the actual aggregation. + */ +class ProjectAggregationBufferProcessor( + partIndex: Int, + groupingExpressions: Seq[NamedExpression], + originalInputAttributes: Seq[Attribute], + aggregateExpressions: Seq[AggregateExpression], + aggregateAttributes: Seq[Attribute], + initialInputBufferOffset: Int, + resultExpressions: Seq[NamedExpression], + newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection) + extends GenericBufferAggregationIterator( + partIndex, + groupingExpressions, + originalInputAttributes, + aggregateExpressions, + aggregateAttributes, + initialInputBufferOffset, + resultExpressions, + newMutableProjection) { + + def hasNext: Boolean = + throw new UnsupportedOperationException( + "hasNext is not supported in ProjectAggregationBufferProcessor") + + def next(): UnsafeRow = + throw new UnsupportedOperationException( + "next is not supported in ProjectAggregationBufferProcessor") + + def process(newInput: InternalRow): UnsafeRow = { + val groupingKey = groupingProjection.apply(newInput) + val buffer = newAggregationBuffer() + processRow(buffer, newInput) + generateOutput(groupingKey, buffer) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala new file mode 100644 index 0000000000000..18d7d1a68b5fc --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StatefulStreamlineAggregateExec.scala @@ -0,0 +1,494 @@ +/* + * 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.execution.streaming + +import java.util.concurrent.TimeUnit.NANOSECONDS + +import scala.util.control.NonFatal + +import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache, RemovalNotification} +import org.apache.hadoop.conf.Configuration + +import org.apache.spark.TaskContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.analysis.WidenStatefulOpNullability +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, MutableProjection, NamedExpression, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.streaming.InternalOutputModes.{Append, Complete, Update} +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, SortAggregateExec} +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorCustomMetric, StatefulOperatorCustomSumMetric, StatefulOperatorStateInfo, StatefulOperatorsUtils, StateStoreWriter, StreamingAggregationStateManager, WatermarkSupport} +import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, StateSchemaCompatibilityChecker, StateSchemaValidationResult, StateStore, StateStoreColFamilySchema, StateStoreOps} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.OutputMode +import org.apache.spark.sql.types.StructType +import org.apache.spark.util.CompletionIterator + +/** + * The physical plan of streaming aggregation which "streamlines" the process of aggregation. + * (Here the term "streamline" represents the loop of "read-process-output" for each input.) + * + * Refer to the classdoc of [[StatefulStreamlineAggregationProcessor]] for more details. + * + * For producing result table as output, it follows the semantic of output mode. + * + * - Append mode: accumulated result for the grouping key will be produced once the watermark + * passes and there will be no further update against the grouping key. + * - Update mode: each input will produce intermediate accumulated result as an output. + * Note that this is different from streaming aggregation in microbatch mode + * ([[StateStoreSaveExec]]) which produces the final intermediate accumulated result for + * each grouping key only in this microbatch. + * - Complete mode: the entire result table will be produced per each microbatch. + */ +case class StatefulStreamlineAggregateExec( + requiredChildDistributionExpressions: Option[Seq[Expression]], + numShufflePartitions: Option[Int], + groupingExpressions: Seq[NamedExpression], + aggregateExpressions: Seq[AggregateExpression], + aggregateAttributes: Seq[Attribute], + initialInputBufferOffset: Int, + resultExpressions: Seq[NamedExpression], + isFinalAggregate: Boolean, + outputMode: Option[OutputMode] = None, + stateFormatVersion: Int, + child: SparkPlan, + stateInfo: Option[StatefulOperatorStateInfo] = None, + eventTimeWatermarkForLateEvents: Option[Long] = None, + eventTimeWatermarkForEviction: Option[Long] = None) + extends BaseAggregateExec with StateStoreWriter with WatermarkSupport { + + override val isStreaming: Boolean = true + + override def shortName: String = + StatefulOperatorsUtils.STATEFUL_STREAMLINE_AGGREGATE_EXEC_OP_NAME + + override def keyExpressions: Seq[Attribute] = groupingExpressions.map(_.toAttribute) + + // SPARK-57003 component (b): widen StatefulStreamlineAggregateExec output. + override def output: Seq[Attribute] = + WidenStatefulOpNullability.widenOutputForStatefulOp( + resultExpressions.map(_.toAttribute)) + + override def customStatefulOperatorMetrics: Seq[StatefulOperatorCustomMetric] = { + Seq( + StatefulOperatorCustomSumMetric( + "numRowsReadDuringEviction", "number of state rows read during state eviction" + ), + StatefulOperatorCustomSumMetric( + "numRowsIncrementallyRemoved", "number of state rows removed during incremental eviction" + ) + ) + } + + private[sql] val stateManager = StreamingAggregationStateManager.createStateManager( + keyExpressions, child.output, stateFormatVersion) + + // SPARK-57003 component (a): widen state schemas to nullable at construction. Both the + // schema-check site (`validateAndMaybeEvolveStateSchema`) and the runtime + // `mapPartitionsWithStateStore` site read these. + private val stateKeySchema: StructType = + WidenStatefulOpNullability.widenStateSchema(keyExpressions.toStructType) + private val stateValueSchema: StructType = + WidenStatefulOpNullability.widenStateSchema(stateManager.getStateValueSchema) + + private val incrementalCleanupFactor = session.sessionState.conf.getConf( + SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR) + + private def doIncrementalCleanup = incrementalCleanupFactor > 0 + + override def validateAndMaybeEvolveStateSchema( + hadoopConf: Configuration, batchId: Long, stateSchemaVersion: Int): + List[StateSchemaValidationResult] = { + val newStateSchema = List(StateStoreColFamilySchema(StateStore.DEFAULT_COL_FAMILY_NAME, + 0, stateKeySchema, 0, stateValueSchema)) + List(StateSchemaCompatibilityChecker.validateAndMaybeEvolveStateSchema(getStateInfo, + hadoopConf, newStateSchema, session.sessionState, stateSchemaVersion)) + } + + override protected def doExecute(): RDD[InternalRow] = { + metrics // force lazy init at driver + + val numOutputRows = longMetric("numOutputRows") + val numUpdatedStateRows = longMetric("numUpdatedStateRows") + val allUpdatesTimeMs = longMetric("allUpdatesTimeMs") + val numRowsReadDuringEviction = longMetric("numRowsReadDuringEviction") + val numRemovedStateRows = longMetric("numRemovedStateRows") + val allRemovalsTimeMs = longMetric("allRemovalsTimeMs") + val commitTimeMs = longMetric("commitTimeMs") + val numRowsIncrementallyRemoved = longMetric("numRowsIncrementallyRemoved") + + assert(outputMode.nonEmpty, + "Incorrect planning in IncrementalExecution, outputMode has not been set") + + child.execute().mapPartitionsWithStateStore( + getStateInfo, + stateKeySchema, + stateValueSchema, + NoPrefixKeyStateEncoderSpec(stateKeySchema), + session.sessionState, + Some(session.streams.stateStoreCoordinator)) { (store, iter) => + + // It's feasible to overload the method to provide a partition index, but now it's too + // many... + val partIdx = TaskContext.get().partitionId() + + // Filter late date using watermark if specified + val baseIterator = watermarkPredicateForDataForLateEvents match { + case Some(predicate) => applyRemovingRowsOlderThanWatermark(iter, predicate) + case None => iter + } + + val aggProcessor = new StatefulStreamlineAggregationProcessor( + partIdx, + groupingExpressions, + inputAttributes, + aggregateExpressions, + aggregateAttributes, + initialInputBufferOffset, + resultExpressions, + (expressions, inputSchema) => + MutableProjection.create(expressions, inputSchema), + stateManager, + store, + numUpdatedStateRows) + + // Each output row from aggIter references the same underlying row. + // It is the caller's responsibility to ensure that each row is consumed before the next + // one is produced. + // + // flushDirtyWrites is aggIter's completion action, so it runs when aggIter is exhausted and + // BEFORE store.commit() in every mode below: Complete drains aggIter (line ~191) ahead of the + // commit iterator; Append drains it (~222) ahead of its own; Update chains it inside the + // outer CompletionIterator whose completion commits, and draining the outer drains aggIter + // first. This ordering is what lets flushDirtyWrites surface a state-write failure (see its + // rethrow) in time to fail the task before the batch commits -- a partial write can never be + // committed. + var tmpRow: UnsafeRow = null + val aggIter = CompletionIterator[UnsafeRow, Iterator[UnsafeRow]]( + baseIterator.map { row => + allUpdatesTimeMs += timeTakenMs { + tmpRow = aggProcessor.process(row) + } + tmpRow + }, + aggProcessor.flushDirtyWrites() // Lazily evaluated + ) + + // Remaining logic performs the same thing with StateStoreSaveExec. It's mostly duplicated, + // with slight modification to cover the usage of aggregation iterator. + + outputMode match { + // Update and output all rows in the StateStore. + case Some(Complete) => + // consume iterator fully to process all inputs and save the result into state store. + aggIter.foreach(_ => ()) + + // SPARK-45582 - Ensure that store instance is not used after commit is called + // to invoke the iterator. + val rangeIter = stateManager.values(store) + + CompletionIterator[UnsafeRow, Iterator[UnsafeRow]]( + rangeIter.map { valueRow => + numOutputRows += 1 + valueRow + }, { + allRemovalsTimeMs += 0 + commitTimeMs += timeTakenMs { + store.commit() + } + setStoreMetrics(store) + setOperatorMetrics() + } + ) + + // Update and output only rows being evicted from the StateStore + // Assumption: watermark predicates must be non-empty if append mode is allowed + case Some(Append) => + assert(watermarkPredicateForDataForLateEvents.isDefined, + "Watermark needs to be defined for streaming aggregation query in append mode") + + assert(watermarkPredicateForKeysForEviction.isDefined, + "Watermark needs to be defined for streaming aggregation query in append mode") + + allUpdatesTimeMs += timeTakenMs { + // consume iterator fully to process all inputs and save the result into state store. + while (aggIter.hasNext) { + aggIter.next() + } + } + + val removalStartTimeNs = System.nanoTime + val evictionIterator = + stateManager.evictionIterator(store, eventTimeWatermarkForEviction) + + CompletionIterator[UnsafeRow, Iterator[UnsafeRow]]( + evictionIterator.map(_.value), { + numRowsReadDuringEviction += evictionIterator.numRowsReadDuringEvictionSoFar + numRemovedStateRows += evictionIterator.numRowsRemovedSoFar + numOutputRows += evictionIterator.numRowsRemovedSoFar + + // Note: Due to the iterator lazy exec, this metric also captures the time taken + // by the consumer operators in addition to the processing in this operator. + allRemovalsTimeMs += NANOSECONDS.toMillis(System.nanoTime - removalStartTimeNs) + commitTimeMs += timeTakenMs { + store.commit() + } + setStoreMetrics(store) + setOperatorMetrics() + }) + + // Update and output modified rows from the StateStore. + case Some(Update) => + /** + * When doing incremental cleanup, we have to be careful what watermark to use. Because + * the late events timestamp is less than the eviction timestamp, within a batch it is + * possible for us to receive events whose timestamps is less than the eviction + * timestamp. Thus, it is possible to evict a record at timestamp t, such that + * t < evictionTimestamp, and, within the same batch, receive a record at timestamp t. + * + * Thus, when using incremental eviction, we have to make sure to clean up records + * up to the timestamp before which we will _never_ receive new records. This would be + * the event time watermark for late events. + */ + val incrementalAwareEvictionWatermark = if (doIncrementalCleanup) { + eventTimeWatermarkForLateEvents + } else { + eventTimeWatermarkForEviction + } + + // Only create it if we are doing incremental cleanup. If we instantiate it and are + // not doing incremental cleanup, then the iterator will not iterate through records + // inserted into the store during the batch. + val incrementalEvictionIter = if (doIncrementalCleanup) { + Some(stateManager.evictionIterator(store, incrementalAwareEvictionWatermark)) + } else { + None + } + + val updateIter = aggIter.map { row => + incrementalEvictionIter.foreach { evictionIter => + allRemovalsTimeMs += timeTakenMs { + var numRemovalsCurrRecord = 0 + // NOTE: EvictionIterator removes (and counts) a row inside hasNext, not next(), so + // a bare hasNext already deletes and bumps the metrics. To stop at exactly + // incrementalCleanupFactor removals per input we must re-check the count before + // each hasNext rather than draining the iterator. + while (numRemovalsCurrRecord < incrementalCleanupFactor + && evictionIter.hasNext) { + // The removal happens inside of the iterator; in Update mode, + // we don't need the result. + evictionIter.next() + numRemovalsCurrRecord += 1 + } + } + } + numOutputRows += 1 + row + } + + CompletionIterator[UnsafeRow, Iterator[UnsafeRow]](updateIter, { + // Anything removed so far must have been part of incremental eviction + incrementalEvictionIter.foreach { iter => + numRowsIncrementallyRemoved += iter.numRowsRemovedSoFar + } + + // If the incremental eviction iterator is defined, we'll finish eviction here + // if any records remain. If it's not, we'll construct an eviction iterator + // and also use it up here. + val evictionIter = incrementalEvictionIter.getOrElse( + stateManager.evictionIterator(store, eventTimeWatermarkForEviction)) + + allRemovalsTimeMs += timeTakenMs { + while (evictionIter.hasNext) { + evictionIter.next() + } + } + + numRowsReadDuringEviction += evictionIter.numRowsReadDuringEvictionSoFar + numRemovedStateRows += evictionIter.numRowsRemovedSoFar + + commitTimeMs += timeTakenMs { + store.commit() + } + setStoreMetrics(store) + setOperatorMetrics() + }) + + case _ => throw QueryExecutionErrors.unsupportedOutputModeForStreamingOperationError( + outputMode.get, "streaming aggregations") + } + } + } + + override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = { + (outputMode.contains(Append) || outputMode.contains(Update)) && + eventTimeWatermarkForEviction.isDefined && + newInputWatermark > eventTimeWatermarkForEviction.get + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = { + copy(child = newChild) + } + + // FIXME: How we can prevent this to be called? + override def toSortAggregate: SortAggregateExec = { + throw new IllegalStateException("This class cannot be replaced with SortAggregate!") + } +} + +/** + * This class is an implementation of GenericBufferAggregationIterator which performs the + * aggregation against state store instead of maintaining aggregation hash table. + * + * For each input, do the following + * - Read the previous value for grouping key in state store + * - Merge the input and previous value (if any) + * - Produce the merged result + * - Store the new value to the dirty writes. + * + * This class maintains dirty writes which represent the cache of state store, performing both + * caching and deferred writes. Note that there is a size limit of dirty writes which should be + * considered carefully in both 1) memory usage and 2) latency spike on flushing writes. + * + * TODO: dirty writes can be implemented via LRU, which will optimize the ability of cache and + * also less spike on flushing writes (as evicted entry would be small portion of LRU). + */ +class StatefulStreamlineAggregationProcessor( + partIndex: Int, + groupingExpressions: Seq[NamedExpression], + originalInputAttributes: Seq[Attribute], + aggregateExpressions: Seq[AggregateExpression], + aggregateAttributes: Seq[Attribute], + initialInputBufferOffset: Int, + resultExpressions: Seq[NamedExpression], + newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection, + stateManager: StreamingAggregationStateManager, + stateStore: StateStore, + numUpdatedStateRows: SQLMetric) + extends GenericBufferAggregationIterator( + partIndex, + groupingExpressions, + originalInputAttributes, + aggregateExpressions, + aggregateAttributes, + initialInputBufferOffset, + resultExpressions, + newMutableProjection) { + + // The value for unsafe buffer is just an conservative arbitrary value - it should be probably + // safer to increase the value. + // The value for safe buffer is picked from the default value of fallback threshold of object + // hash aggregate, 128. + // That said, it's conservative, but should we still make this be configurable? + // NOTE: We need to consider the latency spike on flushing, so the number should not be too high + // even though there is more available memory to cache more entries. + protected val flushThresholdNumKeys: Int = if (useUnsafeBuffer) 1000 else 100 + + // Holds the first failure thrown while a removal listener writes an entry to the state store. + // Guava logs and SWALLOWS any exception a removal listener throws (see the + // CacheBuilder.removalListener scaladoc), so the put below cannot fail the task on its own -- + // without this capture a failed write would be silently dropped and the batch would still commit, + // losing that key's update. flushDirtyWrites rethrows it so the task fails instead of committing + // partial state. + private var writeFailure: Option[Throwable] = None + + // The writes which did not go into state store yet. We also leverage this dirty writes to the + // cache of state store. + private val dirtyWrites: LoadingCache[UnsafeRow, UnsafeRowReference] = CacheBuilder.newBuilder() + .maximumSize(flushThresholdNumKeys) + .removalListener((notification: RemovalNotification[UnsafeRow, UnsafeRowReference]) => { + // Skip once a write has already failed: the task is going to fail anyway, and further puts + // into a broken store are pointless. + if (writeFailure.isEmpty) { + try { + val value = notification.getValue + assert(value.getRow != null, "dirty writes should contain the valid row to update, " + + "but found null.") + stateManager.put(stateStore, value.getRow) + numUpdatedStateRows += 1 + } catch { + case NonFatal(e) => writeFailure = Some(e) + } + } + }) + .build(new CacheLoader[UnsafeRow, UnsafeRowReference] { + override def load(key: UnsafeRow): UnsafeRowReference = { + val newRef = new UnsafeRowReference + val existingValueInState = stateManager.get(stateStore, key) + // NOTE: newRef.getRow could be still null after this line + newRef.putRow(existingValueInState) + newRef + } + }) + + def hasNext: Boolean = + throw new UnsupportedOperationException( + "hasNext is not supported for StatefulStreamlineAggregationProcessor") + + def next(): UnsafeRow = + throw new UnsupportedOperationException( + "next is not supported for StatefulStreamlineAggregationProcessor") + + /** + * Flush the dirty writes to state store. This should be called at the end of each microbatch. + * + * Rethrows the first state-store write failure a removal listener captured, so that a failed put + * fails the task rather than committing state that is missing an update. + */ + def flushDirtyWrites(): Unit = { + dirtyWrites.invalidateAll() + writeFailure.foreach { e => + throw new IllegalStateException( + "Failed to write an aggregation update to the state store", e) + } + } + + def process(newInput: InternalRow): UnsafeRow = { + // groupingProjection hands back the same UnsafeRow on every call, overwriting its bytes, so the + // key has to be copied before the cache can store it. Without the copy, one shared row becomes + // the key of every entry: each entry is then only findable through the hash captured when it + // was inserted, and the equality check that should discriminate between keys degenerates to a + // reference match against that one row. + val groupingKey = groupingProjection.apply(newInput).copy() + val buffer = newAggregationBuffer() + + val existingValueRef = dirtyWrites.get(groupingKey) + + if (existingValueRef.getRow != null) { + processRow(buffer, existingValueRef.getRow) + } + + processRow(buffer, newInput) + + val output = generateOutput(groupingKey, buffer) + existingValueRef.putRow(output.copy()) + output + } +} + +class UnsafeRowReference { + private var row: UnsafeRow = _ + + def getRow: UnsafeRow = row + + def putRow(r: UnsafeRow): Unit = { + row = r + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala index 8357053cdc46c..63dab2c392c0f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/StreamingAggregationStateManager.scala @@ -21,7 +21,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.expressions.{Attribute, UnsafeRow} import org.apache.spark.sql.catalyst.expressions.codegen.{GenerateUnsafeProjection, GenerateUnsafeRowJoiner} import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.execution.streaming.state.{NoopStatePartitionKeyExtractor, ReadStateStore, StateStore, UnsafeRowPair} +import org.apache.spark.sql.execution.streaming.state.{EvictionIterator, NoopStatePartitionKeyExtractor, ReadStateStore, StateStore, UnsafeRowPair} import org.apache.spark.sql.types.StructType /** @@ -61,6 +61,15 @@ sealed trait StreamingAggregationStateManager extends Serializable { /** Return an iterator containing all the values in target state store. */ def values(store: ReadStateStore): Iterator[UnsafeRow] + + /** + * Return an iterator over the rows of the target state store whose event time is older than + * `evictionTimestamp`, removing each row as it is returned. See [[EvictionIterator]]. + */ + def evictionIterator( + store: StateStore, + evictionTimestamp: Option[Long], + allowMultipleEventTimeColumns: Boolean = false): EvictionIterator } object StreamingAggregationStateManager extends Logging { @@ -96,6 +105,18 @@ abstract class StreamingAggregationStateManagerBaseImpl( // discard and don't convert values to avoid computation store.iterator().map(_.key) } + + override def evictionIterator( + store: StateStore, + evictionTimestamp: Option[Long], + allowMultipleEventTimeColumns: Boolean = false): EvictionIterator = { + EvictionIterator( + store, + iterator(store), + keyExpressions, + allowMultipleEventTimeColumns, + evictionTimestamp) + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala index 022fa3469eea5..799f6326f7fd3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/statefulOperators.scala @@ -631,18 +631,6 @@ trait WatermarkSupport extends SparkPlan { watermarkExpression.map(Predicate.create(_, child.output)) } - protected def removeKeysOlderThanWatermark(store: StateStore): Unit = { - if (watermarkPredicateForKeysForEviction.nonEmpty) { - val numRemovedStateRows = longMetric("numRemovedStateRows") - store.iterator().foreach { rowPair => - if (watermarkPredicateForKeysForEviction.get.eval(rowPair.key)) { - store.remove(rowPair.key) - numRemovedStateRows += 1 - } - } - } - } - protected def removeKeysOlderThanWatermark( storeManager: StreamingAggregationStateManager, store: StateStore): Unit = { @@ -1385,6 +1373,11 @@ abstract class BaseStreamingDeduplicateExec protected lazy val stateValueSchema: StructType = WidenStatefulOpNullability.widenStateSchema(schemaForValueRow) + // Initialize it to 0 so that operators explicitly have to opt-in (i.e. set it to non-zero) to + // enable incremental cleanup. + protected val incrementalCleanupFactor: Long = 0 + protected def doIncrementalCleanup: Boolean = incrementalCleanupFactor > 0 + override protected def doExecute(): RDD[InternalRow] = { metrics // force lazy init at driver @@ -1406,6 +1399,7 @@ abstract class BaseStreamingDeduplicateExec val allRemovalsTimeMs = longMetric("allRemovalsTimeMs") val commitTimeMs = longMetric("commitTimeMs") val numDroppedDuplicateRows = longMetric("numDroppedDuplicateRows") + val numRowsIncrementallyRemoved = longMetric("numRowsIncrementallyRemoved") val baseIterator = watermarkPredicateForDataForLateEvents match { case Some(predicate) => applyRemovingRowsOlderThanWatermark(iter, predicate) @@ -1416,10 +1410,37 @@ abstract class BaseStreamingDeduplicateExec val updatesStartTimeNs = System.nanoTime + // Only create the eviction iterator if we are doing incremental cleanup. It is opened over + // the store before the batch processes any rows, so it may not observe records inserted into + // the store later in the batch (the guarantee depends on the store provider's iterator + // semantics). When cleanup is disabled we defer building it to batch end, after all inserts. + val incrementalEvictionIter = if (doIncrementalCleanup) { + Some(iteratorForEviction(store)) + } else { + None + } + val result = baseIterator.filter { r => val row = r.asInstanceOf[UnsafeRow] val key = getKey(row) val keyExists = store.keyExists(key) + + incrementalEvictionIter.foreach { evictionIter => + allRemovalsTimeMs += timeTakenMs { + // Remove up to incrementalCleanupFactor eligible rows for every input record. + var numRemovalsCurrRecord = 0 + // NOTE: the eviction iterator removes a row inside hasNext, not next(), so a bare + // hasNext already deletes it. To stop at exactly incrementalCleanupFactor removals per + // input we must re-check the count before each hasNext rather than draining the + // iterator. + while (numRemovalsCurrRecord < incrementalCleanupFactor && evictionIter.hasNext) { + evictionIter.next() + numRemovalsCurrRecord += 1 + numRowsIncrementallyRemoved += 1 + } + } + } + if (!keyExists) { putDupInfoIntoState(store, row, key, reusedDupInfoRow) numUpdatedStateRows += 1 @@ -1434,7 +1455,17 @@ abstract class BaseStreamingDeduplicateExec CompletionIterator[InternalRow, Iterator[InternalRow]](result, { allUpdatesTimeMs += NANOSECONDS.toMillis(System.nanoTime - updatesStartTimeNs) - allRemovalsTimeMs += timeTakenMs { evictDupInfoFromState(store) } + + // Finish eviction: if incremental cleanup ran, drain whatever eligible rows it did not get + // to during record processing; otherwise (incremental cleanup disabled) build the eviction + // iterator now and drain it fully -- the batch-end-only behavior. + allRemovalsTimeMs += timeTakenMs { + val evictionIter = incrementalEvictionIter.getOrElse(iteratorForEviction(store)) + while (evictionIter.hasNext) { + evictionIter.next() + } + } + commitTimeMs += timeTakenMs { store.commit() } setStoreMetrics(store) setOperatorMetrics() @@ -1450,15 +1481,46 @@ abstract class BaseStreamingDeduplicateExec key: UnsafeRow, reusedDupInfoRow: Option[UnsafeRow]): Unit - protected def evictDupInfoFromState(store: StateStore): Unit + /** + * Creates an iterator that removes the state rows eligible for eviction (older than the eviction + * watermark), updating the eviction metrics as it goes. Each `next()` corresponds to an actual + * removal, so a caller can stop early -- after `incrementalCleanupFactor` removals per input row + * -- and the store stays consistent with what was surfaced. + * + * Note the returned iterator is not necessarily an [[EvictionIterator]]: that helper reads the + * eviction timestamp from the state store key, which is not how every subclass stores it (e.g. + * [[StreamingDeduplicateWithinWatermarkExec]] keeps `expiresAtMicros` in the value row). The + * contents are unused by the caller, hence `Iterator[Any]`. + */ + protected def iteratorForEviction(store: StateStore): Iterator[Any] override def output: Seq[Attribute] = WidenStatefulOpNullability.widenOutputForStatefulOp(child.output) override def outputPartitioning: Partitioning = child.outputPartitioning + /** + * Three eviction metrics are exposed: + * + * - numRemovedStateRows: the total number of state rows removed. + * - numRowsIncrementallyRemoved: the subset of those removed during incremental eviction (i.e. + * spread across input-record processing rather than at batch end). + * - numRowsReadDuringEviction: the number of state rows the eviction iterator scanned, whether + * or not they were removed. + * + * Since incrementally removed rows are a subset of all removed rows, and every removed row was + * read, the relationship is: + * + * numRowsReadDuringEviction >= numRemovedStateRows >= numRowsIncrementallyRemoved + */ override def customStatefulOperatorMetrics: Seq[StatefulOperatorCustomMetric] = { - Seq(StatefulOperatorCustomSumMetric("numDroppedDuplicateRows", "number of duplicates dropped")) + Seq( + StatefulOperatorCustomSumMetric("numDroppedDuplicateRows", "number of duplicates dropped"), + StatefulOperatorCustomSumMetric( + "numRowsReadDuringEviction", "number of state rows read during state eviction"), + StatefulOperatorCustomSumMetric( + "numRowsIncrementallyRemoved", "number of state rows removed during incremental eviction") + ) } override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = { @@ -1484,6 +1546,12 @@ case class StreamingDeduplicateExec( protected val extraOptionOnStateStore: Map[String, String] = Map(StateStoreConf.FORMAT_VALIDATION_CHECK_VALUE_CONFIG -> "false") + // Read via the null-safe `conf` accessor rather than `session.sessionState.conf`: on a + // session-less thread (e.g. during canonicalization) `session` is null, and `conf` falls back to + // SparkPlan's active/default conf instead of throwing. + override protected val incrementalCleanupFactor: Long = + conf.getConf(SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR) + protected def initializeReusedDupInfoRow(): Option[UnsafeRow] = None protected def putDupInfoIntoState( @@ -1494,8 +1562,45 @@ case class StreamingDeduplicateExec( store.put(key, StreamingDeduplicateExec.EMPTY_ROW) } - protected def evictDupInfoFromState(store: StateStore): Unit = { - removeKeysOlderThanWatermark(store) + override protected def iteratorForEviction(store: StateStore): Iterator[Any] = { + val numRemovedStateRows = longMetric("numRemovedStateRows") + val numRowsReadDuringEviction = longMetric("numRowsReadDuringEviction") + + // When doing incremental cleanup, we have to be careful which watermark to use. The late-events + // watermark is less than or equal to the eviction watermark, so within a batch it is possible + // to receive an event whose timestamp is below the eviction watermark. Evicting a key at + // timestamp t and then, in the same batch, receiving another record at t would let that record + // through as if it were new. Thus, with incremental eviction we may only clean up records up to + // the timestamp before which we will never receive new records -- the event time watermark for + // late events. Without incremental cleanup, all input has been processed by the time we evict, + // so the eviction watermark is safe. + val incrementalAwareEvictionWatermark = if (doIncrementalCleanup) { + eventTimeWatermarkForLateEvents + } else { + eventTimeWatermarkForEviction + } + + val evictionIterator = EvictionIterator( + store, + store.iterator(), + keyExpressions, + // The previous full-eviction predicate (watermarkPredicateForKeysForEviction) compiled + // against `keyExpressions` and only produced a predicate when `keyExpressions` carried an + // event-time column -- so resolving the event-time column from `keyExpressions` here is + // equivalent for the dedup key, and matches the aggregation eviction path and the runtime. + // The `allowMultipleStatefulOperators` knob is carried over unchanged, hence + // !allowMultipleStatefulOperators. (The old path found the event-time *column* from + // child.output, which could throw MULTIPLE_EVENT_TIME_COLUMNS when child.output -- not the + // dedup key -- carried several event-time columns; resolving from the key here does not, but + // that only differs under the non-default allowMultiple=false with 2+ event-time columns + // outside the dedup key.) + allowMultipleEventTimeColumns = !allowMultipleStatefulOperators, + incrementalAwareEvictionWatermark) + + CompletionIterator[UnsafeRowPair, Iterator[UnsafeRowPair]](evictionIterator, { + numRowsReadDuringEviction += evictionIterator.numRowsReadDuringEvictionSoFar + numRemovedStateRows += evictionIterator.numRowsRemovedSoFar + }) } override def shortName: String = StatefulOperatorsUtils.DEDUPLICATE_EXEC_OP_NAME @@ -1538,6 +1643,17 @@ case class StreamingDeduplicateWithinWatermarkExec( protected val extraOptionOnStateStore: Map[String, String] = Map.empty + // Incremental cleanup is deliberately NOT enabled for dropDuplicatesWithinWatermark. Its dedup + // key is only the user's columns; the expiry (`expiresAtMicros`) lives in the value row and is + // decoupled from a record's event time. So two records that share a dedup key can have different + // event times, and one can be non-late while the other's stored entry has already expired. + // Evicting that entry mid-batch (as incremental cleanup would) makes a later non-late record with + // the same key miss the existing entry and be emitted as new -- an output that depends on the + // cleanup factor and store iteration order. Evicting only at batch end (the inherited factor-0 + // path) keeps the entry visible to every record in the batch, so the output is deterministic. + // (StreamingDeduplicateExec is immune because its key includes the event time: any record sharing + // an evictable key is itself late and is dropped by the late-events filter first.) + // Below three variables are defined as lazy, as evaluating these variables does not work with // canonicalized plan. Specifically, attributes in child won't have an event time column in // the canonicalized plan. These variables are NOT referenced in canonicalized plan, hence @@ -1571,18 +1687,29 @@ case class StreamingDeduplicateWithinWatermarkExec( store.put(key, timeoutRow) } - protected def evictDupInfoFromState(store: StateStore): Unit = { + override protected def iteratorForEviction(store: StateStore): Iterator[Any] = { val numRemovedStateRows = longMetric("numRemovedStateRows") + val numRowsReadDuringEviction = longMetric("numRowsReadDuringEviction") - // Convert watermark value to micros. + // Incremental cleanup is not enabled for this operator (this class does not override the base + // incrementalCleanupFactor, which defaults to 0 -- see the class-level comment for why), so + // eviction always runs once at batch end against the eviction watermark. val watermarkForEviction = DateTimeUtils.millisToMicros(eventTimeWatermarkForEviction.get) - store.iterator().foreach { rowPair => - val valueRow = rowPair.value - val expiresAt = valueRow.getLong(0) + // We cannot reuse `EvictionIterator` here because it reads the eviction timestamp from the + // state store key, whereas this operator stores `expiresAtMicros` in the value row (the key is + // only the dedup key columns). We still match EvictionIterator's semantics by returning only + // the rows that are actually evicted, so each next() corresponds to a real removal and the + // caller's numRemovedStateRows metric counts removals rather than state rows scanned. + store.iterator().filter { rowPair => + numRowsReadDuringEviction += 1 + val expiresAt = rowPair.value.getLong(0) if (watermarkForEviction >= expiresAt) { store.remove(rowPair.key) numRemovedStateRows += 1 + true + } else { + false } } } @@ -1626,7 +1753,8 @@ trait SchemaValidationUtils extends Logging { stateSchemaDir: Path, session: SparkSession, operatorStateMetadataVersion: Int = 2, - stateStoreEncodingFormat: String = StateStoreEncoding.UnsafeRow.toString + stateStoreEncodingFormat: String = StateStoreEncoding.UnsafeRow.toString, + isRealTimeMode: Boolean = false ): List[StateSchemaValidationResult] = { assert(stateSchemaVersion >= 3) val usingAvro = stateStoreEncodingFormat == StateStoreEncoding.Avro.toString @@ -1634,8 +1762,11 @@ trait SchemaValidationUtils extends Logging { val newStateSchemaFilePath = new Path(stateSchemaDir, s"${batchId}_${UUID.randomUUID().toString}") val metadataPath = new Path(info.checkpointLocation, s"${info.operatorId}") + // RTM can leave metadata for an uncommitted attempt of the current batch, so only metadata + // from the previous committed batch is safe to use for schema validation. + val metadataBatchId = if (isRealTimeMode) batchId - 1 else batchId val metadataReader = OperatorStateMetadataReader.createReader( - metadataPath, hadoopConf, operatorStateMetadataVersion, batchId) + metadataPath, hadoopConf, operatorStateMetadataVersion, metadataBatchId) val operatorStateMetadata = try { metadataReader.read() } catch { @@ -1677,6 +1808,7 @@ object StatefulOperatorsUtils { ) val SYMMETRIC_HASH_JOIN_EXEC_OP_NAME = "symmetricHashJoin" val STATE_STORE_SAVE_EXEC_OP_NAME = "stateStoreSave" + val STATEFUL_STREAMLINE_AGGREGATE_EXEC_OP_NAME = "StatefulStreamlineAggregate" val DEDUPLICATE_EXEC_OP_NAME = "dedupe" val DEDUPLICATE_WITHIN_WATERMARK_EXEC_OP_NAME = "dedupeWithinWatermark" val SESSION_WINDOW_STATE_STORE_SAVE_EXEC_OP_NAME = "sessionWindowStateStoreSaveExec" diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateTypesEncoderUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateTypesEncoderUtils.scala index d147ad66c2460..480aebfb351fa 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateTypesEncoderUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/StateTypesEncoderUtils.scala @@ -117,14 +117,14 @@ class StateTypesEncoder[V]( throw StateStoreErrors.implicitKeyNotFound(stateName) } - keySerializer.apply(keyOption.get).asInstanceOf[UnsafeRow] + keySerializer(keyOption.get).asInstanceOf[UnsafeRow] } /** * Encode the specified value in Spark UnsafeRow with no ttl. */ def encodeValue(value: V): UnsafeRow = { - objToRowSerializer.apply(value).asInstanceOf[UnsafeRow] + objToRowSerializer(value).asInstanceOf[UnsafeRow] } /** @@ -132,15 +132,15 @@ class StateTypesEncoder[V]( * with provided ttl expiration. */ def encodeValue(value: V, expirationMs: Long): UnsafeRow = { - val objRow: InternalRow = objToRowSerializer.apply(value) - valueTTLProjection.apply(InternalRow(objRow, expirationMs)) + val objRow: InternalRow = objToRowSerializer(value) + valueTTLProjection(InternalRow(objRow, expirationMs)) } def decodeValue(row: UnsafeRow): V = { if (hasTtl) { - rowToObjDeserializer.apply(row.getStruct(0, valEncoder.schema.length)) + rowToObjDeserializer(row.getStruct(0, valEncoder.schema.length)) } else { - rowToObjDeserializer.apply(row) + rowToObjDeserializer(row) } } @@ -213,14 +213,14 @@ class CompositeKeyStateEncoder[K, V]( throw StateStoreErrors.implicitKeyNotFound(stateName) } val groupingKey = keyOption.get - val groupingKeyRow = groupingKeySerializer.apply(groupingKey) + val groupingKeyRow = groupingKeySerializer(groupingKey) // Create the final unsafeRow mapping column name "key" to the keyRow groupingKeyProjection(InternalRow(groupingKeyRow)) } def encodeUserKey(userKey: K): UnsafeRow = { - val userKeyRow = userKeySerializer.apply(userKey) + val userKeyRow = userKeySerializer(userKey) // Create the final unsafeRow mapping column name "userKey" to the userKeyRow userKeyProjection(InternalRow(userKeyRow)) @@ -237,8 +237,8 @@ class CompositeKeyStateEncoder[K, V]( } val groupingKey = keyOption.get - val keyRow = groupingKeySerializer.apply(groupingKey) - val userKeyRow = userKeySerializer.apply(userKey) + val keyRow = groupingKeySerializer(groupingKey) + val userKeyRow = userKeySerializer(userKey) // Create the final unsafeRow combining the keyRow and userKeyRow compositeKeyProjection(InternalRow(keyRow, userKeyRow)) @@ -253,7 +253,7 @@ class CompositeKeyStateEncoder[K, V]( * Only user key is returned though grouping key also exist in the row. */ def decodeCompositeKey(row: UnsafeRow): K = { - userKeyRowToObjDeserializer.apply(row.getStruct(1, userKeyEnc.schema.length)) + userKeyRowToObjDeserializer(row.getStruct(1, userKeyEnc.schema.length)) } } @@ -264,7 +264,7 @@ class TTLEncoder(schema: StructType) { // Take a groupingKey UnsafeRow and turn it into a (expirationMs, groupingKey) UnsafeRow. def encodeTTLRow(expirationMs: Long, elementKey: UnsafeRow): UnsafeRow = { - ttlKeyProjection.apply( + ttlKeyProjection( InternalRow(expirationMs, elementKey.asInstanceOf[InternalRow])) } } @@ -294,22 +294,22 @@ class TimerKeyEncoder(keyExprEnc: ExpressionEncoder[Any]) { private val secIndexKeyProjection = UnsafeProjection.create(keySchemaForSecIndex) def encodedKey(groupingKey: Any, expiryTimestampMs: Long): UnsafeRow = { - val keyRow = keySerializer.apply(groupingKey) - keyRowProjection.apply(InternalRow(keyRow, expiryTimestampMs)) + val keyRow = keySerializer(groupingKey) + keyRowProjection(InternalRow(keyRow, expiryTimestampMs)) } def encodeSecIndexKey(groupingKey: Any, expiryTimestampMs: Long): UnsafeRow = { - val keyRow = keySerializer.apply(groupingKey) - secIndexKeyProjection.apply(InternalRow(expiryTimestampMs, keyRow)) + val keyRow = keySerializer(groupingKey) + secIndexKeyProjection(InternalRow(expiryTimestampMs, keyRow)) } def encodePrefixKey(groupingKey: Any): UnsafeRow = { - val keyRow = keySerializer.apply(groupingKey) - prefixKeyProjection.apply(InternalRow(keyRow)) + val keyRow = keySerializer(groupingKey) + prefixKeyProjection(InternalRow(keyRow)) } def decodePrefixKey(retUnsafeRow: UnsafeRow): Any = { - keyDeserializer.apply(retUnsafeRow) + keyDeserializer(retUnsafeRow) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExec.scala index f0e3003b2b710..1966347ea8571 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExec.scala @@ -27,8 +27,10 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.WidenStatefulOpNullability import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeProjection import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorStateInfo, StatefulOperatorsUtils} import org.apache.spark.sql.execution.streaming.operators.stateful.join.StreamingSymmetricHashJoinHelper.StateStoreAwareZipPartitionsHelper import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.statefulprocessor.{DriverStatefulProcessorHandleImpl, ImplicitGroupingKeyTracker, StatefulProcessorHandleImpl, StatefulProcessorHandleState} @@ -55,6 +57,7 @@ import org.apache.spark.util.{CompletionIterator, SerializableConfiguration, Uti * @param eventTimeWatermarkForLateEvents event time watermark for filtering late events * @param eventTimeWatermarkForEviction event time watermark for state eviction * @param isStreaming defines whether the query is streaming or batch + * @param isRealTimeMode defines whether the query is running in Real-Time Mode * @param child the physical plan for the underlying data */ case class TransformWithStateExec( @@ -74,6 +77,7 @@ case class TransformWithStateExec( eventTimeWatermarkForEviction: Option[Long], child: SparkPlan, isStreaming: Boolean = true, + isRealTimeMode: Boolean = false, hasInitialState: Boolean = false, initialStateGroupingAttrs: Seq[Attribute], initialStateDataAttrs: Seq[Attribute], @@ -87,7 +91,8 @@ case class TransformWithStateExec( eventTimeWatermarkForEviction, child, initialStateGroupingAttrs, - initialState) + initialState, + isRealTimeMode) with ObjectProducerExec { override def output: Seq[Attribute] = @@ -200,7 +205,10 @@ case class TransformWithStateExec( } } - private def handleInputRows(keyRow: UnsafeRow, valueRowIter: Iterator[InternalRow]): + private def handleInputRows( + keyRow: UnsafeRow, + currentProcessingTimeMs: Option[Long], + valueRowIter: Iterator[InternalRow]): Iterator[InternalRow] = { val getOutputRow = ObjectOperator.wrapObjectToRow(outputObjectType) @@ -218,7 +226,7 @@ case class TransformWithStateExec( statefulProcessor.handleInputRows( keyObj, valueObjIter, - new TimerValuesImpl(batchTimestampMs, eventTimeWatermarkForEviction)).map { obj => + new TimerValuesImpl(currentProcessingTimeMs, eventTimeWatermarkForEviction)).map { obj => getOutputRow(obj) } } @@ -254,20 +262,86 @@ case class TransformWithStateExec( val groupedIter = GroupedIterator(dataIter, groupingAttributes, child.output) groupedIter.flatMap { case (keyRow, valueRowIter) => val keyUnsafeRow = keyRow.asInstanceOf[UnsafeRow] - handleInputRows(keyUnsafeRow, valueRowIter) + handleInputRows(keyUnsafeRow, batchTimestampMs, valueRowIter) + } + } + + // This is used in Real-time mode to process the data and expire timers interleaved; we also + // can do incremental cleanup for every record that we process. + private def processNewDataAndTimersAndTTL( + dataIter: Iterator[InternalRow], + processorHandle: StatefulProcessorHandleImpl): Iterator[InternalRow] = { + val keyProjection = GenerateUnsafeProjection.generate(groupingAttributes, child.output) + val ttlEvictionIntervalMs = conf.getConf( + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS) + var lastTTLEvictionTriggeredMillis = -1L + + dataIter.flatMap { row => + val keyRow = keyProjection(row) + val valueRowIter = watermarkPredicateForDataForLateEvents match { + case Some(predicate) if timeMode == TimeMode.EventTime() => + applyRemovingRowsOlderThanWatermark(Iterator.single(row), predicate) + case _ => + Iterator.single(row) + } + + val outputDataItr = if (valueRowIter.isEmpty) { + Iterator.empty + } else { + handleInputRows( + keyRow, Some(LowLatencyClock.getClock.getTimeMillis()), valueRowIter) + } + + // handle expired timer rows + // late bind for the expired timers to deal with lazy iterators + val expiredTimersOutputItr = new Iterator[InternalRow] { + private lazy val itr = getIterator() + override def hasNext: Boolean = itr.hasNext + override def next(): InternalRow = itr.next() + private def getIterator(): Iterator[InternalRow] = { + processTimers( + timeMode, + Some(LowLatencyClock.getClock.getTimeMillis()), + processorHandle, + useReusableIterator = true) + } + } + + // handle expired state values via ttl cleanup + val cleanupTtlIter = new Iterator[InternalRow] { + private var yetEvaluated = true + + override def hasNext: Boolean = { + if (yetEvaluated) { + yetEvaluated = false + val currentTimeMs = LowLatencyClock.getClock.getTimeMillis() + if (currentTimeMs - lastTTLEvictionTriggeredMillis > ttlEvictionIntervalMs) { + processorHandle.doTtlCleanup(currentTimeMs) + lastTTLEvictionTriggeredMillis = currentTimeMs + } + } + false + } + + override def next(): InternalRow = + throw new IllegalStateException("next() should not be called on this iterator") + } + + outputDataItr ++ expiredTimersOutputItr ++ cleanupTtlIter } } private def handleTimerRows( keyObj: Any, expiryTimestampMs: Long, + currentProcessingTimeMs: Option[Long], processorHandle: StatefulProcessorHandleImpl): Iterator[InternalRow] = { val getOutputRow = ObjectOperator.wrapObjectToRow(outputObjectType) ImplicitGroupingKeyTracker.setImplicitKey(keyObj) val mappedIterator = withStatefulProcessorErrorHandling("handleExpiredTimer") { statefulProcessor.handleExpiredTimer( keyObj, - new TimerValuesImpl(batchTimestampMs, eventTimeWatermarkForEviction), + new TimerValuesImpl(currentProcessingTimeMs, eventTimeWatermarkForEviction), new ExpiredTimerInfoImpl(Some(expiryTimestampMs))).map { obj => getOutputRow(obj) } @@ -281,31 +355,37 @@ case class TransformWithStateExec( private def processTimers( timeMode: TimeMode, - processorHandle: StatefulProcessorHandleImpl): Iterator[InternalRow] = { + currentProcessingTimeMs: Option[Long], + processorHandle: StatefulProcessorHandleImpl, + useReusableIterator: Boolean = false): Iterator[InternalRow] = { val numExpiredTimers = longMetric("numExpiredTimers") - // SPARK-56566: Timers are always scanned without a lower bound (full scan up to the current - // batch timestamp / eviction watermark). We intentionally do not pass - // prevBatchTimestampMs / lateEventsWatermark as the exclusive lower bound here: - // registerTimer has no guard on the registered expiry, so a user-registered timer with expiry - // at or below the previous batch's lower bound would be silently dropped by a bounded scan. - // Revisit once registerTimer enforces ts > currentBatchTimestamp / watermark. + def getExpiredTimers(expiryTimestampMs: Long): Iterator[(Any, Long)] = { + if (useReusableIterator) { + processorHandle.getExpiredTimersReusableIterator(expiryTimestampMs) + } else { + processorHandle.getExpiredTimers(expiryTimestampMs) + } + } + // The final batch scan has no lower bound. RTM's reusable per-row scan resumes from its + // previous expiration threshold to avoid repeatedly scanning older timer entries. timeMode match { case ProcessingTime => - assert(batchTimestampMs.isDefined) - val batchTimestamp = batchTimestampMs.get - processorHandle.getExpiredTimers(batchTimestamp) + assert(currentProcessingTimeMs.isDefined) + getExpiredTimers(currentProcessingTimeMs.get) .flatMap { case (keyObj, expiryTimestampMs) => numExpiredTimers += 1 - handleTimerRows(keyObj, expiryTimestampMs, processorHandle) + handleTimerRows( + keyObj, expiryTimestampMs, currentProcessingTimeMs, processorHandle) } case EventTime => assert(eventTimeWatermarkForEviction.isDefined) val watermark = eventTimeWatermarkForEviction.get - processorHandle.getExpiredTimers(watermark) + getExpiredTimers(watermark) .flatMap { case (keyObj, expiryTimestampMs) => numExpiredTimers += 1 - handleTimerRows(keyObj, expiryTimestampMs, processorHandle) + handleTimerRows( + keyObj, expiryTimestampMs, currentProcessingTimeMs, processorHandle) } case _ => Iterator.empty @@ -334,17 +414,21 @@ case class TransformWithStateExec( val updatesStartTimeNs = currentTimeNs var timerProcessingStartTimeNs = currentTimeNs - // If timeout is based on event time, then filter late data based on watermark - val filteredIter = watermarkPredicateForDataForLateEvents match { - case Some(predicate) if timeMode == TimeMode.EventTime() => - applyRemovingRowsOlderThanWatermark(iter, predicate) - case _ => - iter + val dataOutputIter = if (isRealTimeMode) { + processNewDataAndTimersAndTTL(iter, processorHandle) + } else { + // If timeout is based on event time, then filter late data based on watermark. + val filteredIter = watermarkPredicateForDataForLateEvents match { + case Some(predicate) if timeMode == TimeMode.EventTime() => + applyRemovingRowsOlderThanWatermark(iter, predicate) + case _ => + iter + } + processNewData(filteredIter) } - val newDataProcessorIter = - CompletionIterator[InternalRow, Iterator[InternalRow]]( - processNewData(filteredIter), { + val newDataProcessorIter = CompletionIterator[InternalRow, Iterator[InternalRow]]( + dataOutputIter, { // Note: Due to the iterator lazy execution, this metric also captures the time taken // by the upstream (consumer) operators in addition to the processing in this operator. allUpdatesTimeMs += NANOSECONDS.toMillis(System.nanoTime - updatesStartTimeNs) @@ -364,7 +448,14 @@ case class TransformWithStateExec( override def next() = itr.next() private def getIterator(): Iterator[InternalRow] = CompletionIterator[InternalRow, Iterator[InternalRow]]( - processTimers(timeMode, processorHandle), { + processTimers( + timeMode, + if (isRealTimeMode) { + Some(LowLatencyClock.getClock.getTimeMillis()) + } else { + batchTimestampMs + }, + processorHandle), { // Note: `timerProcessingTimeMs` also includes the time the parent operators take for // processing output returned from the timers that fire. timerProcessingTimeMs += @@ -409,7 +500,8 @@ case class TransformWithStateExec( val info = getStateInfo val stateSchemaDir = stateSchemaDirPath() validateAndWriteStateSchema(hadoopConf, batchId, stateSchemaVersion, - info, stateSchemaDir, session, operatorStateMetadataVersion, conf.stateStoreEncodingFormat) + info, stateSchemaDir, session, operatorStateMetadataVersion, conf.stateStoreEncodingFormat, + isRealTimeMode = isRealTimeMode) } override protected def doExecute(): RDD[InternalRow] = { @@ -533,9 +625,10 @@ case class TransformWithStateExec( */ private def processData(store: StateStore, singleIterator: Iterator[InternalRow]): CompletionIterator[InternalRow, Iterator[InternalRow]] = { + val currentTimestampMs = if (isRealTimeMode) Some(currentTimestampMsFn) else None val processorHandle = new StatefulProcessorHandleImpl( store, getStateInfo.queryRunId, keyEncoder, timeMode, - isStreaming, batchTimestampMs, prevBatchTimestampMs, metrics) + isStreaming, batchTimestampMs, prevBatchTimestampMs, metrics, currentTimestampMs) assert(processorHandle.getHandleState == StatefulProcessorHandleState.CREATED) statefulProcessor.setHandle(processorHandle) withStatefulProcessorErrorHandling("init") { @@ -550,8 +643,10 @@ case class TransformWithStateExec( childDataIterator: Iterator[InternalRow], initStateIterator: Iterator[InternalRow]): CompletionIterator[InternalRow, Iterator[InternalRow]] = { + val currentTimestampMs = if (isRealTimeMode) Some(currentTimestampMsFn) else None val processorHandle = new StatefulProcessorHandleImpl(store, getStateInfo.queryRunId, - keyEncoder, timeMode, isStreaming, batchTimestampMs, prevBatchTimestampMs, metrics) + keyEncoder, timeMode, isStreaming, batchTimestampMs, prevBatchTimestampMs, metrics, + currentTimestampMs) assert(processorHandle.getHandleState == StatefulProcessorHandleState.CREATED) statefulProcessor.setHandle(processorHandle) withStatefulProcessorErrorHandling("init") { @@ -626,6 +721,7 @@ object TransformWithStateExec { None, child, isStreaming = false, + isRealTimeMode = false, hasInitialState, initialStateGroupingAttrs, initialStateDataAttrs, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExecBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExecBase.scala index f5abe333d0f35..82fce42870d4a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExecBase.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/TransformWithStateExecBase.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, SortOrde import org.apache.spark.sql.catalyst.plans.logical.{EventTime, ProcessingTime} import org.apache.spark.sql.catalyst.plans.physical.Distribution import org.apache.spark.sql.execution.{BinaryExecNode, SparkPlan} +import org.apache.spark.sql.execution.datasources.v2.LowLatencyClock import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorCustomMetric, StatefulOperatorCustomSumMetric, StatefulOperatorPartitioning, StateStoreWriter, WatermarkSupport} import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.statefulprocessor.ImplicitGroupingKeyTracker import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, RocksDBStateStoreProvider, StateStoreErrors, TransformWithStateUserFunctionException} @@ -46,7 +47,8 @@ abstract class TransformWithStateExecBase( eventTimeWatermarkForEviction: Option[Long], child: SparkPlan, initialStateGroupingAttrs: Seq[Attribute], - initialState: SparkPlan) + initialState: SparkPlan, + isRealTimeMode: Boolean = false) extends BinaryExecNode with StateStoreWriter with WatermarkSupport @@ -69,6 +71,22 @@ abstract class TransformWithStateExecBase( // The keys that may have a watermark attribute. override def keyExpressions: Seq[Attribute] = groupingAttributes + @inline + protected lazy val currentTimestampMsFn: () => Long = () => { + assert( + batchTimestampMs.isDefined, + "batchTimestampMs should be set when " + + "invoking the currentTimestampMs function. This function must have been " + + "eagerly created; ensure it is lazy and never invoked before physical planning.") + // In real-time mode, we use the local wall time as the current processing time. + // Skew between executors is possible, but we assume it is acceptable for now. + if (isRealTimeMode) { + LowLatencyClock.getClock.getTimeMillis() + } else { + batchTimestampMs.get + } + } + /** * Distribute by grouping attributes - We need the underlying data and the initial state data to * have the same grouping so that the data are co-located on the same task. @@ -89,9 +107,18 @@ abstract class TransformWithStateExecBase( * We need the initial state to also use the ordering as the data so that we can co-locate the * keys from the underlying data and the initial state. */ - override def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq( - groupingAttributes.map(SortOrder(_, Ascending)), - initialStateGroupingAttrs.map(SortOrder(_, Ascending))) + override def requiredChildOrdering: Seq[Seq[SortOrder]] = { + if (isRealTimeMode) { + // In real-time mode, we don't need to order the initial state data since we produce the + // (key, value) pair for every single data. Also, streaming shuffle does not support + // sorting by nature. + Seq.fill(children.size)(Nil) + } else { + Seq( + groupingAttributes.map(SortOrder(_, Ascending)), + initialStateGroupingAttrs.map(SortOrder(_, Ascending))) + } + } override def shouldRunAnotherBatch(newInputWatermark: Long): Boolean = { if (timeMode == ProcessingTime) { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/statefulprocessor/StatefulProcessorHandleImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/statefulprocessor/StatefulProcessorHandleImpl.scala index 291cc02ea989b..25bdc58262faf 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/statefulprocessor/StatefulProcessorHandleImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/statefulprocessor/StatefulProcessorHandleImpl.scala @@ -106,6 +106,8 @@ class QueryInfoImpl( * @param isStreaming - defines whether the query is streaming or batch * @param batchTimestampMs - timestamp for the current batch if available * @param metrics - metrics to be updated as part of stateful processing + * @param currentTimestampMs - optional function for obtaining the current processing time. Real- + * time mode supplies a live clock; other modes use batchTimestampMs. */ class StatefulProcessorHandleImpl( store: StateStore, @@ -115,7 +117,8 @@ class StatefulProcessorHandleImpl( isStreaming: Boolean = true, batchTimestampMs: Option[Long] = None, prevBatchTimestampMs: Option[Long] = None, - metrics: Map[String, SQLMetric] = Map.empty) + metrics: Map[String, SQLMetric] = Map.empty, + currentTimestampMs: Option[() => Long] = None) extends StatefulProcessorHandleImplBase(timeMode, keyEncoder) with Logging { import StatefulProcessorHandleState._ @@ -125,6 +128,10 @@ class StatefulProcessorHandleImpl( */ private[sql] val ttlStates: util.List[TTLState] = new util.ArrayList[TTLState]() + private lazy val ttlTimestampMs = currentTimestampMs.orElse { + batchTimestampMs.map(timestamp => () => timestamp) + } + private val BATCH_QUERY_ID = "00000000-0000-0000-0000-000000000000" currState = CREATED @@ -187,6 +194,15 @@ class StatefulProcessorHandleImpl( timerState.getExpiredTimers(expiryTimestampMs, prevExpiryTimestampMs) } + /** + * Return expired timers through a cached native iterator. RTM invokes this once per input row, + * so reusing the iterator avoids allocating native scan resources for every row. + */ + def getExpiredTimersReusableIterator(expiryTimestampMs: Long): Iterator[(Any, Long)] = { + verifyTimerOperations("get_expired_timers") + timerState.getExpiredTimersReusable(expiryTimestampMs) + } + /** * Function to list all the registered timers for given implicit key * Note: calling listTimers() within the `handleInputRows` method of the StatefulProcessor @@ -204,9 +220,13 @@ class StatefulProcessorHandleImpl( * which is expired will be cleaned up from StateStore. */ def doTtlCleanup(): Unit = { + ttlTimestampMs.foreach(currentTimestampMs => doTtlCleanup(currentTimestampMs())) + } + + def doTtlCleanup(evictionTimestampMs: Long): Unit = { val numValuesRemovedDueToTTLExpiry = metrics.get("numValuesRemovedDueToTTLExpiry").get ttlStates.forEach { s => - numValuesRemovedDueToTTLExpiry += s.clearExpiredStateForAllKeys() + numValuesRemovedDueToTTLExpiry += s.clearExpiredStateForAllKeys(evictionTimestampMs) } } @@ -242,9 +262,9 @@ class StatefulProcessorHandleImpl( val stateEncoder = encoderFor[T].asInstanceOf[ExpressionEncoder[Any]] val result = if (ttlEnabled) { validateTTLConfig(ttlConfig, stateName) - assert(batchTimestampMs.isDefined) + assert(ttlTimestampMs.isDefined) val valueStateWithTTL = new ValueStateImplWithTTL[T](store, stateName, - keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, + keyEncoder, stateEncoder, ttlConfig, ttlTimestampMs.get, prevBatchTimestampMs, metrics) ttlStates.add(valueStateWithTTL) TWSMetricsUtils.incrementMetric(metrics, "numValueStateWithTTLVars") @@ -292,9 +312,9 @@ class StatefulProcessorHandleImpl( val stateEncoder = encoderFor[T].asInstanceOf[ExpressionEncoder[Any]] val result = if (ttlEnabled) { validateTTLConfig(ttlConfig, stateName) - assert(batchTimestampMs.isDefined) + assert(ttlTimestampMs.isDefined) val listStateWithTTL = new ListStateImplWithTTL[T](store, stateName, - keyEncoder, stateEncoder, ttlConfig, batchTimestampMs.get, + keyEncoder, stateEncoder, ttlConfig, ttlTimestampMs.get, prevBatchTimestampMs, metrics) TWSMetricsUtils.incrementMetric(metrics, "numListStateWithTTLVars") ttlStates.add(listStateWithTTL) @@ -331,9 +351,9 @@ class StatefulProcessorHandleImpl( val valEncoder = encoderFor[V].asInstanceOf[ExpressionEncoder[Any]] val result = if (ttlEnabled) { validateTTLConfig(ttlConfig, stateName) - assert(batchTimestampMs.isDefined) + assert(ttlTimestampMs.isDefined) val mapStateWithTTL = new MapStateImplWithTTL[K, V](store, stateName, keyEncoder, userKeyEnc, - valEncoder, ttlConfig, batchTimestampMs.get, + valEncoder, ttlConfig, ttlTimestampMs.get, prevBatchTimestampMs, metrics) TWSMetricsUtils.incrementMetric(metrics, "numMapStateWithTTLVars") ttlStates.add(mapStateWithTTL) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala index 977e15e38e66a..c93f434359368 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/timers/TimerStateImpl.scala @@ -16,6 +16,8 @@ */ package org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.timers +import java.io.Closeable + import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{EXPIRY_TIMESTAMP, KEY} import org.apache.spark.sql.catalyst.InternalRow @@ -114,6 +116,15 @@ class TimerStateImpl( private val secIndexProjection = UnsafeProjection.create(keySchemaForSecIndex) + private var reusableIterator: Option[ReusableIterator[UnsafeRowPair]] = None + private var lastScannedExpiryTimestampMs = 0L + + private lazy val expiryTimestampProjection = UnsafeProjection.create( + new StructType().add("expiryTimestampMs", LongType, nullable = false)) + + private lazy val lastScannedExpiryTimestampRow = + expiryTimestampProjection.apply(InternalRow(lastScannedExpiryTimestampMs)) + // Placeholder grouping-key struct used in range-scan boundary rows; see // [[RangeScanBoundaryUtils]] for rationale. Correctness relies on real stored // entries never having a null grouping-key struct, which is preserved by @@ -236,6 +247,37 @@ class TimerStateImpl( val endKey = encodeTimestampAsKey(expiryTimestampMs) val iter = store.rangeScan(startKey, endKey, tsToKeyCFName) + getExpiredTimersIterator(iter, expiryTimestampMs, isReusable = false) + } + + /** + * Return expired timers using one cached native iterator. After the first scan, refresh and + * resume from the previous scan's expiration threshold. + */ + private[sql] def getExpiredTimersReusable( + expiryTimestampMs: Long): Iterator[(Any, Long)] = { + store match { + case reusableStore: SupportsReusableIterator => + val iter = reusableIterator match { + case Some(existingIterator) => + lastScannedExpiryTimestampRow.setLong(0, lastScannedExpiryTimestampMs) + existingIterator.refreshAndSeekToPrefix(lastScannedExpiryTimestampRow) + existingIterator + case None => + val newIterator = reusableStore.reusableIterator(tsToKeyCFName) + reusableIterator = Some(newIterator) + newIterator + } + lastScannedExpiryTimestampMs = expiryTimestampMs + getExpiredTimersIterator(iter, expiryTimestampMs, isReusable = true) + case _ => getExpiredTimers(expiryTimestampMs) + } + } + + private def getExpiredTimersIterator( + iter: Iterator[UnsafeRowPair] with Closeable, + expiryTimestampMs: Long, + isReusable: Boolean): Iterator[(Any, Long)] = { new NextIterator[(Any, Long)] { override protected def getNext(): (Any, Long) = { if (iter.hasNext) { @@ -255,7 +297,9 @@ class TimerStateImpl( } override protected def close(): Unit = { - iter.close() + if (!isReusable) { + iter.close() + } } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ListStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ListStateImplWithTTL.scala index 10ec3a58500af..8c69dd7f33816 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ListStateImplWithTTL.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ListStateImplWithTTL.scala @@ -34,7 +34,7 @@ import org.apache.spark.util.NextIterator * @param keyExprEnc - Spark SQL encoder for key * @param valEncoder - Spark SQL encoder for value * @param ttlConfig - TTL configuration for values stored in this state - * @param batchTimestampMs - current batch processing timestamp. + * @param currentTimestampMs - function to get the current processing time timestamp. * @param prevBatchTimestampMs - batch timestamp from the previous micro-batch (exclusive). * Entries with expiration at or below this timestamp are assumed * to have been already cleaned up and will be skipped during @@ -48,11 +48,11 @@ class ListStateImplWithTTL[S]( keyExprEnc: ExpressionEncoder[Any], valEncoder: ExpressionEncoder[Any], ttlConfig: TTLConfig, - batchTimestampMs: Long, + currentTimestampMs: () => Long, prevBatchTimestampMs: Option[Long] = None, metrics: Map[String, SQLMetric]) extends OneToManyTTLState( - stateName, store, keyExprEnc.schema, ttlConfig, batchTimestampMs, + stateName, store, keyExprEnc.schema, ttlConfig, currentTimestampMs, prevBatchTimestampMs, metrics) with ListState[S] { private lazy val stateTypesEncoder = StateTypesEncoder(keyExprEnc, valEncoder, @@ -83,7 +83,7 @@ class ListStateImplWithTTL[S]( override protected def getNext(): S = { val iter = unsafeRowValuesIterator.dropWhile { row => - stateTypesEncoder.isExpired(row, batchTimestampMs) + stateTypesEncoder.isExpired(row, currentTimestampMs()) } if (iter.hasNext) { @@ -167,7 +167,7 @@ class ListStateImplWithTTL[S]( var newMinExpirationMsOpt: Option[Long] = None var isFirst = true unsafeRowValuesIterator.foreach { encodedValue => - if (!stateTypesEncoder.isExpired(encodedValue, batchTimestampMs)) { + if (!stateTypesEncoder.isExpired(encodedValue, currentTimestampMs())) { if (isFirst) { isFirst = false store.put(elementKey, encodedValue, stateName) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/MapStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/MapStateImplWithTTL.scala index 03aa8aaa6ace2..6eb6801700ed1 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/MapStateImplWithTTL.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/MapStateImplWithTTL.scala @@ -35,7 +35,7 @@ import org.apache.spark.util.NextIterator * @param userKeyEnc - Spark SQL encoder for the map key * @param valEncoder - SQL encoder for state variable * @param ttlConfig - the ttl configuration (time to live duration etc.) - * @param batchTimestampMs - current batch processing timestamp. + * @param currentTimestampMs - function to get the current processing time timestamp. * @param prevBatchTimestampMs - batch timestamp from the previous micro-batch (exclusive). * Entries with expiration at or below this timestamp are assumed * to have been already cleaned up and will be skipped during @@ -52,12 +52,12 @@ class MapStateImplWithTTL[K, V]( userKeyEnc: ExpressionEncoder[Any], valEncoder: ExpressionEncoder[Any], ttlConfig: TTLConfig, - batchTimestampMs: Long, + currentTimestampMs: () => Long, prevBatchTimestampMs: Option[Long] = None, metrics: Map[String, SQLMetric]) extends OneToOneTTLState( stateName, store, getCompositeKeySchema(keyExprEnc.schema, userKeyEnc.schema), ttlConfig, - batchTimestampMs, prevBatchTimestampMs, metrics) with MapState[K, V] with Logging { + currentTimestampMs, prevBatchTimestampMs, metrics) with MapState[K, V] with Logging { private val stateTypesEncoder = new CompositeKeyStateEncoder( keyExprEnc, userKeyEnc, valEncoder, stateName, hasTtl = true) @@ -84,7 +84,7 @@ class MapStateImplWithTTL[K, V]( val retRow = store.get(encodedCompositeKey, stateName) if (retRow != null) { - if (!stateTypesEncoder.isExpired(retRow, batchTimestampMs)) { + if (!stateTypesEncoder.isExpired(retRow, currentTimestampMs())) { stateTypesEncoder.decodeValue(retRow).asInstanceOf[V] } else { null.asInstanceOf[V] @@ -107,7 +107,7 @@ class MapStateImplWithTTL[K, V]( val encodedCompositeKey = stateTypesEncoder.encodeCompositeKey(key) val ttlExpirationMs = StateTTL - .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs) + .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, currentTimestampMs()) val encodedValue = stateTypesEncoder.encodeValue(value, ttlExpirationMs) updatePrimaryAndSecondaryIndices(encodedCompositeKey, encodedValue, ttlExpirationMs) @@ -120,7 +120,7 @@ class MapStateImplWithTTL[K, V]( new NextIterator[(K, V)] { override protected def getNext(): (K, V) = { val iter = unsafeRowPairIterator.dropWhile { rowPair => - stateTypesEncoder.isExpired(rowPair.value, batchTimestampMs) + stateTypesEncoder.isExpired(rowPair.value, currentTimestampMs()) } if (iter.hasNext) { val currentRowPair = iter.next() diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/TTLState.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/TTLState.scala index cab6aa9630e9c..bcc4b122ae693 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/TTLState.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/TTLState.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.execution.metric.SQLMetric import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateKeyValueRowSchemaUtils._ import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TTLEncoder import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.statefulprocessor.TWSMetricsUtils -import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, RangeKeyScanStateEncoderSpec, RangeScanBoundaryUtils, StateStore} +import org.apache.spark.sql.execution.streaming.state.{NoPrefixKeyStateEncoderSpec, RangeKeyScanStateEncoderSpec, RangeScanBoundaryUtils, ReusableIterator, StateStore, SupportsReusableIterator, UnsafeRowPair} import org.apache.spark.sql.streaming.TTLConfig import org.apache.spark.sql.types._ @@ -84,9 +84,9 @@ trait TTLState { // a map key. private[sql] def elementKeySchema: StructType - // The timestamp at which the batch is being processed. All state variables that have - // an expiration at or before this timestamp must be cleaned up. - private[sql] def batchTimestampMs: Long + // Returns the current processing-time timestamp. It is fixed to the journaled batch timestamp + // in micro-batch mode and reads the executor clock in Real-Time Mode. + private[sql] def currentTimestampMs: () => Long // The batch timestamp from the previous micro-batch, used to derive the startKey // for scan-based TTL eviction. Entries at or below prevBatchTimestampMs were already @@ -125,7 +125,7 @@ trait TTLState { UnsafeProjection.create(Array[DataType](NullType)).apply(InternalRow.apply(null)) private[sql] final def ttlExpirationMs = StateTTL - .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs) + .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, currentTimestampMs()) store.createColFamilyIfAbsent( TTL_INDEX, @@ -170,35 +170,67 @@ trait TTLState { store.iterator(TTL_INDEX).map(kv => toTTLRow(kv.key)) } - // Returns an Iterator over the keys in the TTL index that have expired. Uses a bounded - // range scan over [prevBatchTimestampMs+1, batchTimestampMs+1) to skip entries that - // were already evicted in previous batches. + // fields for reusable iterator. + private var reusableIteratorOpt: Option[ReusableIterator[UnsafeRowPair]] = None + private var lastTimerDeletedExpiryTimeMs: Long = 0L + + private lazy val secondaryIndexExpiryTsProjection = UnsafeProjection.create( + new StructType() + .add("expiryTimestampMs", LongType, nullable = false) + ) + + private lazy val lastTimerDeletedExpiryTimeMsRow: UnsafeRow = + secondaryIndexExpiryTsProjection.apply(InternalRow(lastTimerDeletedExpiryTimeMs)) + + // Returns an Iterator over the keys in the TTL index that have expired. When the store + // does not support reusable iterators, uses a bounded range scan over + // [prevBatchTimestampMs+1, evictionTimestampMs+1) to skip entries that were already + // evicted (and are now tombstones) in previous batches. // // This method does not delete the keys from the TTL index; it is the responsibility of // the caller to do so. // // The schema of the UnsafeRow returned by this iterator is (expirationMs, elementKey). - private[sql] def ttlEvictionIterator(): Iterator[UnsafeRow] = { - val startKey = prevBatchTimestampMs.flatMap { prevTs => - if (prevTs < Long.MaxValue) { - Some(TTL_ENCODER.encodeTTLRow(prevTs + 1, DEFAULT_ELEMENT_KEY).copy()) - } else { - None - } - } - val endKey = if (batchTimestampMs < Long.MaxValue) { - Some(TTL_ENCODER.encodeTTLRow(batchTimestampMs + 1, DEFAULT_ELEMENT_KEY).copy()) - } else { - None + private[sql] def ttlEvictionIterator(evictionTimestampMs: Long): Iterator[UnsafeRow] = { + val ttlIterator = store match { + case s: SupportsReusableIterator => + if (reusableIteratorOpt.isEmpty) { + reusableIteratorOpt = Some(s.reusableIterator(TTL_INDEX)) + } else { + val reusableIter = reusableIteratorOpt.get + + // Seek to the last deleted timer's expiry time + // Please note that we do not need to specify the grouping key here because + // we are only interested in seeking to the position of the last expiry time + lastTimerDeletedExpiryTimeMsRow.setLong(0, lastTimerDeletedExpiryTimeMs) + reusableIter.refreshAndSeekToPrefix(lastTimerDeletedExpiryTimeMsRow) + } + + lastTimerDeletedExpiryTimeMs = evictionTimestampMs + + reusableIteratorOpt.get + case _ => + val startKey = prevBatchTimestampMs.flatMap { prevTs => + if (prevTs < Long.MaxValue) { + Some(TTL_ENCODER.encodeTTLRow(prevTs + 1, DEFAULT_ELEMENT_KEY).copy()) + } else { + None + } + } + val endKey = if (evictionTimestampMs < Long.MaxValue) { + Some(TTL_ENCODER.encodeTTLRow(evictionTimestampMs + 1, DEFAULT_ELEMENT_KEY).copy()) + } else { + None + } + store.rangeScan(startKey, endKey, TTL_INDEX) } - val ttlIterator = store.rangeScan(startKey, endKey, TTL_INDEX) // Recall that the format is (expirationMs, elementKey) -> TTL_EMPTY_VALUE_ROW, so // kv.value doesn't ever need to be used. // Safety filter: keep only truly expired entries ttlIterator.takeWhile { kv => val expirationMs = kv.key.getLong(0) - StateTTL.isExpired(expirationMs, batchTimestampMs) + StateTTL.isExpired(expirationMs, evictionTimestampMs) }.map(_.key) } @@ -218,7 +250,7 @@ trait TTLState { * * @return number of values cleaned up. */ - private[sql] def clearExpiredStateForAllKeys(): Long + private[sql] def clearExpiredStateForAllKeys(evictionTimestampMs: Long): Long /** * When a user calls clear() on a stateful variable, this method is invoked to @@ -253,14 +285,14 @@ abstract class OneToOneTTLState( storeArg: StateStore, elementKeySchemaArg: StructType, ttlConfigArg: TTLConfig, - batchTimestampMsArg: Long, + currentTimestampMsArg: () => Long, prevBatchTimestampMsArg: Option[Long], metricsArg: Map[String, SQLMetric]) extends TTLState { override private[sql] def stateName: String = stateNameArg override private[sql] def store: StateStore = storeArg override private[sql] def elementKeySchema: StructType = elementKeySchemaArg override private[sql] def ttlConfig: TTLConfig = ttlConfigArg - override private[sql] def batchTimestampMs: Long = batchTimestampMsArg + override private[sql] def currentTimestampMs: () => Long = currentTimestampMsArg override private[sql] def prevBatchTimestampMs: Option[Long] = prevBatchTimestampMsArg override private[sql] def metrics: Map[String, SQLMetric] = metricsArg @@ -310,10 +342,10 @@ abstract class OneToOneTTLState( } } - override private[sql] def clearExpiredStateForAllKeys(): Long = { + override private[sql] def clearExpiredStateForAllKeys(evictionTimestampMs: Long): Long = { var numValuesExpired = 0L - ttlEvictionIterator().foreach { ttlKey => + ttlEvictionIterator(evictionTimestampMs).foreach { ttlKey => // Delete from secondary index deleteFromTTLIndex(ttlKey) // Delete from primary index @@ -372,14 +404,14 @@ abstract class OneToManyTTLState( storeArg: StateStore, elementKeySchemaArg: StructType, ttlConfigArg: TTLConfig, - batchTimestampMsArg: Long, + currentTimestampMsArg: () => Long, prevBatchTimestampMsArg: Option[Long], metricsArg: Map[String, SQLMetric]) extends TTLState { override private[sql] def stateName: String = stateNameArg override private[sql] def store: StateStore = storeArg override private[sql] def elementKeySchema: StructType = elementKeySchemaArg override private[sql] def ttlConfig: TTLConfig = ttlConfigArg - override private[sql] def batchTimestampMs: Long = batchTimestampMsArg + override private[sql] def currentTimestampMs: () => Long = currentTimestampMsArg override private[sql] def prevBatchTimestampMs: Option[Long] = prevBatchTimestampMsArg override private[sql] def metrics: Map[String, SQLMetric] = metricsArg @@ -517,10 +549,10 @@ abstract class OneToManyTTLState( // Clears all the expired values for the given elementKey. protected def clearExpiredValues(elementKey: UnsafeRow): ValueExpirationResult - override private[sql] def clearExpiredStateForAllKeys(): Long = { + override private[sql] def clearExpiredStateForAllKeys(evictionTimestampMs: Long): Long = { var totalNumValuesExpired = 0L - ttlEvictionIterator().foreach { ttlKey => + ttlEvictionIterator(evictionTimestampMs).foreach { ttlKey => val ttlRow = toTTLRow(ttlKey) val elementKey = ttlRow.elementKey diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ValueStateImplWithTTL.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ValueStateImplWithTTL.scala index 1559acf7222cf..7dce677df69ac 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ValueStateImplWithTTL.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/operators/stateful/transformwithstate/ttl/ValueStateImplWithTTL.scala @@ -32,7 +32,7 @@ import org.apache.spark.sql.streaming.{TTLConfig, ValueState} * @param keyExprEnc - Spark SQL encoder for key * @param valEncoder - Spark SQL encoder for value * @param ttlConfig - TTL configuration for values stored in this state - * @param batchTimestampMs - current batch processing timestamp. + * @param currentTimestampMs - function to get the current processing time timestamp. * @param prevBatchTimestampMs - batch timestamp from the previous micro-batch (exclusive). * Entries with expiration at or below this timestamp are assumed * to have been already cleaned up and will be skipped during @@ -46,11 +46,11 @@ class ValueStateImplWithTTL[S]( keyExprEnc: ExpressionEncoder[Any], valEncoder: ExpressionEncoder[Any], ttlConfig: TTLConfig, - batchTimestampMs: Long, + currentTimestampMs: () => Long, prevBatchTimestampMs: Option[Long] = None, metrics: Map[String, SQLMetric] = Map.empty) extends OneToOneTTLState( - stateName, store, keyExprEnc.schema, ttlConfig, batchTimestampMs, + stateName, store, keyExprEnc.schema, ttlConfig, currentTimestampMs, prevBatchTimestampMs, metrics) with ValueState[S] { private val stateTypesEncoder = @@ -80,7 +80,7 @@ class ValueStateImplWithTTL[S]( // Getting the 0th ordinal of the struct using valEncoder val resState = stateTypesEncoder.decodeValue(retRow) - if (!stateTypesEncoder.isExpired(retRow, batchTimestampMs)) { + if (!stateTypesEncoder.isExpired(retRow, currentTimestampMs())) { resState.asInstanceOf[S] } else { null.asInstanceOf[S] @@ -95,7 +95,7 @@ class ValueStateImplWithTTL[S]( val encodedKey = stateTypesEncoder.encodeGroupingKey() val ttlExpirationMs = StateTTL - .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, batchTimestampMs) + .calculateExpirationTimeForDuration(ttlConfig.ttlDuration, currentTimestampMs()) val encodedValue = stateTypesEncoder.encodeValue(newState, ttlExpirationMs) updatePrimaryAndSecondaryIndices(encodedKey, encodedValue, ttlExpirationMs) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncProgressTrackingMicroBatchExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncProgressTrackingMicroBatchExecution.scala index d779f48f5e6e1..a87c4c3574d18 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncProgressTrackingMicroBatchExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncProgressTrackingMicroBatchExecution.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.streaming.WriteToStream import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.streaming.{AvailableNowTrigger, OneTimeTrigger, ProcessingTimeTrigger, RealTimeTrigger, StreamingErrors} -import org.apache.spark.sql.execution.streaming.checkpointing.{AsyncCommitLog, AsyncOffsetSeqLog, CommitMetadata, OffsetSeqBase, OffsetSeqLog} +import org.apache.spark.sql.execution.streaming.checkpointing.{AsyncCommitLog, AsyncOffsetSeqLog, CommitMetadataBase, OffsetSeqBase, OffsetSeqLog} import org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreWriter import org.apache.spark.sql.streaming.Trigger import org.apache.spark.util.{Clock, ThreadUtils} @@ -261,7 +261,7 @@ class AsyncProgressTrackingMicroBatchExecution( || isFirstBatch) { isFirstBatch = false commitLog - .addAsync(execCtx.batchId, CommitMetadata(watermarkTracker.currentWatermark)) + .addAsync(execCtx.batchId, createAsyncCommitMetadata(execCtx)) .thenAccept((batchId: Long) => { logInfo(log"Committed async commit log to disk for batch ${MDC(BATCH_ID, batchId)}.") }) @@ -273,7 +273,7 @@ class AsyncProgressTrackingMicroBatchExecution( }) } else { if (!commitLog.addInMemory( - execCtx.batchId, CommitMetadata(watermarkTracker.currentWatermark))) { + execCtx.batchId, createAsyncCommitMetadata(execCtx))) { throw QueryExecutionErrors.concurrentStreamLogUpdate(execCtx.batchId) } logInfo( @@ -285,6 +285,21 @@ class AsyncProgressTrackingMicroBatchExecution( committedOffsets ++= execCtx.endOffsets } + /** + * Builds the commit log entry for an async batch at the version resolved for this run, so that a + * Real-Time Mode query writes the commit log format its state store checkpoint format implies + * (v2 for the RTM default) rather than always VERSION_1. Async progress tracking does not support + * stateful queries, so there are no state store checkpoint ids to persist. + */ + private def createAsyncCommitMetadata( + execCtx: MicroBatchExecutionContext): CommitMetadataBase = { + commitLog.createMetadata( + nextBatchWatermarkMs = watermarkTracker.currentWatermark, + stateUniqueIds = None, + commitLogFormatVersion = execCtx.commitLogFormatVersionOpt.getOrElse( + sparkSessionForStream.sessionState.conf.streamingCommitLogFormatVersion)) + } + /** * Categorize a raw IO failure surfaced via an async log-write future, then route it through * the standard async error handling path. CompletableFuture wraps the underlying cause in diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala index 0d2e4a6941a00..9efee984457dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala @@ -36,17 +36,18 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec, QueryExecution, SerializeFromObjectExec, SparkPlan, SparkPlanner, SparkStrategy => Strategy, UnaryExecNode} import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec, UpdatingSessionsExec} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec import org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader -import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike +import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike} import org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec, TransformWithStateInPySparkExec} -import org.apache.spark.sql.execution.streaming.{StreamingErrors, StreamingQueryPlanTraverseHelper} +import org.apache.spark.sql.execution.streaming.{ProjectAggregationBufferExec, StatefulStreamlineAggregateExec, StreamingErrors, StreamingQueryPlanTraverseHelper} import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, OffsetSeqMetadata, OffsetSeqMetadataBase} import org.apache.spark.sql.execution.streaming.operators.stateful.{SessionWindowStateStoreRestoreExec, SessionWindowStateStoreSaveExec, StatefulOperator, StatefulOperatorStateInfo, StateStoreRestoreExec, StateStoreSaveExec, StateStoreWriter, StreamingDeduplicateExec, StreamingDeduplicateWithinWatermarkExec, StreamingGlobalLimitExec, StreamingLocalLimitExec, UpdateEventTimeColumnExec} import org.apache.spark.sql.execution.streaming.operators.stateful.flatmapgroupswithstate.FlatMapGroupsWithStateExec import org.apache.spark.sql.execution.streaming.operators.stateful.join.{StreamingSymmetricHashJoinExec, StreamingSymmetricHashJoinHelper} import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateExec import org.apache.spark.sql.execution.streaming.sources.WriteToMicroBatchDataSourceV1 -import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadataReader, OperatorStateMetadataV1, OperatorStateMetadataV2, OperatorStateMetadataWriter, StateSchemaBroadcast, StateSchemaMetadata} +import org.apache.spark.sql.execution.streaming.state.{OperatorStateMetadata, OperatorStateMetadataReader, OperatorStateMetadataV1, OperatorStateMetadataV2, OperatorStateMetadataWriter, StateSchemaBroadcast, StateSchemaMetadata} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.OutputMode import org.apache.spark.util.{SerializableConfiguration, Utils} @@ -82,7 +83,8 @@ class IncrementalExecution( val stateSchemaMetadatas: MutableMap[Long, StateSchemaBroadcast] = MutableMap[Long, StateSchemaBroadcast](), mode: CommandExecutionMode.Value = CommandExecutionMode.ALL, - val isTerminatingTrigger: Boolean = false) + val isTerminatingTrigger: Boolean = false, + val isRealTimeMode: Boolean = false) extends QueryExecution(sparkSession, logicalPlan, mode = mode, shuffleCleanupModeOpt = Some(QueryExecution.determineShuffleCleanupMode(sparkSession.sessionState.conf)), @@ -110,6 +112,12 @@ class IncrementalExecution( private lazy val hadoopConf = sparkSession.sessionState.newHadoopConf() + // Populated while state schemas are validated during planning. In micro-batch mode each entry + // is also written immediately, preserving the existing behavior. Real-Time Mode writes these + // entries only after its delayed offset log entry is durable. + private var stateStoreWritersWithMetadata: + Option[Seq[(StateStoreWriter, OperatorStateMetadata)]] = None + private[sql] val numStateStores = OffsetSeqMetadata.readValueOpt(offsetSeqMetadata, SQLConf.STATEFUL_SHUFFLE_PARTITIONS_INTERNAL) .map(SQLConf.SHUFFLE_PARTITIONS.valueConverter) @@ -194,6 +202,9 @@ class IncrementalExecution( case a: UpdatingSessionsExec if a.isStreaming => a.copy(numShufflePartitions = Some(numStateStores)) + + case a: ProjectAggregationBufferExec if a.isStreaming => + a.copy(numShufflePartitions = Some(numStateStores)) } } @@ -280,12 +291,11 @@ class IncrementalExecution( ssw.validateNewMetadata(oldMetadata, metadata) case None => } - val metadataWriter = OperatorStateMetadataWriter.createWriter( - new Path(checkpointLocation, ssw.getStateInfo.operatorId.toString), - hadoopConf, - ssw.operatorStateMetadataVersion, - Some(currentBatchId)) - metadataWriter.write(metadata) + stateStoreWritersWithMetadata = Some( + stateStoreWritersWithMetadata.getOrElse(Seq.empty) :+ (ssw -> metadata)) + if (!isRealTimeMode) { + writeStateMetadata(ssw, metadata) + } if (ssw.supportsSchemaEvolution) { val stateSchemaMetadata = StateSchemaMetadata .createStateSchemaMetadata(checkpointLocation, hadoopConf, stateSchemaList.head) @@ -314,8 +324,37 @@ class IncrementalExecution( } } + private def writeStateMetadata( + stateStoreWriter: StateStoreWriter, + metadata: OperatorStateMetadata): Unit = { + val metadataWriter = OperatorStateMetadataWriter.createWriter( + new Path(checkpointLocation, stateStoreWriter.getStateInfo.operatorId.toString), + hadoopConf, + metadata.version, + Some(currentBatchId)) + metadataWriter.write(metadata) + } + + /** Write state metadata recorded during planning. Used after the delayed RTM offset WAL write. */ + private[streaming] def writeRecordedStateMetadata(): Unit = { + assert(stateStoreWritersWithMetadata.isDefined, + "stateStoreWritersWithMetadata must be defined before writing state metadata") + stateStoreWritersWithMetadata.get.foreach { case (stateStoreWriter, metadata) => + writeStateMetadata(stateStoreWriter, metadata) + } + } + object StateOpIdRule extends SparkPlanPartialRule { override val rule: PartialFunction[SparkPlan, SparkPlan] = { + case a: StatefulStreamlineAggregateExec => + val aggStateInfo = nextStatefulOperationStateInfo() + a.copy( + numShufflePartitions = Some(aggStateInfo.numPartitions), + stateInfo = Some(aggStateInfo), + outputMode = Some(outputMode), + eventTimeWatermarkForLateEvents = None, + eventTimeWatermarkForEviction = None) + case StateStoreSaveExec(keys, None, None, None, None, stateFormatVersion, UnaryExecNode(agg, StateStoreRestoreExec(_, None, _, child))) => @@ -392,6 +431,7 @@ class IncrementalExecution( prevBatchTimestampMs = prevOffsetSeqMetadata.map(_.batchTimestampMs), eventTimeWatermarkForLateEvents = None, eventTimeWatermarkForEviction = None, + isRealTimeMode = IncrementalExecution.this.isRealTimeMode, hasInitialState = hasInitialState ) @@ -440,6 +480,12 @@ class IncrementalExecution( } override val rule: PartialFunction[SparkPlan, SparkPlan] = { + case a: StatefulStreamlineAggregateExec if a.stateInfo.isDefined => + a.copy( + eventTimeWatermarkForLateEvents = inputWatermarkForLateEvents(a.stateInfo.get), + eventTimeWatermarkForEviction = inputWatermarkForEviction(a.stateInfo.get) + ) + case s: StateStoreSaveExec if s.stateInfo.isDefined => s.copy( eventTimeWatermarkForLateEvents = inputWatermarkForLateEvents(s.stateInfo.get), @@ -584,6 +630,23 @@ class IncrementalExecution( rulesToCompose.reduceLeft { (ruleA, ruleB) => ruleA orElse ruleB } } + /** + * Returns true if a checkpoint whose metadata records operator `oldOpName` may be reopened by a + * plan whose operator is `newOpName`, without tripping the operator-mismatch guard. + * + * A streaming aggregation may move between the micro-batch operator ([[StateStoreSaveExec]], + * "stateStoreSave") and the streamline operator ([[StatefulStreamlineAggregateExec]]) in either + * direction: the two share [[StreamingAggregationStateManager]] and the same state format + * version, so the on-disk state written by one is readable by the other. + */ + private def isOperatorMetadataConvertible(oldOpName: String, newOpName: String): Boolean = { + oldOpName == newOpName || ((oldOpName, newOpName) match { + case ("stateStoreSave", "StatefulStreamlineAggregate") => true + case ("StatefulStreamlineAggregate", "stateStoreSave") => true + case _ => false + }) + } + private def checkOperatorValidWithMetadata( planWithStateOpId: SparkPlan, batchId: Long): Unit = { @@ -639,7 +702,7 @@ class IncrementalExecution( (opMapInMetadata.keySet ++ opMapInPhysicalPlan.keySet).foreach { opId => val opInMetadata = opMapInMetadata.getOrElse(opId, "not found") val opInCurBatch = opMapInPhysicalPlan.getOrElse(opId, "not found") - if (opInMetadata != opInCurBatch) { + if (!isOperatorMetadataConvertible(opInMetadata, opInCurBatch)) { throw QueryExecutionErrors.statefulOperatorNotMatchInStateMetadataError( opMapInMetadata, opMapInPhysicalPlan) @@ -656,6 +719,7 @@ class IncrementalExecution( checkOperatorValidWithMetadata(planWithStateOpId, currentBatchId - 1) } + stateStoreWritersWithMetadata = Some(Seq.empty) val planWithSchemas = planWithStateOpId transform StateSchemaAndOperatorMetadataRule.rule simulateWatermarkPropagation(planWithSchemas) @@ -663,7 +727,123 @@ class IncrementalExecution( } } - override def preparations: Seq[Rule[SparkPlan]] = state +: super.preparations + private def isTransformWithStateInitialStateBootstrap: Boolean = { + isRealTimeMode && currentBatchId == 0 && logicalPlan.exists { + case tws: TransformWithState => tws.hasInitialState + case _ => false + } + } + + /** + * The initial state is loaded in a finite batch before the Real-Time Mode source starts its + * first long-running batch. This lets the initial-state shuffle materialize and prevents input + * that was already available when the query started from being processed before initialization. + * The pipelined-shuffle rule also skips this batch because the DAGScheduler does not support a + * job that mixes the initial state's regular shuffle with a pipelined streaming shuffle. + */ + object PrepareTransformWithStateInitialStateForRealTimeMode extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (isTransformWithStateInitialStateBootstrap) { + plan.transformUp { + case scan: RealTimeStreamScanExec => scan.copy(batchDurationMs = 0L) + } + } else { + plan + } + } + } + + /** + * For a Real-Time Mode batch, mark the shuffle exchanges as pipelined so the DAGScheduler + * co-schedules a stateful query's producer (source scan) and consumer (stateful operator) stages + * as one pipelined group -- records stream through a transient shuffle instead of the consumer + * waiting for the producer to fully materialize. The exchange carries the decision as a field + * (see ShuffleExchangeExec.pipelined); the PipelinedShuffleDependency it then builds is the whole + * opt-in -- routing to the streaming shuffle manager and pipelined-group scheduling both follow + * from that dependency type. + * + * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf (there is no + * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary microbatch path is + * unchanged. + * + * Marks every eligible shuffle exchange on the streaming path, so a plan with several pipelined + * shuffles in a chain (e.g. two repartitions, or a repartition feeding a keyed stateful operator) + * is handled: each becomes a PipelinedShuffleDependency and the whole all-pipelined job is + * co-scheduled as one pipelined group (the DAGScheduler treats an all-pipelined job's stage graph + * as a single group). There is no shuffle-count restriction. An exchange whose subtree does not + * reach the real-time scan is skipped -- the static side of a broadcast stream-static join must + * materialize, because it runs to completion rather than streaming. A partitioning the pipelined + * path cannot serve, such as range partitioning, is rejected up front by RealTimeModeAllowlist + * rather than being handled here. + * + * The walk does not descend into a ReusedExchangeExec (a leaf whose wrapped exchange is a field, + * not a tree child), so a REUSED shuffle exchange would keep pipelined=false while its standalone + * twin flips to true. That divergence is not reachable: a reused shuffle requires + * referencing the same streaming source more than once (self-join / self-union / CTE read twice), + * which Real-Time Mode rejects when the query starts (MicroBatchExecution, + * IDENTICAL_SOURCES_IN_UNION_NOT_SUPPORTED) before this rule runs. The only ReusedExchangeExec + * that reaches an RTM plan wraps a BROADCAST exchange (multiple broadcast joins on the same + * static table, SC-209926), which this rule does not match. + * + * Fan-out -- one shuffle read by more than one consumer -- is a limitation that marking a shuffle + * here introduces, not one that was already there. A regular materialized shuffle serves any + * number of consumers; a pipelined one is transient and read once, so the DAGScheduler rejects + * fan-out for a PipelinedShuffleDependency specifically (checkPipelinedGroupsSupportedInRDDGraph, + * inert for a regular ShuffleDependency). So a fan-out query that runs fine unmarked is rejected + * once this rule marks it. The check runs in handleJobSubmitted, so it rejects the batch's job + * rather than the query: such a query fails the same way on every batch instead of failing once + * when it is planned. A plan-time guard here would improve only where the failure is reported, + * not whether the query can run. + */ + object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec]) + if (!isRealTimeMode || isTransformWithStateInitialStateBootstrap) { + plan + } else { + markStreamingPath(plan)._1 + } + } + + /** + * Marks the shuffles that are on the streaming path -- those whose subtree reaches a + * [[RealTimeStreamScanExec]] -- and returns the rewritten plan along with whether this + * subtree reaches one. + * + * A plan can hold a static subtree alongside the streaming one: the static side of a + * broadcast stream-static join is planned in the same physical plan and may contain its own + * shuffle. That shuffle materializes normally and is not part of the pipelined group -- and + * cannot be, since a static side runs to completion rather than streaming. Marking it + * pipelined would pull it into the group and demand slots for stages that must instead + * finish, which fails admission (CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT). This mirrors the + * streaming-path detection the operator allowlist uses (RealTimeModeAllowlist), which only + * inspects nodes whose subtree reaches the real-time scan; marking a wider set than the + * allowlist checks would flip shuffles it never validated. + */ + private def markStreamingPath(plan: SparkPlan): (SparkPlan, Boolean) = plan match { + case rts: RealTimeStreamScanExec => (rts, true) + case p if p.children.isEmpty => (p, false) + case p => + val results = p.children.map(markStreamingPath) + val onStreamingPath = results.exists(_._2) + val newPlan = p.withNewChildren(results.map(_._1)) + newPlan match { + case s: ShuffleExchangeExec if onStreamingPath && !s.pipelined => + // A bare case-class copy does not carry the node's tags, which is where the logical + // link lives, so copy them over the way the tree transforms do. + val marked = s.copy(pipelined = true) + marked.copyTagsFrom(s) + (marked, true) + case other => (other, onStreamingPath) + } + } + + } + + override def preparations: Seq[Rule[SparkPlan]] = + state +: (super.preparations :+ + PrepareTransformWithStateInitialStateForRealTimeMode :+ + MarkPipelinedShuffleForRealTimeMode) /** no need to try-catch again as this is already done once */ override def assertAnalyzed(): Unit = analyzed diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala index de84c9d15a6e4..e7500e6d634d6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala @@ -31,7 +31,7 @@ import org.apache.spark.internal.LogKeys import org.apache.spark.internal.LogKeys._ import org.apache.spark.sql.catalyst.analysis.{ResolveDeduplicate, V2TableReference} import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, CurrentBatchTimestamp, CurrentDate, CurrentTimestamp, FileSourceMetadataAttribute, LocalTimestamp} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, CurrentBatchTimestamp, CurrentDate, CurrentTimestamp, CurrentTimestampNanos, FileSourceMetadataAttribute, LocalTimestamp, LocalTimestampNanos} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Deduplicate, DeduplicateWithinWatermark, Distinct, FlatMapGroupsInPandasWithState, FlatMapGroupsWithState, GlobalLimit, Join, LeafNode, LocalRelation, LogicalPlan, Project, StreamSourceAwareLogicalPlan, TransformWithState, TransformWithStateInPySpark} import org.apache.spark.sql.catalyst.streaming.{StreamingRelationV2, Unassigned, WriteToStream} import org.apache.spark.sql.catalyst.trees.TreePattern.CURRENT_LIKE @@ -46,7 +46,7 @@ import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, RealTimeStreamScanExec, StreamingDataSourceV2Relation, StreamingDataSourceV2ScanRelation, StreamWriterCommitProgress, WriteToDataSourceV2Exec} import org.apache.spark.sql.execution.streaming.{AvailableNowTrigger, Offset, OneTimeTrigger, ProcessingTimeTrigger, RealTimeTrigger, Sink, Source, StreamingQueryPlanTraverseHelper} -import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, CheckpointVersionManager, CommitLog, CommitMetadataV3, OffsetLogType, OffsetSeqBase, OffsetSeqLog, OffsetSeqMetadata, OffsetSeqMetadataV2, SinkMetadataInfo} +import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, CheckpointVersionManager, CommitLog, CommitLogType, CommitMetadataV3, OffsetLogType, OffsetSeqBase, OffsetSeqLog, OffsetSeqMetadata, OffsetSeqMetadataV2, SinkMetadataInfo} import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperatorStateInfo, StatefulOpStateStoreCheckpointInfo, StateStoreWriter} import org.apache.spark.sql.execution.streaming.runtime.StreamingCheckpointConstants.{DIR_NAME_COMMITS, DIR_NAME_OFFSETS, DIR_NAME_STATE} import org.apache.spark.sql.execution.streaming.sources.{ForeachBatchSink, WriteToMicroBatchDataSource, WriteToMicroBatchDataSourceV1} @@ -82,6 +82,7 @@ class MicroBatchExecution( -1, sparkSession, offsetLogFormatVersionOpt = None, + commitLogFormatVersionOpt = None, previousContext = None) override def getLatestExecutionContext(): StreamExecutionContext = latestExecutionContext @@ -554,9 +555,47 @@ class MicroBatchExecution( CheckpointVersionManager.setFormatVersion( sparkSessionForStream, OffsetLogType, offsetLogFormatVersion) + // Resolve the commit log format version the same way: an existing checkpoint keeps the version + // it was created with, and only a fresh one takes the version from the session config. Persist + // the implied state store checkpoint format so it agrees with what is actually being written. + val commitLogFormatVersion = CheckpointVersionManager.resolveCommitLogVersion( + sparkSessionForStream, latestCommittedBatch) + CheckpointVersionManager.setFormatVersion( + sparkSessionForStream, + CommitLogType, + commitLogFormatVersion, + latestCommittedBatch.map(_._2)) + val stateStoreCheckpointFormatVersion = + sparkSessionForStream.sessionState.conf.stateStoreCheckpointFormatVersion + + // Real-Time Mode requires state store checkpoint format v2. It writes the offset log at + // batch end + // (markMicroBatchStart is a no-op for it), so a mid-batch failure can leave durable state at a + // version that was never logged; the re-execution then rewrites that same state version. With + // checkpoint format v1 the rewritten files reuse the same names as the orphaned ones, so a load + // can pick up a stale file (see the checksum hazard documented on + // StateStoreConf.skipChecksumOnFileMissingChecksum). Format v2 avoids this because each batch + // run generates unique state store checkpoint ids. Commit log v2 persists those ids; a v3 + // commit may or may not contain them, so the resolved state store format is the authoritative + // check. Reject v1 with an escape hatch. + if (trigger.isInstanceOf[RealTimeTrigger] && stateStoreCheckpointFormatVersion < 2) { + if (!sparkSessionForStream.sessionState.conf + .getConf(SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1)) { + throw new SparkIllegalArgumentException( + errorClass = "STREAMING_REAL_TIME_MODE.CHECKPOINT_FORMAT_V1_NOT_SUPPORTED", + messageParameters = Map( + "config" -> SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1.key)) + } + logWarning(log"Starting a Real-Time Mode query on state store checkpoint format version " + + log"${MDC(LogKeys.FILE_VERSION, stateStoreCheckpointFormatVersion)} because " + + log"${MDC(LogKeys.CONFIG, SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1 + .key)} is set. A failed batch may lose data on rerun.") + } + val execCtx = new MicroBatchExecutionContext(id, runId, name, triggerClock, sources, sink, progressReporter, -1, sparkSession, offsetLogFormatVersionOpt = Some(offsetLogFormatVersion), + commitLogFormatVersionOpt = Some(commitLogFormatVersion), previousContext = None) execCtx.offsetSeqMetadata = offsetLogFormatVersion match { @@ -868,8 +907,13 @@ class MicroBatchExecution( private def verifyNewCheckpointDirectory(): Unit = { val fileManager = CheckpointFileManager.create(new Path(resolvedCheckpointRoot), sparkSession.sessionState.newHadoopConf()) - val dirNamesThatShouldNotHaveFiles = Array[String]( - DIR_NAME_OFFSETS, DIR_NAME_STATE, DIR_NAME_COMMITS) + var dirNamesThatShouldNotHaveFiles = Array[String](DIR_NAME_OFFSETS, DIR_NAME_COMMITS) + + // Since real-time mode writes the offset log after the batch is committed, the state directory + // may contain files so we want to allow the streaming query to retry. + if (!trigger.isInstanceOf[RealTimeTrigger]) { + dirNamesThatShouldNotHaveFiles :+= DIR_NAME_STATE + } dirNamesThatShouldNotHaveFiles.foreach { dirName => val path = new Path(resolvedCheckpointRoot, dirName) @@ -1190,9 +1234,19 @@ class MicroBatchExecution( // dummy string to prevent UnresolvedException and to prevent to be used in the future. CurrentBatchTimestamp(execCtx.offsetSeqMetadata.batchTimestampMs, ct.dataType, Some("Dummy TimeZoneId")) + case ct: CurrentTimestampNanos => + // Like CurrentTimestamp, the nanosecond current_timestamp(p) is not + // TimeZoneAwareExpression, so supply the dummy time zone. The batch timestamp is + // millisecond resolution, so the folded TIMESTAMP_LTZ(p) literal has zero + // nanos-within-micro. + CurrentBatchTimestamp(execCtx.offsetSeqMetadata.batchTimestampMs, + ct.dataType, Some("Dummy TimeZoneId")) case lt: LocalTimestamp => CurrentBatchTimestamp(execCtx.offsetSeqMetadata.batchTimestampMs, lt.dataType, lt.timeZoneId) + case lt: LocalTimestampNanos => + CurrentBatchTimestamp(execCtx.offsetSeqMetadata.batchTimestampMs, + lt.dataType, lt.timeZoneId) case cd: CurrentDate => CurrentBatchTimestamp(execCtx.offsetSeqMetadata.batchTimestampMs, cd.dataType, cd.timeZoneId) @@ -1228,7 +1282,8 @@ class MicroBatchExecution( execCtx.previousContext.isEmpty, currentStateStoreCkptId, stateSchemaMetadatas, - isTerminatingTrigger = trigger.isInstanceOf[AvailableNowTrigger.type]) + isTerminatingTrigger = trigger.isInstanceOf[AvailableNowTrigger.type], + isRealTimeMode = trigger.isInstanceOf[RealTimeTrigger]) execCtx.executionPlan.executedPlan // Force the lazy generation of execution plan } // Set up StateStore commit tracking before execution begins @@ -1453,6 +1508,11 @@ class MicroBatchExecution( log"Committed offsets for batch ${MDC(LogKeys.BATCH_ID, execCtx.batchId)}. Metadata " + log"${MDC(LogKeys.OFFSET_SEQUENCE_METADATA, execCtx.offsetSeqMetadata)}" ) + + // State schema validation and broadcast creation still happen during planning, but RTM + // defers operator metadata until its end offset is durable. This prevents a failed batch + // from leaving metadata that has no corresponding offset log entry. + execCtx.executionPlan.writeRecordedStateMetadata() } execCtx.reportTimeTaken("commitOffsets") { @@ -1489,7 +1549,9 @@ class MicroBatchExecution( } else { commitLog.createMetadata( nextBatchWatermarkMs = watermarkTracker.currentWatermark, - stateUniqueIds = stateStoreCkptId) + stateUniqueIds = stateStoreCkptId, + commitLogFormatVersion = execCtx.commitLogFormatVersionOpt.getOrElse( + sparkSessionForStream.sessionState.conf.streamingCommitLogFormatVersion)) } if (!commitLog.add(execCtx.batchId, metadata)) { throw QueryExecutionErrors.concurrentStreamLogUpdate(execCtx.batchId) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala index 7ae557797f36c..8b7853b7f03d6 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala @@ -19,10 +19,13 @@ package org.apache.spark.sql.execution.streaming.runtime import org.apache.spark.SparkIllegalArgumentException import org.apache.spark.internal.{Logging, LogKeys, MessageWithContext} +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, RangePartitioning} import org.apache.spark.sql.connector.catalog.Table import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.execution.streaming.operators.stateful._ +import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateExec object RealTimeModeAllowlist extends Logging { private val allowedSinks = Set( @@ -55,8 +58,27 @@ object RealTimeModeAllowlist extends Logging { "org.apache.spark.sql.execution.datasources.v2.WriteToDataSourceV2Exec", "org.apache.spark.sql.execution.exchange.BroadcastExchangeExec", "org.apache.spark.sql.execution.exchange.ReusedExchangeExec", + // A pipelined (streaming) shuffle repartitions a stateful query's input by key. In Real-Time + // Mode the DAGScheduler co-schedules the shuffle's producer and consumer stages via a + // PipelinedShuffleDependency (see IncrementalExecution's pipelined-shuffle rule), so the + // exchange is a supported member of a pipelined group rather than a materialization barrier. + "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec", "org.apache.spark.sql.execution.joins.BroadcastHashJoinExec", - classOf[EventTimeWatermarkExec].getName + // Streaming aggregation. A Real-Time Mode batch does not end when its input is exhausted, so + // an aggregation is planned as the streamline operator, which merges each input row against + // state and emits immediately, wrapped by the two buffer-projection stages that initialize + // the aggregation buffer and produce the result columns (see + // AggUtils.planStreamlineStreamingAggregation). + "org.apache.spark.sql.execution.streaming.ProjectAggregationBufferExec", + "org.apache.spark.sql.execution.streaming.StatefulStreamlineAggregateExec", + // Streaming deduplication and the state-store access operators it plans into. These run in the + // pipelined-shuffle consumer stage, keyed by the same columns the shuffle repartitions on. + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec", + "org.apache.spark.sql.execution.streaming.operators.stateful.StreamingDeduplicateExec", + classOf[EventTimeWatermarkExec].getName, + classOf[TransformWithStateExec].getName, + classOf[UpdateEventTimeColumnExec].getName ) private def classNamesString(classNames: Seq[String]): MessageWithContext = { @@ -120,12 +142,37 @@ object RealTimeModeAllowlist extends Logging { collectNodesWhoseSubtreeHasRTS(root)._2 } + /** + * Whether this operator is supported in Real-Time Mode. + * + * Membership in `allowedOperators` is the general rule, but a shuffle needs a second look: only + * some partitionings can run in Real-Time Mode. A `RangePartitioning` cannot, because building + * its `RangePartitioner` runs a separate job that samples the input to compute range bounds (see + * `ShuffleExchangeExec.prepareShuffleDependency`), and that job cannot complete while the source + * keeps producing. Rejecting it here fails the query fast with the usual allowlist error instead + * of letting it stall on a sampling job that never finishes. + */ + private def isAllowed(node: SparkPlan): Boolean = { + allowedOperators.contains(node.getClass.getName) && (node match { + case e: ShuffleExchangeExec => supportsPartitioning(e.outputPartitioning) + case _ => true + }) + } + + /** + * Whether a shuffle with this partitioning is supported in Real-Time Mode. See `isAllowed` for + * why range partitioning is not. + */ + private def supportsPartitioning(partitioning: Partitioning): Boolean = { + !partitioning.isInstanceOf[RangePartitioning] + } + def checkAllowedPhysicalOperator(operator: SparkPlan, throwException: Boolean): Unit = { val nodesToCheck = collectRealtimeNodes(operator) val violations = nodesToCheck .collect { case node => - if (allowedOperators.contains(node.getClass.getName)) { + if (isAllowed(node)) { None } else { Some(node.getClass.getName) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecution.scala index 38b4a58385ebb..cef636cc58b19 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecution.scala @@ -34,7 +34,7 @@ import org.apache.logging.log4j.CloseableThreadContext import org.apache.spark.{JobArtifactSet, SparkContext, SparkException, SparkThrowable} import org.apache.spark.internal.Logging -import org.apache.spark.internal.LogKeys.{CHECKPOINT_PATH, CHECKPOINT_ROOT, LOGICAL_PLAN, PATH, PRETTY_ID_STRING, QUERY_ID, RUN_ID, SPARK_DATA_STREAM} +import org.apache.spark.internal.LogKeys.{CHECKPOINT_PATH, CHECKPOINT_ROOT, CONFIG, LOGICAL_PLAN, PATH, PRETTY_ID_STRING, QUERY_ID, RUN_ID, SPARK_DATA_STREAM} import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.streaming.InternalOutputModes._ import org.apache.spark.sql.classic.{SparkSession, StreamingQuery} @@ -43,10 +43,11 @@ import org.apache.spark.sql.connector.read.streaming.{Offset => OffsetV2, ReadLi import org.apache.spark.sql.connector.write.{LogicalWriteInfoImpl, SupportsTruncate, Write} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.command.StreamingExplainCommand -import org.apache.spark.sql.execution.streaming.ContinuousTrigger +import org.apache.spark.sql.execution.streaming.{ContinuousTrigger, RealTimeTrigger} import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager, CommitLog, OffsetSeqLog, OffsetSeqMetadata} import org.apache.spark.sql.execution.streaming.operators.stateful.{StatefulOperator, StateStoreWriter} import org.apache.spark.sql.execution.streaming.sources.{ForeachBatchUserFuncException, ForeachUserFuncException} +import org.apache.spark.sql.execution.streaming.state.{RocksDBConf, RocksDBStateStoreProvider} import org.apache.spark.sql.execution.streaming.state.OperatorStateMetadataV2FileManager import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.SupportsStreamingUpdateAsAppend @@ -325,6 +326,10 @@ abstract class StreamExecution( sparkSessionForStream.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") } + if (trigger.isInstanceOf[RealTimeTrigger]) { + setSparkSessionConfigsForRealTimeMode(sparkSessionForStream) + } + sparkSessionForStream.conf.get(SQLConf.STATEFUL_SHUFFLE_PARTITIONS_INTERNAL) match { case Some(_) => // no-op case None => @@ -458,6 +463,86 @@ abstract class StreamExecution( } } + /** + * Applies the configuration a Real-Time Mode query needs but that is not the engine-wide default, + * because it is only the right choice for a low-latency, long-running batch. Runs once at query + * start, before the logical plan is forced, so a config read during planning sees the final + * value. + * + * Every setting here is a SOFT DEFAULT: applied only when the user has not set the key, so an + * explicit choice always wins. The state store settings, `changelogCheckpointing`, and + * `sortBeforeRepartition` are these. Except for checkpoint format v1 when its escape hatch is + * enabled, an explicit value that is incompatible with Real-Time Mode is rejected up front by the + * preflight in StreamingQueryManager (throwIfConfsAreRealTimeModeIncompatible). Therefore, an + * explicit value that reaches this method is safe to keep or was explicitly allowed. + * + * Every default applied here is logged, so those changes are recoverable from the driver log. + * + * Deliberately not set, though Databricks Runtime does default them for Real-Time Mode: + * - The incremental state-cleanup factor: that mechanism does not exist in OSS yet. + * - The Python/Pandas UDF latency knobs: those configs do not exist in OSS. + */ + private def setSparkSessionConfigsForRealTimeMode(sparkSessionForStream: SparkSession): Unit = { + val conf = sparkSessionForStream.conf + + // SOFT DEFAULTS. Real-Time Mode benefits from state store checkpoint format v2: it gives each + // batch run its own state store checkpoint ids, which prevents a re-executed batch from reusing + // the state file names of a partially-written failed batch (see the v1 hazard described at the + // fail-fast in MicroBatchExecution.initializeExecution). v2 requires the RocksDB state store + // provider. These are defaulted with two INDEPENDENT guards, matching the Databricks runtime: + // each key is set only if the user has not set that key. An explicit non-RocksDB provider has + // already been rejected by the pre-flight in StreamingQueryManager + // (throwIfConfsAreRealTimeModeIncompatible), so by the time this runs the provider is either + // unset (defaulted to RocksDB just below) or already RocksDB; the version default is applied + // independently of the provider default. + val checkpointVersionKey = SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key + if (!conf.contains(checkpointVersionKey)) { + logInfo(log"Real-Time Mode: defaulting ${MDC(CONFIG, checkpointVersionKey)}=2") + conf.set(checkpointVersionKey, CommitLog.VERSION_2.toString) + } + val providerKey = SQLConf.STATE_STORE_PROVIDER_CLASS.key + if (!conf.contains(providerKey)) { + logInfo(log"Real-Time Mode: defaulting " + + log"${MDC(CONFIG, providerKey)}=RocksDBStateStoreProvider") + conf.set(providerKey, classOf[RocksDBStateStoreProvider].getName) + } + + // SOFT DEFAULT. Changelog checkpointing writes a changelog rather than a full snapshot on each + // commit, which shortens the state-commit step at a batch boundary. That step is on the + // critical path between batches in Real-Time Mode, so this is a latency optimization. It is + // only meaningful with RocksDB, hence the check against the provider actually in force. + val changelogKey = s"${RocksDBConf.ROCKSDB_SQL_CONF_NAME_PREFIX}.changelogCheckpointing.enabled" + val usingRocksDb = + conf.get(SQLConf.STATE_STORE_PROVIDER_CLASS.key, + SQLConf.STATE_STORE_PROVIDER_CLASS.defaultValueString) == + classOf[RocksDBStateStoreProvider].getName + if (usingRocksDb && !conf.contains(changelogKey)) { + logInfo(log"Real-Time Mode: defaulting ${MDC(CONFIG, changelogKey)}=true") + conf.set(changelogKey, "true") + } + + // SOFT DEFAULT. `repartition(n)` uses RoundRobinPartitioning, and by default (SPARK-23207) + // Spark inserts a local sort before a round-robin shuffle so the row-to-partition assignment is + // deterministic -- otherwise a retried task could emit rows in a different order and lose data. + // That sort is fully blocking: it drains its whole input before emitting a row. A Real-Time + // Mode task reads an unbounded stream, so the input never ends, the producer never emits, and + // the query makes no progress at all. Note this sort is not a SortExec node (it lives inside + // ShuffleExchangeExec's RDD), so the Real-Time Mode operator allowlist cannot catch it. + // + // Determinism is not needed here: Real-Time Mode does not retry tasks (TaskSetManager caps a + // pipelined task set at one attempt, and its group is aborted as a unit rather than + // recomputed), so the hazard the sort guards against does not arise. + // + // This is a soft default rather than an override: an explicit `true` is rejected up front by + // the pre-flight in StreamingQueryManager (throwIfConfsAreRealTimeModeIncompatible), so by the + // time we get here an explicit value can only be `false`. Matches the Databricks runtime. + if (!conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key)) { + logInfo(log"Real-Time Mode: defaulting " + + log"${MDC(CONFIG, SQLConf.SORT_BEFORE_REPARTITION.key)}=false") + conf.set(SQLConf.SORT_BEFORE_REPARTITION.key, "false") + } + } + private def isInterruptedByStop(e: Throwable, sc: SparkContext): Boolean = { if (state.get == TERMINATED) { StreamExecution.isInterruptionException(e, sc) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecutionContext.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecutionContext.scala index a7061d14d3347..1c2fa064ca88f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecutionContext.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamExecutionContext.scala @@ -128,6 +128,7 @@ class MicroBatchExecutionContext( var _batchId: Long, sparkSession: SparkSession, val offsetLogFormatVersionOpt: Option[Int], + val commitLogFormatVersionOpt: Option[Int], var previousContext: Option[MicroBatchExecutionContext]) extends StreamExecutionContext( id, @@ -192,6 +193,7 @@ class MicroBatchExecutionContext( batchId + 1, sparkSession, offsetLogFormatVersionOpt, + commitLogFormatVersionOpt, Some(this)) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala new file mode 100644 index 0000000000000..e8b71adeb7466 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/EvictionIterator.scala @@ -0,0 +1,105 @@ +/* + * 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.execution.streaming.state + +import org.apache.spark.sql.catalyst.expressions.{Attribute, Predicate} +import org.apache.spark.sql.execution.streaming.operators.stateful.WatermarkSupport + +/** + * An iterator over the evictable rows of a [[StateStore]], which removes each row it returns and + * reports how far it has progressed. + * + * Only the rows that are actually evicted are returned, so every row this iterator yields has been + * removed from the store. The removal happens in `hasNext` rather than `next` (see the note on + * `pending` below) precisely so that a caller which stops early still leaves the store consistent + * with what it observed. That lets a caller both emit the evicted rows -- streaming aggregation in + * append mode outputs a grouping key once the watermark passes it -- and count real removals rather + * than state rows scanned. + */ +trait EvictionIterator extends Iterator[UnsafeRowPair] { + /** Number of state rows examined so far, whether or not they were evicted. */ + def numRowsReadDuringEvictionSoFar: Long + + /** Number of state rows removed so far. */ + def numRowsRemovedSoFar: Long +} + +object EvictionIterator { + + /** + * Returns an [[EvictionIterator]] over the rows of `store` whose event time is older than + * `evictionTimestamp`, removing each row as it is returned. + * + * The event time is read from the state store key, using the watermark metadata on + * `keyExpressions`. If those attributes carry no event time column, or `evictionTimestamp` is + * empty, nothing can be evicted and the iterator is empty. + * + * Note `evictionTimestamp` is not necessarily the current watermark: a caller doing incremental + * cleanup may pass an earlier timestamp, before which no further input can arrive. + */ + def apply( + store: StateStore, + storeIterator: Iterator[UnsafeRowPair], + keyExpressions: Seq[Attribute], + allowMultipleEventTimeColumns: Boolean, + evictionTimestamp: Option[Long]): EvictionIterator = { + + val evictionPredicate = WatermarkSupport.watermarkExpression( + WatermarkSupport.findEventTimeColumn(keyExpressions, allowMultipleEventTimeColumns), + evictionTimestamp).map { expr => + Predicate.create(expr, keyExpressions) + } + + new EvictionIterator { + private var rowsRead = 0L + private var rowsRemoved = 0L + + override def numRowsReadDuringEvictionSoFar: Long = rowsRead + override def numRowsRemovedSoFar: Long = rowsRemoved + + // The row hasNext has advanced to and already removed from the store, held so next() can + // return it. Removal happens in hasNext (not next()) so that a caller which stops iterating + // after hasNext -- without the matching next() -- still leaves the store consistent with the + // rows it was told are evicted; every row this iterator surfaces has already been removed. + private var pending: Option[UnsafeRowPair] = None + + override def hasNext: Boolean = evictionPredicate match { + case Some(predicate) => + while (pending.isEmpty && storeIterator.hasNext) { + val rowPair = storeIterator.next() + rowsRead += 1 + if (predicate.eval(rowPair.key)) { + store.remove(rowPair.key) + rowsRemoved += 1 + pending = Some(rowPair) + } + } + pending.isDefined + case None => false + } + + override def next(): UnsafeRowPair = { + if (!hasNext) { + throw new NoSuchElementException("End of the iterator") + } + val rowPair = pending.get + pending = None + rowPair + } + } + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala index 39ac1331c6103..b72959f93bcaa 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala @@ -488,12 +488,27 @@ private[sql] class HDFSBackedStateStoreProvider extends StateStoreProvider with /** Do maintenance backing data files, including creating snapshots and cleaning up old files */ override def doMaintenance(): Unit = { + doSnapshotMaintenance() + doCleanupMaintenance() + } + + /** Run only the snapshot upload portion of maintenance. */ + override def doSnapshotMaintenance(): Unit = { try { doSnapshot("maintenance") + } catch { + case NonFatal(e) => + logWarning(log"Error performing snapshot maintenance", e) + } + } + + /** Run only the cleanup portion of maintenance. */ + override def doCleanupMaintenance(): Unit = { + try { cleanup() } catch { case NonFatal(e) => - logWarning(log"Error performing snapshot and cleaning up") + logWarning(log"Error performing cleanup maintenance", e) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala index 8b23c96284506..eeb8ce298c41f 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDB.scala @@ -1658,6 +1658,36 @@ class RocksDB( } } + /** + * Return an iterator that remains open when exhausted so it can be refreshed and reused. + */ + private[state] def reusableIterator( + cfName: String = StateStore.DEFAULT_COL_FAMILY_NAME): RocksDBIterator = { + updateMemoryUsageIfNeeded() + val virtualColumnFamilyId = if (useColumnFamilies) { + Some(getColumnFamilyInfo(cfName).cfId) + } else { + None + } + val reusableIterator = new RocksDBIterator( + db.newIterator(), + useColumnFamilies, + virtualColumnFamilyId, + conf.rowChecksumEnabled, + readVerifier, + delimiterSize) + if (useColumnFamilies) { + reusableIterator.seek(Array.emptyByteArray) + } else { + reusableIterator.seekToFirst() + } + + Option(TaskContext.get()).foreach { tc => + tc.addTaskCompletionListener[Unit] { _ => reusableIterator.close() } + } + reusableIterator + } + private def countKeys(): (Long, Long) = { val iter = db.newIterator() @@ -2051,7 +2081,8 @@ class RocksDB( logInfo(log"Rolled back to ${MDC(LogKeys.VERSION_NUM, loadedVersion)}") } - def doMaintenance(): Unit = { + /** Run only the snapshot upload portion of maintenance. */ + def doSnapshotMaintenance(): Unit = { if (enableChangelogCheckpointing) { var mostRecentSnapshot: Option[RocksDBSnapshot] = None @@ -2082,6 +2113,10 @@ class RocksDB( uploadSnapshot(snapshotToUpload) } } + } + + /** Run only the cleanup portion of maintenance. */ + def doCleanupMaintenance(): Unit = { val cleanupTime = timeTakenMs { fileManager.deleteOldVersions( numVersionsToRetain = conf.minVersionsToRetain, @@ -2091,6 +2126,11 @@ class RocksDB( logInfo(log"Cleaned old data, time taken: ${MDC(LogKeys.TIME_UNITS, cleanupTime)} ms") } + def doMaintenance(): Unit = { + doSnapshotMaintenance() + doCleanupMaintenance() + } + /** * This replaces stale reused files in the snapshot with new ones to be uploaded. * Stale means they are potential candidates for deletion by another diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBIterator.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBIterator.scala new file mode 100644 index 0000000000000..f0fcd45b25589 --- /dev/null +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBIterator.scala @@ -0,0 +1,73 @@ +/* + * 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.execution.streaming.state + +import org.rocksdb.{RocksIterator => NativeRocksIterator} + +/** + * A RocksDB iterator that can be refreshed and repositioned without recreating its native + * resources. + */ +private[state] class RocksDBIterator( + iter: NativeRocksIterator, + useColumnFamilies: Boolean, + virtualColumnFamilyId: Option[Short], + rowChecksumEnabled: Boolean, + readVerifier: Option[KeyValueIntegrityVerifier], + delimiterSize: Int) extends Iterator[ByteArrayPair] with AutoCloseable { + + private val byteArrayPair = new ByteArrayPair() + + def refresh(): Unit = iter.refresh() + + def seek(key: Array[Byte]): Unit = { + val encodedKey = virtualColumnFamilyId match { + case Some(id) => RocksDBStateStoreProvider.encodeStateRowWithPrefix(key, id) + case None => key + } + iter.seek(encodedKey) + } + + def seekToFirst(): Unit = iter.seekToFirst() + + override def hasNext: Boolean = { + iter.isValid && virtualColumnFamilyId.forall { id => + RocksDBStateStoreProvider.getColumnFamilyBytesAsId(iter.key()) == id + } + } + + override def next(): ByteArrayPair = { + val key = if (useColumnFamilies) { + RocksDBStateStoreProvider.decodeStateRowWithPrefix(iter.key()) + } else { + iter.key() + } + val value = if (rowChecksumEnabled) { + KeyValueChecksumEncoder.decodeAndVerifyValueRowWithChecksum( + readVerifier, iter.key(), iter.value(), delimiterSize) + } else { + iter.value() + } + + byteArrayPair.set(key, value) + iter.next() + byteArrayPair + } + + override def close(): Unit = iter.close() +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala index 52e7ce296567a..79c69634509bb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala @@ -47,7 +47,8 @@ private[sql] class RocksDBStateStoreProvider lastVersion: Long, private[RocksDBStateStoreProvider] val stamp: Long, private[RocksDBStateStoreProvider] var readOnly: Boolean, - private[RocksDBStateStoreProvider] var forceSnapshotOnCommit: Boolean) extends StateStore { + private[RocksDBStateStoreProvider] var forceSnapshotOnCommit: Boolean) + extends StateStore with SupportsReusableIterator { private sealed trait OPERATION private case object UPDATE extends OPERATION @@ -148,10 +149,10 @@ private[sql] class RocksDBStateStoreProvider case Some(nextState) => nextState case None => val errorMsg = operation match { - case UPDATE => s"Cannot update after ${oldState.toString}" - case ABORT => s"Cannot abort after ${oldState.toString}" - case RELEASE => s"Cannot release after ${oldState.toString}" - case COMMIT => s"Cannot commit after ${oldState.toString}" + case UPDATE => s"Cannot update after ${oldState}" + case ABORT => s"Cannot abort after ${oldState}" + case RELEASE => s"Cannot release after ${oldState}" + case COMMIT => s"Cannot commit after ${oldState}" case METRICS => s"Cannot get metrics in ${oldState} state" } throw StateStoreErrors.stateStoreOperationOutOfOrder(errorMsg) @@ -490,6 +491,49 @@ private[sql] class RocksDBStateStoreProvider } } + private def wrapReusableIterator( + rocksDbIter: RocksDBIterator, + keyEncoder: RocksDBKeyStateEncoder): ReusableIterator[ByteArrayPair] = { + new ReusableIterator[ByteArrayPair] { + override def refreshAndSeekToPrefix(prefixRow: UnsafeRow): Unit = { + rocksDbIter.refresh() + val encoded = keyEncoder match { + case rangeKeyScanStateEncoder: RangeKeyScanStateEncoder => + rangeKeyScanStateEncoder.encodePrefixKey(prefixRow) + case _ => + throw new StateStoreUnsupportedOperationException( + "refreshAndSeekToPrefix", keyEncoder.getClass.getName) + } + rocksDbIter.seek(encoded) + } + + override def hasNext: Boolean = rocksDbIter.hasNext + + override def next(): ByteArrayPair = rocksDbIter.next() + + override def close(): Unit = rocksDbIter.close() + } + } + + override def reusableIterator(colFamilyName: String): ReusableIterator[UnsafeRowPair] = { + validateAndTransitionState(UPDATE) + verifyColFamilyOperations("iterator", colFamilyName) + + val kvEncoder = keyValueEncoderMap.get(colFamilyName) + val rowPair = new UnsafeRowPair() + wrapReusableIterator(rocksDB.reusableIterator(colFamilyName), kvEncoder._1).map { kv => + rowPair.withRows( + kvEncoder._1.decodeKey(kv.key), + kvEncoder._2.decodeValue(kv.value)) + if (!isValidated && rowPair.value != null && !useColumnFamilies) { + StateStoreProvider.validateStateRowFormat( + rowPair.key, keySchema, rowPair.value, valueSchema, stateStoreId, storeConf) + isValidated = true + } + rowPair + } + } + override def iteratorWithMultiValues( colFamilyName: String): StateStoreIterator[UnsafeRowPair] = { validateAndTransitionState(UPDATE) @@ -1072,15 +1116,32 @@ private[sql] class RocksDBStateStoreProvider } override def doMaintenance(): Unit = { + doSnapshotMaintenance() + doCleanupMaintenance() + } + + /** Run only the snapshot upload portion of maintenance. */ + override def doSnapshotMaintenance(): Unit = { + doMaintenanceOp(rocksDB.doSnapshotMaintenance(), "snapshot maintenance") + } + + /** Run only the cleanup portion of maintenance. */ + override def doCleanupMaintenance(): Unit = { + doMaintenanceOp(rocksDB.doCleanupMaintenance(), "cleanup maintenance") + } + + /** + * Common wrapper for maintenance operations: verifies the state machine and swallows non-fatal + * exceptions (SPARK-46547) to avoid deadlock between the maintenance thread and the streaming + * aggregation operator. + */ + private def doMaintenanceOp(op: => Unit, opName: String): Unit = { stateMachine.verifyForMaintenance() try { - rocksDB.doMaintenance() + op } catch { - // SPARK-46547 - Swallow non-fatal exception in maintenance task to avoid deadlock between - // maintenance thread and streaming aggregation operator case NonFatal(ex) => - logWarning(s"Ignoring error while performing maintenance operations with exception=", - ex) + logWarning(s"Ignoring error while performing $opName with exception=", ex) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateRewriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateRewriter.scala index 546a9a6019647..3280671851e3e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateRewriter.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateRewriter.scala @@ -381,7 +381,8 @@ class StateRewriter( // StateRewriter from writing state files in a format that disagrees with the source // checkpoint. Using the read batch commit since the latest commit could be a skipped batch. readCheckpoint.commitLog.get(readBatchId).foreach { metadata => - val configuredVersion = readCheckpoint.commitLog.defaultVersion + val configuredVersion = + sparkSession.sessionState.conf.getConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION) if (metadata.version != configuredVersion) { throw StateRewriterErrors.stateCheckpointFormatVersionMismatchError( checkpointLocationForRead, diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala index 961d2f963ce85..9c969efe93191 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala @@ -20,6 +20,8 @@ package org.apache.spark.sql.execution.streaming.state import java.io.Closeable import java.util.UUID import java.util.concurrent.{ConcurrentLinkedQueue, ScheduledFuture, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.locks.ReentrantReadWriteLock import javax.annotation.concurrent.GuardedBy import scala.collection.mutable @@ -66,6 +68,34 @@ class StateStoreIterator[A]( override def close(): Unit = onClose() } +/** + * An iterator that can be refreshed and repositioned instead of recreated for every scan. + */ +private[sql] abstract class ReusableIterator[A] extends Iterator[A] with Closeable { + /** Refresh this iterator and seek to the encoded prefix row. */ + def refreshAndSeekToPrefix(prefixRow: UnsafeRow): Unit + + override def map[B](f: A => B): ReusableIterator[B] = { + val self = this + new ReusableIterator[B] { + override def refreshAndSeekToPrefix(prefixRow: UnsafeRow): Unit = + self.refreshAndSeekToPrefix(prefixRow) + + override def hasNext: Boolean = self.hasNext + + override def next(): B = f(self.next()) + + override def close(): Unit = self.close() + } + } +} + +/** State stores that can return an iterator whose native resources are reused across scans. */ +private[sql] trait SupportsReusableIterator { + def reusableIterator( + colFamilyName: String = StateStore.DEFAULT_COL_FAMILY_NAME): ReusableIterator[UnsafeRowPair] +} + sealed trait StateStoreEncoding { override def toString: String = this match { case StateStoreEncoding.UnsafeRow => "unsaferow" @@ -82,10 +112,36 @@ sealed trait MaintenanceTaskType object MaintenanceTaskType { case object FromUnloadedProvidersQueue extends MaintenanceTaskType - case object FromTaskThread extends MaintenanceTaskType case object FromLoadedProviders extends MaintenanceTaskType } +/** + * Tracks which maintenance operations still need to run before a provider can be closed. + * Used as a tag on queue entries in `unloadedProvidersToClose`. + */ +sealed trait MaintenanceOpRequest + +object MaintenanceOpRequest { + /** All maintenance operations still need to run (e.g. query-thread-initiated unload). */ + case object All extends MaintenanceOpRequest + /** Only snapshot still needs to run (cleanup already ran as the triggering op). */ + case object Snapshot extends MaintenanceOpRequest + /** Only cleanup still needs to run (snapshot already ran as the triggering op). */ + case object Cleanup extends MaintenanceOpRequest +} + +/** + * Specifies which maintenance operation a single pool task should perform. + * Unlike MaintenanceOpRequest (which tracks remaining ops before close), + * this is the concrete op assigned to one pool thread submission. + */ +sealed trait MaintenanceOpType + +object MaintenanceOpType { + case object Snapshot extends MaintenanceOpType + case object Cleanup extends MaintenanceOpType +} + /** * Base trait for a versioned key-value store which provides read operations. Each instance of a * `ReadStateStore` represents a specific version of state data, and such instances are created @@ -831,6 +887,27 @@ case class TimestampAsPostfixKeyStateEncoderSpec(keySchema: StructType) */ trait StateStoreProvider { + // Whether this provider has been unloaded from the executor. It is read on the maintenance + // thread and set when the provider is unloaded, so maintenance does not run on a provider that + // is already being torn down. Volatile because it can be set on a query execution thread while + // being read on a maintenance thread. + @volatile var unloaded: Boolean = false + + /** + * Read-write lock for coordinating maintenance and close operations. + * - Read lock: held during snapshot/cleanup work (allows concurrent maintenance ops) + * - Write lock: held during close (waits for all maintenance to finish) + * This prevents close from racing with in-flight maintenance on the same provider. + * + * Lock ordering: maintenanceLock must be acquired before + * loadedProviders.synchronized to avoid ABBA deadlock. + * + * Passing fair=true to ensure fairness across reads and writes, + * so the write lock (close) is not starved by continuous read lock + * acquisitions (maintenance ops). + */ + val maintenanceLock: ReentrantReadWriteLock = new ReentrantReadWriteLock(true) + /** * Initialize the provide with more contextual information from the SQL operator. * This method will be called first after creating an instance of the StateStoreProvider by @@ -873,6 +950,11 @@ trait StateStoreProvider { */ def close(): Unit + /** Marks this provider as unloaded so maintenance threads stop processing it. */ + def setUnloaded(): Unit = { + unloaded = true + } + /** * Return an instance of [[StateStore]] representing state data of the given version. * If `stateStoreCkptId` is provided, the instance also needs to match the ID. @@ -923,6 +1005,12 @@ trait StateStoreProvider { /** Optional method for providers to allow for background maintenance (e.g. compactions) */ def doMaintenance(): Unit = { } + /** Run only the snapshot upload portion of maintenance. */ + def doSnapshotMaintenance(): Unit = { } + + /** Run only the cleanup portion of maintenance. */ + def doCleanupMaintenance(): Unit = { } + /** * Optional custom metrics that the implementation may want to report. * @note The StateStore objects created by this provider must report the same custom metrics @@ -1267,16 +1355,18 @@ object StateStore extends Logging { @GuardedBy("loadedProviders") private val loadedProviders = new mutable.HashMap[StateStoreProviderId, StateStoreProvider]() - private val maintenanceThreadPoolLock = new Object - private val unloadedProvidersToClose = - new ConcurrentLinkedQueue[(StateStoreProviderId, StateStoreProvider)] - - // This set is to keep track of the partitions that are queued - // for maintenance or currently have maintenance running on them - // to prevent the same partition from being processed concurrently. - @GuardedBy("maintenanceThreadPoolLock") - private val maintenancePartitions = new mutable.HashSet[StateStoreProviderId] + new ConcurrentLinkedQueue[(StateStoreProviderId, StateStoreProvider, MaintenanceOpRequest)] + + // These sets track which providers currently have maintenance tasks in-flight, + // one per operation type, to prevent concurrent same-type operations on the same provider. + // Each set has its own lock. + private val snapshotPartitionsLock = new Object + @GuardedBy("snapshotPartitionsLock") + private val snapshotPartitions = new mutable.HashSet[StateStoreProviderId] + private val cleanupPartitionsLock = new Object + @GuardedBy("cleanupPartitionsLock") + private val cleanupPartitions = new mutable.HashSet[StateStoreProviderId] /** Reports to the coordinator that a StateStore has committed */ def reportCommitToCoordinator( @@ -1307,30 +1397,68 @@ object StateStore extends Logging { * StateStoreProvider is also unloaded. Any exception that happens in the MaintenanceTask * is indeed exceptional and thus we let it propagate. */ - class MaintenanceTask(periodMs: Long, task: => Unit) { + class MaintenanceTask(periodMs: Long, task: Boolean => Unit) { private val executor = ThreadUtils.newDaemonSingleThreadScheduledExecutor("state-store-maintenance-task") + private def runTask(processUnloadedOnly: Boolean = false): Unit = { + try { + task(processUnloadedOnly) + } catch { + case NonFatal(e) => + logWarning(s"Error running maintenance task, " + + s"processUnloadedOnly=$processUnloadedOnly", e) + throw e + } + } + private val runnable = new Runnable { - override def run(): Unit = { + override def run(): Unit = runTask() + } + + private val future: ScheduledFuture[_] = executor.scheduleAtFixedRate( + runnable, periodMs, periodMs, TimeUnit.MILLISECONDS) + + private val triggerPending = new AtomicBoolean(false) + + /** + * Submit a maintenance cycle to the scheduler executor. If the scheduler is + * idle, it runs immediately. If a cycle is already running, this queues behind + * it. If a triggered run is already queued, this is a no-op. The AtomicBoolean + * ensures at most one triggered run is pending at a time. The flag resets before + * execution so a new trigger can be queued while this one is running. + * + * @param processUnloadedOnly when true (default), only drains the unload + * queue without iterating loadedProviders. This avoids submitting + * unnecessary maintenance work for all providers. + */ + def triggerNow(processUnloadedOnly: Boolean = true): Unit = { + if (triggerPending.compareAndSet(false, true)) { try { - task + executor.execute(() => { + triggerPending.set(false) + runTask(processUnloadedOnly) + }) } catch { - case NonFatal(e) => - logWarning("Error running maintenance thread", e) - throw e + // Executor already shut down by stop(). Reset the flag for completeness. + case _: java.util.concurrent.RejectedExecutionException => + logWarning("triggerNow called after scheduler maintenance task stopped") + triggerPending.set(false) } } } - private val future: ScheduledFuture[_] = executor.scheduleAtFixedRate( - runnable, periodMs, periodMs, TimeUnit.MILLISECONDS) - def stop(): Unit = { future.cancel(false) executor.shutdown() } + /** Stops the scheduler and waits for any in-flight cycle to finish. */ + def stopAndAwait(): Unit = { + stop() + executor.awaitTermination(10, TimeUnit.SECONDS) + } + def isRunning: Boolean = !future.isDone } @@ -1341,18 +1469,22 @@ object StateStore extends Logging { class MaintenanceThreadPool( numThreads: Int, shutdownTimeout: Long, - forceShutdownTimeout: Long) { - private val threadPool = ThreadUtils.newDaemonFixedThreadPool( - numThreads, "state-store-maintenance-thread") + forceShutdownTimeout: Long, + name: String) { + private val threadPool = ThreadUtils.newDaemonFixedThreadPool(numThreads, name) def execute(runnable: Runnable): Unit = { threadPool.execute(runnable) } - def stop(): Unit = { - logInfo("Shutting down MaintenanceThreadPool") + /** Initiate shutdown without waiting. Call awaitStop() to wait. */ + def shutdown(): Unit = { + logInfo(log"Shutting down MaintenanceThreadPool") threadPool.shutdown() // Disable new tasks from being submitted + } + /** Wait for threads to finish after shutdown() was called. */ + def awaitStop(): Unit = { // Wait a while for existing tasks to terminate if (!threadPool.awaitTermination(shutdownTimeout, TimeUnit.SECONDS)) { logWarning( @@ -1367,13 +1499,21 @@ object StateStore extends Logging { } } } + + def stop(): Unit = { + shutdown() + awaitStop() + } } @GuardedBy("loadedProviders") private var maintenanceTask: MaintenanceTask = null @GuardedBy("loadedProviders") - private var maintenanceThreadPool: MaintenanceThreadPool = null + private var highPriorityThreadPool: MaintenanceThreadPool = null + + @GuardedBy("loadedProviders") + private var lowPriorityThreadPool: MaintenanceThreadPool = null @GuardedBy("loadedProviders") private var _coordRef: StateStoreCoordinatorRef = null @@ -1525,16 +1665,25 @@ object StateStore extends Logging { }.getOrElse(log"") providerStatus.providerIdsToUnload.foreach(id => { loadedProviders.remove(id).foreach( provider => { - // Trigger maintenance thread to immediately do maintenance on and close the provider. - // Doing maintenance first allows us to do maintenance for a constantly-moving state - // store. - logInfo(log"Submitted maintenance from task thread to close " + + // Queue provider for maintenance + close. The scheduler will drain the queue + // and submit tasks. remove() returning non-null ensures only one queuer. + logInfo(log"Queuing provider from task thread for maintenance and close " + log"provider=${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}." + taskContextIdLogLine + log"Removed provider from loadedProviders") - submitMaintenanceWorkForProvider( - id, provider, storeConf, MaintenanceTaskType.FromTaskThread) + unloadedProvidersToClose.add((id, provider, MaintenanceOpRequest.All)) }) }) + + // Submit a scheduler cycle so queued providers are processed without + // waiting for the next periodic tick, minimizing the time stale + // providers wait to be closed. Without this, we would wait up to + // 2 maintenance cycles for both operations to finish and the + // provider to be closed from the time it is queued. At most one + // triggered cycle can be pending at a time. + if (providerStatus.providerIdsToUnload.nonEmpty && maintenanceTask != null) { + maintenanceTask.triggerNow() + } + providerStatus.shouldForceSnapshotUpload } else { false @@ -1555,29 +1704,41 @@ object StateStore extends Logging { } /** - * Unload a state store provider. - * If alreadyRemovedFromLoadedProviders is None, provider will be - * removed from loadedProviders and closed. - * If alreadyRemovedFromLoadedProviders is Some, provider will be closed - * using passed in provider. + * Close a provider and release its resources. No-op if already unloaded. * WARNING: CAN ONLY BE CALLED FROM MAINTENANCE THREAD! */ - def removeFromLoadedProvidersAndClose( + def closeProvider( storeProviderId: StateStoreProviderId, - alreadyRemovedProvider: Option[StateStoreProvider] = None): Unit = { - val providerToClose = alreadyRemovedProvider.orElse { - loadedProviders.synchronized { - loadedProviders.remove(storeProviderId) - } + provider: StateStoreProvider): Unit = { + if (provider.unloaded) { + logInfo(log"Skipping close for ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, storeProviderId)}" + + log" because provider is already unloaded") + return } - providerToClose.foreach { provider => + logInfo(log"Closing ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, storeProviderId)}") + try { provider.close() + } finally { + provider.setUnloaded() + } + logInfo(log"Closed ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, storeProviderId)}") + } + + /** + * Remove a provider from loadedProviders by key and close it. + * WARNING: CAN ONLY BE CALLED FROM MAINTENANCE THREAD! + */ + def removeFromLoadedProvidersAndClose(storeProviderId: StateStoreProviderId): Unit = { + loadedProviders.synchronized { + loadedProviders.remove(storeProviderId) + }.foreach { provider => + closeProvider(storeProviderId, provider) } } /** Unload all state store providers: unit test purpose */ private[sql] def unloadAll(): Unit = loadedProviders.synchronized { - loadedProviders.keySet.foreach { key => removeFromLoadedProvidersAndClose(key) } + loadedProviders.foreach { case (id, provider) => closeProvider(id, provider) } loadedProviders.clear() } @@ -1601,71 +1762,101 @@ object StateStore extends Logging { * it can work-around a deadlock condition where a maintenance task is waiting for the lock * */ private[streaming] def stopMaintenanceTaskWithoutLock(): Unit = { - if (maintenanceThreadPool != null) { - maintenanceThreadPoolLock.synchronized { - maintenancePartitions.clear() - } - maintenanceThreadPool.stop() - maintenanceThreadPool = null - } + // Stop the scheduler first so no new work is submitted to the pools. if (maintenanceTask != null) { - maintenanceTask.stop() + maintenanceTask.stopAndAwait() maintenanceTask = null } + // Shut down both pools concurrently, then await both, so we don't + // double the blocking time. + if (highPriorityThreadPool != null) highPriorityThreadPool.shutdown() + if (lowPriorityThreadPool != null) lowPriorityThreadPool.shutdown() + if (highPriorityThreadPool != null) { + highPriorityThreadPool.awaitStop() + highPriorityThreadPool = null + } + if (lowPriorityThreadPool != null) { + lowPriorityThreadPool.awaitStop() + lowPriorityThreadPool = null + } + snapshotPartitionsLock.synchronized { snapshotPartitions.clear() } + cleanupPartitionsLock.synchronized { cleanupPartitions.clear() } } /** Unload and stop all state store providers */ - def stop(): Unit = loadedProviders.synchronized { - loadedProviders.keySet.foreach { key => removeFromLoadedProvidersAndClose(key) } - loadedProviders.clear() - _coordRef = null - stopMaintenanceTask() + def stop(): Unit = { + // Stop scheduler and pools outside loadedProviders lock. Pool threads + // acquire maintenanceLock then loadedProviders.synchronized, so holding + // loadedProviders while awaiting termination would deadlock. + stopMaintenanceTaskWithoutLock() + loadedProviders.synchronized { + loadedProviders.foreach { case (id, provider) => closeProvider(id, provider) } + loadedProviders.clear() + _coordRef = null + } + // Drain after stopping the pool to catch anything queued during shutdown. + while (!unloadedProvidersToClose.isEmpty) { + val (id, provider, _) = unloadedProvidersToClose.poll() + closeProvider(id, provider) + } logInfo("StateStore stopped") } + /** + * Determines the number of threads for the snapshot and cleanup pools + * using the configured ratio. Snapshot gets the rounded value, clamped + * to [1, total - 1]. Cleanup gets the remainder. Each pool gets at + * least 1 thread and the total is never exceeded. + * @return (snapshotThreads, cleanupThreads) + */ + private[streaming] def getPoolSizes(storeConf: StateStoreConf): (Int, Int) = { + val total = storeConf.numStateStoreMaintenanceThreads + val ratio = storeConf.snapshotToCleanupThreadRatio + val snapshotBeforeClamp = math.round(total * ratio).toInt + // Clamp to [1, total - 1] so each pool gets at least 1 thread + // and total is never exceeded. + val snapshot = math.max(1, math.min(total - 1, snapshotBeforeClamp)) + val cleanup = total - snapshot + (snapshot, cleanup) + } + /** Start the periodic maintenance task if not already started and if Spark active */ private def startMaintenanceIfNeeded(storeConf: StateStoreConf): Unit = { - val numMaintenanceThreads = storeConf.numStateStoreMaintenanceThreads val maintenanceShutdownTimeout = storeConf.stateStoreMaintenanceShutdownTimeout val maintenanceForceShutdownTimeout = storeConf.stateStoreMaintenanceForceShutdownTimeout loadedProviders.synchronized { if (SparkEnv.get != null && !isMaintenanceRunning && !storeConf.unloadOnCommit) { maintenanceTask = new MaintenanceTask( storeConf.maintenanceInterval, - task = { doMaintenance(storeConf) } + task = { processUnloadedOnly => doMaintenance(storeConf, processUnloadedOnly) } ) - maintenanceThreadPool = new MaintenanceThreadPool(numMaintenanceThreads, - maintenanceShutdownTimeout, maintenanceForceShutdownTimeout) + // Separate pools for snapshot and cleanup to prevent one operation type + // from starving the other when pool threads are saturated. + val (snapshotThreads, cleanupThreads) = getPoolSizes(storeConf) + highPriorityThreadPool = new MaintenanceThreadPool(snapshotThreads, + maintenanceShutdownTimeout, maintenanceForceShutdownTimeout, + "state-store-maintenance-high-priority") + lowPriorityThreadPool = new MaintenanceThreadPool(cleanupThreads, + maintenanceShutdownTimeout, maintenanceForceShutdownTimeout, + "state-store-maintenance-low-priority") logInfo("State Store maintenance task started") } } } - // Wait until this partition can be processed - private def awaitProcessThisPartition( - id: StateStoreProviderId, - timeoutMs: Long): Boolean = maintenanceThreadPoolLock synchronized { - val startTime = System.currentTimeMillis() - val endTime = startTime + timeoutMs - - // If immediate processing fails, wait with timeout - var canProcessThisPartition = processThisPartition(id) - while (!canProcessThisPartition && System.currentTimeMillis() < endTime) { - maintenanceThreadPoolLock.wait(timeoutMs) - canProcessThisPartition = processThisPartition(id) - } - val elapsedTime = System.currentTimeMillis() - startTime - logInfo(log"Waited for ${MDC(LogKeys.TOTAL_TIME, elapsedTime)} ms to be able to process " + - log"maintenance for partition ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}") - canProcessThisPartition - } - private def doMaintenance(): Unit = doMaintenance(StateStoreConf.empty) - private def processThisPartition(id: StateStoreProviderId): Boolean = { - maintenanceThreadPoolLock.synchronized { - if (!maintenancePartitions.contains(id)) { - maintenancePartitions.add(id) + /** Claim a single partition set slot. Returns true if claimed. */ + private[streaming] def tryClaimPartition( + id: StateStoreProviderId, + opType: MaintenanceOpType): Boolean = { + val (partitionSet, lock) = opType match { + case MaintenanceOpType.Snapshot => (snapshotPartitions, snapshotPartitionsLock) + case MaintenanceOpType.Cleanup => (cleanupPartitions, cleanupPartitionsLock) + } + lock.synchronized { + if (!partitionSet.contains(id)) { + partitionSet.add(id) true } else { false @@ -1677,26 +1868,56 @@ object StateStore extends Logging { * Execute background maintenance task in all the loaded store providers if they are still * the active instances according to the coordinator. */ - private def doMaintenance(storeConf: StateStoreConf): Unit = { + private def doMaintenance( + storeConf: StateStoreConf, + processUnloadedOnly: Boolean = false): Unit = { logDebug("Doing maintenance") if (SparkEnv.get == null) { throw new IllegalStateException("SparkEnv not active, cannot do maintenance on StateStores") } // Providers that couldn't be processed now and need to be added back to the queue - val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]() - - // unloadedProvidersToClose are StateStoreProviders that have been removed from - // loadedProviders, and can now be processed for maintenance. This queue contains - // providers for which we weren't able to process for maintenance on the previous iteration + val providersToRequeue = + new ArrayBuffer[(StateStoreProviderId, StateStoreProvider, MaintenanceOpRequest)]() + + // Phase 1: Drain unloadedProvidersToClose queue. + // These are providers removed from loadedProviders that need maintenance before close. + // opRequest determines which task to submit: + // All: pick one available op, nextOp enqueues the other after completion + // Snapshot: submit snapshot, nextOp = None so provider is closed after + // Cleanup: submit cleanup, nextOp = None so provider is closed after while (!unloadedProvidersToClose.isEmpty) { - val (providerId, provider) = unloadedProvidersToClose.poll() + val (providerId, provider, opRequest) = unloadedProvidersToClose.poll() + + val submitted = opRequest match { + case MaintenanceOpRequest.All => + // All ops should have run recently before the provider can be + // closed. We serialize them by submitting one op now with + // nextOp pointing to the other. When the first completes, + // the pool thread enqueues the remaining op with nextOp = None, + // which closes the provider after finishing. + // Pick whichever partition set is available with short circuit + // evaluation. + tryClaimAndSubmit(providerId, provider, storeConf, + MaintenanceOpType.Snapshot, MaintenanceTaskType.FromUnloadedProvidersQueue, + nextOp = Some(otherMaintenanceOpRequest(MaintenanceOpType.Snapshot))) || + tryClaimAndSubmit(providerId, provider, storeConf, + MaintenanceOpType.Cleanup, MaintenanceTaskType.FromUnloadedProvidersQueue, + nextOp = Some(otherMaintenanceOpRequest(MaintenanceOpType.Cleanup))) + case MaintenanceOpRequest.Snapshot => + tryClaimAndSubmit( + providerId, provider, storeConf, + MaintenanceOpType.Snapshot, MaintenanceTaskType.FromUnloadedProvidersQueue) + case MaintenanceOpRequest.Cleanup => + tryClaimAndSubmit( + providerId, provider, storeConf, + MaintenanceOpType.Cleanup, MaintenanceTaskType.FromUnloadedProvidersQueue) + } - if (processThisPartition(providerId)) { - submitMaintenanceWorkForProvider( - providerId, provider, storeConf, MaintenanceTaskType.FromUnloadedProvidersQueue) - } else { - providersToRequeue += ((providerId, provider)) + // If the partition set is occupied, buffer for requeue. These will be + // added back to the queue after draining and retried on the next cycle. + if (!submitted) { + providersToRequeue += ((providerId, provider, opRequest)) } } @@ -1707,117 +1928,231 @@ object StateStore extends Logging { providersToRequeue.foreach(unloadedProvidersToClose.offer) - loadedProviders.synchronized { - loadedProviders.toSeq - }.foreach { case (id, provider) => - if (processThisPartition(id)) { - submitMaintenanceWorkForProvider( - id, provider, storeConf, MaintenanceTaskType.FromLoadedProviders) - } else { - logInfo(log"Not processing partition ${MDC(LogKeys.PARTITION_ID, id)} " + - log"for maintenance because it is currently " + - log"being processed") + // Phase 2: Submit separate snapshot and cleanup tasks for loaded providers. + // Skipped when processUnloadedOnly is true to avoid submitting + // unnecessary work for all providers. + if (!processUnloadedOnly) { + loadedProviders.synchronized { + loadedProviders.toSeq + }.foreach { case (id, provider) => + tryClaimAndSubmit( + id, provider, storeConf, + MaintenanceOpType.Snapshot, MaintenanceTaskType.FromLoadedProviders) + tryClaimAndSubmit( + id, provider, storeConf, + MaintenanceOpType.Cleanup, MaintenanceTaskType.FromLoadedProviders) } } } + /** + * Attempts to claim a partition set slot and submit maintenance work for a provider. + * Returns true if the work was submitted, false if the partition set was occupied. + */ + private def tryClaimAndSubmit( + providerId: StateStoreProviderId, + provider: StateStoreProvider, + storeConf: StateStoreConf, + opType: MaintenanceOpType, + source: MaintenanceTaskType, + // Only used when source is FromUnloadedProvidersQueue. + nextOp: Option[MaintenanceOpRequest] = None): Boolean = { + if (tryClaimPartition(providerId, opType)) { + submitMaintenanceWorkForProvider( + providerId, provider, storeConf, source, opType, nextOp) + logDebug(s"Submitted $providerId with source $source" + + s" for $opType, nextOp=$nextOp") + true + } else { + logInfo(log"Not processing partition " + + log"${MDC(LogKeys.STATE_STORE_PROVIDER_ID, providerId)} " + + log"with source ${MDC(LogKeys.MAINTENANCE_TASK_TYPE, source)} " + + log"for ${MDC(LogKeys.OP_TYPE, opType)} " + + log"because partition set is occupied") + false + } + } + + /** + * Determines the MaintenanceOpRequest for the "other" operation given the current opType. + * Used when a task needs to queue the provider for the remaining operation before close. + */ + private[streaming] def otherMaintenanceOpRequest( + opType: MaintenanceOpType): MaintenanceOpRequest = opType match { + case MaintenanceOpType.Snapshot => MaintenanceOpRequest.Cleanup + case MaintenanceOpType.Cleanup => MaintenanceOpRequest.Snapshot + } + /** * Submits maintenance work for a provider to the maintenance thread pool. * * @param id The StateStore provider ID to perform maintenance on * @param provider The StateStore provider instance + * @param storeConf The StateStore configuration + * @param source Where this request originated from + * @param opType Which maintenance operation to perform + * @param nextOp If set, the remaining op to enqueue after this one + * completes. Only used when source is FromUnloadedProvidersQueue. */ private def submitMaintenanceWorkForProvider( id: StateStoreProviderId, provider: StateStoreProvider, storeConf: StateStoreConf, - source: MaintenanceTaskType = FromLoadedProviders): Unit = { - maintenanceThreadPool.execute(() => { + source: MaintenanceTaskType, + opType: MaintenanceOpType, + // Only used when source is FromUnloadedProvidersQueue. + nextOp: Option[MaintenanceOpRequest] = None): Unit = { + val pool = opType match { + case MaintenanceOpType.Snapshot => highPriorityThreadPool + case MaintenanceOpType.Cleanup => lowPriorityThreadPool + } + pool.execute(() => { + logDebug(s"Starting $opType maintenance for $id, source=$source") val startTime = System.currentTimeMillis() - // Determine if we can process this partition based on the source - val canProcessThisPartition = source match { - case FromTaskThread => - // Provider from task thread needs to wait for lock - // We potentially need to wait for ongoing maintenance to finish processing - // this partition - val timeoutMs = storeConf.stateStoreMaintenanceProcessingTimeout * 1000 - val ableToProcessNow = awaitProcessThisPartition(id, timeoutMs) - if (!ableToProcessNow) { - // Add to queue for later processing if we can't process now - // This will be resubmitted for maintenance later by the background maintenance task - unloadedProvidersToClose.add((id, provider)) - } - ableToProcessNow - - case FromUnloadedProvidersQueue => - // Provider from queue can be processed immediately - // (we've already removed it from loadedProviders) - true - - case FromLoadedProviders => - // Provider from loadedProviders can be processed immediately - // as it's in maintenancePartitions - true - } - - if (canProcessThisPartition) { - val awaitingPartitionDuration = System.currentTimeMillis() - startTime + try { + // We use a var instead of early return because `return` inside + // a closure (pool.execute) throws NonLocalReturnControl in Scala. + var canProcess = false + // If we can't acquire the lock, the write lock is held, which + // means another thread is closing this provider. The entire + // maintenance task is a no-op in that case, so we skip and free + // the pool thread rather than blocking. The zero timeout honors + // fair ordering so readers do not starve a queued writer. + val lockAcquired = provider.maintenanceLock.readLock().tryLock(0, TimeUnit.SECONDS) try { - provider.doMaintenance() - // Handle unloading based on source - source match { - case FromTaskThread | FromUnloadedProvidersQueue => - // Provider already removed from loadedProviders, just close it - removeFromLoadedProvidersAndClose(id, Some(provider)) - - case FromLoadedProviders => - // Check if provider should be unloaded - if (!verifyIfStoreInstanceActive(id)) { - removeFromLoadedProvidersAndClose(id) + if (lockAcquired) { + canProcess = source match { + case FromLoadedProviders => + // Checks that the ID is still in loadedProviders and that the instance + // matches the one we were given. The scheduler submits from a stale copy + // of loadedProviders, so the provider may have been removed and replaced + // by a new instance under the same key. + loadedProviders.synchronized { loadedProviders.get(id).contains(provider) } && + !provider.unloaded + case _ => + // FromUnloadedProvidersQueue: provider already + // removed, reference passed directly, no containsKey needed + !provider.unloaded + } + if (canProcess) { + // Do maintenance work + opType match { + case MaintenanceOpType.Snapshot => provider.doSnapshotMaintenance() + case MaintenanceOpType.Cleanup => provider.doCleanupMaintenance() } - } - logInfo(log"Unloaded ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}") - } catch { - case NonFatal(e) => - logWarning(log"Error doing maintenance on provider:" + - log" ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}. " + - log"Could not unload state store provider", e) - // When we get a non-fatal exception, we just unload the provider. - // - // By not bubbling the exception to the maintenance task thread or the query execution - // thread, it's possible for a maintenance thread pool task to continue failing on - // the same partition. Additionally, if there is some global issue that will cause - // all maintenance thread pool tasks to fail, then bubbling the exception and - // stopping the pool is faster than waiting for all tasks to see the same exception. - // - // However, we assume that repeated failures on the same partition and global issues - // are rare. The benefit to unloading just the partition with an exception is that - // transient issues on a given provider do not affect any other providers; so, in - // most cases, this should be a more performant solution. - source match { - case FromTaskThread | FromUnloadedProvidersQueue => - removeFromLoadedProvidersAndClose(id, Some(provider)) - case FromLoadedProviders => - removeFromLoadedProvidersAndClose(id) + // Dispatch based on source. FromLoadedProviders runs inside the + // read lock so no close can interleave between work and enqueue. + // FromUnloadedProvidersQueue uses nextOp to decide whether to + // enqueue the remaining op or release the read lock and acquire + // the write lock to close the provider. + source match { + case FromLoadedProviders => + // Check if provider should be unloaded + if (!verifyIfStoreInstanceActive(id)) { + // Only remove if the map still holds the same provider instance + // we were given. Between verifyIfStoreInstanceActive and this + // remove, a concurrent get() may have loaded a new provider + // under the same key. Removing by key alone would incorrectly + // remove the new provider. + val removed = loadedProviders.synchronized { + if (loadedProviders.get(id).contains(provider)) { + loadedProviders.remove(id) + } else { + None + } + } + if (removed.isDefined) { + val remaining = otherMaintenanceOpRequest(opType) + unloadedProvidersToClose.add((id, provider, remaining)) + logInfo(log"${MDC(LogKeys.MAINTENANCE_TASK_TYPE, source)}: " + + log"${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)} verified inactive, " + + log"queued for close with ${MDC(LogKeys.OP_TYPE, remaining)}") + } else { + logInfo(log"${MDC(LogKeys.MAINTENANCE_TASK_TYPE, source)}: " + + log"${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)} verified inactive " + + log"but provider instance differs, skipping removal") + } + } + + case FromUnloadedProvidersQueue => nextOp match { + case Some(remainingOp) => + // Enqueue the remaining op. It will run with + // nextOp = None and close the provider after. + unloadedProvidersToClose.add((id, provider, remainingOp)) + logInfo(log"${MDC(LogKeys.MAINTENANCE_TASK_TYPE, source)}: queued " + + log"${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)} for close with " + + log"${MDC(LogKeys.OP_TYPE, remainingOp)}") + if (maintenanceTask != null) maintenanceTask.triggerNow() + case None => + // Release read lock, then acquire write lock to wait + // for any in-flight maintenance to finish. + provider.maintenanceLock.readLock().unlock() + provider.maintenanceLock.writeLock().lock() + try { + closeProvider(id, provider) + } finally { + // Downgrade: reacquire read lock while holding write + // lock, then release write lock. The outer finally + // unconditionally releases the read lock. + provider.maintenanceLock.readLock().lock() + provider.maintenanceLock.writeLock().unlock() + } + } + } + } else { + logInfo(log"Skipping maintenance for " + + log"${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}, " + + log"provider was removed from loadedProviders or already unloaded") } - } finally { - val duration = System.currentTimeMillis() - startTime - val logMsg = - log"Finished maintenance task for " + - log"provider=${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}" + - log" in elapsed_time=${MDC(LogKeys.TIME_UNITS, duration)}" + - log" and awaiting_partition_time=" + - log"${MDC(LogKeys.TIME_UNITS, awaitingPartitionDuration)}\n" - if (duration > 5000) { - logInfo(logMsg) } else { - logDebug(logMsg) + logDebug(s"Skipping $opType maintenance for $id, could not acquire read lock") + } + } finally { + if (lockAcquired) provider.maintenanceLock.readLock().unlock() + } + } catch { + case NonFatal(e) => + logWarning(log"Error doing maintenance on provider:" + + log" ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}. " + + log"Closing provider due to error", e) + if (source == FromLoadedProviders) { + // Only remove if the map still holds the same provider instance. + // A concurrent get() may have loaded a new provider under the same key. + loadedProviders.synchronized { + if (loadedProviders.get(id).contains(provider)) { + loadedProviders.remove(id) + } + } } - maintenanceThreadPoolLock.synchronized { - maintenancePartitions.remove(id) - maintenanceThreadPoolLock.notifyAll() + // Acquire write lock before close to wait for any concurrent + // maintenance on the other pool thread to finish. + provider.maintenanceLock.writeLock().lock() + try { + // Always close this provider instance regardless of whether we + // removed from the map. Maintenance failed, so we must clean up + // this provider's resources. We cannot rely on the queue to close + // the provider because maintenance may error again. + closeProvider(id, provider) + } finally { + provider.maintenanceLock.writeLock().unlock() } + } finally { + val duration = System.currentTimeMillis() - startTime + val logMsg = + log"Finished maintenance task for " + + log"provider=${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}" + + log" in elapsed_time=${MDC(LogKeys.TIME_UNITS, duration)}\n" + if (duration > 5000) { + logInfo(logMsg) + } else { + logDebug(logMsg) + } + opType match { + case MaintenanceOpType.Snapshot => + snapshotPartitionsLock.synchronized { snapshotPartitions.remove(id) } + case MaintenanceOpType.Cleanup => + cleanupPartitionsLock.synchronized { cleanupPartitions.remove(id) } } } }) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala index e19ac06732fa1..5da62611fda98 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStoreConf.scala @@ -31,10 +31,16 @@ class StateStoreConf( def this() = this(new SQLConf) /** - * Size of MaintenanceThreadPool to perform maintenance tasks for StateStore + * Total number of maintenance threads. Split evenly between the snapshot + * and cleanup thread pools. Each pool needs at least 1 thread, so the + * minimum is 2. */ val numStateStoreMaintenanceThreads: Int = sqlConf.numStateStoreMaintenanceThreads + /** Ratio of threads for snapshot pool. Remainder goes to cleanup. + * Each pool gets at least 1 thread and the total is never exceeded. */ + val snapshotToCleanupThreadRatio: Double = sqlConf.snapshotToCleanupThreadRatio + /** * Timeout for state store maintenance operations to complete on shutdown */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/subquery.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/subquery.scala index d5f258a8084be..25b87f7aab6fe 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/subquery.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/subquery.scala @@ -127,23 +127,36 @@ case class InSubqueryExec( override def nullable: Boolean = child.nullable override def toString: String = s"$child IN ${plan.name}" - override def withNewPlan(plan: BaseSubqueryExec): InSubqueryExec = copy(plan = plan) + override def withNewPlan(plan: BaseSubqueryExec): InSubqueryExec = + copy(plan = plan, result = null) final override def nodePatternsInternal(): Seq[TreePattern] = Seq(IN_SUBQUERY_EXEC) def updateResult(): Unit = { - val rows = plan.executeCollect() - result = if (plan.output.length > 1) { + val (rows, unavailable) = ProjectedBroadcastValueSubqueryExec.resultOf(plan) match { + case Some(BroadcastValueResult.Available(values)) => (values, false) + case Some(BroadcastValueResult.Unavailable) => (Array.empty[InternalRow], true) + case None => (plan.executeCollect(), false) + } + result = if (unavailable) { + assert(isDynamicPruning, + "An unavailable projected broadcast value domain is only supported for " + + "dynamic partition pruning.") + InSubqueryExecResultState.unavailableResult + } else if (plan.output.length > 1) { rows.asInstanceOf[Array[Any]] } else { rows.map(_.get(0, child.dataType)) } - if (!isDynamicPruning) { + if (!isDynamicPruning && !isResultUnavailable) { resultBroadcast = plan.session.sparkContext.broadcast(result) } } // This is used only by DPP where we don't need broadcast the result. - def values(): Option[Array[Any]] = Option(result) + def values(): Option[Array[Any]] = if (isResultUnavailable) None else Option(result) + + private[sql] def isResultUnavailable: Boolean = + InSubqueryExecResultState.isUnavailable(result) private def prepareResult(): Unit = { require(result != null || resultBroadcast != null, s"$this has not finished") @@ -154,12 +167,12 @@ case class InSubqueryExec( override def eval(input: InternalRow): Any = { prepareResult() - inSet.eval(input) + if (isResultUnavailable) true else inSet.eval(input) } override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { prepareResult() - inSet.doGenCode(ctx, ev) + if (isResultUnavailable) Literal.TrueLiteral.doGenCode(ctx, ev) else inSet.doGenCode(ctx, ev) } override lazy val canonicalized: InSubqueryExec = { @@ -175,6 +188,17 @@ case class InSubqueryExec( copy(child = newChild) } +private[execution] object InSubqueryExecResultState { + private case object UnavailableMarker + + def unavailableResult: Array[Any] = Array(UnavailableMarker) + + def isUnavailable(result: Array[Any]): Boolean = { + result != null && result.length == 1 && + (result(0).asInstanceOf[AnyRef] eq UnavailableMarker) + } +} + /** * Plans subqueries that are present in the given [[SparkPlan]]. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala index cdc6556f18629..e40267bd74489 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala @@ -434,7 +434,7 @@ private[window] class WindowSegmentTree( if (!evictEldest() || !acquireBlockMemory()) { // scalastyle:off throwerror throw QueryExecutionErrors.cannotAcquireMemoryForWindowAggregateError( - blockBytes, 0L) + blockBytes, 0L, taskMemoryManager.getMemoryConsumptionBreakdown()) // scalastyle:on throwerror } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala index 52c6821d00011..1ee7699fff429 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala @@ -22,7 +22,7 @@ import org.apache.spark.sql.artifact.ArtifactManager import org.apache.spark.sql.catalyst.analysis.{Analyzer, EvalSubqueriesForTimeTravel, FunctionRegistry, InvokeProcedures, ReplaceCharWithVarchar, ResolveDataSource, ResolveEventTimeWatermark, ResolveExecuteImmediate, ResolveMetricView, ResolveSessionCatalog, ResolveSetCatalogCommand, ResolveTranspose, TableFunctionRegistry} import org.apache.spark.sql.catalyst.analysis.resolver.ResolverExtension import org.apache.spark.sql.catalyst.catalog.{FunctionExpressionBuilder, SessionCatalog} -import org.apache.spark.sql.catalyst.expressions.{Expression, ExtractSemiStructuredFields} +import org.apache.spark.sql.catalyst.expressions.{Expression, ExtractSemiStructuredFields, ParseSql} import org.apache.spark.sql.catalyst.normalizer.NormalizeCTEIds import org.apache.spark.sql.catalyst.optimizer.Optimizer import org.apache.spark.sql.catalyst.parser.ParserInterface @@ -96,7 +96,12 @@ abstract class BaseSessionStateBuilder( */ protected lazy val functionRegistry: FunctionRegistry = { parentState.map(_.functionRegistry.clone()) - .getOrElse(extensions.registerFunctions(FunctionRegistry.builtin.clone())) + .getOrElse { + val registry = FunctionRegistry.builtin.clone() + // sql/core-only builtins that need SparkSqlParser. + ParseSql.register(registry) + extensions.registerFunctions(registry) + } } /** diff --git a/sql/core/src/main/scala/org/apache/spark/sql/internal/SharedState.scala b/sql/core/src/main/scala/org/apache/spark/sql/internal/SharedState.scala index 8e641294bf8cc..3271f0e1f1001 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/internal/SharedState.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/internal/SharedState.scala @@ -101,7 +101,8 @@ private[sql] class SharedState( * A relation cache backed by the cache manager. */ private[sql] val relationCache: RelationCache = { - (nameParts, resolver) => cacheManager.lookupCachedTable(nameParts, resolver) + (catalog, ident, tableId, stateOptions, resolver) => + cacheManager.lookupCachedTable(catalog, ident, tableId, stateOptions, resolver) } /** A global lock for all streaming query lifecycle tracking and management. */ diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala index bfc01f406dd97..b20faf52300ea 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala @@ -587,7 +587,7 @@ abstract class JdbcDialect extends Serializable with Logging { def schemasExists(conn: Connection, options: JDBCOptions, schema: String): Boolean = { val rs = conn.getMetaData.getSchemas(null, schema) while (rs.next()) { - if (rs.getString(1) == schema) return true; + if (rs.getString(1) == schema) return true } false } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala index b301c0c0bd5bc..ed533795329f9 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala @@ -118,6 +118,16 @@ private case class MySQLDialect() extends JdbcDialect with SQLConfHelper with No } else { super.visitAggregateFunction(funcName, isDistinct, inputs) } + + override def visitCast(expr: String, exprDataType: DataType, dataType: DataType): String = { + dataType match { + case DoubleType => + // The common JDBC mapping is DOUBLE PRECISION, which is not a portable CAST target for + // MySQL-compatible databases. In particular, MariaDB rejects it with a syntax error. + throw new UnsupportedOperationException("Cannot cast to double type") + case _ => super.visitCast(expr, exprDataType, dataType) + } + } } override def compileExpression(expr: Expression): Option[String] = { @@ -350,7 +360,7 @@ private case class MySQLDialect() extends JdbcDialect with SQLConfHelper with No } else { // The only property we are building here is `COMMENT` because it's the only one // we can get from `SHOW INDEXES`. - val properties = new util.Properties(); + val properties = new util.Properties() if (indexComment.nonEmpty) properties.put("COMMENT", indexComment) val index = new TableIndex(indexName, indexType, Array(FieldReference(colName)), new util.HashMap[NamedReference, util.Properties](), properties) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala index f5eb1fa6d7a06..46080dd57b1dd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/jdbc/OracleDialect.scala @@ -121,7 +121,7 @@ private case class OracleDialect() extends JdbcDialect with SQLConfHelper with N case (_, lit: Literal[_]) if lit.dataType == BinaryType => compareBlob(le, name, lit) case _ => - super.visitBinaryComparison(name, le, re); + super.visitBinaryComparison(name, le, re) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala index 110889563a2be..82cd74a736180 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/sources/interfaces.scala @@ -192,8 +192,8 @@ trait CreatableRelationProvider { case udt: UserDefinedType[_] => supportsDataType(udt.sqlType) case BinaryType | BooleanType | ByteType | _: CharType | DateType | _: DecimalType | DoubleType | FloatType | IntegerType | LongType | NullType | ObjectType(_) | ShortType | - _: StringType | _: TimeType | TimestampNTZType | TimestampType | - _: VarcharType => true + _: StringType | _: TimeType | _: TimestampNTZNanosType | _: TimestampLTZNanosType | + TimestampNTZType | TimestampType | _: VarcharType => true case _ => false } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala b/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala index bb8f04a8a5565..f56a6f882eda0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/util/PartitionKeyedAccumulator.scala @@ -34,12 +34,9 @@ import org.apache.spark.util.AccumulatorV2 * failed/interrupted tasks are dropped by the accumulator framework (it is not * `countFailedValues`), so only complete per-partition values are ever merged. * - * Backed by a `ConcurrentHashMap`, whose per-entry atomicity is sufficient here: `add` and the - * `putAll` in `merge` are last-write-wins per key, and the reads (`value`, - * `accumulatedNumPartitions`, `foldValues`) only require thread-safety and eventual consistency - * -- they are weakly consistent during concurrent updates but exact once all updates have been - * merged. This avoids any explicit locking (and the nested-lock pattern a two-map `merge` would - * otherwise need). + * Backed by a `ConcurrentHashMap`. Mutations and folds are synchronized so a caller can atomically + * verify that every partition completed and read its statistics from the same stable snapshot. + * Framework-facing reads remain safe, weakly consistent views. * * @tparam T the per-partition value type. Must be non-null (`ConcurrentHashMap` forbids nulls). */ @@ -52,23 +49,29 @@ class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T), java.util.Map override def copyAndReset(): PartitionKeyedAccumulator[T] = new PartitionKeyedAccumulator[T] - override def copy(): PartitionKeyedAccumulator[T] = { + override def copy(): PartitionKeyedAccumulator[T] = synchronized { val newAcc = new PartitionKeyedAccumulator[T] newAcc.byPartition.putAll(byPartition) newAcc } - override def reset(): Unit = byPartition.clear() + override def reset(): Unit = synchronized { + byPartition.clear() + } - override def add(v: (Int, T)): Unit = byPartition.put(v._1, v._2) + override def add(v: (Int, T)): Unit = synchronized { + byPartition.put(v._1, v._2) + } - override def merge(other: AccumulatorV2[(Int, T), java.util.Map[Int, T]]): Unit = other match { - case o: PartitionKeyedAccumulator[T] => - // Last-write-wins per partition id: a partition recorded by more than one task replaces - // rather than accumulates, keeping any caller-derived aggregate exact. - byPartition.putAll(o.byPartition) - case _ => throw new UnsupportedOperationException( - s"Cannot merge ${this.getClass.getName} with ${other.getClass.getName}") + override def merge(other: AccumulatorV2[(Int, T), java.util.Map[Int, T]]): Unit = synchronized { + other match { + case o: PartitionKeyedAccumulator[T] => + // Last-write-wins per partition id: a partition recorded by more than one task replaces + // rather than accumulates, keeping any caller-derived aggregate exact. + byPartition.putAll(o.byPartition) + case _ => throw new UnsupportedOperationException( + s"Cannot merge ${this.getClass.getName} with ${other.getClass.getName}") + } } // A read-only VIEW over the live map -- no copy. Only the accumulator framework calls `value` @@ -80,11 +83,27 @@ class PartitionKeyedAccumulator[T] extends AccumulatorV2[(Int, T), java.util.Map /** Number of distinct partitions that have been recorded. */ def accumulatedNumPartitions: Long = byPartition.size().toLong - /** Folds the per-partition values (each partition counted once) into a single aggregate. */ - def foldValues[A](zero: A)(op: (A, T) => A): A = { + private def foldValuesUnsafe[A](zero: A)(op: (A, T) => A): A = { var result = zero val it = byPartition.values().iterator() while (it.hasNext) result = op(result, it.next()) result } + + /** Folds the per-partition values (each partition counted once) into a single aggregate. */ + def foldValues[A](zero: A)(op: (A, T) => A): A = synchronized { + foldValuesUnsafe(zero)(op) + } + + /** Atomically checks that all partitions completed and folds their statistics. */ + def foldValuesIfComplete[A]( + expectedNumPartitions: Int, + zero: A)( + op: (A, T) => A): Option[A] = synchronized { + if (byPartition.size() == expectedNumPartitions) { + Some(foldValuesUnsafe(zero)(op)) + } else { + None + } + } } diff --git a/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/SqlResource.scala b/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/SqlResource.scala index 223b19f09dda1..42685189865fa 100644 --- a/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/SqlResource.scala +++ b/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/SqlResource.scala @@ -28,6 +28,7 @@ import jakarta.ws.rs.core.{Context, MediaType, UriInfo} import org.apache.spark.JobExecutionStatus import org.apache.spark.internal.config.UI.UI_SQL_GROUP_SUB_EXECUTION_ENABLED import org.apache.spark.sql.execution.ui.{SparkPlanGraph, SparkPlanGraphCluster, SparkPlanGraphNode, SQLAppStatusStore, SQLExecutionUIData} +import org.apache.spark.status.AppStatusStore import org.apache.spark.status.api.v1.{BaseAppResource, NotFoundException} import org.apache.spark.ui.UIUtils @@ -51,7 +52,7 @@ private[v1] class SqlResource extends BaseAppResource { } execs.map { exec => val graph = sqlStore.planGraph(exec.executionId) - prepareExecutionData(exec, graph, details, planDescription) + prepareExecutionData(exec, graph, details, planDescription, ui.store) } } } @@ -67,7 +68,10 @@ private[v1] class SqlResource extends BaseAppResource { val sqlStore = new SQLAppStatusStore(ui.store.store) sqlStore .execution(execId) - .map(prepareExecutionData(_, sqlStore.planGraph(execId), details, planDescription)) + .map { exec => + prepareExecutionData(exec, sqlStore.planGraph(execId), details, planDescription, + ui.store) + } .getOrElse(throw new NotFoundException("unknown query execution id: " + execId)) } } @@ -146,18 +150,31 @@ private[v1] class SqlResource extends BaseAppResource { val start = Option(uriParams.getFirst("start")).map(_.toInt).getOrElse(0) val length = Option(uriParams.getFirst("length")).map(_.toInt).getOrElse(20) - val sortedRoots = sortExecs(rootRows, sortCol, sortDir) + // Precompute the total task time of every root row only when the list is + // sorted by it, so the sort and the page rows reuse the same values + // instead of recomputing per stage attempt. When sorting by another + // column, `execToRow` computes it only for the rows on the current page. + val totalTaskTimeMap: Map[Long, Long] = + if (sortCol == "totalTaskTime") { + rootRows.iterator.map(e => e.executionId -> totalTaskTime(e, ui.store)).toMap + } else { + Map.empty + } + + val sortedRoots = sortExecs(rootRows, sortCol, sortDir, totalTaskTimeMap) val page = if (length > 0) sortedRoots.slice(start, start + length) else sortedRoots // Convert to Java-compatible row data; embed sub-executions when grouping. // Always emit a `subExecutions` field (possibly empty) in grouped mode so // JSON consumers see a consistent schema; flat mode never includes it. val aaData = page.map { exec => - val row = execToRow(exec) + val row = execToRow(exec, totalTaskTimeMap, ui.store) if (groupSubExec) { val subs = subsByRoot.getOrElse(exec.executionId, Seq.empty) // Sort subs by id ascending so they appear in chronological order - row.put("subExecutions", sortExecs(subs, "id", "asc").map(execToRow).asJava) + row.put("subExecutions", + sortExecs(subs, "id", "asc", totalTaskTimeMap) + .map(execToRow(_, totalTaskTimeMap, ui.store)).asJava) } row } @@ -191,7 +208,8 @@ private[v1] class SqlResource extends BaseAppResource { private def sortExecs( execs: Seq[SQLExecutionUIData], sortCol: String, - sortDir: String): Seq[SQLExecutionUIData] = { + sortDir: String, + totalTaskTimeMap: Map[Long, Long]): Seq[SQLExecutionUIData] = { val sorted = sortCol match { case "id" => execs.sortBy(_.executionId) case "status" => execs.sortBy(_.executionStatus) @@ -200,12 +218,36 @@ private[v1] class SqlResource extends BaseAppResource { case "duration" => execs.sortBy(e => e.completionTime.getOrElse(new Date()).getTime - e.submissionTime) + case "totalTaskTime" => + execs.sortBy(e => totalTaskTimeMap.getOrElse(e.executionId, -1L)) case _ => execs.sortBy(_.executionId) } if (sortDir == "asc") sorted else sorted.reverse } - private def execToRow(exec: SQLExecutionUIData): java.util.LinkedHashMap[String, Object] = { + /** + * Total task time of an execution, in milliseconds, aggregated across all + * stages of the execution. Sums `executorRunTime` (the cumulative time + * executors spent running tasks, which is the "Total Time Across All Tasks" + * stage-level metric) of every attempt of every stage: each attempt + * genuinely consumed task time, including failed attempts that were + * retried. Returns -1 when the execution has no stages to aggregate, so + * callers can distinguish "no task time information" from a genuine zero. + */ + private def totalTaskTime(exec: SQLExecutionUIData, store: AppStatusStore): Long = { + if (exec.stages.isEmpty) { + -1L + } else { + exec.stages.iterator.flatMap { stageId => + store.stageData(stageId).map(_.executorRunTime) + }.sum + } + } + + private def execToRow( + exec: SQLExecutionUIData, + totalTaskTimeMap: Map[Long, Long], + store: AppStatusStore): java.util.LinkedHashMap[String, Object] = { val duration = exec.completionTime.getOrElse(new Date()).getTime - exec.submissionTime val jobIds = exec.jobs.collect { case (id, JobExecutionStatus.SUCCEEDED) => id @@ -216,6 +258,8 @@ private[v1] class SqlResource extends BaseAppResource { row.put("description", exec.description) row.put("submissionTime", new Date(exec.submissionTime)) row.put("duration", java.lang.Long.valueOf(duration)) + row.put("totalTaskTime", java.lang.Long.valueOf( + totalTaskTimeMap.getOrElse(exec.executionId, totalTaskTime(exec, store)))) row.put("jobIds", jobIds) row.put("queryId", if (exec.queryId != null) exec.queryId.toString else null) row.put("errorMessage", exec.errorMessage.orNull) @@ -227,7 +271,8 @@ private[v1] class SqlResource extends BaseAppResource { exec: SQLExecutionUIData, graph: SparkPlanGraph, details: Boolean, - planDescription: Boolean): ExecutionData = { + planDescription: Boolean, + store: AppStatusStore): ExecutionData = { var running = Seq[Int]() var completed = Seq[Int]() @@ -267,7 +312,8 @@ private[v1] class SqlResource extends BaseAppResource { if (exec.queryId != null) exec.queryId.toString else null, exec.errorMessage.orNull, exec.rootExecutionId, - exec.modifiedConfigs) + exec.modifiedConfigs, + totalTaskTime(exec, store)) } private def printableMetrics(allNodes: collection.Seq[SparkPlanGraphNode], diff --git a/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/api.scala b/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/api.scala index 9eee17b4c1299..9bdea2ae6e69c 100644 --- a/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/api.scala +++ b/sql/core/src/main/scala/org/apache/spark/status/api/v1/sql/api.scala @@ -44,4 +44,5 @@ class ExecutionData private[spark] ( val queryId: String = null, val errorMessage: String = null, val rootExecutionId: Long = -1, - val modifiedConfigs: Map[String, String] = Map.empty) + val modifiedConfigs: Map[String, String] = Map.empty, + val totalTaskTime: Long = -1L) diff --git a/sql/core/src/test/resources/script.py b/sql/core/src/test/resources/script.py index 4fcd483f44d43..8e347f74568ab 100644 --- a/sql/core/src/test/resources/script.py +++ b/sql/core/src/test/resources/script.py @@ -17,6 +17,7 @@ # import sys + for line in sys.stdin: (a, b, c, d, e) = line.split('\t') sys.stdout.write('\t'.join([a, b, c, d, e])) diff --git a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md index 2637fbc6f077f..6f11ca0435ae9 100644 --- a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md +++ b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md @@ -44,16 +44,22 @@ | org.apache.spark.sql.catalyst.expressions.Atan2 | atan2 | SELECT atan2(0, 0) | struct<ATAN2(0, 0):double> | | org.apache.spark.sql.catalyst.expressions.Atanh | atanh | SELECT atanh(0) | struct<ATANH(0):double> | | org.apache.spark.sql.catalyst.expressions.BRound | bround | SELECT bround(2.5, 0) | struct<bround(2.5, 0):decimal(2,0)> | +| org.apache.spark.sql.catalyst.expressions.Base32 | to_base32 | SELECT to_base32('foobar') | struct<to_base32(foobar):string> | | org.apache.spark.sql.catalyst.expressions.Base64 | base64 | SELECT base64('Spark SQL') | struct<base64(Spark SQL):string> | | org.apache.spark.sql.catalyst.expressions.Between | between | SELECT 0.5 between 0.1 AND 1.0 | struct<between(0.5, 0.1, 1.0):boolean> | | org.apache.spark.sql.catalyst.expressions.Bin | bin | SELECT bin(13) | struct<bin(13):string> | | org.apache.spark.sql.catalyst.expressions.BitLength | bit_length | SELECT bit_length('Spark SQL') | struct<bit_length(Spark SQL):int> | +| org.apache.spark.sql.catalyst.expressions.BitmapAnd | bitmap_and | SELECT substring(hex(bitmap_and(X 'F0', X '70')), 0, 2) | struct<substring(hex(bitmap_and(X'F0', X'70')), 0, 2):string> | | org.apache.spark.sql.catalyst.expressions.BitmapAndAgg | bitmap_and_agg | SELECT substring(hex(bitmap_and_agg(col)), 0, 6) FROM VALUES (X 'F0'), (X '70'), (X '30') AS tab(col) | struct<substring(hex(bitmap_and_agg(col)), 0, 6):string> | +| org.apache.spark.sql.catalyst.expressions.BitmapAndNot | bitmap_andnot | SELECT substring(hex(bitmap_andnot(X 'F0', X '70')), 0, 2) | struct<substring(hex(bitmap_andnot(X'F0', X'70')), 0, 2):string> | | org.apache.spark.sql.catalyst.expressions.BitmapBitPosition | bitmap_bit_position | SELECT bitmap_bit_position(1) | struct<bitmap_bit_position(1):bigint> | | org.apache.spark.sql.catalyst.expressions.BitmapBucketNumber | bitmap_bucket_number | SELECT bitmap_bucket_number(123) | struct<bitmap_bucket_number(123):bigint> | | org.apache.spark.sql.catalyst.expressions.BitmapConstructAgg | bitmap_construct_agg | SELECT substring(hex(bitmap_construct_agg(bitmap_bit_position(col))), 0, 6) FROM VALUES (1), (2), (3) AS tab(col) | struct<substring(hex(bitmap_construct_agg(bitmap_bit_position(col))), 0, 6):string> | | org.apache.spark.sql.catalyst.expressions.BitmapCount | bitmap_count | SELECT bitmap_count(X '1010') | struct<bitmap_count(X'1010'):bigint> | +| org.apache.spark.sql.catalyst.expressions.BitmapOr | bitmap_or | SELECT substring(hex(bitmap_or(X '10', X '20')), 0, 2) | struct<substring(hex(bitmap_or(X'10', X'20')), 0, 2):string> | | org.apache.spark.sql.catalyst.expressions.BitmapOrAgg | bitmap_or_agg | SELECT substring(hex(bitmap_or_agg(col)), 0, 6) FROM VALUES (X '10'), (X '20'), (X '40') AS tab(col) | struct<substring(hex(bitmap_or_agg(col)), 0, 6):string> | +| org.apache.spark.sql.catalyst.expressions.BitmapXor | bitmap_xor | SELECT substring(hex(bitmap_xor(X 'F0', X '70')), 0, 2) | struct<substring(hex(bitmap_xor(X'F0', X'70')), 0, 2):string> | +| org.apache.spark.sql.catalyst.expressions.BitmapXorAgg | bitmap_xor_agg | SELECT substring(hex(bitmap_xor_agg(col)), 0, 6) FROM VALUES (X'10'), (X'30'), (X'40') AS tab(col) | struct<substring(hex(bitmap_xor_agg(col)), 0, 6):string> | | org.apache.spark.sql.catalyst.expressions.BitwiseAnd | & | SELECT 3 & 5 | struct<(3 & 5):int> | | org.apache.spark.sql.catalyst.expressions.BitwiseCount | bit_count | SELECT bit_count(0) | struct<bit_count(0):int> | | org.apache.spark.sql.catalyst.expressions.BitwiseGet | bit_get | SELECT bit_get(11, 0) | struct<bit_get(11, 0):tinyint> | @@ -112,7 +118,7 @@ | org.apache.spark.sql.catalyst.expressions.CurrentTime | current_time | SELECT current_time() | struct<current_time(6):time(6)> | | org.apache.spark.sql.catalyst.expressions.CurrentTime | localtime | SELECT localtime() | struct<current_time(6):time(6)> | | org.apache.spark.sql.catalyst.expressions.CurrentTimeZone | current_timezone | SELECT current_timezone() | struct<current_timezone():string> | -| org.apache.spark.sql.catalyst.expressions.CurrentTimestamp | current_timestamp | SELECT current_timestamp() | struct<current_timestamp():timestamp> | +| org.apache.spark.sql.catalyst.expressions.CurrentTimestampExpressionBuilder | current_timestamp | SELECT current_timestamp() | struct<current_timestamp():timestamp> | | org.apache.spark.sql.catalyst.expressions.CurrentUser | current_user | SELECT current_user() | struct<current_user():string> | | org.apache.spark.sql.catalyst.expressions.CurrentUser | session_user | SELECT session_user() | struct<session_user():string> | | org.apache.spark.sql.catalyst.expressions.CurrentUser | user | SELECT user() | struct<user():string> | @@ -187,6 +193,7 @@ | org.apache.spark.sql.catalyst.expressions.JsonObjectKeys | json_object_keys | SELECT json_object_keys('{}') | struct<json_object_keys({}):array<string>> | | org.apache.spark.sql.catalyst.expressions.JsonToStructs | from_json | SELECT from_json('{"a":1, "b":0.8}', 'a INT, b DOUBLE') | struct<from_json({"a":1, "b":0.8}):struct<a:int,b:double>> | | org.apache.spark.sql.catalyst.expressions.JsonTuple | json_tuple | SELECT json_tuple('{"a":1, "b":2}', 'a', 'b') | struct<c0:string,c1:string> | +| org.apache.spark.sql.catalyst.expressions.JsonTypeof | json_typeof | SELECT json_typeof('{"a": 1}') | struct<json_typeof({"a": 1}):string> | | org.apache.spark.sql.catalyst.expressions.KllSketchGetNBigint | kll_sketch_get_n_bigint | SELECT kll_sketch_get_n_bigint(kll_sketch_agg_bigint(col)) FROM VALUES (1), (2), (3), (4), (5) tab(col) | struct<kll_sketch_get_n_bigint(kll_sketch_agg_bigint(col)):bigint> | | org.apache.spark.sql.catalyst.expressions.KllSketchGetNDouble | kll_sketch_get_n_double | SELECT kll_sketch_get_n_double(kll_sketch_agg_double(col)) FROM VALUES (CAST(1.0 AS DOUBLE)), (CAST(2.0 AS DOUBLE)), (CAST(3.0 AS DOUBLE)), (CAST(4.0 AS DOUBLE)), (CAST(5.0 AS DOUBLE)) tab(col) | struct<kll_sketch_get_n_double(kll_sketch_agg_double(col)):bigint> | | org.apache.spark.sql.catalyst.expressions.KllSketchGetNFloat | kll_sketch_get_n_float | SELECT kll_sketch_get_n_float(kll_sketch_agg_float(col)) FROM VALUES (CAST(1.0 AS FLOAT)), (CAST(2.0 AS FLOAT)), (CAST(3.0 AS FLOAT)), (CAST(4.0 AS FLOAT)), (CAST(5.0 AS FLOAT)) tab(col) | struct<kll_sketch_get_n_float(kll_sketch_agg_float(col)):bigint> | @@ -217,7 +224,7 @@ | org.apache.spark.sql.catalyst.expressions.LessThanOrEqual | <= | SELECT 2 <= 2 | struct<(2 <= 2):boolean> | | org.apache.spark.sql.catalyst.expressions.Levenshtein | levenshtein | SELECT levenshtein('kitten', 'sitting') | struct<levenshtein(kitten, sitting):int> | | org.apache.spark.sql.catalyst.expressions.Like | like | SELECT like('Spark', '_park') | struct<Spark LIKE _park:boolean> | -| org.apache.spark.sql.catalyst.expressions.LocalTimestamp | localtimestamp | SELECT localtimestamp() | struct<localtimestamp():timestamp_ntz> | +| org.apache.spark.sql.catalyst.expressions.LocalTimestampExpressionBuilder | localtimestamp | SELECT localtimestamp() | struct<localtimestamp():timestamp_ntz> | | org.apache.spark.sql.catalyst.expressions.Log | ln | SELECT ln(1) | struct<ln(1):double> | | org.apache.spark.sql.catalyst.expressions.Log10 | log10 | SELECT log10(10) | struct<LOG10(10):double> | | org.apache.spark.sql.catalyst.expressions.Log1p | log1p | SELECT log1p(0) | struct<LOG1P(0):double> | @@ -259,9 +266,10 @@ | org.apache.spark.sql.catalyst.expressions.NaNvl | nanvl | SELECT nanvl(cast('NaN' as double), 123) | struct<nanvl(CAST(NaN AS DOUBLE), 123):double> | | org.apache.spark.sql.catalyst.expressions.NanosToTimestamp | timestamp_nanos | SELECT timestamp_nanos(1230219000123456789) | struct<timestamp_nanos(1230219000123456789):timestamp_ltz(9)> | | org.apache.spark.sql.catalyst.expressions.NextDay | next_day | SELECT next_day('2015-01-14', 'TU') | struct<next_day(2015-01-14, TU):date> | +| org.apache.spark.sql.catalyst.expressions.Normalize | normalize | SELECT normalize('fi', 'NFKC') | struct<normalize(fi, NFKC):string> | | org.apache.spark.sql.catalyst.expressions.Not | ! | SELECT ! true | struct<(NOT true):boolean> | | org.apache.spark.sql.catalyst.expressions.Not | not | SELECT not true | struct<(NOT true):boolean> | -| org.apache.spark.sql.catalyst.expressions.Now | now | SELECT now() | struct<now():timestamp> | +| org.apache.spark.sql.catalyst.expressions.NowExpressionBuilder | now | SELECT now() | struct<now():timestamp> | | org.apache.spark.sql.catalyst.expressions.NthValue | nth_value | SELECT a, b, nth_value(b, 2) OVER (PARTITION BY a ORDER BY b) FROM VALUES ('A1', 2), ('A1', 1), ('A2', 3), ('A1', 1) tab(a, b) | struct<a:string,b:int,nth_value(b, 2) OVER (PARTITION BY a ORDER BY b ASC NULLS FIRST RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW):int> | | org.apache.spark.sql.catalyst.expressions.NullIf | nullif | SELECT nullif(2, 2) | struct<nullif(2, 2):int> | | org.apache.spark.sql.catalyst.expressions.NullIfZero | nullifzero | SELECT nullifzero(0) | struct<nullifzero(0):int> | @@ -271,6 +279,7 @@ | org.apache.spark.sql.catalyst.expressions.OctetLength | octet_length | SELECT octet_length('Spark SQL') | struct<octet_length(Spark SQL):int> | | org.apache.spark.sql.catalyst.expressions.Or | or | SELECT true or false | struct<(true OR false):boolean> | | org.apache.spark.sql.catalyst.expressions.Overlay | overlay | SELECT overlay('Spark SQL' PLACING '_' FROM 6) | struct<overlay(Spark SQL, _, 6, -1):string> | +| org.apache.spark.sql.catalyst.expressions.ParseSql | parse_sql | SELECT parse_sql('SELECT a, b FROM t') | struct<parse_sql(SELECT a, b FROM t):string> | | org.apache.spark.sql.catalyst.expressions.ParseToDate | to_date | SELECT to_date('2009-07-30 04:17:52') | struct<to_date(2009-07-30 04:17:52):date> | | org.apache.spark.sql.catalyst.expressions.ParseToTimestamp | to_timestamp | SELECT to_timestamp('2016-12-31 00:12:00') | struct<to_timestamp(2016-12-31 00:12:00):timestamp> | | org.apache.spark.sql.catalyst.expressions.ParseToTimestampLTZExpressionBuilder | to_timestamp_ltz | SELECT to_timestamp_ltz('2016-12-31 00:12:00') | struct<to_timestamp_ltz(2016-12-31 00:12:00):timestamp> | @@ -389,8 +398,10 @@ | org.apache.spark.sql.catalyst.expressions.ToUnixTimestamp | to_unix_timestamp | SELECT to_unix_timestamp('2016-04-08', 'yyyy-MM-dd') | struct<to_unix_timestamp(2016-04-08, yyyy-MM-dd):bigint> | | org.apache.spark.sql.catalyst.expressions.TransformKeys | transform_keys | SELECT transform_keys(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> k + 1) | struct<transform_keys(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), lambdafunction((namedlambdavariable() + 1), namedlambdavariable(), namedlambdavariable())):map<int,int>> | | org.apache.spark.sql.catalyst.expressions.TransformValues | transform_values | SELECT transform_values(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), (k, v) -> v + 1) | struct<transform_values(map_from_arrays(array(1, 2, 3), array(1, 2, 3)), lambdafunction((namedlambdavariable() + 1), namedlambdavariable(), namedlambdavariable())):map<int,int>> | +| org.apache.spark.sql.catalyst.expressions.TrimArray | trim_array | SELECT trim_array(array(1, 2, 3, 4, 5), 2) | struct<trim_array(array(1, 2, 3, 4, 5), 2):array<int>> | | org.apache.spark.sql.catalyst.expressions.TruncDate | trunc | SELECT trunc('2019-08-04', 'week') | struct<trunc(2019-08-04, week):date> | | org.apache.spark.sql.catalyst.expressions.TruncTimestamp | date_trunc | SELECT date_trunc('YEAR', '2015-03-05T09:32:05.359') | struct<date_trunc(YEAR, 2015-03-05T09:32:05.359):timestamp> | +| org.apache.spark.sql.catalyst.expressions.Truncate | truncate | SELECT truncate(1234.5678, 2) | struct<truncate(1234.5678, 2):decimal(7,2)> | | org.apache.spark.sql.catalyst.expressions.TryAdd | try_add | SELECT try_add(1, 2) | struct<try_add(1, 2):int> | | org.apache.spark.sql.catalyst.expressions.TryAesDecrypt | try_aes_decrypt | SELECT try_aes_decrypt(unhex('6E7CA17BBB468D3084B5744BCA729FB7B2B7BCB8E4472847D02670489D95FA97DBBA7D3210'), '0000111122223333', 'GCM') | struct<try_aes_decrypt(unhex(6E7CA17BBB468D3084B5744BCA729FB7B2B7BCB8E4472847D02670489D95FA97DBBA7D3210), 0000111122223333, GCM, DEFAULT, ):binary> | | org.apache.spark.sql.catalyst.expressions.TryDivide | try_divide | SELECT try_divide(3, 2) | struct<try_divide(3, 2):double> | @@ -430,6 +441,7 @@ | org.apache.spark.sql.catalyst.expressions.TupleUnionThetaDoubleExpressionBuilder | tuple_union_theta_double | SELECT tuple_sketch_estimate_double(tuple_union_theta_double(tuple_sketch_agg_double(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1.0D, 4), (2, 2.0D, 5), (3, 3.0D, 6) tab(col1, val1, col2) | struct<tuple_sketch_estimate_double(tuple_union_theta_double(tuple_sketch_agg_double(col1, val1, 12, sum), theta_sketch_agg(col2, 12), 12, sum)):double> | | org.apache.spark.sql.catalyst.expressions.TupleUnionThetaIntegerExpressionBuilder | tuple_union_theta_integer | SELECT tuple_sketch_estimate_integer(tuple_union_theta_integer(tuple_sketch_agg_integer(col1, val1), theta_sketch_agg(col2))) FROM VALUES (1, 1, 4), (2, 2, 5), (3, 3, 6) tab(col1, val1, col2) | struct<tuple_sketch_estimate_integer(tuple_union_theta_integer(tuple_sketch_agg_integer(col1, val1, 12, sum), theta_sketch_agg(col2, 12), 12, sum)):double> | | org.apache.spark.sql.catalyst.expressions.TypeOf | typeof | SELECT typeof(1) | struct<typeof(1):string> | +| org.apache.spark.sql.catalyst.expressions.UnBase32 | from_base32 | SELECT from_base32('MZXW6YTBOI======') | struct<from_base32(MZXW6YTBOI======):binary> | | org.apache.spark.sql.catalyst.expressions.UnBase64 | unbase64 | SELECT unbase64('U3BhcmsgU1FM') | struct<unbase64(U3BhcmsgU1FM):binary> | | org.apache.spark.sql.catalyst.expressions.UnaryMinus | negative | SELECT negative(1) | struct<negative(1):int> | | org.apache.spark.sql.catalyst.expressions.UnaryPositive | positive | SELECT positive(1) | struct<(+ 1):int> | @@ -460,6 +472,8 @@ | org.apache.spark.sql.catalyst.expressions.WindowTime | window_time | SELECT a, window.start as start, window.end as end, window_time(window), cnt FROM (SELECT a, window, count(*) as cnt FROM VALUES ('A1', '2021-01-01 00:00:00'), ('A1', '2021-01-01 00:04:30'), ('A1', '2021-01-01 00:06:00'), ('A2', '2021-01-01 00:01:00') AS tab(a, b) GROUP by a, window(b, '5 minutes') ORDER BY a, window.start) | struct<a:string,start:timestamp,end:timestamp,window_time(window):timestamp,cnt:bigint> | | org.apache.spark.sql.catalyst.expressions.XmlToStructs | from_xml | SELECT from_xml('<p><a>1</a><b>0.8</b></p>', 'a INT, b DOUBLE') | struct<from_xml(<p><a>1</a><b>0.8</b></p>):struct<a:int,b:double>> | | org.apache.spark.sql.catalyst.expressions.XxHash64 | xxhash64 | SELECT xxhash64('Spark', array(123), 2) | struct<xxhash64(Spark, array(123), 2):bigint> | +| org.apache.spark.sql.catalyst.expressions.Xxh3128 | xxh3_128 | SELECT xxh3_128('Spark') | struct<xxh3_128(Spark):string> | +| org.apache.spark.sql.catalyst.expressions.Xxh364 | xxh3_64 | SELECT xxh3_64('Spark') | struct<xxh3_64(Spark):bigint> | | org.apache.spark.sql.catalyst.expressions.Year | year | SELECT year('2016-07-30') | struct<year(2016-07-30):int> | | org.apache.spark.sql.catalyst.expressions.ZeroIfNull | zeroifnull | SELECT zeroifnull(NULL) | struct<zeroifnull(NULL):int> | | org.apache.spark.sql.catalyst.expressions.ZipWith | zip_with | SELECT zip_with(array(1, 2, 3), array('a', 'b', 'c'), (x, y) -> (y, x)) | struct<zip_with(array(1, 2, 3), array(a, b, c), lambdafunction(named_struct(y, namedlambdavariable(), x, namedlambdavariable()), namedlambdavariable(), namedlambdavariable())):array<struct<y:string,x:int>>> | @@ -482,6 +496,7 @@ | org.apache.spark.sql.catalyst.expressions.aggregate.CollectList | array_agg | SELECT array_agg(col) FROM VALUES (1), (2), (1) AS tab(col) | struct<collect_list(col):array<int>> | | org.apache.spark.sql.catalyst.expressions.aggregate.CollectList | collect_list | SELECT collect_list(col) FROM VALUES (1), (2), (1) AS tab(col) | struct<collect_list(col):array<int>> | | org.apache.spark.sql.catalyst.expressions.aggregate.CollectSet | collect_set | SELECT collect_set(col) FROM VALUES (1), (2), (1) AS tab(col) | struct<collect_set(col):array<int>> | +| org.apache.spark.sql.catalyst.expressions.aggregate.CollectUnion | collect_union | SELECT collect_union(col) FROM VALUES (array(1, 2)), (array(2, 3)), (array(1)) AS tab(col) | struct<collect_union(col):array<int>> | | org.apache.spark.sql.catalyst.expressions.aggregate.Corr | corr | SELECT corr(c1, c2) FROM VALUES (3, 2), (3, 3), (6, 4) as tab(c1, c2) | struct<corr(c1, c2):double> | | org.apache.spark.sql.catalyst.expressions.aggregate.Count | count | SELECT count(*) FROM VALUES (NULL), (5), (5), (20) AS tab(col) | struct<count(1):bigint> | | org.apache.spark.sql.catalyst.expressions.aggregate.CountIf | count_if | SELECT count_if(col % 2 = 0) FROM VALUES (NULL), (0), (1), (2), (3) AS tab(col) | struct<count_if(((col % 2) = 0)):bigint> | @@ -561,9 +576,12 @@ | org.apache.spark.sql.catalyst.expressions.variant.TryVariantSetExpressionBuilder | try_variant_set | SELECT try_variant_set(parse_json('{"a": 1}'), '$.a', 2) | struct<try_variant_set(parse_json({"a": 1}), $.a, 2, true):variant> | | org.apache.spark.sql.catalyst.expressions.variant.VariantArrayAppendExpressionBuilder | variant_array_append | SELECT variant_array_append(parse_json('[1, 2, 3]'), '$', 4) | struct<variant_array_append(parse_json([1, 2, 3]), $, 4):variant> | | org.apache.spark.sql.catalyst.expressions.variant.VariantDelete | variant_delete | SELECT variant_delete(parse_json('{"a": 1, "b": 2, "c": 3, "items": [1, 2, 3]}'), NULL, '$.a', '$.c') | struct<variant_delete(parse_json({"a": 1, "b": 2, "c": 3, "items": [1, 2, 3]}), NULL, $.a, $.c):variant> | +| org.apache.spark.sql.catalyst.expressions.variant.VariantFromArrays | variant_from_arrays | SELECT variant_from_arrays(array('a', 'b'), array(1, 2)) | struct<variant_from_arrays(array(a, b), array(1, 2)):variant> | +| org.apache.spark.sql.catalyst.expressions.variant.VariantFromEntries | variant_from_entries | SELECT variant_from_entries(array(struct('a', 1), struct('b', 2))) | struct<variant_from_entries(array(struct(a, 1), struct(b, 2))):variant> | | org.apache.spark.sql.catalyst.expressions.variant.VariantGetExpressionBuilder | variant_get | SELECT variant_get(parse_json('{"a": 1}'), '$.a', 'int') | struct<variant_get(parse_json({"a": 1}), $.a):int> | | org.apache.spark.sql.catalyst.expressions.variant.VariantInsertExpressionBuilder | variant_insert | SELECT variant_insert(parse_json('{"a": 1}'), '$.b', 2) | struct<variant_insert(parse_json({"a": 1}), $.b, 2):variant> | | org.apache.spark.sql.catalyst.expressions.variant.VariantSetExpressionBuilder | variant_set | SELECT variant_set(parse_json('{"a": 1}'), '$.a', 2) | struct<variant_set(parse_json({"a": 1}), $.a, 2, true):variant> | +| org.apache.spark.sql.catalyst.expressions.variant.VariantStripNullsExpressionBuilder | variant_strip_nulls | SELECT variant_strip_nulls(parse_json('{"a": 1, "b": null, "c": 3}')) | struct<variant_strip_nulls(parse_json({"a": 1, "b": null, "c": 3}), true):variant> | | org.apache.spark.sql.catalyst.expressions.xml.XPathBoolean | xpath_boolean | SELECT xpath_boolean('<a><b>1</b></a>','a/b') | struct<xpath_boolean(<a><b>1</b></a>, a/b):boolean> | | org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_double | SELECT xpath_double('<a><b>1</b><b>2</b></a>', 'sum(a/b)') | struct<xpath_double(<a><b>1</b><b>2</b></a>, sum(a/b)):double> | | org.apache.spark.sql.catalyst.expressions.xml.XPathDouble | xpath_number | SELECT xpath_number('<a><b>1</b><b>2</b></a>', 'sum(a/b)') | struct<xpath_number(<a><b>1</b><b>2</b></a>, sum(a/b)):double> | diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out index ede479100dcb4..aa2795c16f85e 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/array.sql.out @@ -951,3 +951,73 @@ select array_distinct(array(0.0, -0.0, -0.0, DOUBLE("NaN"), DOUBLE("NaN"))) -- !query analysis Project [array_distinct(array(cast(0.0 as double), cast(0.0 as double), cast(0.0 as double), cast(NaN as double), cast(NaN as double))) AS array_distinct(array(0.0, 0.0, 0.0, NaN, NaN))#x] +- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 0) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 0) AS trim_array(array(1, 2, 3, 4, 5), 0)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 2) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 2) AS trim_array(array(1, 2, 3, 4, 5), 2)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 5) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 5) AS trim_array(array(1, 2, 3, 4, 5), 5)#x] ++- OneRowRelation + + +-- !query +select trim_array(array('a', 'b', 'c'), 1) +-- !query analysis +Project [trim_array(array(a, b, c), 1) AS trim_array(array(a, b, c), 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, null, 4), 1) +-- !query analysis +Project [trim_array(array(1, 2, cast(null as int), 4), 1) AS trim_array(array(1, 2, NULL, 4), 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(), 0) +-- !query analysis +Project [trim_array(array(), 0) AS trim_array(array(), 0)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), -1) +-- !query analysis +Project [trim_array(array(1, 2, 3), -1) AS trim_array(array(1, 2, 3), -1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), 4) +-- !query analysis +Project [trim_array(array(1, 2, 3), 4) AS trim_array(array(1, 2, 3), 4)#x] ++- OneRowRelation + + +-- !query +select trim_array(CAST(null AS ARRAY<INT>), 1) +-- !query analysis +Project [trim_array(cast(null as array<int>), 1) AS trim_array(NULL, 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), CAST(null AS INT)) +-- !query analysis +Project [trim_array(array(1, 2, 3), cast(null as int)) AS trim_array(array(1, 2, 3), CAST(NULL AS INT))#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar-standard-semantics.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar-standard-semantics.sql.out new file mode 100644 index 0000000000000..8081f9cecaaf0 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/charvarchar-standard-semantics.sql.out @@ -0,0 +1,1036 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT typeof(CAST('ab' AS CHAR(5))) +-- !query analysis +Project [typeof(cast(ab as char(5))) AS typeof(CAST(ab AS CHAR(5)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(CAST('hello' AS VARCHAR(5))) +-- !query analysis +Project [typeof(cast(hello as varchar(5))) AS typeof(CAST(hello AS VARCHAR(5)))#x] ++- OneRowRelation + + +-- !query +SELECT 'X' || CAST('5' AS CHAR(5)) || 'X' +-- !query analysis +Project [concat(concat(X, cast(cast(5 as char(5)) as string)), X) AS concat(concat(X, CAST(5 AS CHAR(5))), X)#x] ++- OneRowRelation + + +-- !query +SELECT CAST('ab ' AS CHAR(2)) +-- !query analysis +Project [cast(ab as char(2)) AS CAST(ab AS CHAR(2))#x] ++- OneRowRelation + + +-- !query +SELECT CAST('abcdef' AS CHAR(2)) +-- !query analysis +Project [cast(abcdef as char(2)) AS CAST(abcdef AS CHAR(2))#x] ++- OneRowRelation + + +-- !query +SELECT CAST('abcdef' AS VARCHAR(2)) +-- !query analysis +Project [cast(abcdef as varchar(2)) AS CAST(abcdef AS VARCHAR(2))#x] ++- OneRowRelation + + +-- !query +SELECT try_cast('abcdef' AS CHAR(2)) +-- !query analysis +Project [try_cast(abcdef as char(2)) AS TRY_CAST(abcdef AS CHAR(2))#x] ++- OneRowRelation + + +-- !query +SELECT try_cast('abcdef' AS VARCHAR(2)) +-- !query analysis +Project [try_cast(abcdef as varchar(2)) AS TRY_CAST(abcdef AS VARCHAR(2))#x] ++- OneRowRelation + + +-- !query +SELECT CAST(12345 AS VARCHAR(4)) +-- !query analysis +Project [cast(12345 as varchar(4)) AS CAST(12345 AS VARCHAR(4))#x] ++- OneRowRelation + + +-- !query +SELECT CAST(12345 AS VARCHAR(5)) +-- !query analysis +Project [cast(12345 as varchar(5)) AS CAST(12345 AS VARCHAR(5))#x] ++- OneRowRelation + + +-- !query +SELECT try_cast(12345 AS VARCHAR(4)) +-- !query analysis +Project [try_cast(12345 as varchar(4)) AS TRY_CAST(12345 AS VARCHAR(4))#x] ++- OneRowRelation + + +-- !query +SELECT coalesce(CAST('abcdef' AS VARCHAR(2)), CAST('x' AS VARCHAR(4))) +-- !query analysis +Project [coalesce(cast(cast(abcdef as varchar(2)) as varchar(4)), cast(x as varchar(4))) AS coalesce(CAST(abcdef AS VARCHAR(2)), CAST(x AS VARCHAR(4)))#x] ++- OneRowRelation + + +-- !query +SELECT CASE WHEN true THEN CAST('abcdef' AS VARCHAR(2)) ELSE CAST('x' AS VARCHAR(4)) END +-- !query analysis +Project [CASE WHEN true THEN cast(cast(abcdef as varchar(2)) as varchar(4)) ELSE cast(x as varchar(4)) END AS CASE WHEN true THEN CAST(abcdef AS VARCHAR(2)) ELSE CAST(x AS VARCHAR(4)) END#x] ++- OneRowRelation + + +-- !query +SELECT CAST('abcdef' AS VARCHAR(2)) IN (CAST('ab' AS VARCHAR(4))) +-- !query analysis +Project [cast(cast(abcdef as varchar(2)) as varchar(4)) IN (cast(cast(ab as varchar(4)) as varchar(4))) AS (CAST(abcdef AS VARCHAR(2)) IN (CAST(ab AS VARCHAR(4))))#x] ++- OneRowRelation + + +-- !query +SELECT coalesce( + CAST('abcdef' AS VARCHAR(2) COLLATE UTF8_LCASE), + CAST('x' AS VARCHAR(4) COLLATE UTF8_LCASE)) +-- !query analysis +Project [coalesce(cast(cast(abcdef as varchar(2) collate UTF8_LCASE) as varchar(4) collate UTF8_LCASE), cast(x as varchar(4) collate UTF8_LCASE)) AS coalesce(CAST(abcdef AS VARCHAR(2) COLLATE UTF8_LCASE), CAST(x AS VARCHAR(4) COLLATE UTF8_LCASE))#x] ++- OneRowRelation + + +-- !query +SELECT coalesce(try_cast(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))) +-- !query analysis +Project [coalesce(cast(try_cast(12345 as varchar(4)) as varchar(5)), cast(x as varchar(5))) AS coalesce(TRY_CAST(12345 AS VARCHAR(4)), CAST(x AS VARCHAR(5)))#x] ++- OneRowRelation + + +-- !query +SELECT coalesce(CAST(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))) +-- !query analysis +Project [coalesce(cast(cast(12345 as varchar(4)) as varchar(5)), cast(x as varchar(5))) AS coalesce(CAST(12345 AS VARCHAR(4)), CAST(x AS VARCHAR(5)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world' AS VARCHAR(10)))) +-- !query analysis +Project [typeof(coalesce(cast(cast(hello as varchar(5)) as varchar(10)), cast(world as varchar(10)))) AS typeof(coalesce(CAST(hello AS VARCHAR(5)), CAST(world AS VARCHAR(10))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world!' AS CHAR(6)))) +-- !query analysis +Project [typeof(coalesce(cast(cast(hello as varchar(5)) as varchar(6)), cast(cast(world! as char(6)) as varchar(6)))) AS typeof(coalesce(CAST(hello AS VARCHAR(5)), CAST(world! AS CHAR(6))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), cast('world!' AS CHAR(6)))) +-- !query analysis +Project [typeof(coalesce(cast(cast(hello as char(5)) as char(6)), cast(world! as char(6)))) AS typeof(coalesce(CAST(hello AS CHAR(5)), CAST(world! AS CHAR(6))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), 'world')) +-- !query analysis +Project [typeof(coalesce(cast(cast(hello as varchar(5)) as string collate UTF8_BINARY), world)) AS typeof(coalesce(CAST(hello AS VARCHAR(5)), world))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), NULL)) +-- !query analysis +Project [typeof(coalesce(cast(hello as char(5)), cast(null as char(5)))) AS typeof(coalesce(CAST(hello AS CHAR(5)), NULL))#x] ++- OneRowRelation + + +-- !query +SELECT typeof( + CASE WHEN true THEN cast('a' AS CHAR(2)) ELSE cast('bb' AS CHAR(4)) END) +-- !query analysis +Project [typeof(CASE WHEN true THEN cast(cast(a as char(2)) as char(4)) ELSE cast(bb as char(4)) END) AS typeof(CASE WHEN true THEN CAST(a AS CHAR(2)) ELSE CAST(bb AS CHAR(4)) END)#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS CHAR(2)), cast('bbb' AS VARCHAR(3))) +-- !query analysis +Project [cast(cast(a as char(2)) as varchar(3)) IN (cast(cast(a as char(2)) as varchar(3)),cast(cast(bbb as varchar(3)) as varchar(3))) AS (CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(2)), CAST(bbb AS VARCHAR(3))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(c) FROM (SELECT cast('a' AS CHAR(2)) AS c) t WHERE c IN ('a ', 'b') +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- Filter cast(c#x as string collate UTF8_BINARY) IN (cast(a as string collate UTF8_BINARY),cast(b as string collate UTF8_BINARY)) + +- SubqueryAlias t + +- Project [cast(a as char(2)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(upper(cast('ab' AS CHAR(2)))) +-- !query analysis +Project [typeof(upper(cast(cast(ab as char(2)) as string))) AS typeof(upper(CAST(ab AS CHAR(2))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(lower(cast('AB' AS VARCHAR(2)))) +-- !query analysis +Project [typeof(lower(cast(cast(AB as varchar(2)) as string))) AS typeof(lower(CAST(AB AS VARCHAR(2))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(cast('a' AS CHAR(1)) || cast('b' AS VARCHAR(1))) +-- !query analysis +Project [typeof(concat(cast(cast(a as char(1)) as string), cast(cast(b as varchar(1)) as string))) AS typeof(concat(CAST(a AS CHAR(1)), CAST(b AS VARCHAR(1))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(substr(cast('hello' AS VARCHAR(5)), 1, 2)) +-- !query analysis +Project [typeof(substr(cast(cast(hello as varchar(5)) as string), 1, 2)) AS typeof(substr(CAST(hello AS VARCHAR(5)), 1, 2))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(upper(coalesce(cast('a' AS CHAR(2)), cast('b' AS CHAR(4))))) +-- !query analysis +Project [typeof(upper(cast(coalesce(cast(cast(a as char(2)) as char(4)), cast(b as char(4))) as string))) AS typeof(upper(coalesce(CAST(a AS CHAR(2)), CAST(b AS CHAR(4)))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(concat(cast('a' AS CHAR(2)), cast('b' AS CHAR(3)))) +-- !query analysis +Project [typeof(concat(cast(cast(a as char(2)) as string), cast(cast(b as char(3)) as string))) AS typeof(concat(CAST(a AS CHAR(2)), CAST(b AS CHAR(3))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(concat(cast(cast(cast(a as char(2) collate UTF8_LCASE) as char(3) collate UTF8_LCASE) as string collate UTF8_LCASE), cast(cast(b as char(3) collate UTF8_LCASE) as string collate UTF8_LCASE))) AS typeof(concat(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(b AS CHAR(3) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT concat('<', concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE)), '>') +-- !query analysis +Project [concat(<, concat(cast(cast(cast(a as char(2) collate UTF8_LCASE) as char(3) collate UTF8_LCASE) as string collate UTF8_LCASE), cast(cast(b as char(3) collate UTF8_LCASE) as string collate UTF8_LCASE)), >) AS concat('<' collate UTF8_LCASE, concat(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(b AS CHAR(3) COLLATE UTF8_LCASE)), '>' collate UTF8_LCASE)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(elt(1, cast(cast(ab as char(5) collate UTF8_LCASE) as string collate UTF8_LCASE), cast(cast(cast(x as char(1) collate UTF8_LCASE) as char(5) collate UTF8_LCASE) as string collate UTF8_LCASE), true)) AS typeof(elt(1, CAST(ab AS CHAR(5) COLLATE UTF8_LCASE), CAST(x AS CHAR(1) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT concat('<', elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE)), '>') +-- !query analysis +Project [concat(<, elt(1, cast(cast(ab as char(5) collate UTF8_LCASE) as string collate UTF8_LCASE), cast(cast(cast(x as char(1) collate UTF8_LCASE) as char(5) collate UTF8_LCASE) as string collate UTF8_LCASE), true), >) AS concat('<' collate UTF8_LCASE, elt(1, CAST(ab AS CHAR(5) COLLATE UTF8_LCASE), CAST(x AS CHAR(1) COLLATE UTF8_LCASE)), '>' collate UTF8_LCASE)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(trim(cast('ab ' AS CHAR(4)))) +-- !query analysis +Project [typeof(trim(cast(cast(ab as char(4)) as string), None)) AS typeof(trim(CAST(ab AS CHAR(4))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(lpad(cast('ab' AS CHAR(2)), 5, 'x')) +-- !query analysis +Project [typeof(lpad(cast(cast(ab as char(2)) as string), 5, x)) AS typeof(lpad(CAST(ab AS CHAR(2)), 5, x))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(regexp_replace(cast('ab' AS CHAR(2)), 'a', 'x')) +-- !query analysis +Project [typeof(regexp_replace(cast(cast(ab as char(2)) as string), a, x, 1)) AS typeof(regexp_replace(CAST(ab AS CHAR(2)), a, x, 1))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(regexp_extract(cast('ab' AS VARCHAR(2)), '(a)', 1)) +-- !query analysis +Project [typeof(regexp_extract(cast(cast(ab as varchar(2)) as string), (a), 1)) AS typeof(regexp_extract(CAST(ab AS VARCHAR(2)), (a), 1))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(regexp_extract_all(cast('aab' AS VARCHAR(3)), '(a)', 1)) +-- !query analysis +Project [typeof(regexp_extract_all(cast(cast(aab as varchar(3)) as string), (a), 1)) AS typeof(regexp_extract_all(CAST(aab AS VARCHAR(3)), (a), 1))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(split(cast('a,b' AS CHAR(3)), ',')) +-- !query analysis +Project [typeof(split(cast(cast(a,b as char(3)) as string), ,, -1)) AS typeof(split(CAST(a,b AS CHAR(3)), ,, -1))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(mask(cast('ab' AS CHAR(2)))) +-- !query analysis +Project [typeof(mask(cast(cast(ab as char(2)) as string), X, x, n, null)) AS typeof(mask(CAST(ab AS CHAR(2)), X, x, n, NULL))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1)) +-- !query analysis +Project [typeof(overlay(cast(cast(ab as char(5)) as string), x, 1, -1)) AS typeof(overlay(CAST(ab AS CHAR(5)), x, 1, -1))#x] ++- OneRowRelation + + +-- !query +SELECT concat('<', overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1), '>') +-- !query analysis +Project [concat(<, overlay(cast(cast(ab as char(5)) as string), x, 1, -1), >) AS concat(<, overlay(CAST(ab AS CHAR(5)), x, 1, -1), >)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(elt(1, cast('ab' AS CHAR(5)), 'x')) +-- !query analysis +Project [typeof(elt(1, cast(cast(ab as char(5)) as string), x, true)) AS typeof(elt(1, CAST(ab AS CHAR(5)), x))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(right(cast('ab' AS CHAR(5)), 2)) +-- !query analysis +Project [typeof(right(cast(cast(ab as char(5)) as string), 2)) AS typeof(right(CAST(ab AS CHAR(5)), 2))#x] ++- OneRowRelation + + +-- !query +SELECT concat('<', right(cast('ab' AS CHAR(5)), 2), '>') +-- !query analysis +Project [concat(<, right(cast(cast(ab as char(5)) as string), 2), >) AS concat(<, right(CAST(ab AS CHAR(5)), 2), >)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(left(cast('ab' AS CHAR(5)), 2)) +-- !query analysis +Project [typeof(left(cast(cast(ab as char(5)) as string), 2)) AS typeof(left(CAST(ab AS CHAR(5)), 2))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(reverse(cast('ab' AS CHAR(5)))) +-- !query analysis +Project [typeof(reverse(cast(cast(ab as char(5)) as string))) AS typeof(reverse(CAST(ab AS CHAR(5))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(hex(cast('ab' AS CHAR(5)))) +-- !query analysis +Project [typeof(hex(cast(cast(ab as char(5)) as string))) AS typeof(hex(CAST(ab AS CHAR(5))))#x] ++- OneRowRelation + + +-- !query +SELECT hex(cast('ab' AS CHAR(5))) +-- !query analysis +Project [hex(cast(cast(ab as char(5)) as string)) AS hex(CAST(ab AS CHAR(5)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-')) +-- !query analysis +Project [typeof(array_join(cast(array(cast(ab as char(5)), cast(cd as char(5))) as array<string>), -, None)) AS typeof(array_join(array(CAST(ab AS CHAR(5)), CAST(cd AS CHAR(5))), -))#x] ++- OneRowRelation + + +-- !query +SELECT concat('<', array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-'), '>') +-- !query analysis +Project [concat(<, array_join(cast(array(cast(ab as char(5)), cast(cd as char(5))) as array<string>), -, None), >) AS concat(<, array_join(array(CAST(ab AS CHAR(5)), CAST(cd AS CHAR(5))), -), >)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(reverse(array(1, 2))) +-- !query analysis +Project [typeof(reverse(array(1, 2))) AS typeof(reverse(array(1, 2)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(str_to_map(cast('a:1,b:2' AS CHAR(7)))) +-- !query analysis +Project [typeof(str_to_map(cast(cast(a:1,b:2 as char(7)) as string), ,, :)) AS typeof(str_to_map(CAST(a:1,b:2 AS CHAR(7)), ,, :))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(c0) FROM (SELECT json_tuple(cast('{"a":"1"}' AS CHAR(9)), 'a') AS c0) +-- !query analysis +Project [typeof(c0#x) AS typeof(c0)#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [c0#x] + +- Generate json_tuple(cast(cast({"a":"1"} as char(9)) as string), a), false, [c0#x] + +- OneRowRelation + + +-- !query +SELECT typeof(cast('a' AS CHAR(2) COLLATE UTF8_LCASE)) +-- !query analysis +Project [typeof(cast(a as char(2) collate UTF8_LCASE)) AS typeof(CAST(a AS CHAR(2) COLLATE UTF8_LCASE))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(2) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(coalesce(cast(a as char(2) collate UTF8_LCASE), cast(bb as char(2) collate UTF8_LCASE))) AS typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(2) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(coalesce(cast(cast(a as char(2) collate UTF8_LCASE) as char(4) collate UTF8_LCASE), cast(bb as char(4) collate UTF8_LCASE))) AS typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(4) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query analysis +Project [hex(cast(coalesce(cast(cast(a as char(2) collate UTF8_LCASE) as char(4) collate UTF8_LCASE), cast(bb as char(4) collate UTF8_LCASE)) as string collate UTF8_LCASE)) AS hex(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(4) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS VARCHAR(4) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(coalesce(cast(cast(a as char(2) collate UTF8_LCASE) as varchar(4) collate UTF8_LCASE), cast(bb as varchar(4) collate UTF8_LCASE))) AS typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS VARCHAR(4) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query analysis +Project [typeof(coalesce(cast(cast(a as char(2) collate UTF8_LCASE) as char(4) collate UTF8_LCASE), cast(1 as char(4) collate UTF8_LCASE))) AS typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(1 AS CHAR(4) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query analysis +Project [hex(cast(coalesce(cast(cast(a as char(2) collate UTF8_LCASE) as char(4) collate UTF8_LCASE), cast(1 as char(4) collate UTF8_LCASE)) as string collate UTF8_LCASE)) AS hex(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(1 AS CHAR(4) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS VARCHAR(3)) AS c + UNION ALL + SELECT cast('abcd' AS VARCHAR(8)) AS c +) t LIMIT 1 +-- !query analysis +GlobalLimit 1 ++- LocalLimit 1 + +- Project [typeof(c#x) AS typeof(c)#x] + +- SubqueryAlias t + +- Union false, false + :- Project [cast(c#x as varchar(8)) AS c#x] + : +- Project [cast(a as varchar(3)) AS c#x] + : +- OneRowRelation + +- Project [cast(abcd as varchar(8)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t LIMIT 1 +-- !query analysis +GlobalLimit 1 ++- LocalLimit 1 + +- Project [typeof(c#x) AS typeof(c)#x] + +- SubqueryAlias t + +- Union false, false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(a as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(bb as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t +-- !query analysis +Project [concat(<, cast(c#x as string), >) AS concat(<, c, >)#x] ++- SubqueryAlias t + +- Union false, false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(a as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(bb as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- SubqueryAlias t + +- Distinct + +- Union false, false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(a as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(a as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t +-- !query analysis +Project [concat(<, cast(c#x as string), >) AS concat(<, c, >)#x] ++- SubqueryAlias t + +- Distinct + +- Union false, false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(a as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(a as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- SubqueryAlias t + +- Intersect false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(ab as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(ab as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t +-- !query analysis +Project [concat(<, cast(c#x as string), >) AS concat(<, c, >)#x] ++- SubqueryAlias t + +- Intersect false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(ab as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(ab as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- SubqueryAlias t + +- Except false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(ab as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(xy as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t +-- !query analysis +Project [concat(<, cast(c#x as string), >) AS concat(<, c, >)#x] ++- SubqueryAlias t + +- Except false + :- Project [cast(c#x as char(4)) AS c#x] + : +- Project [cast(ab as char(2)) AS c#x] + : +- OneRowRelation + +- Project [cast(xy as char(4)) AS c#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c) FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c) +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- SubqueryAlias t + +- Project [col1#x AS c#x] + +- LocalRelation [col1#x] + + +-- !query +SELECT concat('<', c, '>') FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c) +-- !query analysis +Project [concat(<, cast(c#x as string), >) AS concat(<, c, >)#x] ++- SubqueryAlias t + +- Project [col1#x AS c#x] + +- LocalRelation [col1#x] + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a' AS CHAR(4)) +-- !query analysis +Project [(cast(cast(a as char(2)) as char(4)) = cast(a as char(4))) AS (CAST(a AS CHAR(2)) = CAST(a AS CHAR(4)))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a' AS VARCHAR(2)) +-- !query analysis +Project [(cast(cast(a as char(2)) as varchar(2)) = cast(a as varchar(2))) AS (CAST(a AS CHAR(2)) = CAST(a AS VARCHAR(2)))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a ' AS VARCHAR(2)) +-- !query analysis +Project [(cast(cast(a as char(2)) as varchar(2)) = cast(a as varchar(2))) AS (CAST(a AS CHAR(2)) = CAST(a AS VARCHAR(2)))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) = 'a' +-- !query analysis +Project [(cast(cast(a as char(2)) as string collate UTF8_BINARY) = a) AS (CAST(a AS CHAR(2)) = a)#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) = 'a ' +-- !query analysis +Project [(cast(cast(a as char(2)) as string collate UTF8_BINARY) = a ) AS (CAST(a AS CHAR(2)) = a )#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = 'a' +-- !query analysis +Project [(cast(a as char(2) collate UTF8_BINARY_RTRIM) = cast(a as char(2) collate UTF8_BINARY_RTRIM)) AS (CAST(a AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = a)#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = + cast('a' AS CHAR(4) COLLATE UTF8_BINARY_RTRIM) +-- !query analysis +Project [(cast(cast(a as char(2) collate UTF8_BINARY_RTRIM) as char(4) collate UTF8_BINARY_RTRIM) = cast(a as char(4) collate UTF8_BINARY_RTRIM)) AS (CAST(a AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = CAST(a AS CHAR(4) COLLATE UTF8_BINARY_RTRIM))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4))) +-- !query analysis +Project [cast(cast(a as char(2)) as char(4)) IN (cast(cast(a as char(4)) as char(4))) AS (CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(4))))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS VARCHAR(2))) +-- !query analysis +Project [cast(cast(a as char(2)) as varchar(2)) IN (cast(cast(a as varchar(2)) as varchar(2))) AS (CAST(a AS CHAR(2)) IN (CAST(a AS VARCHAR(2))))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS VARCHAR(2))) +-- !query analysis +Project [cast(cast(a as char(2)) as varchar(2)) IN (cast(cast(a as varchar(2)) as varchar(2))) AS (CAST(a AS CHAR(2)) IN (CAST(a AS VARCHAR(2))))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN ('a', 'b') +-- !query analysis +Project [cast(cast(a as char(2)) as string collate UTF8_BINARY) IN (cast(a as string collate UTF8_BINARY),cast(b as string collate UTF8_BINARY)) AS (CAST(a AS CHAR(2)) IN (a, b))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN ('a ', 'b') +-- !query analysis +Project [cast(cast(a as char(2)) as string collate UTF8_BINARY) IN (cast(a as string collate UTF8_BINARY),cast(b as string collate UTF8_BINARY)) AS (CAST(a AS CHAR(2)) IN (a , b))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4)), cast('b' AS VARCHAR(3))) +-- !query analysis +Project [cast(cast(a as char(2)) as varchar(4)) IN (cast(cast(a as char(4)) as varchar(4)),cast(cast(b as varchar(3)) as varchar(4))) AS (CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(4)), CAST(b AS VARCHAR(3))))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) = cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE) +-- !query analysis +Project [(cast(cast(a as char(2) collate UTF8_LCASE) as varchar(2) collate UTF8_LCASE) = cast(a as varchar(2) collate UTF8_LCASE)) AS (CAST(a AS CHAR(2) COLLATE UTF8_LCASE) = CAST(a AS VARCHAR(2) COLLATE UTF8_LCASE))#x] ++- OneRowRelation + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) IN (cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE)) +-- !query analysis +Project [cast(cast(a as char(2) collate UTF8_LCASE) as varchar(2) collate UTF8_LCASE) IN (cast(a as varchar(2) collate UTF8_LCASE)) AS (CAST(a AS CHAR(2) COLLATE UTF8_LCASE) IN (CAST(a AS VARCHAR(2) COLLATE UTF8_LCASE)))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('123' AS CHAR(3)), 1)), + typeof(coalesce(cast('123' AS VARCHAR(3)), 1)), + typeof(coalesce(cast('123' AS STRING), 1)) +-- !query analysis +Project [typeof(coalesce(cast(cast(123 as char(3)) as bigint), cast(1 as bigint))) AS typeof(coalesce(CAST(123 AS CHAR(3)), 1))#x, typeof(coalesce(cast(cast(123 as varchar(3)) as bigint), cast(1 as bigint))) AS typeof(coalesce(CAST(123 AS VARCHAR(3)), 1))#x, typeof(coalesce(cast(cast(123 as string) as bigint), cast(1 as bigint))) AS typeof(coalesce(CAST(123 AS STRING), 1))#x] ++- OneRowRelation + + +-- !query +SELECT coalesce(cast('123' AS CHAR(3)), 1), + coalesce(cast('123' AS VARCHAR(3)), 1), + coalesce(cast('123' AS STRING), 1) +-- !query analysis +Project [coalesce(cast(cast(123 as char(3)) as bigint), cast(1 as bigint)) AS coalesce(CAST(123 AS CHAR(3)), 1)#xL, coalesce(cast(cast(123 as varchar(3)) as bigint), cast(1 as bigint)) AS coalesce(CAST(123 AS VARCHAR(3)), 1)#xL, coalesce(cast(cast(123 as string) as bigint), cast(1 as bigint)) AS coalesce(CAST(123 AS STRING), 1)#xL] ++- OneRowRelation + + +-- !query +SELECT cast('123' AS CHAR(3)) = 123, + cast('123' AS VARCHAR(3)) = 123, + cast('123' AS STRING) = 123 +-- !query analysis +Project [(cast(cast(123 as char(3)) as bigint) = cast(123 as bigint)) AS (CAST(123 AS CHAR(3)) = 123)#x, (cast(cast(123 as varchar(3)) as bigint) = cast(123 as bigint)) AS (CAST(123 AS VARCHAR(3)) = 123)#x, (cast(cast(123 as string) as bigint) = cast(123 as bigint)) AS (CAST(123 AS STRING) = 123)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('1.5' AS CHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS VARCHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS STRING), 1.5)) +-- !query analysis +Project [typeof(coalesce(cast(cast(1.5 as char(3)) as double), cast(1.5 as double))) AS typeof(coalesce(CAST(1.5 AS CHAR(3)), 1.5))#x, typeof(coalesce(cast(cast(1.5 as varchar(3)) as double), cast(1.5 as double))) AS typeof(coalesce(CAST(1.5 AS VARCHAR(3)), 1.5))#x, typeof(coalesce(cast(cast(1.5 as string) as double), cast(1.5 as double))) AS typeof(coalesce(CAST(1.5 AS STRING), 1.5))#x] ++- OneRowRelation + + +-- !query +SELECT cast('2020-01-02' AS CHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS VARCHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS STRING) = date'2020-01-02' +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT cast('true' AS CHAR(4)) = true, + cast('true' AS VARCHAR(4)) = true, + cast('true' AS STRING) = true +-- !query analysis +Project [(cast(cast(true as char(4)) as boolean) = true) AS (CAST(true AS CHAR(4)) = true)#x, (cast(cast(true as varchar(4)) as boolean) = true) AS (CAST(true AS VARCHAR(4)) = true)#x, (cast(cast(true as string) as boolean) = true) AS (CAST(true AS STRING) = true)#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('true' AS CHAR(4)), true)) +-- !query analysis +Project [typeof(coalesce(cast(cast(true as char(4)) as boolean), true)) AS typeof(coalesce(CAST(true AS CHAR(4)), true))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('true' AS VARCHAR(4)), true)) +-- !query analysis +Project [typeof(coalesce(cast(cast(true as varchar(4)) as boolean), true)) AS typeof(coalesce(CAST(true AS VARCHAR(4)), true))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(coalesce(cast('true' AS STRING), true)) +-- !query analysis +Project [typeof(coalesce(cast(cast(true as string) as boolean), true)) AS typeof(coalesce(CAST(true AS STRING), true))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(array(cast('a' AS CHAR(2)), cast('bb' AS CHAR(3)))) +-- !query analysis +Project [typeof(array(cast(cast(a as char(2)) as char(3)), cast(bb as char(3)))) AS typeof(array(CAST(a AS CHAR(2)), CAST(bb AS CHAR(3))))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(struct(cast('a' AS CHAR(2)) AS f)) +-- !query analysis +Project [typeof(struct(f, cast(a as char(2)))) AS typeof(struct(CAST(a AS CHAR(2)) AS f))#x] ++- OneRowRelation + + +-- !query +SELECT typeof(map('k', cast('a' AS VARCHAR(2)))) +-- !query analysis +Project [typeof(map(k, cast(a as varchar(2)))) AS typeof(map(k, CAST(a AS VARCHAR(2))))#x] ++- OneRowRelation + + +-- !query +CREATE TABLE char_varchar_std (c CHAR(5), v VARCHAR(5)) USING parquet +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`char_varchar_std`, false + + +-- !query +INSERT INTO char_varchar_std VALUES ('ab', 'ab') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/char_varchar_std, false, Parquet, [path=file:[not included in comparison]/{warehouse_dir}/char_varchar_std], Append, `spark_catalog`.`default`.`char_varchar_std`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/char_varchar_std), [c, v] ++- Project [static_invoke(CharVarcharCodegenUtils.charTypeWriteSideCheck(cast(col1#x as char(5)), 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeWriteSideCheck(cast(col2#x as varchar(5)), 5)) AS v#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x, typeof(v#x) AS typeof(v)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std +-- !query analysis +Project [concat([, cast(c#x as string), ]) AS concat([, c, ])#x, concat([, cast(v#x as string), ]) AS concat([, v, ])#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +SELECT length(c), length(v) FROM char_varchar_std +-- !query analysis +Project [length(cast(c#x as string)) AS length(c)#x, length(cast(v#x as string)) AS length(v)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +CREATE TABLE char_varchar_std_ctas USING parquet AS SELECT c, v FROM char_varchar_std +-- !query analysis +CreateDataSourceTableAsSelectCommand `spark_catalog`.`default`.`char_varchar_std_ctas`, ErrorIfExists, [c, v] + +- Project [c#x, v#x] + +- SubqueryAlias spark_catalog.default.char_varchar_std + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std_ctas +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x, typeof(v#x) AS typeof(v)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std_ctas + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std_ctas[c#x,v#x] parquet + + +-- !query +CREATE VIEW char_varchar_std_view AS SELECT c FROM char_varchar_std +-- !query analysis +CreateViewCommand `spark_catalog`.`default`.`char_varchar_std_view`, SELECT c FROM char_varchar_std, false, false, PersistedView, COMPENSATION, true + +- Project [c#x] + +- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +SELECT typeof(c) FROM char_varchar_std_view +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std_view + +- View (`spark_catalog`.`default`.`char_varchar_std_view`, [c#x]) + +- Project [cast(c#x as char(5)) AS c#x] + +- Project [c#x] + +- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +CREATE VIEW char_varchar_std_view_v AS SELECT v FROM char_varchar_std +-- !query analysis +CreateViewCommand `spark_catalog`.`default`.`char_varchar_std_view_v`, SELECT v FROM char_varchar_std, false, false, PersistedView, COMPENSATION, true + +- Project [v#x] + +- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +SELECT typeof(v) FROM char_varchar_std_view_v +-- !query analysis +Project [typeof(v#x) AS typeof(v)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std_view_v + +- View (`spark_catalog`.`default`.`char_varchar_std_view_v`, [v#x]) + +- Project [cast(v#x as varchar(5)) AS v#x] + +- Project [v#x] + +- SubqueryAlias spark_catalog.default.char_varchar_std + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet + + +-- !query +WITH t AS (SELECT c, v FROM char_varchar_std) SELECT typeof(c), typeof(v) FROM t +-- !query analysis +WithCTE +:- CTERelationDef xxxx, false +: +- SubqueryAlias t +: +- Project [c#x, v#x] +: +- SubqueryAlias spark_catalog.default.char_varchar_std +: +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] +: +- Relation spark_catalog.default.char_varchar_std[c#x,v#x] parquet ++- Project [typeof(c#x) AS typeof(c)#x, typeof(v#x) AS typeof(v)#x] + +- SubqueryAlias t + +- CTERelationRef xxxx, true, [c#x, v#x], false, false + + +-- !query +CREATE TABLE char_varchar_std_orc (c CHAR(5), v VARCHAR(5)) USING orc +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`char_varchar_std_orc`, false + + +-- !query +INSERT INTO char_varchar_std_orc VALUES ('ab', 'cd') +-- !query analysis +InsertIntoHadoopFsRelationCommand file:[not included in comparison]/{warehouse_dir}/char_varchar_std_orc, false, ORC, [path=file:[not included in comparison]/{warehouse_dir}/char_varchar_std_orc], Append, `spark_catalog`.`default`.`char_varchar_std_orc`, org.apache.spark.sql.execution.datasources.InMemoryFileIndex(file:[not included in comparison]/{warehouse_dir}/char_varchar_std_orc), [c, v] ++- Project [static_invoke(CharVarcharCodegenUtils.charTypeWriteSideCheck(cast(col1#x as char(5)), 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeWriteSideCheck(cast(col2#x as varchar(5)), 5)) AS v#x] + +- LocalRelation [col1#x, col2#x] + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std_orc +-- !query analysis +Project [typeof(c#x) AS typeof(c)#x, typeof(v#x) AS typeof(v)#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std_orc + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std_orc[c#x,v#x] orc + + +-- !query +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std_orc +-- !query analysis +Project [concat([, cast(c#x as string), ]) AS concat([, c, ])#x, concat([, cast(v#x as string), ]) AS concat([, v, ])#x] ++- SubqueryAlias spark_catalog.default.char_varchar_std_orc + +- Project [static_invoke(CharVarcharCodegenUtils.charTypeReadSideCheck(c#x, 5)) AS c#x, static_invoke(CharVarcharCodegenUtils.varcharTypeReadSideCheck(v#x, 5)) AS v#x] + +- Relation spark_catalog.default.char_varchar_std_orc[c#x,v#x] orc + + +-- !query +DROP VIEW char_varchar_std_view_v +-- !query analysis +DropTableCommand `spark_catalog`.`default`.`char_varchar_std_view_v`, false, true, false + + +-- !query +DROP VIEW char_varchar_std_view +-- !query analysis +DropTableCommand `spark_catalog`.`default`.`char_varchar_std_view`, false, true, false + + +-- !query +DROP TABLE char_varchar_std_ctas +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.char_varchar_std_ctas + + +-- !query +DROP TABLE char_varchar_std_orc +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.char_varchar_std_orc + + +-- !query +DROP TABLE char_varchar_std +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.char_varchar_std diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/distinct-map-aggregates.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/distinct-map-aggregates.sql.out new file mode 100644 index 0000000000000..50f9153b7e3f8 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/distinct-map-aggregates.sql.out @@ -0,0 +1,363 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +CREATE OR REPLACE TEMPORARY VIEW distinct_map_data AS SELECT * FROM VALUES + (2, map('a', 1, 'b', 2), 1, true), + (2, map('b', 2, 'a', 1), 1, true), + (1, map('a', 1, 'b', 2), 1, true), + (1, map('a', 3), 2, false) +AS distinct_map_data(g, m, id, should_keep) +-- !query analysis +CreateViewCommand `distinct_map_data`, SELECT * FROM VALUES + (2, map('a', 1, 'b', 2), 1, true), + (2, map('b', 2, 'a', 1), 1, true), + (1, map('a', 1, 'b', 2), 1, true), + (1, map('a', 3), 2, false) +AS distinct_map_data(g, m, id, should_keep), false, true, LocalTempView, UNSUPPORTED, true + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT m) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct m#x) AS count(DISTINCT m)#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT SIZE(COLLECT_LIST(DISTINCT m)) FROM distinct_map_data +-- !query analysis +Aggregate [size(collect_list(distinct m#x, 0, 0, true), false) AS size(collect_list(DISTINCT m))#x] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a') +-- !query analysis +Project [map_entries(m)#x] ++- Sort [element_at(m#x, a, None, true) ASC NULLS FIRST], true + +- Project [map_entries(m#x) AS map_entries(m)#x, m#x] + +- SubqueryAlias collected_maps + +- Project [m#x] + +- Generate explode(_gen_input_0#x), false, [m#x] + +- Aggregate [collect_list(distinct m#x, 0, 0, true) AS _gen_input_0#x] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT map_entries(FIRST(DISTINCT m)), map_entries(LAST(DISTINCT m)), COUNT(DISTINCT m) +FROM VALUES (map('b', 2, 'a', 1)) AS single_map_data(m) +-- !query analysis +Aggregate [map_entries(first(distinct m#x, false)) AS map_entries(first(DISTINCT m))#x, map_entries(last(distinct m#x, false)) AS map_entries(last(DISTINCT m))#x, count(distinct m#x) AS count(DISTINCT m)#xL] ++- SubqueryAlias single_map_data + +- LocalRelation [m#x] + + +-- !query +SELECT COUNT(DISTINCT m, id) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct m#x, id#x) AS count(DISTINCT m, id)#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT m), COUNT(DISTINCT id) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct m#x) AS count(DISTINCT m)#xL, count(distinct id#x) AS count(DISTINCT id)#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT g, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY g +ORDER BY g +-- !query analysis +Sort [g#x ASC NULLS FIRST], true ++- Aggregate [g#x], [g#x, count(distinct m#x) AS count(DISTINCT m)#xL] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT m, COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a') +-- !query analysis +Sort [element_at(m#x, a, None, true) ASC NULLS FIRST], true ++- Aggregate [m#x], [m#x, count(distinct m#x) AS count(DISTINCT m)#xL, collect_list(distinct m#x, 0, 0, true) AS collect_list(DISTINCT m)#x] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT m) FILTER (WHERE should_keep) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct m#x) FILTER (WHERE should_keep#x) AS count(DISTINCT m) FILTER (WHERE should_keep)#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT MAX(map_values(m)[0]) +FROM distinct_map_data +WHERE id = 1 +-- !query analysis +Aggregate [max(map_values(m#x)[0]) AS max(map_values(m)[0])#x] ++- Filter (id#x = 1) + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT MAX(map_values(m)[0]), COUNT(DISTINCT m) +FROM distinct_map_data +WHERE id = 1 +-- !query analysis +Aggregate [max(map_values(m#x)[0]) AS max(map_values(m)[0])#x, count(distinct m#x) AS count(DISTINCT m)#xL] ++- Filter (id#x = 1) + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT g +FROM distinct_map_data +GROUP BY g +ORDER BY COUNT(DISTINCT m), g +-- !query analysis +Project [g#x] ++- Sort [count(DISTINCT m)#xL ASC NULLS FIRST, g#x ASC NULLS FIRST], true + +- Aggregate [g#x], [g#x, count(distinct m#x) AS count(DISTINCT m)#xL] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT g +FROM distinct_map_data +GROUP BY g +HAVING COUNT(DISTINCT m) = 1 +ORDER BY g +-- !query analysis +Sort [g#x ASC NULLS FIRST], true ++- Project [g#x] + +- Filter (count(DISTINCT m)#xL = cast(1 as bigint)) + +- Aggregate [g#x], [g#x, count(distinct m#x) AS count(DISTINCT m)#xL] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct named_struct(m, m#x)) AS count(DISTINCT named_struct(m, m))#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT array(m)) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct array(m#x)) AS count(DISTINCT array(m))#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT map('m', m)) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct map(m, m#x)) AS count(DISTINCT map(m, m))#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM VALUES + (CAST(map() AS MAP<STRING, INT>)), + (CAST(map() AS MAP<STRING, INT>)), + (CAST(NULL AS MAP<STRING, INT>)) +AS null_and_empty_map_data(m) +-- !query analysis +Aggregate [count(distinct m#x) AS count(DISTINCT m)#xL, collect_list(distinct m#x, 0, 0, true) AS collect_list(DISTINCT m)#x] ++- SubqueryAlias null_and_empty_map_data + +- LocalRelation [m#x] + + +-- !query +SELECT g, GROUPING(g), COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY GROUPING SETS ((g), ()) +ORDER BY GROUPING(g), g +-- !query analysis +Sort [grouping(g)#x ASC NULLS FIRST, g#x ASC NULLS FIRST], true ++- Aggregate [g#x, spark_grouping_id#xL], [g#x, cast((shiftright(spark_grouping_id#xL, 0) & 1) as tinyint) AS grouping(g)#x, count(distinct m#x) AS count(DISTINCT m)#xL] + +- Expand [[g#x, m#x, id#x, should_keep#x, g#x, 0], [g#x, m#x, id#x, should_keep#x, null, 1]], [g#x, m#x, id#x, should_keep#x, g#x, spark_grouping_id#xL] + +- Project [g#x, m#x, id#x, should_keep#x, g#x AS g#x] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m, 'n', n)) +FROM VALUES + (map('a', 1, 'b', 2), map('x', 1, 'y', 2)), + (map('b', 2, 'a', 1), map('y', 2, 'x', 1)) +AS grouped_distinct_map_data(m, n) +GROUP BY m +-- !query analysis +Aggregate [m#x], [count(distinct named_struct(m, m#x, n, n#x)) AS count(DISTINCT named_struct(m, m, n, n))#xL] ++- SubqueryAlias grouped_distinct_map_data + +- LocalRelation [m#x, n#x] + + +-- !query +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=false +-- !query analysis +SetCommand (spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled,Some(false)) + + +-- !query +SELECT COUNT(DISTINCT m) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct m#x) AS count(DISTINCT m)#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a'), map_entries(m)[0].key +-- !query analysis +Project [map_entries(m)#x] ++- Sort [element_at(m#x, a, None, true) ASC NULLS FIRST, map_entries(m#x)[0].key ASC NULLS FIRST], true + +- Project [map_entries(m#x) AS map_entries(m)#x, m#x] + +- SubqueryAlias collected_maps + +- Project [m#x] + +- Generate explode(_gen_input_0#x), false, [m#x] + +- Aggregate [collect_list(distinct m#x, 0, 0, true) AS _gen_input_0#x] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data +-- !query analysis +Aggregate [count(distinct named_struct(m, m#x)) AS count(DISTINCT named_struct(m, m))#xL] ++- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SELECT m, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a') +-- !query analysis +Sort [element_at(m#x, a, None, true) ASC NULLS FIRST], true ++- Aggregate [m#x], [m#x, count(distinct m#x) AS count(DISTINCT m)#xL] + +- SubqueryAlias distinct_map_data + +- View (`distinct_map_data`, [g#x, m#x, id#x, should_keep#x]) + +- Project [cast(g#x as int) AS g#x, cast(m#x as map<string,int>) AS m#x, cast(id#x as int) AS id#x, cast(should_keep#x as boolean) AS should_keep#x] + +- Project [g#x, m#x, id#x, should_keep#x] + +- SubqueryAlias distinct_map_data + +- LocalRelation [g#x, m#x, id#x, should_keep#x] + + +-- !query +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=true +-- !query analysis +SetCommand (spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled,Some(true)) diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/generators-resolution-edge-cases.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/generators-resolution-edge-cases.sql.out index 0b21d6e6b85d5..00930b19cd417 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/generators-resolution-edge-cases.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/generators-resolution-edge-cases.sql.out @@ -781,3 +781,63 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "fragment" : "val" } ] } + + +-- !query +SELECT explode(array(a)) AS col, count(*) OVER () AS cnt, a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL +-- !query analysis +Project [col#x, cnt#xL, a#x] ++- Generate explode(_gen_input_0#x), false, [col#x] + +- Project [_gen_input_0#x, cnt#xL, a#x] + +- Project [_gen_input_0#x, a#x, cnt#xL, cnt#xL] + +- Window [count(1) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS cnt#xL] + +- Filter isnotnull(a#x) + +- Aggregate [a#x], [array(a#x) AS _gen_input_0#x, a#x] + +- SubqueryAlias t + +- LocalRelation [a#x] + + +-- !query +SELECT explode(array(a)) AS col, + count(*) OVER () AS group_count, + count(*) AS row_count +FROM VALUES (1), (1), (2), (3), (3) AS t(a) +GROUP BY a +HAVING row_count > 1 +-- !query analysis +Project [col#x, group_count#xL, row_count#xL] ++- Generate explode(_gen_input_0#x), false, [col#x] + +- Project [_gen_input_0#x, group_count#xL, row_count#xL] + +- Project [_gen_input_0#x, row_count#xL, group_count#xL, group_count#xL] + +- Window [count(1) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS group_count#xL] + +- Filter (row_count#xL > cast(1 as bigint)) + +- Aggregate [a#x], [array(a#x) AS _gen_input_0#x, count(1) AS row_count#xL] + +- SubqueryAlias t + +- LocalRelation [a#x] + + +-- !query +SELECT explode(array(a)) AS col, + count(*) OVER () AS cnt, + count(*) OVER ( + ORDER BY a + ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + ) AS ordered_cnt, + a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL +-- !query analysis +Project [col#x, cnt#xL, ordered_cnt#xL, a#x] ++- Generate explode(_gen_input_0#x), false, [col#x] + +- Project [_gen_input_0#x, cnt#xL, ordered_cnt#xL, a#x] + +- Project [_gen_input_0#x, a#x, cnt#xL, ordered_cnt#xL, cnt#xL, ordered_cnt#xL] + +- Window [count(1) windowspecdefinition(a#x ASC NULLS FIRST, specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS ordered_cnt#xL], [a#x ASC NULLS FIRST] + +- Window [count(1) windowspecdefinition(specifiedwindowframe(RowFrame, unboundedpreceding$(), unboundedfollowing$())) AS cnt#xL] + +- Filter isnotnull(a#x) + +- Aggregate [a#x], [array(a#x) AS _gen_input_0#x, a#x] + +- SubqueryAlias t + +- LocalRelation [a#x] diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause-legacy.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause-legacy.sql.out index 6164ff65dabc8..24f34afc6f091 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause-legacy.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause-legacy.sql.out @@ -61,6 +61,14 @@ Project [c1#x] +- LocalRelation [c1#x] +-- !query +SELECT IDENTIFIER(concat('a', 'b')) FROM VALUES(1) AS T(ab) +-- !query analysis +Project [ab#x] ++- SubqueryAlias T + +- LocalRelation [ab#x] + + -- !query CREATE SCHEMA IF NOT EXISTS s -- !query analysis @@ -163,6 +171,14 @@ Project [c1#x] +- Relation spark_catalog.s.tab[c1#x] csv +-- !query +SELECT * FROM IDENTIFIER(concat('t', 'ab')) +-- !query analysis +Project [c1#x] ++- SubqueryAlias spark_catalog.s.tab + +- Relation spark_catalog.s.tab[c1#x] csv + + -- !query USE SCHEMA default -- !query analysis @@ -190,6 +206,13 @@ Project [coalesce(cast(null as int), 1) AS coalesce(NULL, 1)#x] +- OneRowRelation +-- !query +SELECT IDENTIFIER(concat('COAL', 'ESCE'))(NULL, 1) +-- !query analysis +Project [coalesce(cast(null as int), 1) AS coalesce(NULL, 1)#x] ++- OneRowRelation + + -- !query SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) -- !query analysis @@ -205,6 +228,483 @@ Project [id#xL] +- Range (0, 1, step=1) +-- !query +VALUES(IDENTIFIER(abs(1))) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.WRONG_TYPE", + "sqlState" : "42601", + "messageParameters" : { + "dataType" : "int", + "expr" : "abs(1)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "abs(1)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(nullif('a', 'a')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "nullif('a', 'a')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 41, + "fragment" : "nullif('a', 'a')" + } ] +} + + +-- !query +SELECT IDENTIFIER(max('c1')) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('c1')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 27, + "fragment" : "max('c1')" + } ] +} + + +-- !query +SELECT IDENTIFIER(array_join(transform(array('c', '1'), element -> element), '')) +FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('c', '1'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 80, + "fragment" : "array_join(transform(array('c', '1'), element -> element), '')" + } ] +} + + +-- !query +SELECT IDENTIFIER(rand()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "rand()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "rand()" + } ] +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER ()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WINDOW_FUNCTION_FRAME_NOT_ORDERED", + "sqlState" : "42601", + "messageParameters" : { + "wf_expr" : "row_number()", + "wf_name" : "row_number" + } +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER (ORDER BY 'x')) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 50, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT IDENTIFIER(explode(array('c1'))) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "explode(array('c1'))", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 38, + "fragment" : "explode(array('c1'))" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(max('identifier_function_table')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('identifier_function_table')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "max('identifier_function_table')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(row_number() OVER (ORDER BY 'x')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + array_join(transform(array('identifier', '_function_table'), element -> element), '') +) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('identifier', '_function_table'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 113, + "fragment" : "array_join(transform(array('identifier', '_function_table'), element -> element), '')" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand identifier_name, STRING, 'c1', false, true, false, false + + +-- !query +SELECT IDENTIFIER(identifier_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 35, + "fragment" : "identifier_name()" + } ] +} + + +-- !query +DROP TEMPORARY FUNCTION identifier_name +-- !query analysis +DropFunctionCommand identifier_name, false, true + + +-- !query +CREATE FUNCTION persistent_identifier_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_name, STRING, 'c1', false, false, false, false + + +-- !query +SELECT IDENTIFIER(persistent_identifier_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 46, + "fragment" : "persistent_identifier_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_name, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_function(value INT) +RETURNS INT +RETURN value + 1 +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_function, value INT, INT, value + 1, false, false, false, false + + +-- !query +SELECT IDENTIFIER(concat('persistent_identifier_', 'function'))(1) +-- !query analysis +Project [spark_catalog.default.persistent_identifier_function(value#x) AS spark_catalog.default.persistent_identifier_function(1)#x] ++- Project [cast(1 as int) AS value#x] + +- OneRowRelation + + +-- !query +DROP FUNCTION persistent_identifier_function +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_function, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_base_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_base_name, STRING, 'c1', false, false, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_nested_name() +RETURNS STRING +RETURN persistent_identifier_base_name() +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_nested_name, STRING, persistent_identifier_base_name(), false, false, false, false + + +-- !query +SELECT IDENTIFIER(persistent_identifier_nested_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_nested_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 53, + "fragment" : "persistent_identifier_nested_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_nested_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_nested_name, false, false + + +-- !query +DROP FUNCTION persistent_identifier_base_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_base_name, false, false + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' +-- !query analysis +CreateSQLFunctionCommand identifier_relation_name, STRING, 'identifier_function_table', false, true, false, false + + +-- !query +CREATE TABLE identifier_function_table(c1 INT) USING csv +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`identifier_function_table`, false + + +-- !query +SELECT * FROM IDENTIFIER(identifier_relation_name()) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 51, + "fragment" : "identifier_relation_name()" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_count() +RETURNS BIGINT +RETURN SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')) +-- !query analysis +CreateSQLFunctionCommand identifier_relation_count, BIGINT, SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')), false, true, false, false + + +-- !query +SELECT identifier_relation_count() +-- !query analysis +Project [identifier_relation_count() AS identifier_relation_count()#xL] +: +- Aggregate [count(1) AS count(1)#xL] +: +- SubqueryAlias spark_catalog.default.identifier_function_table +: +- Relation spark_catalog.default.identifier_function_table[c1#x] csv ++- Project + +- OneRowRelation + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_count +-- !query analysis +DropFunctionCommand identifier_relation_count, false, true + + +-- !query +CREATE FUNCTION persistent_identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_relation_name, STRING, 'identifier_function_table', false, false, false, false + + +-- !query +SELECT * FROM IDENTIFIER(persistent_identifier_relation_name()) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 62, + "fragment" : "persistent_identifier_relation_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_relation_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_relation_name, false, false + + +-- !query +DROP TABLE identifier_function_table +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.identifier_function_table + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_name +-- !query analysis +DropFunctionCommand identifier_relation_name, false, true + + -- !query CREATE TABLE IDENTIFIER('tab')(c1 INT) USING CSV -- !query analysis @@ -743,6 +1243,12 @@ CREATE TABLE t(col1 INT) CreateDataSourceTableCommand `spark_catalog`.`default`.`t`, false +-- !query +CREATE TABLE identifier_name_source(name STRING) USING csv +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`identifier_name_source`, false + + -- !query SELECT * FROM IDENTIFIER((SELECT 't')) -- !query analysis @@ -764,6 +1270,50 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +SELECT * FROM IDENTIFIER((SELECT max(name) FROM identifier_name_source)) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 71, + "fragment" : "(SELECT max(name) FROM identifier_name_source)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + (SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1) +) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 84, + "fragment" : "(SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1)" + } ] +} + + -- !query SELECT * FROM (SELECT IDENTIFIER((SELECT 'col1')) FROM IDENTIFIER((SELECT 't'))) -- !query analysis @@ -848,6 +1398,13 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +DROP TABLE identifier_name_source +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.identifier_name_source + + -- !query DROP TABLE t -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out index 35ceffe4ac9f1..31df295777932 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/identifier-clause.sql.out @@ -61,6 +61,14 @@ Project [c1#x] +- LocalRelation [c1#x] +-- !query +SELECT IDENTIFIER(concat('a', 'b')) FROM VALUES(1) AS T(ab) +-- !query analysis +Project [ab#x] ++- SubqueryAlias T + +- LocalRelation [ab#x] + + -- !query CREATE SCHEMA IF NOT EXISTS s -- !query analysis @@ -163,6 +171,14 @@ Project [c1#x] +- Relation spark_catalog.s.tab[c1#x] csv +-- !query +SELECT * FROM IDENTIFIER(concat('t', 'ab')) +-- !query analysis +Project [c1#x] ++- SubqueryAlias spark_catalog.s.tab + +- Relation spark_catalog.s.tab[c1#x] csv + + -- !query USE SCHEMA default -- !query analysis @@ -190,6 +206,13 @@ Project [coalesce(cast(null as int), 1) AS coalesce(NULL, 1)#x] +- OneRowRelation +-- !query +SELECT IDENTIFIER(concat('COAL', 'ESCE'))(NULL, 1) +-- !query analysis +Project [coalesce(cast(null as int), 1) AS coalesce(NULL, 1)#x] ++- OneRowRelation + + -- !query SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) -- !query analysis @@ -205,6 +228,483 @@ Project [id#xL] +- Range (0, 1, step=1) +-- !query +VALUES(IDENTIFIER(abs(1))) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.WRONG_TYPE", + "sqlState" : "42601", + "messageParameters" : { + "dataType" : "int", + "expr" : "abs(1)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "abs(1)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(nullif('a', 'a')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "nullif('a', 'a')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 41, + "fragment" : "nullif('a', 'a')" + } ] +} + + +-- !query +SELECT IDENTIFIER(max('c1')) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('c1')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 27, + "fragment" : "max('c1')" + } ] +} + + +-- !query +SELECT IDENTIFIER(array_join(transform(array('c', '1'), element -> element), '')) +FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('c', '1'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 80, + "fragment" : "array_join(transform(array('c', '1'), element -> element), '')" + } ] +} + + +-- !query +SELECT IDENTIFIER(rand()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "rand()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "rand()" + } ] +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER ()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WINDOW_FUNCTION_FRAME_NOT_ORDERED", + "sqlState" : "42601", + "messageParameters" : { + "wf_expr" : "row_number()", + "wf_name" : "row_number" + } +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER (ORDER BY 'x')) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 50, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT IDENTIFIER(explode(array('c1'))) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "explode(array('c1'))", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 38, + "fragment" : "explode(array('c1'))" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(max('identifier_function_table')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('identifier_function_table')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "max('identifier_function_table')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(row_number() OVER (ORDER BY 'x')) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + array_join(transform(array('identifier', '_function_table'), element -> element), '') +) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('identifier', '_function_table'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 113, + "fragment" : "array_join(transform(array('identifier', '_function_table'), element -> element), '')" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand identifier_name, STRING, 'c1', false, true, false, false + + +-- !query +SELECT IDENTIFIER(identifier_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 35, + "fragment" : "identifier_name()" + } ] +} + + +-- !query +DROP TEMPORARY FUNCTION identifier_name +-- !query analysis +DropFunctionCommand identifier_name, false, true + + +-- !query +CREATE FUNCTION persistent_identifier_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_name, STRING, 'c1', false, false, false, false + + +-- !query +SELECT IDENTIFIER(persistent_identifier_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 46, + "fragment" : "persistent_identifier_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_name, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_function(value INT) +RETURNS INT +RETURN value + 1 +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_function, value INT, INT, value + 1, false, false, false, false + + +-- !query +SELECT IDENTIFIER(concat('persistent_identifier_', 'function'))(1) +-- !query analysis +Project [spark_catalog.default.persistent_identifier_function(value#x) AS spark_catalog.default.persistent_identifier_function(1)#x] ++- Project [cast(1 as int) AS value#x] + +- OneRowRelation + + +-- !query +DROP FUNCTION persistent_identifier_function +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_function, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_base_name() +RETURNS STRING +RETURN 'c1' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_base_name, STRING, 'c1', false, false, false, false + + +-- !query +CREATE FUNCTION persistent_identifier_nested_name() +RETURNS STRING +RETURN persistent_identifier_base_name() +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_nested_name, STRING, persistent_identifier_base_name(), false, false, false, false + + +-- !query +SELECT IDENTIFIER(persistent_identifier_nested_name()) FROM VALUES(1) AS T(c1) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_nested_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 53, + "fragment" : "persistent_identifier_nested_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_nested_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_nested_name, false, false + + +-- !query +DROP FUNCTION persistent_identifier_base_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_base_name, false, false + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' +-- !query analysis +CreateSQLFunctionCommand identifier_relation_name, STRING, 'identifier_function_table', false, true, false, false + + +-- !query +CREATE TABLE identifier_function_table(c1 INT) USING csv +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`identifier_function_table`, false + + +-- !query +SELECT * FROM IDENTIFIER(identifier_relation_name()) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 51, + "fragment" : "identifier_relation_name()" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_count() +RETURNS BIGINT +RETURN SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')) +-- !query analysis +CreateSQLFunctionCommand identifier_relation_count, BIGINT, SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')), false, true, false, false + + +-- !query +SELECT identifier_relation_count() +-- !query analysis +Project [identifier_relation_count() AS identifier_relation_count()#xL] +: +- Aggregate [count(1) AS count(1)#xL] +: +- SubqueryAlias spark_catalog.default.identifier_function_table +: +- Relation spark_catalog.default.identifier_function_table[c1#x] csv ++- Project + +- OneRowRelation + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_count +-- !query analysis +DropFunctionCommand identifier_relation_count, false, true + + +-- !query +CREATE FUNCTION persistent_identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' +-- !query analysis +CreateSQLFunctionCommand spark_catalog.default.persistent_identifier_relation_name, STRING, 'identifier_function_table', false, false, false, false + + +-- !query +SELECT * FROM IDENTIFIER(persistent_identifier_relation_name()) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 62, + "fragment" : "persistent_identifier_relation_name()" + } ] +} + + +-- !query +DROP FUNCTION persistent_identifier_relation_name +-- !query analysis +DropFunctionCommand spark_catalog.default.persistent_identifier_relation_name, false, false + + +-- !query +DROP TABLE identifier_function_table +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.identifier_function_table + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_name +-- !query analysis +DropFunctionCommand identifier_relation_name, false, true + + -- !query CREATE TABLE IDENTIFIER('tab')(c1 INT) USING CSV -- !query analysis @@ -743,6 +1243,12 @@ CREATE TABLE t(col1 INT) CreateDataSourceTableCommand `spark_catalog`.`default`.`t`, false +-- !query +CREATE TABLE identifier_name_source(name STRING) USING csv +-- !query analysis +CreateDataSourceTableCommand `spark_catalog`.`default`.`identifier_name_source`, false + + -- !query SELECT * FROM IDENTIFIER((SELECT 't')) -- !query analysis @@ -764,6 +1270,50 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +SELECT * FROM IDENTIFIER((SELECT max(name) FROM identifier_name_source)) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 71, + "fragment" : "(SELECT max(name) FROM identifier_name_source)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + (SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1) +) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 84, + "fragment" : "(SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1)" + } ] +} + + -- !query SELECT * FROM (SELECT IDENTIFIER((SELECT 'col1')) FROM IDENTIFIER((SELECT 't'))) -- !query analysis @@ -848,6 +1398,13 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +DROP TABLE identifier_name_source +-- !query analysis +DropTable false, false ++- ResolvedIdentifier V2SessionCatalog(spark_catalog), default.identifier_name_source + + -- !query DROP TABLE t -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out index 4fb1f0f04231a..688014da6f2c1 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/json-functions.sql.out @@ -742,6 +742,172 @@ Project [json_object_keys([1, 2, 3]) AS json_object_keys([1, 2, 3])#x] +- OneRowRelation +-- !query +select json_typeof() +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "0", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "1", + "functionName" : "`json_typeof`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 20, + "fragment" : "json_typeof()" + } ] +} + + +-- !query +select json_typeof(null) +-- !query analysis +Project [json_typeof(null) AS json_typeof(NULL)#x] ++- OneRowRelation + + +-- !query +select json_typeof(200) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"200\"", + "inputType" : "\"INT\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"json_typeof(200)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "json_typeof(200)" + } ] +} + + +-- !query +select json_typeof('') +-- !query analysis +Project [json_typeof() AS json_typeof()#x] ++- OneRowRelation + + +-- !query +select json_typeof('{}') +-- !query analysis +Project [json_typeof({}) AS json_typeof({})#x] ++- OneRowRelation + + +-- !query +select json_typeof('{"key": 1, "arr": [1, 2]}') +-- !query analysis +Project [json_typeof({"key": 1, "arr": [1, 2]}) AS json_typeof({"key": 1, "arr": [1, 2]})#x] ++- OneRowRelation + + +-- !query +select json_typeof('[]') +-- !query analysis +Project [json_typeof([]) AS json_typeof([])#x] ++- OneRowRelation + + +-- !query +select json_typeof('[1, 2, 3]') +-- !query analysis +Project [json_typeof([1, 2, 3]) AS json_typeof([1, 2, 3])#x] ++- OneRowRelation + + +-- !query +select json_typeof('"hello"') +-- !query analysis +Project [json_typeof("hello") AS json_typeof("hello")#x] ++- OneRowRelation + + +-- !query +select json_typeof('123') +-- !query analysis +Project [json_typeof(123) AS json_typeof(123)#x] ++- OneRowRelation + + +-- !query +select json_typeof('1.5') +-- !query analysis +Project [json_typeof(1.5) AS json_typeof(1.5)#x] ++- OneRowRelation + + +-- !query +select json_typeof('-123') +-- !query analysis +Project [json_typeof(-123) AS json_typeof(-123)#x] ++- OneRowRelation + + +-- !query +select json_typeof('-1.5') +-- !query analysis +Project [json_typeof(-1.5) AS json_typeof(-1.5)#x] ++- OneRowRelation + + +-- !query +select json_typeof('true') +-- !query analysis +Project [json_typeof(true) AS json_typeof(true)#x] ++- OneRowRelation + + +-- !query +select json_typeof('false') +-- !query analysis +Project [json_typeof(false) AS json_typeof(false)#x] ++- OneRowRelation + + +-- !query +select json_typeof('null') +-- !query analysis +Project [json_typeof(null) AS json_typeof(null)#x] ++- OneRowRelation + + +-- !query +select json_typeof('bad') +-- !query analysis +Project [json_typeof(bad) AS json_typeof(bad)#x] ++- OneRowRelation + + +-- !query +select json_typeof('{"key": 45, "random_string"}') +-- !query analysis +Project [json_typeof({"key": 45, "random_string"}) AS json_typeof({"key": 45, "random_string"})#x] ++- OneRowRelation + + +-- !query +select json_typeof('123 true') +-- !query analysis +Project [json_typeof(123 true) AS json_typeof(123 true)#x] ++- OneRowRelation + + -- !query DROP VIEW IF EXISTS jsonTable -- !query analysis @@ -979,3 +1145,583 @@ GlobalLimit 1 +- LocalLimit 1 +- Project [from_json(StructField(time,TimeType(6),true), {"time": "14:30:45"}, Some(America/Los_Angeles), false) AS from_json({"time": "14:30:45"})#x] +- OneRowRelation + + +-- !query +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.name') +-- !query analysis +Project [json_value({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, $.name, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.name')#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.id' RETURNING INT) +-- !query analysis +Project [json_value({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, $.id, IntegerType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.id' RETURNING INT)#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7,"name":"Ada"}', '$.id' RETURNING INT) + 1 +-- !query analysis +Project [(json_value({"id":7,"name":"Ada"}, $.id, IntegerType, Null, Null, None, None, Some(America/Los_Angeles), true) + 1) AS (JSON_VALUE({"id":7,"name":"Ada"}, '$.id' RETURNING INT) + 1)#x] ++- OneRowRelation + + +-- !query +select json_value('{"score":null}', '$.score') +-- !query analysis +Project [json_value({"score":null}, $.score, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"score":null}, '$.score')#x] ++- OneRowRelation + + +-- !query +select json_value('{"addr":{"city":"NYC"}}', '$.addr') +-- !query analysis +Project [json_value({"addr":{"city":"NYC"}}, $.addr, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"addr":{"city":"NYC"}}, '$.addr')#x] ++- OneRowRelation + + +-- !query +select json_value('{"tags":["x","y"]}', '$.tags') +-- !query analysis +Project [json_value({"tags":["x","y"]}, $.tags, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"tags":["x","y"]}, '$.tags')#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7}', '$.missing') +-- !query analysis +Project [json_value({"id":7}, $.missing, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7}, '$.missing')#x] ++- OneRowRelation + + +-- !query +select json_value(cast(null as string), '$.a') +-- !query analysis +Project [json_value(cast(null as string), $.a, StringType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE(CAST(NULL AS STRING), '$.a')#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7}', '$.missing' DEFAULT '?' ON EMPTY) +-- !query analysis +Project [json_value({"id":7}, $.missing, StringType, Default, Null, Some(?), None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7}, '$.missing' DEFAULT ? ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7}', '$.missing' RETURNING INT DEFAULT 42 ON EMPTY) +-- !query analysis +Project [json_value({"id":7}, $.missing, IntegerType, Default, Null, Some(42), None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7}, '$.missing' RETURNING INT DEFAULT 42 ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_value('{"id":7}', '$.missing' ERROR ON EMPTY) +-- !query analysis +Project [json_value({"id":7}, $.missing, StringType, Error, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"id":7}, '$.missing' ERROR ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_value('{"addr":{"city":"NYC"}}', '$.addr' DEFAULT 'n/a' ON ERROR) +-- !query analysis +Project [json_value({"addr":{"city":"NYC"}}, $.addr, StringType, Null, Default, None, Some(n/a), Some(America/Los_Angeles), true) AS JSON_VALUE({"addr":{"city":"NYC"}}, '$.addr' DEFAULT n/a ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('not json', '$.a' DEFAULT 'bad' ON ERROR) +-- !query analysis +Project [json_value(not json, $.a, StringType, Null, Default, None, Some(bad), Some(America/Los_Angeles), true) AS JSON_VALUE(not json, '$.a' DEFAULT bad ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('not json', '$.a' ERROR ON ERROR) +-- !query analysis +Project [json_value(not json, $.a, StringType, Null, Error, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE(not json, '$.a' ERROR ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT) +-- !query analysis +Project [json_value({"name":"Ada"}, $.name, IntegerType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"name":"Ada"}, '$.name' RETURNING INT)#x] ++- OneRowRelation + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT ERROR ON ERROR) +-- !query analysis +Project [json_value({"name":"Ada"}, $.name, IntegerType, Null, Error, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"name":"Ada"}, '$.name' RETURNING INT ERROR ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT DEFAULT -1 ON ERROR) +-- !query analysis +Project [json_value({"name":"Ada"}, $.name, IntegerType, Null, Default, None, Some(-1), Some(America/Los_Angeles), true) AS JSON_VALUE({"name":"Ada"}, '$.name' RETURNING INT DEFAULT -1 ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('{"a":"x"}', '$.b' DEFAULT 'e' ON EMPTY DEFAULT 'r' ON ERROR) +-- !query analysis +Project [json_value({"a":"x"}, $.b, StringType, Default, Default, Some(e), Some(r), Some(America/Los_Angeles), true) AS JSON_VALUE({"a":"x"}, '$.b' DEFAULT e ON EMPTY DEFAULT r ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_value('{"v":"3.14"}', '$.v' RETURNING DOUBLE) +-- !query analysis +Project [json_value({"v":"3.14"}, $.v, DoubleType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"v":"3.14"}, '$.v' RETURNING DOUBLE)#x] ++- OneRowRelation + + +-- !query +select json_value('{"v":"true"}', '$.v' RETURNING BOOLEAN) +-- !query analysis +Project [json_value({"v":"true"}, $.v, BooleanType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"v":"true"}, '$.v' RETURNING BOOLEAN)#x] ++- OneRowRelation + + +-- !query +select json_value('{"v":"2020-01-02"}', '$.v' RETURNING DATE) +-- !query analysis +Project [json_value({"v":"2020-01-02"}, $.v, DateType, Null, Null, None, None, Some(America/Los_Angeles), true) AS JSON_VALUE({"v":"2020-01-02"}, '$.v' RETURNING DATE)#x] ++- OneRowRelation + + +-- !query +select json_value('{"a":[1,2]}', '$.a[*]') +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_value`", + "path" : "'$.a[*]'", + "sqlExpr" : "\"JSON_VALUE({\"a\":[1,2]}, '$.a[*]')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 42, + "fragment" : "json_value('{\"a\":[1,2]}', '$.a[*]')" + } ] +} + + +-- !query +select json_value('{"a":1}', '$.a' RETURNING STRUCT<x:INT>) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_SCALAR_RETURNING_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_value`", + "returningType" : "\"STRUCT<x: INT>\"", + "sqlExpr" : "\"JSON_VALUE({\"a\":1}, '$.a' RETURNING STRUCT<x: INT>)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 59, + "fragment" : "json_value('{\"a\":1}', '$.a' RETURNING STRUCT<x:INT>)" + } ] +} + + +-- !query +select json_value('{}', '$.x' RETURNING INT DEFAULT array(1) ON EMPTY) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"JSON_VALUE({}, '$.x' RETURNING INT DEFAULT array(1) ON EMPTY)\"", + "srcType" : "\"ARRAY<INT>\"", + "targetType" : "\"INT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 70, + "fragment" : "json_value('{}', '$.x' RETURNING INT DEFAULT array(1) ON EMPTY)" + } ] +} + + +-- !query +select json_exists('{"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}', '$.addr.city') +-- !query analysis +Project [json_exists({"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}, $.addr.city, False) AS JSON_EXISTS({"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}, '$.addr.city')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"score":null}', '$.score') +-- !query analysis +Project [json_exists({"score":null}, $.score, False) AS JSON_EXISTS({"score":null}, '$.score')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.addr.zip') +-- !query analysis +Project [json_exists({"addr":{"city":"NYC"}}, $.addr.zip, False) AS JSON_EXISTS({"addr":{"city":"NYC"}}, '$.addr.zip')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.addr') +-- !query analysis +Project [json_exists({"addr":{"city":"NYC"}}, $.addr, False) AS JSON_EXISTS({"addr":{"city":"NYC"}}, '$.addr')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"tags":["x","y"]}', '$.tags[0]') +-- !query analysis +Project [json_exists({"tags":["x","y"]}, $.tags[0], False) AS JSON_EXISTS({"tags":["x","y"]}, '$.tags[0]')#x] ++- OneRowRelation + + +-- !query +select json_exists(cast(null as string), '$.a') +-- !query analysis +Project [json_exists(cast(null as string), $.a, False) AS JSON_EXISTS(CAST(NULL AS STRING), '$.a')#x] ++- OneRowRelation + + +-- !query +select json_exists('not json', '$.a') +-- !query analysis +Project [json_exists(not json, $.a, False) AS JSON_EXISTS(not json, '$.a')#x] ++- OneRowRelation + + +-- !query +select json_exists('not json', '$.a' TRUE ON ERROR) +-- !query analysis +Project [json_exists(not json, $.a, True) AS JSON_EXISTS(not json, '$.a' TRUE ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_exists('not json', '$.a' FALSE ON ERROR) +-- !query analysis +Project [json_exists(not json, $.a, False) AS JSON_EXISTS(not json, '$.a')#x] ++- OneRowRelation + + +-- !query +select json_exists('not json', '$.a' UNKNOWN ON ERROR) +-- !query analysis +Project [json_exists(not json, $.a, Unknown) AS JSON_EXISTS(not json, '$.a' UNKNOWN ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_exists('not json', '$.a' ERROR ON ERROR) +-- !query analysis +Project [json_exists(not json, $.a, Error) AS JSON_EXISTS(not json, '$.a' ERROR ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":[1,2]}', '$.a[*]') +-- !query analysis +Project [json_exists({"a":[1,2]}, $.a[*], False) AS JSON_EXISTS({"a":[1,2]}, '$.a[*]')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":[]}', '$.a[*]') +-- !query analysis +Project [json_exists({"a":[]}, $.a[*], False) AS JSON_EXISTS({"a":[]}, '$.a[*]')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":5}', '$.a[*]') +-- !query analysis +Project [json_exists({"a":5}, $.a[*], False) AS JSON_EXISTS({"a":5}, '$.a[*]')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":[{"b":1},{"c":2}]}', '$.a[*].b') +-- !query analysis +Project [json_exists({"a":[{"b":1},{"c":2}]}, $.a[*].b, False) AS JSON_EXISTS({"a":[{"b":1},{"c":2}]}, '$.a[*].b')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":[1,2]}', '$.a[5]') +-- !query analysis +Project [json_exists({"a":[1,2]}, $.a[5], False) AS JSON_EXISTS({"a":[1,2]}, '$.a[5]')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":[{"b":1},{"b":2}]}', '$.a.b') +-- !query analysis +Project [json_exists({"a":[{"b":1},{"b":2}]}, $.a.b, False) AS JSON_EXISTS({"a":[{"b":1},{"b":2}]}, '$.a.b')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.*') +-- !query analysis +Project [json_exists({"addr":{"city":"NYC"}}, $.*, False) AS JSON_EXISTS({"addr":{"city":"NYC"}}, '$.*')#x] ++- OneRowRelation + + +-- !query +select json_exists('{"a":1}', '$[') +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_exists`", + "path" : "'$['", + "sqlExpr" : "\"JSON_EXISTS({\"a\":1}, '$[')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 35, + "fragment" : "json_exists('{\"a\":1}', '$[')" + } ] +} + + +-- !query +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.addr') +-- !query analysis +Project [json_query({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, $.addr, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.addr')#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.tags') +-- !query analysis +Project [json_query({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, $.tags, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.tags')#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.id') +-- !query analysis +Project [json_query({"id":7}, $.id, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"id":7}, '$.id')#x] ++- OneRowRelation + + +-- !query +select json_query('{"name":"Ada"}', '$.name') +-- !query analysis +Project [json_query({"name":"Ada"}, $.name, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"name":"Ada"}, '$.name')#x] ++- OneRowRelation + + +-- !query +select json_query('{"score":null}', '$.score') +-- !query analysis +Project [json_query({"score":null}, $.score, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"score":null}, '$.score')#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.missing') +-- !query analysis +Project [json_query({"id":7}, $.missing, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"id":7}, '$.missing')#x] ++- OneRowRelation + + +-- !query +select json_query(cast(null as string), '$.a') +-- !query analysis +Project [json_query(cast(null as string), $.a, StringType, Without, Keep, Null, Null) AS JSON_QUERY(CAST(NULL AS STRING), '$.a')#x] ++- OneRowRelation + + +-- !query +select json_query('{"tags":["x","y"]}', '$.tags[0]' WITH ARRAY WRAPPER) +-- !query analysis +Project [json_query({"tags":["x","y"]}, $.tags[0], StringType, Unconditional, Keep, Null, Null) AS JSON_QUERY({"tags":["x","y"]}, '$.tags[0]' WITH UNCONDITIONAL ARRAY WRAPPER)#x] ++- OneRowRelation + + +-- !query +select json_query('{"tags":["x","y"]}', '$.tags' WITH UNCONDITIONAL ARRAY WRAPPER) +-- !query analysis +Project [json_query({"tags":["x","y"]}, $.tags, StringType, Unconditional, Keep, Null, Null) AS JSON_QUERY({"tags":["x","y"]}, '$.tags' WITH UNCONDITIONAL ARRAY WRAPPER)#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.id' WITH ARRAY WRAPPER) +-- !query analysis +Project [json_query({"id":7}, $.id, StringType, Unconditional, Keep, Null, Null) AS JSON_QUERY({"id":7}, '$.id' WITH UNCONDITIONAL ARRAY WRAPPER)#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.id' WITH CONDITIONAL ARRAY WRAPPER) +-- !query analysis +Project [json_query({"id":7}, $.id, StringType, Conditional, Keep, Null, Null) AS JSON_QUERY({"id":7}, '$.id' WITH CONDITIONAL ARRAY WRAPPER)#x] ++- OneRowRelation + + +-- !query +select json_query('{"addr":{"city":"NYC"}}', '$.addr' WITH CONDITIONAL ARRAY WRAPPER) +-- !query analysis +Project [json_query({"addr":{"city":"NYC"}}, $.addr, StringType, Conditional, Keep, Null, Null) AS JSON_QUERY({"addr":{"city":"NYC"}}, '$.addr' WITH CONDITIONAL ARRAY WRAPPER)#x] ++- OneRowRelation + + +-- !query +select json_query('{"name":"Ada"}', '$.name' OMIT QUOTES) +-- !query analysis +Project [json_query({"name":"Ada"}, $.name, StringType, Without, Omit, Null, Null) AS JSON_QUERY({"name":"Ada"}, '$.name' OMIT QUOTES)#x] ++- OneRowRelation + + +-- !query +select json_query('{"name":"Ada"}', '$.name' KEEP QUOTES) +-- !query analysis +Project [json_query({"name":"Ada"}, $.name, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"name":"Ada"}, '$.name')#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.missing' EMPTY ARRAY ON EMPTY) +-- !query analysis +Project [json_query({"id":7}, $.missing, StringType, Without, Keep, EmptyArray, Null) AS JSON_QUERY({"id":7}, '$.missing' EMPTY ARRAY ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.missing' EMPTY OBJECT ON EMPTY) +-- !query analysis +Project [json_query({"id":7}, $.missing, StringType, Without, Keep, EmptyObject, Null) AS JSON_QUERY({"id":7}, '$.missing' EMPTY OBJECT ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_query('{"id":7}', '$.missing' ERROR ON EMPTY) +-- !query analysis +Project [json_query({"id":7}, $.missing, StringType, Without, Keep, Error, Null) AS JSON_QUERY({"id":7}, '$.missing' ERROR ON EMPTY)#x] ++- OneRowRelation + + +-- !query +select json_query('not json', '$.a') +-- !query analysis +Project [json_query(not json, $.a, StringType, Without, Keep, Null, Null) AS JSON_QUERY(not json, '$.a')#x] ++- OneRowRelation + + +-- !query +select json_query('not json', '$.a' EMPTY ARRAY ON ERROR) +-- !query analysis +Project [json_query(not json, $.a, StringType, Without, Keep, Null, EmptyArray) AS JSON_QUERY(not json, '$.a' EMPTY ARRAY ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_query('not json', '$.a' EMPTY OBJECT ON ERROR) +-- !query analysis +Project [json_query(not json, $.a, StringType, Without, Keep, Null, EmptyObject) AS JSON_QUERY(not json, '$.a' EMPTY OBJECT ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_query('not json', '$.a' ERROR ON ERROR) +-- !query analysis +Project [json_query(not json, $.a, StringType, Without, Keep, Null, Error) AS JSON_QUERY(not json, '$.a' ERROR ON ERROR)#x] ++- OneRowRelation + + +-- !query +select json_query('{"addr":{"city":"NYC"}}', '$.addr' RETURNING STRING) +-- !query analysis +Project [json_query({"addr":{"city":"NYC"}}, $.addr, StringType, Without, Keep, Null, Null) AS JSON_QUERY({"addr":{"city":"NYC"}}, '$.addr')#x] ++- OneRowRelation + + +-- !query +select json_query('{"a":[1,2]}', '$.a[*]') +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "path" : "'$.a[*]'", + "sqlExpr" : "\"JSON_QUERY({\"a\":[1,2]}, '$.a[*]')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 42, + "fragment" : "json_query('{\"a\":[1,2]}', '$.a[*]')" + } ] +} + + +-- !query +select json_query('{"a":1}', '$.a' RETURNING INT) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_QUERY_RETURNING_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "returningType" : "\"INT\"", + "sqlExpr" : "\"JSON_QUERY({\"a\":1}, '$.a' RETURNING INT)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 49, + "fragment" : "json_query('{\"a\":1}', '$.a' RETURNING INT)" + } ] +} + + +-- !query +select json_query('{"name":"Ada"}', '$.name' WITH ARRAY WRAPPER OMIT QUOTES) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_QUERY_WRAPPER_AND_QUOTES", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "sqlExpr" : "\"JSON_QUERY({\"name\":\"Ada\"}, '$.name' WITH UNCONDITIONAL ARRAY WRAPPER OMIT QUOTES)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 76, + "fragment" : "json_query('{\"name\":\"Ada\"}', '$.name' WITH ARRAY WRAPPER OMIT QUOTES)" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out index fa87a63e7f13d..38fac614df706 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/linear-regression.sql.out @@ -433,3 +433,50 @@ Aggregate [regr_r2(cast(y#x as double), cast(k#x as double)) AS regr_r2(y, k)#x] +- Project [k#x, y#x, x#x] +- SubqueryAlias testRegression +- LocalRelation [k#x, y#x, x#x] + + +-- !query +CREATE OR REPLACE TEMPORARY VIEW testCorrConstant AS SELECT * FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y) +-- !query analysis +CreateViewCommand `testCorrConstant`, SELECT * FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y), false, true, LocalTempView, UNSUPPORTED, true + +- Project [x#x, y#x] + +- SubqueryAlias t + +- LocalRelation [x#x, y#x] + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM testCorrConstant +-- !query analysis +Aggregate [corr(cast(x#x as double), cast(y#x as double)) AS corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE))#x] ++- SubqueryAlias testcorrconstant + +- View (`testCorrConstant`, [x#x, y#x]) + +- Project [cast(x#x as int) AS x#x, cast(y#x as int) AS y#x] + +- Project [x#x, y#x] + +- SubqueryAlias t + +- LocalRelation [x#x, y#x] + + +-- !query +DROP VIEW testCorrConstant +-- !query analysis +DropTempViewCommand testCorrConstant, false + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y) +-- !query analysis +Aggregate [corr(cast(x#x as double), cast(y#x as double)) AS corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE))#x] ++- SubqueryAlias t + +- LocalRelation [x#x, y#x] + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (2, 1), (3, 1) AS t(x, y) +-- !query analysis +Aggregate [corr(cast(x#x as double), cast(y#x as double)) AS corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE))#x] ++- SubqueryAlias t + +- LocalRelation [x#x, y#x] diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/math.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/math.sql.out index 1fa7b7513993d..053740acf1914 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/math.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/math.sql.out @@ -391,6 +391,307 @@ Project [bround(-9223372036854775808, -1) AS bround(-9223372036854775808, -1)#xL +- OneRowRelation +-- !query +SELECT truncate(25y, 1) +-- !query analysis +Project [truncate(25, 1) AS truncate(25, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, 0) +-- !query analysis +Project [truncate(25, 0) AS truncate(25, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -1) +-- !query analysis +Project [truncate(25, -1) AS truncate(25, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -2) +-- !query analysis +Project [truncate(25, -2) AS truncate(25, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -3) +-- !query analysis +Project [truncate(25, -3) AS truncate(25, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, 1) +-- !query analysis +Project [truncate(-25, 1) AS truncate(-25, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, 0) +-- !query analysis +Project [truncate(-25, 0) AS truncate(-25, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -1) +-- !query analysis +Project [truncate(-25, -1) AS truncate(-25, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -2) +-- !query analysis +Project [truncate(-25, -2) AS truncate(-25, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -3) +-- !query analysis +Project [truncate(-25, -3) AS truncate(-25, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(127y, -1) +-- !query analysis +Project [truncate(127, -1) AS truncate(127, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-128y, -1) +-- !query analysis +Project [truncate(-128, -1) AS truncate(-128, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(1234.5678) +-- !query analysis +Project [truncate(1234.5678, 0) AS truncate(1234.5678, 0)#x] ++- OneRowRelation + + -- !query SELECT conv('100', 2, 10) -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/misc-functions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/misc-functions.sql.out index a470d33d98931..08d901f949f8f 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/misc-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/misc-functions.sql.out @@ -296,3 +296,45 @@ SELECT hmac('key', 'message', 'SHA-3') -- !query analysis Project [hmac(cast(key as binary), cast(message as binary), SHA-3) AS hmac(key, message, SHA-3)#x] +- OneRowRelation + + +-- !query +SELECT xxh3_64('Spark') +-- !query analysis +Project [xxh3_64(cast(Spark as binary)) AS xxh3_64(Spark)#xL] ++- OneRowRelation + + +-- !query +SELECT xxh3_64(CAST('Spark' AS BINARY)) +-- !query analysis +Project [xxh3_64(cast(Spark as binary)) AS xxh3_64(CAST(Spark AS BINARY))#xL] ++- OneRowRelation + + +-- !query +SELECT xxh3_128('Spark') +-- !query analysis +Project [xxh3_128(cast(Spark as binary)) AS xxh3_128(Spark)#x] ++- OneRowRelation + + +-- !query +SELECT xxh3_128(CAST('Spark' AS BINARY)) +-- !query analysis +Project [xxh3_128(cast(Spark as binary)) AS xxh3_128(CAST(Spark AS BINARY))#x] ++- OneRowRelation + + +-- !query +SELECT xxh3_64(CAST(NULL AS STRING)) +-- !query analysis +Project [xxh3_64(cast(cast(null as string) as binary)) AS xxh3_64(CAST(NULL AS STRING))#xL] ++- OneRowRelation + + +-- !query +SELECT xxh3_128(CAST(NULL AS BINARY)) +-- !query analysis +Project [xxh3_128(cast(null as binary)) AS xxh3_128(CAST(NULL AS BINARY))#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/array.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/array.sql.out index 5f2ea9b475ea6..1e5f15dff8294 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/array.sql.out @@ -951,3 +951,73 @@ select array_distinct(array(0.0, -0.0, -0.0, DOUBLE("NaN"), DOUBLE("NaN"))) -- !query analysis Project [array_distinct(array(cast(0.0 as double), cast(0.0 as double), cast(0.0 as double), cast(NaN as double), cast(NaN as double))) AS array_distinct(array(0.0, 0.0, 0.0, NaN, NaN))#x] +- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 0) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 0) AS trim_array(array(1, 2, 3, 4, 5), 0)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 2) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 2) AS trim_array(array(1, 2, 3, 4, 5), 2)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 5) +-- !query analysis +Project [trim_array(array(1, 2, 3, 4, 5), 5) AS trim_array(array(1, 2, 3, 4, 5), 5)#x] ++- OneRowRelation + + +-- !query +select trim_array(array('a', 'b', 'c'), 1) +-- !query analysis +Project [trim_array(array(a, b, c), 1) AS trim_array(array(a, b, c), 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, null, 4), 1) +-- !query analysis +Project [trim_array(array(1, 2, cast(null as int), 4), 1) AS trim_array(array(1, 2, NULL, 4), 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(), 0) +-- !query analysis +Project [trim_array(array(), 0) AS trim_array(array(), 0)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), -1) +-- !query analysis +Project [trim_array(array(1, 2, 3), -1) AS trim_array(array(1, 2, 3), -1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), 4) +-- !query analysis +Project [trim_array(array(1, 2, 3), 4) AS trim_array(array(1, 2, 3), 4)#x] ++- OneRowRelation + + +-- !query +select trim_array(CAST(null AS ARRAY<INT>), 1) +-- !query analysis +Project [trim_array(cast(null as array<int>), 1) AS trim_array(NULL, 1)#x] ++- OneRowRelation + + +-- !query +select trim_array(array(1, 2, 3), CAST(null AS INT)) +-- !query analysis +Project [trim_array(array(1, 2, 3), cast(null as int)) AS trim_array(array(1, 2, 3), CAST(NULL AS INT))#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/math.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/math.sql.out index 5fe1b69352f57..161477a44b5fe 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/math.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/math.sql.out @@ -391,6 +391,307 @@ Project [bround(-9223372036854775808, -1) AS bround(-9223372036854775808, -1)#xL +- OneRowRelation +-- !query +SELECT truncate(25y, 1) +-- !query analysis +Project [truncate(25, 1) AS truncate(25, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, 0) +-- !query analysis +Project [truncate(25, 0) AS truncate(25, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -1) +-- !query analysis +Project [truncate(25, -1) AS truncate(25, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -2) +-- !query analysis +Project [truncate(25, -2) AS truncate(25, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(25y, -3) +-- !query analysis +Project [truncate(25, -3) AS truncate(25, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, 1) +-- !query analysis +Project [truncate(-25, 1) AS truncate(-25, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, 0) +-- !query analysis +Project [truncate(-25, 0) AS truncate(-25, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -1) +-- !query analysis +Project [truncate(-25, -1) AS truncate(-25, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -2) +-- !query analysis +Project [truncate(-25, -2) AS truncate(-25, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-25y, -3) +-- !query analysis +Project [truncate(-25, -3) AS truncate(-25, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(127y, -1) +-- !query analysis +Project [truncate(127, -1) AS truncate(127, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-128y, -1) +-- !query analysis +Project [truncate(-128, -1) AS truncate(-128, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525s, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525s, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(-525, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#x] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, 1) +-- !query analysis +Project [truncate(525, 1) AS truncate(525, 1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, 0) +-- !query analysis +Project [truncate(525, 0) AS truncate(525, 0)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -1) +-- !query analysis +Project [truncate(525, -1) AS truncate(525, -1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -2) +-- !query analysis +Project [truncate(525, -2) AS truncate(525, -2)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(525L, -3) +-- !query analysis +Project [truncate(525, -3) AS truncate(525, -3)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, 1) +-- !query analysis +Project [truncate(-525, 1) AS truncate(-525, 1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, 0) +-- !query analysis +Project [truncate(-525, 0) AS truncate(-525, 0)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -1) +-- !query analysis +Project [truncate(-525, -1) AS truncate(-525, -1)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -2) +-- !query analysis +Project [truncate(-525, -2) AS truncate(-525, -2)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(-525L, -3) +-- !query analysis +Project [truncate(-525, -3) AS truncate(-525, -3)#xL] ++- OneRowRelation + + +-- !query +SELECT truncate(1234.5678) +-- !query analysis +Project [truncate(1234.5678, 0) AS truncate(1234.5678, 0)#x] ++- OneRowRelation + + -- !query SELECT conv('100', 2, 10) -- !query analysis diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/string-functions.sql.out index 15e6f3ada2668..64b21c7114a08 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/string-functions.sql.out @@ -2167,3 +2167,38 @@ select instr(null, null, cast(null as int), cast(null as int)) -- !query analysis Project [instr(cast(null as string), cast(null as string), cast(null as int), cast(null as int)) AS instr(NULL, NULL, CAST(NULL AS INT), CAST(NULL AS INT))#x] +- OneRowRelation + + +-- !query +select normalize('hello') +-- !query analysis +Project [normalize(hello, NFC) AS normalize(hello, NFC)#x] ++- OneRowRelation + + +-- !query +select normalize('hello', 'NFD') +-- !query analysis +Project [normalize(hello, NFD) AS normalize(hello, NFD)#x] ++- OneRowRelation + + +-- !query +select normalize('fi', 'NFKC') +-- !query analysis +Project [normalize(fi, NFKC) AS normalize(fi, NFKC)#x] ++- OneRowRelation + + +-- !query +select normalize(null, 'NFC') +-- !query analysis +Project [normalize(cast(null as string), NFC) AS normalize(NULL, NFC)#x] ++- OneRowRelation + + +-- !query +select normalize('hello', 'not_a_form') +-- !query analysis +Project [normalize(hello, not_a_form) AS normalize(hello, not_a_form)#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/timestamp.sql.out index a6ebe17e8b317..a1191af8f6c64 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/nonansi/timestamp.sql.out @@ -875,7 +875,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:10\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP '2011-11-11 11:11:11' - 2011-11-11 11:11:10)\"" }, "queryContext" : [ { @@ -899,7 +899,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:11\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(2011-11-11 11:11:11 - TIMESTAMP '2011-11-11 11:11:10')\"" }, "queryContext" : [ { @@ -943,7 +943,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(str - TIMESTAMP '2011-11-11 11:11:11')\"" }, "queryContext" : [ { @@ -967,7 +967,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP '2011-11-11 11:11:11' - str)\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql-gating.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql-gating.sql.out new file mode 100644 index 0000000000000..64ac040ca2e81 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql-gating.sql.out @@ -0,0 +1,14 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT parse_sql('SELECT 1') +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "FEATURE_NOT_ENABLED", + "sqlState" : "56038", + "messageParameters" : { + "configKey" : "spark.sql.function.parseSql.enabled", + "configValue" : "true", + "featureName" : "parse_sql" + } +} diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql.sql.out new file mode 100644 index 0000000000000..bc9bbba7c3c71 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/parse-sql.sql.out @@ -0,0 +1,599 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT parse_sql(NULL) +-- !query analysis +Project [parse_sql(cast(null as string)) AS parse_sql(NULL)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT a, b FROM t') +-- !query analysis +Project [parse_sql(SELECT a, b FROM t) AS parse_sql(SELECT a, b FROM t)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2') +-- !query analysis +Project [parse_sql(SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2) AS parse_sql(SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2)#x] ++- OneRowRelation + + +-- !query +SELECT + get_json_object(result, '$.statement_identifier') AS statement_identifier, + get_json_object(result, '$.source_table_references[0][0]') AS first_table, + get_json_object(result, '$.select_list[1].name[0]') AS second_column +FROM (SELECT parse_sql('SELECT a, b FROM t') AS result) +-- !query analysis +Project [get_json_object(result#x, $.statement_identifier) AS statement_identifier#x, get_json_object(result#x, $.source_table_references[0][0]) AS first_table#x, get_json_object(result#x, $.select_list[1].name[0]) AS second_column#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [parse_sql(SELECT a, b FROM t) AS result#x] + +- OneRowRelation + + +-- !query +SELECT parse_sql('INSERT INTO t SELECT 1') +-- !query analysis +Project [parse_sql(INSERT INTO t SELECT 1) AS parse_sql(INSERT INTO t SELECT 1)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('DELETE FROM t WHERE a = 1') +-- !query analysis +Project [parse_sql(DELETE FROM t WHERE a = 1) AS parse_sql(DELETE FROM t WHERE a = 1)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('UPDATE t SET a = 1 WHERE b = 2') +-- !query analysis +Project [parse_sql(UPDATE t SET a = 1 WHERE b = 2) AS parse_sql(UPDATE t SET a = 1 WHERE b = 2)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE') +-- !query analysis +Project [parse_sql(MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE) AS parse_sql(MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('CREATE TABLE t (a INT)') +-- !query analysis +Project [parse_sql(CREATE TABLE t (a INT)) AS parse_sql(CREATE TABLE t (a INT))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('CREATE TABLE t AS SELECT 1 AS a') +-- !query analysis +Project [parse_sql(CREATE TABLE t AS SELECT 1 AS a) AS parse_sql(CREATE TABLE t AS SELECT 1 AS a)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('DROP TABLE t') +-- !query analysis +Project [parse_sql(DROP TABLE t) AS parse_sql(DROP TABLE t)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('CACHE TABLE t') +-- !query analysis +Project [parse_sql(CACHE TABLE t) AS parse_sql(CACHE TABLE t)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('TABLE t') +-- !query analysis +Project [parse_sql(TABLE t) AS parse_sql(TABLE t)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('VALUES (1), (2)') +-- !query analysis +Project [parse_sql(VALUES (1), (2)) AS parse_sql(VALUES (1), (2))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('CREATE FUNCTION f AS ''x'' USING JAR ''y.jar''') +-- !query analysis +Project [parse_sql(CREATE FUNCTION f AS 'x' USING JAR 'y.jar') AS parse_sql(CREATE FUNCTION f AS 'x' USING JAR 'y.jar')#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('DECLARE VARIABLE x INT') +-- !query analysis +Project [parse_sql(DECLARE VARIABLE x INT) AS parse_sql(DECLARE VARIABLE x INT)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT * FROM t WHERE a = :foo AND b = ?') +-- !query analysis +Project [parse_sql(SELECT * FROM t WHERE a = :foo AND b = ?) AS parse_sql(SELECT * FROM t WHERE a = :foo AND b = ?)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte') +-- !query analysis +Project [parse_sql(WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte) AS parse_sql(WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)') +-- !query analysis +Project [parse_sql(SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)) AS parse_sql(SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a') +-- !query analysis +Project [parse_sql(WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a) AS parse_sql(WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id)') +-- !query analysis +Project [parse_sql(SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id)) AS parse_sql(SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql( +'SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, ''x'') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1)') +-- !query analysis +Project [parse_sql(SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, 'x') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1)) AS parse_sql(SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, 'x') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql( +'MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name))') +-- !query analysis +Project [parse_sql(MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name))) AS parse_sql(MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name)))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql( +'CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper(''x'') + )') +-- !query analysis +Project [parse_sql(CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper('x') + )) AS parse_sql(CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper('x') + ))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELEC FROM t') +-- !query analysis +Project [parse_sql(SELEC FROM t) AS parse_sql(SELEC FROM t)#x] ++- OneRowRelation + + +-- !query +SELECT + get_json_object(result, '$.parse_success') AS parse_success, + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM (SELECT parse_sql('SELEC FROM t') AS result) +-- !query analysis +Project [get_json_object(result#x, $.parse_success) AS parse_success#x, get_json_object(result#x, $.error.errorClass) AS error_class#x, get_json_object(result#x, $.error.queryContext[0].fragment) AS fragment#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [parse_sql(SELEC FROM t) AS result#x] + +- OneRowRelation + + +-- !query +SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b') +-- !query analysis +Project [parse_sql(SELECT * + FROM t + ORDER BY a + CLUSTER BY b) AS parse_sql(SELECT * + FROM t + ORDER BY a + CLUSTER BY b)#x] ++- OneRowRelation + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].startIndex') AS start_index +FROM ( + SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b') AS result +) +-- !query analysis +Project [get_json_object(result#x, $.error.errorClass) AS error_class#x, get_json_object(result#x, $.error.line) AS line#x, get_json_object(result#x, $.error.position) AS position#x, get_json_object(result#x, $.error.queryContext[0].startIndex) AS start_index#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [parse_sql(SELECT * + FROM t + ORDER BY a + CLUSTER BY b) AS result#x] + +- OneRowRelation + + +-- !query +SELECT parse_sql('') +-- !query analysis +Project [parse_sql() AS parse_sql()#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('USE bad-name') +-- !query analysis +Project [parse_sql(USE bad-name) AS parse_sql(USE bad-name)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c') +-- !query analysis +Project [parse_sql(WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c) AS parse_sql(WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('MERGE INTO target USING source ON target.id = source.id') +-- !query analysis +Project [parse_sql(MERGE INTO target USING source ON target.id = source.id) AS parse_sql(MERGE INTO target USING source ON target.id = source.id)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('EXPLAIN SELECT 1') +-- !query analysis +Project [parse_sql(EXPLAIN SELECT 1) AS parse_sql(EXPLAIN SELECT 1)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SET spark.sql.adaptive.enabled=true') +-- !query analysis +Project [parse_sql(SET spark.sql.adaptive.enabled=true) AS parse_sql(SET spark.sql.adaptive.enabled=true)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('ADD JAR /tmp/x.jar') +-- !query analysis +Project [parse_sql(ADD JAR /tmp/x.jar) AS parse_sql(ADD JAR /tmp/x.jar)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('CREATE VIEW v AS SELECT a, b FROM t') +-- !query analysis +Project [parse_sql(CREATE VIEW v AS SELECT a, b FROM t) AS parse_sql(CREATE VIEW v AS SELECT a, b FROM t)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT 1 AS IDENTIFIER(''alias.field'')') +-- !query analysis +Project [parse_sql(SELECT 1 AS IDENTIFIER('alias.field')) AS parse_sql(SELECT 1 AS IDENTIFIER('alias.field'))#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('SELECT DATE ''not-a-date''') +-- !query analysis +Project [parse_sql(SELECT DATE 'not-a-date') AS parse_sql(SELECT DATE 'not-a-date')#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END') +-- !query analysis +Project [parse_sql(BEGIN + SELECT 1; + SELEC 2; + END) AS parse_sql(BEGIN + SELECT 1; + SELEC 2; + END)#x] ++- OneRowRelation + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END') AS result +) +-- !query analysis +Project [get_json_object(result#x, $.error.errorClass) AS error_class#x, get_json_object(result#x, $.error.line) AS line#x, get_json_object(result#x, $.error.position) AS position#x, get_json_object(result#x, $.error.queryContext[0].fragment) AS fragment#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [parse_sql(BEGIN + SELECT 1; + SELEC 2; + END) AS result#x] + +- OneRowRelation + + +-- !query +SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END') +-- !query analysis +Project [parse_sql(BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END) AS parse_sql(BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END)#x] ++- OneRowRelation + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END') AS result +) +-- !query analysis +Project [get_json_object(result#x, $.error.errorClass) AS error_class#x, get_json_object(result#x, $.error.line) AS line#x, get_json_object(result#x, $.error.position) AS position#x, get_json_object(result#x, $.error.queryContext[0].fragment) AS fragment#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [parse_sql(BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END) AS result#x] + +- OneRowRelation + + +-- !query +SELECT sql_text, parse_sql(sql_text) FROM VALUES + ('SELECT 1'), + ('INSERT INTO t SELECT 1'), + ('CACHE TABLE t') +AS t(sql_text) +-- !query analysis +Project [sql_text#x, parse_sql(sql_text#x) AS parse_sql(sql_text)#x] ++- SubqueryAlias t + +- LocalRelation [sql_text#x] + + +-- !query +SELECT parse_sql('BEGIN SELECT 1; END') +-- !query analysis +Project [parse_sql(BEGIN SELECT 1; END) AS parse_sql(BEGIN SELECT 1; END)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('BEGIN SELECT count(a) FROM script_t WHERE c = :p; END') +-- !query analysis +Project [parse_sql(BEGIN SELECT count(a) FROM script_t WHERE c = :p; END) AS parse_sql(BEGIN SELECT count(a) FROM script_t WHERE c = :p; END)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('BEGIN SELECT * FROM t WHERE a = ?; END') +-- !query analysis +Project [parse_sql(BEGIN SELECT * FROM t WHERE a = ?; END) AS parse_sql(BEGIN SELECT * FROM t WHERE a = ?; END)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END') +-- !query analysis +Project [parse_sql(BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END) AS parse_sql(BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql('BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END') +-- !query analysis +Project [parse_sql(BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END) AS parse_sql(BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END)#x] ++- OneRowRelation + + +-- !query +SELECT parse_sql( +'BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string(''%s'', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END') +-- !query analysis +Project [parse_sql(BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string('%s', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END) AS parse_sql(BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string('%s', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END)#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out index 96a4b2ec91c7a..0f19b148db17b 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/random.sql.out @@ -709,7 +709,24 @@ org.apache.spark.sql.AnalysisException -- !query SELECT randstr(-1, 0) AS result -- !query analysis -[Analyzer test output redacted due to nondeterminism] +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.VALUE_OUT_OF_RANGE", + "sqlState" : "42K09", + "messageParameters" : { + "currentValue" : "-1", + "exprName" : "`length`", + "sqlExpr" : "\"randstr(-1, 0)\"", + "valueRange" : "[0, 2147483647]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 21, + "fragment" : "randstr(-1, 0)" + } ] +} -- !query diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/string-functions.sql.out index 15e6f3ada2668..64b21c7114a08 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/string-functions.sql.out @@ -2167,3 +2167,38 @@ select instr(null, null, cast(null as int), cast(null as int)) -- !query analysis Project [instr(cast(null as string), cast(null as string), cast(null as int), cast(null as int)) AS instr(NULL, NULL, CAST(NULL AS INT), CAST(NULL AS INT))#x] +- OneRowRelation + + +-- !query +select normalize('hello') +-- !query analysis +Project [normalize(hello, NFC) AS normalize(hello, NFC)#x] ++- OneRowRelation + + +-- !query +select normalize('hello', 'NFD') +-- !query analysis +Project [normalize(hello, NFD) AS normalize(hello, NFD)#x] ++- OneRowRelation + + +-- !query +select normalize('fi', 'NFKC') +-- !query analysis +Project [normalize(fi, NFKC) AS normalize(fi, NFKC)#x] ++- OneRowRelation + + +-- !query +select normalize(null, 'NFC') +-- !query analysis +Project [normalize(cast(null as string), NFC) AS normalize(NULL, NFC)#x] ++- OneRowRelation + + +-- !query +select normalize('hello', 'not_a_form') +-- !query analysis +Project [normalize(hello, not_a_form) AS normalize(hello, not_a_form)#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/in-subquery/in-limit.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/in-subquery/in-limit.sql.out index 1820f23cfc793..8179080bba1ba 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/in-subquery/in-limit.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/subquery/in-subquery/in-limit.sql.out @@ -201,6 +201,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a LIMIT 10 OFFSET 2) LIMIT 2 @@ -214,14 +215,15 @@ GlobalLimit 2 : +- GlobalLimit 10 : +- LocalLimit 10 : +- Offset 2 - : +- Project [t2a#x] - : +- Filter (outer(t1d#xL) = t2d#xL) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2a#x ASC NULLS FIRST], true + : +- Project [t2a#x] + : +- Filter (outer(t1d#xL) = t2d#xL) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -263,6 +265,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a OFFSET 2) OFFSET 1 -- !query analysis @@ -270,14 +273,15 @@ Offset 1 +- Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] +- Filter t1a#x IN (list#x [t1d#xL]) : +- Offset 2 - : +- Project [t2a#x] - : +- Filter (outer(t1d#xL) = t2d#xL) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2a#x ASC NULLS FIRST], true + : +- Project [t2a#x] + : +- Filter (outer(t1d#xL) = t2d#xL) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -292,6 +296,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4 -- !query analysis @@ -301,14 +306,15 @@ GlobalLimit 4 +- Filter t1c#x IN (list#x []) : +- GlobalLimit 2 : +- LocalLimit 2 - : +- Project [t2c#x] - : +- Filter (cast(t2b#x as int) >= 8) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2c#x ASC NULLS LAST], true + : +- Project [t2c#x] + : +- Filter (cast(t2b#x as int) >= 8) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -359,6 +365,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b <= t1d + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4 -- !query analysis @@ -368,14 +375,15 @@ GlobalLimit 4 +- Filter t1c#x IN (list#x [t1d#xL]) : +- GlobalLimit 2 : +- LocalLimit 2 - : +- Project [t2c#x] - : +- Filter (cast(t2b#x as bigint) <= outer(t1d#xL)) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2c#x ASC NULLS LAST], true + : +- Project [t2c#x] + : +- Filter (cast(t2b#x as bigint) <= outer(t1d#xL)) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -390,20 +398,22 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2) -- !query analysis Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] +- Filter NOT t1b#x IN (list#x []) : +- GlobalLimit 2 : +- LocalLimit 2 - : +- Project [t2b#x] - : +- Filter (cast(t2b#x as int) > 6) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2b#x ASC NULLS FIRST], true + : +- Project [t2b#x] + : +- Filter (cast(t2b#x as int) > 6) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -557,6 +567,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2 OFFSET 2) LIMIT 4 @@ -570,14 +581,15 @@ GlobalLimit 4 : +- GlobalLimit 2 : +- LocalLimit 2 : +- Offset 2 - : +- Project [t2c#x] - : +- Filter (cast(t2b#x as int) >= 8) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2c#x ASC NULLS LAST], true + : +- Project [t2c#x] + : +- Filter (cast(t2b#x as int) >= 8) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -630,6 +642,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2 OFFSET 2) -- !query analysis @@ -638,14 +651,15 @@ Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] : +- GlobalLimit 2 : +- LocalLimit 2 : +- Offset 2 - : +- Project [t2b#x] - : +- Filter (cast(t2b#x as int) > 6) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2b#x ASC NULLS FIRST], true + : +- Project [t2b#x] + : +- Filter (cast(t2b#x as int) > 6) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -660,20 +674,22 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b <= t1d + ORDER BY t2b LIMIT 2) -- !query analysis Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] +- Filter NOT t1b#x IN (list#x [t1d#xL]) : +- GlobalLimit 2 : +- LocalLimit 2 - : +- Project [t2b#x] - : +- Filter (cast(t2b#x as bigint) <= outer(t1d#xL)) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2b#x ASC NULLS FIRST], true + : +- Project [t2b#x] + : +- Filter (cast(t2b#x as bigint) <= outer(t1d#xL)) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -858,6 +874,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c DESC NULLS LAST OFFSET 2) OFFSET 4 -- !query analysis @@ -865,14 +882,15 @@ Offset 4 +- Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] +- Filter t1c#x IN (list#x []) : +- Offset 2 - : +- Project [t2c#x] - : +- Filter (cast(t2b#x as int) >= 8) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2c#x DESC NULLS LAST], true + : +- Project [t2c#x] + : +- Filter (cast(t2b#x as int) >= 8) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] @@ -952,19 +970,21 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b OFFSET 2) -- !query analysis Project [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x] +- Filter NOT t1b#x IN (list#x []) : +- Offset 2 - : +- Project [t2b#x] - : +- Filter (cast(t2b#x as int) > 6) - : +- SubqueryAlias t2 - : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) - : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] - : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] - : +- SubqueryAlias t2 - : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- Sort [t2b#x ASC NULLS FIRST], true + : +- Project [t2b#x] + : +- Filter (cast(t2b#x as int) > 6) + : +- SubqueryAlias t2 + : +- View (`t2`, [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x]) + : +- Project [cast(t2a#x as string) AS t2a#x, cast(t2b#x as smallint) AS t2b#x, cast(t2c#x as int) AS t2c#x, cast(t2d#xL as bigint) AS t2d#xL, cast(t2e#x as float) AS t2e#x, cast(t2f#x as double) AS t2f#x, cast(t2g#x as decimal(4,0)) AS t2g#x, cast(t2h#x as timestamp) AS t2h#x, cast(t2i#x as date) AS t2i#x] + : +- Project [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] + : +- SubqueryAlias t2 + : +- LocalRelation [t2a#x, t2b#x, t2c#x, t2d#xL, t2e#x, t2f#x, t2g#x, t2h#x, t2i#x] +- SubqueryAlias t1 +- View (`t1`, [t1a#x, t1b#x, t1c#x, t1d#xL, t1e#x, t1f#x, t1g#x, t1h#x, t1i#x]) +- Project [cast(t1a#x as string) AS t1a#x, cast(t1b#x as smallint) AS t1b#x, cast(t1c#x as int) AS t1c#x, cast(t1d#xL as bigint) AS t1d#xL, cast(t1e#x as float) AS t1e#x, cast(t1f#x as double) AS t1f#x, cast(t1g#x as decimal(4,0)) AS t1g#x, cast(t1h#x as timestamp) AS t1h#x, cast(t1i#x as date) AS t1i#x] diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out index 3b0cece05ea6d..3173aa5e2f935 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/time.sql.out @@ -1940,7 +1940,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"TIME '12:30:41.123'\"", "inputType" : "\"TIME(6)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIME '12:30:41.123' - TIMESTAMP '2025-07-11 10:00:01')\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out index cdfb8a78bb22c..368ddf703b928 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out @@ -603,6 +603,41 @@ Project [cast(1960-01-01 19:04:05.123456789 + INTERVAL '0 00:00:00.000001' DAY T +- OneRowRelation +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' YEAR +-- !query analysis +Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR#x] ++- OneRowRelation + + +-- !query +SELECT INTERVAL '1' YEAR + TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' +-- !query analysis +Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +-- !query analysis +Project [2020-01-01 19:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1-2' YEAR TO MONTH +-- !query analysis +Project [2020-01-01 19:04:05.123456789 - INTERVAL '1-2' YEAR TO MONTH AS TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '1960-01-31 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +-- !query analysis +Project [1960-01-30 19:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_LTZ '1960-01-30 19:04:05.123456789' + INTERVAL '1' MONTH#x] ++- OneRowRelation + + -- !query SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0) -- !query analysis @@ -628,23 +663,80 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query -SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' +-- !query analysis +Project [(2020-01-01 19:04:05.123456789 - 2019-12-31 19:04:05.000000111) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456001 UTC' +-- !query analysis +Project [(2020-01-01 19:04:05.123456789 - 2020-01-01 19:04:05.123456001) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456001')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' +-- !query analysis +Project [(2019-12-31 19:04:05.000000111 - 2020-01-01 19:04:05.123456789) AS (TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789')#x] ++- OneRowRelation + + +-- !query +SELECT ('2020-01-02 03:04:05.1234567 UTC' :: timestamp_ltz(7)) - ('2020-01-01 03:04:05.000000009 UTC' :: timestamp_ltz(9)) +-- !query analysis +Project [(cast(cast(2020-01-02 03:04:05.1234567 UTC as timestamp_ltz(7)) as timestamp_ltz(9)) - cast(2020-01-01 03:04:05.000000009 UTC as timestamp_ltz(9))) AS (CAST(2020-01-02 03:04:05.1234567 UTC AS TIMESTAMP_LTZ(7)) - CAST(2020-01-01 03:04:05.000000009 UTC AS TIMESTAMP_LTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05 UTC' +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01' +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC' - TIMESTAMP_LTZ '1960-01-01 00:00:00.000000999 UTC' +-- !query analysis +Project [(2019-12-31 16:00:00.123456789 - 1959-12-31 16:00:00.000000999) AS (TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789' - TIMESTAMP_LTZ '1959-12-31 16:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - CAST(NULL AS timestamp_ltz(9)) +-- !query analysis +Project [(2020-01-01 19:04:05.123456789 - cast(null as timestamp_ltz(9))) AS (TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - CAST(NULL AS TIMESTAMP_LTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.123456789 UTC' :: timestamp_ltz(9)) -- !query analysis org.apache.spark.sql.catalyst.ExtendedAnalysisException { - "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES", + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", "sqlState" : "42K09", "messageParameters" : { - "left" : "\"TIMESTAMP_LTZ(9)\"", - "right" : "\"INTERVAL MONTH\"", - "sqlExpr" : "\"(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH)\"" + "inputSql" : "\"CAST(2022-03-27 03:00:00.123456789 UTC AS TIMESTAMP_LTZ(9))\"", + "inputType" : "\"TIMESTAMP_LTZ(9)\"", + "paramIndex" : "third", + "requiredType" : "\"(TIMESTAMP_NTZ OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\"", + "sqlExpr" : "\"convert_timezone(Europe/Brussels, Europe/Moscow, CAST(2022-03-27 03:00:00.123456789 UTC AS TIMESTAMP_LTZ(9)))\"" }, "queryContext" : [ { "objectType" : "", "objectName" : "", "startIndex" : 8, - "stopIndex" : 77, - "fragment" : "TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH" + "stopIndex" : 120, + "fragment" : "convert_timezone('Europe/Brussels', 'Europe/Moscow',\n '2022-03-27 03:00:00.123456789 UTC' :: timestamp_ltz(9))" } ] } @@ -673,6 +765,55 @@ Sort [c#x ASC NULLS FIRST], true +- LocalRelation [c#x] +-- !query +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3), + (CAST(NULL AS timestamp_ltz(9)), 4), + (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Aggregate [k#x], [k#x, count(1) AS count(1)#xL, sum(v#x) AS sum(v)#xL] + +- SubqueryAlias t + +- LocalRelation [k#x, v#x] + + +-- !query +SELECT mode(c) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c) +-- !query analysis +Aggregate [mode(c#x, 0, 0, None) AS mode(c)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c) +-- !query analysis +Aggregate [sort_array(collect_set(c#x, 0, 0, true), true) AS sort_array(collect_set(c), true)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c) +-- !query analysis +Aggregate [sort_array(collect_list(c#x, 0, 0, true), true) AS sort_array(collect_list(c), true)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + -- !query SELECT unix_timestamp(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789') -- !query analysis @@ -734,6 +875,21 @@ Aggregate [max_by(v#x, k#x) AS max_by(v, k)#x, min_by(v#x, k#x) AS min_by(v, k)# +- LocalRelation [v#x, k#x] +-- !query +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c) + ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Distinct + +- Project [c#x] + +- SubqueryAlias t + +- LocalRelation [c#x] + + -- !query SELECT unix_nanos(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789 UTC') -- !query analysis @@ -1036,6 +1192,63 @@ Sort [v#x ASC NULLS FIRST], true +- OneRowRelation +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query analysis +Project [(c#x = cast(2020-01-02 03:04:05.123456789 as timestamp_ltz(9))) AS (c = 2020-01-02 03:04:05.123456789)#x, (c#x = cast(2020-01-02 03:04:05.123456788 as timestamp_ltz(9))) AS (c = 2020-01-02 03:04:05.123456788)#x, (c#x < cast(2020-01-02 03:04:05.123456790 as timestamp_ltz(9))) AS (c < 2020-01-02 03:04:05.123456790)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005' +-- !query analysis +Project [c#x] ++- Filter between(c#x, 2020-01-02 03:04:05.000000001, 2020-01-02 03:04:05.000000005) + +- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SET spark.sql.ansi.enabled=false +-- !query analysis +SetCommand (spark.sql.ansi.enabled,Some(false)) + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true +-- !query analysis +SetCommand (spark.sql.legacy.typeCoercion.datetimeToString.enabled,Some(true)) + + +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query analysis +Project [(c#x = cast(2020-01-02 03:04:05.123456789 as timestamp_ltz(9))) AS (c = 2020-01-02 03:04:05.123456789)#x, (cast(c#x as string) < 2020-01-02 03:04:05.123456790) AS (c < 2020-01-02 03:04:05.123456790)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false +-- !query analysis +SetCommand (spark.sql.legacy.typeCoercion.datetimeToString.enabled,Some(false)) + + +-- !query +SET spark.sql.ansi.enabled=true +-- !query analysis +SetCommand (spark.sql.ansi.enabled,Some(true)) + + -- !query SELECT unix_seconds(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789 UTC') -- !query analysis @@ -1174,3 +1387,237 @@ SELECT date_trunc('NANOSECOND', TIMESTAMP_LTZ '2020-01-01 12:34:56.123456789 UTC -- !query analysis Project [date_trunc(NANOSECOND, 2020-01-01 04:34:56.123456789, Some(America/Los_Angeles)) AS date_trunc(NANOSECOND, TIMESTAMP_LTZ '2020-01-01 04:34:56.123456789')#x] +- OneRowRelation + + +-- !query +SELECT typeof(current_timestamp(9)), typeof(current_timestamp(8)), typeof(current_timestamp(7)) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT typeof(now(9)), typeof(now(6)) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT typeof(current_timestamp()), typeof(current_timestamp(6)) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT typeof(current_timestamp(7 + 2)) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT current_timestamp(9) = current_timestamp(9), now(9) = current_timestamp(9) +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT current_timestamp(3) +-- !query analysis +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "3", + "type" : "TIMESTAMP_LTZ" + } +} + + +-- !query +SELECT current_timestamp(10) +-- !query analysis +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "10", + "type" : "TIMESTAMP_LTZ" + } +} + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + INTERSECT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Intersect false + :- Project [c#x] + : +- SubqueryAlias __auto_generated_subquery_name + : +- Union false, false + : :- Project [2019-12-31 16:00:00.000000001 AS c#x] + : : +- OneRowRelation + : +- Project [2019-12-31 16:00:00.000000999 AS TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'#x] + : +- OneRowRelation + +- Project [2019-12-31 16:00:00.000000001 AS TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001'#x] + +- OneRowRelation + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + EXCEPT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Except false + :- Project [c#x] + : +- SubqueryAlias __auto_generated_subquery_name + : +- Union false, false + : :- Project [2019-12-31 16:00:00.000000001 AS c#x] + : : +- OneRowRelation + : +- Project [2019-12-31 16:00:00.000000999 AS TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'#x] + : +- OneRowRelation + +- Project [2019-12-31 16:00:00.000000001 AS TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001'#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009 UTC' :: timestamp_ltz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900 UTC' :: timestamp_ltz(9))) ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Project [typeof(c#x) AS typeof(c)#x, c#x] + +- SubqueryAlias __auto_generated_subquery_name + +- Intersect false + :- Project [cast(c#x as timestamp_ltz(9)) AS c#x] + : +- Project [cast(2020-01-01 00:00:00.0000009 UTC as timestamp_ltz(7)) AS c#x] + : +- OneRowRelation + +- Project [cast(2020-01-01 00:00:00.000000900 UTC as timestamp_ltz(9)) AS CAST(2020-01-01 00:00:00.000000900 UTC AS TIMESTAMP_LTZ(9))#x] + +- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000500 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC' +-- !query analysis +Project [between(2019-12-31 16:00:00.0000005, 2019-12-31 16:00:00.000000001, 2019-12-31 16:00:00.000000999) AS between(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000500', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000001000 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC' +-- !query analysis +Project [between(2019-12-31 16:00:00.000001, 2019-12-31 16:00:00.000000001, 2019-12-31 16:00:00.000000999) AS between(TIMESTAMP_LTZ '2019-12-31 16:00:00.000001000', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT '2020-01-01 00:00:00.000000500 UTC' :: timestamp_ltz(9) + BETWEEN '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7) + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000001 UTC' +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7), + TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC') AS v) +-- !query analysis +Project [typeof(v#x) AS typeof(v)#x, v#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [if (true) cast(cast(2020-01-01 00:00:00.0000001 UTC as timestamp_ltz(7)) as timestamp_ltz(9)) else 2019-12-31 16:00:00.123456789 AS v#x] + +- OneRowRelation + + +-- !query +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ltz(9)), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS v) +-- !query analysis +Project [typeof(v#x) AS typeof(v)#x, v#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [nvl(cast(null as timestamp_ltz(9)), 2019-12-31 16:00:00.000000999) AS v#x] + +- OneRowRelation + + +-- !query +SELECT ifnull(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', CAST(NULL AS timestamp_ltz(9))) +-- !query analysis +Project [ifnull(2019-12-31 16:00:00.000000001, cast(null as timestamp_ltz(9))) AS ifnull(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', CAST(NULL AS TIMESTAMP_LTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS t(k) + WHERE k IN (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter k#x IN (list#x []) + : +- Project [2019-12-31 16:00:00.000000999 AS TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'))) ORDER BY col +-- !query analysis +Sort [col#x ASC NULLS FIRST], true ++- Project [typeof(col#x) AS typeof(col)#x, col#x] + +- SubqueryAlias __auto_generated_subquery_name + +- Project [col#x] + +- Generate explode(array(2019-12-31 16:00:00.000000001, 2019-12-31 16:00:00.000000999)), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT element_at(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), 2) +-- !query analysis +Project [element_at(array(2019-12-31 16:00:00.000000001, 2019-12-31 16:00:00.000000999), 2, None, true) AS element_at(array(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'), 2)#x] ++- OneRowRelation + + +-- !query +SELECT (named_struct('f', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')).f +-- !query analysis +Project [named_struct(f, 2019-12-31 16:00:00.123456789).f AS named_struct(f, TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789').f#x] ++- OneRowRelation + + +-- !query +SELECT map('k', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')['k'] +-- !query analysis +Project [map(k, 2019-12-31 16:00:00.123456789)[k] AS map(k, TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789')[k]#x] ++- OneRowRelation + + +-- !query +SELECT map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b')[ + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'] +-- !query analysis +Project [map(2019-12-31 16:00:00.000000001, a, 2019-12-31 16:00:00.000000999, b)[2019-12-31 16:00:00.000000999] AS map(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', a, TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999', b)[TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999']#x] ++- OneRowRelation + + +-- !query +SELECT element_at(map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b'), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') +-- !query analysis +Project [element_at(map(2019-12-31 16:00:00.000000001, a, 2019-12-31 16:00:00.000000999, b), 2019-12-31 16:00:00.000000001, None, true) AS element_at(map(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', a, TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999', b), TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001')#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out index 7b03461ea2ac1..8fa51bf5ff5c1 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out @@ -530,6 +530,41 @@ Project [cast(1960-01-02 03:04:05.123456789 + INTERVAL '0 00:00:00.000001' DAY T +- OneRowRelation +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR +-- !query analysis +Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR#x] ++- OneRowRelation + + +-- !query +SELECT INTERVAL '1' YEAR + TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' +-- !query analysis +Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' YEAR AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH +-- !query analysis +Project [2020-01-02 03:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH +-- !query analysis +Project [2020-01-02 03:04:05.123456789 - INTERVAL '1-2' YEAR TO MONTH AS TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH +-- !query analysis +Project [1960-01-31 03:04:05.123456789 + INTERVAL '1' MONTH AS TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH#x] ++- OneRowRelation + + -- !query SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + make_interval(0, 1, 0, 2, 0, 0, 0) -- !query analysis @@ -555,25 +590,81 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query -SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' -- !query analysis -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES", - "sqlState" : "42K09", - "messageParameters" : { - "left" : "\"TIMESTAMP_NTZ(9)\"", - "right" : "\"INTERVAL MONTH\"", - "sqlExpr" : "\"(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH)\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 73, - "fragment" : "TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH" - } ] -} +Project [(2020-01-02 03:04:05.123456789 - 2020-01-01 03:04:05.000000111) AS (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456001' +-- !query analysis +Project [(2020-01-02 03:04:05.123456789 - 2020-01-02 03:04:05.123456001) AS (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456001')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' +-- !query analysis +Project [(2020-01-01 03:04:05.000000111 - 2020-01-02 03:04:05.123456789) AS (TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789')#x] ++- OneRowRelation + + +-- !query +SELECT ('2020-01-02 03:04:05.1234567' :: timestamp_ntz(7)) - ('2020-01-01 03:04:05.000000009' :: timestamp_ntz(9)) +-- !query analysis +Project [(cast(cast(2020-01-02 03:04:05.1234567 as timestamp_ntz(7)) as timestamp_ntz(9)) - cast(2020-01-01 03:04:05.000000009 as timestamp_ntz(9))) AS (CAST(2020-01-02 03:04:05.1234567 AS TIMESTAMP_NTZ(7)) - CAST(2020-01-01 03:04:05.000000009 AS TIMESTAMP_NTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05' +-- !query analysis +Project [(2020-01-02 03:04:05.123456789 - cast(2020-01-02 03:04:05 as timestamp_ntz(9))) AS (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01' +-- !query analysis +[Analyzer test output redacted due to nondeterminism] + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789' - TIMESTAMP_NTZ '1960-01-01 00:00:00.000000999' +-- !query analysis +Project [(2020-01-01 00:00:00.123456789 - 1960-01-01 00:00:00.000000999) AS (TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789' - TIMESTAMP_NTZ '1960-01-01 00:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - CAST(NULL AS timestamp_ntz(9)) +-- !query analysis +Project [(2020-01-02 03:04:05.123456789 - cast(null as timestamp_ntz(9))) AS (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - CAST(NULL AS TIMESTAMP_NTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + TIMESTAMP_NTZ '2022-03-27 03:00:00.123456789') +-- !query analysis +Project [convert_timezone(Europe/Brussels, Europe/Moscow, 2022-03-27 03:00:00.123456789) AS convert_timezone(Europe/Brussels, Europe/Moscow, TIMESTAMP_NTZ '2022-03-27 03:00:00.123456789')#x] ++- OneRowRelation + + +-- !query +SELECT typeof(convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.1234567' :: timestamp_ntz(7))) +-- !query analysis +Project [typeof(convert_timezone(Europe/Brussels, Europe/Moscow, cast(2022-03-27 03:00:00.1234567 as timestamp_ntz(7)))) AS typeof(convert_timezone(Europe/Brussels, Europe/Moscow, CAST(2022-03-27 03:00:00.1234567 AS TIMESTAMP_NTZ(7))))#x] ++- OneRowRelation + + +-- !query +SELECT convert_timezone('America/Los_Angeles', 'UTC', CAST(NULL AS timestamp_ntz(9))) +-- !query analysis +Project [convert_timezone(America/Los_Angeles, UTC, cast(null as timestamp_ntz(9))) AS convert_timezone(America/Los_Angeles, UTC, CAST(NULL AS TIMESTAMP_NTZ(9)))#x] ++- OneRowRelation -- !query @@ -600,6 +691,55 @@ Sort [c#x ASC NULLS FIRST], true +- LocalRelation [c#x] +-- !query +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3), + (CAST(NULL AS timestamp_ntz(9)), 4), + (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Aggregate [k#x], [k#x, count(1) AS count(1)#xL, sum(v#x) AS sum(v)#xL] + +- SubqueryAlias t + +- LocalRelation [k#x, v#x] + + +-- !query +SELECT mode(c) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c) +-- !query analysis +Aggregate [mode(c#x, 0, 0, None) AS mode(c)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c) +-- !query analysis +Aggregate [sort_array(collect_set(c#x, 0, 0, true), true) AS sort_array(collect_set(c), true)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c) +-- !query analysis +Aggregate [sort_array(collect_list(c#x, 0, 0, true), true) AS sort_array(collect_list(c), true)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + -- !query SELECT unix_timestamp(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query analysis @@ -654,6 +794,21 @@ Aggregate [max_by(v#x, k#x) AS max_by(v, k)#x, min_by(v#x, k#x) AS min_by(v, k)# +- LocalRelation [v#x, k#x] +-- !query +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c) + ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Distinct + +- Project [c#x] + +- SubqueryAlias t + +- LocalRelation [c#x] + + -- !query SELECT unix_nanos(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query analysis @@ -870,6 +1025,63 @@ Sort [v#x ASC NULLS FIRST], true +- OneRowRelation +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query analysis +Project [(c#x = cast(2020-01-02 03:04:05.123456789 as timestamp_ntz(9))) AS (c = 2020-01-02 03:04:05.123456789)#x, (c#x = cast(2020-01-02 03:04:05.123456788 as timestamp_ntz(9))) AS (c = 2020-01-02 03:04:05.123456788)#x, (c#x < cast(2020-01-02 03:04:05.123456790 as timestamp_ntz(9))) AS (c < 2020-01-02 03:04:05.123456790)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SELECT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005' +-- !query analysis +Project [c#x] ++- Filter between(c#x, 2020-01-02 03:04:05.000000001, 2020-01-02 03:04:05.000000005) + +- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SET spark.sql.ansi.enabled=false +-- !query analysis +SetCommand (spark.sql.ansi.enabled,Some(false)) + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true +-- !query analysis +SetCommand (spark.sql.legacy.typeCoercion.datetimeToString.enabled,Some(true)) + + +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query analysis +Project [(c#x = cast(2020-01-02 03:04:05.123456789 as timestamp_ntz(9))) AS (c = 2020-01-02 03:04:05.123456789)#x, (c#x < cast(2020-01-02 03:04:05.123456790 as timestamp_ntz(9))) AS (c < 2020-01-02 03:04:05.123456790)#x] ++- SubqueryAlias t + +- LocalRelation [c#x] + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false +-- !query analysis +SetCommand (spark.sql.legacy.typeCoercion.datetimeToString.enabled,Some(false)) + + +-- !query +SET spark.sql.ansi.enabled=true +-- !query analysis +SetCommand (spark.sql.ansi.enabled,Some(true)) + + -- !query SELECT unix_seconds(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query analysis @@ -1001,3 +1213,365 @@ SELECT date_trunc('NANOSECOND', TIMESTAMP_NTZ '2020-01-01 12:34:56.123456789') -- !query analysis Project [date_trunc(NANOSECOND, 2020-01-01 12:34:56.123456789, Some(America/Los_Angeles)) AS date_trunc(NANOSECOND, TIMESTAMP_NTZ '2020-01-01 12:34:56.123456789')#x] +- OneRowRelation + + +-- !query +SELECT typeof(localtimestamp(9)), typeof(localtimestamp(8)), typeof(localtimestamp(7)) +-- !query analysis +Project [typeof(localtimestamp(9, Some(America/Los_Angeles))) AS typeof(localtimestamp())#x, typeof(localtimestamp(8, Some(America/Los_Angeles))) AS typeof(localtimestamp())#x, typeof(localtimestamp(7, Some(America/Los_Angeles))) AS typeof(localtimestamp())#x] ++- OneRowRelation + + +-- !query +SELECT typeof(localtimestamp()), typeof(localtimestamp(6)) +-- !query analysis +Project [typeof(localtimestamp(Some(America/Los_Angeles))) AS typeof(localtimestamp())#x, typeof(localtimestamp(Some(America/Los_Angeles))) AS typeof(localtimestamp())#x] ++- OneRowRelation + + +-- !query +SELECT typeof(localtimestamp(8 + 1)) +-- !query analysis +Project [typeof(localtimestamp(9, Some(America/Los_Angeles))) AS typeof(localtimestamp())#x] ++- OneRowRelation + + +-- !query +SELECT localtimestamp(9) = localtimestamp(9) +-- !query analysis +Project [(localtimestamp(9, Some(America/Los_Angeles)) = localtimestamp(9, Some(America/Los_Angeles))) AS (localtimestamp() = localtimestamp())#x] ++- OneRowRelation + + +-- !query +SELECT localtimestamp(3) +-- !query analysis +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "3", + "type" : "TIMESTAMP_NTZ" + } +} + + +-- !query +SELECT localtimestamp(10) +-- !query analysis +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "10", + "type" : "TIMESTAMP_NTZ" + } +} + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + INTERSECT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Intersect false + :- Project [c#x] + : +- SubqueryAlias __auto_generated_subquery_name + : +- Union false, false + : :- Project [2020-01-01 00:00:00.000000001 AS c#x] + : : +- OneRowRelation + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- Project [2020-01-01 00:00:00.000000001 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'#x] + +- OneRowRelation + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + EXCEPT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Except false + :- Project [c#x] + : +- SubqueryAlias __auto_generated_subquery_name + : +- Union false, false + : :- Project [2020-01-01 00:00:00.000000001 AS c#x] + : : +- OneRowRelation + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- Project [2020-01-01 00:00:00.000000001 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'#x] + +- OneRowRelation + + +-- !query +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009' :: timestamp_ntz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900' :: timestamp_ntz(9))) ORDER BY c +-- !query analysis +Sort [c#x ASC NULLS FIRST], true ++- Project [typeof(c#x) AS typeof(c)#x, c#x] + +- SubqueryAlias __auto_generated_subquery_name + +- Intersect false + :- Project [cast(c#x as timestamp_ntz(9)) AS c#x] + : +- Project [cast(2020-01-01 00:00:00.0000009 as timestamp_ntz(7)) AS c#x] + : +- OneRowRelation + +- Project [cast(2020-01-01 00:00:00.000000900 as timestamp_ntz(9)) AS CAST(2020-01-01 00:00:00.000000900 AS TIMESTAMP_NTZ(9))#x] + +- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' +-- !query analysis +Project [between(2020-01-01 00:00:00.0000005, 2020-01-01 00:00:00.000000001, 2020-01-01 00:00:00.000000999) AS between(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000001000' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' +-- !query analysis +Project [between(2020-01-01 00:00:00.000001, 2020-01-01 00:00:00.000000001, 2020-01-01 00:00:00.000000999) AS between(TIMESTAMP_NTZ '2020-01-01 00:00:00.000001000', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')#x] ++- OneRowRelation + + +-- !query +SELECT '2020-01-01 00:00:00.000000500' :: timestamp_ntz(9) + BETWEEN '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7) + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000001' +-- !query analysis +Project [between(cast(2020-01-01 00:00:00.000000500 as timestamp_ntz(9)), cast(2020-01-01 00:00:00.0000001 as timestamp_ntz(7)), 2020-01-01 00:00:00.000001) AS between(CAST(2020-01-01 00:00:00.000000500 AS TIMESTAMP_NTZ(9)), CAST(2020-01-01 00:00:00.0000001 AS TIMESTAMP_NTZ(7)), TIMESTAMP_NTZ '2020-01-01 00:00:00.000001')#x] ++- OneRowRelation + + +-- !query +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7), + TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789') AS v) +-- !query analysis +Project [typeof(v#x) AS typeof(v)#x, v#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [if (true) cast(cast(2020-01-01 00:00:00.0000001 as timestamp_ntz(7)) as timestamp_ntz(9)) else 2020-01-01 00:00:00.123456789 AS v#x] + +- OneRowRelation + + +-- !query +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ntz(9)), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS v) +-- !query analysis +Project [typeof(v#x) AS typeof(v)#x, v#x] ++- SubqueryAlias __auto_generated_subquery_name + +- Project [nvl(cast(null as timestamp_ntz(9)), 2020-01-01 00:00:00.000000999) AS v#x] + +- OneRowRelation + + +-- !query +SELECT ifnull(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', CAST(NULL AS timestamp_ntz(9))) +-- !query analysis +Project [ifnull(2020-01-01 00:00:00.000000001, cast(null as timestamp_ntz(9))) AS ifnull(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', CAST(NULL AS TIMESTAMP_NTZ(9)))#x] ++- OneRowRelation + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter k#x IN (list#x []) + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') +-- !query analysis +Project [scalar-subquery#x [] AS scalarsubquery()#x] +: +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] +: +- OneRowRelation ++- OneRowRelation + + +-- !query +SELECT typeof((SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')) +-- !query analysis +Project [typeof(scalar-subquery#x []) AS typeof(scalarsubquery())#x] +: +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] +: +- OneRowRelation ++- OneRowRelation + + +-- !query +SELECT typeof((SELECT CAST(NULL AS timestamp_ntz(9)))) +-- !query analysis +Project [typeof(scalar-subquery#x []) AS typeof(scalarsubquery())#x] +: +- Project [cast(null as timestamp_ntz(9)) AS CAST(NULL AS TIMESTAMP_NTZ(9))#x] +: +- OneRowRelation ++- OneRowRelation + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k = (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter (k#x = scalar-subquery#x []) + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS s(v) + WHERE s.v = t.k) ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter exists#x [k#x] + : +- Project [1 AS 1#x] + : +- Filter (v#x = outer(k#x)) + : +- SubqueryAlias s + : +- LocalRelation [v#x] + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE NOT EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS s(v) + WHERE s.v = t.k) ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter NOT exists#x [k#x] + : +- Project [1 AS 1#x] + : +- Filter (v#x = outer(k#x)) + : +- SubqueryAlias s + : +- LocalRelation [v#x] + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter NOT k#x IN (list#x []) + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT k FROM VALUES + ('2020-01-01 00:00:00.0000009' :: timestamp_ntz(7)) AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter NOT cast(k#x as timestamp_ntz(9)) IN (list#x []) + : +- Project [TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' + UNION ALL SELECT CAST(NULL AS timestamp_ntz(9))) ORDER BY k +-- !query analysis +Sort [k#x ASC NULLS FIRST], true ++- Project [k#x] + +- Filter NOT k#x IN (list#x []) + : +- Union false, false + : :- Project [2020-01-01 00:00:00.000000999 AS TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'#x] + : : +- OneRowRelation + : +- Project [cast(null as timestamp_ntz(9)) AS CAST(NULL AS TIMESTAMP_NTZ(9))#x] + : +- OneRowRelation + +- SubqueryAlias t + +- LocalRelation [k#x] + + +-- !query +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'))) ORDER BY col +-- !query analysis +Sort [col#x ASC NULLS FIRST], true ++- Project [typeof(col#x) AS typeof(col)#x, col#x] + +- SubqueryAlias __auto_generated_subquery_name + +- Project [col#x] + +- Generate explode(array(2020-01-01 00:00:00.000000001, 2020-01-01 00:00:00.000000999)), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT element_at(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), 2) +-- !query analysis +Project [element_at(array(2020-01-01 00:00:00.000000001, 2020-01-01 00:00:00.000000999), 2, None, true) AS element_at(array(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), 2)#x] ++- OneRowRelation + + +-- !query +SELECT (named_struct('f', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')).f +-- !query analysis +Project [named_struct(f, 2020-01-01 00:00:00.123456789).f AS named_struct(f, TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789').f#x] ++- OneRowRelation + + +-- !query +SELECT map('k', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')['k'] +-- !query analysis +Project [map(k, 2020-01-01 00:00:00.123456789)[k] AS map(k, TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')[k]#x] ++- OneRowRelation + + +-- !query +SELECT map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b')[ + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'] +-- !query analysis +Project [map(2020-01-01 00:00:00.000000001, a, 2020-01-01 00:00:00.000000999, b)[2020-01-01 00:00:00.000000999] AS map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', a, TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', b)[TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999']#x] ++- OneRowRelation + + +-- !query +SELECT element_at(map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b'), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') +-- !query analysis +Project [element_at(map(2020-01-01 00:00:00.000000001, a, 2020-01-01 00:00:00.000000999, b), 2020-01-01 00:00:00.000000001, None, true) AS element_at(map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', a, TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', b), TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001')#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out index af47dffedcb12..f66c0b84e8829 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/timestampNTZ/timestamp.sql.out @@ -877,7 +877,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:10\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP_NTZ '2011-11-11 11:11:11' - 2011-11-11 11:11:10)\"" }, "queryContext" : [ { @@ -901,7 +901,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:11\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(2011-11-11 11:11:11 - TIMESTAMP_NTZ '2011-11-11 11:11:10')\"" }, "queryContext" : [ { @@ -947,7 +947,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(str - TIMESTAMP_NTZ '2011-11-11 11:11:11')\"" }, "queryContext" : [ { @@ -971,7 +971,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP_NTZ '2011-11-11 11:11:11' - str)\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/decimalPrecision.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/decimalPrecision.sql.out index 4458e15e53cf7..45e987583c9ea 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/decimalPrecision.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/typeCoercion/native/decimalPrecision.sql.out @@ -1706,7 +1706,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(3,0))\"", "inputType" : "\"DECIMAL(3,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(3,0)))\"" }, "queryContext" : [ { @@ -1730,7 +1730,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(5,0))\"", "inputType" : "\"DECIMAL(5,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(5,0)))\"" }, "queryContext" : [ { @@ -1754,7 +1754,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(10,0))\"", "inputType" : "\"DECIMAL(10,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(10,0)))\"" }, "queryContext" : [ { @@ -1778,7 +1778,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(20,0))\"", "inputType" : "\"DECIMAL(20,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(20,0)))\"" }, "queryContext" : [ { @@ -2426,7 +2426,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(3,0))\"", "inputType" : "\"DECIMAL(3,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(3,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2450,7 +2450,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(5,0))\"", "inputType" : "\"DECIMAL(5,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(5,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2474,7 +2474,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(10,0))\"", "inputType" : "\"DECIMAL(10,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(10,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2498,7 +2498,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(20,0))\"", "inputType" : "\"DECIMAL(20,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(20,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/unnest.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/unnest.sql.out new file mode 100644 index 0000000000000..a67c021b79424 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/unnest.sql.out @@ -0,0 +1,278 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +CREATE OR REPLACE TEMPORARY VIEW nested AS SELECT * FROM VALUES +(1, array(10, 20, 30), array('a', 'b')), +(2, array(40), array('c', 'd', 'e')), +(3, array(), array()), +(4, cast(null as array<int>), array('f')) +AS nested(id, xs, ys) +-- !query analysis +CreateViewCommand `nested`, SELECT * FROM VALUES +(1, array(10, 20, 30), array('a', 'b')), +(2, array(40), array('c', 'd', 'e')), +(3, array(), array()), +(4, cast(null as array<int>), array('f')) +AS nested(id, xs, ys), false, true, LocalTempView, UNSUPPORTED, true + +- Project [id#x, xs#x, ys#x] + +- SubqueryAlias nested + +- LocalRelation [id#x, xs#x, ys#x] + + +-- !query +SELECT * FROM UNNEST(array(10, 20, 30)) +-- !query analysis +Project [col#x] ++- Generate unnest(array(10, 20, 30)), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT v FROM UNNEST(array(10, 20, 30)) AS t(v) +-- !query analysis +Project [v#x] ++- SubqueryAlias t + +- Project [col#x AS v#x] + +- Generate unnest(array(10, 20, 30)), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array()) +-- !query analysis +Project [col#x] ++- Generate unnest(array()), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(cast(null as array<int>)) +-- !query analysis +Project [col#x] ++- Generate unnest(cast(null as array<int>)), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(10, 20, 30)) WITH ORDINALITY +-- !query analysis +Project [col#x, ordinality#xL] ++- Generate unnest(array(10, 20, 30), WITH ORDINALITY), false, [col#x, ordinality#xL] + +- OneRowRelation + + +-- !query +SELECT val, pos FROM UNNEST(array('x', 'y')) WITH ORDINALITY AS t(val, pos) +-- !query analysis +Project [val#x, pos#xL] ++- SubqueryAlias t + +- Project [col#x AS val#x, ordinality#xL AS pos#xL] + +- Generate unnest(array(x, y), WITH ORDINALITY), false, [col#x, ordinality#xL] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) AS t(a, b) +-- !query analysis +Project [a#x, b#x] ++- SubqueryAlias t + +- Project [col0#x AS a#x, col1#x AS b#x] + +- Generate unnest(array(1, 2), array(10, 20, 30)), false, [col0#x, col1#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) WITH ORDINALITY AS t(a, b, ord) +-- !query analysis +Project [a#x, b#x, ord#xL] ++- SubqueryAlias t + +- Project [col0#x AS a#x, col1#x AS b#x, ordinality#xL AS ord#xL] + +- Generate unnest(array(1, 2), array(10, 20, 30), WITH ORDINALITY), false, [col0#x, col1#x, ordinality#xL] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(struct(1, 'a'), struct(2, 'b'))) AS t(s) +-- !query analysis +Project [s#x] ++- SubqueryAlias t + +- Project [col#x AS s#x] + +- Generate unnest(array(struct(col1, 1, col2, a), struct(col1, 2, col2, b))), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT id, elem FROM nested, LATERAL UNNEST(xs) AS t(elem) ORDER BY id, elem +-- !query analysis +Sort [id#x ASC NULLS FIRST, elem#x ASC NULLS FIRST], true ++- Project [id#x, elem#x] + +- LateralJoin lateral-subquery#x [xs#x], Inner + : +- SubqueryAlias t + : +- Project [col#x AS elem#x] + : +- Generate unnest(outer(xs#x)), false, [col#x] + : +- OneRowRelation + +- SubqueryAlias nested + +- View (`nested`, [id#x, xs#x, ys#x]) + +- Project [cast(id#x as int) AS id#x, cast(xs#x as array<int>) AS xs#x, cast(ys#x as array<string>) AS ys#x] + +- Project [id#x, xs#x, ys#x] + +- SubqueryAlias nested + +- LocalRelation [id#x, xs#x, ys#x] + + +-- !query +SELECT id, x, y, ord +FROM nested, LATERAL UNNEST(xs, ys) WITH ORDINALITY AS t(x, y, ord) +ORDER BY id, ord +-- !query analysis +Sort [id#x ASC NULLS FIRST, ord#xL ASC NULLS FIRST], true ++- Project [id#x, x#x, y#x, ord#xL] + +- LateralJoin lateral-subquery#x [xs#x && ys#x], Inner + : +- SubqueryAlias t + : +- Project [col0#x AS x#x, col1#x AS y#x, ordinality#xL AS ord#xL] + : +- Generate unnest(outer(xs#x), outer(ys#x), WITH ORDINALITY), false, [col0#x, col1#x, ordinality#xL] + : +- OneRowRelation + +- SubqueryAlias nested + +- View (`nested`, [id#x, xs#x, ys#x]) + +- Project [cast(id#x as int) AS id#x, cast(xs#x as array<int>) AS xs#x, cast(ys#x as array<string>) AS ys#x] + +- Project [id#x, xs#x, ys#x] + +- SubqueryAlias nested + +- LocalRelation [id#x, xs#x, ys#x] + + +-- !query +SELECT id, elem +FROM nested LEFT JOIN LATERAL UNNEST(xs) AS t(elem) ON true +ORDER BY id, elem +-- !query analysis +Sort [id#x ASC NULLS FIRST, elem#x ASC NULLS FIRST], true ++- Project [id#x, elem#x] + +- LateralJoin lateral-subquery#x [xs#x], LeftOuter, true + : +- SubqueryAlias t + : +- Project [col#x AS elem#x] + : +- Generate unnest(outer(xs#x)), false, [col#x] + : +- OneRowRelation + +- SubqueryAlias nested + +- View (`nested`, [id#x, xs#x, ys#x]) + +- Project [cast(id#x as int) AS id#x, cast(xs#x as array<int>) AS xs#x, cast(ys#x as array<string>) AS ys#x] + +- Project [id#x, xs#x, ys#x] + +- SubqueryAlias nested + +- LocalRelation [id#x, xs#x, ys#x] + + +-- !query +SELECT * FROM UNNEST(array(array(1, 2), array(3))) AS t(inner) +-- !query analysis +Project [inner#x] ++- SubqueryAlias t + +- Project [col#x AS inner#x] + +- Generate unnest(array(array(1, 2), array(3))), false, [col#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(1, cast(null as int), 3)) WITH ORDINALITY +-- !query analysis +Project [col#x, ordinality#xL] ++- Generate unnest(array(1, cast(null as int), 3), WITH ORDINALITY), false, [col#x, ordinality#xL] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(array(1), array(10, 20, 30)) AS t(a, b) +-- !query analysis +Project [a#x, b#x] ++- SubqueryAlias t + +- Project [col0#x AS a#x, col1#x AS b#x] + +- Generate unnest(array(1), array(10, 20, 30)), false, [col0#x, col1#x] + +- OneRowRelation + + +-- !query +SELECT * FROM UNNEST(42) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"42\"", + "inputType" : "\"INT\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(42)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 24, + "fragment" : "UNNEST(42)" + } ] +} + + +-- !query +SELECT * FROM UNNEST(map('a', 1)) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"map(a, 1)\"", + "inputType" : "\"MAP<STRING, INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(map(a, 1))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 33, + "fragment" : "UNNEST(map('a', 1))" + } ] +} + + +-- !query +SELECT * FROM UNNEST(array(1, 2), 3) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"3\"", + "inputType" : "\"INT\"", + "paramIndex" : "second", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(array(1, 2), 3)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 36, + "fragment" : "UNNEST(array(1, 2), 3)" + } ] +} + + +-- !query +SELECT * FROM `unnest`(array(1, 2)) +-- !query analysis +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "UNRESOLVABLE_TABLE_VALUED_FUNCTION", + "sqlState" : "42883", + "messageParameters" : { + "name" : "`unnest`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 35, + "fragment" : "`unnest`(array(1, 2))" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/variant/variant-from-arrays-entries.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/variant/variant-from-arrays-entries.sql.out new file mode 100644 index 0000000000000..0f159016698c2 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/variant/variant-from-arrays-entries.sql.out @@ -0,0 +1,244 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +select cast(variant_from_arrays(array('z', 'a'), array(1, 2)) as string) +-- !query analysis +Project [cast(variant_from_arrays(array(z, a), array(1, 2)) as string) AS CAST(variant_from_arrays(array(z, a), array(1, 2)) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_arrays(cast(array() as array<string>), cast(array() as array<int>)) as string) +-- !query analysis +Project [cast(variant_from_arrays(cast(array() as array<string>), cast(array() as array<int>)) as string) AS CAST(variant_from_arrays(array(), array()) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_arrays(array('a', 'b'), array(1, cast(null as int))) as string) +-- !query analysis +Project [cast(variant_from_arrays(array(a, b), array(1, cast(null as int))) as string) AS CAST(variant_from_arrays(array(a, b), array(1, CAST(NULL AS INT))) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_arrays(array('a'), array(array(1, 2, 3))) as string) +-- !query analysis +Project [cast(variant_from_arrays(array(a), array(array(1, 2, 3))) as string) AS CAST(variant_from_arrays(array(a), array(array(1, 2, 3))) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_arrays(cast(null as array<string>), array(1)) as string) +-- !query analysis +Project [cast(variant_from_arrays(cast(null as array<string>), array(1)) as string) AS CAST(variant_from_arrays(NULL, array(1)) AS STRING)#x] ++- OneRowRelation + + +-- !query +select variant_from_arrays(array('a', cast(null as string)), array(1, 2)) +-- !query analysis +Project [variant_from_arrays(array(a, cast(null as string)), array(1, 2)) AS variant_from_arrays(array(a, CAST(NULL AS STRING)), array(1, 2))#x] ++- OneRowRelation + + +-- !query +select variant_from_arrays(array('a', 'a'), array(1, 2)) +-- !query analysis +Project [variant_from_arrays(array(a, a), array(1, 2)) AS variant_from_arrays(array(a, a), array(1, 2))#x] ++- OneRowRelation + + +-- !query +select variant_from_arrays(array('a', 'b'), array(1)) +-- !query analysis +Project [variant_from_arrays(array(a, b), array(1)) AS variant_from_arrays(array(a, b), array(1))#x] ++- OneRowRelation + + +-- !query +select variant_from_arrays(array(1, 2), array('a', 'b')) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(1, 2)\"", + "inputType" : "\"ARRAY<INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY<STRING>\"", + "sqlExpr" : "\"variant_from_arrays(array(1, 2), array(a, b))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 56, + "fragment" : "variant_from_arrays(array(1, 2), array('a', 'b'))" + } ] +} + + +-- !query +select variant_from_arrays(array('a'), array(map(1, 2))) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"variant_from_arrays(array(a), array(map(1, 2)))\"", + "srcType" : "\"MAP<INT, INT>\"", + "targetType" : "\"VARIANT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 56, + "fragment" : "variant_from_arrays(array('a'), array(map(1, 2)))" + } ] +} + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'b', 'v', 2))) as string) +-- !query analysis +Project [cast(variant_from_entries(array(named_struct(k, a, v, 1), named_struct(k, b, v, 2))) as string) AS CAST(variant_from_entries(array(named_struct(k, a, v, 1), named_struct(k, b, v, 2))) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_entries(cast(array() as array<struct<k:string,v:int>>)) as string) +-- !query analysis +Project [cast(variant_from_entries(cast(array() as array<struct<k:string,v:int>>)) as string) AS CAST(variant_from_entries(array()) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', cast(null as int)))) as string) +-- !query analysis +Project [cast(variant_from_entries(array(named_struct(k, a, v, cast(null as int)))) as string) AS CAST(variant_from_entries(array(named_struct(k, a, v, CAST(NULL AS INT)))) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), cast(null as struct<k:string,v:int>))) as string) +-- !query analysis +Project [cast(variant_from_entries(array(named_struct(k, a, v, 1), cast(null as struct<k:string,v:int>))) as string) AS CAST(variant_from_entries(array(named_struct(k, a, v, 1), NULL)) AS STRING)#x] ++- OneRowRelation + + +-- !query +select cast(variant_from_entries(cast(null as array<struct<k:string,v:int>>)) as string) +-- !query analysis +Project [cast(variant_from_entries(cast(null as array<struct<k:string,v:int>>)) as string) AS CAST(variant_from_entries(NULL) AS STRING)#x] ++- OneRowRelation + + +-- !query +select variant_from_entries(array(named_struct('k', cast(null as string), 'v', 1))) +-- !query analysis +Project [variant_from_entries(array(named_struct(k, cast(null as string), v, 1))) AS variant_from_entries(array(named_struct(k, CAST(NULL AS STRING), v, 1)))#x] ++- OneRowRelation + + +-- !query +select variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'a', 'v', 2))) +-- !query analysis +Project [variant_from_entries(array(named_struct(k, a, v, 1), named_struct(k, a, v, 2))) AS variant_from_entries(array(named_struct(k, a, v, 1), named_struct(k, a, v, 2)))#x] ++- OneRowRelation + + +-- !query +select variant_from_entries(array(1, 2)) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(1, 2)\"", + "inputType" : "\"ARRAY<INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(1, 2))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 40, + "fragment" : "variant_from_entries(array(1, 2))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 'a'))) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(named_struct(k, a))\"", + "inputType" : "\"ARRAY<STRUCT<k: STRING NOT NULL>>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, a)))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 58, + "fragment" : "variant_from_entries(array(named_struct('k', 'a')))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 1, 'v', 'a'))) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(named_struct(k, 1, v, a))\"", + "inputType" : "\"ARRAY<STRUCT<k: INT NOT NULL, v: STRING NOT NULL>>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, 1, v, a)))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 66, + "fragment" : "variant_from_entries(array(named_struct('k', 1, 'v', 'a')))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 'a', 'v', map(1, 2)))) +-- !query analysis +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, a, v, map(1, 2))))\"", + "srcType" : "\"MAP<INT, INT>\"", + "targetType" : "\"VARIANT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 74, + "fragment" : "variant_from_entries(array(named_struct('k', 'a', 'v', map(1, 2))))" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/vector-distance.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/vector-distance.sql.out index 42239b3297305..85b570926cbfd 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/vector-distance.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/vector-distance.sql.out @@ -654,3 +654,66 @@ SELECT vector_l2_distance( -- !query analysis Project [vector_l2_distance(array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), array(16.0, 15.0, 14.0, 13.0, 12.0, 11.0, 10.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0)) AS vector_l2_distance(array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), array(16.0, 15.0, 14.0, 13.0, 12.0, 11.0, 10.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0))#x] +- OneRowRelation + + +-- !query +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(3.0e19F, 4.0e19F)) +-- !query analysis +Project [vector_cosine_similarity(array(3.0E19, 4.0E19), array(3.0E19, 4.0E19)) AS vector_cosine_similarity(array(3.0E19, 4.0E19), array(3.0E19, 4.0E19))#x] ++- OneRowRelation + + +-- !query +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(-3.0e19F, -4.0e19F)) +-- !query analysis +Project [vector_cosine_similarity(array(3.0E19, 4.0E19), array(-3.0E19, -4.0E19)) AS vector_cosine_similarity(array(3.0E19, 4.0E19), array(-3.0E19, -4.0E19))#x] ++- OneRowRelation + + +-- !query +SELECT vector_inner_product(array(1.0e20F, 1.0e20F), array(1.0e20F, -1.0e20F)) +-- !query analysis +Project [vector_inner_product(array(1.0E20, 1.0E20), array(1.0E20, -1.0E20)) AS vector_inner_product(array(1.0E20, 1.0E20), array(1.0E20, -1.0E20))#x] ++- OneRowRelation + + +-- !query +SELECT vector_l2_distance(array(3.0e19F, 4.0e19F), array(0.0F, 0.0F)) +-- !query analysis +Project [vector_l2_distance(array(3.0E19, 4.0E19), array(0.0, 0.0)) AS vector_l2_distance(array(3.0E19, 4.0E19), array(0.0, 0.0))#x] ++- OneRowRelation + + +-- !query +SELECT vector_cosine_similarity(array(1.0e-23F, 0.0F), array(1.0e-23F, 0.0F)) +-- !query analysis +Project [vector_cosine_similarity(array(1.0E-23, 0.0), array(1.0E-23, 0.0)) AS vector_cosine_similarity(array(1.0E-23, 0.0), array(1.0E-23, 0.0))#x] ++- OneRowRelation + + +-- !query +SELECT vector_cosine_similarity(array(1.0e-23F, 1.0e-23F), array(1.0e-23F, -1.0e-23F)) +-- !query analysis +Project [vector_cosine_similarity(array(1.0E-23, 1.0E-23), array(1.0E-23, -1.0E-23)) AS vector_cosine_similarity(array(1.0E-23, 1.0E-23), array(1.0E-23, -1.0E-23))#x] ++- OneRowRelation + + +-- !query +SELECT vector_cosine_similarity(array(float('inf'), 1.0F), array(1.0F, 1.0F)) +-- !query analysis +Project [vector_cosine_similarity(array(cast(inf as float), 1.0), array(1.0, 1.0)) AS vector_cosine_similarity(array(inf, 1.0), array(1.0, 1.0))#x] ++- OneRowRelation + + +-- !query +SELECT vector_inner_product(array(float('inf'), 1.0F), array(1.0F, 1.0F)) +-- !query analysis +Project [vector_inner_product(array(cast(inf as float), 1.0), array(1.0, 1.0)) AS vector_inner_product(array(inf, 1.0), array(1.0, 1.0))#x] ++- OneRowRelation + + +-- !query +SELECT vector_l2_distance(array(float('inf'), 1.0F), array(0.0F, 0.0F)) +-- !query analysis +Project [vector_l2_distance(array(cast(inf as float), 1.0), array(0.0, 0.0)) AS vector_l2_distance(array(inf, 1.0), array(0.0, 0.0))#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/vector-norm.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/vector-norm.sql.out index 035ad40cc1222..287c039e6c942 100644 --- a/sql/core/src/test/resources/sql-tests/analyzer-results/vector-norm.sql.out +++ b/sql/core/src/test/resources/sql-tests/analyzer-results/vector-norm.sql.out @@ -572,3 +572,66 @@ SELECT vector_norm( -- !query analysis Project [vector_norm(array(1.0, 2.0, 3.0, 4.0, 5.0, cast(null as float), 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), 2.0) AS vector_norm(array(1.0, 2.0, 3.0, 4.0, 5.0, CAST(NULL AS FLOAT), 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), 2.0)#x] +- OneRowRelation + + +-- !query +SELECT vector_norm(array(3.0e19F, 4.0e19F), 2.0F) +-- !query analysis +Project [vector_norm(array(3.0E19, 4.0E19), 2.0) AS vector_norm(array(3.0E19, 4.0E19), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_normalize(array(3.0e19F, 4.0e19F), 2.0F) +-- !query analysis +Project [vector_normalize(array(3.0E19, 4.0E19), 2.0) AS vector_normalize(array(3.0E19, 4.0E19), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_norm(array(3.0e38F, 3.0e38F), 1.0F) +-- !query analysis +Project [vector_norm(array(3.0E38, 3.0E38), 1.0) AS vector_norm(array(3.0E38, 3.0E38), 1.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_normalize(array(3.0e38F, 3.0e38F), 1.0F) +-- !query analysis +Project [vector_normalize(array(3.0E38, 3.0E38), 1.0) AS vector_normalize(array(3.0E38, 3.0E38), 1.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_norm(array(1.0e-23F, 0.0F), 2.0F) +-- !query analysis +Project [vector_norm(array(1.0E-23, 0.0), 2.0) AS vector_norm(array(1.0E-23, 0.0), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_normalize(array(1.0e-23F, 0.0F), 2.0F) +-- !query analysis +Project [vector_normalize(array(1.0E-23, 0.0), 2.0) AS vector_normalize(array(1.0E-23, 0.0), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_normalize(array(1.0e-23F, 1.0e-23F), 2.0F) +-- !query analysis +Project [vector_normalize(array(1.0E-23, 1.0E-23), 2.0) AS vector_normalize(array(1.0E-23, 1.0E-23), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_norm(array(float('inf'), 1.0F), 2.0F) +-- !query analysis +Project [vector_norm(array(cast(inf as float), 1.0), 2.0) AS vector_norm(array(inf, 1.0), 2.0)#x] ++- OneRowRelation + + +-- !query +SELECT vector_normalize(array(float('inf'), 1.0F), 2.0F) +-- !query analysis +Project [vector_normalize(array(cast(inf as float), 1.0), 2.0) AS vector_normalize(array(inf, 1.0), 2.0)#x] ++- OneRowRelation diff --git a/sql/core/src/test/resources/sql-tests/inputs/array.sql b/sql/core/src/test/resources/sql-tests/inputs/array.sql index 27fb3b3d07774..cf642cbbfb10f 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/array.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/array.sql @@ -192,3 +192,17 @@ select array_prepend(array(CAST(NULL AS String)), CAST(NULL AS String)); -- SPARK-45599: Confirm 0.0, -0.0, and NaN are handled appropriately. select array_union(array(0.0, -0.0, DOUBLE("NaN")), array(0.0, -0.0, DOUBLE("NaN"))); select array_distinct(array(0.0, -0.0, -0.0, DOUBLE("NaN"), DOUBLE("NaN"))); + +-- function trim_array +select trim_array(array(1, 2, 3, 4, 5), 0); +select trim_array(array(1, 2, 3, 4, 5), 2); +select trim_array(array(1, 2, 3, 4, 5), 5); +select trim_array(array('a', 'b', 'c'), 1); +select trim_array(array(1, 2, null, 4), 1); +select trim_array(array(), 0); +-- trim_array errors: n negative or greater than the array cardinality +select trim_array(array(1, 2, 3), -1); +select trim_array(array(1, 2, 3), 4); +-- trim_array with NULL array or NULL n returns NULL +select trim_array(CAST(null AS ARRAY<INT>), 1); +select trim_array(array(1, 2, 3), CAST(null AS INT)); diff --git a/sql/core/src/test/resources/sql-tests/inputs/charvarchar-standard-semantics.sql b/sql/core/src/test/resources/sql-tests/inputs/charvarchar-standard-semantics.sql new file mode 100644 index 0000000000000..23baf68d2a98b --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/charvarchar-standard-semantics.sql @@ -0,0 +1,252 @@ +--SET spark.sql.charVarchar.standardSemantics.enabled=true + +-- R3: CAST introduces CHAR/VARCHAR +SELECT typeof(CAST('ab' AS CHAR(5))); +SELECT typeof(CAST('hello' AS VARCHAR(5))); +SELECT 'X' || CAST('5' AS CHAR(5)) || 'X'; + +-- CAST length: character-to-character truncates (ISO 6.13); numeric overflow errors +SELECT CAST('ab ' AS CHAR(2)); +SELECT CAST('abcdef' AS CHAR(2)); +SELECT CAST('abcdef' AS VARCHAR(2)); +SELECT try_cast('abcdef' AS CHAR(2)); +SELECT try_cast('abcdef' AS VARCHAR(2)); +SELECT CAST(12345 AS VARCHAR(4)); +SELECT CAST(12345 AS VARCHAR(5)); +SELECT try_cast(12345 AS VARCHAR(4)); + +-- Explicit CAST inside LCT must keep inner truncation / overflow (do not retarget). +SELECT coalesce(CAST('abcdef' AS VARCHAR(2)), CAST('x' AS VARCHAR(4))); +SELECT CASE WHEN true THEN CAST('abcdef' AS VARCHAR(2)) ELSE CAST('x' AS VARCHAR(4)) END; +SELECT CAST('abcdef' AS VARCHAR(2)) IN (CAST('ab' AS VARCHAR(4))); +SELECT coalesce( + CAST('abcdef' AS VARCHAR(2) COLLATE UTF8_LCASE), + CAST('x' AS VARCHAR(4) COLLATE UTF8_LCASE)); +SELECT coalesce(try_cast(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))); +SELECT coalesce(CAST(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))); + +-- R2: least common type (COALESCE / CASE) +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world' AS VARCHAR(10)))); +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world!' AS CHAR(6)))); +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), cast('world!' AS CHAR(6)))); +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), 'world')); +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), NULL)); +SELECT typeof( + CASE WHEN true THEN cast('a' AS CHAR(2)) ELSE cast('bb' AS CHAR(4)) END); + +-- R2: least common type for IN lists +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS CHAR(2)), cast('bbb' AS VARCHAR(3))); +SELECT typeof(c) FROM (SELECT cast('a' AS CHAR(2)) AS c) t WHERE c IN ('a ', 'b'); + +-- R1: transforming functions return STRING +SELECT typeof(upper(cast('ab' AS CHAR(2)))); +SELECT typeof(lower(cast('AB' AS VARCHAR(2)))); +SELECT typeof(cast('a' AS CHAR(1)) || cast('b' AS VARCHAR(1))); +SELECT typeof(substr(cast('hello' AS VARCHAR(5)), 1, 2)); +SELECT typeof(upper(coalesce(cast('a' AS CHAR(2)), cast('b' AS CHAR(4))))); +SELECT typeof(concat(cast('a' AS CHAR(2)), cast('b' AS CHAR(3)))); +SELECT typeof(concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE))); +SELECT concat('<', concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE)), '>'); +SELECT typeof(elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE))); +SELECT concat('<', elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE)), '>'); +SELECT typeof(trim(cast('ab ' AS CHAR(4)))); +SELECT typeof(lpad(cast('ab' AS CHAR(2)), 5, 'x')); + +-- R1: regexp / mask / split family +SELECT typeof(regexp_replace(cast('ab' AS CHAR(2)), 'a', 'x')); +SELECT typeof(regexp_extract(cast('ab' AS VARCHAR(2)), '(a)', 1)); +SELECT typeof(regexp_extract_all(cast('aab' AS VARCHAR(3)), '(a)', 1)); +SELECT typeof(split(cast('a,b' AS CHAR(3)), ',')); +SELECT typeof(mask(cast('ab' AS CHAR(2)))); + +-- R1: CHAR/VARCHAR promote to STRING where a plain string is expected, so expressions that +-- require all their string inputs to share one type accept them alongside a STRING argument. +-- Values are wrapped in sentinels because the golden format trims trailing blanks, which would +-- otherwise hide the CHAR padding these expressions operate on. +SELECT typeof(overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1)); +SELECT concat('<', overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1), '>'); +SELECT typeof(elt(1, cast('ab' AS CHAR(5)), 'x')); +SELECT typeof(right(cast('ab' AS CHAR(5)), 2)); +SELECT concat('<', right(cast('ab' AS CHAR(5)), 2), '>'); +SELECT typeof(left(cast('ab' AS CHAR(5)), 2)); + +-- R1: transforms whose result length differs from the input must not inherit the constraint. +SELECT typeof(reverse(cast('ab' AS CHAR(5)))); +SELECT typeof(hex(cast('ab' AS CHAR(5)))); +SELECT hex(cast('ab' AS CHAR(5))); +SELECT typeof(array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-')); +SELECT concat('<', array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-'), '>'); +-- reverse() on a non-string input is unaffected. +SELECT typeof(reverse(array(1, 2))); +-- Expressions that pull values out of a wrapper do not inherit the wrapper's constraint. These +-- take their inputs through ExpectsInputTypes, so no implicit cast strips the length for them. +SELECT typeof(str_to_map(cast('a:1,b:2' AS CHAR(7)))); +SELECT typeof(c0) FROM (SELECT json_tuple(cast('{"a":"1"}' AS CHAR(9)), 'a') AS c0); + +-- Collation survives CAST and LCT. Mixed lengths with the same collation widen to max(n, m); +-- they must not collapse to an indeterminate collation. +SELECT typeof(cast('a' AS CHAR(2) COLLATE UTF8_LCASE)); +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(2) COLLATE UTF8_LCASE))); +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))); +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))); +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS VARCHAR(4) COLLATE UTF8_LCASE))); +-- Mixed strength, same collation: Implicit string CAST CHAR(2) vs Default +-- non-string CAST CHAR(4). Length still widens to max(n, m); the COLLATE +-- operator itself is STRING, so it is not used here. +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))); +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))); + +-- Set operations and multi-row VALUES share the same LCT as COALESCE. +SELECT typeof(c) FROM ( + SELECT cast('a' AS VARCHAR(3)) AS c + UNION ALL + SELECT cast('abcd' AS VARCHAR(8)) AS c +) t LIMIT 1; +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t LIMIT 1; +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t; +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t; +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t; +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t; +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t; +-- Non-empty EXCEPT: after widening, 'ab ' is not 'xy '. +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t; +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t; +SELECT typeof(c) FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c); +SELECT concat('<', c, '>') FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c); + +-- Comparison and IN: both sides (including the IN left-hand side) are cast to the LCT of all +-- participants. Casting to CHAR pads, so CHAR vs CHAR of different lengths compares equal after +-- widening; casting to VARCHAR/STRING keeps the CHAR pad, so CHAR 'a' (stored as 'a ') is not equal +-- to VARCHAR/STRING 'a' unless the other side carries the same trailing blank. Trailing-blank +-- ignoring is a collation concern (RTRIM), not a type-level PAD SPACE policy. +SELECT cast('a' AS CHAR(2)) = cast('a' AS CHAR(4)); +SELECT cast('a' AS CHAR(2)) = cast('a' AS VARCHAR(2)); +SELECT cast('a' AS CHAR(2)) = cast('a ' AS VARCHAR(2)); +SELECT cast('a' AS CHAR(2)) = 'a'; +SELECT cast('a' AS CHAR(2)) = 'a '; +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = 'a'; +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = + cast('a' AS CHAR(4) COLLATE UTF8_BINARY_RTRIM); +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4))); +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS VARCHAR(2))); +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS VARCHAR(2))); +SELECT cast('a' AS CHAR(2)) IN ('a', 'b'); +SELECT cast('a' AS CHAR(2)) IN ('a ', 'b'); +-- Three-part IN: LHS CHAR(2) and list CHAR(4)/VARCHAR(3) all widen to VARCHAR(4). +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4)), cast('b' AS VARCHAR(3))); +-- Same CHAR-pad rule under a non-RTRIM collation: UTF8_LCASE must not make +-- CHAR 'a' equal VARCHAR 'a'. The analyzer must nest CHAR then VARCHAR, not +-- retarget CAST('a' AS CHAR(2) COLLATE UTF8_LCASE) to VARCHAR(2). +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) = cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE); +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) IN (cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE)); + +-- CHAR/VARCHAR vs non-string follow STRING: compare promotes the string side to the other +-- atomic type; COALESCE uses the same STRING promotion (including vs BOOLEAN). +SELECT typeof(coalesce(cast('123' AS CHAR(3)), 1)), + typeof(coalesce(cast('123' AS VARCHAR(3)), 1)), + typeof(coalesce(cast('123' AS STRING), 1)); +SELECT coalesce(cast('123' AS CHAR(3)), 1), + coalesce(cast('123' AS VARCHAR(3)), 1), + coalesce(cast('123' AS STRING), 1); +SELECT cast('123' AS CHAR(3)) = 123, + cast('123' AS VARCHAR(3)) = 123, + cast('123' AS STRING) = 123; +SELECT typeof(coalesce(cast('1.5' AS CHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS VARCHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS STRING), 1.5)); +SELECT cast('2020-01-02' AS CHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS VARCHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS STRING) = date'2020-01-02'; +SELECT cast('true' AS CHAR(4)) = true, + cast('true' AS VARCHAR(4)) = true, + cast('true' AS STRING) = true; +SELECT typeof(coalesce(cast('true' AS CHAR(4)), true)); +SELECT typeof(coalesce(cast('true' AS VARCHAR(4)), true)); +SELECT typeof(coalesce(cast('true' AS STRING), true)); + +-- Nested types keep CHAR/VARCHAR +SELECT typeof(array(cast('a' AS CHAR(2)), cast('bb' AS CHAR(3)))); +SELECT typeof(struct(cast('a' AS CHAR(2)) AS f)); +SELECT typeof(map('k', cast('a' AS VARCHAR(2)))); + +-- R3: bare column references keep the declared type; write and read pad CHAR +CREATE TABLE char_varchar_std (c CHAR(5), v VARCHAR(5)) USING parquet; +INSERT INTO char_varchar_std VALUES ('ab', 'ab'); +SELECT typeof(c), typeof(v) FROM char_varchar_std; +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std; +SELECT length(c), length(v) FROM char_varchar_std; + +-- Language surfaces: CTAS / VIEW inherit CHAR/VARCHAR; ORC catalog round-trip +CREATE TABLE char_varchar_std_ctas USING parquet AS SELECT c, v FROM char_varchar_std; +SELECT typeof(c), typeof(v) FROM char_varchar_std_ctas; +CREATE VIEW char_varchar_std_view AS SELECT c FROM char_varchar_std; +SELECT typeof(c) FROM char_varchar_std_view; +CREATE VIEW char_varchar_std_view_v AS SELECT v FROM char_varchar_std; +SELECT typeof(v) FROM char_varchar_std_view_v; +WITH t AS (SELECT c, v FROM char_varchar_std) SELECT typeof(c), typeof(v) FROM t; +CREATE TABLE char_varchar_std_orc (c CHAR(5), v VARCHAR(5)) USING orc; +INSERT INTO char_varchar_std_orc VALUES ('ab', 'cd'); +SELECT typeof(c), typeof(v) FROM char_varchar_std_orc; +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std_orc; + +DROP VIEW char_varchar_std_view_v; +DROP VIEW char_varchar_std_view; +DROP TABLE char_varchar_std_ctas; +DROP TABLE char_varchar_std_orc; +DROP TABLE char_varchar_std; diff --git a/sql/core/src/test/resources/sql-tests/inputs/distinct-map-aggregates.sql b/sql/core/src/test/resources/sql-tests/inputs/distinct-map-aggregates.sql new file mode 100644 index 0000000000000..523ed0caee89e --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/distinct-map-aggregates.sql @@ -0,0 +1,102 @@ +-- Test DISTINCT aggregates with MapType arguments. + +CREATE OR REPLACE TEMPORARY VIEW distinct_map_data AS SELECT * FROM VALUES + (2, map('a', 1, 'b', 2), 1, true), + (2, map('b', 2, 'a', 1), 1, true), + (1, map('a', 1, 'b', 2), 1, true), + (1, map('a', 3), 2, false) +AS distinct_map_data(g, m, id, should_keep); + +SELECT COUNT(DISTINCT m) FROM distinct_map_data; + +SELECT SIZE(COLLECT_LIST(DISTINCT m)) FROM distinct_map_data; + +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a'); + +SELECT map_entries(FIRST(DISTINCT m)), map_entries(LAST(DISTINCT m)), COUNT(DISTINCT m) +FROM VALUES (map('b', 2, 'a', 1)) AS single_map_data(m); + +SELECT COUNT(DISTINCT m, id) FROM distinct_map_data; + +SELECT COUNT(DISTINCT m), COUNT(DISTINCT id) FROM distinct_map_data; + +SELECT g, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY g +ORDER BY g; + +SELECT m, COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a'); + +SELECT COUNT(DISTINCT m) FILTER (WHERE should_keep) FROM distinct_map_data; + +SELECT MAX(map_values(m)[0]) +FROM distinct_map_data +WHERE id = 1; + +SELECT MAX(map_values(m)[0]), COUNT(DISTINCT m) +FROM distinct_map_data +WHERE id = 1; + +SELECT g +FROM distinct_map_data +GROUP BY g +ORDER BY COUNT(DISTINCT m), g; + +SELECT g +FROM distinct_map_data +GROUP BY g +HAVING COUNT(DISTINCT m) = 1 +ORDER BY g; + +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data; + +SELECT COUNT(DISTINCT array(m)) FROM distinct_map_data; + +SELECT COUNT(DISTINCT map('m', m)) FROM distinct_map_data; + +SELECT COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM VALUES + (CAST(map() AS MAP<STRING, INT>)), + (CAST(map() AS MAP<STRING, INT>)), + (CAST(NULL AS MAP<STRING, INT>)) +AS null_and_empty_map_data(m); + +SELECT g, GROUPING(g), COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY GROUPING SETS ((g), ()) +ORDER BY GROUPING(g), g; + +SELECT COUNT(DISTINCT named_struct('m', m, 'n', n)) +FROM VALUES + (map('a', 1, 'b', 2), map('x', 1, 'y', 2)), + (map('b', 2, 'a', 1), map('y', 2, 'x', 1)) +AS grouped_distinct_map_data(m, n) +GROUP BY m; + +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=false; + +SELECT COUNT(DISTINCT m) FROM distinct_map_data; + +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a'), map_entries(m)[0].key; + +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data; + +SELECT m, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a'); + +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=true; diff --git a/sql/core/src/test/resources/sql-tests/inputs/generators-resolution-edge-cases.sql b/sql/core/src/test/resources/sql-tests/inputs/generators-resolution-edge-cases.sql index 67b5ae77c3369..5928ccea56e76 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/generators-resolution-edge-cases.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/generators-resolution-edge-cases.sql @@ -215,3 +215,29 @@ SELECT posexplode(array('x', 'y')) as (pos, val), pos, val FROM (VALUES (42)) AS -- generator's multi-alias does not shadow table column with aggregate SELECT posexplode(array('x', 'y')) as (pos, val), pos, val, count(*) FROM (VALUES (42)) AS t(pos) GROUP BY pos; + +-- HAVING on a grouping key should be evaluated before a window function with a generator +SELECT explode(array(a)) AS col, count(*) OVER () AS cnt, a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL; + +-- HAVING on an aggregate should be evaluated before a window function with a generator +SELECT explode(array(a)) AS col, + count(*) OVER () AS group_count, + count(*) AS row_count +FROM VALUES (1), (1), (2), (3), (3) AS t(a) +GROUP BY a +HAVING row_count > 1; + +-- HAVING should be evaluated before multiple window functions with a generator +SELECT explode(array(a)) AS col, + count(*) OVER () AS cnt, + count(*) OVER ( + ORDER BY a + ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + ) AS ordered_cnt, + a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL; diff --git a/sql/core/src/test/resources/sql-tests/inputs/identifier-clause.sql b/sql/core/src/test/resources/sql-tests/inputs/identifier-clause.sql index 1425004e9fcca..3c4c4f73e3bc0 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/identifier-clause.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/identifier-clause.sql @@ -14,6 +14,7 @@ SELECT IDENTIFIER('`t`.c1') FROM VALUES(1) AS T(c1); SELECT IDENTIFIER('`c 1`') FROM VALUES(1) AS T(`c 1`); SELECT IDENTIFIER('``') FROM VALUES(1) AS T(``); SELECT IDENTIFIER('c' || '1') FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(concat('a', 'b')) FROM VALUES(1) AS T(ab); -- Table references CREATE SCHEMA IF NOT EXISTS s; @@ -29,6 +30,7 @@ SELECT * FROM IDENTIFIER('tab'); SELECT * FROM IDENTIFIER('s.tab'); SELECT * FROM IDENTIFIER('`s`.`tab`'); SELECT * FROM IDENTIFIER('t' || 'a' || 'b'); +SELECT * FROM IDENTIFIER(concat('t', 'ab')); USE SCHEMA default; DROP TABLE s.tab; @@ -36,9 +38,73 @@ DROP SCHEMA s; -- Function reference SELECT IDENTIFIER('COAL' || 'ESCE')(NULL, 1); +SELECT IDENTIFIER(concat('COAL', 'ESCE'))(NULL, 1); SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1); SELECT * FROM IDENTIFIER('ra' || 'nge')(0, 1); +VALUES(IDENTIFIER(abs(1))); +SELECT * FROM IDENTIFIER(nullif('a', 'a')); + +-- Function resolution while computing an identifier name +SELECT IDENTIFIER(max('c1')) FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(array_join(transform(array('c', '1'), element -> element), '')) +FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(rand()) FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(row_number() OVER ()) FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(row_number() OVER (ORDER BY 'x')) FROM VALUES(1) AS T(c1); +SELECT IDENTIFIER(explode(array('c1'))) FROM VALUES(1) AS T(c1); +SELECT * FROM IDENTIFIER(max('identifier_function_table')); +SELECT * FROM IDENTIFIER(row_number() OVER (ORDER BY 'x')); +SELECT * FROM IDENTIFIER( + array_join(transform(array('identifier', '_function_table'), element -> element), '') +); + +CREATE TEMPORARY FUNCTION identifier_name() +RETURNS STRING +RETURN 'c1'; +SELECT IDENTIFIER(identifier_name()) FROM VALUES(1) AS T(c1); +DROP TEMPORARY FUNCTION identifier_name; + +CREATE FUNCTION persistent_identifier_name() +RETURNS STRING +RETURN 'c1'; +SELECT IDENTIFIER(persistent_identifier_name()) FROM VALUES(1) AS T(c1); +DROP FUNCTION persistent_identifier_name; + +CREATE FUNCTION persistent_identifier_function(value INT) +RETURNS INT +RETURN value + 1; +SELECT IDENTIFIER(concat('persistent_identifier_', 'function'))(1); +DROP FUNCTION persistent_identifier_function; + +CREATE FUNCTION persistent_identifier_base_name() +RETURNS STRING +RETURN 'c1'; +CREATE FUNCTION persistent_identifier_nested_name() +RETURNS STRING +RETURN persistent_identifier_base_name(); +SELECT IDENTIFIER(persistent_identifier_nested_name()) FROM VALUES(1) AS T(c1); +DROP FUNCTION persistent_identifier_nested_name; +DROP FUNCTION persistent_identifier_base_name; + +CREATE TEMPORARY FUNCTION identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table'; +CREATE TABLE identifier_function_table(c1 INT) USING csv; +SELECT * FROM IDENTIFIER(identifier_relation_name()); +CREATE TEMPORARY FUNCTION identifier_relation_count() +RETURNS BIGINT +RETURN SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')); +SELECT identifier_relation_count(); +DROP TEMPORARY FUNCTION identifier_relation_count; +CREATE FUNCTION persistent_identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table'; +SELECT * FROM IDENTIFIER(persistent_identifier_relation_name()); +DROP FUNCTION persistent_identifier_relation_name; +DROP TABLE identifier_function_table; +DROP TEMPORARY FUNCTION identifier_relation_name; + -- Table DDL CREATE TABLE IDENTIFIER('tab')(c1 INT) USING CSV; DROP TABLE IF EXISTS IDENTIFIER('ta' || 'b'); @@ -120,11 +186,17 @@ VALUES(IDENTIFIER(SUBSTR('HELLO', 1, RAND() + 1))); SELECT `IDENTIFIER`('abs')(c1) FROM VALUES(-1) AS T(c1); CREATE TABLE t(col1 INT); +CREATE TABLE identifier_name_source(name STRING) USING csv; SELECT * FROM IDENTIFIER((SELECT 't')); +SELECT * FROM IDENTIFIER((SELECT max(name) FROM identifier_name_source)); +SELECT * FROM IDENTIFIER( + (SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1) +); SELECT * FROM (SELECT IDENTIFIER((SELECT 'col1')) FROM IDENTIFIER((SELECT 't'))); SELECT IDENTIFIER((SELECT 'col1')) FROM VALUES(1); SELECT col1, IDENTIFIER((SELECT col1)) FROM VALUES(1); SELECT IDENTIFIER((SELECT 'col1', 'col2')) FROM VALUES(1,2); +DROP TABLE identifier_name_source; DROP TABLE t; CREATE TABLE IDENTIFIER(1)(c1 INT) USING csv; diff --git a/sql/core/src/test/resources/sql-tests/inputs/json-functions.sql b/sql/core/src/test/resources/sql-tests/inputs/json-functions.sql index 66134545107a3..43a0a5b62db91 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/json-functions.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/json-functions.sql @@ -108,6 +108,27 @@ select json_object_keys('{[1,2]}'); select json_object_keys('{"key": 45, "random_string"}'); select json_object_keys('[1, 2, 3]'); +-- json_typeof +select json_typeof(); +select json_typeof(null); +select json_typeof(200); +select json_typeof(''); +select json_typeof('{}'); +select json_typeof('{"key": 1, "arr": [1, 2]}'); +select json_typeof('[]'); +select json_typeof('[1, 2, 3]'); +select json_typeof('"hello"'); +select json_typeof('123'); +select json_typeof('1.5'); +select json_typeof('-123'); +select json_typeof('-1.5'); +select json_typeof('true'); +select json_typeof('false'); +select json_typeof('null'); +select json_typeof('bad'); +select json_typeof('{"key": 45, "random_string"}'); +select json_typeof('123 true'); + -- Clean up DROP VIEW IF EXISTS jsonTable; @@ -152,3 +173,115 @@ select to_json(from_json('{"time":"23:59:59.999999"}', 'time TIME(6)')); select schema_of_json('{"time": "14:30:45"}'); select schema_of_json('{"time": "14:30:45.123456"}'); select from_json('{"time": "14:30:45"}', 'time TIME') LIMIT 1; + +-- JSON_VALUE: extract a scalar value (ANSI SQL:2016) +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.name'); +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.id' RETURNING INT); +select json_value('{"id":7,"name":"Ada"}', '$.id' RETURNING INT) + 1; +-- present but JSON null -> SQL NULL +select json_value('{"score":null}', '$.score'); +-- non-scalar (object / array) -> NULL ON ERROR (default) +select json_value('{"addr":{"city":"NYC"}}', '$.addr'); +select json_value('{"tags":["x","y"]}', '$.tags'); +-- missing path -> NULL ON EMPTY (default) +select json_value('{"id":7}', '$.missing'); +-- NULL input propagates to NULL +select json_value(cast(null as string), '$.a'); +-- ON EMPTY behaviors +select json_value('{"id":7}', '$.missing' DEFAULT '?' ON EMPTY); +select json_value('{"id":7}', '$.missing' RETURNING INT DEFAULT 42 ON EMPTY); +select json_value('{"id":7}', '$.missing' ERROR ON EMPTY); +-- ON ERROR behaviors +select json_value('{"addr":{"city":"NYC"}}', '$.addr' DEFAULT 'n/a' ON ERROR); +select json_value('not json', '$.a' DEFAULT 'bad' ON ERROR); +select json_value('not json', '$.a' ERROR ON ERROR); +-- failed cast -> ON ERROR +select json_value('{"name":"Ada"}', '$.name' RETURNING INT); +select json_value('{"name":"Ada"}', '$.name' RETURNING INT ERROR ON ERROR); +select json_value('{"name":"Ada"}', '$.name' RETURNING INT DEFAULT -1 ON ERROR); +-- combined ON EMPTY + ON ERROR +select json_value('{"a":"x"}', '$.b' DEFAULT 'e' ON EMPTY DEFAULT 'r' ON ERROR); +-- RETURNING types +select json_value('{"v":"3.14"}', '$.v' RETURNING DOUBLE); +select json_value('{"v":"true"}', '$.v' RETURNING BOOLEAN); +select json_value('{"v":"2020-01-02"}', '$.v' RETURNING DATE); +-- invalid: wildcard path +select json_value('{"a":[1,2]}', '$.a[*]'); +-- invalid: non-scalar RETURNING type +select json_value('{"a":1}', '$.a' RETURNING STRUCT<x:INT>); +-- invalid: a DEFAULT that cannot cast to the RETURNING type +select json_value('{}', '$.x' RETURNING INT DEFAULT array(1) ON EMPTY); + +-- JSON_EXISTS: test path presence (ANSI SQL:2016) +select json_exists('{"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}', '$.addr.city'); +-- present but JSON null -> true +select json_exists('{"score":null}', '$.score'); +-- absent -> false +select json_exists('{"addr":{"city":"NYC"}}', '$.addr.zip'); +-- matches an object / array -> true +select json_exists('{"addr":{"city":"NYC"}}', '$.addr'); +select json_exists('{"tags":["x","y"]}', '$.tags[0]'); +-- NULL input -> NULL (unknown) +select json_exists(cast(null as string), '$.a'); +-- malformed input -> FALSE ON ERROR (default) +select json_exists('not json', '$.a'); +-- ON ERROR behaviors +select json_exists('not json', '$.a' TRUE ON ERROR); +select json_exists('not json', '$.a' FALSE ON ERROR); +select json_exists('not json', '$.a' UNKNOWN ON ERROR); +select json_exists('not json', '$.a' ERROR ON ERROR); +-- lax wildcard [*]: true iff the array has elements +select json_exists('{"a":[1,2]}', '$.a[*]'); +select json_exists('{"a":[]}', '$.a[*]'); +-- lax auto-wrap: [*] over a non-array treats it as a single-element array +select json_exists('{"a":5}', '$.a[*]'); +-- embedded wildcard: any element has the field +select json_exists('{"a":[{"b":1},{"c":2}]}', '$.a[*].b'); +-- out-of-range index -> false +select json_exists('{"a":[1,2]}', '$.a[5]'); +-- lax auto-unwrap: a member step over an array applies to each element +select json_exists('{"a":[{"b":1},{"b":2}]}', '$.a.b'); +-- member wildcard .* matches any member +select json_exists('{"addr":{"city":"NYC"}}', '$.*'); +-- invalid: an unparseable path is rejected at analysis +select json_exists('{"a":1}', '$['); + +-- JSON_QUERY: extract an object or array as JSON text (ANSI SQL:2016) +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.addr'); +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.tags'); +-- a scalar result is emitted as JSON text (not an error) under the default WITHOUT ARRAY WRAPPER +select json_query('{"id":7}', '$.id'); +select json_query('{"name":"Ada"}', '$.name'); +-- present but JSON null -> the JSON text null +select json_query('{"score":null}', '$.score'); +-- missing path -> NULL ON EMPTY (default) +select json_query('{"id":7}', '$.missing'); +-- NULL input propagates to NULL +select json_query(cast(null as string), '$.a'); +-- ARRAY WRAPPER +select json_query('{"tags":["x","y"]}', '$.tags[0]' WITH ARRAY WRAPPER); +select json_query('{"tags":["x","y"]}', '$.tags' WITH UNCONDITIONAL ARRAY WRAPPER); +select json_query('{"id":7}', '$.id' WITH ARRAY WRAPPER); +-- CONDITIONAL wraps only a scalar; an object/array is left as is +select json_query('{"id":7}', '$.id' WITH CONDITIONAL ARRAY WRAPPER); +select json_query('{"addr":{"city":"NYC"}}', '$.addr' WITH CONDITIONAL ARRAY WRAPPER); +-- OMIT QUOTES strips the quotes from a scalar string result +select json_query('{"name":"Ada"}', '$.name' OMIT QUOTES); +select json_query('{"name":"Ada"}', '$.name' KEEP QUOTES); +-- ON EMPTY behaviors +select json_query('{"id":7}', '$.missing' EMPTY ARRAY ON EMPTY); +select json_query('{"id":7}', '$.missing' EMPTY OBJECT ON EMPTY); +select json_query('{"id":7}', '$.missing' ERROR ON EMPTY); +-- ON ERROR behaviors (malformed input) +select json_query('not json', '$.a'); +select json_query('not json', '$.a' EMPTY ARRAY ON ERROR); +select json_query('not json', '$.a' EMPTY OBJECT ON ERROR); +select json_query('not json', '$.a' ERROR ON ERROR); +-- RETURNING STRING is allowed (the result is JSON text) +select json_query('{"addr":{"city":"NYC"}}', '$.addr' RETURNING STRING); +-- invalid: wildcard path +select json_query('{"a":[1,2]}', '$.a[*]'); +-- invalid: non-string RETURNING type +select json_query('{"a":1}', '$.a' RETURNING INT); +-- invalid: OMIT QUOTES combined with an array wrapper +select json_query('{"name":"Ada"}', '$.name' WITH ARRAY WRAPPER OMIT QUOTES); diff --git a/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql b/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql index a3fa6d4c4cd49..7939b2b9f4c70 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/linear-regression.sql @@ -54,3 +54,17 @@ SELECT k, regr_intercept(y, x) FROM testRegression WHERE x IS NOT NULL AND y IS -- SPARK-55969: regr_r2 should treat first param as dependent variable SELECT regr_r2(k, x) FROM testRegression where k=2; SELECT regr_r2(y, k) FROM testRegression where k=2; + +-- SPARK-58213: corr should return NULL for constant columns (zero variance) +CREATE OR REPLACE TEMPORARY VIEW testCorrConstant AS SELECT * FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y); +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM testCorrConstant; +DROP VIEW testCorrConstant; + +-- x-constant (y varies): corr should return NULL +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y); + +-- y-constant (x varies): corr should return NULL +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (2, 1), (3, 1) AS t(x, y); diff --git a/sql/core/src/test/resources/sql-tests/inputs/math.sql b/sql/core/src/test/resources/sql-tests/inputs/math.sql index 14a647a610cc3..bd9a6b2c3de3c 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/math.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/math.sql @@ -70,6 +70,64 @@ SELECT bround(525L, -3); SELECT bround(9223372036854775807L, -1); SELECT bround(-9223372036854775808L, -1); +-- Truncate with Byte input +SELECT truncate(25y, 1); +SELECT truncate(25y, 0); +SELECT truncate(25y, -1); +SELECT truncate(25y, -2); +SELECT truncate(25y, -3); +-- Truncate with negative Byte input: truncation is toward zero, unlike floor. +SELECT truncate(-25y, 1); +SELECT truncate(-25y, 0); +SELECT truncate(-25y, -1); +SELECT truncate(-25y, -2); +SELECT truncate(-25y, -3); +-- Truncate never overflows, unlike round: truncate(-128y, -1) is -120, not an overflow. +SELECT truncate(127y, -1); +SELECT truncate(-128y, -1); + +-- Truncate with short integer input +SELECT truncate(525s, 1); +SELECT truncate(525s, 0); +SELECT truncate(525s, -1); +SELECT truncate(525s, -2); +SELECT truncate(525s, -3); +-- Truncate with negative short integer input: truncation is toward zero, unlike floor. +SELECT truncate(-525s, 1); +SELECT truncate(-525s, 0); +SELECT truncate(-525s, -1); +SELECT truncate(-525s, -2); +SELECT truncate(-525s, -3); + +-- Truncate with integer input +SELECT truncate(525, 1); +SELECT truncate(525, 0); +SELECT truncate(525, -1); +SELECT truncate(525, -2); +SELECT truncate(525, -3); +-- Truncate with negative integer input: truncation is toward zero, unlike floor. +SELECT truncate(-525, 1); +SELECT truncate(-525, 0); +SELECT truncate(-525, -1); +SELECT truncate(-525, -2); +SELECT truncate(-525, -3); + +-- Truncate with big integer input +SELECT truncate(525L, 1); +SELECT truncate(525L, 0); +SELECT truncate(525L, -1); +SELECT truncate(525L, -2); +SELECT truncate(525L, -3); +-- Truncate with negative big integer input: truncation is toward zero, unlike floor. +SELECT truncate(-525L, 1); +SELECT truncate(-525L, 0); +SELECT truncate(-525L, -1); +SELECT truncate(-525L, -2); +SELECT truncate(-525L, -3); + +-- Truncate with the scale argument omitted; it defaults to 0. +SELECT truncate(1234.5678); + -- Conv SELECT conv('100', 2, 10); SELECT conv(-10, 16, -10); diff --git a/sql/core/src/test/resources/sql-tests/inputs/misc-functions.sql b/sql/core/src/test/resources/sql-tests/inputs/misc-functions.sql index e0b7b6ac3f88f..232f16963c84a 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/misc-functions.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/misc-functions.sql @@ -58,3 +58,12 @@ SELECT hmac('key', CAST(NULL AS BINARY)); SELECT hmac('key', 'message', CAST(NULL AS STRING)); -- Unsupported algorithm. SELECT hmac('key', 'message', 'SHA-3'); + +-- xxh3_64 and xxh3_128 +SELECT xxh3_64('Spark'); +SELECT xxh3_64(CAST('Spark' AS BINARY)); +SELECT xxh3_128('Spark'); +SELECT xxh3_128(CAST('Spark' AS BINARY)); +-- Null propagation. +SELECT xxh3_64(CAST(NULL AS STRING)); +SELECT xxh3_128(CAST(NULL AS BINARY)); diff --git a/sql/core/src/test/resources/sql-tests/inputs/parse-sql-gating.sql b/sql/core/src/test/resources/sql-tests/inputs/parse-sql-gating.sql new file mode 100644 index 0000000000000..ae2a650f08592 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/parse-sql-gating.sql @@ -0,0 +1,4 @@ +-- parse_sql is off by default while the JSON contract is still evolving. +--SET spark.sql.function.parseSql.enabled=false + +SELECT parse_sql('SELECT 1'); diff --git a/sql/core/src/test/resources/sql-tests/inputs/parse-sql.sql b/sql/core/src/test/resources/sql-tests/inputs/parse-sql.sql new file mode 100644 index 0000000000000..b4eb15ffbc38e --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/parse-sql.sql @@ -0,0 +1,249 @@ +-- End-to-end coverage for parse_sql (SPARK-58738). +-- Returns compact JSON for parse-only statement analysis via SparkSqlParser. +-- Off by default while the JSON contract is still evolving. +--SET spark.sql.function.parseSql.enabled=true + +-- null input +SELECT parse_sql(NULL); + +-- basic SELECT classification and references +SELECT parse_sql('SELECT a, b FROM t'); +SELECT parse_sql('SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2'); + +-- JSON-path access over one shared successful parse result +SELECT + get_json_object(result, '$.statement_identifier') AS statement_identifier, + get_json_object(result, '$.source_table_references[0][0]') AS first_table, + get_json_object(result, '$.select_list[1].name[0]') AS second_column +FROM (SELECT parse_sql('SELECT a, b FROM t') AS result); + +-- DML +SELECT parse_sql('INSERT INTO t SELECT 1'); +SELECT parse_sql('DELETE FROM t WHERE a = 1'); +SELECT parse_sql('UPDATE t SET a = 1 WHERE b = 2'); +SELECT parse_sql('MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE'); + +-- DDL / CTAS +SELECT parse_sql('CREATE TABLE t (a INT)'); +SELECT parse_sql('CREATE TABLE t AS SELECT 1 AS a'); +SELECT parse_sql('DROP TABLE t'); + +-- Spark-only statements (negative Table 39 codes) +SELECT parse_sql('CACHE TABLE t'); + +-- TABLE / VALUES are SELECT-shaped (not Unrecognized) +SELECT parse_sql('TABLE t'); +SELECT parse_sql('VALUES (1), (2)'); + +-- function / variable names are not target or source table references +SELECT parse_sql('CREATE FUNCTION f AS ''x'' USING JAR ''y.jar'''); +SELECT parse_sql('DECLARE VARIABLE x INT'); + +-- parameter markers +SELECT parse_sql('SELECT * FROM t WHERE a = :foo AND b = ?'); + +-- CTE: lineage excludes CTE names; still walks CTE bodies for real tables +SELECT parse_sql('WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte'); + +-- CTE shadowing is scoped: the inner CTE real_t does not hide the outer table +SELECT parse_sql('SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)'); + +-- a CTE definition sees only preceding aliases, so b below is the real table +SELECT parse_sql('WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a'); + +-- nested subqueries +SELECT parse_sql('SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id)'); + +-- functions in projection, window, join, TVF, predicates, subquery, grouping, and ordering +SELECT parse_sql( +'SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, ''x'') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1)'); + +-- functions and tables throughout a multiline MERGE +SELECT parse_sql( +'MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name))'); + +-- functions embedded in DDL column defaults +SELECT parse_sql( +'CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper(''x'') + )'); + +-- syntax error: dump the complete STANDARD error, including query context +SELECT parse_sql('SELEC FROM t'); + +-- JSON-path access over one shared parse result +SELECT + get_json_object(result, '$.parse_success') AS parse_success, + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM (SELECT parse_sql('SELEC FROM t') AS result); + +-- full multiline parse-time validation error, including context and location +SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b'); + +-- JSON-path access over one shared multiline parse result +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].startIndex') AS start_index +FROM ( + SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b') AS result +); + +-- parse-only validation errors beyond PARSE_SYNTAX_ERROR +SELECT parse_sql(''); +SELECT parse_sql('USE bad-name'); +SELECT parse_sql('WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c'); +SELECT parse_sql('MERGE INTO target USING source ON target.id = source.id'); +SELECT parse_sql('EXPLAIN SELECT 1'); +SELECT parse_sql('SET spark.sql.adaptive.enabled=true'); +SELECT parse_sql('ADD JAR /tmp/x.jar'); +SELECT parse_sql('CREATE VIEW v AS SELECT a, b FROM t'); +SELECT parse_sql('SELECT 1 AS IDENTIFIER(''alias.field'')'); +SELECT parse_sql('SELECT DATE ''not-a-date'''); + +-- location for an error inside a multiline script +--QUERY-DELIMITER-START +SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END'); +--QUERY-DELIMITER-END + +-- JSON-path access over one shared scripting parse result +--QUERY-DELIMITER-START +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END') AS result +); +--QUERY-DELIMITER-END + +-- location for a SQL scripting semantic validation error +--QUERY-DELIMITER-START +SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END'); +--QUERY-DELIMITER-END + +-- JSON-path access over one shared scripting validation result +--QUERY-DELIMITER-START +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END') AS result +); +--QUERY-DELIMITER-END + +-- batch over a column of SQL text +SELECT sql_text, parse_sql(sql_text) FROM VALUES + ('SELECT 1'), + ('INSERT INTO t SELECT 1'), + ('CACHE TABLE t') +AS t(sql_text); + +-- BEGIN END scripts contain ';' inside the string literal; use query delimiters +-- so the test harness does not split on those semicolons. +--QUERY-DELIMITER-START +SELECT parse_sql('BEGIN SELECT 1; END'); +--QUERY-DELIMITER-END + +--QUERY-DELIMITER-START +SELECT parse_sql('BEGIN SELECT count(a) FROM script_t WHERE c = :p; END'); +--QUERY-DELIMITER-END + +-- Positional markers under SingleStatement must not be double-counted. +--QUERY-DELIMITER-START +SELECT parse_sql('BEGIN SELECT * FROM t WHERE a = ?; END'); +--QUERY-DELIMITER-END + +--QUERY-DELIMITER-START +SELECT parse_sql('BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END'); +--QUERY-DELIMITER-END + +--QUERY-DELIMITER-START +SELECT parse_sql('BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END'); +--QUERY-DELIMITER-END + +-- Complex, genuinely multiline script: dump the complete JSON result. +--QUERY-DELIMITER-START +SELECT parse_sql( +'BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string(''%s'', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END'); +--QUERY-DELIMITER-END diff --git a/sql/core/src/test/resources/sql-tests/inputs/string-functions.sql b/sql/core/src/test/resources/sql-tests/inputs/string-functions.sql index c432093e1ae88..4e4cb6a401f5b 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/string-functions.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/string-functions.sql @@ -370,4 +370,11 @@ select instr('a', 'b', 1, cast(null as int)); select instr(null, 'b', 1); select instr('a', null, 1); select instr('a', 'b', cast(null as int), 2); -select instr(null, null, cast(null as int), cast(null as int)); \ No newline at end of file +select instr(null, null, cast(null as int), cast(null as int)); + +-- normalize +select normalize('hello'); +select normalize('hello', 'NFD'); +select normalize('fi', 'NFKC'); +select normalize(null, 'NFC'); +select normalize('hello', 'not_a_form'); \ No newline at end of file diff --git a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql index 7c816d8a41672..a094dcb81157a 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/subquery/in-subquery/in-limit.sql @@ -88,6 +88,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a LIMIT 10 OFFSET 2) LIMIT 2 @@ -109,6 +110,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a OFFSET 2) OFFSET 1; @@ -118,6 +120,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4; @@ -138,6 +141,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b <= t1d + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4; @@ -148,6 +152,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2); -- TC 01.05 @@ -200,6 +205,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2 OFFSET 2) LIMIT 4 @@ -225,6 +231,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2 OFFSET 2); @@ -233,6 +240,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b <= t1d + ORDER BY t2b LIMIT 2); SELECT * @@ -296,6 +304,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c DESC NULLS LAST OFFSET 2) OFFSET 4; @@ -327,6 +336,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b OFFSET 2); -- OFFSET with NOT IN correlated diff --git a/sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql b/sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql index 62e9b8eb93b5f..1eba70d4f4b2b 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql @@ -168,11 +168,48 @@ SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1 00:04:00.000321' DAY TO SECOND; SELECT TIMESTAMP_LTZ '1960-01-02 03:04:05.123456789 UTC' + INTERVAL '0 00:00:00.000001' DAY TO SECOND; --- SPARK-57501: nanos timestamps support only ANSI day-time intervals. A (legacy) calendar interval --- is rejected by TimestampAddInterval's type check, and a year-month interval has no supported --- operator overload. -SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0); +-- SPARK-57825: TIMESTAMP_LTZ(p) +/- ANSI year-month interval keeps the nanos type/precision and +-- carries the whole fraction (including the sub-microsecond digits) through unchanged; the month +-- shift is applied on the session-local wall clock. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' YEAR; +-- The interval-first operand order resolves to the same addition. +SELECT INTERVAL '1' YEAR + TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC'; SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH; +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1-2' YEAR TO MONTH; +-- Jan-31 -> Feb-29 day clamp on a pre-epoch (leap-year) value. +SELECT TIMESTAMP_LTZ '1960-01-31 03:04:05.123456789 UTC' + INTERVAL '1' MONTH; +-- SPARK-57501, SPARK-57825: nanos timestamps support ANSI day-time and year-month intervals; the +-- legacy calendar interval is still rejected by TimestampAddInterval's type check. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0); + +-- SPARK-57832: TIMESTAMP_LTZ(p) - TIMESTAMP_LTZ(p) yields a microsecond-grid DayTimeIntervalType. +-- Only each operand's epochMicros participates, so the sub-microsecond remainder is truncated: the +-- 789/111 sub-micro digits drop out and the difference is exactly 1 day + 0.123456 s. Like the +-- micro TIMESTAMP_LTZ case, the subtraction runs on session-zone local date-times, so a DST +-- transition between the two instants could shift the interval; there is none in the UTC span here. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC'; +-- Two values inside the same microsecond subtract to zero once the remainder is truncated. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456001 UTC'; +-- The subtraction is antisymmetric. +SELECT TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC'; +-- Mixed precision (7 vs 9) widens to the common nanos type before subtracting. +SELECT ('2020-01-02 03:04:05.1234567 UTC' :: timestamp_ltz(7)) - ('2020-01-01 03:04:05.000000009 UTC' :: timestamp_ltz(9)); +-- Mixed with a micro TIMESTAMP_LTZ operand. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05 UTC'; +-- A DATE operand is cast to the nanos LTZ type (midnight in the session time zone); the +-- fraction below the micro grid drops. The bare LTZ literal and the DATE both read in the session +-- zone (America/Los_Angeles), so the difference is exactly 1 day. +SELECT TIMESTAMP_LTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01'; +-- Pre-epoch operand exercises the negative-epoch path. +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC' - TIMESTAMP_LTZ '1960-01-01 00:00:00.000000999 UTC'; +-- NULL operand propagates. +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - CAST(NULL AS timestamp_ltz(9)); + +-- SPARK-57818: convert_timezone is NTZ-only, so a nanosecond LTZ(p) source is rejected rather +-- than silently reinterpreted as NTZ (the positive TIMESTAMP_NTZ(p) path is covered in +-- timestamp-ntz-nanos.sql). +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.123456789 UTC' :: timestamp_ltz(9)); -- SPARK-57103: MAX / MIN over nanosecond-precision TIMESTAMP_LTZ. The aggregate preserves the -- nanosecond type and orders by the sub-microsecond remainder; NULLs are ignored. Values are @@ -188,6 +225,47 @@ SELECT c, count(*) FROM VALUES (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c) GROUP BY c ORDER BY c; +-- GROUP BY a nanosecond key with aggregates and a NULL group: exact-duplicate keys collapse, two +-- keys sharing epochMicros but differing within the microsecond stay in separate groups, and all +-- NULL keys group together (unlike an equi-join). Three groups: .000000001 (count 2, sum 3), +-- .000000999 (count 1, sum 3), NULL (count 2, sum 9). Values render in the session time zone. +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3), + (CAST(NULL AS timestamp_ltz(9)), 4), + (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k; + +-- SPARK-56822: mode over nanosecond-precision TIMESTAMP_LTZ. Frequencies are counted on the full +-- nanos value, so the most-frequent value is selected down to the sub-microsecond and the result +-- type stays TIMESTAMP_LTZ(9); the value renders in the session time zone (America/Los_Angeles). +-- .000000001 appears twice, .000000999 once. +SELECT mode(c) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c); + +-- SPARK-56822: collect_set over nanosecond-precision TIMESTAMP_LTZ. It deduplicates on the full +-- sub-microsecond value: the two .000000001 rows collapse to one, the .000000999 row stays, so the +-- sorted set has two distinct elements and the element type stays TIMESTAMP_LTZ(9); values render +-- in the session time zone (America/Los_Angeles). collect_set order is non-deterministic, so the +-- output is stabilized with sort_array. +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c); + +-- SPARK-56822: collect_list over nanosecond-precision TIMESTAMP_LTZ. The buffer holds the full +-- nanos value, so the sub-microsecond remainder survives and the result element type stays +-- TIMESTAMP_LTZ(9); values render in the session time zone (America/Los_Angeles). collect_list +-- order is non-deterministic, so the output is stabilized with sort_array; duplicates are kept and +-- NULLs are dropped. +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c); -- SPARK-57528: unix_timestamp / to_unix_timestamp over nanosecond-precision values. The result is -- whole-second BIGINT; the sub-second digits are dropped. A literal without an explicit zone is @@ -211,6 +289,15 @@ SELECT max_by(v, k), min_by(v, k) FROM VALUES (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3), (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000500 UTC', 2), (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000007 UTC', CAST(NULL AS INT)) AS t(v, k); +-- DISTINCT over a nanosecond column: exact duplicates are removed, two values sharing epochMicros +-- but differing within the microsecond are both kept, and NULL survives as a single row. Three +-- rows: .000000001, .000000999, NULL. Values render in the session time zone. +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c) + ORDER BY c; -- SPARK-57527: unix_nanos over nanosecond-precision values returns DECIMAL(21, 0) nanoseconds since -- the epoch. The explicit-zone literals below fix the instant directly, independent of the session @@ -329,6 +416,36 @@ SELECT v, lead(v) OVER (ORDER BY v) AS next_v FROM ( UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000100' UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000500') ORDER BY v; +-- SPARK-57811: a string operand is coerced to the nanosecond timestamp type in comparisons and +-- predicates (not truncated to micros, not promoted to string). The 9th fractional digit is +-- significant, so an off-by-one-nanosecond literal does not compare equal. +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c); + +-- BETWEEN over nanosecond timestamps: only the value inside the sub-microsecond range qualifies. +SELECT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005'; + +-- SPARK-57811: this is the config that exercises the new non-ANSI production arms. The arms live +-- in TypeCoercion (not AnsiTypeCoercion), so ANSI must be disabled; and only under legacy +-- castDatetimeToString does the range path differ. With both flags set, the range comparison +-- promotes BOTH operands to string (Cast(c AS STRING), matching the micro TIMESTAMP type), while +-- equality still casts the string to the nanos type (the Equality arm fires before the range arm). +-- In every other config (ANSI on, or ANSI off without the legacy flag) the string coerces to the +-- nanos type, identical to the parent commit, which is why the default-config cases above do not +-- distinguish this change. +SET spark.sql.ansi.enabled=false; +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true; +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c); +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false; +SET spark.sql.ansi.enabled=true; + -- SPARK-57814: unix_seconds / unix_millis / unix_micros over nanosecond-precision values. The result -- is a whole BIGINT count of the unit; sub-unit digits are dropped. Explicit-zone literals fix the -- instant directly, independent of the session time zone. @@ -370,3 +487,84 @@ SELECT date_trunc('HOUR', '2020-01-01 12:34:56.123456789 UTC' :: timestamp_ltz(9 SELECT date_trunc('DAY', '2020-01-01 04:00:00.000000123 UTC' :: timestamp_ltz(7)); -- An unsupported (sub-microsecond) unit yields NULL; the result still carries the nanos type. SELECT date_trunc('NANOSECOND', TIMESTAMP_LTZ '2020-01-01 12:34:56.123456789 UTC'); + +-- SPARK-57837: current_timestamp(p) / now(p) with a nanosecond precision return TIMESTAMP_LTZ(p). +-- The values are non-deterministic, so only the (deterministic) result type and query-stable +-- self-equality are checked. Precision 6 keeps the standard microsecond TIMESTAMP. +SELECT typeof(current_timestamp(9)), typeof(current_timestamp(8)), typeof(current_timestamp(7)); +SELECT typeof(now(9)), typeof(now(6)); +SELECT typeof(current_timestamp()), typeof(current_timestamp(6)); +-- A foldable (constant) precision expression is accepted. +SELECT typeof(current_timestamp(7 + 2)); +-- All references to current_timestamp(p) within a query see the same value. +SELECT current_timestamp(9) = current_timestamp(9), now(9) = current_timestamp(9); +-- Out-of-range precision is rejected. +SELECT current_timestamp(3); +SELECT current_timestamp(10); + +-- SPARK-57841: end-to-end coverage for operators that ride on the resolved widening (SPARK-57454) +-- and complex-type access over nanosecond values, mirroring timestamp-ntz-nanos.sql for the LTZ +-- family. Every case turns on the SUB-MICROSECOND remainder or on cross-precision widening. LTZ is +-- zone-dependent, so literals carry an explicit UTC zone and render in the session zone +-- (America/Los_Angeles, UTC-08:00). Multi-row queries end in a top-level ORDER BY. + +-- INTERSECT / EXCEPT distinguish the sub-microsecond remainder. +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + INTERSECT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c; +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + EXCEPT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c; +-- Mixed-precision set op widens to the wider precision. +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009 UTC' :: timestamp_ltz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900 UTC' :: timestamp_ltz(9))) ORDER BY c; + +-- BETWEEN on a sub-microsecond boundary. +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000500 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'; +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000001000 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'; +-- Mixed-precision BETWEEN widens the bounds to the probe's precision. +SELECT '2020-01-01 00:00:00.000000500 UTC' :: timestamp_ltz(9) + BETWEEN '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7) + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000001 UTC'; + +-- if / nvl / ifnull preserve the nanos type and widen mixed-precision branches to the wider type. +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7), + TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC') AS v); +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ltz(9)), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS v); +SELECT ifnull(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', CAST(NULL AS timestamp_ltz(9))); + +-- IN (subquery): the semi-join matches on the full nanos key. +SELECT k FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS t(k) + WHERE k IN (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') ORDER BY k; + +-- explode(array<ts_nanos>) yields one row per element. +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'))) ORDER BY col; + +-- element_at over array<ts_nanos> (1-based). +SELECT element_at(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), 2); + +-- struct-field extraction. +SELECT (named_struct('f', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')).f; + +-- map lookup by string key and by nanosecond key. +SELECT map('k', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')['k']; +SELECT map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b')[ + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC']; +SELECT element_at(map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b'), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'); diff --git a/sql/core/src/test/resources/sql-tests/inputs/timestamp-ntz-nanos.sql b/sql/core/src/test/resources/sql-tests/inputs/timestamp-ntz-nanos.sql index 971a2f8e148fd..266829fe7ab26 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/timestamp-ntz-nanos.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/timestamp-ntz-nanos.sql @@ -144,11 +144,49 @@ SELECT named_struct('f', DATE '2020-01-01') :: struct<f: timestamp_ntz(9)>; SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '2 00:03:00.000456' DAY TO SECOND; SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1 00:04:00.000321' DAY TO SECOND; SELECT TIMESTAMP_NTZ '1960-01-02 03:04:05.123456789' + INTERVAL '0 00:00:00.000001' DAY TO SECOND; --- SPARK-57501: nanos timestamps support only ANSI day-time intervals. A (legacy) calendar interval --- is rejected by TimestampAddInterval's type check, and a year-month interval has no supported --- operator overload. -SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + make_interval(0, 1, 0, 2, 0, 0, 0); +-- SPARK-57825: TIMESTAMP_NTZ(p) +/- ANSI year-month interval keeps the nanos type/precision and +-- carries the whole fraction (including the sub-microsecond digits) through unchanged; a month +-- shift never touches the time of day. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR; +-- The interval-first operand order resolves to the same addition. +SELECT INTERVAL '1' YEAR + TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789'; SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH; +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH; +-- Jan-31 -> Feb-29 day clamp on a pre-epoch (leap-year) value. +SELECT TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH; +-- SPARK-57501, SPARK-57825: nanos timestamps support ANSI day-time and year-month intervals; the +-- legacy calendar interval is still rejected by TimestampAddInterval's type check. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + make_interval(0, 1, 0, 2, 0, 0, 0); + +-- SPARK-57832: TIMESTAMP_NTZ(p) - TIMESTAMP_NTZ(p) yields a microsecond-grid DayTimeIntervalType. +-- Only each operand's epochMicros participates, so the sub-microsecond remainder is truncated: the +-- 789/111 sub-micro digits drop out and the difference is exactly 1 day + 0.123456 s. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111'; +-- Two values inside the same microsecond subtract to zero once the remainder is truncated. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456001'; +-- The subtraction is antisymmetric. +SELECT TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789'; +-- Mixed precision (7 vs 9) widens to the common nanos type before subtracting; the result stays on +-- the micros grid. +SELECT ('2020-01-02 03:04:05.1234567' :: timestamp_ntz(7)) - ('2020-01-01 03:04:05.000000009' :: timestamp_ntz(9)); +-- Mixed with a micro TIMESTAMP_NTZ operand. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05'; +-- A DATE operand is cast to the nanos type (midnight); the fraction below the micro grid drops. +SELECT TIMESTAMP_NTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01'; +-- Pre-epoch operand exercises the negative-epoch path. +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789' - TIMESTAMP_NTZ '1960-01-01 00:00:00.000000999'; +-- NULL operand propagates. +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - CAST(NULL AS timestamp_ntz(9)); + +-- SPARK-57818: convert_timezone over nanosecond-precision TIMESTAMP_NTZ. The sub-microsecond +-- remainder is carried through unchanged; only the whole-microsecond part shifts with the zone +-- offset, and the result keeps the source's exact precision. +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + TIMESTAMP_NTZ '2022-03-27 03:00:00.123456789'); +SELECT typeof(convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.1234567' :: timestamp_ntz(7))); +-- NULL nanosecond timestamp. +SELECT convert_timezone('America/Los_Angeles', 'UTC', CAST(NULL AS timestamp_ntz(9))); -- SPARK-57103: MAX / MIN over nanosecond-precision TIMESTAMP_NTZ. The aggregate preserves the -- nanosecond type and orders by the sub-microsecond remainder (two values share the same @@ -164,6 +202,44 @@ SELECT c, count(*) FROM VALUES (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c) GROUP BY c ORDER BY c; +-- GROUP BY a nanosecond key with aggregates and a NULL group: exact-duplicate keys collapse, two +-- keys sharing epochMicros but differing within the microsecond stay in separate groups, and all +-- NULL keys group together (unlike an equi-join). Three groups: .000000001 (count 2, sum 3), +-- .000000999 (count 1, sum 3), NULL (count 2, sum 9). +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3), + (CAST(NULL AS timestamp_ntz(9)), 4), + (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k; + +-- SPARK-56822: mode over nanosecond-precision TIMESTAMP_NTZ. Frequencies are counted on the full +-- nanos value, so the most-frequent value is selected down to the sub-microsecond and the result +-- type stays TIMESTAMP_NTZ(9). .000000001 appears twice, .000000999 once. +SELECT mode(c) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c); + +-- SPARK-56822: collect_set over nanosecond-precision TIMESTAMP_NTZ. It deduplicates on the full +-- sub-microsecond value: the two .000000001 rows collapse to one, the .000000999 row stays, so the +-- sorted set has two distinct elements and the element type stays TIMESTAMP_NTZ(9). collect_set +-- order is non-deterministic, so the output is stabilized with sort_array. +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c); + +-- SPARK-56822: collect_list over nanosecond-precision TIMESTAMP_NTZ. The buffer holds the full +-- nanos value, so the sub-microsecond remainder survives and the result element type stays +-- TIMESTAMP_NTZ(9). collect_list order is non-deterministic, so the output is stabilized with +-- sort_array; duplicates are kept and NULLs are dropped. +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c); -- SPARK-57528: unix_timestamp / to_unix_timestamp over nanosecond-precision values. The result is -- whole-second BIGINT; the sub-second digits are dropped and NTZ applies no zone shift, so the @@ -185,6 +261,15 @@ SELECT max_by(v, k), min_by(v, k) FROM VALUES (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3), (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500', 2), (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000007', CAST(NULL AS INT)) AS t(v, k); +-- DISTINCT over a nanosecond column: exact duplicates are removed, two values sharing epochMicros +-- but differing within the microsecond are both kept, and NULL survives as a single row. Three +-- rows: .000000001, .000000999, NULL. +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c) + ORDER BY c; -- SPARK-57527: unix_nanos over nanosecond-precision values returns DECIMAL(21, 0) nanoseconds since -- the epoch; NTZ applies no zone shift, so the wall-clock value is read as the epoch instant. The @@ -269,6 +354,34 @@ SELECT v, lead(v) OVER (ORDER BY v) AS next_v FROM ( UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000100' UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500') ORDER BY v; +-- SPARK-57811: a string operand is coerced to the nanosecond timestamp type in comparisons and +-- predicates (not truncated to micros, not promoted to string). The 9th fractional digit is +-- significant, so an off-by-one-nanosecond literal does not compare equal. +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c); + +-- BETWEEN over nanosecond timestamps: only the value inside the sub-microsecond range qualifies. +SELECT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005'; + +-- SPARK-57811: TIMESTAMP_NTZ(p) mirrors micros TimestampNTZType under legacy castDatetimeToString. +-- Micros TimestampNTZType has no arm in findCommonTypeForBinaryComparison, so it stays config-blind +-- and casts the string to the timestamp type even under the legacy flag; nanos NTZ has no arm +-- either and does the same. (Only the LTZ family, like micros TimestampType, promotes the range +-- comparison to string under this flag -- see timestamp-ltz-nanos.sql.) So with both flags set the +-- NTZ range comparison still casts the string to the nanos type, identical to the default config. +SET spark.sql.ansi.enabled=false; +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true; +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c); +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false; +SET spark.sql.ansi.enabled=true; + -- SPARK-57814: unix_seconds / unix_millis / unix_micros over nanosecond-precision values. The result -- is a whole BIGINT count of the unit; sub-unit digits (incl. the sub-microsecond remainder) are -- dropped and NTZ applies no zone shift, so the wall-clock value is read as the epoch instant. @@ -305,3 +418,144 @@ SELECT date_trunc('HOUR', '2020-01-01 12:34:56.123456789' :: timestamp_ntz(9)); SELECT date_trunc('DAY', '2020-06-21 23:30:00.000000123' :: timestamp_ntz(7)); -- An unsupported (sub-microsecond) unit yields NULL; the result still carries the nanos type. SELECT date_trunc('NANOSECOND', TIMESTAMP_NTZ '2020-01-01 12:34:56.123456789'); + +-- SPARK-57837: localtimestamp(p) with a nanosecond precision returns TIMESTAMP_NTZ(p). The values +-- are non-deterministic, so only the (deterministic) result type and query-stable self-equality +-- are checked. Precision 6 keeps the standard microsecond TIMESTAMP_NTZ. +SELECT typeof(localtimestamp(9)), typeof(localtimestamp(8)), typeof(localtimestamp(7)); +SELECT typeof(localtimestamp()), typeof(localtimestamp(6)); +-- A foldable (constant) precision expression is accepted. +SELECT typeof(localtimestamp(8 + 1)); +-- All references to localtimestamp(p) within a query see the same value. +SELECT localtimestamp(9) = localtimestamp(9); +-- Out-of-range precision is rejected. +SELECT localtimestamp(3); +SELECT localtimestamp(10); + +-- SPARK-57841: end-to-end coverage for operators that ride on the resolved widening (SPARK-57454) +-- and complex-type access over nanosecond values. Every case turns on the SUB-MICROSECOND remainder +-- (.000000001 vs .000000999 share a microsecond; only the full nanos value tells them apart) or on +-- cross-precision widening. Multi-row queries end in a top-level ORDER BY so the golden output order +-- is meaningful (SQLQueryTestSuite re-sorts otherwise). NTZ is zone-independent; the LTZ file mirrors +-- these in the session zone. + +-- INTERSECT / EXCEPT distinguish the sub-microsecond remainder. A micro-only set op would wrongly +-- merge .000000001 and .000000999; here INTERSECT keeps only the common value and EXCEPT removes it. +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + INTERSECT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c; +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + EXCEPT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c; +-- Mixed-precision set op widens to the wider precision; the equal instant matches after widening. +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009' :: timestamp_ntz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900' :: timestamp_ntz(9))) ORDER BY c; + +-- BETWEEN on a sub-microsecond boundary: the bounds share the microsecond with the probe, so only +-- the full nanos value decides inclusivity. .000000500 is inside [.000000001, .000000999]; +-- .000001000 (next microsecond) is outside. +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'; +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000001000' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'; +-- Mixed-precision BETWEEN widens the bounds to the probe's precision. +SELECT '2020-01-01 00:00:00.000000500' :: timestamp_ntz(9) + BETWEEN '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7) + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000001'; + +-- if / nvl / ifnull preserve the nanos type and widen mixed-precision branches to the wider type. +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7), + TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789') AS v); +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ntz(9)), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS v); +SELECT ifnull(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', CAST(NULL AS timestamp_ntz(9))); + +-- IN (subquery): the semi-join matches on the full nanos key, so only the .000000999 row qualifies. +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k; + +-- Scalar subquery in projection returns the nanos value and carries the nanos type; +-- the sub-microsecond precision survives scalar-subquery result boxing. +SELECT (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'); +SELECT typeof((SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')); +-- A NULL scalar subquery still carries the nanos type. +SELECT typeof((SELECT CAST(NULL AS timestamp_ntz(9)))); +-- Scalar subquery in a WHERE comparison: the sub-microsecond value decides the match, so only +-- the .000000999 row qualifies (not the .000000001 row). +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k = (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k; + +-- EXISTS (correlated on a nanos equality): the outer row is kept iff a matching nanos key exists +-- in the subquery relation. Only the .000000999 row correlates to s.v. +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS s(v) + WHERE s.v = t.k) ORDER BY k; + +-- NOT EXISTS (correlated on a nanos equality): the opposite; the outer row is kept iff no matching +-- nanos key exists. The subquery holds only .000000001, so the .000000999 row survives. +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE NOT EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS s(v) + WHERE s.v = t.k) ORDER BY k; + +-- NOT IN (subquery): anti-semi-join on the full nanos key. The .000000999 row is in the subquery +-- set (excluded); the .000000001 row is not (kept). Sub-microsecond precision decides membership -- +-- the two values differ in the nanosecond digit, not by rounding error. +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k; + +-- Mixed-precision NOT IN widens the probe to p=9 before the anti-join. The p=7 value .0000009 +-- becomes .000000900 at p=9, which is not .000000999, so the row is not in the set and is kept. +SELECT k FROM VALUES + ('2020-01-01 00:00:00.0000009' :: timestamp_ntz(7)) AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k; + +-- NOT IN with a NULL in the subquery set: three-valued logic. For the row that does not equal the +-- non-null member, the comparison against NULL is UNKNOWN, so NOT IN is UNKNOWN and the row is +-- filtered out; the row that equals the non-null member is a definite match and also excluded. +-- The result is therefore empty. +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' + UNION ALL SELECT CAST(NULL AS timestamp_ntz(9))) ORDER BY k; + +-- explode(array<ts_nanos>) yields one row per element, each keeping the nanos type and value. +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'))) ORDER BY col; + +-- element_at over array<ts_nanos> (1-based) returns the addressed element unchanged. +SELECT element_at(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), 2); + +-- struct-field extraction reads the nanos value back out of a struct. +SELECT (named_struct('f', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')).f; + +-- map lookup by string key and by nanosecond key (GetMapValue / element_at over a nanos-keyed map). +-- The nanos-keyed lookup must consult the full sub-microsecond value: looking up .000000999 returns +-- 'b', not 'a', even though both keys share the microsecond. +SELECT map('k', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')['k']; +SELECT map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b')[ + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999']; +SELECT element_at(map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b'), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'); diff --git a/sql/core/src/test/resources/sql-tests/inputs/unnest.sql b/sql/core/src/test/resources/sql-tests/inputs/unnest.sql new file mode 100644 index 0000000000000..3758fda053080 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/unnest.sql @@ -0,0 +1,72 @@ +-- Tests for the ANSI SQL UNNEST collection derived table in the FROM clause. + +CREATE OR REPLACE TEMPORARY VIEW nested AS SELECT * FROM VALUES +(1, array(10, 20, 30), array('a', 'b')), +(2, array(40), array('c', 'd', 'e')), +(3, array(), array()), +(4, cast(null as array<int>), array('f')) +AS nested(id, xs, ys); + +-- Single array: one row per element, default column name `col`. +SELECT * FROM UNNEST(array(10, 20, 30)); + +-- Single array with a table and column alias. +SELECT v FROM UNNEST(array(10, 20, 30)) AS t(v); + +-- Empty array produces no rows. +SELECT * FROM UNNEST(array()); + +-- NULL array is treated as empty and produces no rows. +SELECT * FROM UNNEST(cast(null as array<int>)); + +-- WITH ORDINALITY appends a 1-based BIGINT position column. +SELECT * FROM UNNEST(array(10, 20, 30)) WITH ORDINALITY; + +-- WITH ORDINALITY with column aliases. +SELECT val, pos FROM UNNEST(array('x', 'y')) WITH ORDINALITY AS t(val, pos); + +-- Multiple arrays are expanded in parallel and padded with NULLs to the longest length. +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) AS t(a, b); + +-- Multiple arrays with WITH ORDINALITY. +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) WITH ORDINALITY AS t(a, b, ord); + +-- An array of structs keeps the struct as a single column (unlike inline). +SELECT * FROM UNNEST(array(struct(1, 'a'), struct(2, 'b'))) AS t(s); + +-- Correlated UNNEST over a table column, via LATERAL. +SELECT id, elem FROM nested, LATERAL UNNEST(xs) AS t(elem) ORDER BY id, elem; + +-- Correlated UNNEST of two arrays with ordinality, via LATERAL. +SELECT id, x, y, ord +FROM nested, LATERAL UNNEST(xs, ys) WITH ORDINALITY AS t(x, y, ord) +ORDER BY id, ord; + +-- LEFT JOIN LATERAL preserves outer rows when the array is empty or NULL. +SELECT id, elem +FROM nested LEFT JOIN LATERAL UNNEST(xs) AS t(elem) ON true +ORDER BY id, elem; + +-- Nested arrays: the element type is preserved as-is (a single array-typed column). +SELECT * FROM UNNEST(array(array(1, 2), array(3))) AS t(inner); + +-- Array elements that are themselves NULL are emitted as NULL rows (distinct from a NULL array). +SELECT * FROM UNNEST(array(1, cast(null as int), 3)) WITH ORDINALITY; + +-- The first array being shorter still pads it (not just trailing arrays). +SELECT * FROM UNNEST(array(1), array(10, 20, 30)) AS t(a, b); + +-- Non-array argument is rejected. +SELECT * FROM UNNEST(42); + +-- A MAP argument is rejected (unlike explode, UNNEST is array-only per the SQL standard). +SELECT * FROM UNNEST(map('a', 1)); + +-- A mix of array and non-array arguments is rejected. +SELECT * FROM UNNEST(array(1, 2), 3); + +-- UNNEST is a non-reserved keyword: a table-valued function named `unnest` can still be invoked +-- by quoting the name, which bypasses the dedicated UNNEST relation syntax. Here it resolves as a +-- generic (unregistered) TVF and fails at analysis, proving the name is not swallowed by the +-- grammar. +SELECT * FROM `unnest`(array(1, 2)); diff --git a/sql/core/src/test/resources/sql-tests/inputs/variant/variant-from-arrays-entries.sql b/sql/core/src/test/resources/sql-tests/inputs/variant/variant-from-arrays-entries.sql new file mode 100644 index 0000000000000..d4db88252d1fa --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/inputs/variant/variant-from-arrays-entries.sql @@ -0,0 +1,47 @@ +-- variant_from_arrays + +-- Basic object construction; keys are sorted in the resulting variant object. +select cast(variant_from_arrays(array('z', 'a'), array(1, 2)) as string); +-- Empty input produces an empty object. +select cast(variant_from_arrays(cast(array() as array<string>), cast(array() as array<int>)) as string); +-- Null values are kept as variant null. +select cast(variant_from_arrays(array('a', 'b'), array(1, cast(null as int))) as string); +-- Nested values are converted recursively. +select cast(variant_from_arrays(array('a'), array(array(1, 2, 3))) as string); +-- A null array input produces null. +select cast(variant_from_arrays(cast(null as array<string>), array(1)) as string); +-- A null key is rejected. +select variant_from_arrays(array('a', cast(null as string)), array(1, 2)); +-- Duplicate keys are rejected. +select variant_from_arrays(array('a', 'a'), array(1, 2)); +-- Mismatched array lengths are rejected. +select variant_from_arrays(array('a', 'b'), array(1)); +-- A non-string key type is rejected. +select variant_from_arrays(array(1, 2), array('a', 'b')); +-- A value type that cannot be cast to variant is rejected. +select variant_from_arrays(array('a'), array(map(1, 2))); + +-- variant_from_entries + +-- Basic object construction from key/value struct entries. +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'b', 'v', 2))) as string); +-- Empty input produces an empty object. +select cast(variant_from_entries(cast(array() as array<struct<k:string,v:int>>)) as string); +-- Null values are kept as variant null. +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', cast(null as int)))) as string); +-- A null entry makes the whole result null. +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), cast(null as struct<k:string,v:int>))) as string); +-- A null array input produces null. +select cast(variant_from_entries(cast(null as array<struct<k:string,v:int>>)) as string); +-- A null key is rejected. +select variant_from_entries(array(named_struct('k', cast(null as string), 'v', 1))); +-- Duplicate keys are rejected. +select variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'a', 'v', 2))); +-- A non-array-of-pair-struct input is rejected. +select variant_from_entries(array(1, 2)); +-- A struct with the wrong number of fields is rejected. +select variant_from_entries(array(named_struct('k', 'a'))); +-- A non-string key type is rejected. +select variant_from_entries(array(named_struct('k', 1, 'v', 'a'))); +-- A value type that cannot be cast to variant is rejected. +select variant_from_entries(array(named_struct('k', 'a', 'v', map(1, 2)))); diff --git a/sql/core/src/test/resources/sql-tests/inputs/vector-distance.sql b/sql/core/src/test/resources/sql-tests/inputs/vector-distance.sql index 24035963260b4..ce05649a7d5ea 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/vector-distance.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/vector-distance.sql @@ -130,3 +130,38 @@ SELECT vector_l2_distance( array(1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F, 13.0F, 14.0F, 15.0F, 16.0F), array(16.0F, 15.0F, 14.0F, 13.0F, 12.0F, 11.0F, 10.0F, 9.0F, 8.0F, 7.0F, 6.0F, 5.0F, 4.0F, 3.0F, 2.0F, 1.0F) ); + +-- SPARK-58544: large magnitudes, intermediate sums of squares/products must not overflow the +-- float range + +-- vector_cosine_similarity is scale invariant: identical vectors have similarity 1.0 +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(3.0e19F, 4.0e19F)); + +-- vector_cosine_similarity: opposite vectors have similarity -1.0 +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(-3.0e19F, -4.0e19F)); + +-- vector_inner_product: individual products overflow the float range but cancel out +SELECT vector_inner_product(array(1.0e20F, 1.0e20F), array(1.0e20F, -1.0e20F)); + +-- vector_l2_distance: sqrt((3e19)^2 + (4e19)^2) = 5e19 +SELECT vector_l2_distance(array(3.0e19F, 4.0e19F), array(0.0F, 0.0F)); + +-- SPARK-58544: small magnitudes, intermediate sums of squares/products must not underflow to zero + +-- vector_cosine_similarity of identical tiny vectors is 1.0, not NULL +SELECT vector_cosine_similarity(array(1.0e-23F, 0.0F), array(1.0e-23F, 0.0F)); + +-- vector_cosine_similarity of orthogonal tiny vectors is 0.0, not NULL +SELECT vector_cosine_similarity(array(1.0e-23F, 1.0e-23F), array(1.0e-23F, -1.0e-23F)); + +-- SPARK-58544: elements that are already infinite still propagate to NaN or Infinity; the wider +-- accumulators do not change that + +-- vector_cosine_similarity: the norm is infinite, so the similarity is NaN +SELECT vector_cosine_similarity(array(float('inf'), 1.0F), array(1.0F, 1.0F)); + +-- vector_inner_product: the dot product is infinite +SELECT vector_inner_product(array(float('inf'), 1.0F), array(1.0F, 1.0F)); + +-- vector_l2_distance: the distance is infinite +SELECT vector_l2_distance(array(float('inf'), 1.0F), array(0.0F, 0.0F)); diff --git a/sql/core/src/test/resources/sql-tests/inputs/vector-norm.sql b/sql/core/src/test/resources/sql-tests/inputs/vector-norm.sql index 13eacf854b78a..cab7548a628dd 100644 --- a/sql/core/src/test/resources/sql-tests/inputs/vector-norm.sql +++ b/sql/core/src/test/resources/sql-tests/inputs/vector-norm.sql @@ -153,3 +153,36 @@ SELECT vector_norm( array(1.0F, 2.0F, 3.0F, 4.0F, 5.0F, CAST(NULL AS FLOAT), 7.0F, 8.0F, 9.0F, 10.0F, 11.0F, 12.0F, 13.0F, 14.0F, 15.0F, 16.0F), 2.0F ); + +-- SPARK-58544: large magnitudes, the intermediate sum of squares must not overflow the float range + +-- vector_norm: sqrt((3e19)^2 + (4e19)^2) = 5e19 +SELECT vector_norm(array(3.0e19F, 4.0e19F), 2.0F); + +-- vector_normalize: [3e19, 4e19] normalizes to [0.6, 0.8] +SELECT vector_normalize(array(3.0e19F, 4.0e19F), 2.0F); + +-- vector_norm: the L1 norm itself is not representable as a float, so it stays infinite +SELECT vector_norm(array(3.0e38F, 3.0e38F), 1.0F); + +-- vector_normalize: normalization still succeeds when the norm exceeds the float range +SELECT vector_normalize(array(3.0e38F, 3.0e38F), 1.0F); + +-- SPARK-58544: small magnitudes, the intermediate sum of squares must not underflow to zero + +-- vector_norm: the L2 norm of a tiny vector is not zero +SELECT vector_norm(array(1.0e-23F, 0.0F), 2.0F); + +-- vector_normalize: normalization is scale invariant, so the result is a unit vector, not NULL +SELECT vector_normalize(array(1.0e-23F, 0.0F), 2.0F); + +SELECT vector_normalize(array(1.0e-23F, 1.0e-23F), 2.0F); + +-- SPARK-58544: elements that are already infinite still produce an infinite norm; the wider +-- accumulators do not change that + +-- vector_norm: an infinite element makes the L2 norm infinite +SELECT vector_norm(array(float('inf'), 1.0F), 2.0F); + +-- vector_normalize: dividing by an infinite norm yields NaN for the infinite element +SELECT vector_normalize(array(float('inf'), 1.0F), 2.0F); diff --git a/sql/core/src/test/resources/sql-tests/results/array.sql.out b/sql/core/src/test/resources/sql-tests/results/array.sql.out index 2113ff5ee4913..7de11e3adeef4 100644 --- a/sql/core/src/test/resources/sql-tests/results/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/array.sql.out @@ -1152,3 +1152,103 @@ select array_distinct(array(0.0, -0.0, -0.0, DOUBLE("NaN"), DOUBLE("NaN"))) struct<array_distinct(array(0.0, 0.0, 0.0, NaN, NaN)):array<double>> -- !query output [0.0,NaN] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 0) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 0):array<int>> +-- !query output +[1,2,3,4,5] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 2) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 2):array<int>> +-- !query output +[1,2,3] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 5) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 5):array<int>> +-- !query output +[] + + +-- !query +select trim_array(array('a', 'b', 'c'), 1) +-- !query schema +struct<trim_array(array(a, b, c), 1):array<string>> +-- !query output +["a","b"] + + +-- !query +select trim_array(array(1, 2, null, 4), 1) +-- !query schema +struct<trim_array(array(1, 2, NULL, 4), 1):array<int>> +-- !query output +[1,2,null] + + +-- !query +select trim_array(array(), 0) +-- !query schema +struct<trim_array(array(), 0):array<void>> +-- !query output +[] + + +-- !query +select trim_array(array(1, 2, 3), -1) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + "sqlState" : "22023", + "messageParameters" : { + "functionName" : "`trim_array`", + "length" : "-1", + "numElements" : "3", + "parameter" : "`n`" + } +} + + +-- !query +select trim_array(array(1, 2, 3), 4) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + "sqlState" : "22023", + "messageParameters" : { + "functionName" : "`trim_array`", + "length" : "4", + "numElements" : "3", + "parameter" : "`n`" + } +} + + +-- !query +select trim_array(CAST(null AS ARRAY<INT>), 1) +-- !query schema +struct<trim_array(NULL, 1):array<int>> +-- !query output +NULL + + +-- !query +select trim_array(array(1, 2, 3), CAST(null AS INT)) +-- !query schema +struct<trim_array(array(1, 2, 3), CAST(NULL AS INT)):array<int>> +-- !query output +NULL diff --git a/sql/core/src/test/resources/sql-tests/results/charvarchar-standard-semantics.sql.out b/sql/core/src/test/resources/sql-tests/results/charvarchar-standard-semantics.sql.out new file mode 100644 index 0000000000000..c9592e8361c32 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/charvarchar-standard-semantics.sql.out @@ -0,0 +1,1068 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT typeof(CAST('ab' AS CHAR(5))) +-- !query schema +struct<typeof(CAST(ab AS CHAR(5))):string> +-- !query output +char(5) + + +-- !query +SELECT typeof(CAST('hello' AS VARCHAR(5))) +-- !query schema +struct<typeof(CAST(hello AS VARCHAR(5))):string> +-- !query output +varchar(5) + + +-- !query +SELECT 'X' || CAST('5' AS CHAR(5)) || 'X' +-- !query schema +struct<concat(concat(X, CAST(5 AS CHAR(5))), X):string> +-- !query output +X5 X + + +-- !query +SELECT CAST('ab ' AS CHAR(2)) +-- !query schema +struct<CAST(ab AS CHAR(2)):char(2)> +-- !query output +ab + + +-- !query +SELECT CAST('abcdef' AS CHAR(2)) +-- !query schema +struct<CAST(abcdef AS CHAR(2)):char(2)> +-- !query output +ab + + +-- !query +SELECT CAST('abcdef' AS VARCHAR(2)) +-- !query schema +struct<CAST(abcdef AS VARCHAR(2)):varchar(2)> +-- !query output +ab + + +-- !query +SELECT try_cast('abcdef' AS CHAR(2)) +-- !query schema +struct<TRY_CAST(abcdef AS CHAR(2)):char(2)> +-- !query output +ab + + +-- !query +SELECT try_cast('abcdef' AS VARCHAR(2)) +-- !query schema +struct<TRY_CAST(abcdef AS VARCHAR(2)):varchar(2)> +-- !query output +ab + + +-- !query +SELECT CAST(12345 AS VARCHAR(4)) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "EXCEED_LIMIT_LENGTH", + "sqlState" : "54006", + "messageParameters" : { + "limit" : "4" + } +} + + +-- !query +SELECT CAST(12345 AS VARCHAR(5)) +-- !query schema +struct<CAST(12345 AS VARCHAR(5)):varchar(5)> +-- !query output +12345 + + +-- !query +SELECT try_cast(12345 AS VARCHAR(4)) +-- !query schema +struct<TRY_CAST(12345 AS VARCHAR(4)):varchar(4)> +-- !query output +NULL + + +-- !query +SELECT coalesce(CAST('abcdef' AS VARCHAR(2)), CAST('x' AS VARCHAR(4))) +-- !query schema +struct<coalesce(CAST(abcdef AS VARCHAR(2)), CAST(x AS VARCHAR(4))):varchar(4)> +-- !query output +ab + + +-- !query +SELECT CASE WHEN true THEN CAST('abcdef' AS VARCHAR(2)) ELSE CAST('x' AS VARCHAR(4)) END +-- !query schema +struct<CASE WHEN true THEN CAST(abcdef AS VARCHAR(2)) ELSE CAST(x AS VARCHAR(4)) END:varchar(4)> +-- !query output +ab + + +-- !query +SELECT CAST('abcdef' AS VARCHAR(2)) IN (CAST('ab' AS VARCHAR(4))) +-- !query schema +struct<(CAST(abcdef AS VARCHAR(2)) IN (CAST(ab AS VARCHAR(4)))):boolean> +-- !query output +true + + +-- !query +SELECT coalesce( + CAST('abcdef' AS VARCHAR(2) COLLATE UTF8_LCASE), + CAST('x' AS VARCHAR(4) COLLATE UTF8_LCASE)) +-- !query schema +struct<coalesce(CAST(abcdef AS VARCHAR(2) COLLATE UTF8_LCASE), CAST(x AS VARCHAR(4) COLLATE UTF8_LCASE)):varchar(4) collate UTF8_LCASE> +-- !query output +ab + + +-- !query +SELECT coalesce(try_cast(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))) +-- !query schema +struct<coalesce(TRY_CAST(12345 AS VARCHAR(4)), CAST(x AS VARCHAR(5))):varchar(5)> +-- !query output +x + + +-- !query +SELECT coalesce(CAST(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "EXCEED_LIMIT_LENGTH", + "sqlState" : "54006", + "messageParameters" : { + "limit" : "4" + } +} + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world' AS VARCHAR(10)))) +-- !query schema +struct<typeof(coalesce(CAST(hello AS VARCHAR(5)), CAST(world AS VARCHAR(10)))):string> +-- !query output +varchar(10) + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), cast('world!' AS CHAR(6)))) +-- !query schema +struct<typeof(coalesce(CAST(hello AS VARCHAR(5)), CAST(world! AS CHAR(6)))):string> +-- !query output +varchar(6) + + +-- !query +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), cast('world!' AS CHAR(6)))) +-- !query schema +struct<typeof(coalesce(CAST(hello AS CHAR(5)), CAST(world! AS CHAR(6)))):string> +-- !query output +char(6) + + +-- !query +SELECT typeof(coalesce(cast('hello' AS VARCHAR(5)), 'world')) +-- !query schema +struct<typeof(coalesce(CAST(hello AS VARCHAR(5)), world)):string> +-- !query output +string collate UTF8_BINARY + + +-- !query +SELECT typeof(coalesce(cast('hello' AS CHAR(5)), NULL)) +-- !query schema +struct<typeof(coalesce(CAST(hello AS CHAR(5)), NULL)):string> +-- !query output +char(5) + + +-- !query +SELECT typeof( + CASE WHEN true THEN cast('a' AS CHAR(2)) ELSE cast('bb' AS CHAR(4)) END) +-- !query schema +struct<typeof(CASE WHEN true THEN CAST(a AS CHAR(2)) ELSE CAST(bb AS CHAR(4)) END):string> +-- !query output +char(4) + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS CHAR(2)), cast('bbb' AS VARCHAR(3))) +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(2)), CAST(bbb AS VARCHAR(3)))):boolean> +-- !query output +true + + +-- !query +SELECT typeof(c) FROM (SELECT cast('a' AS CHAR(2)) AS c) t WHERE c IN ('a ', 'b') +-- !query schema +struct<typeof(c):string> +-- !query output +char(2) + + +-- !query +SELECT typeof(upper(cast('ab' AS CHAR(2)))) +-- !query schema +struct<typeof(upper(CAST(ab AS CHAR(2)))):string> +-- !query output +string + + +-- !query +SELECT typeof(lower(cast('AB' AS VARCHAR(2)))) +-- !query schema +struct<typeof(lower(CAST(AB AS VARCHAR(2)))):string> +-- !query output +string + + +-- !query +SELECT typeof(cast('a' AS CHAR(1)) || cast('b' AS VARCHAR(1))) +-- !query schema +struct<typeof(concat(CAST(a AS CHAR(1)), CAST(b AS VARCHAR(1)))):string> +-- !query output +string + + +-- !query +SELECT typeof(substr(cast('hello' AS VARCHAR(5)), 1, 2)) +-- !query schema +struct<typeof(substr(CAST(hello AS VARCHAR(5)), 1, 2)):string> +-- !query output +string + + +-- !query +SELECT typeof(upper(coalesce(cast('a' AS CHAR(2)), cast('b' AS CHAR(4))))) +-- !query schema +struct<typeof(upper(coalesce(CAST(a AS CHAR(2)), CAST(b AS CHAR(4))))):string> +-- !query output +string + + +-- !query +SELECT typeof(concat(cast('a' AS CHAR(2)), cast('b' AS CHAR(3)))) +-- !query schema +struct<typeof(concat(CAST(a AS CHAR(2)), CAST(b AS CHAR(3)))):string> +-- !query output +string + + +-- !query +SELECT typeof(concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(concat(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(b AS CHAR(3) COLLATE UTF8_LCASE))):string> +-- !query output +string collate UTF8_LCASE + + +-- !query +SELECT concat('<', concat( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast('b' AS CHAR(3) COLLATE UTF8_LCASE)), '>') +-- !query schema +struct<concat('<' collate UTF8_LCASE, concat(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(b AS CHAR(3) COLLATE UTF8_LCASE)), '>' collate UTF8_LCASE):string collate UTF8_LCASE> +-- !query output +<a b > + + +-- !query +SELECT typeof(elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(elt(1, CAST(ab AS CHAR(5) COLLATE UTF8_LCASE), CAST(x AS CHAR(1) COLLATE UTF8_LCASE))):string> +-- !query output +string collate UTF8_LCASE + + +-- !query +SELECT concat('<', elt( + 1, + cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + cast('x' AS CHAR(1) COLLATE UTF8_LCASE)), '>') +-- !query schema +struct<concat('<' collate UTF8_LCASE, elt(1, CAST(ab AS CHAR(5) COLLATE UTF8_LCASE), CAST(x AS CHAR(1) COLLATE UTF8_LCASE)), '>' collate UTF8_LCASE):string collate UTF8_LCASE> +-- !query output +<ab > + + +-- !query +SELECT typeof(trim(cast('ab ' AS CHAR(4)))) +-- !query schema +struct<typeof(trim(CAST(ab AS CHAR(4)))):string> +-- !query output +string + + +-- !query +SELECT typeof(lpad(cast('ab' AS CHAR(2)), 5, 'x')) +-- !query schema +struct<typeof(lpad(CAST(ab AS CHAR(2)), 5, x)):string> +-- !query output +string + + +-- !query +SELECT typeof(regexp_replace(cast('ab' AS CHAR(2)), 'a', 'x')) +-- !query schema +struct<typeof(regexp_replace(CAST(ab AS CHAR(2)), a, x, 1)):string> +-- !query output +string + + +-- !query +SELECT typeof(regexp_extract(cast('ab' AS VARCHAR(2)), '(a)', 1)) +-- !query schema +struct<typeof(regexp_extract(CAST(ab AS VARCHAR(2)), (a), 1)):string> +-- !query output +string + + +-- !query +SELECT typeof(regexp_extract_all(cast('aab' AS VARCHAR(3)), '(a)', 1)) +-- !query schema +struct<typeof(regexp_extract_all(CAST(aab AS VARCHAR(3)), (a), 1)):string> +-- !query output +array<string> + + +-- !query +SELECT typeof(split(cast('a,b' AS CHAR(3)), ',')) +-- !query schema +struct<typeof(split(CAST(a,b AS CHAR(3)), ,, -1)):string> +-- !query output +array<string> + + +-- !query +SELECT typeof(mask(cast('ab' AS CHAR(2)))) +-- !query schema +struct<typeof(mask(CAST(ab AS CHAR(2)), X, x, n, NULL)):string> +-- !query output +string + + +-- !query +SELECT typeof(overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1)) +-- !query schema +struct<typeof(overlay(CAST(ab AS CHAR(5)), x, 1, -1)):string> +-- !query output +string + + +-- !query +SELECT concat('<', overlay(cast('ab' AS CHAR(5)) PLACING 'x' FROM 1), '>') +-- !query schema +struct<concat(<, overlay(CAST(ab AS CHAR(5)), x, 1, -1), >):string> +-- !query output +<xb > + + +-- !query +SELECT typeof(elt(1, cast('ab' AS CHAR(5)), 'x')) +-- !query schema +struct<typeof(elt(1, CAST(ab AS CHAR(5)), x)):string> +-- !query output +string + + +-- !query +SELECT typeof(right(cast('ab' AS CHAR(5)), 2)) +-- !query schema +struct<typeof(right(CAST(ab AS CHAR(5)), 2)):string> +-- !query output +string + + +-- !query +SELECT concat('<', right(cast('ab' AS CHAR(5)), 2), '>') +-- !query schema +struct<concat(<, right(CAST(ab AS CHAR(5)), 2), >):string> +-- !query output +< > + + +-- !query +SELECT typeof(left(cast('ab' AS CHAR(5)), 2)) +-- !query schema +struct<typeof(left(CAST(ab AS CHAR(5)), 2)):string> +-- !query output +string + + +-- !query +SELECT typeof(reverse(cast('ab' AS CHAR(5)))) +-- !query schema +struct<typeof(reverse(CAST(ab AS CHAR(5)))):string> +-- !query output +string + + +-- !query +SELECT typeof(hex(cast('ab' AS CHAR(5)))) +-- !query schema +struct<typeof(hex(CAST(ab AS CHAR(5)))):string> +-- !query output +string + + +-- !query +SELECT hex(cast('ab' AS CHAR(5))) +-- !query schema +struct<hex(CAST(ab AS CHAR(5))):string> +-- !query output +6162202020 + + +-- !query +SELECT typeof(array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-')) +-- !query schema +struct<typeof(array_join(array(CAST(ab AS CHAR(5)), CAST(cd AS CHAR(5))), -)):string> +-- !query output +string + + +-- !query +SELECT concat('<', array_join(array(cast('ab' AS CHAR(5)), cast('cd' AS CHAR(5))), '-'), '>') +-- !query schema +struct<concat(<, array_join(array(CAST(ab AS CHAR(5)), CAST(cd AS CHAR(5))), -), >):string> +-- !query output +<ab -cd > + + +-- !query +SELECT typeof(reverse(array(1, 2))) +-- !query schema +struct<typeof(reverse(array(1, 2))):string> +-- !query output +array<int> + + +-- !query +SELECT typeof(str_to_map(cast('a:1,b:2' AS CHAR(7)))) +-- !query schema +struct<typeof(str_to_map(CAST(a:1,b:2 AS CHAR(7)), ,, :)):string> +-- !query output +map<string,string> + + +-- !query +SELECT typeof(c0) FROM (SELECT json_tuple(cast('{"a":"1"}' AS CHAR(9)), 'a') AS c0) +-- !query schema +struct<typeof(c0):string> +-- !query output +string + + +-- !query +SELECT typeof(cast('a' AS CHAR(2) COLLATE UTF8_LCASE)) +-- !query schema +struct<typeof(CAST(a AS CHAR(2) COLLATE UTF8_LCASE)):string> +-- !query output +char(2) collate UTF8_LCASE + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(2) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(2) COLLATE UTF8_LCASE))):string> +-- !query output +char(2) collate UTF8_LCASE + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(4) COLLATE UTF8_LCASE))):string> +-- !query output +char(4) collate UTF8_LCASE + + +-- !query +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query schema +struct<hex(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS CHAR(4) COLLATE UTF8_LCASE))):string collate UTF8_LCASE> +-- !query output +61202020 + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), cast('bb' AS VARCHAR(4) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(bb AS VARCHAR(4) COLLATE UTF8_LCASE))):string> +-- !query output +varchar(4) collate UTF8_LCASE + + +-- !query +SELECT typeof(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query schema +struct<typeof(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(1 AS CHAR(4) COLLATE UTF8_LCASE))):string> +-- !query output +char(4) collate UTF8_LCASE + + +-- !query +SELECT hex(coalesce( + cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + cast(1 AS CHAR(4) COLLATE UTF8_LCASE))) +-- !query schema +struct<hex(coalesce(CAST(a AS CHAR(2) COLLATE UTF8_LCASE), CAST(1 AS CHAR(4) COLLATE UTF8_LCASE))):string collate UTF8_LCASE> +-- !query output +61202020 + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS VARCHAR(3)) AS c + UNION ALL + SELECT cast('abcd' AS VARCHAR(8)) AS c +) t LIMIT 1 +-- !query schema +struct<typeof(c):string> +-- !query output +varchar(8) + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t LIMIT 1 +-- !query schema +struct<typeof(c):string> +-- !query output +char(4) + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION ALL + SELECT cast('bb' AS CHAR(4)) AS c +) t +-- !query schema +struct<concat(<, c, >):string> +-- !query output +<a > +<bb > + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t +-- !query schema +struct<typeof(c):string> +-- !query output +char(4) + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('a' AS CHAR(2)) AS c + UNION + SELECT cast('a' AS CHAR(4)) AS c +) t +-- !query schema +struct<concat(<, c, >):string> +-- !query output +<a > + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t +-- !query schema +struct<typeof(c):string> +-- !query output +char(4) + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + INTERSECT + SELECT cast('ab' AS CHAR(4)) AS c +) t +-- !query schema +struct<concat(<, c, >):string> +-- !query output +<ab > + + +-- !query +SELECT typeof(c) FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t +-- !query schema +struct<typeof(c):string> +-- !query output +char(4) + + +-- !query +SELECT concat('<', c, '>') FROM ( + SELECT cast('ab' AS CHAR(2)) AS c + EXCEPT + SELECT cast('xy' AS CHAR(4)) AS c +) t +-- !query schema +struct<concat(<, c, >):string> +-- !query output +<ab > + + +-- !query +SELECT typeof(c) FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c) +-- !query schema +struct<typeof(c):string> +-- !query output +char(4) +char(4) + + +-- !query +SELECT concat('<', c, '>') FROM (VALUES + (cast('a' AS CHAR(2))), + (cast('bb' AS CHAR(4))) +) t(c) +-- !query schema +struct<concat(<, c, >):string> +-- !query output +<a > +<bb > + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a' AS CHAR(4)) +-- !query schema +struct<(CAST(a AS CHAR(2)) = CAST(a AS CHAR(4))):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a' AS VARCHAR(2)) +-- !query schema +struct<(CAST(a AS CHAR(2)) = CAST(a AS VARCHAR(2))):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2)) = cast('a ' AS VARCHAR(2)) +-- !query schema +struct<(CAST(a AS CHAR(2)) = CAST(a AS VARCHAR(2))):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) = 'a' +-- !query schema +struct<(CAST(a AS CHAR(2)) = a):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2)) = 'a ' +-- !query schema +struct<(CAST(a AS CHAR(2)) = a ):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = 'a' +-- !query schema +struct<(CAST(a AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = a):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = + cast('a' AS CHAR(4) COLLATE UTF8_BINARY_RTRIM) +-- !query schema +struct<(CAST(a AS CHAR(2) COLLATE UTF8_BINARY_RTRIM) = CAST(a AS CHAR(4) COLLATE UTF8_BINARY_RTRIM)):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4))) +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(4)))):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS VARCHAR(2))) +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (CAST(a AS VARCHAR(2)))):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a ' AS VARCHAR(2))) +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (CAST(a AS VARCHAR(2)))):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) IN ('a', 'b') +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (a, b)):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2)) IN ('a ', 'b') +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (a , b)):boolean> +-- !query output +true + + +-- !query +SELECT cast('a' AS CHAR(2)) IN (cast('a' AS CHAR(4)), cast('b' AS VARCHAR(3))) +-- !query schema +struct<(CAST(a AS CHAR(2)) IN (CAST(a AS CHAR(4)), CAST(b AS VARCHAR(3)))):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) = cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE) +-- !query schema +struct<(CAST(a AS CHAR(2) COLLATE UTF8_LCASE) = CAST(a AS VARCHAR(2) COLLATE UTF8_LCASE)):boolean> +-- !query output +false + + +-- !query +SELECT cast('a' AS CHAR(2) COLLATE UTF8_LCASE) IN (cast('a' AS VARCHAR(2) COLLATE UTF8_LCASE)) +-- !query schema +struct<(CAST(a AS CHAR(2) COLLATE UTF8_LCASE) IN (CAST(a AS VARCHAR(2) COLLATE UTF8_LCASE))):boolean> +-- !query output +false + + +-- !query +SELECT typeof(coalesce(cast('123' AS CHAR(3)), 1)), + typeof(coalesce(cast('123' AS VARCHAR(3)), 1)), + typeof(coalesce(cast('123' AS STRING), 1)) +-- !query schema +struct<typeof(coalesce(CAST(123 AS CHAR(3)), 1)):string,typeof(coalesce(CAST(123 AS VARCHAR(3)), 1)):string,typeof(coalesce(CAST(123 AS STRING), 1)):string> +-- !query output +bigint bigint bigint + + +-- !query +SELECT coalesce(cast('123' AS CHAR(3)), 1), + coalesce(cast('123' AS VARCHAR(3)), 1), + coalesce(cast('123' AS STRING), 1) +-- !query schema +struct<coalesce(CAST(123 AS CHAR(3)), 1):bigint,coalesce(CAST(123 AS VARCHAR(3)), 1):bigint,coalesce(CAST(123 AS STRING), 1):bigint> +-- !query output +123 123 123 + + +-- !query +SELECT cast('123' AS CHAR(3)) = 123, + cast('123' AS VARCHAR(3)) = 123, + cast('123' AS STRING) = 123 +-- !query schema +struct<(CAST(123 AS CHAR(3)) = 123):boolean,(CAST(123 AS VARCHAR(3)) = 123):boolean,(CAST(123 AS STRING) = 123):boolean> +-- !query output +true true true + + +-- !query +SELECT typeof(coalesce(cast('1.5' AS CHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS VARCHAR(3)), 1.5)), + typeof(coalesce(cast('1.5' AS STRING), 1.5)) +-- !query schema +struct<typeof(coalesce(CAST(1.5 AS CHAR(3)), 1.5)):string,typeof(coalesce(CAST(1.5 AS VARCHAR(3)), 1.5)):string,typeof(coalesce(CAST(1.5 AS STRING), 1.5)):string> +-- !query output +double double double + + +-- !query +SELECT cast('2020-01-02' AS CHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS VARCHAR(10)) = date'2020-01-02', + cast('2020-01-02' AS STRING) = date'2020-01-02' +-- !query schema +struct<(CAST(2020-01-02 AS CHAR(10)) = DATE '2020-01-02'):boolean,(CAST(2020-01-02 AS VARCHAR(10)) = DATE '2020-01-02'):boolean,(CAST(2020-01-02 AS STRING) = DATE '2020-01-02'):boolean> +-- !query output +true true true + + +-- !query +SELECT cast('true' AS CHAR(4)) = true, + cast('true' AS VARCHAR(4)) = true, + cast('true' AS STRING) = true +-- !query schema +struct<(CAST(true AS CHAR(4)) = true):boolean,(CAST(true AS VARCHAR(4)) = true):boolean,(CAST(true AS STRING) = true):boolean> +-- !query output +true true true + + +-- !query +SELECT typeof(coalesce(cast('true' AS CHAR(4)), true)) +-- !query schema +struct<typeof(coalesce(CAST(true AS CHAR(4)), true)):string> +-- !query output +boolean + + +-- !query +SELECT typeof(coalesce(cast('true' AS VARCHAR(4)), true)) +-- !query schema +struct<typeof(coalesce(CAST(true AS VARCHAR(4)), true)):string> +-- !query output +boolean + + +-- !query +SELECT typeof(coalesce(cast('true' AS STRING), true)) +-- !query schema +struct<typeof(coalesce(CAST(true AS STRING), true)):string> +-- !query output +boolean + + +-- !query +SELECT typeof(array(cast('a' AS CHAR(2)), cast('bb' AS CHAR(3)))) +-- !query schema +struct<typeof(array(CAST(a AS CHAR(2)), CAST(bb AS CHAR(3)))):string> +-- !query output +array<char(3)> + + +-- !query +SELECT typeof(struct(cast('a' AS CHAR(2)) AS f)) +-- !query schema +struct<typeof(struct(CAST(a AS CHAR(2)) AS f)):string> +-- !query output +struct<f:char(2)> + + +-- !query +SELECT typeof(map('k', cast('a' AS VARCHAR(2)))) +-- !query schema +struct<typeof(map(k, CAST(a AS VARCHAR(2)))):string> +-- !query output +map<string,varchar(2)> + + +-- !query +CREATE TABLE char_varchar_std (c CHAR(5), v VARCHAR(5)) USING parquet +-- !query schema +struct<> +-- !query output + + + +-- !query +INSERT INTO char_varchar_std VALUES ('ab', 'ab') +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std +-- !query schema +struct<typeof(c):string,typeof(v):string> +-- !query output +char(5) varchar(5) + + +-- !query +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std +-- !query schema +struct<concat([, c, ]):string,concat([, v, ]):string> +-- !query output +[ab ] [ab] + + +-- !query +SELECT length(c), length(v) FROM char_varchar_std +-- !query schema +struct<length(c):int,length(v):int> +-- !query output +5 2 + + +-- !query +CREATE TABLE char_varchar_std_ctas USING parquet AS SELECT c, v FROM char_varchar_std +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std_ctas +-- !query schema +struct<typeof(c):string,typeof(v):string> +-- !query output +char(5) varchar(5) + + +-- !query +CREATE VIEW char_varchar_std_view AS SELECT c FROM char_varchar_std +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT typeof(c) FROM char_varchar_std_view +-- !query schema +struct<typeof(c):string> +-- !query output +char(5) + + +-- !query +CREATE VIEW char_varchar_std_view_v AS SELECT v FROM char_varchar_std +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT typeof(v) FROM char_varchar_std_view_v +-- !query schema +struct<typeof(v):string> +-- !query output +varchar(5) + + +-- !query +WITH t AS (SELECT c, v FROM char_varchar_std) SELECT typeof(c), typeof(v) FROM t +-- !query schema +struct<typeof(c):string,typeof(v):string> +-- !query output +char(5) varchar(5) + + +-- !query +CREATE TABLE char_varchar_std_orc (c CHAR(5), v VARCHAR(5)) USING orc +-- !query schema +struct<> +-- !query output + + + +-- !query +INSERT INTO char_varchar_std_orc VALUES ('ab', 'cd') +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT typeof(c), typeof(v) FROM char_varchar_std_orc +-- !query schema +struct<typeof(c):string,typeof(v):string> +-- !query output +char(5) varchar(5) + + +-- !query +SELECT concat('[', c, ']'), concat('[', v, ']') FROM char_varchar_std_orc +-- !query schema +struct<concat([, c, ]):string,concat([, v, ]):string> +-- !query output +[ab ] [cd] + + +-- !query +DROP VIEW char_varchar_std_view_v +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP VIEW char_varchar_std_view +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP TABLE char_varchar_std_ctas +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP TABLE char_varchar_std_orc +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP TABLE char_varchar_std +-- !query schema +struct<> +-- !query output + diff --git a/sql/core/src/test/resources/sql-tests/results/distinct-map-aggregates.sql.out b/sql/core/src/test/resources/sql-tests/results/distinct-map-aggregates.sql.out new file mode 100644 index 0000000000000..57f53db7e3f54 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/distinct-map-aggregates.sql.out @@ -0,0 +1,265 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +CREATE OR REPLACE TEMPORARY VIEW distinct_map_data AS SELECT * FROM VALUES + (2, map('a', 1, 'b', 2), 1, true), + (2, map('b', 2, 'a', 1), 1, true), + (1, map('a', 1, 'b', 2), 1, true), + (1, map('a', 3), 2, false) +AS distinct_map_data(g, m, id, should_keep) +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT COUNT(DISTINCT m) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT m):bigint> +-- !query output +2 + + +-- !query +SELECT SIZE(COLLECT_LIST(DISTINCT m)) FROM distinct_map_data +-- !query schema +struct<size(collect_list(DISTINCT m)):int> +-- !query output +2 + + +-- !query +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a') +-- !query schema +struct<map_entries(m):array<struct<key:string,value:int>>> +-- !query output +[{"key":"a","value":1},{"key":"b","value":2}] +[{"key":"a","value":3}] + + +-- !query +SELECT map_entries(FIRST(DISTINCT m)), map_entries(LAST(DISTINCT m)), COUNT(DISTINCT m) +FROM VALUES (map('b', 2, 'a', 1)) AS single_map_data(m) +-- !query schema +struct<map_entries(first(DISTINCT m)):array<struct<key:string,value:int>>,map_entries(last(DISTINCT m)):array<struct<key:string,value:int>>,count(DISTINCT m):bigint> +-- !query output +[{"key":"a","value":1},{"key":"b","value":2}] [{"key":"a","value":1},{"key":"b","value":2}] 1 + + +-- !query +SELECT COUNT(DISTINCT m, id) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT m, id):bigint> +-- !query output +2 + + +-- !query +SELECT COUNT(DISTINCT m), COUNT(DISTINCT id) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT m):bigint,count(DISTINCT id):bigint> +-- !query output +2 2 + + +-- !query +SELECT g, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY g +ORDER BY g +-- !query schema +struct<g:int,count(DISTINCT m):bigint> +-- !query output +1 2 +2 1 + + +-- !query +SELECT m, COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a') +-- !query schema +struct<m:map<string,int>,count(DISTINCT m):bigint,collect_list(DISTINCT m):array<map<string,int>>> +-- !query output +{"a":1,"b":2} 1 [{"a":1,"b":2}] +{"a":3} 1 [{"a":3}] + + +-- !query +SELECT COUNT(DISTINCT m) FILTER (WHERE should_keep) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT m) FILTER (WHERE should_keep):bigint> +-- !query output +1 + + +-- !query +SELECT MAX(map_values(m)[0]) +FROM distinct_map_data +WHERE id = 1 +-- !query schema +struct<max(map_values(m)[0]):int> +-- !query output +2 + + +-- !query +SELECT MAX(map_values(m)[0]), COUNT(DISTINCT m) +FROM distinct_map_data +WHERE id = 1 +-- !query schema +struct<max(map_values(m)[0]):int,count(DISTINCT m):bigint> +-- !query output +2 1 + + +-- !query +SELECT g +FROM distinct_map_data +GROUP BY g +ORDER BY COUNT(DISTINCT m), g +-- !query schema +struct<g:int> +-- !query output +2 +1 + + +-- !query +SELECT g +FROM distinct_map_data +GROUP BY g +HAVING COUNT(DISTINCT m) = 1 +ORDER BY g +-- !query schema +struct<g:int> +-- !query output +2 + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT named_struct(m, m)):bigint> +-- !query output +2 + + +-- !query +SELECT COUNT(DISTINCT array(m)) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT array(m)):bigint> +-- !query output +2 + + +-- !query +SELECT COUNT(DISTINCT map('m', m)) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT map(m, m)):bigint> +-- !query output +2 + + +-- !query +SELECT COUNT(DISTINCT m), COLLECT_LIST(DISTINCT m) +FROM VALUES + (CAST(map() AS MAP<STRING, INT>)), + (CAST(map() AS MAP<STRING, INT>)), + (CAST(NULL AS MAP<STRING, INT>)) +AS null_and_empty_map_data(m) +-- !query schema +struct<count(DISTINCT m):bigint,collect_list(DISTINCT m):array<map<string,int>>> +-- !query output +1 [{}] + + +-- !query +SELECT g, GROUPING(g), COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY GROUPING SETS ((g), ()) +ORDER BY GROUPING(g), g +-- !query schema +struct<g:int,grouping(g):tinyint,count(DISTINCT m):bigint> +-- !query output +1 0 2 +2 0 1 +NULL 1 2 + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m, 'n', n)) +FROM VALUES + (map('a', 1, 'b', 2), map('x', 1, 'y', 2)), + (map('b', 2, 'a', 1), map('y', 2, 'x', 1)) +AS grouped_distinct_map_data(m, n) +GROUP BY m +-- !query schema +struct<count(DISTINCT named_struct(m, m, n, n)):bigint> +-- !query output +1 + + +-- !query +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=false +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled false + + +-- !query +SELECT COUNT(DISTINCT m) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT m):bigint> +-- !query output +3 + + +-- !query +SELECT map_entries(m) +FROM ( + SELECT EXPLODE(COLLECT_LIST(DISTINCT m)) AS m + FROM distinct_map_data +) AS collected_maps +ORDER BY element_at(m, 'a'), map_entries(m)[0].key +-- !query schema +struct<map_entries(m):array<struct<key:string,value:int>>> +-- !query output +[{"key":"a","value":1},{"key":"b","value":2}] +[{"key":"b","value":2},{"key":"a","value":1}] +[{"key":"a","value":3}] + + +-- !query +SELECT COUNT(DISTINCT named_struct('m', m)) FROM distinct_map_data +-- !query schema +struct<count(DISTINCT named_struct(m, m)):bigint> +-- !query output +3 + + +-- !query +SELECT m, COUNT(DISTINCT m) +FROM distinct_map_data +GROUP BY m +ORDER BY element_at(m, 'a') +-- !query schema +struct<m:map<string,int>,count(DISTINCT m):bigint> +-- !query output +{"a":1,"b":2} 1 +{"a":3} 1 + + +-- !query +SET spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled=true +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.optimizer.insertMapSortInDistinctAggregates.enabled true diff --git a/sql/core/src/test/resources/sql-tests/results/generators-resolution-edge-cases.sql.out b/sql/core/src/test/resources/sql-tests/results/generators-resolution-edge-cases.sql.out index 7dc50cfb4aac4..7876b772b720d 100644 --- a/sql/core/src/test/resources/sql-tests/results/generators-resolution-edge-cases.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/generators-resolution-edge-cases.sql.out @@ -781,3 +781,49 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "fragment" : "val" } ] } + + +-- !query +SELECT explode(array(a)) AS col, count(*) OVER () AS cnt, a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL +-- !query schema +struct<col:int,cnt:bigint,a:int> +-- !query output +1 3 1 +2 3 2 +3 3 3 + + +-- !query +SELECT explode(array(a)) AS col, + count(*) OVER () AS group_count, + count(*) AS row_count +FROM VALUES (1), (1), (2), (3), (3) AS t(a) +GROUP BY a +HAVING row_count > 1 +-- !query schema +struct<col:int,group_count:bigint,row_count:bigint> +-- !query output +1 2 2 +3 2 2 + + +-- !query +SELECT explode(array(a)) AS col, + count(*) OVER () AS cnt, + count(*) OVER ( + ORDER BY a + ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + ) AS ordered_cnt, + a +FROM VALUES (1), (2), (3), (NULL) AS t(a) +GROUP BY a +HAVING a IS NOT NULL +-- !query schema +struct<col:int,cnt:bigint,ordered_cnt:bigint,a:int> +-- !query output +1 3 3 1 +2 3 3 2 +3 3 3 3 diff --git a/sql/core/src/test/resources/sql-tests/results/identifier-clause-legacy.sql.out b/sql/core/src/test/resources/sql-tests/results/identifier-clause-legacy.sql.out index 3b3bc11999c7f..f881ebd53c72a 100644 --- a/sql/core/src/test/resources/sql-tests/results/identifier-clause-legacy.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/identifier-clause-legacy.sql.out @@ -63,16 +63,445 @@ struct<c1:int> 1 +-- !query +SELECT IDENTIFIER(concat('a', 'b')) FROM VALUES(1) AS T(ab) +-- !query schema +struct<ab:int> +-- !query output +1 + + -- !query CREATE SCHEMA IF NOT EXISTS s -- !query schema struct<> -- !query output - + + + +-- !query +CREATE TABLE s.tab(c1 INT) USING CSV +-- !query schema +struct<> +-- !query output + + + +-- !query +USE SCHEMA s +-- !query schema +struct<> +-- !query output + + + +-- !query +INSERT INTO IDENTIFIER('ta' || 'b') VALUES(1) +-- !query schema +struct<> +-- !query output + + + +-- !query +DELETE FROM IDENTIFIER('ta' || 'b') WHERE 1=0 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "DELETE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +UPDATE IDENTIFIER('ta' || 'b') SET c1 = 2 +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkUnsupportedOperationException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "UPDATE TABLE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +MERGE INTO IDENTIFIER('ta' || 'b') AS t USING IDENTIFIER('ta' || 'b') AS s ON s.c1 = t.c1 + WHEN MATCHED THEN UPDATE SET c1 = 3 +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkUnsupportedOperationException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "MERGE INTO TABLE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +SELECT * FROM IDENTIFIER('tab') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('s.tab') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('`s`.`tab`') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('t' || 'a' || 'b') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER(concat('t', 'ab')) +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +USE SCHEMA default +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP TABLE s.tab +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP SCHEMA s +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT IDENTIFIER('COAL' || 'ESCE')(NULL, 1) +-- !query schema +struct<coalesce(NULL, 1):int> +-- !query output +1 + + +-- !query +SELECT IDENTIFIER(concat('COAL', 'ESCE'))(NULL, 1) +-- !query schema +struct<coalesce(NULL, 1):int> +-- !query output +1 + + +-- !query +SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) +-- !query schema +struct<abs(c1):int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('ra' || 'nge')(0, 1) +-- !query schema +struct<id:bigint> +-- !query output +0 + + +-- !query +VALUES(IDENTIFIER(abs(1))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.WRONG_TYPE", + "sqlState" : "42601", + "messageParameters" : { + "dataType" : "int", + "expr" : "abs(1)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "abs(1)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(nullif('a', 'a')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "nullif('a', 'a')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 41, + "fragment" : "nullif('a', 'a')" + } ] +} + + +-- !query +SELECT IDENTIFIER(max('c1')) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('c1')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 27, + "fragment" : "max('c1')" + } ] +} + + +-- !query +SELECT IDENTIFIER(array_join(transform(array('c', '1'), element -> element), '')) +FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('c', '1'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 80, + "fragment" : "array_join(transform(array('c', '1'), element -> element), '')" + } ] +} + + +-- !query +SELECT IDENTIFIER(rand()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "rand()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "rand()" + } ] +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER ()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WINDOW_FUNCTION_FRAME_NOT_ORDERED", + "sqlState" : "42601", + "messageParameters" : { + "wf_expr" : "row_number()", + "wf_name" : "row_number" + } +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER (ORDER BY 'x')) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 50, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT IDENTIFIER(explode(array('c1'))) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "explode(array('c1'))", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 38, + "fragment" : "explode(array('c1'))" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(max('identifier_function_table')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('identifier_function_table')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "max('identifier_function_table')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(row_number() OVER (ORDER BY 'x')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + array_join(transform(array('identifier', '_function_table'), element -> element), '') +) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('identifier', '_function_table'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 113, + "fragment" : "array_join(transform(array('identifier', '_function_table'), element -> element), '')" + } ] +} -- !query -CREATE TABLE s.tab(c1 INT) USING CSV +CREATE TEMPORARY FUNCTION identifier_name() +RETURNS STRING +RETURN 'c1' -- !query schema struct<> -- !query output @@ -80,7 +509,30 @@ struct<> -- !query -USE SCHEMA s +SELECT IDENTIFIER(identifier_name()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 35, + "fragment" : "identifier_name()" + } ] +} + + +-- !query +DROP TEMPORARY FUNCTION identifier_name -- !query schema struct<> -- !query output @@ -88,7 +540,9 @@ struct<> -- !query -INSERT INTO IDENTIFIER('ta' || 'b') VALUES(1) +CREATE FUNCTION persistent_identifier_name() +RETURNS STRING +RETURN 'c1' -- !query schema struct<> -- !query output @@ -96,88 +550,166 @@ struct<> -- !query -DELETE FROM IDENTIFIER('ta' || 'b') WHERE 1=0 +SELECT IDENTIFIER(persistent_identifier_name()) FROM VALUES(1) AS T(c1) -- !query schema struct<> -- !query output org.apache.spark.sql.AnalysisException { - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", "messageParameters" : { - "operation" : "DELETE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } + "expr" : "spark_catalog.default.persistent_identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 46, + "fragment" : "persistent_identifier_name()" + } ] } -- !query -UPDATE IDENTIFIER('ta' || 'b') SET c1 = 2 +DROP FUNCTION persistent_identifier_name -- !query schema struct<> -- !query output -org.apache.spark.SparkUnsupportedOperationException -{ - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", - "messageParameters" : { - "operation" : "UPDATE TABLE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } -} + -- !query -MERGE INTO IDENTIFIER('ta' || 'b') AS t USING IDENTIFIER('ta' || 'b') AS s ON s.c1 = t.c1 - WHEN MATCHED THEN UPDATE SET c1 = 3 +CREATE FUNCTION persistent_identifier_function(value INT) +RETURNS INT +RETURN value + 1 -- !query schema struct<> -- !query output -org.apache.spark.SparkUnsupportedOperationException + + + +-- !query +SELECT IDENTIFIER(concat('persistent_identifier_', 'function'))(1) +-- !query schema +struct<spark_catalog.default.persistent_identifier_function(1):int> +-- !query output +2 + + +-- !query +DROP FUNCTION persistent_identifier_function +-- !query schema +struct<> +-- !query output + + + +-- !query +CREATE FUNCTION persistent_identifier_base_name() +RETURNS STRING +RETURN 'c1' +-- !query schema +struct<> +-- !query output + + + +-- !query +CREATE FUNCTION persistent_identifier_nested_name() +RETURNS STRING +RETURN persistent_identifier_base_name() +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT IDENTIFIER(persistent_identifier_nested_name()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException { - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", "messageParameters" : { - "operation" : "MERGE INTO TABLE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } + "expr" : "spark_catalog.default.persistent_identifier_nested_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 53, + "fragment" : "persistent_identifier_nested_name()" + } ] } -- !query -SELECT * FROM IDENTIFIER('tab') +DROP FUNCTION persistent_identifier_nested_name -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('s.tab') +DROP FUNCTION persistent_identifier_base_name -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('`s`.`tab`') +CREATE TEMPORARY FUNCTION identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('t' || 'a' || 'b') +CREATE TABLE identifier_function_table(c1 INT) USING csv -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -USE SCHEMA default +SELECT * FROM IDENTIFIER(identifier_relation_name()) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 51, + "fragment" : "identifier_relation_name()" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_count() +RETURNS BIGINT +RETURN SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')) -- !query schema struct<> -- !query output @@ -185,7 +717,15 @@ struct<> -- !query -DROP TABLE s.tab +SELECT identifier_relation_count() +-- !query schema +struct<identifier_relation_count():bigint> +-- !query output +0 + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_count -- !query schema struct<> -- !query output @@ -193,7 +733,9 @@ struct<> -- !query -DROP SCHEMA s +CREATE FUNCTION persistent_identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' -- !query schema struct<> -- !query output @@ -201,27 +743,50 @@ struct<> -- !query -SELECT IDENTIFIER('COAL' || 'ESCE')(NULL, 1) +SELECT * FROM IDENTIFIER(persistent_identifier_relation_name()) -- !query schema -struct<coalesce(NULL, 1):int> +struct<> -- !query output -1 +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 62, + "fragment" : "persistent_identifier_relation_name()" + } ] +} -- !query -SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) +DROP FUNCTION persistent_identifier_relation_name -- !query schema -struct<abs(c1):int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('ra' || 'nge')(0, 1) +DROP TABLE identifier_function_table -- !query schema -struct<id:bigint> +struct<> +-- !query output + + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_name +-- !query schema +struct<> -- !query output -0 + -- !query @@ -854,6 +1419,14 @@ struct<> +-- !query +CREATE TABLE identifier_name_source(name STRING) USING csv +-- !query schema +struct<> +-- !query output + + + -- !query SELECT * FROM IDENTIFIER((SELECT 't')) -- !query schema @@ -877,6 +1450,54 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +SELECT * FROM IDENTIFIER((SELECT max(name) FROM identifier_name_source)) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 71, + "fragment" : "(SELECT max(name) FROM identifier_name_source)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + (SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1) +) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 84, + "fragment" : "(SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1)" + } ] +} + + -- !query SELECT * FROM (SELECT IDENTIFIER((SELECT 'col1')) FROM IDENTIFIER((SELECT 't'))) -- !query schema @@ -969,6 +1590,14 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +DROP TABLE identifier_name_source +-- !query schema +struct<> +-- !query output + + + -- !query DROP TABLE t -- !query schema diff --git a/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out b/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out index a45c2da46ada6..045348cf7ea0a 100644 --- a/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/identifier-clause.sql.out @@ -63,16 +63,445 @@ struct<c1:int> 1 +-- !query +SELECT IDENTIFIER(concat('a', 'b')) FROM VALUES(1) AS T(ab) +-- !query schema +struct<ab:int> +-- !query output +1 + + -- !query CREATE SCHEMA IF NOT EXISTS s -- !query schema struct<> -- !query output - + + + +-- !query +CREATE TABLE s.tab(c1 INT) USING CSV +-- !query schema +struct<> +-- !query output + + + +-- !query +USE SCHEMA s +-- !query schema +struct<> +-- !query output + + + +-- !query +INSERT INTO IDENTIFIER('ta' || 'b') VALUES(1) +-- !query schema +struct<> +-- !query output + + + +-- !query +DELETE FROM IDENTIFIER('ta' || 'b') WHERE 1=0 +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "DELETE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +UPDATE IDENTIFIER('ta' || 'b') SET c1 = 2 +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkUnsupportedOperationException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "UPDATE TABLE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +MERGE INTO IDENTIFIER('ta' || 'b') AS t USING IDENTIFIER('ta' || 'b') AS s ON s.c1 = t.c1 + WHEN MATCHED THEN UPDATE SET c1 = 3 +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkUnsupportedOperationException +{ + "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", + "sqlState" : "0A000", + "messageParameters" : { + "operation" : "MERGE INTO TABLE", + "tableName" : "`spark_catalog`.`s`.`tab`" + } +} + + +-- !query +SELECT * FROM IDENTIFIER('tab') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('s.tab') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('`s`.`tab`') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('t' || 'a' || 'b') +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER(concat('t', 'ab')) +-- !query schema +struct<c1:int> +-- !query output +1 + + +-- !query +USE SCHEMA default +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP TABLE s.tab +-- !query schema +struct<> +-- !query output + + + +-- !query +DROP SCHEMA s +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT IDENTIFIER('COAL' || 'ESCE')(NULL, 1) +-- !query schema +struct<coalesce(NULL, 1):int> +-- !query output +1 + + +-- !query +SELECT IDENTIFIER(concat('COAL', 'ESCE'))(NULL, 1) +-- !query schema +struct<coalesce(NULL, 1):int> +-- !query output +1 + + +-- !query +SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) +-- !query schema +struct<abs(c1):int> +-- !query output +1 + + +-- !query +SELECT * FROM IDENTIFIER('ra' || 'nge')(0, 1) +-- !query schema +struct<id:bigint> +-- !query output +0 + + +-- !query +VALUES(IDENTIFIER(abs(1))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.WRONG_TYPE", + "sqlState" : "42601", + "messageParameters" : { + "dataType" : "int", + "expr" : "abs(1)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "abs(1)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(nullif('a', 'a')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "nullif('a', 'a')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 41, + "fragment" : "nullif('a', 'a')" + } ] +} + + +-- !query +SELECT IDENTIFIER(max('c1')) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('c1')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 27, + "fragment" : "max('c1')" + } ] +} + + +-- !query +SELECT IDENTIFIER(array_join(transform(array('c', '1'), element -> element), '')) +FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('c', '1'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 80, + "fragment" : "array_join(transform(array('c', '1'), element -> element), '')" + } ] +} + + +-- !query +SELECT IDENTIFIER(rand()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "rand()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 24, + "fragment" : "rand()" + } ] +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER ()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WINDOW_FUNCTION_FRAME_NOT_ORDERED", + "sqlState" : "42601", + "messageParameters" : { + "wf_expr" : "row_number()", + "wf_name" : "row_number" + } +} + + +-- !query +SELECT IDENTIFIER(row_number() OVER (ORDER BY 'x')) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 50, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT IDENTIFIER(explode(array('c1'))) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "explode(array('c1'))", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 38, + "fragment" : "explode(array('c1'))" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(max('identifier_function_table')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "max('identifier_function_table')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "max('identifier_function_table')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER(row_number() OVER (ORDER BY 'x')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "row_number() OVER (ORDER BY 'x' ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 57, + "fragment" : "row_number() OVER (ORDER BY 'x')" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + array_join(transform(array('identifier', '_function_table'), element -> element), '') +) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "array_join(transform(array('identifier', '_function_table'), lambdafunction(namedlambdavariable(), namedlambdavariable())), '')", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 113, + "fragment" : "array_join(transform(array('identifier', '_function_table'), element -> element), '')" + } ] +} -- !query -CREATE TABLE s.tab(c1 INT) USING CSV +CREATE TEMPORARY FUNCTION identifier_name() +RETURNS STRING +RETURN 'c1' -- !query schema struct<> -- !query output @@ -80,7 +509,30 @@ struct<> -- !query -USE SCHEMA s +SELECT IDENTIFIER(identifier_name()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 35, + "fragment" : "identifier_name()" + } ] +} + + +-- !query +DROP TEMPORARY FUNCTION identifier_name -- !query schema struct<> -- !query output @@ -88,7 +540,9 @@ struct<> -- !query -INSERT INTO IDENTIFIER('ta' || 'b') VALUES(1) +CREATE FUNCTION persistent_identifier_name() +RETURNS STRING +RETURN 'c1' -- !query schema struct<> -- !query output @@ -96,88 +550,166 @@ struct<> -- !query -DELETE FROM IDENTIFIER('ta' || 'b') WHERE 1=0 +SELECT IDENTIFIER(persistent_identifier_name()) FROM VALUES(1) AS T(c1) -- !query schema struct<> -- !query output org.apache.spark.sql.AnalysisException { - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", "messageParameters" : { - "operation" : "DELETE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } + "expr" : "spark_catalog.default.persistent_identifier_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 46, + "fragment" : "persistent_identifier_name()" + } ] } -- !query -UPDATE IDENTIFIER('ta' || 'b') SET c1 = 2 +DROP FUNCTION persistent_identifier_name -- !query schema struct<> -- !query output -org.apache.spark.SparkUnsupportedOperationException -{ - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", - "messageParameters" : { - "operation" : "UPDATE TABLE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } -} + -- !query -MERGE INTO IDENTIFIER('ta' || 'b') AS t USING IDENTIFIER('ta' || 'b') AS s ON s.c1 = t.c1 - WHEN MATCHED THEN UPDATE SET c1 = 3 +CREATE FUNCTION persistent_identifier_function(value INT) +RETURNS INT +RETURN value + 1 -- !query schema struct<> -- !query output -org.apache.spark.SparkUnsupportedOperationException + + + +-- !query +SELECT IDENTIFIER(concat('persistent_identifier_', 'function'))(1) +-- !query schema +struct<spark_catalog.default.persistent_identifier_function(1):int> +-- !query output +2 + + +-- !query +DROP FUNCTION persistent_identifier_function +-- !query schema +struct<> +-- !query output + + + +-- !query +CREATE FUNCTION persistent_identifier_base_name() +RETURNS STRING +RETURN 'c1' +-- !query schema +struct<> +-- !query output + + + +-- !query +CREATE FUNCTION persistent_identifier_nested_name() +RETURNS STRING +RETURN persistent_identifier_base_name() +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT IDENTIFIER(persistent_identifier_nested_name()) FROM VALUES(1) AS T(c1) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException { - "errorClass" : "UNSUPPORTED_FEATURE.TABLE_OPERATION", - "sqlState" : "0A000", + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", "messageParameters" : { - "operation" : "MERGE INTO TABLE", - "tableName" : "`spark_catalog`.`s`.`tab`" - } + "expr" : "spark_catalog.default.persistent_identifier_nested_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 19, + "stopIndex" : 53, + "fragment" : "persistent_identifier_nested_name()" + } ] } -- !query -SELECT * FROM IDENTIFIER('tab') +DROP FUNCTION persistent_identifier_nested_name -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('s.tab') +DROP FUNCTION persistent_identifier_base_name -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('`s`.`tab`') +CREATE TEMPORARY FUNCTION identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('t' || 'a' || 'b') +CREATE TABLE identifier_function_table(c1 INT) USING csv -- !query schema -struct<c1:int> +struct<> -- !query output -1 + -- !query -USE SCHEMA default +SELECT * FROM IDENTIFIER(identifier_relation_name()) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 51, + "fragment" : "identifier_relation_name()" + } ] +} + + +-- !query +CREATE TEMPORARY FUNCTION identifier_relation_count() +RETURNS BIGINT +RETURN SELECT count(*) FROM IDENTIFIER(concat('identifier_function_', 'table')) -- !query schema struct<> -- !query output @@ -185,7 +717,15 @@ struct<> -- !query -DROP TABLE s.tab +SELECT identifier_relation_count() +-- !query schema +struct<identifier_relation_count():bigint> +-- !query output +0 + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_count -- !query schema struct<> -- !query output @@ -193,7 +733,9 @@ struct<> -- !query -DROP SCHEMA s +CREATE FUNCTION persistent_identifier_relation_name() +RETURNS STRING +RETURN 'identifier_function_table' -- !query schema struct<> -- !query output @@ -201,27 +743,50 @@ struct<> -- !query -SELECT IDENTIFIER('COAL' || 'ESCE')(NULL, 1) +SELECT * FROM IDENTIFIER(persistent_identifier_relation_name()) -- !query schema -struct<coalesce(NULL, 1):int> +struct<> -- !query output -1 +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "spark_catalog.default.persistent_identifier_relation_name()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 62, + "fragment" : "persistent_identifier_relation_name()" + } ] +} -- !query -SELECT IDENTIFIER('abs')(c1) FROM VALUES(-1) AS T(c1) +DROP FUNCTION persistent_identifier_relation_name -- !query schema -struct<abs(c1):int> +struct<> -- !query output -1 + -- !query -SELECT * FROM IDENTIFIER('ra' || 'nge')(0, 1) +DROP TABLE identifier_function_table -- !query schema -struct<id:bigint> +struct<> +-- !query output + + + +-- !query +DROP TEMPORARY FUNCTION identifier_relation_name +-- !query schema +struct<> -- !query output -0 + -- !query @@ -854,6 +1419,14 @@ struct<> +-- !query +CREATE TABLE identifier_name_source(name STRING) USING csv +-- !query schema +struct<> +-- !query output + + + -- !query SELECT * FROM IDENTIFIER((SELECT 't')) -- !query schema @@ -877,6 +1450,54 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException } +-- !query +SELECT * FROM IDENTIFIER((SELECT max(name) FROM identifier_name_source)) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 26, + "stopIndex" : 71, + "fragment" : "(SELECT max(name) FROM identifier_name_source)" + } ] +} + + +-- !query +SELECT * FROM IDENTIFIER( + (SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1) +) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "NOT_A_CONSTANT_STRING.NOT_CONSTANT", + "sqlState" : "42601", + "messageParameters" : { + "expr" : "scalarsubquery()", + "name" : "IDENTIFIER" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 29, + "stopIndex" : 84, + "fragment" : "(SELECT 'x' FROM (SELECT 1) WHERE explode(array(1)) = 1)" + } ] +} + + -- !query SELECT * FROM (SELECT IDENTIFIER((SELECT 'col1')) FROM IDENTIFIER((SELECT 't'))) -- !query schema @@ -969,6 +1590,14 @@ org.apache.spark.sql.catalyst.parser.ParseException } +-- !query +DROP TABLE identifier_name_source +-- !query schema +struct<> +-- !query output + + + -- !query DROP TABLE t -- !query schema diff --git a/sql/core/src/test/resources/sql-tests/results/json-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/json-functions.sql.out index 96c6dab19dc7c..71b94295d37d1 100644 --- a/sql/core/src/test/resources/sql-tests/results/json-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/json-functions.sql.out @@ -825,6 +825,193 @@ struct<json_object_keys([1, 2, 3]):array<string>> NULL +-- !query +select json_typeof() +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "WRONG_NUM_ARGS.WITHOUT_SUGGESTION", + "sqlState" : "42605", + "messageParameters" : { + "actualNum" : "0", + "docroot" : "https://spark.apache.org/docs/latest", + "expectedNum" : "1", + "functionName" : "`json_typeof`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 20, + "fragment" : "json_typeof()" + } ] +} + + +-- !query +select json_typeof(null) +-- !query schema +struct<json_typeof(NULL):string> +-- !query output +NULL + + +-- !query +select json_typeof(200) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"200\"", + "inputType" : "\"INT\"", + "paramIndex" : "first", + "requiredType" : "\"STRING\"", + "sqlExpr" : "\"json_typeof(200)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 23, + "fragment" : "json_typeof(200)" + } ] +} + + +-- !query +select json_typeof('') +-- !query schema +struct<json_typeof():string> +-- !query output +NULL + + +-- !query +select json_typeof('{}') +-- !query schema +struct<json_typeof({}):string> +-- !query output +object + + +-- !query +select json_typeof('{"key": 1, "arr": [1, 2]}') +-- !query schema +struct<json_typeof({"key": 1, "arr": [1, 2]}):string> +-- !query output +object + + +-- !query +select json_typeof('[]') +-- !query schema +struct<json_typeof([]):string> +-- !query output +array + + +-- !query +select json_typeof('[1, 2, 3]') +-- !query schema +struct<json_typeof([1, 2, 3]):string> +-- !query output +array + + +-- !query +select json_typeof('"hello"') +-- !query schema +struct<json_typeof("hello"):string> +-- !query output +string + + +-- !query +select json_typeof('123') +-- !query schema +struct<json_typeof(123):string> +-- !query output +number + + +-- !query +select json_typeof('1.5') +-- !query schema +struct<json_typeof(1.5):string> +-- !query output +number + + +-- !query +select json_typeof('-123') +-- !query schema +struct<json_typeof(-123):string> +-- !query output +number + + +-- !query +select json_typeof('-1.5') +-- !query schema +struct<json_typeof(-1.5):string> +-- !query output +number + + +-- !query +select json_typeof('true') +-- !query schema +struct<json_typeof(true):string> +-- !query output +boolean + + +-- !query +select json_typeof('false') +-- !query schema +struct<json_typeof(false):string> +-- !query output +boolean + + +-- !query +select json_typeof('null') +-- !query schema +struct<json_typeof(null):string> +-- !query output +null + + +-- !query +select json_typeof('bad') +-- !query schema +struct<json_typeof(bad):string> +-- !query output +NULL + + +-- !query +select json_typeof('{"key": 45, "random_string"}') +-- !query schema +struct<json_typeof({"key": 45, "random_string"}):string> +-- !query output +NULL + + +-- !query +select json_typeof('123 true') +-- !query schema +struct<json_typeof(123 true):string> +-- !query output +NULL + + -- !query DROP VIEW IF EXISTS jsonTable -- !query schema @@ -1095,3 +1282,706 @@ select from_json('{"time": "14:30:45"}', 'time TIME') LIMIT 1 struct<from_json({"time": "14:30:45"}):struct<time:time(6)>> -- !query output {"time":14:30:45} + + +-- !query +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.name') +-- !query schema +struct<JSON_VALUE({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.name'):string> +-- !query output +Ada + + +-- !query +select json_value('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.id' RETURNING INT) +-- !query schema +struct<JSON_VALUE({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.id' RETURNING INT):int> +-- !query output +7 + + +-- !query +select json_value('{"id":7,"name":"Ada"}', '$.id' RETURNING INT) + 1 +-- !query schema +struct<(JSON_VALUE({"id":7,"name":"Ada"}, '$.id' RETURNING INT) + 1):int> +-- !query output +8 + + +-- !query +select json_value('{"score":null}', '$.score') +-- !query schema +struct<JSON_VALUE({"score":null}, '$.score'):string> +-- !query output +NULL + + +-- !query +select json_value('{"addr":{"city":"NYC"}}', '$.addr') +-- !query schema +struct<JSON_VALUE({"addr":{"city":"NYC"}}, '$.addr'):string> +-- !query output +NULL + + +-- !query +select json_value('{"tags":["x","y"]}', '$.tags') +-- !query schema +struct<JSON_VALUE({"tags":["x","y"]}, '$.tags'):string> +-- !query output +NULL + + +-- !query +select json_value('{"id":7}', '$.missing') +-- !query schema +struct<JSON_VALUE({"id":7}, '$.missing'):string> +-- !query output +NULL + + +-- !query +select json_value(cast(null as string), '$.a') +-- !query schema +struct<JSON_VALUE(CAST(NULL AS STRING), '$.a'):string> +-- !query output +NULL + + +-- !query +select json_value('{"id":7}', '$.missing' DEFAULT '?' ON EMPTY) +-- !query schema +struct<JSON_VALUE({"id":7}, '$.missing' DEFAULT ? ON EMPTY):string> +-- !query output +? + + +-- !query +select json_value('{"id":7}', '$.missing' RETURNING INT DEFAULT 42 ON EMPTY) +-- !query schema +struct<JSON_VALUE({"id":7}, '$.missing' RETURNING INT DEFAULT 42 ON EMPTY):int> +-- !query output +42 + + +-- !query +select json_value('{"id":7}', '$.missing' ERROR ON EMPTY) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_VALUE_ON_ERROR.EMPTY", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_value`", + "path" : "'$.missing'" + } +} + + +-- !query +select json_value('{"addr":{"city":"NYC"}}', '$.addr' DEFAULT 'n/a' ON ERROR) +-- !query schema +struct<JSON_VALUE({"addr":{"city":"NYC"}}, '$.addr' DEFAULT n/a ON ERROR):string> +-- !query output +n/a + + +-- !query +select json_value('not json', '$.a' DEFAULT 'bad' ON ERROR) +-- !query schema +struct<JSON_VALUE(not json, '$.a' DEFAULT bad ON ERROR):string> +-- !query output +bad + + +-- !query +select json_value('not json', '$.a' ERROR ON ERROR) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_VALUE_ON_ERROR.ERROR", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_value`", + "path" : "'$.a'" + } +} + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT) +-- !query schema +struct<JSON_VALUE({"name":"Ada"}, '$.name' RETURNING INT):int> +-- !query output +NULL + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT ERROR ON ERROR) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_VALUE_ON_ERROR.ERROR", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_value`", + "path" : "'$.name'" + } +} + + +-- !query +select json_value('{"name":"Ada"}', '$.name' RETURNING INT DEFAULT -1 ON ERROR) +-- !query schema +struct<JSON_VALUE({"name":"Ada"}, '$.name' RETURNING INT DEFAULT -1 ON ERROR):int> +-- !query output +-1 + + +-- !query +select json_value('{"a":"x"}', '$.b' DEFAULT 'e' ON EMPTY DEFAULT 'r' ON ERROR) +-- !query schema +struct<JSON_VALUE({"a":"x"}, '$.b' DEFAULT e ON EMPTY DEFAULT r ON ERROR):string> +-- !query output +e + + +-- !query +select json_value('{"v":"3.14"}', '$.v' RETURNING DOUBLE) +-- !query schema +struct<JSON_VALUE({"v":"3.14"}, '$.v' RETURNING DOUBLE):double> +-- !query output +3.14 + + +-- !query +select json_value('{"v":"true"}', '$.v' RETURNING BOOLEAN) +-- !query schema +struct<JSON_VALUE({"v":"true"}, '$.v' RETURNING BOOLEAN):boolean> +-- !query output +true + + +-- !query +select json_value('{"v":"2020-01-02"}', '$.v' RETURNING DATE) +-- !query schema +struct<JSON_VALUE({"v":"2020-01-02"}, '$.v' RETURNING DATE):date> +-- !query output +2020-01-02 + + +-- !query +select json_value('{"a":[1,2]}', '$.a[*]') +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_value`", + "path" : "'$.a[*]'", + "sqlExpr" : "\"JSON_VALUE({\"a\":[1,2]}, '$.a[*]')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 42, + "fragment" : "json_value('{\"a\":[1,2]}', '$.a[*]')" + } ] +} + + +-- !query +select json_value('{"a":1}', '$.a' RETURNING STRUCT<x:INT>) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_SCALAR_RETURNING_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_value`", + "returningType" : "\"STRUCT<x: INT>\"", + "sqlExpr" : "\"JSON_VALUE({\"a\":1}, '$.a' RETURNING STRUCT<x: INT>)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 59, + "fragment" : "json_value('{\"a\":1}', '$.a' RETURNING STRUCT<x:INT>)" + } ] +} + + +-- !query +select json_value('{}', '$.x' RETURNING INT DEFAULT array(1) ON EMPTY) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"JSON_VALUE({}, '$.x' RETURNING INT DEFAULT array(1) ON EMPTY)\"", + "srcType" : "\"ARRAY<INT>\"", + "targetType" : "\"INT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 70, + "fragment" : "json_value('{}', '$.x' RETURNING INT DEFAULT array(1) ON EMPTY)" + } ] +} + + +-- !query +select json_exists('{"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}', '$.addr.city') +-- !query schema +struct<JSON_EXISTS({"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}, '$.addr.city'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"score":null}', '$.score') +-- !query schema +struct<JSON_EXISTS({"score":null}, '$.score'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.addr.zip') +-- !query schema +struct<JSON_EXISTS({"addr":{"city":"NYC"}}, '$.addr.zip'):boolean> +-- !query output +false + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.addr') +-- !query schema +struct<JSON_EXISTS({"addr":{"city":"NYC"}}, '$.addr'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"tags":["x","y"]}', '$.tags[0]') +-- !query schema +struct<JSON_EXISTS({"tags":["x","y"]}, '$.tags[0]'):boolean> +-- !query output +true + + +-- !query +select json_exists(cast(null as string), '$.a') +-- !query schema +struct<JSON_EXISTS(CAST(NULL AS STRING), '$.a'):boolean> +-- !query output +NULL + + +-- !query +select json_exists('not json', '$.a') +-- !query schema +struct<JSON_EXISTS(not json, '$.a'):boolean> +-- !query output +false + + +-- !query +select json_exists('not json', '$.a' TRUE ON ERROR) +-- !query schema +struct<JSON_EXISTS(not json, '$.a' TRUE ON ERROR):boolean> +-- !query output +true + + +-- !query +select json_exists('not json', '$.a' FALSE ON ERROR) +-- !query schema +struct<JSON_EXISTS(not json, '$.a'):boolean> +-- !query output +false + + +-- !query +select json_exists('not json', '$.a' UNKNOWN ON ERROR) +-- !query schema +struct<JSON_EXISTS(not json, '$.a' UNKNOWN ON ERROR):boolean> +-- !query output +NULL + + +-- !query +select json_exists('not json', '$.a' ERROR ON ERROR) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_EXISTS_ON_ERROR", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_exists`", + "path" : "'$.a'" + } +} + + +-- !query +select json_exists('{"a":[1,2]}', '$.a[*]') +-- !query schema +struct<JSON_EXISTS({"a":[1,2]}, '$.a[*]'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"a":[]}', '$.a[*]') +-- !query schema +struct<JSON_EXISTS({"a":[]}, '$.a[*]'):boolean> +-- !query output +false + + +-- !query +select json_exists('{"a":5}', '$.a[*]') +-- !query schema +struct<JSON_EXISTS({"a":5}, '$.a[*]'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"a":[{"b":1},{"c":2}]}', '$.a[*].b') +-- !query schema +struct<JSON_EXISTS({"a":[{"b":1},{"c":2}]}, '$.a[*].b'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"a":[1,2]}', '$.a[5]') +-- !query schema +struct<JSON_EXISTS({"a":[1,2]}, '$.a[5]'):boolean> +-- !query output +false + + +-- !query +select json_exists('{"a":[{"b":1},{"b":2}]}', '$.a.b') +-- !query schema +struct<JSON_EXISTS({"a":[{"b":1},{"b":2}]}, '$.a.b'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"addr":{"city":"NYC"}}', '$.*') +-- !query schema +struct<JSON_EXISTS({"addr":{"city":"NYC"}}, '$.*'):boolean> +-- !query output +true + + +-- !query +select json_exists('{"a":1}', '$[') +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_exists`", + "path" : "'$['", + "sqlExpr" : "\"JSON_EXISTS({\"a\":1}, '$[')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 35, + "fragment" : "json_exists('{\"a\":1}', '$[')" + } ] +} + + +-- !query +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.addr') +-- !query schema +struct<JSON_QUERY({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.addr'):string> +-- !query output +{"city":"NYC"} + + +-- !query +select json_query('{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}', '$.tags') +-- !query schema +struct<JSON_QUERY({"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}, '$.tags'):string> +-- !query output +["x","y"] + + +-- !query +select json_query('{"id":7}', '$.id') +-- !query schema +struct<JSON_QUERY({"id":7}, '$.id'):string> +-- !query output +7 + + +-- !query +select json_query('{"name":"Ada"}', '$.name') +-- !query schema +struct<JSON_QUERY({"name":"Ada"}, '$.name'):string> +-- !query output +"Ada" + + +-- !query +select json_query('{"score":null}', '$.score') +-- !query schema +struct<JSON_QUERY({"score":null}, '$.score'):string> +-- !query output +null + + +-- !query +select json_query('{"id":7}', '$.missing') +-- !query schema +struct<JSON_QUERY({"id":7}, '$.missing'):string> +-- !query output +NULL + + +-- !query +select json_query(cast(null as string), '$.a') +-- !query schema +struct<JSON_QUERY(CAST(NULL AS STRING), '$.a'):string> +-- !query output +NULL + + +-- !query +select json_query('{"tags":["x","y"]}', '$.tags[0]' WITH ARRAY WRAPPER) +-- !query schema +struct<JSON_QUERY({"tags":["x","y"]}, '$.tags[0]' WITH UNCONDITIONAL ARRAY WRAPPER):string> +-- !query output +["x"] + + +-- !query +select json_query('{"tags":["x","y"]}', '$.tags' WITH UNCONDITIONAL ARRAY WRAPPER) +-- !query schema +struct<JSON_QUERY({"tags":["x","y"]}, '$.tags' WITH UNCONDITIONAL ARRAY WRAPPER):string> +-- !query output +[["x","y"]] + + +-- !query +select json_query('{"id":7}', '$.id' WITH ARRAY WRAPPER) +-- !query schema +struct<JSON_QUERY({"id":7}, '$.id' WITH UNCONDITIONAL ARRAY WRAPPER):string> +-- !query output +[7] + + +-- !query +select json_query('{"id":7}', '$.id' WITH CONDITIONAL ARRAY WRAPPER) +-- !query schema +struct<JSON_QUERY({"id":7}, '$.id' WITH CONDITIONAL ARRAY WRAPPER):string> +-- !query output +[7] + + +-- !query +select json_query('{"addr":{"city":"NYC"}}', '$.addr' WITH CONDITIONAL ARRAY WRAPPER) +-- !query schema +struct<JSON_QUERY({"addr":{"city":"NYC"}}, '$.addr' WITH CONDITIONAL ARRAY WRAPPER):string> +-- !query output +{"city":"NYC"} + + +-- !query +select json_query('{"name":"Ada"}', '$.name' OMIT QUOTES) +-- !query schema +struct<JSON_QUERY({"name":"Ada"}, '$.name' OMIT QUOTES):string> +-- !query output +Ada + + +-- !query +select json_query('{"name":"Ada"}', '$.name' KEEP QUOTES) +-- !query schema +struct<JSON_QUERY({"name":"Ada"}, '$.name'):string> +-- !query output +"Ada" + + +-- !query +select json_query('{"id":7}', '$.missing' EMPTY ARRAY ON EMPTY) +-- !query schema +struct<JSON_QUERY({"id":7}, '$.missing' EMPTY ARRAY ON EMPTY):string> +-- !query output +[] + + +-- !query +select json_query('{"id":7}', '$.missing' EMPTY OBJECT ON EMPTY) +-- !query schema +struct<JSON_QUERY({"id":7}, '$.missing' EMPTY OBJECT ON EMPTY):string> +-- !query output +{} + + +-- !query +select json_query('{"id":7}', '$.missing' ERROR ON EMPTY) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_QUERY_ON_ERROR.EMPTY", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_query`", + "path" : "'$.missing'" + } +} + + +-- !query +select json_query('not json', '$.a') +-- !query schema +struct<JSON_QUERY(not json, '$.a'):string> +-- !query output +NULL + + +-- !query +select json_query('not json', '$.a' EMPTY ARRAY ON ERROR) +-- !query schema +struct<JSON_QUERY(not json, '$.a' EMPTY ARRAY ON ERROR):string> +-- !query output +[] + + +-- !query +select json_query('not json', '$.a' EMPTY OBJECT ON ERROR) +-- !query schema +struct<JSON_QUERY(not json, '$.a' EMPTY OBJECT ON ERROR):string> +-- !query output +{} + + +-- !query +select json_query('not json', '$.a' ERROR ON ERROR) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "JSON_QUERY_ON_ERROR.ERROR", + "sqlState" : "2203G", + "messageParameters" : { + "functionName" : "`json_query`", + "path" : "'$.a'" + } +} + + +-- !query +select json_query('{"addr":{"city":"NYC"}}', '$.addr' RETURNING STRING) +-- !query schema +struct<JSON_QUERY({"addr":{"city":"NYC"}}, '$.addr'):string> +-- !query output +{"city":"NYC"} + + +-- !query +select json_query('{"a":[1,2]}', '$.a[*]') +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_PATH", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "path" : "'$.a[*]'", + "sqlExpr" : "\"JSON_QUERY({\"a\":[1,2]}, '$.a[*]')\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 42, + "fragment" : "json_query('{\"a\":[1,2]}', '$.a[*]')" + } ] +} + + +-- !query +select json_query('{"a":1}', '$.a' RETURNING INT) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_QUERY_RETURNING_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "returningType" : "\"INT\"", + "sqlExpr" : "\"JSON_QUERY({\"a\":1}, '$.a' RETURNING INT)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 49, + "fragment" : "json_query('{\"a\":1}', '$.a' RETURNING INT)" + } ] +} + + +-- !query +select json_query('{"name":"Ada"}', '$.name' WITH ARRAY WRAPPER OMIT QUOTES) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.INVALID_JSON_QUERY_WRAPPER_AND_QUOTES", + "sqlState" : "42K09", + "messageParameters" : { + "functionName" : "`json_query`", + "sqlExpr" : "\"JSON_QUERY({\"name\":\"Ada\"}, '$.name' WITH UNCONDITIONAL ARRAY WRAPPER OMIT QUOTES)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 76, + "fragment" : "json_query('{\"name\":\"Ada\"}', '$.name' WITH ARRAY WRAPPER OMIT QUOTES)" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/results/keywords-enforced.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords-enforced.sql.out index 7a6608885e3c0..84dbbaf041d2a 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords-enforced.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords-enforced.sql.out @@ -77,6 +77,7 @@ COMPENSATION false COMPUTE false CONCATENATE false CONDITION false +CONDITIONAL false CONSTRAINT true CONTAINS false CONTINUE false @@ -130,8 +131,10 @@ DOUBLE false DROP false ELSE true ELSEIF false +EMPTY false END true ENFORCED false +ERROR false ESCAPE true ESCAPED false EVOLUTION false @@ -210,6 +213,11 @@ ITEMS false ITERATE false JOIN true JSON false +JSON_EXISTS false +JSON_QUERY false +JSON_TABLE false +JSON_VALUE false +KEEP false KEY false KEYS false LANGUAGE false @@ -268,8 +276,10 @@ NOT true NULL true NULLS false NUMERIC false +OBJECT false OF false OFFSET true +OMIT false ON true ONLY true OPEN false @@ -277,6 +287,7 @@ OPTION false OPTIONS false OR true ORDER true +ORDINALITY false OUT false OUTER true OUTPUTFORMAT false @@ -302,6 +313,7 @@ PURGE false QUALIFY false QUARTER false QUERY false +QUOTES false RANGE false READ false READS false @@ -324,6 +336,7 @@ RESET false RESPECT false RESTRICT false RETURN false +RETURNING false RETURNS false REVOKE false RIGHT true @@ -407,11 +420,13 @@ TYPE false UNARCHIVE false UNBOUNDED false UNCACHE false +UNCONDITIONAL false UNIFORM false UNION true UNIQUE true UNKNOWN true UNLOCK false +UNNEST false UNPIVOT false UNSET false UNTIL false @@ -440,6 +455,7 @@ WINDOW false WITH true WITHIN true WITHOUT false +WRAPPER false X false YEAR false YEARS false diff --git a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out index db8bfa0e205c9..4b57e254ef964 100644 --- a/sql/core/src/test/resources/sql-tests/results/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/keywords.sql.out @@ -77,6 +77,7 @@ COMPENSATION false COMPUTE false CONCATENATE false CONDITION false +CONDITIONAL false CONSTRAINT false CONTAINS false CONTINUE false @@ -130,8 +131,10 @@ DOUBLE false DROP false ELSE false ELSEIF false +EMPTY false END false ENFORCED false +ERROR false ESCAPE false ESCAPED false EVOLUTION false @@ -210,6 +213,11 @@ ITEMS false ITERATE false JOIN false JSON false +JSON_EXISTS false +JSON_QUERY false +JSON_TABLE false +JSON_VALUE false +KEEP false KEY false KEYS false LANGUAGE false @@ -268,8 +276,10 @@ NOT false NULL false NULLS false NUMERIC false +OBJECT false OF false OFFSET false +OMIT false ON false ONLY false OPEN false @@ -277,6 +287,7 @@ OPTION false OPTIONS false OR false ORDER false +ORDINALITY false OUT false OUTER false OUTPUTFORMAT false @@ -302,6 +313,7 @@ PURGE false QUALIFY false QUARTER false QUERY false +QUOTES false RANGE false READ false READS false @@ -324,6 +336,7 @@ RESET false RESPECT false RESTRICT false RETURN false +RETURNING false RETURNS false REVOKE false RIGHT false @@ -407,11 +420,13 @@ TYPE false UNARCHIVE false UNBOUNDED false UNCACHE false +UNCONDITIONAL false UNIFORM false UNION false UNIQUE false UNKNOWN false UNLOCK false +UNNEST false UNPIVOT false UNSET false UNTIL false @@ -440,6 +455,7 @@ WINDOW false WITH false WITHIN false WITHOUT false +WRAPPER false X false YEAR false YEARS false diff --git a/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out b/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out index 96b2aa08884ef..410d7192e11d3 100644 --- a/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/linear-regression.sql.out @@ -290,3 +290,46 @@ SELECT regr_r2(y, k) FROM testRegression where k=2 struct<regr_r2(y, k):double> -- !query output NULL + + +-- !query +CREATE OR REPLACE TEMPORARY VIEW testCorrConstant AS SELECT * FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y) +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM testCorrConstant +-- !query schema +struct<corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)):double> +-- !query output +NULL + + +-- !query +DROP VIEW testCorrConstant +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (1, 2), (1, 3) AS t(x, y) +-- !query schema +struct<corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)):double> +-- !query output +NULL + + +-- !query +SELECT corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)) FROM VALUES +(1, 1), (2, 1), (3, 1) AS t(x, y) +-- !query schema +struct<corr(CAST(x AS DOUBLE), CAST(y AS DOUBLE)):double> +-- !query output +NULL diff --git a/sql/core/src/test/resources/sql-tests/results/math.sql.out b/sql/core/src/test/resources/sql-tests/results/math.sql.out index e2abcb099130a..1f027d362cdf8 100644 --- a/sql/core/src/test/resources/sql-tests/results/math.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/math.sql.out @@ -703,6 +703,350 @@ org.apache.spark.SparkArithmeticException } +-- !query +SELECT truncate(25y, 1) +-- !query schema +struct<truncate(25, 1):tinyint> +-- !query output +25 + + +-- !query +SELECT truncate(25y, 0) +-- !query schema +struct<truncate(25, 0):tinyint> +-- !query output +25 + + +-- !query +SELECT truncate(25y, -1) +-- !query schema +struct<truncate(25, -1):tinyint> +-- !query output +20 + + +-- !query +SELECT truncate(25y, -2) +-- !query schema +struct<truncate(25, -2):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(25y, -3) +-- !query schema +struct<truncate(25, -3):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(-25y, 1) +-- !query schema +struct<truncate(-25, 1):tinyint> +-- !query output +-25 + + +-- !query +SELECT truncate(-25y, 0) +-- !query schema +struct<truncate(-25, 0):tinyint> +-- !query output +-25 + + +-- !query +SELECT truncate(-25y, -1) +-- !query schema +struct<truncate(-25, -1):tinyint> +-- !query output +-20 + + +-- !query +SELECT truncate(-25y, -2) +-- !query schema +struct<truncate(-25, -2):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(-25y, -3) +-- !query schema +struct<truncate(-25, -3):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(127y, -1) +-- !query schema +struct<truncate(127, -1):tinyint> +-- !query output +120 + + +-- !query +SELECT truncate(-128y, -1) +-- !query schema +struct<truncate(-128, -1):tinyint> +-- !query output +-120 + + +-- !query +SELECT truncate(525s, 1) +-- !query schema +struct<truncate(525, 1):smallint> +-- !query output +525 + + +-- !query +SELECT truncate(525s, 0) +-- !query schema +struct<truncate(525, 0):smallint> +-- !query output +525 + + +-- !query +SELECT truncate(525s, -1) +-- !query schema +struct<truncate(525, -1):smallint> +-- !query output +520 + + +-- !query +SELECT truncate(525s, -2) +-- !query schema +struct<truncate(525, -2):smallint> +-- !query output +500 + + +-- !query +SELECT truncate(525s, -3) +-- !query schema +struct<truncate(525, -3):smallint> +-- !query output +0 + + +-- !query +SELECT truncate(-525s, 1) +-- !query schema +struct<truncate(-525, 1):smallint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525s, 0) +-- !query schema +struct<truncate(-525, 0):smallint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525s, -1) +-- !query schema +struct<truncate(-525, -1):smallint> +-- !query output +-520 + + +-- !query +SELECT truncate(-525s, -2) +-- !query schema +struct<truncate(-525, -2):smallint> +-- !query output +-500 + + +-- !query +SELECT truncate(-525s, -3) +-- !query schema +struct<truncate(-525, -3):smallint> +-- !query output +0 + + +-- !query +SELECT truncate(525, 1) +-- !query schema +struct<truncate(525, 1):int> +-- !query output +525 + + +-- !query +SELECT truncate(525, 0) +-- !query schema +struct<truncate(525, 0):int> +-- !query output +525 + + +-- !query +SELECT truncate(525, -1) +-- !query schema +struct<truncate(525, -1):int> +-- !query output +520 + + +-- !query +SELECT truncate(525, -2) +-- !query schema +struct<truncate(525, -2):int> +-- !query output +500 + + +-- !query +SELECT truncate(525, -3) +-- !query schema +struct<truncate(525, -3):int> +-- !query output +0 + + +-- !query +SELECT truncate(-525, 1) +-- !query schema +struct<truncate(-525, 1):int> +-- !query output +-525 + + +-- !query +SELECT truncate(-525, 0) +-- !query schema +struct<truncate(-525, 0):int> +-- !query output +-525 + + +-- !query +SELECT truncate(-525, -1) +-- !query schema +struct<truncate(-525, -1):int> +-- !query output +-520 + + +-- !query +SELECT truncate(-525, -2) +-- !query schema +struct<truncate(-525, -2):int> +-- !query output +-500 + + +-- !query +SELECT truncate(-525, -3) +-- !query schema +struct<truncate(-525, -3):int> +-- !query output +0 + + +-- !query +SELECT truncate(525L, 1) +-- !query schema +struct<truncate(525, 1):bigint> +-- !query output +525 + + +-- !query +SELECT truncate(525L, 0) +-- !query schema +struct<truncate(525, 0):bigint> +-- !query output +525 + + +-- !query +SELECT truncate(525L, -1) +-- !query schema +struct<truncate(525, -1):bigint> +-- !query output +520 + + +-- !query +SELECT truncate(525L, -2) +-- !query schema +struct<truncate(525, -2):bigint> +-- !query output +500 + + +-- !query +SELECT truncate(525L, -3) +-- !query schema +struct<truncate(525, -3):bigint> +-- !query output +0 + + +-- !query +SELECT truncate(-525L, 1) +-- !query schema +struct<truncate(-525, 1):bigint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525L, 0) +-- !query schema +struct<truncate(-525, 0):bigint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525L, -1) +-- !query schema +struct<truncate(-525, -1):bigint> +-- !query output +-520 + + +-- !query +SELECT truncate(-525L, -2) +-- !query schema +struct<truncate(-525, -2):bigint> +-- !query output +-500 + + +-- !query +SELECT truncate(-525L, -3) +-- !query schema +struct<truncate(-525, -3):bigint> +-- !query output +0 + + +-- !query +SELECT truncate(1234.5678) +-- !query schema +struct<truncate(1234.5678, 0):decimal(5,0)> +-- !query output +1234 + + -- !query SELECT conv('100', 2, 10) -- !query schema diff --git a/sql/core/src/test/resources/sql-tests/results/misc-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/misc-functions.sql.out index 9373dc63289aa..5d9438d3a5234 100644 --- a/sql/core/src/test/resources/sql-tests/results/misc-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/misc-functions.sql.out @@ -402,3 +402,51 @@ org.apache.spark.SparkRuntimeException "parameter" : "`algorithm`" } } + + +-- !query +SELECT xxh3_64('Spark') +-- !query schema +struct<xxh3_64(Spark):bigint> +-- !query output +80997306238743657 + + +-- !query +SELECT xxh3_64(CAST('Spark' AS BINARY)) +-- !query schema +struct<xxh3_64(CAST(Spark AS BINARY)):bigint> +-- !query output +80997306238743657 + + +-- !query +SELECT xxh3_128('Spark') +-- !query schema +struct<xxh3_128(Spark):string> +-- !query output +7d57dd84c60c86ca1f4e82ab91a12b5e + + +-- !query +SELECT xxh3_128(CAST('Spark' AS BINARY)) +-- !query schema +struct<xxh3_128(CAST(Spark AS BINARY)):string> +-- !query output +7d57dd84c60c86ca1f4e82ab91a12b5e + + +-- !query +SELECT xxh3_64(CAST(NULL AS STRING)) +-- !query schema +struct<xxh3_64(CAST(NULL AS STRING)):bigint> +-- !query output +NULL + + +-- !query +SELECT xxh3_128(CAST(NULL AS BINARY)) +-- !query schema +struct<xxh3_128(CAST(NULL AS BINARY)):string> +-- !query output +NULL diff --git a/sql/core/src/test/resources/sql-tests/results/nonansi/array.sql.out b/sql/core/src/test/resources/sql-tests/results/nonansi/array.sql.out index 3bd7671ae386f..b5c3ee79929bd 100644 --- a/sql/core/src/test/resources/sql-tests/results/nonansi/array.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/nonansi/array.sql.out @@ -1010,3 +1010,103 @@ select array_distinct(array(0.0, -0.0, -0.0, DOUBLE("NaN"), DOUBLE("NaN"))) struct<array_distinct(array(0.0, 0.0, 0.0, NaN, NaN)):array<double>> -- !query output [0.0,NaN] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 0) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 0):array<int>> +-- !query output +[1,2,3,4,5] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 2) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 2):array<int>> +-- !query output +[1,2,3] + + +-- !query +select trim_array(array(1, 2, 3, 4, 5), 5) +-- !query schema +struct<trim_array(array(1, 2, 3, 4, 5), 5):array<int>> +-- !query output +[] + + +-- !query +select trim_array(array('a', 'b', 'c'), 1) +-- !query schema +struct<trim_array(array(a, b, c), 1):array<string>> +-- !query output +["a","b"] + + +-- !query +select trim_array(array(1, 2, null, 4), 1) +-- !query schema +struct<trim_array(array(1, 2, NULL, 4), 1):array<int>> +-- !query output +[1,2,null] + + +-- !query +select trim_array(array(), 0) +-- !query schema +struct<trim_array(array(), 0):array<void>> +-- !query output +[] + + +-- !query +select trim_array(array(1, 2, 3), -1) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + "sqlState" : "22023", + "messageParameters" : { + "functionName" : "`trim_array`", + "length" : "-1", + "numElements" : "3", + "parameter" : "`n`" + } +} + + +-- !query +select trim_array(array(1, 2, 3), 4) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.TRIM_ARRAY_LENGTH", + "sqlState" : "22023", + "messageParameters" : { + "functionName" : "`trim_array`", + "length" : "4", + "numElements" : "3", + "parameter" : "`n`" + } +} + + +-- !query +select trim_array(CAST(null AS ARRAY<INT>), 1) +-- !query schema +struct<trim_array(NULL, 1):array<int>> +-- !query output +NULL + + +-- !query +select trim_array(array(1, 2, 3), CAST(null AS INT)) +-- !query schema +struct<trim_array(array(1, 2, 3), CAST(NULL AS INT)):array<int>> +-- !query output +NULL diff --git a/sql/core/src/test/resources/sql-tests/results/nonansi/keywords.sql.out b/sql/core/src/test/resources/sql-tests/results/nonansi/keywords.sql.out index db8bfa0e205c9..4b57e254ef964 100644 --- a/sql/core/src/test/resources/sql-tests/results/nonansi/keywords.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/nonansi/keywords.sql.out @@ -77,6 +77,7 @@ COMPENSATION false COMPUTE false CONCATENATE false CONDITION false +CONDITIONAL false CONSTRAINT false CONTAINS false CONTINUE false @@ -130,8 +131,10 @@ DOUBLE false DROP false ELSE false ELSEIF false +EMPTY false END false ENFORCED false +ERROR false ESCAPE false ESCAPED false EVOLUTION false @@ -210,6 +213,11 @@ ITEMS false ITERATE false JOIN false JSON false +JSON_EXISTS false +JSON_QUERY false +JSON_TABLE false +JSON_VALUE false +KEEP false KEY false KEYS false LANGUAGE false @@ -268,8 +276,10 @@ NOT false NULL false NULLS false NUMERIC false +OBJECT false OF false OFFSET false +OMIT false ON false ONLY false OPEN false @@ -277,6 +287,7 @@ OPTION false OPTIONS false OR false ORDER false +ORDINALITY false OUT false OUTER false OUTPUTFORMAT false @@ -302,6 +313,7 @@ PURGE false QUALIFY false QUARTER false QUERY false +QUOTES false RANGE false READ false READS false @@ -324,6 +336,7 @@ RESET false RESPECT false RESTRICT false RETURN false +RETURNING false RETURNS false REVOKE false RIGHT false @@ -407,11 +420,13 @@ TYPE false UNARCHIVE false UNBOUNDED false UNCACHE false +UNCONDITIONAL false UNIFORM false UNION false UNIQUE false UNKNOWN false UNLOCK false +UNNEST false UNPIVOT false UNSET false UNTIL false @@ -440,6 +455,7 @@ WINDOW false WITH false WITHIN false WITHOUT false +WRAPPER false X false YEAR false YEARS false diff --git a/sql/core/src/test/resources/sql-tests/results/nonansi/math.sql.out b/sql/core/src/test/resources/sql-tests/results/nonansi/math.sql.out index 09f4383933288..5ab0677bfd7c2 100644 --- a/sql/core/src/test/resources/sql-tests/results/nonansi/math.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/nonansi/math.sql.out @@ -447,6 +447,350 @@ struct<bround(-9223372036854775808, -1):bigint> 9223372036854775806 +-- !query +SELECT truncate(25y, 1) +-- !query schema +struct<truncate(25, 1):tinyint> +-- !query output +25 + + +-- !query +SELECT truncate(25y, 0) +-- !query schema +struct<truncate(25, 0):tinyint> +-- !query output +25 + + +-- !query +SELECT truncate(25y, -1) +-- !query schema +struct<truncate(25, -1):tinyint> +-- !query output +20 + + +-- !query +SELECT truncate(25y, -2) +-- !query schema +struct<truncate(25, -2):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(25y, -3) +-- !query schema +struct<truncate(25, -3):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(-25y, 1) +-- !query schema +struct<truncate(-25, 1):tinyint> +-- !query output +-25 + + +-- !query +SELECT truncate(-25y, 0) +-- !query schema +struct<truncate(-25, 0):tinyint> +-- !query output +-25 + + +-- !query +SELECT truncate(-25y, -1) +-- !query schema +struct<truncate(-25, -1):tinyint> +-- !query output +-20 + + +-- !query +SELECT truncate(-25y, -2) +-- !query schema +struct<truncate(-25, -2):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(-25y, -3) +-- !query schema +struct<truncate(-25, -3):tinyint> +-- !query output +0 + + +-- !query +SELECT truncate(127y, -1) +-- !query schema +struct<truncate(127, -1):tinyint> +-- !query output +120 + + +-- !query +SELECT truncate(-128y, -1) +-- !query schema +struct<truncate(-128, -1):tinyint> +-- !query output +-120 + + +-- !query +SELECT truncate(525s, 1) +-- !query schema +struct<truncate(525, 1):smallint> +-- !query output +525 + + +-- !query +SELECT truncate(525s, 0) +-- !query schema +struct<truncate(525, 0):smallint> +-- !query output +525 + + +-- !query +SELECT truncate(525s, -1) +-- !query schema +struct<truncate(525, -1):smallint> +-- !query output +520 + + +-- !query +SELECT truncate(525s, -2) +-- !query schema +struct<truncate(525, -2):smallint> +-- !query output +500 + + +-- !query +SELECT truncate(525s, -3) +-- !query schema +struct<truncate(525, -3):smallint> +-- !query output +0 + + +-- !query +SELECT truncate(-525s, 1) +-- !query schema +struct<truncate(-525, 1):smallint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525s, 0) +-- !query schema +struct<truncate(-525, 0):smallint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525s, -1) +-- !query schema +struct<truncate(-525, -1):smallint> +-- !query output +-520 + + +-- !query +SELECT truncate(-525s, -2) +-- !query schema +struct<truncate(-525, -2):smallint> +-- !query output +-500 + + +-- !query +SELECT truncate(-525s, -3) +-- !query schema +struct<truncate(-525, -3):smallint> +-- !query output +0 + + +-- !query +SELECT truncate(525, 1) +-- !query schema +struct<truncate(525, 1):int> +-- !query output +525 + + +-- !query +SELECT truncate(525, 0) +-- !query schema +struct<truncate(525, 0):int> +-- !query output +525 + + +-- !query +SELECT truncate(525, -1) +-- !query schema +struct<truncate(525, -1):int> +-- !query output +520 + + +-- !query +SELECT truncate(525, -2) +-- !query schema +struct<truncate(525, -2):int> +-- !query output +500 + + +-- !query +SELECT truncate(525, -3) +-- !query schema +struct<truncate(525, -3):int> +-- !query output +0 + + +-- !query +SELECT truncate(-525, 1) +-- !query schema +struct<truncate(-525, 1):int> +-- !query output +-525 + + +-- !query +SELECT truncate(-525, 0) +-- !query schema +struct<truncate(-525, 0):int> +-- !query output +-525 + + +-- !query +SELECT truncate(-525, -1) +-- !query schema +struct<truncate(-525, -1):int> +-- !query output +-520 + + +-- !query +SELECT truncate(-525, -2) +-- !query schema +struct<truncate(-525, -2):int> +-- !query output +-500 + + +-- !query +SELECT truncate(-525, -3) +-- !query schema +struct<truncate(-525, -3):int> +-- !query output +0 + + +-- !query +SELECT truncate(525L, 1) +-- !query schema +struct<truncate(525, 1):bigint> +-- !query output +525 + + +-- !query +SELECT truncate(525L, 0) +-- !query schema +struct<truncate(525, 0):bigint> +-- !query output +525 + + +-- !query +SELECT truncate(525L, -1) +-- !query schema +struct<truncate(525, -1):bigint> +-- !query output +520 + + +-- !query +SELECT truncate(525L, -2) +-- !query schema +struct<truncate(525, -2):bigint> +-- !query output +500 + + +-- !query +SELECT truncate(525L, -3) +-- !query schema +struct<truncate(525, -3):bigint> +-- !query output +0 + + +-- !query +SELECT truncate(-525L, 1) +-- !query schema +struct<truncate(-525, 1):bigint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525L, 0) +-- !query schema +struct<truncate(-525, 0):bigint> +-- !query output +-525 + + +-- !query +SELECT truncate(-525L, -1) +-- !query schema +struct<truncate(-525, -1):bigint> +-- !query output +-520 + + +-- !query +SELECT truncate(-525L, -2) +-- !query schema +struct<truncate(-525, -2):bigint> +-- !query output +-500 + + +-- !query +SELECT truncate(-525L, -3) +-- !query schema +struct<truncate(-525, -3):bigint> +-- !query output +0 + + +-- !query +SELECT truncate(1234.5678) +-- !query schema +struct<truncate(1234.5678, 0):decimal(5,0)> +-- !query output +1234 + + -- !query SELECT conv('100', 2, 10) -- !query schema diff --git a/sql/core/src/test/resources/sql-tests/results/nonansi/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/nonansi/string-functions.sql.out index 1a8521a37cfce..c6655b01d0e65 100644 --- a/sql/core/src/test/resources/sql-tests/results/nonansi/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/nonansi/string-functions.sql.out @@ -2701,3 +2701,52 @@ select instr(null, null, cast(null as int), cast(null as int)) struct<instr(NULL, NULL, CAST(NULL AS INT), CAST(NULL AS INT)):int> -- !query output NULL + + +-- !query +select normalize('hello') +-- !query schema +struct<normalize(hello, NFC):string> +-- !query output +hello + + +-- !query +select normalize('hello', 'NFD') +-- !query schema +struct<normalize(hello, NFD):string> +-- !query output +hello + + +-- !query +select normalize('fi', 'NFKC') +-- !query schema +struct<normalize(fi, NFKC):string> +-- !query output +fi + + +-- !query +select normalize(null, 'NFC') +-- !query schema +struct<normalize(NULL, NFC):string> +-- !query output +NULL + + +-- !query +select normalize('hello', 'not_a_form') +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.NORMALIZE_FORM", + "sqlState" : "22023", + "messageParameters" : { + "form" : "'not_a_form'", + "functionName" : "`normalize`", + "parameter" : "`form`" + } +} diff --git a/sql/core/src/test/resources/sql-tests/results/nonansi/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/results/nonansi/timestamp.sql.out index 4d7786512e556..c0aaec386bf9b 100644 --- a/sql/core/src/test/resources/sql-tests/results/nonansi/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/nonansi/timestamp.sql.out @@ -1035,7 +1035,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:10\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP '2011-11-11 11:11:11' - 2011-11-11 11:11:10)\"" }, "queryContext" : [ { @@ -1061,7 +1061,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:11\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(2011-11-11 11:11:11 - TIMESTAMP '2011-11-11 11:11:10')\"" }, "queryContext" : [ { @@ -1111,7 +1111,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(str - TIMESTAMP '2011-11-11 11:11:11')\"" }, "queryContext" : [ { @@ -1137,7 +1137,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP '2011-11-11 11:11:11' - str)\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/results/parse-sql-gating.sql.out b/sql/core/src/test/resources/sql-tests/results/parse-sql-gating.sql.out new file mode 100644 index 0000000000000..aac0288887045 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/parse-sql-gating.sql.out @@ -0,0 +1,16 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT parse_sql('SELECT 1') +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "FEATURE_NOT_ENABLED", + "sqlState" : "56038", + "messageParameters" : { + "configKey" : "spark.sql.function.parseSql.enabled", + "configValue" : "true", + "featureName" : "parse_sql" + } +} diff --git a/sql/core/src/test/resources/sql-tests/results/parse-sql.sql.out b/sql/core/src/test/resources/sql-tests/results/parse-sql.sql.out new file mode 100644 index 0000000000000..986ce7c4ae945 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/parse-sql.sql.out @@ -0,0 +1,571 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +SELECT parse_sql(NULL) +-- !query schema +struct<parse_sql(NULL):string> +-- !query output +NULL + + +-- !query +SELECT parse_sql('SELECT a, b FROM t') +-- !query schema +struct<parse_sql(SELECT a, b FROM t):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["t"]],"select_list":[{"name":["a"]},{"name":["b"]}]} + + +-- !query +SELECT parse_sql('SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2') +-- !query schema +struct<parse_sql(SELECT db.my_func(a), count(b) FROM cat.ns.t1 JOIN t2):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["cat","ns","t1"],["t2"]],"function_references":[["db","my_func"],["count"]],"select_list":[{"name":[]},{"name":[]}]} + + +-- !query +SELECT + get_json_object(result, '$.statement_identifier') AS statement_identifier, + get_json_object(result, '$.source_table_references[0][0]') AS first_table, + get_json_object(result, '$.select_list[1].name[0]') AS second_column +FROM (SELECT parse_sql('SELECT a, b FROM t') AS result) +-- !query schema +struct<statement_identifier:string,first_table:string,second_column:string> +-- !query output +SELECT t b + + +-- !query +SELECT parse_sql('INSERT INTO t SELECT 1') +-- !query schema +struct<parse_sql(INSERT INTO t SELECT 1):string> +-- !query output +{"parse_success":true,"statement_identifier":"INSERT","statement_code":50,"target_table_references":[["t"]],"select_list":[{"name":[]}]} + + +-- !query +SELECT parse_sql('DELETE FROM t WHERE a = 1') +-- !query schema +struct<parse_sql(DELETE FROM t WHERE a = 1):string> +-- !query output +{"parse_success":true,"statement_identifier":"DELETE WHERE","statement_code":19,"target_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('UPDATE t SET a = 1 WHERE b = 2') +-- !query schema +struct<parse_sql(UPDATE t SET a = 1 WHERE b = 2):string> +-- !query output +{"parse_success":true,"statement_identifier":"UPDATE WHERE","statement_code":82,"target_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE') +-- !query schema +struct<parse_sql(MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE):string> +-- !query output +{"parse_success":true,"statement_identifier":"MERGE","statement_code":128,"target_table_references":[["t"]],"source_table_references":[["s"]]} + + +-- !query +SELECT parse_sql('CREATE TABLE t (a INT)') +-- !query schema +struct<parse_sql(CREATE TABLE t (a INT)):string> +-- !query output +{"parse_success":true,"statement_identifier":"CREATE TABLE","statement_code":77,"target_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('CREATE TABLE t AS SELECT 1 AS a') +-- !query schema +struct<parse_sql(CREATE TABLE t AS SELECT 1 AS a):string> +-- !query output +{"parse_success":true,"statement_identifier":"CREATE TABLE","statement_code":77,"target_table_references":[["t"]],"select_list":[{"name":["a"]}]} + + +-- !query +SELECT parse_sql('DROP TABLE t') +-- !query schema +struct<parse_sql(DROP TABLE t):string> +-- !query output +{"parse_success":true,"statement_identifier":"DROP TABLE","statement_code":32,"target_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('CACHE TABLE t') +-- !query schema +struct<parse_sql(CACHE TABLE t):string> +-- !query output +{"parse_success":true,"statement_identifier":"CACHE TABLE","statement_code":-1,"target_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('TABLE t') +-- !query schema +struct<parse_sql(TABLE t):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["t"]]} + + +-- !query +SELECT parse_sql('VALUES (1), (2)') +-- !query schema +struct<parse_sql(VALUES (1), (2)):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21} + + +-- !query +SELECT parse_sql('CREATE FUNCTION f AS ''x'' USING JAR ''y.jar''') +-- !query schema +struct<parse_sql(CREATE FUNCTION f AS 'x' USING JAR 'y.jar'):string> +-- !query output +{"parse_success":true,"statement_identifier":"CREATE ROUTINE","statement_code":14} + + +-- !query +SELECT parse_sql('DECLARE VARIABLE x INT') +-- !query schema +struct<parse_sql(DECLARE VARIABLE x INT):string> +-- !query output +{"parse_success":true,"statement_identifier":"DECLARE VARIABLE","statement_code":-8} + + +-- !query +SELECT parse_sql('SELECT * FROM t WHERE a = :foo AND b = ?') +-- !query schema +struct<parse_sql(SELECT * FROM t WHERE a = :foo AND b = ?):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["t"]],"select_list":[{"name":["*"]}],"parameter_markers":{"named":["foo"],"unnamed_count":1}} + + +-- !query +SELECT parse_sql('WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte') +-- !query schema +struct<parse_sql(WITH cte AS (SELECT a FROM hidden_base) SELECT a FROM cte):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["hidden_base"]],"select_list":[{"name":["a"]}]} + + +-- !query +SELECT parse_sql('SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)') +-- !query schema +struct<parse_sql(SELECT * FROM real_t WHERE EXISTS (WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["inner_base"],["real_t"]],"select_list":[{"name":["*"]}]} + + +-- !query +SELECT parse_sql('WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a') +-- !query schema +struct<parse_sql(WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["b"]],"select_list":[{"name":["*"]}]} + + +-- !query +SELECT parse_sql('SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id)') +-- !query schema +struct<parse_sql(SELECT (SELECT max(v) FROM scalar_src) AS m, t.a FROM outer_t t WHERE EXISTS (SELECT 1 FROM exists_src e WHERE e.id = t.id)):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["scalar_src"],["exists_src"],["outer_t"]],"function_references":[["max"]],"select_list":[{"name":["m"]},{"name":["t","a"]}]} + + +-- !query +SELECT parse_sql( +'SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, ''x'') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1)') +-- !query schema +struct<parse_sql(SELECT coalesce(t.a, 0), sum(abs(t.b)) OVER ( + PARTITION BY lower(t.c) ORDER BY length(t.d)) + FROM left_t t + JOIN right_t r ON hash(t.id) = hash(r.id) + JOIN LATERAL range(cast(t.n AS BIGINT)) rng + WHERE startswith(t.c, 'x') + AND EXISTS (SELECT max(s.v) FROM scalar_t s WHERE s.id = t.id) + GROUP BY coalesce(t.a, 0), t.b, t.c, t.d + HAVING count_if(t.b > 0) > 0 + ORDER BY greatest(t.a, 1)):string> +-- !query output +{"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"source_table_references":[["scalar_t"],["left_t"],["right_t"]],"function_references":[["greatest"],["count_if"],["coalesce"],["sum"],["abs"],["lower"],["length"],["startswith"],["max"],["range"],["hash"]],"select_list":[{"name":[]},{"name":[]}]} + + +-- !query +SELECT parse_sql( +'MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name))') +-- !query schema +struct<parse_sql(MERGE INTO target t + USING ( + SELECT id, normalize_name(name) AS name + FROM source + WHERE is_valid(id) + ) s + ON hash(t.id) = hash(s.id) + WHEN MATCHED AND should_update(t.name, s.name) THEN + UPDATE SET name = coalesce(s.name, upper(t.name)) + WHEN NOT MATCHED THEN + INSERT (id, name) VALUES (s.id, lower(s.name))):string> +-- !query output +{"parse_success":true,"statement_identifier":"MERGE","statement_code":128,"target_table_references":[["target"]],"source_table_references":[["source"]],"function_references":[["hash"],["should_update"],["coalesce"],["upper"],["lower"],["normalize_name"],["is_valid"]]} + + +-- !query +SELECT parse_sql( +'CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper(''x'') + )') +-- !query schema +struct<parse_sql(CREATE TABLE defaults ( + created DATE DEFAULT current_date(), + normalized STRING DEFAULT upper('x') + )):string> +-- !query output +{"parse_success":true,"statement_identifier":"CREATE TABLE","statement_code":77,"target_table_references":[["defaults"]],"function_references":[["current_date"],["upper"]]} + + +-- !query +SELECT parse_sql('SELEC FROM t') +-- !query schema +struct<parse_sql(SELEC FROM t):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"PARSE_SYNTAX_ERROR","messageTemplate":"Syntax error at or near <error><hint>.","sqlState":"42601","messageParameters":{"error":"'SELEC'","hint":""},"queryContext":[{"objectType":"","objectName":"","startIndex":1,"stopIndex":12,"fragment":"SELEC FROM t"}],"line":1,"position":0}} + + +-- !query +SELECT + get_json_object(result, '$.parse_success') AS parse_success, + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM (SELECT parse_sql('SELEC FROM t') AS result) +-- !query schema +struct<parse_success:string,error_class:string,fragment:string> +-- !query output +false PARSE_SYNTAX_ERROR SELEC FROM t + + +-- !query +SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b') +-- !query schema +struct<parse_sql(SELECT * + FROM t + ORDER BY a + CLUSTER BY b):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"UNSUPPORTED_FEATURE.COMBINATION_QUERY_RESULT_CLAUSES","messageTemplate":"The feature is not supported: Combination of ORDER BY/SORT BY/DISTRIBUTE BY/CLUSTER BY.","sqlState":"0A000","queryContext":[{"objectType":"","objectName":"","startIndex":19,"stopIndex":42,"fragment":"ORDER BY a\n CLUSTER BY b"}],"line":3,"position":1}} + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].startIndex') AS start_index +FROM ( + SELECT parse_sql( +'SELECT * + FROM t + ORDER BY a + CLUSTER BY b') AS result +) +-- !query schema +struct<error_class:string,line:string,position:string,start_index:string> +-- !query output +UNSUPPORTED_FEATURE.COMBINATION_QUERY_RESULT_CLAUSES 3 1 19 + + +-- !query +SELECT parse_sql('') +-- !query schema +struct<parse_sql():string> +-- !query output +{"parse_success":false,"error":{"errorClass":"PARSE_EMPTY_STATEMENT","messageTemplate":"Syntax error, unexpected empty statement.","sqlState":"42617","line":1,"position":0}} + + +-- !query +SELECT parse_sql('USE bad-name') +-- !query schema +struct<parse_sql(USE bad-name):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"INVALID_IDENTIFIER","messageTemplate":"The unquoted identifier <ident> is invalid and must be back quoted as: `<ident>`.\nUnquoted identifiers can only contain ASCII letters ('a' - 'z', 'A' - 'Z'), digits ('0' - '9'), and underbar ('_').\nUnquoted identifiers must also not start with a digit.\nDifferent data sources and meta stores may impose additional restrictions on valid identifiers.","sqlState":"42602","messageParameters":{"ident":"bad-name"},"queryContext":[{"objectType":"","objectName":"","startIndex":1,"stopIndex":12,"fragment":"USE bad-name"}],"line":1,"position":7}} + + +-- !query +SELECT parse_sql('WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c') +-- !query schema +struct<parse_sql(WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"DUPLICATED_CTE_NAMES","messageTemplate":"CTE definition can't have duplicate names: <duplicateNames>.","sqlState":"42602","messageParameters":{"duplicateNames":"`c`"},"queryContext":[{"objectType":"","objectName":"","startIndex":1,"stopIndex":53,"fragment":"WITH c AS (SELECT 1), c AS (SELECT 2) SELECT * FROM c"}],"line":1,"position":0}} + + +-- !query +SELECT parse_sql('MERGE INTO target USING source ON target.id = source.id') +-- !query schema +struct<parse_sql(MERGE INTO target USING source ON target.id = source.id):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"MERGE_WITHOUT_WHEN","messageTemplate":"There must be at least one WHEN clause in a MERGE statement.","sqlState":"42601","queryContext":[{"objectType":"","objectName":"","startIndex":1,"stopIndex":55,"fragment":"MERGE INTO target USING source ON target.id = source.id"}],"line":1,"position":0}} + + +-- !query +SELECT parse_sql('EXPLAIN SELECT 1') +-- !query schema +struct<parse_sql(EXPLAIN SELECT 1):string> +-- !query output +{"parse_success":true,"statement_identifier":"EXPLAIN","statement_code":-23,"select_list":[{"name":[]}]} + + +-- !query +SELECT parse_sql('SET spark.sql.adaptive.enabled=true') +-- !query schema +struct<parse_sql(SET spark.sql.adaptive.enabled=true):string> +-- !query output +{"parse_success":true,"statement_identifier":"SET","statement_code":-24} + + +-- !query +SELECT parse_sql('ADD JAR /tmp/x.jar') +-- !query schema +struct<parse_sql(ADD JAR /tmp/x.jar):string> +-- !query output +{"parse_success":true,"statement_identifier":"ADD JAR","statement_code":-26} + + +-- !query +SELECT parse_sql('CREATE VIEW v AS SELECT a, b FROM t') +-- !query schema +struct<parse_sql(CREATE VIEW v AS SELECT a, b FROM t):string> +-- !query output +{"parse_success":true,"statement_identifier":"CREATE VIEW","statement_code":84,"target_table_references":[["v"]],"source_table_references":[["t"]],"select_list":[{"name":["a"]},{"name":["b"]}]} + + +-- !query +SELECT parse_sql('SELECT 1 AS IDENTIFIER(''alias.field'')') +-- !query schema +struct<parse_sql(SELECT 1 AS IDENTIFIER('alias.field')):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"IDENTIFIER_TOO_MANY_NAME_PARTS","messageTemplate":"<identifier> is not a valid identifier as it has more than <limit> name parts.","sqlState":"42601","messageParameters":{"identifier":"`alias`.`field`","limit":"1"},"queryContext":[{"objectType":"","objectName":"","startIndex":8,"stopIndex":37,"fragment":"1 AS IDENTIFIER('alias.field')"}],"line":1,"position":12}} + + +-- !query +SELECT parse_sql('SELECT DATE ''not-a-date''') +-- !query schema +struct<parse_sql(SELECT DATE 'not-a-date'):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"INVALID_TYPED_LITERAL","messageTemplate":"The value of the typed literal <valueType> is invalid: <value>.","sqlState":"42604","messageParameters":{"value":"'not-a-date'","valueType":"\"DATE\""},"queryContext":[{"objectType":"","objectName":"","startIndex":8,"stopIndex":24,"fragment":"DATE 'not-a-date'"}],"line":1,"position":7}} + + +-- !query +SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END') +-- !query schema +struct<parse_sql(BEGIN + SELECT 1; + SELEC 2; + END):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"PARSE_SYNTAX_ERROR","messageTemplate":"Syntax error at or near <error><hint>.","sqlState":"42601","messageParameters":{"error":"'2'","hint":""},"queryContext":[{"objectType":"","objectName":"","startIndex":1,"stopIndex":35,"fragment":"BEGIN\n SELECT 1;\n SELEC 2;\n END"}],"line":3,"position":9}} + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + SELECT 1; + SELEC 2; + END') AS result +) +-- !query schema +struct<error_class:string,line:string,position:string,fragment:string> +-- !query output +PARSE_SYNTAX_ERROR 3 9 BEGIN + SELECT 1; + SELEC 2; + END + + +-- !query +SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END') +-- !query schema +struct<parse_sql(BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END):string> +-- !query output +{"parse_success":false,"error":{"errorClass":"LABELS_MISMATCH","messageTemplate":"Begin label <beginLabel> does not match the end label <endLabel>.","sqlState":"42K0L","messageParameters":{"beginLabel":"`lbl_begin`","endLabel":"`lbl_end`"},"queryContext":[{"objectType":"","objectName":"","startIndex":10,"stopIndex":19,"fragment":"lbl_begin:"}],"line":2,"position":3}} + + +-- !query +SELECT + get_json_object(result, '$.error.errorClass') AS error_class, + get_json_object(result, '$.error.line') AS line, + get_json_object(result, '$.error.position') AS position, + get_json_object(result, '$.error.queryContext[0].fragment') AS fragment +FROM ( + SELECT parse_sql( +'BEGIN + lbl_begin: BEGIN + SELECT 1; + END lbl_end; + END') AS result +) +-- !query schema +struct<error_class:string,line:string,position:string,fragment:string> +-- !query output +LABELS_MISMATCH 2 3 lbl_begin: + + +-- !query +SELECT sql_text, parse_sql(sql_text) FROM VALUES + ('SELECT 1'), + ('INSERT INTO t SELECT 1'), + ('CACHE TABLE t') +AS t(sql_text) +-- !query schema +struct<sql_text:string,parse_sql(sql_text):string> +-- !query output +CACHE TABLE t {"parse_success":true,"statement_identifier":"CACHE TABLE","statement_code":-1,"target_table_references":[["t"]]} +INSERT INTO t SELECT 1 {"parse_success":true,"statement_identifier":"INSERT","statement_code":50,"target_table_references":[["t"]],"select_list":[{"name":[]}]} +SELECT 1 {"parse_success":true,"statement_identifier":"SELECT","statement_code":21,"select_list":[{"name":[]}]} + + +-- !query +SELECT parse_sql('BEGIN SELECT 1; END') +-- !query schema +struct<parse_sql(BEGIN SELECT 1; END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22} + + +-- !query +SELECT parse_sql('BEGIN SELECT count(a) FROM script_t WHERE c = :p; END') +-- !query schema +struct<parse_sql(BEGIN SELECT count(a) FROM script_t WHERE c = :p; END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22,"source_table_references":[["script_t"]],"function_references":[["count"]],"parameter_markers":{"named":["p"]}} + + +-- !query +SELECT parse_sql('BEGIN SELECT * FROM t WHERE a = ?; END') +-- !query schema +struct<parse_sql(BEGIN SELECT * FROM t WHERE a = ?; END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22,"source_table_references":[["t"]],"parameter_markers":{"unnamed_count":1}} + + +-- !query +SELECT parse_sql('BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END') +-- !query schema +struct<parse_sql(BEGIN IF (SELECT flag FROM gate) THEN INSERT INTO dest SELECT * FROM src_if; ELSE DELETE FROM src_else; END IF; END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22,"target_table_references":[["dest"],["src_else"]],"source_table_references":[["gate"],["src_if"]]} + + +-- !query +SELECT parse_sql('BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END') +-- !query schema +struct<parse_sql(BEGIN DECLARE EXIT HANDLER FOR SQLEXCEPTION BEGIN INSERT INTO err_log SELECT * FROM failing_row; END; SELECT a FROM main_t; END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22,"target_table_references":[["err_log"]],"source_table_references":[["failing_row"],["main_t"]]} + + +-- !query +SELECT parse_sql( +'BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string(''%s'', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END') +-- !query schema +struct<parse_sql(BEGIN + DECLARE EXIT HANDLER FOR SQLEXCEPTION + BEGIN + INSERT INTO error_log + SELECT format_string('%s', message) FROM error_source; + END; + + WITH prepared AS ( + SELECT id, normalize_name(name) AS name + FROM input_names + WHERE is_valid(id) + ) + INSERT INTO output_names + SELECT id, upper(name) FROM prepared; + + IF EXISTS (SELECT 1 FROM control_flags WHERE enabled()) THEN + UPDATE update_target + SET value = coalesce((SELECT max(value) FROM update_source), 0) + WHERE should_update(id); + ELSE + DELETE FROM delete_target + WHERE id IN (SELECT id FROM delete_source WHERE expired(ts)); + END IF; + + FOR row AS + SELECT id FROM loop_source WHERE ready(id) + DO + SELECT audit(row.id), count(*) FROM loop_body; + END FOR; + END):string> +-- !query output +{"parse_success":true,"statement_identifier":"BEGIN END","statement_code":-22,"target_table_references":[["error_log"],["output_names"],["update_target"],["delete_target"]],"source_table_references":[["error_source"],["input_names"],["control_flags"],["update_source"],["delete_source"],["loop_source"],["loop_body"]],"function_references":[["format_string"],["normalize_name"],["is_valid"],["upper"],["enabled"],["coalesce"],["should_update"],["max"],["expired"],["ready"],["audit"],["count"]]} diff --git a/sql/core/src/test/resources/sql-tests/results/random.sql.out b/sql/core/src/test/resources/sql-tests/results/random.sql.out index 049d134003cd4..5439683037b7e 100644 --- a/sql/core/src/test/resources/sql-tests/results/random.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/random.sql.out @@ -857,15 +857,23 @@ SELECT randstr(-1, 0) AS result -- !query schema struct<> -- !query output -org.apache.spark.SparkRuntimeException +org.apache.spark.sql.catalyst.ExtendedAnalysisException { - "errorClass" : "INVALID_PARAMETER_VALUE.LENGTH", - "sqlState" : "22023", + "errorClass" : "DATATYPE_MISMATCH.VALUE_OUT_OF_RANGE", + "sqlState" : "42K09", "messageParameters" : { - "functionName" : "`randstr`", - "length" : "-1", - "parameter" : "`length`" - } + "currentValue" : "-1", + "exprName" : "`length`", + "sqlExpr" : "\"randstr(-1, 0)\"", + "valueRange" : "[0, 2147483647]" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 21, + "fragment" : "randstr(-1, 0)" + } ] } diff --git a/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out b/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out index b0497277c59bc..e6be16af67848 100644 --- a/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/string-functions.sql.out @@ -2769,3 +2769,52 @@ select instr(null, null, cast(null as int), cast(null as int)) struct<instr(NULL, NULL, CAST(NULL AS INT), CAST(NULL AS INT)):int> -- !query output NULL + + +-- !query +select normalize('hello') +-- !query schema +struct<normalize(hello, NFC):string> +-- !query output +hello + + +-- !query +select normalize('hello', 'NFD') +-- !query schema +struct<normalize(hello, NFD):string> +-- !query output +hello + + +-- !query +select normalize('fi', 'NFKC') +-- !query schema +struct<normalize(fi, NFKC):string> +-- !query output +fi + + +-- !query +select normalize(null, 'NFC') +-- !query schema +struct<normalize(NULL, NFC):string> +-- !query output +NULL + + +-- !query +select normalize('hello', 'not_a_form') +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "INVALID_PARAMETER_VALUE.NORMALIZE_FORM", + "sqlState" : "22023", + "messageParameters" : { + "form" : "'not_a_form'", + "functionName" : "`normalize`", + "parameter" : "`form`" + } +} diff --git a/sql/core/src/test/resources/sql-tests/results/subquery/in-subquery/in-limit.sql.out b/sql/core/src/test/resources/sql-tests/results/subquery/in-subquery/in-limit.sql.out index d501c93973a31..67aa738ee36b9 100644 --- a/sql/core/src/test/resources/sql-tests/results/subquery/in-subquery/in-limit.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/subquery/in-subquery/in-limit.sql.out @@ -113,6 +113,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a LIMIT 10 OFFSET 2) LIMIT 2 @@ -145,6 +146,7 @@ FROM t1 WHERE t1a IN (SELECT t2a FROM t2 WHERE t1d = t2d + ORDER BY t2a OFFSET 2) OFFSET 1 -- !query schema @@ -161,6 +163,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4 -- !query schema @@ -168,8 +171,6 @@ struct<t1a:string,t1b:smallint,t1c:int,t1d:bigint,t1e:float,t1f:double,t1g:decim -- !query output val1a 16 12 10 15.0 20.0 2000 2014-07-04 01:01:00 2014-07-04 val1a 16 12 21 15.0 20.0 2000 2014-06-04 01:02:00.001 2014-06-04 -val1b 8 16 19 17.0 25.0 2600 2014-05-04 01:01:00 2014-05-04 -val1c 8 16 19 17.0 25.0 2600 2014-05-04 01:02:00.001 2014-05-05 -- !query @@ -195,6 +196,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b <= t1d + ORDER BY t2c NULLS LAST LIMIT 2) LIMIT 4 -- !query schema @@ -210,6 +212,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2) -- !query schema struct<t1a:string,t1b:smallint,t1c:int,t1d:bigint,t1e:float,t1f:double,t1g:decimal(4,0),t1h:timestamp,t1i:date> @@ -218,6 +221,10 @@ val1a 16 12 10 15.0 20.0 2000 2014-07-04 01:01:00 2014-07-04 val1a 16 12 21 15.0 20.0 2000 2014-06-04 01:02:00.001 2014-06-04 val1a 6 8 10 15.0 20.0 2000 2014-04-04 01:00:00 2014-04-04 val1a 6 8 10 15.0 20.0 2000 2014-04-04 01:02:00.001 2014-04-04 +val1d 10 NULL 12 17.0 25.0 2600 2015-05-04 01:01:00 2015-05-04 +val1e 10 NULL 19 17.0 25.0 2600 2014-05-04 01:01:00 2014-05-04 +val1e 10 NULL 19 17.0 25.0 2600 2014-09-04 01:02:00.001 2014-09-04 +val1e 10 NULL 25 17.0 25.0 2600 2014-08-04 01:01:00 2014-08-04 -- !query @@ -292,6 +299,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c NULLS LAST LIMIT 2 OFFSET 2) LIMIT 4 @@ -327,6 +335,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b LIMIT 2 OFFSET 2) -- !query schema @@ -348,14 +357,17 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b <= t1d + ORDER BY t2b LIMIT 2) -- !query schema struct<t1a:string,t1b:smallint,t1c:int,t1d:bigint,t1e:float,t1f:double,t1g:decimal(4,0),t1h:timestamp,t1i:date> -- !query output val1a 16 12 10 15.0 20.0 2000 2014-07-04 01:01:00 2014-07-04 val1a 16 12 21 15.0 20.0 2000 2014-06-04 01:02:00.001 2014-06-04 -val1b 8 16 19 17.0 25.0 2600 2014-05-04 01:01:00 2014-05-04 -val1c 8 16 19 17.0 25.0 2600 2014-05-04 01:02:00.001 2014-05-05 +val1d 10 NULL 12 17.0 25.0 2600 2015-05-04 01:01:00 2015-05-04 +val1e 10 NULL 19 17.0 25.0 2600 2014-05-04 01:01:00 2014-05-04 +val1e 10 NULL 19 17.0 25.0 2600 2014-09-04 01:02:00.001 2014-09-04 +val1e 10 NULL 25 17.0 25.0 2600 2014-08-04 01:01:00 2014-08-04 -- !query @@ -456,6 +468,7 @@ FROM t1 WHERE t1c IN (SELECT t2c FROM t2 WHERE t2b >= 8 + ORDER BY t2c DESC NULLS LAST OFFSET 2) OFFSET 4 -- !query schema @@ -504,6 +517,7 @@ FROM t1 WHERE t1b NOT IN (SELECT t2b FROM t2 WHERE t2b > 6 + ORDER BY t2b OFFSET 2) -- !query schema struct<t1a:string,t1b:smallint,t1c:int,t1d:bigint,t1e:float,t1f:double,t1g:decimal(4,0),t1h:timestamp,t1i:date> diff --git a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out index b37fe614e2376..9b1322dcb2e81 100644 --- a/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/subquery/scalar-subquery/scalar-subquery-predicate.sql.out @@ -598,7 +598,7 @@ WHERE t1c = (SELECT t2c -- !query schema struct<t1a:string,t1b:smallint> -- !query output -val1a 16 + -- !query @@ -625,7 +625,7 @@ WHERE t1c = (SELECT DISTINCT t2c -- !query schema struct<t1a:string,t1b:smallint> -- !query output -val1a 16 + -- !query diff --git a/sql/core/src/test/resources/sql-tests/results/time.sql.out b/sql/core/src/test/resources/sql-tests/results/time.sql.out index 58e72aec919e9..5bb6fb64c7485 100644 --- a/sql/core/src/test/resources/sql-tests/results/time.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/time.sql.out @@ -2358,7 +2358,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"TIME '12:30:41.123'\"", "inputType" : "\"TIME(6)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIME '12:30:41.123' - TIMESTAMP '2025-07-11 10:00:01')\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/results/timestamp-ltz-nanos.sql.out b/sql/core/src/test/resources/sql-tests/results/timestamp-ltz-nanos.sql.out index 388ce96f73967..5fae7cc5041e6 100644 --- a/sql/core/src/test/resources/sql-tests/results/timestamp-ltz-nanos.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/timestamp-ltz-nanos.sql.out @@ -678,6 +678,46 @@ struct<TIMESTAMP_LTZ '1960-01-01 19:04:05.123456789' + INTERVAL '0 00:00:00.0000 1960-01-01 19:04:05.123457789 +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' YEAR +-- !query schema +struct<TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR:timestamp_ltz(9)> +-- !query output +2021-01-01 19:04:05.123456789 + + +-- !query +SELECT INTERVAL '1' YEAR + TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' +-- !query schema +struct<TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' YEAR:timestamp_ltz(9)> +-- !query output +2021-01-01 19:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +-- !query schema +struct<TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH:timestamp_ltz(9)> +-- !query output +2020-02-01 19:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - INTERVAL '1-2' YEAR TO MONTH +-- !query schema +struct<TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH:timestamp_ltz(9)> +-- !query output +2018-11-01 19:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_LTZ '1960-01-31 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +-- !query schema +struct<TIMESTAMP_LTZ '1960-01-30 19:04:05.123456789' + INTERVAL '1' MONTH:timestamp_ltz(9)> +-- !query output +1960-02-29 19:04:05.123456789 + + -- !query SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + make_interval(0, 1, 0, 2, 0, 0, 0) -- !query schema @@ -705,25 +745,92 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query -SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' +-- !query schema +struct<(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111'):interval day to second> +-- !query output +1 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456001 UTC' +-- !query schema +struct<(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456001'):interval day to second> +-- !query output +0 00:00:00.000000000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 03:04:05.000000111 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' +-- !query schema +struct<(TIMESTAMP_LTZ '2019-12-31 19:04:05.000000111' - TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789'):interval day to second> +-- !query output +-1 00:00:00.123456000 + + +-- !query +SELECT ('2020-01-02 03:04:05.1234567 UTC' :: timestamp_ltz(7)) - ('2020-01-01 03:04:05.000000009 UTC' :: timestamp_ltz(9)) +-- !query schema +struct<(CAST(2020-01-02 03:04:05.1234567 UTC AS TIMESTAMP_LTZ(7)) - CAST(2020-01-01 03:04:05.000000009 UTC AS TIMESTAMP_LTZ(9))):interval day to second> +-- !query output +1 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - TIMESTAMP_LTZ '2020-01-02 03:04:05 UTC' +-- !query schema +struct<(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - TIMESTAMP '2020-01-01 19:04:05'):interval day to second> +-- !query output +0 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01' +-- !query schema +struct<(TIMESTAMP_LTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01'):interval day to second> +-- !query output +1 00:00:00.000000000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC' - TIMESTAMP_LTZ '1960-01-01 00:00:00.000000999 UTC' +-- !query schema +struct<(TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789' - TIMESTAMP_LTZ '1959-12-31 16:00:00.000000999'):interval day to second> +-- !query output +21915 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' - CAST(NULL AS timestamp_ltz(9)) +-- !query schema +struct<(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' - CAST(NULL AS TIMESTAMP_LTZ(9))):interval day to second> +-- !query output +NULL + + +-- !query +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.123456789 UTC' :: timestamp_ltz(9)) -- !query schema struct<> -- !query output org.apache.spark.sql.catalyst.ExtendedAnalysisException { - "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES", + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", "sqlState" : "42K09", "messageParameters" : { - "left" : "\"TIMESTAMP_LTZ(9)\"", - "right" : "\"INTERVAL MONTH\"", - "sqlExpr" : "\"(TIMESTAMP_LTZ '2020-01-01 19:04:05.123456789' + INTERVAL '1' MONTH)\"" + "inputSql" : "\"CAST(2022-03-27 03:00:00.123456789 UTC AS TIMESTAMP_LTZ(9))\"", + "inputType" : "\"TIMESTAMP_LTZ(9)\"", + "paramIndex" : "third", + "requiredType" : "\"(TIMESTAMP_NTZ OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\"", + "sqlExpr" : "\"convert_timezone(Europe/Brussels, Europe/Moscow, CAST(2022-03-27 03:00:00.123456789 UTC AS TIMESTAMP_LTZ(9)))\"" }, "queryContext" : [ { "objectType" : "", "objectName" : "", "startIndex" : 8, - "stopIndex" : 77, - "fragment" : "TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789 UTC' + INTERVAL '1' MONTH" + "stopIndex" : 120, + "fragment" : "convert_timezone('Europe/Brussels', 'Europe/Moscow',\n '2022-03-27 03:00:00.123456789 UTC' :: timestamp_ltz(9))" } ] } @@ -752,6 +859,56 @@ struct<c:timestamp_ltz(9),count(1):bigint> 2019-12-31 16:00:00.000000999 1 +-- !query +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3), + (CAST(NULL AS timestamp_ltz(9)), 4), + (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k +-- !query schema +struct<k:timestamp_ltz(9),count(1):bigint,sum(v):bigint> +-- !query output +NULL 2 9 +2019-12-31 16:00:00.000000001 2 3 +2019-12-31 16:00:00.000000999 1 3 + + +-- !query +SELECT mode(c) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c) +-- !query schema +struct<mode(c):timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000001 + + +-- !query +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c) +-- !query schema +struct<sort_array(collect_set(c), true):array<timestamp_ltz(9)>> +-- !query output +[2019-12-31 16:00:00.000000001,2019-12-31 16:00:00.000000999] + + +-- !query +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c) +-- !query schema +struct<sort_array(collect_list(c), true):array<timestamp_ltz(9)>> +-- !query output +[2019-12-31 16:00:00.000000001,2019-12-31 16:00:00.000000001,2019-12-31 16:00:00.000000999] + + -- !query SELECT unix_timestamp(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789') -- !query schema @@ -820,6 +977,21 @@ struct<max_by(v, k):timestamp_ltz(9),min_by(v, k):timestamp_ltz(9)> 2019-12-31 16:00:00.000000999 2019-12-31 16:00:00.000000001 +-- !query +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), + (CAST(NULL AS timestamp_ltz(9))) AS t(c) + ORDER BY c +-- !query schema +struct<c:timestamp_ltz(9)> +-- !query output +NULL +2019-12-31 16:00:00.000000001 +2019-12-31 16:00:00.000000999 + + -- !query SELECT unix_nanos(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789 UTC') -- !query schema @@ -1127,6 +1299,70 @@ struct<v:timestamp_ltz(9),next_v:timestamp_ltz(9)> 2020-01-01 00:00:00.0000009 NULL +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query schema +struct<(c = 2020-01-02 03:04:05.123456789):boolean,(c = 2020-01-02 03:04:05.123456788):boolean,(c < 2020-01-02 03:04:05.123456790):boolean> +-- !query output +true false true + + +-- !query +SELECT c FROM VALUES + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_LTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005' +-- !query schema +struct<c:timestamp_ltz(9)> +-- !query output +2020-01-02 03:04:05.000000001 + + +-- !query +SET spark.sql.ansi.enabled=false +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.ansi.enabled false + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.legacy.typeCoercion.datetimeToString.enabled true + + +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_LTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query schema +struct<(c = 2020-01-02 03:04:05.123456789):boolean,(c < 2020-01-02 03:04:05.123456790):boolean> +-- !query output +true true + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.legacy.typeCoercion.datetimeToString.enabled false + + +-- !query +SET spark.sql.ansi.enabled=true +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.ansi.enabled true + + -- !query SELECT unix_seconds(TIMESTAMP_LTZ '2020-01-01 13:24:35.123456789 UTC') -- !query schema @@ -1285,3 +1521,231 @@ SELECT date_trunc('NANOSECOND', TIMESTAMP_LTZ '2020-01-01 12:34:56.123456789 UTC struct<date_trunc(NANOSECOND, TIMESTAMP_LTZ '2020-01-01 04:34:56.123456789'):timestamp_ltz(9)> -- !query output NULL + + +-- !query +SELECT typeof(current_timestamp(9)), typeof(current_timestamp(8)), typeof(current_timestamp(7)) +-- !query schema +struct<typeof(current_timestamp()):string,typeof(current_timestamp()):string,typeof(current_timestamp()):string> +-- !query output +timestamp_ltz(9) timestamp_ltz(8) timestamp_ltz(7) + + +-- !query +SELECT typeof(now(9)), typeof(now(6)) +-- !query schema +struct<typeof(current_timestamp()):string,typeof(now()):string> +-- !query output +timestamp_ltz(9) timestamp + + +-- !query +SELECT typeof(current_timestamp()), typeof(current_timestamp(6)) +-- !query schema +struct<typeof(current_timestamp()):string,typeof(current_timestamp()):string> +-- !query output +timestamp timestamp + + +-- !query +SELECT typeof(current_timestamp(7 + 2)) +-- !query schema +struct<typeof(current_timestamp()):string> +-- !query output +timestamp_ltz(9) + + +-- !query +SELECT current_timestamp(9) = current_timestamp(9), now(9) = current_timestamp(9) +-- !query schema +struct<(current_timestamp() = current_timestamp()):boolean,(current_timestamp() = current_timestamp()):boolean> +-- !query output +true true + + +-- !query +SELECT current_timestamp(3) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "3", + "type" : "TIMESTAMP_LTZ" + } +} + + +-- !query +SELECT current_timestamp(10) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "10", + "type" : "TIMESTAMP_LTZ" + } +} + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + INTERSECT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c +-- !query schema +struct<c:timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000001 + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' AS c + UNION ALL SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') + EXCEPT SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' ORDER BY c +-- !query schema +struct<c:timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000999 + + +-- !query +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009 UTC' :: timestamp_ltz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900 UTC' :: timestamp_ltz(9))) ORDER BY c +-- !query schema +struct<typeof(c):string,c:timestamp_ltz(9)> +-- !query output +timestamp_ltz(9) 2019-12-31 16:00:00.0000009 + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000500 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC' +-- !query schema +struct<between(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000500', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'):boolean> +-- !query output +true + + +-- !query +SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000001000 UTC' + BETWEEN TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC' + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC' +-- !query schema +struct<between(TIMESTAMP_LTZ '2019-12-31 16:00:00.000001000', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'):boolean> +-- !query output +false + + +-- !query +SELECT '2020-01-01 00:00:00.000000500 UTC' :: timestamp_ltz(9) + BETWEEN '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7) + AND TIMESTAMP_LTZ '2020-01-01 00:00:00.000001 UTC' +-- !query schema +struct<between(CAST(2020-01-01 00:00:00.000000500 UTC AS TIMESTAMP_LTZ(9)), CAST(2020-01-01 00:00:00.0000001 UTC AS TIMESTAMP_LTZ(7)), TIMESTAMP '2019-12-31 16:00:00.000001'):boolean> +-- !query output +true + + +-- !query +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001 UTC' :: timestamp_ltz(7), + TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC') AS v) +-- !query schema +struct<typeof(v):string,v:timestamp_ltz(9)> +-- !query output +timestamp_ltz(9) 2019-12-31 16:00:00.0000001 + + +-- !query +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ltz(9)), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS v) +-- !query schema +struct<typeof(v):string,v:timestamp_ltz(9)> +-- !query output +timestamp_ltz(9) 2019-12-31 16:00:00.000000999 + + +-- !query +SELECT ifnull(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', CAST(NULL AS timestamp_ltz(9))) +-- !query schema +struct<ifnull(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', CAST(NULL AS TIMESTAMP_LTZ(9))):timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000001 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'), + (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') AS t(k) + WHERE k IN (SELECT TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC') ORDER BY k +-- !query schema +struct<k:timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000999 + + +-- !query +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'))) ORDER BY col +-- !query schema +struct<typeof(col):string,col:timestamp_ltz(9)> +-- !query output +timestamp_ltz(9) 2019-12-31 16:00:00.000000001 +timestamp_ltz(9) 2019-12-31 16:00:00.000000999 + + +-- !query +SELECT element_at(array( + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'), 2) +-- !query schema +struct<element_at(array(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999'), 2):timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.000000999 + + +-- !query +SELECT (named_struct('f', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')).f +-- !query schema +struct<named_struct(f, TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789').f:timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.123456789 + + +-- !query +SELECT map('k', TIMESTAMP_LTZ '2020-01-01 00:00:00.123456789 UTC')['k'] +-- !query schema +struct<map(k, TIMESTAMP_LTZ '2019-12-31 16:00:00.123456789')[k]:timestamp_ltz(9)> +-- !query output +2019-12-31 16:00:00.123456789 + + +-- !query +SELECT map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b')[ + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'] +-- !query schema +struct<map(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', a, TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999', b)[TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999']:string> +-- !query output +b + + +-- !query +SELECT element_at(map(TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 'a', + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 'b'), + TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') +-- !query schema +struct<element_at(map(TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001', a, TIMESTAMP_LTZ '2019-12-31 16:00:00.000000999', b), TIMESTAMP_LTZ '2019-12-31 16:00:00.000000001'):string> +-- !query output +a diff --git a/sql/core/src/test/resources/sql-tests/results/timestamp-ntz-nanos.sql.out b/sql/core/src/test/resources/sql-tests/results/timestamp-ntz-nanos.sql.out index 821aad995bc65..46fb523e378aa 100644 --- a/sql/core/src/test/resources/sql-tests/results/timestamp-ntz-nanos.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/timestamp-ntz-nanos.sql.out @@ -596,6 +596,46 @@ struct<TIMESTAMP_NTZ '1960-01-02 03:04:05.123456789' + INTERVAL '0 00:00:00.0000 1960-01-02 03:04:05.123457789 +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR +-- !query schema +struct<TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR:timestamp_ntz(9)> +-- !query output +2021-01-02 03:04:05.123456789 + + +-- !query +SELECT INTERVAL '1' YEAR + TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' +-- !query schema +struct<TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' YEAR:timestamp_ntz(9)> +-- !query output +2021-01-02 03:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH +-- !query schema +struct<TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH:timestamp_ntz(9)> +-- !query output +2020-02-02 03:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH +-- !query schema +struct<TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - INTERVAL '1-2' YEAR TO MONTH:timestamp_ntz(9)> +-- !query output +2018-11-02 03:04:05.123456789 + + +-- !query +SELECT TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH +-- !query schema +struct<TIMESTAMP_NTZ '1960-01-31 03:04:05.123456789' + INTERVAL '1' MONTH:timestamp_ntz(9)> +-- !query output +1960-02-29 03:04:05.123456789 + + -- !query SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + make_interval(0, 1, 0, 2, 0, 0, 0) -- !query schema @@ -623,27 +663,93 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException -- !query -SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' -- !query schema -struct<> +struct<(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111'):interval day to second> -- !query output -org.apache.spark.sql.catalyst.ExtendedAnalysisException -{ - "errorClass" : "DATATYPE_MISMATCH.BINARY_OP_DIFF_TYPES", - "sqlState" : "42K09", - "messageParameters" : { - "left" : "\"TIMESTAMP_NTZ(9)\"", - "right" : "\"INTERVAL MONTH\"", - "sqlExpr" : "\"(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH)\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 8, - "stopIndex" : 73, - "fragment" : "TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' + INTERVAL '1' MONTH" - } ] -} +1 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456001' +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456001'):interval day to second> +-- !query output +0 00:00:00.000000000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-01 03:04:05.000000111' - TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789'):interval day to second> +-- !query output +-1 00:00:00.123456000 + + +-- !query +SELECT ('2020-01-02 03:04:05.1234567' :: timestamp_ntz(7)) - ('2020-01-01 03:04:05.000000009' :: timestamp_ntz(9)) +-- !query schema +struct<(CAST(2020-01-02 03:04:05.1234567 AS TIMESTAMP_NTZ(7)) - CAST(2020-01-01 03:04:05.000000009 AS TIMESTAMP_NTZ(9))):interval day to second> +-- !query output +1 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05' +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - TIMESTAMP_NTZ '2020-01-02 03:04:05'):interval day to second> +-- !query output +0 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01' +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-02 00:00:00.000000789' - DATE '2020-01-01'):interval day to second> +-- !query output +1 00:00:00.000000000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789' - TIMESTAMP_NTZ '1960-01-01 00:00:00.000000999' +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789' - TIMESTAMP_NTZ '1960-01-01 00:00:00.000000999'):interval day to second> +-- !query output +21915 00:00:00.123456000 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - CAST(NULL AS timestamp_ntz(9)) +-- !query schema +struct<(TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789' - CAST(NULL AS TIMESTAMP_NTZ(9))):interval day to second> +-- !query output +NULL + + +-- !query +SELECT convert_timezone('Europe/Brussels', 'Europe/Moscow', + TIMESTAMP_NTZ '2022-03-27 03:00:00.123456789') +-- !query schema +struct<convert_timezone(Europe/Brussels, Europe/Moscow, TIMESTAMP_NTZ '2022-03-27 03:00:00.123456789'):timestamp_ntz(9)> +-- !query output +2022-03-27 04:00:00.123456789 + + +-- !query +SELECT typeof(convert_timezone('Europe/Brussels', 'Europe/Moscow', + '2022-03-27 03:00:00.1234567' :: timestamp_ntz(7))) +-- !query schema +struct<typeof(convert_timezone(Europe/Brussels, Europe/Moscow, CAST(2022-03-27 03:00:00.1234567 AS TIMESTAMP_NTZ(7)))):string> +-- !query output +timestamp_ntz(7) + + +-- !query +SELECT convert_timezone('America/Los_Angeles', 'UTC', CAST(NULL AS timestamp_ntz(9))) +-- !query schema +struct<convert_timezone(America/Los_Angeles, UTC, CAST(NULL AS TIMESTAMP_NTZ(9))):timestamp_ntz(9)> +-- !query output +NULL -- !query @@ -670,6 +776,56 @@ struct<c:timestamp_ntz(9),count(1):bigint> 2020-01-01 00:00:00.000000999 1 +-- !query +SELECT k, count(*), sum(v) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3), + (CAST(NULL AS timestamp_ntz(9)), 4), + (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v) + GROUP BY k ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9),count(1):bigint,sum(v):bigint> +-- !query output +NULL 2 9 +2020-01-01 00:00:00.000000001 2 3 +2020-01-01 00:00:00.000000999 1 3 + + +-- !query +SELECT mode(c) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c) +-- !query schema +struct<mode(c):timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000001 + + +-- !query +SELECT sort_array(collect_set(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c) +-- !query schema +struct<sort_array(collect_set(c), true):array<timestamp_ntz(9)>> +-- !query output +[2020-01-01 00:00:00.000000001,2020-01-01 00:00:00.000000999] + + +-- !query +SELECT sort_array(collect_list(c)) FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c) +-- !query schema +struct<sort_array(collect_list(c), true):array<timestamp_ntz(9)>> +-- !query output +[2020-01-01 00:00:00.000000001,2020-01-01 00:00:00.000000001,2020-01-01 00:00:00.000000999] + + -- !query SELECT unix_timestamp(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query schema @@ -730,6 +886,21 @@ struct<max_by(v, k):timestamp_ntz(9),min_by(v, k):timestamp_ntz(9)> 2020-01-01 00:00:00.000000999 2020-01-01 00:00:00.000000001 +-- !query +SELECT DISTINCT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), + (CAST(NULL AS timestamp_ntz(9))) AS t(c) + ORDER BY c +-- !query schema +struct<c:timestamp_ntz(9)> +-- !query output +NULL +2020-01-01 00:00:00.000000001 +2020-01-01 00:00:00.000000999 + + -- !query SELECT unix_nanos(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query schema @@ -924,6 +1095,70 @@ struct<v:timestamp_ntz(9),next_v:timestamp_ntz(9)> 2020-01-01 00:00:00.0000009 NULL +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c = '2020-01-02 03:04:05.123456788', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query schema +struct<(c = 2020-01-02 03:04:05.123456789):boolean,(c = 2020-01-02 03:04:05.123456788):boolean,(c < 2020-01-02 03:04:05.123456790):boolean> +-- !query output +true false true + + +-- !query +SELECT c FROM VALUES + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000001'), + (TIMESTAMP_NTZ '2020-01-02 03:04:05.000000009') AS t(c) + WHERE c BETWEEN '2020-01-02 03:04:05.000000001' AND '2020-01-02 03:04:05.000000005' +-- !query schema +struct<c:timestamp_ntz(9)> +-- !query output +2020-01-02 03:04:05.000000001 + + +-- !query +SET spark.sql.ansi.enabled=false +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.ansi.enabled false + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=true +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.legacy.typeCoercion.datetimeToString.enabled true + + +-- !query +SELECT c = '2020-01-02 03:04:05.123456789', + c < '2020-01-02 03:04:05.123456790' + FROM VALUES (TIMESTAMP_NTZ '2020-01-02 03:04:05.123456789') AS t(c) +-- !query schema +struct<(c = 2020-01-02 03:04:05.123456789):boolean,(c < 2020-01-02 03:04:05.123456790):boolean> +-- !query output +true true + + +-- !query +SET spark.sql.legacy.typeCoercion.datetimeToString.enabled=false +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.legacy.typeCoercion.datetimeToString.enabled false + + +-- !query +SET spark.sql.ansi.enabled=true +-- !query schema +struct<key:string,value:string> +-- !query output +spark.sql.ansi.enabled true + + -- !query SELECT unix_seconds(TIMESTAMP_NTZ '2020-01-01 13:24:35.123456789') -- !query schema @@ -1074,3 +1309,317 @@ SELECT date_trunc('NANOSECOND', TIMESTAMP_NTZ '2020-01-01 12:34:56.123456789') struct<date_trunc(NANOSECOND, TIMESTAMP_NTZ '2020-01-01 12:34:56.123456789'):timestamp_ntz(9)> -- !query output NULL + + +-- !query +SELECT typeof(localtimestamp(9)), typeof(localtimestamp(8)), typeof(localtimestamp(7)) +-- !query schema +struct<typeof(localtimestamp()):string,typeof(localtimestamp()):string,typeof(localtimestamp()):string> +-- !query output +timestamp_ntz(9) timestamp_ntz(8) timestamp_ntz(7) + + +-- !query +SELECT typeof(localtimestamp()), typeof(localtimestamp(6)) +-- !query schema +struct<typeof(localtimestamp()):string,typeof(localtimestamp()):string> +-- !query output +timestamp_ntz timestamp_ntz + + +-- !query +SELECT typeof(localtimestamp(8 + 1)) +-- !query schema +struct<typeof(localtimestamp()):string> +-- !query output +timestamp_ntz(9) + + +-- !query +SELECT localtimestamp(9) = localtimestamp(9) +-- !query schema +struct<(localtimestamp() = localtimestamp()):boolean> +-- !query output +true + + +-- !query +SELECT localtimestamp(3) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "3", + "type" : "TIMESTAMP_NTZ" + } +} + + +-- !query +SELECT localtimestamp(10) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkException +{ + "errorClass" : "INVALID_TIMESTAMP_PRECISION", + "sqlState" : "22023", + "messageParameters" : { + "precision" : "10", + "type" : "TIMESTAMP_NTZ" + } +} + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + INTERSECT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c +-- !query schema +struct<c:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000001 + + +-- !query +SELECT c FROM (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' AS c + UNION ALL SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') + EXCEPT SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' ORDER BY c +-- !query schema +struct<c:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT typeof(c), c FROM ( + (SELECT '2020-01-01 00:00:00.0000009' :: timestamp_ntz(7) AS c) + INTERSECT (SELECT '2020-01-01 00:00:00.000000900' :: timestamp_ntz(9))) ORDER BY c +-- !query schema +struct<typeof(c):string,c:timestamp_ntz(9)> +-- !query output +timestamp_ntz(9) 2020-01-01 00:00:00.0000009 + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' +-- !query schema +struct<between(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000500', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'):boolean> +-- !query output +true + + +-- !query +SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000001000' + BETWEEN TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001' + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' +-- !query schema +struct<between(TIMESTAMP_NTZ '2020-01-01 00:00:00.000001000', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'):boolean> +-- !query output +false + + +-- !query +SELECT '2020-01-01 00:00:00.000000500' :: timestamp_ntz(9) + BETWEEN '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7) + AND TIMESTAMP_NTZ '2020-01-01 00:00:00.000001' +-- !query schema +struct<between(CAST(2020-01-01 00:00:00.000000500 AS TIMESTAMP_NTZ(9)), CAST(2020-01-01 00:00:00.0000001 AS TIMESTAMP_NTZ(7)), TIMESTAMP_NTZ '2020-01-01 00:00:00.000001'):boolean> +-- !query output +true + + +-- !query +SELECT typeof(v), v FROM (SELECT if(true, + '2020-01-01 00:00:00.0000001' :: timestamp_ntz(7), + TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789') AS v) +-- !query schema +struct<typeof(v):string,v:timestamp_ntz(9)> +-- !query output +timestamp_ntz(9) 2020-01-01 00:00:00.0000001 + + +-- !query +SELECT typeof(v), v FROM (SELECT nvl( + CAST(NULL AS timestamp_ntz(9)), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS v) +-- !query schema +struct<typeof(v):string,v:timestamp_ntz(9)> +-- !query output +timestamp_ntz(9) 2020-01-01 00:00:00.000000999 + + +-- !query +SELECT ifnull(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', CAST(NULL AS timestamp_ntz(9))) +-- !query schema +struct<ifnull(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', CAST(NULL AS TIMESTAMP_NTZ(9))):timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000001 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') +-- !query schema +struct<scalarsubquery():timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT typeof((SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')) +-- !query schema +struct<typeof(scalarsubquery()):string> +-- !query output +timestamp_ntz(9) + + +-- !query +SELECT typeof((SELECT CAST(NULL AS timestamp_ntz(9)))) +-- !query schema +struct<typeof(scalarsubquery()):string> +-- !query output +timestamp_ntz(9) + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k = (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS s(v) + WHERE s.v = t.k) ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE NOT EXISTS (SELECT 1 FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS s(v) + WHERE s.v = t.k) ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000001 + + +-- !query +SELECT k FROM VALUES + ('2020-01-01 00:00:00.0000009' :: timestamp_ntz(7)) AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') ORDER BY k +-- !query schema +struct<k:timestamp_ntz(7)> +-- !query output +2020-01-01 00:00:00.0000009 + + +-- !query +SELECT k FROM VALUES + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'), + (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS t(k) + WHERE k NOT IN (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999' + UNION ALL SELECT CAST(NULL AS timestamp_ntz(9))) ORDER BY k +-- !query schema +struct<k:timestamp_ntz(9)> +-- !query output + + + +-- !query +SELECT typeof(col), col FROM (SELECT explode(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'))) ORDER BY col +-- !query schema +struct<typeof(col):string,col:timestamp_ntz(9)> +-- !query output +timestamp_ntz(9) 2020-01-01 00:00:00.000000001 +timestamp_ntz(9) 2020-01-01 00:00:00.000000999 + + +-- !query +SELECT element_at(array( + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), 2) +-- !query schema +struct<element_at(array(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'), 2):timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.000000999 + + +-- !query +SELECT (named_struct('f', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')).f +-- !query schema +struct<named_struct(f, TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789').f:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.123456789 + + +-- !query +SELECT map('k', TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')['k'] +-- !query schema +struct<map(k, TIMESTAMP_NTZ '2020-01-01 00:00:00.123456789')[k]:timestamp_ntz(9)> +-- !query output +2020-01-01 00:00:00.123456789 + + +-- !query +SELECT map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b')[ + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'] +-- !query schema +struct<map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', a, TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', b)[TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999']:string> +-- !query output +b + + +-- !query +SELECT element_at(map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 'a', + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 'b'), + TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') +-- !query schema +struct<element_at(map(TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', a, TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', b), TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'):string> +-- !query output +a diff --git a/sql/core/src/test/resources/sql-tests/results/timestampNTZ/timestamp.sql.out b/sql/core/src/test/resources/sql-tests/results/timestampNTZ/timestamp.sql.out index 904cffe03b447..c0a154facec2b 100644 --- a/sql/core/src/test/resources/sql-tests/results/timestampNTZ/timestamp.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/timestampNTZ/timestamp.sql.out @@ -1035,7 +1035,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:10\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP_NTZ '2011-11-11 11:11:11' - 2011-11-11 11:11:10)\"" }, "queryContext" : [ { @@ -1061,7 +1061,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"2011-11-11 11:11:11\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(2011-11-11 11:11:11 - TIMESTAMP_NTZ '2011-11-11 11:11:10')\"" }, "queryContext" : [ { @@ -1111,7 +1111,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(str - TIMESTAMP_NTZ '2011-11-11 11:11:11')\"" }, "queryContext" : [ { @@ -1137,7 +1137,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"str\"", "inputType" : "\"STRING\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(TIMESTAMP_NTZ '2011-11-11 11:11:11' - str)\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/results/typeCoercion/native/decimalPrecision.sql.out b/sql/core/src/test/resources/sql-tests/results/typeCoercion/native/decimalPrecision.sql.out index 54e26851ba57e..d0a62b06e1c4b 100644 --- a/sql/core/src/test/resources/sql-tests/results/typeCoercion/native/decimalPrecision.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/typeCoercion/native/decimalPrecision.sql.out @@ -1508,7 +1508,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(3,0))\"", "inputType" : "\"DECIMAL(3,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(3,0)))\"" }, "queryContext" : [ { @@ -1534,7 +1534,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(5,0))\"", "inputType" : "\"DECIMAL(5,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(5,0)))\"" }, "queryContext" : [ { @@ -1560,7 +1560,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(10,0))\"", "inputType" : "\"DECIMAL(10,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(10,0)))\"" }, "queryContext" : [ { @@ -1586,7 +1586,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(20,0))\"", "inputType" : "\"DECIMAL(20,0)\"", "paramIndex" : "second", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(2017-12-11 09:30:00.0 AS TIMESTAMP) - CAST(1 AS DECIMAL(20,0)))\"" }, "queryContext" : [ { @@ -2164,7 +2164,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(3,0))\"", "inputType" : "\"DECIMAL(3,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(3,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2190,7 +2190,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(5,0))\"", "inputType" : "\"DECIMAL(5,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(5,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2216,7 +2216,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(10,0))\"", "inputType" : "\"DECIMAL(10,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(10,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { @@ -2242,7 +2242,7 @@ org.apache.spark.sql.catalyst.ExtendedAnalysisException "inputSql" : "\"CAST(1 AS DECIMAL(20,0))\"", "inputType" : "\"DECIMAL(20,0)\"", "paramIndex" : "first", - "requiredType" : "\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\"", + "requiredType" : "(\"(TIMESTAMP OR TIMESTAMP WITHOUT TIME ZONE)\" or \"(TIMESTAMP_LTZ(P) OR TIMESTAMP_NTZ(P) WITH P IN [7, 9])\")", "sqlExpr" : "\"(CAST(1 AS DECIMAL(20,0)) - CAST(2017-12-11 09:30:00.0 AS TIMESTAMP))\"" }, "queryContext" : [ { diff --git a/sql/core/src/test/resources/sql-tests/results/udf/udf-window.sql.out b/sql/core/src/test/resources/sql-tests/results/udf/udf-window.sql.out index 20bb65d1d3e31..4d60e6003d6a4 100644 --- a/sql/core/src/test/resources/sql-tests/results/udf/udf-window.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/udf/udf-window.sql.out @@ -377,23 +377,17 @@ FROM testData WINDOW w AS (PARTITION BY udf(cate) ORDER BY udf(val)) ORDER BY cate, udf(val) -- !query schema -struct<> +struct<udf(val):int,cate:string,max:int,min:int,min:int,count:bigint,sum:bigint,avg:double,stddev:double,first_value:int,first_value_ignore_null:int,first_value_contain_null:int,any_value:int,any_value_ignore_null:int,any_value_contain_null:int,last_value:int,last_value_ignore_null:int,last_value_contain_null:int,rank:int,dense_rank:int,cume_dist:double,percent_rank:double,ntile:int,row_number:int,var_pop:double,var_samp:double,approx_count_distinct:bigint,covar_pop:double,corr:double,stddev_samp:double,stddev_pop:double,collect_list:array<int>,collect_set:array<int>,skewness:double,kurtosis:double> -- !query output -org.apache.spark.SparkArithmeticException -{ - "errorClass" : "DIVIDE_BY_ZERO", - "sqlState" : "22012", - "messageParameters" : { - "config" : "\"spark.sql.ansi.enabled\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 1126, - "stopIndex" : 1161, - "fragment" : "corr(udf(val), udf(val_long)) OVER w" - } ] -} +NULL NULL NULL NULL NULL 0 NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL 1 1 0.5 0.0 1 1 NULL NULL 0 NULL NULL NULL NULL [] [] NULL NULL +3 NULL 3 3 3 1 3 3.0 NULL NULL 3 NULL NULL 3 NULL 3 3 3 2 2 1.0 1.0 2 2 0.0 NULL 1 0.0 NULL NULL 0.0 [3] [3] NULL NULL +NULL a NULL NULL NULL 0 NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL 1 1 0.25 0.0 1 1 NULL NULL 0 NULL NULL NULL NULL [] [] NULL NULL +1 a 1 1 1 2 2 1.0 0.0 NULL 1 NULL NULL 1 NULL 1 1 1 2 2 0.75 0.3333333333333333 1 2 0.0 0.0 1 0.0 NULL 0.0 0.0 [1,1] [1] 0.7071067811865476 -1.5 +1 a 1 1 1 2 2 1.0 0.0 NULL 1 NULL NULL 1 NULL 1 1 1 2 2 0.75 0.3333333333333333 2 3 0.0 0.0 1 0.0 NULL 0.0 0.0 [1,1] [1] 0.7071067811865476 -1.5 +2 a 2 1 1 3 4 1.3333333333333333 0.5773502691896258 NULL 1 NULL NULL 1 NULL 2 2 2 4 3 1.0 1.0 2 4 0.22222222222222224 0.33333333333333337 2 4.772185885555555E8 1.0 0.5773502691896258 0.4714045207910317 [1,1,2] [1,2] 1.1539890888012805 -0.6672217220327235 +1 b 1 1 1 1 1 1.0 NULL 1 1 1 1 1 1 1 1 1 1 1 0.3333333333333333 0.0 1 1 0.0 NULL 1 NULL NULL NULL 0.0 [1] [1] NULL NULL +2 b 2 1 1 2 3 1.5 0.7071067811865476 1 1 1 1 1 1 2 2 2 2 2 0.6666666666666666 0.5 1 2 0.25 0.5 2 0.0 NULL 0.7071067811865476 0.5 [1,2] [1,2] 0.0 -2.0000000000000013 +3 b 3 1 1 3 6 2.0 1.0 1 1 1 1 1 1 3 3 3 3 3 1.0 1.0 2 3 0.6666666666666666 1.0 3 5.3687091175E8 1.0 1.0 0.816496580927726 [1,2,3] [1,2,3] 0.7057890433107311 -1.4999999999999984 -- !query diff --git a/sql/core/src/test/resources/sql-tests/results/unnest.sql.out b/sql/core/src/test/resources/sql-tests/results/unnest.sql.out new file mode 100644 index 0000000000000..e04dc7ecc531d --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/unnest.sql.out @@ -0,0 +1,267 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +CREATE OR REPLACE TEMPORARY VIEW nested AS SELECT * FROM VALUES +(1, array(10, 20, 30), array('a', 'b')), +(2, array(40), array('c', 'd', 'e')), +(3, array(), array()), +(4, cast(null as array<int>), array('f')) +AS nested(id, xs, ys) +-- !query schema +struct<> +-- !query output + + + +-- !query +SELECT * FROM UNNEST(array(10, 20, 30)) +-- !query schema +struct<col:int> +-- !query output +10 +20 +30 + + +-- !query +SELECT v FROM UNNEST(array(10, 20, 30)) AS t(v) +-- !query schema +struct<v:int> +-- !query output +10 +20 +30 + + +-- !query +SELECT * FROM UNNEST(array()) +-- !query schema +struct<col:void> +-- !query output + + + +-- !query +SELECT * FROM UNNEST(cast(null as array<int>)) +-- !query schema +struct<col:int> +-- !query output + + + +-- !query +SELECT * FROM UNNEST(array(10, 20, 30)) WITH ORDINALITY +-- !query schema +struct<col:int,ordinality:bigint> +-- !query output +10 1 +20 2 +30 3 + + +-- !query +SELECT val, pos FROM UNNEST(array('x', 'y')) WITH ORDINALITY AS t(val, pos) +-- !query schema +struct<val:string,pos:bigint> +-- !query output +x 1 +y 2 + + +-- !query +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) AS t(a, b) +-- !query schema +struct<a:int,b:int> +-- !query output +1 10 +2 20 +NULL 30 + + +-- !query +SELECT * FROM UNNEST(array(1, 2), array(10, 20, 30)) WITH ORDINALITY AS t(a, b, ord) +-- !query schema +struct<a:int,b:int,ord:bigint> +-- !query output +1 10 1 +2 20 2 +NULL 30 3 + + +-- !query +SELECT * FROM UNNEST(array(struct(1, 'a'), struct(2, 'b'))) AS t(s) +-- !query schema +struct<s:struct<col1:int,col2:string>> +-- !query output +{"col1":1,"col2":"a"} +{"col1":2,"col2":"b"} + + +-- !query +SELECT id, elem FROM nested, LATERAL UNNEST(xs) AS t(elem) ORDER BY id, elem +-- !query schema +struct<id:int,elem:int> +-- !query output +1 10 +1 20 +1 30 +2 40 + + +-- !query +SELECT id, x, y, ord +FROM nested, LATERAL UNNEST(xs, ys) WITH ORDINALITY AS t(x, y, ord) +ORDER BY id, ord +-- !query schema +struct<id:int,x:int,y:string,ord:bigint> +-- !query output +1 10 a 1 +1 20 b 2 +1 30 NULL 3 +2 40 c 1 +2 NULL d 2 +2 NULL e 3 +4 NULL f 1 + + +-- !query +SELECT id, elem +FROM nested LEFT JOIN LATERAL UNNEST(xs) AS t(elem) ON true +ORDER BY id, elem +-- !query schema +struct<id:int,elem:int> +-- !query output +1 10 +1 20 +1 30 +2 40 +3 NULL +4 NULL + + +-- !query +SELECT * FROM UNNEST(array(array(1, 2), array(3))) AS t(inner) +-- !query schema +struct<inner:array<int>> +-- !query output +[1,2] +[3] + + +-- !query +SELECT * FROM UNNEST(array(1, cast(null as int), 3)) WITH ORDINALITY +-- !query schema +struct<col:int,ordinality:bigint> +-- !query output +1 1 +3 3 +NULL 2 + + +-- !query +SELECT * FROM UNNEST(array(1), array(10, 20, 30)) AS t(a, b) +-- !query schema +struct<a:int,b:int> +-- !query output +1 10 +NULL 20 +NULL 30 + + +-- !query +SELECT * FROM UNNEST(42) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"42\"", + "inputType" : "\"INT\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(42)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 24, + "fragment" : "UNNEST(42)" + } ] +} + + +-- !query +SELECT * FROM UNNEST(map('a', 1)) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"map(a, 1)\"", + "inputType" : "\"MAP<STRING, INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(map(a, 1))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 33, + "fragment" : "UNNEST(map('a', 1))" + } ] +} + + +-- !query +SELECT * FROM UNNEST(array(1, 2), 3) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"3\"", + "inputType" : "\"INT\"", + "paramIndex" : "second", + "requiredType" : "\"ARRAY\"", + "sqlExpr" : "\"unnest(array(1, 2), 3)\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 36, + "fragment" : "UNNEST(array(1, 2), 3)" + } ] +} + + +-- !query +SELECT * FROM `unnest`(array(1, 2)) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.AnalysisException +{ + "errorClass" : "UNRESOLVABLE_TABLE_VALUED_FUNCTION", + "sqlState" : "42883", + "messageParameters" : { + "name" : "`unnest`" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 15, + "stopIndex" : 35, + "fragment" : "`unnest`(array(1, 2))" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/results/variant/variant-from-arrays-entries.sql.out b/sql/core/src/test/resources/sql-tests/results/variant/variant-from-arrays-entries.sql.out new file mode 100644 index 0000000000000..c51ef9d03a0b8 --- /dev/null +++ b/sql/core/src/test/resources/sql-tests/results/variant/variant-from-arrays-entries.sql.out @@ -0,0 +1,296 @@ +-- Automatically generated by SQLQueryTestSuite +-- !query +select cast(variant_from_arrays(array('z', 'a'), array(1, 2)) as string) +-- !query schema +struct<CAST(variant_from_arrays(array(z, a), array(1, 2)) AS STRING):string> +-- !query output +{"a":2,"z":1} + + +-- !query +select cast(variant_from_arrays(cast(array() as array<string>), cast(array() as array<int>)) as string) +-- !query schema +struct<CAST(variant_from_arrays(array(), array()) AS STRING):string> +-- !query output +{} + + +-- !query +select cast(variant_from_arrays(array('a', 'b'), array(1, cast(null as int))) as string) +-- !query schema +struct<CAST(variant_from_arrays(array(a, b), array(1, CAST(NULL AS INT))) AS STRING):string> +-- !query output +{"a":1,"b":null} + + +-- !query +select cast(variant_from_arrays(array('a'), array(array(1, 2, 3))) as string) +-- !query schema +struct<CAST(variant_from_arrays(array(a), array(array(1, 2, 3))) AS STRING):string> +-- !query output +{"a":[1,2,3]} + + +-- !query +select cast(variant_from_arrays(cast(null as array<string>), array(1)) as string) +-- !query schema +struct<CAST(variant_from_arrays(NULL, array(1)) AS STRING):string> +-- !query output +NULL + + +-- !query +select variant_from_arrays(array('a', cast(null as string)), array(1, 2)) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "NULL_MAP_KEY", + "sqlState" : "2200E" +} + + +-- !query +select variant_from_arrays(array('a', 'a'), array(1, 2)) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "VARIANT_DUPLICATE_KEY", + "sqlState" : "22023", + "messageParameters" : { + "key" : "a" + } +} + + +-- !query +select variant_from_arrays(array('a', 'b'), array(1)) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "_LEGACY_ERROR_TEMP_2128" +} + + +-- !query +select variant_from_arrays(array(1, 2), array('a', 'b')) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(1, 2)\"", + "inputType" : "\"ARRAY<INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY<STRING>\"", + "sqlExpr" : "\"variant_from_arrays(array(1, 2), array(a, b))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 56, + "fragment" : "variant_from_arrays(array(1, 2), array('a', 'b'))" + } ] +} + + +-- !query +select variant_from_arrays(array('a'), array(map(1, 2))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"variant_from_arrays(array(a), array(map(1, 2)))\"", + "srcType" : "\"MAP<INT, INT>\"", + "targetType" : "\"VARIANT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 56, + "fragment" : "variant_from_arrays(array('a'), array(map(1, 2)))" + } ] +} + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'b', 'v', 2))) as string) +-- !query schema +struct<CAST(variant_from_entries(array(named_struct(k, a, v, 1), named_struct(k, b, v, 2))) AS STRING):string> +-- !query output +{"a":1,"b":2} + + +-- !query +select cast(variant_from_entries(cast(array() as array<struct<k:string,v:int>>)) as string) +-- !query schema +struct<CAST(variant_from_entries(array()) AS STRING):string> +-- !query output +{} + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', cast(null as int)))) as string) +-- !query schema +struct<CAST(variant_from_entries(array(named_struct(k, a, v, CAST(NULL AS INT)))) AS STRING):string> +-- !query output +{"a":null} + + +-- !query +select cast(variant_from_entries(array(named_struct('k', 'a', 'v', 1), cast(null as struct<k:string,v:int>))) as string) +-- !query schema +struct<CAST(variant_from_entries(array(named_struct(k, a, v, 1), NULL)) AS STRING):string> +-- !query output +NULL + + +-- !query +select cast(variant_from_entries(cast(null as array<struct<k:string,v:int>>)) as string) +-- !query schema +struct<CAST(variant_from_entries(NULL) AS STRING):string> +-- !query output +NULL + + +-- !query +select variant_from_entries(array(named_struct('k', cast(null as string), 'v', 1))) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "NULL_MAP_KEY", + "sqlState" : "2200E" +} + + +-- !query +select variant_from_entries(array(named_struct('k', 'a', 'v', 1), named_struct('k', 'a', 'v', 2))) +-- !query schema +struct<> +-- !query output +org.apache.spark.SparkRuntimeException +{ + "errorClass" : "VARIANT_DUPLICATE_KEY", + "sqlState" : "22023", + "messageParameters" : { + "key" : "a" + } +} + + +-- !query +select variant_from_entries(array(1, 2)) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(1, 2)\"", + "inputType" : "\"ARRAY<INT>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(1, 2))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 40, + "fragment" : "variant_from_entries(array(1, 2))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 'a'))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(named_struct(k, a))\"", + "inputType" : "\"ARRAY<STRUCT<k: STRING NOT NULL>>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, a)))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 58, + "fragment" : "variant_from_entries(array(named_struct('k', 'a')))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 1, 'v', 'a'))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + "sqlState" : "42K09", + "messageParameters" : { + "inputSql" : "\"array(named_struct(k, 1, v, a))\"", + "inputType" : "\"ARRAY<STRUCT<k: INT NOT NULL, v: STRING NOT NULL>>\"", + "paramIndex" : "first", + "requiredType" : "\"ARRAY\" of pair \"STRUCT\"", + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, 1, v, a)))\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 66, + "fragment" : "variant_from_entries(array(named_struct('k', 1, 'v', 'a')))" + } ] +} + + +-- !query +select variant_from_entries(array(named_struct('k', 'a', 'v', map(1, 2)))) +-- !query schema +struct<> +-- !query output +org.apache.spark.sql.catalyst.ExtendedAnalysisException +{ + "errorClass" : "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + "sqlState" : "42K09", + "messageParameters" : { + "sqlExpr" : "\"variant_from_entries(array(named_struct(k, a, v, map(1, 2))))\"", + "srcType" : "\"MAP<INT, INT>\"", + "targetType" : "\"VARIANT\"" + }, + "queryContext" : [ { + "objectType" : "", + "objectName" : "", + "startIndex" : 8, + "stopIndex" : 74, + "fragment" : "variant_from_entries(array(named_struct('k', 'a', 'v', map(1, 2))))" + } ] +} diff --git a/sql/core/src/test/resources/sql-tests/results/vector-distance.sql.out b/sql/core/src/test/resources/sql-tests/results/vector-distance.sql.out index ccc39c23fb53c..c75be87615616 100644 --- a/sql/core/src/test/resources/sql-tests/results/vector-distance.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/vector-distance.sql.out @@ -4,7 +4,7 @@ SELECT vector_cosine_similarity(array(1.0F, 2.0F, 3.0F), array(4.0F, 5.0F, 6.0F) -- !query schema struct<vector_cosine_similarity(array(1.0, 2.0, 3.0), array(4.0, 5.0, 6.0)):float> -- !query output -0.9746319 +0.97463185 -- !query @@ -744,3 +744,75 @@ SELECT vector_l2_distance( struct<vector_l2_distance(array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), array(16.0, 15.0, 14.0, 13.0, 12.0, 11.0, 10.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0)):float> -- !query output 36.878178 + + +-- !query +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(3.0e19F, 4.0e19F)) +-- !query schema +struct<vector_cosine_similarity(array(3.0E19, 4.0E19), array(3.0E19, 4.0E19)):float> +-- !query output +1.0 + + +-- !query +SELECT vector_cosine_similarity(array(3.0e19F, 4.0e19F), array(-3.0e19F, -4.0e19F)) +-- !query schema +struct<vector_cosine_similarity(array(3.0E19, 4.0E19), array(-3.0E19, -4.0E19)):float> +-- !query output +-1.0 + + +-- !query +SELECT vector_inner_product(array(1.0e20F, 1.0e20F), array(1.0e20F, -1.0e20F)) +-- !query schema +struct<vector_inner_product(array(1.0E20, 1.0E20), array(1.0E20, -1.0E20)):float> +-- !query output +0.0 + + +-- !query +SELECT vector_l2_distance(array(3.0e19F, 4.0e19F), array(0.0F, 0.0F)) +-- !query schema +struct<vector_l2_distance(array(3.0E19, 4.0E19), array(0.0, 0.0)):float> +-- !query output +5.0E19 + + +-- !query +SELECT vector_cosine_similarity(array(1.0e-23F, 0.0F), array(1.0e-23F, 0.0F)) +-- !query schema +struct<vector_cosine_similarity(array(1.0E-23, 0.0), array(1.0E-23, 0.0)):float> +-- !query output +1.0 + + +-- !query +SELECT vector_cosine_similarity(array(1.0e-23F, 1.0e-23F), array(1.0e-23F, -1.0e-23F)) +-- !query schema +struct<vector_cosine_similarity(array(1.0E-23, 1.0E-23), array(1.0E-23, -1.0E-23)):float> +-- !query output +0.0 + + +-- !query +SELECT vector_cosine_similarity(array(float('inf'), 1.0F), array(1.0F, 1.0F)) +-- !query schema +struct<vector_cosine_similarity(array(inf, 1.0), array(1.0, 1.0)):float> +-- !query output +NaN + + +-- !query +SELECT vector_inner_product(array(float('inf'), 1.0F), array(1.0F, 1.0F)) +-- !query schema +struct<vector_inner_product(array(inf, 1.0), array(1.0, 1.0)):float> +-- !query output +Infinity + + +-- !query +SELECT vector_l2_distance(array(float('inf'), 1.0F), array(0.0F, 0.0F)) +-- !query schema +struct<vector_l2_distance(array(inf, 1.0), array(0.0, 0.0)):float> +-- !query output +Infinity diff --git a/sql/core/src/test/resources/sql-tests/results/vector-norm.sql.out b/sql/core/src/test/resources/sql-tests/results/vector-norm.sql.out index 8ebd583aa7e47..31a2aa2fb08ee 100644 --- a/sql/core/src/test/resources/sql-tests/results/vector-norm.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/vector-norm.sql.out @@ -668,7 +668,7 @@ SELECT vector_normalize( -- !query schema struct<vector_normalize(array(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), 2.0):array<float>> -- !query output -[0.025854385,0.05170877,0.07756316,0.10341754,0.12927192,0.15512632,0.1809807,0.20683508,0.23268947,0.25854385,0.28439823,0.31025264,0.33610702,0.3619614,0.38781577,0.41367015] +[0.025854385,0.05170877,0.07756315,0.10341754,0.12927192,0.1551263,0.1809807,0.20683508,0.23268946,0.25854385,0.28439823,0.3102526,0.336107,0.3619614,0.38781577,0.41367015] -- !query @@ -680,3 +680,75 @@ SELECT vector_norm( struct<vector_norm(array(1.0, 2.0, 3.0, 4.0, 5.0, CAST(NULL AS FLOAT), 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0), 2.0):float> -- !query output NULL + + +-- !query +SELECT vector_norm(array(3.0e19F, 4.0e19F), 2.0F) +-- !query schema +struct<vector_norm(array(3.0E19, 4.0E19), 2.0):float> +-- !query output +5.0E19 + + +-- !query +SELECT vector_normalize(array(3.0e19F, 4.0e19F), 2.0F) +-- !query schema +struct<vector_normalize(array(3.0E19, 4.0E19), 2.0):array<float>> +-- !query output +[0.6,0.8] + + +-- !query +SELECT vector_norm(array(3.0e38F, 3.0e38F), 1.0F) +-- !query schema +struct<vector_norm(array(3.0E38, 3.0E38), 1.0):float> +-- !query output +Infinity + + +-- !query +SELECT vector_normalize(array(3.0e38F, 3.0e38F), 1.0F) +-- !query schema +struct<vector_normalize(array(3.0E38, 3.0E38), 1.0):array<float>> +-- !query output +[0.5,0.5] + + +-- !query +SELECT vector_norm(array(1.0e-23F, 0.0F), 2.0F) +-- !query schema +struct<vector_norm(array(1.0E-23, 0.0), 2.0):float> +-- !query output +1.0E-23 + + +-- !query +SELECT vector_normalize(array(1.0e-23F, 0.0F), 2.0F) +-- !query schema +struct<vector_normalize(array(1.0E-23, 0.0), 2.0):array<float>> +-- !query output +[1.0,0.0] + + +-- !query +SELECT vector_normalize(array(1.0e-23F, 1.0e-23F), 2.0F) +-- !query schema +struct<vector_normalize(array(1.0E-23, 1.0E-23), 2.0):array<float>> +-- !query output +[0.70710677,0.70710677] + + +-- !query +SELECT vector_norm(array(float('inf'), 1.0F), 2.0F) +-- !query schema +struct<vector_norm(array(inf, 1.0), 2.0):float> +-- !query output +Infinity + + +-- !query +SELECT vector_normalize(array(float('inf'), 1.0F), 2.0F) +-- !query schema +struct<vector_normalize(array(inf, 1.0), 2.0):array<float>> +-- !query output +[NaN,0.0] diff --git a/sql/core/src/test/resources/sql-tests/results/window.sql.out b/sql/core/src/test/resources/sql-tests/results/window.sql.out index a6e76c599731c..daa91aab30703 100644 --- a/sql/core/src/test/resources/sql-tests/results/window.sql.out +++ b/sql/core/src/test/resources/sql-tests/results/window.sql.out @@ -555,23 +555,17 @@ FROM testData WINDOW w AS (PARTITION BY cate ORDER BY val) ORDER BY cate, val -- !query schema -struct<> --- !query output -org.apache.spark.SparkArithmeticException -{ - "errorClass" : "DIVIDE_BY_ZERO", - "sqlState" : "22012", - "messageParameters" : { - "config" : "\"spark.sql.ansi.enabled\"" - }, - "queryContext" : [ { - "objectType" : "", - "objectName" : "", - "startIndex" : 1016, - "stopIndex" : 1041, - "fragment" : "corr(val, val_long) OVER w" - } ] -} +struct<val:int,cate:string,max:int,min:int,min:int,count:bigint,sum:bigint,avg:double,stddev:double,first_value:int,first_value_ignore_null:int,first_value_contain_null:int,any_value:int,any_value_ignore_null:int,any_value_contain_null:int,last_value:int,last_value_ignore_null:int,last_value_contain_null:int,rank:int,dense_rank:int,cume_dist:double,percent_rank:double,ntile:int,row_number:int,var_pop:double,var_samp:double,approx_count_distinct:bigint,covar_pop:double,corr:double,stddev_samp:double,stddev_pop:double,collect_list:array<int>,collect_set:array<int>,skewness:double,kurtosis:double> +-- !query output +NULL NULL NULL NULL NULL 0 NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL 1 1 0.5 0.0 1 1 NULL NULL 0 NULL NULL NULL NULL [] [] NULL NULL +3 NULL 3 3 3 1 3 3.0 NULL NULL 3 NULL NULL 3 NULL 3 3 3 2 2 1.0 1.0 2 2 0.0 NULL 1 0.0 NULL NULL 0.0 [3] [3] NULL NULL +NULL a NULL NULL NULL 0 NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL 1 1 0.25 0.0 1 1 NULL NULL 0 NULL NULL NULL NULL [] [] NULL NULL +1 a 1 1 1 2 2 1.0 0.0 NULL 1 NULL NULL 1 NULL 1 1 1 2 2 0.75 0.3333333333333333 1 2 0.0 0.0 1 0.0 NULL 0.0 0.0 [1,1] [1] 0.7071067811865476 -1.5 +1 a 1 1 1 2 2 1.0 0.0 NULL 1 NULL NULL 1 NULL 1 1 1 2 2 0.75 0.3333333333333333 2 3 0.0 0.0 1 0.0 NULL 0.0 0.0 [1,1] [1] 0.7071067811865476 -1.5 +2 a 2 1 1 3 4 1.3333333333333333 0.5773502691896258 NULL 1 NULL NULL 1 NULL 2 2 2 4 3 1.0 1.0 2 4 0.22222222222222224 0.33333333333333337 2 4.772185885555555E8 1.0 0.5773502691896258 0.4714045207910317 [1,1,2] [1,2] 1.1539890888012805 -0.6672217220327235 +1 b 1 1 1 1 1 1.0 NULL 1 1 1 1 1 1 1 1 1 1 1 0.3333333333333333 0.0 1 1 0.0 NULL 1 NULL NULL NULL 0.0 [1] [1] NULL NULL +2 b 2 1 1 2 3 1.5 0.7071067811865476 1 1 1 1 1 1 2 2 2 2 2 0.6666666666666666 0.5 1 2 0.25 0.5 2 0.0 NULL 0.7071067811865476 0.5 [1,2] [1,2] 0.0 -2.0000000000000013 +3 b 3 1 1 3 6 2.0 1.0 1 1 1 1 1 1 3 3 3 3 3 1.0 1.0 2 3 0.6666666666666666 1.0 3 5.3687091175E8 1.0 1.0 0.816496580927726 [1,2,3] [1,2,3] 0.7057890433107311 -1.4999999999999984 -- !query diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala index 1d462c75f20f5..dbccbcce73933 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.ExtendedAnalysisException import org.apache.spark.sql.catalyst.expressions.aggregate.{ApproxTopK, ApproxTopKAggregateBuffer, CombineInternal} import org.apache.spark.sql.errors.DataTypeErrors.toSQLType +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{BooleanType, ByteType, DataType, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, TimestampNTZType, TimestampType, TimeType} @@ -319,6 +320,52 @@ class ApproxTopKSuite extends SharedSparkSession { } } + // scalastyle:off nonascii + test("SPARK-58096: approx_top_k keys on raw sort-key bytes, not a lossy-decoded String, " + + "so ICU-collation-distinct non-ASCII values are not over-merged") { + // U+4E14, U+4E15 and U+4E16 are distinct under UNICODE_CI, but their ICU sort keys are + // arbitrary bytes that decode to the same String via a lossy UTF-8 conversion. Keying the + // sketch on that decoded String would collapse them into a single item with an inflated + // count; keying on the raw sort-key bytes keeps them separate. + val res = sql( + """SELECT approx_top_k(c, 3) + |FROM (SELECT CAST(col AS STRING COLLATE UNICODE_CI) AS c + | FROM VALUES ('且'), ('且'), ('且'), + | ('丕'), ('丕'), ('世') AS t(col)) + |""".stripMargin) + checkAnswer(res, Row(Seq(Row("且", 3), Row("丕", 2), Row("世", 1)))) + } + + test("SPARK-58096: approx_top_k_accumulate/estimate keeps ICU-collation-distinct " + + "non-ASCII values separate through the serde round-trip") { + val res = sql( + """SELECT approx_top_k_estimate(approx_top_k_accumulate(c), 3) + |FROM (SELECT CAST(col AS STRING COLLATE UNICODE_CI) AS c + | FROM VALUES ('且'), ('且'), ('且'), + | ('丕'), ('丕'), ('世') AS t(col)) + |""".stripMargin) + checkAnswer(res, Row(Seq(Row("且", 3), Row("丕", 2), Row("世", 1)))) + } + + test("SPARK-58096: approx_top_k_combine keeps ICU-collation-distinct non-ASCII values " + + "separate across sketches and a shuffle") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val sketches = sql( + """SELECT approx_top_k_accumulate(CAST(col AS STRING COLLATE UNICODE_CI)) AS sketch + | FROM VALUES ('且'), ('且'), ('且'), ('丕') AS t(col) + |UNION ALL + |SELECT approx_top_k_accumulate(CAST(col AS STRING COLLATE UNICODE_CI)) AS sketch + | FROM VALUES ('丕'), ('世') AS t(col) + |""".stripMargin).repartition(2) + sketches.createOrReplaceTempView("approx_top_k_sketches") + val res = sql( + "SELECT approx_top_k_estimate(approx_top_k_combine(sketch, 100), 3) " + + "FROM approx_top_k_sketches") + checkAnswer(res, Row(Seq(Row("且", 3), Row("丕", 2), Row("世", 1)))) + } + } + // scalastyle:on nonascii + test("SPARK-52588: accumulate and estimate of Decimal(4, 1)") { val res = sql("SELECT approx_top_k_estimate(approx_top_k_accumulate(expr, 10)) " + "FROM VALUES CAST(0.0 AS DECIMAL(4, 1)), CAST(0.0 AS DECIMAL(4, 1)), " + diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala index ae24a21538f15..a7705de0b4b76 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ApproximatePercentileQuerySuite.scala @@ -20,12 +20,19 @@ package org.apache.spark.sql import java.sql.{Date, Timestamp} import java.time.{Duration, LocalDateTime, LocalTime, Period} -import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile +import org.apache.spark.SparkArithmeticException +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, + ApproximatePercentile, Final, Partial} import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile.DEFAULT_PERCENTILE_ACCURACY import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile.PercentileDigest +import org.apache.spark.sql.catalyst.plans.logical.Expand import org.apache.spark.sql.catalyst.util.DateTimeUtils +import org.apache.spark.sql.execution.ReusedSubqueryExec +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec +import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.sql.types.TimeType +import org.apache.spark.sql.types.{ArrayType, DoubleType, TimeType} import org.apache.spark.tags.SlowSQLTest /** @@ -36,6 +43,461 @@ class ApproximatePercentileQuerySuite extends SharedSparkSession { import testImplicits._ private val table = "percentile_approx" + private val constantFoldingRule = + "org.apache.spark.sql.catalyst.optimizer.ConstantFolding" + + private def fusionTest(name: String)(body: => Unit): Unit = test(name) { + withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true") { + body + } + } + + private def excludedRules: Seq[String] = { + spark.sessionState.conf.optimizerExcludedRules.toSeq + .flatMap(_.split(",")) + .map(_.trim) + .filter(_.nonEmpty) + } + + private def assertPercentileDigestCount(query: DataFrame, expected: Int): Unit = { + val counts = query.queryExecution.sparkPlan.collect { + case aggregate: ObjectHashAggregateExec => + aggregate.aggregateExpressions.count( + _.aggregateFunction.isInstanceOf[ApproximatePercentile]) + } + assert(counts.nonEmpty) + assert(counts.forall(_ == expected), counts) + } + + private def checkMatchesUnfusedBaseline(sql: String, expectedDigests: Int): Unit = { + val withoutConstantFolding = (excludedRules :+ constantFoldingRule).distinct + val baseline = withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> withoutConstantFolding.mkString(","), + SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") { + spark.sql(sql).collect().toSeq + } + + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> withoutConstantFolding.mkString(",")) { + val query = spark.sql(sql) + checkAnswer(query, baseline) + assertPercentileDigestCount(query, expectedDigests) + } + } + + test("approximate percentile fusion can be disabled") { + withSQLConf(SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") { + val query = spark.sql( + "SELECT percentile_approx(id, 0.5D), percentile_approx(id, 0.9D) FROM range(10)") + checkAnswer(query, Row(4L, 8L)) + assertPercentileDigestCount(query, 2) + } + } + + fusionTest("compatible scalar percentiles share one physical percentile digest") { + withTempView(table) { + (1 to 1000).toDF("col").createOrReplaceTempView(table) + val query = spark.sql( + s"""SELECT + | approx_percentile(col, 0.5, 10000), + | approx_percentile(col, 0.9, 10000), + | approx_percentile(col, 0.95, 10000) + |FROM $table + |""".stripMargin) + + checkAnswer(query, Row(500, 900, 950)) + assertPercentileDigestCount(query, 1) + val optimizedPercentiles = query.queryExecution.optimizedPlan.expressions.flatMap { + _.collect { case percentile: ApproximatePercentile => percentile } + } + assert(optimizedPercentiles.nonEmpty) + assert(optimizedPercentiles.forall(_.prettyName == "approx_percentile")) + val modes = query.queryExecution.sparkPlan.collect { + case aggregate: ObjectHashAggregateExec => + aggregate.aggregateExpressions.collect { + case expression @ AggregateExpression(_: ApproximatePercentile, _, _, _, _) => + expression.mode + } + }.flatten.toSet + assert(modes == Set(Partial, Final)) + } + } + + fusionTest("do not fuse duplicate percentages already shared by physical planning") { + withTempView(table) { + (1 to 1000).toDF("col").createOrReplaceTempView(table) + val query = spark.sql( + s"""SELECT + | percentile_approx(col, 0.5D), + | percentile_approx(col, 0.25D + 0.25D) + |FROM $table + |""".stripMargin) + + checkAnswer(query, Row(500, 500)) + val percentiles = query.queryExecution.sparkPlan.collect { + case aggregate: ObjectHashAggregateExec => + aggregate.aggregateExpressions.collect { + case AggregateExpression( + percentile: ApproximatePercentile, _, _, _, _) => percentile + } + }.flatten + assert(percentiles.nonEmpty) + assert(percentiles.forall(_.percentageExpression.dataType == DoubleType)) + } + } + + fusionTest("preserve structural input and filter evaluation") { + checkAnswer( + spark.sql( + """SELECT + | percentile_approx((a + b) + c, 0.5D), + | percentile_approx(a + (b + c), 0.9D) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(a, b, c) + |""".stripMargin), + Row(1.0d, 0.0d)) + + checkAnswer( + spark.sql( + """SELECT + | percentile_approx(v, 0.5D) + | FILTER (WHERE (a + b) + c = 1D), + | percentile_approx(v, 0.9D) + | FILTER (WHERE a + (b + c) = 1D) + |FROM VALUES ( + | 7, + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(v, a, b, c) + |""".stripMargin), + Row(7, null)) + + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + val exception = intercept[SparkArithmeticException] { + spark.sql( + """SELECT + | percentile_approx(a + (b + c), 0.5D), + | percentile_approx((a + b) + c, 0.9D) + |FROM VALUES ( + | CAST(2147483647 AS INT), + | CAST(1 AS INT), + | CAST(-1 AS INT) + |) AS t(a, b, c) + |""".stripMargin).collect() + } + assert(exception.getCondition == "ARITHMETIC_OVERFLOW") + } + } + + fusionTest("do not fuse canonical input or filter collisions") { + checkMatchesUnfusedBaseline( + """SELECT + | percentile_approx(a + (b + c), 0.5D), + | percentile_approx((a + b) + c, 0.5D), + | percentile_approx((a + b) + c, 0.9D), + | percentile_approx(a + (b + c), 0.9D) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(a, b, c) + |""".stripMargin, + expectedDigests = 2) + + val crossDistinctCollision = + """SELECT + | percentile_approx(DISTINCT a + (b + c), 0.5D), + | percentile_approx(DISTINCT a + (b + c), 0.9D), + | percentile_approx((a + b) + c, 0.5D), + | percentile_approx((a + b) + c, 0.9D) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(a, b, c) + |""".stripMargin + val unfusedBaseline = withSQLConf( + SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") { + spark.sql(crossDistinctCollision).collect().toSeq + } + checkAnswer(spark.sql(crossDistinctCollision), unfusedBaseline) + + val filteredArrayCollision = spark.sql( + """SELECT + | percentile_approx(v, array(0.5D, 0.9D)) + | FILTER (WHERE a + (b + c) = 0D), + | percentile_approx(v, 0.5D) + | FILTER (WHERE (a + b) + c = 0D), + | percentile_approx(v, 0.9D) + | FILTER (WHERE (a + b) + c = 0D) + |FROM VALUES ( + | 7, + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(v, a, b, c) + |""".stripMargin) + checkAnswer(filteredArrayCollision, Row(Seq(7, 7), null, null)) + assertPercentileDigestCount(filteredArrayCollision, 2) + } + + fusionTest("preserve canonically colliding parameters") { + val cases = Seq( + ( + """SELECT + | percentile_approx( + | v, 0.5D, CAST((1e16D + -1e16D) + 3D AS INT)), + | percentile_approx( + | v, 0.5D, CAST(1e16D + (-1e16D + 3D) AS INT)), + | percentile_approx( + | v, 0.9D, CAST(1e16D + (-1e16D + 3D) AS INT)), + | percentile_approx( + | v, 0.9D, CAST((1e16D + -1e16D) + 3D AS INT)) + |FROM range(100) AS t(v) + |""".stripMargin, + 2), + ( + """SELECT + | percentile_approx( + | id, array((1e16D + -1e16D) + 0.5D, 0.9D)), + | percentile_approx( + | id, 1e16D + (-1e16D + 0.5D)), + | percentile_approx(id, 0.9D) + |FROM range(100) + |""".stripMargin, + 2)) + + cases.foreach { case (sql, expectedDigests) => + checkMatchesUnfusedBaseline(sql, expectedDigests) + } + } + + fusionTest("preserve existing arrays that collide after distinct removal") { + val query = spark.sql( + """SELECT + | percentile_approx( + | DISTINCT a + (b + c), array(0.5D, 0.9D)), + | percentile_approx(DISTINCT a + (b + c), 0.1D), + | percentile_approx((a + b) + c, 0.5D), + | percentile_approx((a + b) + c, 0.9D) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t(a, b, c) + |""".stripMargin) + checkAnswer(query, Row(Seq(0.0d, 0.0d), 0.0d, 1.0d, 1.0d)) + assertPercentileDigestCount(query, 3) + } + + fusionTest("fused percentiles use fresh result IDs across CTE references") { + val query = spark.sql( + """WITH c AS ( + | SELECT + | percentile_approx(v, 0.5D) AS a, + | percentile_approx(v, 0.9D) AS b + | FROM range(1, 6) AS t(v) + |) + |SELECT c1.a, c1.b, c2.a + |FROM c AS c1 JOIN c AS c2 + |""".stripMargin) + checkAnswer(query, Row(3L, 5L, 3L)) + val optimizedPlan = query.queryExecution.optimizedPlan + assert(optimizedPlan.subqueriesAll.exists(_.exists(_.expressions.exists(_.exists { + case percentile: ApproximatePercentile => + percentile.percentageExpression.dataType.isInstanceOf[ArrayType] + case _ => false + }))), optimizedPlan.numberedTreeString) + } + + fusionTest("preserve percentile inputs across scalar subquery reuse") { + val query = + """SELECT + | (SELECT named_struct( + | 'p50', percentile_approx((a + b) + c, 0.5D), + | 'p90', percentile_approx((a + b) + c, 0.9D)) + | FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + | ) AS t(a, b, c)), + | (SELECT named_struct( + | 'p50', get(percentile_approx( + | a + (b + c), array(0.5D, 0.9D)), 0), + | 'p90', get(percentile_approx( + | a + (b + c), array(0.5D, 0.9D)), 1)) + | FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + | ) AS t(a, b, c)) + |""".stripMargin + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "false") { + checkAnswer( + spark.sql(query), + Row(Row(1.0d, 1.0d), Row(0.0d, 0.0d))) + } + } + + fusionTest("preserve pre-fusion identity across exchange reuse") { + val parameterPairs = Seq( + ("(1e16D + -1e16D) + 0.5D", "100"), + ("0.5D", "CAST((1e16D + -1e16D) + 100D AS INT)"), + ("0.5D", "100L")) + val activeRules = (excludedRules :+ constantFoldingRule).distinct.mkString(",") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "false", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> activeRules) { + parameterPairs.foreach { case (firstPercentage, secondAccuracy) => + val query = spark.sql( + s"""SELECT + | (a + b) + c, + | array( + | percentile_approx(1D, $firstPercentage, 100), + | percentile_approx(1D, 0.9D, $secondAccuracy)) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t1(a, b, c) + |GROUP BY (a + b) + c + |UNION ALL + |SELECT + | a + (b + c), + | array( + | percentile_approx(1D, 0.5D, 100), + | percentile_approx(1D, 0.9D, 100)) + |FROM VALUES ( + | CAST(10000000000000000 AS DOUBLE), + | CAST(-10000000000000000 AS DOUBLE), + | CAST(1 AS DOUBLE) + |) AS t2(a, b, c) + |GROUP BY a + (b + c) + |""".stripMargin) + + checkAnswer( + query, + Seq( + Row(1.0d, Seq(1.0d, 1.0d)), + Row(0.0d, Seq(1.0d, 1.0d)))) + assert(query.queryExecution.executedPlan.collect { + case _: ReusedExchangeExec => true + }.isEmpty) + } + } + } + + fusionTest("identical fused percentiles remain reusable") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true") { + val subquery = + """(SELECT array( + | percentile_approx(v, 0D), + | percentile_approx(v, 1D)) + | FROM VALUES (1), (2), (3) AS t(v)) + |""".stripMargin + val query = spark.sql(s"SELECT $subquery, $subquery") + + checkAnswer(query, Row(Seq(1, 3), Seq(1, 3))) + assert(query.queryExecution.executedPlan.collectWithSubqueries { + case _: ReusedSubqueryExec => true + }.nonEmpty) + } + } + + fusionTest("fused distinct percentiles keep a single distinct group") { + val query = spark.sql( + """SELECT + | count(DISTINCT id), + | percentile_approx(DISTINCT id, 0.5D), + | percentile_approx(DISTINCT id, 0.9D) + |FROM range(10) + |""".stripMargin) + + checkAnswer(query, Row(10L, 4L, 8L)) + assertPercentileDigestCount(query, 1) + assert(query.queryExecution.optimizedPlan.collect { + case _: Expand => true + }.isEmpty) + } + + fusionTest("combined scalar percentiles preserve empty-input nulls") { + withTempView(table) { + Seq.empty[Int].toDF("col").createOrReplaceTempView(table) + val query = spark.sql( + s"""SELECT + | percentile_approx(col, 0.5D), + | percentile_approx(col, 0.9D) + |FROM $table + |""".stripMargin) + checkAnswer(query, Row(null, null)) + assertPercentileDigestCount(query, 1) + } + } + + fusionTest("preserve compressed and merged low-accuracy percentile digests") { + val sql = + """SELECT + | percentile_approx(id, 0.1D, 100), + | percentile_approx(id, 0.5D, 100), + | percentile_approx(id, 0.9D, 100) + |FROM range(0, 50000, 1, 4) + |""".stripMargin + val baseline = withSQLConf( + SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false") { + spark.sql(sql).collect().toSeq + } + + val query = spark.sql(sql) + checkAnswer(query, baseline) + assertPercentileDigestCount(query, 1) + } + + fusionTest("fuse compatible percentiles introduced by merged scalar subqueries") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SUBQUERY_REUSE_ENABLED.key -> "true") { + val query = spark.sql( + """SELECT + | (SELECT percentile_approx(id, 0.5D) FROM range(10)), + | (SELECT percentile_approx(id, 0.9D) FROM range(10)) + |""".stripMargin) + + checkAnswer(query, Row(4L, 8L)) + val digestCounts = query.queryExecution.executedPlan.collectWithSubqueries { + case aggregate: ObjectHashAggregateExec => + aggregate.aggregateExpressions.count( + _.aggregateFunction.isInstanceOf[ApproximatePercentile]) + } + assert(digestCounts.nonEmpty && digestCounts.forall(_ == 1), digestCounts) + } + } + + fusionTest("fuse canonical input groups with disjoint percentages independently") { + val query = spark.sql( + """SELECT + | percentile_approx(a + b, 0.5D), + | percentile_approx(a + b, 0.9D), + | percentile_approx(b + a, 0.25D), + | percentile_approx(b + a, 0.75D) + |FROM VALUES (1D, 2D), (3D, 4D), (5D, 6D) AS t(a, b) + |""".stripMargin) + + checkAnswer(query, Row(7.0d, 11.0d, 3.0d, 11.0d)) + assertPercentileDigestCount(query, 2) + } test("percentile_approx, single percentile value") { withTempView(table) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala index 7ec687e3d9534..9a40015036192 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala @@ -460,4 +460,111 @@ class BinBySuite extends QueryTest with SharedSparkSession { Row(2, null, null, null, null))) } } + + test("ORDER BY on a column BIN BY re-outputs without its qualifier is rejected") { + // BIN BY re-outputs its input columns without the original qualifier, so `ORDER BY t.value` + // cannot resolve against them; hidden-output insertion appends the column to an operator whose + // child does not produce it, and analysis fails with MISSING_ATTRIBUTES. + withSQLConf( + SQLConf.BIN_BY_ENABLED.key -> "true", + SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "false") { + checkError( + exception = intercept[AnalysisException] { + spark.sql( + """SELECT * FROM VALUES + | (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP '2024-01-01 00:05:00', 100.0D) + | AS t(ts_start, ts_end, value) + |BIN BY ( + | RANGE ts_start TO ts_end + | BIN WIDTH INTERVAL '5' MINUTE + | ALIGN TO TIMESTAMP '2024-01-01 00:00:00' + | DISTRIBUTE UNIFORM (value)) + |ORDER BY t.value""".stripMargin) + }, + condition = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION", + parameters = Map( + "missingAttributes" -> "\"value\"", + "input" -> ("\"ts_start\", \"ts_end\", \"value\", \"bin_start\", \"bin_end\", " + + "\"bin_distribute_ratio\""), + "operator" -> "!Sort \\[value#\\d+ ASC NULLS FIRST\\], true", + "operation" -> "\"value\""), + matchPVals = true, + queryContext = + Array(ExpectedContext(fragment = "ORDER BY t.value", start = 271, stop = 286))) + } + } + + test("BIN BY resolves a bare ORDER BY on a re-output DISTRIBUTE column that is not projected") { + withSQLConf( + SQLConf.BIN_BY_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + // A bare ORDER BY on the re-output DISTRIBUTE column resolves even when it is not projected. + val df = spark.sql( + """SELECT bin_start + |FROM VALUES + | (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP '2024-01-01 00:10:00', 100.0D) + | AS metrics(ts_start, ts_end, value) + |BIN BY ( + | RANGE ts_start TO ts_end BIN WIDTH INTERVAL '5' MINUTE + | ALIGN TO TIMESTAMP '2024-01-01 00:00:00' DISTRIBUTE UNIFORM (value)) + |ORDER BY value""".stripMargin) + checkAnswer(df, Seq( + Row(tsAt("2024-01-01 00:00:00")), + Row(tsAt("2024-01-01 00:05:00")))) + } + } + + test("BIN BY rejects a source-qualified reference to a re-output DISTRIBUTE column") { + withSQLConf( + SQLConf.BIN_BY_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + // The re-output DISTRIBUTE column drops its source qualifier, so `metrics.value` fails. + checkError( + exception = intercept[AnalysisException] { + spark.sql( + """SELECT metrics.value + |FROM VALUES + | (TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP '2024-01-01 00:10:00', 100.0D) + | AS metrics(ts_start, ts_end, value) + |BIN BY ( + | RANGE ts_start TO ts_end BIN WIDTH INTERVAL '5' MINUTE + | ALIGN TO TIMESTAMP '2024-01-01 00:00:00' + | DISTRIBUTE UNIFORM (value))""".stripMargin) + }, + condition = "UNRESOLVED_COLUMN.WITH_SUGGESTION", + parameters = Map( + "objectName" -> "`metrics`.`value`", + "proposal" -> ("`metrics`.`ts_end`, `value`, `bin_start`, " + + "`metrics`.`ts_start`, `bin_end`")), + queryContext = + Array(ExpectedContext(fragment = "metrics.value", start = 7, stop = 19))) + } + } + + test("BIN BY keeps a file-source metadata column reachable in ORDER BY and SELECT") { + withSQLConf( + SQLConf.BIN_BY_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + withTempDir { dir => + val path = new java.io.File(dir, "metrics").getCanonicalPath + spark.sql( + """SELECT TIMESTAMP '2024-01-01 00:00:00' AS ts_start, + | TIMESTAMP '2024-01-01 00:10:00' AS ts_end, 100.0D AS value""".stripMargin) + .write.mode("overwrite").parquet(path) + val binBy = + s"""parquet.`$path` BIN BY (RANGE ts_start TO ts_end BIN WIDTH INTERVAL '5' MINUTE + | ALIGN TO TIMESTAMP '2024-01-01 00:00:00' DISTRIBUTE UNIFORM (value))""".stripMargin + + // `_metadata` is a hidden metadata column; BIN BY carries it through for a bare ORDER BY. + checkAnswer( + spark.sql(s"SELECT bin_start FROM $binBy ORDER BY _metadata.file_size"), + Seq(Row(tsAt("2024-01-01 00:00:00")), Row(tsAt("2024-01-01 00:05:00")))) + + // Also selectable through BIN BY. + checkAnswer( + spark.sql(s"SELECT bin_start, _metadata.file_size > 0 FROM $binBy"), + Seq(Row(tsAt("2024-01-01 00:00:00"), true), Row(tsAt("2024-01-01 00:05:00"), true))) + } + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala index 458688de3fbde..7d079fb894e7d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/BitmapExpressionsQuerySuite.scala @@ -17,7 +17,27 @@ package org.apache.spark.sql -import org.apache.spark.sql.functions.{bitmap_and_agg, bitmap_bit_position, bitmap_bucket_number, bitmap_construct_agg, bitmap_count, bitmap_or_agg, col, expr, hex, lit, substring, to_binary} +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.functions.{ + bitmap_and, + bitmap_and_agg, + bitmap_andnot, + bitmap_bit_position, + bitmap_bucket_number, + bitmap_construct_agg, + bitmap_count, + bitmap_or, + bitmap_or_agg, + bitmap_xor, + bitmap_xor_agg, + col, + expr, + hex, + lit, + substring, + sum, + to_binary} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class BitmapExpressionsQuerySuite extends SharedSparkSession { @@ -297,6 +317,167 @@ class BitmapExpressionsQuerySuite extends SharedSparkSession { } } + test("scalar bitmap binary operations") { + def checkResults(): Unit = { + val df = Seq(("F00F", "70")).toDF("left", "right") + checkAnswer( + df.selectExpr( + "substring(hex(bitmap_and(to_binary(left, 'hex'), to_binary(right, 'hex'))), 0, 4)", + "substring(hex(bitmap_or(to_binary(left, 'hex'), to_binary(right, 'hex'))), 0, 4)", + "substring(hex(bitmap_andnot(to_binary(left, 'hex'), to_binary(right, 'hex'))), 0, 4)", + "substring(hex(bitmap_xor(to_binary(left, 'hex'), to_binary(right, 'hex'))), 0, 4)"), + Seq(Row("7000", "F00F", "800F", "800F"))) + + val leftBitmap = to_binary(col("left"), lit("hex")) + val rightBitmap = to_binary(col("right"), lit("hex")) + checkAnswer( + df.select( + substring(hex(bitmap_and(leftBitmap, rightBitmap)), 0, 4), + substring(hex(bitmap_or(leftBitmap, rightBitmap)), 0, 4), + substring(hex(bitmap_andnot(leftBitmap, rightBitmap)), 0, 4), + substring(hex(bitmap_xor(leftBitmap, rightBitmap)), 0, 4)), + Seq(Row("7000", "F00F", "800F", "800F"))) + + val shortBitmap = Seq(("", "FF")).toDF("left", "right") + checkAnswer( + shortBitmap.selectExpr( + "length(bitmap_and(to_binary(left, 'hex'), to_binary(right, 'hex')))", + "length(bitmap_or(to_binary(left, 'hex'), to_binary(right, 'hex')))", + "length(bitmap_andnot(to_binary(left, 'hex'), to_binary(right, 'hex')))", + "length(bitmap_xor(to_binary(left, 'hex'), to_binary(right, 'hex')))"), + Seq(Row(4096, 4096, 4096, 4096))) + + checkAnswer( + spark.sql(""" + |SELECT bitmap_count(bitmap_and(X'', X'FF')), + | bitmap_count(bitmap_or(X'', X'FF')), + | bitmap_count(bitmap_andnot(X'', X'FF')), + | bitmap_count(bitmap_xor(X'', X'FF')) + |""".stripMargin), + Seq(Row(0, 8, 0, 8))) + + val boundaryBitmap = + Seq((Array.fill[Byte](4096)(0), Array.fill[Byte](4096)(0))).toDF("left", "right") + checkAnswer( + boundaryBitmap.selectExpr( + "length(bitmap_and(left, right))", + "length(bitmap_or(left, right))", + "length(bitmap_andnot(left, right))", + "length(bitmap_xor(left, right))"), + Seq(Row(4096, 4096, 4096, 4096))) + + checkAnswer( + spark.sql("SELECT bitmap_count(bitmap_or(bitmap_and(X 'F0', X '70'), X '0F'))"), + Seq(Row(7))) + } + + Seq("CODEGEN_ONLY" -> "true", "NO_CODEGEN" -> "false").foreach { + case (codegenMode, wholeStageEnabled) => + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStageEnabled, + SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) { + checkResults() + } + } + } + + test("scalar bitmap binary operations with aggregate bitmaps") { + val constructed = Seq((1, 2), (2, 3), (3, 4)).toDF("left", "right") + val leftBitmap = bitmap_construct_agg(bitmap_bit_position(col("left"))) + val rightBitmap = bitmap_construct_agg(bitmap_bit_position(col("right"))) + checkAnswer( + constructed.agg(bitmap_count(bitmap_and(leftBitmap, rightBitmap))), + Seq(Row(2))) + + val precomputed = Seq("F0", "70").toDF("bitmap") + val orBitmap = bitmap_or_agg(to_binary(col("bitmap"), lit("hex"))) + val andBitmap = bitmap_and_agg(to_binary(col("bitmap"), lit("hex"))) + checkAnswer( + precomputed.agg( + bitmap_count(bitmap_and(orBitmap, andBitmap)), + bitmap_count(bitmap_or(orBitmap, andBitmap)), + bitmap_count(bitmap_andnot(orBitmap, andBitmap)), + bitmap_count(bitmap_xor(orBitmap, andBitmap))), + Seq(Row(3, 4, 1, 1))) + } + + test("scalar bitmap binary operations in grouped query") { + val df = Seq((1, "F0", "70"), (1, "10", "20"), (2, "FF", "0F")) + .toDF("group_id", "left", "right") + checkAnswer( + df.selectExpr( + "group_id", + "bitmap_count(bitmap_and(to_binary(left, 'hex'), to_binary(right, 'hex'))) AS count") + .groupBy("group_id") + .agg(sum("count")) + .orderBy("group_id"), + Seq(Row(1, 3), Row(2, 4))) + } + + test("scalar bitmap binary operations with nulls") { + val df = Seq[(String, String)](("F0", null), (null, "70"), (null, null)).toDF("left", "right") + checkAnswer( + df.selectExpr( + "bitmap_and(to_binary(left, 'hex'), to_binary(right, 'hex'))", + "bitmap_or(to_binary(left, 'hex'), to_binary(right, 'hex'))", + "bitmap_andnot(to_binary(left, 'hex'), to_binary(right, 'hex'))", + "bitmap_xor(to_binary(left, 'hex'), to_binary(right, 'hex'))"), + Seq( + Row(null, null, null, null), + Row(null, null, null, null), + Row(null, null, null, null))) + } + + test("scalar bitmap binary operations reject oversized inputs") { + val oversizedInputs = Seq( + "left" -> Seq((Array.fill[Byte](4097)(0), Array[Byte](0))).toDF("left", "right"), + "right" -> Seq((Array[Byte](0), Array.fill[Byte](4097)(0))).toDF("left", "right")) + + Seq("CODEGEN_ONLY" -> "true", "NO_CODEGEN" -> "false").foreach { + case (codegenMode, wholeStageEnabled) => + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStageEnabled, + SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) { + oversizedInputs.foreach { case (inputSide, oversized) => + Seq("bitmap_and", "bitmap_or", "bitmap_andnot", "bitmap_xor").foreach { functionName => + withClue(s"$functionName with oversized $inputSide input in $codegenMode: ") { + checkError( + exception = intercept[SparkRuntimeException] { + oversized.selectExpr(s"$functionName(left, right)").collect() + }, + condition = "BITMAP_INPUT_TOO_LARGE", + parameters = Map("inputNumBytes" -> "4097", "maxNumBytes" -> "4096")) + } + } + } + } + } + } + + test("scalar bitmap binary operations called with non-binary types") { + val invalidInputs = Seq( + ("first", "left", Seq((12, Array[Byte](0))).toDF("left", "right")), + ("second", "right", Seq((Array[Byte](0), 13)).toDF("left", "right"))) + + invalidInputs.foreach { case (paramIndex, inputColumn, df) => + Seq("bitmap_and", "bitmap_or", "bitmap_andnot", "bitmap_xor").foreach { functionName => + val sqlExpr = s"$functionName(left, right)" + checkError( + exception = intercept[AnalysisException] { + df.selectExpr(sqlExpr) + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> s""""$sqlExpr"""", + "paramIndex" -> paramIndex, + "requiredType" -> "\"BINARY\"", + "inputSql" -> s""""$inputColumn"""", + "inputType" -> "\"INT\""), + context = ExpectedContext(fragment = sqlExpr, start = 0, stop = sqlExpr.length - 1)) + } + } + } + test("bitmap_count called with non-binary type") { val df = Seq(12).toDF("a") checkError( @@ -356,4 +537,100 @@ class BitmapExpressionsQuerySuite extends SharedSparkSession { "inputType" -> "\"INT\""), context = ExpectedContext(fragment = "bitmap_and_agg(a)", start = 0, stop = 16)) } + + test("bitmap_xor_agg") { + // Test basic XOR functionality: 10 ^ 30 ^ 40 = 60 + val df = Seq("10", "30", "40").toDF("a") + checkAnswer( + df.selectExpr("substring(hex(bitmap_xor_agg(to_binary(a, 'hex'))), 0, 6)"), + Seq(Row("600000"))) + checkAnswer( + df.select(substring(hex(bitmap_xor_agg(to_binary(col("a"), lit("hex")))), 0, 6)), + Seq(Row("600000"))) + + // Test with same values - XOR should result in zero: 10 ^ 10 = 00 + val df2 = Seq("10", "10").toDF("a") + checkAnswer( + df2.selectExpr("substring(hex(bitmap_xor_agg(to_binary(a, 'hex'))), 0, 6)"), + Seq(Row("000000"))) + + // Test with zero - X XOR 0 = X: A0 ^ 00 = A0 + val df3 = Seq("A0", "00").toDF("a") + checkAnswer( + df3.selectExpr("substring(hex(bitmap_xor_agg(to_binary(a, 'hex'))), 0, 6)"), + Seq(Row("A00000"))) + + // Test with binary values of different lengths: 0A0B ^ 0A = 000B + val df4 = Seq("0A0B", "0A").toDF("a") + checkAnswer( + df4.selectExpr("substring(hex(bitmap_xor_agg(to_binary(a, 'hex'))), 0, 6)"), + Seq(Row("000B00"))) + + // Test empty result (no rows) - should return all zeros as XOR identity + val emptyDf = Seq.empty[String].toDF("a") + checkAnswer( + emptyDf.selectExpr("substring(hex(bitmap_xor_agg(to_binary(a, 'hex'))), 0, 6)"), + Seq(Row("000000"))) + + val emptyDf2 = Seq.empty[(String, Int)].toDF("a", "b") + checkAnswer( + emptyDf2 + .selectExpr("bitmap_xor_agg(to_binary(a, 'hex')) " + + "filter (where b = 1) as xor_agg") + .select(substring(hex(col("xor_agg")), 0, 6)), + Seq(Row("000000"))) + + // Test empty result (no rows) with GROUP BY - should return empty DataFrame + val emptyDf3 = Seq.empty[(String, Int, Int)].toDF("a", "b", "c") + checkAnswer( + emptyDf3 + .groupBy("c") + .agg(expr("bitmap_xor_agg(to_binary(a, 'hex')) " + + "filter (where b = 1) as xor_agg").alias("xor_agg")) + .select(substring(hex(col("xor_agg")), 0, 6)), + Seq()) + } + + test("bitmap_xor_agg with complex bitmaps from bitmap_construct_agg") { + val table = "bitmap_xor_test_table" + withTable(table) { + // Create test data: Group 1 has bits 1,2,4; Group 2 has bits 1,2,3. + spark.sql(s""" + | CREATE TABLE $table (group_id INT, bit_pos LONG) + | """.stripMargin) + spark.sql(s""" + | INSERT INTO $table VALUES + | (1, 1), (1, 2), (1, 4), + | (2, 1), (2, 2), (2, 3) + | """.stripMargin) + // XOR of the group bitmaps cancels bits 1 and 2, leaving bits 3 and 4. + val xorResult = spark.sql(s""" + | SELECT bitmap_count( + | bitmap_xor_agg(group_bitmap) + | ) as xor_count + | FROM ( + | SELECT bitmap_construct_agg(bitmap_bit_position(bit_pos)) as group_bitmap + | FROM $table + | GROUP BY group_id + | ) + | """.stripMargin) + checkAnswer(xorResult, Seq(Row(2))) + } + } + + test("bitmap_xor_agg called with non-binary type") { + val df = Seq(12).toDF("a") + checkError( + exception = intercept[AnalysisException] { + df.selectExpr("bitmap_xor_agg(a)") + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "\"bitmap_xor_agg(a)\"", + "paramIndex" -> "first", + "requiredType" -> "\"BINARY\"", + "inputSql" -> "\"a\"", + "inputType" -> "\"INT\""), + context = ExpectedContext(fragment = "bitmap_xor_agg(a)", start = 0, stop = 16)) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CharVarcharTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CharVarcharTestSuite.scala index 8d22575f5d09b..39694ad3d6869 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CharVarcharTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CharVarcharTestSuite.scala @@ -17,16 +17,24 @@ package org.apache.spark.sql -import org.apache.spark.{SparkConf, SparkRuntimeException} -import org.apache.spark.sql.catalyst.expressions.{Attribute, EqualTo, GreaterThan, ScalarSubquery, StringRPad} +import scala.util.Try + +import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException, SparkThrowable} +import org.apache.spark.sql.catalyst.analysis.FunctionRegistry +import org.apache.spark.sql.catalyst.analysis.resolver.ResolverGuard +import org.apache.spark.sql.catalyst.expressions.{ + ArrayJoin, Attribute, Concat, EqualTo, Expression, GreaterThan, Literal, ScalarSubquery, + StringRPad, StringToMap, Upper +} import org.apache.spark.sql.catalyst.expressions.Cast.toSQLId -import org.apache.spark.sql.catalyst.parser.CatalystSqlParser -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Project} +import org.apache.spark.sql.catalyst.parser.{CatalystSqlParser, ParseException} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, Project} import org.apache.spark.sql.catalyst.util.CharVarcharUtils import org.apache.spark.sql.connector.SchemaRequiredDataSource import org.apache.spark.sql.connector.catalog.{CatalogV2Util, InMemoryPartitionTableCatalog} import org.apache.spark.sql.execution.datasources.LogicalRelation import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.functions import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.sources.SimpleInsertSource import org.apache.spark.sql.test.SharedSparkSession @@ -854,6 +862,1173 @@ class BasicCharVarcharTestSuite extends SharedSparkSession { } } + test("SPARK-58797: CAST to CHAR/VARCHAR with standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charDf = sql("SELECT CAST('ab' AS CHAR(5)) AS c") + assert(charDf.schema.head.dataType === CharType(5)) + checkAnswer(charDf, Row("ab ")) + + val varcharDf = sql("SELECT CAST('hello' AS VARCHAR(5)) AS v") + assert(varcharDf.schema.head.dataType === VarcharType(5)) + checkAnswer(varcharDf, Row("hello")) + + // ISO 6.13: character-to-character CAST truncates rather than erroring. + checkAnswer(sql("SELECT CAST('hello!' AS VARCHAR(5)) AS v"), Row("hello")) + checkAnswer(sql("SELECT CAST('abcdef' AS CHAR(2)) AS c"), Row("ab")) + checkAnswer(sql("SELECT CAST('abcdef' AS VARCHAR(2)) AS v"), Row("ab")) + checkAnswer(sql("SELECT try_cast('abcdef' AS CHAR(2)) AS c"), Row("ab")) + checkAnswer(sql("SELECT try_cast('abcdef' AS VARCHAR(2)) AS v"), Row("ab")) + + // Multi-byte characters: length is in characters, not octets. + // scalastyle:off nonascii + checkAnswer(sql("SELECT CAST('你好' AS VARCHAR(2)) AS v"), Row("你好")) + checkAnswer(sql("SELECT CAST('你好啊' AS VARCHAR(2)) AS v"), Row("你好")) + // scalastyle:on nonascii + + // ISO 6.13 numeric-to-character CAST still errors when the literal does not fit. + checkError( + exception = intercept[SparkRuntimeException] { + sql("SELECT CAST(12345 AS VARCHAR(4))").collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "4") + ) + checkAnswer(sql("SELECT CAST(12345 AS VARCHAR(5)) AS v"), Row("12345")) + checkAnswer(sql("SELECT try_cast(12345 AS VARCHAR(4)) AS v"), Row(null)) + + // LCT must wrap the inner CAST, not retarget it (truncation / overflow stay). + checkAnswer( + sql("SELECT coalesce(CAST('abcdef' AS VARCHAR(2)), CAST('x' AS VARCHAR(4))) AS c"), + Row("ab")) + checkAnswer( + sql("""SELECT CASE WHEN true THEN CAST('abcdef' AS VARCHAR(2)) + |ELSE CAST('x' AS VARCHAR(4)) END AS c""".stripMargin), + Row("ab")) + checkAnswer( + sql("SELECT CAST('abcdef' AS VARCHAR(2)) IN (CAST('ab' AS VARCHAR(4)))"), + Row(true)) + checkAnswer( + sql("""SELECT coalesce( + | CAST('abcdef' AS VARCHAR(2) COLLATE UTF8_LCASE), + | CAST('x' AS VARCHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin), + Row("ab")) + checkAnswer( + sql("SELECT coalesce(try_cast(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5))) AS c"), + Row("x")) + checkError( + exception = intercept[SparkRuntimeException] { + sql("SELECT coalesce(CAST(12345 AS VARCHAR(4)), CAST('x' AS VARCHAR(5)))").collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "4") + ) + } + } + + test("SPARK-58797: store assignment with standardSemantics and charVarcharAsString") { + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true", + SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key -> "true") { + val wide = new StructType().add("c", CharType(10)) + val df = spark.createDataFrame(java.util.Arrays.asList(Row("spark")), wide) + checkError( + exception = intercept[SparkRuntimeException] { + df.to(new StructType().add("c", CharType(3))).collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + withTable("std_and_as_string") { + sql("CREATE TABLE std_and_as_string (v VARCHAR(2)) USING parquet") + checkError( + exception = intercept[SparkRuntimeException] { + sql("INSERT INTO std_and_as_string VALUES ('abc')") + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "2")) + } + } + } + + test("SPARK-58798: least common type for COALESCE/CASE with CHAR/VARCHAR") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assert(sql( + "SELECT coalesce(cast('hello' AS VARCHAR(5)), cast('world' AS VARCHAR(10))) AS c") + .schema.head.dataType === VarcharType(10)) + assert(sql( + "SELECT coalesce(cast('hello' AS VARCHAR(5)), cast('world!' AS CHAR(6))) AS c") + .schema.head.dataType === VarcharType(6)) + assert(sql( + "SELECT coalesce(cast('hello' AS CHAR(5)), cast('world!' AS CHAR(6))) AS c") + .schema.head.dataType === CharType(6)) + assert(sql( + "SELECT coalesce(cast('hello' AS VARCHAR(5)), 'world') AS c") + .schema.head.dataType === StringType) + assert(sql( + """SELECT CASE WHEN true THEN cast('a' AS CHAR(2)) + |ELSE cast('bb' AS CHAR(4)) END AS c""".stripMargin) + .schema.head.dataType === CharType(4)) + // LCT(NULL, T) = T + assert(sql("SELECT coalesce(null, cast('a' AS CHAR(5))) AS c") + .schema.head.dataType === CharType(5)) + } + } + + test("SPARK-58799: transforming string functions return STRING") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assert(sql("SELECT upper(cast('ab' AS CHAR(2))) AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT lower(cast('AB' AS VARCHAR(2))) AS c") + .schema.head.dataType === StringType) + assert(sql( + "SELECT cast('a' AS CHAR(1)) || cast('b' AS VARCHAR(1)) AS c") + .schema.head.dataType === StringType) + // Pads from CHAR participate in the concatenated value. + checkAnswer( + sql("SELECT cast('he' AS CHAR(4)) || cast('llo' AS CHAR(3)) AS c"), + Row("he llo")) + // Collated CHAR must promote to the same collation, not UTF8_BINARY STRING. + assert(sql( + """SELECT concat( + | cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + | cast('b' AS CHAR(3) COLLATE UTF8_LCASE)) AS c""".stripMargin) + .schema.head.dataType === StringType("UTF8_LCASE")) + checkAnswer( + sql("""SELECT concat( + | cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + | cast('b' AS CHAR(3) COLLATE UTF8_LCASE)) AS c""".stripMargin), + Row("a b ")) + assert(sql( + """SELECT elt( + | 1, + | cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + | cast('x' AS CHAR(1) COLLATE UTF8_LCASE)) AS c""".stripMargin) + .schema.head.dataType === StringType("UTF8_LCASE")) + checkAnswer( + sql("""SELECT elt( + | 1, + | cast('ab' AS CHAR(5) COLLATE UTF8_LCASE), + | cast('x' AS CHAR(1) COLLATE UTF8_LCASE)) AS c""".stripMargin), + Row("ab ")) + assert(sql("SELECT substr(cast('hello' AS VARCHAR(5)), 1, 2) AS c") + .schema.head.dataType === StringType) + assert(sql( + "SELECT upper(coalesce(cast('a' AS CHAR(2)), cast('b' AS CHAR(4)))) AS c") + .schema.head.dataType === StringType) + } + } + + test("SPARK-58798: LCT preserves collation on CHAR/VARCHAR") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val c1 = CharType(2, "UTF8_LCASE") + val c2 = CharType(4, "UTF8_LCASE") + assert(StringHelper.tightestCommonString(c1, c2).contains(CharType(4, "UTF8_LCASE"))) + val v1 = VarcharType(3, "UTF8_LCASE") + val v2 = VarcharType(5, "UTF8_LCASE") + assert(StringHelper.tightestCommonString(v1, v2).contains(VarcharType(5, "UTF8_LCASE"))) + assert(StringHelper.tightestCommonString(c1, v2).contains(VarcharType(5, "UTF8_LCASE"))) + } + } + + test("SPARK-58799: regexp/mask/split return STRING under standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assert(sql("SELECT regexp_replace(cast('ab' AS CHAR(2)), 'a', 'x') AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT regexp_extract(cast('ab' AS VARCHAR(2)), '(a)', 1) AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT split(cast('a,b' AS CHAR(3)), ',') AS c") + .schema.head.dataType === ArrayType(StringType, containsNull = false)) + assert(sql("SELECT mask(cast('ab' AS CHAR(2))) AS c") + .schema.head.dataType === StringType) + } + } + + test("SPARK-58796: preserve vs standardSemantics result-type matrix") { + // preserve-only: transforming ops may keep Char/Varchar (leaky experimental path). + withSQLConf(SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true") { + assert(sql("SELECT upper(cast('ab' AS CHAR(2))) AS c") + .schema.head.dataType === CharType(2)) + } + // standardSemantics: transforming ops return STRING even if preserve is also on. + withSQLConf( + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true", + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assert(sql("SELECT upper(cast('ab' AS CHAR(2))) AS c") + .schema.head.dataType === StringType) + } + } + + test("SPARK-58797: standardSemantics wins over charVarcharAsString") { + withSQLConf( + SQLConf.LEGACY_CHAR_VARCHAR_AS_STRING.key -> "true", + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val df = sql("SELECT CAST('ab' AS CHAR(5)) AS c") + assert(df.schema.head.dataType === CharType(5)) + checkAnswer(df, Row("ab ")) + } + } + + test("SPARK-58796: createDataFrame allows CHAR/VARCHAR when standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val df = spark.range(1).map(_.toString).toDF() + val schema = new StructType().add("id", CharType(5)) + val created = spark.createDataFrame(df.collectAsList(), schema) + assert(created.schema.head.dataType === CharType(5)) + checkAnswer(created, Row("0 ")) + + // RowEncoder must retain a declared collation on the constrained type, not rebuild + // CharType(length) / VarcharType(length) with the default collation. + val collated = new StructType() + .add("c", CharType(5, "UTF8_LCASE")) + .add("v", VarcharType(5, "UTF8_LCASE")) + val collatedDf = spark.createDataFrame( + java.util.Arrays.asList(Row("ab", "cd")), collated) + assert(collatedDf.schema("c").dataType === CharType(5, "UTF8_LCASE")) + assert(collatedDf.schema("v").dataType === VarcharType(5, "UTF8_LCASE")) + checkAnswer(collatedDf, Row("ab ", "cd")) + } + } + + test("SPARK-58803: Dataset/encoder/UDF CHAR/VARCHAR under standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + // createDataFrame / RowEncoder write-side: pad CHAR, reject oversize. + val charSchema = new StructType().add("c", CharType(3)) + checkAnswer( + spark.createDataFrame(java.util.Arrays.asList(Row("ab")), charSchema), + Row("ab ")) + checkError( + exception = intercept[SparkRuntimeException] { + spark.createDataFrame(java.util.Arrays.asList(Row("abcd")), charSchema).collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + val varcharSchema = new StructType().add("v", VarcharType(3)) + checkAnswer( + spark.createDataFrame(java.util.Arrays.asList(Row("ab")), varcharSchema), + Row("ab")) + // Oversize by trailing blanks only: trim just enough to fit the limit. + checkAnswer( + spark.createDataFrame(java.util.Arrays.asList(Row("abc ")), varcharSchema), + Row("abc")) + checkError( + exception = intercept[SparkRuntimeException] { + spark.createDataFrame(java.util.Arrays.asList(Row("abcd")), varcharSchema).collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + + // Explicit Encoders.CHAR / VARCHAR: typed Dataset write-side checks. + val charDs = spark.createDataset(Seq("ab"))(Encoders.CHAR(4)) + assert(charDs.schema.head.dataType === CharType(4)) + checkAnswer(charDs.toDF(), Row("ab ")) + checkError( + exception = intercept[SparkRuntimeException] { + spark.createDataset(Seq("abcde"))(Encoders.VARCHAR(3)).collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + + // UDF register: return type stays CHAR/VARCHAR; write-side pad / length apply. + spark.udf.register("std_char_udf", () => "B", CharType(3)) + spark.udf.register("std_varchar_udf", (x: String) => x, VarcharType(3)) + val charUdf = sql("SELECT std_char_udf() AS c") + assert(charUdf.schema.head.dataType === CharType(3)) + checkAnswer(charUdf, Row("B ")) + val varcharUdf = sql("SELECT std_varchar_udf('ab') AS v") + assert(varcharUdf.schema.head.dataType === VarcharType(3)) + checkAnswer(varcharUdf, Row("ab")) + checkError( + exception = intercept[SparkException] { + sql("SELECT std_varchar_udf('abcd')").collect() + }.getCause.asInstanceOf[SparkRuntimeException], + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + + // Java udf(..., returnType) path and Dataset.encoder from CHAR result schema. + val javaUdf = functions.udf( + new org.apache.spark.sql.api.java.UDF0[String] { + override def call(): String = "a" + }, + CharType(5)) + val javaUdfDf = spark.range(1).select(javaUdf().as("c")) + assert(javaUdfDf.schema.head.dataType === CharType(5)) + checkAnswer(javaUdfDf, Row("a ")) + assert(javaUdfDf.encoder.schema.head.dataType === CharType(5)) + + // Dataset.to: CHAR/VARCHAR target schema allowed; Cast applies store assignment. + withTable("std_cv_to") { + sql("CREATE TABLE std_cv_to (c CHAR(10), v VARCHAR(255)) USING parquet") + sql("INSERT INTO std_cv_to VALUES ('spark', 'awesome')") + val df = sql("SELECT * FROM std_cv_to") + assert(df.schema("c").dataType === CharType(10)) + assert(df.schema("v").dataType === VarcharType(255)) + val reordered = StructType.fromDDL("v VARCHAR(255), c CHAR(10)") + val toDf = df.to(reordered) + assert(toDf.schema.map(_.dataType) === Seq(VarcharType(255), CharType(10))) + checkAnswer(toDf, Row("awesome", "spark ")) + // Narrowing CHAR length is store assignment and must enforce length. + checkError( + exception = intercept[SparkRuntimeException] { + df.select($"c").to(new StructType().add("c", CharType(3))).collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3")) + } + + // DataFrameReader / DataStreamReader user schemas keep CHAR/VARCHAR. + val readerSchema = new StructType().add("id", CharType(5)) + val csvInput = spark.range(1).map(_.toString) + val csvDf = spark.read.schema(readerSchema).csv(csvInput) + assert(csvDf.schema.head.dataType === CharType(5)) + checkAnswer(csvDf, Row("0 ")) + val csvDfDdl = spark.read.schema("id VARCHAR(5)").csv(csvInput) + assert(csvDfDdl.schema.head.dataType === VarcharType(5)) + withTempPath { dir => + spark.range(1).write.save(dir.toString) + val streamDf = spark.readStream.schema(readerSchema).load(dir.toString) + assert(streamDf.schema.head.dataType === CharType(5)) + val streamDdl = spark.readStream.schema("id VARCHAR(5)").load(dir.toString) + assert(streamDdl.schema.head.dataType === VarcharType(5)) + } + } + } + + test("SPARK-58794: promotion unifies CHAR/VARCHAR with STRING at plain-string inputs") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + withTable("std_promote") { + sql("CREATE TABLE std_promote (c CHAR(5), v VARCHAR(5)) USING parquet") + sql("INSERT INTO std_promote VALUES ('ab', 'ab')") + + // Expressions requiring all their string inputs to share one type must accept a + // CHAR/VARCHAR argument alongside a STRING one by promoting it to STRING. + Seq( + "overlay(c PLACING 'x' FROM 1)" -> "xb ", + "overlay(v PLACING 'x' FROM 1)" -> "xb", + "string_agg(c, '-')" -> "ab ", + "listagg(c, '-')" -> "ab ", + "elt(1, c, 'x')" -> "ab ", + // right() is RuntimeReplaceable; its literal branches must agree with the substring + // branch, which promotion has already reduced to STRING. + "right(c, 2)" -> " ", + "left(c, 2)" -> "ab").foreach { case (expr, expected) => + val df = sql(s"SELECT $expr AS r FROM std_promote") + assert(df.schema.head.dataType === StringType, s"$expr should return STRING") + checkAnswer(df, Row(expected)) + } + + // Transforming expressions must not inherit the input's length constraint: each of these + // produces a value whose length differs from the CHAR(5) input. + Seq( + "reverse(c)" -> " ba", + "hex(c)" -> "6162202020", + "array_join(array(c, c), '-')" -> "ab -ab ").foreach { case (expr, expected) => + val df = sql(s"SELECT $expr AS r FROM std_promote") + assert(df.schema.head.dataType === StringType, s"$expr should return STRING") + checkAnswer(df, Row(expected)) + } + + // Promotion must not reach pass-through / least-common-type sites, which preserve + // CHAR/VARCHAR. + Seq( + "c", "coalesce(c, c)", "case when true then c else c end", "max(c)", + "element_at(array(c), 1)", "transform(array(c), x -> x)[0]", + "first_value(c) over (order by 1)").foreach { expr => + val df = sql(s"SELECT $expr AS r FROM std_promote") + assert(df.schema.head.dataType === CharType(5), s"$expr should stay CHAR(5)") + } + + // reverse() on non-string inputs is unaffected by the promotion. + assert(sql("SELECT reverse(array(1, 2)) AS r").schema.head.dataType === + ArrayType(IntegerType, containsNull = false)) + } + } + } + + // Pass-through and container functions that may keep CHAR(n)/VARCHAR(n): aggregates + // and ordering that return an input unchanged, null-handling, element access, + // array/map/struct constructors, and collection rearrangements that keep element types. + // Legitimacy is still per shape: reverse(array(c)) may keep CHAR, reverse(c) must not. + private val charVarcharPassThroughFunctions = Set( + "any_value", "approx_top_k", "approx_top_k_accumulate", "array", "array_agg", "array_compact", + "array_distinct", "array_max", "array_min", "array_repeat", "array_sort", "arrays_zip", + "coalesce", "collect_list", "collect_set", "collect_union", "concat", "explode", + "explode_outer", "first", "first_value", "flatten", "get", "greatest", "ifnull", "last", + "last_value", "least", "map", "map_concat", "map_entries", "map_keys", "map_values", "max", + "max_by", "measure", "min", "min_by", "mode", "named_struct", "nullif", "nullifzero", "nvl", + "nvl2", "reverse", "shuffle", "sort_array", "struct", "trim_array", "when") + + // String-transforming shapes of otherwise pass-through functions. These must reduce to + // unconstrained STRING even though the same function keeps CHAR on collection inputs. + private val charVarcharTransformingCalls = Set( + "concat(c)", "concat(c, c)", "concat(c, c, c)", "concat(c, 'x')", "concat('x', c)", + "reverse(c)") + + private val inventoryScalarShapes = Seq( + "%s(c)", "%s(c, c)", "%s(c, 'x')", "%s('x', c)", "%s(c, 1)", "%s(array(c))", + "%s(array(c), '-')") + + private val inventoryNestedShapes = Seq( + "%s(c, c, c)", "%s(array(array(c)))", "%s(named_struct('x', c))", + "%s(map(c, 1))", "%s(map(1, c))") + + private def assertNoInventoriedCharVarcharLeaks(argumentShapes: Seq[String]): Unit = { + val (passThroughLeaks, transformingLeaks) = + FunctionRegistry.functionSet.map(_.funcName).toSeq.sorted.flatMap { name => + argumentShapes.map(_.format(name)).flatMap { call => + // Most shapes do not typecheck for a given function; those are simply not evidence. + val keepsCharVarchar = + Try(sql(s"SELECT $call AS r FROM std_inventory").schema.head.dataType) + .toOption + .exists(CharVarcharUtils.hasCharVarchar) + Option.when(keepsCharVarchar && + (!charVarcharPassThroughFunctions.contains(name) || + charVarcharTransformingCalls.contains(call)))((name, call)) + } + }.partition { case (_, call) => !charVarcharTransformingCalls.contains(call) } + + assert(passThroughLeaks.isEmpty, + "these inventoried calls returned a CHAR/VARCHAR type; if they legitimately pass through " + + "their input type, add the function name to charVarcharPassThroughFunctions: " + + passThroughLeaks.map(_._2).mkString(", ")) + assert(transformingLeaks.isEmpty, + "these transforming calls returned a CHAR/VARCHAR type; fix the expression to return " + + "plain STRING: " + transformingLeaks.map(_._2).mkString(", ")) + } + + test("SPARK-58794: inventoried shapes do not leak CHAR/VARCHAR under standardSemantics") { + withTable("std_inventory") { + sql("CREATE TABLE std_inventory (c CHAR(5)) USING parquet") + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assertNoInventoriedCharVarcharLeaks(inventoryScalarShapes) + } + } + } + + test("SPARK-59016: nested inventoried shapes do not leak CHAR/VARCHAR") { + withTable("std_inventory") { + sql("CREATE TABLE std_inventory (c CHAR(5)) USING parquet") + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + assertNoInventoriedCharVarcharLeaks(inventoryNestedShapes) + } + } + } + + test("SPARK-58794: collated mixed-length LCT ignores collation strength for length") { + val mixedLength = + """SELECT coalesce( + | cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + | cast('bb' AS CHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin + val mixedStrength = + """SELECT coalesce( + | cast('a' AS CHAR(2) COLLATE UTF8_LCASE), + | cast(1 AS CHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin + + Seq( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true", + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true").foreach { case (key, value) => + withSQLConf(key -> value) { + assert(sql(mixedLength).schema.head.dataType === CharType(4, "UTF8_LCASE"), + s"$key=$value same-strength mixed CHAR lengths") + assert(sql(mixedStrength).schema.head.dataType === CharType(4, "UTF8_LCASE"), + s"$key=$value Implicit CHAR(2) vs Default CHAR(4) must widen, not narrow") + } + } + } + + test("SPARK-58794: typed CHAR Literal is re-padded when LCT widens the length") { + // CollationTypeCoercion.changeType used to `copy(dataType)` on Literal, which would + // leave CHAR(2) "a " as a CHAR(4) value without the extra pad. SQL CAST is a Cast + // node so goldens do not cover this; Literal.create does. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val c2 = Column(Literal.create("a", CharType(2, "UTF8_LCASE"))) + val c4 = Column(Literal.create("bb", CharType(4, "UTF8_LCASE"))) + val coalesced = functions.coalesce(c2, c4) + val df = spark.range(1).select(coalesced.as("c")) + assert(df.schema.head.dataType === CharType(4, "UTF8_LCASE")) + checkAnswer(df, Row("a ")) + } + } + + test("SPARK-58794: CHAR/VARCHAR vs non-string follow STRING for compare and COALESCE") { + // Comparisons promote the string side to the other atomic type; COALESCE uses the same + // STRING promotion (and the same BOOLEAN path). CHAR/VARCHAR extend StringType, so they must + // match the STRING analogue on schema, values, and analysis errors. Exact-length CHAR keeps + // padding from confounding comparisons with the unpadded STRING value. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + def followString(cvSql: String, strSql: String): Unit = { + val strResult = Try { + val df = sql(strSql) + (df.schema.map(_.dataType), df.collect().toSeq) + } + val cvResult = Try { + val df = sql(cvSql) + (df.schema.map(_.dataType), df.collect().toSeq) + } + (strResult, cvResult) match { + case (scala.util.Success((st, sr)), scala.util.Success((ct, cr))) => + assert(ct === st, s"schema:\n $cvSql -> $ct\n $strSql -> $st") + assert(cr === sr, s"rows:\n $cvSql -> $cr\n $strSql -> $sr") + case (scala.util.Failure(eStr: SparkThrowable), + scala.util.Failure(eCv: SparkThrowable)) => + assert(eCv.getCondition === eStr.getCondition, + s"$cvSql vs $strSql: ${eCv.getCondition} != ${eStr.getCondition}") + case (s, c) => + fail(s"$cvSql vs $strSql: STRING success=${s.isSuccess} CV success=${c.isSuccess}") + } + } + + val combos = Seq( + ("cast('123' AS CHAR(3))", "cast('123' AS VARCHAR(3))", "cast('123' AS STRING)", "123"), + ("cast('1.5' AS CHAR(3))", "cast('1.5' AS VARCHAR(3))", "cast('1.5' AS STRING)", "1.5"), + ("cast('2020-01-02' AS CHAR(10))", "cast('2020-01-02' AS VARCHAR(10))", + "cast('2020-01-02' AS STRING)", "date'2020-01-02'"), + ("cast('2020-01-02 03:04:05' AS CHAR(19))", "cast('2020-01-02 03:04:05' AS VARCHAR(19))", + "cast('2020-01-02 03:04:05' AS STRING)", "timestamp'2020-01-02 03:04:05'"), + ("cast('true' AS CHAR(4))", "cast('true' AS VARCHAR(4))", "cast('true' AS STRING)", "true"), + // Padded CHAR vs the same padded STRING / VARCHAR bytes. + ("cast('123' AS CHAR(5))", "cast('123 ' AS VARCHAR(5))", "cast('123 ' AS STRING)", "123")) + + val templates: Seq[(String, String) => String] = Seq( + (cv, other) => s"SELECT typeof(coalesce($cv, $other))", + (cv, other) => s"SELECT coalesce($cv, $other)", + (cv, other) => s"SELECT $cv = $other", + (cv, other) => s"SELECT $cv < $other", + (cv, other) => s"SELECT $cv IN ($other)") + + for ((charExpr, varcharExpr, stringExpr, other) <- combos) { + for (mk <- templates) { + followString(mk(charExpr, other), mk(stringExpr, other)) + followString(mk(varcharExpr, other), mk(stringExpr, other)) + } + } + } + } + + test("SPARK-58794: ImplicitTypeCasts promotes CHAR/VARCHAR to STRING") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + def analyzed(sqlText: String): LogicalPlan = sql(sqlText).queryExecution.analyzed + def exprs(plan: LogicalPlan): Seq[Expression] = + plan.flatMap(_.expressions.flatMap(e => e +: e.collect { case c => c })) + + val upper = exprs(analyzed("SELECT upper(cast('ab' AS CHAR(2)))")) + .collect { case u: Upper => u }.head + assert(upper.child.dataType === StringType) + + val concat = exprs(analyzed( + "SELECT concat(cast('a' AS CHAR(2)), cast('b' AS CHAR(3)))")) + .collect { case c: Concat => c }.head + assert(concat.children.forall(_.dataType == StringType)) + + val arrayJoin = exprs(analyzed( + "SELECT array_join(array(cast('ab' AS CHAR(5))), '-')")) + .collect { case a: ArrayJoin => a }.head + assert(arrayJoin.array.dataType === ArrayType(StringType, containsNull = false)) + + val stringToMap = exprs(analyzed("SELECT str_to_map(cast('a:1' AS CHAR(5)))")) + .collect { case s: StringToMap => s }.head + assert(stringToMap.first.dataType === StringType) + assert(stringToMap.dataType === MapType(StringType, StringType, valueContainsNull = true)) + } + } + + test("SPARK-58794: parameterized CHAR/VARCHAR lengths under standardSemantics") { + // Length positions accept parameter markers (`integerValue` -> `parameterMarker`). Under + // standardSemantics the bound type stays first-class: CAST keeps CHAR/VARCHAR, pads and + // enforces length, and DDL schemas retain the substituted n. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + val charDf = spark.sql("SELECT cast('ab' AS CHAR(:n)) AS c", Map("n" -> 5)) + assert(charDf.schema.head.dataType === CharType(5)) + checkAnswer( + spark.sql("SELECT concat('<', cast('ab' AS CHAR(:n)), '>')", Map("n" -> 5)), + Row("<ab >")) + + val varcharDf = spark.sql("SELECT cast('hello' AS VARCHAR(?)) AS c", Array(5)) + assert(varcharDf.schema.head.dataType === VarcharType(5)) + checkAnswer( + spark.sql("SELECT cast('abcdef' AS VARCHAR(?))", Array(2)), + Row("ab")) + + withTable("param_varchar", "param_char") { + spark.sql( + "CREATE TABLE param_varchar (c VARCHAR(:n)) USING parquet", Map("n" -> 7)) + assert(spark.table("param_varchar").schema.head.dataType === VarcharType(7)) + spark.sql("CREATE TABLE param_char (c CHAR(?)) USING parquet", Array(4)) + assert(spark.table("param_char").schema.head.dataType === CharType(4)) + } + + // Non-integral / negative lengths fail when substituted into the length position. + checkError( + exception = intercept[ParseException] { + spark.sql("SELECT cast('a' AS CHAR(:n))", Map("n" -> -1)) + }, + condition = "PARSE_SYNTAX_ERROR", + parameters = Map("error" -> "'-'", "hint" -> ""), + context = ExpectedContext( + fragment = "SELECT cast('a' AS CHAR(:n))", + start = 0, + stop = 27)) + checkError( + exception = intercept[ParseException] { + spark.sql("SELECT cast('a' AS CHAR(:n))", Map("n" -> 1.5)) + }, + condition = "PARSE_SYNTAX_ERROR", + parameters = Map("error" -> "'1.5D'", "hint" -> ""), + context = ExpectedContext( + fragment = "SELECT cast('a' AS CHAR(:n))", + start = 0, + stop = 27)) + } + } + + test("SPARK-58794: language surfaces keep CHAR/VARCHAR under standardSemantics") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + // CTAS / CREATE VIEW inherit projected CHAR/VARCHAR. + withTable("std_src", "std_ctas") { + sql("CREATE TABLE std_src (c CHAR(5), v VARCHAR(5)) USING parquet") + sql("INSERT INTO std_src VALUES ('ab', 'ab')") + sql("CREATE TABLE std_ctas USING parquet AS SELECT c, v FROM std_src") + assert(spark.table("std_ctas").schema.map(_.dataType) === + Seq(CharType(5), VarcharType(5))) + checkAnswer( + sql("SELECT concat('<', c, '>'), concat('<', v, '>') FROM std_ctas"), + Row("<ab >", "<ab>")) + val cteDf = sql( + "WITH t AS (SELECT c, v FROM std_src) SELECT typeof(c), typeof(v) FROM t") + checkAnswer(cteDf, Row("char(5)", "varchar(5)")) + } + withTable("std_view_src") { + withView("std_cv_view", "std_cv_view_v") { + sql("CREATE TABLE std_view_src (c CHAR(4), v VARCHAR(4)) USING parquet") + sql("INSERT INTO std_view_src VALUES ('xy', 'ab')") + sql("CREATE VIEW std_cv_view AS SELECT c FROM std_view_src") + sql("CREATE VIEW std_cv_view_v AS SELECT v FROM std_view_src") + assert(spark.table("std_cv_view").schema.head.dataType === CharType(4)) + assert(spark.table("std_cv_view_v").schema.head.dataType === VarcharType(4)) + checkAnswer( + sql("SELECT concat('<', c, '>') FROM std_cv_view"), Row("<xy >")) + checkAnswer(sql("SELECT v FROM std_cv_view_v"), Row("ab")) + } + } + + // ALTER COLUMN equal-length CHAR/VARCHAR remains supported with first-class types. + // VARCHAR widen / CHAR->VARCHAR are allowed by CheckAnalysis on V2 tables + // (see DSV2CharVarcharDDLTestSuite); V1 file-source ALTER only evolves collation + // (same StringConstraint), so length changes stay rejected there. + withTable("std_alter") { + sql("CREATE TABLE std_alter (c CHAR(4), v VARCHAR(4)) USING parquet") + sql("ALTER TABLE std_alter CHANGE COLUMN c TYPE CHAR(4)") + sql("ALTER TABLE std_alter CHANGE COLUMN v TYPE VARCHAR(4)") + assert(spark.table("std_alter").schema.map(_.dataType) === + Seq(CharType(4), VarcharType(4))) + intercept[AnalysisException] { + sql("ALTER TABLE std_alter CHANGE COLUMN c TYPE CHAR(5)") + } + intercept[AnalysisException] { + sql("ALTER TABLE std_alter CHANGE COLUMN v TYPE VARCHAR(5)") + } + } + + // Session variables: DECLARE / SET keep the type and apply CAST assignment. + sql("DECLARE OR REPLACE VARIABLE std_char_var CHAR(4)") + sql("DECLARE OR REPLACE VARIABLE std_varchar_var VARCHAR(4)") + try { + sql("SET VARIABLE std_char_var = 'ab'") + val charVarDf = sql("SELECT std_char_var AS c") + assert(charVarDf.schema.head.dataType === CharType(4)) + checkAnswer(sql("SELECT concat('<', std_char_var, '>')"), Row("<ab >")) + // Oversize by trailing blanks only is trimmed to fit CHAR(n). + sql("SET VARIABLE std_char_var = 'abcd '") + checkAnswer(sql("SELECT concat('<', std_char_var, '>')"), Row("<abcd>")) + intercept[SparkRuntimeException] { + sql("SET VARIABLE std_char_var = 'abcde'").collect() + } + + sql("SET VARIABLE std_varchar_var = 'ab'") + val varcharVarDf = sql("SELECT std_varchar_var AS v") + assert(varcharVarDf.schema.head.dataType === VarcharType(4)) + checkAnswer(sql("SELECT concat('<', std_varchar_var, '>')"), Row("<ab>")) + // Oversize by trailing blanks only is trimmed to fit. + sql("SET VARIABLE std_varchar_var = 'abcd '") + checkAnswer(sql("SELECT std_varchar_var"), Row("abcd")) + intercept[SparkRuntimeException] { + sql("SET VARIABLE std_varchar_var = 'abcde'").collect() + } + } finally { + sql("DROP TEMPORARY VARIABLE IF EXISTS std_char_var") + sql("DROP TEMPORARY VARIABLE IF EXISTS std_varchar_var") + } + + // SQL scripting local variables keep CHAR/VARCHAR inside a compound statement. + val localVarScript = + """ + |BEGIN + | DECLARE c CHAR(4); + | DECLARE v VARCHAR(4); + | SET c = 'ab'; + | SET v = 'cd'; + | SELECT typeof(c), concat('<', c, '>'), typeof(v), concat('<', v, '>'); + |END + |""".stripMargin + val localVarDf = sql(localVarScript) + assert(localVarDf.schema.map(_.dataType) === + Seq(StringType, StringType, StringType, StringType)) + checkAnswer(localVarDf, Row("char(4)", "<ab >", "varchar(4)", "<cd>")) + // Trailing-blank trim on local SET into CHAR/VARCHAR. + checkAnswer( + sql( + """ + |BEGIN + | DECLARE c CHAR(4); + | DECLARE v VARCHAR(4); + | SET c = 'abcd '; + | SET v = 'abcd '; + | SELECT concat('<', c, '>'), v; + |END + |""".stripMargin), + Row("<abcd>", "abcd")) + intercept[SparkRuntimeException] { + sql( + """ + |BEGIN + | DECLARE c CHAR(4); + | SET c = 'abcde'; + |END + |""".stripMargin).collect() + } + intercept[SparkRuntimeException] { + sql( + """ + |BEGIN + | DECLARE v VARCHAR(4); + | SET v = 'abcde'; + |END + |""".stripMargin).collect() + } + + // Cursor FETCH INTO CHAR/VARCHAR locals applies store assignment (pad / length). + withSQLConf(SQLConf.SQL_SCRIPTING_CURSOR_ENABLED.key -> "true") { + val cursorScript = + """ + |BEGIN + | DECLARE fetched_c CHAR(4); + | DECLARE fetched_v VARCHAR(4); + | DECLARE cur CURSOR FOR + | SELECT cast('ab' AS CHAR(4)) AS c, cast('cd' AS VARCHAR(4)) AS v; + | OPEN cur; + | FETCH cur INTO fetched_c, fetched_v; + | SELECT typeof(fetched_c), concat('<', fetched_c, '>'), + | typeof(fetched_v), concat('<', fetched_v, '>'); + | CLOSE cur; + |END + |""".stripMargin + checkAnswer( + sql(cursorScript), + Row("char(4)", "<ab >", "varchar(4)", "<cd>")) + + // FETCH plain STRING into CHAR pads via assignment cast. + checkAnswer( + sql( + """ + |BEGIN + | DECLARE fetched CHAR(4); + | DECLARE cur CURSOR FOR SELECT 'ab' AS c; + | OPEN cur; + | FETCH cur INTO fetched; + | SELECT typeof(fetched), concat('<', fetched, '>'); + | CLOSE cur; + |END + |""".stripMargin), + Row("char(4)", "<ab >")) + + // FETCH into a wider CHAR pads; trailing blanks trim into a shorter target. + checkAnswer( + sql( + """ + |BEGIN + | DECLARE fetched CHAR(5); + | DECLARE cur CURSOR FOR SELECT cast('xy' AS CHAR(2)) AS c; + | OPEN cur; + | FETCH cur INTO fetched; + | SELECT concat('<', fetched, '>'); + | CLOSE cur; + |END + |""".stripMargin), + Row("<xy >")) + checkAnswer( + sql( + """ + |BEGIN + | DECLARE fetched_c CHAR(4); + | DECLARE fetched_v VARCHAR(4); + | DECLARE cur CURSOR FOR + | SELECT cast('abcd ' AS CHAR(5)) AS c, cast('abcd ' AS VARCHAR(5)) AS v; + | OPEN cur; + | FETCH cur INTO fetched_c, fetched_v; + | SELECT concat('<', fetched_c, '>'), fetched_v; + | CLOSE cur; + |END + |""".stripMargin), + Row("<abcd>", "abcd")) + intercept[SparkRuntimeException] { + sql( + """ + |BEGIN + | DECLARE fetched VARCHAR(2); + | DECLARE cur CURSOR FOR SELECT cast('abcd' AS VARCHAR(4)) AS v; + | OPEN cur; + | FETCH cur INTO fetched; + | CLOSE cur; + |END + |""".stripMargin).collect() + } + } + + // SQL FUNCTION params/RETURNS apply store assignment (pad, blank-trim, overflow). + sql("CREATE OR REPLACE TEMPORARY FUNCTION std_char_fn() RETURNS CHAR(3) RETURN 'a'") + sql( + """CREATE OR REPLACE TEMPORARY FUNCTION std_varchar_ret() + |RETURNS VARCHAR(3) RETURN 'ab'""".stripMargin) + sql( + """CREATE OR REPLACE TEMPORARY FUNCTION std_char_param(x CHAR(3)) + |RETURNS CHAR(3) RETURN x""".stripMargin) + sql( + """CREATE OR REPLACE TEMPORARY FUNCTION std_varchar_param(x VARCHAR(3)) + |RETURNS VARCHAR(3) RETURN x""".stripMargin) + try { + val fnDf = sql("SELECT std_char_fn() AS c") + assert(fnDf.schema.head.dataType === CharType(3)) + checkAnswer(sql("SELECT concat('<', std_char_fn(), '>')"), Row("<a >")) + val retVarcharDf = sql("SELECT std_varchar_ret() AS v") + assert(retVarcharDf.schema.head.dataType === VarcharType(3)) + checkAnswer(retVarcharDf, Row("ab")) + + // STRING -> CHAR(n) param: pad; trailing blanks trim; non-blank overflow errors. + val charParamDf = sql("SELECT std_char_param('a') AS c") + assert(charParamDf.schema.head.dataType === CharType(3)) + checkAnswer(sql("SELECT concat('<', std_char_param('a'), '>')"), Row("<a >")) + checkAnswer( + sql("SELECT concat('<', std_char_param('abc '), '>')"), + Row("<abc>")) + intercept[SparkRuntimeException] { + sql("SELECT std_char_param('abcd')").collect() + } + + val paramDf = sql("SELECT std_varchar_param('ab') AS v") + assert(paramDf.schema.head.dataType === VarcharType(3)) + checkAnswer(paramDf, Row("ab")) + checkAnswer(sql("SELECT std_varchar_param('abc ')"), Row("abc")) + intercept[SparkRuntimeException] { + sql("SELECT std_varchar_param('abcd')").collect() + } + } finally { + sql("DROP TEMPORARY FUNCTION IF EXISTS std_char_fn") + sql("DROP TEMPORARY FUNCTION IF EXISTS std_varchar_ret") + sql("DROP TEMPORARY FUNCTION IF EXISTS std_char_param") + sql("DROP TEMPORARY FUNCTION IF EXISTS std_varchar_param") + } + + // ORC catalog tables stamp the catalyst type so typeof survives write/read. + withTable("std_orc") { + sql("CREATE TABLE std_orc (c CHAR(5), v VARCHAR(5)) USING orc") + sql("INSERT INTO std_orc VALUES ('ab', 'cd')") + assert(spark.table("std_orc").schema.map(_.dataType) === + Seq(CharType(5), VarcharType(5))) + checkAnswer( + sql("SELECT concat('<', c, '>'), concat('<', v, '>') FROM std_orc"), + Row("<ab >", "<cd>")) + } + + // File-only ORC inference recovers the catalyst type stamped on write. + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(1).selectExpr("cast('ab' AS CHAR(4)) AS c") + .write.mode("overwrite").orc(path) + val orcDf = spark.read.orc(path) + assert(orcDf.schema.head.dataType === CharType(4)) + checkAnswer(orcDf.selectExpr("concat('<', c, '>')"), Row("<ab >")) + // Reading with first-class types off replaces CHAR with STRING even if the + // file was stamped under standardSemantics. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") { + val readOff = spark.read.orc(path) + assert(readOff.schema.head.dataType === StringType) + } + } + // First-class types off: CAST CHAR is STRING before the writer, so ORC does not stamp CHAR. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(1).selectExpr("cast('ab' AS CHAR(4)) AS c") + .write.mode("overwrite").orc(path) + assert(spark.read.orc(path).schema.head.dataType === StringType) + } + } + // preserveCharVarcharTypeInfo also keeps first-class types, so write still stamps CHAR. + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false", + SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> "true") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(1).selectExpr("cast('ab' AS CHAR(4)) AS c") + .write.mode("overwrite").orc(path) + assert(spark.read.orc(path).schema.head.dataType === CharType(4)) + } + } + // ORC stamps collated unbounded STRING as plain "string"; the inferred type is the + // same as Avro, which omits the catalyst property on unbounded STRING. + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(1).selectExpr("cast('ab' AS STRING COLLATE UTF8_LCASE) AS c") + .write.mode("overwrite").orc(path) + assert(spark.read.orc(path).schema.head.dataType === StringType) + } + + // Avro conversion stamps CHAR/VARCHAR (including nested fields and CHAR map keys). + // The avro data source lives in connector/avro; sql/core still owns toAvroType, + // serializer, and deserializer, so round-trip CHAR/VARCHAR here with a DataFileWriter. + val converters = org.apache.spark.sql.avro.SchemaConverters + Seq(CharType(5), VarcharType(7)).foreach { dt => + val avro = converters.toAvroType(dt, nullable = false) + val back = converters.toSqlType(avro).dataType + assert(back === dt, s"Avro round-trip lost $dt, got $back") + } + val nestedSchema = new StructType() + .add("c", CharType(4)) + .add("s", new StructType().add("f", VarcharType(3))) + .add("m", MapType(CharType(2), VarcharType(3))) + val nestedAvro = converters.toAvroType(nestedSchema, nullable = false) + assert(converters.toSqlType(nestedAvro).dataType === nestedSchema) + val badStamp = org.apache.avro.SchemaBuilder.builder().stringType() + badStamp.addProp("spark.sql.catalyst.type", "int") + val badStampErr = intercept[Exception] { + converters.toSqlType(badStamp) + } + assert(badStampErr.getMessage.contains("STRING subtype") || + Option(badStampErr.getCause).exists(_.getMessage.contains("STRING subtype"))) + // Flag-off replace happens before Avro sees the type, so CHAR is not stamped. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "false") { + val replaced = CharVarcharUtils.replaceCharVarcharWithString(CharType(4)) + assert(replaced === StringType) + val replacedAvro = converters.toAvroType(replaced, nullable = false) + assert(replacedAvro.getProp("spark.sql.catalyst.type") === null) + assert(converters.toSqlType(replacedAvro).dataType === StringType) + } + + withTempPath { file => + val ser = new org.apache.spark.sql.avro.AvroSerializer( + nestedSchema, nestedAvro, nullable = false) + val row = org.apache.spark.sql.catalyst.InternalRow( + org.apache.spark.unsafe.types.UTF8String.fromString("ab "), + org.apache.spark.sql.catalyst.InternalRow( + org.apache.spark.unsafe.types.UTF8String.fromString("xy")), + new org.apache.spark.sql.catalyst.util.ArrayBasedMapData( + new org.apache.spark.sql.catalyst.util.GenericArrayData(Array( + org.apache.spark.unsafe.types.UTF8String.fromString("k "))), + new org.apache.spark.sql.catalyst.util.GenericArrayData(Array( + org.apache.spark.unsafe.types.UTF8String.fromString("v"))))) + val record = ser.serialize(row) + .asInstanceOf[org.apache.avro.generic.GenericRecord] + val writer = new org.apache.avro.file.DataFileWriter( + new org.apache.avro.generic.GenericDatumWriter[org.apache.avro.generic.GenericRecord]( + nestedAvro)) + writer.create(nestedAvro, file) + writer.append(record) + writer.close() + val reader = new org.apache.avro.file.DataFileReader( + file, + new org.apache.avro.generic.GenericDatumReader[org.apache.avro.generic.GenericRecord]()) + try { + assert(converters.toSqlType(reader.getSchema).dataType === nestedSchema) + val deser = new org.apache.spark.sql.avro.AvroDeserializer( + nestedAvro, nestedSchema, "CORRECTED", false, "", -1) + val back = deser.deserialize(reader.next()).get + .asInstanceOf[org.apache.spark.sql.catalyst.InternalRow] + assert(back.getUTF8String(0).toString === "ab ") + val nestedField = back.getStruct(1, 1) + assert(nestedField.getUTF8String(0).toString === "xy") + val mapData = back.getMap(2) + assert(mapData.numElements() === 1) + assert(mapData.keyArray().getUTF8String(0).toString === "k ") + assert(mapData.valueArray().getUTF8String(0).toString === "v") + } finally { + reader.close() + } + } + + // JSON / CSV keep a user-specified CHAR/VARCHAR schema under the flag. + withTempPath { dir => + val path = dir.getCanonicalPath + spark.range(1).selectExpr("cast(id AS STRING) AS c").write.mode("overwrite") + .json(s"$path/json") + val jsonDf = spark.read.schema("c CHAR(5)").json(s"$path/json") + assert(jsonDf.schema.head.dataType === CharType(5)) + checkAnswer(jsonDf.selectExpr("concat('<', c, '>')"), Row("<0 >")) + + spark.range(1).selectExpr("cast(id AS STRING) AS c").write.mode("overwrite") + .option("header", "true").csv(s"$path/csv") + val csvDf = spark.read.schema("c VARCHAR(5)").option("header", "true") + .csv(s"$path/csv") + assert(csvDf.schema.head.dataType === VarcharType(5)) + checkAnswer(csvDf, Row("0")) + } + } + } + + test("SPARK-58802: single-pass resolver agrees with fixed-point under standardSemantics") { + // Dual run defaults to on under tests, but pin it explicitly so this coverage cannot be + // silently lost: the HybridAnalyzer compares output schema and normalized plan across the + // two analyzers and fails with HYBRID_ANALYZER_EXCEPTION on any divergence. Resolver has no + // Char/Varchar-specific logic; it inherits Expression.dataType and shared TypeCoercion, so + // this matrix is the proof that least common type, CAST, and promotion stay aligned. + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true", + SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "true", + SQLConf.ANALYZER_DUAL_RUN_SAMPLE_RATE.key -> "1.0", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_EXPOSE_RESOLVER_GUARD_FAILURE.key -> "true") { + def sql(sqlText: String): DataFrame = { + val parsed = spark.sessionState.sqlParser.parsePlan(sqlText) + val unsupportedReason = + new ResolverGuard(spark.sessionState.catalogManager).apply(parsed).planUnsupportedReason + assert( + unsupportedReason.isEmpty, + s"ResolverGuard skipped dual-run for [$sqlText]: $unsupportedReason") + spark.sql(sqlText) + } + + // CAST / try_cast introduce the type. + assert(sql("SELECT CAST('ab' AS CHAR(5)) AS c").schema.head.dataType === CharType(5)) + assert(sql("SELECT CAST('hello' AS VARCHAR(5)) AS c").schema.head.dataType === + VarcharType(5)) + assert(sql("SELECT try_cast('abcdef' AS CHAR(2)) AS c").schema.head.dataType === + CharType(2)) + checkAnswer(sql("SELECT try_cast('abcdef' AS VARCHAR(2)) AS c"), Row("ab")) + checkAnswer( + sql("SELECT coalesce(CAST('abcdef' AS VARCHAR(2)), CAST('x' AS VARCHAR(4))) AS c"), + Row("ab")) + checkAnswer( + sql("SELECT CAST('abcdef' AS VARCHAR(2)) IN (CAST('ab' AS VARCHAR(4)))"), + Row(true)) + checkAnswer( + sql("""SELECT coalesce( + | CAST('abcdef' AS VARCHAR(2) COLLATE UTF8_LCASE), + | CAST('x' AS VARCHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin), + Row("ab")) + + // Least common type: COALESCE / CASE / NULL / CHAR+VARCHAR / CHAR+STRING. + assert(sql( + "SELECT coalesce(CAST('a' AS VARCHAR(3)), CAST('bb' AS VARCHAR(7))) AS c") + .schema.head.dataType === VarcharType(7)) + assert(sql( + "SELECT coalesce(CAST('a' AS CHAR(2)), CAST('bb' AS VARCHAR(4))) AS c") + .schema.head.dataType === VarcharType(4)) + assert(sql( + "SELECT coalesce(CAST('a' AS CHAR(2)), 'bb') AS c") + .schema.head.dataType === StringType) + assert(sql( + "SELECT coalesce(CAST('a' AS CHAR(5)), CAST(NULL AS CHAR(5))) AS c") + .schema.head.dataType === CharType(5)) + assert(sql( + "SELECT CASE WHEN true THEN CAST('a' AS CHAR(2)) ELSE CAST('bb' AS CHAR(4)) END AS c") + .schema.head.dataType === CharType(4)) + assert(sql( + "SELECT CASE WHEN false THEN CAST('a' AS VARCHAR(2)) ELSE CAST('bb' AS CHAR(4)) END AS c") + .schema.head.dataType === VarcharType(4)) + + // IN-list common type (side condition uses LCT; result is boolean). + checkAnswer( + sql("SELECT CAST('a' AS CHAR(2)) IN (CAST('a ' AS CHAR(2)), CAST('bbb' AS VARCHAR(3)))"), + Row(true)) + + // Transforming operators return STRING. + assert(sql("SELECT upper(CAST('ab' AS CHAR(2))) AS c").schema.head.dataType === StringType) + assert(sql("SELECT lower(CAST('AB' AS VARCHAR(2))) AS c").schema.head.dataType === + StringType) + assert(sql("SELECT CAST('a' AS CHAR(1)) || CAST('b' AS VARCHAR(1)) AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT concat(CAST('a' AS CHAR(2)), CAST('b' AS CHAR(3))) AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT substr(CAST('hello' AS VARCHAR(5)), 1, 2) AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT trim(CAST('ab ' AS CHAR(4))) AS c").schema.head.dataType === StringType) + assert(sql("SELECT regexp_replace(CAST('ab' AS CHAR(2)), 'a', 'x') AS c") + .schema.head.dataType === StringType) + assert(sql("SELECT mask(CAST('ab' AS CHAR(2))) AS c").schema.head.dataType === StringType) + assert(sql("SELECT split(CAST('a,b' AS CHAR(3)), ',') AS c").schema.head.dataType === + ArrayType(StringType, containsNull = false)) + // Promotion after least common type: coalesce stays CHAR, upper widens to STRING. + assert(sql( + "SELECT upper(coalesce(CAST('a' AS CHAR(2)), CAST('b' AS CHAR(4)))) AS c") + .schema.head.dataType === StringType) + + // Set-operation LCT. + val union = sql( + """SELECT CAST('a' AS VARCHAR(3)) AS c + |UNION ALL + |SELECT CAST('abcd' AS VARCHAR(8)) AS c""".stripMargin) + assert(union.schema.head.dataType === VarcharType(8)) + checkAnswer(union, Seq(Row("a"), Row("abcd"))) + + val intersect = sql( + """SELECT CAST('ab' AS CHAR(2)) AS c + |INTERSECT + |SELECT CAST('ab' AS CHAR(4)) AS c""".stripMargin) + assert(intersect.schema.head.dataType === CharType(4)) + checkAnswer(intersect, Seq(Row("ab "))) + + // Nested types keep CHAR/VARCHAR through analysis. + assert(sql("SELECT array(CAST('a' AS CHAR(2)), CAST('bb' AS CHAR(3))) AS c") + .schema.head.dataType === ArrayType(CharType(3), containsNull = false)) + assert(sql("SELECT struct(CAST('a' AS CHAR(2)) AS f) AS c") + .schema.head.dataType === + StructType(Seq(StructField("f", CharType(2), nullable = false)))) + + // Collated mixed-length LCT (the CollationTypeCoercion equal-strength and + // mixed-strength paths). Set ops have a separate resolver path. + assert(sql( + """SELECT coalesce( + | CAST('a' AS CHAR(2) COLLATE UTF8_LCASE), + | CAST('bb' AS CHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin) + .schema.head.dataType === CharType(4, "UTF8_LCASE")) + assert(sql( + """SELECT coalesce( + | CAST('a' AS CHAR(2) COLLATE UTF8_LCASE), + | CAST(1 AS CHAR(4) COLLATE UTF8_LCASE)) AS c""".stripMargin) + .schema.head.dataType === CharType(4, "UTF8_LCASE")) + checkAnswer( + sql("SELECT CAST('a' AS CHAR(2) COLLATE UTF8_LCASE) = " + + "CAST('a' AS CHAR(4) COLLATE UTF8_LCASE)"), + Row(true)) + checkAnswer( + sql("SELECT CAST('a' AS CHAR(2) COLLATE UTF8_LCASE) IN " + + "(CAST('a' AS CHAR(4) COLLATE UTF8_LCASE))"), + Row(true)) + checkAnswer( + sql("SELECT CAST('a' AS CHAR(2) COLLATE UTF8_LCASE) = " + + "CAST('a' AS VARCHAR(2) COLLATE UTF8_LCASE)"), + Row(false)) + checkAnswer( + sql("SELECT CAST('a' AS CHAR(2) COLLATE UTF8_LCASE) IN " + + "(CAST('a' AS VARCHAR(2) COLLATE UTF8_LCASE))"), + Row(false)) + + val mixedCharUnion = sql( + """SELECT CAST('a' AS CHAR(2)) AS c + |UNION ALL + |SELECT CAST('bb' AS CHAR(4)) AS c""".stripMargin) + assert(mixedCharUnion.schema.head.dataType === CharType(4)) + checkAnswer(mixedCharUnion, Seq(Row("a "), Row("bb "))) + + // Bare column references keep the declared type through dual-run analysis. + withTable("char_varchar_dual_run") { + spark.sql("CREATE TABLE char_varchar_dual_run (c CHAR(5), v VARCHAR(5)) USING parquet") + spark.sql("INSERT INTO char_varchar_dual_run VALUES ('ab', 'ab')") + val df = sql("SELECT c, v FROM char_varchar_dual_run") + assert(df.schema("c").dataType === CharType(5)) + assert(df.schema("v").dataType === VarcharType(5)) + checkAnswer(df, Row("ab ", "ab")) + } + } + } + test("invalidate char/varchar in functions") { checkError( exception = intercept[AnalysisException] { @@ -1010,6 +2185,47 @@ class BasicCharVarcharTestSuite extends SharedSparkSession { } } } + + test("SPARK-59001: empty CHAR/VARCHAR partition values become null like STRING") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + // CHAR(n>0) pads '' to spaces; CHAR(0) is the empty CHAR that empty2null should treat + // the same as VARCHAR/STRING. + Seq("CHAR(0)", "VARCHAR(5)").foreach { typ => + withTempPath { path => + sql(s"SELECT 0 AS id, CAST('' AS $typ) AS p UNION ALL SELECT 1, CAST(NULL AS $typ)") + .write.mode("overwrite").partitionBy("p").parquet(path.getCanonicalPath) + val df = spark.read.parquet(path.getCanonicalPath) + checkAnswer(df.where("p IS NULL").select("id"), Seq(Row(0), Row(1))) + val dirs = path.listFiles().filterNot( + f => f.getName.startsWith(".") || f.getName.startsWith("_")) + assert(dirs.length === 1, dirs.map(_.getName).mkString(",")) + } + } + } + } + + test("SPARK-59001: text datasource accepts CHAR/VARCHAR as a string family type") { + Seq("text", "").foreach { useV1SourceList => + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true", + SQLConf.USE_V1_SOURCE_LIST.key -> useV1SourceList) { + withTempPath { dir => + val path = dir.getCanonicalPath + sql("SELECT CAST('ab' AS CHAR(4)) AS value").write.mode("overwrite").text(path) + val df = spark.read.schema("value CHAR(4)").text(path) + assert(df.schema.head.dataType === CharType(4)) + checkAnswer(df.selectExpr("concat('<', value, '>')"), Row("<ab >")) + } + withTempPath { dir => + val path = dir.getCanonicalPath + sql("SELECT CAST('cd' AS VARCHAR(5)) AS value").write.mode("overwrite").text(path) + val df = spark.read.schema("value VARCHAR(5)").text(path) + assert(df.schema.head.dataType === VarcharType(5)) + checkAnswer(df, Row("cd")) + } + } + } + } } class FileSourceCharVarcharTestSuite extends CharVarcharTestSuite with SharedSparkSession { @@ -1040,6 +2256,90 @@ class FileSourceCharVarcharTestSuite extends CharVarcharTestSuite with SharedSpa } } + test("SPARK-58801: standardSemantics scan pads CHAR and errors on oversize") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + withTempPath { dir => + withTable("t") { + sql("SELECT '12' as col").write.format(format).save(dir.toString) + sql(s"CREATE TABLE t (col CHAR(3)) using $format LOCATION '$dir'") + checkAnswer(sql("SELECT * FROM t"), Row("12 ")) + } + } + Seq("CHAR", "VARCHAR").foreach { typ => + withTempPath { dir => + withTable("t") { + sql("SELECT '123456' as col").write.format(format).save(dir.toString) + sql(s"CREATE TABLE t (col $typ(2)) using $format LOCATION '$dir'") + checkError( + exception = intercept[SparkRuntimeException] { + sql("SELECT * FROM t").collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "2") + ) + } + } + } + // An oversized value consisting only of trailing blanks is trimmed successfully. + withTempPath { dir => + withTable("t") { + sql("SELECT '12 ' as col").write.format(format).save(dir.toString) + sql(s"CREATE TABLE t (col VARCHAR(2)) using $format LOCATION '$dir'") + checkAnswer(sql("SELECT * FROM t"), Row("12")) + } + } + } + } + + test("SPARK-58794: EXTERNAL TABLE CHAR/VARCHAR pad, assignment, and oversize") { + // External file tables use the same store assignment on write and pad + + // length-check on scan. Hive TRANSFORM is out of scope for this epic. + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + withTempPath { dir => + val path = dir.getCanonicalPath + // Pre-existing short file values: CHAR pads on scan; types stay first-class. + sql("SELECT 'ab' AS c, 'cd' AS v").write.format(format).save(path) + withTable("std_ext") { + sql( + s"""CREATE EXTERNAL TABLE std_ext (c CHAR(5), v VARCHAR(5)) + |USING $format LOCATION '$path'""".stripMargin) + assert(spark.table("std_ext").schema.map(_.dataType) === + Seq(CharType(5), VarcharType(5))) + checkAnswer( + sql("SELECT concat('<', c, '>'), concat('<', v, '>') FROM std_ext"), + Row("<ab >", "<cd>")) + + // INSERT into EXTERNAL applies write-side assignment (pad / length). + sql("INSERT INTO std_ext VALUES ('x', 'yz')") + checkAnswer( + sql("SELECT concat('<', c, '>'), concat('<', v, '>') FROM std_ext " + + "WHERE c LIKE 'x%'"), + Row("<x >", "<yz>")) + intercept[SparkRuntimeException] { + sql("INSERT INTO std_ext VALUES ('too-long', 'ok')").collect() + } + } + } + // Bypass-writer oversize non-blank values fail on scan. + withTempPath { dir => + val path = dir.getCanonicalPath + sql("SELECT 'abcdef' AS c").write.format(format).save(path) + withTable("std_ext_oversize") { + sql( + s"""CREATE EXTERNAL TABLE std_ext_oversize (c CHAR(3)) + |USING $format LOCATION '$path'""".stripMargin) + checkError( + exception = intercept[SparkRuntimeException] { + sql("SELECT * FROM std_ext_oversize").collect() + }, + condition = "EXCEED_LIMIT_LENGTH", + parameters = Map("limit" -> "3") + ) + } + } + } + } + test("alter table set location w/ fit length values") { withTempPath { dir => withTable("t") { @@ -1171,6 +2471,27 @@ class DSV2CharVarcharTestSuite extends CharVarcharTestSuite .set(SQLConf.DEFAULT_CATALOG.key, "testcat") } + test("SPARK-58794: VARCHAR widen and CHAR->VARCHAR under standardSemantics") { + // V2 CheckAnalysis allows length-preserving CHAR->VARCHAR and VARCHAR widen; + // V1 file-source ALTER does not (collation-only evolution). + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + withTable("std_v2_alter") { + sql(s"CREATE TABLE std_v2_alter (c CHAR(4), v VARCHAR(4)) USING $format") + sql("ALTER TABLE std_v2_alter CHANGE COLUMN v TYPE VARCHAR(5)") + assert(spark.table("std_v2_alter").schema("v").dataType === VarcharType(5)) + sql("ALTER TABLE std_v2_alter CHANGE COLUMN c TYPE VARCHAR(5)") + assert(spark.table("std_v2_alter").schema.map(_.dataType) === + Seq(VarcharType(5), VarcharType(5))) + intercept[AnalysisException] { + sql("ALTER TABLE std_v2_alter CHANGE COLUMN v TYPE VARCHAR(4)") + } + intercept[AnalysisException] { + sql("ALTER TABLE std_v2_alter CHANGE COLUMN c TYPE CHAR(5)") + } + } + } + } + test("char/varchar type values length check: partitioned columns of other types") { Seq("CHAR(5)", "VARCHAR(5)").foreach { typ => withTable("t") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 32517af68c916..5300c0f3ad4a5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala @@ -573,6 +573,56 @@ class DataFrameAggregateSuite extends SharedSparkSession } } + testWithWholeStageCodegenOnAndOff("SPARK-58213: corr returns NULL for zero variance") { _ => + val input = Seq( + ("both-constant", Some(1.0), Some(1.0)), + ("both-constant", Some(1.0), Some(1.0)), + ("empty", None, None), + ("normal", Some(1.0), Some(1.0)), + ("normal", Some(2.0), Some(2.0)), + ("normal", Some(3.0), Some(3.0)), + ("partial-null", None, Some(0.0)), + ("partial-null", Some(1.0), Some(1.0)), + ("partial-null", Some(2.0), None), + ("partial-null", Some(3.0), Some(3.0)), + ("single", Some(1.0), Some(2.0)), + ("x-constant", Some(1.0), Some(1.0)), + ("x-constant", Some(1.0), Some(2.0)), + ("x-constant", Some(1.0), Some(3.0)), + ("y-constant", Some(1.0), Some(1.0)), + ("y-constant", Some(2.0), Some(1.0)), + ("y-constant", Some(3.0), Some(1.0)), + ("zero-correlation", Some(0.0), Some(1.0)), + ("zero-correlation", Some(1.0), Some(-2.0)), + ("zero-correlation", Some(2.0), Some(1.0))).toDF("case", "x", "y") + + Seq(true, false).foreach { ansiEnabled => + Seq(true, false).foreach { legacyStatisticalAggregate => + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString, + SQLConf.LEGACY_STATISTICAL_AGGREGATE.key -> + legacyStatisticalAggregate.toString) { + val singleRow = if (legacyStatisticalAggregate) { + Row("single", Double.NaN) + } else { + Row("single", null) + } + checkAnswer( + input.groupBy($"case").agg(corr($"x", $"y")).orderBy($"case"), + Seq( + Row("both-constant", null), + Row("empty", null), + Row("normal", 1.0), + Row("partial-null", 1.0), + singleRow, + Row("x-constant", null), + Row("y-constant", null), + Row("zero-correlation", 0.0))) + } + } + } + } + test("null moments") { val emptyTableData = Seq.empty[(Int, Int)].toDF("a", "b") checkAnswer(emptyTableData.agg( @@ -676,6 +726,72 @@ class DataFrameAggregateSuite extends SharedSparkSession ) } + test("collect_union function") { + // Distinct union of array elements across rows. + val df = Seq(Seq(1, 2), Seq(2, 3), Seq(1)).toDF("arr") + checkDataset( + df.select(collect_union($"arr").as("u")).as[Set[Int]], + Set(1, 2, 3)) + checkAnswer( + df.select(sort_array(collect_union($"arr"))), + Seq(Row(Seq(1, 2, 3)))) + checkAnswer( + df.selectExpr("sort_array(collect_union(arr))"), + Seq(Row(Seq(1, 2, 3)))) + + // NULL array inputs are always skipped. + val dfNulls = Seq(Seq(1, 2), null, Seq(2, 3)).toDF("arr") + checkAnswer( + dfNulls.select(sort_array(collect_union($"arr"))), + Seq(Row(Seq(1, 2, 3)))) + + // NULL elements: dropped by default (IGNORE NULLS, matching collect_set) ... + val dfNullElem = Seq(Seq(Integer.valueOf(1), null), Seq(Integer.valueOf(2))).toDF("arr") + checkAnswer( + dfNullElem.select(sort_array(collect_union($"arr"))), + Seq(Row(Seq(1, 2)))) + // ... and kept (a single null) with RESPECT NULLS, matching + // array_distinct(flatten(collect_list(...))). sort_array puts null first in asc order. + checkAnswer( + dfNullElem.selectExpr("sort_array(collect_union(arr) RESPECT NULLS)"), + Seq(Row(Seq(null, 1, 2)))) + // Equivalent to array_distinct(flatten(collect_list(...))) under RESPECT NULLS. Compare + // order-insensitively via sort_array, since element order in either result is unspecified. + checkAnswer( + dfNullElem.selectExpr("sort_array(collect_union(arr) RESPECT NULLS)"), + dfNullElem.selectExpr("sort_array(array_distinct(flatten(collect_list(arr))))")) + + // Per-group union. + val g = Seq(("a", Seq(1, 2)), ("a", Seq(2, 3)), ("b", Seq(4))).toDF("k", "arr") + checkAnswer( + g.groupBy("k").agg(sort_array(collect_union($"arr"))).orderBy("k"), + Seq(Row("a", Seq(1, 2, 3)), Row("b", Seq(4)))) + + // Empty result: only-NULL arrays produce an empty array, not null. + val dfEmpty = Seq[Seq[Int]](null, null).toDF("arr") + checkAnswer( + dfEmpty.select(collect_union($"arr")), + Seq(Row(Seq.empty[Int]))) + } + + test("collect_union requires an array input") { + // A non-array (scalar) input is rejected at analysis time. + val df = Seq(1, 2, 3).toDF("a") + checkError( + exception = intercept[AnalysisException] { + df.select(collect_union($"a")).collect() + }, + condition = "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE", + parameters = Map( + "sqlExpr" -> "\"collect_union(a)\"", + "paramIndex" -> "first", + "inputSql" -> "\"a\"", + "inputType" -> "\"INT\"", + "requiredType" -> "\"ARRAY\""), + context = ExpectedContext( + fragment = "collect_union", callSitePattern = getCurrentClassCallSitePattern)) + } + test("SPARK-55256: array_agg and collect_list skip nulls by default") { val df = Seq((1, Some(2)), (2, None), (3, Some(4))).toDF("a", "b") @@ -2741,6 +2857,70 @@ class DataFrameAggregateSuite extends SharedSparkSession } } + // Byte 3 of the serialized sketch holds lgConfigK, so SUBSTRING(hex(...), 7, 2) reads it as a + // hex pair: 0x0F = 15, 0x0C = 12. + private val allNullGroupSketches = + """ + |with sketches as ( + | select 'has_data' as grp, hll_sketch_agg(cast(id as string), 15) as sketch + | from (select explode(sequence(1, 1000)) as id) + | union all + | select 'all_null' as grp, cast(null as binary) as sketch + |) + |""".stripMargin + + test("hll_union_agg reports the default lgConfigK for a group with no non-NULL sketch") { + // hll_union_agg has no lgConfigK parameter, and the all-NULL group holds no sketch to take one + // from, so it cannot report the 15 the other group was built at. The estimate is correct either + // way; the mismatched precision is what the next test has to cope with. + checkAnswer( + sql(allNullGroupSketches + + """ + |select + | grp, + | substring(hex(hll_union_agg(sketch)), 7, 2) as lg_config_k, + | hll_sketch_estimate(hll_union_agg(sketch)) as estimate + |from sketches + |group by grp + |""".stripMargin), + Seq(Row("all_null", "0C", 0L), Row("has_data", "0F", 1000L))) + } + + test("hll_union_agg and hll_union re-merge the empty sketch produced for an all-NULL group") { + // The stored sketches are now an empty lgConfigK=12 one beside a populated lgConfigK=15 one. + // Rolling them back up must not fail under default settings, because an empty sketch holds no + // coupons and so cannot cost precision at any lgConfigK. + val stored = sql(allNullGroupSketches + + "select grp, hll_union_agg(sketch) as sketch from sketches group by grp") + .collect() + .map(row => row.getString(0) -> row.getAs[Array[Byte]]("sketch")) + .toMap + + Seq( + "empty sketch first" -> Seq(stored("all_null"), stored("has_data")), + "empty sketch last" -> Seq(stored("has_data"), stored("all_null")) + ).foreach { case (order, sketches) => + // One partition exercises update(); two exercise the partial-to-final merge() as well. + Seq(1, 2).foreach { numPartitions => + withClue(s"hll_union_agg, $order, $numPartitions partition(s): ") { + checkAnswer( + sketches.toDF("sketch").repartition(numPartitions).selectExpr( + "substring(hex(hll_union_agg(sketch)), 7, 2) as lg_config_k", + "hll_sketch_estimate(hll_union_agg(sketch)) as estimate"), + Seq(Row("0F", 1000L))) + } + } + + withClue(s"hll_union, $order: ") { + checkAnswer( + Seq((sketches.head, sketches.last)).toDF("left", "right").selectExpr( + "substring(hex(hll_union(left, right)), 7, 2) as lg_config_k", + "hll_sketch_estimate(hll_union(left, right)) as estimate"), + Seq(Row("0F", 1000L))) + } + } + } + test("hll_sketch_agg") { val df = Seq(1, 1, 2, 2, 3).toDF("col") checkAnswer( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala index f86613bf6684d..f20c6bd7a88ea 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala @@ -84,6 +84,7 @@ class DataFrameFunctionsSuite extends SharedSparkSession { "bucket", "days", "hours", "months", "years", // Datasource v2 partition transformations "product", // Discussed in https://github.com/apache/spark/pull/30745 "unwrap_udt", + "wrap_udt", "timestamp_add", "timestamp_diff" ) @@ -499,6 +500,23 @@ class DataFrameFunctionsSuite extends SharedSparkSession { callSitePattern = "", startIndex = 0, stopIndex = 0)) + expr = randstr(lit(-1), lit(0)) + checkError( + intercept[AnalysisException](df.select(expr)), + condition = "DATATYPE_MISMATCH.VALUE_OUT_OF_RANGE", + parameters = Map( + "sqlExpr" -> "\"randstr(-1, 0)\"", + "exprName" -> "`length`", + "valueRange" -> "[0, 2147483647]", + "currentValue" -> "-1"), + context = ExpectedContext( + contextType = QueryContextType.DataFrame, + fragment = "randstr", + objectType = "", + objectName = "", + callSitePattern = "", + startIndex = 0, + stopIndex = 0)) } test("uniform function") { @@ -664,6 +682,23 @@ class DataFrameFunctionsSuite extends SharedSparkSession { Row(2743272264L, 2180413220L)) } + test("misc xxh3_64 and xxh3_128 function") { + val df = Seq(("ABC", Array[Byte](1, 2, 3, 4, 5, 6))).toDF("a", "b") + checkAnswer( + df.select(xxh3_64($"a"), xxh3_64($"b")), + Row(2615927343983396622L, -4044731995552965649L)) + checkAnswer( + df.select(xxh3_128($"a"), xxh3_128($"b")), + Row("9e947f00ecd6acb2244da40f405c870e", "866737830f560dbf3e1f439d2d785f44")) + + checkAnswer( + df.selectExpr("xxh3_64(a)", "xxh3_64(b)"), + Row(2615927343983396622L, -4044731995552965649L)) + checkAnswer( + df.selectExpr("xxh3_128(a)", "xxh3_128(b)"), + Row("9e947f00ecd6acb2244da40f405c870e", "866737830f560dbf3e1f439d2d785f44")) + } + test("misc aes function") { val key32 = "abcdefghijklmnop12345678ABCDEFGH" val encryptedEcb = "9J3iZbIxnmaG+OIA9Amd+A==" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala index 129d6bb686762..1d93006342abc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameJoinSuite.scala @@ -630,4 +630,13 @@ class DataFrameJoinSuite extends SharedSparkSession checkAnswer(df1, df2) }) } + + test("SPARK-58384: preserve null semantics under NOT in join condition") { + val left = Seq((Some(0), 10), (None, 11)).toDF("a", "b").as("left") + val right = Seq((None, 20), (Some(0), 21), (Some(1), 22)).toDF("x", "y").as("right") + val condition = !(($"left.a" === $"right.x") || + ($"left.a".isNull && $"right.x".isNull)) + + checkAnswer(left.join(right, condition), Row(0, 10, 1, 22)) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala index 270d74cf5ef54..81e69fd09d972 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFramePivotSuite.scala @@ -476,4 +476,45 @@ class DataFramePivotSuite extends SharedSparkSession { Row(10, 30)) } } + + test("ORDER BY on a column dropped by PIVOT fails name resolution") { + // PIVOT drops the pivoted measure column, so `t.v` fails name resolution with + // UNRESOLVED_COLUMN: the dropped column is never appended to a lower operator. + withSQLConf(SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "false") { + checkError( + exception = intercept[AnalysisException] { + spark.sql( + """SELECT * FROM VALUES (1, 1, 100), (2, 2, 200) AS t(id, m, v) + |PIVOT (SUM(v) FOR m IN (1, 2)) + |ORDER BY t.v""".stripMargin) + }, + condition = "UNRESOLVED_COLUMN.WITH_SUGGESTION", + parameters = Map( + "objectName" -> "`t`.`v`", + "proposal" -> "`1`, `2`, `t`.`id`"), + queryContext = + Array(ExpectedContext(fragment = "t.v", start = 101, stop = 103))) + } + } + + test("DataFrame sort on a column dropped by PIVOT is rejected") { + // The Column reference carries the dropped column's id, so unlike the SQL case above it + // reaches hidden-output insertion rather than failing name resolution, and analysis fails + // with MISSING_ATTRIBUTES. + withSQLConf(SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "false") { + val df = Seq((1, 1, 100), (2, 2, 200)).toDF("id", "m", "v") + val pivoted = df.groupBy("id").pivot("m", Seq(1, 2)).agg(sum($"v")) + + checkError( + exception = intercept[AnalysisException] { + pivoted.sort(df("v")) + }, + condition = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", + parameters = Map( + "missingAttributes" -> "\"v\"", + "input" -> "\"id\", \"1\", \"2\"", + "operator" -> "!Sort \\[v#\\d+ ASC NULLS FIRST\\], true"), + matchPVals = true) + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala index faefbab6c8ca9..32b6da3a30b15 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSetOperationsSuite.scala @@ -20,6 +20,7 @@ package org.apache.spark.sql import java.sql.{Date, Timestamp} import java.util.Locale +import org.apache.spark.sql.catalyst.expressions.AttributeReference import org.apache.spark.sql.catalyst.optimizer.RemoveNoopUnion import org.apache.spark.sql.catalyst.plans.logical.Union import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, KeyedPartitioning, PartitioningCollection, UnknownPartitioning} @@ -1589,6 +1590,159 @@ class DataFrameSetOperationsSuite extends SharedSparkSession with AdaptiveSparkP } } + test("SPARK-58819: union outputPartitioning compares children in the union's attribute space") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTempView("t1", "t2") { + // `DISTRIBUTE BY id` resolves its key with the `t1` qualifier inside the subquery, but + // the outer `Project` exposes `id` with the subquery qualifier, so child 0's partitioning + // differs from its output in qualifier, not nullability. The branches overlap on `id = 1`, + // so a shared key's grouping depends on the union's co-location claim being honored + // (guarded by `checkAnswer(grouped, ...)` below). + Seq((Option(1), 10), (Option(2), 20)).toDF("id", "v").createOrReplaceTempView("t1") + Seq((Option(1), 30), (Option(3), 40)).toDF("id", "v").createOrReplaceTempView("t2") + + val sqlText = + """ + |SELECT id, v FROM (SELECT id, v FROM t1 DISTRIBUTE BY id) WHERE id IS NOT NULL + |UNION ALL + |SELECT id, sum(v) AS v FROM t2 GROUP BY id + |""".stripMargin + val union = spark.sql(sqlText) + val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + + // Pin the discriminator to qualifier: child 0's partitioning references the same column + // (`id`) as its output, differing only in qualifier (`[t1]` vs the subquery qualifier), + // never nullability. The explicit nullability check rules out the nullability cause the + // DataFrame test below exercises; a future refactor leaking a nullability difference + // fails loudly instead of silently changing what the test covers. + val child0 = unionExec.head.children.head + val outputId = child0.output.head.asInstanceOf[AttributeReference] + val partitioningId = + child0.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head + .asInstanceOf[AttributeReference] + assert(outputId.exprId == partitioningId.exprId, + s"expected the same column: $outputId vs $partitioningId") + assert(outputId.nullable == partitioningId.nullable, + s"nullability must match, qualifier is the sole difference: " + + s"$outputId vs $partitioningId") + assert(outputId.qualifier != partitioningId.qualifier, + s"expected a qualifier difference: $outputId vs $partitioningId") + + // Child 0's `HashPartitioning` references `id` with the `t1` qualifier while its output + // `id` carries the subquery qualifier; remapping both to the union's output attributes + // still propagates the hash partitioning despite the qualifier difference. The propagated + // partitioning must be expressed in the union's own output attribute (the subquery + // qualifier), not child 0's `[t1]` attribute, since `toUnionOutput` was removed. + assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning], + s"expected a HashPartitioning pass-through but got ${unionExec.head.outputPartitioning}") + val hashPartitioning = + unionExec.head.outputPartitioning.asInstanceOf[HashPartitioning] + assert(hashPartitioning.expressions == Seq(unionExec.head.output.head)) + + // The two branches contribute one shuffle each (DISTRIBUTE BY and GROUP BY). The propagated + // HashPartitioning lets the downstream group-by reuse them instead of adding a third. + val unionShuffles = union.queryExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + }.size + val grouped = union.groupBy($"id").count() + val groupedShuffles = grouped.queryExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + }.size + assert(unionShuffles == 2, s"union should have 2 shuffles but got $unionShuffles") + assert(groupedShuffles == 2, + s"group-by should reuse the union's partitioning (expect 2 shuffles) but got " + + s"$groupedShuffles\n${grouped.queryExecution.executedPlan}") + + // `UNION_OUTPUT_PARTITIONING=false` drops the pass-through so the group-by adds its own + // shuffle; that freshly-planned path is the oracle for both the raw union rows and the + // grouped result. + val (correctResult, correctGrouped) = + withSQLConf(SQLConf.UNION_OUTPUT_PARTITIONING.key -> "false") { + val baseline = spark.sql(sqlText) + (baseline.collect(), baseline.groupBy($"id").count().collect()) + } + checkAnswer(union, correctResult) + checkAnswer(grouped, correctGrouped) + } + } + } + + test("SPARK-58819: union outputPartitioning ignores partition key nullability") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + // `.sample(1.0)` keeps `FilterExec` above the shuffle without disabling any rule: + // `PushPredicateThroughNonJoin.canPushThrough` has no `Sample` case, so `IsNotNull` stays + // above the `Sample`, which passes its child's output and partitioning through verbatim. + // The filter's output nullability is thus narrowed while its passed-through partitioning + // keeps the nullable key. `fraction = 1.0` passes every row, keeping counts deterministic. + val df1 = Seq((Option(1), 10L), (Option(2), 20L)).toDF("id", "v") + .repartition($"id").sample(withReplacement = false, fraction = 1.0, seed = 42) + .filter($"id".isNotNull) + val df2 = Seq((Option(1), 30L), (Option(3), 40L)).toDF("id", "v") + .groupBy($"id").agg(sum($"v").as("v")) + + val union = df1.union(df2) + val unionExec = union.queryExecution.executedPlan.collect { case u: UnionExec => u } + assert(unionExec.size == 1) + + // Pin the discriminator to nullability: child 0 is a `FilterExec` whose partitioning + // (passed through verbatim) references the same column as its `output`, differing only in + // nullability. The explicit qualifier check rules out the qualifier cause the SQL-variant + // test above exercises; a future refactor leaking a qualifier difference fails loudly + // instead of silently changing what the test covers. + val child0 = unionExec.head.children.head + val outputId = child0.output.head.asInstanceOf[AttributeReference] + val partitioningId = + child0.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head + .asInstanceOf[AttributeReference] + assert(outputId.exprId == partitioningId.exprId, + s"expected the same column: $outputId vs $partitioningId") + assert(outputId.qualifier == partitioningId.qualifier, + s"qualifier must match (nullability is the sole difference): $outputId vs $partitioningId") + assert(outputId.nullable != partitioningId.nullable, + s"expected a nullability difference: $outputId vs $partitioningId") + assert(outputId.withNullability(partitioningId.nullable) == partitioningId, + s"only nullability should differ: $outputId vs $partitioningId") + + // child1 (groupBy) is the well-behaved branch: its output and partitioning reference the + // same nullable group key, so the nullability mismatch is confined to child0's Filter. + val child1 = unionExec.head.children(1) + val child1OutputId = child1.output.head.asInstanceOf[AttributeReference] + val child1PartitioningId = + child1.outputPartitioning.asInstanceOf[HashPartitioning].expressions.head + .asInstanceOf[AttributeReference] + assert(child1OutputId.nullable, s"group key should be nullable: $child1OutputId") + assert(child1OutputId.exprId == child1PartitioningId.exprId, + s"child1 output and partitioning should reference the same column") + assert(child1OutputId.nullable == child1PartitioningId.nullable, + s"child1 has no nullability mismatch: $child1OutputId vs $child1PartitioningId") + + // The union propagates the co-located HashPartitioning in its own output attribute space + // with the merged nullability (true, from child1's nullable group key). This also pins that + // the partitioning attribute is the union's output attribute, not a leaked child attribute. + assert(unionExec.head.outputPartitioning.isInstanceOf[HashPartitioning], + s"expected HashPartitioning pass-through but got ${unionExec.head.outputPartitioning}") + val unionPartitioning = unionExec.head.outputPartitioning.asInstanceOf[HashPartitioning] + assert(unionPartitioning.expressions == Seq(unionExec.head.output.head)) + + // Without the fix the union reports UnknownPartitioning and the downstream group-by adds a + // third shuffle; with it, the pass-through lets the group-by reuse the two existing ones. + val grouped = union.groupBy($"id").count() + val groupedShuffles = grouped.queryExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + }.size + assert(groupedShuffles == 2, + s"group-by should reuse the union's partitioning (expect 2 shuffles) but got " + + s"$groupedShuffles\n${grouped.queryExecution.executedPlan}") + + // The two branches overlap on `id = 1`, so the correct count proves the co-location claim + // is actually honored (a false claim would split id = 1 into duplicate groups). + checkAnswer(grouped, Row(1, 2L) :: Row(2, 1L) :: Row(3, 1L) :: Nil) + } + } + test("SPARK-52921: union partitioning - range partitioning") { val df1 = Seq((1, 2, 4), (1, 3, 5), (2, 2, 3), (2, 4, 5)).toDF("a", "b", "c") val df2 = Seq((4, 1, 5), (2, 4, 6), (1, 4, 2), (3, 5, 1)).toDF("d", "e", "f") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala index 99d4a66195c85..1154863a2dcda 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameStatSuite.scala @@ -264,7 +264,7 @@ class DataFrameStatSuite extends SharedSparkSession { val q1 = 0.5 val q2 = 0.8 - val epsilons = List(2.0, 5.0, 100.0) + val epsilons = List(2.0, 5.0, 100.0, Double.PositiveInfinity) val Array(single1_1) = df.stat.approxQuantile("singles", Array(q1), 1.0) val Array(s1_1, s2_1) = df.stat.approxQuantile("singles", Array(q1, q2), 1.0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala index 577ae025f1dfb..423dcc83dbed0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala @@ -636,6 +636,18 @@ class DataFrameSuite extends SharedSparkSession assert(df.schema.map(_.name) === Seq("key", "value", "newCol1", "newCol2")) } + test("withColumns: internal method with dependent columns") { + val df = testData.toDF().withColumns( + Seq("newCol1", "newCol2"), + Seq(col("key") + 1, col("newCol1") + 1)) + checkAnswer( + df, + testData.collect().map { case Row(key: Int, value: String) => + Row(key, value, key + 1, key + 2) + }.toSeq) + assert(df.schema.map(_.name) === Seq("key", "value", "newCol1", "newCol2")) + } + test("withColumns: internal method") { val df = testData.toDF().withColumns(Seq("newCol1", "newCol2"), Seq(col("key") + 1, col("key") + 2)) @@ -2822,5 +2834,3 @@ case class OutputListAwareConstraintsTestPlan( override def newInstance(): LogicalPlan = copy(outputList = outputList.map(_.newInstance())) } - - diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala index c1d1f13915a6e..3530ccbb6ecfa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala @@ -79,6 +79,28 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession parameters = Map("wf_name" -> "row_number", "wf_expr" -> "row_number()")) } + test("SPARK-58757: collapse window with an empty order spec into an ordered sibling") { + val df = Seq( + (0, 0), (0, 2), (0, 4), + (1, 1), (1, 3), (1, 5)).toDF("k", "v") + val ordered = Window.partitionBy("k").orderBy("v") + val unordered = Window.partitionBy("k") + val res = df.select( + $"k", + $"v", + row_number().over(ordered).as("rn"), + collect_list($"v").over(unordered).as("vs"), + first($"v").over(unordered).as("f")) + assert(res.queryExecution.optimizedPlan.collect { case w: LogicalWindow => w }.size === 1) + checkAnswer(res, Seq( + Row(0, 0, 1, Array(0, 2, 4), 0), + Row(0, 2, 2, Array(0, 2, 4), 0), + Row(0, 4, 3, Array(0, 2, 4), 0), + Row(1, 1, 1, Array(1, 3, 5), 1), + Row(1, 3, 2, Array(1, 3, 5), 1), + Row(1, 5, 3, Array(1, 3, 5), 1))) + } + test("corr, covar_pop, stddev_pop functions in specific window") { withSQLConf(SQLConf.LEGACY_STATISTICAL_AGGREGATE.key -> "true", SQLConf.ANSI_ENABLED.key -> "false") { @@ -1360,7 +1382,7 @@ class DataFrameWindowFunctionsSuite extends SharedSparkSession def isShuffleExecByRequirement( plan: ShuffleExchangeExec, desiredClusterColumns: Seq[String]): Boolean = plan match { - case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _) => + case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS, _, _) => partitionExpressionsColumns(op.expressions) === desiredClusterColumns case _ => false } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala index 3b28cae31a134..096eee137b320 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala @@ -1942,7 +1942,7 @@ class DatasetSuite extends SharedSparkSession val agg = cp.groupBy($"id" % 2).agg(count($"id")) agg.queryExecution.executedPlan.collectFirst { - case ShuffleExchangeExec(_, _: RDDScanExec, _, _) => + case ShuffleExchangeExec(_, _: RDDScanExec, _, _, _) => case BroadcastExchangeExec(_, _: RDDScanExec) => }.foreach { _ => fail( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetUnpivotSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DatasetUnpivotSuite.scala index f83eef5fb0ce3..97ab31c3405a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetUnpivotSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetUnpivotSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql import org.apache.spark.sql.functions.{length, struct, sum} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ import org.apache.spark.util.ArrayImplicits._ @@ -666,6 +667,50 @@ class DatasetUnpivotSuite extends SharedSparkSession { } } + test("ORDER BY on a column dropped by UNPIVOT is rejected") { + // UNPIVOT drops the unpivoted value columns, so `ORDER BY t.a` cannot resolve against them; + // hidden-output insertion appends the column to an operator whose child does not produce it, + // and analysis fails with MISSING_ATTRIBUTES. + withSQLConf(SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "false") { + checkError( + exception = intercept[AnalysisException] { + spark.sql( + """SELECT * FROM VALUES (1, 10, 20) AS t(id, a, b) + |UNPIVOT (val FOR name IN (a, b)) + |ORDER BY t.a""".stripMargin) + }, + condition = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", + parameters = Map( + "missingAttributes" -> "\"a\"", + "input" -> "\"id\", \"name\", \"val\"", + "operator" -> "!Sort \\[a#\\d+ ASC NULLS FIRST\\], true"), + matchPVals = true, + queryContext = + Array(ExpectedContext(fragment = "ORDER BY t.a", start = 81, stop = 92))) + } + } + + test("DataFrame sort on a column dropped by UNPIVOT is rejected") { + // The Column reference carries the dropped column's id, so it reaches hidden-output insertion + // rather than failing name resolution as a bare name would, and analysis fails with + // MISSING_ATTRIBUTES. + withSQLConf(SQLConf.ANALYZER_DUAL_RUN_LEGACY_AND_SINGLE_PASS_RESOLVER.key -> "false") { + val df = Seq((1, 10, 20)).toDF("id", "a", "b") + val unpivoted = df.unpivot(Array($"id"), Array($"a", $"b"), "name", "val") + + checkError( + exception = intercept[AnalysisException] { + unpivoted.sort(df("a")) + }, + condition = "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_MISSING_FROM_INPUT", + parameters = Map( + "missingAttributes" -> "\"a\"", + "input" -> "\"id\", \"name\", \"val\"", + "operator" -> "!Sort \\[a#\\d+ ASC NULLS FIRST\\], true"), + matchPVals = true) + } + } + } case class WideData(id: Int, str1: String, str2: String, int1: Option[Int], long1: Option[Long]) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala index 4aa4a8870b770..9f497680b7fdb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala @@ -17,22 +17,30 @@ package org.apache.spark.sql +import java.lang.reflect.Modifier +import java.sql.Date import java.util.concurrent.atomic.AtomicInteger import scala.collection.concurrent.TrieMap +import scala.collection.mutable.ArrayBuffer import org.scalatest.GivenWhenThen -import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression} +import org.apache.spark.SparkException +import org.apache.spark.sql.catalyst.expressions.{BroadcastValueProjection, + DynamicPruningExpression, DynamicPruningSubquery, Expression} import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode._ -import org.apache.spark.sql.catalyst.plans.ExistenceJoin +import org.apache.spark.sql.catalyst.optimizer.BuildRight +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, Inner, LeftAnti} import org.apache.spark.sql.catalyst.plans.logical.LocalRelation -import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableWithV2FilterCatalog} +import org.apache.spark.sql.connector.catalog.{InMemoryTableCatalog, InMemoryTableCatalystRuntimeFilterCatalog, InMemoryTableWithV2FilterCatalog} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive._ import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.execution.dynamicpruning.PlanDynamicPruningFilters import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ReusedExchangeExec} import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec +import org.apache.spark.sql.execution.metric.SQLLastAttemptMetrics import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryWrapper} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -195,8 +203,9 @@ abstract class DynamicPartitionPruningSuiteBase case InSubqueryExec(_, _: SubqueryExec, _, _, _, _) => true case _ => false } - val subqueryBroadcast = dpExprs.collect { + val subqueryBroadcast: Seq[BaseSubqueryExec with UnaryExecNode] = dpExprs.collect { case InSubqueryExec(_, b: SubqueryBroadcastExec, _, _, _, _) => b + case InSubqueryExec(_, b: ProjectedBroadcastValueSubqueryExec, _, _, _, _) => b } val hasFilter = if (withSubquery) "Should" else "Shouldn't" @@ -255,6 +264,348 @@ abstract class DynamicPartitionPruningSuiteBase assert(buf.distinct.size == n) } + test("reuse ancestor broadcast value rows for composite-key date partition pruning") { + withTable("dpp_projection_dates", "dpp_projection_users", "dpp_projection_target") { + Seq( + (Date.valueOf("2024-01-01"), "2024-01-01", "us"), + (Date.valueOf("2024-01-01"), "2024-01-01", "us"), + (Date.valueOf("2024-01-08"), "2024-01-08", "eu"), + (Date.valueOf("2024-01-15"), null, "us"), + (null.asInstanceOf[Date], "2024-01-22", "apac"), + (Date.valueOf("2024-01-29"), "2024-01-29", "latam")) + .toDF("cohort_date", "cohort_ds", "region") + .write.format(tableFormat).saveAsTable("dpp_projection_dates") + + Seq( + ("2024-01-01", "us", "u1", true), + ("2024-01-08", "eu", "u2", true), + (null, "us", "null-key", true), + ("2024-01-22", "apac", "null-value", true), + ("2024-01-08", "us", "inactive", false)) + .toDF("cohort_ds", "region", "user_id", "active") + .write.format(tableFormat).saveAsTable("dpp_projection_users") + + Seq( + ("u1", "2024-01-07", "p1"), + ("u2", "2024-01-14", "p2"), + ("ghost", "2024-02-04", "extra-domain"), + ("u1", "2024-02-01", "unselected")) + .toDF("user_id", "ds", "payload") + .write.partitionBy("ds").format(tableFormat).saveAsTable("dpp_projection_target") + + val query = + """ + |WITH cohorts AS ( + | SELECT /*+ BROADCAST(d) */ d.cohort_date, u.user_id + | FROM dpp_projection_dates d + | JOIN dpp_projection_users u + | ON d.cohort_ds = u.cohort_ds AND d.region = u.region + | WHERE u.active AND u.user_id <> 'filtered' + |) + |SELECT c.user_id, w.payload + |FROM cohorts c LEFT JOIN dpp_projection_target w + | ON w.user_id = c.user_id + | AND w.ds = date_format(date_add(c.cohort_date, 6), 'yyyy-MM-dd') + |ORDER BY c.user_id, w.payload + |""".stripMargin + val expected = Seq( + Row("null-value", null), + Row("u1", "p1"), + Row("u1", "p1"), + Row("u2", "p2")) + + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "false", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql(query) + checkAnswer(df, expected) + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.isEmpty, df.queryExecution.executedPlan) + } + + Seq("true", "false").foreach { ansiEnabled => + withSQLConf( + SQLConf.ANSI_ENABLED.key -> ansiEnabled, + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql(query) + val candidates = ArrayBuffer.empty[DynamicPruningSubquery] + df.queryExecution.optimizedPlan.transformAllExpressionsWithSubqueries { + case candidate: DynamicPruningSubquery => + candidates += candidate + candidate + } + val projections = candidates.flatMap(_.broadcastValueProjection) + assert(projections.size === 1, df.queryExecution.optimizedPlan) + assert(projections.head.sourceHashKeys.map(_.references.head.name) === + Seq("cohort_ds", "region")) + + checkAnswer(df, expected) + + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + val lastAttemptMetric = SQLLastAttemptMetrics.createMetric( + spark.sparkContext, "projected broadcast pruning") + assert(lastAttemptMetric.lastAttemptValueForDataset(df) === Some(0L)) + val pruning = collectDynamicPruningExpressions(df.queryExecution.executedPlan) + .collectFirst { case in: InSubqueryExec => in }.get + assert(pruning.values().get.map(String.valueOf).toSet === + Set("2024-01-07", "2024-01-14", "2024-02-04")) + val projectedMetrics = projected.head.metrics.map { + case (name, metric) => name -> metric.value + } + assert(projected.head.metrics("numInputRows").value === 5, projectedMetrics) + assert(projected.head.metrics("numOutputRows").value === 3, projectedMetrics) + } + } + + Seq( + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_ROWS.key, + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_BYTES.key, + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_SOURCE_BYTES.key + ).foreach { limitKey => + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "true", + limitKey -> "0", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql(query) + checkAnswer(df, expected) + + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + assert(projected.head.metrics("projectionDisabled").value === 1) + + val pruning = collectDynamicPruningExpressions(df.queryExecution.executedPlan) + .collectFirst { case in: InSubqueryExec => in }.get + assert(pruning.isResultUnavailable) + assert(pruning.values().isEmpty) + + val collectError = intercept[SparkException] { + projected.head.executeCollect() + } + assert(collectError.getCondition === "INTERNAL_ERROR") + assert(collectError.getMessage.contains("executeCollectResult")) + + val ordinaryIn = pruning.copy(isDynamicPruning = false) + val ordinaryInError = intercept[AssertionError] { + ordinaryIn.updateResult() + } + assert(ordinaryInError.getMessage.contains("dynamic partition pruning")) + } + } + } + } + + test("reuse ancestor broadcast string values with a residual join predicate") { + withTable("dpp_projection_products", "dpp_projection_promotions", + "dpp_projection_sales") { + Seq( + (101, "books", 10), + (202, "toys", 20), + (303, "clothing", 30), + (404, null.asInstanceOf[String], 40)) + .toDF("product_id", "category", "price") + .write.format(tableFormat).saveAsTable("dpp_projection_products") + + Seq((101, true, 15), (202, true, 25), (303, false, 35), (404, true, 45)) + .toDF("product_id", "active", "max_price") + .write.format(tableFormat).saveAsTable("dpp_projection_promotions") + + Seq( + (101, "books"), + (202, "toys"), + (303, "clothing"), + (909, "electronics")) + .toDF("product_id", "category") + .write.partitionBy("category").format(tableFormat) + .saveAsTable("dpp_projection_sales") + + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql( + """ + |WITH eligible AS ( + | SELECT /*+ BROADCAST(p) */ p.product_id, p.category + | FROM dpp_projection_products p + | JOIN dpp_projection_promotions promo + | ON p.product_id = promo.product_id + | AND p.price <= promo.max_price + | WHERE promo.active AND promo.product_id > 0 + |) + |SELECT s.product_id, s.category + |FROM dpp_projection_sales s + |JOIN eligible e + | ON s.product_id = e.product_id + | AND s.category = e.category + |ORDER BY s.product_id + |""".stripMargin) + + checkAnswer(df, Seq(Row(101, "books"), Row(202, "toys"))) + + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + + val pruning = collectDynamicPruningExpressions(df.queryExecution.executedPlan) + .collectFirst { case in: InSubqueryExec => in }.get + assert(pruning.values().get.map(String.valueOf).toSet === + Set("books", "toys", "clothing")) + } + } + } + + test("bind reordered broadcast value columns for partition pruning") { + withTable("dpp_reordered_dates", "dpp_reordered_users", "dpp_reordered_target") { + Seq( + (Date.valueOf("2024-01-01"), "us", "2024-01-01"), + (Date.valueOf("2024-01-08"), "eu", "2024-01-08"), + (Date.valueOf("2024-01-15"), "us", "2024-01-15")) + .toDF("cohort_date", "region", "cohort_ds") + .write.format(tableFormat).saveAsTable("dpp_reordered_dates") + + Seq( + ("2024-01-01", "us", "u1", true), + ("2024-01-08", "eu", "u2", true), + ("2024-01-15", "us", "inactive", false)) + .toDF("cohort_ds", "region", "user_id", "active") + .write.format(tableFormat).saveAsTable("dpp_reordered_users") + + Seq( + ("u1", "20240107", "p1"), + ("u2", "20240114", "p2"), + ("inactive", "20240121", "extra-domain")) + .toDF("user_id", "ds", "payload") + .write.partitionBy("ds").format(tableFormat).saveAsTable("dpp_reordered_target") + + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql( + """ + |WITH reordered_dates AS ( + | SELECT cohort_date, region, cohort_ds + | FROM dpp_reordered_dates + |), cohorts AS ( + | SELECT /*+ BROADCAST(d) */ u.user_id, d.cohort_date + | FROM reordered_dates d + | JOIN dpp_reordered_users u + | ON d.cohort_ds = u.cohort_ds AND d.region = u.region + | WHERE u.active AND u.user_id <> 'filtered' + |) + |SELECT c.user_id, w.payload + |FROM cohorts c + |LEFT JOIN dpp_reordered_target w + | ON w.user_id = c.user_id + | AND w.ds = date_format(date_add(c.cohort_date, 6), 'yyyyMMdd') + |ORDER BY c.user_id, w.payload + |""".stripMargin) + + checkAnswer(df, Seq(Row("u1", "p1"), Row("u2", "p2"))) + + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + assert(projected.head.child.output.map(_.name) === + Seq("cohort_date", "region", "cohort_ds")) + + val pruning = collectDynamicPruningExpressions(df.queryExecution.executedPlan) + .collectFirst { case in: InSubqueryExec => in }.get + assert(pruning.values().get.map(String.valueOf).toSet === + Set("20240107", "20240114", "20240121")) + } + } + } + + test("project every value row from dense and sparse long-key hash broadcasts") { + Seq(2L, 1000000000L).foreach { otherKey => + withTable("dpp_long_dates", "dpp_long_users", "dpp_long_target") { + Seq( + (1L, Date.valueOf("2024-01-01")), + (1L, Date.valueOf("2024-01-08")), + (otherKey, Date.valueOf("2024-01-15"))) + .toDF("cohort_id", "cohort_date") + .write.format(tableFormat).saveAsTable("dpp_long_dates") + + Seq((1L, "u1", true), (otherKey, "u2", true)) + .toDF("cohort_id", "user_id", "active") + .write.format(tableFormat).saveAsTable("dpp_long_users") + + Seq( + ("u1", "2024-01-07", "p1"), + ("u1", "2024-01-14", "p2"), + ("u2", "2024-01-21", "p3")) + .toDF("user_id", "ds", "payload") + .write.partitionBy("ds").format(tableFormat).saveAsTable("dpp_long_target") + + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql( + """ + |WITH cohorts AS ( + | SELECT /*+ BROADCAST(d) */ d.cohort_date, u.user_id + | FROM dpp_long_dates d + | JOIN dpp_long_users u ON d.cohort_id = u.cohort_id + | WHERE u.active AND u.user_id <> 'filtered' + |) + |SELECT c.user_id, w.payload + |FROM cohorts c + |LEFT JOIN dpp_long_target w + | ON w.user_id = c.user_id + | AND w.ds = date_format(date_add(c.cohort_date, 6), 'yyyy-MM-dd') + |ORDER BY c.user_id, w.payload + |""".stripMargin) + + checkAnswer(df, Seq( + Row("u1", "p1"), + Row("u1", "p2"), + Row("u2", "p3"))) + + val projected = collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + val pruning = collectDynamicPruningExpressions(df.queryExecution.executedPlan) + .collectFirst { case in: InSubqueryExec => in }.get + assert(pruning.values().get.map(String.valueOf).toSet === + Set("2024-01-07", "2024-01-14", "2024-01-21")) + val projectedMetrics = projected.head.metrics.map { + case (name, metric) => name -> metric.value + } + assert(projected.head.metrics("numInputRows").value === 3, projectedMetrics) + assert(projected.head.metrics("numOutputRows").value === 3, projectedMetrics) + } + } + } + } + /** * Collect the children of all correctly pushed down dynamic pruning expressions in a spark plan. */ @@ -2181,11 +2532,95 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat } class DynamicPartitionPruningV1SuiteAEOff extends DynamicPartitionPruningV1Suite - with DisableAdaptiveExecutionSuite + with DisableAdaptiveExecutionSuite { + + test("preserve direct broadcast pruning for existing hash and null-aware modes") { + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> "false", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val source = spark.table("product").queryExecution.optimizedPlan + val build = QueryExecution.createSparkPlan(spark.sessionState.planner, source) + val probe = spark.range(4).queryExecution.sparkPlan + val mismatchedHashJoin = BroadcastHashJoinExec( + Seq(probe.output.head), + Seq(build.output(1)), + Inner, + BuildRight, + None, + probe, + build) + assert(!mismatchedHashJoin.rightKeys.head.semanticEquals(source.output.head)) + + val nullAwareJoin = BroadcastHashJoinExec( + Seq(probe.output.head), + Seq(build.output.head), + LeftAnti, + BuildRight, + None, + probe, + build, + isNullAwareAntiJoin = true) + assert(nullAwareJoin.isNullAwareAntiJoin) + + Seq( + "mismatched hash mode" -> mismatchedHashJoin, + "null-aware anti join" -> nullAwareJoin).foreach { case (name, existingJoin) => + val pruning = DynamicPruningSubquery( + probe.output.head, + source, + Seq(source.output.head), + Seq(0), + onlyInBroadcast = true)() + val plan = FilterExec(DynamicPruningExpression(pruning), existingJoin) + val rewritten = PlanDynamicPruningFilters(spark).apply(plan) + val directFilters = rewritten.expressions.flatMap(_.collect { + case in: InSubqueryExec => in + }) + + assert(directFilters.size === 1, s"$name:\n$rewritten") + assert(directFilters.head.plan.isInstanceOf[SubqueryBroadcastExec], + s"$name:\n$rewritten") + } + } + } +} class DynamicPartitionPruningV1SuiteAEOn extends DynamicPartitionPruningV1Suite with EnableAdaptiveExecutionSuite { + test("adaptive broadcast value projection survives Catalyst copies") { + val source = spark.table("product").queryExecution.optimizedPlan + val key = source.output.head + val child = QueryExecution.createSparkPlan(spark.sessionState.planner, source) + val projection = BroadcastValueProjection(source, Seq(key), key) + val adaptive = SubqueryAdaptiveBroadcastExec( + "dynamicpruning", Seq(0), true, source, Seq(key), child)(Some(projection)) + + assert(adaptive.productArity === 6) + assert(SubqueryAdaptiveBroadcastExec.unapply(adaptive).exists(_.productArity == 6)) + assert(Modifier.isTransient( + classOf[SubqueryAdaptiveBroadcastExec].getDeclaredField("broadcastValueProjection") + .getModifiers)) + + Seq( + adaptive.copy()(adaptive.broadcastValueProjection), + adaptive.withNewChildren(Seq(ProjectExec(child.output, child))) + .asInstanceOf[SubqueryAdaptiveBroadcastExec], + adaptive.makeCopy(adaptive.productIterator.map(_.asInstanceOf[AnyRef]).toArray) + .asInstanceOf[SubqueryAdaptiveBroadcastExec] + ).foreach { rewritten => + assert(rewritten.broadcastValueProjection.contains(projection)) + } + + val unprojected = adaptive.copy()(None) + assert(adaptive === unprojected) + assert(adaptive.canonicalized.asInstanceOf[SubqueryAdaptiveBroadcastExec] + .broadcastValueProjection.isEmpty) + assert(adaptive.canonicalized === unprojected.canonicalized) + } + test("SPARK-39447: Avoid AssertionError in AdaptiveSparkPlanExec.doExecuteBroadcast") { val df = sql( """ @@ -2255,6 +2690,32 @@ class DynamicPartitionPruningV2FilterSuiteAEOn extends DynamicPartitionPruningV2FilterSuite with EnableAdaptiveExecutionSuite +/** + * Runs the DSv2 dynamic partition pruning tests against scans that receive runtime filters as + * Catalyst expressions, via + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], rather than as + * connector predicates. This is the counterpart of [[DynamicPartitionPruningV2FilterSuite]], + * which covers the same tests for + * [[org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering]]. + */ +abstract class DynamicPartitionPruningV2CatalystFilterSuite + extends DynamicPartitionPruningV2Suite { + + override protected def initState(): Unit = { + super.initState() + spark.conf.set("spark.sql.catalog.testcat", + classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) + } +} + +class DynamicPartitionPruningV2CatalystFilterSuiteAEOff + extends DynamicPartitionPruningV2CatalystFilterSuite + with DisableAdaptiveExecutionSuite + +class DynamicPartitionPruningV2CatalystFilterSuiteAEOn + extends DynamicPartitionPruningV2CatalystFilterSuite + with EnableAdaptiveExecutionSuite + private object DppMaterializedInputTestState { private val counters = TrieMap.empty[String, AtomicInteger] diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ExpressionsSchemaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ExpressionsSchemaSuite.scala index 4550ad76d43b8..5e5a5bb8d5805 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ExpressionsSchemaSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ExpressionsSchemaSuite.scala @@ -23,6 +23,7 @@ import java.nio.file.Files import scala.collection.mutable.ArrayBuffer import org.apache.spark.sql.catalyst.util.stringToFile +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils @@ -127,12 +128,15 @@ class ExpressionsSchemaSuite extends SharedSparkSession { // AvroDataToCatalyst or CatalystDataToAvro classes which are not available in this // test. case exampleRe(sql, _) => - val df = spark.sql(sql) - val escapedSql = sql.replaceAll("\\|", "|") - val schema = df.schema.catalogString.replaceAll("\\|", "|") - val queryOutput = QueryOutput(className, funcName, escapedSql, schema) - outputBuffer += queryOutput.toString - outputs += queryOutput + // parse_sql examples require the experimental feature flag. + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> true.toString) { + val df = spark.sql(sql) + val escapedSql = sql.replaceAll("\\|", "|") + val schema = df.schema.catalogString.replaceAll("\\|", "|") + val queryOutput = QueryOutput(className, funcName, escapedSql, schema) + outputBuffer += queryOutput.toString + outputs += queryOutput + } case _ => } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala index f045a92b6dc94..beabadd178295 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/InjectRuntimeFilterSuite.scala @@ -19,15 +19,18 @@ package org.apache.spark.sql import java.io.File -import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal} +import org.apache.spark.sql.catalyst.expressions.{Alias, BloomFilterMightContain, Literal, ScalarSubquery} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} -import org.apache.spark.sql.catalyst.optimizer.MergeSubplans import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan} +import org.apache.spark.sql.columnar.CachedBatch import org.apache.spark.sql.execution.{ReusedSubqueryExec, SubqueryExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEPropagateEmptyRelation} +import org.apache.spark.sql.execution.columnar.InMemoryRelation +import org.apache.spark.sql.execution.planmerging.MergeSubplans import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.storage.StorageLevel class InjectRuntimeFilterSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { @@ -299,6 +302,592 @@ class InjectRuntimeFilterSuite extends SharedSparkSession checkWithAndWithoutFeatureEnabled(query, shouldReplace = false) } + test("SPARK-58272: safely use fully materialized selectively filtered caches") { + val cacheName = "cached_bloom_filter_keys" + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + + def cachedRelation: InMemoryRelation = { + spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + } + + val before = cachedRelation + assert(before.materializedMetadata.isEmpty) + assert(!before.statsAvailable) + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + + val buffers = before.cacheBuilder.cachedColumnBuffers + assert(buffers.getNumPartitions > 1) + spark.sparkContext.runJob( + buffers, + (iterator: Iterator[CachedBatch]) => iterator.size, + Seq(0)) + assert(before.materializedMetadata.isEmpty) + assert(!before.statsAvailable) + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + + assert(spark.table(cacheName).count() == 1) + val materialized = cachedRelation + val metadata = materialized.materializedMetadata.get + assert(metadata.rowCount == 1L) + assert(metadata.statsAvailable) + assert(materialized.statsAvailable) + assert(materialized.isOutputRepeatable) + assert(materialized.hasSelectivePredicate) + val cachedSize = materialized.stats.sizeInBytes + assert(cachedSize > 0) + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + (cachedSize - 1).toString) { + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + cachedSize.toString) { + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 1) + checkAnswer(actual, expected) + + val leftOuter = + s"SELECT * FROM bf1 LEFT OUTER JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(leftOuter).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + + test("SPARK-58272: bound materialized caches by Bloom filter capacity") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MAX_NUM_ITEMS.key -> "4", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq( + ("at_bloom_capacity", 4L, 1), + ("over_bloom_capacity", 5L, 0) + ).foreach { case (cacheName, rowCount, expectedBloomFilters) => + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 8, 1, numPartitions = 4) + .where(s"id < $rowCount") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == rowCount) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + assert(relation.stats.rowCount.contains(rowCount)) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == + expectedBloomFilters) + } + } + } + } + } + + test("SPARK-58272: size projected cached runtime filters using exact materialized rows") { + val cacheName = "projected_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_EXPECTED_NUM_ITEMS.key -> "1", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "false") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 8, 1, numPartitions = 2) + .where("id < 4") + .selectExpr("CAST(id AS INT) AS c2", "CAST(id * 2 AS INT) AS payload") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 4) + + val cachedRelation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case relation: InMemoryRelation => relation + }.get + assert(cachedRelation.stats.rowCount.contains(BigInt(4))) + assert(cachedRelation.hasSelectivePredicate) + + val projectedKeys = spark.table(cacheName).selectExpr("c2 + 1 AS projected_key") + val projectedPlan = projectedKeys.queryExecution.optimizedPlan + assert(projectedPlan.stats.rowCount.isEmpty) + assert(projectedPlan.stats.sizeInBytes < cachedRelation.stats.sizeInBytes) + + val fact = spark.table("bf1") + val query = fact.join(projectedKeys, fact("c1") === projectedKeys("projected_key")) + val plan = withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + projectedPlan.stats.sizeInBytes.toString) { + query.queryExecution.optimizedPlan + } + assert(getNumBloomFilters(plan) == 1) + + val bloomAggregates = plan.collect { + case Filter(condition, _) => condition.collect { + case subquery: ScalarSubquery => subquery.plan.collect { + case Aggregate(_, aggregateExpressions, _, _) => aggregateExpressions.collect { + case Alias(AggregateExpression(aggregate: BloomFilterAggregate, _, _, _, _), _) => + aggregate + } + }.flatten + }.flatten + }.flatten + assert(bloomAggregates.size == 1) + assert(bloomAggregates.head.estimatedNumItemsExpression.eval() == 4L) + } + } + } + } + + test("SPARK-58272: preserve selective runtime filters over all cache states") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq( + ("memory_only_filtered_bloom_keys", StorageLevel.MEMORY_ONLY, true), + ("cold_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, false), + ("materialized_filtered_bloom_keys", StorageLevel.MEMORY_AND_DISK, true)).foreach { + case (cacheName, storageLevel, materialize) => + withTempView(cacheName) { + withCache(cacheName) { + val cached = spark.range(0, 40, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(storageLevel) + cached.createOrReplaceTempView(cacheName) + + if (materialize) { + assert(cached.count() == 40) + } + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cachedRelation: InMemoryRelation => cachedRelation + }.get + assert(relation.statsAvailable == (storageLevel.useDisk && materialize)) + assert(relation.isOutputRepeatable == materialize) + + val query = s"SELECT * FROM bf1 JOIN " + + s"(SELECT * FROM $cacheName WHERE c2 = 8) filtered_keys " + + "ON bf1.c1 = filtered_keys.c2" + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 1) + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + checkAnswer(actual, expected) + } + } + } + } + } + + test("SPARK-58272: reject materialized runtime filters for nonbinary collated join keys") { + val cacheName = "collated_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 1") + .selectExpr("IF(id = 1, 'a', 'z') COLLATE UTF8_LCASE AS k") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + + val actual = sql( + s"SELECT fact.k FROM " + + "(SELECT (CASE WHEN a1 = 73 THEN 'A' ELSE 'X' END) " + + "COLLATE UTF8_LCASE AS k FROM bf1) fact " + + s"JOIN $cacheName ON fact.k = $cacheName.k COLLATE UTF8_LCASE") + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == 0) + checkAnswer(actual, Row("A")) + } + } + } + } + + test("SPARK-58272: use materialized caches without predicates only with pruning statistics") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + val factPlan = spark.table("bf1").queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c1").get + val applicationDistinctCount = + factPlan.stats.attributeStats.get(factKey).flatMap(_.distinctCount).get + assert(applicationDistinctCount > 5 && applicationDistinctCount < 100) + + Seq(("small_cached_bloom_keys", 8L, 13L, true), + ("large_cached_bloom_keys", 0L, 100L, false)).foreach { + case (cacheName, start, end, shouldInject) => + withTempView(cacheName) { + withCache(cacheName) { + spark.range(start, end, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + spark.table(cacheName).count() + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(!relation.hasSelectivePredicate) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + val actual = sql(query) + assert(getNumBloomFilters(actual.queryExecution.optimizedPlan) == + (if (shouldInject) 1 else 0)) + + if (shouldInject) { + val derivedJoin = sql(s"SELECT * FROM bf1 JOIN $cacheName " + + s"ON bf1.c1 % 2 = $cacheName.c2 % 2") + assert(getNumBloomFilters(derivedJoin.queryExecution.optimizedPlan) == 0) + + val projectedFact = spark.table("bf1").selectExpr("c1 AS fact_key") + val projectedKeys = spark.table(cacheName).selectExpr("c2 AS cached_key") + val projectedJoin = projectedFact.join( + projectedKeys, projectedFact("fact_key") === projectedKeys("cached_key")) + assert(getNumBloomFilters(projectedJoin.queryExecution.optimizedPlan) == 1) + + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + checkAnswer(actual, expected) + + withSQLConf(SQLConf.CBO_ENABLED.key -> "false") { + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + } + } + + test("SPARK-58272: use filtered application-side distinct counts") { + val cacheName = "filtered_application_bloom_keys" + val nullFilteredCacheName = "null_filtered_application_bloom_keys" + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + withTempView(cacheName, nullFilteredCacheName) { + withCache(cacheName, nullFilteredCacheName) { + spark.range(8, 13, 1, numPartitions = 4) + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 5) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(!relation.hasSelectivePredicate) + + val filteredFact = spark.table("bf1").where("a1 = 77") + val factPlan = filteredFact.queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c1").get + assert(factPlan.stats.attributeStats.get(factKey) + .flatMap(_.distinctCount).contains(BigInt(1))) + + Seq( + s"SELECT * FROM (SELECT * FROM bf1 WHERE a1 = 77) fact " + + s"JOIN $cacheName keys ON fact.c1 = keys.c2", + s"SELECT * FROM $cacheName keys " + + "JOIN (SELECT * FROM bf1 WHERE a1 = 77) fact ON keys.c2 = fact.c1" + ).foreach { query => + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + + val actual = sql(query) + val plan = actual.queryExecution.optimizedPlan + assert(getNumBloomFilters(plan) == 1) + val applicationSides = plan.collect { + case Filter(condition, child) + if condition.exists(_.isInstanceOf[BloomFilterMightContain]) => child + } + assert(applicationSides.size == 1) + assert(applicationSides.head.exists(_.isInstanceOf[InMemoryRelation]), plan.treeString) + checkAnswer(actual, expected) + } + + val windowedQuery = + "SELECT * FROM (" + + "SELECT *, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY a1) AS rn " + + "FROM bf1 WHERE a1 = 77) fact " + + s"JOIN $cacheName keys ON fact.c1 = keys.c2" + assert(getNumBloomFilters(sql(windowedQuery).queryExecution.optimizedPlan) == 0) + + spark.range(0, 6, 1, numPartitions = 4) + .selectExpr( + "CAST(CASE id WHEN 0 THEN 0 WHEN 1 THEN 8 WHEN 2 THEN 23 " + + "WHEN 3 THEN 58 WHEN 4 THEN 74 ELSE 86 END AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(nullFilteredCacheName) + assert(spark.table(nullFilteredCacheName).count() == 6) + + val nullFilteredFact = spark.table("bf1").where("a1 IS NOT NULL") + val nullFilteredPlan = nullFilteredFact.queryExecution.optimizedPlan + val nullFilteredKey = nullFilteredPlan.output.find(_.name == "c1").get + assert(nullFilteredPlan.stats.attributeStats.get(nullFilteredKey) + .flatMap(_.distinctCount).contains(BigInt(6))) + + val query = s"SELECT * FROM (SELECT * FROM bf1 WHERE a1 IS NOT NULL) fact " + + s"JOIN $nullFilteredCacheName keys ON fact.c1 = keys.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + + test("SPARK-58272: preserve transitive selective filters past materialized caches") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "1MB", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.CBO_ENABLED.key -> "true") { + val factPlan = spark.table("bf3").queryExecution.optimizedPlan + val factKey = factPlan.output.find(_.name == "c3").get + val applicationDistinctCount = + factPlan.stats.attributeStats.get(factKey).flatMap(_.distinctCount).get + assert(applicationDistinctCount > 5 && applicationDistinctCount < 100) + + Seq( + ("large_transitive_bloom_keys", 0L, 100L, "1MB", false), + ("oversized_transitive_bloom_keys", 8L, 13L, "0", false), + ("selective_transitive_bloom_keys", 0L, 1000L, "1MB", true) + ).foreach { + case (cacheName, start, end, materializedThreshold, hasHiddenSelectivePredicate) => + withTempView(cacheName) { + withCache(cacheName) { + val keys = spark.range(start, end, 1, numPartitions = 4) + val cachedKeys = if (hasHiddenSelectivePredicate) { + keys.where("id < 100") + } else { + keys + } + cachedKeys + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + val expectedCount = if (hasHiddenSelectivePredicate) 100L else end - start + assert(spark.table(cacheName).count() == expectedCount) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate == hasHiddenSelectivePredicate) + + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_MATERIALIZED_CREATION_SIDE_THRESHOLD.key -> + materializedThreshold) { + Seq( + s"SELECT /*+ BROADCAST(selective) */ keys.c2 FROM $cacheName keys " + + "JOIN (SELECT c2 FROM bf2 WHERE a2 = 5) selective " + + "ON keys.c2 = selective.c2", + s"SELECT /*+ BROADCAST(selective) */ keys.c2 " + + "FROM (SELECT c2 FROM bf2 WHERE a2 = 5) selective " + + s"JOIN $cacheName keys ON selective.c2 = keys.c2" + ).foreach { creationSide => + val query = s"SELECT * FROM bf3 fact JOIN ($creationSide) creation " + + "ON fact.c3 = creation.c2" + var expected: Array[Row] = null + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "false") { + expected = sql(query).collect() + } + + val actual = sql(query) + val plan = actual.queryExecution.optimizedPlan + assert(getNumBloomFilters(plan) == 1) + val applicationSides = plan.collect { + case Filter(condition, child) + if condition.exists(_.isInstanceOf[BloomFilterMightContain]) => child + } + assert(applicationSides.size == 1) + assert(applicationSides.head.output.exists(_.name == "c3"), plan.treeString) + val creationPlans = plan.collect { + case Filter(condition, _) => condition.collect { + case BloomFilterMightContain(subquery: ScalarSubquery, _) => subquery.plan + } + }.flatten + assert(creationPlans.size == 1) + assert(!creationPlans.head.exists(_.isInstanceOf[InMemoryRelation]), + plan.treeString) + checkAnswer(actual, expected) + } + } + } + } + } + } + } + + test("SPARK-58272: reject memory-only and non-repeatable materialized caches") { + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + Seq(("memory_only_bloom_keys", false), + ("nondeterministic_bloom_keys", true)).foreach { + case (cacheName, nondeterministic) => + withTempView(cacheName) { + withCache(cacheName) { + val keys = if (nondeterministic) { + spark.range(0, 40, 1, numPartitions = 4) + .selectExpr("CAST(rand() * 20 AS INT) AS c2") + .where("c2 < 10") + .persist(StorageLevel.MEMORY_AND_DISK) + } else { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_ONLY) + } + keys.createOrReplaceTempView(cacheName) + keys.count() + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(!relation.statsAvailable) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + } + + test("SPARK-58272: do not trust cached runtime-replaceable encryption") { + val cacheName = "encrypted_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true") { + withTempView(cacheName) { + withCache(cacheName) { + spark.range(0, 40, 1, numPartitions = 4) + .where("id = 8") + .selectExpr( + "CAST(id AS INT) AS c2", + "aes_encrypt(CAST(id AS STRING), '0000111122223333') AS encrypted") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.hasSelectivePredicate) + assert(!relation.statsAvailable) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 0) + } + } + } + } + + test("SPARK-58272: use materialized selectively filtered file-backed caches") { + val cacheName = "parquet_cached_bloom_filter_keys" + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.IGNORE_MISSING_FILES.key -> "false", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + withTempPath { path => + spark.range(0, 40, 1, numPartitions = 4).write.parquet(path.getCanonicalPath) + withTempView(cacheName) { + withCache(cacheName) { + spark.read.parquet(path.getCanonicalPath) + .where("id = 8") + .selectExpr("CAST(id AS INT) AS c2") + .persist(StorageLevel.MEMORY_AND_DISK) + .createOrReplaceTempView(cacheName) + assert(spark.table(cacheName).count() == 1) + + val relation = spark.table(cacheName).queryExecution.withCachedData.collectFirst { + case cached: InMemoryRelation => cached + }.get + assert(relation.statsAvailable) + assert(relation.hasSelectivePredicate) + + val query = s"SELECT * FROM bf1 JOIN $cacheName ON bf1.c1 = $cacheName.c2" + assert(getNumBloomFilters(sql(query).queryExecution.optimizedPlan) == 1) + } + } + } + } + } + + test("SPARK-58272: materialized creation-side threshold exceeds broadcast threshold") { + val conf = new SQLConf + assert(conf.runtimeFilterMaterializedCreationSideThreshold > conf.autoBroadcastJoinThreshold) + } + test("Runtime bloom filter join: simple") { withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "3000", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "2000") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala index 7f695e90df884..3c0b00793ca5e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/JoinSuite.scala @@ -83,6 +83,14 @@ class JoinSuite extends SharedSparkSession with AdaptiveSparkPlanHelper canPlanAsBroadcastHashJoin(optimized.asInstanceOf[Join], conf) === operators.head.isInstanceOf[BroadcastHashJoinExec], "canPlanAsBroadcastHashJoin not in sync with join selection codepath!") + operators.head match { + case bhj: BroadcastHashJoinExec => + assert( + getBroadcastHashJoinBuildSide(optimized.asInstanceOf[Join], conf) + .contains(bhj.buildSide), + "getBroadcastHashJoinBuildSide not in sync with join selection codepath!") + case _ => + } operators.head } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonExistsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonExistsSuite.scala new file mode 100644 index 0000000000000..a4bac2c628a54 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonExistsSuite.scala @@ -0,0 +1,188 @@ +/* + * 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 + +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.catalyst.expressions.{JsonExists, JsonExistsBehavior, Literal} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.BooleanType + +/** + * End-to-end tests for the SQL:2016 `JSON_EXISTS` predicate. + */ +class JsonExistsSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val doc = + """{"id":7,"addr":{"city":"NYC"},"score":null,"tags":["x","y"]}""" + + test("returns BOOLEAN") { + assert(sql(s"SELECT json_exists('$doc', '$$.id')").schema.head.dataType === BooleanType) + } + + test("path present -> true, absent -> false") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr.city')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr.zip')"), Row(false)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.missing')"), Row(false)) + } + + test("present but JSON null -> true (distinguishes present-null from absent)") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.score')"), Row(true)) + } + + test("a matched object or array -> true") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.tags')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.tags[1]')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.tags[5]')"), Row(false)) + } + + test("NULL input yields NULL (Unknown), not the ON ERROR path") { + checkAnswer( + sql("SELECT json_exists(CAST(NULL AS STRING), '$.a' ERROR ON ERROR)"), Row(null)) + } + + test("malformed input: FALSE ON ERROR by default, and each ON ERROR behavior") { + checkAnswer(sql("SELECT json_exists('not json', '$.a')"), Row(false)) + checkAnswer(sql("SELECT json_exists('not json', '$.a' FALSE ON ERROR)"), Row(false)) + checkAnswer(sql("SELECT json_exists('not json', '$.a' TRUE ON ERROR)"), Row(true)) + checkAnswer(sql("SELECT json_exists('not json', '$.a' UNKNOWN ON ERROR)"), Row(null)) + val e = intercept[SparkRuntimeException] { + sql("SELECT json_exists('not json', '$.a' ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_EXISTS_ON_ERROR") + } + + test("empty or whitespace-only input is malformed -> ON ERROR") { + checkAnswer(sql("SELECT json_exists('', '$.a')"), Row(false)) + checkAnswer(sql("SELECT json_exists('', '$.a' TRUE ON ERROR)"), Row(true)) + checkAnswer(sql("SELECT json_exists(' ', '$.a')"), Row(false)) + checkAnswer(sql("SELECT json_exists(' ', '$.a' TRUE ON ERROR)"), Row(true)) + } + + test("a partial/prefix-garbage input is malformed -> ON ERROR") { + // A valid prefix followed by trailing garbage is not a single well-formed value. This holds + // whether the path matches (drainToRootEnd runs after the match) or is absent (drainToRootEnd + // runs after the miss) -- both surface the trailing content as malformed input. + checkAnswer(sql("""SELECT json_exists('{"a":1} trailing', '$.a')"""), Row(false)) + checkAnswer(sql("""SELECT json_exists('{"a":1} trailing', '$.a' TRUE ON ERROR)"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"a":1} trailing', '$.missing')"""), Row(false)) + checkAnswer( + sql("""SELECT json_exists('{"a":1} trailing', '$.missing' TRUE ON ERROR)"""), Row(true)) + // A second well-formed root value (not just garbage) is also more than one value -> malformed; + // drainToRootEnd must reject the trailing root, even though the path matched in the first one. + checkAnswer(sql("""SELECT json_exists('{"a":1} {"b":2}', '$.a')"""), Row(false)) + checkAnswer(sql("""SELECT json_exists('{"a":1} {"b":2}', '$.a' TRUE ON ERROR)"""), Row(true)) + } + + test("works over a column of JSON documents") { + withTempView("docs") { + Seq( + (1, """{"a":1}"""), + (2, """{"b":2}"""), // absent -> false + (3, """{"a":null}"""), // present null -> true + (4, "not json")) // malformed -> false (default) + .toDF("k", "j").createOrReplaceTempView("docs") + checkAnswer( + sql("SELECT k, json_exists(j, '$.a') AS e FROM docs ORDER BY k"), + Seq(Row(1, true), Row(2, false), Row(3, true), Row(4, false))) + } + } + + test("bracket and nested path syntax") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$[\\'addr\\'][\\'city\\']')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr.city')"), Row(true)) + } + + test("sql escapes a quoted path literal so the rendering re-parses") { + // A bracket-quoted path contains single quotes; the `sql` rendering must escape them, otherwise + // it would emit invalid SQL such as JSON_EXISTS('{}', '$['a']['b']'). + val e = JsonExists(Literal("{}"), "$['a']['b']", JsonExistsBehavior.False) + val parsed = spark.sessionState.sqlParser.parseExpression(e.sql) + assert(parsed.isInstanceOf[JsonExists]) + assert(parsed.asInstanceOf[JsonExists].path === "$['a']['b']") + } + + test("lax wildcard [*]: true iff the array has elements; auto-wraps a non-array") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.tags[*]')"), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"tags":[]}', '$.tags[*]')"""), Row(false)) + checkAnswer(sql("""SELECT json_exists('{"a":[1,2]}', '$.a[*]')"""), Row(true)) + // lax auto-wrap: a non-array value (scalar or object) is a single-element array. + checkAnswer(sql("""SELECT json_exists('{"a":5}', '$.a[*]')"""), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr[*]')"), Row(true)) + // lax auto-wrap for an explicit index: over a non-array, [0] matches the wrapped value and + // any [i>0] does not. + checkAnswer(sql("""SELECT json_exists('{"a":5}', '$.a[0]')"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"a":5}', '$.a[1]')"""), Row(false)) + } + + test("lax embedded wildcard $.a[*].b matches when any element has the field") { + checkAnswer(sql("""SELECT json_exists('{"a":[{"b":1},{"c":2}]}', '$.a[*].b')"""), Row(true)) + // Match only in a later element (guards the short-circuit loop). + checkAnswer(sql("""SELECT json_exists('{"a":[{"c":1},{"b":2}]}', '$.a[*].b')"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"a":[{"c":1}]}', '$.a[*].b')"""), Row(false)) + } + + test("lax member wildcard .* and ['*'] match any member") { + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr.*')"), Row(true)) + checkAnswer(sql("""SELECT json_exists('{}', '$.*')"""), Row(false)) + // ['*'] is the bracket-quoted spelling of the same member wildcard. + checkAnswer(sql(s"SELECT json_exists('$doc', '$$[\\'*\\']')"), Row(true)) + checkAnswer(sql(s"SELECT json_exists('$doc', '$$.addr[\\'*\\']')"), Row(true)) + checkAnswer(sql("SELECT json_exists('{}', '$[\\'*\\']')"), Row(false)) + // lax auto-unwrap: a member wildcard over an array applies to each element, matching iff some + // element has a member. + checkAnswer(sql("""SELECT json_exists('[{"a":1}]', '$.*')"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('[1]', '$.*')"""), Row(false)) + } + + test("index step over a non-array member is skipped without corrupting the traversal") { + // $.*[1] visits each member; [1] over the object member `a` must be fully consumed (return + // false) so the traversal advances to member `b`, where b[1] matches. + checkAnswer(sql("""SELECT json_exists('{"a":{"x":1},"b":[0,1]}', '$.*[1]')"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"a":{"x":1}}', '$.*[1]')"""), Row(false)) + } + + test("lax auto-unwrap: a member step over an array applies to each element") { + checkAnswer(sql("""SELECT json_exists('{"a":[{"b":1},{"b":2}]}', '$.a.b')"""), Row(true)) + checkAnswer(sql("""SELECT json_exists('{"a":[{"c":1}]}', '$.a.b')"""), Row(false)) + } + + test("duplicate object keys: first-match, consistent with JSON_VALUE / JSON_TABLE") { + // A named-key step follows only the first member with that name; a later duplicate is ignored. + checkAnswer(sql("""SELECT json_exists('{"a":{},"a":{"b":1}}', '$.a.b')"""), Row(false)) + checkAnswer(sql("""SELECT json_exists('{"a":{"b":1},"a":{}}', '$.a.b')"""), Row(true)) + // The duplicate key itself still exists, so a path that stops at it matches. + checkAnswer(sql("""SELECT json_exists('{"a":{},"a":{"b":1}}', '$.a')"""), Row(true)) + } + + test("invalid: an unparseable path is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_exists('$doc', '$$[')").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_PATH") + } + + test("usable in a WHERE predicate") { + withTempView("docs") { + Seq((1, """{"a":1}"""), (2, """{"b":2}""")).toDF("k", "j").createOrReplaceTempView("docs") + checkAnswer( + sql("SELECT k FROM docs WHERE json_exists(j, '$.a')"), Seq(Row(1))) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala index 4bbd3b533d34a..c0ae3fc157f71 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonFunctionsSuite.scala @@ -1798,6 +1798,37 @@ class JsonFunctionsSuite extends SharedSparkSession { } } + test("SPARK-58707: json pruning keeps the corrupt record column populated") { + Seq("true", "false").foreach { enabled => + withSQLConf(SQLConf.JSON_EXPRESSION_OPTIMIZATION.key -> enabled) { + // `b` is a type mismatch rather than a structural malformation, so it only fails when the + // parser is asked for it. Pruning the schema down to the corrupt record column alone left + // the parser nothing to convert, so the record was reported as clean. + val df = Seq("""{"a": 1, "b": "bad"}""").toDS() + .selectExpr("from_json(value, 'a int, b int, _corrupt_record string') as p") + .selectExpr("p._corrupt_record") + + checkAnswer(df, Row("""{"a": 1, "b": "bad"}""")) + } + } + } + + test("SPARK-58707: named_struct keeps the corrupt record column populated") { + Seq("true", "false").foreach { enabled => + withSQLConf(SQLConf.JSON_EXPRESSION_OPTIMIZATION.key -> enabled) { + // The two from_json calls are written inline so that CollapseProject does not keep them + // behind an alias, which would stop the rule from firing at all. + val fromJson = "from_json(value, 'a int, b int, _corrupt_record string')" + val df = Seq("""{"a": 1, "b": "bad"}""").toDS() + .selectExpr( + s"named_struct('a', $fromJson.a, '_corrupt_record', $fromJson._corrupt_record) as s") + .selectExpr("s.a", "s._corrupt_record") + + checkAnswer(df, Row(1, """{"a": 1, "b": "bad"}""")) + } + } + } + test("SPARK-33907: json pruning optimization with corrupt record field") { Seq("true", "false").foreach { enabled => withSQLConf(SQLConf.JSON_EXPRESSION_OPTIMIZATION.key -> enabled) { @@ -1806,6 +1837,9 @@ class JsonFunctionsSuite extends SharedSparkSession { .add("b", IntegerType) val badRec = """{"a" 1, "b": 11}""" + // Since SPARK-58707 the rule no longer prunes to the corrupt record column, so both + // iterations run the same plan. The record here is structurally malformed, which fails at + // tokenization whatever schema is requested, so the answer never depended on the pruning. val df = Seq(badRec, """{"a": 2, "b": 12}""").toDS() .selectExpr("from_json(value, 'a int, b int, _corrupt_record string') as parsed") .selectExpr("parsed._corrupt_record") @@ -1815,6 +1849,31 @@ class JsonFunctionsSuite extends SharedSparkSession { } } + test("SPARK-58373: bad json input with json pruning optimization: named_struct") { + Seq("true", "false").foreach { enabled => + withSQLConf(SQLConf.JSON_EXPRESSION_OPTIMIZATION.key -> enabled) { + // `c` is a type mismatch rather than a structural malformation, so it only fails + // when the parser is actually asked for it. A structurally broken record would + // fail at tokenization regardless of the requested schema and hide the pruning. + // The two from_json calls are written inline: an aliased subquery reference would + // not be inlined by CollapseProject, so the named_struct would not see them. + val fromJson = "from_json(value, 'a int, b int, c int', map('mode', 'FAILFAST'))" + val df = Seq("""{"a": 1, "b": 2, "c": "bad"}""").toDS() + .selectExpr(s"named_struct('a', $fromJson.a, 'b', $fromJson.b)") + + checkError( + exception = intercept[SparkException] { + df.collect() + }, + condition = "MALFORMED_RECORD_IN_PARSING.WITHOUT_SUGGESTION", + parameters = Map( + "badRecord" -> "[1,2,null]", + "failFastMode" -> "FAILFAST") + ) + } + } + } + test("SPARK-47670: separately projected from_json fields preserve malformed-input semantics") { withSQLConf( SQLConf.JSON_EXPRESSION_OPTIMIZATION.key -> "true", @@ -1951,6 +2010,15 @@ class JsonFunctionsSuite extends SharedSparkSession { checkAnswer(df.select(json_object_keys($"a")), expected) } + test("json_typeof function") { + val df = Seq(null, "{}", "[1, 2, 3]", "\"str\"", "123", "true", "null", "", "bad") + .toDF("a") + val expected = Seq(Row(null), Row("object"), Row("array"), Row("string"), + Row("number"), Row("boolean"), Row("null"), Row(null), Row(null)) + checkAnswer(df.selectExpr("json_typeof(a)"), expected) + checkAnswer(df.select(json_typeof($"a")), expected) + } + test("function get_json_object - Codegen Support") { withTempView("GetJsonObjectTable") { val data = Seq(("1", """{"f1": "value1", "f5": 5.23}""")).toDF("key", "jstring") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonQuerySuite.scala new file mode 100644 index 0000000000000..083d319f6c279 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonQuerySuite.scala @@ -0,0 +1,228 @@ +/* + * 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 + +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch +import org.apache.spark.sql.catalyst.expressions.{JsonQuery, JsonQueryBehavior, JsonQueryQuotes, + JsonQueryWrapper, Literal} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{CharType, StringType, VarcharType} + +/** + * End-to-end tests for the SQL:2016 `JSON_QUERY` function. + */ +class JsonQuerySuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val doc = + """{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null}""" + + test("extract an object or array as verbatim JSON text") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.addr')"), Row("""{"city":"NYC"}""")) + checkAnswer(sql(s"SELECT json_query('$doc', '$$.tags')"), Row("""["x","y"]""")) + } + + test("default RETURNING type is STRING") { + assert(sql(s"SELECT json_query('$doc', '$$.addr')").schema.head.dataType === StringType) + } + + test("RETURNING VARCHAR/CHAR is normalized to STRING and does not truncate") { + // JSON_QUERY returns the fragment verbatim (no length-enforcing cast), so a CHAR/VARCHAR + // RETURNING must not advertise a length it cannot enforce. The result type is STRING and the + // value is not truncated -- including when char/varchar type info is otherwise preserved. + Seq("false", "true").foreach { preserve => + withSQLConf(SQLConf.PRESERVE_CHAR_VARCHAR_TYPE_INFO.key -> preserve) { + // VARCHAR(n) and CHAR(n) exercise the two separate normalization branches. + Seq("VARCHAR(2)", "CHAR(2)").foreach { returning => + val df = sql(s"SELECT json_query('$doc', '$$.addr' RETURNING $returning)") + assert(df.schema.head.dataType === StringType, s"$returning preserve=$preserve") + checkAnswer(df, Row("""{"city":"NYC"}""")) + } + } + } + } + + test("a directly-constructed JsonQuery with a CHAR/VARCHAR RETURNING is rejected") { + // The parser normalizes CHAR/VARCHAR to STRING, but a raw CharType/VarcharType supplied by + // direct Catalyst construction would otherwise advertise a length JSON_QUERY does not enforce. + // isValidReturningType rejects it, so checkInputDataTypes fails. + Seq(VarcharType(2), CharType(2)).foreach { returning => + val expr = JsonQuery(Literal("{}"), "$.a", returning, JsonQueryWrapper.Without, + JsonQueryQuotes.Keep, JsonQueryBehavior.Null, JsonQueryBehavior.Null) + expr.checkInputDataTypes() match { + case DataTypeMismatch(errorSubClass, _) => + assert(errorSubClass == "INVALID_JSON_QUERY_RETURNING_TYPE", s"for $returning") + case other => fail(s"expected DataTypeMismatch for $returning, got $other") + } + } + } + + test("RETURNING STRING is allowed (the result is JSON text)") { + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.addr' RETURNING STRING)"), Row("""{"city":"NYC"}""")) + } + + test("a scalar result is emitted as JSON text under the default WITHOUT ARRAY WRAPPER") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.id')"), Row("7")) + // A string scalar keeps its surrounding quotes by default (KEEP QUOTES). + checkAnswer(sql(s"SELECT json_query('$doc', '$$.name')"), Row("\"Ada\"")) + } + + test("a present JSON null yields the JSON text null") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.score')"), Row("null")) + } + + test("a missing path is an ON EMPTY case, NULL by default") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.missing')"), Row(null)) + } + + test("NULL JSON input propagates to NULL (not ON EMPTY / ON ERROR)") { + checkAnswer(sql("SELECT json_query(CAST(NULL AS STRING), '$.a')"), Row(null)) + } + + test("WITH [UNCONDITIONAL] ARRAY WRAPPER wraps the result in a one-element array") { + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.tags[0]' WITH ARRAY WRAPPER)"), Row("""["x"]""")) + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.tags' WITH UNCONDITIONAL ARRAY WRAPPER)"), + Row("""[["x","y"]]""")) + checkAnswer(sql(s"SELECT json_query('$doc', '$$.id' WITH ARRAY WRAPPER)"), Row("[7]")) + } + + test("WITH CONDITIONAL ARRAY WRAPPER wraps only a scalar") { + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.id' WITH CONDITIONAL ARRAY WRAPPER)"), Row("[7]")) + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.addr' WITH CONDITIONAL ARRAY WRAPPER)"), + Row("""{"city":"NYC"}""")) + checkAnswer( + sql(s"SELECT json_query('$doc', '$$.tags' WITH CONDITIONAL ARRAY WRAPPER)"), + Row("""["x","y"]""")) + } + + test("OMIT QUOTES strips the quotes from a scalar string result") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.name' OMIT QUOTES)"), Row("Ada")) + // KEEP QUOTES is the default and keeps them. + checkAnswer(sql(s"SELECT json_query('$doc', '$$.name' KEEP QUOTES)"), Row("\"Ada\"")) + // OMIT QUOTES is a no-op for a non-string scalar and for structural results. + checkAnswer(sql(s"SELECT json_query('$doc', '$$.id' OMIT QUOTES)"), Row("7")) + checkAnswer(sql(s"SELECT json_query('$doc', '$$.addr' OMIT QUOTES)"), Row("""{"city":"NYC"}""")) + } + + test("OMIT QUOTES unescapes an escaped string scalar") { + // Pass the JSON via a column so SQL string-literal escaping does not rewrite it first. The + // stored JSON is {"s":"a\"b\\c\n"}; s decodes to a, quote, b, backslash, c, newline. + val df = Seq("""{"s":"a\"b\\c\n"}""").toDF("j") + // KEEP QUOTES (default) returns the verbatim, re-escaped JSON string. + checkAnswer(df.selectExpr("json_query(j, '$.s')"), Row(""""a\"b\\c\n"""")) + // OMIT QUOTES returns the raw unescaped content, which is intentionally no longer valid JSON. + checkAnswer(df.selectExpr("json_query(j, '$.s' OMIT QUOTES)"), Row("a\"b\\c\n")) + } + + test("the JSON_QUERY keyword is non-reserved and usable as an identifier") { + // JSON_QUERY and the OBJECT keyword introduced for the ON EMPTY / ON ERROR clause are + // non-reserved in both modes, so they remain usable as column names. + withTable("t") { + sql("CREATE TABLE t (json_query INT, object STRING) USING parquet") + sql("INSERT INTO t VALUES (1, 'x')") + checkAnswer(sql("SELECT json_query, object FROM t"), Row(1, "x")) + } + } + + test("EMPTY ARRAY / EMPTY OBJECT ON EMPTY") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.missing' EMPTY ARRAY ON EMPTY)"), Row("[]")) + checkAnswer(sql(s"SELECT json_query('$doc', '$$.missing' EMPTY OBJECT ON EMPTY)"), Row("{}")) + } + + test("ERROR ON EMPTY raises for a missing path") { + val e = intercept[SparkRuntimeException] { + sql(s"SELECT json_query('$doc', '$$.missing' ERROR ON EMPTY)").collect() + } + assert(e.getCondition == "JSON_QUERY_ON_ERROR.EMPTY") + } + + test("malformed input is an ON ERROR case, NULL by default") { + checkAnswer(sql("SELECT json_query('not json', '$.a')"), Row(null)) + checkAnswer(sql("SELECT json_query('not json', '$.a' EMPTY ARRAY ON ERROR)"), Row("[]")) + checkAnswer(sql("SELECT json_query('not json', '$.a' EMPTY OBJECT ON ERROR)"), Row("{}")) + } + + test("ERROR ON ERROR raises for malformed input") { + val e = intercept[SparkRuntimeException] { + sql("SELECT json_query('not json', '$.a' ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_QUERY_ON_ERROR.ERROR") + } + + test("a valid JSON prefix followed by trailing content is an ON ERROR case") { + checkAnswer(sql("""SELECT json_query('{"a":{"b":1}} trailing', '$.a')"""), Row(null)) + checkAnswer(sql("""SELECT json_query('{"a":1}{"a":2}', '$.a')"""), Row(null)) + val e = intercept[SparkRuntimeException] { + sql("""SELECT json_query('{"a":{"b":1}} trailing', '$.a' ERROR ON ERROR)""").collect() + } + assert(e.getCondition == "JSON_QUERY_ON_ERROR.ERROR") + } + + test("nested path into an object and array index") { + checkAnswer(sql(s"SELECT json_query('$doc', '$$.addr.city')"), Row("\"NYC\"")) + checkAnswer(sql(s"SELECT json_query('$doc', '$$.tags[1]')"), Row("\"y\"")) + } + + test("invalid: wildcard path is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_query('$doc', '$$.tags[*]')").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_PATH") + } + + test("invalid: non-string RETURNING type is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_query('$doc', '$$.id' RETURNING INT)").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_QUERY_RETURNING_TYPE") + } + + test("invalid: OMIT QUOTES combined with an array wrapper is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_query('$doc', '$$.name' WITH ARRAY WRAPPER OMIT QUOTES)").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_QUERY_WRAPPER_AND_QUOTES") + } + + test("sql renders a bracket-quoted path as a valid, re-parseable string literal") { + val df = sql(s"SELECT json_query('$doc', '$$[\\'addr\\']')") + val jsonQuery = df.queryExecution.analyzed.expressions + .flatMap(_.collect { case jq: JsonQuery => jq }).head + val rendered = jsonQuery.sql + assert(rendered.contains("\\'addr\\'"), s"path was not escaped in: $rendered") + checkAnswer(sql(s"SELECT $rendered"), Row("""{"city":"NYC"}""")) + } + + test("works over a column of JSON documents") { + val df = Seq( + """{"a":{"x":1}}""", + """{"a":[1,2]}""", + """{"b":2}""", + "not json").toDF("j") + checkAnswer( + df.selectExpr("json_query(j, '$.a')"), + Seq(Row("""{"x":1}"""), Row("[1,2]"), Row(null), Row(null))) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonTableSuite.scala new file mode 100644 index 0000000000000..3772cffc39d19 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonTableSuite.scala @@ -0,0 +1,800 @@ +/* + * 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 + +import org.apache.spark.{SparkException, SparkThrowable} +import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, LongType, StringType, TimestampType} + +/** + * End-to-end tests for the SQL:2016 `JSON_TABLE` table-valued function (flat, non-nested subset). + */ +class JsonTableSuite extends QueryTest with SharedSparkSession { + + test("expand a JSON array into rows with typed columns and ordinality") { + val json = """{"items":[{"id":1,"n":"a"},{"id":2,"n":"b"},{"id":3,"n":"c"}]}""" + val df = sql( + s""" + |SELECT t.* FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS ( + | seq FOR ORDINALITY, + | id INT PATH '$$.id', + | name STRING PATH '$$.n' + | ) + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1L, 1, "a"), Row(2L, 2, "b"), Row(3L, 3, "c"))) + // Ordinality is a BIGINT, typed columns take their declared types. + assert(df.schema.map(_.dataType) === Seq(LongType, IntegerType, StringType)) + } + + test("implicit column path derived from column name") { + val json = """{"rows":[{"id":10,"name":"x"},{"id":20,"name":"y"}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.rows[*]' + | COLUMNS (id INT, name STRING) + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(10, "x"), Row(20, "y"))) + } + + test("row path matching a single object yields one row") { + val json = """{"a":1,"b":"hello"}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$' + | COLUMNS (a INT PATH '$$.a', b STRING PATH '$$.b') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1, "hello"))) + } + + test("missing column path yields null") { + val json = """{"items":[{"id":1},{"id":2,"n":"b"}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (id INT PATH '$$.id', name STRING PATH '$$.n') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1, null), Row(2, "b"))) + } + + test("EXISTS column reports presence as boolean") { + val json = """{"items":[{"id":1,"opt":9},{"id":2}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (id INT PATH '$$.id', hasOpt BOOLEAN EXISTS PATH '$$.opt') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1, true), Row(2, false))) + } + + test("empty array yields no rows") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[]}', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq.empty) + } + + test("row path matching nothing yields no rows") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]}', + | '$.absent[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq.empty) + } + + test("NULL ON ERROR (default) yields no rows for malformed JSON") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{ this is not valid json', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq.empty) + + // Explicit NULL ON ERROR behaves the same. + val df2 = sql( + """ + |SELECT * FROM json_table( + | '{ this is not valid json', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + | NULL ON ERROR + |) AS t + """.stripMargin) + checkAnswer(df2, Seq.empty) + } + + test("ERROR ON ERROR raises for malformed JSON") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{ this is not valid json', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + | ERROR ON ERROR + |) AS t + """.stripMargin) + intercept[SparkException] { + df.collect() + } + } + + test("null JSON input yields no rows") { + val df = sql( + """ + |SELECT * FROM json_table( + | CAST(NULL AS STRING), + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq.empty) + } + + test("untyped NULL input is coerced and yields no rows (NULL ON ERROR)") { + // An untyped SQL NULL (NullType) must be coerced to STRING and apply NULL ON ERROR, not be + // rejected during analysis. + val df = sql( + """ + |SELECT * FROM json_table( + | NULL, + | '$.items[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq.empty) + } + + test("join JSON_TABLE output against a base table") { + import testImplicits._ + withTempView("docs") { + Seq( + (1, """{"tags":[{"k":"a"},{"k":"b"}]}"""), + (2, """{"tags":[{"k":"c"}]}""") + ).toDF("id", "doc").createOrReplaceTempView("docs") + + val df = sql( + """ + |SELECT d.id, t.k + |FROM docs d, + |LATERAL json_table(d.doc, '$.tags[*]' COLUMNS (k STRING PATH '$.k')) AS t + |ORDER BY d.id, t.k + """.stripMargin) + checkAnswer(df, Seq(Row(1, "a"), Row(1, "b"), Row(2, "c"))) + } + } + + test("nested field extraction within a row item") { + val json = """{"items":[{"meta":{"score":7}},{"meta":{"score":8}}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (score INT PATH '$$.meta.score') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(7), Row(8))) + } + + test("column path that is a prefix of another is resolved correctly") { + // `$.meta` both terminates a column and is a prefix of `$.meta.score` and `$.meta.name`, so + // all three (plus an EXISTS on the prefix) must resolve from the same item. + val json = """{"items":[{"meta":{"score":7,"name":"a"}},{"meta":{"score":8,"name":"b"}}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS ( + | meta STRING PATH '$$.meta', + | has_meta BOOLEAN EXISTS PATH '$$.meta', + | score INT PATH '$$.meta.score', + | name STRING PATH '$$.meta.name') + |) AS t + """.stripMargin) + checkAnswer(df, Seq( + Row("""{"score":7,"name":"a"}""", true, 7, "a"), + Row("""{"score":8,"name":"b"}""", true, 8, "b"))) + } + + test("ordinality-only table produces one numbered row per array element") { + // No path columns: every element still yields a row, numbered by ordinality, even when the + // element itself is a scalar (the item is never inspected for a value). + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[10,20,30]}', + | '$.items[*]' + | COLUMNS (seq FOR ORDINALITY) + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1L), Row(2L), Row(3L))) + } + + test("column alias list from the table alias") { + val json = """{"items":[{"id":1,"n":"a"}]}""" + val df = sql( + s""" + |SELECT renamed_id, renamed_name FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (id INT PATH '$$.id', name STRING PATH '$$.n') + |) AS t(renamed_id, renamed_name) + """.stripMargin) + checkAnswer(df, Seq(Row(1, "a"))) + } + + test("duplicate column names are rejected at parse time") { + val e = intercept[ParseException] { + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]}', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id', id STRING PATH '$.id') + |) AS t + """.stripMargin) + } + assert(e.getCondition == "INVALID_SQL_SYNTAX.DUPLICATE_JSON_TABLE_COLUMN") + assert(e.getMessageParameters.get("columnName") == "`id`") + } + + test("duplicate column name detection follows spark.sql.caseSensitive") { + // Names differing only in case: distinct under a case-sensitive resolver, colliding otherwise. + val mixedCase = + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1,"ID":2}]}', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id', ID INT PATH '$.ID') + |) AS t + """.stripMargin + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + checkAnswer(sql(mixedCase), Seq(Row(1, 2))) + } + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val e = intercept[ParseException](sql(mixedCase)) + assert(e.getCondition == "INVALID_SQL_SYNTAX.DUPLICATE_JSON_TABLE_COLUMN") + assert(e.getMessageParameters.get("columnName") == "`id`") + } + // An exact-case duplicate is rejected in both modes. + val exactDuplicate = + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]}', + | '$.items[*]' + | COLUMNS (id INT PATH '$.id', id INT PATH '$.id') + |) AS t + """.stripMargin + Seq("true", "false").foreach { caseSensitive => + withSQLConf(SQLConf.CASE_SENSITIVE.key -> caseSensitive) { + val e = intercept[ParseException](sql(exactDuplicate)) + assert(e.getCondition == "INVALID_SQL_SYNTAX.DUPLICATE_JSON_TABLE_COLUMN") + } + } + } + + test("value cast honors ANSI mode") { + val json = """{"items":[{"v":"not_a_number"}]}""" + val query = + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (v INT PATH '$$.v') + |) AS t + """.stripMargin + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + checkAnswer(sql(query), Seq(Row(null))) + } + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + // ANSI cast failures surface as a SparkThrowable (e.g. SparkNumberFormatException). + intercept[SparkThrowable] { + sql(query).collect() + } + } + } + + test("EXISTS distinguishes a present JSON null from a missing key") { + // A key present with a JSON null value EXISTS (true); a truly absent key does not (false). + val json = """{"items":[{"a":null},{"b":1}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (hasA BOOLEAN EXISTS PATH '$$.a') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(true), Row(false))) + } + + test("value column returns SQL NULL (not the string 'null') for a JSON null") { + // JSON null must become SQL NULL, while a JSON string "null" must remain the string. + val json = """{"items":[{"v":null},{"v":"null"},{"v":"x"}]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (v STRING PATH '$$.v') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(null), Row("null"), Row("x"))) + } + + test("value column returns SQL NULL for a JSON null reached via an array index") { + val json = """{"arr":[null, "kept"]}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$' + | COLUMNS (first STRING PATH '$$.arr[0]', second STRING PATH '$$.arr[1]') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(null, "kept"))) + } + + test("mid-path wildcard in a column or row path is rejected") { + // A wildcard anywhere except a single trailing '[*]' on the row path is unsupported. + val e1 = intercept[AnalysisException] { + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"x":1}]}', + | '$.items[*].x' + | COLUMNS (x INT PATH '$.x') + |) AS t + """.stripMargin).collect() + } + assert(e1.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_TABLE_PATH") + assert(e1.getMessageParameters.get("location") == "row path") + assert(e1.getMessageParameters.get("path") == "'$.items[*].x'") + // The rendered expression carries the full JSON_TABLE syntax (row path + columns + ON ERROR), + // not just the JSON input, so the diagnostic is actionable. + assert(e1.getMessage.contains( + "JSON_TABLE({\"items\":[{\"x\":1}]}, '$.items[*].x' " + + "COLUMNS (x INT PATH '$.x') NULL ON ERROR)")) + + val e2 = intercept[AnalysisException] { + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"vals":[1,2]}]}', + | '$.items[*]' + | COLUMNS (v INT PATH '$.vals[*]') + |) AS t + """.stripMargin).collect() + } + assert(e2.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_TABLE_PATH") + assert(e2.getMessageParameters.get("location") == "column 'v'") + assert(e2.getMessageParameters.get("path") == "'$.vals[*]'") + } + + test("[*] over a non-array fires ERROR ON ERROR and is empty under NULL ON ERROR") { + val json = """{"items":{"a":1}}""" + // NULL ON ERROR (default): a non-array under [*] yields no rows. + checkAnswer( + sql( + s""" + |SELECT * FROM json_table( + | '$json', '$$.items[*]' COLUMNS (a INT PATH '$$.a') + |) AS t + """.stripMargin), + Seq.empty) + // ERROR ON ERROR: the same input raises. + intercept[SparkException] { + sql( + s""" + |SELECT * FROM json_table( + | '$json', '$$.items[*]' COLUMNS (a INT PATH '$$.a') ERROR ON ERROR + |) AS t + """.stripMargin).collect() + } + } + + test("non-explode row path resolving to a string value") { + // The matched row item is a JSON string; a value column reading '$' must get its content. + val df = sql( + """ + |SELECT * FROM json_table( + | '{"name":"hello world"}', + | '$.name' + | COLUMNS (c STRING PATH '$', present BOOLEAN EXISTS PATH '$') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row("hello world", true))) + } + + test("row path resolving to a top-level string") { + val df = sql( + """ + |SELECT * FROM json_table('"just a string"', '$' COLUMNS (c STRING PATH '$')) AS t + """.stripMargin) + checkAnswer(df, Seq(Row("just a string"))) + } + + test("array of scalars as the row source") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"nums":[1,2,3]}', + | '$.nums[*]' + | COLUMNS (seq FOR ORDINALITY, v INT PATH '$') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1L, 1), Row(2L, 2), Row(3L, 3))) + } + + test("high-precision fractional value is preserved exactly, not rounded to a double") { + // The matched fragment is reserialized before JSON_TABLE casts it. The number carries more + // significant digits than a double can hold, so serializing via a lossy float copy would round + // it before the DECIMAL/STRING cast sees it. Exact structure copying keeps the digits verbatim. + val json = """{"v":123456789.123456789123456789}""" + val df = sql( + s""" + |SELECT * FROM json_table( + | '$json', + | '$$' + | COLUMNS (d DECIMAL(38, 18) PATH '$$.v', s STRING PATH '$$.v') + |) AS t + """.stripMargin) + checkAnswer(df, + Seq(Row(BigDecimal("123456789.123456789123456789"), "123456789.123456789123456789"))) + } + + test("array of strings as the row source keeps string content") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"tags":["a","b c"]}', + | '$.tags[*]' + | COLUMNS (v STRING PATH '$') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row("a"), Row("b c"))) + } + + test("duplicate keys within an item resolve to the first value") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"a":1,"a":2}]}', + | '$.items[*]' + | COLUMNS (a INT PATH '$.a') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1))) + } + + test("structure-valued column serialized to STRING") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"o":{"x":1}}]}', + | '$.items[*]' + | COLUMNS (o STRING PATH '$.o') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row("""{"x":1}"""))) + } + + test("unsupported declared column type is rejected at analysis") { + // A value column cannot be declared as a complex type that STRING cannot be cast to. + val e = intercept[AnalysisException] { + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"v":1}]}', + | '$.items[*]' + | COLUMNS (v STRUCT<a: INT> PATH '$.v') + |) AS t + """.stripMargin).collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION") + } + + test("oversized numeric path index is rejected as an invalid path") { + val e = intercept[AnalysisException] { + sql( + """ + |SELECT * FROM json_table( + | '{"a":[1]}', + | '$' + | COLUMNS (v INT PATH '$.a[999999999999999999999999]') + |) AS t + """.stripMargin).collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_TABLE_PATH") + } + + test("ERROR ON ERROR rejects trailing garbage and empty input") { + // A valid JSON prefix followed by garbage must raise, not silently produce rows. + intercept[SparkException] { + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]} trailing garbage', + | '$.items[*]' COLUMNS (id INT PATH '$.id') ERROR ON ERROR + |) AS t + """.stripMargin).collect() + } + // Empty input is malformed under ERROR ON ERROR. + intercept[SparkException] { + sql( + """ + |SELECT * FROM json_table('', '$.items[*]' COLUMNS (id INT PATH '$.id') ERROR ON ERROR) + |AS t + """.stripMargin).collect() + } + // Under the default NULL ON ERROR the same inputs simply yield no rows. + checkAnswer( + sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]} trailing garbage', + | '$.items[*]' COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin), + Seq.empty) + } + + test("large array row source is expanded correctly") { + // A big array is fully and correctly expanded. Row emission is streamed element by element + // from the source parser (the whole expanded payload is not materialized at once); note the + // input is still scanned once up front to validate it is a single well-formed JSON document. + val n = 5000 + val arr = (1 to n).map(i => s"""{"id":$i}""").mkString(",") + val json = s"""{"items":[$arr]}""" + val df = sql( + s""" + |SELECT count(*) AS c, sum(id) AS s FROM json_table( + | '$json', + | '$$.items[*]' + | COLUMNS (id INT PATH '$$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(n.toLong, (n.toLong * (n + 1)) / 2))) + + // A bare LIMIT (no ORDER BY) stops pulling generator rows early, so the streaming iterator is + // abandoned before exhaustion -- this exercises the task-completion-listener parser cleanup. + // Elements are in document order, so the first three ids are 1, 2, 3. + checkAnswer( + sql( + s""" + |SELECT id FROM json_table('$json', '$$.items[*]' COLUMNS (id INT PATH '$$.id')) AS t + |LIMIT 3 + """.stripMargin), + Seq(Row(1), Row(2), Row(3))) + } + + test("deep container path to an array row source") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"a":{"b":[{"id":1},{"id":2}]}}', + | '$.a.b[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1), Row(2))) + } + + test("indexed container path then wildcard") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"a":[{"vals":[{"id":10},{"id":20}]}]}', + | '$.a[0].vals[*]' + | COLUMNS (id INT PATH '$.id') + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(10), Row(20))) + } + + test("non-explode row path resolving to JSON null yields one row of nulls") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"x":null}', + | '$.x' + | COLUMNS (v STRING PATH '$', present BOOLEAN EXISTS PATH '$') + |) AS t + """.stripMargin) + // The row source is a JSON null: one row; a value column is SQL NULL, EXISTS is true. + checkAnswer(df, Seq(Row(null, true))) + } + + test("[*] over a JSON null container yields no rows / errors per ON ERROR") { + val json = """{"items":null}""" + checkAnswer( + sql( + s""" + |SELECT * FROM json_table('$json', '$$.items[*]' COLUMNS (id INT PATH '$$.id')) AS t + """.stripMargin), + Seq.empty) + intercept[SparkException] { + sql( + s""" + |SELECT * FROM json_table( + | '$json', '$$.items[*]' COLUMNS (id INT PATH '$$.id') ERROR ON ERROR + |) AS t + """.stripMargin).collect() + } + } + + test("implicit path for a column name containing a dot reads the literal key") { + // `a.b` with no PATH must read the JSON key "a.b", not the nested path a -> b. + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"a.b":1, "a":{"b":2}}]}', + | '$.items[*]' + | COLUMNS (`a.b` INT) + |) AS t + """.stripMargin) + checkAnswer(df, Seq(Row(1))) + } + + test("column cast eval mode is captured at plan construction, not execution") { + // Build the DataFrame with ANSI off (bad cast -> NULL); enabling ANSI afterwards must not + // change the already-planned generator's behavior. + val query = + """ + |SELECT * FROM json_table( + | '{"items":[{"v":"not_a_number"}]}', + | '$.items[*]' + | COLUMNS (v INT PATH '$.v') + |) AS t + """.stripMargin + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + val df = sql(query) + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + checkAnswer(df, Seq(Row(null))) + } + } + } + + test("FOR ORDINALITY column is non-nullable in the output schema") { + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"id":1}]}', + | '$.items[*]' + | COLUMNS (seq FOR ORDINALITY, id INT PATH '$.id') + |) AS t + """.stripMargin) + assert(!df.schema("seq").nullable) + assert(df.schema("id").nullable) + } + + test("castability check uses the session ANSI mode") { + // BOOLEAN -> TIMESTAMP is castable in non-ANSI mode but not in ANSI mode, so an + // `EXISTS ... TIMESTAMP` column must be accepted under non-ANSI and rejected under ANSI, + // matching the eval mode of the actual per-column Cast. + val query = + """ + |SELECT * FROM json_table( + | '{"items":[{"a":1}]}', + | '$.items[*]' + | COLUMNS (hasA TIMESTAMP EXISTS PATH '$.a') + |) AS t + """.stripMargin + withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") { + // Accepted at analysis; BOOLEAN true casts to a timestamp value. + assert(sql(query).schema("hasA").dataType == TimestampType) + } + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + val e = intercept[AnalysisException](sql(query).collect()) + assert(e.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION") + } + } + + test("CHAR/VARCHAR column types are normalized to STRING") { + // A raw CHAR/VARCHAR target has no runtime encoder; like a normal CAST, the declared type is + // normalized to STRING so the column is produced and queried without error. + val df = sql( + """ + |SELECT * FROM json_table( + | '{"items":[{"c":"hi","v":"world"}]}', + | '$.items[*]' + | COLUMNS (c CHAR(5) PATH '$.c', v VARCHAR(10) PATH '$.v') + |) AS t + """.stripMargin) + assert(df.schema("c").dataType == StringType) + assert(df.schema("v").dataType == StringType) + checkAnswer(df, Seq(Row("hi", "world"))) + } + + test("wide projection interleaving all column kinds keeps every column in its own slot") { + // The per-row projection reads column kinds, paths, casts and trie slots from parallel arrays, + // so a wide projection that interleaves the three kinds (and leaves some paths unmatched) + // pins each column to its own slot. + val n = 60 + val fields = (1 to n).map(i => s""""f$i":$i""").mkString(",") + val columns = (1 to n).map { i => + i % 3 match { + case 0 => s"ord$i FOR ORDINALITY" + case 1 => s"val$i INT PATH '$$.f$i'" + // Some EXISTS columns point at a key the item does not have, so both outcomes are covered. + case _ if i % 4 == 0 => s"ex$i BOOLEAN EXISTS PATH '$$.f$i'" + case _ => s"ex$i BOOLEAN EXISTS PATH '$$.missing$i'" + } + }.mkString(", ") + val df = sql( + s""" + |SELECT * FROM json_table('{"items":[{$fields}]}', '$$.items[*]' COLUMNS ($columns)) AS t + """.stripMargin) + val expected = (1 to n).map { i => + i % 3 match { + case 0 => 1L // single row, so ordinality is 1 + case 1 => i + case _ => i % 4 == 0 + } + } + checkAnswer(df, Seq(Row.fromSeq(expected))) + } + + test("array row source over many input rows streams without leaking parsers") { + // Exercises the per-task (not per-row) parser cleanup: many input rows, each with a `[*]` + // array row source, must all shred correctly. + withTempView("docs") { + spark.range(0, 200) + .selectExpr("concat('{\"items\":[{\"v\":', id, '},{\"v\":', id, '}]}') AS j") + .createOrReplaceTempView("docs") + val df = sql( + """ + |SELECT t.v + |FROM docs d, + |LATERAL json_table(d.j, '$.items[*]' COLUMNS (v INT PATH '$.v')) AS t + """.stripMargin) + // Two rows per input document. + assert(df.count() == 400) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/JsonValueSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/JsonValueSuite.scala new file mode 100644 index 0000000000000..d3106ac33ddbe --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/JsonValueSuite.scala @@ -0,0 +1,263 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql + +import org.apache.spark.SparkRuntimeException +import org.apache.spark.sql.catalyst.expressions.JsonValue +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StringType} + +/** + * End-to-end tests for the SQL:2016 `JSON_VALUE` scalar function. + */ +class JsonValueSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val doc = + """{"id":7,"name":"Ada","tags":["x","y"],"addr":{"city":"NYC"},"score":null,"f":"3.14"}""" + + test("extract a scalar value as STRING by default") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.name')"), Row("Ada")) + // Numbers and booleans come back as their JSON text under the default STRING RETURNING. + checkAnswer(sql(s"SELECT json_value('$doc', '$$.id')"), Row("7")) + } + + test("RETURNING casts the scalar to the requested type") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.id' RETURNING INT)"), Row(7)) + assert(sql(s"SELECT json_value('$doc', '$$.id' RETURNING INT)").schema.head.dataType + === IntegerType) + checkAnswer(sql(s"SELECT json_value('$doc', '$$.f' RETURNING DOUBLE)"), Row(3.14d)) + checkAnswer(sql("SELECT json_value('{\"v\":\"true\"}', '$.v' RETURNING BOOLEAN)"), Row(true)) + checkAnswer( + sql("SELECT json_value('{\"v\":\"2020-01-02\"}', '$.v' RETURNING DATE)"), + Row(java.sql.Date.valueOf("2020-01-02"))) + } + + test("default RETURNING type is STRING") { + assert(sql(s"SELECT json_value('$doc', '$$.name')").schema.head.dataType === StringType) + } + + test("a raw JSON number keeps its exact source digits (no double rounding)") { + // The matched scalar is read straight from the parser, so a fraction with more digits than a + // double can represent reaches the DECIMAL cast (and the default STRING form) intact. + val big = """{"v":0.123456789012345678}""" + checkAnswer( + sql(s"SELECT json_value('$big', '$$.v' RETURNING DECIMAL(38,18))"), + Row(new java.math.BigDecimal("0.123456789012345678"))) + checkAnswer(sql(s"SELECT json_value('$big', '$$.v')"), Row("0.123456789012345678")) + } + + test("a present JSON null yields SQL NULL") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.score')"), Row(null)) + } + + test("a non-scalar (object/array) match is an ON ERROR case, NULL by default") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.addr')"), Row(null)) + checkAnswer(sql(s"SELECT json_value('$doc', '$$.tags')"), Row(null)) + } + + test("a missing path is an ON EMPTY case, NULL by default") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.missing')"), Row(null)) + checkAnswer(sql(s"SELECT json_value('$doc', '$$.addr.zip')"), Row(null)) + } + + test("NULL JSON input propagates to NULL (not ON EMPTY / ON ERROR)") { + checkAnswer(sql("SELECT json_value(CAST(NULL AS STRING), '$.a' ERROR ON EMPTY ERROR ON ERROR)"), + Row(null)) + } + + test("DEFAULT ON EMPTY") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.missing' DEFAULT '?' ON EMPTY)"), Row("?")) + checkAnswer( + sql(s"SELECT json_value('$doc', '$$.missing' RETURNING INT DEFAULT 42 ON EMPTY)"), Row(42)) + } + + test("ERROR ON EMPTY raises for a missing path") { + val e = intercept[SparkRuntimeException] { + sql(s"SELECT json_value('$doc', '$$.missing' ERROR ON EMPTY)").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.EMPTY") + } + + test("DEFAULT ON ERROR for a non-scalar match and for malformed input") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.addr' DEFAULT 'n/a' ON ERROR)"), Row("n/a")) + checkAnswer(sql("SELECT json_value('not json', '$.a' DEFAULT 'bad' ON ERROR)"), Row("bad")) + } + + test("ERROR ON ERROR raises for malformed input") { + val e = intercept[SparkRuntimeException] { + sql("SELECT json_value('not json', '$.a' ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.ERROR") + } + + test("a valid JSON prefix followed by trailing content is an ON ERROR case") { + // The whole input must be a single well-formed JSON value: a value that parses but is trailed + // by garbage (or a second root value) is malformed, even when the path matches within the + // prefix. NULL ON ERROR by default; ERROR ON ERROR raises. + checkAnswer(sql("""SELECT json_value('{"a":1} trailing', '$.a')"""), Row(null)) + checkAnswer(sql("""SELECT json_value('{"a":1}{"a":2}', '$.a')"""), Row(null)) + val e = intercept[SparkRuntimeException] { + sql("""SELECT json_value('{"a":1} trailing', '$.a' ERROR ON ERROR)""").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.ERROR") + } + + test("ERROR ON ERROR raises for a non-scalar match") { + val e = intercept[SparkRuntimeException] { + sql(s"SELECT json_value('$doc', '$$.addr' ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.ERROR") + } + + test("a failed cast is an ON ERROR case") { + // NULL ON ERROR default. + checkAnswer(sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT)"), Row(null)) + // DEFAULT ON ERROR. + checkAnswer( + sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT DEFAULT -1 ON ERROR)"), Row(-1)) + // ERROR ON ERROR. + val e = intercept[SparkRuntimeException] { + sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.ERROR") + } + + test("independent ON EMPTY and ON ERROR behaviors") { + // Missing path -> ON EMPTY branch; malformed input / non-scalar value -> ON ERROR branch. + checkAnswer( + sql(s"SELECT json_value('$doc', '$$.missing' DEFAULT 'e' ON EMPTY DEFAULT 'r' ON ERROR)"), + Row("e")) + checkAnswer( + sql(s"SELECT json_value('$doc', '$$.addr' DEFAULT 'e' ON EMPTY DEFAULT 'r' ON ERROR)"), + Row("r")) + } + + test("ON ERROR governs a failed cast in BOTH ANSI and non-ANSI mode") { + // The extracted-scalar cast is always an ANSI (throwing) cast, so a bad conversion routes to ON + // ERROR identically regardless of the session's ANSI setting -- a non-ANSI session must not + // silently return NULL and bypass DEFAULT / ERROR ON ERROR. (A NULL-ON-ERROR-only check would + // pass vacuously in non-ANSI mode, where a lenient cast already yields NULL.) + Seq("true", "false").foreach { ansi => + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi) { + withClue(s"ansi=$ansi ") { + // A valid conversion still succeeds in both modes (the ANSI cast is not stricter here). + checkAnswer(sql(s"SELECT json_value('$doc', '$$.id' RETURNING INT)"), Row(7)) + // NULL ON ERROR (default). + checkAnswer(sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT)"), Row(null)) + // DEFAULT ON ERROR. + checkAnswer( + sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT DEFAULT -1 ON ERROR)"), + Row(-1)) + // ERROR ON ERROR. + val e = intercept[SparkRuntimeException] { + sql(s"SELECT json_value('$doc', '$$.name' RETURNING INT ERROR ON ERROR)").collect() + } + assert(e.getCondition == "JSON_VALUE_ON_ERROR.ERROR") + } + } + } + } + + test("works over a column of JSON documents") { + withTempView("docs") { + Seq( + (1, """{"a":10}"""), + (2, """{"a":20}"""), + (3, """{"b":30}"""), // missing -> NULL ON EMPTY + (4, "not json")) // malformed -> NULL ON ERROR + .toDF("k", "j").createOrReplaceTempView("docs") + checkAnswer( + sql("SELECT k, json_value(j, '$.a' RETURNING INT) AS a FROM docs ORDER BY k"), + Seq(Row(1, 10), Row(2, 20), Row(3, null), Row(4, null))) + } + } + + test("invalid: wildcard path is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_value('$doc', '$$.tags[*]')").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_PATH") + } + + test("invalid: non-scalar RETURNING type is rejected at analysis") { + val e = intercept[AnalysisException] { + sql(s"SELECT json_value('$doc', '$$.addr' RETURNING STRUCT<x:INT>)").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.INVALID_JSON_SCALAR_RETURNING_TYPE") + } + + test("invalid: a DEFAULT that cannot cast to the RETURNING type is rejected at analysis") { + // The DEFAULT-to-RETURNING cast is validated up front, so an uncastable default (here an + // ARRAY DEFAULT with RETURNING INT) fails at analysis rather than late in the fallback branch. + Seq("ON EMPTY", "ON ERROR").foreach { clause => + val e = intercept[AnalysisException] { + sql(s"SELECT json_value('{}', '$$.x' RETURNING INT DEFAULT array(1) $clause)").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION", + s"unexpected condition for $clause: ${e.getCondition}") + } + } + + test("bracket path syntax and array index") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.tags[0]')"), Row("x")) + checkAnswer(sql(s"SELECT json_value('$doc', '$$[\\'name\\']')"), Row("Ada")) + } + + test("sql renders a bracket-quoted path as a valid, re-parseable string literal") { + // A path containing single quotes (e.g. `$['name']`) must be escaped when rendered back to + // SQL, or the round-tripped statement is malformed. Extract the JsonValue expression, render + // its `.sql`, and confirm the rendered form both parses and evaluates to the same result. + val df = sql(s"SELECT json_value('$doc', '$$[\\'name\\']')") + val jsonValue = df.queryExecution.analyzed.expressions + .flatMap(_.collect { case jv: JsonValue => jv }).head + val rendered = jsonValue.sql + assert(rendered.contains("\\'name\\'"), s"path was not escaped in: $rendered") + checkAnswer(sql(s"SELECT $rendered"), Row("Ada")) + } + + test("nested path into an object") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.addr.city')"), Row("NYC")) + } + + test("DEFAULT NULL behaves like NULL ON ERROR/EMPTY") { + checkAnswer(sql(s"SELECT json_value('$doc', '$$.missing' DEFAULT NULL ON EMPTY)"), Row(null)) + checkAnswer(sql(s"SELECT json_value('$doc', '$$.addr' DEFAULT NULL ON ERROR)"), Row(null)) + } + + test("DEFAULT is a real child expression, not a constant") { + // A DEFAULT that references a column proves the DEFAULT clause is resolved as a child, and that + // the value flows per row. + withTempView("t") { + Seq((1, "a"), (2, "b")).toDF("k", "fallback").createOrReplaceTempView("t") + checkAnswer( + sql("SELECT json_value('{}', '$.x' DEFAULT fallback ON EMPTY) FROM t ORDER BY k"), + Seq(Row("a"), Row("b"))) + } + } + + test("a non-string JSON input is rejected at analysis") { + // JSON_VALUE takes a STRING JSON input; an INT input is a type mismatch (matching the existing + // JSON functions' ExpectsInputTypes behavior), not a silent coercion. + val e = intercept[AnalysisException] { + sql("SELECT json_value(123, '$.a')").collect() + } + assert(e.getCondition == "DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/MathFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/MathFunctionsSuite.scala index f7b27c2052935..3ce1940e10e7a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/MathFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/MathFunctionsSuite.scala @@ -324,6 +324,46 @@ class MathFunctionsSuite extends SharedSparkSession { testOneToOneMathFunction(rint, math.rint) } + test("truncate") { + val df = Seq(5, 55, 555).map(Tuple1(_)).toDF("a") + checkAnswer( + df.select(truncate($"a", lit(-1)), truncate($"a", lit(-2))), + Seq(Row(0, 0), Row(50, 0), Row(550, 500)) + ) + // Truncation rounds toward zero, unlike floor for negative values. + val df2 = Seq(1234.5678, -1234.5678).map(Tuple1(_)).toDF("a") + checkAnswer( + df2.select(truncate($"a", lit(2))), + Seq(Row(1234.56), Row(-1234.56)) + ) + checkAnswer( + df2.selectExpr("truncate(a, 2)"), + Seq(Row(1234.56), Row(-1234.56)) + ) + // The truncate(Column, Int) and truncate(Column) overloads. + checkAnswer( + df2.select(truncate($"a", 2)), + Seq(Row(1234.56), Row(-1234.56)) + ) + checkAnswer( + df2.select(truncate($"a")), + Seq(Row(1234.0), Row(-1234.0)) + ) + // Decimal input, the type that motivated the new Decimal.changePrecision rounding mode. + val df3 = Seq(BigDecimal("1234.5678"), BigDecimal("-1234.5678")).map(Tuple1(_)).toDF("a") + checkAnswer( + df3.select(truncate($"a", lit(2))), + Seq(Row(BigDecimal("1234.56")), Row(BigDecimal("-1234.56"))) + ) + // NaN and Infinity are returned unchanged, matching RoundBase's other rounding modes. + val df4 = Seq(Double.NaN, Double.PositiveInfinity, Double.NegativeInfinity) + .map(Tuple1(_)).toDF("a") + checkAnswer( + df4.select(truncate($"a", lit(2))), + Seq(Row(Double.NaN), Row(Double.PositiveInfinity), Row(Double.NegativeInfinity)) + ) + } + test("round/bround/ceil/floor") { val df = Seq(5, 55, 555).map(Tuple1(_)).toDF("a") checkAnswer( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/MetadataCacheSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/MetadataCacheSuite.scala index e9d272728c944..a8b68ab773000 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/MetadataCacheSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/MetadataCacheSuite.scala @@ -22,6 +22,7 @@ import java.io.File import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.tags.ExtendedSQLTest /** * Test suite to handle metadata cache related. @@ -129,6 +130,7 @@ class MetadataCacheV1Suite extends MetadataCacheSuite { } } +@ExtendedSQLTest class MetadataCacheV2Suite extends MetadataCacheSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/NestedDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/NestedDataSourceSuite.scala index 4f2794ff87dce..017fe22c666ea 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/NestedDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/NestedDataSourceSuite.scala @@ -20,6 +20,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{LongType, StructType} +import org.apache.spark.tags.ExtendedSQLTest // Datasource tests for nested schemas trait NestedDataSourceSuiteBase extends SharedSparkSession { @@ -83,6 +84,7 @@ class NestedDataSourceV1Suite extends NestedDataSourceSuiteBase { .set(SQLConf.USE_V1_SOURCE_LIST, nestedDataSources.mkString(",")) } +@ExtendedSQLTest class NestedDataSourceV2Suite extends NestedDataSourceSuiteBase { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala index 6cd49948630da..120679f661133 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/PlanStabilitySuite.scala @@ -189,6 +189,8 @@ trait PlanStabilitySuite extends DisableAdaptiveExecutionSuite { subqueriesMap.getOrElseUpdate(subquery.id, subqueriesMap.size + 1) case subquery: SubqueryBroadcastExec => subqueriesMap.getOrElseUpdate(subquery.id, subqueriesMap.size + 1) + case subquery: ProjectedBroadcastValueSubqueryExec => + subqueriesMap.getOrElseUpdate(subquery.id, subqueriesMap.size + 1) case ReusedSubqueryExec(subquery) => subqueriesMap.getOrElseUpdate(subquery.id, subqueriesMap.size + 1) case _ => -1 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PushDownJoinThroughUnionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/PushDownJoinThroughUnionSuite.scala index 3b38e7d11fd5b..3cd59f8738ab2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/PushDownJoinThroughUnionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/PushDownJoinThroughUnionSuite.scala @@ -18,13 +18,18 @@ package org.apache.spark.sql import org.apache.spark.SparkConf +import org.apache.spark.sql.catalyst.optimizer.BuildLeft +import org.apache.spark.sql.catalyst.plans.logical.Join +import org.apache.spark.sql.execution.FileSourceScanExec import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, BroadcastQueryStageExec} import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class PushDownJoinThroughUnionSuite - extends SharedSparkSession + extends QueryTest + with SharedSparkSession with AdaptiveSparkPlanHelper { import testImplicits._ @@ -35,8 +40,10 @@ class PushDownJoinThroughUnionSuite withTempView("fact1", "fact2", "dim") { withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val fact1 = Seq((1, "a"), (2, "b")).toDF("id", "val1") - val fact2 = Seq((3, "c"), (4, "d")).toDF("id", "val1") + // Every branch must stay larger than the dimension table in bytes, otherwise the planner + // builds from the branch and the rule declines to push the join down. + val fact1 = Seq((1, "a"), (2, "b"), (5, "e"), (6, "f")).toDF("id", "val1") + val fact2 = Seq((3, "c"), (4, "d"), (7, "g"), (8, "h")).toDF("id", "val1") val dim = Seq((1, "x"), (2, "y"), (3, "z")).toDF("id", "label") fact1.createOrReplaceTempView("fact1") @@ -54,6 +61,7 @@ class PushDownJoinThroughUnionSuite Row(2, "b", "y"), Row(3, "c", "z") )) + assertJoinCount(result, 2) } } } @@ -62,9 +70,11 @@ class PushDownJoinThroughUnionSuite withTempView("fact1", "fact2", "fact3", "dim") { withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val fact1 = Seq((1, 10), (2, 20)).toDF("id", "amount") - val fact2 = Seq((3, 30), (4, 40)).toDF("id", "amount") - val fact3 = Seq((1, 50), (5, 60)).toDF("id", "amount") + // Twelve rows per branch keep every branch strictly larger than the dimension table in + // bytes; the extra ids join with nothing and leave the expected rows unchanged. + val fact1 = ((1, 10) +: (2, 20) +: (21 to 30).map(i => (i, i * 10))).toDF("id", "amount") + val fact2 = ((3, 30) +: (4, 40) +: (31 to 40).map(i => (i, i * 10))).toDF("id", "amount") + val fact3 = ((1, 50) +: (5, 60) +: (41 to 50).map(i => (i, i * 10))).toDF("id", "amount") val dim = Seq((1, "web"), (2, "store"), (3, "catalog"), (5, "other")) .toDF("id", "channel") @@ -90,6 +100,7 @@ class PushDownJoinThroughUnionSuite Row(1, 50, "web"), Row(5, 60, "other") )) + assertJoinCount(result, 3) } } } @@ -118,14 +129,17 @@ class PushDownJoinThroughUnionSuite Row(3, "c", "z"), Row(99, "d", null) )) + assertJoinCount(result, 2) } } } test("Optimization disabled produces same results") { withTempView("fact1", "fact2", "dim") { - val fact1 = Seq((1, "a"), (2, "b")).toDF("id", "val1") - val fact2 = Seq((3, "c"), (4, "d")).toDF("id", "val1") + // Branches larger than the dimension table in bytes, so the rule fires in the first run and + // the comparison against the excluded-rule run is meaningful. + val fact1 = Seq((1, "a"), (2, "b"), (5, "e"), (6, "f")).toDF("id", "val1") + val fact2 = Seq((3, "c"), (4, "d"), (7, "g"), (8, "h")).toDF("id", "val1") val dim = Seq((1, "x"), (2, "y"), (3, "z")).toDF("id", "label") fact1.createOrReplaceTempView("fact1") @@ -146,14 +160,18 @@ class PushDownJoinThroughUnionSuite withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - checkAnswer(sql(query), expected) + val enabled = sql(query) + checkAnswer(enabled, expected) + assertJoinCount(enabled, 2) } withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760", SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> "org.apache.spark.sql.catalyst.optimizer.PushDownJoinThroughUnion") { - checkAnswer(sql(query), expected) + val excluded = sql(query) + checkAnswer(excluded, expected) + assertJoinCount(excluded, 1) } } } @@ -162,8 +180,12 @@ class PushDownJoinThroughUnionSuite withTempView("fact1", "fact2", "dim") { withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val fact1 = Seq((1, "a", 100), (2, "b", 200)).toDF("id", "val1", "val2") - val fact2 = Seq((3, "c", 300), (4, "d", 400)).toDF("id", "val1", "val2") + // Column pruning shrinks both sides before the rule reads their stats, so the branches need + // enough rows to stay strictly larger than the dimension table once pruned. + val fact1 = ((1, "a", 100) +: (2, "b", 200) +: (21 to 30).map(i => (i, "x", i))) + .toDF("id", "val1", "val2") + val fact2 = ((3, "c", 300) +: (4, "d", 400) +: (31 to 40).map(i => (i, "y", i))) + .toDF("id", "val1", "val2") val dim = Seq((1, "x", "extra1"), (2, "y", "extra2"), (3, "z", "extra3")) .toDF("id", "label", "info") @@ -182,6 +204,7 @@ class PushDownJoinThroughUnionSuite Row(2, "y"), Row(3, "z") )) + assertJoinCount(result, 2) } } } @@ -225,8 +248,12 @@ class PushDownJoinThroughUnionSuite withTempView("fact1", "fact2", "dim") { withSQLConf( SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { - val fact1 = Seq((1, "a"), (2, "b")).toDF("id", "val1") - val fact2 = Seq((3, "c"), (4, "d")).toDF("id", "val1") + // Each branch must stay strictly larger than the dimension table in bytes: at equal sizes + // the rule fires only through the `right <= left` tie-break in `getSmallerSide`. + val fact1 = Seq((1, "a"), (2, "b"), (5, "e"), (6, "f"), (9, "i"), (10, "j")) + .toDF("id", "val1") + val fact2 = Seq((3, "c"), (4, "d"), (7, "g"), (8, "h"), (11, "k"), (12, "l")) + .toDF("id", "val1") val dim = Seq((1, "x"), (2, "y"), (3, "z"), (4, "w")).toDF("id", "label") fact1.createOrReplaceTempView("fact1") @@ -244,7 +271,109 @@ class PushDownJoinThroughUnionSuite Row(1, "a", "x"), Row(3, "c", "z") )) + assertJoinCount(result, 2) + } + } + } + + test("SPARK-58449: right side is scanned once when only the Union side is broadcastable") { + withTable("fact1", "fact2", "dim") { + // The threshold sits between the two Union branches and the right side, so the Union is + // broadcastable but the right side is not. An inner join can broadcast either side, so the + // join still plans as a broadcast hash join and the rule used to fire, cloning the right side + // once per branch. Nothing reuses a bare probe-side scan, so the right side would be read + // twice. + spark.range(0, 4).selectExpr("id", "id AS v").write.format("parquet").saveAsTable("fact1") + spark.range(4, 8).selectExpr("id", "id AS v").write.format("parquet").saveAsTable("fact2") + spark.range(0, 2000).selectExpr("id AS did", "id AS label").write + .format("parquet").saveAsTable("dim") + + val unionSize = Seq("fact1", "fact2").map(t => + spark.table(t).queryExecution.optimizedPlan.stats.sizeInBytes).sum + val rightSize = spark.table("dim").queryExecution.optimizedPlan.stats.sizeInBytes + assert(unionSize < rightSize, "test setup: the Union must be smaller than the right side") + + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> unionSize.toString) { + val df = sql( + """SELECT f.id, d.label + |FROM (SELECT * FROM fact1 UNION ALL SELECT * FROM fact2) f + |JOIN dim d ON f.id = d.did + """.stripMargin) + checkAnswer(df, (0 until 8).map(i => Row(i.toLong, i.toLong))) + + assertBuildsFromTheLeft(df) + assertJoinCount(df, 1) + assert(rightScansOf(df) == 1, + s"the right side must be scanned once, found ${rightScansOf(df)} scans of it") } } } + + test("SPARK-58449: right side is scanned once when it is larger than each Union branch") { + withTable("fact1", "fact2", "dim") { + // Both sides are under the threshold, so an inner join can build from either one and the + // planner picks the smaller side, which is the left one. The rule used to fire anyway, + // because a broadcast hash join was plannable, and every branch then probed its own copy of + // the right side. + spark.range(0, 4).selectExpr("id", "id AS v").write.format("parquet").saveAsTable("fact1") + spark.range(4, 8).selectExpr("id", "id AS v").write.format("parquet").saveAsTable("fact2") + spark.range(0, 2000).selectExpr("id AS did", "id AS label").write + .format("parquet").saveAsTable("dim") + + val branchSizes = Seq("fact1", "fact2").map(t => + spark.table(t).queryExecution.optimizedPlan.stats.sizeInBytes) + val rightSize = spark.table("dim").queryExecution.optimizedPlan.stats.sizeInBytes + assert(branchSizes.forall(_ < rightSize) && branchSizes.sum < rightSize, + "test setup: the right side must be larger than the whole Union") + + withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> (rightSize * 2).toString) { + val df = sql( + """SELECT f.id, d.label + |FROM (SELECT * FROM fact1 UNION ALL SELECT * FROM fact2) f + |JOIN dim d ON f.id = d.did + """.stripMargin) + checkAnswer(df, (0 until 8).map(i => Row(i.toLong, i.toLong))) + + assertBuildsFromTheLeft(df) + assertJoinCount(df, 1) + assert(rightScansOf(df) == 1, + s"the right side must be scanned once, found ${rightScansOf(df)} scans of it") + } + } + } + + /** + * Asserts how many `Join` nodes the optimized plan holds. The rule turns one join into one per + * Union branch, so this pins whether it fired. Without it, a fixture whose branches stopped being + * larger than the dimension table would still produce the right rows and pass silently. + */ + private def assertJoinCount(df: DataFrame, expected: Int): Unit = { + val joins = df.queryExecution.optimizedPlan.collect { case j: Join => j } + assert(joins.size == expected, + s"expected $expected Join nodes in the optimized plan, found ${joins.size}") + } + + /** + * Asserts that every broadcast hash join in the plan builds from the left, the condition under + * which the rule used to fire and duplicate the probe side. Without it the scan count would also + * be satisfied by a plan that never became a broadcast hash join, which is a different reason for + * the rule not to fire. + * + * This is a premise of the two tests above rather than a property of the rule: the number of + * joins is left to their own assertion. Should a planner change stop picking a build-left + * broadcast hash join here, retune the table sizes and the threshold. + */ + private def assertBuildsFromTheLeft(df: DataFrame): Unit = { + val joins = collectWithSubqueries(df.queryExecution.executedPlan) { + case j: BroadcastHashJoinExec => j + } + assert(joins.nonEmpty && joins.forall(_.buildSide == BuildLeft), + s"expected every broadcast hash join to build from the left, found ${joins.map(_.buildSide)}") + } + + private def rightScansOf(df: DataFrame): Int = { + collectWithSubqueries(df.queryExecution.executedPlan) { + case s: FileSourceScanExec if s.tableIdentifier.exists(_.table == "dim") => s + }.size + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala index 02cb61787abc0..04d633698604d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala @@ -282,9 +282,6 @@ trait QueryTestBase protected def sparkContext = spark.sparkContext - // Shorthand for running a query using our SparkSession - protected lazy val sql: String => DataFrame = spark.sql _ - /** * A helper object for importing SQL implicits. * diff --git a/sql/core/src/test/scala/org/apache/spark/sql/ReplaceIntegerLiteralsWithOrdinalsDataframeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/ReplaceIntegerLiteralsWithOrdinalsDataframeSuite.scala index 5c3b6a6ec696e..92f54b5530536 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/ReplaceIntegerLiteralsWithOrdinalsDataframeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/ReplaceIntegerLiteralsWithOrdinalsDataframeSuite.scala @@ -24,7 +24,6 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class ReplaceIntegerLiteralsWithOrdinalsDataframeSuite extends SharedSparkSession { - import testImplicits._ test("Group by ordinal - Dataframe") { val query = "SELECT * FROM VALUES(1,2),(1,3),(2,4)" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala new file mode 100644 index 0000000000000..f99429cd1649b --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.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.spark.sql + +import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +class RewriteDistinctAggregatesConditionalQuerySuite extends QueryTest with SharedSparkSession { + + private def checkRewriteAndResult( + conditionalSql: String, + filterSql: String): Unit = { + val expectedRows = spark.sql(filterSql).collect() + + // Verify the rewrite produces the same result as the explicit FILTER form. + withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + checkAnswer(spark.sql(conditionalSql), expectedRows) + } + + // Verify the non-rewritten form also matches the explicit FILTER form. + withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + checkAnswer(spark.sql(conditionalSql), expectedRows) + } + } + + private def hasUserConditionFilter(plan: LogicalPlan): Boolean = { + val aggregateExpressions = plan.collect { + case a: Aggregate => a.aggregateExpressions + }.flatten + val aggregateExprs = aggregateExpressions.flatMap { + _.collect { case ae: AggregateExpression => ae } + } + // The Expand rewrite adds gid-routing filters to every distinct aggregate; we only + // care about user-condition filters introduced by the IF/CASE canonicalization. + aggregateExprs.flatMap(_.filter).exists(_.references.exists(_.name != "gid")) + } + + test("rewrite COUNT(DISTINCT IF(cond, col, NULL)) correctness") { + withTempView("t") { + spark.range(7) + .selectExpr( + "cast(id % 3 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "case when id % 4 = 0 then null else cast(id * 100 as int) end as col2") + .createOrReplaceTempView("t") + + // Two conditional distinct counts on the same base so the rewrite fires (size > 1). + checkRewriteAndResult( + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)), + | COUNT(DISTINCT IF(col1 > 20, col2, NULL)) + |FROM t GROUP BY key""".stripMargin, + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t GROUP BY key""".stripMargin) + } + } + + test("rewrite COUNT(DISTINCT CASE WHEN cond THEN col END) correctness") { + withTempView("t") { + spark.range(7) + .selectExpr( + "cast(id % 3 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "case when id % 4 = 0 then null else cast(id * 100 as string) end as col2") + .createOrReplaceTempView("t") + + // Two CASE WHEN conditional distinct counts on the same base so the rewrite fires. + checkRewriteAndResult( + """SELECT key, + | COUNT(DISTINCT CASE WHEN col1 > 10 THEN col2 END), + | COUNT(DISTINCT CASE WHEN col1 > 20 THEN col2 END) + |FROM t GROUP BY key""".stripMargin, + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t GROUP BY key""".stripMargin) + } + } + + test("rewrite COUNT(DISTINCT CASE WHEN cond THEN col ELSE NULL END) correctness") { + withTempView("t") { + spark.range(6) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "case when id % 4 = 0 then null else cast(id * 1.0 as double) end as col2") + .createOrReplaceTempView("t") + + // Two CASE WHEN ... ELSE NULL counts on the same base so the rewrite fires. + checkRewriteAndResult( + """SELECT key, + | COUNT(DISTINCT CASE WHEN col1 > 10 THEN col2 ELSE NULL END), + | COUNT(DISTINCT CASE WHEN col1 > 20 THEN col2 ELSE NULL END) + |FROM t GROUP BY key""".stripMargin, + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t GROUP BY key""".stripMargin) + } + } + + test("rewrite with no GROUP BY") { + withTempView("t") { + spark.range(5) + .selectExpr( + "cast(id * 10 as int) as col1", + "case when id % 3 = 0 then null else cast(id * 100 as int) end as col2") + .createOrReplaceTempView("t") + + // Two counts so the rewrite fires without GROUP BY. + checkRewriteAndResult( + """SELECT + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)), + | COUNT(DISTINCT IF(col1 > 20, col2, NULL)) + |FROM t""".stripMargin, + """SELECT + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t""".stripMargin) + } + } + + test("rewrite with all NULLs in conditional branch") { + withTempView("t") { + spark.range(3) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id * 5 as int) as col1", + "cast(id * 100 as int) as col2") + .createOrReplaceTempView("t") + + // col1 values are 0, 5, 10 - both thresholds (> 10 and > 20) yield zero matches, + // so both counts return 0. Two counts so the rewrite fires. + checkRewriteAndResult( + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)), + | COUNT(DISTINCT IF(col1 > 20, col2, NULL)) + |FROM t GROUP BY key""".stripMargin, + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t GROUP BY key""".stripMargin) + } + } + + test("rewrite with duplicates in base column") { + withTempView("t") { + spark.range(6) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "case when id % 3 = 0 then 100 when id % 3 = 1 then 100 else 200 end as col2") + .createOrReplaceTempView("t") + + // Two counts on the same base (with duplicates) so the rewrite fires. + checkRewriteAndResult( + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)), + | COUNT(DISTINCT IF(col1 > 20, col2, NULL)) + |FROM t GROUP BY key""".stripMargin, + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10), + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20) + |FROM t GROUP BY key""".stripMargin) + } + } + + test("multiple conditional distinct counts collapse and produce correct results") { + withTempView("t") { + spark.range(5) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "case when id % 3 = 0 then null else cast(id * 100 as int) end as col2", + "case when id % 4 = 0 then null else cast(id * 10 as string) end as col3") + .createOrReplaceTempView("t") + + val conditionalSql = + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)) as cnt1, + | COUNT(DISTINCT IF(col1 > 5, col3, NULL)) as cnt2 + |FROM t GROUP BY key""".stripMargin + + val filterSql = + """SELECT key, + | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10) as cnt1, + | COUNT(DISTINCT col3) FILTER (WHERE col1 > 5) as cnt2 + |FROM t GROUP BY key""".stripMargin + + checkRewriteAndResult(conditionalSql, filterSql) + } + } + + test("rewrite does not affect COUNT(DISTINCT IF(cond, col, non_null))") { + withTempView("t") { + spark.range(3) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id * 10 as int) as col1", + "cast(id * 100 as int) as col2") + .createOrReplaceTempView("t") + + // Two counts so mayNeedtoRewrite fires and the non-null-else guard is exercised. + val sqlText = + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, 0)), + | COUNT(DISTINCT IF(col1 > 20, col2, 0)) + |FROM t GROUP BY key""".stripMargin + + // Collect the conf-off result as the semantic baseline; the conf-on query must match it. + val withoutRewriteRows = withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + spark.sql(sqlText).collect() + } + withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + checkAnswer(spark.sql(sqlText), withoutRewriteRows) + } + + // When the switch is enabled, the non-null ELSE branch must not be canonicalized + // into a user-condition FILTER; only the internal gid-routing filters introduced + // by the Expand rewrite may be present. + val hasFilter = withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + hasUserConditionFilter(spark.sql(sqlText).queryExecution.optimizedPlan) + } + assert(!hasFilter, + "Non-null ELSE branch should not produce a user-condition FILTER clause") + } + } + + test("rewrite is not applied to error-prone base expressions (ANSI divide by zero)") { + withTempView("t") { + // Rows with x = 0 make both conditions false, so the IF branches short-circuit + // and 10 / x is never evaluated for them. + spark.range(5) + .selectExpr( + "cast(id % 2 + 1 as int) as key", + "cast(id % 2 as int) as x") + .createOrReplaceTempView("t") + + // Two conditional distinct counts on the same error-prone base so that rewrite() + // is reached; the canonicalization must NOT fire for such a base. + val sqlText = + """SELECT key, + | COUNT(DISTINCT IF(x <> 0, 10 / x, NULL)), + | COUNT(DISTINCT IF(x > 0, 10 / x, NULL)) + |FROM t GROUP BY key""".stripMargin + + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + // Must not raise DIVIDE_BY_ZERO: the canonicalization must leave the IFs intact + // so that their branch short-circuiting is preserved. + checkAnswer(spark.sql(sqlText), Seq(Row(1, 0L, 0L), Row(2, 1L, 1L))) + assert(!hasUserConditionFilter(spark.sql(sqlText).queryExecution.optimizedPlan), + "error-prone base should not produce a user-condition FILTER clause") + } + } + } + + test("rewrite is present in optimized plan") { + withTempView("t") { + spark.range(2) + .selectExpr( + "cast(id + 1 as int) as key", + "cast(id * 10 as int) as col1", + "cast(id * 100 as int) as col2") + .createOrReplaceTempView("t") + + // Two conditional distinct counts on the same base column trigger canonicalization, + // so the optimized plan must contain user-condition FILTER clauses. + val hasFilter = withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + val df = spark.sql( + """SELECT key, + | COUNT(DISTINCT IF(col1 > 10, col2, NULL)) as cnt1, + | COUNT(DISTINCT IF(col1 > 5, col2, NULL)) as cnt2 + |FROM t GROUP BY key""".stripMargin) + hasUserConditionFilter(df.queryExecution.optimizedPlan) + } + + assert(hasFilter, "Optimized plan should contain user-condition FILTER clause") + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLInsertTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLInsertTestSuite.scala index 7452917f57d73..7b1fb2e7dcc3e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLInsertTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLInsertTestSuite.scala @@ -591,6 +591,73 @@ trait SQLInsertTestSuite extends QueryTest with AdaptiveSparkPlanHelper { } } } + + test("SPARK-58816: insert with column list resolves structs inside arrays positionally") { + withTable("t") { + createTable("t", Seq("arr"), Seq("ARRAY<STRUCT<x: INT, y: INT>>")) + // Source has fields in (y, x) order; target expects (x, y). + // Positional resolution must rename y->x and x->y so the cast maps by position. + sql("INSERT INTO t (arr) SELECT array(named_struct('y', 20, 'x', 10))") + checkAnswer(spark.table("t"), Row(Seq(Row(20, 10)))) + + // Contrast: INSERT BY NAME must continue to resolve by name, giving {x:10, y:20}. + sql("INSERT INTO t BY NAME SELECT array(named_struct('y', 20, 'x', 10)) AS arr") + checkAnswer(spark.table("t"), Seq(Row(Seq(Row(20, 10))), Row(Seq(Row(10, 20))))) + } + } + + test("SPARK-58816: insert with column list resolves structs inside maps positionally") { + withTable("t") { + // Map value: MAP<STRING, STRUCT<x: INT, y: INT>> + createTable("t", Seq("m"), Seq("MAP<STRING, STRUCT<x: INT, y: INT>>")) + sql("INSERT INTO t (m) SELECT map('k', named_struct('y', 20, 'x', 10))") + checkAnswer(spark.table("t"), Row(Map("k" -> Row(20, 10)))) + } + // Map key: MAP<STRUCT<x: INT, y: INT>, STRING> -- exercises the key recursion branch. + withTable("t") { + createTable("t", Seq("m"), Seq("MAP<STRUCT<x: INT, y: INT>, STRING>")) + sql("INSERT INTO t (m) SELECT map(named_struct('y', 20, 'x', 10), 'v')") + checkAnswer(spark.table("t"), Row(Map(Row(20, 10) -> "v"))) + } + } + + test("SPARK-58816: insert with column list resolves structs positionally at all nesting levels") { + // All three nesting modes must behave identically: direct struct, array-of-struct, + // map-of-struct. Source fields are (y, x); target fields are (x, y). + // Positional resolution writes y->x slot and x->y slot for each case. + withTable("t") { + createTable( + "t", + Seq("s", "arr", "m"), + Seq( + "STRUCT<x: INT, y: INT>", + "ARRAY<STRUCT<x: INT, y: INT>>", + "MAP<STRING, STRUCT<x: INT, y: INT>>")) + sql( + """INSERT INTO t (s, arr, m) + |SELECT + | named_struct('y', 20, 'x', 10), + | array(named_struct('y', 20, 'x', 10)), + | map('k', named_struct('y', 20, 'x', 10)) + |""".stripMargin) + // After positional resolution: first field slot gets 20, second gets 10 in all three cases. + checkAnswer( + spark.table("t"), + Row(Row(20, 10), Seq(Row(20, 10)), Map("k" -> Row(20, 10)))) + } + // Deep mixed nesting: ARRAY<STRUCT<nested: ARRAY<STRUCT<x: INT, y: INT>>>>. + // Two recursive transitions (Array->Struct->Array->Struct) must all be renamed positionally. + withTable("t") { + createTable( + "t", Seq("col"), + Seq("ARRAY<STRUCT<nested: ARRAY<STRUCT<x: INT, y: INT>>>>")) + sql( + """INSERT INTO t (col) + |SELECT array(named_struct( + | 'nested', array(named_struct('y', 20, 'x', 10))))""".stripMargin) + checkAnswer(spark.table("t"), Row(Seq(Row(Seq(Row(20, 10)))))) + } + } } @SQLUserDefinedType(udt = classOf[MyIntUDT]) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala index 03d3a122da16d..8075a8a0776bc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala @@ -5158,6 +5158,163 @@ class SQLQuerySuite extends SharedSparkSession with AdaptiveSparkPlanHelper sql("SELECT col1 - rand() FROM VALUES(1) GROUP BY ALL") } } + + test("SPARK-57353: CUBE with ORDER BY - single pass resolver (tentative fallback)") { + // ORDER BY + grouping analytics throws ExplicitlyUnsupportedResolverFeature in pure + // single-pass (SPARK-57346). In tentative mode, the HybridAnalyzer falls back to legacy. + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) as s FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY s""".stripMargin), + Row(1, 30) :: Row(2, 30) :: Row(null, 60) :: Nil) + } + } + + test("SPARK-57353: ROLLUP with HAVING - single pass resolver (tentative fallback)") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a) HAVING SUM(b) > 30""".stripMargin), + Row(null, 60) :: Nil) + } + } + + test("SPARK-57353: GROUPING SETS with ORDER BY - single pass resolver (tentative fallback)") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, b, SUM(b) as s FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY a, b GROUPING SETS ((a, b), (a)) ORDER BY s""".stripMargin), + Row(1, 10, 10) :: Row(1, 20, 20) :: Row(1, null, 30) :: + Row(2, 30, 30) :: Row(2, null, 30) :: Nil) + } + } + + test("SPARK-57353: CUBE with NULL grouping columns - single pass resolver (tentative fallback)") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) as s FROM VALUES (1,10),(null,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY s""".stripMargin), + Row(1, 10) :: Row(null, 20) :: Row(2, 30) :: Row(null, 60) :: Nil) + } + } + + test("SPARK-57353: ROLLUP with HAVING filtering all rows - single pass resolver" + + " (tentative fallback)") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a) HAVING SUM(b) > 100""".stripMargin), + Nil) + } + } + + test("SPARK-57353: CUBE with multiple aggregates in ORDER BY - single pass resolver" + + " (tentative fallback)") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) as s, COUNT(b) as c + |FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY COUNT(b), SUM(b)""".stripMargin), + Row(2, 30, 1) :: Row(1, 30, 2) :: Row(null, 60, 3) :: Nil) + } + } + + test("SPARK-57353: multi-column ROLLUP with HAVING - single pass resolver (tentative fallback" + + ", SPARK-57346)") { + // SPARK-57346: multi-column ROLLUP with HAVING previously produced wrong results (1 row + // instead of 4) in pure single-pass. Now throws ExplicitlyUnsupportedResolverFeature so + // tentative mode falls back to legacy for correct results. + val query = + """SELECT a, b, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a, b) HAVING SUM(b) > 25""".stripMargin + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + val legacyResult = withSQLConf( + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "false") { + sql(query).collect().toSeq + } + checkAnswer(sql(query), legacyResult) + } + } + + test("SPARK-57353: missing-aggregation still detected after GROUPING SETS expansion") { + // Confirms aggregate-expression validation still fires: column 'b' is not in GROUP BY + // ROLLUP(a) and is not aggregated, so MISSING_AGGREGATION should be raised. + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "true") { + val ex = intercept[AnalysisException] { + sql( + """SELECT a, b, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a) HAVING SUM(b) > 10""".stripMargin).collect() + } + assert(ex.getCondition === "MISSING_AGGREGATION") + } + } + + test("SPARK-57353: LCA with GROUPING SETS - single pass resolver") { + // LCA + grouping analytics throws ExplicitlyUnsupportedResolverFeature in single-pass mode. + // In tentative mode, the HybridAnalyzer catches it and falls back to legacy for correct + // results. + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED_TENTATIVELY.key -> "true") { + checkAnswer( + sql( + """SELECT a, SUM(b) as total, total + 1 + |FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY total""".stripMargin), + Row(1, 30, 31) :: Row(2, 30, 31) :: Row(null, 60, 61) :: Nil) + } + } + + test("SPARK-57353: legacy analyzer handles CUBE/ROLLUP/GROUPING SETS correctly") { + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false") { + checkAnswer( + sql( + """SELECT a, SUM(b) as s FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY s""".stripMargin), + Row(1, 30) :: Row(2, 30) :: Row(null, 60) :: Nil) + checkAnswer( + sql( + """SELECT a, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a) HAVING SUM(b) > 30""".stripMargin), + Row(null, 60) :: Nil) + } + } + + test("SPARK-57353: nested aggregate with GROUPING SETS is rejected") { + // Confirms aggregate-expression validation (nested agg check) still fires with grouping sets. + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "true") { + val ex = intercept[AnalysisException] { + sql( + """SELECT a, SUM(COUNT(b)) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a)""".stripMargin).collect() + } + assert(ex.getCondition === "NESTED_AGGREGATE_FUNCTION") + } + } + + test("SPARK-57353: GROUPING SETS with empty grouping list - all rows") { + // GROUPING SETS (()) produces a single grand-total row. + withSQLConf(SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "true") { + QueryTest.checkAnswer( + sql( + """SELECT SUM(b) as s FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY a GROUPING SETS (())""".stripMargin), + Row(60) :: Nil, + checkToRDD = false) + } + } } case class Foo(bar: Option[String]) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTestBeforeAfterHooksSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTestBeforeAfterHooksSuite.scala new file mode 100644 index 0000000000000..f3253def7113a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTestBeforeAfterHooksSuite.scala @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql + +class SessionQueryTestBeforeAfterHooksSuite extends SessionQueryTest { + + override protected def beforeAll(): Unit = { + super.beforeAll() + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + } + + override protected def beforeEach(): Unit = { + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + super.beforeEach() + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + } + + override protected def afterEach(): Unit = { + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + super.afterEach() + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + } + + override protected def afterAll(): Unit = { + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + super.afterAll() + } + + test("assert spark is available in BeforeAndAfter hooks") { + checkAnswer(spark.sql("SELECT 1"), Seq(Row(1))) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala index d7b2511eac2a3..85de26b7a567e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionJobTaggingAndCancellationSuite.scala @@ -81,6 +81,27 @@ class SparkSessionJobTaggingAndCancellationSuite assert(session.getTags() == Set("one")) } + test("SPARK-58358: addTag and removeTag reject invalid tags") { + val session = classic.SparkSession.builder().master("local").getOrCreate() + + // A tag cannot be empty or contain the ',' separator. Both addTag and removeTag must + // reject such tags (removeTag previously skipped this validation). + Seq("", "a,b", ",").foreach { invalidTag => + intercept[IllegalArgumentException] { + session.addTag(invalidTag) + } + intercept[IllegalArgumentException] { + session.removeTag(invalidTag) + } + } + // A rejected tag must not leak into the tag set. + assert(session.getTags() == Set()) + + // Removing a valid, absent tag stays a no-op. + session.removeTag("absent") + assert(session.getTags() == Set()) + } + test("Tags set from session are prefixed with session UUID") { sc = new SparkContext("local[2]", "test") val session = classic.SparkSession.builder().sparkContext(sc).getOrCreate() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionProvider.scala b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionProvider.scala index 67ff122efec87..1f0aaa1f39bf4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionProvider.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionProvider.scala @@ -26,4 +26,9 @@ package org.apache.spark.sql */ trait SparkSessionProvider { protected def spark: SparkSession + + /** + * Shorthand for running a query using the [[SparkSession]] + */ + protected def sql(query: String): DataFrame = spark.sql(query) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala index e5fa22b05ecb2..00b6c926b92bf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/StringFunctionsSuite.scala @@ -308,6 +308,18 @@ class StringFunctionsSuite extends SharedSparkSession { Row("AQIDBA==", bytes)) } + test("string to_base32/from_base32 function") { + val bytes = "foobar".getBytes("UTF-8") + val df = Seq((bytes, "MZXW6YTBOI======")).toDF("a", "b") + checkAnswer( + df.select(to_base32($"a"), from_base32($"b")), + Row("MZXW6YTBOI======", bytes)) + + checkAnswer( + df.selectExpr("to_base32(a)", "from_base32(b)"), + Row("MZXW6YTBOI======", bytes)) + } + test("string overlay function") { // scalastyle:off // non ascii characters are not allowed in the code, so we disable the scalastyle here. @@ -407,6 +419,16 @@ class StringFunctionsSuite extends SharedSparkSession { // scalastyle:on } + test("string normalize") { + // scalastyle:off + val df = Seq("\uFB01").toDF("s") + checkAnswer(df.select(normalize($"s")), Row("\uFB01")) + checkAnswer(df.select(normalize($"s", lit("NFKC"))), Row("fi")) + checkAnswer(df.selectExpr("normalize(s, 'NFKC')"), Row("fi")) + checkAnswer(df.select(normalize(lit(null), lit("NFC"))), Row(null)) + // scalastyle:on + } + test("string translate") { val df = Seq(("translate", "")).toDF("a", "b") checkAnswer(df.select(translate($"a", "rnlt", "123")), Row("1a2s3ae")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala index d8d885c9b927e..d5cc3e6c0c00d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SubquerySuite.scala @@ -20,7 +20,8 @@ package org.apache.spark.sql import scala.collection.mutable.ArrayBuffer import org.apache.spark.SparkRuntimeException -import org.apache.spark.sql.catalyst.expressions.{EqualTo, NamedExpression, OuterReference, SubqueryExpression} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{EqualTo, Literal, NamedExpression, OuterReference, SubqueryExpression} import org.apache.spark.sql.catalyst.plans.{LeftAnti, LeftSemi} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LogicalPlan, Project, Sort, Union} import org.apache.spark.sql.execution._ @@ -264,6 +265,17 @@ class SubquerySuite extends SharedSparkSession Row(3, 3.0) :: Nil) } + test("IN predicate subquery preserves its broadcast when replacing its plan") { + val subqueryPlan = SubqueryExec("subquery", spark.range(3).queryExecution.executedPlan) + val subquery = InSubqueryExec( + Literal(1L), subqueryPlan, NamedExpression.newExprId, isDynamicPruning = false) + subquery.updateResult() + + val updatedSubquery = subquery.withNewPlan(subqueryPlan) + assert(updatedSubquery.values().isEmpty) + assert(updatedSubquery.eval(InternalRow.empty) == true) + } + test("NOT IN predicate subquery") { checkAnswer( sql("select * from l where a not in (select c from r)"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosDistinctSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosDistinctSuiteBase.scala new file mode 100644 index 0000000000000..5574344520ba5 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosDistinctSuiteBase.scala @@ -0,0 +1,183 @@ +/* + * 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 + +import java.time.{Instant, LocalDateTime} + +import org.apache.spark.SparkConf +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ + +/** + * End-to-end `DISTINCT` correctness tests over the nanosecond-precision timestamp types + * `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in `[7, 9]`), part of the nanosecond timestamp + * preview (SPARK-56822). Set-distinctness over these types executes today with no production + * change -- it rides on the nanos hashing and equality the types already implement (SPARK-57103 + * extended `Murmur3Hash` / `XxHash64` / `HiveHash` to hash the carrier's two fields, + * `epochMicros: Long` and `nanosWithinMicro: Short in [0, 999]`). There is no coverage organized + * around distinctness by column type, so this per-type suite adds it and locks the regression. + * + * Headline risk: SUB-MICROSECOND distinctness, mirroring the join suite. The carrier is + * `TimestampNanosVal = (epochMicros: Long, nanosWithinMicro: Short in [0, 999])`. Every value in + * these tables shares the SAME epochMicros (1577836800000000, = 2020-01-01T00:00:00Z) and differs + * ONLY in nanosWithinMicro. So the micro-level path alone cannot tell them apart -- correct + * distinctness MUST be driven by the full nanos value: + * - two values equal in epochMicros but DIFFERENT in nanosWithinMicro must NOT be deduplicated, + * - exact duplicates (equal incl. the sub-microsecond remainder) MUST be deduplicated, + * - NULL, unlike an equi-join, IS a single distinct value -- `DISTINCT` keeps exactly one NULL. + * If the sub-microsecond remainder were ignored, the two distinct sub-microsecond values here + * would fold into one, dropping the distinct-row count from 3 to 2 and failing these tests loudly. + * + * Precision-safety: all sub-microsecond remainders are multiples of 100ns (100 / 900), which are + * exact at every p in [7, 9] (`createDataFrame` floors nanosWithinMicro to (n/100)*100 at p=7 and + * (n/10)*10 at p=8). So the SAME inputs and the SAME expected results are valid verbatim at all + * three precisions, and the two distinct remainders never collide even at the coarsest p=7. + * + * Each test runs under whole-stage codegen on and off, so the same sub-microsecond distinctness is + * proven on the nanos hash path in both codegen modes, for NTZ and LTZ. The mixed-precision test + * additionally pins that a `UNION` of two different nanos precisions widens the column to the + * higher precision (`findWiderDateTimeType`) and preserves the distinction. + * + * The nanosecond timestamp types are gated behind a preview flag enabled by default under tests + * (`Utils.isTesting`), so it is not set here. The session time zone is fixed so the + * `TIMESTAMP_LTZ` (`Instant`) values are deterministic. The two subclasses run every test with + * ANSI mode on and off. + */ +abstract class TimestampNanosDistinctSuiteBase extends QueryTest with SharedSparkSession { + + override def sparkConf: SparkConf = super.sparkConf + .set(SQLConf.SESSION_LOCAL_TIMEZONE.key, "America/Los_Angeles") + + // Whole-stage codegen on (CODEGEN_ONLY) vs off (NO_CODEGEN). The hash-aggregate exec backing + // DISTINCT is the same in both modes; only the WholeStageCodegenExec wrapper differs. Mirrors the + // join and functions suites. + protected val codegenModes: Seq[Seq[(String, String)]] = Seq( + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY"), + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")) + + private def cgLabel(cgConf: Seq[(String, String)]): String = + if (cgConf.contains(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true")) "codegen on" + else "codegen off" + + // ---- relation builders: single nanos column "k" ---- + // + // Four rows, all sharing epochMicros (2020-01-01T00:00:00Z) and differing only within the + // microsecond: + // - two rows with the 100ns value (an exact duplicate): collapse to one distinct row, + // - one row with the 900ns value: same microsecond, distinct sub-microsecond remainder, + // - one NULL row: a single distinct NULL. + // Expected three distinct values: 100ns, 900ns, NULL. + + private def ntzDF(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100")), + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100")), + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000900")), + Row(null))), + new StructType().add("k", TimestampNTZNanosType(p))) + + private def ltzDF(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(Instant.parse("2020-01-01T00:00:00.000000100Z")), + Row(Instant.parse("2020-01-01T00:00:00.000000100Z")), + Row(Instant.parse("2020-01-01T00:00:00.000000900Z")), + Row(null))), + new StructType().add("k", TimestampLTZNanosType(p))) + + // The external nanosecond remainder of a collected value, whichever family it came back as + // (NTZ -> LocalDateTime, LTZ -> Instant). Every value here shares epochMicros, so the remainder + // alone identifies it. + private def nanoOf(v: Any): Int = v match { + case ldt: LocalDateTime => ldt.getNano + case i: Instant => i.getNano + } + + // Each family as (label, builder(p), widerType(p)). DISTINCT is type-agnostic, so the family only + // picks the data builder and (for the mixed-precision widening assertion) the expected widened + // column type. + protected val families: Seq[(String, Int => DataFrame, Int => DataType)] = Seq( + ("NTZ", ntzDF, (p: Int) => TimestampNTZNanosType(p)), + ("LTZ", ltzDF, (p: Int) => TimestampLTZNanosType(p))) + + // Mixed-precision UNION pairs; the column widens to max(pl, pr). + private val mixedPrecisionPairs: Seq[(Int, Int)] = Seq((7, 9), (7, 8), (8, 9)) + + // Asserts that `distinct` holds exactly the three expected values: the 100ns and 900ns remainders + // and a single NULL. Checking the actual remainders (not just the row count) pins that the + // surviving rows carry the full sub-microsecond value. + private def assertDistinctKeys(distinct: DataFrame, p: Int): Unit = { + assert(distinct.count() == 3, s"expected 3 distinct values (100ns, 900ns, NULL) at p=$p") + val rows = distinct.collect() + assert(rows.count(_.isNullAt(0)) == 1, s"expected exactly one NULL at p=$p") + val remainders = rows.filterNot(_.isNullAt(0)).map(r => nanoOf(r.get(0))).toSet + assert(remainders == Set(100, 900), + s"expected distinct remainders {100, 900} at p=$p, got $remainders") + } + + for { + (family, builder, widerType) <- families + cgConf <- codegenModes + } { + // ======================================================================================== + // DISTINCT over a nanos column: exact duplicates are removed, sub-microsecond-distinct values + // are kept, and NULL survives as exactly one row. Expected three distinct rows: 100ns, 900ns, + // NULL. If the sub-microsecond remainder were ignored, 100ns and 900ns would collapse to one, + // dropping the count to 2 and failing this assertion. + // ======================================================================================== + test(s"$family nanos DISTINCT keeps sub-microsecond-distinct values - ${cgLabel(cgConf)}") { + withSQLConf(cgConf: _*) { + Seq(7, 8, 9).foreach { p => + assertDistinctKeys(builder(p).select("k").distinct(), p) + } + } + } + + // ======================================================================================== + // DISTINCT over a mixed-precision column: a UNION of the same data at two different nanos + // precisions widens the column to max(pl, pr) (findWiderDateTimeType). DISTINCT over the + // widened column still separates the two sub-microsecond values and still collapses to three + // rows (the union only adds more exact duplicates; all remainders are multiples of 100ns, + // exact at every p in [7, 9]). + // ======================================================================================== + test(s"$family nanos DISTINCT across mixed precisions widens column - ${cgLabel(cgConf)}") { + withSQLConf(cgConf: _*) { + mixedPrecisionPairs.foreach { case (pl, pr) => + val union = builder(pl).union(builder(pr)) + assert(union.schema("k").dataType == widerType(math.max(pl, pr)), + s"expected column widened to ${widerType(math.max(pl, pr))} for ($pl, $pr)") + assertDistinctKeys(union.distinct(), math.max(pl, pr)) + } + } + } + } +} + +// Runs the nanosecond timestamp DISTINCT tests with ANSI mode enabled explicitly. +class TimestampNanosDistinctAnsiOnSuite extends TimestampNanosDistinctSuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "true") +} + +// Runs the nanosecond timestamp DISTINCT tests with ANSI mode disabled explicitly. +class TimestampNanosDistinctAnsiOffSuite extends TimestampNanosDistinctSuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala index adfeb902cc03e..4a73b2c448dfb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosFunctionsSuiteBase.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql import java.time.{Instant, LocalDateTime} -import org.apache.spark.{SparkConf, SparkRuntimeException, SparkUnsupportedOperationException} +import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException, SparkUnsupportedOperationException} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -710,6 +710,117 @@ abstract class TimestampNanosFunctionsSuiteBase extends SharedSparkSession { } } + // mode over nanosecond-precision timestamps (SPARK-56822). `Mode` counts frequencies in an + // `OpenHashMap` keyed on the physical `TimestampNanosVal` (its `equals`/`hashCode` cover the full + // `(epochMicros, nanosWithinMicro)` pair) and returns `child.dataType`, so the most-frequent + // value is selected on the full nanos value and its precision and family (NTZ/LTZ) are preserved. + + test("SPARK-56822: mode over nanosecond-precision timestamps returns the most frequent value") { + Seq(7, 8, 9).foreach { p => + val schema = new StructType() + .add("ntz", TimestampNTZNanosType(p)) + .add("ltz", TimestampLTZNanosType(p)) + // The frequent value (3 rows) and the rare one (1 row) differ only within the microsecond, so + // frequency counting must key on the full nanos value; a NULL row is ignored. The fractions + // are multiples of 100ns, exact at every p in [7, 9]. There is a unique most-frequent value, + // so the result is deterministic without a WITHIN GROUP / deterministic argument. + val ldtHot = LocalDateTime.parse("2020-01-01T00:00:00.000000100") + val ldtCold = LocalDateTime.parse("2020-01-01T00:00:00.000000900") + val insHot = Instant.parse("2020-01-01T00:00:00.000000100Z") + val insCold = Instant.parse("2020-01-01T00:00:00.000000900Z") + val data = Seq( + Row(ldtHot, insHot), Row(ldtHot, insHot), Row(ldtHot, insHot), + Row(ldtCold, insCold), Row(null, null)) + val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema) + + val res = df.selectExpr("mode(ntz)", "mode(ltz)") + // The result keeps the family (NTZ/LTZ) and precision of the input. + assert(res.schema.map(_.dataType) === Seq(TimestampNTZNanosType(p), TimestampLTZNanosType(p))) + checkAnswer(res, Row(ldtHot, insHot)) + } + } + + // collect_set over nanosecond-precision timestamps (SPARK-56822). `CollectSet` deduplicates via a + // `HashSet` keyed on the physical `TimestampNanosVal`, whose `equals`/`hashCode` cover the full + // `(epochMicros, nanosWithinMicro)` pair, so sub-microsecond-distinct values are kept distinct; + // the result element type is exactly `child.dataType`. + + test("SPARK-56822: collect_set over nanos deduplicates on the full sub-microsecond value") { + // Two values share the microsecond and differ only in the last nanosecond digit. Flooring the + // input to precision `p` collapses them when p < 9 (the distinguishing digit is below the + // grid), but keeps them apart at p = 9. collect_set must dedup on the full stored value, + // not on micros. + Seq(7, 8, 9).foreach { p => + val schema = new StructType().add("ntz", TimestampNTZNanosType(p)) + val data = Seq( + Row(LocalDateTime.parse("2020-01-01T12:34:56.123456780")), + Row(LocalDateTime.parse("2020-01-01T12:34:56.123456789")), + Row(LocalDateTime.parse("2020-01-01T12:34:56.123456780"))) + val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema) + + val res = df.selectExpr("collect_set(ntz)") + assert(res.schema.head.dataType === ArrayType(TimestampNTZNanosType(p), containsNull = false)) + + // Values are floored to `p` on ingestion, so the distinct set depends on `p`. + val expected = p match { + case 7 => Set(LocalDateTime.parse("2020-01-01T12:34:56.123456700")) + case 8 => Set(LocalDateTime.parse("2020-01-01T12:34:56.123456780")) + case _ => Set( + LocalDateTime.parse("2020-01-01T12:34:56.123456780"), + LocalDateTime.parse("2020-01-01T12:34:56.123456789")) + } + val collected = res.collect().head.getSeq[LocalDateTime](0).toSet + assert(collected === expected, s"collect_set(p=$p) expected $expected, got $collected") + } + } + + // collect_list over nanosecond-precision timestamps (SPARK-56822). `CollectList` is + // type-agnostic: its `ArrayBuffer` buffer holds the physical `TimestampNanosVal` and the + // result element type is exactly `child.dataType`, so the input precision, family (NTZ/LTZ) + // and sub-microsecond remainder survive with no truncation to micros. The collection order + // after aggregation is not deterministic, so the contents are compared as a set. + + test("SPARK-56822: collect_list over nanosecond-precision timestamps preserves type and nanos " + + "remainder") { + Seq(7, 8, 9).foreach { p => + val schema = new StructType() + .add("ntz", TimestampNTZNanosType(p)) + .add("ltz", TimestampLTZNanosType(p)) + // The sub-microsecond parts are multiples of 100ns, so they are exact at every p in [7, 9] + // (no flooring) yet non-zero -- a value truncated to micros would be visibly wrong. The NULL + // row is dropped (collect_list ignores nulls by default). + val data = Seq( + Row(LocalDateTime.parse("2020-01-01T12:34:56.000000100"), + Instant.parse("2020-01-01T12:34:56.000000100Z")), + Row(LocalDateTime.parse("2020-01-02T00:00:00.000000900"), + Instant.parse("2020-01-02T00:00:00.000000900Z")), + Row(null, null)) + val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema) + + val sqlRes = df.selectExpr("collect_list(ntz)", "collect_list(ltz)") + // The result keeps the family (NTZ/LTZ) and precision; nulls are dropped so + // containsNull=false. + assert(sqlRes.schema.map(_.dataType) === Seq( + ArrayType(TimestampNTZNanosType(p), containsNull = false), + ArrayType(TimestampLTZNanosType(p), containsNull = false))) + + val expectedNtz = Set( + LocalDateTime.parse("2020-01-01T12:34:56.000000100"), + LocalDateTime.parse("2020-01-02T00:00:00.000000900")) + val expectedLtz = Set( + Instant.parse("2020-01-01T12:34:56.000000100Z"), + Instant.parse("2020-01-02T00:00:00.000000900Z")) + // The Scala Column API path agrees with the SQL path; both drop the null and keep + // both values. + val sqlRow = sqlRes.collect().head + val colRow = df.select(collect_list(col("ntz")), collect_list(col("ltz"))).collect().head + assert(sqlRow.getSeq[LocalDateTime](0).toSet === expectedNtz) + assert(sqlRow.getSeq[Instant](1).toSet === expectedLtz) + assert(colRow.getSeq[LocalDateTime](0).toSet === expectedNtz) + assert(colRow.getSeq[Instant](1).toSet === expectedLtz) + } + } + test("SPARK-57816: date_format / to_char / to_varchar over nanosecond-precision timestamps") { // The 9-`S` pattern is a fixed-width fraction field, so it always emits 9 digits; truncating to // precision `p` zeros the low digits (floor); it does not drop them. The session zone is @@ -824,6 +935,88 @@ abstract class TimestampNanosFunctionsSuiteBase extends SharedSparkSession { } } } + + test("SPARK-57837: current_timestamp(p) / now(p) return TIMESTAMP_LTZ(p)") { + val df = spark.range(1) + Seq("current_timestamp", "now").foreach { fn => + Seq(7, 8, 9).foreach { p => + val schema = df.selectExpr(s"$fn($p) AS c").schema + assert(schema.head.dataType === TimestampLTZNanosType(p), + s"$fn($p) should be TIMESTAMP_LTZ($p)") + } + // Precision 6 and the no-arg form keep the micro TIMESTAMP type. + assert(df.selectExpr(s"$fn(6) AS c").schema.head.dataType === TimestampType) + assert(df.selectExpr(s"$fn() AS c").schema.head.dataType === TimestampType) + } + } + + test("SPARK-57837: localtimestamp(p) returns TIMESTAMP_NTZ(p)") { + val df = spark.range(1) + Seq(7, 8, 9).foreach { p => + val schema = df.selectExpr(s"localtimestamp($p) AS c").schema + assert(schema.head.dataType === TimestampNTZNanosType(p), + s"localtimestamp($p) should be TIMESTAMP_NTZ($p)") + } + assert(df.selectExpr("localtimestamp(6) AS c").schema.head.dataType === TimestampNTZType) + assert(df.selectExpr("localtimestamp() AS c").schema.head.dataType === TimestampNTZType) + } + + test("SPARK-57837: nanos current-timestamp values are query-stable and floored to precision") { + val df = spark.range(1) + // All references within a query see the same value. + checkAnswer( + df.selectExpr( + "current_timestamp(9) = current_timestamp(9)", + "now(9) = current_timestamp(9)", + "localtimestamp(9) = localtimestamp(9)"), + Row(true, true, true)) + + // Sub-precision digits are floored: the extracted DECIMAL(11, 9) second has zeros below `p`. + Seq(7, 8, 9).foreach { p => + val secExpr = s"extract(SECOND FROM current_timestamp($p))" + val secondDecimal = df.selectExpr(secExpr).head().getDecimal(0) + val scaled = secondDecimal.movePointRight(9).longValueExact() + val step = math.pow(10, 9 - p).toLong + assert(scaled % step == 0, + s"extract(SECOND) of current_timestamp($p) = $secondDecimal not floored to precision $p") + } + } + + test("SPARK-57837: current_timestamp(p) / localtimestamp(p) reject out-of-range precision") { + val df = spark.range(1) + Seq( + "current_timestamp(3)" -> "TIMESTAMP_LTZ", + "now(3)" -> "TIMESTAMP_LTZ", + "localtimestamp(10)" -> "TIMESTAMP_NTZ").foreach { case (expr, typeName) => + val precision = expr.replaceAll("\\D", "") + checkError( + exception = intercept[SparkException] { + df.selectExpr(expr).collect() + }, + condition = "INVALID_TIMESTAMP_PRECISION", + parameters = Map("precision" -> precision, "type" -> typeName)) + } + } + + test("SPARK-57837: current_timestamp(p) / localtimestamp(p) require the nanos preview flag") { + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "false") { + val df = spark.range(1) + Seq("current_timestamp(9)", "now(8)", "localtimestamp(7)").foreach { expr => + checkError( + exception = intercept[SparkException] { + df.selectExpr(expr).collect() + }, + condition = "FEATURE_NOT_ENABLED", + parameters = Map( + "featureName" -> "Nanosecond-precision timestamp types", + "configKey" -> "spark.sql.timestampNanosTypes.enabled", + "configValue" -> "true")) + } + // With the flag off, precision 6 and the no-arg forms still work (micro types). + assert(df.selectExpr("current_timestamp(6) AS c").schema.head.dataType === TimestampType) + assert(df.selectExpr("localtimestamp() AS c").schema.head.dataType === TimestampNTZType) + } + } } // Runs the nanosecond timestamp function tests with ANSI mode enabled explicitly. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosGroupBySuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosGroupBySuiteBase.scala new file mode 100644 index 0000000000000..731db201836b6 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosGroupBySuiteBase.scala @@ -0,0 +1,190 @@ +/* + * 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 + +import java.time.{Instant, LocalDateTime} + +import org.apache.spark.SparkConf +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ + +/** + * End-to-end `GROUP BY` correctness tests over the nanosecond-precision timestamp types + * `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in `[7, 9]`), part of the nanosecond timestamp + * preview (SPARK-56822). Grouping over these types executes today with no production change -- it + * rides on the nanos hashing and equality the types already implement (SPARK-57103 extended + * `Murmur3Hash` / `XxHash64` / `HiveHash` to hash the carrier's two fields, `epochMicros: Long` + * and `nanosWithinMicro: Short in [0, 999]`). There is no coverage organized around grouping by + * column type, so this per-type suite adds it and locks the regression. + * + * Headline risk: SUB-MICROSECOND key correctness, mirroring the join suite. The carrier is + * `TimestampNanosVal = (epochMicros: Long, nanosWithinMicro: Short in [0, 999])`. Every key in + * these tables shares the SAME epochMicros (1577836800000000, = 2020-01-01T00:00:00Z) and differs + * ONLY in nanosWithinMicro. So the micro-level path alone cannot tell the keys apart -- correct + * grouping MUST be driven by the full nanos value: + * - two rows equal in epochMicros but DIFFERENT in nanosWithinMicro must NOT collapse into one + * group, + * - exact duplicates (equal incl. the sub-microsecond remainder) MUST collapse, + * - NULL keys, unlike an equi-join, DO group together -- all NULL rows land in a single group. + * If the sub-microsecond remainder were ignored, the two distinct sub-microsecond keys here would + * fold into one, dropping the group count from 3 to 2 and failing these tests loudly. + * + * Precision-safety: all sub-microsecond remainders are multiples of 100ns (100 / 900), which are + * exact at every p in [7, 9] (`createDataFrame` floors nanosWithinMicro to (n/100)*100 at p=7 and + * (n/10)*10 at p=8). So the SAME inputs and the SAME expected results are valid verbatim at all + * three precisions, and the two distinct remainders never collide even at the coarsest p=7. + * + * Each test runs under whole-stage codegen on and off, so the same sub-microsecond grouping is + * proven on the nanos hash path in both codegen modes, for NTZ and LTZ. The mixed-precision test + * additionally pins that a `UNION` of two different nanos precisions widens the key to the higher + * precision (`findWiderDateTimeType`) and preserves the distinction. + * + * The nanosecond timestamp types are gated behind a preview flag enabled by default under tests + * (`Utils.isTesting`), so it is not set here. The session time zone is fixed so the + * `TIMESTAMP_LTZ` (`Instant`) values are deterministic. The two subclasses run every test with + * ANSI mode on and off. + */ +abstract class TimestampNanosGroupBySuiteBase extends QueryTest with SharedSparkSession { + + override def sparkConf: SparkConf = super.sparkConf + .set(SQLConf.SESSION_LOCAL_TIMEZONE.key, "America/Los_Angeles") + + // Whole-stage codegen on (CODEGEN_ONLY) vs off (NO_CODEGEN). The hash-aggregate exec is the same + // in both modes; only the WholeStageCodegenExec wrapper differs. Mirrors the join and functions + // suites. + protected val codegenModes: Seq[Seq[(String, String)]] = Seq( + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY"), + Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")) + + private def cgLabel(cgConf: Seq[(String, String)]): String = + if (cgConf.contains(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true")) "codegen on" + else "codegen off" + + // ---- relation builders: key column "k" of the given nanos type + an Int value column "v" ---- + // + // Five rows, all sharing epochMicros (2020-01-01T00:00:00Z) and differing only within the + // microsecond: + // - two rows with the 100ns key (an exact duplicate): must collapse into one group, + // - one row with the 900ns key: same microsecond, distinct sub-microsecond remainder, + // - two NULL-key rows: group together into a single NULL group. + // The value column distinguishes the group signatures: sum(v) is 1+2=3 for the 100ns group, + // 3 for the 900ns group, and 4+5=9 for the NULL group, so the three groups are all identifiable + // by their (count, sum) pair even without inspecting the key. + + private def ntzDF(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100"), 1), + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100"), 2), + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000900"), 3), + Row(null, 4), + Row(null, 5))), + new StructType().add("k", TimestampNTZNanosType(p)).add("v", IntegerType)) + + private def ltzDF(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(Instant.parse("2020-01-01T00:00:00.000000100Z"), 1), + Row(Instant.parse("2020-01-01T00:00:00.000000100Z"), 2), + Row(Instant.parse("2020-01-01T00:00:00.000000900Z"), 3), + Row(null, 4), + Row(null, 5))), + new StructType().add("k", TimestampLTZNanosType(p)).add("v", IntegerType)) + + // Each family as (label, builder(p), widerType(p)). GROUP BY is type-agnostic, so the family only + // picks the data builder and (for the mixed-precision widening assertion) the expected widened + // key type. + protected val families: Seq[(String, Int => DataFrame, Int => DataType)] = Seq( + ("NTZ", ntzDF, (p: Int) => TimestampNTZNanosType(p)), + ("LTZ", ltzDF, (p: Int) => TimestampLTZNanosType(p))) + + // Mixed-precision UNION pairs; the key widens to max(pl, pr). + private val mixedPrecisionPairs: Seq[(Int, Int)] = Seq((7, 9), (7, 8), (8, 9)) + + for { + (family, builder, widerType) <- families + cgConf <- codegenModes + } { + // ======================================================================================== + // GROUP BY a nanos key: exact duplicates collapse, but two keys sharing epochMicros and + // differing only within the microsecond stay in separate groups. Aggregates (count, sum) are + // computed per group. Expected three groups: 100ns (count 2, sum 3), 900ns (count 1, sum 3), + // NULL (count 2, sum 9). If the sub-microsecond remainder were ignored, 100ns and 900ns would + // merge into a single (count 3, sum 6) group and this checkAnswer would fail. + // ======================================================================================== + test(s"$family nanos GROUP BY distinguishes the sub-microsecond remainder - " + + s"${cgLabel(cgConf)}") { + withSQLConf(cgConf: _*) { + Seq(7, 8, 9).foreach { p => + val grouped = builder(p).groupBy("k") + .agg(count(lit(1)).as("c"), sum(col("v")).as("s")) + checkAnswer(grouped.select("c", "s"), + Seq(Row(2L, 3L), Row(1L, 3L), Row(2L, 9L))) + } + } + } + + // ======================================================================================== + // GROUP BY over a mixed-precision key: a UNION of the same data at two different nanos + // precisions widens the key to max(pl, pr) (findWiderDateTimeType). Grouping the widened column + // still separates the two sub-microsecond keys (all remainders are multiples of 100ns, exact at + // every p in [7, 9]). The union doubles every row, so counts/sums double: 100ns (4, 6), + // 900ns (2, 6), NULL (4, 18). + // ======================================================================================== + test(s"$family nanos GROUP BY across mixed precisions widens key - ${cgLabel(cgConf)}") { + withSQLConf(cgConf: _*) { + mixedPrecisionPairs.foreach { case (pl, pr) => + val union = builder(pl).union(builder(pr)) + assert(union.schema("k").dataType == widerType(math.max(pl, pr)), + s"expected key widened to ${widerType(math.max(pl, pr))} for ($pl, $pr)") + val grouped = union.groupBy("k") + .agg(count(lit(1)).as("c"), sum(col("v")).as("s")) + checkAnswer(grouped.select("c", "s"), + Seq(Row(4L, 6L), Row(2L, 6L), Row(4L, 18L))) + } + } + } + + // ======================================================================================== + // NULL grouping: unlike an equi-join (where NULL never matches NULL), GROUP BY treats NULL as a + // single group. Both NULL-key rows land in one group (count 2). + // ======================================================================================== + test(s"$family nanos GROUP BY collapses NULL keys into a single group - ${cgLabel(cgConf)}") { + withSQLConf(cgConf: _*) { + Seq(7, 8, 9).foreach { p => + val nullGroup = builder(p).groupBy("k").count().filter(col("k").isNull) + checkAnswer(nullGroup.select("count"), Seq(Row(2L))) + } + } + } + } +} + +// Runs the nanosecond timestamp GROUP BY tests with ANSI mode enabled explicitly. +class TimestampNanosGroupByAnsiOnSuite extends TimestampNanosGroupBySuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "true") +} + +// Runs the nanosecond timestamp GROUP BY tests with ANSI mode disabled explicitly. +class TimestampNanosGroupByAnsiOffSuite extends TimestampNanosGroupBySuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosJoinSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosJoinSuiteBase.scala index b6e7b9aef28c4..79762085c7272 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosJoinSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosJoinSuiteBase.scala @@ -20,6 +20,8 @@ package org.apache.spark.sql import java.time.{Instant, LocalDateTime} import org.apache.spark.SparkConf +import org.apache.spark.sql.catalyst.expressions.EqualTo +import org.apache.spark.sql.catalyst.plans.logical.Join import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} @@ -112,6 +114,27 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt s"Executed plan:\n$plan") } + /** + * Asserts that the equi-join key in `df`'s analyzed plan was coerced to `expected` on BOTH sides + * (type coercion inserts a `Cast` around the narrower side). This pins that a mixed-precision or + * micro-vs-nanos join actually widens the key via `findWiderDateTimeType` rather than, say, + * comparing at the narrower precision and dropping the distinguishing digit. + */ + protected def assertJoinKeyType(df: DataFrame, expected: DataType): Unit = { + val equalTos = df.queryExecution.analyzed.collect { case j: Join => j.condition }.flatten + .flatMap(_.collect { case e: EqualTo => e }) + assert(equalTos.nonEmpty, + s"No EqualTo join condition found in analyzed plan:\n${df.queryExecution.analyzed}") + equalTos.foreach { e => + assert(e.left.dataType == expected && e.right.dataType == expected, + s"Join key not coerced to $expected: ${e.left.dataType} === ${e.right.dataType}") + } + } + + private def cgLabel(cgConf: Seq[(String, String)]): String = + if (cgConf.contains(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true")) "codegen on" + else "codegen off" + // ---- relation builders: key column "k" of the given nanos type + an Int id column ---- private def ntzLeft(p: Int): DataFrame = @@ -150,6 +173,45 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt Row(null, 40))), new StructType().add("k", TimestampLTZNanosType(p)).add("rid", IntegerType)) + // Microsecond-typed sides (TIMESTAMP_NTZ / TIMESTAMP): whole-microsecond keys (nanosWithinMicro + // implicitly 0). Joining these against a nanos side widens the key to the nanos type, so only a + // nanos row whose sub-microsecond remainder is also 0 can match. lid=1 (a whole micro) matches + // the nanos .000000000 row; lid=2 is another whole micro with no match; lid=4 is a NULL key. + private def ntzMicroLeft(): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(LocalDateTime.parse("2020-01-01T00:00:00"), 1), + Row(LocalDateTime.parse("2020-01-01T00:00:01"), 2), + Row(null, 4))), + new StructType().add("k", TimestampNTZType).add("lid", IntegerType)) + + private def ltzMicroLeft(): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(Instant.parse("2020-01-01T00:00:00Z"), 1), + Row(Instant.parse("2020-01-01T00:00:01Z"), 2), + Row(null, 4))), + new StructType().add("k", TimestampType).add("lid", IntegerType)) + + // Nanos right side for the micro-vs-nanos join: rid=10 is exactly on a microsecond boundary + // (.000000000, matches the micro lid=1); rid=20 shares that microsecond but carries a 500ns + // remainder, so it must NOT match the micro key (the widened key keeps the sub-micro bit). + private def ntzNanosRightForMicro(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000000"), 10), + Row(LocalDateTime.parse("2020-01-01T00:00:00.000000500"), 20), + Row(null, 40))), + new StructType().add("k", TimestampNTZNanosType(p)).add("rid", IntegerType)) + + private def ltzNanosRightForMicro(p: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize(Seq( + Row(Instant.parse("2020-01-01T00:00:00.000000000Z"), 10), + Row(Instant.parse("2020-01-01T00:00:00.000000500Z"), 20), + Row(null, 40))), + new StructType().add("k", TimestampLTZNanosType(p)).add("rid", IntegerType)) + // Expected join outputs (order-insensitive). Selected as (lid, rid) so each row is identifiable. // INNER: only the fully-equal sub-microsecond pair (500ns == 500ns). private val expectedInner: Seq[Row] = Seq(Row(1, 10)) @@ -159,6 +221,12 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt private val expectedLeftOuter: Seq[Row] = Seq(Row(1, 10), Row(2, null), Row(3, null), Row(4, null)) + // micro-vs-nanos: only the whole-microsecond left row (lid=1) matches the nanos .000000000 row + // (rid=10). The nanos .000000500 row (rid=20) shares the microsecond but must NOT match, and the + // NULL key (lid=4) never matches. + private val expectedMicroInner: Seq[Row] = Seq(Row(1, 10)) + private val expectedMicroLeftOuter: Seq[Row] = Seq(Row(1, 10), Row(2, null), Row(4, null)) + // ========================================================================================== // NTZ: inner + left-outer over a sub-microsecond key, every strategy x codegen mode x p. // ========================================================================================== @@ -166,14 +234,8 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt (stratName, execClass, stratConf) <- joinStrategies cgConf <- codegenModes } { - val cgName = if (cgConf.exists(_ == (SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true"))) { - "codegen on" - } else { - "codegen off" - } - test(s"NTZ nanos join distinguishes the sub-microsecond remainder - " + - s"$stratName - $cgName") { + s"$stratName - ${cgLabel(cgConf)}") { withSQLConf((stratConf ++ cgConf): _*) { Seq(7, 8, 9).foreach { p => val left = ntzLeft(p) @@ -200,14 +262,8 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt (stratName, execClass, stratConf) <- joinStrategies cgConf <- codegenModes } { - val cgName = if (cgConf.exists(_ == (SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true"))) { - "codegen on" - } else { - "codegen off" - } - test(s"LTZ nanos join distinguishes the sub-microsecond remainder - " + - s"$stratName - $cgName") { + s"$stratName - ${cgLabel(cgConf)}") { withSQLConf((stratConf ++ cgConf): _*) { Seq(7, 8, 9).foreach { p => val left = ltzLeft(p) @@ -226,6 +282,121 @@ abstract class TimestampNanosJoinSuiteBase extends SharedSparkSession with Adapt } } } + + // Mixed-precision pairs. Both sides carry a nanos type but at DIFFERENT precisions; the join key + // is widened to the higher precision (findWiderDateTimeType). Because every remainder is a + // multiple of 100ns (exact at every p in [7, 9]), the same inputs and expected rows hold for + // each pair, and the sub-microsecond equal/distinct relationship is preserved after widening. + private val mixedPrecisionPairs: Seq[(Int, Int)] = Seq((7, 9), (7, 8), (8, 9)) + + // ========================================================================================== + // NTZ mixed precision: p_left != p_right, key widens to max(p_left, p_right). + // ========================================================================================== + for { + (stratName, execClass, stratConf) <- joinStrategies + cgConf <- codegenModes + } { + test(s"NTZ nanos join across mixed precisions widens key - $stratName - ${cgLabel(cgConf)}") { + withSQLConf((stratConf ++ cgConf): _*) { + mixedPrecisionPairs.foreach { case (pl, pr) => + val left = ntzLeft(pl) + val right = ntzRight(pr) + val wider = TimestampNTZNanosType(math.max(pl, pr)) + + val inner = left.join(right, left("k") === right("k"), "inner") + .select(left("lid"), right("rid")) + assertJoinUsed(inner, execClass) + assertJoinKeyType(inner, wider) + checkAnswer(inner, expectedInner) + + val leftOuter = left.join(right, left("k") === right("k"), "left_outer") + .select(left("lid"), right("rid")) + assertJoinUsed(leftOuter, execClass) + checkAnswer(leftOuter, expectedLeftOuter) + } + } + } + } + + // ========================================================================================== + // LTZ mixed precision: p_left != p_right, key widens to max(p_left, p_right). + // ========================================================================================== + for { + (stratName, execClass, stratConf) <- joinStrategies + cgConf <- codegenModes + } { + test(s"LTZ nanos join across mixed precisions widens key - $stratName - ${cgLabel(cgConf)}") { + withSQLConf((stratConf ++ cgConf): _*) { + mixedPrecisionPairs.foreach { case (pl, pr) => + val left = ltzLeft(pl) + val right = ltzRight(pr) + val wider = TimestampLTZNanosType(math.max(pl, pr)) + + val inner = left.join(right, left("k") === right("k"), "inner") + .select(left("lid"), right("rid")) + assertJoinUsed(inner, execClass) + assertJoinKeyType(inner, wider) + checkAnswer(inner, expectedInner) + + val leftOuter = left.join(right, left("k") === right("k"), "left_outer") + .select(left("lid"), right("rid")) + assertJoinUsed(leftOuter, execClass) + checkAnswer(leftOuter, expectedLeftOuter) + } + } + } + } + + // ========================================================================================== + // Microsecond timestamp JOIN nanosecond timestamp: the key widens from the micro type to the + // nanos type, and only a whole-microsecond (zero-remainder) nanos row can match the micro key. + // ========================================================================================== + for { + (stratName, execClass, stratConf) <- joinStrategies + cgConf <- codegenModes + } { + test(s"NTZ micro join nanos widens to the nanos key - $stratName - ${cgLabel(cgConf)}") { + withSQLConf((stratConf ++ cgConf): _*) { + Seq(7, 8, 9).foreach { p => + val left = ntzMicroLeft() + val right = ntzNanosRightForMicro(p) + val wider = TimestampNTZNanosType(p) + + val inner = left.join(right, left("k") === right("k"), "inner") + .select(left("lid"), right("rid")) + assertJoinUsed(inner, execClass) + assertJoinKeyType(inner, wider) + checkAnswer(inner, expectedMicroInner) + + val leftOuter = left.join(right, left("k") === right("k"), "left_outer") + .select(left("lid"), right("rid")) + assertJoinUsed(leftOuter, execClass) + checkAnswer(leftOuter, expectedMicroLeftOuter) + } + } + } + + test(s"LTZ micro join nanos widens to the nanos key - $stratName - ${cgLabel(cgConf)}") { + withSQLConf((stratConf ++ cgConf): _*) { + Seq(7, 8, 9).foreach { p => + val left = ltzMicroLeft() + val right = ltzNanosRightForMicro(p) + val wider = TimestampLTZNanosType(p) + + val inner = left.join(right, left("k") === right("k"), "inner") + .select(left("lid"), right("rid")) + assertJoinUsed(inner, execClass) + assertJoinKeyType(inner, wider) + checkAnswer(inner, expectedMicroInner) + + val leftOuter = left.join(right, left("k") === right("k"), "left_outer") + .select(left("lid"), right("rid")) + assertJoinUsed(leftOuter, execClass) + checkAnswer(leftOuter, expectedMicroLeftOuter) + } + } + } + } } // Runs the nanosecond timestamp join tests with ANSI mode enabled explicitly. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosRenderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosRenderingSuiteBase.scala new file mode 100644 index 0000000000000..f3525de4af0e8 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosRenderingSuiteBase.scala @@ -0,0 +1,177 @@ +/* + * 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 + +import java.time.{Instant, LocalDateTime} + +import org.apache.spark.SparkConf +// castToImpl: `spark.createDataFrame` is typed as the public sql.DataFrame, but `showString` is a +// classic-only method; this implicit narrows the receiver to classic.Dataset for that call. +import org.apache.spark.sql.classic.ClassicConversions.castToImpl +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ + +/** + * End-to-end `show()` / `collect()` rendering tests over the nanosecond-precision timestamp types + * `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in `[7, 9]`). These ride on the nanos cast-to-string + * (`ToPrettyString` -> `ToStringBase`, which routes nanos through the Types Framework fraction + * formatter) and the nanos external encoders, so no production change is required. + * + * Two distinct surfaces are pinned, neither of which the golden SQL files can express: + * - `Dataset.show()` string rendering: the fractional second is rendered to the type's precision + * with sub-precision digits floored to zero and an all-zero fraction trimmed entirely. A value + * that only differs below the type's precision therefore renders identically -- the display is + * lossy at `p`, exactly as the stored value is. (The golden `.out` files render via + * `hiveResultString`, a different entry point onto the same formatter; `HiveResultSuite` covers + * that one. `show()` is what a user actually sees at the console.) + * - `collect()` external round-trip: NTZ comes back as `java.time.LocalDateTime` and LTZ as + * `java.time.Instant` (see `RowEncoder`), and `.getNano` preserves the sub-microsecond + * remainder floored to the type's precision -- so two values that share `epochMicros` but + * differ in `nanosWithinMicro` are distinguishable after a full driver round-trip. + * + * The nanosecond timestamp types are gated behind a preview flag enabled by default under tests + * (`Utils.isTesting`), so it is not set here. The session time zone is fixed to America/Los_Angeles + * (UTC-08:00, no DST on 2020-01-01) so the `TIMESTAMP_LTZ` (`Instant`) values render + * deterministically in wall-clock time. The two subclasses run every test with ANSI mode on/off. + */ +abstract class TimestampNanosRenderingSuiteBase extends QueryTest with SharedSparkSession { + + override def sparkConf: SparkConf = super.sparkConf + .set(SQLConf.SESSION_LOCAL_TIMEZONE.key, "America/Los_Angeles") + + // Single nanosecond TIMESTAMP_NTZ(p) column "c"; a null element becomes a NULL row. + private def ntzDF(values: Seq[String], precision: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize( + values.map(s => Row(if (s == null) null else LocalDateTime.parse(s)))), + new StructType().add("c", TimestampNTZNanosType(precision))) + + // Single nanosecond TIMESTAMP_LTZ(p) column "c"; a null element becomes a NULL row. + private def ltzDF(values: Seq[String], precision: Int): DataFrame = + spark.createDataFrame( + spark.sparkContext.parallelize( + values.map(s => Row(if (s == null) null else Instant.parse(s)))), + new StructType().add("c", TimestampLTZNanosType(precision))) + + /** + * The data cells of `Dataset.show(truncate = 0)`, one string per (single-column) row, in row + * order with the alignment padding stripped. `showString` frames the table as a top border, + * a header row, a separator, one line per data row, then a bottom border; dropping the first + * three lines and the last isolates the data rows, and `.trim` removes the right-alignment + * padding so the exact rendered value can be compared (a prefix `contains` check could not tell + * a precision-floored fraction from a longer one). + */ + private def shownCells(df: DataFrame): Seq[String] = { + val lines = df.showString(100, truncate = 0, vertical = false).split("\n").toSeq + lines.drop(3).dropRight(1).map(_.stripPrefix("|").stripSuffix("|").trim) + } + + // The per-precision floored rendering of ".123456789": p=7 -> 7 digits, p=8 -> 8, p=9 -> 9. + private val frac: Map[Int, String] = Map( + 7 -> ".1234567", 8 -> ".12345678", 9 -> ".123456789") + + // ========================================================================================== + // show() renders the fraction to the type's precision, flooring sub-precision digits and + // trimming an all-zero fraction. NTZ is zone-independent. + // ========================================================================================== + test("show() renders nanosecond TIMESTAMP_NTZ to the type precision") { + Seq(7, 8, 9).foreach { p => + val df = ntzDF(Seq( + "2020-01-01T00:00:00.123456789", // floored to p digits + "2020-01-01T00:00:00.000000001", // non-zero only at digit 9 -> survives only at p=9 + "2020-01-01T00:00:00", // all-zero fraction -> no fraction + null), p) + assert(shownCells(df) === Seq( + "2020-01-01 00:00:00" + frac(p), + if (p == 9) "2020-01-01 00:00:00.000000001" else "2020-01-01 00:00:00", + "2020-01-01 00:00:00", + "NULL")) + } + } + + test("show() renders nanosecond TIMESTAMP_LTZ in the session zone to the type precision") { + Seq(7, 8, 9).foreach { p => + // UTC instants; the session zone is UTC-08:00, so 08:00:00Z renders as 00:00:00 wall-clock. + val df = ltzDF(Seq( + "2020-01-01T08:00:00.123456789Z", + "2020-01-01T08:00:00.000000001Z", + "2020-01-01T08:00:00Z", + null), p) + assert(shownCells(df) === Seq( + "2020-01-01 00:00:00" + frac(p), + if (p == 9) "2020-01-01 00:00:00.000000001" else "2020-01-01 00:00:00", + "2020-01-01 00:00:00", + "NULL")) + } + } + + test("show() renders nanosecond timestamps nested in array / struct") { + // A 9-digit fraction inside a complex type still renders to the element precision (p=9 here). + val ntz = ntzDF(Seq("2020-01-01T00:00:00.123456789"), 9) + assert(shownCells(ntz.selectExpr("array(c)")) === + Seq("[2020-01-01 00:00:00.123456789]")) + assert(shownCells(ntz.selectExpr("named_struct('f', c)")) === + Seq("{2020-01-01 00:00:00.123456789}")) + } + + // ========================================================================================== + // collect() returns the external LocalDateTime / Instant, preserving the sub-microsecond + // remainder floored to the type precision. Two values sharing epochMicros are distinguishable. + // ========================================================================================== + // .000000123 floors to 100ns at p=7, 120ns at p=8, 123ns at p=9; .000000999 -> 900 / 990 / 999. + private def flooredNano(base: Int, p: Int): Int = p match { + case 7 => base / 100 * 100 + case 8 => base / 10 * 10 + case 9 => base + } + + test("collect() over nanosecond TIMESTAMP_NTZ preserves the precision-floored remainder") { + Seq(7, 8, 9).foreach { p => + val df = ntzDF(Seq( + "2020-01-01T00:00:00.000000123", + "2020-01-01T00:00:00.000000999", + null), p) + val got = df.collect().map(r => Option(r.getAs[LocalDateTime]("c")).map(_.getNano)).toSet + assert(got === Set(Some(flooredNano(123, p)), Some(flooredNano(999, p)), None)) + // The two non-null values share epochMicros yet stay distinct at every supported precision. + assert(flooredNano(123, p) != flooredNano(999, p)) + } + } + + test("collect() over nanosecond TIMESTAMP_LTZ preserves the precision-floored remainder") { + Seq(7, 8, 9).foreach { p => + val df = ltzDF(Seq( + "2020-01-01T00:00:00.000000123Z", + "2020-01-01T00:00:00.000000999Z", + null), p) + val got = df.collect().map(r => Option(r.getAs[Instant]("c")).map(_.getNano)).toSet + assert(got === Set(Some(flooredNano(123, p)), Some(flooredNano(999, p)), None)) + } + } +} + +// Runs the nanosecond timestamp rendering tests with ANSI mode enabled explicitly. +class TimestampNanosRenderingAnsiOnSuite extends TimestampNanosRenderingSuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "true") +} + +// Runs the nanosecond timestamp rendering tests with ANSI mode disabled explicitly. +class TimestampNanosRenderingAnsiOffSuite extends TimestampNanosRenderingSuiteBase { + override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false") +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosWideningSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosWideningSuiteBase.scala index 22bb69fdfc6b4..91bfcb16f9d6e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosWideningSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosWideningSuiteBase.scala @@ -120,6 +120,40 @@ abstract class TimestampNanosWideningSuiteBase extends SharedSparkSession { checkAnswer(ntz.selectExpr("a IN (b)"), Row(false)) } + test("SPARK-56822: NOT IN / NOT EXISTS / scalar subquery over nanosecond timestamps") { + val ntzA = LocalDateTime.parse("2020-01-01T00:00:00.000000001") + val ntzB = LocalDateTime.parse("2020-01-01T00:00:00.000000999") + // Two sub-microsecond-distinct nanos values in the outer relation. + val outer = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(ntzA), Row(ntzB))), + new StructType().add("c", TimestampNTZNanosType(9))) + + withTempView("outer_t") { + outer.createOrReplaceTempView("outer_t") + + // NOT IN over the nanos key: only the value absent from the subquery set survives, and the + // nanosecond digit -- not just the microsecond -- decides membership. + checkAnswer( + spark.sql( + "SELECT c FROM outer_t WHERE c NOT IN " + + "(SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999')"), + Row(ntzA)) + + // NOT EXISTS correlated on the nanos key: the outer row survives iff no matching key exists. + checkAnswer( + spark.sql( + "SELECT o.c FROM outer_t o WHERE NOT EXISTS " + + "(SELECT 1 FROM outer_t i WHERE i.c = o.c AND " + + "i.c = TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001')"), + Row(ntzB)) + } + + // A scalar subquery carries the nanos type into its result column. + val scalar = spark.sql("SELECT (SELECT TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999') AS c") + assert(scalar.schema("c").dataType === TimestampNTZNanosType(9)) + checkAnswer(scalar, Row(ntzB)) + } + test("SPARK-57454: binary comparison widens nanosecond timestamps") { // Equal absolute instants stored at different precisions compare equal. val ltzEq = twoCols(TimestampType, instantA, TimestampLTZNanosType(9), instantA) @@ -132,6 +166,29 @@ abstract class TimestampNanosWideningSuiteBase extends SharedSparkSession { val ntzEq = twoCols(TimestampNTZType, ldtA, TimestampNTZNanosType(9), ldtA) checkAnswer(ntzEq.selectExpr("a = b", "a < b"), Row(true, false)) } + + test("SPARK-57811: string operand coerces to the nanosecond timestamp type") { + // Equality on the exact sub-microsecond value: the string is parsed at nanos precision, so a + // string that differs only in the 9th digit does NOT compare equal. The paired + // `= '...789' => true` / `= '...788' => false` is the load-bearing pair: it rules out + // micro-truncation (truncating the operand to micros would make both strings equal the column). + // The nanos-vs-string-promotion distinction is locked separately by the analyzer goldens + // (explicit `cast as timestamp_ntz(9)`) and the unit tests in TypeCoercionSuite. + val ntz = + single(TimestampNTZNanosType(9), LocalDateTime.parse("2020-01-02T03:04:05.123456789")) + checkAnswer( + ntz.selectExpr( + "c = '2020-01-02 03:04:05.123456789'", + "c = '2020-01-02 03:04:05.123456788'", + "c < '2020-01-02 03:04:05.123456790'", + "c > '2020-01-02 03:04:05.123456788'"), + Row(true, false, true, true)) + + // LTZ operand: the string literal is parsed in the session zone (America/Los_Angeles), so the + // 2020-01-01T00:00:00Z instant equals the local wall-clock time eight hours earlier. + val ltz = single(TimestampLTZNanosType(9), Instant.parse("2020-01-01T00:00:00Z")) + checkAnswer(ltz.selectExpr("c = '2019-12-31 16:00:00'"), Row(true)) + } } // Runs the nanosecond timestamp widening tests with ANSI mode enabled explicitly. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala index 4e39f01bf5374..21384d7c4b917 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala @@ -1252,6 +1252,21 @@ class UDFSuite extends SharedSparkSession { parameters = Map("dataType" -> s"\"${dt.sql}\"") ) } + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + Seq(CharType(5), VarcharType(5)).foreach { dt => + val f = udf( + new UDF0[String] { + override def call(): String = "a" + }, + dt + ) + val df = spark.range(1).select(f().as("c")) + assert(df.schema.head.dataType === dt) + assert(df.encoder.schema.head.dataType === dt) + val expected = if (dt.isInstanceOf[CharType]) "a " else "a" + checkAnswer(df, Row(expected)) + } + } } test("SPARK-47927: ScalaUDF null handling") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala index 18adb45c1234e..04be47523c588 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/UserDefinedTypeSuite.scala @@ -293,8 +293,33 @@ class UserDefinedTypeSuite extends SharedSparkSession with ParquetTest val unwrappedFeaturesArrays: Array[Array[Double]] = unwrappedFeatures.collect() assert(unwrappedFeaturesArrays.length === 2) - java.util.Arrays.equals(unwrappedFeaturesArrays(0), Array(0.1, 1.0)) - java.util.Arrays.equals(unwrappedFeaturesArrays(1), Array(0.2, 2.0)) + assert(Arrays.equals(unwrappedFeaturesArrays(0), Array(0.1, 1.0))) + assert(Arrays.equals(unwrappedFeaturesArrays(1), Array(0.2, 2.0))) + } + + test("Test wrap_udt function") { + val udt = new TestUDT.MyDenseVectorUDT() + val df = Seq(Array(0.1, 1.0), Array(0.2, 2.0)).toDF("features") + val wrappedFeatures = df.select(wrap_udt(col("features"), udt).as("features")) + + assert(wrappedFeatures.schema("features").dataType === udt) + checkAnswer( + wrappedFeatures, + Seq( + Row(new TestUDT.MyDenseVector(Array(0.1, 1.0))), + Row(new TestUDT.MyDenseVector(Array(0.2, 2.0))))) + } + + test("Test unwrap_udt and wrap_udt round trip") { + val udt = new TestUDT.MyDenseVectorUDT() + val roundTrip = pointsRDD.select(wrap_udt(unwrap_udt(col("features")), udt).as("features")) + + assert(roundTrip.schema("features").dataType === udt) + checkAnswer( + roundTrip, + Seq( + Row(new TestUDT.MyDenseVector(Array(0.1, 1.0))), + Row(new TestUDT.MyDenseVector(Array(0.2, 2.0))))) } test("SPARK-46289: UDT ordering") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala index 2d26356890d28..30675e759428d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala @@ -246,6 +246,31 @@ class VariantEndToEndSuite extends SharedSparkSession { } } + test("variant_from_arrays and variant_from_entries - Codegen Support") { + Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { codegenMode => + withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) { + val entryType = StructType(Array( + StructField("k", StringType), StructField("v", IntegerType))) + val schema = StructType(Array( + StructField("keys", ArrayType(StringType)), + StructField("values", ArrayType(IntegerType)), + StructField("entries", ArrayType(entryType)))) + // Source non-foldable rows so the operators' doGenCode is actually exercised under codegen. + val data = Seq(Row(Seq("a", "b"), Seq(1, 2), Seq(Row("a", 1), Row("b", 2)))) + val df = spark.createDataFrame(spark.sparkContext.parallelize(data), schema) + val arraysDF = df.select(variant_from_arrays(col("keys"), col("values")).cast("string")) + val entriesDF = df.select(variant_from_entries(col("entries")).cast("string")) + val wholeStage = codegenMode == "CODEGEN_ONLY" + assert(arraysDF.queryExecution.executedPlan.exists( + _.isInstanceOf[WholeStageCodegenExec]) == wholeStage) + assert(entriesDF.queryExecution.executedPlan.exists( + _.isInstanceOf[WholeStageCodegenExec]) == wholeStage) + checkAnswer(arraysDF, Row("""{"a":1,"b":2}""")) + checkAnswer(entriesDF, Row("""{"a":1,"b":2}""")) + } + } + } + test("schema_of_variant") { def check(json: String, expected: String): Unit = { val df = Seq(json).toDF("j").selectExpr("schema_of_variant(parse_json(j))") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala index 864d0c9b631bd..9c2bfaca03a2a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala @@ -486,6 +486,11 @@ class VariantSuite extends SharedSparkSession with ExpressionEvalHelper { assert(intercept[AnalysisException] { sql("SELECT variant_set(parse_json('{}'), '$.a', named_struct('x', 1))") }.getCondition == "DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION") + + // A non-constant create_if_missing is rejected at analysis. + assert(intercept[AnalysisException] { + sql("SELECT variant_set(parse_json('{\"a\": 1}'), '$.a', 2, c) FROM VALUES (true) AS t(c)") + }.getCondition == "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT") } test("variant_set with dynamic arguments") { @@ -680,6 +685,82 @@ class VariantSuite extends SharedSparkSession with ExpressionEvalHelper { } } + test("variant_strip_nulls with literal arguments") { + def rows(results: Any*): Seq[Row] = results.map(Row(_)) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('{\"a\": 1, \"b\": null, \"c\": 3}')))"), + rows("""{"a":1,"c":3}""")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('[1, null, 3]')))"), + rows("[1,3]")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json(" + + "'{\"a\": [null, 3, {\"b\": null, \"c\": [null, 1]}], \"d\": null, " + + "\"e\": {\"f\": null, \"g\": 2}}')))"), + rows("""{"a":[3,{"c":[1]}],"e":{"g":2}}""")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json(" + + "'{\"a\": 100000, \"b\": null, \"c\": 10000000000, \"d\": \"hello world\"}')))"), + rows("""{"a":100000,"c":10000000000,"d":"hello world"}""")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('{\"a\": [1, null], \"b\": null}'), " + + "include_arrays => false))"), + rows("""{"a":[1,null]}""")) + + // include_arrays = false keeps array null elements but still strips null fields of objects. + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(" + + "parse_json('[{\"a\": 1, \"b\": null}, null, {\"c\": null, \"d\": 4}]'), false))"), + rows("""[{"a":1},null,{"d":4}]""")) + + // Empty containers are preserved. + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('{\"a\": null}')))"), + rows("{}")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('[null, null]')))"), + rows("[]")) + + // Top-level variant null is unchanged. + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(parse_json('null')))"), + rows("null")) + + checkAnswer( + sql("SELECT to_json(variant_strip_nulls(CAST(NULL AS VARIANT)))"), + rows(null)) + + assert(intercept[AnalysisException] { + sql("SELECT variant_strip_nulls(parse_json('{\"a\": null}'), c) FROM VALUES (true) AS t(c)") + }.getCondition == "DATATYPE_MISMATCH.NON_FOLDABLE_INPUT") + } + + test("variant_strip_nulls with dynamic arguments") { + def rows(results: Any*): Seq[Row] = results.map(Row(_)) + val df = Seq( + """{"a": [1, null], "b": null}""", + """{"x": null, "y": 2}""", + null + ).toDF("json") + val v = parse_json(col("json")) + + // Single-argument overload defaults include_arrays to true. + checkAnswer( + df.select(to_json(variant_strip_nulls(v)).alias("r")), + rows("""{"a":[1]}""", """{"y":2}""", null)) + + // Boolean overload with include_arrays = false preserves array null elements. + checkAnswer( + df.select(to_json(variant_strip_nulls(v, false)).alias("r")), + rows("""{"a":[1,null]}""", """{"y":2}""", null)) + } + test("round trip tests") { withSQLConf(SQLConf.VARIANT_INFER_SHREDDING_SCHEMA.key -> "false") { val rand = new Random(42) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/XmlFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/XmlFunctionsSuite.scala index 73cdc7df139ad..bfe46e681261f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/XmlFunctionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/XmlFunctionsSuite.scala @@ -361,6 +361,69 @@ class XmlFunctionsSuite extends SharedSparkSession { checkAnswer(dfTwo, readBackTwo) } + test("SPARK-57458: to_xml with nanos timestamp types") { + withSQLConf( + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + foreachNanosPrecision { p => + // The pattern must carry `p` fractional digits to emit the full declared precision; the + // stored value is truncated to precision `p`, so the rendered fraction is the first + // `p` digits of 123456789. + val fracPat = "S" * p + val frac = "123456789".take(p) + val ldt = LocalDateTime.of(2020, 1, 1, 0, 0, 0, 123456789) + Seq( + (TimestampNTZNanosType(p): DataType, "timestampNTZFormat", + s"yyyy-MM-dd'T'HH:mm:ss.$fracPat", ldt: Any, + s"2020-01-01T00:00:00.$frac"), + (TimestampLTZNanosType(p): DataType, "timestampFormat", + s"yyyy-MM-dd'T'HH:mm:ss.${fracPat}XXX", ldt.toInstant(ZoneOffset.UTC): Any, + s"2020-01-01T00:00:00.${frac}Z")).foreach { + case (nanosType, optKey, fmt, value, expectedTs) => + val schema = new StructType().add("ts", nanosType) + val df = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(value))), schema) + val expectedXml = + s"""|<ROW> + | <ts>$expectedTs</ts> + |</ROW>""".stripMargin + checkAnswer( + df.select(to_xml(struct($"ts"), Map(optKey -> fmt).asJava)), + Row(expectedXml)) + } + } + } + } + + test("SPARK-57458: roundtrip in to_xml and from_xml - nanos timestamps") { + withSQLConf( + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true", + SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + foreachNanosPrecision { p => + val fracPat = "S" * p + val ldt = LocalDateTime.of(2020, 1, 1, 0, 0, 0, 123456789) + Seq( + (TimestampNTZNanosType(p): DataType, "timestampNTZFormat", + s"yyyy-MM-dd'T'HH:mm:ss.$fracPat", ldt: Any), + (TimestampLTZNanosType(p): DataType, "timestampFormat", + s"yyyy-MM-dd'T'HH:mm:ss.${fracPat}XXX", ldt.toInstant(ZoneOffset.UTC): Any)).foreach { + case (nanosType, optKey, fmt, value) => + val schema = new StructType().add("ts", nanosType) + val df = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(value))), schema) + val options = Map(optKey -> fmt).asJava + // The input column already carries precision `p`, so the to_xml -> from_xml + // round-trip with a `p`-digit pattern is loss-free. + val readBack = df + .select(to_xml(struct($"ts"), options).as("xml")) + .select(from_xml($"xml", schema, options).as("data")) + .select($"data.ts".as("ts")) + checkAnswer(readBack, df.select($"ts")) + } + } + } + } + test("Support to_xml in SQL") { val schemaOne = StructType(StructField("a", IntegerType, nullable = false) :: Nil) val dataOne = Seq(Row(1)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/DataFrameAnalyzerTestGapsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/DataFrameAnalyzerTestGapsSuite.scala index 001ec8f9cb0ed..a533d1525b4a4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/DataFrameAnalyzerTestGapsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/DataFrameAnalyzerTestGapsSuite.scala @@ -75,4 +75,12 @@ class DataFrameAnalyzerTestGapsSuite extends SharedSparkSession { Seq(Row(4), Row(6), Row(8)) ) } + + test("Stacked orderBy resolving hidden columns across levels") { + // Each orderBy resolves a dropped column from hidden output, widening an intermediate Project. + // The missing-input check runs on each such Project; a valid query must not trip it. + val table = Seq((1, 2, 3), (4, 5, 6)).toDF("col1", "col2", "col3") + val query = table.select($"col1").orderBy($"col2").orderBy($"col1").orderBy($"col2") + checkAnswer(query, Seq(Row(1), Row(4))) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ExplicitlyUnsupportedResolverFeatureSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ExplicitlyUnsupportedResolverFeatureSuite.scala index 0a9c1c07d1bda..a8c1041509cc7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ExplicitlyUnsupportedResolverFeatureSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ExplicitlyUnsupportedResolverFeatureSuite.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.analysis.resolver.{ Resolver } import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class ExplicitlyUnsupportedResolverFeatureSuite extends SharedSparkSession { @@ -44,6 +45,98 @@ class ExplicitlyUnsupportedResolverFeatureSuite extends SharedSparkSession { } } + test("SPARK-57353: HAVING with grouping analytics is unsupported (SPARK-57346)") { + checkResolution( + """SELECT a, SUM(b) FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY ROLLUP(a) HAVING SUM(b) > 30""".stripMargin, + shouldPass = false, + expectedMessage = Some("HAVING with grouping analytics (SPARK-57346)") + ) + } + + test("SPARK-57353: ORDER BY with grouping analytics is unsupported (SPARK-57346)") { + checkResolution( + """SELECT a, SUM(b) as s FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a) ORDER BY s""".stripMargin, + shouldPass = false, + expectedMessage = Some("ORDER BY with grouping analytics (SPARK-57346)") + ) + } + + test("SPARK-57353: LCA with grouping analytics is unsupported") { + checkResolution( + """SELECT a, SUM(b) as total, total + 1 + |FROM VALUES (1,10),(1,20),(2,30) AS t(a,b) + |GROUP BY CUBE(a)""".stripMargin, + shouldPass = false, + expectedMessage = Some("lateral column alias with grouping analytics") + ) + } + + test("ASOF JOIN MATCH_CONDITION operands requiring element-wise ordering") { + withSQLConf(SQLConf.SQL_ASOF_JOIN_ENABLED.key -> "true") { + checkResolution( + """SELECT t.a, r.a FROM VALUES (ARRAY(1, 3)) AS t(a) + |ASOF JOIN VALUES (ARRAY(1, 2)) AS r(a) MATCH_CONDITION (t.a >= r.a);""".stripMargin, + expectedMessage = Some("MATCH_CONDITION with a lambda-based ordering expression") + ) + checkResolution( + """SELECT t.a, r.a FROM VALUES (ARRAY(named_struct('seq', 1))) AS t(a) + |ASOF JOIN VALUES (ARRAY(named_struct('seq', 2))) AS r(a) + |MATCH_CONDITION (t.a >= r.a);""".stripMargin, + expectedMessage = Some("MATCH_CONDITION with a lambda-based ordering expression") + ) + checkResolution( + """SELECT t.k, r.k FROM VALUES (10) AS t(k) + |ASOF JOIN VALUES (5) AS r(k) MATCH_CONDITION (t.k >= r.k);""".stripMargin, + shouldPass = true + ) + } + } + + test("SPARK-57353: ORDER BY with grouping analytics in subquery does not reject outer sort") { + // Grouping analytics inside a scalar subquery must not leak hasGroupingAnalytics to the + // outer query. The outer ORDER BY is unrelated and should resolve successfully. + checkResolution( + """SELECT (SELECT SUM(x) FROM VALUES (1),(2) t(x) + | GROUP BY GROUPING SETS (())) AS s ORDER BY s""".stripMargin, + shouldPass = true + ) + } + + test("SPARK-57353: HAVING with grouping analytics in subquery does not reject outer having") { + // Same boundary isolation for HAVING: grouping analytics confined to a subquery must not + // cause the outer HAVING to be rejected. + checkResolution( + """SELECT a, (SELECT SUM(x) FROM VALUES (1),(2) t(x) + | GROUP BY GROUPING SETS (())) AS s + |FROM VALUES (1),(2),(3) v(a) + |GROUP BY a + |HAVING a > 1""".stripMargin, + shouldPass = true + ) + } + + test("SPARK-57353: HAVING with grouping analytics in derived table does not reject outer") { + // Grouping analytics inside a derived table (SubqueryAlias) must not leak + // hasGroupingAnalytics to the outer query. The outer HAVING is unrelated. + checkResolution( + """SELECT a FROM (SELECT a FROM VALUES (1) t(a) GROUP BY CUBE(a)) s + |GROUP BY a HAVING a > 0""".stripMargin, + shouldPass = true + ) + } + + test("SPARK-57353: ORDER BY with grouping analytics in derived table does not reject outer") { + // Same boundary isolation for ORDER BY: grouping analytics confined to a derived table + // must not cause the outer ORDER BY to be rejected. + checkResolution( + """SELECT a FROM (SELECT a FROM VALUES (1) t(a) GROUP BY CUBE(a)) s + |ORDER BY a""".stripMargin, + shouldPass = true + ) + } + private def checkResolution( sqlText: String, shouldPass: Boolean = false, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ResolverGuardSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ResolverGuardSuite.scala index 6c61c7797a6ca..26e95313e5941 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ResolverGuardSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/analysis/resolver/ResolverGuardSuite.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.analysis.resolver.{ } import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession class ResolverGuardSuite extends ResolverGuardSuiteBase { @@ -455,6 +456,27 @@ class ResolverGuardSuite extends ResolverGuardSuiteBase { } } + gridTest("ASOF JOIN")(Seq(true, false)) { enabled => + val asOfJoinQuery = + "SELECT t.symbol, q.bid_price FROM " + + "VALUES (TIMESTAMP '2026-06-29 10:00:05', 'AAPL') AS t(trade_time, symbol) " + + "ASOF JOIN VALUES (TIMESTAMP '2026-06-29 10:00:00', 'AAPL', 180.10) " + + "AS q(quote_time, symbol, bid_price) " + + "MATCH_CONDITION (t.trade_time >= q.quote_time) ON t.symbol = q.symbol" + val expectedReason = if (enabled) { + None + } else { + Some("class org.apache.spark.sql.catalyst.plans.logical.AsOfJoin operator resolution") + } + withSQLConf( + SQLConf.SQL_ASOF_JOIN_ENABLED.key -> "true", + SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLE_ASOF_JOIN_RESOLUTION.key -> + enabled.toString + ) { + checkResolverGuard(asOfJoinQuery, expectedReason) + } + } + } trait ResolverGuardSuiteBase extends SharedSparkSession { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/artifact/ArtifactManagerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/artifact/ArtifactManagerSuite.scala index a1678335b7755..0391103aa283d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/artifact/ArtifactManagerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/artifact/ArtifactManagerSuite.scala @@ -22,10 +22,10 @@ import java.nio.file.{Files, Path, Paths} import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException} import org.apache.spark.metrics.source.CodegenMetrics -import org.apache.spark.sql.Artifact +import org.apache.spark.sql.{AnalysisException, Artifact} import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.functions.col -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.DataTypes import org.apache.spark.storage.CacheId @@ -253,6 +253,30 @@ class ArtifactManagerSuite extends SharedSparkSession { assert(copiedClassFile.exists()) } + test("SPARK-58531: allowDestLocal cannot be set from a session") { + // The conf gates writes to a local filesystem destination on the driver, so it must stay a + // static conf: a session that could turn it on would be able to write to arbitrary paths on + // the driver. Guard against it being made session-settable again. + val key = StaticSQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL.key + assert(SQLConf.isStaticConfigKey(key)) + checkError( + exception = intercept[AnalysisException](spark.conf.set(key, "true")), + condition = "CANNOT_MODIFY_STATIC_CONFIG", + parameters = Map("key" -> s""""$key"""")) + } + + test("SPARK-58531: StaticSQLConf can initialize before SQLConf") { + // Use a fresh process because both objects may already be initialized in this test JVM. + val sparkHome = sys.props.getOrElse("spark.test.home", fail("spark.test.home is not set!")) + val process = Utils.executeCommand( + Seq( + s"$sparkHome/bin/spark-class", + StaticSQLConfInitializationTestApp.getClass.getCanonicalName.stripSuffix("$")), + new File(sparkHome), + Map("SPARK_TESTING" -> "1", "SPARK_HOME" -> sparkHome)) + assert(process.waitFor() === 0) + } + test("Removal of resources") { withTempPath { path => @@ -714,3 +738,10 @@ class ArtifactManagerSuite extends SharedSparkSession { } } } + +object StaticSQLConfInitializationTestApp { + def main(args: Array[String]): Unit = { + val key = StaticSQLConf.ARTIFACT_COPY_FROM_LOCAL_TO_FS_ALLOW_DEST_LOCAL.key + require(key == "spark.sql.artifact.copyFromLocalToFs.allowDestLocal") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/catalyst/expressions/ParseSqlSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/catalyst/expressions/ParseSqlSuite.scala new file mode 100644 index 0000000000000..f44a409a0d641 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/catalyst/expressions/ParseSqlSuite.scala @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.catalyst.expressions + +import org.json4s._ +import org.json4s.jackson.JsonMethods.parse + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.catalyst.analysis.TypeCheckResult +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.plans.SQLHelper +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StringType +import org.apache.spark.unsafe.types.UTF8String + +class ParseSqlSuite extends SparkFunSuite with ExpressionEvalHelper with SQLHelper { + + private def evalJson(sql: String): JValue = { + val result = ParseSql(Literal(sql)).eval().asInstanceOf[UTF8String].toString + parse(result) + } + + test("parse_sql is disabled by default") { + assert(!SQLConf.get.parseSqlEnabled) + checkError( + exception = intercept[AnalysisException] { + ParseSql(Literal("SELECT 1")).checkInputDataTypes() + }, + condition = "FEATURE_NOT_ENABLED", + parameters = Map( + "featureName" -> "parse_sql", + "configKey" -> SQLConf.PARSE_SQL_ENABLED.key, + "configValue" -> "true")) + } + + test("parse_sql type check succeeds when enabled") { + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> "true") { + assert(ParseSql(Literal("SELECT 1")).checkInputDataTypes() === + TypeCheckResult.TypeCheckSuccess) + } + } + + test("parse_sql returns JSON for a valid SELECT") { + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> "true") { + val j = evalJson("SELECT 1 AS a") + assert(j \ "parse_success" === JBool(true)) + assert(j \ "statement_identifier" === JString("SELECT")) + assert(j \ "statement_code" === JInt(21)) + } + } + + test("parse_sql returns null for null input") { + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> "true") { + checkEvaluation(ParseSql(Literal.create(null, StringType)), null) + } + } + + test("parse_sql does not throw on syntax error") { + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> "true") { + val j = evalJson("NOT A STATEMENT !!!") + assert(j \ "parse_success" === JBool(false)) + assert(j \ "error" \ "errorClass" === JString("PARSE_SYNTAX_ERROR")) + } + } + + test("parse_sql works with CodegenFallback path") { + withSQLConf(SQLConf.PARSE_SQL_ENABLED.key -> "true") { + val expr = ParseSql(Literal("INSERT INTO t SELECT 1")) + assert(expr.isInstanceOf[CodegenFallback]) + val j = evalJson("INSERT INTO t SELECT 1") + assert(j \ "statement_identifier" === JString("INSERT")) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResultSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResultSuite.scala new file mode 100644 index 0000000000000..2dd031fb2c20d --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/catalyst/parser/ParseSqlResultSuite.scala @@ -0,0 +1,146 @@ +/* + * 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.catalyst.parser + +import org.json4s._ +import org.json4s.jackson.JsonMethods.parse + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.internal.SQLConf + +/** + * Pin Table 39 codes and contracts that goldens do not cover. + * Behavioral coverage lives in sql-tests/inputs/parse-sql.sql. + */ +class ParseSqlResultSuite extends SparkFunSuite { + + private def obj(sql: String): JObject = + parse(ParseSqlResult.fromSql(sql)).asInstanceOf[JObject] + + private def tableRefs(sql: String, field: String): Set[Seq[String]] = + obj(sql) \ field match { + case JNothing => Set.empty + case JArray(arr) => arr.map { + case JArray(parts) => parts.map(_.asInstanceOf[JString].s) + case other => fail(s"unexpected $field entry: $other") + }.toSet + case other => fail(s"unexpected $field: $other") + } + + private def sourceTableRefs(sql: String): Set[Seq[String]] = + tableRefs(sql, "source_table_references") + + private def targetTableRefs(sql: String): Set[Seq[String]] = + tableRefs(sql, "target_table_references") + + test("Table 39 standard and Spark code pairs are pinned") { + assert(SqlStatementCodes.Select.statementCode === 21) + assert(SqlStatementCodes.Insert.statementCode === 50) + assert(SqlStatementCodes.DeleteWhere.statementCode === 19) + assert(SqlStatementCodes.UpdateWhere.statementCode === 82) + assert(SqlStatementCodes.Merge.statementCode === 128) + assert(SqlStatementCodes.CreateTable.statementCode === 77) + assert(SqlStatementCodes.CreateView.statementCode === 84) + assert(SqlStatementCodes.DropTable.statementCode === 32) + assert(SqlStatementCodes.AlterTable.statementCode === 4) + assert(SqlStatementCodes.TruncateTable.statementCode === 139) + assert(SqlStatementCodes.Unrecognized.statementCode === 0) + assert(SqlStatementCodes.CacheTable.statementCode < 0) + assert(SqlStatementCodes.BeginEnd.statementCode === -22) + assert(SqlStatementCodes.Explain.statementCode === -23) + assert(SqlStatementCodes.Set.statementCode === -24) + assert(SqlStatementCodes.CreateMetricViewStmt.statementCode === -37) + } + + test("TABLE and VALUES classify as SELECT") { + val table = obj("TABLE t") + assert(table \ "statement_identifier" === JString("SELECT")) + assert(table \ "statement_code" === JInt(21)) + assert(sourceTableRefs("TABLE t") === Set(Seq("t"))) + + // Eager inlining must not flip VALUES between SELECT and Unrecognized. + Seq(true, false).foreach { eager => + SQLConf.withExistingConf(new SQLConf) { + SQLConf.get.setConf(SQLConf.EAGER_EVAL_OF_UNRESOLVED_INLINE_TABLE_ENABLED, eager) + val values = obj("VALUES (1), (2)") + assert(values \ "statement_identifier" === JString("SELECT"), + s"eager=$eager") + assert(values \ "statement_code" === JInt(21), s"eager=$eager") + } + } + } + + test("CREATE FUNCTION and DECLARE VARIABLE are not table references") { + assert(sourceTableRefs("CREATE FUNCTION f AS 'x' USING JAR 'y.jar'").isEmpty) + assert(targetTableRefs("CREATE FUNCTION f AS 'x' USING JAR 'y.jar'").isEmpty) + assert(sourceTableRefs("DECLARE VARIABLE x INT").isEmpty) + assert(targetTableRefs("DECLARE VARIABLE x INT").isEmpty) + // Contrast: CREATE VIEW still reports the view target and query source. + assert(targetTableRefs("CREATE VIEW v AS SELECT 1 AS a") === Set(Seq("v"))) + assert(sourceTableRefs("CREATE VIEW v AS SELECT 1 AS a").isEmpty) + } + + test("DML and DDL split target and source table references") { + assert(targetTableRefs("INSERT INTO t SELECT 1") === Set(Seq("t"))) + assert(sourceTableRefs("INSERT INTO t SELECT 1").isEmpty) + + assert(targetTableRefs("DELETE FROM t WHERE a = 1") === Set(Seq("t"))) + assert(sourceTableRefs("DELETE FROM t WHERE a = 1").isEmpty) + + assert(targetTableRefs("UPDATE t SET a = 1 WHERE b = 2") === Set(Seq("t"))) + assert(sourceTableRefs("UPDATE t SET a = 1 WHERE b = 2").isEmpty) + + assert(targetTableRefs( + "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE") === Set(Seq("t"))) + assert(sourceTableRefs( + "MERGE INTO t USING s ON t.id = s.id WHEN MATCHED THEN DELETE") === Set(Seq("s"))) + + assert(targetTableRefs("CREATE TABLE x AS SELECT a FROM src") === Set(Seq("x"))) + assert(sourceTableRefs("CREATE TABLE x AS SELECT a FROM src") === Set(Seq("src"))) + + assert(targetTableRefs("DROP TABLE t") === Set(Seq("t"))) + assert(sourceTableRefs("DROP TABLE t").isEmpty) + } + + test("CTE aliases only shadow references within their own scope") { + // The inner CTE named real_t must not hide the outer real table real_t. + assert(sourceTableRefs( + "SELECT * FROM real_t WHERE EXISTS (" + + "WITH real_t AS (SELECT * FROM inner_base) SELECT * FROM real_t)") === + Set(Seq("real_t"), Seq("inner_base"))) + // A definition sees only preceding aliases, so b here is the real table. + assert(sourceTableRefs( + "WITH a AS (SELECT * FROM b), b AS (SELECT 1 AS x) SELECT * FROM a") === + Set(Seq("b"))) + } + + test("positional markers inside BEGIN END are counted once") { + val j = obj("BEGIN SELECT * FROM t WHERE a = ?; END") + assert(j \ "parse_success" === JBool(true)) + assert(j \ "statement_identifier" === JString("BEGIN END")) + assert(j \ "parameter_markers" \ "unnamed_count" === JInt(1)) + assert(j \ "parameter_markers" \ "named" === JNothing) + assert(sourceTableRefs("BEGIN SELECT * FROM t WHERE a = ?; END") === Set(Seq("t"))) + } + + test("syntax error returns STANDARD error JSON without throwing") { + val j = obj("SELEC FROM t") + assert(j \ "parse_success" === JBool(false)) + assert(j \ "error" \ "errorClass" === JString("PARSE_SYNTAX_ERROR")) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionBinder.scala b/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionBinder.scala index 2f79876d841d8..71c44cfd05b11 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionBinder.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionBinder.scala @@ -22,6 +22,9 @@ import org.apache.spark.{sql, SparkFunSuite} /** * Overrides [[spark]] to provide a [[SparkSession classic.SparkSession]] */ -trait SparkSessionBinder extends sql.SparkSessionBinder { self: SparkFunSuite => +trait SparkSessionBinder + extends sql.SparkSessionBinder + with SparkSessionProvider { self: SparkFunSuite => + override protected def spark: SparkSession = super.spark.asInstanceOf[SparkSession] } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionProvider.scala b/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionProvider.scala index 77de0db4bf68b..288c1fe879581 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionProvider.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/classic/SparkSessionProvider.scala @@ -21,4 +21,6 @@ import org.apache.spark.sql trait SparkSessionProvider extends sql.SparkSessionProvider { override protected def spark: SparkSession + + override protected def sql(query: String): DataFrame = spark.sql(query) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollatedFilterPushDownToParquetSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollatedFilterPushDownToParquetSuite.scala index 19e60f6c6276a..bdbbf8057cf18 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollatedFilterPushDownToParquetSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollatedFilterPushDownToParquetSuite.scala @@ -31,6 +31,7 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf} import org.apache.spark.sql.sources.{EqualTo, Filter, IsNotNull} import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.tags.ExtendedSQLTest abstract class CollatedFilterPushDownToParquetSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { @@ -248,6 +249,7 @@ class CollatedFilterPushDownToParquetV1Suite extends CollatedFilterPushDownToPar } } +@ExtendedSQLTest class CollatedFilterPushDownToParquetV2Suite extends CollatedFilterPushDownToParquetSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala index dd22f647ab31d..ec5be2c602613 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala @@ -34,7 +34,6 @@ import org.apache.spark.util.Utils * collations */ class CollationExpressionWalkerSuite extends SharedSparkSession { - import testImplicits._ // Trait to distinguish different cases for generation sealed trait CollationType @@ -381,9 +380,13 @@ class CollationExpressionWalkerSuite extends SharedSparkSession { "sha1", "unbase64", "base64", + "to_base32", + "from_base32", "sha2", "sha", "crc32", + "xxh3_64", + "xxh3_128", "ascii", "time_trunc", // The result/sketch embeds the original item value, which now preserves the @@ -391,7 +394,11 @@ class CollationExpressionWalkerSuite extends SharedSparkSession { "approx_top_k", "approx_top_k_accumulate", "approx_top_k_combine", - "approx_top_k_estimate" + "approx_top_k_estimate", + // The variant object embeds the key string, which preserves the input case for collated + // strings, so the result is not comparable across collations. + "variant_from_arrays", + "variant_from_entries" ) logInfo("Total number of expression: " + expressionCounter) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala index 045748ad7cb13..76a28e1af51e2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationSuite.scala @@ -41,7 +41,9 @@ import org.apache.spark.sql.execution.joins._ import org.apache.spark.sql.functions.col import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, Metadata, MetadataBuilder, StringType, StructField, StructType} +import org.apache.spark.tags.ExtendedSQLTest +@ExtendedSQLTest class CollationSuite extends DatasourceV2SQLBase with AdaptiveSparkPlanHelper { protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName diff --git a/sql/core/src/test/scala/org/apache/spark/sql/collation/DefaultCollationTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/collation/DefaultCollationTestSuite.scala index 82bb616480233..1c0c0c965dac9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/collation/DefaultCollationTestSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/collation/DefaultCollationTestSuite.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAM import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{BooleanType, StringType, StructType} +import org.apache.spark.tags.ExtendedSQLTest abstract class DefaultCollationTestSuite extends SharedSparkSession { @@ -1935,6 +1936,7 @@ abstract class DefaultCollationTestSuiteV2 } } +@ExtendedSQLTest class DefaultCollationStringTestSuiteV1 extends DefaultCollationTestSuiteV1 { override protected def testDataType(testName: String)(testFn: String => Unit): Unit = { test(s"$testName [STRING]") { @@ -1943,6 +1945,7 @@ class DefaultCollationStringTestSuiteV1 extends DefaultCollationTestSuiteV1 { } } +@ExtendedSQLTest class DefaultCollationStringTestSuiteV2 extends DefaultCollationTestSuiteV2 { override protected def testDataType(testName: String)(testFn: String => Unit): Unit = { test(s"$testName [STRING]") { @@ -1974,6 +1977,7 @@ class DefaultCollationStringTestSuiteV2 extends DefaultCollationTestSuiteV2 { } } +@ExtendedSQLTest class DefaultCollationCharVarcharTestSuiteV1 extends DefaultCollationTestSuiteV1 { override protected def excluded: Seq[String] = super.excluded ++ stringTestNames ++ stringTestNamesV1 @@ -1988,6 +1992,7 @@ class DefaultCollationCharVarcharTestSuiteV1 extends DefaultCollationTestSuiteV1 } } +@ExtendedSQLTest class DefaultCollationCharVarcharTestSuiteV2 extends DefaultCollationTestSuiteV2 { override protected def excluded: Seq[String] = super.excluded ++ stringTestNames diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/AppendDataTransactionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/AppendDataTransactionSuite.scala index aef9c65550fc4..e48d99061671b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/AppendDataTransactionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/AppendDataTransactionSuite.scala @@ -19,14 +19,70 @@ package org.apache.spark.sql.connector import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.Row -import org.apache.spark.sql.connector.catalog.{Aborted, Committed} +import org.apache.spark.sql.connector.catalog.{ + Aborted, + Committed, + TableContext, + TableWritePrivilege, + TimeTravel, + Txn, + TxnTableCatalog} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode import org.apache.spark.sql.sources +import org.apache.spark.sql.util.CaseInsensitiveStringMap class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { + private val targetLoadOption = "targetLoadOption" + private val targetLoadValue = "loadValue" + private val targetWriteOption = "targetWriteOption" + private val targetWriteValue = "writeValue" + private val targetOptionsClause = + s"WITH (`$targetLoadOption` = '$targetLoadValue', " + + s"`$targetWriteOption` = '$targetWriteValue')" + + private def assertTargetLoadAndWriteOptions( + txn: Txn, + expectedPrivileges: java.util.Set[TableWritePrivilege], + minTargetLoads: Int = 1): Unit = { + val targetLoads = txn.catalog.loadTableCalls.filter { + case (context, _) => context.writePrivileges() == expectedPrivileges + } + assert(targetLoads.size >= minTargetLoads, + s"expected at least $minTargetLoads target loads with write options") + targetLoads.foreach { case (context, options) => + assert(context.writePrivileges() === expectedPrivileges) + assert(options.get(targetLoadOption) === targetLoadValue) + assert(options.asCaseSensitiveMap().containsKey(targetLoadOption)) + assert(options.get(targetWriteOption) === null) + assert(options.size() === 1) + } + + assert(table.lastWriteInfo != null, "the V2 table did not receive LogicalWriteInfo") + assert(table.lastWriteInfo.options().get(targetLoadOption) === targetLoadValue) + assert(table.lastWriteInfo.options().asCaseSensitiveMap().containsKey(targetLoadOption)) + assert(table.lastWriteInfo.options().get(targetWriteOption) === targetWriteValue) + } + + test("SPARK-58389: transaction catalog honors time travel context") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" }""") + val pinnedVersion = catalog.loadTable(ident).version() + catalog.pinTable(ident, "pinned") + append("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 2, "salary": 200, "dep": "software" }""") + assert(catalog.loadTable(ident).version() !== pinnedVersion) + + val txnCatalog = new TxnTableCatalog(catalog) + val context = new TableContext( + new TimeTravel.AsOfVersion("pinned"), java.util.Set.of[TableWritePrivilege]()) + val loaded = txnCatalog.loadTable(ident, context, CaseInsensitiveStringMap.empty()) + + assert(loaded.version() === pinnedVersion) + } + test("writeTo append with transactional checks") { // create table with initial data createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", @@ -35,14 +91,17 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { |""".stripMargin) // create a source on top of itself that will be fully resolved and analyzed - val sourceDF = spark.table(tableNameAsString) + val sourceDF = spark.read.option("customReadOption", "customValue").table(tableNameAsString) .where("pk == 1") .select(col("pk") + 10 as "pk", col("salary"), col("dep")) sourceDF.queryExecution.assertAnalyzed() // append data using the DataFrame API val (txn, txnTables) = executeTransaction { - sourceDF.writeTo(tableNameAsString).append() + sourceDF.writeTo(tableNameAsString) + .option(targetLoadOption, targetLoadValue) + .option(targetWriteOption, targetWriteValue) + .append() } // check txn was properly committed and closed @@ -50,6 +109,22 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { assert(txn.isClosed) assert(txnTables.size === 1) assert(table.version() === "2") + val sourceLoads = txn.catalog.loadTableCalls.filter { + case (context, _) => context.writePrivileges().isEmpty + } + assert(sourceLoads.nonEmpty, "transaction re-resolution did not reload the source") + sourceLoads.foreach { case (context, options) => + assert(context.timeTravel().isEmpty) + assert(context.writePrivileges().isEmpty) + assert(options.isEmpty) + } + assertTargetLoadAndWriteOptions(txn, java.util.Set.of(TableWritePrivilege.INSERT)) + txn.catalog.loadTableCalls.filter { + case (context, _) => !context.writePrivileges().isEmpty + }.foreach { case (_, options) => + assert(options.get("customReadOption") === null) + } + assert(table.lastWriteInfo.options().get("customReadOption") === null) // check the source scan was tracked via the transaction catalog val targetTxnTable = txnTables(tableNameAsString) @@ -77,7 +152,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { // SQL INSERT INTO using VALUES val (txn, txnTables) = executeTransaction { - sql(s"INSERT INTO $tableNameAsString VALUES (3, 300, 'hr'), (4, 400, 'finance')") + sql(s"INSERT INTO $tableNameAsString $targetOptionsClause " + + "VALUES (3, 300, 'hr'), (4, 400, 'finance')") } // check txn was properly committed and closed @@ -87,6 +163,7 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { // VALUES literal - No catalog tables were scanned assert(txnTables.isEmpty) + assertTargetLoadAndWriteOptions(txn, java.util.Set.of(TableWritePrivilege.INSERT)) // check data was inserted correctly checkAnswer( @@ -109,12 +186,12 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { val insertOverwrite = if (isDynamic) { // OverwritePartitionsDynamic - s"""INSERT OVERWRITE $tableNameAsString + s"""INSERT OVERWRITE $tableNameAsString $targetOptionsClause |SELECT pk + 10, salary, dep FROM $tableNameAsString WHERE dep = 'hr' |""".stripMargin } else { // OverwriteByExpression - s"""INSERT OVERWRITE $tableNameAsString + s"""INSERT OVERWRITE $tableNameAsString $targetOptionsClause |PARTITION (dep = 'hr') |SELECT pk + 10, salary FROM $tableNameAsString WHERE dep = 'hr' |""".stripMargin @@ -137,6 +214,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { case sources.EqualTo("dep", "hr") => true case _ => false }) + assertTargetLoadAndWriteOptions( + txn, java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -159,7 +238,10 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { toDF("pk", "salary", "dep") val (txn, txnTables) = executeTransaction { - sourceDF.writeTo(tableNameAsString).overwrite(col("dep") === "hr") + sourceDF.writeTo(tableNameAsString) + .option(targetLoadOption, targetLoadValue) + .option(targetWriteOption, targetWriteValue) + .overwrite(col("dep") === "hr") } assert(txn.currentState === Committed) @@ -168,6 +250,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { // literal DataFrame source - no catalog tables were scanned assert(txnTables.isEmpty) + assertTargetLoadAndWriteOptions( + txn, java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -190,7 +274,10 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { toDF("pk", "salary", "dep") val (txn, txnTables) = executeTransaction { - sourceDF.writeTo(tableNameAsString).overwritePartitions() + sourceDF.writeTo(tableNameAsString) + .option(targetLoadOption, targetLoadValue) + .option(targetWriteOption, targetWriteValue) + .overwritePartitions() } assert(txn.currentState === Committed) @@ -199,6 +286,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { // literal DataFrame source - no catalog tables were scanned assert(txnTables.isEmpty) + assertTargetLoadAndWriteOptions( + txn, java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -412,7 +501,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { sql(s"INSERT INTO $sourceNameAsString VALUES (3, 300, 'hr', true), (4, 400, 'software', false)") val (txn, txnTables) = executeTransaction { - sql(s"INSERT WITH SCHEMA EVOLUTION INTO $tableNameAsString SELECT * FROM $sourceNameAsString") + sql(s"INSERT WITH SCHEMA EVOLUTION INTO $tableNameAsString $targetOptionsClause " + + s"SELECT * FROM $sourceNameAsString") } assert(txn.currentState === Committed) @@ -420,6 +510,8 @@ class AppendDataTransactionSuite extends RowLevelOperationSuiteBase { // the new column must be visible in the committed delegate's schema assert(table.schema.fieldNames.toSeq === Seq("pk", "salary", "dep", "active")) + assertTargetLoadAndWriteOptions( + txn, java.util.Set.of(TableWritePrivilege.INSERT), minTargetLoads = 2) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..c0a1faa032635 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -0,0 +1,708 @@ +/* + * 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.connector + +import org.apache.spark.{SparkConf, SparkException} +import org.apache.spark.sql.{AnalysisException, DataFrame, Row} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GetStructField, GreaterThan, Literal, RLike} +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper +import org.apache.spark.sql.connector.catalog.{ + Column, + Identifier, + InMemoryCatalystRuntimeFilterTable, + InMemoryTable, + InMemoryTableCatalystRuntimeFilterCatalog, + TableCatalog} +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{Batch, HasPartitionKey, InputPartition, PartitionReaderFactory, Scan, SupportsRuntimeV2Filtering} +import org.apache.spark.sql.execution.{FilterExec, ScalarSubquery => ExecScalarSubquery} +import org.apache.spark.sql.execution.ExplainUtils.stripAQEPlan +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, DataSourceV2Strategy, PushDownUtils} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StringType, StructType} + +/** + * Tests for scans that implement + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], + * where runtime filters are pushed once as Catalyst expressions instead of connector + * predicates. + */ +class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { + + protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + protected val catalogName = "testcatalystruntimefilter" + + override def sparkConf: SparkConf = super.sparkConf + .set(s"spark.sql.catalog.$catalogName", + classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) + + private def withDPPConf(f: => Unit): Unit = { + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10")(f) + } + + test("scalar subquery on partition column -> pushed as Catalyst expression") { + val tbl = s"$catalogName.tbl1" + val dim = s"$catalogName.dim1" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3)) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3))) + // `part` is not declared fully pushed, so Spark still evaluates the filter after the scan. + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + + // The answer alone would be right even without pruning, since that post-scan filter drops + // the extra rows. + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") + } + } + + test("predicate on fully pushed filter attributes -> not evaluated after the scan") { + val tbl = s"$catalogName.tbl_fully_pushed" + val dim = s"$catalogName.dim_fully_pushed" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')") + // Matching and nonmatching partitions: the scan must prune nonmatching ones itself + // because Spark drops the post-scan FilterExec for fully pushed attributes. + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3)) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3))) + assertScalarSubqueryEvaluatedAfterScan(df, expected = false) + } + } + + test("non-deterministic predicate on fully pushed attributes -> evaluated after the scan") { + val tbl = s"$catalogName.tbl_nondeterministic" + val dim = s"$catalogName.dim_nondeterministic" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + // A non-deterministic filter is never pushed, so it must keep its post-scan FilterExec even + // though it only references a fully pushed attribute. Dropping it there would leave nothing + // to evaluate it and the scan would return the nonmatching partitions too. + val df = sql( + s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim) OR rand() < 0.5") + // The row in the matching partition always qualifies, the others qualify at random. + assert(df.collect().contains(Row(3, 3))) + + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + assertScalarSubqueryRuntimeFilters(df, expectedCount = 0) + assertPushedCatalystPredicates(df, 0) + } + } + + test("predicate on partly fully pushed filter attributes -> evaluated after the scan") { + val tbl = s"$catalogName.tbl_partly_pushed" + val dim = s"$catalogName.dim_partly_pushed" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + + "PARTITIONED BY (p1, p2) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'p1')") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, 1, 2)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + // The predicate also references p2, which is not declared fully pushed, so it is not + // considered fully pushed and Spark keeps evaluating it after the scan. + val df = sql(s"SELECT * FROM $tbl WHERE p1 + p2 = (SELECT max(val) FROM $dim)") + checkAnswer(df, (0 until 5).map(i => Row(i, 1, 2))) + + assertScalarSubqueryRuntimeFilters(df) + val p1 = AttributeReference("p1", IntegerType, nullable = false)() + val p2 = AttributeReference("p2", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(Add(p1, p2), Literal(3))) + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + } + } + + test("arithmetic around a scalar subquery -> subquery literalized, expression pushed intact") { + val tbl = s"$catalogName.tbl2" + val dim = s"$catalogName.dim2" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2)") + + // The scalar subquery is literalized on the way to the scan, and the arithmetic around it + // survives. This one would also translate to a V2 predicate (`part > 3`, after the + // literalized operands are folded); see the next test for one that would not. + val df = sql(s"SELECT * FROM $tbl WHERE part > (SELECT max(val) FROM $dim) + 1") + checkAnswer(df, Row(4, 4)) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual( + df, GreaterThan(part, Add(Literal(2), Literal(1)))) + } + } + + test("filter with no V2 translation -> pushed instead of dropped") { + val tbl = s"$catalogName.tbl_untranslatable" + val dim = s"$catalogName.dim_untranslatable" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part STRING) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, '$i')") + } + sql(s"CREATE TABLE $dim (val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES ('3')") + + // `V2ExpressionBuilder` has no RLike case, so this filter has no V2 predicate at all and + // the V2 interfaces would never see it. The pattern is still a subquery when the optimizer + // runs, so nothing rewrites the RLIKE into a translatable Contains beforehand either. + val df = sql(s"SELECT * FROM $tbl WHERE part RLIKE (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, "3")) + + assertScalarSubqueryRuntimeFilters(df) + val runtimeFilter = collectBatchScan(df).runtimeFilters.head + assert(DataSourceV2Strategy.translateScalarSubqueryFilterV2(runtimeFilter).isEmpty, + s"Expected no V2 translation for $runtimeFilter") + val part = AttributeReference("part", StringType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, RLike(part, Literal("3"))) + } + } + + test("DPP filter -> pushed as InSubqueryExec expression") { + val fact = s"$catalogName.fact3" + val dim = s"$catalogName.dim3" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (dim_id INT, dim_val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2, 'two')") + + withDPPConf { + val df = sql( + s"""SELECT f.id, f.part FROM $fact f JOIN $dim d + |ON f.part = d.dim_id WHERE d.dim_val = 'two'""".stripMargin) + checkAnswer(df, Row(2, 2)) + + assertDPPRuntimeFilters(df) + val dppPredicate = collectBatchScan(df).runtimeFilters.collectFirst { + case DynamicPruningExpression(e) => e + }.get + assertPushedCatalystPredicatesEqual(df, dppPredicate) + } + } + } + + test("DPP filter on a nested partition field -> pushed with the nested access intact") { + val fact = s"$catalogName.fact_nested_dpp" + val dim = s"$catalogName.dim_nested_dpp" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, s STRUCT<part: INT>) USING $v2Source " + + "PARTITIONED BY (s.part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($i, named_struct('part', $i))") + } + sql(s"CREATE TABLE $dim (dim_id INT, dim_val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2, 'two')") + + withDPPConf { + val df = sql( + s"""SELECT f.id FROM $fact f JOIN $dim d + |ON f.s.part = d.dim_id WHERE d.dim_val = 'two'""".stripMargin) + checkAnswer(df, Row(2)) + + assertDPPRuntimeFilters(df) + // The scan only reports `s`, the struct holding the partition field, but the predicate it + // receives keeps the nested access, so it can still tell which partition to keep. + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + assert(pushed.head.exists(_.isInstanceOf[GetStructField]), + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") + } + } + } + + test("scan implementing both runtime filtering interfaces -> rejected") { + val tbl = s"$catalogName.tbl_both_interfaces" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + sql(s"INSERT INTO $tbl VALUES (1, 1)") + + val scanRelation = sql(s"SELECT * FROM $tbl").queryExecution.optimizedPlan.collectFirst { + case r: DataSourceV2ScanRelation => r + }.getOrElse(fail("Expected a DataSourceV2ScanRelation")) + + // every runtime filtering path starts at runtimeFilterAttrs + val e = intercept[SparkException] { + scanRelation.copy(scan = new BothRuntimeFilteringInterfacesScan).runtimeFilterAttrs + } + assert(e.getMessage.contains("A scan must not implement both SupportsRuntimeV2Filtering " + + "and SupportsRuntimeCatalystFiltering")) + } + } + + test("filter on column outside filterAttributes -> not pushed, even if declared fully pushed") { + val tbl = s"$catalogName.tbl4" + val dim = s"$catalogName.dim4" + withTable(tbl, dim) { + // p2 is a partition column but is not declared filterable, so no runtime filter is derived + // for it. Declaring it fully pushed as well, which the interface forbids for an attribute + // that is not filterable, must not cost it the post-scan filter: nothing was pushed, so the + // scan prunes nothing and the nonmatching rows would come back. + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + + "PARTITIONED BY (p1, p2) " + + "TBLPROPERTIES('filter-attributes' = 'p1', 'fully-pushed-filter-attributes' = 'p2')") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3, 3)) + + assert(collectBatchScan(df).runtimeFilters.isEmpty, + "Expected no runtime filters for a column outside filterAttributes") + assertPushedCatalystPredicates(df, 0) + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + } + } + + test("two predicates on filter attributes -> pushed together in a single filter() call") { + val tbl = s"$catalogName.tbl_two_predicates" + val dim1 = s"$catalogName.dim_two_predicates1" + val dim2 = s"$catalogName.dim_two_predicates2" + withTable(tbl, dim1, dim2) { + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source PARTITIONED BY (p1, p2)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i, ${i * 10})") + } + sql(s"CREATE TABLE $dim1 (val INT) USING $v2Source") + sql(s"INSERT INTO $dim1 VALUES (3)") + sql(s"CREATE TABLE $dim2 (val INT) USING $v2Source") + sql(s"INSERT INTO $dim2 VALUES (30)") + + val df = sql(s"SELECT * FROM $tbl WHERE p1 = (SELECT max(val) FROM $dim1) " + + s"AND p2 = (SELECT max(val) FROM $dim2)") + checkAnswer(df, Row(3, 3, 30)) + + assertScalarSubqueryRuntimeFilters(df, expectedCount = 2) + val p1 = AttributeReference("p1", IntegerType, nullable = false)() + val p2 = AttributeReference("p2", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual( + df, EqualTo(p1, Literal(3)), EqualTo(p2, Literal(30))) + assert(getCatalystScan(df).filterCallCount === 1, + "expected both predicates pushed in a single filter() call") + } + } + + test("nested field of a filter attribute -> pushed with the nested access intact") { + val tbl = s"$catalogName.tbl_nested" + val dim = s"$catalogName.dim_nested" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, s STRUCT<tz: STRING>) USING $v2Source " + + "PARTITIONED BY (s.tz)") + for (i <- 0 until 3) { + sql(s"INSERT INTO $tbl VALUES ($i, named_struct('tz', 'tz$i'))") + } + sql(s"CREATE TABLE $dim (val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES ('tz1')") + + // The scan declares the top-level struct column `s` as its filter attribute, so the + // predicate qualifies for pushdown even though it reaches into `s.tz`. Matching the nested + // access against the partition layout is left to the scan, which this fixture does. + val df = sql(s"SELECT * FROM $tbl WHERE s.tz = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(1, Row("tz1"))) + + assertScalarSubqueryRuntimeFilters(df) + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + val nestedAccesses = pushed.head.collect { case g: GetStructField => g } + assert(nestedAccesses.size === 1, + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + assert(nestedAccesses.head.childSchema.fieldNames.contains("tz")) + } + } + + test("dotted top-level and nested partition columns -> bound to the correct partition slot") { + val tbl = s"$catalogName.tbl_dotted_collision" + val dim = s"$catalogName.dim_dotted_collision" + withTable(tbl, dim) { + // Two partition columns whose dotted names collide: the quoted top-level column `x.y` and + // the nested field `x`.`y`. They carry different values in each row, so a predicate bound + // to the wrong slot would prune the wrong partitions. `x.y` is 3 exactly where `x`.`y` is + // 30, so binding a filter on the nested field to the top-level slot would find nothing. + sql(s"CREATE TABLE $tbl (id INT, `x.y` INT, x STRUCT<y: INT>) USING $v2Source " + + "PARTITIONED BY (`x.y`, x.y)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i, named_struct('y', ${i * 10}))") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (30)") + + // Alias the table so `f.x.y` unambiguously reads the nested field, not the column `x.y`. + val df = sql(s"SELECT * FROM $tbl f WHERE f.x.y = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3, Row(30))) + + assertScalarSubqueryRuntimeFilters(df) + // The pushed predicate keeps the nested access, and the scan prunes to the single partition + // whose nested `x`.`y` is 30 rather than binding to the colliding top-level `x.y` slot. + val pushed = getPushedCatalystPredicates(df) + assert(pushed.size === 1, s"expected a single pushed predicate, got $pushed") + assert(pushed.head.exists(_.isInstanceOf[GetStructField]), + s"expected the pushed predicate to keep the nested access, got ${pushed.head}") + val batchScan = collectBatchScan(df) + assert(batchScan.inputPartitions.size === 5) + assert(batchScan.filteredPartitions.flatten.size === 1, + s"expected 1 partition after pruning, got ${batchScan.filteredPartitions.flatten.size}") + } + } + + test("filterAttributes that is not a top-level scan attribute") { + val tbl = s"$catalogName.tbl_unresolvable_attr" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT, s STRUCT<tz: STRING>) USING $v2Source " + + "PARTITIONED BY (part)") + sql(s"INSERT INTO $tbl VALUES (1, 1, named_struct('tz', 'a'))") + + val scanRelation = sql(s"SELECT * FROM $tbl").queryExecution.optimizedPlan.collectFirst { + case r: DataSourceV2ScanRelation => r + }.getOrElse(fail("Expected a DataSourceV2ScanRelation")) + + // An attribute the read schema does not carry, such as one pruned out of the projection. + val missing = intercept[AnalysisException] { + scanRelation.copy(scan = new MissingFilterAttributeScan).runtimeFilterAttrs + } + checkError( + exception = missing, + condition = "_LEGACY_ERROR_TEMP_1137", + parameters = Map("name" -> "missing", "outputStr" -> "id,part,s")) + + // A nested reference is rejected up front, since `filterAttributes()` must return + // top-level read-schema attributes. This holds even over an int column that could never + // carry a nested field. + val nested = intercept[SparkException] { + scanRelation.copy(scan = new NestedFilterAttributeScan).runtimeFilterAttrs + } + assert(nested.getMessage.contains("must be a top-level attribute"), + s"expected the nested reference to be rejected, got ${nested.getMessage}") + + // Over a struct it is rejected the same way, rather than widening to the struct column: + // accepting `s.tz` would make runtime filters over every field of `s` eligible, not just + // `s.tz`. + val struct = intercept[SparkException] { + scanRelation.copy(scan = new StructNestedFilterAttributeScan).runtimeFilterAttrs + } + assert(struct.getMessage.contains("must be a top-level attribute"), + s"expected the nested struct reference to be rejected, got ${struct.getMessage}") + } + } + + test("no runtime filter -> filter() is never called") { + val tbl = s"$catalogName.tbl5" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + + val df = sql(s"SELECT * FROM $tbl WHERE part = 3") + checkAnswer(df, Row(3, 3)) + + assert(collectBatchScan(df).runtimeFilters.isEmpty) + assertPushedCatalystPredicates(df, 0) + } + } + + test("ALTER TABLE keeps the Catalyst runtime-filter table type") { + val tbl = s"$catalogName.tbl_alter" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + sql(s"ALTER TABLE $tbl ADD COLUMNS (extra INT)") + + val table = spark.sessionState.catalogManager.catalog(catalogName) + .asInstanceOf[TableCatalog] + .loadTable(Identifier.of(Array.empty, "tbl_alter")) + assert(table.isInstanceOf[InMemoryCatalystRuntimeFilterTable], + s"ALTER TABLE reconstructed ${table.getClass.getName}") + + sql(s"INSERT INTO $tbl VALUES (0, 0, 10), (1, 1, 11)") + checkAnswer(sql(s"SELECT id, part FROM $tbl WHERE extra = 11"), Row(1, 1)) + } + } + + /** + * While SPJ is active the scan's partitioning has to survive runtime filtering, so the + * post-filter partitions still line up with the other side of the join: splits may be pruned, + * but the source may not drop a partition key, invent one, or grow a key's split count. + */ + test("data source that breaks the partitioning it reported -> rejected") { + val partAttr = AttributeReference("part", IntegerType)() + val table = new InMemoryTable("t", Array(Column.create("part", IntegerType)), + Array.empty[Transform], java.util.Collections.emptyMap[String, String]) + val partitioning = KeyedPartitioning( + Seq(partAttr), + Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))), + isGrouped = false) + + def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = { + val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter) + PushDownUtils.replanWithRuntimeFilters(scan, Seq(EqualTo(partAttr, Literal(1))), table, + Seq(partAttr), partitioning, originalPartitions = Seq.empty) + } + + val keyDropped = intercept[SparkException](replanAfterFiltering(Seq(new InputPartition {}))) + assert(keyDropped.getMessage.contains("must have preserved the original partitioning")) + + val keyInvented = intercept[SparkException](replanAfterFiltering(Seq(KeyedInputPartition(99)))) + assert(keyInvented.getMessage.contains("must not report new partition keys")) + + val splitsGrown = intercept[SparkException] { + replanAfterFiltering(Seq(KeyedInputPartition(1), KeyedInputPartition(1))) + } + assert(splitsGrown.getMessage.contains("must not report new partitions for a given key")) + } + + // --------------------------------------------------------------------------- + // Helper methods + // --------------------------------------------------------------------------- + + private def assertDPPRuntimeFilters( + df: DataFrame, expectedCount: Int = 1): Unit = { + val batchScan = collectBatchScan(df) + val dppFilters = batchScan.runtimeFilters.collect { + case d: DynamicPruningExpression => d + } + assert(dppFilters.size === expectedCount, + s"Expected $expectedCount DynamicPruningExpression(s) " + + s"in runtimeFilters, got ${dppFilters.size}") + } + + private def assertScalarSubqueryRuntimeFilters( + df: DataFrame, expectedCount: Int = 1): Unit = { + val batchScan = collectBatchScan(df) + val scalarFilters = batchScan.runtimeFilters.collect { + case f if !f.isInstanceOf[DynamicPruning] => f + } + val dppFilters = batchScan.runtimeFilters.collect { + case d: DynamicPruning => d + } + assert(scalarFilters.size === expectedCount, + s"Expected $expectedCount scalar subquery runtime filter(s), " + + s"got ${scalarFilters.size}") + assert(dppFilters.isEmpty, + "Expected non-DPP runtime filters (scalar subquery)") + } + + /** + * Checks whether a scalar subquery runtime filter is still evaluated by a [[FilterExec]] above + * the scan. Filters that only reference `fullyPushedFilterAttributes` are dropped from it. + */ + private def assertScalarSubqueryEvaluatedAfterScan( + df: DataFrame, + expected: Boolean): Unit = { + val postScanConditions = stripAQEPlan(df.queryExecution.executedPlan).collect { + case f: FilterExec => f.condition + } + val evaluated = postScanConditions.exists(_.exists(_.isInstanceOf[ExecScalarSubquery])) + assert(evaluated === expected, + s"Expected scalar subquery evaluated after scan to be $expected, " + + s"post-scan filter conditions: $postScanConditions") + } + + private def collectBatchScan(df: DataFrame): BatchScanExec = { + stripAQEPlan(df.queryExecution.executedPlan).collectFirst { + case b: BatchScanExec => b + }.getOrElse(fail("Expected BatchScanExec in plan")) + } + + private type CatalystScan = + InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan + + private def getCatalystScan(df: DataFrame): CatalystScan = { + collectBatchScan(df).scan match { + case s: CatalystScan => s + case other => fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other") + } + } + + private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { + getCatalystScan(df).pushedCatalystPredicates + } + + private def assertPushedCatalystPredicates(df: DataFrame, expected: Int): Unit = { + val preds = getPushedCatalystPredicates(df) + assert(preds.size === expected, + s"Expected $expected pushed Catalyst runtime predicate(s), got ${preds.size}: $preds") + } + + /** + * Binds [[AttributeReference]]s in `expected` to the scan output (by name) and checks that the + * pushed Catalyst runtime predicates match exactly via [[Expression.semanticEquals]]. + */ + private def assertPushedCatalystPredicatesEqual( + df: DataFrame, + expected: Expression*): Unit = { + val batchScan = collectBatchScan(df) + val actual = getPushedCatalystPredicates(df) + val normalizedExpected = expected.map(bindToScanOutput(_, batchScan.output)) + assert(actual.size === normalizedExpected.size, + s"Expected ${normalizedExpected.size} pushed Catalyst predicate(s), " + + s"got ${actual.size}: $actual") + actual.zip(normalizedExpected).foreach { case (a, e) => + assert(a.semanticEquals(e), + s"Pushed Catalyst predicate mismatch.\nExpected: $e\nActual: $a") + } + } + + private def bindToScanOutput( + expr: Expression, + output: Seq[AttributeReference]): Expression = { + val resolver = SQLConf.get.resolver + expr.transformUp { + case a: AttributeReference => + output.find(o => resolver(o.name, a.name)) + .map(_.withNullability(a.nullable).withQualifier(a.qualifier)) + .getOrElse(a) + } + } +} + +/** A scan violating the rule that only one runtime filtering interface may be implemented. */ +private class BothRuntimeFilteringInterfacesScan + extends Scan with SupportsRuntimeV2Filtering with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference("part")) + + override def filter(predicates: Array[Predicate]): Unit = {} + + override def filter(expressions: Array[Expression]): Unit = {} +} + +/** A scan declaring a filter attribute the read schema does not carry. */ +private class MissingFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference("missing")) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +/** + * A scan breaking the rule that a filter attribute must be a top level read schema column: it + * reports `part.nested` over the int column `part`, so it is rejected as a nested reference. + */ +private class NestedFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def filterAttributes(): Array[NamedReference] = + Array(FieldReference(Seq("part", "nested"))) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +/** + * A scan breaking the same rule over a struct column: it reports `s.tz` where `s` is a struct. + * The nested reference is rejected rather than widening to the struct column `s`. + */ +private class StructNestedFilterAttributeScan extends Scan with SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = + new StructType().add("s", new StructType().add("tz", StringType)) + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference(Seq("s", "tz"))) + + override def filter(expressions: Array[Expression]): Unit = {} +} + +private case class KeyedInputPartition(key: Int) extends InputPartition with HasPartitionKey { + override def partitionKey(): InternalRow = InternalRow(key) +} + +/** + * A scan reporting one set of partitions before filtering and another after, so it can break the + * requirement to preserve the partitioning it originally reported. + */ +private class PartitioningBreakingScan( + initialPartitions: Seq[InputPartition], + afterFilter: Seq[InputPartition]) + extends Scan with Batch with SupportsRuntimeCatalystFiltering { + + private var filtered = false + + override def readSchema(): StructType = new StructType().add("part", IntegerType) + + override def toBatch: Batch = this + + override def planInputPartitions(): Array[InputPartition] = + if (filtered) afterFilter.toArray else initialPartitions.toArray + + override def createReaderFactory(): PartitionReaderFactory = + throw new UnsupportedOperationException() + + override def filterAttributes(): Array[NamedReference] = Array(FieldReference("part")) + + override def filter(expressions: Array[Expression]): Unit = { + filtered = true + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSessionCatalogSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSessionCatalogSuite.scala index 97cdebe2d32df..d104601813298 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSessionCatalogSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSessionCatalogSuite.scala @@ -98,6 +98,29 @@ class DataSourceV2DataFrameSessionCatalogSuite verifyTable("t", df) } } + + test("SPARK-58389: time travel options are ignored for V1 table writes") { + withTable("t") { + sql("CREATE TABLE t(c BIGINT) USING csv") + val df = spark.range(1).toDF("c") + + df.write + .format(v2Format) + .option("versionAsOf", "1") + .insertInto("t") + verifyTable("t", df) + + // SaveMode.Ignore leaves the existing V1 table unchanged, but the write must still + // reach the V1 fallback before Spark rejects the time-travel option. + df.write + .format(v2Format) + .option("versionAsOf", "1") + .mode(SaveMode.Ignore) + .saveAsTable("t") + + verifyTable("t", df) + } + } } class InMemoryTableSessionCatalog extends TestV2SessionCatalogBase[InMemoryTable] { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala index 5c5a6210f29d8..07a214c24fbc3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala @@ -20,11 +20,11 @@ package org.apache.spark.sql.connector import org.scalatest.BeforeAndAfter import org.apache.spark.sql.{DataFrame, Row} -import org.apache.spark.sql.catalyst.expressions.{DynamicPruning, DynamicPruningExpression} +import org.apache.spark.sql.catalyst.expressions.{DynamicPruning, DynamicPruningExpression, Literal} import org.apache.spark.sql.connector.catalog.{BufferedRows, InMemoryEnhancedRuntimePartitionFilterTable, InMemoryTableEnhancedRuntimePartitionFilterCatalog} import org.apache.spark.sql.connector.expressions.PartitionFieldReference import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate -import org.apache.spark.sql.execution.ExplainUtils.stripAQEPlan +import org.apache.spark.sql.execution.{ExplainUtils, ProjectedBroadcastValueSubqueryExec} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -75,6 +75,93 @@ class DataSourceV2EnhancedRuntimePartitionFilterSuite SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10")(f) } + private def withProjectedBroadcastRuntimeFiltering( + maxRows: Int, + projectionEnabled: Boolean = true)(f: DataFrame => Unit): Unit = { + val fact = s"$catalogName.projection_fact" + val codes = "broadcast_projection_codes" + val selectors = "broadcast_projection_selectors" + + withTable(fact, codes, selectors) { + sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (part <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($part, $part)") + } + + sql(s"CREATE TABLE $codes (join_id INT, part INT) USING parquet") + sql(s"INSERT INTO $codes VALUES (1, 2), (2, 3), (3, 4)") + sql(s"CREATE TABLE $selectors (join_id INT, keep BOOLEAN) USING parquet") + sql(s"INSERT INTO $selectors VALUES (1, true), (2, true), (3, false)") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_ENABLED.key -> + projectionEnabled.toString, + SQLConf.DYNAMIC_PARTITION_PRUNING_BROADCAST_PROJECTION_MAX_ROWS.key -> + maxRows.toString, + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = sql( + s""" + |WITH selected AS ( + | SELECT /*+ BROADCAST(c) */ c.part + | FROM $codes c + | JOIN $selectors x ON c.join_id = x.join_id + | WHERE x.keep AND x.join_id > 0 + |) + |SELECT /*+ MERGE(f, s) */ f.id, f.part + |FROM $fact f + |JOIN selected s ON f.part = s.part + |ORDER BY f.id + |""".stripMargin) + f(df) + } + } + } + + test("disabled broadcast projection preserves existing no-op iterative V2 filtering") { + withProjectedBroadcastRuntimeFiltering(maxRows = 10, projectionEnabled = false) { df => + checkAnswer(df, Seq(Row(2, 2), Row(3, 3))) + assertDPPRuntimeFilters(df) + assert(collectBatchScan(df).runtimeFilters.contains( + DynamicPruningExpression(Literal.TrueLiteral))) + assertPushedPartitionPredicates(df, expectedCount = 1) + assert(getPushedPartitionPredicates(df).head.references().isEmpty) + assertScanReturnsPartitionKeys(df, Set("0", "1", "2", "3", "4")) + } + } + + test("unavailable broadcast values do not reach iterative V2 partition filtering") { + withProjectedBroadcastRuntimeFiltering(maxRows = 0) { df => + checkAnswer(df, Seq(Row(2, 2), Row(3, 3))) + val projected = ExplainUtils.collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + assert(projected.head.metrics("projectionDisabled").value === 1) + assertDPPRuntimeFilters(df) + assertPushedPartitionPredicates(df, expectedCount = 0) + assertScanReturnsPartitionKeys(df, Set("0", "1", "2", "3", "4")) + } + } + + test("projected broadcast values prune iterative V2 partitions with a safe superset") { + withProjectedBroadcastRuntimeFiltering(maxRows = 10) { df => + checkAnswer(df, Seq(Row(2, 2), Row(3, 3))) + val projected = ExplainUtils.collectWithSubqueries(df.queryExecution.executedPlan) { + case subquery: ProjectedBroadcastValueSubqueryExec => subquery + } + assert(projected.size === 1, df.queryExecution.executedPlan) + assert(projected.head.metrics("numInputRows").value === 3) + assert(projected.head.metrics("numOutputRows").value === 3) + assertDPPRuntimeFilters(df) + assertPushedPartitionPredicates(df, expectedCount = 1) + assertScanReturnsPartitionKeys(df, Set("2", "3", "4")) + } + } + // --------------------------------------------------------------------------- // PartitionPredicate IS created // --------------------------------------------------------------------------- @@ -520,7 +607,7 @@ class DataSourceV2EnhancedRuntimePartitionFilterSuite } private def collectBatchScan(df: DataFrame): BatchScanExec = { - stripAQEPlan(df.queryExecution.executedPlan).collectFirst { + ExplainUtils.collectFirst(df.queryExecution.executedPlan) { case b: BatchScanExec => b }.getOrElse(fail("Expected BatchScanExec in plan")) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala index f785cac9c0124..1490871e4440e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala @@ -17,20 +17,183 @@ package org.apache.spark.sql.connector -import org.apache.spark.sql.{AnalysisException, Row} +import java.util +import java.util.concurrent.atomic.AtomicInteger + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.{AnalysisException, DataFrame, Row} import org.apache.spark.sql.QueryTest.withQueryExecutionsCaptured +import org.apache.spark.sql.catalyst.analysis.{ + AnalysisContext, + AsOfVersion, + RelationCache, + RelationResolution, + UnresolvedRelation, + V2TableReference} import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 -import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTableCatalog} -import org.apache.spark.sql.execution.CommandResultExec +import org.apache.spark.sql.connector.catalog.{ + Identifier, + InMemoryBaseTable, + InMemoryCatalog, + InMemoryRowLevelOperationTableCatalog, + InMemoryTableCatalog, + StagedTable, + StagingTableCatalog, + Table, + TableChange, + TableInfo, + TableWritePrivilege, + TimeTravel} +import org.apache.spark.sql.connector.write.Write +import org.apache.spark.sql.execution.{CommandResultExec, QueryExecution, SparkPlan} import org.apache.spark.sql.execution.datasources.v2._ import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +class LoadCountingInMemoryCatalog extends InMemoryCatalog { + val singleArgLoads = new AtomicInteger(0) + + override def loadTable(ident: Identifier): Table = { + singleArgLoads.incrementAndGet() + super.loadTable(ident) + } +} + +class StateAwareInMemoryCatalog extends LoadCountingInMemoryCatalog { + // Include Spark's internal marker so the write-context test detects if it leaks before state + // option projection. Production catalogs should declare only raw user option keys. + override def tableStateOptionKeys(): util.Set[String] = + util.Set.of("snapshot", UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) +} + +class NullReturningInMemoryCatalog extends InMemoryCatalog { + override def createTable(ident: Identifier, tableInfo: TableInfo): Table = { + super.createTable(ident, tableInfo) + null + } +} + +class NullReturningStagingInMemoryCatalog extends InMemoryCatalog with StagingTableCatalog { + override def stageCreate(ident: Identifier, tableInfo: TableInfo): StagedTable = { + createTable(ident, tableInfo) + null + } + + override def stageReplace(ident: Identifier, tableInfo: TableInfo): StagedTable = { + dropTable(ident) + createTable(ident, tableInfo) + null + } + + override def stageCreateOrReplace(ident: Identifier, tableInfo: TableInfo): StagedTable = { + if (tableExists(ident)) { + dropTable(ident) + } + createTable(ident, tableInfo) + null + } +} class DataSourceV2OptionSuite extends DatasourceV2SQLBase { import testImplicits._ private val catalogAndNamespace = "testcat.ns1.ns2." + private def inMemoryCatalog: InMemoryCatalog = + catalog("testcat").asInstanceOf[InMemoryCatalog] + + private def withStateAwareTable( + f: (StateAwareInMemoryCatalog, String) => Unit): Unit = { + withSQLConf( + "spark.sql.catalog.statecat" -> classOf[StateAwareInMemoryCatalog].getName, + "spark.sql.catalog.statecat.copyOnLoad" -> "true") { + val tableName = "statecat.ns.table" + withTable(tableName) { + sql(s"CREATE TABLE $tableName (id bigint, data string)") + sql(s"INSERT INTO $tableName VALUES (1, 'a'), (2, 'b')") + f(catalog("statecat").asInstanceOf[StateAwareInMemoryCatalog], tableName) + } + } + } + + private def assertOnlySnapshotOptions( + catalog: StateAwareInMemoryCatalog, + expectedSnapshot: String): Unit = { + val loadOptions = catalog.loadTableCalls.map(_._2) + assert(loadOptions.nonEmpty, "expected at least one options-aware table load") + assert(loadOptions.forall { options => + options.size() == 1 && options.get("snapshot") == expectedSnapshot + }, s"expected only snapshot=$expectedSnapshot to be forwarded, got: $loadOptions") + } + + private val loadOption = "load-Option" + private val loadOptionValue = "load-value" + private val writeOption = "write-option" + private val writeOptionValue = "write-value" + + private def testWithLoadOptionAsTableState(testName: String)(f: => Unit): Unit = { + test(testName) { + withSQLConf("spark.sql.catalog.testcat.tableStateOptionKeys" -> loadOption)(f) + } + } + + private def assertTargetOptions( + options: CaseInsensitiveStringMap): org.scalatest.Assertion = { + assert(options.get(loadOption) === loadOptionValue) + assert(options.get(writeOption) === writeOptionValue) + } + + private def assertWriteLoad( + tableCatalog: InMemoryTableCatalog, + expectedPrivileges: Set[TableWritePrivilege]): Unit = { + val matchingCalls = tableCatalog.loadTableCalls.filter { + case (context, _) => context.writePrivileges() == expectedPrivileges.asJava + } + assert(matchingCalls.nonEmpty, "loadTable did not receive the expected write privileges") + matchingCalls.foreach { case (context, options) => + assert(context.writePrivileges() === expectedPrivileges.asJava) + assert(options.get(loadOption) === loadOptionValue) + assert(options.asCaseSensitiveMap().containsKey(loadOption)) + assert(options.get(writeOption) === null) + assert(options.size() === 1) + } + } + + private def assertWriteLoad(expectedPrivileges: Set[TableWritePrivilege]): Unit = { + assertWriteLoad(inMemoryCatalog, expectedPrivileges) + } + + private def inMemoryWriteOptions(write: Write): CaseInsensitiveStringMap = { + write.toBatch match { + case append: InMemoryBaseTable#Append => append.info.options + case overwrite: InMemoryBaseTable#TruncateAndAppend => overwrite.info.options + case dynamic: InMemoryBaseTable#DynamicOverwrite => dynamic.info.options + case other => fail(s"expected a V2 in-memory batch write, got ${other.getClass.getName}") + } + } + + private def collectInMemoryWriteOptions(plan: SparkPlan): Seq[CaseInsensitiveStringMap] = { + val direct = plan.collect { + case AppendDataExec(_, _, write, _, _) => inMemoryWriteOptions(write) + case OverwriteByExpressionExec(_, _, write, _, _) => inMemoryWriteOptions(write) + case OverwritePartitionsDynamicExec(_, _, write, _, _) => inMemoryWriteOptions(write) + } + val commandResults = plan.collect { + case CommandResultExec(_, commandPhysicalPlan, _) => + collectInMemoryWriteOptions(commandPhysicalPlan) + }.flatten + direct ++ commandResults + } + + private def assertV2Write(captured: Seq[QueryExecution]): Unit = { + val writeOptions = captured.flatMap(qe => collectInMemoryWriteOptions(qe.executedPlan)) + assert(writeOptions.nonEmpty, "expected a V2 in-memory batch write") + writeOptions.foreach(assertTargetOptions) + } + test("SPARK-36680: Supports Dynamic Table Options for SQL Select") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { @@ -96,15 +259,19 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { } } - test("SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert") { + testWithLoadOptionAsTableState( + "SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") - val df = sql(s"INSERT INTO $t1 WITH (`write.split-size` = 10) VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() + val df = sql(s"INSERT INTO $t1 WITH (`$loadOption` = '$loadOptionValue', " + + s"`$writeOption` = '$writeOptionValue') VALUES (1, 'a'), (2, 'b')") var collected = df.queryExecution.optimizedPlan.collect { case CommandResult(_, AppendData(relation: DataSourceV2Relation, _, _, _, _, _, _), _, _) => - assert(relation.options.get("write.split-size") == "10") + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assertTargetOptions(relation.options) } assert (collected.size == 1) @@ -113,9 +280,10 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { _, AppendDataExec(_, _, write, _, _), _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#Append] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT)) val insertResult = sql(s"SELECT * FROM $t1") checkAnswer(insertResult, Seq(Row(1, "a"), Row(2, "b"))) @@ -223,74 +391,87 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { } } - test("SPARK-50286: Propagate options for DataFrameWriter Append") { + testWithLoadOptionAsTableState("SPARK-50286: Propagate options for DataFrameWriter Append") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") + inMemoryCatalog.resetLoadTableCalls() val captured = withQueryExecutionsCaptured(spark) { Seq(1 -> "a", 2 -> "b").toDF("id", "data") .write - .option("write.split-size", "10") + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) .mode("append") .insertInto(t1) } assert(captured.size === 1) val qe = captured.head var collected = qe.optimizedPlan.collect { - case AppendData(_: DataSourceV2Relation, _, writeOptions, _, _, _, _) => - assert(writeOptions("write.split-size") == "10") + case AppendData(relation: DataSourceV2Relation, _, writeOptions, _, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) } assert (collected.size == 1) collected = qe.executedPlan.collect { case AppendDataExec(_, _, write, _, _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#Append] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT)) } } - test("SPARK-50286: Propagate options for DataFrameWriterV2 Append") { + testWithLoadOptionAsTableState("SPARK-50286: Propagate options for DataFrameWriterV2 Append") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") + inMemoryCatalog.resetLoadTableCalls() val captured = withQueryExecutionsCaptured(spark) { Seq(1 -> "a", 2 -> "b").toDF("id", "data") .writeTo(t1) - .option("write.split-size", "10") + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) .append() } assert(captured.size === 1) val qe = captured.head var collected = qe.optimizedPlan.collect { - case AppendData(_: DataSourceV2Relation, _, writeOptions, _, _, _, _) => - assert(writeOptions("write.split-size") == "10") + case AppendData(relation: DataSourceV2Relation, _, writeOptions, _, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) } assert (collected.size == 1) collected = qe.executedPlan.collect { case AppendDataExec(_, _, write, _, _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#Append] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT)) } } - test("SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert Overwrite") { + testWithLoadOptionAsTableState( + "SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert Overwrite") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() - val df = sql(s"INSERT OVERWRITE $t1 WITH (`write.split-size` = 10) " + - s"VALUES (3, 'c'), (4, 'd')") + val df = sql(s"INSERT OVERWRITE $t1 WITH (`$loadOption` = '$loadOptionValue', " + + s"`$writeOption` = '$writeOptionValue') VALUES (3, 'c'), (4, 'd')") var collected = df.queryExecution.optimizedPlan.collect { case CommandResult(_, OverwriteByExpression(relation: DataSourceV2Relation, _, _, _, _, _, _, _), _, _) => - assert(relation.options.get("write.split-size") === "10") + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assertTargetOptions(relation.options) } assert (collected.size == 1) @@ -299,58 +480,69 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { _, OverwriteByExpressionExec(_, _, write, _, _), _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#TruncateAndAppend] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) val insertResult = sql(s"SELECT * FROM $t1") checkAnswer(insertResult, Seq(Row(3, "c"), Row(4, "d"))) } } - test("SPARK-50286: Propagate options for DataFrameWriterV2 OverwritePartitions") { + testWithLoadOptionAsTableState( + "SPARK-50286: Propagate options for DataFrameWriterV2 OverwritePartitions") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() val captured = withQueryExecutionsCaptured(spark) { Seq(3 -> "c", 4 -> "d").toDF("id", "data") .writeTo(t1) - .option("write.split-size", "10") + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) .overwritePartitions() } assert(captured.size === 1) val qe = captured.head var collected = qe.optimizedPlan.collect { - case OverwritePartitionsDynamic(_: DataSourceV2Relation, _, writeOptions, _, _, _) => - assert(writeOptions("write.split-size") === "10") + case OverwritePartitionsDynamic( + relation: DataSourceV2Relation, _, writeOptions, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) } assert (collected.size == 1) collected = qe.executedPlan.collect { case OverwritePartitionsDynamicExec(_, _, write, _, _) => val dynOverwrite = write.toBatch.asInstanceOf[InMemoryBaseTable#DynamicOverwrite] - assert(dynOverwrite.info.options.get("write.split-size") === "10") + assertTargetOptions(dynOverwrite.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) } } - test("SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert Replace") { + testWithLoadOptionAsTableState( + "SPARK-49098, SPARK-50286: Supports Dynamic Table Options for SQL Insert Replace") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() - val df = sql(s"INSERT INTO $t1 WITH (`write.split-size` = 10) " + - s"REPLACE WHERE TRUE " + + val df = sql(s"INSERT INTO $t1 WITH (`$loadOption` = '$loadOptionValue', " + + s"`$writeOption` = '$writeOptionValue') REPLACE WHERE TRUE " + s"VALUES (3, 'c'), (4, 'd')") var collected = df.queryExecution.optimizedPlan.collect { case CommandResult(_, OverwriteByExpression(relation: DataSourceV2Relation, _, _, _, _, _, _, _), _, _) => - assert(relation.options.get("write.split-size") == "10") + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assertTargetOptions(relation.options) } assert (collected.size == 1) @@ -359,23 +551,26 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { _, OverwriteByExpressionExec(_, _, write, _, _), _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#TruncateAndAppend] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) val insertResult = sql(s"SELECT * FROM $t1") checkAnswer(insertResult, Seq(Row(3, "c"), Row(4, "d"))) } } - test("SPARK-50286: Propagate options for DataFrameWriter Overwrite") { + testWithLoadOptionAsTableState("SPARK-50286: Propagate options for DataFrameWriter Overwrite") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") + inMemoryCatalog.resetLoadTableCalls() val captured = withQueryExecutionsCaptured(spark) { Seq(1 -> "a", 2 -> "b").toDF("id", "data") .write - .option("write.split-size", "10") + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) .mode("overwrite") .insertInto(t1) } @@ -383,47 +578,1110 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { val qe = captured.head var collected = qe.optimizedPlan.collect { - case OverwriteByExpression(_: DataSourceV2Relation, _, _, writeOptions, _, _, _, _) => - assert(writeOptions("write.split-size") === "10") + case OverwriteByExpression( + relation: DataSourceV2Relation, _, _, writeOptions, _, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) } assert (collected.size == 1) collected = qe.executedPlan.collect { case OverwriteByExpressionExec(_, _, write, _, _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#TruncateAndAppend] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + } + } + + testWithLoadOptionAsTableState( + "SPARK-58389: dynamic partition overwrite separates load and write options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string) PARTITIONED BY (id)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() + + val captured = withSQLConf( + SQLConf.PARTITION_OVERWRITE_MODE.key -> + SQLConf.PartitionOverwriteMode.DYNAMIC.toString) { + withQueryExecutionsCaptured(spark) { + Seq(2 -> "updated", 3 -> "new").toDF("id", "data") + .write + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .mode("overwrite") + .insertInto(t1) + } + } + assert(captured.size === 1) + + val qe = captured.head + val logicalWrites = qe.optimizedPlan.collect { + case OverwritePartitionsDynamic( + relation: DataSourceV2Relation, _, writeOptions, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) + } + assert(logicalWrites.size === 1) + + val physicalWrites = qe.executedPlan.collect { + case OverwritePartitionsDynamicExec(_, _, write, _, _) => + val dynamicOverwrite = write.toBatch.asInstanceOf[InMemoryBaseTable#DynamicOverwrite] + assertTargetOptions(dynamicOverwrite.info.options) + } + assert(physicalWrites.size === 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + checkAnswer(sql(s"SELECT * FROM $t1"), + Seq(Row(1, "a"), Row(2, "updated"), Row(3, "new"))) } } - test("SPARK-50286: Propagate options for DataFrameWriterV2 Overwrite") { + testWithLoadOptionAsTableState( + "SPARK-50286: Propagate options for DataFrameWriterV2 Overwrite") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() val captured = withQueryExecutionsCaptured(spark) { Seq(3 -> "c", 4 -> "d").toDF("id", "data") .writeTo(t1) - .option("write.split-size", "10") + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) .overwrite(lit(true)) } assert(captured.size === 1) val qe = captured.head var collected = qe.optimizedPlan.collect { - case OverwriteByExpression(_: DataSourceV2Relation, _, _, writeOptions, _, _, _, _) => - assert(writeOptions("write.split-size") === "10") + case OverwriteByExpression( + relation: DataSourceV2Relation, _, _, writeOptions, _, _, _, _) => + assert(relation.table.isInstanceOf[InMemoryBaseTable]) + assert(writeOptions(loadOption) === loadOptionValue) + assert(writeOptions(writeOption) === writeOptionValue) } assert (collected.size == 1) collected = qe.executedPlan.collect { case OverwriteByExpressionExec(_, _, write, _, _) => val append = write.toBatch.asInstanceOf[InMemoryBaseTable#TruncateAndAppend] - assert(append.info.options.get("write.split-size") === "10") + assertTargetOptions(append.info.options) } assert (collected.size == 1) + assertWriteLoad(Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + } + } + + testWithLoadOptionAsTableState( + "SPARK-58389: DataFrameWriter saveAsTable separates load and write options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + Seq( + ("append", Set(TableWritePrivilege.INSERT)), + ("overwrite", Set(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + ).foreach { case (mode, expectedPrivileges) => + inMemoryCatalog.resetLoadTableCalls() + + val captured = withQueryExecutionsCaptured(spark) { + Seq(1 -> "a").toDF("id", "data") + .write + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .mode(mode) + .saveAsTable(t1) + } + + assertWriteLoad(expectedPrivileges) + assertV2Write(captured) + } + } + } + + testWithLoadOptionAsTableState( + "SPARK-58389: schema evolution reload separates table-state and write options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint)") + inMemoryCatalog.resetLoadTableCalls() + + val captured = withQueryExecutionsCaptured(spark) { + Seq(1L -> "a").toDF("id", "data") + .writeTo(t1) + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .withSchemaEvolution() + .append() + } + + val matchingCalls = inMemoryCatalog.loadTableCalls.filter { + case (_, options) => options.get(loadOption) == loadOptionValue + } + assert(matchingCalls.size >= 2, "expected initial target load and post-evolution reload") + matchingCalls.foreach { case (context, options) => + assert(context.writePrivileges() === java.util.Set.of(TableWritePrivilege.INSERT)) + assert(options.get(writeOption) === null) + assert(options.size() === 1) + } + assertV2Write(captured) + } + } + + Seq( + "non-staging" -> classOf[NullReturningInMemoryCatalog], + "staging" -> classOf[NullReturningStagingInMemoryCatalog] + ).foreach { case (catalogType, catalogClass) => + test(s"SPARK-58389: $catalogType CTAS/RTAS fallback separates load and write options") { + val catalogName = s"${catalogType.replace('-', '_')}_null_catalog" + registerCatalog(catalogName, catalogClass) + spark.conf.set(s"spark.sql.catalog.$catalogName.tableStateOptionKeys", loadOption) + val fallbackCatalog = catalog(catalogName).asInstanceOf[InMemoryCatalog] + val t1 = s"$catalogName.table" + withTable(t1) { + fallbackCatalog.resetLoadTableCalls() + val createExecutions = withQueryExecutionsCaptured(spark) { + spark.range(1).writeTo(t1) + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .create() + } + assertWriteLoad(fallbackCatalog, Set(TableWritePrivilege.INSERT)) + assertV2Write(createExecutions) + + fallbackCatalog.resetLoadTableCalls() + val replaceExecutions = withQueryExecutionsCaptured(spark) { + spark.range(1, 2).writeTo(t1) + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .replace() + } + assertWriteLoad(fallbackCatalog, Set(TableWritePrivilege.INSERT)) + assertV2Write(replaceExecutions) + + fallbackCatalog.resetLoadTableCalls() + val createOrReplaceExecutions = withQueryExecutionsCaptured(spark) { + spark.range(2, 3).writeTo(t1) + .option(loadOption, loadOptionValue) + .option(writeOption, writeOptionValue) + .createOrReplace() + } + assertWriteLoad(fallbackCatalog, Set(TableWritePrivilege.INSERT)) + assertV2Write(createOrReplaceExecutions) + } + } + } + + test("only table-state options are forwarded to loadTable - DataFrame API") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + spark.read + .option("snapshot", "s1") + .option("customOption", "customValue") + .table(tableName) + .collect() + + assertOnlySnapshotOptions(stateCatalog, "s1") + } + } + + test("only table-state options are forwarded to loadTable - SQL") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + sql(s"SELECT * FROM $tableName " + + "WITH ('snapshot' = 's1', 'customOption' = 'customValue')").collect() + + assertOnlySnapshotOptions(stateCatalog, "s1") + } + } + + test("only table-state options are forwarded to loadTable - DataStreamReader") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + // Trigger analysis of the streaming relation. + spark.readStream + .option("snapshot", "s1") + .option("customOption", "customValue") + .table(tableName) + .queryExecution + .analyzed + + assertOnlySnapshotOptions(stateCatalog, "s1") + } + } + + test("time travel is passed in TableContext and excluded from load options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + // versionAsOf is the default time-travel version option key + // (SQLConf.TIME_TRAVEL_VERSION_KEY). Pin a versioned copy so the versioned load succeeds. + inMemoryCatalog.pinTable(Identifier.of(Array("ns1", "ns2"), "table"), "v1") + + spark.read + .option("versionAsOf", "v1") + .option("customOption", "customValue") + .table(t1) + .collect() + + val ctx = inMemoryCatalog.lastTableContext + assert(ctx.isDefined) + assert(ctx.get.timeTravel().isPresent) + assert(ctx.get.timeTravel().get() === new TimeTravel.AsOfVersion("v1")) + + val opts = inMemoryCatalog.lastLoadTableOptions + assert(opts.isDefined) + assert(opts.get.isEmpty) + } + } + + test("write privileges are carried in TableContext, internal key stripped") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + sql(s"INSERT INTO $tableName WITH ('snapshot' = 's1') VALUES (3, 'c')") + + val ctx = stateCatalog.lastTableContext + assert(ctx.isDefined) + assert(!ctx.get.writePrivileges().isEmpty) + + val opts = stateCatalog.lastLoadTableOptions + assert(opts.isDefined) + assert(opts.get.get("snapshot") === "s1") + // The internal write-privileges marker must not leak to the connector as a user option. + assert(opts.get.get(UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) === null) + } + } + + test("execution refresh filters load options when a catalog declares no state options") { + registerCatalog("loadcounting", classOf[LoadCountingInMemoryCatalog]) + val loadCountingCatalog = + catalog("loadcounting").asInstanceOf[LoadCountingInMemoryCatalog] + val t1 = "loadcounting.ns1.ns2.table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + loadCountingCatalog.resetLoadTableCalls() + loadCountingCatalog.singleArgLoads.set(0) + + spark.read.option("split-size", "5").table(t1).collect() + + // Both analysis and the execution-time refresh enter through the options-aware overload, + // but this catalog declares no table-state options, so each receives an empty option map. + val optionAwareLoads = loadCountingCatalog.loadTableCalls + assert(optionAwareLoads.size === 2, + s"expected one analysis load and one refresh load, got: $optionAwareLoads") + assert(optionAwareLoads.forall(_._2.isEmpty), + s"expected analysis and refresh to filter split-size, got: $optionAwareLoads") + assert(loadCountingCatalog.singleArgLoads.get() === 2, + s"expected two delegated single-argument loads, got: " + + loadCountingCatalog.singleArgLoads.get()) + } + } + + test("catalogs with no declared state options share table state across option bags") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() + + // The catalog declares no table-state option keys, so none of them affect table state. + val df = sql(s"SELECT a.id FROM $t1 WITH (`split-size` = 5) a " + + s"JOIN $t1 WITH (`split-size` = 9) b ON a.id = b.id") + df.queryExecution.analyzed + + // The first reference loads the table and the second reuses that table state. + val analysisLoads = inMemoryCatalog.loadTableCalls + assert(analysisLoads.size === 1, + s"expected the second option bag to reuse the loaded table state, got: $analysisLoads") + assert(analysisLoads.head._2.isEmpty, + s"expected no table-state options to be forwarded, got: $analysisLoads") + + // The execution-time refresh loads the first reference and reuses its table state for the + // second reference. + inMemoryCatalog.resetLoadTableCalls() + df.collect() + val refreshLoads = inMemoryCatalog.loadTableCalls + assert(refreshLoads.size === 1, + s"expected refresh to share table state across option bags, got: $refreshLoads") + assert(refreshLoads.head._2.isEmpty, + s"expected refresh to filter scan options, got: $refreshLoads") + + // Each scan also keeps its own option end-to-end (neither reference inherits the other's). + val splitSizes = df.queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation => s.relation.options.get("split-size") + }.sorted + assert(splitSizes === Seq("5", "9")) + } + } + + test("same table state shares one Table while preserving each reference's options") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`SnApShOt` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + + val analyzedRelations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("split-size") => r + } + assert(analyzedRelations.size == 2) + assert(analyzedRelations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(analyzedRelations.map(_.options.get("snapshot")).distinct == Seq("s1")) + assert(analyzedRelations.head.table eq analyzedRelations.last.table) + + val analysisLoads = stateCatalog.loadTableCalls.filter(_._2.get("snapshot") == "s1") + assert(analysisLoads.size == 1, + s"expected one catalog load for state s1 during analysis, got: $analysisLoads") + assert(analysisLoads.head._2.size() == 1, + s"expected scan options to be filtered during analysis, got: $analysisLoads") + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + assert(df.collect().toSeq == Seq(Row(1), Row(2))) + + val refreshLoads = stateCatalog.loadTableCalls.filter(_._2.get("snapshot") == "s1") + assert(refreshLoads.size == 1, + s"expected one catalog load for state s1 during refresh, got: $refreshLoads") + assert(refreshLoads.head._2.size() == 1, + s"expected scan options to be filtered during refresh, got: $refreshLoads") + val refreshedRelations = df.queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation if s.relation.options.containsKey("split-size") => + s.relation + } + assert(refreshedRelations.size == 2) + assert(refreshedRelations.head.table eq refreshedRelations.last.table) + } + } + + test("shared relation cache selects by state and preserves each reference's full options") { + withStateAwareTable { (stateCatalog, tableName) => + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + val otherStateCached = spark.read + .option("snapshot", "s2") + .option("split-size", "7") + .table(tableName) + cached.cache() + otherStateCached.cache() + try { + cached.collect() + // Cache this relation last so it is searched first. An s1 lookup must scan past it. + otherStateCached.collect() + val cachedTable = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + + def relations(df: DataFrame): Seq[DataSourceV2Relation] = { + df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + } + + stateCatalog.resetLoadTableCalls() + val cachedFirst = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + val cachedFirstRelations = relations(cachedFirst) + assert(cachedFirstRelations.size == 2) + assert(cachedFirstRelations.forall(_.table eq cachedTable)) + assert(cachedFirstRelations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") == 1) + + stateCatalog.resetLoadTableCalls() + assert(cachedFirst.collect().map(_.getLong(0)).sorted.toSeq == Seq(1L, 2L)) + assert(stateCatalog.loadTableCalls.isEmpty, + s"shared relation cache pin should avoid refresh reloads, got: " + + stateCatalog.loadTableCalls) + + stateCatalog.resetLoadTableCalls() + val uncachedFirst = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) b ON a.id = b.id") + val uncachedFirstRelations = relations(uncachedFirst) + assert(uncachedFirstRelations.size == 2) + assert(uncachedFirstRelations.forall(_.table eq cachedTable)) + assert(uncachedFirstRelations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") == 1) + + stateCatalog.resetLoadTableCalls() + assert(uncachedFirst.collect().map(_.getLong(0)).sorted.toSeq == Seq(1L, 2L)) + assert(stateCatalog.loadTableCalls.isEmpty, + s"shared relation cache pin should avoid refresh reloads, got: " + + stateCatalog.loadTableCalls) + } finally { + otherStateCached.unpersist() + cached.unpersist() + } + } + } + + test("shared relation cache makes table selection independent of reference order") { + withStateAwareTable { (stateCatalog, tableName) => + val ident = Identifier.of(Array("ns"), "table") + stateCatalog.alterTable(ident, TableChange.setProperty("version", "X")) + + val cached = spark.read.option("split-size", "5").table(tableName) + cached.cache() + try { + cached.collect() + val versionX = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + assert(versionX.properties().get("version") == "X") + + // Simulate an external catalog update that does not invalidate Spark's relation cache. + val versionY = stateCatalog.alterTable( + ident, + TableChange.setProperty("version", "Y")) + assert(versionY.properties().get("version") == "Y") + assert(versionY.id == versionX.id) + assert(versionY ne versionX) + + def relations(df: DataFrame): Seq[DataSourceV2Relation] = { + df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("split-size") => r + } + } + + val readA = spark.read.option("split-size", "5").table(tableName) + val readB = spark.read.option("split-size", "9").table(tableName) + assert(relations(readA).map(_.table) == Seq(versionX)) + assert(relations(readB).map(_.table) == Seq(versionX)) + + val aThenB = sql(s"SELECT a.id FROM $tableName WITH (`split-size` = 5) a " + + s"JOIN $tableName WITH (`split-size` = 9) b ON a.id = b.id") + val bThenA = sql(s"SELECT a.id FROM $tableName WITH (`split-size` = 9) a " + + s"JOIN $tableName WITH (`split-size` = 5) b ON a.id = b.id") + assert(relations(aThenB).map(_.table) == Seq(versionX, versionX)) + assert(relations(bThenA).map(_.table) == Seq(versionX, versionX)) + } finally { + cached.unpersist() + } + } + } + + test("shared relation cache rejects a different declared table state") { + withStateAwareTable { (stateCatalog, tableName) => + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + cached.cache() + try { + cached.collect() + val cachedTable = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + + stateCatalog.resetLoadTableCalls() + val differentState = spark.read + .option("snapshot", "s2") + .option("split-size", "5") + .table(tableName) + val relation = differentState.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r + }.getOrElse(fail("expected a v2 relation")) + + assert(relation.table ne cachedTable) + assert(relation.options.get("snapshot") == "s2") + assert(relation.options.get("split-size") == "5") + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s2") == 1) + } finally { + cached.unpersist() + } + } + } + + test("streaming references participate in the query table-state cache") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + val df = sql(s"SELECT a.id FROM STREAM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN STREAM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: StreamingRelationV2 if r.extraOptions.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.extraOptions.get("split-size")).sorted == Seq("5", "9")) + assert(relations.head.table eq relations.last.table) + val stateLoads = stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") + assert(stateLoads == 1, s"expected one streaming table load, got: $stateLoads") + } + } + + test("different table-state option values establish separate table pins") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's2', `split-size` = 9) b ON a.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("snapshot")).sorted == Seq("s1", "s2")) + assert(relations.head.table ne relations.last.table) + val loadedStates = stateCatalog.loadTableCalls + .map(_._2.get("snapshot")) + .filter(_ != null) + .sorted + assert(loadedStates == Seq("s1", "s2")) + } + } + + test("persistent write targets bypass query-scoped cache lookups") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + RelationCache.empty) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "5")) + val read = UnresolvedRelation(tableName.split("\\.").toSeq, options) + val write = read.requireWritePrivileges(Set(TableWritePrivilege.INSERT)) + + def resolve(relation: UnresolvedRelation): DataSourceV2Relation = { + resolver.resolveRelation(relation).flatMap(_.collectFirst { + case r: DataSourceV2Relation => r + }).getOrElse(fail(s"failed to resolve ${relation.name} as a v2 relation")) + } + + AnalysisContext.withNewAnalysisContext { + val readRelation = resolve(read) + val writeRelation = resolve(write) + + assert(readRelation.table ne writeRelation.table) + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 1) + } + + assert(stateCatalog.loadTableCalls.size == 2) + assert(stateCatalog.loadTableCalls.count(_._1.writePrivileges().isEmpty) == 1) + assert(stateCatalog.loadTableCalls.count( + _._1.writePrivileges().contains(TableWritePrivilege.INSERT)) == 1) + assert(stateCatalog.loadTableCalls.forall(_._2.get("snapshot") == "s1")) + assert(stateCatalog.loadTableCalls.forall(_._2.size() == 1)) + } + } + + test("SPARK-58389: explicit time travel specs on internal write targets use qualified names") { + withStateAwareTable { (_, tableName) => + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + RelationCache.empty) + val write = UnresolvedRelation(Seq("ns", "table")) + .requireWritePrivileges(Set(TableWritePrivilege.INSERT)) + val previousCatalog = spark.catalog.currentCatalog() + + try { + spark.catalog.setCurrentCatalog("statecat") + val e = AnalysisContext.withNewAnalysisContext { + intercept[AnalysisException] { + resolver.resolveRelation(write, Some(AsOfVersion("v1"))) + } + } + checkError( + exception = e, + condition = "UNSUPPORTED_FEATURE.TIME_TRAVEL", + parameters = Map("relationId" -> "`statecat`.`ns`.`table`")) + } finally { + spark.catalog.setCurrentCatalog(previousCatalog) + } + } + } + + test("persistent write targets establish table pins for subsequent reads") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + RelationCache.empty) + val writeOptions = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "5")) + val readOptions = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "9")) + val write = UnresolvedRelation(tableName.split("\\.").toSeq, writeOptions) + .requireWritePrivileges(Set(TableWritePrivilege.INSERT)) + val read = UnresolvedRelation(tableName.split("\\.").toSeq, readOptions) + + def resolve(relation: UnresolvedRelation): DataSourceV2Relation = { + resolver.resolveRelation(relation).flatMap(_.collectFirst { + case r: DataSourceV2Relation => r + }).getOrElse(fail(s"failed to resolve ${relation.name} as a v2 relation")) + } + + AnalysisContext.withNewAnalysisContext { + val writeRelation = resolve(write) + val readRelation = resolve(read) + + assert(writeRelation.table eq readRelation.table) + assert(writeRelation.options.get("split-size") == "5") + assert(readRelation.options.get("split-size") == "9") + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + } + + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._1.writePrivileges().contains( + TableWritePrivilege.INSERT)) + assertOnlySnapshotOptions(stateCatalog, "s1") + } + } + + test("transaction V2TableReference skips shared lookup and writes bypass query caches") { + withStateAwareTable { (stateCatalog, tableName) => + val original = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + .queryExecution + .analyzed + .collectFirst { case r: DataSourceV2Relation => r } + .getOrElse(fail("expected a v2 relation")) + val readRef = V2TableReference.createForTransaction(original) + val otherReadRef = V2TableReference.createForTransaction(original.copy( + options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "9")))) + val writeRef = V2TableReference.createForWriteTarget(original) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = (_, _, _, _, _) => { + sharedRelationCacheLookups += 1 + None + } + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + sharedRelationCache) + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + AnalysisContext.withNewAnalysisContext { + val readRelation = resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + val cachedReadRelation = + resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + val otherReadRelation = + resolver.resolveReference(otherReadRef).asInstanceOf[DataSourceV2Relation] + val writeRelation = resolver.resolveReference(writeRef).asInstanceOf[DataSourceV2Relation] + val readAfterWriteRelation = + resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + + assert(readRelation.table eq cachedReadRelation.table) + assert(readRelation.table eq otherReadRelation.table) + assert(readRelation.table eq readAfterWriteRelation.table) + assert(readRelation.table ne writeRelation.table) + assert(readRelation.options.get("split-size") == "5") + assert(otherReadRelation.options.get("split-size") == "9") + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + } + + assert(sharedRelationCacheLookups == 0) + assert(stateCatalog.singleArgLoads.get() == 2) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.size() == 1) + } + } + + test("nested view resolution shares the query table-state cache") { + withStateAwareTable { (stateCatalog, tableName) => + withView("state_nested_view") { + sql(s"CREATE VIEW state_nested_view AS SELECT * FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5)") + stateCatalog.resetLoadTableCalls() + + val df = sql(s"SELECT v.id FROM state_nested_view v JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON v.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(relations.head.table eq relations.last.table) + val stateLoads = stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") + assert(stateLoads == 1, s"expected one nested-view table load, got: $stateLoads") + } + } + } + + test("temporary-view V2TableReference consults shared cache only for the initial pin") { + withStateAwareTable { (stateCatalog, tableName) => + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + .queryExecution + .analyzed + .collectFirst { case r: DataSourceV2Relation => r } + .getOrElse(fail("expected a v2 relation")) + val initialRef = V2TableReference.createForTempView(cached, Seq("state_view")) + val otherOptionsRef = V2TableReference.createForTempView( + cached.copy(options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "9"))), + Seq("state_view")) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = (_, _, _, _, _) => { + sharedRelationCacheLookups += 1 + Some(cached) + } + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + sharedRelationCache) + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + AnalysisContext.withNewAnalysisContext { + val initialRelation = + resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + val cachedRelation = + resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + val otherOptionsRelation = + resolver.resolveReference(otherOptionsRef).asInstanceOf[DataSourceV2Relation] + + assert(initialRelation.table eq cached.table) + assert(cachedRelation.table eq cached.table) + assert(otherOptionsRelation.table eq cached.table) + assert(initialRelation.options.get("split-size") == "5") + assert(otherOptionsRelation.options.get("split-size") == "9") + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + } + + assert(sharedRelationCacheLookups == 1) + assert(stateCatalog.singleArgLoads.get() == 1) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.size() == 1) + } + } + + test("temporary-view re-resolution preserves the CacheManager table pin") { + withStateAwareTable { (stateCatalog, tableName) => + withTempView("state_view") { + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + cached.cache() + try { + cached.collect() + val cachedTable = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + cached.createOrReplaceTempView("state_view") + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT v.id FROM state_view v JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON v.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(relations.forall(_.table eq cachedTable)) + assert(stateCatalog.singleArgLoads.get() == 1, + s"expected one table load while re-resolving the query, got: " + + stateCatalog.singleArgLoads.get()) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.size() == 1) + } finally { + cached.unpersist() + } + } + } + } + + test("SPARK-58389: repeated references with the same options load the table once") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + inMemoryCatalog.resetLoadTableCalls() + + // Both references carry the same options, so they share one relation-cache entry: the table + // is loaded once (resolve-once-per-query is preserved for identical option bags). + val df = sql(s"SELECT a.id FROM $t1 WITH (`split-size` = 5) a " + + s"JOIN $t1 WITH (`split-size` = 5) b ON a.id = b.id") + df.queryExecution.analyzed + + val analysisLoads = inMemoryCatalog.loadTableCalls + assert(analysisLoads.size === 1, + s"expected a single loadTable for identical option bags, got: $analysisLoads") + assert(analysisLoads.head._2.isEmpty, + s"expected scan options to be filtered during analysis, got: $analysisLoads") + + // The refresh phase also reuses one load for repeated references with identical options. + inMemoryCatalog.resetLoadTableCalls() + df.collect() + val refreshLoads = inMemoryCatalog.loadTableCalls + assert(refreshLoads.size === 1, + s"expected refresh to load identical option bags once, got: $refreshLoads") + assert(refreshLoads.head._2.isEmpty, + s"expected scan options to be filtered during refresh, got: $refreshLoads") + } + } + + test("SPARK-58389: time travel is rejected for resolved and newly created V2 write targets") { + val t1 = s"${catalogAndNamespace}table" + val newTable = s"${catalogAndNamespace}new_table" + withTable(t1, newTable) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + + val input = Seq(1L -> "a").toDF("id", "data") + val existingRelationId = "`testcat`.`ns1`.`ns2`.`table`" + val newRelationId = "`testcat`.`ns1`.`ns2`.`new_table`" + Seq( + existingRelationId -> (() => sql(s"INSERT INTO $t1 WITH " + + "('versionAsOf' = 'v1', 'timestampAsOf' = '2021-01-01') VALUES (1, 'a')")), + existingRelationId -> (() => input.writeTo(t1) + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .append()), + existingRelationId -> (() => input.write + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .insertInto(t1)), + existingRelationId -> (() => input.write + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .mode("append") + .saveAsTable(t1)), + newRelationId -> (() => input.writeTo(newTable) + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .create()), + existingRelationId -> (() => input.writeTo(t1) + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .replace()), + existingRelationId -> (() => input.writeTo(t1) + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .createOrReplace()) + ).foreach { case (relationId, writeToTimeTravelTarget) => + checkError( + exception = intercept[AnalysisException](writeToTimeTravelTarget()), + condition = "UNSUPPORTED_FEATURE.TIME_TRAVEL", + parameters = Map("relationId" -> relationId)) + } + + withTempView("temp_view") { + input.createOrReplaceTempView("temp_view") + checkError( + exception = intercept[AnalysisException] { + input.writeTo("temp_view") + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .append() + }, + condition = "UNSUPPORTED_FEATURE.TIME_TRAVEL", + parameters = Map("relationId" -> "`temp_view`")) + } + } + } + + test("SPARK-58389: a missing insertInto target reports table not found before time travel") { + val missingTable = s"${catalogAndNamespace}missing" + val input = Seq(1L -> "a").toDF("id", "data") + + checkError( + exception = intercept[AnalysisException] { + input.write.option("versionAsOf", "v1").insertInto(missingTable) + }, + condition = "TABLE_OR_VIEW_NOT_FOUND", + parameters = Map("relationName" -> "`ns1`.`ns2`.`missing`")) + } + + test("SPARK-58389: SQL CTAS and RTAS options remain table properties") { + val t1 = s"${catalogAndNamespace}table" + val ident = Identifier.of(Array("ns1", "ns2"), "table") + withTable(t1) { + sql(s"CREATE TABLE $t1 USING foo OPTIONS ('versionAsOf' = 'v1') " + + "AS SELECT 1L AS id, 'a' AS data") + val createdProperties = inMemoryCatalog.loadTable(ident).properties() + assert(createdProperties.get("versionAsOf") === "v1") + assert(createdProperties.get("option.versionAsOf") === "v1") + + sql(s"REPLACE TABLE $t1 USING foo OPTIONS ('timestampAsOf' = '2021-01-01') " + + "AS SELECT 2L AS id, 'b' AS data") + val replacedProperties = inMemoryCatalog.loadTable(ident).properties() + assert(replacedProperties.get("timestampAsOf") === "2021-01-01") + assert(replacedProperties.get("option.timestampAsOf") === "2021-01-01") + checkAnswer(spark.table(t1), Row(2L, "b")) + } + } + + test("SPARK-58389: CACHE TABLE result is not reused for a read carrying different options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + + // Cache the option-free read. The CacheManager keys entries on the query plan, and a DSv2 + // relation's plan carries its `options`, so the cached entry's fingerprint is "no options". + val cached = spark.table(t1) + cached.cache() + try { + val cacheManager = spark.sharedState.cacheManager + + // An option-free read has the same plan fingerprint, so it reuses the cached result. + assert(cacheManager.lookupCachedData(spark.table(t1)).isDefined, + "an option-free read should hit the cached result") + + // A read carrying options has a different fingerprint, so it must NOT reuse the cached + // result -- otherwise the connector's options would be silently ignored on a cache hit. + assert( + cacheManager.lookupCachedData(spark.read.option("split-size", "5").table(t1)).isEmpty, + "a read carrying options must not reuse the option-free cached result") + } finally { + cached.unpersist() + } + } + } + + test("execution refresh reuses cached table state with different scan options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + + val cached = spark.table(t1) + cached.cache() + try { + assert(cached.count() === 2) + + // This catalog declares no state options, so analysis and refresh may reuse the cached + // table while the relation keeps split-size=5 as its complete option bag. + val df = spark.read.option("split-size", "5").table(t1).filter("id > 0") + val analyzedRelation = df.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r + }.getOrElse(fail("expected a v2 relation")) + assert(analyzedRelation.options.get("split-size") == "5") + inMemoryCatalog.resetLoadTableCalls() + df.collect() + + assert(inMemoryCatalog.loadTableCalls.isEmpty, + s"expected refresh to reuse the matching table state, got: " + + inMemoryCatalog.loadTableCalls) + } finally { + cached.unpersist() + } + } + } + + test("recaching preserves relation options and filters table load options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + + val cached = spark.read.option("split-size", "5").table(t1) + cached.cache() + try { + assert(cached.count() === 2) + + // Refreshing a cached table rebuilds its CacheManager entry. The rebuilt relation retains + // the original options, but this catalog declares no state options for the table reload. + inMemoryCatalog.resetLoadTableCalls() + spark.catalog.refreshTable(t1) + + val recacheLoads = inMemoryCatalog.loadTableCalls + assert(recacheLoads.nonEmpty, "expected recache to reload the table") + assert(recacheLoads.forall(_._2.isEmpty), + s"expected recache to filter split-size, got: $recacheLoads") + + val cacheManager = spark.sharedState.cacheManager + val sameOptions = spark.read.option("split-size", "5").table(t1) + val recached = cacheManager.lookupCachedData(sameOptions) + assert(recached.isDefined, "a read with the original options should hit after recache") + + val recachedOptions = recached.get.plan.collect { + case r: DataSourceV2Relation => r.options.get("split-size") + } + assert(recachedOptions === Seq("5"), + s"expected the recached relation to retain split-size=5, got: $recachedOptions") + + assert(cacheManager.lookupCachedData(spark.table(t1)).isEmpty, + "an option-free read must not reuse the recached option-carrying result") + } finally { + spark.catalog.clearCache() + } + } + } + + test("recaching a filtered plan filters table load options") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + + // The filter forces CacheManager.tryRefreshPlan through V2TableRefreshUtil instead of the + // bare-relation fast path covered by the preceding test. + val cached = spark.read.option("split-size", "5").table(t1).filter("id > 0") + cached.cache() + try { + assert(cached.count() === 2) + inMemoryCatalog.resetLoadTableCalls() + + spark.catalog.refreshTable(t1) + + val recacheLoads = inMemoryCatalog.loadTableCalls + assert(recacheLoads.nonEmpty, "expected filtered recache to reload the table") + assert(recacheLoads.forall(_._2.isEmpty), + s"expected filtered recache to filter split-size, got: $recacheLoads") + + val samePlan = spark.read.option("split-size", "5").table(t1).filter("id > 0") + assert(spark.sharedState.cacheManager.lookupCachedData(samePlan).isDefined, + "the filtered plan should remain cached after refresh") + } finally { + spark.catalog.clearCache() + } + } + } + + test("SPARK-58389: a DataFrame temp view's options do not leak to a later reference") { + val t1 = s"${catalogAndNamespace}table" + withTable(t1) { + sql(s"CREATE TABLE $t1 (id bigint, data string)") + sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") + + // The temp view resolves via the V2TableReference path, whose relation carries the view's + // options. A later option-free reference to the same table must not inherit them. + withTempView("v") { + spark.read.option("split-size", "5").table(t1).createOrReplaceTempView("v") + inMemoryCatalog.resetLoadTableCalls() + val df = sql(s"SELECT v.id FROM v JOIN $t1 b ON v.id = b.id") + + val splitSizes = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation => Option(r.options.get("split-size")) + } + // Exactly one reference (`v`) keeps its option; `b` (option-free) must not inherit it. + assert(splitSizes.flatten === Seq("5"), + s"option leaked to the option-free reference, got: $splitSizes") + assert(splitSizes.contains(None), + s"expected an option-free reference, got: $splitSizes") + assert(inMemoryCatalog.loadTableCalls.nonEmpty, + "V2TableReference temp-view reload did not load the table") + assert(inMemoryCatalog.loadTableCalls.forall(_._2.isEmpty), + s"expected temp-view reloads to filter scan options, got: " + + inMemoryCatalog.loadTableCalls) + } } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index 79c46bdb4ccf7..76c7902188ed3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -5478,6 +5478,47 @@ class DataSourceV2SQLSuiteV2Filter extends DataSourceV2SQLSuite { s"Expected 1 partition after scalar subquery pruning, got $numPartitions") } } + + test("SPARK-58207: non-deterministic scalar subquery filters are not pushed into " + + "runtimeFilters") { + val tbl = s"${catalogAndNamespace}tbl" + val dim = s"${catalogAndNamespace}dim" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Format PARTITIONED BY (part)") + for (i <- 0 until 10) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + + sql(s"CREATE TABLE $dim (val INT) USING $v2Format") + sql(s"INSERT INTO $dim VALUES (3)") + + // `part = (subquery) OR rand() < 0.5` references only the partition column and holds a + // scalar subquery, so it is a candidate for runtime pushdown, but it is non-deterministic. + // Routing it would push it to the source for pruning while the FilterExec above the scan + // re-evaluates it, and a partition the source dropped on its own evaluation could not be + // recovered. + val df = sql( + s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim) OR rand() < 0.5") + df.collect() + + val batchScan = collect(df.queryExecution.executedPlan) { + case b: BatchScanExec => b + }.head + assert(batchScan.runtimeFilters.isEmpty, + s"Expected no runtime filters for a non-deterministic filter, " + + s"got ${batchScan.runtimeFilters}") + + // No pruning at the source, and the filter is still evaluated by Spark after the scan. + val numPartitions = batchScan.filteredPartitions.count(_.isDefined) + assert(numPartitions == 10, + s"Expected all 10 partitions to be retained, got $numPartitions") + val postScanConditions = collect(df.queryExecution.executedPlan) { + case f: FilterExec => f.condition + } + assert(postScanConditions.exists(!_.deterministic), + s"Expected the non-deterministic filter above the scan, got $postScanConditions") + } + } } class ReserveSchemaNullabilityCatalog extends InMemoryCatalog { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala index cc784e6f73a07..eb4e6b53dd23d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2Suite.scala @@ -19,37 +19,43 @@ package org.apache.spark.sql.connector import java.io.File import java.util +import java.util.Optional import java.util.OptionalLong import scala.jdk.CollectionConverters._ import test.org.apache.spark.sql.connector._ -import org.apache.spark.SparkUnsupportedOperationException -import org.apache.spark.sql.{AnalysisException, DataFrame, Row} +import org.apache.spark.{SparkException, SparkUnsupportedOperationException} +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SQLContext} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{ AttributeReference, Expression => CatalystExpression, GreaterThan => CatalystGreaterThan, - Literal => CatalystLiteral, ScalarSubquery} + LessThan => CatalystLessThan, Literal => CatalystLiteral, ScalarSubquery} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter => LogicalFilter, Project} -import org.apache.spark.sql.connector.catalog.{PartitionInternalRow, SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils +import org.apache.spark.sql.connector.catalog.{PartitionInternalRow, SupportsRead, SupportsWrite, Table, TableCapability, TableProvider} import org.apache.spark.sql.connector.catalog.TableCapability._ import org.apache.spark.sql.connector.expressions.{Expression, FieldReference, Literal, NamedReference, NullOrdering, SortDirection, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read._ import org.apache.spark.sql.connector.read.Scan.ColumnarSupportMode +import org.apache.spark.sql.connector.read.colstats.ColumnStatistics import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning, Partitioning, UnknownPartitioning} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, LogicalWriteInfo, PhysicalWriteInfo, Write, WriteBuilder, WriterCommitMessage} import org.apache.spark.sql.execution.SortExec import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper -import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering} +import org.apache.spark.sql.execution.datasources.v2.{ + BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation, PushedDownOperators, + V1ScanWrapper, V2ScanPartitioningAndOrdering} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ import org.apache.spark.sql.execution.exchange.{Exchange, ShuffleExchangeExec} import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters -import org.apache.spark.sql.sources.{Filter, GreaterThan} +import org.apache.spark.sql.internal.connector.{SimpleTableProvider, SupportsPushDownCatalystFilters} +import org.apache.spark.sql.sources.{BaseRelation, Filter, GreaterThan, TableScan} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructField, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -387,16 +393,30 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper Seq(classOf[ReportStatisticsDataSource], classOf[JavaReportStatisticsDataSource]).foreach { cls => withClue(cls.getName) { - val df = spark.read.format(cls.getName).load() - val logical = df.queryExecution.optimizedPlan.collect { - case d: DataSourceV2ScanRelation => d - }.head - - val statics = logical.computeStats() - assert(statics.rowCount.isDefined && statics.rowCount.get === 10, - "Row count statics should be reported by data source") - assert(statics.sizeInBytes === 80, - "Size in bytes statics should be reported by data source") + def scanStats: org.apache.spark.sql.catalyst.plans.logical.Statistics = { + val df = spark.read.format(cls.getName).load() + df.queryExecution.optimizedPlan.collect { + case d: DataSourceV2ScanRelation => d + }.head.computeStats() + } + + withSQLConf( + SQLConf.CBO_ENABLED.key -> "false", + SQLConf.PLAN_STATS_ENABLED.key -> "false") { + val statics = scanStats + assert(statics.rowCount.isEmpty, + "Row count statics should not be reported when Spark only needs size") + assert(statics.sizeInBytes === 80, + "Size in bytes statics should be reported by data source") + } + + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val statics = scanStats + assert(statics.rowCount.isDefined && statics.rowCount.get === 10, + "Row count statics should be reported by data source") + assert(statics.sizeInBytes === 80, + "Size in bytes statics should be reported by data source") + } } } } @@ -485,6 +505,17 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper } } + test("SPARK-58352: WRITING_JOB_FAILED when batch write commit and abort both fail") { + val cls = classOf[CommitAndAbortFailingDataSource] + checkError( + exception = intercept[SparkException] { + spark.range(1).select($"id" as Symbol("i"), -$"id" as Symbol("j")) + .write.format(cls.getName).mode("append").save() + }, + condition = "WRITING_JOB_FAILED", + parameters = Map.empty[String, String]) + } + test("simple counter in writer with onDataWriterCommit") { Seq(classOf[SimpleWritableDataSource]).foreach { cls => withTempPath { file => @@ -1168,6 +1199,228 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper }.head } + private def hasIGt3(condition: CatalystExpression): Boolean = { + condition.exists { + case CatalystGreaterThan(attr: AttributeReference, CatalystLiteral(value: Int, _)) => + attr.name == "i" && value == 3 + case _ => false + } + } + + private def hasJLtNeg5(condition: CatalystExpression): Boolean = { + condition.exists { + case CatalystLessThan(attr: AttributeReference, CatalystLiteral(value: Int, _)) => + attr.name == "j" && value == -5 + case _ => false + } + } + + test("Spark post-pushdown adjustments re-add fully pushed predicates") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"Expected i > 3 in a post-scan Filter when Spark owns post-pushdown adjustments:\n$plan") + val scan = getScanRelation(q) + assert(scan.pushedFilters.exists(hasIGt3), + "scan.pushedFilters should still record the pushed predicate") + } + + test("Spark post-pushdown adjustments use delegated stats end-to-end") { + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val scan = getScanRelation(q) + val stats = scan.stats + val expectedSize = EstimationUtils.getOutputSize( + scan.output, BigInt(10), stats.attributeStats) + assert(stats.rowCount.contains(BigInt(10)), + "fake connector should delegate numRows through scan stats") + assert(stats.sizeInBytes === expectedSize, + "fake connector should use Spark's projection-aware scan-delegated stats size") + assert(q.queryExecution.optimizedPlan.stats.rowCount.contains(BigInt(7)), + "re-added pushed predicate should adjust plan row count from the scan's pre-filter stats") + } + } + + test("scan stats alone do not imply Spark post-pushdown adjustments") { + withSQLConf(SQLConf.CBO_ENABLED.key -> "true") { + val df = spark.read.format(classOf[AdvancedDataSourceV2WithScanStats].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(!plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"i > 3 should be stripped from the post-scan Filter by default:\n$plan") + + val scan = getScanRelation(q) + val stats = scan.stats + assert(scan.scan.asInstanceOf[SupportsReportStatistics].reflectsFullyPushedDownFilters(), + "this fake connector reports scan stats as post-pushdown by default") + assert(stats.rowCount.contains(BigInt(2)), + "connector should use scan rowCount") + assert(stats.sizeInBytes === BigInt(32), + "connector should use scan sizeInBytes") + } + } + + test("V1ScanWrapper delegates optional reported statistics") { + val v1Scan = new V1Scan with SupportsReportStatistics { + override def readSchema(): StructType = TestingV2Source.schema + override def toV1TableScan[T <: BaseRelation with TableScan](context: SQLContext): T = + throw new UnsupportedOperationException("not used") + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(128L) + override def numRows(): OptionalLong = OptionalLong.of(8L) + } + override def estimateSizeInBytes(): OptionalLong = OptionalLong.of(64L) + override def reflectsFullyPushedDownFilters(): Boolean = false + } + + val wrapper = V1ScanWrapper(v1Scan, Nil, + PushedDownOperators(None, None, None, None, Nil, Nil, Nil, None)) + + val stats = wrapper.estimateStatistics() + assert(stats.sizeInBytes().getAsLong === 128L) + assert(stats.numRows().getAsLong === 8L) + assert(wrapper.estimateSizeInBytes().getAsLong === 64L) + assert(!wrapper.reflectsFullyPushedDownFilters()) + } + + test("V1ScanWrapper uses empty/default statistics when V1 scan does not report them") { + val v1Scan = new V1Scan { + override def readSchema(): StructType = TestingV2Source.schema + override def toV1TableScan[T <: BaseRelation with TableScan](context: SQLContext): T = + throw new UnsupportedOperationException("not used") + } + + val wrapper = V1ScanWrapper(v1Scan, Nil, + PushedDownOperators(None, None, None, None, Nil, Nil, Nil, None)) + + val stats = wrapper.estimateStatistics() + assert(!stats.sizeInBytes().isPresent) + assert(!stats.numRows().isPresent) + assert(!wrapper.estimateSizeInBytes().isPresent) + assert(wrapper.reflectsFullyPushedDownFilters()) + } + + test("Spark post-pushdown adjustments are not added for scans without reported stats") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() + val q = df.filter($"i" > 3) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val plan = q.queryExecution.optimizedPlan + assert(!plan.collect { case f: LogicalFilter => f }.exists(f => hasIGt3(f.condition)), + s"i > 3 should not be re-added without SupportsReportStatistics opt-in:\n$plan") + + val scan = getScanRelation(q) + assert(scan.pushedFilters.exists(hasIGt3), + "scan.pushedFilters should still record the pushed predicate") + } + } + + test("Spark post-pushdown adjustments merge pushed and unpushed predicates") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3 && $"j" < -5) + checkAnswer(q, (6 until 10).map(i => Row(i, -i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed and unpushed predicates should be merged into one Filter above the scan:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head)) + assert(hasJLtNeg5(directScanFilters.head)) + } + + test("Spark post-pushdown adjustments preserve projection when re-adding predicates") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).select($"j") + checkAnswer(q, (4 until 10).map(i => Row(-i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed predicate should be re-added directly above the scan before projection:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head)) + assert(optimizedPlan.output.map(_.name) == Seq("j")) + + val scan = getScanRelation(q) + assert(scan.output.exists(_.name == "i"), + "scan output should retain the pushed-filter column needed for the re-added predicate") + assert(scan.output.exists(_.name == "j"), + "scan output should retain the projected column") + } + } + + test("Spark post-pushdown adjustments skip pushed predicates pruned from scan output") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithBestEffortSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).select($"j") + checkAnswer(q, (4 until 10).map(i => Row(-i))) + + val scan = getScanRelation(q) + assert(!scan.output.exists(_.name == "i"), + "column i should be pruned from the scan output") + assert(!scan.scan.asInstanceOf[SupportsReportStatistics].reflectsFullyPushedDownFilters(), + "fake connector should request Spark post-pushdown adjustments") + // SPARK-40259 made pushedFilters the complete set: it keeps a fully-pushed filter even when + // the filter's column is pruned from the scan output (the adjustment below uses a separate, + // pruned-output-remapped set, so it still does not re-add i > 3). + assert(scan.pushedFilters.exists(hasIGt3), + "pushedFilters keeps i > 3 even though column i was pruned (complete-set semantics)") + + val optimizedPlan = q.queryExecution.optimizedPlan + assert(!optimizedPlan.collect { case f: LogicalFilter => f }.exists(f => + hasIGt3(f.condition)), + s"pruned pushed predicate should not be re-added above the scan:\n$optimizedPlan") + } + } + + test("Spark post-pushdown adjustments re-add pushed predicates below stacked residual filters") { + withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { + // Non-deterministic, always-true residual filters: they cannot be pushed down and cannot be + // combined with each other, so two Filter nodes stay stacked above the scan. This forces + // addToScan to recurse through the outer Filter to re-add the fully pushed i > 3. + val alwaysTrue = udf((_: Int) => true).asNondeterministic() + val df = spark.read.format( + classOf[AdvancedDataSourceV2WithSparkPostPushdownAdjustments].getName).load() + val q = df.filter($"i" > 3).filter(alwaysTrue($"i")).filter(alwaysTrue($"j")) + checkAnswer(q, (4 until 10).map(i => Row(i, -i))) + + val optimizedPlan = q.queryExecution.optimizedPlan + val filters = optimizedPlan.collect { case f: LogicalFilter => f } + assert(filters.length == 2, + s"the two non-deterministic residual filters should stay stacked above the scan:\n" + + optimizedPlan) + val directScanFilters = optimizedPlan.collect { + case LogicalFilter(condition, _: DataSourceV2ScanRelation) => condition + } + assert(directScanFilters.length == 1, + s"pushed predicate should be re-added into the filter directly above the scan:\n" + + optimizedPlan) + assert(hasIGt3(directScanFilters.head), + s"addToScan should recurse through the outer residual filter to re-add i > 3:\n" + + optimizedPlan) + } + } + test("pushedFilters are set for fully pushed filters") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() // AdvancedDataSourceV2 only supports pushing GreaterThan on column "i". @@ -1271,33 +1524,34 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "pushedFilters should not contain the unsupported filter on column j") } - test("pushedFilters are remapped by ProjectionOverSchema after nested schema pruning") { + test("pushedFilters keep the relation-level struct type after nested schema pruning") { val df = spark.read.format(classOf[NestedSchemaDataSourceV2].getName).load() // NestedSchemaScanBuilder pushes GreaterThan on "s.a". - // Selecting only s.a triggers nested schema pruning: s goes from struct<a,b> to struct<a>. + // Selecting only s.a triggers nested schema pruning: the scan output narrows s to struct<a>, + // but pushedFilters records the fully-pushed filter against the relation's pre-pruning schema, + // so its struct column keeps the full type struct<a, b>. val q = df.select($"s.a").filter($"s.a" > 3) checkAnswer(q, (4 until 10).map(i => Row(i))) val scanRelation = getScanRelation(q) assert(scanRelation.pushedFilters.nonEmpty, "pushedFilters should be non-empty") - // Find the struct attribute referenced by the pushed filter. - // Before remapping it would have type struct<a,b>; after remapping, struct<a>. val structAttrs = scanRelation.pushedFilters .flatMap(_.collect { case a: AttributeReference if a.name == "s" => a }) assert(structAttrs.nonEmpty, "pushed filter should reference struct column s") - val prunedStructType = structAttrs.head.dataType.asInstanceOf[StructType] - assert(prunedStructType.fieldNames.toSeq == Seq("a"), - s"struct column in pushed filter should be pruned to struct<a> but was $prunedStructType") + val structType = structAttrs.head.dataType.asInstanceOf[StructType] + assert(structType.fieldNames.toSeq == Seq("a", "b"), + s"struct column in pushed filter should keep type struct<a, b> but was $structType") } - test("pushedFilters drops filters referencing pruned nested struct fields") { + test("pushedFilters keep filters referencing pruned nested struct fields") { // Disable constraint propagation so IsNotNull(s.a) is not added as a post-scan // filter (it would keep field a alive in the struct). withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { val df = spark.read.format(classOf[NestedSchemaDataSourceV2].getName).load() - // Filter on s.a but select only s.b. Column pruning narrows s to struct<b>, - // so the pushed filter on s.a can't be remapped and should be dropped. + // Filter on s.a but select only s.b. Column pruning narrows the scan output's s to struct<b>, + // but pushedFilters keeps the fully-pushed filter on s.a (against the relation schema) so a + // later scan merge can re-enforce it. val q = df.filter($"s.a" > 3).select($"s.b") checkAnswer(q, (4 until 10).map(i => Row(-i))) @@ -1306,8 +1560,8 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper filter.collect { case a: AttributeReference if a.name == "s" => a } .flatMap(_.dataType.asInstanceOf[StructType].fieldNames) } - assert(!referencedStructFields.contains("a"), - "pushedFilters should not reference pruned nested field a") + assert(referencedStructFields.contains("a"), + "pushedFilters should keep the filter referencing nested field a") } } @@ -1336,6 +1590,32 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "Canonicalized instances with equivalent pushedFilters should be equal") } + test("scan canonicalization distinguishes different pushedFilters") { + // Negative counterpart to the test above: two scans of the same relation whose pushedFilters + // differ must NOT canonicalize equal. This inequality is what lets MergeSubplans tell the two + // scans apart; if they compared equal it could treat them as identical and apply one side's + // filter to both -- a wrong-answer bug. + val table = new SimpleDataSourceV2().getTable(CaseInsensitiveStringMap.empty()) + + val relation1 = DataSourceV2Relation.create( + table, None, None, CaseInsensitiveStringMap.empty()) + val relation2 = DataSourceV2Relation.create( + table, None, None, CaseInsensitiveStringMap.empty()) + val scan1 = relation1.table.asReadable.newScanBuilder(relation1.options).build() + val scan2 = relation2.table.asReadable.newScanBuilder(relation2.options).build() + + val filter1 = CatalystGreaterThan(relation1.output.head, CatalystLiteral(3)) + val filter2 = CatalystGreaterThan(relation2.output.head, CatalystLiteral(5)) + + val scanRelation1 = DataSourceV2ScanRelation(relation1, scan1, relation1.output, + pushedFilters = Seq(filter1)) + val scanRelation2 = DataSourceV2ScanRelation(relation2, scan2, relation2.output, + pushedFilters = Seq(filter2)) + + assert(scanRelation1.canonicalized != scanRelation2.canonicalized, + "Canonicalized instances with different pushedFilters must not be equal") + } + test("pushedFilters excludes non-deterministic filters") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() // i > 3 is pushable and deterministic; rand() > 0.5 is non-deterministic and not pushable. @@ -1399,21 +1679,66 @@ class DataSourceV2Suite extends SharedSparkSession with AdaptiveSparkPlanHelper "non-deterministic filter should be retained as a post-scan Filter") } - test("pushedFilters drops filters referencing pruned columns") { + test("pushedFilters keep filters referencing pruned columns") { // Disable constraint propagation so IsNotNull(i) is not added (it would keep // column i in the scan output). This simulates a connector that pushes IsNotNull. withSQLConf(SQLConf.CONSTRAINT_PROPAGATION_ENABLED.key -> "false") { val df = spark.read.format(classOf[AdvancedDataSourceV2].getName).load() - // i > 3 is fully pushed; selecting only j causes column pruning to drop i. + // i > 3 is fully pushed; selecting only j causes column pruning to drop i from the scan + // output, but pushedFilters keeps the fully-pushed filter on i (against the relation schema) + // so a later scan merge can re-enforce it. val q = df.filter($"i" > 3).select($"j") checkAnswer(q, (4 until 10).map(i => Row(-i))) val scanRelation = getScanRelation(q) assert(!scanRelation.output.exists(_.name == "i"), "column i should be pruned from scan output") - assert(scanRelation.pushedFilters.isEmpty, - "pushedFilters should drop filters referencing pruned columns") + val referencedCols = scanRelation.pushedFilters.flatMap(_.references.map(_.name)).toSet + assert(referencedCols.contains("i"), + "pushedFilters should keep the filter referencing the pruned column i") + } + } + + test("SPARK-57225: DSv2 source without batch write capability throws clear error") { + val cls = classOf[ReadOnlyV2DataSource].getName + val df = spark.range(1).toDF("i") + checkError( + exception = intercept[AnalysisException] { + df.write.format(cls).mode("append").save() + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`read_only_v2_test`", + "operation" -> "batch write" + ) + ) + + checkError( + exception = intercept[AnalysisException] { + df.write.format(cls).save() + }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`read_only_v2_test`", + "operation" -> "batch write" + ) + ) + } + + test("SPARK-57225: DSv2 source with V1_BATCH_WRITE still falls back to V1 path") { + val cls = classOf[V1BatchWriteV2DataSource].getName + val df = spark.range(1).toDF("i") + // V1_BATCH_WRITE sources should fall through to saveToV1SourceCommand, which calls + // DataSource.planForWriting. Since our test source doesn't implement + // CreatableRelationProvider or FileFormat, planForWriting hits the case _ branch + // and throws INTERNAL_ERROR with "does not allow create table as select". + // This proves V1_BATCH_WRITE reached the V1 write path. + val ex = intercept[SparkException] { + df.write.format(cls).mode("append").save() } + assert(ex.getCondition == "INTERNAL_ERROR", + "V1_BATCH_WRITE source should reach DataSource.planForWriting (V1 path)") + assert(ex.getMessage.contains("does not allow create table as select")) } } @@ -1633,6 +1958,89 @@ class AdvancedBatch(val filters: Array[Filter], val requiredSchema: StructType) } } +class AdvancedDataSourceV2WithSparkPostPushdownAdjustments extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithSparkPostPushdownAdjustments() + } + } +} + +class AdvancedScanBuilderWithSparkPostPushdownAdjustments + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def reflectsFullyPushedDownFilters(): Boolean = false + + override def pruneColumns(requiredSchema: StructType): Unit = { + this.requiredSchema = schemaWithPushedFilterColumns(requiredSchema) + } + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(10L) + override def columnStats(): java.util.Map[NamedReference, ColumnStatistics] = { + val stats = new java.util.HashMap[NamedReference, ColumnStatistics]() + stats.put(FieldReference.column("i"), new ColumnStatistics { + override def distinctCount(): OptionalLong = OptionalLong.of(10L) + override def min(): Optional[AnyRef] = Optional.of(Int.box(0)) + override def max(): Optional[AnyRef] = Optional.of(Int.box(9)) + }) + stats + } + } + + private def schemaWithPushedFilterColumns(prunedSchema: StructType): StructType = { + val existingFields = prunedSchema.fieldNames.toSet + val pushedFilterFields = filters.collect { + case GreaterThan(columnName: String, _) => columnName + }.distinct + .filterNot(existingFields.contains) + .flatMap { columnName => + TestingV2Source.schema.fields.find(_.name == columnName) + } + StructType(prunedSchema.fields ++ pushedFilterFields) + } +} + +class AdvancedDataSourceV2WithBestEffortSparkPostPushdownAdjustments extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithBestEffortSparkPostPushdownAdjustments() + } + } +} + +class AdvancedScanBuilderWithBestEffortSparkPostPushdownAdjustments + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def reflectsFullyPushedDownFilters(): Boolean = false + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.empty() + override def numRows(): OptionalLong = OptionalLong.of(10L) + } +} + +class AdvancedDataSourceV2WithScanStats extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new AdvancedScanBuilderWithReportedStatistics() + } + } +} + +class AdvancedScanBuilderWithReportedStatistics + extends AdvancedScanBuilder with SupportsReportStatistics { + + override def estimateStatistics(): Statistics = new Statistics { + override def sizeInBytes(): OptionalLong = OptionalLong.of(32L) + override def numRows(): OptionalLong = OptionalLong.of(2L) + } +} + class AdvancedDataSourceV2WithV2Filter extends TestingV2Source { override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable { @@ -2110,6 +2518,56 @@ object SpecificReaderFactory extends PartitionReaderFactory { class SchemaReadAttemptException(m: String) extends RuntimeException(m) +/** + * A writable data source whose batch write both fails to commit and then fails to abort. This + * drives the exact branch in `WriteToDataSourceV2Exec` that raises `WRITING_JOB_FAILED`: the + * write's `commit` throws, and the follow-up `abort` also throws, so the original failure is + * wrapped rather than re-thrown. The per-task `DataWriter` succeeds so the failure is driver-side + * and deterministic (no dependence on task scheduling). + */ +class CommitAndAbortFailingDataSource extends TestingV2Source { + + override def getTable(options: CaseInsensitiveStringMap): Table = new SimpleBatchTable + with SupportsWrite { + + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, + TableCapability.TRUNCATE) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new SimpleScanBuilder { + override def planInputPartitions(): Array[InputPartition] = Array.empty + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = new WriteBuilder { + override def build(): Write = new Write { + override def toBatch: BatchWrite = new BatchWrite { + override def createBatchWriterFactory(info: PhysicalWriteInfo): DataWriterFactory = + CommitAndAbortFailingDataSource.WriterFactory + + override def commit(messages: Array[WriterCommitMessage]): Unit = + throw new RuntimeException("commit failed") + + override def abort(messages: Array[WriterCommitMessage]): Unit = + throw new RuntimeException("abort failed") + } + } + } + } +} + +object CommitAndAbortFailingDataSource { + object WriterFactory extends DataWriterFactory { + override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = + new DataWriter[InternalRow] { + override def write(record: InternalRow): Unit = {} + override def commit(): WriterCommitMessage = null + override def abort(): Unit = {} + override def close(): Unit = {} + } + } +} + class SimpleWriteOnlyDataSource extends SimpleWritableDataSource { override def getTable(options: CaseInsensitiveStringMap): Table = { @@ -2191,3 +2649,41 @@ class InvalidDataSource extends TestingV2Source { override def getTable(options: CaseInsensitiveStringMap): Table = null } + +/** + * A read-only DSv2 data source that only declares BATCH_READ capability. + * Used to test that write attempts produce a clear error instead of falling through to V1. + */ +class ReadOnlyV2DataSource extends SimpleTableProvider { + override def getTable(options: CaseInsensitiveStringMap): Table = { + new ReadOnlyV2Table + } +} + +class ReadOnlyV2Table extends Table with SupportsRead { + override def name(): String = "read_only_v2_test" + override def schema(): StructType = new StructType().add("i", "long") + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(TableCapability.BATCH_READ) + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + throw new UnsupportedOperationException("scan not needed for write tests") +} + +/** + * A DSv2 data source that declares V1_BATCH_WRITE, simulating the JDBC pattern. + * Used to verify that V1 fallback is preserved for sources that opt into it. + */ +class V1BatchWriteV2DataSource extends SimpleTableProvider { + override def getTable(options: CaseInsensitiveStringMap): Table = { + new V1BatchWriteV2Table + } +} + +class V1BatchWriteV2Table extends Table with SupportsRead { + override def name(): String = "v1_batch_write_test" + override def schema(): StructType = new StructType().add("i", "long") + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(TableCapability.BATCH_READ, TableCapability.V1_BATCH_WRITE) + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + throw new UnsupportedOperationException("scan not needed for write tests") +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeleteFromTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeleteFromTableSuiteBase.scala index b75c0fc1c9474..f8cc82eb8a89c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeleteFromTableSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeleteFromTableSuiteBase.scala @@ -22,7 +22,7 @@ import org.apache.spark.sql.{AnalysisException, Row} import org.apache.spark.sql.QueryTest.withQueryExecutionsCaptured import org.apache.spark.sql.catalyst.expressions.CheckInvariant import org.apache.spark.sql.catalyst.plans.logical.Filter -import org.apache.spark.sql.connector.catalog.{Aborted, Committed, InMemoryRowLevelOperationTable, InMemoryTable, InMemoryTruncatableOnlyTable, InMemoryTruncatableOnlyTableCatalog, TableCatalog} +import org.apache.spark.sql.connector.catalog.{Aborted, Committed, InMemoryRowLevelOperationTable, InMemoryTable, InMemoryTruncatableOnlyTable, InMemoryTruncatableOnlyTableCatalog, TableWritePrivilege} import org.apache.spark.sql.connector.write.DeleteSummary import org.apache.spark.sql.execution.datasources.v2.{DeleteFromTableExec, ReplaceDataExec, TruncateTableExec, WriteDeltaExec} import org.apache.spark.sql.internal.SQLConf @@ -1047,8 +1047,10 @@ abstract class DeleteFromTableSuiteBase extends RowLevelOperationSuiteBase { // replace the row-level plan with a filter-only deleteWhere call via // OptimizeMetadataOnlyDeleteFromTable (canDeleteWhere returns false for LessThan). checkRowLevelOperationOptions( - sql(s"DELETE FROM $tableNameAsString WITH (`write.split-size` = 10) WHERE salary < 200"), - "write.split-size" -> "10") + sql(s"DELETE FROM $tableNameAsString WITH " + + s"(`load-option` = 'load-value', `write-option` = 'write-value') WHERE salary < 200"), + "load-option" -> "load-value", + "write-option" -> "write-value") checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -1101,16 +1103,20 @@ abstract class DeleteFromTableSuiteBase extends RowLevelOperationSuiteBase { // dep = 'hr' is an EqualTo on the partition column. OptimizeMetadataOnlyDeleteFromTable // converts the row-level plan to a deleteWhere call. Verify options flow into the exec. val Seq(qe) = withQueryExecutionsCaptured(spark) { - sql(s"DELETE FROM $tableNameAsString WITH (`write.split-size` = 10) WHERE dep = 'hr'") + sql(s"DELETE FROM $tableNameAsString WITH " + + s"(`load-option` = 'load-value', `write-option` = 'write-value') WHERE dep = 'hr'") } val exec = qe.executedPlan.collectFirst { case e: DeleteFromTableExec => e }.getOrElse(fail("expected DeleteFromTableExec for the metadata-only deleteWhere path")) - assert(exec.options.get("write.split-size") === "10", - "options must reach DeleteFromTableExec on the deleteWhere path") - assert(exec.table.asInstanceOf[InMemoryRowLevelOperationTable].lastDeleteOptions - .get("write.split-size") === "10", - "options must be forwarded to SupportsDeleteV2.deleteWhere(predicates, options)") + assert(exec.options.get("load-option") === "load-value") + assert(exec.options.get("write-option") === "write-value") + val v2Table = exec.table.asInstanceOf[InMemoryRowLevelOperationTable] + assert(v2Table.lastDeleteOptions.get("load-option") === "load-value") + assert(v2Table.lastDeleteOptions.get("write-option") === "write-value") + assertLastTransactionWriteLoadOptions( + "load-option" -> "load-value", + "write-option" -> "write-value") checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -1146,27 +1152,40 @@ abstract class DeleteFromTableSuiteBase extends RowLevelOperationSuiteBase { // is planned as TruncateTableExec (not DeleteFromTableExec). withSQLConf( "spark.sql.catalog.trunccat" -> - classOf[InMemoryTruncatableOnlyTableCatalog].getName) { + classOf[InMemoryTruncatableOnlyTableCatalog].getName, + "spark.sql.catalog.trunccat.tableStateOptionKeys" -> "load-option") { withTable("trunccat.ns.tbl") { sql("CREATE TABLE trunccat.ns.tbl (pk INT NOT NULL, dep STRING) USING foo") sql("INSERT INTO trunccat.ns.tbl VALUES (1, 'hr'), (2, 'software')") + val truncCatalog = spark.sessionState.catalogManager + .catalog("trunccat").asInstanceOf[InMemoryTruncatableOnlyTableCatalog] + truncCatalog.resetLoadTableCalls() val Seq(qe) = withQueryExecutionsCaptured(spark) { - sql("DELETE FROM trunccat.ns.tbl WITH (`write.split-size` = 10)") + sql("DELETE FROM trunccat.ns.tbl WITH " + + "(`load-option` = 'load-value', `write-option` = 'write-value')") } val exec = qe.executedPlan.collectFirst { case e: TruncateTableExec => e }.getOrElse(fail("expected TruncateTableExec for the truncate path")) - assert(exec.options.get("write.split-size") === "10", - "options must reach TruncateTableExec") + assert(exec.options.get("load-option") === "load-value") + assert(exec.options.get("write-option") === "write-value") + + val targetLoads = truncCatalog.loadTableCalls.filter { + case (context, options) => + context.writePrivileges() === java.util.Set.of(TableWritePrivilege.DELETE) && + options.get("load-option") == "load-value" && + options.get("write-option") == null && + options.size() == 1 + } + assert(targetLoads.nonEmpty, "target loadTable did not receive truncate state options") - val truncTable = spark.sessionState.catalogManager - .catalog("trunccat").asInstanceOf[TableCatalog] + val truncTable = truncCatalog .loadTable(org.apache.spark.sql.connector.catalog.Identifier.of( Array("ns"), "tbl")) .asInstanceOf[InMemoryTruncatableOnlyTable] - assert(truncTable.lastTruncateOptions.get("write.split-size") === "10", - "options must be forwarded to TruncatableTable.truncateTable(options)") + assert(truncTable.lastTruncateOptions.get("load-option") === "load-value") + assert(truncTable.lastTruncateOptions.get("write-option") === "write-value") checkAnswer(spark.table("trunccat.ns.tbl"), Nil) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..ce104f5631f11 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -0,0 +1,50 @@ +/* + * 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.connector + +import org.apache.spark.sql.Row + +class DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = super.extraTableProps + props.put("supports-deltas", "true") + props + } + + test("delete does not use group filtering when the group key is not scanned") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + // a delta-based delete scans the row ID, the condition columns and the metadata columns, so + // `dep` is not read and the scan cannot declare it as a filter attribute + assertNoCatalystGroupFilter(executedPlan) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala index c1741cac8ad3c..98ed8bf0b3d0d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala @@ -33,9 +33,12 @@ abstract class DistributionAndOrderingSuiteBase extends SharedSparkSession with BeforeAndAfter with AdaptiveSparkPlanHelper { import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._ + /** The catalog implementation `testcat` is registered with. */ + protected def catalogClassName: String = classOf[InMemoryCatalog].getName + override def beforeAll(): Unit = { super.beforeAll() - spark.conf.set("spark.sql.catalog.testcat", classOf[InMemoryCatalog].getName) + spark.conf.set("spark.sql.catalog.testcat", catalogClassName) } override def afterAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala index 81b2ea59cad9e..c963d91005f10 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GeneratedColumnWriteSuite.scala @@ -18,10 +18,12 @@ package org.apache.spark.sql.connector import org.apache.spark.SparkRuntimeException -import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.{AnalysisException, DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.QueryPlanningTracker import org.apache.spark.sql.catalyst.expressions.CheckInvariant -import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog +import org.apache.spark.sql.catalyst.util.GeneratedColumn +import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryRowLevelOperationTableCatalog} +import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.internal.SQLConf @@ -97,6 +99,17 @@ class GeneratedColumnWriteSuite extends QueryTest with DatasourceV2SQLBase { } } + private def generationExpressionOf(table: String, column: String): Option[String] = { + val loaded = catalog("testcat").asTableCatalog.loadTable(Identifier.of(Array(), table)) + Option(loaded.columns().find(_.name == column).get.generationExpression()) + } + + private def assertNoGenerationMetadata(df: DataFrame, column: String): Unit = { + val metadata = df.schema(column).metadata + assert(!metadata.contains(GeneratedColumn.GENERATION_EXPRESSION_METADATA_KEY), + s"generation expression leaked into the schema of $column: ${metadata.json}") + } + testGeneratedColumnWrite("append_data_v2") { table => import testImplicits._ Seq(Tuple5(1L, "foo", "2020-10-11 12:30:30", 100, "2020-11-12")) @@ -979,6 +992,68 @@ class GeneratedColumnWriteSuite extends QueryTest with DatasourceV2SQLBase { } } + test("generation expression is not exposed in the read schema") { + val tblName = "my_tab" + withTable(s"testcat.$tblName") { + sql(s"""CREATE TABLE testcat.$tblName( + | id INT, + | doubled INT GENERATED ALWAYS AS (id * 2) + |) USING foo""".stripMargin) + // The table still declares the generated column; only the internal metadata is hidden. + assert(generationExpressionOf(tblName, "doubled").contains("id * 2")) + assertNoGenerationMetadata(spark.table(s"testcat.$tblName"), "doubled") + assertNoGenerationMetadata(spark.read.table(s"testcat.$tblName"), "doubled") + assertNoGenerationMetadata(sql(s"SELECT * FROM testcat.$tblName"), "doubled") + assertNoGenerationMetadata(spark.readStream.table(s"testcat.$tblName"), "doubled") + } + } + + test("CTAS from a table with generated columns does not create generated columns") { + val srcName = "src_tab" + val dstName = "dst_tab" + withTable(s"testcat.$srcName", s"testcat.$dstName") { + sql(s"""CREATE TABLE testcat.$srcName( + | id INT, + | doubled INT GENERATED ALWAYS AS (id * 2) + |) USING foo""".stripMargin) + sql(s"INSERT INTO testcat.$srcName (id) VALUES (1)") + sql(s"CREATE TABLE testcat.$dstName USING foo AS SELECT * FROM testcat.$srcName") + // The query output is plain data, so the new table must not inherit the generated column. + assert(generationExpressionOf(dstName, "doubled").isEmpty) + // Confirms it behaves as an ordinary column: a value the expression would not produce is + // written through instead of failing a generated column constraint. + sql(s"INSERT INTO testcat.$dstName VALUES (5, 999)") + checkAnswer(spark.table(s"testcat.$dstName"), Row(1, 2) :: Row(5, 999) :: Nil) + } + } + + test("streaming write to a new table does not create generated columns") { + val srcName = "src_tab" + val dstName = "dst_tab" + withTable(s"testcat.$srcName", s"testcat.$dstName") { + withTempDir { checkpointDir => + sql(s"""CREATE TABLE testcat.$srcName( + | id INT, + | doubled INT GENERATED ALWAYS AS (id * 2) + |) USING foo""".stripMargin) + sql(s"INSERT INTO testcat.$srcName (id) VALUES (1)") + // toTable creates the target from the streaming DataFrame's schema, which must not carry + // the source's generation expression over: a new table with a generated column would make + // this very write unsupported. + val query = spark.readStream.table(s"testcat.$srcName").writeStream + .option("checkpointLocation", checkpointDir.getAbsolutePath) + .toTable(s"testcat.$dstName") + try { + query.processAllAvailable() + } finally { + query.stop() + } + assert(generationExpressionOf(dstName, "doubled").isEmpty) + checkAnswer(spark.table(s"testcat.$dstName"), Row(1, 2)) + } + } + } + test("streaming write with generated columns is blocked") { import testImplicits._ val tblName = "my_tab" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..34b43e2354abc --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -0,0 +1,60 @@ +/* + * 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.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.connector.catalog.InMemoryTable +import org.apache.spark.sql.connector.write.DeleteSummary + +class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + test("delete runtime group filtering with SupportsRuntimeCatalystFiltering") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // only pk 1 matches, so hr is rewritten and its other row (pk 3) is copied over + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "salary INT, dep STRING", groups = Seq("hr"))) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + + checkReplacedPartitions(Seq("hr")) + checkDeleteMetrics(numDeletedRows = 1, numCopiedRows = 1) + } + + private def checkDeleteMetrics(numDeletedRows: Long, numCopiedRows: Long): Unit = { + val t = catalog.loadTable(ident).asInstanceOf[InMemoryTable] + val summary = t.commits.last.writeSummary.get.asInstanceOf[DeleteSummary] + assert(summary.numDeletedRows() === numDeletedRows, + s"Expected numDeletedRows=$numDeletedRows, got ${summary.numDeletedRows()}") + assert(summary.numCopiedRows() === numCopiedRows, + s"Expected numCopiedRows=$numCopiedRows, got ${summary.numCopiedRows()}") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index ceaac9729ddd1..5d564b2136557 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Literal, TransformExpression} import org.apache.spark.sql.catalyst.plans.physical -import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog} +import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryCatalystRuntimeFilterCatalog, InMemoryTableCatalog} import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.connector.distributions.Distributions import org.apache.spark.sql.connector.expressions._ @@ -45,8 +45,287 @@ import org.apache.spark.sql.functions.{col, max} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf._ import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest -class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with ExplainSuiteHelper { +abstract class KeyGroupedPartitioningSuiteBase extends DistributionAndOrderingSuiteBase { + + protected val emptyProps: java.util.Map[String, String] = { + Collections.emptyMap[String, String] + } + + protected val items: String = "items" + protected val itemsColumns: Array[Column] = Array( + Column.create("id", LongType), + Column.create("name", StringType), + Column.create("price", FloatType), + Column.create("arrive_time", TimestampType)) + + protected val purchases: String = "purchases" + protected val purchasesColumns: Array[Column] = Array( + Column.create("item_id", LongType), + Column.create("price", FloatType), + Column.create("time", TimestampType)) + + protected def createTable( + table: String, + columns: Array[Column], + partitions: Array[Transform], + ordering: Array[SortOrder] = Array.empty, + catalog: InMemoryTableCatalog = catalog): Unit = { + catalog.createTable(Identifier.of(Array("ns"), table), + columns, partitions, emptyProps, Distributions.unspecified(), ordering, None, None, + numRowsPerSplit = 1) + } + + protected def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { + // here we skip collecting shuffle operators that are not associated with SMJ + collect(plan) { + case s: SortMergeJoinExec => s + }.flatMap(smj => + collect(smj) { + case s: ShuffleExchangeExec => s + }) + }.toSet.toSeq + + protected def collectGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { + // here we skip collecting group-partition operators that are not associated with SMJ + collect(plan) { + case s: SortMergeJoinExec => s + }.flatMap(smj => + collect(smj) { + case g: GroupPartitionsExec => g + }) + }.toSet.toSeq + + protected def collectScans(plan: SparkPlan): Seq[BatchScanExec] = { + collect(plan) { case s: BatchScanExec => s } + } + +} + +/** + * Tests for runtime filtering under a storage-partitioned join, whose outcome depends on how the + * scan takes runtime filters. + */ +trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSuiteBase { + + /** + * Helper method to verify that filteredPartitions contains the expected number of + * Some and None values. This is used to verify that dynamic partition filtering + * properly fills filtered-out partitions with None. + */ + private def assertFilteredPartitions( + scans: Seq[BatchScanExec], + expectedTotalPartitions: Seq[Int], + expectedFilteredOutPartitions: Seq[Int]): Unit = { + assert(scans.size === expectedTotalPartitions.size, + s"Expected ${expectedTotalPartitions.size} scans but got ${scans.size}") + + scans.zip(expectedTotalPartitions).zip(expectedFilteredOutPartitions).foreach { + case ((scan, expectedTotal), expectedFiltered) => + val filtered = scan.filteredPartitions + assert(filtered.size === expectedTotal, + s"Expected $expectedTotal total partitions but got ${filtered.size}") + + val noneCount = filtered.count(_.isEmpty) + assert(noneCount === expectedFiltered, + s"Expected $expectedFiltered None values but got $noneCount") + + val someCount = filtered.count(_.isDefined) + assert(someCount === (expectedTotal - expectedFiltered), + s"Expected ${expectedTotal - expectedFiltered} Some values but got $someCount") + } + } + + test("data source partitioning + dynamic partition filtering") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { + // number of unique partitions changed after dynamic filtering - the gap should be filled + // with empty partitions and the job should still succeed + var df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price > 40.0") + + var shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + var scans = collectScans(df.queryExecution.executedPlan) + assert(scans.forall(_.outputPartitioning.numPartitions === 5)) + var groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) + + checkAnswer(df, Seq(Row(131))) + + // Verify that filteredPartitions contains None for filtered-out partitions. + // After DPF with filter i.price > 40.0, only id=1 survives on items side. + // The purchases side should be pruned to only item_id=1. + // purchases: 5 total partitions (3 for id=1, 1 for id=2, 1 for id=3) + // After DPF: 3 Some (id=1), 2 None (id=2, id=3) + assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 2)) + + // dynamic filtering doesn't change partitioning so storage-partitioned join should kick + // in + df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price >= 10.0") + + shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + scans = collectScans(df.queryExecution.executedPlan) + assert(scans.forall(_.outputPartitioning.numPartitions === 5)) + groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) + + checkAnswer(df, Seq(Row(303.5))) + + // With filter i.price >= 10.0, all ids (1, 2, 3) survive, + // so no partitions should be filtered out + assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 0)) + } + } + } + } + + test("SPARK-42038: partially clustered: with dynamic partition filtering") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp)), " + + s"(4, 'dd', 18.0, cast('2023-01-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(1, 50.0, cast('2020-01-15' as timestamp)), " + + s"(1, 55.0, cast('2020-01-15' as timestamp)), " + + s"(1, 60.0, cast('2020-01-15' as timestamp)), " + + s"(1, 65.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp)), " + + s"(5, 25.0, cast('2023-01-01' as timestamp)), " + + s"(5, 26.0, cast('2023-01-01' as timestamp)), " + + s"(5, 28.0, cast('2023-01-01' as timestamp)), " + + s"(6, 50.0, cast('2023-02-01' as timestamp)), " + + s"(6, 50.0, cast('2023-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + Seq(("true", 15), ("false", 6)).foreach { + case (enable, expected) => + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { + + // When partition values are pushed down, storage-partitioned join fills the missing + // partitions & splits after dynamic filtering with empty partitions & splits. + val df = sql(s"SELECT sum(p.price) from " + + s"testcat.ns.$purchases p, testcat.ns.$items i WHERE " + + s"p.item_id = i.id AND p.price < 45.0") + + checkAnswer(df, Seq(Row(213.5))) + val shuffles = collectShuffles(df.queryExecution.executedPlan) + val scans = collectScans(df.queryExecution.executedPlan) + val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(scans.map(_.outputPartitioning.numPartitions) === Seq(14, 6)) + if (pushDownValues) { + assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") + assert(groupPartitions.forall(_.outputPartitioning.numPartitions === expected)) + } else { + assert(shuffles.nonEmpty, + "should contain shuffle when not pushing down partition values") + assert(groupPartitions.isEmpty) + } + + // Verify filteredPartitions for DPF. + // After filter p.price < 45.0, purchases has item_ids {1, 2, 3, 5}. + // Items side should be pruned to these ids. Since items has {1, 2, 3, 4}, + // id=4 should be filtered out. + // purchases: 14 total, all kept (0 None) - no DPF on probe side + // items: 6 total, id=4 filtered (1 None) + assertFilteredPartitions(scans, Seq(14, 6), Seq(0, 1)) + } + } + } + } + + test("SPARK-45652: SPJ should handle empty partition after dynamic filtering") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") { + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-15' as timestamp)), " + + s"(1, 45.0, cast('2020-01-15' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + Seq(true, false).foreach { pushDownValues => + Seq(true, false).foreach { partiallyClustered => { + withSQLConf( + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> + partiallyClustered.toString, + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { + // The dynamic filtering effectively filtered out all the partitions + val df = sql(s"SELECT p.price from testcat.ns.$items i, testcat.ns.$purchases p " + + "WHERE i.id = p.item_id AND i.price > 50.0") + checkAnswer(df, Seq.empty) + } + } + } + } + } + } +} + +@ExtendedSQLTest +class KeyGroupedPartitioningSuite + extends KeyGroupedPartitioningSuiteBase with ExplainSuiteHelper { private val functions = Seq( UnboundYearsFunction, UnboundDaysFunction, @@ -68,9 +347,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with catalog.clearFunctions() } - private val emptyProps: java.util.Map[String, String] = { - Collections.emptyMap[String, String] - } private val table: String = "tbl" private val columns: Array[Column] = Array( @@ -256,17 +532,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with assert(expectedPartitioning == scan.outputPartitioning) } - private def createTable( - table: String, - columns: Array[Column], - partitions: Array[Transform], - ordering: Array[SortOrder] = Array.empty, - catalog: InMemoryTableCatalog = catalog): Unit = { - catalog.createTable(Identifier.of(Array("ns"), table), - columns, partitions, emptyProps, Distributions.unspecified(), ordering, None, None, - numRowsPerSplit = 1) - } - private val customers: String = "customers" private val customersColumns: Array[Column] = Array( Column.create("customer_name", StringType), @@ -341,59 +606,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } - protected def collectShuffles(plan: SparkPlan): Seq[ShuffleExchangeLike] = { - // here we skip collecting shuffle operators that are not associated with SMJ - collect(plan) { - case s: SortMergeJoinExec => s - }.flatMap(smj => - collect(smj) { - case s: ShuffleExchangeExec => s - }) - }.toSet.toSeq - - protected def collectGroupPartitions(plan: SparkPlan): Seq[GroupPartitionsExec] = { - // here we skip collecting shuffle operators that are not associated with SMJ - collect(plan) { - case s: SortMergeJoinExec => s - }.flatMap(smj => - collect(smj) { - case g: GroupPartitionsExec => g - }) - }.toSet.toSeq - - private def collectScans(plan: SparkPlan): Seq[BatchScanExec] = { - collect(plan) { case s: BatchScanExec => s } - } - - /** - * Helper method to verify that filteredPartitions contains the expected number of - * Some and None values. This is used to verify that dynamic partition filtering - * properly fills filtered-out partitions with None. - */ - private def assertFilteredPartitions( - scans: Seq[BatchScanExec], - expectedTotalPartitions: Seq[Int], - expectedFilteredOutPartitions: Seq[Int]): Unit = { - assert(scans.size === expectedTotalPartitions.size, - s"Expected ${expectedTotalPartitions.size} scans but got ${scans.size}") - - scans.zip(expectedTotalPartitions).zip(expectedFilteredOutPartitions).foreach { - case ((scan, expectedTotal), expectedFiltered) => - val filtered = scan.filteredPartitions - assert(filtered.size === expectedTotal, - s"Expected $expectedTotal total partitions but got ${filtered.size}") - - val noneCount = filtered.count(_.isEmpty) - assert(noneCount === expectedFiltered, - s"Expected $expectedFiltered None values but got $noneCount") - - val someCount = filtered.count(_.isDefined) - assert(someCount === (expectedTotal - expectedFiltered), - s"Expected ${expectedTotal - expectedFiltered} Some values but got $someCount") - } - } - - test("partitioned join: exact distribution (same number of buckets) from both sides") { val customers_partitions = Array(bucket(4, "customer_id")) val orders_partitions = Array(bucket(4, "customer_id")) @@ -415,19 +627,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with testWithCustomersAndOrders(customers_partitions, Array.empty, 2, 0) } - private val items: String = "items" - private val itemsColumns: Array[Column] = Array( - Column.create("id", LongType), - Column.create("name", StringType), - Column.create("price", FloatType), - Column.create("arrive_time", TimestampType)) - - private val purchases: String = "purchases" - private val purchasesColumns: Array[Column] = Array( - Column.create("item_id", LongType), - Column.create("price", FloatType), - Column.create("time", TimestampType)) - private val details: String = "details" private val detailsColumns: Array[Column] = Array( Column.create("item_id", LongType), @@ -1199,7 +1398,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with enable <- Seq("true", "false") } yield { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> false.toString, SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { // The left side uses a key-grouped partitioning to satisfy the WINDOW function's @@ -1320,206 +1518,63 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with |""".stripMargin) checkAnswer(df, Seq(Row(1, 41.0f), Row(2, 10.0f), Row(3, 15.5f))) - // One GroupPartitionsExec per join child to align the partially-clustered - // partitions, and one above the join to group for the window. - val joinGP = collectGroupPartitions(df.queryExecution.executedPlan) - assert(joinGP.size === 2, - "expected 2 GroupPartitionsExec under the join") - val allGP = collectAllGroupPartitions(df.queryExecution.executedPlan) - assert(allGP.size === 3, - "expected 3 GroupPartitionsExec total (2 under join + 1 above for window)") - } - } - - test("SPARK-55848: checkpointed partially-clustered join with dedup") { - withTempDir { dir => - spark.sparkContext.setCheckpointDir(dir.getPath) - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - "(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - "(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - "(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") - - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - "(1, 42.0, cast('2020-01-01' as timestamp)), " + - "(1, 50.0, cast('2020-01-02' as timestamp)), " + - "(2, 11.0, cast('2020-01-01' as timestamp)), " + - "(3, 19.5, cast('2020-02-01' as timestamp))") - - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> true.toString) { - // Checkpoint the JOIN result (not the scan) so the checkpoint node carries the - // partially-clustered KeyGroupedPartitioning. The dedup on top must still insert - // the required grouping operator because partially-clustered partitioning does not - // satisfy ClusteredDistribution. - val joinedDf = sql( - s"""${selectWithMergeJoinHint("i", "p")} i.id, i.name, i.price - |FROM testcat.ns.$items i - |JOIN testcat.ns.$purchases p ON i.id = p.item_id""".stripMargin) - val checkpointedDf = joinedDf.checkpoint() - val df = checkpointedDf.select("id").distinct() - checkAnswer(df, Seq(Row(1), Row(2), Row(3))) - - val checkpointScans = collect(df.queryExecution.executedPlan) { - case r: RDDScanExec => r - } - assert(checkpointScans.exists(_.outputPartitioning match { - case kp: physical.KeyedPartitioning => !kp.isGrouped - case _ => false - }), "checkpoint (RDDScanExec) should have ungrouped KeyedPartitioning") - - val allGroupPartitions = collectAllGroupPartitions(df.queryExecution.executedPlan) - assert(allGroupPartitions.size === 1, - "expected 1 GroupPartitionsExec above the checkpointed join for dedup") - } - } - } - - test("data source partitioning + dynamic partition filtering") { - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", - SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") - - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp))") - - Seq(true, false).foreach { pushDownValues => - withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { - // number of unique partitions changed after dynamic filtering - the gap should be filled - // with empty partitions and the job should still succeed - var df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price > 40.0") - - var shuffles = collectShuffles(df.queryExecution.executedPlan) - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - var scans = collectScans(df.queryExecution.executedPlan) - assert(scans.forall(_.outputPartitioning.numPartitions === 5)) - var groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) - - checkAnswer(df, Seq(Row(131))) - - // Verify that filteredPartitions contains None for filtered-out partitions. - // After DPF with filter i.price > 40.0, only id=1 survives on items side. - // The purchases side should be pruned to only item_id=1. - // purchases: 5 total partitions (3 for id=1, 1 for id=2, 1 for id=3) - // After DPF: 3 Some (id=1), 2 None (id=2, id=3) - assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 2)) - - // dynamic filtering doesn't change partitioning so storage-partitioned join should kick - // in - df = sql(s"SELECT sum(p.price) from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price >= 10.0") - - shuffles = collectShuffles(df.queryExecution.executedPlan) - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - scans = collectScans(df.queryExecution.executedPlan) - assert(scans.forall(_.outputPartitioning.numPartitions === 5)) - groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === 3)) - - checkAnswer(df, Seq(Row(303.5))) - - // With filter i.price >= 10.0, all ids (1, 2, 3) survive, - // so no partitions should be filtered out - assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 0)) - } - } - } - } - - test("SPARK-42038: partially clustered: with dynamic partition filtering") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp)), " + - s"(4, 'dd', 18.0, cast('2023-01-01' as timestamp))") + // One GroupPartitionsExec per join child to align the partially-clustered + // partitions, and one above the join to group for the window. + val joinGP = collectGroupPartitions(df.queryExecution.executedPlan) + assert(joinGP.size === 2, + "expected 2 GroupPartitionsExec under the join") + val allGP = collectAllGroupPartitions(df.queryExecution.executedPlan) + assert(allGP.size === 3, + "expected 3 GroupPartitionsExec total (2 under join + 1 above for window)") + } + } - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(1, 50.0, cast('2020-01-15' as timestamp)), " + - s"(1, 55.0, cast('2020-01-15' as timestamp)), " + - s"(1, 60.0, cast('2020-01-15' as timestamp)), " + - s"(1, 65.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp)), " + - s"(5, 25.0, cast('2023-01-01' as timestamp)), " + - s"(5, 26.0, cast('2023-01-01' as timestamp)), " + - s"(5, 28.0, cast('2023-01-01' as timestamp)), " + - s"(6, 50.0, cast('2023-02-01' as timestamp)), " + - s"(6, 50.0, cast('2023-02-01' as timestamp))") + test("SPARK-55848: checkpointed partially-clustered join with dedup") { + withTempDir { dir => + spark.sparkContext.setCheckpointDir(dir.getPath) + val items_partitions = Array(identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + + "(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") - Seq(true, false).foreach { pushDownValues => - Seq(("true", 15), ("false", 6)).foreach { - case (enable, expected) => - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", - SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10", - SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, - SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> enable) { + val purchases_partitions = Array(identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(1, 42.0, cast('2020-01-01' as timestamp)), " + + "(1, 50.0, cast('2020-01-02' as timestamp)), " + + "(2, 11.0, cast('2020-01-01' as timestamp)), " + + "(3, 19.5, cast('2020-02-01' as timestamp))") - // storage-partitioned join should kick in and fill the missing partitions & splits - // after dynamic filtering with empty partitions & splits, respectively. - val df = sql(s"SELECT sum(p.price) from " + - s"testcat.ns.$purchases p, testcat.ns.$items i WHERE " + - s"p.item_id = i.id AND p.price < 45.0") + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> true.toString) { + // Checkpoint the JOIN result (not the scan) so the checkpoint node carries the + // partially-clustered KeyGroupedPartitioning. The dedup on top must still insert + // the required grouping operator because partially-clustered partitioning does not + // satisfy ClusteredDistribution. + val joinedDf = sql( + s"""${selectWithMergeJoinHint("i", "p")} i.id, i.name, i.price + |FROM testcat.ns.$items i + |JOIN testcat.ns.$purchases p ON i.id = p.item_id""".stripMargin) + val checkpointedDf = joinedDf.checkpoint() + val df = checkpointedDf.select("id").distinct() + checkAnswer(df, Seq(Row(1), Row(2), Row(3))) - checkAnswer(df, Seq(Row(213.5))) - val shuffles = collectShuffles(df.queryExecution.executedPlan) - val scans = collectScans(df.queryExecution.executedPlan) - val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) - assert(scans.map(_.outputPartitioning.numPartitions) === Seq(14, 6)) - if (pushDownValues) { - assert(shuffles.isEmpty, "should not add shuffle for both sides of the join") - assert(groupPartitions.forall(_.outputPartitioning.numPartitions === expected)) - } else { - assert(shuffles.nonEmpty, - "should contain shuffle when not pushing down partition values") - assert(groupPartitions.isEmpty) - } + val checkpointScans = collect(df.queryExecution.executedPlan) { + case r: RDDScanExec => r + } + assert(checkpointScans.exists(_.outputPartitioning match { + case kp: physical.KeyedPartitioning => !kp.isGrouped + case _ => false + }), "checkpoint (RDDScanExec) should have ungrouped KeyedPartitioning") - // Verify filteredPartitions for DPF. - // After filter p.price < 45.0, purchases has item_ids {1, 2, 3, 5}. - // Items side should be pruned to these ids. Since items has {1, 2, 3, 4}, - // id=4 should be filtered out. - // purchases: 14 total, all kept (0 None) - no DPF on probe side - // items: 6 total, id=4 filtered (1 None) - assertFilteredPartitions(scans, Seq(14, 6), Seq(0, 1)) - } + val allGroupPartitions = collectAllGroupPartitions(df.queryExecution.executedPlan) + assert(allGroupPartitions.size === 1, + "expected 1 GroupPartitionsExec above the checkpointed join for dedup") } } } @@ -1830,7 +1885,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { partiallyClustered => Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -1972,7 +2026,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> @@ -2132,7 +2185,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> @@ -2310,7 +2362,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", @@ -2375,7 +2426,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach{ allowPushDown => Seq(true, false).foreach{ partiallyClustered => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> allowPushDown.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -2433,7 +2483,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { partiallyClustered => Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> @@ -2496,7 +2545,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with Seq(true, false).foreach { allowKeysSubsetOfPartitionKeys => withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString, SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> partiallyClustered.toString, @@ -2578,45 +2626,56 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } - test("SPARK-45652: SPJ should handle empty partition after dynamic filtering") { - withSQLConf( - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", - SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", - SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") { - val items_partitions = Array(identity("id")) - createTable(items, itemsColumns, items_partitions) - sql(s"INSERT INTO testcat.ns.$items VALUES " + - s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + - s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " + - s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + - s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " + - s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + test("SPARK-59025: shuffle one side and join keys are less than partition keys " + + "when the keyed side reports a PartitioningCollection") { + val items_partitions = Array(identity("id"), identity("name")) + createTable(items, itemsColumns, items_partitions) - val purchases_partitions = Array(identity("item_id")) - createTable(purchases, purchasesColumns, purchases_partitions) - sql(s"INSERT INTO testcat.ns.$purchases VALUES " + - s"(1, 42.0, cast('2020-01-01' as timestamp)), " + - s"(1, 44.0, cast('2020-01-15' as timestamp)), " + - s"(1, 45.0, cast('2020-01-15' as timestamp)), " + - s"(2, 11.0, cast('2020-01-01' as timestamp)), " + - s"(3, 19.5, cast('2020-02-01' as timestamp))") + // 4 distinct (id, name) partition keys but only 3 distinct ids, so grouping by the join + // key must reduce the keyed side from 4 partitions to 3. + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(1, 'ab', 30.0, cast('2020-01-02' as timestamp)), " + + "(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + "(4, 'cc', 15.5, cast('2020-02-01' as timestamp))") - Seq(true, false).foreach { pushDownValues => - Seq(true, false).foreach { partiallyClustered => { - withSQLConf( - SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> - partiallyClustered.toString, - SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> pushDownValues.toString) { - // The dynamic filtering effectively filtered out all the partitions - val df = sql(s"SELECT p.price from testcat.ns.$items i, testcat.ns.$purchases p " + - "WHERE i.id = p.item_id AND i.price > 50.0") - checkAnswer(df, Seq.empty) - } - } - } - } + createTable(purchases, purchasesColumns, Array.empty) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(1, 42.0, cast('2020-01-01' as timestamp)), " + + "(1, 89.0, cast('2020-01-03' as timestamp)), " + + "(3, 19.5, cast('2020-02-01' as timestamp)), " + + "(5, 26.0, cast('2023-01-01' as timestamp))") + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + // Duplicating `id` under two aliases makes the projection report a + // `PartitioningCollection` of `KeyedPartitioning`s, so the keyed side's shuffle spec + // is a `ShuffleSpecCollection` wrapping a `KeyedShuffleSpec` with join key positions. + val df = sql( + s""" + |${selectWithMergeJoinHint("i", "p")} + |id1, id2, name, i.price AS purchase_price, p.price AS sale_price + |FROM (SELECT id AS id1, id AS id2, name, price FROM testcat.ns.$items) i + |JOIN testcat.ns.$purchases p ON i.id1 = p.item_id + |ORDER BY id1, purchase_price, sale_price + |""".stripMargin) + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.size == 1, "only the non-keyed side should be shuffled") + val groupPartitions = collectGroupPartitions(df.queryExecution.executedPlan) + assert(groupPartitions.size == 1 && groupPartitions.head.joinKeyPositions.isDefined, + "the keyed side should be grouped by the join keys") + assert(groupPartitions.head.outputPartitioning.numPartitions == 3, + "the keyed side should be grouped down to 3 partitions") + assert(shuffles.head.outputPartitioning.numPartitions == 3, + "the shuffled side should match the 3 grouped partitions") + checkAnswer(df, Seq( + Row(1, 1, "ab", 30.0, 42.0), + Row(1, 1, "ab", 30.0, 89.0), + Row(1, 1, "aa", 40.0, 42.0), + Row(1, 1, "aa", 40.0, 89.0), + Row(3, 3, "bb", 10.0, 19.5))) } } @@ -2708,7 +2767,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with "(6, 50.0, cast('2023-02-01' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", @@ -3210,7 +3268,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with test("SPARK-55411: Fix ArrayIndexOutOfBoundsException when join keys " + "are less than cluster keys") { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", @@ -3286,7 +3343,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with test("SPARK-55535: Multi table join granular partition grouping") { withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { val items_partitions = Array(identity("id"), years("arrive_time")) @@ -3482,7 +3538,6 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with createTable(purchases, purchasesColumns, purchases_partitions) sql(s"INSERT INTO testcat.ns.$purchases VALUES (2, 10.0, cast('2021-01-01' as timestamp))") withSQLConf( - SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { val df = sql( @@ -4217,6 +4272,173 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } + test("SPARK-58988: v2 bucketed table with subset join keys joining v1 table") { + // The v2 table is partitioned by an extra identity key `dt` plus `bucket(16, c1)`, while the + // join is only on `c1`. allowKeysSubsetOfPartitionKeys lets the operation key `c1` be a subset + // of the partition keys `[dt, bucket(16, c1)]`, so EnsureRequirements projects the keyed side + // to `[bucket(16, c1)]`. v2BucketingShuffleEnabled then re-shuffles only the v1 side using that + // projected KeyedPartitioning. ShuffledJoin wraps the two output partitionings into a + // PartitioningCollection, which requires all KeyedPartitionings to share equal partitionKeys. + // The v2 side's keys are sorted by GroupPartitionsExec, while the keys re-used for the v1 side + // keep their first-occurrence order from createShuffleSpec, so the two sequences disagree and + // the collection construction used to fail. + val cols = Array( + Column.create("c1", LongType), + Column.create("c2", StringType), + Column.create("dt", StringType)) + val partitions = Array(identity("dt"), bucket(16, "c1")) + + createTable("iceberg_t2", cols, partitions) + sql("INSERT INTO testcat.ns.iceberg_t2 VALUES (2, 'cc', '2020'), (1, 'aa', '2021')") + + withTable("t1") { + sql("CREATE TABLE t1 (c1 BIGINT, c2 STRING) USING parquet") + sql("INSERT INTO t1 VALUES (1, 'aa'), (2, 'cc')") + + withSQLConf( + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql("SELECT * FROM testcat.ns.iceberg_t2 t0 JOIN t1 ON t0.c1 = t1.c1") + val plan = df.queryExecution.executedPlan + // Only the v1 side is re-shuffled; the v2 side is regrouped onto the join key instead. + assert(collectShuffles(plan).length == 1) + assert(collectGroupPartitions(plan).length == 1) + checkAnswer(df, Seq( + Row(1L, "aa", "2021", 1L, "aa"), + Row(2L, "cc", "2020", 2L, "cc"))) + } + } + } + + test("SPARK-59022: keyed shuffle follows the declared partition key order") { + val cols = Array( + Column.create("id", LongType), + Column.create("dt", StringType)) + val partitions = Array[Transform](identity("dt"), identity("id")) + + createTable("nt", cols, partitions) + // The scan reports its keys sorted on the full key: [('2020', 2), ('2021', 1)]. Narrowing them + // to `[id]` gives [2, 1], which is not sorted -- projecting a sorted sequence onto a subset of + // key positions does not preserve sortedness. + sql("INSERT INTO testcat.ns.nt VALUES (2, '2020'), (1, '2021')") + + withTable("t1") { + sql("CREATE TABLE t1 (id BIGINT, data STRING) USING parquet") + sql("INSERT INTO t1 VALUES (1, 'a'), (2, 'b')") + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + // `length(dt) > 2` is not pushed down, so `dt` reaches the scan while the projection above + // it drops the column, narrowing the KeyedPartitioning to `[identity(id)]`. + val df = sql( + """ + |SELECT nt.id, t1.data + |FROM (SELECT id FROM testcat.ns.nt WHERE length(dt) > 2) nt + |JOIN t1 ON nt.id = t1.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + // Only the v1 side is shuffled, onto the narrowed KeyedPartitioning of the v2 side, and + // nothing re-groups the v2 side -- so the shuffle is the only thing that can align them. + assert(collectShuffles(plan).length == 1) + assert(collectGroupPartitions(plan).isEmpty) + checkAnswer(df, Seq(Row(1L, "a"), Row(2L, "b"))) + } + } + } + + test("SPARK-59022: keyed shuffle follows the declared partition key order over a union") { + val cols = Array( + Column.create("id", LongType), + Column.create("data", StringType)) + val partitions = Array[Transform](identity("id")) + + createTable("nt1", cols, partitions) + createTable("nt2", cols, partitions) + // `UnionExec` concatenates its children's partition keys in child order, so the merged keys are + // [3, 4] ++ [1, 2]. They are unique, hence grouped, and no projection is involved, hence not + // narrowed -- but they are not sorted. + sql("INSERT INTO testcat.ns.nt1 VALUES (3, 'c'), (4, 'd')") + sql("INSERT INTO testcat.ns.nt2 VALUES (1, 'a'), (2, 'b')") + + withTable("t1") { + sql("CREATE TABLE t1 (id BIGINT, x STRING) USING parquet") + sql("INSERT INTO t1 VALUES (1, 'x'), (2, 'x'), (3, 'x'), (4, 'x')") + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql( + """ + |SELECT u.id, u.data, t1.x + |FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u + |JOIN t1 ON u.id = t1.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + // Only the v1 side is shuffled, onto the union's KeyedPartitioning, and nothing re-groups + // the union -- so the shuffle is the only thing that can align them. + assert(collectShuffles(plan).length == 1) + assert(collectGroupPartitions(plan).isEmpty) + checkAnswer(df, Seq( + Row(1L, "a", "x"), Row(2L, "b", "x"), Row(3L, "c", "x"), Row(4L, "d", "x"))) + } + + // The plan-shape assertions above need AQE off, but the wrong results were live in the + // default configuration, so pin the answer with AQE on as well. + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + checkAnswer( + sql( + """ + |SELECT u.id, u.data, t1.x + |FROM (SELECT * FROM testcat.ns.nt1 UNION ALL SELECT * FROM testcat.ns.nt2) u + |JOIN t1 ON u.id = t1.id + |""".stripMargin), + Seq(Row(1L, "a", "x"), Row(2L, "b", "x"), Row(3L, "c", "x"), Row(4L, "d", "x"))) + } + } + } + + test("SPARK-59027: v2 bucketed table with subset join keys left-outer joining v1 table") { + // Same shape as the SPARK-58988 test above, but LEFT OUTER: `ShuffledJoin` then exposes only + // the left side's partitioning, so no `PartitioningCollection` invariant compares the two + // sides' declared keys at planning time. If the order declared by `createShuffleSpec` and the + // physical layouts of the two sides (`GroupPartitionsExec` on the keyed side, the shuffle + // partitioner on the other) ever diverged, this join would silently lose matches instead of + // failing planning, so the answer check is the guard here. The unmatched row distinguishes a + // legitimate outer-join null from a lost match. + val cols = Array( + Column.create("c1", LongType), + Column.create("c2", StringType), + Column.create("dt", StringType)) + val partitions = Array(identity("dt"), bucket(16, "c1")) + + createTable("iceberg_t3", cols, partitions) + sql("INSERT INTO testcat.ns.iceberg_t3 VALUES " + + "(2, 'cc', '2020'), (1, 'aa', '2021'), (3, 'ee', '2022')") + + withTable("t1") { + sql("CREATE TABLE t1 (c1 BIGINT, c2 STRING) USING parquet") + sql("INSERT INTO t1 VALUES (1, 'aa'), (2, 'cc')") + + withSQLConf( + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql("SELECT * FROM testcat.ns.iceberg_t3 t0 LEFT JOIN t1 ON t0.c1 = t1.c1") + val plan = df.queryExecution.executedPlan + // Only the v1 side is re-shuffled; the v2 side is regrouped onto the join key instead. + assert(collectShuffles(plan).length == 1) + assert(collectGroupPartitions(plan).length == 1) + checkAnswer(df, Seq( + Row(1L, "aa", "2021", 1L, "aa"), + Row(2L, "cc", "2020", 2L, "cc"), + Row(3L, "ee", "2022", null, null))) + } + } + } + test("SPARK-57881: storage-partitioned join leverages union output KeyedPartitioning to " + "avoid shuffle") { val cols = Array( @@ -4495,4 +4717,56 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } } + + test("SPARK-58558: SPJ on a join key column partitioned by multiple transforms") { + // Partition expressions outnumber the join keys because `id` is partitioned twice, but + // every join key is covered, so SPJ works with default configs. + val items_partitions = Array(bucket(8, "id"), identity("id")) + createTable(items, itemsColumns, items_partitions) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + "(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + val purchases_partitions = Array(bucket(8, "item_id"), identity("item_id")) + createTable(purchases, purchasesColumns, purchases_partitions) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(1, 42.0, cast('2020-01-01' as timestamp)), " + + "(2, 19.5, cast('2020-02-01' as timestamp))") + + val df = createJoinTestDF(Seq("id" -> "item_id")) + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, "should not contain any shuffle") + checkAnswer(df, Seq(Row(1, "aa", 40.0, 42.0), Row(2, "bb", 10.0, 19.5))) + } +} + +/** + * Runs the runtime filtering tests against a catalog whose scans take runtime filters as connector + * predicates, via [[org.apache.spark.sql.connector.read.SupportsRuntimeFiltering]]. + */ +@ExtendedSQLTest +class KeyGroupedPartitioningRuntimeFilterSuite + extends KeyGroupedPartitioningSuiteBase with KeyGroupedPartitioningRuntimeFilterTests { + + after { + catalog.clearTables() + } +} + +/** + * Runs the runtime filtering tests against a catalog whose scans take runtime filters as Catalyst + * expressions, via + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]]. + */ +@ExtendedSQLTest +class KeyGroupedPartitioningCatalystRuntimeFilterSuite + extends KeyGroupedPartitioningSuiteBase with KeyGroupedPartitioningRuntimeFilterTests { + + override protected def catalogClassName: String = + classOf[InMemoryCatalystRuntimeFilterCatalog].getName + + after { + catalog.clearTables() + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala index de6c450069115..ea15f13c225cf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala @@ -2866,13 +2866,15 @@ abstract class MergeIntoTableSuiteBase extends RowLevelOperationSuiteBase // DataSourceV2Relation, the RowLevelOperationInfo, and the write builder's LogicalWriteInfo checkRowLevelOperationOptions( sql( - s"""MERGE INTO $tableNameAsString t WITH (`write.split-size` = 10) + s"""MERGE INTO $tableNameAsString t WITH + | (`load-option` = 'load-value', `write-option` = 'write-value') |USING $sourceNameAsString s |ON t.pk = s.pk |WHEN MATCHED THEN UPDATE SET t.salary = s.salary |WHEN NOT MATCHED THEN INSERT * |""".stripMargin), - "write.split-size" -> "10") + "load-option" -> "load-value", + "write-option" -> "write-value") checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -3016,7 +3018,8 @@ abstract class MergeIntoTableSuiteBase extends RowLevelOperationSuiteBase // write from the target's options at a different site than the row-level rewrite path. val executedPlan = executeAndKeepPlan { sql( - s"""MERGE INTO $tableNameAsString t WITH (`write.split-size` = 10) + s"""MERGE INTO $tableNameAsString t WITH + | (`load-option` = 'load-value', `write-option` = 'write-value') |USING $sourceNameAsString s |ON t.pk = s.pk |WHEN NOT MATCHED THEN INSERT * @@ -3026,7 +3029,11 @@ abstract class MergeIntoTableSuiteBase extends RowLevelOperationSuiteBase case e: InsertOnlyMergeExec => e.write }.getOrElse(fail("expected an InsertOnlyMergeExec in the executed plan")) val append = write.toBatch.asInstanceOf[InMemoryBaseTable#Append] - assert(append.info.options.get("write.split-size") === "10") + assert(append.info.options.get("load-option") === "load-value") + assert(append.info.options.get("write-option") === "write-value") + assertLastTransactionWriteLoadOptions( + "load-option" -> "load-value", + "write-option" -> "write-value") checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala new file mode 100644 index 0000000000000..810131eb54cef --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala @@ -0,0 +1,229 @@ +/* + * 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.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.{BufferedRows, InMemoryRowLevelOperationTable} +import org.apache.spark.sql.execution.InSubqueryExec +import org.apache.spark.sql.execution.ReusedSubqueryExec +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.types.StructType +import org.apache.spark.unsafe.types.UTF8String + +/** + * Verifies that row-level runtime group filtering injects filters for scans that implement + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], where the filter + * reaches the connector as a Catalyst expression instead of a connector predicate. + * + * The tests here apply to both group-based and delta-based row-level operations. DELETE is left + * to the concrete suites because delta-based deletes only scan the row ID and the condition + * columns, so the group key is not available to filter on. + */ +abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase + extends RowLevelOperationSuiteBase { + + import testImplicits._ + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("use-catalyst-runtime-filtering", "true") + props + } + + test("update runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("updated_id") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // the subquery blocks planning-time pushdown, leaving group filtering to do the pruning; + // only id 1 matches, and it lives in hr + val updatedIdDF = Seq(Some(1), None).toDF() + updatedIdDF.createOrReplaceTempView("updated_id") + + val executedPlan = executeAndKeepPlan { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM updated_id)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "id INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 1, -1, "hr") :: Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } + } + + test("merge runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 200, "dep": "hr" } + |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" } + |{ "pk": 4, "id": 4, "salary": 400, "dep": "software" } + |{ "pk": 5, "id": 5, "salary": 500, "dep": "software" } + |""".stripMargin) + + // pk 1 to 3 match rows in hr, pk 6 matches nothing and becomes an insert, so hr is the + // only group that has to be rewritten + val sourceDF = Seq(1, 2, 3, 6).toDF("pk") + sourceDF.createOrReplaceTempView("source") + + val executedPlan = executeAndKeepPlan { + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = t.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, id, salary, dep) VALUES (s.pk, 0, 0, 'hr') + |""".stripMargin) + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "pk INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq( + Row(1, 1, 101, "hr"), + Row(2, 2, 201, "hr"), + Row(3, 3, 301, "hr"), + Row(4, 4, 400, "software"), + Row(5, 5, 500, "software"), + Row(6, 0, 0, "hr"))) + } + } + + /** + * Asserts the injected group filter down to its contents: the scan declares + * `expectedFilterAttrs` in `filterAttributes`, every scan node carries one dynamic pruning + * filter matching `expectedFilter`, the connector received that same filter as a Catalyst + * expression, and the scan then read only `expectedFilter.groups`. + * + * A group-based UPDATE is rewritten as a union of two branches sharing one scan, so the plan + * can hold more than one scan node. Each carries its own copy of the filter, keyed on that + * branch's own attributes and with its own expr IDs, and pushes it separately (see the UPDATE + * case in RowLevelOperationRuntimeGroupFiltering.buildMatchingRowsPlan). Every copy is checked + * against `expectedFilter` rather than against one another, so the expectation stays explicit + * and does not depend on how many copies the rewrite happens to produce. + */ + protected def assertCatalystGroupFilter( + executedPlan: SparkPlan, + expectedFilterAttrs: Seq[String], + expectedFilter: GroupFilter): Unit = { + val batchScans = collect(executedPlan) { case s: BatchScanExec => s } + assert(batchScans.nonEmpty, "expected a batch scan for the row-level operation") + val scan = catalystScan(batchScans.head) + assert(batchScans.forall(_.scan eq scan), + s"expected all ${batchScans.size} scan nodes to share one scan") + + val filterAttrs = scan.filterAttributes().map(_.fieldNames.mkString(".")).toSeq + assert(filterAttrs === expectedFilterAttrs, + s"expected the scan to declare $expectedFilterAttrs as filter attributes, got $filterAttrs") + + batchScans.foreach { batchScan => + batchScan.runtimeFilters match { + case Seq(DynamicPruningExpression(inSubquery: InSubqueryExec)) => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => fail(s"expected a single dynamic pruning group filter, got $other") + } + } + + // the scan must receive the Catalyst subquery Spark planned, not a translated connector + // predicate, once per scan node + val pushed = scan.pushedCatalystPredicates + assert(pushed.size === batchScans.size, + s"expected each of the ${batchScans.size} scan node(s) to push the filter once, got $pushed") + pushed.foreach { + case inSubquery: InSubqueryExec => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => + fail(s"expected the group filter pushed as an InSubqueryExec, got $other") + } + + val scannedGroups = scan.data.map(_.asInstanceOf[BufferedRows].keyString()).distinct + assert(scannedGroups.sorted === expectedFilter.groups.sorted, + s"scan must read only the filtered groups, got ${scannedGroups.mkString(", ")}") + } + + /** + * The expected shape of a group filter subquery: the columns it reads, which must be only those + * needed to evaluate the row-level condition, and the groups it resolves to at runtime. + */ + protected case class GroupFilter(scanSchema: String, groups: Seq[String]) + + private def assertGroupFilter( + filter: InSubqueryExec, + expectedFilterAttrs: Seq[String], + expectedFilter: GroupFilter): Unit = { + assert(filter.child.references.toSeq.map(_.name) === expectedFilterAttrs, + s"expected the group filter keyed on $expectedFilterAttrs, got ${filter.child}") + + // the second branch of a group-based UPDATE reuses the first branch's subquery, and + // ReusedSubqueryExec is a leaf node, so unwrap it to reach the plan underneath + val subqueryPlan = filter.plan match { + case reused: ReusedSubqueryExec => reused.child + case plan => plan + } + val subqueryScan = find(subqueryPlan) { case _: BatchScanExec => true; case _ => false } + .getOrElse(fail(s"could not find the scan of group filter subquery ${filter.plan.name}")) + assert( + DataTypeUtils.sameType(subqueryScan.schema, StructType.fromDDL(expectedFilter.scanSchema)), + s"unexpected group filter subquery scan schema ${subqueryScan.schema.sql}") + + val groups = filter.values() + .getOrElse(fail("group filter subquery produced no values")) + .map(_.asInstanceOf[UTF8String].toString) + assert(groups.toSeq.sorted === expectedFilter.groups.sorted, + s"group filter must select the groups holding matching rows, got ${groups.mkString(", ")}") + } + + /** Asserts no group filter was injected, e.g. because the scan does not read the group key. */ + protected def assertNoCatalystGroupFilter(executedPlan: SparkPlan): Unit = { + val batchScan = collect(executedPlan) { case s: BatchScanExec => s }.head + val scan = catalystScan(batchScan) + assert(scan.filterAttributes().isEmpty, + s"expected no filter attributes, got ${scan.filterAttributes().mkString(", ")}") + assert(batchScan.runtimeFilters.isEmpty, + s"expected no runtime filters, got ${batchScan.runtimeFilters}") + assert(scan.pushedCatalystPredicates.isEmpty, + s"expected no pushed predicates, got ${scan.pushedCatalystPredicates}") + } + + private type CatalystRowLevelScan = + InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan + + private def catalystScan(batchScan: BatchScanExec): CatalystRowLevelScan = { + batchScan.scan match { + case s: InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan => s + case other => fail(s"expected InMemoryCatalystRowLevelBatchScan, got ${other.getClass}") + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala index 5d52325da818a..21e27e9b6c1cb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala @@ -19,6 +19,8 @@ package org.apache.spark.sql.connector import java.util.Collections +import scala.jdk.CollectionConverters._ + import org.scalatest.BeforeAndAfter import org.apache.spark.sql.{DataFrame, Encoders, Row} @@ -48,12 +50,15 @@ abstract class RowLevelOperationSuiteBase before { spark.conf.set("spark.sql.catalog.cat", classOf[InMemoryRowLevelOperationTableCatalog].getName) + spark.conf.set( + "spark.sql.catalog.cat.tableStateOptionKeys", "load-option,targetLoadOption") } after { catalog.nextTxnRejectRegisteredScansAttempt = false spark.sessionState.catalogManager.reset() spark.sessionState.conf.unsetConf("spark.sql.catalog.cat") + spark.sessionState.conf.unsetConf("spark.sql.catalog.cat.tableStateOptionKeys") } protected final val PK_FIELD = StructField("pk", IntegerType, nullable = false) @@ -179,9 +184,9 @@ abstract class RowLevelOperationSuiteBase }.getOrElse(fail("couldn't find row-level operation in optimized plan")) } - // asserts the given SQL options reached every layer that should carry them: the rewritten - // DataSourceV2Relation, the RowLevelOperationInfo passed to the operation builder, and the - // write builder's LogicalWriteInfo + // Asserts the given SQL options reached every V2 layer that should carry them: the target + // catalog load, the rewritten DataSourceV2Relation, the RowLevelOperationInfo passed to the + // operation builder, and the write builder's LogicalWriteInfo. protected def checkRowLevelOperationOptions( func: => Unit, expectedOptions: (String, String)*): Unit = { @@ -191,8 +196,12 @@ abstract class RowLevelOperationSuiteBase case wd: WriteDelta => wd.table }.getOrElse(fail("couldn't find row-level operation in optimized plan")) .asInstanceOf[DataSourceV2Relation] + assert(writeRelation.catalog.nonEmpty, "expected a catalog-backed V2 write relation") + assert(writeRelation.identifier.nonEmpty, "expected a catalog-backed V2 write identifier") val operation = writeRelation.table.asInstanceOf[RowLevelOperationTable].operation .asInstanceOf[RowLevelOperationWithOptions] + assertLastTransactionWriteLoadOptions(expectedOptions: _*) + expectedOptions.foreach { case (key, value) => assert(writeRelation.options.get(key) === value, s"relation option '$key'") assert(operation.options.get(key) === value, s"row-level operation option '$key'") @@ -200,6 +209,27 @@ abstract class RowLevelOperationSuiteBase } } + protected def assertLastTransactionWriteLoadOptions( + expectedOptions: (String, String)*): Unit = { + val stateKeys = catalog.tableStateOptionKeys().asScala + val expectedStateOptions = expectedOptions.filter { case (key, _) => + stateKeys.exists(_.equalsIgnoreCase(key)) + } + val targetLoads = catalog.lastTransaction.catalog.loadTableCalls.filter { + case (context, _) => !context.writePrivileges().isEmpty + } + assert(targetLoads.nonEmpty, "target loadTable did not receive write privileges") + targetLoads.foreach { case (_, options) => + assert(options.size() === expectedStateOptions.size) + expectedStateOptions.foreach { case (key, value) => + assert(options.get(key) === value, s"table-state option '$key'") + } + expectedOptions.diff(expectedStateOptions).foreach { case (key, _) => + assert(options.get(key) === null, s"non-state load option '$key'") + } + } + } + protected def assertNoScanPlanning(plan: LogicalPlan): Unit = { val relations = plan.collect { case r: DataSourceV2Relation => r } assert(relations.nonEmpty, "plan must contain relations") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/SupportsCatalogOptionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/SupportsCatalogOptionsSuite.scala index bcd8ba185d1dd..1c293bfa268d5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/SupportsCatalogOptionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/SupportsCatalogOptionsSuite.scala @@ -27,14 +27,15 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.SparkException import org.apache.spark.sql.{AnalysisException, DataFrame, Dataset, SaveMode} +import org.apache.spark.sql.QueryTest.withQueryExecutionsCaptured import org.apache.spark.sql.catalyst.analysis.{AsOfTimestamp, AsOfVersion, NoSuchTableException, TableAlreadyExistsException, TimeTravelSpec} import org.apache.spark.sql.catalyst.plans.logical.{AppendData, LogicalPlan, OverwriteByExpression} import org.apache.spark.sql.catalyst.util.DateTimeUtils -import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Column, Identifier, InMemoryTableCatalog, SupportsCatalogOptions, TableCatalog} +import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Column, Identifier, InMemoryBaseTable, InMemoryTableCatalog, SupportsCatalogOptions, TableCatalog, TableWritePrivilege} import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME import org.apache.spark.sql.connector.expressions.{FieldReference, IdentityTransform} import org.apache.spark.sql.execution.QueryExecution -import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.execution.datasources.v2.{AppendDataExec, DataSourceV2Relation, OverwriteByExpressionExec} import org.apache.spark.sql.internal.SQLConf.V2_SESSION_CATALOG_IMPLEMENTATION import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.LongType @@ -65,6 +66,7 @@ class SupportsCatalogOptionsSuite extends SharedSparkSession with BeforeAndAfter V2_SESSION_CATALOG_IMPLEMENTATION.key, classOf[InMemoryTableSessionCatalog].getName) spark.conf.set( s"spark.sql.catalog.$catalogName", classOf[InMemoryTableCatalog].getName) + spark.conf.set(s"spark.sql.catalog.$catalogName.tableStateOptionKeys", "load-option") } override def afterEach(): Unit = { @@ -75,6 +77,7 @@ class SupportsCatalogOptionsSuite extends SharedSparkSession with BeforeAndAfter catalog(catalogName).dropTable(_)) spark.conf.unset(V2_SESSION_CATALOG_IMPLEMENTATION.key) spark.conf.unset(s"spark.sql.catalog.$catalogName") + spark.conf.unset(s"spark.sql.catalog.$catalogName.tableStateOptionKeys") } private def testCreateAndRead( @@ -385,6 +388,94 @@ class SupportsCatalogOptionsSuite extends SharedSparkSession with BeforeAndAfter assert(relation.timeTravelSpec.contains(expectedTimeTravelSpec)) } + test("read options are preserved for scans but filtered from loadTable") { + sql(s"create table $catalogName.t1 (id bigint) using $format") + val cat = catalog(catalogName).asInstanceOf[InMemoryTableCatalog] + cat.resetLoadTableCalls() + + // The provider uses the options to identify the table, but this catalog declares no + // table-state options. The relation still retains the complete option map for scan planning. + val df = load("t1", Some(catalogName)) + df.collect() + val relation = df.logicalPlan.asInstanceOf[DataSourceV2Relation] + assert(relation.options.get("name") === "t1") + + val opts = cat.lastLoadTableOptions + assert(opts.isDefined, "loadTable(context, options) was not invoked") + assert(opts.get.isEmpty) + } + + test("SPARK-58389: SupportsCatalogOptions separates table-state and write options") { + sql(s"create table $catalogName.t1 (id bigint) using $format") + val cat = catalog(catalogName).asInstanceOf[InMemoryTableCatalog] + val loadOption = "load-Option" + val loadValue = "load-value" + val writeOption = "write-option" + + Seq( + (SaveMode.Append, "append", java.util.Set.of(TableWritePrivilege.INSERT)), + (SaveMode.Overwrite, "overwrite", + java.util.Set.of(TableWritePrivilege.INSERT, TableWritePrivilege.DELETE)) + ).foreach { case (mode, optionValue, expectedPrivileges) => + cat.resetLoadTableCalls() + val Seq(qe) = withQueryExecutionsCaptured(spark) { + spark.range(1).write + .format(format) + .option("name", "t1") + .option("catalog", catalogName) + .option(loadOption, loadValue) + .option(writeOption, optionValue) + .mode(mode) + .save() + } + + val matchingCalls = cat.loadTableCalls.filter { + case (_, options) => options.get(loadOption) == loadValue + } + assert(matchingCalls.nonEmpty, "loadTable(context, options) was not invoked for the write") + matchingCalls.foreach { case (context, options) => + assert(context.writePrivileges() === expectedPrivileges) + assert(options.size() === 1) + assert(options.asCaseSensitiveMap().containsKey(loadOption)) + assert(options.get(writeOption) === null) + } + + val actualWriteOptions = mode match { + case SaveMode.Append => + qe.executedPlan.collectFirst { + case AppendDataExec(_, _, write, _, _) => + write.toBatch.asInstanceOf[InMemoryBaseTable#Append].info.options + } + case SaveMode.Overwrite => + qe.executedPlan.collectFirst { + case OverwriteByExpressionExec(_, _, write, _, _) => + write.toBatch.asInstanceOf[InMemoryBaseTable#TruncateAndAppend].info.options + } + case other => fail(s"unexpected save mode: $other") + } + assert(actualWriteOptions.isDefined, "expected a V2 in-memory batch write") + assert(actualWriteOptions.get.get(loadOption) === loadValue) + assert(actualWriteOptions.get.get(writeOption) === optionValue) + assert(actualWriteOptions.get.get("name") === "t1") + assert(actualWriteOptions.get.get("catalog") === catalogName) + } + } + + test("SPARK-58389: SupportsCatalogOptions rejects time travel options for table creation") { + checkError( + exception = intercept[AnalysisException] { + spark.range(1).write + .format(format) + .option("name", "t1") + .option("catalog", catalogName) + .option("versionAsOf", "v1") + .option("timestampAsOf", "2021-01-01") + .save() + }, + condition = "UNSUPPORTED_FEATURE.TIME_TRAVEL", + parameters = Map("relationId" -> "`testcat`.`t1`")) + } + private def load( name: String, catalogOpt: Option[String], diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala index b54445466d7ef..c122713c1873f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala @@ -1243,8 +1243,11 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase { |""".stripMargin) checkRowLevelOperationOptions( - sql(s"UPDATE $tableNameAsString WITH (`write.split-size` = 10) SET salary = -1 WHERE pk = 1"), - "write.split-size" -> "10") + sql(s"UPDATE $tableNameAsString WITH " + + s"(`load-option` = 'load-value', `write-option` = 'write-value') " + + s"SET salary = -1 WHERE pk = 1"), + "load-option" -> "load-value", + "write-option" -> "write-value") checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala index ce2f7c232a961..7ec699353e687 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryCompilationErrorsSuite.scala @@ -22,6 +22,8 @@ import java.util.IllegalFormatException import org.apache.spark.{SPARK_DOC_ROOT, SparkIllegalArgumentException, SparkUnsupportedOperationException} import org.apache.spark.sql._ import org.apache.spark.sql.api.java.{UDF1, UDF2, UDF23Test} +import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType} import org.apache.spark.sql.catalyst.expressions.{Coalesce, Literal, UnsafeRow} import org.apache.spark.sql.catalyst.parser.ParseException import org.apache.spark.sql.execution.datasources.SaveIntoDataSourceCommand @@ -1133,6 +1135,22 @@ class QueryCompilationErrorsSuite parameters = Map("database" -> s"`$globalTempDB`") ) } + + test("SPARK-58349: TABLE_LOCATION_URI_NOT_SPECIFIED: table does not specify locationUri") { + val identifier = TableIdentifier("t", Some("db")) + val table = CatalogTable( + identifier = identifier, + tableType = CatalogTableType.MANAGED, + storage = CatalogStorageFormat.empty, + schema = new StructType()) + checkError( + exception = intercept[AnalysisException] { + table.location + }, + condition = "TABLE_LOCATION_URI_NOT_SPECIFIED", + parameters = Map("identifier" -> identifier.toString) + ) + } } class MyCastToString extends SparkUserDefinedFunction( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionAnsiErrorsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionAnsiErrorsSuite.scala index d5185289e021d..1ce04c8064ab1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionAnsiErrorsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/errors/QueryExecutionAnsiErrorsSuite.scala @@ -195,6 +195,14 @@ class QueryExecutionAnsiErrorsSuite extends SharedSparkSession { context = ExpectedContext( fragment = "apply", callSitePattern = getCurrentClassCallSitePattern)) + + checkError( + exception = intercept[SparkArrayIndexOutOfBoundsException] { + sql("select array(id, 2, 3)[5] from range(1)").collect() + }, + condition = "INVALID_ARRAY_INDEX", + parameters = Map("indexValue" -> "5", "arraySize" -> "3"), + context = ExpectedContext(fragment = "array(id, 2, 3)[5]", start = 7, stop = 24)) } test("INVALID_ARRAY_INDEX_IN_ELEMENT_AT: element_at from array") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala index fbc7ba71d6ce9..566f3cf64836d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala @@ -148,7 +148,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper spark.range(0, 20).selectExpr("id", "50 as cnt").collect().toImmutableArraySeq) // Then, let's look at the number of post-shuffle partitions estimated - // by the ExchangeCoordinator. + // by the CoalesceShufflePartitions rule. val finalPlan = stripAQEPlan(agg.queryExecution.executedPlan) val shuffleReads = finalPlan.collect { case r @ CoalescedShuffleRead() => r @@ -193,7 +193,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper expectedAnswer.collect().toImmutableArraySeq) // Then, let's look at the number of post-shuffle partitions estimated - // by the ExchangeCoordinator. + // by the CoalesceShufflePartitions rule. val finalPlan = stripAQEPlan(join.queryExecution.executedPlan) val shuffleReads = finalPlan.collect { case r @ CoalescedShuffleRead() => r @@ -248,7 +248,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper expectedAnswer.collect().toImmutableArraySeq) // Then, let's look at the number of post-shuffle partitions estimated - // by the ExchangeCoordinator. + // by the CoalesceShufflePartitions rule. val finalPlan = stripAQEPlan(join.queryExecution.executedPlan) val shuffleReads = finalPlan.collect { case r @ CoalescedShuffleRead() => r @@ -298,7 +298,7 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite with SQLConfHelper expectedAnswer.collect().toImmutableArraySeq) // Then, let's look at the number of post-shuffle partitions estimated - // by the ExchangeCoordinator. + // by the CoalesceShufflePartitions rule. val finalPlan = stripAQEPlan(join.queryExecution.executedPlan) val shuffleReads = finalPlan.collect { case r @ CoalescedShuffleRead() => r diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala index 17675705060d6..bd970f1c34d8a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/CombineAdjacentAggregationSuite.scala @@ -18,8 +18,9 @@ package org.apache.spark.sql.execution import org.apache.spark.sql.{QueryTest, Row} -import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final} -import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.catalyst.expressions.aggregate.{Complete, Final, PartialMerge} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEShuffleReadExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, SortAggregateExec} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -191,6 +192,89 @@ class CombineAdjacentAggregationSuite extends QueryTest } } + test("Combine adjacent partial merge and final hash aggregates") { + withTempView("t") { + spark.range(20).selectExpr("id % 3 as k", "id % 7 as v").createOrReplaceTempView("t") + val query = "SELECT k, count(*) FROM (SELECT /*+ repartition(k) */ * FROM t) GROUP BY k" + val finalAgg = withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false") { + collect(sql(query).queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.child.isInstanceOf[HashAggregateExec] => agg + }.head + } + val partialAgg = finalAgg.child.asInstanceOf[HashAggregateExec] + val partialMergeInput = InputAdapter(partialAgg) + val partialMergeAgg = partialAgg.copy( + requiredChildDistributionExpressions = Some(partialAgg.groupingExpressions), + isStreaming = true, + numShufflePartitions = Some(5), + initialInputBufferOffset = finalAgg.initialInputBufferOffset, + aggregateExpressions = partialAgg.aggregateExpressions.map(_.copy(mode = PartialMerge)), + child = partialMergeInput) + partialMergeAgg.copyTagsFrom(partialAgg) + val finalOverPartialMerge = finalAgg.copy(child = partialMergeAgg) + finalOverPartialMerge.copyTagsFrom(finalAgg) + + val combined = CombineAdjacentAggregation(finalOverPartialMerge) + .asInstanceOf[HashAggregateExec] + assert(combined.aggregateExpressions.forall(_.mode == Final)) + assert(combined.requiredChildDistributionExpressions == + finalAgg.requiredChildDistributionExpressions) + assert(combined.isStreaming == partialMergeAgg.isStreaming) + assert(combined.numShufflePartitions == partialMergeAgg.numShufflePartitions) + assert(combined.groupingExpressions == partialMergeAgg.groupingExpressions) + assert(combined.aggregateAttributes == finalAgg.aggregateAttributes) + assert(combined.resultExpressions == finalAgg.resultExpressions) + assert(combined.initialInputBufferOffset == partialMergeAgg.initialInputBufferOffset) + assert(combined.child == partialMergeInput) + assert(finalOverPartialMerge.executeCollect().map(_.copy()).toSet == + combined.executeCollect().map(_.copy()).toSet) + + val filteredPartialMerge = partialMergeAgg.copy( + aggregateExpressions = partialMergeAgg.aggregateExpressions.zipWithIndex.map { + case (agg, 0) => agg.copy(filter = Some(Literal.TrueLiteral)) + case (agg, _) => agg + }) + filteredPartialMerge.copyTagsFrom(partialAgg) + val finalOverFilteredPartialMerge = finalAgg.copy(child = filteredPartialMerge) + finalOverFilteredPartialMerge.copyTagsFrom(finalAgg) + assert(collect(CombineAdjacentAggregation(finalOverFilteredPartialMerge)) { + case agg: HashAggregateExec => agg + }.size == 3) + } + } + + test("Combined aggregate retains its distribution requirement under AQE") { + import testImplicits._ + + withTempView("t") { + spark.sparkContext.parallelize( + (1 to 10).map(i => (if (i > 4) 5 else i, i.toString)), 3) + .toDF("k", "v").createOrReplaceTempView("t") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true", + SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1", + SQLConf.SHUFFLE_PARTITIONS.key -> "5", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "150") { + val df = sql( + "SELECT k, count(*) FROM (SELECT /*+ REBALANCE(k) */ * FROM t) GROUP BY k") + checkAnswer(df, Seq(Row(1, 1), Row(2, 1), Row(3, 1), Row(4, 1), Row(5, 6))) + assert(collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec => agg + }.size == 1) + val partitionSpecs = collect(df.queryExecution.executedPlan) { + case read: AQEShuffleReadExec => read + }.flatMap(_.partitionSpecs) + assert(partitionSpecs.nonEmpty) + assert(partitionSpecs.forall(!_.isInstanceOf[PartialReducerPartitionSpec])) + } + } + } + test("Do not combine when a shuffle sits between the partial and final aggregate") { withTempView("t") { spark.range(20).selectExpr("id % 3 as k", "id % 7 as v") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala index b7798b0bde5db..c6c262379d582 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala @@ -238,4 +238,32 @@ class ExchangeSuite extends SharedSparkSession { "ReusedExchangeExec should reuse an existing exchange") } } + + test("ShuffleExchangeExec string args show `pipelined` only when it is set") { + val plan = spark.range(10).selectExpr("id % 4 AS key").queryExecution.executedPlan + val partitioning = HashPartitioning(Seq(Literal(1)), 4) + + // The default: nothing about pipelining is printed, so plan output (and every golden file that + // captures it) is unchanged by the existence of the field. + val nonPipelined = ShuffleExchangeExec(partitioning, plan).simpleString(10) + assert(!nonPipelined.contains("pipelined") && !nonPipelined.contains("false"), + s"a non-pipelined exchange should not mention pipelining, but was: $nonPipelined") + + // When set, it is printed with a name rather than as a bare positional `true`. + val pipelined = ShuffleExchangeExec(partitioning, plan, pipelined = true).simpleString(10) + assert(pipelined.contains("isPipelined=true"), + s"a pipelined exchange should report it, but was: $pipelined") + + // `stringArgs` drops `pipelined` by position, so it must stay the last constructor field -- + // otherwise a newly added field would be dropped instead and the bare `false` would come back. + val exchange = ShuffleExchangeExec(partitioning, plan) + assert(exchange.productElementName(exchange.productArity - 1) == "pipelined", + "`pipelined` must remain the last constructor field of ShuffleExchangeExec, because " + + "stringArgs drops it positionally") + + // Exchange's plan_id suffix is re-appended by the override and must survive. + Seq(nonPipelined, pipelined).foreach { s => + assert(s.contains("[plan_id="), s"the plan_id suffix should be preserved, but was: $s") + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExplainUtilsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExplainUtilsSuite.scala new file mode 100644 index 0000000000000..4830f5d9e78d4 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExplainUtilsSuite.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.spark.sql.execution + +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.logical.Range +import org.apache.spark.sql.catalyst.util.StringConcat +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests for [[ExplainUtils.processPlan]]: operator ID assignment, WholeStageCodegen tag + * propagation, thread-local lifecycle, and subquery handling. + */ +class ExplainUtilsSuite extends QueryTest with SharedSparkSession { + + private def explainOutput(plan: SparkPlan): String = { + val concat = new StringConcat() + ExplainUtils.processPlan(plan, concat.append) + concat.toString + } + + test("processPlan assigns unique operator IDs to all visible plan nodes") { + val df = spark.range(100).filter("id > 10").select("id") + val output = explainOutput(df.queryExecution.executedPlan) + // Each operator ID appears both in the tree header ("Filter (2)") and as the header of its + // verbose section ("(2) Filter"). Anchor to the verbose-section headers, which are the only + // lines that begin with "(N)", so each operator contributes exactly one ID. + val ids = "(?m)^\\((\\d+)\\)".r.findAllMatchIn(output).map(_.group(1).toInt).toSeq + assert(ids.nonEmpty, "processPlan should assign at least one operator ID") + assert(ids == ids.distinct, s"processPlan operator IDs should be unique: $ids") + } + + test("processPlan tags a WholeStageCodegenExec child with its codegen stage id") { + // Build the plan directly so the assertion does not depend on the planner producing a + // WholeStageCodegenExec for any particular query. processPlan should stamp the child of a + // WholeStageCodegenExec with CODEGEN_ID_TAG carrying that node's codegenStageId. + val child = RangeExec(Range(0, 10, 1, 1)) + val plan = WholeStageCodegenExec(child)(codegenStageId = 7) + ExplainUtils.processPlan[SparkPlan](plan, _ => ()) + assert(child.getTagValue(QueryPlan.CODEGEN_ID_TAG).contains(7), + "processPlan should tag a WholeStageCodegenExec child with its codegenStageId") + } + + test("processPlan restores localIdMap to its prior value after completion") { + val prev = ExplainUtils.localIdMap.get() + ExplainUtils.processPlan(spark.range(10).filter("id > 3").queryExecution.executedPlan, + _ => ()) + assert(ExplainUtils.localIdMap.get() eq prev, + "processPlan should restore the thread-local localIdMap on completion") + } + + test("processPlan includes subquery output in the explain string") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + val df = spark.range(100).filter("id > 0") + .filter("id < (SELECT max(id) FROM range(5))") + val output = explainOutput(df.queryExecution.executedPlan) + assert(output.contains("Subqueries"), "processPlan should emit a Subqueries section") + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/MetricViewV2CatalogSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/MetricViewV2CatalogSuite.scala index 9b75e89a1ad8f..5036ab94c133c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/MetricViewV2CatalogSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/MetricViewV2CatalogSuite.scala @@ -930,6 +930,28 @@ class MetricViewV2CatalogSuite extends SharedSparkSession { } } + test("ALTER VIEW <metric_view> SET TBLPROPERTIES preserves the dependency list") { + withTestCatalogTables { + val mv = MetricView( + "0.1", + AssetSource(fullSourceTableName), + where = None, + select = metricViewColumns) + createMetricView(fullMetricViewName, mv) + + sql(s"ALTER VIEW $fullMetricViewName SET TBLPROPERTIES ('k' = 'v')") + + val info = capturedViewInfo() + assert(info.properties().get("k") === "v") + val deps = info.viewDependencies() + assert(deps != null) + assert(deps.dependencies().length === 1) + val tableDep = deps.dependencies()(0).asInstanceOf[TableDependency] + assert(tableDep.nameParts().toSeq === + Seq(testCatalogName, testNamespace, sourceTableName)) + } + } + test("SHOW TABLES on a v2 RelationCatalog lists both tables and metric views") { withTestCatalogTables { val mv = MetricView( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala index 71192e92af0d0..41b74c6d4ff21 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala @@ -733,10 +733,11 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { outputPlan match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, - ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _, _), _), + ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), + _, _, _, _), _), SortExec(_, _, ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), - _, _, _), _), _) => + _, _, _, _), _), _) => assert(leftKeys === smjExec.leftKeys) assert(rightKeys === smjExec.rightKeys) assert(leftKeys === leftPartitioningExpressions) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala index 42f5e217ff4ab..f1f2bb7e0833a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala @@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeRef import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HashPartitioning, KeyedPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning} import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, YearsFunction} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, TimestampType} @@ -701,6 +702,141 @@ class ProjectedOrderingAndPartitioningSuite case other => fail(s"Expected UnknownPartitioning, got $other") } } + + test("SPARK-58405: eliminate redundant shuffle when aggregate groups by an alias of " + + "the window partition key") { + // Register the `testData` view (scoped to just these tests that need it). + testData + // The window output `vset` is consumed downstream so the window is not eliminated. + // The window is hash-partitioned by `key`; the outer aggregate groups by `userid`, which is + // an alias of `key`. Since `key` is already a partition key of the window, the outer + // aggregate's shuffle on `userid` is redundant. + val df = sql( + """ + |SELECT userid, count(*), sum(size(vset)) + |FROM ( + | SELECT key AS userid, + | collect_set(value) OVER (PARTITION BY key) AS vset + | FROM testData + |) u + |GROUP BY 1 + """.stripMargin) + + // Correctness: results must match grouping the same window output by `key` directly. + val expected = sql( + """ + |SELECT key, count(*), sum(size(vset)) + |FROM ( + | SELECT key, + | collect_set(value) OVER (PARTITION BY key) AS vset + | FROM testData + |) + |GROUP BY 1 + """.stripMargin) + checkAnswer(df, expected.collect()) + + // `collect` from `AdaptiveSparkPlanHelper` descends into the finalized AQE plan, so AQE can + // stay enabled. `checkAnswer` above has already materialized the query. + val plan = df.queryExecution.executedPlan + val shuffles = collect(plan) { case e: ShuffleExchangeExec => e } + assert(shuffles.size == 1, + s"Expected 1 shuffle but found ${shuffles.size}:\n$plan") + } + + test("SPARK-58405: eliminate redundant shuffle for a repartition consumer of a window " + + "partitioned by an aliased key") { + // Register the `testData` view (scoped to just these tests that need it). + testData + // Consumer-agnostic: a repartition on `userid` (an alias of the window key `key`) reuses the + // window's shuffle rather than adding its own. + val df = sql( + """ + |SELECT /*+ REPARTITION(userid) */ userid, vset + |FROM ( + | SELECT key AS userid, + | collect_set(value) OVER (PARTITION BY key) AS vset + | FROM testData + |) u + """.stripMargin) + + df.collect() + val plan = df.queryExecution.executedPlan + val shuffles = collect(plan) { case e: ShuffleExchangeExec => e } + assert(shuffles.size == 1, + s"Expected 1 shuffle but found ${shuffles.size}:\n$plan") + } + + test("SPARK-58405: eliminate redundant shuffle and sort for a window consumer partitioned " + + "and ordered by aliased keys") { + // Register the `testData` view (scoped to just these tests that need it). + testData + // Stacked windows: the outer window is PARTITION BY userid ORDER BY tstamp, where `userid` + // and `tstamp` are aliases of the inner window's partition key `key` and order key `value`. + // The inner window's child is already partitioned by `key` and ordered by `[key, value]`; + // pulling both `key AS userid` and `value AS tstamp` above the inner window projects the + // partitioning and the full ordering up through the aliases, so the outer window needs + // neither a redundant shuffle nor a redundant sort. Without the rule the outer window adds + // one of each (2 shuffles, 2 sorts). `size(vset)` keeps the inner window output live so it + // is not pruned away. + val df = sql( + """ + |SELECT userid, tstamp, + | sum(size(vset)) OVER (PARTITION BY userid ORDER BY tstamp) AS s + |FROM ( + | SELECT key AS userid, value AS tstamp, + | collect_set(value) OVER (PARTITION BY key ORDER BY value) AS vset + | FROM testData + |) u + """.stripMargin) + + df.collect() + val plan = df.queryExecution.executedPlan + val shuffles = collect(plan) { case e: ShuffleExchangeExec => e } + val sorts = collect(plan) { case s: SortExec => s } + assert(shuffles.size == 1, + s"Expected 1 shuffle but found ${shuffles.size}:\n$plan") + assert(sorts.size == 1, + s"Expected 1 sort but found ${sorts.size}:\n$plan") + } + + test("SPARK-58405: eliminate redundant shuffle across a chain of windows over an aliased key") { + // Register the `testData` view (scoped to just these tests that need it). + testData + // Two adjacent windows (different order specs, so they are not collapsed and leave no Project + // between them) both partition by `key`; the aggregate downstream groups by `userid`, an alias + // of `key`. Pulling `key AS userid` up across the whole window chain lets the parent project's + // `HashPartitioning(userid)` satisfy the aggregate, so no redundant shuffle is inserted. + // `sum(r1)`/`sum(r2)` keep both window outputs live so the chain is not pruned away. + val df = sql( + """ + |SELECT userid, count(*), sum(r1), sum(r2) + |FROM ( + | SELECT key AS userid, + | row_number() OVER (PARTITION BY key ORDER BY value) AS r1, + | rank() OVER (PARTITION BY key ORDER BY value DESC) AS r2 + | FROM testData + |) u + |GROUP BY 1 + """.stripMargin) + + val expected = sql( + """ + |SELECT key, count(*), sum(r1), sum(r2) + |FROM ( + | SELECT key, + | row_number() OVER (PARTITION BY key ORDER BY value) AS r1, + | rank() OVER (PARTITION BY key ORDER BY value DESC) AS r2 + | FROM testData + |) + |GROUP BY 1 + """.stripMargin) + checkAnswer(df, expected.collect()) + + val plan = df.queryExecution.executedPlan + val shuffles = collect(plan) { case e: ShuffleExchangeExec => e } + assert(shuffles.size == 1, + s"Expected 1 shuffle but found ${shuffles.size}:\n$plan") + } } private case class DummyLeafExecWithPartitioning( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala index f7afdb5e6e537..fd047d54b957f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala @@ -25,7 +25,7 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListe import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, SaveMode} import org.apache.spark.sql.catalyst.{QueryPlanningTracker, QueryPlanningTrackerCallback, TableIdentifier} import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedFunction, UnresolvedRelation} -import org.apache.spark.sql.catalyst.expressions.{Alias, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.{Alias, NamedLambdaVariable, RegExpReplace, UnsafeRow} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{CommandResult, LogicalPlan, OneRowRelation, Project, ShowTables, SubqueryAlias} import org.apache.spark.sql.catalyst.trees.TreeNodeTag @@ -55,6 +55,22 @@ class QueryExecutionSuite extends SharedSparkSession { override protected def sparkConf = super.sparkConf.set(SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key, "0") + private def collectLambdaVariables(plan: LogicalPlan): Seq[NamedLambdaVariable] = { + plan.collect { + case node => node.expressions.flatMap(_.collect { + case variable: NamedLambdaVariable => variable + }) + }.flatten + } + + private def collectRegExpReplaceExpressions(plan: LogicalPlan): Seq[RegExpReplace] = { + plan.collect { + case node => node.expressions.flatMap(_.collect { + case expression: RegExpReplace => expression + }) + }.flatten + } + def checkDumpedPlans(path: String, expected: Int): Unit = Utils.tryWithResource( Source.fromFile(path)) { source => assert(source.getLines().toList @@ -105,6 +121,34 @@ class QueryExecutionSuite extends SharedSparkSession { } } + test("SPARK-58208: optimizedPlan uses fresh stateful expressions") { + val df = spark.range(1).selectExpr("transform(array(id), x -> x + 1) AS v") + val queryExecution = df.queryExecution + + val beforeOptimize = collectLambdaVariables(queryExecution.withCachedData) + val optimized = collectLambdaVariables(queryExecution.optimizedPlan) + + assert(beforeOptimize.nonEmpty) + assert(beforeOptimize.size == optimized.size) + beforeOptimize.zip(optimized).foreach { case (before, after) => + assert(before.exprId == after.exprId) + assert(before.value ne after.value) + } + } + + test("SPARK-58208: optimizedPlan keeps structurally equal fresh stateful expressions") { + val df = spark.range(1).selectExpr( + "regexp_replace(cast(id AS STRING), cast(id AS STRING), 'x') AS v") + val queryExecution = df.queryExecution + + val beforeOptimize = collectRegExpReplaceExpressions(queryExecution.withCachedData) + val optimized = collectRegExpReplaceExpressions(queryExecution.optimizedPlan) + + assert(beforeOptimize.size == 1) + assert(optimized.size == 1) + assert(beforeOptimize.head ne optimized.head) + } + test("dumping query execution info by invalid path") { val path = "1234567890://plans.txt" val exception = intercept[IllegalArgumentException] { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala index 638d69fd7191d..fc1d9099cbaea 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala @@ -351,6 +351,7 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper { event match { case e: SparkListenerSQLExecutionStart => sqlJobTags = e.jobTags + case _ => } } }) @@ -383,6 +384,7 @@ class SQLExecutionSuite extends SparkFunSuite with SQLConfHelper { event match { case e: SparkListenerSQLExecutionStart => sqlJobGroupIdOpt = e.jobGroupId + case _ => } } }) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLFunctionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLFunctionSuite.scala index 7406f953d063e..d7956a9c1217c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLFunctionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLFunctionSuite.scala @@ -79,6 +79,27 @@ class SQLFunctionSuite extends SharedSparkSession { } } + test("SPARK-58779: SQL table function argument referencing a CTE") { + withUserDefinedFunction("my_tvf" -> false) { + sql( + """ + |CREATE FUNCTION my_tvf(a INT, b INT) + |RETURNS TABLE(x INT) + |RETURN SELECT a + b AS x + |""".stripMargin) + // The table function is called with scalar-subquery arguments that reference a CTE. + // ResolveSQLTableFunctions runs checkAnalysis (which runs InlineCTE) on the resolved + // function plan, which contains a CTERelationRef whose CTERelationDef lives in the outer + // WithCTE scope. Previously the unguarded cteMap(ref.cteId) lookup threw + // NoSuchElementException. + checkAnswer(sql( + """ + |WITH cte AS (SELECT 1 AS c1, 2 AS c2) + |SELECT * FROM my_tvf((SELECT c1 FROM cte), (SELECT c2 FROM cte)) + |""".stripMargin), Row(3)) + } + } + test("SQL scalar function with default value") { withUserDefinedFunction("bar" -> false) { sql( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLViewSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLViewSuite.scala index 3fb54d7c43d58..aef97c2a9314f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLViewSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SQLViewSuite.scala @@ -1018,6 +1018,19 @@ abstract class SQLViewSuite extends QueryTest { s"`$SESSION_CATALOG_NAME`.`default`.`view2` -> " + s"`$SESSION_CATALOG_NAME`.`default`.`view1`")) ) + + // Detect cyclic view references from subqueries nested in larger expressions. + checkError( + exception = intercept[AnalysisException] { + sql("ALTER VIEW view1 AS SELECT * FROM jt WHERE id = (SELECT id FROM view2)") + }, + condition = "RECURSIVE_VIEW", + parameters = Map( + "viewIdent" -> s"`$SESSION_CATALOG_NAME`.`default`.`view1`", + "newPath" -> (s"`$SESSION_CATALOG_NAME`.`default`.`view1` -> " + + s"`$SESSION_CATALOG_NAME`.`default`.`view2` -> " + + s"`$SESSION_CATALOG_NAME`.`default`.`view1`")) + ) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala index f6493a8d0e85c..7405aa5fd4de0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/SortSuite.scala @@ -25,11 +25,13 @@ import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest /** * Test sorting. Many of the test cases generate random data and compares the sorted result with one * sorted by a reference implementation ([[ReferenceSort]]). */ +@ExtendedSQLTest class SortSuite extends SharedSparkSession { import testImplicits.newProductEncoder import testImplicits.localSeqToDatasetHolder diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala index 2e73417dd9ed3..6af7d066fccd9 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/WholeStageCodegenSuite.scala @@ -23,18 +23,24 @@ import java.time.Duration import org.apache.spark.SparkException import org.apache.spark.rdd.MapPartitionsWithEvaluatorRDD import org.apache.spark.sql.{Dataset, Row, SaveMode} -import org.apache.spark.sql.catalyst.expressions.{And, Cast, CodegenObjectFactoryMode, Expression, IsNotNull} +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, Cast, CodegenObjectFactoryMode, Expression, IsNotNull} import org.apache.spark.sql.catalyst.expressions.codegen.{ByteCodeStats, CodeAndComment, CodeGenerator} import org.apache.spark.sql.execution.adaptive.DisableAdaptiveExecutionSuite import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec -import org.apache.spark.sql.execution.debug.codegenString +import org.apache.spark.sql.execution.debug.{codegenString, codegenStringSeq} import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{DayTimeIntervalType, DecimalType, DoubleType, FloatType, IntegerType, LongType, StringType, StructField, StructType} +// Nested-struct fixtures for the SPARK-51356 test below. They are top-level so the Dataset +// encoders resolve without an outer scope. +case class Spark51356Inner(d: Int) +case class Spark51356Mid(c: Spark51356Inner = null) +case class Spark51356Outer(b: Spark51356Mid = null) + // Disable AQE because the WholeStageCodegenExec is added when running QueryStageExec class WholeStageCodegenSuite extends SharedSparkSession with DisableAdaptiveExecutionSuite { @@ -753,6 +759,12 @@ class WholeStageCodegenSuite extends SharedSparkSession assert(df.collect() === Array(Row(1), Row(2), Row(3))) } + test("SPARK-58437: SortExec generates assignable iterator type") { + val code = genCode(spark.range(3, 0, -1).toDF().sort(col("id"))).map(_.body).mkString("\n") + assert(!code.contains("scala.collection.Iterator<UnsafeRow>")) + assert(code.contains("scala.collection.Iterator<InternalRow>")) + } + test("MapElements should be included in WholeStageCodegen") { import testImplicits._ @@ -1322,6 +1334,116 @@ class WholeStageCodegenSuite extends SharedSparkSession } } + test("SPARK-51356: FilterExec emits IsNotNull on a nested field before its otherPred") { + // `IsNotNull(b.c)` is null-intolerant and references only `b`, so it is classified as a + // notNullPred -- but its child is a complex expression, so it never matches an otherPred's + // bare attribute reference. It used to be deferred to the trailing leftover block, i.e. + // emitted *after* the UDF that dereferences `b.c`, and `ScalaUDF` hands a null argument to + // its deserializer rather than short-circuiting, so `newInstance(Spark51356Inner)` threw. + // The interpreted path evaluates the conjunction in order and was unaffected. + val data = Seq( + Spark51356Outer(null), + Spark51356Outer(Spark51356Mid(null)), // the row that used to trigger the failure + Spark51356Outer(Spark51356Mid(Spark51356Inner(0))), + Spark51356Outer(Spark51356Mid(Spark51356Inner(1)))) + val isDZero = udf((c: Spark51356Inner) => c.d == 0) + def newDf(): Dataset[Spark51356Mid] = { + // `map(identity)` keeps the input from being folded into a LocalRelation, so the filter + // really goes through whole-stage codegen. + val mids = spark.createDataset(data).map(identity) + .where(col("b").isNotNull).select(col("b").as[Spark51356Mid]) + mids.filter(col("c").isNotNull).filter(not(isDZero(col("c")))) + } + + withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { + val df = newDf() + val plan = df.queryExecution.executedPlan + assert(plan.exists(_.isInstanceOf[WholeStageCodegenExec]), + "Filter should be in whole-stage codegen") + // Guard against optimizer drift: the regression only exists when a single FilterExec + // carries an `IsNotNull` over a complex child together with a non-IsNotNull conjunct + // that consumes the same expression. + def conjuncts(e: Expression): Seq[Expression] = e match { + case And(l, r) => conjuncts(l) ++ conjuncts(r) + case other => Seq(other) + } + val matchingFilter = plan.collect { + case f: FilterExec => + val cs = conjuncts(f.condition) + val nestedIsNotNulls = cs.collect { + case IsNotNull(child) if !child.isInstanceOf[Attribute] => child + } + nestedIsNotNulls.exists { child => + cs.exists { + case _: IsNotNull => false + case other => other.exists(_.semanticEquals(child)) + } + } + }.exists(identity) + assert(matchingFilter, + "expected a FilterExec carrying IsNotNull(<complex>) plus a non-IsNotNull conjunct " + + "over the same expression") + checkAnswer(df.toDF(), Row(Row(1))) + } + // Cross-check the codegen path against the interpreted path. + withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") { + val df = newDf() + assert(!df.queryExecution.executedPlan.exists(_.isInstanceOf[WholeStageCodegenExec]), + "the cross-check must be planned without whole-stage codegen") + checkAnswer(df.toDF(), Row(Row(1))) + } + } + + test("SPARK-51356: FilterExec CSE emits IsNotNull on a nested field before its otherPred") { + // Same defect as above, on the CSE branch of `FilterExec.doConsume`, which inlines its own + // copy of the interleaving. Two otherPreds share the non-cheap `f(c)`, so the branch is + // taken (a bare `b.c` is cheap and would fall back to `generatePredicateCode`). The + // guarding `IsNotNull(b.c)` has to be emitted ahead of the shared CSE precompute, not after + // the predicates that consume it. + val data = Seq( + Spark51356Outer(null), + Spark51356Outer(Spark51356Mid(null)), // the row that used to trigger the failure + Spark51356Outer(Spark51356Mid(Spark51356Inner(0))), + Spark51356Outer(Spark51356Mid(Spark51356Inner(1)))) + val dOf = udf((c: Spark51356Inner) => c.d) + val mids = spark.createDataset(data).map(identity) + .where(col("b").isNotNull).select(col("b").as[Spark51356Mid]) + val df = mids + .filter(col("c").isNotNull) + .filter(dOf(col("c")) > 0) + .filter(dOf(col("c")) < 100) + + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.SUBEXPRESSION_ELIMINATION_FILTER_EXEC_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { + val plan = df.queryExecution.executedPlan + assert(plan.exists(_.isInstanceOf[WholeStageCodegenExec]), + "Filter should be in whole-stage codegen") + checkAnswer(df.toDF(), Row(Row(1))) + } + } + + test("SPARK-51356: FilterExec CSE guards an IsNotNull whose child is itself a subexpression") { + // A dynamic-gap `session_window` guards the session struct with an IsNotNull over the same + // cast the struct is built from, so the guarded expression is itself the common + // subexpression. The check must not reference the CSE state, which is emitted later. + val df = spark.sql( + """ + |SELECT a, count(*) AS cnt + |FROM VALUES ('A1', '2021-01-01 00:00:00'), ('A1', '2021-01-01 00:04:30'), + | ('A2', '2021-01-01 00:01:00') AS tab(a, b) + |GROUP BY a, session_window(b, CASE WHEN a = 'A1' THEN '5 minutes' ELSE '1 minute' END) + """.stripMargin) + + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.SUBEXPRESSION_ELIMINATION_FILTER_EXEC_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") { + checkAnswer(df, Seq(Row("A1", 2), Row("A2", 1))) + } + } + test("SPARK-56032: FilterExec CSE handles shared otherPred refs with guard") { // Filter shape: kind = 'numeric' AND cast(s as int) > 0 AND cast(s as int) < 100. // Exercises invariant (b) on a shape where two cast otherPreds share a ref: the @@ -1550,4 +1672,65 @@ class WholeStageCodegenSuite extends SharedSparkSession s"subexpressionElimination.filterExec.enabled should reduce repeated evaluation: " + s"addExact appears $enabledCount times when enabled vs $disabledCount times when disabled") } + + test("Expand should eliminate common subexpressions across branches") { + // The conditions of the two conditional COUNT DISTINCT aggregates share `sinh(v)`, which + // `RewriteDistinctAggregates` places in different branches of the Expand. Since all the + // branches of an Expand consume the same input row, with subexpression elimination + // `sinh(v)` should be evaluated once per input row (before the branch loop) instead of + // once per branch. + def runQuery(): String = { + val df = spark.range(10) + .selectExpr("id as k", "cast(id as double) as v") + .selectExpr( + "count(DISTINCT IF(sinh(v) > 5.0, k, NULL))", + "count(DISTINCT IF(sinh(v) > 50.0, v, NULL))") + val plan = df.queryExecution.executedPlan + assert(plan.exists(_.isInstanceOf[ExpandExec]), "Expand is expected") + checkAnswer(df, Row(7L, 5L)) + codegenStringSeq(plan).map(_._2).mkString + } + + val sinhPattern = "java\\.lang\\.Math\\.sinh".r + val enabledCode = withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true") { + runQuery() + } + val disabledCode = withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false") { + runQuery() + } + assert(sinhPattern.findAllIn(enabledCode).length == 1, + "sinh(v) should be evaluated only once per input row with subexpression elimination") + assert(sinhPattern.findAllIn(disabledCode).length == 2, + "sinh(v) should be evaluated once per branch without subexpression elimination") + + // Whether the switch/case bodies are split into separate functions (SPARK-35329) + // depends on the amount of code generated for the branches, so force both code paths + // explicitly instead of relying on the default methodSplitThreshold. With a tiny + // threshold the bodies are split and the eliminated subexpression variables have to + // be passed to the split functions as parameters; with a huge threshold the bodies + // stay inline and reference the variables directly. ExpandExec names each split + // function via `ctx.freshName("switchCaseCode")` (e.g. `switchCaseCode_0`) and emits + // both its definition and call site into the generated source, so the substring + // "switchCaseCode" appears in the generated code if and only if a split happened. + val splitCode = withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1") { + runQuery() + } + assert(splitCode.contains("switchCaseCode"), + "switch/case bodies should be split into separate functions with a tiny " + + "methodSplitThreshold") + assert(sinhPattern.findAllIn(splitCode).length == 1, + "sinh(v) should be evaluated only once per input row with function splitting") + + val inlineCode = withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1000000") { + runQuery() + } + assert(!inlineCode.contains("switchCaseCode"), + "switch/case bodies should stay inline with a large methodSplitThreshold") + assert(sinhPattern.findAllIn(inlineCode).length == 1, + "sinh(v) should be evaluated only once per input row without function splitting") + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala index f84f3ad5cb6c9..8e98d7785f6c0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -19,7 +19,10 @@ package org.apache.spark.sql.execution.adaptive import java.io.File import java.net.URI +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} +import scala.concurrent.Future import scala.concurrent.duration._ import org.apache.logging.log4j.Level @@ -31,10 +34,10 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListe import org.apache.spark.shuffle.sort.SortShuffleManager import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, EqualTo, IsNull, Or, SortOrder} -import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight} -import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti} -import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, JoinHint, LocalRelation, LogicalPlan} +import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, EqualTo, IsNull, Literal, Or, SortOrder} +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, EliminateLimits} +import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti, LeftSemi} +import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, EmptyRelation, GlobalLimit, Join, JoinHint, LeafNode, LocalRelation, LogicalPlan, Statistics} import org.apache.spark.sql.catalyst.plans.physical.{CoalescedNullAwareHashPartitioning, SinglePartition} import org.apache.spark.sql.classic.Strategy import org.apache.spark.sql.execution._ @@ -60,8 +63,18 @@ import org.apache.spark.sql.test.SQLTestData.TestData import org.apache.spark.sql.types.{IntegerType, StructType} import org.apache.spark.sql.util.QueryExecutionListener import org.apache.spark.tags.SlowSQLTest +import org.apache.spark.util.{SparkFatalException, ThreadUtils, Utils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.Utils + +/** + * A leaf whose cost estimate under-counts its structural row bound. Used by SPARK-57956 to verify + * that an unmaterialized query stage exposes the structural `maxRows` instead of an estimate. + */ +private case class UnderCountLeaf(output: Seq[Attribute]) extends LeafNode { + override def maxRows: Option[Long] = Some(2L) + override def computeStats(): Statistics = + Statistics(sizeInBytes = BigInt(1), rowCount = Some(BigInt(0))) +} @SlowSQLTest class AdaptiveQueryExecSuite @@ -363,6 +376,1061 @@ class AdaptiveQueryExecSuite } } + test("non-empty global aggregate stage eliminates conditionless semi joins") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val globalAggregate = spark.range(1).where("id < 0").repartition(2) + .agg(count("*").as("c")) + val df = testData.join(globalAggregate, Seq.empty[String], "left_semi") + + val logicalJoin = df.queryExecution.optimizedPlan.collectFirst { + case join: Join => join + }.getOrElse(fail("expected a conditionless semi join before adaptive execution")) + assert(logicalJoin.joinType == LeftSemi) + assert(logicalJoin.condition.isEmpty) + + val initialPlan = df.queryExecution.executedPlan + .asInstanceOf[AdaptiveSparkPlanExec].initialPlan + assert(findTopLevelBaseJoin(initialPlan).size == 1) + + val aggregate = collect(initialPlan) { + case aggregate: BaseAggregateExec if aggregate.groupingExpressions.isEmpty => aggregate + }.headOption.getOrElse(fail("expected a global aggregate in the initial adaptive plan")) + val emptyStage = TestExchangeQueryStageExec( + 0, + aggregate.child, + aggregate.child.canonicalized, + runtimeRowCount = Some(BigInt(0))) + emptyStage.resultOption.set(Some(())) + val aggregateStage = LogicalQueryStage( + logicalJoin.right, aggregate.withNewChildren(Seq(emptyStage))) + val rewrittenJoin = AQEPropagateEmptyRelation(logicalJoin.copy(right = aggregateStage)) + assert(rewrittenJoin.fastEquals(logicalJoin.left), rewrittenJoin) + + checkAnswer(df, testData.collect().toSeq) + + val finalPlan = stripAQEPlan(df.queryExecution.executedPlan) + assert(findTopLevelBaseJoin(finalPlan).isEmpty, finalPlan) + } + } + + test("global aggregate over empty filtered input preserves its single output row") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + withTempView("empty_rows") { + spark.range(1).where("id < 0").createOrReplaceTempView("empty_rows") + + val globalAggregate = sql( + "SELECT count(*) AS c FROM " + + "(SELECT /*+ REPARTITION(2) */ id FROM empty_rows ORDER BY id) sorted_empty_rows") + checkAnswer(globalAggregate, Row(0L)) + val finalAggregatePlan = stripAQEPlan(globalAggregate.queryExecution.executedPlan) + assert(!finalAggregatePlan.isInstanceOf[EmptyRelationExec]) + + val df = sql( + "SELECT l.* FROM testData l LEFT ANTI JOIN " + + "(SELECT count(*) AS c FROM empty_rows HAVING count(*) = 0) r " + + "ON l.key = r.c + 1") + + checkAnswer(df, testData.filter($"key" =!= 1).collect().toSeq) + } + } + } + + test("sort and limit wrappers propagate only provably empty query stages") { + val getEstimatedRowCount = + PrivateMethod[Option[BigInt]](Symbol("getEstimatedRowCount")) + val scan = LocalTableScanExec(Nil, Nil, None) + + val unknownStage = TestExchangeQueryStageExec(0, scan, scan) + val emptyStage = TestExchangeQueryStageExec( + 1, scan, scan, runtimeRowCount = Some(BigInt(0))) + emptyStage.resultOption.set(Some(())) + val nonEmptyStage = TestExchangeQueryStageExec( + 2, scan, scan, runtimeRowCount = Some(BigInt(10))) + nonEmptyStage.resultOption.set(Some(())) + + def estimatedRowCount(plan: SparkPlan): Option[BigInt] = + AQEPropagateEmptyRelation.invokePrivate(getEstimatedRowCount(plan)) + + assert(estimatedRowCount(LocalLimitExec(0, unknownStage)).contains(BigInt(0))) + assert(estimatedRowCount(GlobalLimitExec(0, unknownStage)).contains(BigInt(0))) + + assert(estimatedRowCount(LocalLimitExec(5, emptyStage)).contains(BigInt(0))) + assert(estimatedRowCount(GlobalLimitExec(5, emptyStage)).contains(BigInt(0))) + assert(estimatedRowCount(GlobalLimitExec(5, emptyStage, offset = 2)).contains(BigInt(0))) + assert(estimatedRowCount(GlobalLimitExec(-1, emptyStage, offset = 2)).contains(BigInt(0))) + + assert(estimatedRowCount(LocalLimitExec(5, unknownStage)).isEmpty) + assert(estimatedRowCount(GlobalLimitExec(5, unknownStage)).isEmpty) + assert(estimatedRowCount(LocalLimitExec(1, nonEmptyStage)).isEmpty) + assert(estimatedRowCount(GlobalLimitExec(1, nonEmptyStage, offset = 99)).isEmpty) + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.TOP_K_SORT_FALLBACK_THRESHOLD.key -> "1", + SQLConf.ORDERING_AWARE_LIMIT_OFFSET.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + val sorted = spark.range(0, 10, 1, numPartitions = 2) + .where("id < 0") + .orderBy($"id" % 8) + .limit(2) + .distinct() + val initialPlan = sorted.queryExecution.executedPlan + .asInstanceOf[AdaptiveSparkPlanExec].initialPlan + val (plannedSort, plannedShuffle) = initialPlan.collectFirst { + case GlobalLimitExec(_, sort @ SortExec(_, false, shuffle: ShuffleExchangeExec, _), _) + if sort.logicalLink.exists(logical => + shuffle.logicalLink.exists(_ eq logical)) => (sort, shuffle) + }.getOrElse(fail("expected a logically linked sort above the middle-limit shuffle")) + + val emptySortStage = TestExchangeQueryStageExec( + 3, + plannedShuffle, + plannedShuffle.canonicalized, + runtimeRowCount = Some(BigInt(0))) + emptySortStage.resultOption.set(Some(())) + val sortedEmptyStage = plannedSort.withNewChildren(Seq(emptySortStage)) + + Seq(sortedEmptyStage, LocalLimitExec(5, sortedEmptyStage)).foreach { physicalPlan => + val logicalStage = LogicalQueryStage(plannedSort.logicalLink.get, physicalPlan) + assert(AQEPropagateEmptyRelation(logicalStage).isInstanceOf[EmptyRelation]) + } + + checkAnswer(sorted.toDF(), Seq.empty) + } + } + + test("obsolete stage cancellation preserves shared reused exchange stages") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + val df = spark.sql("SELECT * FROM testData join testData2 ON key = a") + val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val replacementPlan = LocalTableScanExec(Nil, Nil, None) + + val sharedStage = TestExchangeQueryStageExec( + 0, LocalTableScanExec(Nil, Nil, None), LocalTableScanExec(Nil, Nil, None)) + adaptivePlan.context.markSharedStageResult(sharedStage.resultOption) + + val sharedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(sharedStage))) + assert(sharedCancelledIds.isEmpty) + assert(!sharedStage.cancelled) + + val materializedStage = TestExchangeQueryStageExec( + 1, LocalTableScanExec(Nil, Nil, None), LocalTableScanExec(Nil, Nil, None)) + materializedStage.resultOption.set(Some(())) + val materializedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(materializedStage))) + assert(materializedCancelledIds.isEmpty) + assert(!materializedStage.cancelled) + + val retainedStage = TestExchangeQueryStageExec( + 2, LocalTableScanExec(Nil, Nil, None), LocalTableScanExec(Nil, Nil, None)) + val retainedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(retainedStage, Seq(retainedStage))) + assert(retainedCancelledIds.isEmpty) + assert(!retainedStage.cancelled) + + val reusedStage = TestExchangeQueryStageExec( + 3, LocalTableScanExec(Nil, Nil, None), LocalTableScanExec(Nil, Nil, None)) + val reusedReplacementStage = reusedStage.newReuseInstance(4, reusedStage.output) + val reusedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(reusedReplacementStage, Seq(reusedStage))) + assert(reusedCancelledIds.isEmpty) + assert(!reusedStage.cancelled) + + val unsharedStage = TestExchangeQueryStageExec( + 5, LocalTableScanExec(Nil, Nil, None), LocalTableScanExec(Nil, Nil, None)) + adaptivePlan.context.stageCache.put(unsharedStage.plan.canonicalized, unsharedStage) + val unsharedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(unsharedStage))) + assert(unsharedCancelledIds == Seq(unsharedStage.id)) + assert(unsharedStage.cancelled) + assert(!adaptivePlan.context.stageCache.contains(unsharedStage.plan.canonicalized)) + + val failedStage = TestExchangeQueryStageExec( + 6, + LocalTableScanExec(Nil, Nil, None), + LocalTableScanExec(Nil, Nil, None), + cancelFailure = Some(new IllegalStateException("test stage cancellation failed"))) + adaptivePlan.context.stageCache.put(failedStage.plan.canonicalized, failedStage) + val failedCancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(failedStage))) + assert(failedCancelledIds.isEmpty) + assert(!failedStage.cancelled) + assert(adaptivePlan.context.stageCache.get(failedStage.plan.canonicalized) + .exists(_ eq failedStage)) + } + } + + test("obsolete same-plan reused exchange aliases are cancelled only when all are removed") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val adaptivePlan = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + .queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + def acquireStage(): ExchangeQueryStageExec = { + adaptivePlan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ExchangeQueryStageExec] + } + + val originalStage = acquireStage() + val reusedStage = acquireStage() + assert(originalStage.id != reusedStage.id) + assert(originalStage.resultOption.eq(reusedStage.resultOption)) + + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val candidateStages = Seq(originalStage, reusedStage) + + // Retaining either alias keeps the shared physical exchange alive. + assert(adaptivePlan.invokePrivate( + cancelObsoleteStages(reusedStage, candidateStages)).isEmpty) + assert(adaptivePlan.context.stageCache.get(originalStage.plan.canonicalized) + .exists(_.resultOption.eq(originalStage.resultOption))) + + val emptyReplacement = LocalTableScanExec(Nil, Nil, None) + val cancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(emptyReplacement, candidateStages)) + assert(cancelledIds.toSet == Set(originalStage.id, reusedStage.id)) + assert(cancelledIds.distinct.size == candidateStages.size) + assert(!adaptivePlan.context.stageCache.contains(originalStage.plan.canonicalized)) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val cancellationFailure = new IllegalStateException("obsolete reused exchange was cancelled") + candidateStages.foreach { stage => + assert(adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + stage, cancellationFailure, cancelledIds.toSet))) + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + stage, new SparkFatalException(new OutOfMemoryError("fatal obsolete alias")), + cancelledIds.toSet))) + } + + val unrelatedStage = TestExchangeQueryStageExec( + 412, exchange, exchange.canonicalized) + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + unrelatedStage, cancellationFailure, cancelledIds.toSet))) + } + } + + test("obsolete exchange cancellation preserves reuse by another adaptive subquery plan") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val adaptivePlan = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + .queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + def acquireStage(plan: AdaptiveSparkPlanExec): ExchangeQueryStageExec = { + plan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ExchangeQueryStageExec] + } + + val originalStage = acquireStage(adaptivePlan) + val samePlanAlias = acquireStage(adaptivePlan) + val subqueryPlan = adaptivePlan.copy(isSubquery = true) + assert(subqueryPlan ne adaptivePlan) + assert(subqueryPlan == adaptivePlan) + assert(subqueryPlan.context eq adaptivePlan.context) + + val subqueryAlias = acquireStage(subqueryPlan) + assert(originalStage.resultOption.eq(samePlanAlias.resultOption)) + assert(originalStage.resultOption.eq(subqueryAlias.resultOption)) + + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val emptyReplacement = LocalTableScanExec(Nil, Nil, None) + assert(adaptivePlan.invokePrivate(cancelObsoleteStages( + emptyReplacement, Seq(originalStage, samePlanAlias))).isEmpty) + assert(adaptivePlan.context.stageCache.get(originalStage.plan.canonicalized) + .exists(_.resultOption.eq(subqueryAlias.resultOption))) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val requiredStageFailure = new IllegalStateException("required subquery exchange failed") + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + originalStage, requiredStageFailure, Set.empty[Int]))) + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + samePlanAlias, requiredStageFailure, Set.empty[Int]))) + } + } + + test("failed obsolete exchange cancellation is not retried after same-plan reuse") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val adaptivePlan = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + .queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + + val cancellationAttempts = new AtomicInteger() + val originalStage = TestExchangeQueryStageExec( + 413, + exchange, + exchange.canonicalized, + cancelCallback = Some(() => { + if (cancellationAttempts.incrementAndGet() == 1) { + throw new IllegalStateException("first stage cancellation failed") + } + })) + adaptivePlan.context.registerStageOwner(originalStage.resultOption, adaptivePlan) + adaptivePlan.context.stageCache.put(exchange.canonicalized, originalStage) + + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val emptyReplacement = LocalTableScanExec(Nil, Nil, None) + assert(adaptivePlan.invokePrivate( + cancelObsoleteStages(emptyReplacement, Seq(originalStage))).isEmpty) + assert(cancellationAttempts.get() == 1) + assert(!originalStage.cancelled) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized).contains(originalStage)) + + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + val reusedStage = adaptivePlan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ExchangeQueryStageExec] + assert(originalStage.id != reusedStage.id) + assert(originalStage.resultOption.eq(reusedStage.resultOption)) + assert(!adaptivePlan.context.isSharedStageResult(originalStage.resultOption)) + + assert(adaptivePlan.invokePrivate( + cancelObsoleteStages(emptyReplacement, Seq(reusedStage))).isEmpty) + assert(cancellationAttempts.get() == 1, + "a failed exchange cancellation must never be retried for another local alias") + assert(!originalStage.cancelled) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized).contains(originalStage)) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val fatalFailure = new SparkFatalException(new OutOfMemoryError("fatal failed cancellation")) + Seq(originalStage, reusedStage).foreach { stage => + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + stage, fatalFailure, Set.empty[Int]))) + } + + val obsoleteFailure = new IllegalStateException("uncancelled obsolete exchange failed") + Seq(originalStage, reusedStage).foreach { stage => + assert(adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + stage, obsoleteFailure, Set.empty[Int]))) + } + assert(cancellationAttempts.get() == 1) + } + } + + test("obsolete reused exchange cancellation includes aliases dropped by an earlier plan") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val adaptivePlan = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + .queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + def acquireStage(): ExchangeQueryStageExec = { + adaptivePlan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ExchangeQueryStageExec] + } + + val originalStage = acquireStage() + val previouslyDroppedAlias = acquireStage() + val lastRemainingAlias = acquireStage() + val allAliases = Seq(originalStage, previouslyDroppedAlias, lastRemainingAlias) + assert(allAliases.map(_.id).distinct.size == allAliases.size) + assert(allAliases.forall(_.resultOption.eq(originalStage.resultOption))) + + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + assert(adaptivePlan.invokePrivate( + cancelObsoleteStages(lastRemainingAlias, allAliases)).isEmpty) + + val emptyReplacement = LocalTableScanExec(Nil, Nil, None) + val cancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(emptyReplacement, Seq(lastRemainingAlias))) + assert(cancelledIds.toSet == allAliases.map(_.id).toSet) + assert(cancelledIds.distinct.size == allAliases.size) + assert(!adaptivePlan.context.stageCache.contains(originalStage.plan.canonicalized)) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val cancellationFailure = new IllegalStateException("previously dropped exchange failed") + allAliases.foreach { alias => + assert(adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + alias, cancellationFailure, cancelledIds.toSet))) + } + } + } + + test("cancelling a submitted reused shuffle alias preserves its underlying shuffle cleanup ID") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val owner = org.apache.spark.sql.classic.Dataset.ofRows( + spark, spark.range(16).repartition(2).logicalPlan, DoNotCleanup) + val adaptivePlan = owner.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + assert(adaptivePlan.context.qe.shuffleCleanupMode == DoNotCleanup) + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + def acquireStage(): ShuffleQueryStageExec = { + adaptivePlan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ShuffleQueryStageExec] + } + + val originalStage = acquireStage() + val reusedStage = acquireStage() + assert(reusedStage.shuffle eq originalStage.shuffle) + assert(reusedStage.plan.isInstanceOf[ReusedExchangeExec]) + ThreadUtils.awaitResult(originalStage.materialize(), 30.seconds) + assert(originalStage.shuffle.futureAction.get().isDefined) + assert(!originalStage.isMaterialized) + + val shuffleId = originalStage.shuffle.shuffleId + adaptivePlan.context.shuffleIds.remove(shuffleId) + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val emptyReplacement = LocalTableScanExec(Nil, Nil, None) + val cancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(emptyReplacement, Seq(reusedStage))) + + assert(cancelledIds.toSet == Set(originalStage.id, reusedStage.id)) + assert(adaptivePlan.context.shuffleIds.containsKey(shuffleId), + "a reused exchange leaf must not hide its submitted underlying shuffle from cleanup") + assert(!adaptivePlan.context.stageCache.contains(originalStage.plan.canonicalized)) + } + } + + test("obsolete stage cancellation includes stages from previous adaptive plan adoptions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val adaptivePlan = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + .queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchanges = adaptivePlan.initialPlan.collect { + case exchange: ShuffleExchangeExec => exchange + } + assert(exchanges.size >= 2, "expected two shuffle exchanges in the initial adaptive plan") + val previousExchange = exchanges.head + val latestExchange = exchanges.find(exchange => + !exchange.canonicalized.fastEquals(previousExchange.canonicalized)) + .getOrElse(fail("expected independently cached shuffle exchanges")) + + val previouslyRetainedStage = TestExchangeQueryStageExec( + 410, previousExchange, previousExchange.canonicalized) + val latestStage = TestExchangeQueryStageExec( + 411, latestExchange, latestExchange.canonicalized) + val oldPhysicalPlan = UnionExec(Seq(previouslyRetainedStage, latestStage)) + val latestStages = Seq(latestStage) + adaptivePlan.context.stageCache.put(previousExchange.canonicalized, previouslyRetainedStage) + adaptivePlan.context.stageCache.put(latestExchange.canonicalized, latestStage) + + val obsoleteStageCandidates = + PrivateMethod[Seq[QueryStageExec]](Symbol("obsoleteStageCandidates")) + val candidates = adaptivePlan.invokePrivate( + obsoleteStageCandidates(oldPhysicalPlan, latestStages)) + + assert(!latestStages.exists(_.id == previouslyRetainedStage.id)) + assert(candidates.map(_.id).toSet == Set(previouslyRetainedStage.id, latestStage.id)) + assert(candidates.count(_.id == latestStage.id) == 1) + + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val cancelledIds = adaptivePlan.invokePrivate(cancelObsoleteStages(latestStage, candidates)) + assert(cancelledIds == Seq(previouslyRetainedStage.id)) + assert(previouslyRetainedStage.cancelled) + assert(!latestStage.cancelled) + assert(!adaptivePlan.context.stageCache.contains(previousExchange.canonicalized)) + assert(adaptivePlan.context.stageCache.get(latestExchange.canonicalized) + .exists(_ eq latestStage)) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val stageFailure = new IllegalStateException("obsolete stage failed after cancellation") + assert(adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + previouslyRetainedStage, stageFailure, cancelledIds.toSet))) + assert(!adaptivePlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + latestStage, stageFailure, cancelledIds.toSet))) + } + } + + test("obsolete broadcast stages are not cancelled or initialized") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = spark.sql( + "SELECT /*+ BROADCAST(r) */ * FROM testData l JOIN testData2 r ON l.key = r.a") + val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case broadcast: BroadcastExchangeExec => broadcast + }.getOrElse(fail("expected an uninitialized broadcast exchange")) + + val stage = BroadcastQueryStageExec(400, exchange, exchange.canonicalized) + adaptivePlan.context.stageCache.put(exchange.canonicalized, stage) + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val replacementPlan = LocalTableScanExec(Nil, Nil, None) + + assert(adaptivePlan.invokePrivate(cancelObsoleteStages(replacementPlan, Seq(stage))).isEmpty) + assert(stage.resultOption.get().isEmpty) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized).exists(_ eq stage)) + } + } + + test("fatal obsolete stage failures are never ignored") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + val df = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val scan = LocalTableScanExec(Nil, Nil, None) + val stage = TestExchangeQueryStageExec(401, scan, scan) + val cancelledStageIds = Set(stage.id) + val wrappedFatal = new SparkFatalException(new OutOfMemoryError("fatal stage failure")) + val rawFatal = new OutOfMemoryError("fatal stage failure") + val cancellationFailure = new java.util.concurrent.CancellationException("stage cancelled") + + assert(!adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(stage, wrappedFatal, cancelledStageIds))) + assert(!adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(stage, rawFatal, cancelledStageIds))) + assert(adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(stage, cancellationFailure, cancelledStageIds))) + assert(!adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(stage, cancellationFailure, Set.empty[Int]))) + + val failedStage = TestExchangeQueryStageExec( + 402, + scan, + scan, + cancelFailure = Some(new IllegalStateException("stage cancellation failed"))) + adaptivePlan.context.stageCache.put(scan.canonicalized, failedStage) + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + assert(adaptivePlan.invokePrivate(cancelObsoleteStages(scan, Seq(failedStage))).isEmpty) + + failedStage.error.set(Some(wrappedFatal)) + assert(!adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(failedStage, wrappedFatal, Set.empty[Int]))) + assert(!adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(failedStage, rawFatal, Set.empty[Int]))) + assert(adaptivePlan.context.stageCache.get(scan.canonicalized).exists(_ eq failedStage)) + + failedStage.error.set(Some(cancellationFailure)) + assert(adaptivePlan.invokePrivate( + shouldIgnoreObsoleteStageFailure(failedStage, cancellationFailure, Set.empty[Int]))) + assert(!adaptivePlan.context.stageCache.contains(scan.canonicalized)) + } + } + + test("obsolete shuffle stage cancellation records only submitted shuffles for cleanup") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val replacementPlan = LocalTableScanExec(Nil, Nil, None) + + val unsubmitted = spark.range(8).repartition(2) + val unsubmittedPlan = + unsubmitted.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val unsubmittedExchange = unsubmittedPlan.initialPlan.collectFirst { + case exchange: ShuffleExchangeExec => exchange + }.getOrElse(fail("expected an unsubmitted shuffle exchange")) + val unsubmittedStage = TestExchangeQueryStageExec( + 200, unsubmittedExchange, unsubmittedExchange.canonicalized) + unsubmittedPlan.context.stageCache.put(unsubmittedExchange.canonicalized, unsubmittedStage) + + assert(unsubmittedExchange.futureAction.get().isEmpty) + assert(unsubmittedPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(unsubmittedStage))) == Seq(unsubmittedStage.id)) + assert(unsubmittedPlan.context.shuffleIds.isEmpty) + assert(unsubmittedExchange.futureAction.get().isEmpty) + + val submitted = spark.range(8).repartition(2) + submitted.collect() + val submittedPlan = submitted.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val submittedExchange = collect(submittedPlan) { + case stage: ShuffleQueryStageExec if stage.shuffle.futureAction.get().isDefined => + stage.shuffle + }.headOption.getOrElse(fail("expected a submitted shuffle exchange")) + val submittedShuffleId = submittedExchange.shuffleId + submittedPlan.context.shuffleIds.remove(submittedShuffleId) + val submittedStage = TestExchangeQueryStageExec( + 201, submittedExchange, submittedExchange.canonicalized) + submittedPlan.context.stageCache.put(submittedExchange.canonicalized, submittedStage) + + assert(submittedPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(submittedStage))) == Seq(submittedStage.id)) + assert(submittedPlan.context.shuffleIds.containsKey(submittedShuffleId)) + + submittedPlan.context.shuffleIds.remove(submittedShuffleId) + val failedSubmittedStage = TestExchangeQueryStageExec( + 202, + submittedExchange, + submittedExchange.canonicalized, + cancelFailure = Some(new IllegalStateException("submitted stage cancellation failed"))) + submittedPlan.context.stageCache.put(submittedExchange.canonicalized, failedSubmittedStage) + + assert(submittedPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(failedSubmittedStage))).isEmpty) + assert(submittedPlan.context.shuffleIds.containsKey(submittedShuffleId)) + assert(submittedPlan.context.stageCache.get(submittedExchange.canonicalized) + .exists(_ eq failedSubmittedStage)) + + val delegatedExchange = org.apache.spark.sql.MyShuffleExchangeExec( + submittedExchange.asInstanceOf[ShuffleExchangeExec]) + val delegatedStage = ShuffleQueryStageExec( + 203, delegatedExchange, delegatedExchange.canonicalized) + submittedPlan.context.stageCache.put(delegatedExchange.canonicalized, delegatedStage) + + assert(submittedExchange.futureAction.get().isDefined) + assert(delegatedExchange.futureAction.get().isEmpty) + assert(submittedPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(delegatedStage))).isEmpty) + assert(submittedPlan.context.stageCache.get(delegatedExchange.canonicalized) + .exists(_ eq delegatedStage)) + + submittedPlan.context.shuffleIds.remove(submittedShuffleId) + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val delegatedFailure = new IllegalStateException("opaque delegated shuffle failed") + delegatedStage.error.set(Some(delegatedFailure)) + assert(!submittedPlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + delegatedStage, delegatedFailure, Set.empty[Int]))) + assert(!submittedPlan.context.shuffleIds.containsKey(submittedShuffleId)) + assert(submittedExchange.futureAction.get().exists(!_.isCancelled)) + assert(submittedPlan.context.stageCache.get(delegatedExchange.canonicalized) + .exists(_ eq delegatedStage)) + } + } + + test("obsolete submitted shuffle stages are not cancelled before shuffle-file cleanup") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val replacementPlan = LocalTableScanExec(Nil, Nil, None) + + val submitted = spark.range(8).repartition(2) + submitted.collect() + val submittedPlan = submitted.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val submittedExchange = collect(submittedPlan) { + case stage: ShuffleQueryStageExec if stage.shuffle.futureAction.get().isDefined => + stage.shuffle + }.headOption.getOrElse(fail("expected a submitted shuffle exchange")) + val submittedAction = submittedExchange.futureAction.get().get + + val cleanup = org.apache.spark.sql.classic.Dataset.ofRows( + spark, spark.range(16).repartition(2).logicalPlan, RemoveShuffleFiles) + val cleanupPlan = cleanup.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + assert(cleanupPlan.context.qe.shuffleCleanupMode == RemoveShuffleFiles) + + val submittedStage = ShuffleQueryStageExec( + 300, submittedExchange, submittedExchange.canonicalized) + cleanupPlan.context.stageCache.put(submittedExchange.canonicalized, submittedStage) + + assert(!submittedAction.isCancelled) + assert(cleanupPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(submittedStage))).isEmpty) + assert(!submittedAction.isCancelled) + assert(cleanupPlan.context.stageCache.get(submittedExchange.canonicalized) + .exists(_ eq submittedStage)) + + val shouldIgnoreObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("shouldIgnoreObsoleteStageFailure")) + val fatalFailure = new SparkFatalException(new OutOfMemoryError("fatal shuffle failure")) + assert(!cleanupPlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + submittedStage, fatalFailure, Set.empty[Int]))) + + val stageFailure = new IllegalStateException("obsolete submitted shuffle failed") + submittedStage.error.set(Some(stageFailure)) + assert(cleanupPlan.invokePrivate(shouldIgnoreObsoleteStageFailure( + submittedStage, stageFailure, Set.empty[Int]))) + assert(cleanupPlan.context.shuffleIds.containsKey(submittedExchange.shuffleId)) + assert(!cleanupPlan.context.stageCache.contains(submittedExchange.canonicalized)) + assert(!submittedAction.isCancelled) + + val unsubmittedExchange = cleanupPlan.initialPlan.collectFirst { + case exchange: ShuffleExchangeExec => exchange + }.getOrElse(fail("expected an unsubmitted shuffle exchange")) + val unsubmittedStage = ShuffleQueryStageExec( + 301, unsubmittedExchange, unsubmittedExchange.canonicalized) + cleanupPlan.context.stageCache.put(unsubmittedExchange.canonicalized, unsubmittedStage) + + assert(unsubmittedExchange.futureAction.get().isEmpty) + assert(cleanupPlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(unsubmittedStage))) == Seq(unsubmittedStage.id)) + assert(unsubmittedExchange.futureAction.get().isEmpty) + assert(!cleanupPlan.context.stageCache.contains(unsubmittedExchange.canonicalized)) + } + } + + test("concurrent exchange reuse and obsolete cancellation coordinate by stage") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.EXCHANGE_REUSE_ENABLED.key -> "true") { + val createNonResultQueryStages = + PrivateMethod[Any](Symbol("createNonResultQueryStages")) + val cancelObsoleteStages = + PrivateMethod[Seq[Int]](Symbol("cancelObsoleteStages")) + val ignoreFailedObsoleteStageFailure = + PrivateMethod[Boolean](Symbol("ignoreFailedObsoleteStageFailure")) + val replacementPlan = LocalTableScanExec(Nil, Nil, None) + + def newPlanAndExchange(): (AdaptiveSparkPlanExec, ShuffleExchangeExec) = { + val df = spark.sql("SELECT * FROM testData JOIN testData2 ON key = a") + val adaptivePlan = df.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec] + val exchange = adaptivePlan.initialPlan.collectFirst { + case shuffle: ShuffleExchangeExec => shuffle + }.getOrElse(fail("expected a shuffle exchange in the initial adaptive plan")) + (adaptivePlan, exchange) + } + + def acquireStage( + adaptivePlan: AdaptiveSparkPlanExec, + exchange: ShuffleExchangeExec): ExchangeQueryStageExec = { + adaptivePlan.invokePrivate(createNonResultQueryStages(exchange)) + .asInstanceOf[Product].productElement(0).asInstanceOf[ExchangeQueryStageExec] + } + + type LifecycleWorker = (Thread, AtomicReference[Throwable]) + + def startWorker(name: String)(body: => Unit): LifecycleWorker = { + val started = new CountDownLatch(1) + val failure = new AtomicReference[Throwable]() + val worker = new Thread(name) { + override def run(): Unit = { + started.countDown() + try { + spark.withActive { + body + } + } catch { + case error: Throwable => failure.set(error) + } + } + } + worker.setDaemon(true) + worker.start() + assert(started.await(30, TimeUnit.SECONDS), s"$name did not start") + (worker, failure) + } + + def waitForWorker(worker: LifecycleWorker): Unit = { + worker._1.join(TimeUnit.SECONDS.toMillis(30)) + } + + def checkWorker(worker: LifecycleWorker): Unit = { + assert(!worker._1.isAlive, s"${worker._1.getName} did not finish") + Option(worker._2.get()).foreach(throw _) + } + + { + val (adaptivePlan, exchange) = newPlanAndExchange() + val reuseEntered = new CountDownLatch(1) + val releaseReuse = new CountDownLatch(1) + val existingStage = TestExchangeQueryStageExec( + 100, + exchange, + exchange.canonicalized, + reuseCallback = Some(() => { + reuseEntered.countDown() + assert(releaseReuse.await(30, TimeUnit.SECONDS), "stage reuse was not released") + })) + adaptivePlan.context.stageCache.put(exchange.canonicalized, existingStage) + + val acquiredStage = new AtomicReference[ExchangeQueryStageExec]() + val cancelledIds = new AtomicReference[Seq[Int]]() + val reuseWorker = startWorker("aqe-stage-reuse-first") { + acquiredStage.set(acquireStage(adaptivePlan, exchange)) + } + var cancellationWorker: Option[LifecycleWorker] = None + + try { + assert(reuseEntered.await(30, TimeUnit.SECONDS), "exchange reuse did not start") + cancellationWorker = Some(startWorker("aqe-stage-cancellation-after-reuse") { + cancelledIds.set(adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(existingStage)))) + }) + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(cancellationWorker.get._1.getState == Thread.State.BLOCKED) + } + } finally { + releaseReuse.countDown() + waitForWorker(reuseWorker) + cancellationWorker.foreach(waitForWorker) + } + + checkWorker(reuseWorker) + cancellationWorker.foreach(checkWorker) + assert(cancelledIds.get().isEmpty) + assert(!existingStage.cancelled) + assert(acquiredStage.get().resultOption.eq(existingStage.resultOption)) + assert(adaptivePlan.context.isSharedStageResult(existingStage.resultOption)) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized) + .exists(_ eq existingStage)) + } + + { + val (adaptivePlan, exchange) = newPlanAndExchange() + val unrelatedExchange = adaptivePlan.initialPlan.collectFirst { + case candidate: ShuffleExchangeExec + if !candidate.canonicalized.fastEquals(exchange.canonicalized) => candidate + }.getOrElse(fail("expected an unrelated shuffle exchange in the initial adaptive plan")) + val cancellationEntered = new CountDownLatch(1) + val releaseCancellation = new CountDownLatch(1) + val obsoleteStage = TestExchangeQueryStageExec( + 101, + exchange, + exchange.canonicalized, + cancelCallback = Some(() => { + cancellationEntered.countDown() + assert(releaseCancellation.await(30, TimeUnit.SECONDS), + "stage cancellation was not released") + })) + adaptivePlan.context.stageCache.put(exchange.canonicalized, obsoleteStage) + + val acquiredStage = new AtomicReference[ExchangeQueryStageExec]() + val unrelatedStage = new AtomicReference[ExchangeQueryStageExec]() + val unrelatedAcquired = new CountDownLatch(1) + val cancelledIds = new AtomicReference[Seq[Int]]() + val cancellationWorker = startWorker("aqe-stage-cancellation-first") { + cancelledIds.set(adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(obsoleteStage)))) + } + var reuseWorker: Option[LifecycleWorker] = None + var unrelatedWorker: Option[LifecycleWorker] = None + + try { + assert(cancellationEntered.await(30, TimeUnit.SECONDS), + "obsolete stage cancellation did not start") + reuseWorker = Some(startWorker("aqe-stage-reuse-after-cancellation") { + acquiredStage.set(acquireStage(adaptivePlan, exchange)) + }) + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(Set(Thread.State.BLOCKED, Thread.State.WAITING, Thread.State.TIMED_WAITING) + .contains(reuseWorker.get._1.getState)) + assert(acquiredStage.get() == null) + } + + unrelatedWorker = Some(startWorker("aqe-unrelated-stage-during-cancellation") { + unrelatedStage.set(acquireStage(adaptivePlan, unrelatedExchange)) + unrelatedAcquired.countDown() + }) + assert(unrelatedAcquired.await(10, TimeUnit.SECONDS), + "an unrelated exchange was blocked while stage cancellation was running") + assert(unrelatedStage.get().resultOption.ne(obsoleteStage.resultOption)) + assert(acquiredStage.get() == null) + } finally { + releaseCancellation.countDown() + waitForWorker(cancellationWorker) + reuseWorker.foreach(waitForWorker) + unrelatedWorker.foreach(waitForWorker) + } + + checkWorker(cancellationWorker) + reuseWorker.foreach(checkWorker) + unrelatedWorker.foreach(checkWorker) + assert(cancelledIds.get() == Seq(obsoleteStage.id)) + assert(obsoleteStage.cancelled) + assert(acquiredStage.get().resultOption.ne(obsoleteStage.resultOption)) + assert(adaptivePlan.context.stageCache.get(acquiredStage.get().plan.canonicalized) + .exists(_ eq acquiredStage.get())) + assert(adaptivePlan.context.stageCache.get(unrelatedStage.get().plan.canonicalized) + .exists(_ eq unrelatedStage.get())) + } + + { + val (adaptivePlan, obsoleteExchange) = newPlanAndExchange() + val unrelatedExchange = adaptivePlan.initialPlan.collectFirst { + case exchange: ShuffleExchangeExec + if !exchange.canonicalized.fastEquals(obsoleteExchange.canonicalized) => exchange + }.getOrElse(fail("expected an unrelated shuffle exchange in the initial adaptive plan")) + val obsoleteStage = ShuffleQueryStageExec( + 104, obsoleteExchange, obsoleteExchange.canonicalized) + adaptivePlan.context.stageCache.put(obsoleteExchange.canonicalized, obsoleteStage) + + val monitorEntered = new CountDownLatch(1) + val releaseMonitor = new CountDownLatch(1) + val monitorWorker = startWorker("aqe-obsolete-shuffle-monitor-held") { + obsoleteExchange.synchronized { + monitorEntered.countDown() + assert(releaseMonitor.await(30, TimeUnit.SECONDS), + "obsolete shuffle monitor was not released") + } + } + + val cancelledIds = new AtomicReference[Seq[Int]]() + val acquiredStage = new AtomicReference[ExchangeQueryStageExec]() + val unrelatedAcquired = new CountDownLatch(1) + var cancellationWorker: Option[LifecycleWorker] = None + var unrelatedWorker: Option[LifecycleWorker] = None + + try { + assert(monitorEntered.await(30, TimeUnit.SECONDS), + "obsolete shuffle monitor was not acquired") + cancellationWorker = Some(startWorker("aqe-shuffle-cancellation-waits-for-monitor") { + cancelledIds.set(adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(obsoleteStage)))) + }) + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(cancellationWorker.get._1.getState == Thread.State.BLOCKED) + } + + unrelatedWorker = Some(startWorker("aqe-unrelated-stage-acquisition") { + acquiredStage.set(acquireStage(adaptivePlan, unrelatedExchange)) + unrelatedAcquired.countDown() + }) + assert(unrelatedAcquired.await(10, TimeUnit.SECONDS), + "an unrelated exchange was blocked by obsolete shuffle cancellation") + assert(acquiredStage.get().resultOption.ne(obsoleteStage.resultOption)) + assert(adaptivePlan.context.stageCache.get(acquiredStage.get().plan.canonicalized) + .exists(_ eq acquiredStage.get())) + } finally { + releaseMonitor.countDown() + waitForWorker(monitorWorker) + cancellationWorker.foreach(waitForWorker) + unrelatedWorker.foreach(waitForWorker) + } + + checkWorker(monitorWorker) + cancellationWorker.foreach(checkWorker) + unrelatedWorker.foreach(checkWorker) + assert(cancelledIds.get() == Seq(obsoleteStage.id)) + assert(!adaptivePlan.context.stageCache.contains(obsoleteExchange.canonicalized)) + } + + { + val (adaptivePlan, exchange) = newPlanAndExchange() + val cancellationEntered = new CountDownLatch(1) + val releaseCancellation = new CountDownLatch(1) + val obsoleteStage = TestExchangeQueryStageExec( + 105, + exchange, + exchange.canonicalized, + cancelCallback = Some(() => { + cancellationEntered.countDown() + assert(releaseCancellation.await(30, TimeUnit.SECONDS), + "failed stage cancellation was not released") + throw new IllegalStateException("test concurrent stage cancellation failed") + })) + adaptivePlan.context.stageCache.put(exchange.canonicalized, obsoleteStage) + + val cancelledIds = new AtomicReference[Seq[Int]]() + val acquiredStage = new AtomicReference[ExchangeQueryStageExec]() + val cancellationWorker = startWorker("aqe-failed-stage-cancellation") { + cancelledIds.set(adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(obsoleteStage)))) + } + var reuseWorker: Option[LifecycleWorker] = None + + try { + assert(cancellationEntered.await(30, TimeUnit.SECONDS), + "failing stage cancellation did not start") + reuseWorker = Some(startWorker("aqe-stage-reuse-after-failed-cancellation") { + acquiredStage.set(acquireStage(adaptivePlan, exchange)) + }) + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(Set(Thread.State.BLOCKED, Thread.State.WAITING, Thread.State.TIMED_WAITING) + .contains(reuseWorker.get._1.getState)) + assert(acquiredStage.get() == null) + } + } finally { + releaseCancellation.countDown() + waitForWorker(cancellationWorker) + reuseWorker.foreach(waitForWorker) + } + + checkWorker(cancellationWorker) + reuseWorker.foreach(checkWorker) + assert(cancelledIds.get().isEmpty) + assert(!obsoleteStage.cancelled) + assert(acquiredStage.get().resultOption.eq(obsoleteStage.resultOption)) + assert(adaptivePlan.context.isSharedStageResult(obsoleteStage.resultOption)) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized) + .exists(_ eq obsoleteStage)) + } + + { + val (adaptivePlan, exchange) = newPlanAndExchange() + val obsoleteStage = TestExchangeQueryStageExec( + 102, + exchange, + exchange.canonicalized, + cancelFailure = Some(new IllegalStateException("test stage cancellation failed"))) + adaptivePlan.context.stageCache.put(exchange.canonicalized, obsoleteStage) + + val cancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(obsoleteStage))) + assert(cancelledIds.isEmpty) + assert(!obsoleteStage.cancelled) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized) + .exists(_ eq obsoleteStage)) + + val reusedStage = acquireStage(adaptivePlan, exchange) + assert(reusedStage.resultOption.eq(obsoleteStage.resultOption)) + assert(adaptivePlan.context.isSharedStageResult(obsoleteStage.resultOption)) + + obsoleteStage.error.set(Some(new IllegalStateException("shared stage failed"))) + assert(!adaptivePlan.invokePrivate(ignoreFailedObsoleteStageFailure(obsoleteStage))) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized) + .exists(_ eq obsoleteStage)) + } + + { + val (adaptivePlan, exchange) = newPlanAndExchange() + val obsoleteStage = TestExchangeQueryStageExec( + 103, + exchange, + exchange.canonicalized, + cancelFailure = Some(new IllegalStateException("test stage cancellation failed"))) + adaptivePlan.context.stageCache.put(exchange.canonicalized, obsoleteStage) + + val cancelledIds = adaptivePlan.invokePrivate( + cancelObsoleteStages(replacementPlan, Seq(obsoleteStage))) + assert(cancelledIds.isEmpty) + assert(adaptivePlan.context.stageCache.get(exchange.canonicalized) + .exists(_ eq obsoleteStage)) + + obsoleteStage.error.set(Some(new IllegalStateException("obsolete stage failed"))) + assert(adaptivePlan.invokePrivate(ignoreFailedObsoleteStageFailure(obsoleteStage))) + assert(!adaptivePlan.context.stageCache.contains(exchange.canonicalized)) + } + } + } + test("Scalar subquery") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", @@ -416,19 +1484,19 @@ class AdaptiveQueryExecSuite // A possible resulting query plan: // BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastExchange - // +-LocalShuffleReader* + // +-AQEShuffleRead local* // +- ShuffleExchange // After applied the 'OptimizeShuffleWithLocalRead' rule, we can convert all the four @@ -463,20 +1531,20 @@ class AdaptiveQueryExecSuite // A possible resulting query plan: // BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastExchange // +-HashAggregate - // +- CoalescedShuffleReader + // +- AQEShuffleRead coalesced // +- ShuffleExchange // The shuffle added by Aggregate can't apply local read. @@ -508,21 +1576,21 @@ class AdaptiveQueryExecSuite // A possible resulting query plan: // BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin // +- BroadcastExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange - // +- LocalShuffleReader* + // +- AQEShuffleRead local* // +- ShuffleExchange // +- BroadcastHashJoin // +- Filter // +- HashAggregate - // +- CoalescedShuffleReader + // +- AQEShuffleRead coalesced // +- ShuffleExchange // +- BroadcastExchange - // +-LocalShuffleReader* + // +-AQEShuffleRead local* // +- ShuffleExchange // The shuffle added by Aggregate can't apply local read. @@ -3398,6 +4466,32 @@ class AdaptiveQueryExecSuite } } + test("SPARK-57956: unmaterialized query stage exposes structural maxRows, not the estimate") { + // Build a LogicalQueryStage whose underlying stage is never materialized. Its computeStats() + // falls back to the logical plan's cost estimate - here an under-count of 0 rows - which must + // NOT be promoted to a hard maxRows bound. Otherwise EliminateLimits would drop a LIMIT that + // still needs to be applied once the stage runs (SPARK-57956). + val output = Seq(AttributeReference("a", IntegerType)()) + val logical = UnderCountLeaf(output) + val scan = LocalTableScanExec(output, Nil, None) + val exchange = BroadcastExchangeExec( + HashedRelationBroadcastMode(output, isNullAware = false), scan) + val queryStage = LogicalQueryStage(logical, BroadcastQueryStageExec(0, exchange, exchange)) + + assert(!queryStage.isMaterialized) + // The under-counted estimate is what computeStats() surfaces... + assert(queryStage.stats.rowCount.contains(BigInt(0))) + // ...but maxRows must remain the structural bound (2), not the estimate (0). + assert(queryStage.maxRows.contains(2L), + "an unmaterialized stage must expose its structural maxRows, not the row-count estimate") + + // Since the structural bound (2) exceeds the limit (1), the LIMIT must be retained. With the + // estimate wrongly promoted (maxRows = 0), EliminateLimits would instead drop it. + val limited = GlobalLimit(Literal(1), queryStage) + assert(EliminateLimits(limited).isInstanceOf[GlobalLimit], + "LIMIT must be retained when the child's structural row bound exceeds the limit") + } + test("SPARK-48037: Fix SortShuffleWriter lacks shuffle write related metrics " + "resulting in potentially inaccurate data") { withTable("t3") { @@ -4268,6 +5362,41 @@ class AdaptiveQueryExecSuite } } +private case class TestExchangeQueryStageExec( + override val id: Int, + override val plan: SparkPlan, + override val _canonicalized: SparkPlan, + runtimeRowCount: Option[BigInt] = None, + cancelFailure: Option[Throwable] = None, + reuseCallback: Option[() => Unit] = None, + cancelCallback: Option[() => Unit] = None) extends ExchangeQueryStageExec { + @volatile var cancelled: Boolean = false + + override protected def doMaterialize(): Future[Any] = Future.successful(()) + + override def getRuntimeStatistics: org.apache.spark.sql.catalyst.plans.logical.Statistics = + org.apache.spark.sql.catalyst.plans.logical.Statistics( + sizeInBytes = BigInt(0), rowCount = runtimeRowCount) + + override protected def doCancel(reason: String): Unit = { + cancelFailure.foreach { failure => + throw failure + } + cancelCallback.foreach(_()) + cancelled = true + } + + override def newReuseInstance( + newStageId: Int, + newOutput: Seq[Attribute]): ExchangeQueryStageExec = { + reuseCallback.foreach(_()) + val reuse = copy(id = newStageId) + reuse._resultOption = this._resultOption + reuse._error = this._error + reuse + } +} + /** * A minimal leaf plan with a single output attribute, used to build tiny plans for cost tests. */ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala new file mode 100644 index 0000000000000..5eb9b4a8af6f1 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/AdaptivePartialAggregationSuite.scala @@ -0,0 +1,1364 @@ +/* + * 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.execution.aggregate + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.aggregate.{Partial, PartialMerge} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.tags.ExtendedSQLTest + +/** + * Tests for runtime adaptive partial aggregation + * (see [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]). When a partial aggregate is not reducing + * rows, the operator stops aggregating and streams the remaining rows through as single-row partial + * buffers for the Final aggregate to merge. Once pass-through is active the map is frozen, and its + * output always precedes the passed-through rows: a row that collides with a frozen key is held + * behind the map and flushed only after it drains, so every group merges its buffers in the same + * order as a run that never bypasses, including order-sensitive aggregates such as `first`/`last`. + * + * The suite has two halves: + * 1. Correctness: aggregate results are identical to the reference (feature-off) run across the + * full matrix of codegen on/off, two-level map on/off, and spill/no-spill, over a range of + * aggregate shapes, key types, and `Expand`-bearing plans (ROLLUP / CUBE / GROUPING SETS / + * multi-distinct). Order-sensitive aggregates are tested against the reference too, including + * under a fan-out child that queues its whole batch behind the frozen map. + * 2. Triggering: the `numBypassingRows` metric proves the bypass actually fires when (and only + * when) it should -- high-cardinality input bypasses, low-cardinality input keeps aggregating, + * the feature switch and eligibility rules are honored, and both check points work. + */ +@ExtendedSQLTest +class AdaptivePartialAggregationSuite extends QueryTest with SharedSparkSession + with AdaptiveSparkPlanHelper { + + import testImplicits._ + + // A `testFallbackStartsAt` setting ("fastMapCounter, regularMapCounter") that makes the regular + // map fall back (spill) periodically, exercising the spill-check decision path in both the + // codegen and interpreted aggregation paths. Kept moderate so low-cardinality inputs (which are + // never bypassed and therefore really spill) do not open an unbounded number of spill readers. + private val forceSpillFallback = "4, 16" + + // The upstream `CombineAdjacentAggregation` and `ReplaceHashWithSortAgg` rules would change the + // plan of these small single-partition queries away from a Partial+Final `HashAggregateExec`: + // the former merges the two adjacent phases (no shuffle in between) into a single `Complete` + // aggregate, and the latter converts a hash aggregate to a sort aggregate when the input is + // already sorted by the grouping key (a `Range` over an ascending `id` key). The adaptive + // feature lives in the partial hash aggregation, so both rules are disabled to keep that + // structure in the tests. + private val fixedPlanConfs = Seq( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false", + SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") + + // Whether `agg` is a partial or partial-merge phase. This matches every `HashAggregateExec` + // whose modes are all `Partial`/`PartialMerge`, including the DISTINCT intermediate phase + // (`PartialMerge ++ Partial`, the non-distinct aggregates in `PartialMerge` and the distinct + // ones in `Partial`) that the feature now bypasses. It is a superset of the phases the feature + // actually applies to: a pure `PartialMerge` de-duplication phase (excluded by its required + // distribution) and an ineligible group-by-only `Final` (vacuously all-`Partial`) also match, + // but neither registers a `numBypassingRows` metric, so they add 0 to the bypass count below. + private def isPartialPhase(agg: HashAggregateExec): Boolean = + agg.aggregateExpressions.forall(a => a.mode == Partial || a.mode == PartialMerge) + + /** + * Runs `build` with adaptive partial aggregation disabled (the reference) and then across the + * full configuration matrix with it enabled, asserting every enabled run matches the reference. + * + * `build` takes the number of input partitions, which the matrix varies along with everything + * else, because the plan shape decides which parts of the feature run at all. When the two + * aggregates end up in one whole-stage -- no `Exchange` between them -- the partial aggregate's + * output feeds the Final's `doConsume` directly and never reaches + * `BufferedRowIterator.currentRows`, so `shouldStop()` stays false for the whole build and + * neither `needStopCheck` nor the resumed-build path is exercised. Splitting them puts the + * streamed rows through the output buffer and runs both. + * + * More than one input partition is necessary but not sufficient for that split: a `Range` keyed + * directly on `id` already reports an output partitioning that satisfies the Final aggregate's + * `ClusteredDistribution`, so `EnsureRequirements` inserts no `Exchange` however many partitions + * it has. Tests that want the split shape group on a derived key (a cast, say) so the input + * partitioning no longer satisfies the requirement. + * + * `expectBypass` ties the correctness guarantee to the triggering guarantee: beyond matching the + * reference, every cell must either actually stream rows through (when true) or keep + * aggregating (when false). Without it a test could silently stop exercising pass-through if the + * input stopped being bypassable, and only this assertion makes that fail loudly. + */ + private def checkAdaptiveMatchesReference( + build: Int => DataFrame, + expectBypass: Boolean = true): Unit = { + for { + inputPartitions <- Seq(1, 2) + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } { + // The reference is built with the same partitioning, so only the feature differs. + val reference = withSQLConf( + (SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") +: fixedPlanConfs: _*) { + build(inputPartitions).collect().toSeq + } + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + // Small `minRows` so the periodic check runs on modest inputs. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8") ++ + spillConf ++ fixedPlanConfs): _*) { + val msg = s"inputPartitions=$inputPartitions wholeStage=$wholeStage " + + s"twoLevelMap=$twoLevelMap forceSpill=$forceSpill" + withClue(msg) { + // Collect once so the metrics are populated, then check whether the bypass fired for + // this cell. The metric lives on the partial `HashAggregateExec` phases, so those are the + // operators the assertion reads. + val df = build(inputPartitions) + df.collect() + val skipped = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if isPartialPhase(agg) => + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.sum + if (expectBypass) { + assert(skipped > 0, + s"expected rows to bypass partial aggregation, got $skipped bypassed rows") + } else { + assert(skipped == 0, + s"expected no rows to bypass partial aggregation, got $skipped bypassed rows") + } + checkAnswer(df, reference) + } + } + } + } + + /** + * The observable per-run counters we assert on, all read from the partial `HashAggregateExec` in + * a single execution so the metrics are not double-counted: + * - `skipped`: our self-reported `numBypassingRows` metric. + * - `partialOutputRows`: the partial aggregate's own `numOutputRows`. An independent, + * pre-existing counter driven by the normal output path, so it is the ground truth for + * whether rows were streamed through -- it equals the distinct key count when aggregation is + * effective and climbs toward the input row count once the operator bypasses. + * - `spillBytes`: the partial aggregate's `spillSize`. Reliable only when no fallback is + * forced: on the interpreted path this is derived from the task-cumulative memory-spill + * counter, so a forced fallback (or downstream shuffle-write spill) can inflate it. + * Asserted only by the periodic check test, which forces no fallback; use + * `tasksFallBacked` otherwise. + * - `tasksFallBacked`: the partial aggregate's `numTasksFallBacked`, incremented only when the + * regular map actually falls back into sort-based aggregation. When the spill check bypasses + * at the spill boundary the sorter is never created, so this stays 0 -- direct, per-operator + * evidence the bypass replaced the sort fallback. + */ + private case class AggCounters( + skipped: Long, + partialOutputRows: Long, + spillBytes: Long, + tasksFallBacked: Long) + + // Verifies `df` (an already-collected bypassing run) produces the same results as the feature-off + // reference. `build` is re-run for the reference so it gets a genuinely non-adaptive plan rather + // than reusing the bypassing run's cached one. + private def checkAgainstReference(df: DataFrame, build: () => DataFrame): Unit = { + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + build().collect().toSeq + } + checkAnswer(df, reference) + } + + private def runAndReadCounters(build: () => DataFrame): AggCounters = { + // The triggering tests assert on metrics, so also verify the bypassing run produces the same + // results as the feature-off reference. + val df = build() + df.collect() + val partialAggs = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if isPartialPhase(agg) => agg + } + // A partial aggregate is always present for the grouped queries these tests use. + assert(partialAggs.nonEmpty, "expected a partial HashAggregateExec in the plan") + val counters = AggCounters( + // The metric is only registered on aggregates the feature applies to; an aggregate without + // it bypassed nothing. + skipped = partialAggs.map(_.metrics.get("numBypassingRows").map(_.value).getOrElse(0L)).sum, + partialOutputRows = partialAggs.map(_.metrics("numOutputRows").value).sum, + spillBytes = partialAggs.map(_.metrics("spillSize").value).sum, + tasksFallBacked = partialAggs.map(_.metrics("numTasksFallBacked").value).sum) + checkAgainstReference(df, build) + counters + } + + private def numBypassingRows(build: () => DataFrame): Long = runAndReadCounters(build).skipped + + // Returns the bypassed-row count per partial `HashAggregateExec` phase, keyed by the number of + // grouping keys, and verifies the run matches the feature-off reference. A `count(DISTINCT ...)` + // group-by has two such phases -- the de-duplication partial (grouping on key + distinct + // columns) and the distinct partial (grouping on the keys only, whose non-distinct aggregates run + // in `PartialMerge`) -- so their bypasses can be told apart by the grouping key count. + private def bypassRowsByGroupingKeyCount(build: () => DataFrame): Map[Int, Long] = { + val df = build() + df.collect() + val byKeyCount = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if isPartialPhase(agg) => + agg.groupingExpressions.length -> + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.groupBy(_._1).map { case (n, pairs) => n -> pairs.map(_._2).sum } + // The pure `PartialMerge` de-duplication phase (grouping on key + distinct columns) is the one + // phase the `exists(_.mode == Partial)` guard exists to exclude, and only when it carries + // non-distinct aggregates: with none at all its `aggregateExpressions` is empty, so the guard's + // `isEmpty` disjunct admits it and only its required distribution keeps it out, and + // `dedupPhases` below (which requires a non-empty `aggregateExpressions`) does not collect it. + // It shares the 2-key bucket above with the leading `Partial` phase, so pin its ineligibility + // directly. + val dedupPhases = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if agg.aggregateExpressions.nonEmpty && + agg.aggregateExpressions.forall(_.mode == PartialMerge) => agg + } + assert(dedupPhases.forall(!_.metrics.contains("numBypassingRows")), + "the pure-PartialMerge de-duplication phase must stay ineligible") + checkAgainstReference(df, build) + byKeyCount + } + + /** + * Runs `body` once per (wholeStage, twoLevelMap) combination with the feature enabled and a small + * `minRows`, threading a descriptive clue for failure messages. + * + * The fast (first-level) map is append-only and never spills, so only the regular (second-level) + * map can reach a spill boundary. With the default fast-map capacity (2^16) a small + * high-cardinality input would be fully absorbed by the fast map and never reach the regular map, + * so nothing could ever bypass. To make the triggering tests meaningful when the two-level map is + * on, we shrink the fast map via the first field of `testFallbackStartsAt` so rows fall through + * to the regular map. `regularFallback` optionally sets the second field to also force the + * regular map to spill (for the spill check); when 0 the regular map does not spill. + */ + private def forEachCodegenAndMap( + minRows: Long = 8, + regularFallback: Int = 0, + minCompaction: Double = -1.0)( + body: String => Unit): Unit = { + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + // Shrink the fast map to 4 keys when it is on so rows reach the regular map. The second field + // controls regular-map spilling; 0 means "never" (a large sentinel). + val fallbackConf = if (twoLevelMap || regularFallback > 0) { + val fastCap = if (twoLevelMap) 4 else 1 + val regular = if (regularFallback > 0) regularFallback else Int.MaxValue + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, $regular") + } else { + Nil + } + // A negative value means "leave the threshold at its default". + val thresholdConf = if (minCompaction >= 0.0) { + Seq(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.toString) + } else { + Nil + } + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> minRows.toString) ++ + fallbackConf ++ thresholdConf ++ fixedPlanConfs): _*) { + body(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") + } + } + } + + ///////////////////////////////////////////////////////////////////////////// + // Part 1: Correctness -- results identical to the feature-off reference. + ///////////////////////////////////////////////////////////////////////////// + + test("results unchanged for high-cardinality input that bypasses partial aggregation") { + // Every grouping key is distinct, so partial aggregation reduces nothing and should be + // bypassed by the periodic check. + checkAdaptiveMatchesReference { parts => + spark.range(0, 200, 1, parts) + .select($"id".cast("string") as "k", ($"id" * 2) as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", max($"v") as "m") + } + } + + test("results unchanged for low-cardinality input that keeps partial aggregation") { + // Few distinct keys, high reduction: partial aggregation is effective and should be kept, so + // the bypass metric must stay zero in every cell. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 600, 1, parts) + .select(($"id" % 5).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", min($"v") as "mn", max($"v") as "mx") + }) + } + + test("results unchanged for medium-cardinality input near the reduction threshold") { + // Roughly half the rows are distinct keys; exercises the boundary of the ratio checks. The + // overall compaction ratio (~2.0) is above the threshold, but the *first* periodic check still + // sees the leading distinct keys and fires, so the bypass must be observable too. + checkAdaptiveMatchesReference { parts => + spark.range(0, 1000, 1, parts) + .select(($"id" % 500).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with multiple grouping keys and string keys") { + checkAdaptiveMatchesReference { parts => + spark.range(0, 500, 1, parts) + .select( + concat(lit("g"), ($"id" % 300).cast("string")) as "k1", + ($"id" % 7) as "k2", + $"id" as "v") + .groupBy($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with nullable grouping keys") { + // Nulls are sparse enough (1 in 40) that the keys stay close to unique and the input really + // does bypass; a denser null key would lift the compaction ratio above the threshold and the + // test would never engage the feature. + checkAdaptiveMatchesReference { parts => + spark.range(0, 400, 1, parts) + .select( + when($"id" % 40 === 0, lit(null)).otherwise($"id").cast("string") as "k", + $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged with average (multi-slot buffer) aggregate") { + // avg has a two-slot partial buffer (sum, count); pass-through buffers must carry all slots. + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") as "k", ($"id" + 1) as "v") + .groupBy($"k") + .agg(avg($"v") as "a", sum($"v") as "s") + } + } + + test("results unchanged with a mix of many aggregate functions and buffer types") { + // Exercises a wide pass-through buffer spanning several aggregate buffer layouts at once: + // sum (decimal), avg (double), count, min/max, first/last, and stddev (declarative buffer). + // The imperative-buffer case is covered separately (see the `approx_count_distinct` test). + checkAdaptiveMatchesReference { parts => + spark.range(0, 400, 1, parts) + .select( + $"id" as "k", + ($"id" % 97).cast("decimal(10,2)") as "d", + ($"id" % 13).cast("double") as "dbl") + .groupBy($"k") + .agg( + sum($"d") as "sd", + avg($"dbl") as "ad", + count(lit(1)) as "c", + min($"dbl") as "mn", + max($"dbl") as "mx", + first($"dbl") as "f", + last($"dbl") as "l", + stddev($"dbl") as "sd2") + } + } + + test("results unchanged with an imperative-buffer aggregate") { + // `approx_count_distinct` uses `HyperLogLogPlusPlus`, an `ImperativeAggregate` whose buffer + // state is written by `initialize(buffer)` rather than by a projection, so a pass-through + // single-row buffer has to be re-initialized with `copyFrom(initialAggregationBuffer)` for + // every row. No declarative aggregate exercises that reset. It also reports + // `supportCodegen = false`, so the operator only ever runs on `TungstenAggregationIterator`; + // keep it in its own test rather than folding it into a codegen cell that would quietly + // become interpreted. + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(approx_count_distinct($"v") as "c") + } + } + + test("results unchanged with filtered aggregate functions") { + // A `FILTER (WHERE ...)` aggregate is compiled into a per-row guard around the buffer update + // rather than a separate filtering operator: `If(filter, update, buffer)` in the interpreted + // path and an `if (!cond) continue` guard in the generated code. Pass-through reuses those + // exact update expressions, so a bypassed row whose filter is false contributes nothing to its + // single-row buffer. The all-true and all-false filters pin the two extremes, and the fully + // distinct grouping keys ensure rows bypass (in the regular-map-only configurations) so the + // filter guard actually runs in the pass-through path. + withTempView("t") { + spark.range(0, 400, 1, 1) + .select($"id".cast("string") as "k", ($"id" % 100) as "v") + .createOrReplaceTempView("t") + checkAdaptiveMatchesReference { parts => + spark.sql( + """SELECT k, + | sum(v) FILTER (WHERE v % 2 = 0) AS s_even, + | count(1) FILTER (WHERE v > 50) AS c_gt50, + | avg(v) FILTER (WHERE v > 25) AS a_gt25, + | sum(v) FILTER (WHERE true) AS s_all, + | sum(v) FILTER (WHERE false) AS s_none + |FROM t GROUP BY k""".stripMargin) + } + } + } + + test("results unchanged with decimal and date grouping keys") { + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select( + ($"id" % 280).cast("decimal(12,3)") as "k1", + date_add(lit(java.sql.Date.valueOf("2020-01-01")), ($"id" % 250).cast("int")) as "k2", + $"id" as "v") + .groupBy($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + } + } + + test("results unchanged for group-by-only (distinct) with no aggregate functions") { + // No aggregate functions: the pass-through buffer is a zero-column UnsafeRow, so the output is + // just the grouping key. High-cardinality keys should bypass, and the de-duplicated result must + // still match the reference. + checkAdaptiveMatchesReference { parts => + spark.range(0, 400, 1, parts) + .select(($"id" % 350) as "k1", ($"id" % 11) as "k2") + .distinct() + } + } + + test("results unchanged for group-by-only with duplicate keys (Final phase must not bypass)") { + // A group-by-only aggregate has an empty `aggregateExpressions`, so checking the aggregate + // modes alone is vacuously true and could wrongly admit the `Final` phase of the two-phase + // plan. With duplicate keys, a bypassing `Final` would skip its de-duplication and return + // duplicate rows. The two-level map off variants route the rows to the regular map so the + // periodic check fires and the regression would show up. + checkAdaptiveMatchesReference { parts => + spark.range(0, 1000, 1, parts) + .select(($"id" % 10) as "c") + .distinct() + } + } + + test("results unchanged when a large frozen map is output before pass-through streaming") { + // A larger `minRows` lets the map accumulate many keys before the periodic check bypasses, + // so the early map output (which also frees the map) spans several drain cycles and re-enters + // the map-output function; the results must still match the feature-off reference. + val query = () => spark.range(0, 400000, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "200000", + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false") ++ fixedPlanConfs): _*) { + val df = query() + df.collect() + val skipped = collect(df.queryExecution.executedPlan) { + case agg: HashAggregateExec if isPartialPhase(agg) => + agg.metrics.get("numBypassingRows").map(_.value).getOrElse(0L) + }.sum + assert(skipped > 0, + s"expected the large frozen map to eventually bypass, got $skipped bypassed rows") + val reference = withSQLConf( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false") { + query().collect().toSeq + } + checkAnswer(df, reference) + } + } + + test("distinct aggregation stays correct") { + checkAdaptiveMatchesReference { parts => + spark.range(0, 300, 1, parts) + .select($"id".cast("string") as "k", ($"id" % 50) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd", sum($"v") as "s") + } + } + + test("distinct with plain and filtered non-distinct aggregates") { + // One query carries all three shapes through the DISTINCT intermediate phase + // (`PartialMerge ++ Partial`): the distinct aggregate (`count(DISTINCT v)`), a plain + // non-distinct aggregate (`sum(v)`), and a filtered non-distinct aggregate + // (`avg(v) FILTER (...)`, whose `FILTER` is applied in the leading `Partial` phase only). + // Fully distinct keys and values make neither partial phase reduce anything, so both bypass in + // the same execution: asserting the 2-key phase (de-duplication, all `Partial`) and the 1-key + // phase (distinct partial, `PartialMerge ++ Partial`) together proves the two bypasses coexist, + // the plain and filtered non-distinct buffers pass through correctly, and the results still + // match the feature-off reference. + withTempView("t") { + spark.range(0, 400, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .createOrReplaceTempView("t") + forEachCodegenAndMap() { clue => + val df = () => spark.sql( + """SELECT k, + | count(DISTINCT v) AS cd, + | sum(v) AS s, + | avg(v) FILTER (WHERE v > 25) AS a_gt25 + |FROM t GROUP BY k""".stripMargin) + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (PartialMerge++Partial, grouping on k) to bypass, " + + s"got $byKeyCount") + } + } + } + } + + test("distinct with an order-sensitive non-distinct aggregate across partitions") { + // A single-partition `range` fuses the whole four-phase DISTINCT plan into one stage, so the + // split-topology path (the frozen map draining one row per queued row, the queue flush, and + // the `shouldStop()` re-entry) never runs for a `PartialMerge` member, including an + // order-sensitive one. Deriving both grouping columns and varying the partition count forces + // the exchanges; carrying `first`/`last` makes the pass-through's merge-into-an-empty-buffer + // step observable. The forced-spill cells are omitted because the sort-based fallback reorders + // `first`/`last` in the feature-off reference arm too. + forEachCodegenAndMap() { clue => + Seq(1, 2).foreach { parts => + val df = () => spark.range(0, 400, 1, parts) + .select(($"id" % 100).cast("string") as "k", ($"id" % 7) as "v", $"id" as "w") + .groupBy($"k") + .agg(countDistinct($"v") as "cd", sum($"w") as "s", + first($"w") as "f", last($"w") as "l") + withClue(s"$clue parts=$parts") { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (PartialMerge++Partial, grouping on k) to bypass, " + + s"got $byKeyCount") + } + } + } + } + + test("an imperative aggregate stays correct in the distinct intermediate phase") { + // `approx_count_distinct` uses `HyperLogLogPlusPlus`, an `ImperativeAggregate` whose buffer is + // written by `initialize`/`merge` rather than a projection, so the distinct intermediate phase + // runs on `TungstenAggregationIterator` (supportCodegen = false). Asserting both phases bypass + // -- the de-duplication partial (grouping on `k` + `v`) and the distinct partial (grouping on + // `k`, whose `PartialMerge` member is imperative) -- proves the imperative buffer is reset then + // merged with the incoming buffer on pass-through, and the results still match the reference. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 50).cast("string") as "k", ($"id" % 20) as "v", $"id" as "x") + .groupBy($"k") + .agg(approx_count_distinct($"x") as "acd", countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the de-duplication partial (grouping on k, v) to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (imperative PartialMerge, grouping on k) to bypass, " + + s"got $byKeyCount") + } + } + } + + test("distinct aggregation bypasses on high-cardinality input") { + // The `PartialMerge` phase of the multi-phase distinct plan always aggregates (it requires a + // distribution, so it is never eligible), so the rows reaching the distinct `Partial` phase + // are de-duplicated and pass-through carries exactly one distinct value each. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 1000, 1, 1) + .select(($"id" % 100).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + assert(numBypassingRows(df) > 0, + "expected a distinct partial aggregation to bypass for high-cardinality input") + } + } + } + + test("count distinct: the de-duplication partial aggregate bypasses") { + // `count(DISTINCT v) GROUP BY k` plans two `Partial` phases: the de-duplication partial groups + // on (k, v) and the distinct partial groups on (k). Fully distinct (k, v) pairs make the + // de-duplication partial (2 grouping keys) reduce nothing, so it must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 4).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + } + } + } + + test("count distinct: the distinct partial aggregate bypasses") { + // Mirror of the test above for the other phase: with many distinct keys but few distinct + // values per key, the (k, v) de-duplication partial reduces well while the distinct partial + // (1 grouping key) sees a fresh key per row and must bypass. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id".cast("string") as "k", ($"id" % 2) as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("count distinct: both partial aggregates bypass and results stay correct") { + // Fully distinct keys and fully distinct values: neither partial phase reduces anything, so + // both bypass in the same execution. The de-duplication partial keeps the (k, v) pairs unique + // and the distinct partial counts them, so the result must still match the reference. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 400, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(countDistinct($"v") as "cd") + withClue(clue) { + val byKeyCount = bypassRowsByGroupingKeyCount(df) + assert(byKeyCount.get(2).exists(_ > 0), + s"expected the (k, v) de-duplication partial to bypass, got $byKeyCount") + assert(byKeyCount.get(1).exists(_ > 0), + s"expected the distinct partial (grouping on k) to bypass, got $byKeyCount") + } + } + } + + test("global aggregation (no grouping keys) is never bypassed and stays correct") { + withSQLConf(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true") { + checkAnswer( + spark.range(0, 100, 1, 1).agg(sum($"id") as "s", count(lit(1)) as "c"), + Row(4950L, 100L)) + } + } + + test("results unchanged with an empty input") { + // No rows means the check points never fire; the metric must stay zero. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 0, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c") + }) + } + + // The following four tests cover plans where an `ExpandExec` sits below the partial aggregate + // (ROLLUP / CUBE / GROUPING SETS / multi-distinct). PR apache/spark#28804 statically disabled its + // skip-partial-aggregate optimization whenever an Expand was present, but that was a performance + // heuristic guarding its *static* row sampling, not a correctness requirement. Our decision is + // made at runtime from the observed compaction ratio, so we deliberately do not port that + // exclusion. These tests assert results stay correct with the exclusion absent. + // + // The ROLLUP and CUBE cases below use two grouping columns, where the grand-total set keeps the + // compaction ratio high enough that they decline to bypass -- they cover the eligible-but- + // declining side. (Widening the rollup lowers the ratio: with five distinct columns the same + // shape does bypass.) The GROUPING SETS and multi-distinct tests, and `pass-through fires for + // high-cardinality input below an Expand`, cover an Expand that bypasses. + + test("results unchanged for ROLLUP (Expand below partial aggregate)") { + // The grand-total group repeats on every expanded row, so the compaction ratio stays above the + // threshold and the partial aggregate declines to bypass in every cell. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 400, 1, parts) + .select(($"id" % 200) as "k1", ($"id" % 7) as "k2", $"id" as "v") + .rollup($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + }) + } + + test("results unchanged for CUBE (Expand below partial aggregate)") { + // Same as ROLLUP: the grand-total group keeps the ratio above the threshold, so nothing + // bypasses. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 400, 1, parts) + .select(($"id" % 150) as "k1", ($"id" % 5) as "k2", $"id" as "v") + .cube($"k1", $"k2") + .agg(sum($"v") as "s", count(lit(1)) as "c") + }) + } + + test("results unchanged for GROUPING SETS (Expand below partial aggregate)") { + // No `()` grouping set, and both keys distinct, so every expanded row is a fresh key and the + // input genuinely bypasses. A grand-total set would collapse all rows into one group and lift + // the compaction ratio above the threshold (see the ROLLUP and CUBE tests below). + withTempView("t") { + spark.range(0, 400, 1, 1) + .select($"id" as "k1", ($"id" + 1000) as "k2", $"id" as "v") + .createOrReplaceTempView("t") + checkAdaptiveMatchesReference { parts => + spark.sql( + """SELECT k1, k2, sum(v) AS s, count(1) AS c + |FROM t + |GROUP BY k1, k2 GROUPING SETS ((k1, k2), (k1), (k2))""".stripMargin) + } + } + } + + test("results unchanged for multi-distinct (Expand below partial aggregate)") { + checkAdaptiveMatchesReference { parts => + spark.range(0, 400, 1, parts) + .select(($"id" % 100).cast("string") as "k", ($"id" % 30) as "a", ($"id" % 40) as "b") + .groupBy($"k") + .agg(countDistinct($"a") as "da", countDistinct($"b") as "db", sum($"a") as "s") + } + } + + test("results unchanged under a fused Union (child yields from a nested helper)") { + // `UnionExec` wraps each child's produce in its own helper, so a streamed row that fills the + // output buffer returns only as far as the aggregate's build loop. Reaching the end of the + // child's produce therefore does not mean the input is exhausted, and treating it as such + // drops the rest of the partition. + checkAdaptiveMatchesReference { parts => + spark.range(0, 100, 1, parts + 1) + .union(spark.range(100, 200, 1, parts + 1)) + .groupBy("id").count() + } + } + + test("results unchanged when an Exchange separates the two aggregates") { + // Grouping on a derived key stops the `Range`'s output partitioning from satisfying the Final + // aggregate's `ClusteredDistribution`, so the plan keeps an `Exchange` and the two aggregates + // land in separate whole-stages. That is the shape where streamed rows actually pass through + // `BufferedRowIterator.currentRows` -- fused, they go straight into the Final's hash map and + // neither `needStopCheck` nor the resumed build is reached. + checkAdaptiveMatchesReference { parts => + spark.range(0, 200, 1, parts) + .select($"id".cast("string") as "k", ($"id" * 2) as "v") + .groupBy($"k") + .agg(sum($"v") as "s", count(lit(1)) as "c", max($"v") as "m") + } + } + + test("order-sensitive aggregates match a non-bypassed run") { + // Bypassing must merge a group's buffers in the same order as a run that never bypasses: a + // group can straddle the freeze and hold both a map buffer and pass-through buffers, and the + // `Final` merges in emit order. The queue holds each passed-through row behind the frozen map + // and flushes it only after the map drains, so the map buffer always precedes the colliding + // pass-through buffers on both execution paths -- exactly the merge order of a run that never + // bypasses, so every enabled cell must match the feature-off reference and + // `spark.sql.codegen.wholeStage` must not be observable. + // + // The flip fires at the first periodic check: with `minRows=8` it lands on the 8th aggregated + // row, when the map already holds 8 distinct keys (id 0 maps to -1, ids 1-7 to keys 1-7), so + // the compaction ratio 1.0 is below `minCompaction` (1.05) and every remaining row streams; + // id 8 is the first bypassed row. At `dupAt=8` the colliding row is that first bypassed row, + // queued behind the frozen map and flushed only after the map drains, so its buffer still + // reaches the `Final` after the frozen map's and `first`/`last` match the merge order; at + // `dupAt=9` the colliding row streams directly after the map and the group stays in input + // order. Both plan shapes are exercised separately: splitting the aggregates with an + // `Exchange` (or not) changes whether streamed rows pass through + // `BufferedRowIterator.currentRows` or feed the `Final`'s `doConsume` directly. + for { + splits <- Seq(1, 2) + derivedKey <- Seq(false, true) + dupAt <- Seq(8, 9) + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } { + val query = () => { + val base = when($"id" === 0 || $"id" === dupAt, lit(-1L)).otherwise($"id") + spark.range(0, 40, 1, splits) + .select(if (derivedKey) base.cast("string") else base as "k", $"id" as "v") + .toDF("k", "v") + .groupBy($"k") + .agg(first($"v") as "f", last($"v") as "l") + } + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + def run(enabled: Boolean): Seq[String] = withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> enabled.toString, + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8") ++ + spillConf ++ fixedPlanConfs): _*) { + query().collect().toSeq.map(_.toString()).sorted + } + withClue(s"splits=$splits derivedKey=$derivedKey dupAt=$dupAt " + + s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap forceSpill=$forceSpill: ") { + assert(run(enabled = true) == run(enabled = false), + "adaptive partial aggregation changed an order-sensitive result") + } + } + } + + test("fan-out below a split aggregate preserves the merge order") { + // A `GenerateExec` expands a collection without checking `shouldStop()`, so in the + // exchange-split shape the whole fan-out batch of the trigger input row is queued at once + // rather than one row at a time. Each queued row advances the frozen-map output by one row, + // and the queue is flushed only after the map fully drains, so a group whose rows straddle + // the freeze point still merges its map buffer before its pass-through buffers, matching a + // non-bypassed run. + // + // The colliding row is the first row of the trigger batch, so it is the first one queued: a + // design that emitted queued rows ahead of the frozen map would put it before its map buffer. + // Key 2 is inserted early (by id 1) and drained third, so id 4's first exploded element + // collides with it; the other two fan-out tests cover collisions on later batch elements. + checkAdaptiveMatchesReference { parts => + spark.range(0, 20, 1, parts) + .select($"id", explode(array( + when($"id" === 4, lit(2L)).otherwise($"id" * 2), + $"id" * 2 + 1)) as "k") + .select($"k", ($"id" * 100 + $"k") as "v") + .groupBy($"k").agg(first($"v") as "f", last($"v") as "l") + } + } + + test("a wide fan-out batch stays queued behind the frozen map") { + // A colliding key must not merge before its map buffer. This batch is four rows wide and + // collides twice, so the colliding pass-through rows must both wait behind the frozen map for + // the whole map to drain; a design that let any part of the batch escape ahead of the map + // breaks the merge order for at least one collision. + checkAdaptiveMatchesReference { parts => + spark.range(0, 20, 1, parts) + .select($"id", explode(array( + $"id" * 4, + when($"id" === 2, lit(3L)).otherwise($"id" * 4 + 1), + $"id" * 4 + 2, + when($"id" === 2, lit(7L)).otherwise($"id" * 4 + 3))) as "k") + .select($"k", ($"id" * 100 + $"k") as "v") + .groupBy($"k").agg(first($"v") as "f", last($"v") as "l") + } + } + + test("a frozen map larger than the fan-out batch preserves the merge order") { + // The queue bounds the held rows to the batch width regardless of the map size, and the flush + // waits for the whole map to drain. This query grows the map well past the batch width, so a + // design that flushed the held rows early would leave the colliding key's map buffer behind + // them. Every enabled cell must match the feature-off reference, so each compares directly + // against a non-bypassed run, across both the generated and interpreted paths. + for { + splits <- Seq(1, 2) + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + forceSpill <- Seq(true, false) + } { + val query = () => { + spark.range(0, 100, 1, splits) + .select($"id", explode(array( + $"id" * 2, + when($"id" === 16, lit(5L)).otherwise($"id" * 2 + 1))) as "k") + .select($"k", ($"id" * 100 + $"k") as "v") + .groupBy($"k").agg(first($"v") as "f", last($"v") as "l") + } + val spillConf = if (forceSpill) { + Seq("spark.sql.TungstenAggregate.testFallbackStartsAt" -> forceSpillFallback) + } else { + Nil + } + def run(enabled: Boolean): DataFrame = withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> enabled.toString, + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "32") ++ + spillConf ++ fixedPlanConfs): _*) { + query() + } + val reference = run(enabled = false) + val adaptive = run(enabled = true) + withClue(s"splits=$splits wholeStage=$wholeStage " + + s"twoLevelMap=$twoLevelMap forceSpill=$forceSpill: ") { + checkAnswer(adaptive, reference) + } + } + } + + test("results unchanged below a generator that cannot yield mid-fan-out") { + // `GenerateExec` expands a collection without checking `shouldStop()`, so it cannot bound the + // output buffer between the rows of one input row. Each passed-through row queues a copy and + // advances the frozen-map output by one row, so the whole fan-out batch is held behind the map + // and the memory is bounded by the batch width rather than by how many rows the child emits. + checkAdaptiveMatchesReference { parts => + spark.range(0, 1, 1, parts) + .select(explode(sequence(lit(1), lit(500))) as "k") + .groupBy($"k").count() + } + } + + ///////////////////////////////////////////////////////////////////////////// + // Part 2: Triggering -- the bypass fires when, and only when, it should. + ///////////////////////////////////////////////////////////////////////////// + + test("pass-through fires for high-cardinality input, not for low-cardinality input") { + forEachCodegenAndMap() { clue => + // Fully distinct keys: partial aggregation reduces nothing, so rows must bypass. + val highCard = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(highCard) > 0, + "expected some rows to bypass partial aggregation for high-cardinality input") + } + // Few distinct keys, high reduction: partial aggregation is effective, nothing bypasses. + val lowCard = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(lowCard) == 0, + "expected no rows to bypass partial aggregation for low-cardinality input") + } + } + } + + test("group-by-only pass-through fires for high-cardinality input") { + forEachCodegenAndMap() { clue => + val distinctKeys = () => spark.range(0, 200, 1, 1).select($"id" as "k").distinct() + withClue(clue) { + assert(numBypassingRows(distinctKeys) > 0, + "expected group-by-only rows to bypass partial aggregation for high-cardinality input") + } + } + } + + test("filtered aggregate functions are eligible for pass-through") { + // The filter clause does not change eligibility: a partial aggregate over `FILTER (WHERE ...)` + // functions still bypasses on high-cardinality input, and the per-row filter guard runs inside + // the pass-through single-row buffer update. + forEachCodegenAndMap() { clue => + withTempView("t") { + spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", ($"id" % 100) as "v") + .createOrReplaceTempView("t") + val df = () => spark.sql( + """SELECT k, sum(v) FILTER (WHERE v % 2 = 0) AS s + |FROM t GROUP BY k""".stripMargin) + withClue(clue) { + assert(numBypassingRows(df) > 0, + "expected filtered-aggregate rows to bypass partial aggregation for high-cardinality " + + "input") + } + } + } + } + + test("pass-through fires for high-cardinality input below an Expand") { + // The static PR#28804 heuristic would have refused to skip whenever an Expand was present; our + // runtime decision skips because the expanded rows genuinely do not reduce. GROUPING SETS over + // two single-column, fully-distinct sets is used (rather than ROLLUP/CUBE) so there is no + // grand-total group lifting the compaction ratio above the threshold: every expanded row is a + // fresh key, so all configurations bypass. + forEachCodegenAndMap() { clue => + withTempView("t") { + spark.range(0, 200, 1, 1) + .select($"id".as("k1"), ($"id" + 1000).as("k2"), $"id".as("v")) + .createOrReplaceTempView("t") + val gs = () => spark.sql( + """SELECT k1, k2, sum(v) AS s + |FROM t + |GROUP BY k1, k2 GROUPING SETS ((k1), (k2))""".stripMargin) + withClue(clue) { + assert(numBypassingRows(gs) > 0, + "expected rows below an Expand to bypass partial aggregation for high-cardinality " + + "input") + } + } + } + } + + test("no pass-through when the feature is disabled") { + // The metric must stay zero across the whole matrix when the switch is off, even for input that + // would otherwise bypass. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString) ++ fixedPlanConfs): _*) { + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") { + assert(numBypassingRows(df) == 0, + "no rows should bypass partial aggregation when the feature is disabled") + } + } + } + } + + test("no pass-through for a global aggregation with no grouping keys") { + // Global aggregation is ineligible (`groupingExpressions` is empty): there is a single group, + // so there is nothing to stream through. The partial aggregate must never bypass regardless of + // codegen or map settings, even under a forced fallback. + forEachCodegenAndMap(regularFallback = 16) { clue => + val df = () => spark.range(0, 200, 1, 1) + .agg(sum($"id") as "s", count(lit(1)) as "c") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "a global aggregation is not eligible and must never bypass") + } + } + } + + test("no pass-through for a session_window grouping key") { + // A batch `session_window` grouping is not streaming, but its partial aggregate feeds a + // `MergingSessionsExec` that merges overlapping sessions, and passing single-row buffers into + // it is untested. It is gated out in `adaptivePartialAggEnabled`. This input has fully + // distinct sessions (ratio 1.0) and a small `minRows`, so without the gate the periodic check + // would bypass and the `expectBypass = false` assertion would fail. + checkAdaptiveMatchesReference( + expectBypass = false, + build = { parts => + spark.range(0, 40, 1, parts) + .select($"id" as "v", timestamp_seconds($"id" * 60) as "time") + .groupBy(session_window($"time", "10 seconds")) + .agg(sum($"v") as "s") + }) + } + + test("rows absorbed by the fast map count toward the compaction ratio") { + // The compaction ratio is measured at the operator level, so the rows the fast map absorbs + // must count in the numerator just as its keys count in the denominator. This input is + // dominated by a hot-key prefix that the fast map serves without ever reaching the regular + // map, followed by a short distinct tail that does reach it. Counting only the regular map's + // traffic would see the tail alone -- a ratio near 1 -- and bypass an aggregation that is in + // fact collapsing rows heavily. + forEachCodegenAndMap() { clue => + // With the fast map on it holds 4 keys, so `k < 4` is served there and `k >= 4` falls + // through. 400 hot-key rows against 4 hot keys plus 20 tail keys is a ratio of about 17.5, + // far above the default, so nothing may bypass. + val df = () => spark.range(0, 420, 1, 1) + .select(when($"id" < 400, $"id" % 4).otherwise($"id" - 396) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "an aggregation collapsing ~17 rows per key must not bypass; fast-map hits are " + + "missing from the numerator if it does") + } + } + } + + test("a distinct-heavy prefix commits the task even when later rows collapse") { + // The flip is one-way, so the decision is made from the rows seen so far, not the partition as + // a whole. Here the first `minRows` rows are all distinct, so the periodic check bypasses and + // the pass-through stays on permanently: the remaining rows -- which all collapse onto a single + // key and would have aggregated well -- stream through as single-row partial buffers. The + // overall compaction ratio (40 rows / 8 keys = 5.0) would keep aggregation, but the prefix has + // already committed the task. The mirror case (a compacting prefix with a distinct tail) is + // covered by the test above; this pins the other side of the same asymmetry. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 40, 1, 1) + .select(when($"id" < 8, $"id").otherwise(lit(0L)) as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped == 32, + s"expected the 32 rows after the flip to bypass, got ${c.skipped}") + assert(c.partialOutputRows == 40, + s"expected every input row to stream through the partial, got ${c.partialOutputRows}") + } + } + } + + test("the periodic check fires when the ratio shows no reduction, without spilling") { + // No forced regular-map spill: only the periodic check can trigger the bypass. Fully + // distinct keys give a compaction ratio of 1.0, below `minCompaction`, so rows bypass and + // the regular map never spills. + forEachCodegenAndMap() { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, + "the periodic check should bypass fully distinct input") + assert(c.spillBytes == 0, "the periodic check must decide before any spill happens") + assert(c.tasksFallBacked == 0, + "the periodic check must not fall back to sort-based aggregation") + } + } + } + + test("the spill check fires when the map would spill on high-cardinality input") { + // Force the regular map to fall back quickly. High-cardinality input that reaches the fallback + // point should bypass via the spill check rather than spilling. Use a `minRows` larger than the + // input so the periodic check cannot fire first and the spill check is the one exercised. + forEachCodegenAndMap(minRows = 100000, regularFallback = 16) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, + "the spill check should bypass high-cardinality input at the spill boundary") + // The whole point of the spill check is to bypass *instead of* falling back to sort-based + // aggregation, so the sorter is never created. `numTasksFallBacked` is the reliable + // per-operator signal for that (the `spillSize` metric on the interpreted path is derived + // from the task-cumulative memory-spill counter and can be inflated by unrelated spilling + // such as the downstream shuffle write, so it is not asserted here). + assert(c.tasksFallBacked == 0, + "the spill check must replace the sort fallback, not trigger it") + } + } + } + + test("a new in-memory map epoch after a spill can still bypass") { + // A spill starts a new in-memory map epoch and restarts the row counters, so an input whose + // cardinality only turns unfavorable after an early spill is still caught. The first 100 rows + // repeat 5 keys and keep the aggregation effective while the forced fallback makes the map + // spill; the remaining 300 rows are fully distinct, so the new epoch is judged ineffective and + // the rest of the input is passed through. Both the spill and the bypass must be observable. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + val fastCap = if (twoLevelMap) 4 else 1 + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "8", + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, 40") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 400, 1, 1) + .select(when($"id" < 100, $"id" % 5).otherwise($"id") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") { + val c = runAndReadCounters(df) + assert(c.tasksFallBacked > 0, + "the low-cardinality prefix should still spill and fall back to sort") + assert(c.skipped > 0, + "the high-cardinality tail after the spill should be passed through") + } + } + } + } + + test("without the feature the same input really does fall back to sort") { + // Sanity check for the spill check assertion above: with adaptive disabled, the identical + // high-cardinality input under the same forced fallback genuinely falls back to sort-based + // aggregation. This proves the spill check's `tasksFallBacked == 0` reflects the bypass and + // not merely an input that never reached the spill boundary. + for { + wholeStage <- Seq(true, false) + twoLevelMap <- Seq(true, false) + } { + val fastCap = if (twoLevelMap) 4 else 1 + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> twoLevelMap.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> s"$fastCap, 16") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"wholeStage=$wholeStage twoLevelMap=$twoLevelMap") { + val c = runAndReadCounters(df) + assert(c.skipped == 0, "feature disabled: nothing should bypass") + assert(c.tasksFallBacked > 0, + "feature disabled: the forced fallback should trigger sort-based aggregation") + } + } + } + } + + test("the spill check decides identically at the exact ratio boundary with codegen on and off") { + // The check before a spill evaluates the ratio over the rows already aggregated, excluding the + // failed insertion that becomes the first pass-through row, so both execution paths must judge + // the same row set and reach the same decision. `id % 40` over 400 rows gives 40 keys when the + // map fills at 50 aggregated rows, i.e. a compaction ratio of exactly 1.25: demanding 1.25 the + // aggregation is kept (`50 < 40 * 1.25` is false), demanding 1.3 it is bypassed. + Seq(1.25 -> false, 1.3 -> true).foreach { case (minCompaction, shouldBypass) => + val skippedPerCodegen = Seq(true, false).map { wholeStage => + withSQLConf( + (Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> "true", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ENABLE_TWOLEVEL_AGG_MAP.key -> "false", + // A `minRows` larger than the input keeps the periodic check out of the picture. + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "100000", + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key -> minCompaction.toString, + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 50") ++ + fixedPlanConfs): _*) { + val df = () => spark.range(0, 400, 1, 1) + .select(($"id" % 40).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(s"minCompaction=$minCompaction wholeStage=$wholeStage") { + numBypassingRows(df) + } + } + } + withClue(s"minCompaction=$minCompaction skipped=$skippedPerCodegen") { + assert(skippedPerCodegen.forall(_ > 0) == shouldBypass, + s"expected bypass=$shouldBypass at the ratio boundary") + assert(skippedPerCodegen.map(_ > 0).distinct.length == 1, + "codegen and interpreted paths must reach the same decision at the boundary") + } + } + } + + test("a very high minCompaction always bypasses once minRows has been processed") { + // Demanding an unreachable compaction ratio is the most aggressive setting: even a + // low-cardinality input that the default threshold keeps aggregating is bypassed at the first + // check point. The results must still match the feature-off reference. + forEachCodegenAndMap(minCompaction = 1000000.0) { clue => + val df = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) > 0, + "an unreachable compaction ratio must bypass even a low-cardinality input") + } + } + } + + test("minRows = 0 disables the periodic check but keeps the spill check") { + // `minRows = 0` is a sentinel for "never evaluate periodically". Without a forced regular-map + // spill there is no check point at all, so fully distinct input -- which the default settings + // bypass immediately -- must be aggregated all the way through. This is what proves the + // periodic check is genuinely off rather than merely deferred. + forEachCodegenAndMap(minRows = 0) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "minRows = 0 must disable the periodic check, so nothing may bypass without a spill") + } + } + } + + test("minRows = 0 still bypasses at the spill boundary") { + // The other half of the sentinel: with the periodic check off, the spill check alone still + // bypasses instead of paying the spill I/O, which is the spill-only operating mode. + forEachCodegenAndMap(minRows = 0, regularFallback = 16) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(df) + assert(c.skipped > 0, + "the spill check must still bypass when the periodic check is disabled") + assert(c.tasksFallBacked == 0, + "the spill check must replace the sort fallback, not trigger it") + } + } + } + + test("minCompaction rejects values that cannot be generated as a Java literal") { + // The threshold is interpolated straight into the generated source, where a non-finite double + // renders as `InfinityD` -- not a valid Java literal. Reject it at configuration time rather + // than failing to compile the stage (or silently deoptimizing when codegen falls back). + Seq("Infinity", "-Infinity", "NaN", "1e309").foreach { v => + val e = intercept[IllegalArgumentException] { + spark.conf.set(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key, v) + } + assert(e.getMessage.contains("finite"), s"expected a finiteness error for '$v': $e") + } + spark.conf.unset(SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_COMPACTION.key) + } + + test("a larger minRows defers the decision so a small high-cardinality input is not bypassed") { + // With `minRows` larger than the whole input and no regular-map spill forced, the periodic + // check point is never reached, so nothing bypasses even though the keys are fully distinct. + forEachCodegenAndMap(minRows = 100000) { clue => + val df = () => spark.range(0, 200, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + assert(numBypassingRows(df) == 0, + "no bypass expected before `minRows` rows have been processed") + } + } + } + + test("partial aggregate output row count reflects the bypass (independent of the skip metric)") { + // `numOutputRows` on the partial aggregate is the ground truth: it is driven by the normal + // aggregation output path, not by our self-reported `numBypassingRows` metric. This test cross + // checks the two and pins the observable data-side effect of bypassing. + val numRows = 200 + forEachCodegenAndMap() { clue => + // Fully distinct keys: once the bypass fires the operator stops collapsing rows, so the + // partial aggregate emits far more than the handful of keys a real aggregation would. Read + // all counters from a single execution so the metrics are not double-counted. + val highCard = () => spark.range(0, numRows, 1, 1) + .select($"id".cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(highCard) + assert(c.skipped > 0, "high-cardinality input should bypass") + // Every partial output row is either a real (aggregated) group or a bypassed row, so the + // partial output count must be at least the number of bypassed rows, and it climbs toward + // the input row count -- well above the heavy reduction a kept aggregation would give. + assert(c.partialOutputRows >= c.skipped, + s"partial output ${c.partialOutputRows} should be >= bypassed rows ${c.skipped}") + assert(c.partialOutputRows > numRows / 2, + s"partial output (${c.partialOutputRows}) should climb toward the input row count") + } + + // Low-cardinality reference: partial aggregation stays effective, so its output equals the + // small number of distinct keys and nothing is bypassed. + val lowCard = () => spark.range(0, 600, 1, 1) + .select(($"id" % 5).cast("string") as "k", $"id" as "v") + .groupBy($"k") + .agg(sum($"v") as "s") + withClue(clue) { + val c = runAndReadCounters(lowCard) + assert(c.skipped == 0, "low-cardinality input should not bypass") + assert(c.partialOutputRows == 5, + "an effective partial aggregate should emit exactly the distinct key count") + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala new file mode 100644 index 0000000000000..2da4b5e0d66f7 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/aggregate/HashAggregateExecSuite.scala @@ -0,0 +1,137 @@ +/* + * 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.execution.aggregate + +import java.util.Properties + +import scala.util.Random + +import org.mockito.Mockito._ + +import org.apache.spark.{SparkConf, SparkFunSuite, TaskContext, TaskContextImpl} +import org.apache.spark.internal.config.MEMORY_OFFHEAP_ENABLED +import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.execution.{UnsafeFixedWidthAggregationMap, UnsafeKVExternalSorter} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} +import org.apache.spark.unsafe.types.UTF8String + +/** + * Test suite for the static helpers of [[HashAggregateExec]] that its generated code calls into. + */ +class HashAggregateExecSuite extends SparkFunSuite with SharedSparkSession { + + private val groupKeySchema = StructType(StructField("product", StringType) :: Nil) + private val aggBufferSchema = StructType(StructField("salePrice", IntegerType) :: Nil) + private val PAGE_SIZE_BYTES: Long = 1L << 26 // 64 megabytes + + /** + * Runs `f` against a fresh aggregation map, with a task context in place because both the map and + * the sorter it is destructed into allocate through the task memory manager. Asserts on the way + * out that the spilled sorters left nothing behind. + */ + private def withAggregationMap(f: UnsafeFixedWidthAggregationMap => Unit): Unit = { + val conf = new SparkConf().set(MEMORY_OFFHEAP_ENABLED.key, "false") + val taskMemoryManager = new TaskMemoryManager(new TestMemoryManager(conf), 0) + // The map registers a completion listener, which a mock swallows: this test drives the spills + // directly rather than running a task to completion. The sorter, in turn, reads the task memory + // manager off the thread-local context, so a real one has to be in place as well. + val taskContext = mock(classOf[TaskContext]) + when(taskContext.taskMemoryManager()).thenReturn(taskMemoryManager) + TaskContext.setTaskContext(new TaskContextImpl( + stageId = 0, + stageAttemptNumber = 0, + partitionId = 0, + numPartitions = 1, + taskAttemptId = Random.nextInt(10000), + attemptNumber = 0, + taskMemoryManager = taskMemoryManager, + localProperties = new Properties, + metricsSystem = null)) + + val map = new UnsafeFixedWidthAggregationMap( + InternalRow(0), // empty aggregation buffer + aggBufferSchema, + groupKeySchema, + taskContext, + 128, // initial capacity + PAGE_SIZE_BYTES) + + try { + f(map) + } finally { + map.free() + TaskContext.unset() + } + assert(taskMemoryManager.cleanUpAllAllocatedMemory() === 0) + } + + /** Inserts `keys` into the map, using each key's length as its aggregation buffer value. */ + private def insert(map: UnsafeFixedWidthAggregationMap, keys: Seq[String]): Unit = { + keys.foreach { key => + val buffer = map.getAggregationBuffer(InternalRow(UTF8String.fromString(key))) + assert(buffer != null) + buffer.setInt(0, key.length) + } + } + + /** + * Drains the sorter and returns the keys it held, checking along the way that each key still + * carries its own aggregation buffer -- draining to exhaustion also releases the sorter's memory. + */ + private def drainKeys(sorter: UnsafeKVExternalSorter): Seq[String] = { + val keys = Seq.newBuilder[String] + val iter = sorter.sortedIterator() + while (iter.next()) { + assert(iter.getKey.getString(0).length === iter.getValue.getInt(0)) + keys += iter.getKey.getString(0) + } + keys.result() + } + + test("spillHashMapToSorter destructs the map into a new sorter on the first spill") { + withAggregationMap { map => + val keys = Seq("apple", "banana", "cherry") + insert(map, keys) + + val sorter = HashAggregateExec.spillHashMapToSorter(map, null) + + assert(sorter != null) + assert(drainKeys(sorter) === keys.sorted) + } + } + + test("spillHashMapToSorter merges into the existing sorter and returns it") { + withAggregationMap { map => + val firstKeys = Seq("apple", "banana", "cherry") + val secondKeys = Seq("damson", "elderberry", "fig") + insert(map, firstKeys) + val sorter = HashAggregateExec.spillHashMapToSorter(map, null) + + // The map keeps accepting keys after being destructed, so the second spill hits the merge + // branch with a sorter already in hand. + insert(map, secondKeys) + val merged = HashAggregateExec.spillHashMapToSorter(map, sorter) + + // The same sorter comes back, now holding the keys of both spills. + assert(merged eq sorter) + assert(drainKeys(merged) === (firstKeys ++ secondKeys).sorted) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala new file mode 100644 index 0000000000000..0482128b5bf40 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/AdaptivePartialAggregationBenchmark.scala @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.internal.SQLConf + +/** + * Benchmark comparing runtime adaptive partial aggregation (see + * [[SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED]]) against the static pre-shuffle partial + * aggregation. When the partial aggregation is not reducing rows, the operator streams the + * remaining rows through as single-row partial buffers instead of maintaining (and possibly + * spilling) a large aggregation map. + * + * Each scenario runs the query across the full matrix of whole-stage codegen on/off and the + * feature disabled (`adaptive = F`, the pre-change baseline) vs enabled (`adaptive = T`), over a + * {high, low}-cardinality x {no-spill, on-spill} grid: + * - high-cardinality, no spill: the periodic check bypasses, which should win. + * - low-cardinality, no spill: nothing bypasses, and the per-row guard the feature still adds + * shows up as a small overhead to quantify. + * - high-cardinality, forced regular-map spill: the spill check bypasses instead of spilling, + * which should win. + * - low-cardinality, forced regular-map spill: the compaction ratio is far above the threshold, + * so the spill check keeps aggregating and both runs spill identically (no regression). + * + * To run this benchmark: + * {{{ + * 1. build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.AdaptivePartialAggregationBenchmark" + * 2. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.AdaptivePartialAggregationBenchmark" + * Results will be written to "benchmarks/AdaptivePartialAggregationBenchmark-results.txt". + * }}} + */ +object AdaptivePartialAggregationBenchmark extends SqlBasedBenchmark { + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + // The upstream `CombineAdjacentAggregation` and `ReplaceHashWithSortAgg` rules would collapse + // or convert these single-partition hash aggregates, so both are disabled to keep the + // Partial+Final `HashAggregateExec` structure the adaptive feature governs. + val fixedPlanConfs = Seq( + SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "false", + SQLConf.REPLACE_HASH_WITH_SORT_AGG_ENABLED.key -> "false") + + // Adds the (whole-stage codegen, adaptive switch) matrix for `query`. `extraConf` is applied + // to all four cases so the only differences are the two axes. + def addCodegenAdaptiveCases( + benchmark: Benchmark, + query: () => DataFrame, + extraConf: Seq[(String, String)] = Nil): Unit = { + for { + wholeStage <- Seq(true, false) + adaptive <- Seq(false, true) + } { + val adaptiveLabel = if (adaptive) "T" else "F" + val label = s"codegen = $wholeStage, adaptive = $adaptiveLabel" + benchmark.addCase(label) { _ => + withSQLConf( + (Seq( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> wholeStage.toString, + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_ENABLED.key -> adaptive.toString) ++ + fixedPlanConfs ++ extraConf): _*) { + query().noop() + } + } + } + } + + // Fully distinct keys make partial aggregation useless, so the periodic check bypasses: the + // feature should be faster than the baseline that maintains a map entry per row. + runBenchmark("high-cardinality input, pass-through at the periodic check") { + val N = 8L << 20 + val benchmark = new Benchmark("adaptive partial agg, high card, no spill", N, + output = output) + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N)) + benchmark.run() + } + + // 1000 distinct keys over a large input: partial aggregation reduces a lot, so the periodic + // check never activates pass-through. The two runs still differ: the adaptive path pays a + // small per-row overhead (the stop check, the bypass counter and the check-point compare) + // even when nothing bypasses. The measured difference is a few percent, not a functional + // regression. + runBenchmark("low-cardinality input, pass-through at the periodic check") { + val N = 16L << 20 + val benchmark = new Benchmark("adaptive partial agg, low card, no spill", N, + output = output) + addCodegenAdaptiveCases(benchmark, () => + spark.range(N).selectExpr("id % 1000 as k", "id as v").groupBy("k").agg("v" -> "sum")) + benchmark.run() + } + + // Force the regular map to spill quickly and disable the periodic check (huge minRows). With + // fully distinct keys the compaction ratio is 1.0, so at the spill boundary the spill check + // bypasses instead of spilling; the baseline spills repeatedly and falls back to sort-based + // aggregation. + runBenchmark("high-cardinality input, pass-through at the spill check") { + val N = 8L << 20 + val benchmark = new Benchmark("adaptive partial agg, high card, spill", N, output = output) + val spillCheckConf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "0", + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") + addCodegenAdaptiveCases(benchmark, () => distinctKeyedDf(N), extraConf = spillCheckConf) + benchmark.run() + } + + // Force the regular map to spill quickly on low-cardinality input. With only 1000 distinct + // keys the compaction ratio is far above the threshold, so even at the spill boundary the + // spill check correctly keeps aggregating: both runs spill and fall back to sort-based + // aggregation identically (no regression). + runBenchmark("low-cardinality input, pass-through at the spill check") { + val N = 16L << 20 + val benchmark = new Benchmark("adaptive partial agg, low card, spill", N, output = output) + val spillCheckConf = Seq( + SQLConf.ADAPTIVE_PARTIAL_AGGREGATION_MIN_ROWS.key -> "0", + "spark.sql.TungstenAggregate.testFallbackStartsAt" -> "1, 1048576") + addCodegenAdaptiveCases(benchmark, () => + spark.range(N).selectExpr("id % 1000 as k", "id as v").groupBy("k").agg("v" -> "sum"), + extraConf = spillCheckConf) + benchmark.run() + } + } + + private def distinctKeyedDf(N: Long): DataFrame = + spark.range(N).selectExpr("id as k", "id as v").groupBy("k").agg("v" -> "sum") +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala index afb732a00d534..7df7e97ecb255 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ArrowCacheBenchmark.scala @@ -793,6 +793,61 @@ object ArrowCacheBenchmark extends SqlBasedBenchmark { } } + private def columnPruningWideTable(): Unit = { + // 1M rows x 50 long columns is a ~400MB uncompressed cache -- wide enough that pruning 49 of + // 50 columns matters, but small enough to sit comfortably in the benchmark heap so the timing + // reflects the read, not GC pressure from a cache that nearly fills the heap. + val numRows = 1000000 + val numCols = 50 + val cols = (0 until numCols).map(i => s"id + $i as col$i") + + // Measure only the read: summing one column while pruning the other 49. Each case builds and + // fully materializes its cache before starting the timer (via addTimerCase), so the numbers + // reflect the cached-scan read path -- which column pruning speeds up -- not the one-time cache + // materialization. An aggregate is used rather than a projection to a noop sink because the + // latter does not pull the selected column's data through the read path. The cache serializer + // is a JVM-wide static, so each case creates its own fresh session (excluded from timing), + // which also lets the Default and Arrow serializers share one comparison table. + val default = "org.apache.spark.sql.execution.columnar.DefaultCachedBatchSerializer" + val arrow = "org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer" + runBenchmark(s"Cache with column pruning (sum 1 of $numCols columns)") { + val benchmark = new Benchmark( + s"Sum 1 of $numCols columns, $numRows rows", numRows, output = output) + + def addPruningCase(name: String, serializer: String)( + configure: SparkSession => Unit): Unit = { + benchmark.addTimerCase(name) { timer => + val spark = createFreshSession(serializer) + try { + configure(spark) + val df = spark.range(numRows).selectExpr(cols: _*) + df.cache() + df.count() // materialize the cache before timing + timer.startTiming() + df.selectExpr("sum(col0)").collect() + timer.stopTiming() + df.unpersist(blocking = true) + } finally { + spark.stop() + } + } + } + + addPruningCase("Default cache", default)(_ => ()) + addPruningCase("Default cache (uncompressed)", default) { spark => + spark.conf.set("spark.sql.inMemoryColumnarStorage.compressed", "false") + } + addPruningCase("Arrow cache", arrow)(_ => ()) + Seq("-1", "1", "3").foreach { level => + addPruningCase(s"Arrow cache (zstd level $level)", arrow) { spark => + spark.conf.set(SQLConf.ARROW_EXECUTION_COMPRESSION_CODEC.key, "zstd") + spark.conf.set(SQLConf.ARROW_EXECUTION_ZSTD_COMPRESSION_LEVEL.key, level) + } + } + benchmark.run() + } + } + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { runBenchmark("Arrow Cache vs Default Cache") { cachePrimitiveTypes() @@ -800,6 +855,7 @@ object ArrowCacheBenchmark extends SqlBasedBenchmark { cacheColumnarInput() recacheArrowData() columnPruning() + columnPruningWideTable() } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala index 27e67e6b36344..42644708ab5bd 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/ExpandBenchmark.scala @@ -33,6 +33,23 @@ import org.apache.spark.sql.internal.SQLConf * OptimizeExpand rule that pre-aggregates data before the Expand. Controlled by * spark.sql.optimizer.optimizeExpandRatio (default -1 = disabled). * + * It also measures two optimizations that target a shared expensive + * subexpression landing in the branches of the Expand, using a traffic/BI + * "N-day active users" rollup (conditional COUNT(DISTINCT) and conditional SUM + * aggregates whose conditions all share the same datetime subexpression): + * - rewriteCountDistinctConditional collapses N conditional COUNT(DISTINCT) + * on the same base column into one distinct group, shrinking the Expand + * fan-out from Nx to 1x (cuts data amplification, not the per-row + * subexpression evaluation cost); + * - subexpressionElimination (in Expand, via whole-stage codegen) evaluates + * the shared subexpression once per input row instead of once per Expand + * branch occurrence. + * The four cases form a 2x2 matrix over the two optimizations (base, +rewrite, + * +rewrite+CSE, and +CSE alone), so the benefit of each optimization can be + * attributed independently. Controlled by + * spark.sql.optimizer.rewriteCountDistinctConditional.enabled (default true) + * and spark.sql.subexpressionElimination.enabled (default true). + * * To run this benchmark: * {{{ * 1. build/sbt "sql/Test/runMain <this class>" @@ -40,6 +57,14 @@ import org.apache.spark.sql.internal.SQLConf * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain <this class>" * Results will be written to "benchmarks/ExpandBenchmark-results.txt". * }}} + * + * Optional arguments select benchmark sections (default: all): + * {{{ + * <this class> ratio # Expand: varying number of COUNT(DISTINCT) + * <this class> char # Expand: varying data characteristics (pure distinct) + * <this class> subexpr # Expand: subexpression elimination across branches + * <this class> smoke # all sections with tiny data and a single iteration + * }}} */ object ExpandBenchmark extends SqlBasedBenchmark { @@ -74,7 +99,7 @@ object ExpandBenchmark extends SqlBasedBenchmark { private def expandRatioBenchmark( title: String, N: Long, numDistinct: Int, - table: String): Unit = { + table: String, numIters: Int = 5): Unit = { val benchmark = new Benchmark(title, N, output = output) val sqlWithSum = countDistinctQuery( @@ -89,28 +114,28 @@ object ExpandBenchmark extends SqlBasedBenchmark { benchmark.addCase( s"with sum - baseline (ratio $ratioWithSum)", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sqlWithSum).noop() } } benchmark.addCase( s"with sum - optimized (ratio $ratioWithSum)", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sqlWithSum).noop() } } benchmark.addCase( s"pure distinct - baseline (ratio $ratioPure)", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sqlPure).noop() } } benchmark.addCase( s"pure distinct - optimized (ratio $ratioPure)", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sqlPure).noop() } @@ -119,24 +144,136 @@ object ExpandBenchmark extends SqlBasedBenchmark { benchmark.run() } + /** + * Prepares a traffic-like fact table: one row per page view, with a user id, + * a date-string partition column `pt` (format `yyyyMMdd`), and a metric value. + */ + private def prepareTrafficTable(name: String, N: Long, userMod: Int): Unit = { + spark.range(N) + .selectExpr( + s"cast(id % $userMod as bigint) as user_id", + "date_format(date_add(date '2023-06-01', cast(id % 730 as int)), 'yyyyMMdd') as pt", + "cast(id % 3600 as bigint) as value") + .createOrReplaceTempView(name) + } + + /** + * Builds a traffic/BI "N-day active users" rollup query: for each of the + * given day windows, one conditional COUNT(DISTINCT) and one conditional SUM + * whose conditions all share the same expensive datetime subexpression + * (parse `pt`, format it back, then `datediff` against a fixed date). + * `RewriteDistinctAggregates` places those conditions in the branches of an + * Expand, so without subexpression elimination the shared subexpression is + * codegen'd once per occurrence: 9 in the distinct branch(es) and 9 in the + * SUM branch, i.e. 18 evaluations per input row. + */ + private def conditionalAggsQuery(table: String, windows: Seq[Int]): String = { + val daysSince = "datediff(" + + "from_unixtime(unix_timestamp('20250601', 'yyyyMMdd')), " + + "from_unixtime(unix_timestamp(pt, 'yyyyMMdd')))" + val metrics = windows.flatMap { n => + Seq( + s"count(distinct if($daysSince <= $n, pt, null)) as count_${n}d", + s"sum(if($daysSince <= $n, value, null)) as sum_${n}d") + } + s"""SELECT user_id, + | ${metrics.mkString(",\n ")} + |FROM $table + |GROUP BY user_id""".stripMargin + } + + /** + * Measures a traffic/BI "N-day active users" rollup (see [[conditionalAggsQuery]]) + * under all four combinations of the two optimizations that affect the shared + * `daysSince` subexpression in the Expand branches: + * - base: rewriteCountDistinctConditional off, subexpression elimination off + * -> 10 Expand branches (one per conditional COUNT(DISTINCT) plus one for + * the SUMs), daysSince evaluated 18 times per input row; + * - +rewrite: rewriteCountDistinctConditional on, CSE off -> the 9 + * conditional COUNT(DISTINCT) expressions collapse into one distinct group, so the + * Expand has 2 branches (data amplification drops from 10x to 2x), but + * daysSince is still evaluated 18 times per input row; + * - +rewrite+CSE: both on -> the Expand runs the standard whole-stage + * subexpression elimination, evaluating daysSince once per input row + * before the branch loop; + * - +CSE alone: rewriteCountDistinctConditional off, subexpression + * elimination on -> the same 10-branch Expand as the base case, but + * daysSince is evaluated once per input row, isolating the CSE benefit + * without the data-amplification reduction. + */ + private def subExprEliminationBenchmark( + title: String, N: Long, table: String, numIters: Int = 3): Unit = { + val benchmark = new Benchmark(title, N, output = output) + val sql = conditionalAggsQuery(table, Seq(1, 7, 14, 30, 60, 90, 180, 365, 730)) + + benchmark.addCase( + "rewrite off, CSE off (10x amplify, 18 evals/row)", + numIters = numIters) { _ => + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false", + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + spark.sql(sql).noop() + } + } + benchmark.addCase( + "rewrite on, CSE off (2x amplify, 18 evals/row)", + numIters = numIters) { _ => + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "false", + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + spark.sql(sql).noop() + } + } + benchmark.addCase( + "rewrite on, CSE on (2x amplify, 1 eval/row)", + numIters = numIters) { _ => + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + spark.sql(sql).noop() + } + } + benchmark.addCase( + "rewrite off, CSE on (10x amplify, 1 eval/row)", + numIters = numIters) { _ => + withSQLConf( + SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true", + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + spark.sql(sql).noop() + } + } + + benchmark.run() + } + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { - val N = 10L << 20 // ~10M rows + val smoke = mainArgs.contains("smoke") + val sectionFilter = mainArgs.filterNot(_ == "smoke").toSet - runBenchmark("Expand: varying number of COUNT(DISTINCT)") { + def runSection(key: String, name: String)(body: => Unit): Unit = { + if (sectionFilter.isEmpty || sectionFilter.contains(key)) { + runBenchmark(name)(body) + } + } + + val numIters = if (smoke) 1 else 5 + val N = if (smoke) 1000L else 10L << 20 // ~10M rows + + runSection("ratio", "Expand: varying number of COUNT(DISTINCT)") { prepareTable("expand_bench", N, keyMod = 1000, colMods = Seq(100, 200, 300, 400, 500, 600, 700, 800)) expandRatioBenchmark( - "2 distinct aggregates", N, 2, "expand_bench") + "2 distinct aggregates", N, 2, "expand_bench", numIters) expandRatioBenchmark( - "4 distinct aggregates", N, 4, "expand_bench") + "4 distinct aggregates", N, 4, "expand_bench", numIters) expandRatioBenchmark( - "6 distinct aggregates", N, 6, "expand_bench") + "6 distinct aggregates", N, 6, "expand_bench", numIters) expandRatioBenchmark( - "8 distinct aggregates", N, 8, "expand_bench") + "8 distinct aggregates", N, 8, "expand_bench", numIters) } - runBenchmark("Expand: varying data characteristics (pure distinct)") { + runSection("char", "Expand: varying data characteristics (pure distinct)") { // Default: 1K groups, moderate distinct cardinality prepareTable("expand_default", N, keyMod = 1000, colMods = Seq(100, 200, 300, 400, 500, 600)) @@ -162,49 +299,49 @@ object ExpandBenchmark extends SqlBasedBenchmark { "6 pure distinct aggs with varying data", N, output = output) benchDataChar.addCase("1K groups, moderate card - baseline", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sql6d).noop() } } benchDataChar.addCase("1K groups, moderate card - optimized", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sql6d).noop() } } benchDataChar.addCase("100K groups, moderate card - baseline", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sql6dHighKey).noop() } } benchDataChar.addCase("100K groups, moderate card - optimized", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sql6dHighKey).noop() } } benchDataChar.addCase("1K groups, low card (5 vals) - baseline", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sql6dLowCard).noop() } } benchDataChar.addCase("1K groups, low card (5 vals) - optimized", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sql6dLowCard).noop() } } benchDataChar.addCase("no grouping key - baseline", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "-1") { spark.sql(sql6dGlobal).noop() } } benchDataChar.addCase("no grouping key - optimized", - numIters = 5) { _ => + numIters = numIters) { _ => withSQLConf(OPTIMIZE_EXPAND_RATIO -> "2") { spark.sql(sql6dGlobal).noop() } @@ -212,5 +349,13 @@ object ExpandBenchmark extends SqlBasedBenchmark { benchDataChar.run() } + + runSection("subexpr", "Expand: subexpression elimination across branches") { + val numRows = if (smoke) 1000L else 5L << 20 // ~5M rows + prepareTrafficTable("expand_traffic", numRows, userMod = 1000000) + subExprEliminationBenchmark( + "9 conditional COUNT(DISTINCT) + 9 conditional SUM sharing one subexpression", + numRows, "expand_traffic", numIters = if (smoke) 1 else 3) + } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala new file mode 100644 index 0000000000000..23b8e508d0686 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/RuntimeBloomFilterCachedInputBenchmark.scala @@ -0,0 +1,139 @@ +/* + * 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.execution.benchmark + +import scala.collection.mutable +import scala.concurrent.duration._ + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.BloomFilterMightContain +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +/** + * Measures runtime Bloom filtering when a materialized cache hides a selective predicate. + * + * To run this benchmark: + * {{{ + * build/sbt "sql/Test/runMain + * org.apache.spark.sql.execution.benchmark.RuntimeBloomFilterCachedInputBenchmark" + * }}} + * + * The additional shuffle metric reports only the wide fact-side exchange, independently of + * exchanges used to aggregate results or construct the Bloom filter. + */ +object RuntimeBloomFilterCachedInputBenchmark extends SqlBasedBenchmark { + + private val factRows = 500000L + private val filteringStride = 100L + private val partitions = 4 + + override def getSparkSession: SparkSession = { + SparkSession.builder() + .master("local[4]") + .appName(this.getClass.getCanonicalName) + .config(SQLConf.SHUFFLE_PARTITIONS.key, 4) + .config("spark.ui.enabled", false) + .getOrCreate() + } + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.SHUFFLE_PARTITIONS.key -> partitions.toString, + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "0", + SQLConf.RUNTIME_BLOOM_FILTER_CREATION_SIDE_THRESHOLD.key -> "10MB", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false", + SQLConf.IGNORE_MISSING_FILES.key -> "false") { + val fact = spark.range(0, factRows, 1, partitions).selectExpr( + "id AS fact_key", + "sha2(cast(id AS STRING), 256) AS payload") + val dimension = spark.range(0, factRows, 1, partitions) + .filter(s"id % $filteringStride = 0") + .selectExpr("id AS dimension_key") + .persist(StorageLevel.MEMORY_AND_DISK) + + try { + assert(dimension.count() == factRows / filteringStride) + + runBenchmark("Runtime Bloom filter from a materialized selective cache") { + val factShuffleBytes = mutable.Map.empty[Boolean, Long] + val benchmark = new Benchmark( + "Cached input runtime Bloom filter", + factRows, + minNumIters = 2, + warmupTime = 1.second, + minTime = Duration.Zero, + output = output) + + Seq(false, true).foreach { enabled => + val name = if (enabled) { + "Runtime Bloom filter enabled" + } else { + "Runtime Bloom filter disabled" + } + benchmark.addCase(name, numIters = 2) { _ => + withSQLConf(SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> enabled.toString) { + val result = fact.join(dimension, fact("fact_key") === dimension("dimension_key")) + .selectExpr("sum(length(payload)) AS payload_bytes") + val observed = result.collect().head.getLong(0) + assert(observed == factRows / filteringStride * 64) + + val optimizedPlan = result.queryExecution.optimizedPlan + val hasRuntimeBloomFilter = optimizedPlan.exists { node => + node.expressions.exists(_.exists(_.isInstanceOf[BloomFilterMightContain])) + } + assert(hasRuntimeBloomFilter == enabled, + s"Expected runtime Bloom filter enabled=$enabled:\n$optimizedPlan") + + val factExchanges = result.queryExecution.executedPlan.collect { + case exchange: ShuffleExchangeExec + if exchange.output.exists(_.name == "fact_key") => exchange + } + assert(factExchanges.size == 1, + s"Expected one fact-side shuffle exchange, found ${factExchanges.size}") + factShuffleBytes(enabled) = + factExchanges.head.metrics(SQLShuffleWriteMetricsReporter.SHUFFLE_BYTES_WRITTEN) + .value + } + } + } + + benchmark.run() + + val baseline = factShuffleBytes(false) + val optimized = factShuffleBytes(true) + assert(baseline > 0 && optimized < baseline, + s"Expected the Bloom filter to reduce fact shuffle bytes: $baseline -> $optimized") + val reduction = (baseline - optimized) * 100.0 / baseline + // scalastyle:off println + benchmark.out.println(s"Fact-side shuffle bytes without runtime Bloom filter: $baseline") + benchmark.out.println(s"Fact-side shuffle bytes with runtime Bloom filter: $optimized") + benchmark.out.println(f"Fact-side shuffle reduction: $reduction%.2f%%") + // scalastyle:on println + } + } finally { + dimension.unpersist(blocking = true) + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/VariantShreddedPredicatePushdownBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/VariantShreddedPredicatePushdownBenchmark.scala new file mode 100644 index 0000000000000..388a555c7d0c1 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/VariantShreddedPredicatePushdownBenchmark.scala @@ -0,0 +1,166 @@ +/* + * 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.execution.benchmark + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, SaveMode} +import org.apache.spark.sql.internal.SQLConf + +/** + * Synthetic benchmark for row-group skipping on shredded Variant columns (SPARK-55817). + * + * The optimization pushes a predicate on a shredded Variant field to the physical typed_value leaf + * (guarded so residual fallbacks are never skipped), letting Parquet skip row groups the leaf + * min/max cannot match. The lift depends on the layout: the field must be shredded, the predicate a + * literal comparison, the data sorted on that field, and a file must hold many row groups. This + * benchmark writes such a layout (sorted on the shredded field, small block size so a single file + * has many row groups) and compares scan time with the optimization on vs off. + * + * To run this benchmark: + * {{{ + * 1. without sbt: + * bin/spark-submit --class <this class> + * --jars <spark core test jar>,<spark catalyst test jar> <sql core test jar> + * 2. build/sbt "sql/Test/runMain <this class>" + * 3. generate result: + * SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "sql/Test/runMain <this class>" + * Results will be written to + * "benchmarks/VariantShreddedPredicatePushdownBenchmark-results.txt". + * }}} + */ +object VariantShreddedPredicatePushdownBenchmark extends SqlBasedBenchmark { + + private val N = 20 * 1024 * 1024 + private val NUMBER_OF_ITER = 10 + + // A single-column shredded Variant dataset with an object field `a` sorted ascending, so that a + // literal predicate on `a` maps to a contiguous range of row groups. + private val df: DataFrame = spark + .range(0, N, 1, 1) + .selectExpr("parse_json('{\"a\":' || id || '}') AS v") + + // Same, but every row also carries a key `z` outside the shredding schema, so the whole partial + // object lands in the top-level residual `v.value` (non-null on every row). `a` is still fully + // shredded into the typed leaf. This is the normal layout for real Variant data (the inferred + // shredding schema is capped, so extra keys are common), and it is where the flat + // `or(leaf, isNotNull(residual)...)` guard could never skip -- the tighter + // `or(leaf, and(anyResidualNotNull, isNull(leaf)))` guard skips via the "leaf has no nulls" arm. + private val dfPartialObject: DataFrame = spark + .range(0, N, 1, 1) + .selectExpr("parse_json('{\"a\":' || id || ', \"z\":\"outside\"}') AS v") + + // Confs to write the Variant column shredded, forcing `a` to a bigint typed leaf. A small block + // size makes the writer emit many row groups per file. + private val writeConf = Seq( + SQLConf.VARIANT_WRITE_SHREDDING_ENABLED.key -> "true", + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true", + SQLConf.VARIANT_FORCE_SHREDDING_SCHEMA_FOR_TEST.key -> "a bigint") + + private def addCase( + benchmark: Benchmark, + inputPath: String, + enablePushdown: String, + name: String, + withFilter: DataFrame => DataFrame): Unit = { + val loadDF = spark.read.parquet(inputPath).selectExpr("variant_get(v, '$.a', 'bigint') AS a") + benchmark.addCase(name) { _ => + withSQLConf( + SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED.key -> enablePushdown, + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true") { + withFilter(loadDF).noop() + } + } + } + + // A tiny block size makes one file hold many row groups, which is the layout row-group skipping + // needs. `blockSize = None` uses the Parquet default (one row group for this dataset), used to + // show the skip-none overhead disappears when a file is not deliberately dense. + private def createAndRunBenchmark( + name: String, + withFilter: DataFrame => DataFrame, + data: DataFrame = df, + blockSize: Option[Int] = Some(128 * 1024)): Unit = { + withTempPath { tempDir => + val outputPath = tempDir.getCanonicalPath + withSQLConf(writeConf: _*) { + val writer = data.write.mode(SaveMode.Overwrite) + blockSize.foreach(bs => writer.option("parquet.block.size", bs.toString)) + writer.parquet(outputPath) + } + val benchmark = new Benchmark(name, N, NUMBER_OF_ITER, output = output) + addCase(benchmark, outputPath, enablePushdown = "false", + "Without shredded predicate pushdown", withFilter) + addCase(benchmark, outputPath, enablePushdown = "true", + "With shredded predicate pushdown", withFilter) + benchmark.run() + } + } + + /** + * Filter that matches nothing, so the leaf min/max lets Parquet skip every row group when the + * optimization is on. + */ + def runSkipAllRowGroups(): Unit = { + createAndRunBenchmark("Can skip all row groups", _.filter("a < 0")) + } + + /** + * Highly selective filter matching only the last few row groups of the sorted data. + */ + def runSkipSomeRowGroups(): Unit = { + createAndRunBenchmark("Can skip some row groups", _.filter(s"a > ${(N * 0.99).toLong}")) + } + + /** + * Filter that matches the whole range, so no row group can be skipped -- measures the overhead + * of building and evaluating the pushed predicate when it never helps. Written with the tiny + * block size (many row groups), so this is the worst case for the overhead: it is paid per row + * group. Compare with `runSkipNoRowGroupsDefaultBlockSize`. + */ + def runSkipNoRowGroups(): Unit = { + createAndRunBenchmark("Can skip no row groups", _.filter(s"a >= 0 and a <= $N")) + } + + /** + * Same skip-none filter but written at the default Parquet block size (one row group for this + * dataset), the layout a default-configured writer produces. The per-row-group overhead + * effectively vanishes here -- it scales with row-group count, the same knob as the benefit. + */ + def runSkipNoRowGroupsDefaultBlockSize(): Unit = { + createAndRunBenchmark("Can skip no row groups (default block size)", + _.filter(s"a >= 0 and a <= $N"), blockSize = None) + } + + /** + * Same selective filter as `runSkipSomeRowGroups`, but on data whose objects carry a key outside + * the shredding schema (so the top-level residual is non-null on every row). This is the layout + * where the earlier flat OR guard could never skip; the tighter guard still skips here. + */ + def runSkipSomeRowGroupsPartialObject(): Unit = { + createAndRunBenchmark("Can skip some row groups (partial object)", + _.filter(s"a > ${(N * 0.99).toLong}"), data = dfPartialObject) + } + + override def runBenchmarkSuite(mainArgs: Array[String]): Unit = { + runSkipAllRowGroups() + runSkipSomeRowGroups() + runSkipNoRowGroups() + runSkipNoRowGroupsDefaultBlockSize() + runSkipSomeRowGroupsPartialObject() + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala index 73c2b15daf2ff..b3d6c3539a10c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializerSuite.scala @@ -264,6 +264,102 @@ class ArrowCachedBatchSerializerSuite extends QueryTest with SharedSparkSession assert(projected.queryExecution.executedPlan.toString.contains("InMemoryTableScan")) } + test("column projection prunes columns on load for every projection shape") { + // The read path reads only the selected columns' buffers out of the cached bytes. Exercise a + // spread of projection shapes -- reordering, single column at each position, complex columns + // mixed with primitives, and duplicate selection -- under both the row and vectorized read + // paths, so the buffer-span arithmetic (which must skip the exact node/buffer runs of + // unselected columns, including the child buffers of complex columns) is covered end to end. + Seq(false, true).foreach { vectorized => + withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { + val df = (1 to 50).map { i => + (i, s"str$i", Seq(i, i + 1), (i.toLong, s"n$i")) + }.toDF("a", "b", "arr", "st") + df.cache() + try { + def expected(cols: Seq[String]): Seq[Row] = (1 to 50).map { i => + Row.fromSeq(cols.map { + case "a" => i + case "b" => s"str$i" + case "arr" => Seq(i, i + 1) + case "st" => Row(i.toLong, s"n$i") + }) + } + // Selecting a complex column after skipping a var-width one exercises skipping the + // multi-buffer runs (offset + data) of the unselected string. + checkAnswer(df.select("arr"), expected(Seq("arr"))) + checkAnswer(df.select("st"), expected(Seq("st"))) + checkAnswer(df.select("b"), expected(Seq("b"))) + // Reordered projection: the loaded root must be in output (columnIndices) order. + checkAnswer(df.select("st", "a"), expected(Seq("st", "a"))) + checkAnswer(df.select("arr", "b", "a"), expected(Seq("arr", "b", "a"))) + // Duplicate selection maps two output columns to one cached column. + checkAnswer(df.select("a", "a"), (1 to 50).map(i => Row(i, i))) + // Full projection (no pruning) still round-trips. + checkAnswer(df.select("a", "b", "arr", "st"), expected(Seq("a", "b", "arr", "st"))) + } finally { + df.unpersist() + } + } + } + } + + test("column projection prunes deeply nested columns on load") { + // The buffer-span arithmetic must skip (and, when selected, copy) a column's ENTIRE buffer + // subtree, at arbitrary nesting depth. Exercise array<struct>, struct<struct<struct>>, and + // map<int, array<int>>, each selected while pruning neighbours whose own subtrees have varying + // buffer counts, under both read paths, so an off-by-one in the recursive span would surface + // as wrong values in a following column. + Seq(false, true).foreach { vectorized => + withSQLConf(SQLConf.CACHE_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { + val schema = new StructType() + .add("id", IntegerType) + .add("arrOfStruct", ArrayType(new StructType().add("x", LongType).add("y", StringType))) + .add("s", StringType) + .add("deepStruct", new StructType() + .add("l1", new StructType() + .add("l2", new StructType().add("v", LongType).add("t", StringType)))) + .add("mapOfArray", MapType(IntegerType, ArrayType(IntegerType))) + val rows = (1 to 30).map { i => + Row(i, Seq(Row(i.toLong, s"a$i"), Row((i + 1).toLong, s"b$i")), s"s$i", + Row(Row(Row(i.toLong * 10, s"deep$i"))), Map(i -> Seq(i, i + 1, i + 2))) + } + val df = spark.createDataFrame( + spark.sparkContext.parallelize(rows, 1), schema).cache() + try { + def expected(cols: Seq[String]): Seq[Row] = (1 to 30).map { i => + Row.fromSeq(cols.map { + case "id" => i + case "arrOfStruct" => Seq(Row(i.toLong, s"a$i"), Row((i + 1).toLong, s"b$i")) + case "s" => s"s$i" + case "deepStruct" => Row(Row(Row(i.toLong * 10, s"deep$i"))) + case "mapOfArray" => Map(i -> Seq(i, i + 1, i + 2)) + }) + } + // Each nested column selected alone: its whole subtree must be copied, nothing else. + checkAnswer(df.select("arrOfStruct"), expected(Seq("arrOfStruct"))) + checkAnswer(df.select("deepStruct"), expected(Seq("deepStruct"))) + checkAnswer(df.select("mapOfArray"), expected(Seq("mapOfArray"))) + // A primitive after a pruned deep column: the skip must span the whole deep subtree. + checkAnswer(df.select("id", "s"), expected(Seq("id", "s"))) + // Reordered mix of nested and primitive columns, pruning others in between. + checkAnswer( + df.select("mapOfArray", "id", "deepStruct"), + expected(Seq("mapOfArray", "id", "deepStruct"))) + checkAnswer( + df.select("deepStruct", "arrOfStruct"), + expected(Seq("deepStruct", "arrOfStruct"))) + // Full projection round-trips. + checkAnswer( + df.select("id", "arrOfStruct", "s", "deepStruct", "mapOfArray"), + expected(Seq("id", "arrOfStruct", "s", "deepStruct", "mapOfArray"))) + } finally { + df.unpersist() + } + } + } + } + test("caching with multiple batches") { withSQLConf(SQLConf.ARROW_EXECUTION_MAX_RECORDS_PER_BATCH.key -> "10") { val df = (1 to 50).map(i => (i, s"str$i")).toDF("a", "b") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala index fcf7edfcf87b0..91b1b388cb8ca 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryColumnarQuerySuite.scala @@ -28,14 +28,14 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.columnar.CachedBatch -import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} +import org.apache.spark.sql.execution.{FileSourceScanExec, FilterExec, InputAdapter, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.test.SQLTestData._ import org.apache.spark.sql.types._ -import org.apache.spark.storage.StorageLevel +import org.apache.spark.storage.{RDDBlockId, StorageLevel} import org.apache.spark.storage.StorageLevel._ class TestCachedBatchSerializer( @@ -649,4 +649,210 @@ class InMemoryColumnarQuerySuite extends SharedSparkSession with AdaptiveSparkPl assert(exceptionCnt.get == 0) } + + test("SPARK-58272: only fully materialized repeatable disk-backed caches publish exact stats") { + def checkCache(level: StorageLevel, expected: Boolean): Unit = { + val cached = spark.range(0, 20, 1, numPartitions = 2) + .filter($"id" < 10) + .persist(level) + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + + assert(relation.hasSelectivePredicate) + assert(!relation.mayHaveUsableMaterializedStats) + assert(relation.materializedMetadata.isEmpty) + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + relation.cacheBuilder.cachedColumnBuffers.count() + + val metadata = relation.materializedMetadata.get + assert(metadata.rowCount == 10L) + assert(metadata.sizeInBytes == relation.computeStats().sizeInBytes) + assert(metadata.isOutputRepeatable) + assert(metadata.isDurable == expected) + assert(metadata.statsAvailable == expected) + assert(relation.mayHaveUsableMaterializedStats == expected) + assert(relation.isOutputRepeatable) + assert(relation.statsAvailable == expected) + assert(relation.computeStats().rowCount.contains(10L)) + } finally { + cached.unpersist(blocking = true) + } + } + + checkCache(MEMORY_AND_DISK, expected = true) + checkCache(MEMORY_ONLY, expected = false) + } + + test("SPARK-58272: materialized cache stats follow partition recomputation") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val property = "spark.sql.test.cache.recomputedPartitionHasRows" + System.clearProperty(property) + val liveGate = udf(() => java.lang.Boolean.getBoolean(property)).asNondeterministic() + val cached = spark.range(0, 8, 1, numPartitions = 1) + .filter(liveGate()) + .persist(MEMORY_ONLY) + + try { + assert(cached.count() == 0) + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val builder = relation.cacheBuilder + val blockId = RDDBlockId(builder.cachedColumnBuffers.id, 0) + val blockManager = spark.sparkContext.env.blockManager + + assert(blockManager.getStatus(blockId).nonEmpty) + assert(builder.loadedMaterializedStats.exists(_._1 == 0L)) + assert(relation.materializedMetadata.exists(_.rowCount == 0L)) + + System.setProperty(property, "true") + blockManager.removeBlock(blockId) + assert(blockManager.getStatus(blockId).isEmpty) + assert(cached.count() == 8) + assert(blockManager.getStatus(blockId).nonEmpty) + assert(builder.materializedRowCount == 8L) + assert(builder.loadedMaterializedStats.exists(_._1 == 8L)) + assert(relation.materializedMetadata.exists(_.rowCount == 8L)) + + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true") { + checkAnswer(cached.groupBy("id").count(), (0L until 8L).map(id => Row(id, 1L))) + } + } finally { + cached.unpersist(blocking = true) + System.clearProperty(property) + } + } + } + + test("SPARK-58272: materialized caches require trusted strict file reads") { + withTempPath { path => + spark.range(10).write.parquet(path.getCanonicalPath) + + def checkCache(data: DataFrame, expected: Boolean): Unit = { + val cached = data.filter($"id" < 5).persist(MEMORY_AND_DISK) + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + relation.cacheBuilder.cachedColumnBuffers.count() + assert(relation.isOutputRepeatable == expected) + assert(relation.statsAvailable == expected) + assert(relation.hasSelectivePredicate) + } finally { + cached.unpersist(blocking = true) + } + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.IGNORE_MISSING_FILES.key -> "false", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + checkCache(spark.read.parquet(path.getCanonicalPath), expected = true) + + val preinitialized = spark.read.parquet(path.getCanonicalPath) + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + try { + val relation = preinitialized.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val fileScan = relation.cacheBuilder.cachedPlan.collectFirst { + case scan: FileSourceScanExec => scan + }.get + withSQLConf(SQLConf.IGNORE_MISSING_FILES.key -> "true") { + fileScan.inputRDD + } + + relation.cacheBuilder.cachedColumnBuffers.count() + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + } finally { + preinitialized.unpersist(blocking = true) + } + + val rebuildable = spark.read.parquet(path.getCanonicalPath) + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + try { + val relation = rebuildable.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + val builder = relation.cacheBuilder + val fileScan = builder.cachedPlan.collectFirst { + case scan: FileSourceScanExec => scan + }.get + + // Keep the physical file reader strict while making only this cache generation observe + // best-effort session settings. + fileScan.inputRDD + withSQLConf(SQLConf.IGNORE_MISSING_FILES.key -> "true") { + builder.cachedColumnBuffers.count() + } + assert(!relation.materializedMetadata.get.isOutputRepeatable) + assert(!relation.mayHaveUsableMaterializedStats) + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + + builder.clearCache(blocking = true) + assert(relation.materializedMetadata.isEmpty) + builder.cachedColumnBuffers.count() + assert(relation.materializedMetadata.get.statsAvailable) + assert(relation.mayHaveUsableMaterializedStats) + assert(relation.isOutputRepeatable) + assert(relation.statsAvailable) + } finally { + rebuildable.unpersist(blocking = true) + } + + checkCache( + spark.read.option("ignoreMissingFiles", "true").parquet(path.getCanonicalPath), + expected = false) + checkCache( + spark.read.option("ignoreCorruptFiles", "true").parquet(path.getCanonicalPath), + expected = false) + } + + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.IGNORE_MISSING_FILES.key -> "true", + SQLConf.IGNORE_CORRUPT_FILES.key -> "false") { + checkCache(spark.read.parquet(path.getCanonicalPath), expected = false) + } + } + } + + test("SPARK-58272: unsafe cached lineages cannot supply runtime-filter statistics") { + val nondeterministic = spark.range(10).filter(rand() > 0.5).persist(MEMORY_AND_DISK) + val isSmall = udf((value: Long) => value < 5) + val userDefined = spark.range(10) + .filter(isSmall($"id")) + .persist(MEMORY_AND_DISK) + val randomizedEncryption = spark.range(10) + .selectExpr( + "id", + "aes_encrypt(CAST(id AS STRING), '0000111122223333') AS encrypted") + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + val currentTime = spark.range(10) + .selectExpr("id", "current_timestamp() AS observed_at") + .filter($"id" < 5) + .persist(MEMORY_AND_DISK) + + Seq(nondeterministic, userDefined, randomizedEncryption, currentTime).foreach { cached => + try { + val relation = cached.queryExecution.withCachedData.collectFirst { + case plan: InMemoryRelation => plan + }.get + relation.cacheBuilder.cachedColumnBuffers.count() + assert(!relation.isOutputRepeatable) + assert(!relation.statsAvailable) + } finally { + cached.unpersist(blocking = true) + } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlignAssignmentsSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlignAssignmentsSuiteBase.scala index c88ebb0d69ee7..bcc4895616bdc 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlignAssignmentsSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/AlignAssignmentsSuiteBase.scala @@ -30,12 +30,13 @@ import org.apache.spark.sql.catalyst.expressions.objects.AssertNotNull import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogV2Util, Column, ColumnDefaultValue, Identifier, SupportsRowLevelOperations, TableCapability, TableCatalog, TableWritePrivilege} +import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogV2Util, Column, ColumnDefaultValue, Identifier, SupportsRowLevelOperations, TableCapability, TableCatalog, TableContext, TableWritePrivilege} import org.apache.spark.sql.connector.expressions.{LiteralValue, Transform} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.datasources.v2.V2SessionCatalog import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{BooleanType, IntegerType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap abstract class AlignAssignmentsSuiteBase extends AnalysisTest { @@ -163,6 +164,10 @@ abstract class AlignAssignmentsSuiteBase extends AnalysisTest { }) when(newCatalog.loadTable(any(), any[java.util.Set[TableWritePrivilege]]())) .thenCallRealMethod() + // The options-aware overload runs the real default dispatch, which delegates to the + // stubbed overloads above. + when(newCatalog.loadTable(any(), any[TableContext](), any[CaseInsensitiveStringMap]())) + .thenCallRealMethod() when(newCatalog.name()).thenReturn("cat") newCatalog } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala index a924b637a79fa..1065c5baa5858 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala @@ -27,7 +27,7 @@ import org.mockito.invocation.InvocationOnMock import org.apache.spark.SparkUnsupportedOperationException import org.apache.spark.sql.{AnalysisException, SaveMode} import org.apache.spark.sql.catalyst.{AliasIdentifier, TableIdentifier} -import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AnalysisTest, Analyzer, AsOfVersion, EmptyFunctionRegistry, NoSuchTableException, RelationResolution, ResolvedFieldName, ResolvedFieldPosition, ResolvedIdentifier, ResolvedTable, ResolveSessionCatalog, TimeTravelSpec, UnresolvedAttribute, UnresolvedFieldPosition, UnresolvedInlineTable, UnresolvedPartitionSpec, UnresolvedRelation, UnresolvedSubqueryColumnAliases, UnresolvedTable} +import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AnalysisTest, Analyzer, AsOfVersion, EmptyFunctionRegistry, NoSuchTableException, RelationCache, RelationResolution, ResolvedFieldName, ResolvedFieldPosition, ResolvedIdentifier, ResolvedTable, ResolveSessionCatalog, TimeTravelSpec, UnresolvedAttribute, UnresolvedFieldPosition, UnresolvedInlineTable, UnresolvedPartitionSpec, UnresolvedRelation, UnresolvedSubqueryColumnAliases, UnresolvedTable} import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStorageFormat, CatalogTable, CatalogTableType, InMemoryCatalog, SessionCatalog, TempVariableManager} import org.apache.spark.sql.catalyst.expressions.{AttributeReference, Cast, EqualTo, Expression, InSubquery, IntegerLiteral, ListQuery, Literal, StringLiteral} import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke @@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{AlterColumns, AlterColumnSpe import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId import org.apache.spark.sql.connector.FakeV2Provider -import org.apache.spark.sql.connector.catalog.{CatalogManager, Column, ColumnDefaultValue, Identifier, SupportsDelete, Table, TableCapability, TableCatalog, TableChange, TableWritePrivilege, V1Table} +import org.apache.spark.sql.connector.catalog.{CatalogManager, Column, ColumnDefaultValue, Identifier, SupportsDelete, Table, TableCapability, TableCatalog, TableChange, TableContext, TableWritePrivilege, V1Table} import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME import org.apache.spark.sql.connector.expressions.{LiteralValue, Transform} import org.apache.spark.sql.errors.QueryExecutionErrors @@ -47,6 +47,7 @@ import org.apache.spark.sql.internal.SQLConf.{PARTITION_OVERWRITE_MODE, Partitio import org.apache.spark.sql.sources.SimpleScanSource import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{BooleanType, CharType, DoubleType, IntegerType, LongType, StringType, StructField, StructType, VarcharType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.unsafe.types.UTF8String class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { @@ -183,13 +184,17 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { val ident = invocation.getArguments()(0).asInstanceOf[Identifier] val version = invocation.getArguments()(1).asInstanceOf[String] (ident.name, version) match { - case ("tab", "v1") => table + case ("tab", "v1" | "v2") => table case ("tab", _) => throw new RuntimeException("Unknown version: " + version) case _ => throw new NoSuchTableException(Seq(ident.name)) } }) when(newCatalog.loadTable(any(), any[java.util.Set[TableWritePrivilege]]())) .thenCallRealMethod() + // The options-aware overload runs the real default dispatch, which delegates to the + // stubbed overloads above. + when(newCatalog.loadTable(any(), any[TableContext](), any[CaseInsensitiveStringMap]())) + .thenCallRealMethod() when(newCatalog.name()).thenReturn("testcat") newCatalog } @@ -209,6 +214,10 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { }) when(newCatalog.loadTable(any(), any[java.util.Set[TableWritePrivilege]]())) .thenCallRealMethod() + // The options-aware overload runs the real default dispatch, which delegates to the + // stubbed overloads above. + when(newCatalog.loadTable(any(), any[TableContext](), any[CaseInsensitiveStringMap]())) + .thenCallRealMethod() when(newCatalog.name()).thenReturn(CatalogManager.SESSION_CATALOG_NAME) newCatalog } @@ -3383,6 +3392,7 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { AnalysisContext.withNewAnalysisContext { val ctx = AnalysisContext.get assert(ctx.relationCache.isEmpty) + assert(ctx.tableCache.isEmpty) // create two unresolved relations without time travel val unresolved1 = UnresolvedRelation(Seq("testcat", "tab")) @@ -3398,7 +3408,8 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { // after first resolution, cache should have 1 entry (without time travel) assert(ctx.relationCache.size == 1) - assert(ctx.relationCache.keys.head._2.isEmpty) + assert(ctx.relationCache.keys.head.timeTravelSpec.isEmpty) + assert(ctx.tableCache.size == 1) // create unresolved relation with time travel spec val timeTravelSpec = AsOfVersion("v1") @@ -3413,6 +3424,17 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { // after time travel resolution, cache should have 2 entries (with and without time travel) assert(ctx.relationCache.size == 2) + assert(ctx.tableCache.size == 2) + + val otherTimeTravelSpec = AsOfVersion("v2") + val resolved4 = resolve( + UnresolvedRelation(Seq("testcat", "tab")), + Some(otherTimeTravelSpec)) + assert(resolved4.timeTravelSpec.contains(otherTimeTravelSpec)) + + // Distinct parsed time-travel specs are distinct state pins even with identical raw options. + assert(ctx.relationCache.size == 3) + assert(ctx.tableCache.size == 3) } } @@ -3456,6 +3478,194 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { } } + test("option-based time travel bypasses the shared relation cache") { + AnalysisContext.withNewAnalysisContext { + val ident = Identifier.of(Array.empty[String], "tab") + val cachedRelation = DataSourceV2Relation.create(table, Some(testCat), Some(ident)) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = (_, _, _, _, _) => { + sharedRelationCacheLookups += 1 + Some(cachedRelation) + } + val rule = new RelationResolution(catalogManagerWithDefault, sharedRelationCache) + val unresolved = UnresolvedRelation( + Seq("testcat", "tab"), + new CaseInsensitiveStringMap(java.util.Map.of("VeRsIoNaSoF", "v1"))) + + val resolved = rule.resolveRelation(unresolved) match { + case Some(AsDataSourceV2Relation(relation)) => relation + case other => fail(s"failed to resolve as v2 relation: $other") + } + + assert(resolved.timeTravelSpec.contains(AsOfVersion("v1"))) + assert(sharedRelationCacheLookups == 0) + } + } + + test("shared relation cache reuses table state and preserves the read's full options") { + val ident = Identifier.of(Array.empty[String], "tab") + val cachedTable = testCat.loadTable(ident) + + // A shared relation cache entry (as if left by an earlier CACHE TABLE) built with a specific + // set of options. The `id` tag lets the test tell a cache reuse apart from a fresh load, since + // both would otherwise carry the same mock `Table`. + def cachedRelationWith(opts: java.util.Map[String, String]): DataSourceV2Relation = { + val r = DataSourceV2Relation.create( + cachedTable, Some(testCat), Some(ident), new CaseInsensitiveStringMap(opts)) + r.setTagValue(LogicalPlan.PLAN_ID_TAG, 4242L) + r + } + + def resolveWith( + cacheOpts: java.util.Map[String, String], + readOpts: java.util.Map[String, String]): DataSourceV2Relation = { + AnalysisContext.withNewAnalysisContext { + val sharedRelationCache: RelationCache = + (_, _, _, _, _) => Some(cachedRelationWith(cacheOpts)) + val rule = new RelationResolution(catalogManagerWithDefault, sharedRelationCache) + val unresolved = + UnresolvedRelation(Seq("testcat", "tab"), new CaseInsensitiveStringMap(readOpts)) + rule.resolveRelation(unresolved) match { + case Some(AsDataSourceV2Relation(relation)) => relation + case other => fail(s"failed to resolve as v2 relation: $other") + } + } + } + + // Same options as the cached entry: the cache is reused (the tagged cached relation flows + // through, so the tag survives). + val reused = resolveWith( + java.util.Map.of("split-size", "5"), java.util.Map.of("split-size", "5")) + assert(reused.options.get("split-size") === "5") + assert(reused.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(4242L), + "matching options should reuse the cached relation") + + // This catalog does not declare any state options, so a different full option bag still + // reuses the cached table while the returned relation carries this read's complete options. + val reusedWithDifferentOptions = resolveWith( + java.util.Map.of("cachedOnly", "stale"), java.util.Map.of("split-size", "5")) + assert(reusedWithDifferentOptions.options.get("split-size") === "5") + assert(!reusedWithDifferentOptions.options.containsKey("cachedOnly")) + assert(reusedWithDifferentOptions.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(4242L), + "scan-specific options should not prevent shared table reuse") + } + + test("table-state cache consults shared relation cache only while establishing the pin") { + def newTable(id: String): Table = { + val t = mock(classOf[Table]) + when(t.id()).thenReturn(id) + when(t.name()).thenReturn("tab") + when(t.columns()).thenReturn(Array(Column.create("i", IntegerType))) + when(t.capabilities()).thenReturn(java.util.Set.of()) + t + } + + def options(splitSize: String): CaseInsensitiveStringMap = { + new CaseInsensitiveStringMap( + java.util.Map.of("state", "s1", "split-size", splitSize)) + } + + def run( + firstSplitSize: String, + secondSplitSize: String, + sharedRelationCacheEntry: ( + TableCatalog, + Identifier, + Table, + Table) => Option[LogicalPlan]): ( + DataSourceV2Relation, + DataSourceV2Relation, + Table, + Table, + Int, + Int) = { + val currentTable = newTable("table-id") + val cachedTable = newTable("table-id") + val catalog = mock(classOf[TableCatalog]) + when(catalog.name()).thenReturn("statecat") + when(catalog.tableStateOptionKeys()).thenReturn(java.util.Set.of("state")) + var loads = 0 + when(catalog.loadTable( + any[Identifier], + any[TableContext], + any[CaseInsensitiveStringMap])).thenAnswer((_: InvocationOnMock) => { + loads += 1 + currentTable + }) + + val manager = mock(classOf[CatalogManager]) + when(manager.catalog(any())).thenReturn(catalog) + when(manager.v1SessionCatalog).thenReturn(v1SessionCatalog) + val ident = Identifier.of(Array.empty[String], "tab") + val sharedRelationCacheCandidate = + sharedRelationCacheEntry(catalog, ident, currentTable, cachedTable) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = + (_, _, _, _, _) => { + sharedRelationCacheLookups += 1 + sharedRelationCacheCandidate + } + val resolver = new RelationResolution(manager, sharedRelationCache) + + def resolveWith(splitSize: String): DataSourceV2Relation = { + val unresolved = UnresolvedRelation(Seq("statecat", "tab"), options(splitSize)) + resolver.resolveRelation(unresolved) match { + case Some(AsDataSourceV2Relation(relation)) => relation + case other => fail(s"failed to resolve as v2 relation: $other") + } + } + + AnalysisContext.withNewAnalysisContext { + val first = resolveWith(firstSplitSize) + val second = resolveWith(secondSplitSize) + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + (first, second, currentTable, cachedTable, loads, sharedRelationCacheLookups) + } + } + + def cachedRelation( + catalog: TableCatalog, + ident: Identifier, + table: Table, + splitSize: String, + tag: Long): DataSourceV2Relation = { + val relation = DataSourceV2Relation.create( + table, + Some(catalog), + Some(ident), + options(splitSize)) + relation.setTagValue(LogicalPlan.PLAN_ID_TAG, tag) + relation + } + + // Situation A: the state-option cached match establishes the initial pin. The later same-state + // lookup reuses that pin without consulting the shared relation cache. + val situationA = run("5", "9", { (catalog, ident, _, cachedTable) => + Some(cachedRelation(catalog, ident, cachedTable, "5", 2L)) + }) + assert(situationA._1.table eq situationA._4) + assert(situationA._2.table eq situationA._4) + assert(situationA._1.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(2L)) + assert(situationA._2.getTagValue(LogicalPlan.PLAN_ID_TAG).isEmpty) + assert(situationA._5 == 1) + assert(situationA._6 == 1) + + // Situation B: scan-specific options differ on the first reference, but the state options + // still match. The shared cache therefore establishes the same pin regardless of reference + // order, and each returned relation keeps its own complete options. + val situationB = run("9", "5", { (catalog, ident, _, cachedTable) => + Some(cachedRelation(catalog, ident, cachedTable, "5", 6L)) + }) + assert(situationB._1.table eq situationB._4) + assert(situationB._2.table eq situationB._4) + assert(situationB._1.options.get("split-size") == "9") + assert(situationB._2.options.get("split-size") == "5") + assert(situationB._1.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(6L)) + assert(situationB._5 == 1) + assert(situationB._6 == 1) + } + private def resolve( unresolvedRelation: UnresolvedRelation, timeTravelSpec: Option[TimeTravelSpec] = None, diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/ShowFunctionsSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/ShowFunctionsSuiteBase.scala index 19bd830500834..23d8e3d1f0b38 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/ShowFunctionsSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/ShowFunctionsSuiteBase.scala @@ -118,6 +118,9 @@ trait ShowFunctionsSuiteBase extends QueryTest with DDLCommandTestUtils { assert(!systemFuns.filter("function='case'").isEmpty) // Built-in functions assert(!systemFuns.filter("function='substring'").isEmpty) + // sql/core-only builtin registered after FunctionRegistry.builtin is cloned + assert(!systemFuns.filter("function='parse_sql'").isEmpty) + assert(sql(s"SHOW USER FUNCTIONS IN $ns").filter("function='parse_sql'").isEmpty) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/DescribeTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/DescribeTableSuite.scala index 97fa08fef3a64..4e2c4a4c6975b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/DescribeTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/v1/DescribeTableSuite.scala @@ -24,11 +24,13 @@ import org.json4s.jackson.JsonMethods.parse import org.apache.spark.SPARK_VERSION import org.apache.spark.sql.{AnalysisException, QueryTest, Row} +import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType} import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME import org.apache.spark.sql.execution.command import org.apache.spark.sql.execution.command.{DescribeTableJson, Field, SqlPathEntry, TableColumn, Type} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.{IntegerType, StringType, StructType} /** * This base suite contains unified tests for the `DESCRIBE TABLE` command that checks V1 @@ -874,6 +876,40 @@ class DescribeTableSuite extends DescribeTableSuiteBase with CommandSuiteBase { } } + test("DESCRIBE TABLE is resilient to corrupt partition metadata") { + withNamespaceAndTable("ns", "table") { tbl => + // Simulate corrupt metadata where the declared partition column does not match the + // last field in the table schema. + val table = CatalogTable( + identifier = TableIdentifier("table", Some("ns")), + tableType = CatalogTableType.MANAGED, + storage = CatalogStorageFormat.empty, + schema = new StructType() + .add("id", IntegerType) + .add("actual_part", StringType), + provider = Some(getProvider()), + partitionColumnNames = Seq("declared_part")) + spark.sessionState.catalog.createTable(table, ignoreIfExists = false) + + val expectedInvalidInfo = Seq( + Row("# Invalid Partition Information", "", ""), + Row("Declared Partition Columns", "[declared_part]", ""), + Row("Last Columns in Table Schema", "[actual_part]", "")) + + Seq("DESCRIBE TABLE", "DESCRIBE TABLE EXTENDED").foreach { command => + val description = spark.sql(s"$command $tbl").collect().toSeq + assert(description.contains(Row("id", "int", null))) + assert(description.contains(Row("actual_part", "string", null))) + assert(description.containsSlice(expectedInvalidInfo)) + assert(!description.exists(_.getString(0) == "# Partition Information")) + } + + val jsonValue = spark.sql(s"DESCRIBE TABLE EXTENDED $tbl AS JSON").head().getString(0) + val parsedOutput = parse(jsonValue).extract[DescribeTableJson] + assert(parsedOutput.partition_columns === Some(List("declared_part"))) + } + } + test("DESCRIBE TABLE EXTENDED of a table with a default column value") { withTable("t") { spark.sql(s"CREATE TABLE t (id bigint default 42) $defaultUsing") diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ArchiveReadSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ArchiveReadSuiteBase.scala index ae64f5b502930..01a2521a737bf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ArchiveReadSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ArchiveReadSuiteBase.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.{DataFrame, QueryTest, Row} import org.apache.spark.sql.functions.{col, input_file_name, struct} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.sql.types.{StringType, StructType} +import org.apache.spark.sql.types.{DataType, StringType, StructType} import org.apache.spark.util.Utils /** @@ -105,6 +105,19 @@ trait ArchiveReadSuiteBase extends QueryTest with SharedSparkSession { protected def corruptEntryBytes: Array[Byte] = s"This is not a valid $format file".getBytes("UTF-8") + /** + * Whether [[writeArchiveFailingAfterFirstEntry]] is supported, gating the mid-advance + * corrupt-skip regression test. From the container trait. + */ + protected def supportsMidAdvanceFailure: Boolean + + /** + * Writes an archive at `dest` whose first entry (`firstEntry`) reads cleanly but which then fails + * while advancing to a later entry. From the container trait. + */ + protected def writeArchiveFailingAfterFirstEntry( + dest: File, firstEntry: (String, Array[Byte])): Unit + /** An archive extension whose reader fails on corrupt bytes (used by the corrupt-file tests). */ protected def corruptArchiveExtension: String @@ -170,6 +183,19 @@ trait ArchiveReadSuiteBase extends QueryTest with SharedSparkSession { */ protected def localizesEntries: Boolean = false + /** + * Prefix of the per-entry temp dir a [[localizesEntries]] format unpacks each entry into (e.g. + * `parquet-archive`); its inference variant is `<prefix>-infer`. Gates the shared temp-dir + * cleanup tests. + */ + protected def archiveTempDirPrefix: String = "" + + /** + * The config key toggling this format's vectorized/columnar reader, if it has one. When set, the + * shared localize-path read test runs with the reader both on and off. + */ + protected def vectorizedReaderConfKey: Option[String] = None + /** Sample data with a nested struct column, used by the complex-type test. */ protected def complexSampleDf: DataFrame = Seq((1, "NYC", "10001"), (2, "SF", "94105")).toDF("id", "city", "zip") @@ -331,9 +357,74 @@ trait ArchiveReadSuiteBase extends QueryTest with SharedSparkSession { } } + // ----- shared archivePathFilter tests -------------------------------------- + + test("archivePathFilter selects inner entries by full path") { + withArchiveFile() { archive => + writeArchive(archive, Seq( + s"top.$fileExtension" -> encodeFile(sampleDf((1, "top"))), + s"sub/keep.$fileExtension" -> encodeFile(sampleDf((2, "keep"))), + s"other/skip.$fileExtension" -> encodeFile(sampleDf((3, "skip"))))) + // `sub/*` matches the full inner path, so only the entry under `sub/` is ingested. + checkAnswer( + read(archive.getCanonicalPath, Map("archivePathFilter" -> "sub/*")).select("id", "name"), + Seq(Row(2, "keep"))) + } + } + + test("archivePathFilter with an extension glob selects across subdirectories") { + withArchiveFile() { archive => + writeArchive(archive, Seq( + s"top.$fileExtension" -> encodeFile(sampleDf((1, "top"))), + s"sub/nested.$fileExtension" -> encodeFile(sampleDf((2, "nested"))), + "sub/skip.other" -> encodeFile(sampleDf((3, "skip"))))) + // `*` crosses `/`, so the glob keeps both entries of this extension at any depth, while the + // entry with a different extension is filtered out. + checkAnswer( + read(archive.getCanonicalPath, Map("archivePathFilter" -> s"*.$fileExtension")) + .select("id", "name"), + Seq(Row(1, "top"), Row(2, "nested"))) + } + } + + test("archivePathFilter matching no entry yields no rows") { + withArchiveFile() { archive => + writeArchive(archive, Seq(s"data.$fileExtension" -> encodeFile(sampleDf((1, "Alice"))))) + checkAnswer( + read(archive.getCanonicalPath, Map("archivePathFilter" -> "nomatch/*")), Seq.empty[Row]) + } + } + + test("archivePathFilter applies in addition to ignoredPathSegmentRegex") { + withArchiveFile() { archive => + writeArchive(archive, Seq( + s"keep/data.$fileExtension" -> encodeFile(sampleDf((1, "keep"))), + // Matches the glob, but the hidden-file filter still drops the `_`-prefixed entry. + s"keep/_hidden.$fileExtension" -> encodeFile(sampleDf((2, "hidden"))))) + checkAnswer( + read(archive.getCanonicalPath, Map("archivePathFilter" -> "keep/*")).select("id", "name"), + Seq(Row(1, "keep"))) + } + } + // ----- shared schema-inference tests (run when `supportsSchemaInference`) -- if (supportsSchemaInference) { + test("archivePathFilter applies to schema inference, not just the scan") { + // The excluded entry carries a column the kept entry lacks. Inference must skip it, otherwise + // the inferred schema is a superset of what the scan returns and `extra` reads back all-null. + withArchiveFile() { archive => + writeArchive(archive, Seq( + s"keep/data.$fileExtension" -> encodeFile(sampleDf((1, "keep"))), + s"skip/data.$fileExtension" -> + encodeFile(Seq((2, "skip", "x")).toDF("id", "name", "extra")))) + val schema = inferredSchema( + Seq(archive.getCanonicalPath), Map("archivePathFilter" -> "keep/*")) + assert(!schema.fieldNames.contains("extra"), + s"inference read a filtered-out entry; got $schema") + } + } + test("archive infers the same schema as a directory of the same files") { val entries = Seq(sampleDf((1, "Alice"), (2, "Bob")), sampleDf((3, "Carol"))) .zipWithIndex.map { case (p, i) => entryName(i) -> encodeFile(p) } @@ -517,6 +608,135 @@ trait ArchiveReadSuiteBase extends QueryTest with SharedSparkSession { } } } + + // Temp dirs (of `prefix`) surviving under the executor-local dir; the read/inference variants + // must be removed on task completion, so `after -- before` catches a leak. + def archiveTempDirs(prefix: String): Set[String] = { + val localDir = new File(Utils.getLocalDir(spark.sparkContext.getConf)) + Option(localDir.listFiles()).getOrElse(Array.empty) + .filter(_.getName.startsWith(prefix)).map(_.getName).toSet + } + + val vectorizedCases = vectorizedReaderConfKey.map(Seq(_)).getOrElse(Seq.empty) + .flatMap(key => Seq(true, false).map(key -> _)) + (if (vectorizedCases.nonEmpty) vectorizedCases else Seq("" -> false)).foreach { + case (confKey, vectorized) => + val label = if (confKey.isEmpty) "" else s" with vectorized reader = $vectorized" + test(s"archive reads return the same rows$label") { + val conf = if (confKey.isEmpty) Map.empty[String, String] + else Map(confKey -> vectorized.toString) + withSQLConf(conf.toSeq: _*) { + assertArchiveMatchesDir( + Seq(entryName(0) -> encodeFile(sampleDf((1, "Alice"), (2, "Bob"))))) + } + } + } + + test("an abandoned read (LIMIT) over an archive returns partial rows and cleans up") { + withArchiveFile() { archive => + val parts = (0 until 4).map(i => entryName(i) -> encodeFile(sampleDf((i, s"v$i")))) + writeArchive(archive, parts) + val before = archiveTempDirs(archiveTempDirPrefix) + assert(read(archive.getCanonicalPath).limit(2).collect().length == 2) + assert((archiveTempDirs(archiveTempDirPrefix) -- before).isEmpty, + "the read's temp dir was not cleaned up") + } + } + + test("extensionless entries are read and inferred like a directory of part-files") { + val data = sampleDf((1, "Alice"), (2, "Bob")) + withArchiveFile() { archive => + writeArchive(archive, Seq("part-00000" -> encodeFile(data))) + checkAnswer(read(archive.getCanonicalPath), data) + assert(inferredSchema(Seq(archive.getCanonicalPath)).fieldNames.toSet == Set("id", "name"), + "an extensionless entry should be inferred like a directory of part-files") + } + } + + test("a corrupt archive cleans up its read temp dir rather than leaking it") { + // A corrupt archive throws before the read returns an iterator, but must not leak its dir. + withArchiveFile(corruptArchiveExtension) { archive => + writeCorruptArchive(archive) + val before = archiveTempDirs(archiveTempDirPrefix) + intercept[SparkException](read(archive.getCanonicalPath).collect()) + assert((archiveTempDirs(archiveTempDirPrefix) -- before).isEmpty, + "a corrupt archive leaked its read temp dir") + } + } + + test("a corrupt archive cleans up its inference temp dir rather than leaking it") { + // Inference localizes entries too, on a worker without a TaskContext; a corrupt archive + // throws during that eager localize and must not leak the inference temp dir. + withArchiveFile(corruptArchiveExtension) { archive => + writeCorruptArchive(archive) + val before = archiveTempDirs(s"$archiveTempDirPrefix-infer") + intercept[SparkException](inferredSchema(Seq(archive.getCanonicalPath))) + assert((archiveTempDirs(s"$archiveTempDirPrefix-infer") -- before).isEmpty, + "a corrupt archive leaked its inference temp dir") + } + } + + test("archive inference unions differing fields across entries with mergeSchema=true") { + // mergeSchema=true folds every entry's schema; over an archive, one unpacked entry at a time. + // Compare field name/type pairs as a set rather than by exact StructType equality: field + // order across archive entries follows the parallel merge order and is not a guaranteed + // contract, and some formats (ORC) merge entries unordered. Only the name/type union matters. + val withName = sampleDf((1, "Alice"), (2, "Bob")) + val idExtra = Seq((3, 30)).toDF("id", "extra") + val entries = Seq(entryName(0) -> encodeFile(withName), entryName(1) -> encodeFile(idExtra)) + val merge = Map("mergeSchema" -> "true") + def fieldTypes(s: StructType): Set[(String, DataType)] = + s.fields.map(f => f.name -> f.dataType).toSet + withArchiveFile() { archive => + writeArchive(archive, entries) + val archiveSchema = inferredSchema(Seq(archive.getCanonicalPath), merge) + withTempDir { dir => + entries.foreach { case (n, b) => Files.write(new File(dir, n).toPath, b) } + assert(archiveSchema.fieldNames.toSet == Set("id", "name", "extra"), + s"expected the union of entry fields, got $archiveSchema") + assert(fieldTypes(archiveSchema) == + fieldTypes(inferredSchema(Seq(dir.getCanonicalPath), merge)), + s"archive mergeSchema inference diverged from a directory read; got $archiveSchema") + } + } + } + } + + // ----- shared parent-archive _metadata test -------------------------------- + + test("_metadata exposes the parent archive file's values, identical for every row") { + archiveExtensions.foreach { ext => + withArchiveFile(ext) { archive => + // Multiple entries, each multiple rows: the archive is one non-splittable PartitionedFile, + // so every row must carry the same parent-archive metadata (not any inner entry's). + val parts = Seq(sampleDf((1, "Alice"), (2, "Bob")), sampleDf((3, "Carol"), (4, "Dan"))) + writeArchive( + archive, parts.zipWithIndex.map { case (p, i) => entryName(i) -> encodeFile(p) }) + + val rows = read(archive.getCanonicalPath) + .select("_metadata.file_path", "_metadata.file_name", "_metadata.file_size", + "_metadata.file_block_start", "_metadata.file_block_length", + "_metadata.file_modification_time") + .collect() + assert(rows.length == parts.map(_.count()).sum, + s"expected one row per input record, got ${rows.length}") + + val fileSize = archive.length() + rows.foreach { r => + assert(r.getString(0).endsWith(archive.getName) && !r.getString(0).contains(entryName(0)), + s"file_path should be the archive file, got ${r.getString(0)}") + assert(r.getString(1) == archive.getName, s"file_name mismatch: ${r.getString(1)}") + assert(r.getLong(2) == fileSize, s"file_size mismatch: ${r.getLong(2)} != $fileSize") + assert(r.getLong(3) == 0L, s"file_block_start should be 0, got ${r.getLong(3)}") + assert(r.getLong(4) == fileSize, + s"file_block_length should be the archive size, got ${r.getLong(4)}") + assert(r.getAs[java.sql.Timestamp](5).getTime == archive.lastModified(), + "file_modification_time should be the archive's mtime") + } + assert(rows.map(_.toSeq).distinct.length == 1, + "every row must carry the same parent-archive metadata") + } + } } // ----- shared complex-type test (run when `supportsComplexTypes`) ---------- diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/BinaryFileArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/BinaryFileArchiveReadBase.scala index bcbfb89d12415..e8fc1268abb8d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/BinaryFileArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/BinaryFileArchiveReadBase.scala @@ -50,14 +50,6 @@ trait BinaryFileArchiveReadBase extends QueryTest with SharedSparkSession { /** Extension of the archive [[writeCorruptArchive]] produces (corruption is format-specific). */ protected def corruptArchiveExtension: String - /** - * Whether the container reports each entry's size up front. Streaming zip - * (`ZipArchiveInputStream`) does not -- an entry sized only by a trailing data descriptor reads - * back as `-1` -- so the per-entry `length` tests are skipped for zip until it moves to - * `ZipFile`. tar and 7z report real sizes. - */ - protected def entrySizeKnown: Boolean = true - override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ARCHIVE_FORMAT_READER_ENABLED.key, "true") @@ -98,7 +90,6 @@ trait BinaryFileArchiveReadBase extends QueryTest with SharedSparkSession { } test("wholeFile=false sources path and length from each entry, modtime from the parent") { - assume(entrySizeKnown) withArchiveFile() { archive => writeArchive(archive, Seq("a.bin" -> bytes("aaa"), "b.bin" -> bytes("bbbb"))) val rows = read(archive.getCanonicalPath, Map("wholeFile" -> "false")) @@ -116,6 +107,37 @@ trait BinaryFileArchiveReadBase extends QueryTest with SharedSparkSession { } } + test("wholeFile=false _metadata exposes the parent archive's values for every row") { + archiveExtensions.foreach { ext => + withArchiveFile(ext) { archive => + writeArchive(archive, Seq("a.bin" -> bytes("aaa"), "b.bin" -> bytes("bbbb"))) + val rows = read(archive.getCanonicalPath, Map("wholeFile" -> "false")) + .select("_metadata.file_path", "_metadata.file_name", "_metadata.file_size", + "_metadata.file_block_start", "_metadata.file_block_length", + "_metadata.file_modification_time") + .collect() + assert(rows.length == 2) + + val fileSize = archive.length() + rows.foreach { r => + // The `path`/`length` data columns are per entry here, but _metadata stays parent-only: + // it is derived from the single PartitionedFile. + assert(r.getString(0).endsWith(archive.getName) && !r.getString(0).contains("!/"), + s"file_path should be the archive file, got ${r.getString(0)}") + assert(r.getString(1) == archive.getName, s"file_name mismatch: ${r.getString(1)}") + assert(r.getLong(2) == fileSize, s"file_size mismatch: ${r.getLong(2)} != $fileSize") + assert(r.getLong(3) == 0L, s"file_block_start should be 0, got ${r.getLong(3)}") + assert(r.getLong(4) == fileSize, + s"file_block_length should be the archive size, got ${r.getLong(4)}") + assert(r.getAs[java.sql.Timestamp](5).getTime == archive.lastModified(), + "file_modification_time should be the archive's mtime") + } + assert(rows.map(_.toSeq).distinct.length == 1, + "every row must carry the same parent-archive metadata") + } + } + } + test("wholeFile=false on an empty archive yields no rows") { withArchiveFile() { archive => writeArchive(archive, Seq.empty) @@ -135,8 +157,20 @@ trait BinaryFileArchiveReadBase extends QueryTest with SharedSparkSession { } } + test("wholeFile=false honors archivePathFilter, applied on top of hidden-entry filtering") { + withArchiveFile() { archive => + writeArchive(archive, Seq( + "keep/a.bin" -> bytes("a"), + "keep/_hidden.bin" -> bytes("drop"), // matches the glob but hidden + "other/b.bin" -> bytes("drop"))) + checkAnswer( + read(archive.getCanonicalPath, Map("wholeFile" -> "false", "archivePathFilter" -> "keep/*")) + .select("content"), + Seq(Row(bytes("a")))) + } + } + test("wholeFile=false enforces SOURCES_BINARY_FILE_MAX_LENGTH per entry") { - assume(entrySizeKnown) withArchiveFile() { archive => writeArchive(archive, Seq("big.bin" -> bytes("0123456789"))) withSQLConf(SQLConf.SOURCES_BINARY_FILE_MAX_LENGTH.key -> "4") { @@ -149,7 +183,6 @@ trait BinaryFileArchiveReadBase extends QueryTest with SharedSparkSession { } test("wholeFile=false honors length filter pushdown against each entry") { - assume(entrySizeKnown) withArchiveFile() { archive => writeArchive(archive, Seq("a.bin" -> bytes("aaa"), "b.bin" -> bytes("bbbb"))) // Entry lengths are 3 and 4; the filter selects per entry, not against the archive size. @@ -220,11 +253,6 @@ class BinaryFileTarArchiveReadSuite extends BinaryFileArchiveReadBase with TarAr class BinaryFileZipArchiveReadSuite extends BinaryFileArchiveReadBase with ZipArchiveTestUtils { override protected def corruptArchiveExtension: String = "zip" - - // Streaming `ZipArchiveInputStream` cannot report an entry's size before reading it, so the - // per-entry `length` tests are skipped for zip. Remove this override once zip reads move to - // `ZipFile`, which exposes entry sizes from the central directory. - override protected def entrySizeKnown: Boolean = false } class BinaryFileSevenZArchiveReadSuite diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/CSVArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/CSVArchiveReadBase.scala index b3ab43fcd19c1..3665f857edd6d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/CSVArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/CSVArchiveReadBase.scala @@ -131,11 +131,8 @@ trait CSVArchiveReadBase extends ArchiveReadSuiteBase { } test("CSV: the DSv2 path refuses to infer a schema for an archive (UNABLE_TO_INFER_SCHEMA)") { - // Archive scanning is wired into the V1 file source only, so the DSv2 reader cannot read - // archives. On the V2 path inference must keep returning None for an archive input -- raising - // UNABLE_TO_INFER_SCHEMA -- rather than inferring a schema and letting the V2 scan parse the - // raw archive bytes as CSV. Forcing csv off the V1 source list routes the read through - // CSVTable. + // Forcing csv off the V1 source list routes the archive read through the DSv2 CSVTable, which + // cannot read archives and must fail with UNABLE_TO_INFER_SCHEMA, not parse raw bytes. withArchiveFile() { archive => writeArchive(archive, Seq(entryName(0) -> encodeFile(sampleDf((1, "Alice"), (2, "Bob"))))) withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/JSONArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/JSONArchiveReadBase.scala index 24a1873805350..70de9e5e63178 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/JSONArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/JSONArchiveReadBase.scala @@ -184,11 +184,30 @@ trait JSONArchiveReadBase extends ArchiveReadSuiteBase { schema = corruptSchema) } + if (supportsMidAdvanceFailure) { + test("JSON: multiLine inference keeps entries read before a mid-advance failure " + + "(ignoreCorruptFiles)") { + // Entry 0 is read, then advancing to a later entry throws (not at open). A whole-archive drop + // would lose entry 0's `extra`; aborting the traversal would lose the sibling file's `later`. + val opts = Map("multiLine" -> "true") + withArchiveFile() { archive => + writeArchiveFailingAfterFirstEntry(archive, + entryName(0) -> jsonBytes("{\n \"id\": 1,\n \"name\": \"Alice\",\n \"extra\": 9\n}")) + Files.write(new File(archive.getParentFile, s"later.$fileExtension").toPath, + jsonBytes("{\n \"id\": 2,\n \"name\": \"Bob\",\n \"later\": 7\n}")) + withSQLConf(SQLConf.IGNORE_CORRUPT_FILES.key -> "true") { + val schema = inferredSchema(Seq(archive.getParentFile.getCanonicalPath), opts) + assert(schema.fieldNames.toSet == Set("id", "name", "extra", "later"), + "expected `extra` (pre-failure entry) and `later` (sibling file) in the inferred " + + s"schema after the mid-advance skip, got $schema") + } + } + } + } + test("JSON: the DSv2 path refuses to infer a schema for an archive (UNABLE_TO_INFER_SCHEMA)") { - // Archive scanning is wired into the v1 file source only, so the DSv2 reader cannot read - // archives. On the v2 path inference must keep returning None for an archive input -- raising - // UNABLE_TO_INFER_SCHEMA -- rather than inferring a schema the v2 scan would then mis-read as - // raw archive bytes. Forcing json off the v1 source list routes the read through JsonTable. + // Forcing json off the v1 source list routes the archive read through the DSv2 JsonTable, which + // cannot read archives and must fail with UNABLE_TO_INFER_SCHEMA, not parse raw bytes. withArchiveFile() { archive => writeArchive(archive, Seq(entryName(0) -> encodeFile(sampleDf((1, "Alice"), (2, "Bob"))))) withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/OrcArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/OrcArchiveReadBase.scala new file mode 100644 index 0000000000000..a63152caa975b --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/OrcArchiveReadBase.scala @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources + +import org.apache.spark.sql.internal.SQLConf + +/** + * Binds [[ArchiveReadSuiteBase]]'s hooks to ORC (entries unpacked to a local file for footer + * random access). ORC is self-describing, so the base's schema-inference tests run too. + */ +trait OrcArchiveReadBase extends ArchiveReadSuiteBase { + + override protected def format: String = "orc" + + override protected def fileExtension: String = "orc" + + override protected def readOptions: Map[String, String] = Map.empty + + override protected def readSchema: String = "id INT, name STRING" + + // ORC has authoritative per-file schemas and only unions under `mergeSchema`, so it opts out of + // the by-name default-inference union (covered by the shared localize-path tests instead). + override protected def supportsSchemaMerge: Boolean = false + + // ORC samples one part-file for non-merge inference. + override protected def inferenceSamplesOneFile: Boolean = true + + // ORC unpacks each entry to a local temp file for footer random access. + override protected def localizesEntries: Boolean = true + + override protected def archiveTempDirPrefix: String = "orc-archive" + + override protected def vectorizedReaderConfKey: Option[String] = + Some(SQLConf.ORC_VECTORIZED_READER_ENABLED.key) +} + +class OrcTarArchiveReadSuite + extends ArchiveReadSuiteBase + with OrcArchiveReadBase + with TarArchiveReadBase + +class OrcZipArchiveReadSuite + extends ArchiveReadSuiteBase + with OrcArchiveReadBase + with ZipArchiveReadBase + +class OrcSevenZArchiveReadSuite + extends ArchiveReadSuiteBase + with OrcArchiveReadBase + with SevenZArchiveReadBase diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ParquetArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ParquetArchiveReadBase.scala index c09f764ba3d2a..374098789f16c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ParquetArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ParquetArchiveReadBase.scala @@ -18,15 +18,12 @@ package org.apache.spark.sql.execution.datasources import java.io.File -import java.nio.file.Files import org.apache.hadoop.fs.{FileStatus, Path} -import org.apache.spark.SparkException import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType -import org.apache.spark.util.Utils /** * Binds [[ArchiveReadSuiteBase]]'s hooks to Parquet (entries unpacked to a local file for footer @@ -34,8 +31,6 @@ import org.apache.spark.util.Utils */ trait ParquetArchiveReadBase extends ArchiveReadSuiteBase { - import testImplicits._ - override protected def format: String = "parquet" override protected def fileExtension: String = "parquet" @@ -54,89 +49,10 @@ trait ParquetArchiveReadBase extends ArchiveReadSuiteBase { // Parquet unpacks each entry to a local temp file for footer random access. override protected def localizesEntries: Boolean = true - for (vectorized <- Seq(true, false)) { - test(s"archive reads return the same rows with vectorized reader = $vectorized") { - withSQLConf(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> vectorized.toString) { - assertArchiveMatchesDir( - Seq(entryName(0) -> encodeFile(sampleDf((1, "Alice"), (2, "Bob"))))) - } - } - } - - test("an abandoned read (LIMIT) over an archive returns partial rows and cleans up") { - def archiveTempDirs(localDir: File): Set[String] = - Option(localDir.listFiles()).getOrElse(Array.empty) - .filter(_.getName.startsWith("parquet-archive")).map(_.getName).toSet - withArchiveFile() { archive => - val parts = (0 until 4).map(i => entryName(i) -> encodeFile(sampleDf((i, s"v$i")))) - writeArchive(archive, parts) - val localDir = new File(Utils.getLocalDir(spark.sparkContext.getConf)) - val before = archiveTempDirs(localDir) - assert(read(archive.getCanonicalPath).limit(2).collect().length == 2) - // This read's per-entry temp dir (prefix `parquet-archive`) must be removed on task - // completion, so no new one survives. - assert((archiveTempDirs(localDir) -- before).isEmpty, - "the read's temp dir was not cleaned up") - } - } - - test("extensionless entries are read and inferred like a directory of part-files") { - val data = sampleDf((1, "Alice"), (2, "Bob")) - withArchiveFile() { archive => - writeArchive(archive, Seq("part-00000" -> encodeFile(data))) - checkAnswer(read(archive.getCanonicalPath), data) - assert(inferredSchema(Seq(archive.getCanonicalPath)).fieldNames.toSet == Set("id", "name"), - "an extensionless entry should be inferred like a directory of part-files") - } - } + override protected def archiveTempDirPrefix: String = "parquet-archive" - private def parquetArchiveTempDirs(prefix: String): Set[String] = { - val localDir = new File(Utils.getLocalDir(spark.sparkContext.getConf)) - Option(localDir.listFiles()).getOrElse(Array.empty) - .filter(_.getName.startsWith(prefix)).map(_.getName).toSet - } - - test("a corrupt archive cleans up its read temp dir rather than leaking it") { - // A corrupt archive throws before the read returns an iterator, but must not leak the temp dir. - withArchiveFile(corruptArchiveExtension) { archive => - writeCorruptArchive(archive) - val before = parquetArchiveTempDirs("parquet-archive") - intercept[SparkException](read(archive.getCanonicalPath).collect()) - assert((parquetArchiveTempDirs("parquet-archive") -- before).isEmpty, - "a corrupt archive leaked its read temp dir") - } - } - - test("a corrupt archive cleans up its inference temp dir rather than leaking it") { - // Inference localizes entries too (readArchiveFooters), on a worker without a TaskContext; a - // corrupt archive throws during that eager localize and must not leak parquet-archive-infer. - withArchiveFile(corruptArchiveExtension) { archive => - writeCorruptArchive(archive) - val before = parquetArchiveTempDirs("parquet-archive-infer") - intercept[SparkException](inferredSchema(Seq(archive.getCanonicalPath))) - assert((parquetArchiveTempDirs("parquet-archive-infer") -- before).isEmpty, - "a corrupt archive leaked its inference temp dir") - } - } - - test("archive inference unions differing fields across entries with mergeSchema=true") { - // mergeSchema=true folds every entry's footer; over an archive, one unpacked entry at a time. - val withName = sampleDf((1, "Alice"), (2, "Bob")) - val idExtra = Seq((3, 30)).toDF("id", "extra") - val entries = Seq(entryName(0) -> encodeFile(withName), entryName(1) -> encodeFile(idExtra)) - val merge = Map("mergeSchema" -> "true") - withArchiveFile() { archive => - writeArchive(archive, entries) - val archiveSchema = inferredSchema(Seq(archive.getCanonicalPath), merge) - withTempDir { dir => - entries.foreach { case (n, b) => Files.write(new File(dir, n).toPath, b) } - assert(archiveSchema.fieldNames.toSet == Set("id", "name", "extra"), - s"expected the union of entry fields, got $archiveSchema") - assert(archiveSchema == inferredSchema(Seq(dir.getCanonicalPath), merge), - s"archive mergeSchema inference diverged from a directory read; got $archiveSchema") - } - } - } + override protected def vectorizedReaderConfKey: Option[String] = + Some(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key) test("inference skips a missing archive among good ones (ignoreMissingFiles)") { // Exercised on ParquetFileFormat.inferSchema(files) directly: inference now runs on the diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala index d6ed274327a66..73590b039a014 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/PushVariantIntoScanSuite.scala @@ -26,6 +26,7 @@ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest trait PushVariantIntoScanSuiteBase extends SharedSparkSession { override def sparkConf: SparkConf = @@ -34,6 +35,18 @@ trait PushVariantIntoScanSuiteBase extends SharedSparkSession { // Whether the reader-deferral tests should exercise the V2 read path. Subclasses override. protected def useV2: Boolean + test("hoistable variant extractions are not classified as throwable") { + // If either expression becomes throwable, join hoisting requires cast-error deferral. + val v = AttributeReference("v", VariantType)() + assert(!VariantGet( + v, + Literal("$.a"), + IntegerType, + failOnError = true, + timeZoneId = Some(localTimeZone)).throwable) + assert(!Cast(v, IntegerType, timeZoneId = Some(localTimeZone)).throwable) + } + // Write a parquet dataset via V1, then expose it as the temp view `T`. The view's read path is // V2 when `useV2`, V1 otherwise. Use this for tests that need to actually execute a scan and // compare V1 vs V2 behavior. @@ -2707,6 +2720,7 @@ abstract class PushVariantIntoScanV2SuiteBase extends QueryTest with PushVariant } // V2 DataSource tests - Row-based reader +@ExtendedSQLTest class PushVariantIntoScanV2Suite extends PushVariantIntoScanV2SuiteBase { override protected def vectorizedReaderEnabled: Boolean = false override protected def readerName: String = "row-based reader" diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala index 6b8f3495f4a02..e37f004bf1fde 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SchemaPruningSuite.scala @@ -908,7 +908,6 @@ abstract class SchemaPruningSuite testSchemaPruning("select nested field in Expand") { import org.apache.spark.sql.catalyst.dsl.expressions._ - import testImplicits.castToImpl val query1 = Expand( Seq( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SevenZArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SevenZArchiveReadBase.scala index efc6c280bfb65..1c77964e6a5f6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SevenZArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SevenZArchiveReadBase.scala @@ -68,4 +68,12 @@ trait SevenZArchiveTestUtils { protected def writeCorruptArchive(dest: File): Unit = Files.write(dest.toPath, "this is not a valid 7z archive, just some random bytes" .getBytes(StandardCharsets.UTF_8)) + + // 7z validates its whole index when the archive is opened, so this helper cannot construct a + // failure that occurs only while advancing to another entry. + protected def supportsMidAdvanceFailure: Boolean = false + + protected def writeArchiveFailingAfterFirstEntry( + dest: File, firstEntry: (String, Array[Byte])): Unit = + throw new UnsupportedOperationException("7z cannot fail while advancing to a later entry") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormatSuite.scala index b63b9ac310310..d9eae971e975e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormatSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormatSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.execution.datasources import java.io.{ByteArrayOutputStream, Closeable, File, FileOutputStream, InputStream, OutputStream} import java.nio.charset.StandardCharsets +import java.nio.file.Files import java.util.Properties import java.util.regex.Pattern import java.util.zip.GZIPOutputStream @@ -29,9 +30,10 @@ import org.apache.commons.compress.archivers.sevenz.{SevenZArchiveEntry, SevenZO import org.apache.commons.compress.archivers.tar.{TarArchiveEntry, TarArchiveOutputStream} import org.apache.commons.compress.archivers.zip.{ZipArchiveEntry, ZipArchiveOutputStream} import org.apache.hadoop.conf.Configuration -import org.apache.hadoop.fs.Path +import org.apache.hadoop.fs.{GlobPattern, Path} -import org.apache.spark.{SparkFunSuite, TaskContext, TaskContextImpl} +import org.apache.spark.{SparkFunSuite, SparkRuntimeException, TaskContext, TaskContextImpl} +import org.apache.spark.sql.catalyst.FileSourceOptions /** * Unit tests for the streaming [[SupportsArchiveFormat]] engine: `isArchivePath` dispatch and @@ -82,25 +84,29 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { } finally out.close() } + /** Appends `v` as an unsigned 2-byte little-endian integer -- one of ZIP's header fields. */ + private def u16(out: ByteArrayOutputStream, v: Int): Unit = { + out.write(v & 0xFF); out.write((v >>> 8) & 0xFF) + } + + /** Appends `v` as an unsigned 4-byte little-endian integer -- one of ZIP's header fields. */ + private def u32(out: ByteArrayOutputStream, v: Long): Unit = { + out.write((v & 0xFF).toInt); out.write(((v >>> 8) & 0xFF).toInt) + out.write(((v >>> 16) & 0xFF).toInt); out.write(((v >>> 24) & 0xFF).toInt) + } + /** * Writes a zip with one STORED (uncompressed) entry that uses a data descriptor: general-purpose * bit 3 is set and the local header's crc/size fields are zeroed, so the real values live only in - * the trailing data descriptor. `ZipArchiveInputStream` cannot stream such an entry -- it has no - * size to bound the read -- so `read` throws rather than yielding truncated bytes. This is the - * non-streamable case for pure-streaming zip reads; `ZipArchiveOutputStream` cannot produce it - * (it rejects an unsized STORED entry, or rewrites the header when the sink is seekable), so the - * bytes are assembled by hand. + * the trailing data descriptor. */ private def writeStoredEntryWithDataDescriptor(file: File, name: String, body: String): Unit = { val nameBytes = name.getBytes(StandardCharsets.UTF_8) val data = body.getBytes(StandardCharsets.UTF_8) val crc = { val c = new java.util.zip.CRC32(); c.update(data); c.getValue } val out = new ByteArrayOutputStream() - def u16(v: Int): Unit = { out.write(v & 0xFF); out.write((v >>> 8) & 0xFF) } - def u32(v: Long): Unit = { - out.write((v & 0xFF).toInt); out.write(((v >>> 8) & 0xFF).toInt) - out.write(((v >>> 16) & 0xFF).toInt); out.write(((v >>> 24) & 0xFF).toInt) - } + def u16(v: Int): Unit = this.u16(out, v) + def u32(v: Long): Unit = this.u32(out, v) // Local file header: GP bit 3 set (data descriptor), STORED method, sizes zeroed here. val localHeaderOffset = out.size() u32(0x04034b50L); u16(10); u16(0x0008); u16(0); u16(0); u16(0) @@ -120,7 +126,52 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { u32(0x06054b50L); u16(0); u16(0); u16(1); u16(1) u32(cdSize.toLong); u32(cdOffset.toLong); u16(0) val fos = new FileOutputStream(file) - try fos.write(out.toByteArray) finally fos.close() + try { + fos.write(out.toByteArray) + } finally { + fos.close() + } + } + + /** + * Writes a zip whose single STORED entry sets the encryption flag (general-purpose bit 0), so + * `ZipFile#canReadEntryData` returns false. Assembled by hand: no writer we depend on emits an + * encryption-flagged entry (commons-compress's `ZipArchiveOutputStream` cannot). The bytes are + * not actually encrypted -- only the flag is set, which alone drives the read path under test. + * ZIP is a sequence of little-endian records, each led by a 4-byte signature. + */ + private def writeEncryptedEntry(file: File, name: String, body: String): Unit = { + val nameBytes = name.getBytes(StandardCharsets.UTF_8) + val data = body.getBytes(StandardCharsets.UTF_8) + val crc = { val c = new java.util.zip.CRC32(); c.update(data); c.getValue } + val out = new ByteArrayOutputStream() + def u16(v: Int): Unit = this.u16(out, v) + def u32(v: Long): Unit = this.u32(out, v) + val size = data.length.toLong + // Local file header (sig 0x04034b50): version, GP flags = 0x0001 (bit 0 = encrypted), STORED + // method (0), mod time/date, crc, compressed & uncompressed sizes, name & extra lengths. + val localHeaderOffset = out.size() + u32(0x04034b50L); u16(10); u16(0x0001); u16(0); u16(0); u16(0) + u32(crc); u32(size); u32(size) + u16(nameBytes.length); u16(0) + out.write(nameBytes); out.write(data) + // Central directory record (sig 0x02014b50): mirrors the header, same 0x0001 encrypted flag, + // and points back at the local header offset. + val cdOffset = out.size() + u32(0x02014b50L); u16(20); u16(10); u16(0x0001); u16(0); u16(0); u16(0) + u32(crc); u32(size); u32(size) + u16(nameBytes.length); u16(0); u16(0); u16(0); u16(0); u32(0); u32(localHeaderOffset.toLong) + out.write(nameBytes) + val cdSize = out.size() - cdOffset + // End of central directory (sig 0x06054b50): 1 entry, directory size and offset. + u32(0x06054b50L); u16(0); u16(0); u16(1); u16(1) + u32(cdSize.toLong); u32(cdOffset.toLong); u16(0) + val fos = new FileOutputStream(file) + try { + fos.write(out.toByteArray) + } finally { + fos.close() + } } /** Write a 7z archive, used to verify the `.7z` archive path. */ @@ -157,11 +208,20 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { /** Drains every entry into `(name, decodedText)` pairs through `SupportsArchiveFormat`. */ private def collect(file: File): Seq[(String, String)] = - SupportsArchiveFormat.readArchiveEntries(new Path(file.toURI), new Configuration()) { + SupportsArchiveFormat.readArchiveEntries( + new Path(file.toURI), new Configuration(), archivePathFilter = None) { (entry, in) => Iterator.single((entry.getName, new String(readAll(in), StandardCharsets.UTF_8))) }.toList + /** Drains every entry through `SupportsArchiveFormat` under an `archivePathFilter` glob. */ + private def collectFiltered(file: File, glob: String): Seq[String] = + SupportsArchiveFormat.readArchiveEntries( + new Path(file.toURI), new Configuration(), + archivePathFilter = Some(new GlobPattern(glob))) { (entry, _) => + Iterator.single(entry.getName) + }.toList + // ----- isArchivePath ------------------------------------------------------ test("isArchivePath: positive cases") { @@ -256,7 +316,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { // HadoopFSUtils.shouldFilterOutPathName still apply -- mirroring a loose-file listing with // the ignoredPathSegmentRegex option set to the same regex. val entries = SupportsArchiveFormat.readArchiveEntries( - new Path(tar.toURI), new Configuration(), Pattern.compile("(?!)")) { (entry, in) => + new Path(tar.toURI), new Configuration(), Pattern.compile("(?!)"), + archivePathFilter = None) { (entry, in) => Iterator.single((entry.getName, new String(readAll(in), StandardCharsets.UTF_8))) }.toList assert(entries == Seq("_SUCCESS" -> "marker", "real.csv" -> "kept")) @@ -271,7 +332,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { val opened = ArrayBuffer[String]() // parseEntry yields a single element without reading the stream, so each invocation maps to // exactly one consumed output element -- letting us observe when the next entry is opened. - val it = SupportsArchiveFormat.readArchiveEntries(new Path(tar.toURI), new Configuration()) { + val it = SupportsArchiveFormat.readArchiveEntries( + new Path(tar.toURI), new Configuration(), archivePathFilter = None) { (entry, _) => opened += entry.getName Iterator.single(entry.getName) @@ -299,7 +361,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { writeTar(tar, Seq(textEntry("a.csv", "a"), textEntry("b.csv", "b"))) val seen = ArrayBuffer[String]() - val it = SupportsArchiveFormat.readArchiveEntries(new Path(tar.toURI), new Configuration()) { + val it = SupportsArchiveFormat.readArchiveEntries( + new Path(tar.toURI), new Configuration(), archivePathFilter = None) { (entry, in) => val body = new String(readAll(in), StandardCharsets.UTF_8) in.close() // must NOT close the underlying archive @@ -316,7 +379,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { val tar = new File(dir, "closeable.tar") writeTar(tar, Seq(textEntry("a.csv", "a"), textEntry("b.csv", "b"))) - val it = SupportsArchiveFormat.readArchiveEntries(new Path(tar.toURI), new Configuration()) { + val it = SupportsArchiveFormat.readArchiveEntries( + new Path(tar.toURI), new Configuration(), archivePathFilter = None) { (entry, _) => Iterator.single(entry.getName) } assert(it.hasNext) @@ -345,7 +409,7 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { TaskContext.setTaskContext(ctx) try { val it = SupportsArchiveFormat.readArchiveEntries( - new Path(tar.toURI), new Configuration()) { + new Path(tar.toURI), new Configuration(), archivePathFilter = None) { (entry, _) => Iterator.single(entry.getName) } assert(it.hasNext) @@ -359,8 +423,6 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { } // ----- zip ---------------------------------------------------------------- - // The streaming engine is shared with tar (only stream-opening differs), so these cases focus on - // the `.zip` dispatch and the `ZipArchiveInputStream` container behaving like the tar path. test("readArchiveEntries: empty zip yields empty iterator") { withTempDir { dir => @@ -416,7 +478,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { writeZip(zip, Seq(textEntry("a.csv", "a"), textEntry("b.csv", "b"), textEntry("c.csv", "c"))) val opened = ArrayBuffer[String]() - val it = SupportsArchiveFormat.readArchiveEntries(new Path(zip.toURI), new Configuration()) { + val it = SupportsArchiveFormat.readArchiveEntries( + new Path(zip.toURI), new Configuration(), archivePathFilter = None) { (entry, _) => opened += entry.getName Iterator.single(entry.getName) @@ -441,7 +504,8 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { writeZip(zip, Seq(textEntry("a.csv", "a"), textEntry("b.csv", "b"))) val seen = ArrayBuffer[String]() - val it = SupportsArchiveFormat.readArchiveEntries(new Path(zip.toURI), new Configuration()) { + val it = SupportsArchiveFormat.readArchiveEntries( + new Path(zip.toURI), new Configuration(), archivePathFilter = None) { (entry, in) => val body = new String(readAll(in), StandardCharsets.UTF_8) in.close() // must NOT close the underlying archive @@ -453,15 +517,56 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { } } - test("readArchiveEntries: a non-streamable zip entry fails loudly, not with garbled bytes") { + test("readArchiveEntries: a STORED zip entry sized only by a data descriptor reads correctly") { withTempDir { dir => val zip = new File(dir, "stored-dd.zip") writeStoredEntryWithDataDescriptor(zip, "a.csv", "hello") - // A stored entry sized only by a trailing data descriptor is the documented non-streamable - // case: ZipArchiveInputStream throws on read instead of returning truncated/garbled bytes. + assert(collect(zip) == Seq("a.csv" -> "hello")) + } + } + + test("readArchiveEntries: corrupt zip bytes fail loudly") { + withTempDir { dir => + val zip = new File(dir, "corrupt.zip") + Files.write(zip.toPath, + "this is not a valid zip archive, just some random bytes" + .getBytes(StandardCharsets.UTF_8)) + // Not a ZIP (no end-of-central-directory signature), so ZipFile construction fails fast + // rather than reporting an empty archive. The ZipException surfaces wrapped in IOException. val ex = intercept[java.io.IOException](collect(zip)) - assert(ex.getMessage != null && ex.getMessage.contains("data descriptor"), - s"expected a clear unsupported-feature error, got $ex") + assert(ex.getCause.isInstanceOf[java.util.zip.ZipException]) + assert(ex.getCause.getMessage.contains("Archive is not a ZIP archive")) + } + } + + test("readArchiveEntries: zip with duplicate entry names yields both entries") { + withTempDir { dir => + val zip = new File(dir, "dupes.zip") + writeZip(zip, Seq(textEntry("a.csv", "first"), textEntry("a.csv", "second"))) + // ZipFile reads every central-directory record, so both duplicate-named entries surface. + assert(collect(zip) == Seq("a.csv" -> "first", "a.csv" -> "second")) + } + } + + test("readArchiveEntries: encrypted zip entry fails with a clear unsupported error") { + withTempDir { dir => + val zip = new File(dir, "encrypted.zip") + writeEncryptedEntry(zip, "a.csv", "hello") + val ex = intercept[SparkRuntimeException](collect(zip)) + checkError( + exception = ex, + condition = "CANNOT_READ_ZIP_ENTRY", + parameters = Map("entry" -> "a.csv", "path" -> new Path(zip.toURI).toString)) + } + } + + test("readArchiveEntries: an encrypted zip entry that is skipped does not throw") { + withTempDir { dir => + val zip = new File(dir, "encrypted-dotfile.zip") + // The encrypted entry has a dotfile name the engine filters out, so it is never read; its + // readability must not be checked (else it would throw CANNOT_READ_ZIP_ENTRY on a skip). + writeEncryptedEntry(zip, "._skipped.csv", "secret") + assert(collect(zip).isEmpty) } } @@ -518,4 +623,60 @@ class SupportsArchiveFormatSuite extends SparkFunSuite { assert(collect(sevenZ) == Seq("real.csv" -> "kept")) } } + + // ----- archivePathFilter --------------------------------------------------- + + test("readArchiveEntries: archivePathFilter keeps only entries matching the glob") { + withTempDir { dir => + val tar = new File(dir, "filter.tar") + writeTar(tar, Seq( + textEntry("sub/a.csv", "a"), + textEntry("sub/b.csv", "b"), + textEntry("other/c.csv", "c"))) + // The glob matches the entry's full path, so `sub/*` selects only the `sub/` entries. + assert(collectFiltered(tar, "sub/*") == Seq("sub/a.csv", "sub/b.csv")) + } + } + + test("readArchiveEntries: archivePathFilter with a `*` glob crosses directory boundaries") { + withTempDir { dir => + val tar = new File(dir, "filter-ext.tar") + writeTar(tar, Seq( + textEntry("top.csv", "t"), + textEntry("sub/nested.csv", "n"), + textEntry("keep.json", "j"))) + assert(collectFiltered(tar, "*.csv") == Seq("top.csv", "sub/nested.csv")) + } + } + + test("readArchiveEntries: archivePathFilter matching nothing yields an empty iterator") { + withTempDir { dir => + val tar = new File(dir, "filter-none.tar") + writeTar(tar, Seq(textEntry("a.csv", "a"), textEntry("b.csv", "b"))) + assert(collectFiltered(tar, "nomatch/*").isEmpty) + } + } + + test("readArchiveEntries: archivePathFilter applies on top of hidden-entry filtering") { + withTempDir { dir => + val tar = new File(dir, "filter-hidden.tar") + writeTar(tar, Seq( + textEntry("data/real.csv", "kept"), + textEntry("data/_SUCCESS", "marker"))) // matches the glob but hidden by the default regex + assert(collectFiltered(tar, "data/*") == Seq("data/real.csv")) + } + } + + test("archivePathFilter: an invalid glob is rejected with a clear error") { + val ex = intercept[IllegalArgumentException]( + FileSourceOptions.compileArchivePathFilter("[")) + assert(ex.getMessage.contains("archivePathFilter")) + } + + test("archivePathFilter: an empty value disables the filter rather than matching nothing") { + val options = new FileSourceOptions( + Map(FileSourceOptions.ARCHIVE_PATH_FILTER -> "")) + assert(options.archivePathFilter.isEmpty) + assert(options.archivePathFilterPattern.isEmpty) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TableLocationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TableLocationSuite.scala index dacb745f49ae7..5aa165d227873 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TableLocationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TableLocationSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources import org.apache.hadoop.fs.Path -import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.{AnalysisException, Row} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -76,4 +76,25 @@ class TableLocationSuite extends SharedSparkSession { } } } + + test("SPARK-56558: CTAS IF NOT EXISTS should be with non-existent or empty location") { + withSQLConf(SQLConf.ALLOW_NON_EMPTY_LOCATION_IN_CTAS.key -> "false") { + withTempDir { dir => + val tempLocation = dir.getCanonicalPath + withTable("ctas1", "ctas2") { + sql(s"CREATE TABLE ctas1 USING parquet LOCATION '$tempLocation/ctas1' " + + "AS SELECT 1 AS ID") + // Table ctas2 does not exist in the catalog, so IF NOT EXISTS must not skip the + // non-empty location check and overwrite the data of table ctas1. + val m = intercept[AnalysisException] { + sql(s"CREATE TABLE IF NOT EXISTS ctas2 USING parquet LOCATION '$tempLocation/ctas1' " + + "AS SELECT 2 AS ID") + }.getMessage + assert(m.contains("CREATE-TABLE-AS-SELECT cannot create table with location to a " + + "non-empty directory")) + checkAnswer(spark.table("ctas1"), Row(1)) + } + } + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TarArchiveTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TarArchiveTestUtils.scala index 7f1d48e8143b2..a31c2b187eac3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TarArchiveTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TarArchiveTestUtils.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.execution.datasources -import java.io.{File, FileOutputStream, OutputStream} +import java.io.{ByteArrayOutputStream, File, FileOutputStream, OutputStream} import java.nio.charset.StandardCharsets import java.nio.file.Files import java.util.Locale @@ -62,4 +62,34 @@ trait TarArchiveTestUtils { protected def writeCorruptArchive(dest: File): Unit = Files.write(dest.toPath, "this is not a valid gzip-compressed tar archive" .getBytes(StandardCharsets.UTF_8)) + + protected def supportsMidAdvanceFailure: Boolean = true + + /** + * Writes a plain `.tar` (uncompressed, so the first entry survives truncation) with `firstEntry` + * intact followed by a second entry whose header is cut short, so `getNextEntry` throws while + * advancing to it. + */ + protected def writeArchiveFailingAfterFirstEntry( + dest: File, firstEntry: (String, Array[Byte])): Unit = { + require(firstEntry._1.getBytes(StandardCharsets.UTF_8).length <= 100, + "the first entry's name must fit the 100-byte ustar name field so its header is a single " + + s"512-byte block, otherwise the truncation offset below shifts; got '${firstEntry._1}'") + val buf = new ByteArrayOutputStream() + val out = new TarArchiveOutputStream(buf) + Seq(firstEntry, ("part-later.bin", ("x" * 4096).getBytes(StandardCharsets.UTF_8))).foreach { + case (entryName, bytes) => + val entry = new TarArchiveEntry(entryName) + entry.setSize(bytes.length.toLong) + out.putArchiveEntry(entry) + out.write(bytes) + out.closeArchiveEntry() + } + out.finish() + out.close() + // Keep the first entry (512-byte header + block-aligned body) plus half of the second entry's + // 512-byte header, so `getNextEntry` hits a partial header while advancing. + val firstBlocks = 512 + ((firstEntry._2.length + 511) / 512) * 512 + Files.write(dest.toPath, buf.toByteArray.take(firstBlocks + 256)) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TextArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TextArchiveReadBase.scala index 1204c5a71b047..03c8b95d23c84 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TextArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/TextArchiveReadBase.scala @@ -155,6 +155,35 @@ trait TextArchiveReadBase extends QueryTest with SharedSparkSession { } } + test("_metadata exposes the parent archive file's values, identical for every row") { + archiveExtensions.foreach { ext => + withArchiveFile(ext) { archive => + writeArchive(archive, Seq( + "a.txt" -> textBytes("l1\nl2\n"), "b.txt" -> textBytes("l3\nl4\n"))) + val rows = read(archive.getCanonicalPath) + .select("_metadata.file_path", "_metadata.file_name", "_metadata.file_size", + "_metadata.file_block_start", "_metadata.file_block_length", + "_metadata.file_modification_time") + .collect() + assert(rows.length == 4) + val fileSize = archive.length() + rows.foreach { r => + assert(r.getString(0).endsWith(archive.getName) && !r.getString(0).contains("a.txt"), + s"file_path should be the archive file, got ${r.getString(0)}") + assert(r.getString(1) == archive.getName, s"file_name mismatch: ${r.getString(1)}") + assert(r.getLong(2) == fileSize, s"file_size mismatch: ${r.getLong(2)} != $fileSize") + assert(r.getLong(3) == 0L, s"file_block_start should be 0, got ${r.getLong(3)}") + assert(r.getLong(4) == fileSize, + s"file_block_length should be the archive size, got ${r.getLong(4)}") + assert(r.getAs[java.sql.Timestamp](5).getTime == archive.lastModified(), + "file_modification_time should be the archive's mtime") + } + assert(rows.map(_.toSeq).distinct.length == 1, + "every row must carry the same parent-archive metadata") + } + } + } + Seq(true, false).foreach { ignoreCorrupt => test(s"ignoreCorruptFiles=$ignoreCorrupt controls whether a corrupt archive is skipped") { withArchiveFile(corruptArchiveExtension) { archive => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala index bac194255581d..1f71e77bc08a2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala @@ -161,6 +161,24 @@ class V1WriteCommandSuite extends SharedSparkSession with V1WriteCommandSuiteBas } } + test("v1 write with CHAR/VARCHAR partition columns applies empty2null") { + withSQLConf(SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true") { + // The partition values must vary, otherwise the sort on a foldable key is pruned and + // there is no output ordering left to match. + Seq("CHAR(5)", "VARCHAR(5)").foreach { typ => + withPlannedWrite { enabled => + withTable("t") { + sql(s"CREATE TABLE t(i INT) USING PARQUET PARTITIONED BY (p $typ)") + executeAndCheckOrdering( + hasLogicalSort = enabled, orderingMatched = enabled, hasEmpty2Null = enabled) { + sql("INSERT INTO t SELECT i, k FROM t0") + } + } + } + } + } + } + test("v1 write with partition, bucketed and sort columns") { withPlannedWrite { enabled => withTable("t") { @@ -438,4 +456,61 @@ class V1WriteCommandSuite extends SharedSparkSession with V1WriteCommandSuiteBas } } } + + test("SPARK-58444: planned write should not enable concurrent writer when ordering " + + "already matched") { + // The concurrent output writer keeps one open writer per dynamic partition and falls back + // to the sort-based sequential writer once the number of open writers reaches + // `maxConcurrentOutputFileWriters`. When the input is already sorted by the required + // ordering, FileFormatWriter should NOT enable the concurrent writer at all, so that this + // wasteful fall-back never happens. This test pins that behavior via the fall-back log. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.PLANNED_WRITE_ENABLED.key -> "true", + SQLConf.MAX_CONCURRENT_OUTPUT_FILE_WRITERS.key -> "2") { + withTable("t") { + sql("CREATE TABLE t(i INT, k STRING) USING PARQUET PARTITIONED BY (j INT)") + + // t0 has 5 distinct values of `j` (i % 5), which is greater than the + // maxConcurrentOutputFileWriters threshold (2), so a single-task concurrent write over + // all partitions would trigger the fall-back. Use an int partition column to avoid the + // empty2null projection that a string partition column would add. + val fallbackMsg = "Fall back from concurrent writers" + val loggerName = classOf[DynamicPartitionDataConcurrentWriter].getName + val expected = spark.table("t0").select($"i", $"k", $"j") + + // Case 1: input already sorted by the dynamic partition column, collapsed into a single + // task. The output ordering matches the required ordering, so the concurrent writer must + // be disabled -> no fall-back log. + val matchedAppender = new LogAppender("ordering matched, no concurrent writer") + withLogAppender(matchedAppender, Seq(loggerName)) { + expected.repartition(1).sortWithinPartitions("j") + .write.mode("overwrite").insertInto("t") + } + assert(FileFormatWriter.outputOrderingMatched, + "Expected the output ordering to match the required ordering.") + assert(!matchedAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains(fallbackMsg)), + "Concurrent writer should be disabled when ordering already matches, " + + "so no fall-back to the sort-based writer should happen.") + checkAnswer(spark.table("t"), expected) + + // Case 2 (control): input NOT sorted. Ordering does not match, so the concurrent writer + // stays enabled and falls back once open writers exceed the threshold. This proves the + // log assertion above is actually discriminating and not vacuously true. + val unmatchedAppender = new LogAppender("ordering not matched, concurrent writer") + withLogAppender(unmatchedAppender, Seq(loggerName)) { + expected.repartition(1) + .write.mode("overwrite").insertInto("t") + } + assert(!FileFormatWriter.outputOrderingMatched, + "Expected the output ordering NOT to match the required ordering.") + assert(unmatchedAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains(fallbackMsg)), + "Concurrent writer should fall back to the sort-based writer when ordering " + + "does not match.") + checkAnswer(spark.table("t"), expected) + } + } + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/XMLArchiveReadBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/XMLArchiveReadBase.scala index c5e956d4cd38d..ffa5484b3194e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/XMLArchiveReadBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/XMLArchiveReadBase.scala @@ -21,6 +21,7 @@ import java.io.File import java.nio.charset.StandardCharsets import java.nio.file.Files +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StringType /** @@ -132,6 +133,27 @@ trait XMLArchiveReadBase extends ArchiveReadSuiteBase { extraOptions = Map("multiLine" -> "true"), schema = corruptSchema) } + + if (supportsMidAdvanceFailure) { + test("XML: multiLine inference keeps records read before a mid-advance failure " + + "(ignoreCorruptFiles)") { + // Entry 0 is read, then advancing to a later entry throws (not at open). A whole-archive drop + // would lose entry 0's `extra`; aborting the traversal would lose the sibling file's `later`. + val opts = Map("multiLine" -> "true") + withArchiveFile() { archive => + writeArchiveFailingAfterFirstEntry(archive, entryName(0) -> + xmlBytes("<rows><row><id>1</id><name>Alice</name><extra>9</extra></row></rows>")) + Files.write(new File(archive.getParentFile, s"later.$fileExtension").toPath, + xmlBytes("<rows><row><id>2</id><name>Bob</name><later>7</later></row></rows>")) + withSQLConf(SQLConf.IGNORE_CORRUPT_FILES.key -> "true") { + val schema = inferredSchema(Seq(archive.getParentFile.getCanonicalPath), opts) + assert(schema.fieldNames.toSet == Set("id", "name", "extra", "later"), + "expected `extra` (pre-failure entry) and `later` (sibling file) in the inferred " + + s"schema after the mid-advance skip, got $schema") + } + } + } + } } class XMLTarArchiveReadSuite diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ZipArchiveTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ZipArchiveTestUtils.scala index cf9e42262a7b2..51282d1e856db 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ZipArchiveTestUtils.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/ZipArchiveTestUtils.scala @@ -57,4 +57,13 @@ trait ZipArchiveTestUtils { protected def writeCorruptArchive(dest: File): Unit = Files.write(dest.toPath, "this is not a valid zip archive, just some random bytes" .getBytes(StandardCharsets.UTF_8)) + + // Zip can't throw while *advancing*: a cut header reads as clean EOF, and a cut body only throws + // when the entry's bytes are read -- which a lazy reader (JSON) defers downstream. Tar is used + // for the mid-advance regression test instead. + protected def supportsMidAdvanceFailure: Boolean = false + + protected def writeArchiveFailingAfterFirstEntry( + dest: File, firstEntry: (String, Array[Byte])): Unit = + throw new UnsupportedOperationException("zip cannot force a throw while advancing to an entry") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala index 71eb34134920b..3aefe25f6a08c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/csv/CSVSuite.scala @@ -48,6 +48,7 @@ import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf} import org.apache.spark.sql.internal.SQLConf.BinaryOutputStyle import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest abstract class CSVSuite extends SharedSparkSession @@ -4104,6 +4105,7 @@ class CSVv1Suite extends CSVSuite { } } +@ExtendedSQLTest class CSVv2Suite extends CSVSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/ExplodeEmbeddedArrayJsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/ExplodeEmbeddedArrayJsonSuite.scala index 075f61042d909..a21d01af1d311 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/ExplodeEmbeddedArrayJsonSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/ExplodeEmbeddedArrayJsonSuite.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.{AnalysisException, QueryTest, Row} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.StructType +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils class EmbeddedArraySplitterSuite extends SparkFunSuite { @@ -388,6 +389,7 @@ class ExplodeEmbeddedArrayJsonV1Suite extends ExplodeEmbeddedArrayJsonSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "json") } +@ExtendedSQLTest class ExplodeEmbeddedArrayJsonV2Suite extends ExplodeEmbeddedArrayJsonSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala index 308b41d9ec776..4cd6783e1df60 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/json/JsonSuite.scala @@ -53,6 +53,7 @@ import org.apache.spark.sql.types._ import org.apache.spark.sql.types.StructType.fromDDL import org.apache.spark.sql.types.TestUDT.{MyDenseVector, MyDenseVectorUDT} import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils @@ -4356,6 +4357,7 @@ class JsonV1Suite extends JsonSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "json") } +@ExtendedSQLTest class JsonV2Suite extends JsonSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilterSuite.scala index 3168075cc9d9a..a7eb37e95278e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcFilterSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.datasources.orc import java.math.MathContext import java.nio.charset.StandardCharsets import java.sql.{Date, Timestamp} -import java.time.{Duration, LocalDateTime, LocalTime, Period} +import java.time.{Duration, LocalDateTime, LocalTime, Period, ZoneOffset} import scala.jdk.CollectionConverters._ @@ -32,6 +32,7 @@ import org.apache.spark.sql.{AnalysisException, Column, DataFrame, Row} import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.planning.PhysicalOperation +import org.apache.spark.sql.catalyst.util.TimestampNanosTestUtils.foreachNanosPrecision import org.apache.spark.sql.execution.datasources.v2.ExtractV2Scan import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan import org.apache.spark.sql.functions.col @@ -398,6 +399,67 @@ class OrcFilterSuite extends OrcTest with SharedSparkSession { } } + test("SPARK-57823: filter pushdown - nanosecond timestamp") { + // Wall clocks carry sub-microsecond digits. The literal is explicitly nanos-typed, so even a + // microsecond-aligned value would still exercise the nanos pushdown path; the extra digits + // exercise value preservation and comparison beyond microsecond precision. + val wallClocks = Seq( + LocalDateTime.of(1000, 1, 1, 1, 2, 3, 456789123), + LocalDateTime.of(1582, 10, 1, 0, 11, 22, 456789123), + LocalDateTime.of(1900, 1, 1, 23, 59, 59, 456789123), + LocalDateTime.of(2020, 5, 25, 10, 11, 12, 456789123)) + + // Builds a nanos-typed literal so the comparison against the nanos column needs no type + // coercion: NTZ uses the wall clock as a LocalDateTime, LTZ as the same instant at UTC. + def nanosLiteral(nanosType: DataType, wallClock: LocalDateTime): Expression = { + val external: Any = nanosType match { + case _: TimestampNTZNanosType => wallClock + case _: TimestampLTZNanosType => wallClock.toInstant(ZoneOffset.UTC) + } + Literal.create(external, nanosType) + } + + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + foreachNanosPrecision { precision => + Seq(TimestampNTZNanosType(precision), TimestampLTZNanosType(precision)).foreach { + nanosType => + val ts = wallClocks.map(nanosLiteral(nanosType, _)) + withTempPath { dir => + val path = dir.getCanonicalPath + nanosTimestampDf(nanosType, wallClocks).write.orc(path) + // Read back without an explicit schema: the nanos type is recovered from the ORC + // catalyst-type attribute, so the pushed-down column is nanos-typed. + readFile(path) { implicit df => + assert(df("ts").expr.dataType === nanosType) + + checkFilterPredicate($"ts".isNull, PredicateLeaf.Operator.IS_NULL) + + checkFilterPredicate($"ts" === ts(0), PredicateLeaf.Operator.EQUALS) + checkFilterPredicate($"ts" <=> ts(0), PredicateLeaf.Operator.NULL_SAFE_EQUALS) + + checkFilterPredicate($"ts" < ts(1), PredicateLeaf.Operator.LESS_THAN) + checkFilterPredicate($"ts" > ts(2), PredicateLeaf.Operator.LESS_THAN_EQUALS) + checkFilterPredicate($"ts" <= ts(0), PredicateLeaf.Operator.LESS_THAN_EQUALS) + checkFilterPredicate($"ts" >= ts(3), PredicateLeaf.Operator.LESS_THAN) + + checkFilterPredicate(ts(0) === $"ts", PredicateLeaf.Operator.EQUALS) + checkFilterPredicate(ts(0) <=> $"ts", PredicateLeaf.Operator.NULL_SAFE_EQUALS) + checkFilterPredicate(ts(1) > $"ts", PredicateLeaf.Operator.LESS_THAN) + checkFilterPredicate(ts(2) < $"ts", PredicateLeaf.Operator.LESS_THAN_EQUALS) + checkFilterPredicate(ts(0) >= $"ts", PredicateLeaf.Operator.LESS_THAN_EQUALS) + checkFilterPredicate(ts(3) <= $"ts", PredicateLeaf.Operator.LESS_THAN) + + // In covers the per-value castLiteralValue path (values.map(...)) in + // buildLeafSearchArgument, exercising the nanos literal cast for every element. + checkFilterPredicate( + In($"ts", Seq(ts(0), ts(2))), PredicateLeaf.Operator.IN) + } + } + } + } + } + } + test("filter pushdown - combinations with logical operators") { withOrcDataFrame((1 to 4).map(i => Tuple1(Option(i)))) { implicit df => checkFilterPredicate( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala index 909ebf5daf398..de72dd81e7e4f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcQuerySuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.datasources.orc import java.io.File import java.nio.charset.StandardCharsets import java.sql.Timestamp -import java.time.LocalDateTime +import java.time.{LocalDateTime, ZoneOffset} import org.apache.hadoop.conf.Configuration import org.apache.hadoop.fs.{FileSystem, Path} @@ -35,6 +35,7 @@ import org.apache.orc.mapreduce.OrcInputFormat import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.sql._ import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.catalyst.util.DateTimeTestUtils import org.apache.spark.sql.catalyst.util.TimestampNanosTestUtils.foreachNanosPrecision import org.apache.spark.sql.execution.FileSourceScanExec @@ -43,6 +44,7 @@ import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils import org.apache.spark.util.collection.Utils.createArray @@ -1029,6 +1031,101 @@ abstract class OrcQuerySuite extends OrcQueryTest with SharedSparkSession { } } + test("SPARK-57823: ORC predicate pushdown returns correct results for nanos timestamps") { + // One wall clock per second across a minute so each row lands in a distinct stripe below, all + // carrying sub-microsecond digits so the nanosecond fraction participates in the comparison. + val numRows = 60 + val wallClocks = (0 until numRows).map { s => + LocalDateTime.of(2020, 5, 25, 10, 0, s, 123456789) + } + def checkNanosPushdown(): Unit = { + foreachNanosPrecision { precision => + Seq(TimestampNTZNanosType(precision), TimestampLTZNanosType(precision)).foreach { + nanosType => + val inputDf = nanosTimestampDf(nanosType, wallClocks) + // A boundary literal (row 30) expressed as the nanos external type, so a `< threshold` + // predicate keeps exactly the first 30 rows. + val boundary = nanosType match { + case _: TimestampNTZNanosType => Literal.create(wallClocks(30), nanosType) + case _: TimestampLTZNanosType => + Literal.create(wallClocks(30).toInstant(ZoneOffset.UTC), nanosType) + } + // Rebuilt against each DataFrame's own `ts` attribute so the filter resolves cleanly. + def keepFirstHalf(df: DataFrame): Column = df("ts") < Column(boundary) + withTempPath { dir => + val path = dir.getCanonicalPath + // Repartition so the rows are spread over several ORC files/stripes, giving the + // pushed-down search argument something to skip. + inputDf.repartition(numRows).write.orc(path) + Seq(true, false).foreach { vectorized => + withSQLConf( + SQLConf.ORC_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.ORC_FILTER_PUSHDOWN_ENABLED.key -> "true") { + val read = spark.read.schema(new StructType().add("ts", nanosType)).orc(path) + + // Results are correct: the pushdown must not drop or corrupt matching rows. + checkAnswer( + read.where(keepFirstHalf(read)), inputDf.where(keepFirstHalf(inputDf))) + + // Pushdown actually skips data: with the Spark-side filter removed, fewer than + // all rows survive, proving the ORC search argument pruned stripes. + assert(stripSparkFilter(read.where(keepFirstHalf(read))).count() < numRows) + } + } + } + } + } + } + + // The LTZ column stores the UTC instant; the JVM default zone enters only at comparison time, + // when ORC evaluates the pushed predicate against the non-UTC timestamp statistics. So a + // literal built in the wrong frame would only diverge from the stripe stats outside UTC. Run + // under zones with both offset signs so timezone handling is exercised, not just UTC where such + // a bug is masked. + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + Seq(DateTimeTestUtils.UTC, DateTimeTestUtils.LA, DateTimeTestUtils.JST).foreach { zoneId => + DateTimeTestUtils.withDefaultTimeZone(zoneId)(checkNanosPushdown()) + } + } + } + + test("SPARK-57823: LTZ nanos pushdown stays correct across a DST spring-forward gap") { + // Timestamp.valueOf maps a wall clock inside the JVM zone's spring-forward gap one hour ahead, + // so for instants whose UTC wall clock lands in that gap the literal frame is non-monotonic. + // Under America/Los_Angeles the gap is 2020-03-08 02:00..03:00, so lay data one minute apart + // straddling it (UTC wall clocks 01:30..03:29) with a boundary at 02:45. checkAnswer must hold: + // in the gap the pushed predicate may only turn conservative (skip fewer stripes), never prune + // a stripe that holds matching rows. + val numRows = 120 + val base = LocalDateTime.of(2020, 3, 8, 1, 30, 0, 123456789) + val wallClocks = (0 until numRows).map(m => base.plusMinutes(m)) + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + DateTimeTestUtils.withDefaultTimeZone(DateTimeTestUtils.LA) { + foreachNanosPrecision { precision => + val nanosType = TimestampLTZNanosType(precision) + val inputDf = nanosTimestampDf(nanosType, wallClocks) + val boundary = Literal.create( + LocalDateTime.of(2020, 3, 8, 2, 45, 0, 0).toInstant(ZoneOffset.UTC), nanosType) + def keepBeforeGapBoundary(df: DataFrame): Column = df("ts") < Column(boundary) + withTempPath { dir => + val path = dir.getCanonicalPath + inputDf.repartition(numRows).write.orc(path) + Seq(true, false).foreach { vectorized => + withSQLConf( + SQLConf.ORC_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.ORC_FILTER_PUSHDOWN_ENABLED.key -> "true") { + val read = spark.read.schema(new StructType().add("ts", nanosType)).orc(path) + checkAnswer( + read.where(keepBeforeGapBoundary(read)), + inputDf.where(keepBeforeGapBoundary(inputDf))) + } + } + } + } + } + } + } + test("SPARK-57455: ORC round-trips nanos timestamps in nested/complex types") { withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { val wallClocks = Seq(LocalDateTime.of(1970, 1, 1, 0, 20, 34, 567890123)) @@ -1170,6 +1267,7 @@ class OrcV1QuerySuite extends OrcQuerySuite { .set(SQLConf.USE_V1_SOURCE_LIST, "orc") } +@ExtendedSQLTest class OrcV2QuerySuite extends OrcQuerySuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcSourceSuite.scala index 155fd592cf8ac..58b98fe053aea 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/orc/OrcSourceSuite.scala @@ -41,6 +41,7 @@ import org.apache.spark.sql.execution.datasources.orc.OrcCompressionCodec._ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils case class OrcData(intField: Int, stringField: String) @@ -1103,6 +1104,7 @@ class OrcSourceV1Suite extends OrcSourceSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "orc") } +@ExtendedSQLTest class OrcSourceV2Suite extends OrcSourceSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormatSuite.scala index b0ad5ca8cab71..01c8507032ceb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormatSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFileFormatSuite.scala @@ -30,6 +30,7 @@ import org.apache.spark.sql.execution.datasources.CommonFileDataSourceSuite import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.HadoopFSUtils abstract class ParquetFileFormatSuite @@ -269,6 +270,7 @@ class ParquetFileFormatV1Suite extends ParquetFileFormatSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "parquet") } +@ExtendedSQLTest class ParquetFileFormatV2Suite extends ParquetFileFormatSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala index ae977e17755df..f023e8390f155 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetFilterSuite.scala @@ -22,7 +22,7 @@ import java.lang.{Double => JDouble, Float => JFloat, Long => JLong} import java.math.{BigDecimal => JBigDecimal} import java.nio.charset.StandardCharsets import java.sql.{Date, Timestamp} -import java.time.{Duration, LocalDate, LocalDateTime, LocalTime, Period, ZoneId} +import java.time.{Duration, Instant, LocalDate, LocalDateTime, LocalTime, Period, ZoneId} import java.util.HashSet import scala.reflect.ClassTag @@ -34,7 +34,7 @@ import org.apache.parquet.filter2.predicate.FilterApi._ import org.apache.parquet.filter2.predicate.Operators.{Column => _, Eq, Gt, GtEq, In => FilterIn, Lt, LtEq, NotEq, UserDefinedByInstance} import org.apache.parquet.hadoop.{ParquetFileReader, ParquetInputFormat, ParquetOutputFormat} import org.apache.parquet.hadoop.util.HadoopInputFile -import org.apache.parquet.schema.MessageType +import org.apache.parquet.schema.{MessageType, MessageTypeParser} import org.apache.spark.{SparkConf, SparkException, SparkRuntimeException} import org.apache.spark.sql._ @@ -45,7 +45,7 @@ import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.util.RebaseDateTime.RebaseSpec import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.parseColumnPath import org.apache.spark.sql.execution.ExplainMode -import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, HadoopFsRelation, LogicalRelationWithTable, PushableColumnAndNestedColumn} +import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, HadoopFsRelation, LogicalRelationWithTable, PushableColumnAndNestedColumn, VariantMetadata} import org.apache.spark.sql.execution.datasources.v2.ExtractV2Scan import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.functions._ @@ -82,13 +82,15 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession { protected def createParquetFilters( schema: MessageType, caseSensitive: Option[Boolean] = None, - datetimeRebaseSpec: RebaseSpec = RebaseSpec(LegacyBehaviorPolicy.CORRECTED) + datetimeRebaseSpec: RebaseSpec = RebaseSpec(LegacyBehaviorPolicy.CORRECTED), + variantExtractionSchema: Option[StructType] = None ): ParquetFilters = new ParquetFilters(schema, conf.parquetFilterPushDownDate, conf.parquetFilterPushDownTimestamp, conf.parquetFilterPushDownDecimal, conf.parquetFilterPushDownStringPredicate, conf.parquetFilterPushDownInFilterThreshold, caseSensitive.getOrElse(conf.caseSensitiveAnalysis), - datetimeRebaseSpec) + datetimeRebaseSpec, + variantExtractionSchema = variantExtractionSchema) override def beforeEach(): Unit = { super.beforeEach() @@ -2275,6 +2277,556 @@ abstract class ParquetFilterSuite extends ParquetTest with SharedSparkSession { } } } + + test("SPARK-57822: filter pushdown - nanosecond timestamps") { + // The value that reaches `ParquetFilters.createFilter` for a nanosecond timestamp column is + // the externalized Java time value (Instant for LTZ, LocalDateTime for NTZ), NOT a + // TimestampNanosVal. Both convert to the signed INT64 epoch-nanoseconds that + // `TimestampNanosParquetOps` writes, so every operator must produce a `long`-column filter. + // Threshold 1 so a 2-element `In` exceeds it and exercises the `makeInPredicate` nanos arm + // (below the threshold, `In` expands to an `Or` of `Eq`, which is covered by the `Eq` arm). + withSQLConf( + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true", + SQLConf.PARQUET_FILTER_PUSHDOWN_INFILTERTHRESHOLD.key -> "1") { + Seq(7, 8, 9).foreach { p => + // NTZ column: filter values are LocalDateTime. Use a sub-microsecond fraction so the + // epoch-nanos encoding (not truncated micros) is exercised. + val ntzSchema = new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampNTZNanosType(p))) + val ntzFilters = createParquetFilters(ntzSchema) + val ntz = LocalDateTime.parse("2020-01-01T12:34:56.000000123") + val ntzHi = LocalDateTime.parse("2020-01-01T12:34:57.000000123") + assert(ntzFilters.createFilter(sources.IsNull("c")).exists(_.isInstanceOf[Eq[_]])) + assert(ntzFilters.createFilter(sources.IsNotNull("c")).exists(_.isInstanceOf[NotEq[_]])) + assert(ntzFilters.createFilter(sources.EqualTo("c", ntz)).exists(_.isInstanceOf[Eq[_]])) + assert(ntzFilters.createFilter(sources.EqualNullSafe("c", ntz)) + .exists(_.isInstanceOf[Eq[_]])) + assert(ntzFilters.createFilter(sources.Not(sources.EqualTo("c", ntz))) + .exists(_.isInstanceOf[NotEq[_]])) + assert(ntzFilters.createFilter(sources.LessThan("c", ntz)).exists(_.isInstanceOf[Lt[_]])) + assert(ntzFilters.createFilter(sources.LessThanOrEqual("c", ntz)) + .exists(_.isInstanceOf[LtEq[_]])) + assert(ntzFilters.createFilter(sources.GreaterThan("c", ntz)).exists(_.isInstanceOf[Gt[_]])) + assert(ntzFilters.createFilter(sources.GreaterThanOrEqual("c", ntz)) + .exists(_.isInstanceOf[GtEq[_]])) + assert(ntzFilters.createFilter(sources.In("c", Array[Any](ntz, ntzHi))) + .exists(_.isInstanceOf[FilterIn[_]])) + + // LTZ column: filter values are Instant. + val ltzSchema = new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampLTZNanosType(p))) + val ltzFilters = createParquetFilters(ltzSchema) + val ltz = Instant.parse("2020-01-01T12:34:56.000000123Z") + val ltzHi = Instant.parse("2020-01-01T12:34:57.000000123Z") + assert(ltzFilters.createFilter(sources.IsNull("c")).exists(_.isInstanceOf[Eq[_]])) + assert(ltzFilters.createFilter(sources.IsNotNull("c")).exists(_.isInstanceOf[NotEq[_]])) + assert(ltzFilters.createFilter(sources.EqualTo("c", ltz)).exists(_.isInstanceOf[Eq[_]])) + assert(ltzFilters.createFilter(sources.EqualNullSafe("c", ltz)) + .exists(_.isInstanceOf[Eq[_]])) + assert(ltzFilters.createFilter(sources.Not(sources.EqualTo("c", ltz))) + .exists(_.isInstanceOf[NotEq[_]])) + assert(ltzFilters.createFilter(sources.LessThan("c", ltz)).exists(_.isInstanceOf[Lt[_]])) + assert(ltzFilters.createFilter(sources.LessThanOrEqual("c", ltz)) + .exists(_.isInstanceOf[LtEq[_]])) + assert(ltzFilters.createFilter(sources.GreaterThan("c", ltz)).exists(_.isInstanceOf[Gt[_]])) + assert(ltzFilters.createFilter(sources.GreaterThanOrEqual("c", ltz)) + .exists(_.isInstanceOf[GtEq[_]])) + assert(ltzFilters.createFilter(sources.In("c", Array[Any](ltz, ltzHi))) + .exists(_.isInstanceOf[FilterIn[_]])) + } + } + } + + test("SPARK-57822: don't push down nanosecond timestamp filters that would overflow int64") { + // Values outside the int64 epoch-nanoseconds range (~1677-09-21 .. 2262-04-11) must NOT be + // pushed down: encoding them would throw, and a truncated encoding could silently mis-skip a + // row group. Falling back to a full scan is always correct (SPARK-46092 rationale). + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + val ntzSchema = new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampNTZNanosType(9))) + val ntzFilters = createParquetFilters(ntzSchema) + val ltzSchema = new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampLTZNanosType(9))) + val ltzFilters = createParquetFilters(ltzSchema) + + // Out of range: must not push down. + val ntzOverflow = LocalDateTime.parse("2300-01-01T00:00:00") + val ltzOverflow = Instant.parse("2300-01-01T00:00:00Z") + Seq( + sources.LessThan("c", ntzOverflow), + sources.LessThanOrEqual("c", ntzOverflow), + sources.GreaterThan("c", ntzOverflow), + sources.GreaterThanOrEqual("c", ntzOverflow), + sources.EqualTo("c", ntzOverflow), + sources.EqualNullSafe("c", ntzOverflow), + sources.Not(sources.EqualTo("c", ntzOverflow)), + sources.In("c", Array[Any](ntzOverflow)) + ).foreach { filter => + assert(ntzFilters.createFilter(filter).isEmpty, + s"Row group filter $filter shouldn't be pushed down.") + } + Seq( + sources.LessThan("c", ltzOverflow), + sources.LessThanOrEqual("c", ltzOverflow), + sources.GreaterThan("c", ltzOverflow), + sources.GreaterThanOrEqual("c", ltzOverflow), + sources.EqualTo("c", ltzOverflow), + sources.EqualNullSafe("c", ltzOverflow), + sources.Not(sources.EqualTo("c", ltzOverflow)), + sources.In("c", Array[Any](ltzOverflow)) + ).foreach { filter => + assert(ltzFilters.createFilter(filter).isEmpty, + s"Row group filter $filter shouldn't be pushed down.") + } + + // In range: must push down. + assert(ntzFilters.createFilter( + sources.LessThan("c", LocalDateTime.parse("2020-01-01T00:00:00"))).isDefined) + assert(ltzFilters.createFilter( + sources.LessThan("c", Instant.parse("2020-01-01T00:00:00Z"))).isDefined) + } + } + + test("SPARK-57822: an In list mixing in-range and out-of-range nanos values isn't pushed down") { + // A mixed In list (in-range head, out-of-range tail) must fall back to a full scan, not crash + // filter creation. The In arm converts *every* element - per-element `makeEq` under the + // pushdown threshold, `makeInPredicate` above it - so a head-only guard would let the + // out-of-range tail reach the encoder, which throws (SPARK-46092: never throw, full scan + // instead). Both branches are exercised by setting the threshold below and above the list + // size. The list must lead with an in-range value so the fix is what rejects it, not the old + // head check. + Seq("1", "10").foreach { threshold => + withSQLConf( + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true", + SQLConf.PARQUET_FILTER_PUSHDOWN_INFILTERTHRESHOLD.key -> threshold) { + val ntzFilters = createParquetFilters(new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampNTZNanosType(9)))) + val ltzFilters = createParquetFilters(new SparkToParquetSchemaConverter(conf) + .convert(new StructType().add("c", TimestampLTZNanosType(9)))) + val ntzIn = sources.In("c", Array[Any]( + LocalDateTime.parse("2020-01-01T00:00:00"), LocalDateTime.parse("2300-01-01T00:00:00"))) + val ltzIn = sources.In("c", Array[Any]( + Instant.parse("2020-01-01T00:00:00Z"), Instant.parse("2300-01-01T00:00:00Z"))) + // Must not throw, and must not push down (a single out-of-range element disqualifies the + // whole In). + assert(ntzFilters.createFilter(ntzIn).isEmpty, + s"Mixed-range NTZ In (threshold=$threshold) shouldn't be pushed down.") + assert(ltzFilters.createFilter(ltzIn).isEmpty, + s"Mixed-range LTZ In (threshold=$threshold) shouldn't be pushed down.") + // An all-in-range In of the same size still pushes down (the fix rejects only on a + // non-pushable element, it doesn't disable In pushdown wholesale). + assert(ntzFilters.createFilter(sources.In("c", Array[Any]( + LocalDateTime.parse("2020-01-01T00:00:00"), + LocalDateTime.parse("2021-01-01T00:00:00")))).isDefined) + assert(ltzFilters.createFilter(sources.In("c", Array[Any]( + Instant.parse("2020-01-01T00:00:00Z"), + Instant.parse("2021-01-01T00:00:00Z")))).isDefined) + } + } + } + + // ---------------------------------------------------------------------------------------------- + // Shredded-variant filter pushdown (SPARK-55817). + // + // PushVariantIntoScan rewrites variant_get(v, '$.a', 'bigint') > 999 into a struct-field access + // "v.`0`" > 999 where "0" carries VariantMetadata for path "$.a". ParquetFilters maps that + // logical path to the physical shredded leaf v.typed_value.a.typed_value and, for soundness, + // guards it so a row group is skipped only when the leaf cannot match AND every value for the + // path is provably in the leaf (see `makeShreddedFilter`). + // ---------------------------------------------------------------------------------------------- + + /** A variant-extraction StructField named by ordinal, carrying VariantMetadata for `path`. */ + private def variantField(name: String, dt: DataType, path: String): StructField = + StructField(name, dt, metadata = VariantMetadata(path, failOnError = true, "UTC").toMetadata) + + /** + * The variantExtractionSchema PushVariantIntoScan produces for a top-level variant column + * `colName` with the given extraction fields (each an ordinal-named field with VariantMetadata). + */ + private def variantExtractionSchema(colName: String, fields: StructField*): StructType = + StructType(Seq(StructField(colName, StructType(fields)))) + + test("shredded variant filter: single-level bigint resolves to leaf with residual guards") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + val filter = pf.createFilter(sources.GreaterThan("v.`0`", 999L)) + assert(filter.isDefined, "Expected shredded variant predicate to be created") + val s = filter.get.toString + // Guarded shape: or(gt(leaf), and(or(notEq(residual)...), eq(leaf, null))). + assert(s.contains("v.typed_value.a.typed_value"), + s"Expected leaf column v.typed_value.a.typed_value in $s") + assert(s.contains("v.typed_value.a.value"), s"Expected L1 residual guard in $s") + assert(s.contains("v.value"), s"Expected top-level residual guard in $s") + // The leaf-is-null arm must be present (this is what keeps the guard sound and effective). + assert(s.contains("eq(v.typed_value.a.typed_value, null)"), + s"Expected an isNull(leaf) guard arm in $s") + } + + test("shredded variant filter: without variantExtractionSchema the logical path is unknown") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val pf = createParquetFilters(MessageTypeParser.parseMessageType(parquetSchema)) + assert(pf.createFilter(sources.GreaterThan("v.`0`", 999L)).isEmpty) + } + + test("shredded variant filter: string leaf and all comparison operators") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group b { + | optional binary value; + | optional binary typed_value (STRING); + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", StringType, "$.b")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + Seq( + sources.EqualTo("v.`0`", "str"), + sources.LessThan("v.`0`", "str"), + sources.LessThanOrEqual("v.`0`", "str"), + sources.GreaterThan("v.`0`", "str"), + sources.GreaterThanOrEqual("v.`0`", "str")).foreach { f => + val filter = pf.createFilter(f) + assert(filter.isDefined, s"Expected $f to push down") + val s = filter.get.toString + assert(s.contains("v.typed_value.b.typed_value"), s"Expected leaf column in $s for $f") + assert(s.contains("v.typed_value.b.value") && s.contains("v.value"), + s"Expected residual guards in $s for $f") + } + } + + test("shredded variant filter: In pushes OR of guarded equalities") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + val filter = pf.createFilter(sources.In("v.`0`", Array[Any](1L, 2L, 3L))) + assert(filter.isDefined, "Expected In on shredded column to push down") + val s = filter.get.toString + assert(s.contains("v.typed_value.a.typed_value"), s"Expected leaf column in $s") + assert(s.contains("v.typed_value.a.value") && s.contains("v.value"), + s"Expected residual guards in $s") + } + + test("shredded variant filter: multi-level path resolves with a residual per level") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional group typed_value { + | optional group b { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a.b")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + val filter = pf.createFilter(sources.GreaterThan("v.`0`", 5L)) + assert(filter.isDefined, "Expected multi-level shredded predicate to be created") + val s = filter.get.toString + assert(s.contains("v.typed_value.a.typed_value.b.typed_value"), + s"Expected multi-level leaf column in $s") + // Three residual guards: L0, L1 (a), and leaf-level sibling (b). + assert(s.contains("v.value"), s"Expected L0 residual guard in $s") + assert(s.contains("v.typed_value.a.value"), s"Expected L1 residual guard in $s") + assert(s.contains("v.typed_value.a.typed_value.b.value"), + s"Expected leaf-level residual guard in $s") + } + + test("shredded variant filter: array-index path is rejected") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a[0]")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + assert(pf.createFilter(sources.GreaterThan("v.`0`", 999L)).isEmpty, + "Array-index paths must not resolve to a shredded leaf") + } + + test("shredded variant filter: synthetic fields (placeholder / companion) resolve to None") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val placeholder = variantField("0", BooleanType, "$.__placeholder_field__") + val pfPlaceholder = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(variantExtractionSchema("v", placeholder))) + assert(pfPlaceholder.createFilter(sources.EqualTo("v.`0`", true)).isEmpty, + "Placeholder field must not resolve to a shredded leaf") + + // Full-variant passthrough path "$" yields no keys -> None. + val passthrough = variantField("0", LongType, "$") + val pfPassthrough = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(variantExtractionSchema("v", passthrough))) + assert(pfPassthrough.createFilter(sources.GreaterThan("v.`0`", 1L)).isEmpty, + "Full-variant passthrough must not resolve to a shredded leaf") + } + + test("shredded variant filter: absent shredded field resolves to None") { + // The physical schema does not shred `a` (no typed_value.a subtree); the value lives entirely + // in the opaque residual. Nothing should be pushed. + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group c { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + assert(pf.createFilter(sources.GreaterThan("v.`0`", 999L)).isEmpty, + "A path not shredded in this file must not be pushed") + } + + test("shredded variant filter: case-insensitive column name resolves; keys stay exact-case") { + // The top-level variant column name is a Spark identifier, matched case-insensitively. + // The object key is variant data, matched exact-case, so the physical key must equal the + // requested path's key exactly. + val parquetSchema = + """message spark_schema { + | optional group V { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + // Logical column name `v` differs in case from physical `V`; key `a` matches exactly. + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + caseSensitive = Some(false), variantExtractionSchema = Some(extraction)) + val filter = pf.createFilter(sources.GreaterThan("v.`0`", 999L)) + assert(filter.isDefined, "Case-insensitive column matching should resolve the shredded leaf") + assert(filter.get.toString.toLowerCase(java.util.Locale.ROOT).contains("typed_value"), + s"Expected a typed_value leaf predicate, got ${filter.get}") + } + + test("shredded variant filter: object key is matched case-sensitively even when case-" + + "insensitive analysis") { + // Physical schema shreds a key `A` (uppercase). A request for `$.a` (lowercase) must NOT bind + // to `A`, because variant keys are data resolved exact-case by the reader. Binding to `A` + // would be unsound. + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group A { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + caseSensitive = Some(false), variantExtractionSchema = Some(extraction)) + assert(pf.createFilter(sources.GreaterThan("v.`0`", 999L)).isEmpty, + "Key `$.a` must not case-insensitively bind to the physical `A` subtree") + } + + test("shredded variant filter: negated predicate is not pushed") { + // not(or(leaf, and(anyResidualNotNull, isNull(leaf)))) is rewritten by parquet-mr into + // and(not(leaf), or(and(eq(residual, null)...), notEq(leaf, null))), which is droppable once + // some residual has no nulls AND the leaf is entirely NULL -- an all-fallback row group. So a + // negated shredded predicate must not be pushed. + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), variantExtractionSchema = Some(extraction)) + // != arrives as Not(EqualTo(...)); NOT IN as Not(In(...)); Not(GreaterThan(...)) likewise. + assert(pf.createFilter(sources.Not(sources.EqualTo("v.`0`", 700L))).isEmpty, + "Negated EqualTo on a shredded path must not be pushed") + assert(pf.createFilter(sources.Not(sources.In("v.`0`", Array[Any](1L, 2L)))).isEmpty, + "Negated In on a shredded path must not be pushed") + assert(pf.createFilter(sources.Not(sources.GreaterThan("v.`0`", 700L))).isEmpty, + "Negated GreaterThan on a shredded path must not be pushed") + // Still pushed inside an AND alongside a negated shredded predicate? The AND can push the + // non-negated side, but the negated shredded conjunct itself must be dropped. + val andFilter = pf.createFilter( + sources.And(sources.GreaterThan("v.`0`", 700L), sources.Not(sources.EqualTo("v.`0`", 900L)))) + assert(andFilter.isDefined, "AND should still push its non-negated shredded conjunct") + assert(!andFilter.get.toString.contains("not("), + s"AND must not contain a negated shredded predicate, got ${andFilter.get}") + } + + test("shredded variant filter: extraction type narrower than the leaf is not pushed") { + // Physical leaf is a plain int (INT32). A smallint extraction must NOT be pushed: the leaf + // min/max is over int values, so a row group holding only out-of-int16-range values would be + // wrongly skipped, turning an eager INVALID_VARIANT_CAST into an empty result. + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int32 typed_value (INTEGER(32,true)); + | } + | } + | } + |}""".stripMargin + val narrower = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(variantExtractionSchema("v", variantField("0", ShortType, + "$.a")))) + assert(narrower.createFilter(sources.GreaterThan("v.`0`", 5.toShort)).isEmpty, + "A smallint extraction against an int leaf must not be pushed") + // The exact type (int extraction against int leaf) is pushed. + val exact = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(variantExtractionSchema("v", variantField("0", IntegerType, + "$.a")))) + assert(exact.createFilter(sources.GreaterThan("v.`0`", 5)).isDefined, + "An int extraction against an int leaf should be pushed") + } + + test("shredded variant filter: extraction type wider than the leaf is pushed (safe widening)") { + // Physical leaf is a plain int (INT32); a bigint extraction is sound to push because every int + // value casts to bigint losslessly and ordering is preserved. + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int32 typed_value (INTEGER(32,true)); + | } + | } + | } + |}""".stripMargin + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(variantExtractionSchema("v", variantField("0", LongType, + "$.a")))) + // A long literal within int range widens to the int leaf and is pushed. + val filter = pf.createFilter(sources.GreaterThan("v.`0`", 5L)) + assert(filter.isDefined, "A bigint extraction over an int leaf should push (safe widening)") + assert(filter.get.toString.contains("v.typed_value.a.typed_value"), + s"Expected the int leaf column in ${filter.get}") + // A long literal outside int range is still rejected by valueMatchesParquetType. + assert(pf.createFilter(sources.GreaterThan("v.`0`", Int.MaxValue.toLong + 1L)).isEmpty, + "An out-of-int-range literal must not be pushed against an int leaf") + } + + test("shredded variant filter: large In above threshold still pushes") { + val parquetSchema = + """message spark_schema { + | optional group v { + | optional binary value; + | optional group typed_value { + | optional group a { + | optional binary value; + | optional int64 typed_value; + | } + | } + | } + |}""".stripMargin + val extraction = variantExtractionSchema("v", variantField("0", LongType, "$.a")) + withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_INFILTERTHRESHOLD.key -> "2") { + val pf = createParquetFilters( + MessageTypeParser.parseMessageType(parquetSchema), + variantExtractionSchema = Some(extraction)) + // 3 values > threshold 2: the regular path uses FilterApi.in; the shredded path must too, + // OR-ed with the residual guards, rather than returning None. + val filter = pf.createFilter(sources.In("v.`0`", Array[Any](1L, 2L, 3L))) + assert(filter.isDefined, "A large In on a shredded path should still push down") + val s = filter.get.toString + assert(s.contains("v.typed_value.a.typed_value"), s"Expected leaf column in $s") + assert(s.contains("v.typed_value.a.value") && s.contains("v.value"), + s"Expected residual guards in $s") + } + } } @ExtendedSQLTest diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetGeoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetGeoSuite.scala index 057f722f215c8..185ea357fd503 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetGeoSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetGeoSuite.scala @@ -25,7 +25,9 @@ import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.functions.{st_asbinary, st_geogfromwkb, st_geomfromwkb, st_srid} import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{Geography, GeographyType, StructField, StructType} +import org.apache.spark.tags.ExtendedSQLTest +@ExtendedSQLTest class ParquetGeoSuite extends ParquetCompatibilityTest with SharedSparkSession { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala index e36bb50416f31..bdb78b9e105df 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetQuerySuite.scala @@ -39,6 +39,7 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan import org.apache.spark.sql.functions.struct import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils /** @@ -1371,6 +1372,7 @@ class ParquetV1QuerySuite extends ParquetQuerySuite { } } +@ExtendedSQLTest class ParquetV2QuerySuite extends ParquetQuerySuite { import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala index e62d952a9eb91..2ce1566027258 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTimestampNanosSuite.scala @@ -309,6 +309,63 @@ class ParquetTimestampNanosSuite extends QueryTest with ParquetTest with SharedS } } + test("SPARK-57822: nanos timestamp filter pushdown prunes row groups with identical results") { + // End-to-end: write many rows across multiple row groups, filter on the nanos column, and + // assert that (a) the filter is pushed to Parquet and skips row groups (a stripped-Spark-filter + // scan reads more than the matching rows but fewer than all rows) and (b) the full query result + // equals reading the same predicate with pushdown disabled. Record-level filtering is disabled + // so ONLY row-group-level skipping is exercised; the non-vectorized reader is used so + // `stripSparkFilter` reflects exactly what Parquet returned. + withNanosEnabled { + Seq(7, 8, 9).foreach { p => + val frac = "000000123".take(p) + Seq("ntz", "ltz").foreach { kind => + val typ = if (kind == "ntz") "TIMESTAMP_NTZ" else "TIMESTAMP_LTZ" + val expectedType = + if (kind == "ntz") TimestampNTZNanosType(p) else TimestampLTZNanosType(p) + withTempPath { dir => + val path = dir.getCanonicalPath + // 1024 rows one second apart, monotonically increasing so consecutive row groups hold + // disjoint value ranges (statistics-based skipping is effective). A small block size + // forces multiple row groups. + spark.sql( + s"""SELECT + | $typ '2020-01-01 00:00:00.$frac' + make_dt_interval(0,0,0,id) AS c + |FROM range(0, 1024)""".stripMargin) + .coalesce(1) + .write.option("parquet.block.size", 512).parquet(path) + assert(spark.read.parquet(path).schema("c").dataType === expectedType) + + // Matches the single row at id = 1000, which lives only in a late row group. + val predicate = s"c = $typ '2020-01-01 00:16:40.$frac'" + + // Row-group-level skipping: with record filtering off, a stripped scan returns whole + // surviving row groups. Fewer than all rows (skipping happened) but more than the one + // matching row (record filtering is off). + withSQLConf( + SQLConf.PARQUET_RECORD_FILTER_ENABLED.key -> "false", + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "false") { + val df = spark.read.parquet(path).filter(predicate) + val actual = stripSparkFilter(df).collect().length + assert(actual > 1 && actual < 1024, + s"Expected row-group skipping (p=$p, $kind) but scanned $actual of 1024 rows.") + } + + // Results must be identical whether or not the filter is pushed to Parquet. + val expected = withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "false") { + spark.read.parquet(path).filter(predicate).collect().toSeq + } + assert(expected.length === 1) + withSQLConf(SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true") { + checkAnswer(spark.read.parquet(path).filter(predicate), expected) + } + } + } + } + } + } + test("SPARK-57102: nanos timestamps round-trip via the V2 file source") { withNanosEnabled { withSQLConf(SQLConf.USE_V1_SOURCE_LIST.key -> "") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala index 7ccd664f6c7c4..ac33bd7ba3364 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/ParquetTypeWideningSuite.scala @@ -34,7 +34,9 @@ import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ import org.apache.spark.sql.types.DecimalType.{ByteDecimal, IntDecimal, LongDecimal, ShortDecimal} +import org.apache.spark.tags.ExtendedSQLTest +@ExtendedSQLTest class ParquetTypeWideningSuite extends ParquetTest with SharedSparkSession diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala index cbd3d89c36586..588d0352627b8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantInferShreddingSuite.scala @@ -764,6 +764,45 @@ class VariantInferShreddingSuite extends SharedSparkSession with ParquetTest { checkAnswer(spark.read.parquet(dir.getAbsolutePath), df.collect()) } + testWithTempDir( + "SPARK-58949: infer shredding schema from canonical and legacy fields") { dir => + val bmpKey = new String(Character.toChars(65535)) + val supplementaryKey = new String(Character.toChars(0x10000)) + val quote = 34.toChar.toString + val json = "{" + quote + supplementaryKey + quote + ":1," + + quote + bmpKey + quote + ":2}" + val canonical = VariantBuilder.parseJson(json, false) + val legacyValue = canonical.getValue.clone() + VariantUtil.handleObject[Unit](legacyValue, 0, + (size, idSize, offsetSize, idStart, offsetStart, _dataStart) => { + def swap(start: Int, width: Int): Unit = { + val left = start + (size - 2) * width + val right = left + width + val leftValue = VariantUtil.readUnsigned(legacyValue, left, width) + val rightValue = VariantUtil.readUnsigned(legacyValue, right, width) + VariantUtil.writeLong(legacyValue, left, rightValue, width) + VariantUtil.writeLong(legacyValue, right, leftValue, width) + } + swap(idStart, idSize) + swap(offsetStart, offsetSize) + }) + val canonicalValue = canonical.getValue + val metadata = canonical.getMetadata + val rdd = spark.sparkContext.parallelize(0 until 20, 1).map { i => + val value = if (i % 2 == 0) canonicalValue else legacyValue + InternalRow(new VariantVal(value, metadata)) + } + val writeSchema = DataType.fromDDL("struct<v variant>").asInstanceOf[StructType] + val df = Dataset.ofRows(spark, LogicalRDD(DataTypeUtils.toAttributes(writeSchema), rdd)(spark)) + + df.write.mode("overwrite").parquet(dir.getAbsolutePath) + val expected = StructType(Seq( + StructField(supplementaryKey, LongType), + StructField(bmpKey, LongType))) + checkFileSchema(expected, dir) + assert(spark.read.parquet(dir.getAbsolutePath).count() === 20) + } + testWithTempDir("special characters in field names - dots") { dir => val df = spark.sql( """ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantShreddingFilterPushdownSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantShreddingFilterPushdownSuite.scala new file mode 100644 index 0000000000000..9fd3b2f63e88f --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/VariantShreddingFilterPushdownSuite.scala @@ -0,0 +1,445 @@ +/* + * 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.execution.datasources.parquet + +import java.io.File + +import org.apache.spark.SparkException +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.util.AccumulatorContext + +/** + * End-to-end tests for row-group skipping on shredded Variant columns in Parquet (SPARK-55817). + * + * When a Variant column is written with shredding enabled, each extracted scalar field is stored + * as a typed Parquet leaf column (e.g. `v.typed_value.a.typed_value` for `$.a`) carrying min/max + * statistics. On the DSv1 path, PushVariantIntoScan rewrites + * `variant_get(v, '$.a', 'bigint') > 999` into a struct-field access `v.`0` > 999`, and (when + * `spark.sql.variant.shreddedPredicatePushdown.enabled` is true) ParquetFilters maps `v.`0`` to + * the physical leaf and guards it so a row group is skipped only when the leaf cannot match AND + * every value for the path is provably in the leaf (see `makeShreddedFilter`). + * + * Scope: the optimization fires on the DSv1 read path only. DSv2 does rewrite variant extractions + * into `v.`0`` struct accesses, but only after filter pushdown has run, so the filters offered to + * the Parquet scan builder are still `variant_get(v, ...)` predicates and cannot be pushed for + * row-group skipping (see the comment in ParquetScanBuilder). DSv2 reads remain correct -- the + * variant filter is applied post-scan -- they just do not skip row groups. These tests therefore + * assert skipping only on DSv1, and assert correctness on both DSv1 and DSv2. + * + * The central correctness concern is soundness under fallback: shredding is per-row and per-file + * best-effort, so values that don't fit the shredded type (overflow / type mismatch) or that are + * in a file that doesn't shred the path are stored in an opaque residual with `typed_value` NULL. + * Parquet min/max excludes NULLs, so a naive leaf-only predicate could skip a row group that still + * holds a matching row. These tests mix typed and fallback rows in a single row group and assert + * that no matching row is ever dropped and results equal the no-pushdown baseline. + */ +class VariantShreddingFilterPushdownSuite extends QueryTest with ParquetTest + with SharedSparkSession { + + // Base configs to write shredded Variant Parquet files. `annotate` controls whether the physical + // variant group carries the VARIANT logical-type annotation (the production default is true). + private def writeConf(forceSchema: String, annotate: Boolean): Seq[(String, String)] = Seq( + SQLConf.VARIANT_WRITE_SHREDDING_ENABLED.key -> "true", + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true", + SQLConf.VARIANT_FORCE_SHREDDING_SCHEMA_FOR_TEST.key -> forceSchema, + SQLConf.PARQUET_ANNOTATE_VARIANT_LOGICAL_TYPE.key -> annotate.toString) + + /** + * Counts how many Parquet row groups are actually read by the given DataFrame, using the + * accumulator technique from ParquetFilterSuite. Only meaningful with the vectorized reader, + * which reports the row-group count into a registered NumRowGroupsAcc. + */ + private def countRowGroupsRead(df: DataFrame): Int = { + val accu = new NumRowGroupsAcc + sparkContext.register(accu) + try { + df.foreachPartition((it: Iterator[Row]) => it.foreach(_ => accu.add(0))) + accu.value + } finally { + AccumulatorContext.remove(accu.id) + } + } + + /** + * Writes a JSON-per-row Variant Parquet file coalesced to a single partition with a tiny block + * size so the writer emits multiple row groups. `jsonExpr` is the SQL expression producing the + * JSON string per `id` in `range(0, numRows, 1, 1)`. + */ + private def writeShredded( + dir: File, + forceSchema: String, + jsonExpr: String, + numRows: Int, + blockSize: Int = 512, + annotate: Boolean = true): Unit = { + withSQLConf(writeConf(forceSchema, annotate): _*) { + spark.sql( + s"""SELECT parse_json($jsonExpr) AS v + |FROM range(0, $numRows, 1, 1)""".stripMargin) + .coalesce(1) + .write + .option("parquet.block.size", blockSize) + .mode("overwrite") + .parquet(dir.getAbsolutePath) + } + } + + // Run `block` with pushdown enabled, across the {DSv1, DSv2} x {vectorized, non-vectorized} grid. + // `dsv1` is passed so a test can assert row-group skipping only on the DSv1 path. + private def forEachReader(block: (Boolean, Boolean) => Unit): Unit = { + Seq("parquet" -> true, "" -> false).foreach { case (useV1, dsv1) => + Seq(true, false).foreach { vectorized => + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> useV1, + SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true") { + withClue(s"(dsv1=$dsv1, vectorized=$vectorized) ") { + block(dsv1, vectorized) + } + } + } + } + } + + // Read the same query with pushdown disabled: the baseline that must never lose rows. + private def baseline(read: => DataFrame): Seq[Row] = { + withSQLConf( + SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED.key -> "false", + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true") { + read.collect().toSeq + } + } + + // Assert that a row group whose only match for `$.a > 999` is a residual fallback (a value the + // int64 leaf cannot hold, so it lands in v.typed_value.a.value with the leaf NULL) is not + // dropped. `fallbackJson` is the JSON for the fallback row. The leaf min/max is 0..49, so only + // the guard keeps the row group; a leaf-only predicate would drop it and lose the row (#54598). + private def checkResidualFallbackNotDropped(fallbackJson: String): Unit = { + withTempDir { dir => + val jsonExpr = + "case when id = 50 then '" + fallbackJson + "' else '{\"a\":' || id || '}' end" + writeShredded(dir, "a bigint", jsonExpr, numRows = 51, blockSize = 1024 * 1024) + + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a") + .where("a > 999") + val expected = baseline(read) + assert(expected == Seq(Row(1500L)), s"baseline should return the fallback row, got $expected") + + forEachReader { (dsv1, vectorized) => + // The row group's only match is in the residual with a NULL leaf, so the guard must keep + // it: results include the fallback row and the row group is not skipped. + checkAnswer(read, expected) + if (dsv1 && vectorized) { + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a") + assert(countRowGroupsRead(read) == countRowGroupsRead(all), + "Row group whose only match is a residual fallback must NOT be skipped") + } + } + } + } + + test("residual fallback beyond the leaf's min/max is not dropped") { + // Different fallback encodings that all miss the int64 leaf: a non-integral decimal that + // 1500.5 rounds to 1500, and a string "1500". Both are read back as bigint 1500. + checkResidualFallbackNotDropped("{\"a\":1500.5}") + checkResidualFallbackNotDropped("{\"a\":\"1500\"}") + } + + test("negated predicate over an all-fallback row group is not dropped") { + withTempDir { dir => + // `a` shredded as bigint. Both rows are non-integral decimals the int64 leaf cannot hold, so + // both land in the residual (typed leaf entirely NULL, residual has no nulls). The path is + // still pushable (bigint extraction over a bigint leaf). A naive negated push would rewrite + // `!= 700` into and(notEq(leaf), eq(residual, null)) and skip the row group -- losing both + // rows. The negation guard must prevent pushing, so {500, 600} come back. + val jsonExpr = "case when id = 0 then '{\"a\":500.5}' else '{\"a\":600.5}' end" + writeShredded(dir, "a bigint", jsonExpr, numRows = 2, blockSize = 1024 * 1024) + + Seq( + "try_variant_get(v, '$.a', 'bigint') != 700" -> Seq(Row(500L), Row(600L)), + "try_variant_get(v, '$.a', 'bigint') NOT IN (700, 800)" -> Seq(Row(500L), Row(600L)) + ).foreach { case (predicate, want) => + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a") + .where(predicate) + assert(baseline(read).sortBy(_.getLong(0)) == want, s"baseline for $predicate") + forEachReader { (_, _) => + checkAnswer(read, want) + } + } + } + } + + test("type-mismatch fallback: string values in a numeric-shredded field are not dropped") { + withTempDir { dir => + // `a` shredded as bigint. Even rows -> a is a number (typed); odd rows -> a is a string + // (type mismatch -> residual, typed_value NULL). Use try_variant_get so the string rows + // resolve to NULL (filtered out) rather than raising a strict-cast error, and assert the + // matching numeric rows are still returned. + val jsonExpr = + "case when id % 2 = 0 then '{\"a\":' || (id + 1000) || '}' " + + "else '{\"a\":\"str' || id || '\"}' end" + writeShredded(dir, "a bigint", jsonExpr, numRows = 20, blockSize = 1024 * 1024) + + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a") + .where("a > 1005") + val expected = baseline(read) + assert(expected.nonEmpty, "baseline should return the matching numeric rows") + + forEachReader { (_, _) => + checkAnswer(read, expected) + } + } + } + + test("file without the shredded path: value read from residual, predicate not pushed") { + withTempDir { dir => + // Force a shredding schema that does NOT contain `a`; `$.a` lives entirely in the opaque + // top-level residual. Nothing is pushed for `$.a`; results must still be correct. + val jsonExpr = "'{\"a\":' || id || '}'" + writeShredded(dir, "b bigint", jsonExpr, numRows = 20, blockSize = 1024 * 1024) + + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a") + .where("a > 9") + val expected = baseline(read) + assert(expected == (10L to 19L).map(Row(_)), s"unexpected baseline: $expected") + + forEachReader { (_, _) => + checkAnswer(read.orderBy("a"), expected) + } + } + } + + test("residual-null happy path: a row group is skipped (DSv1) and results are correct") { + withTempDir { dir => + // Homogeneous typed data across two row groups. All values shred cleanly (residuals all + // NULL), so the optimization fires and one row group is skipped on DSv1. + val jsonExpr = "'{\"a\":' || id || '}'" + // Small block size -> at least two row groups: [0,999] and [1000,1999]. + writeShredded(dir, "a bigint", jsonExpr, numRows = 2000, blockSize = 512) + + forEachReader { (dsv1, vectorized) => + val filtered = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a") + .where("a > 999") + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a") + checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_))) + if (dsv1 && vectorized) { + assert(countRowGroupsRead(filtered) < countRowGroupsRead(all), + "Expected at least one row group to be skipped by the shredded leaf statistics") + } + } + } + } + + test("partial object with a non-shredded sibling key still skips (leaf has no nulls)") { + withTempDir { dir => + // Every row also carries a key `z` outside the shredding schema, so the whole partial object + // lands in the top-level residual v.value -- it is non-null on every row. `a` is still fully + // shredded into the typed leaf (no nulls). The flat OR guard could never skip here (v.value + // never all-null); the tighter guard skips via the "leaf has no nulls" arm. Sorted on `a` + // across two row groups so `a > 999` can drop the first. + val jsonExpr = "'{\"a\":' || id || ', \"z\":\"outside\"}'" + writeShredded(dir, "a bigint", jsonExpr, numRows = 2000, blockSize = 512) + + forEachReader { (dsv1, vectorized) => + val filtered = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a").where("a > 999") + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a") + checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_))) + if (dsv1 && vectorized) { + assert(countRowGroupsRead(filtered) < countRowGroupsRead(all), + "Expected skipping despite a non-null top-level residual (partial object)") + } + } + } + } + + test("multi-level $.a.b: skip fires on DSv1 and results are correct") { + withTempDir { dir => + // `a` shredded as struct<b bigint>. Homogeneous nested typed data across two row groups so + // the skip fires on the nested leaf `v.typed_value.a.typed_value.b.typed_value`. + val typedJson = "'{\"a\":{\"b\":' || id || '}}'" + writeShredded(dir, "a struct<b bigint>", typedJson, numRows = 2000, blockSize = 512) + + forEachReader { (dsv1, vectorized) => + val filtered = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a.b', 'bigint') AS b") + .where("b > 999") + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a.b', 'bigint') AS b") + checkAnswer(filtered.orderBy("b"), (1000L to 1999L).map(Row(_))) + if (dsv1 && vectorized) { + assert(countRowGroupsRead(filtered) < countRowGroupsRead(all), + "Expected a row group to be skipped by the nested shredded leaf statistics") + } + } + } + } + + test("multi-level $.a.b: nested residual fallback beyond the leaf's min/max is not dropped") { + withTempDir { dir => + // `a` shredded as struct<b bigint>. Rows 0..19 shred cleanly, so the nested leaf + // v.typed_value.a.typed_value.b.typed_value is min 0 / max 19. Row 20 stores `b` as a + // non-integral decimal the int64 leaf cannot hold, so it lands in + // v.typed_value.a.typed_value.b.value with the leaf NULL. `b > 999` matches only that row and + // the leaf min/max cannot match it, so only the guard keeps the row group -- this makes the + // nested leaf-level residual load-bearing (a leaf-only predicate would drop it). + val jsonExpr = + "case when id = 20 then '{\"a\":{\"b\":1500.5}}' else '{\"a\":{\"b\":' || id || '}}' end" + writeShredded(dir, "a struct<b bigint>", jsonExpr, numRows = 21, blockSize = 1024 * 1024) + + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a.b', 'bigint') AS b") + .where("b > 999") + assert(baseline(read) == Seq(Row(1500L)), "baseline should return the nested fallback row") + + forEachReader { (dsv1, vectorized) => + checkAnswer(read, Seq(Row(1500L))) + if (dsv1 && vectorized) { + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a.b', 'bigint') AS b") + assert(countRowGroupsRead(read) == countRowGroupsRead(all), + "Row group whose only match is a nested residual fallback must NOT be skipped") + } + } + } + } + + test("unannotated variant layout: skip and fallback still work") { + // The suite writes the production-default annotated layout everywhere else; this test covers + // the unannotated physical layout explicitly. Both a skip-eligible query and a residual + // fallback must behave correctly. + withTempDir { dir => + writeShredded(dir, "a bigint", "'{\"a\":' || id || '}'", numRows = 2000, blockSize = 512, + annotate = false) + forEachReader { (dsv1, vectorized) => + val filtered = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a").where("a > 999") + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a") + checkAnswer(filtered.orderBy("a"), (1000L to 1999L).map(Row(_))) + if (dsv1 && vectorized) { + assert(countRowGroupsRead(filtered) < countRowGroupsRead(all), + "Expected a row group to be skipped on the unannotated layout") + } + } + } + + withTempDir { dir => + // Residual fallback under the unannotated layout: the row whose only match is in the residual + // (NULL leaf) must survive. + writeShredded(dir, "a bigint", + "case when id = 50 then '{\"a\":1500.5}' else '{\"a\":' || id || '}' end", + numRows = 51, blockSize = 1024 * 1024, annotate = false) + def read: DataFrame = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a").where("a > 999") + assert(baseline(read) == Seq(Row(1500L))) + forEachReader { (_, _) => + checkAnswer(read, Seq(Row(1500L))) + } + } + } + + test("deferCastError=true: strict non-string cast does not fire; try_variant_get still does") { + // With deferCastError, a strict cast to a non-string, non-variant type is rewritten into + // UnwrapVariantCastError, which is not translated to a pushable filter -- so shredded pushdown + // does not fire for it (results still correct). try_variant_get (failOnError=false) and string + // targets are unaffected and still fire. Results must be correct in every combination. + withTempDir { dir => + writeShredded(dir, "a bigint", "'{\"a\":' || id || '}'", numRows = 2000, blockSize = 512) + Seq("true", "false").foreach { defer => + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR.key -> defer, + SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> "true", + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true") { + withClue(s"(deferCastError=$defer) ") { + // Strict cast: correct either way (does not fire when defer=true). + val strict = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'bigint') AS a").where("a > 999") + checkAnswer(strict.orderBy("a"), (1000L to 1999L).map(Row(_))) + // try_variant_get: unaffected by deferCastError and still skips a row group on DSv1. + val tryGet = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a").where("a > 999") + val all = spark.read.parquet(dir.getAbsolutePath) + .selectExpr("try_variant_get(v, '$.a', 'bigint') AS a") + checkAnswer(tryGet.orderBy("a"), (1000L to 1999L).map(Row(_))) + assert(countRowGroupsRead(tryGet) < countRowGroupsRead(all), + "try_variant_get should still skip regardless of deferCastError") + } + } + } + } + } + + test("strict variant_get preserves INVALID_VARIANT_CAST on a residual fallback (not empty)") { + // The worst failure mode: a thrown cast error silently becoming an empty result. `a` shredded + // as int; row 50 stores 3000000000, which overflows int32 so it lands in the residual with the + // leaf NULL (extraction type matches the leaf exactly, so the path is pushed). With + // deferCastError=false (default) the scan casts eagerly, so strict `variant_get(v,'$.a','int')` + // must raise INVALID_VARIANT_CAST -- a leaf-only push would drop the row group (leaf max 49) + // and return empty instead. The guard keeps the row group, so the error is preserved. + withTempDir { dir => + val jsonExpr = "case when id = 50 then '{\"a\":3000000000}' else '{\"a\":' || id || '}' end" + writeShredded(dir, "a int", jsonExpr, numRows = 51, blockSize = 1024 * 1024) + Seq("parquet", "").foreach { useV1 => + Seq(true, false).foreach { vectorized => + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> useV1, + SQLConf.PUSH_VARIANT_INTO_SCAN_DEFER_CAST_ERROR.key -> "false", + SQLConf.VARIANT_SHREDDED_PREDICATE_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key -> vectorized.toString, + SQLConf.VARIANT_ALLOW_READING_SHREDDED.key -> "true") { + withClue(s"(useV1='$useV1', vectorized=$vectorized) ") { + val e = intercept[SparkException] { + spark.read.parquet(dir.getAbsolutePath) + .selectExpr("variant_get(v, '$.a', 'int') AS a").where("a > 999").collect() + } + assert(findCause(e, "INVALID_VARIANT_CAST"), + s"Expected INVALID_VARIANT_CAST to be preserved, got: $e") + } + } + } + } + } + } + + // Walk an exception's cause chain for a Spark error condition (message substring). + private def findCause(e: Throwable, condition: String): Boolean = { + var cur: Throwable = e + while (cur != null) { + if (Option(cur.getMessage).exists(_.contains(condition))) return true + cur = cur.getCause + } + false + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala index febeb7c5fb1ea..f0063f86d1117 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/parquet/types/ops/TimestampNanosParquetOpsSuite.scala @@ -17,11 +17,15 @@ package org.apache.spark.sql.execution.datasources.parquet.types.ops +import java.time.{Instant, LocalDateTime, ZoneOffset} + import org.apache.parquet.column.ColumnDescriptor +import org.apache.parquet.filter2.predicate.FilterApi +import org.apache.parquet.filter2.predicate.SparkFilterApi.longColumn import org.apache.parquet.io.api.PrimitiveConverter import org.apache.parquet.schema.{LogicalTypeAnnotation, Type, Types} import org.apache.parquet.schema.LogicalTypeAnnotation.{TimestampLogicalTypeAnnotation, TimeUnit} -import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT64 +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.{INT32, INT64} import org.apache.parquet.schema.Type.Repetition.REQUIRED import org.apache.spark.{SparkArithmeticException, SparkFunSuite, SparkRuntimeException} @@ -136,6 +140,129 @@ class TimestampNanosParquetOpsSuite extends SparkFunSuite { } } + // ---------- filter-pushdown ops ---------- + + test("filterOps accepts the matching Java time value and rejects others") { + // LTZ pushes down java.time.Instant; NTZ pushes down java.time.LocalDateTime. Each rejects + // the other type (and non-temporal values) so a mismatched literal falls through to no + // pushdown rather than a ClassCastException in the converter. + val ltzOps = TimestampNanosParquetOps.ltzFilterOps + val ntzOps = TimestampNanosParquetOps.ntzFilterOps + assert(ltzOps.acceptsValue(Instant.parse("2020-01-01T00:00:00Z"))) + assert(!ltzOps.acceptsValue(LocalDateTime.parse("2020-01-01T00:00:00"))) + assert(!ltzOps.acceptsValue(java.lang.Long.valueOf(1L))) + assert(ntzOps.acceptsValue(LocalDateTime.parse("2020-01-01T00:00:00"))) + assert(!ntzOps.acceptsValue(Instant.parse("2020-01-01T00:00:00Z"))) + assert(!ntzOps.acceptsValue("2020-01-01")) + } + + test("filterOps rejects values outside the INT64 epoch-nanos range (falls back to full scan)") { + // Year 2300 is past the ~2262 int64 epoch-nanos cutoff: encoding would overflow, so + // acceptsValue must reject it (SPARK-46092-style guard) instead of throwing during filter + // creation. An in-range value is accepted. + val ltzOps = TimestampNanosParquetOps.ltzFilterOps + val ntzOps = TimestampNanosParquetOps.ntzFilterOps + assert(!ltzOps.acceptsValue(Instant.parse("2300-01-01T00:00:00Z"))) + assert(!ntzOps.acceptsValue(LocalDateTime.parse("2300-01-01T00:00:00"))) + assert(ltzOps.acceptsValue(Instant.parse("2020-01-01T00:00:00Z"))) + assert(ntzOps.acceptsValue(LocalDateTime.parse("2020-01-01T00:00:00"))) + } + + test("filterOps make* throws on an out-of-range value - why acceptsValue must gate every In " + + "element, not just the head") { + // The encoder is exact (Math.addExact/multiplyExact), so building a predicate directly from an + // out-of-range value throws rather than clamping. This is the contract behind the + // ParquetFilters In-arm gating every element through acceptsValue: a makeIn set (or per-element + // makeEq) that includes an out-of-range tail would crash filter creation. `acceptsValue` + // (tested above) is the guard that keeps make* off such values on the pushdown path. + val path = Array("c") + val ltzOps = TimestampNanosParquetOps.ltzFilterOps + val ntzOps = TimestampNanosParquetOps.ntzFilterOps + val ltzOverflow = Instant.parse("2300-01-01T00:00:00Z") + val ntzOverflow = LocalDateTime.parse("2300-01-01T00:00:00") + intercept[ArithmeticException](ltzOps.makeEq(path, ltzOverflow)) + intercept[ArithmeticException](ntzOps.makeEq(path, ntzOverflow)) + // A mixed set (in-range head + out-of-range tail) is the crash the head-only guard let through. + val ltzInRange = Instant.parse("2020-01-01T00:00:00Z") + val ntzInRange = LocalDateTime.parse("2020-01-01T00:00:00") + intercept[ArithmeticException](ltzOps.makeIn(path, Array[Any](ltzInRange, ltzOverflow))) + intercept[ArithmeticException](ntzOps.makeIn(path, Array[Any](ntzInRange, ntzOverflow))) + } + + test("filterOps declares the canonical nanos-timestamp Parquet encoding") { + // LTZ is isAdjustedToUTC=true, NTZ is false; both INT64 TIMESTAMP(NANOS). These are the keys + // the ParquetFilters reverse lookup matches against the file schema. + val ltzOps = TimestampNanosParquetOps.ltzFilterOps + val ntzOps = TimestampNanosParquetOps.ntzFilterOps + assert(ltzOps.primitiveTypeName === INT64) + assert(ltzOps.logicalTypeAnnotation === + LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS)) + assert(ntzOps.primitiveTypeName === INT64) + assert(ntzOps.logicalTypeAnnotation === + LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS)) + } + + test("filterOps builds predicates converting to signed INT64 epoch-nanoseconds") { + val path = Array("c") + val col = longColumn(path) + + // LTZ: Instant -> epoch-nanos. Sub-microsecond digits are preserved (not truncated to micros). + val ltzOps = TimestampNanosParquetOps.ltzFilterOps + val instant = Instant.parse("2020-01-01T12:34:56.000000789Z") + val ltzNanos = java.lang.Long.valueOf( + instant.getEpochSecond * DateTimeConstants.NANOS_PER_SECOND + instant.getNano) + assert(ltzOps.makeEq(path, instant) === FilterApi.eq(col, ltzNanos)) + assert(ltzOps.makeNotEq(path, instant) === FilterApi.notEq(col, ltzNanos)) + assert(ltzOps.makeLt(path, instant) === FilterApi.lt(col, ltzNanos)) + assert(ltzOps.makeLtEq(path, instant) === FilterApi.ltEq(col, ltzNanos)) + assert(ltzOps.makeGt(path, instant) === FilterApi.gt(col, ltzNanos)) + assert(ltzOps.makeGtEq(path, instant) === FilterApi.gtEq(col, ltzNanos)) + + // NTZ: LocalDateTime (interpreted at UTC) -> epoch-nanos. + val ntzOps = TimestampNanosParquetOps.ntzFilterOps + val ldt = LocalDateTime.parse("2020-01-01T12:34:56.000000789") + val ntzInstant = ldt.toInstant(ZoneOffset.UTC) + val ntzNanos = java.lang.Long.valueOf( + ntzInstant.getEpochSecond * DateTimeConstants.NANOS_PER_SECOND + ntzInstant.getNano) + assert(ntzOps.makeEq(path, ldt) === FilterApi.eq(col, ntzNanos)) + assert(ntzOps.makeIn(path, Array[Any](ldt)) === { + val set = new java.util.HashSet[java.lang.Long]() + set.add(ntzNanos) + FilterApi.in(col, set) + }) + } + + test("filterOps eq/notEq/in tolerate a null value (IsNull / IsNotNull)") { + val path = Array("c") + val col = longColumn(path) + val nullLong = null.asInstanceOf[java.lang.Long] + // null value -> null Long comparand; used by ParquetFilters for IsNull / IsNotNull. + Seq(TimestampNanosParquetOps.ltzFilterOps, TimestampNanosParquetOps.ntzFilterOps).foreach { + ops => + assert(ops.makeEq(path, null) === FilterApi.eq(col, nullLong)) + assert(ops.makeNotEq(path, null) === FilterApi.notEq(col, nullLong)) + val set = new java.util.HashSet[java.lang.Long]() + set.add(null) + assert(ops.makeIn(path, Array[Any](null)) === FilterApi.in(col, set)) + } + } + + test("ParquetTypeOps.filterOpsFor resolves each nanos encoding and nothing else") { + // The LTZ and NTZ encodings differ only in isAdjustedToUTC; each resolves to its own ops. + assert(ParquetTypeOps.filterOpsFor( + LogicalTypeAnnotation.timestampType(true, TimeUnit.NANOS), INT64) + .contains(TimestampNanosParquetOps.ltzFilterOps)) + assert(ParquetTypeOps.filterOpsFor( + LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS), INT64) + .contains(TimestampNanosParquetOps.ntzFilterOps)) + // MICROS unit, or an INT32 primitive, is not a nanos-timestamp encoding (pushdown falls + // through to no framework ops). + assert(ParquetTypeOps.filterOpsFor( + LogicalTypeAnnotation.timestampType(true, TimeUnit.MICROS), INT64).isEmpty) + assert(ParquetTypeOps.filterOpsFor( + LogicalTypeAnnotation.timestampType(false, TimeUnit.NANOS), INT32).isEmpty) + } + // ---------- helpers ---------- private def nanosField(isAdjustedToUTC: Boolean): Type = diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala index 6778f7109a6e0..38fea013a246e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/TextSuite.scala @@ -33,6 +33,7 @@ import org.apache.spark.sql.execution.datasources.CommonFileDataSourceSuite import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{StringType, StructType} +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils abstract class TextSuite extends SharedSparkSession with CommonFileDataSourceSuite { @@ -349,6 +350,7 @@ class TextV1Suite extends TextSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "text") } +@ExtendedSQLTest class TextV2Suite extends TextSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/WholeTextFileSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/WholeTextFileSuite.scala index 81eddc182539d..08a4a72ec67a8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/WholeTextFileSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/text/WholeTextFileSuite.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.util.HadoopCompressionCodec.GZIP import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{StringType, StructType} +import org.apache.spark.tags.ExtendedSQLTest abstract class WholeTextFileSuite extends SharedSparkSession { @@ -113,6 +114,7 @@ class WholeTextFileV1Suite extends WholeTextFileSuite { .set(SQLConf.USE_V1_SOURCE_LIST, "text") } +@ExtendedSQLTest class WholeTextFileV2Suite extends WholeTextFileSuite { override protected def sparkConf: SparkConf = super diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala index 4af5a32515349..a0ca66183fdf4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.variant.VariantGet +import org.apache.spark.sql.catalyst.optimizer.ConstantFolding import org.apache.spark.sql.catalyst.util.V2ExpressionBuilder import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue, VariantGet => V2VariantGet} import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And => V2And, Not => V2Not, Or => V2Or, Predicate} @@ -1034,6 +1035,25 @@ class DataSourceV2StrategySuite extends SharedSparkSession { } } + test("SPARK-58428: translating an expression that failed to evaluate does not loop forever") { + withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") { + // `coalesce(c, 1 div 0) = 1`. Constant folding defers the divide by zero error because the + // failing expression sits in a conditional branch, so it is tagged FAILED_TO_EVALUATE and + // left as is. `div` returns BIGINT, so `c` is LONG to keep the `coalesce` inputs equal. + val c = AttributeReference("c", LongType)() + val predicate = + EqualTo(Coalesce(Seq(c, IntegralDivide(Literal(1), Literal(0)))), Literal(1L)) + val folded = ConstantFolding.constantFolding(predicate) + assert( + folded.exists(_.containsTag(ConstantFolding.FAILED_TO_EVALUATE)), + "expected the divide by zero branch to be tagged FAILED_TO_EVALUATE") + + // Translating such an expression used to recurse forever. Note that a regression hangs + // this test instead of failing it, as the recursion is in tail position. + assert(new V2ExpressionBuilder(folded, isPredicate = true).build().isEmpty) + } + } + /** * Translate the given Catalyst [[Expression]] into data source V2 [[Predicate]] * then verify against the given [[Predicate]]. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala index cd09fca7c7021..be772bc3f28e3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder} -import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyedShuffleSpec, Partitioning, PartitioningCollection, UnknownPartitioning} import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -255,6 +255,31 @@ class GroupPartitionsExecSuite extends SharedSparkSession { } } + test("SPARK-59027: createShuffleSpec subset-keys spec orders keys the same as this node's " + + "grouping") { + // With `allowKeysSubsetOfPartitionKeys`, `EnsureRequirements` may shuffle the other join side + // onto the spec's projected keys while this side is re-grouped by a `GroupPartitionsExec` + // carrying the spec's `joinKeyPositions`. The two key orders must agree (see + // `KeyedPartitioning.groupedKeyRowOrdering`), or the sides are mis-aligned -- a planning-time + // `PartitioningCollection` invariant failure for inner joins, silent wrong results for join + // types that expose only one side's partitioning. + // First-appearance order of the projected keys ([3], [1], [2]) differs from their sorted + // order ([1], [2], [3]), so the assertion discriminates the sort each side uses. + val partitionKeys = Seq(row(3, 30), row(1, 10), row(2, 20), row(1, 99)) + val partitioning = KeyedPartitioning(Seq(exprA, exprB), partitionKeys) + + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val spec = partitioning.createShuffleSpec(ClusteredDistribution(Seq(exprA))) + .asInstanceOf[KeyedShuffleSpec] + assert(spec.joinKeyPositions === Some(Seq(0))) + + val gpe = GroupPartitionsExec( + DummySparkPlan(outputPartitioning = partitioning), + joinKeyPositions = spec.joinKeyPositions) + assert(gpe.groupedPartitions.map(_._1) === spec.partitioning.partitionKeys) + } + } + test("SPARK-56549: tryEnableSortedMerge returns None when no coalescing occurs") { val partitionKeys = Seq(row(1), row(2), row(3)) val childOrdering = Seq(SortOrder(exprA, Ascending)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalogSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalogSuite.scala index b3b0eb811ead9..1caebf506a3d1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalogSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalogSuite.scala @@ -65,6 +65,7 @@ class JDBCTableCatalogSuite extends SharedSparkSession { jdbcClientType: String): Metadata = new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", dataType.isInstanceOf[NumericType]) .putString("jdbcClientType", jdbcClientType) .build() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceRealTimeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceRealTimeSuite.scala new file mode 100644 index 0000000000000..71714e1eb550c --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/state/StateDataSourceRealTimeSuite.scala @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2.state + +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, + LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.{OutputMode, StreamRealTimeModeManualClockSuiteBase, TimeMode} + +/** + * State data source reads under Real-Time Mode (RTM). + */ +class StateDataSourceRealTimeSuite + extends StreamRealTimeModeManualClockSuiteBase + with StateDataSourceTestBase { + + import testImplicits._ + + private def assertEventuallyBatchCommitted(batchId: Long): StreamAction = { + Execute(s"Assert batch $batchId is committed") { q => + eventually(timeout(1.minute)) { + assert(q.commitLog.getLatest().get._1 === batchId) + } + } + } + + test("transformWithState + RTM: state data source read") { + withSQLConf( + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + withTempDir { tempDir => + val input = LowLatencyMemoryStream[String](2) + val query = input.toDS() + .groupByKey(x => x) + .transformWithState( + new StatefulProcessorWithSingleValueVar(), + TimeMode.ProcessingTime(), + OutputMode.Update()) + val checkpoint = tempDir.getCanonicalPath + + testStream(query, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(trigger = defaultTrigger, checkpointLocation = checkpoint), + AddData(input, "a", "b"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "1"), ("b", "1")), + advanceRealTimeClock, + assertEventuallyBatchCommitted(0), + StopStream + ) + + val stateReaderDf = spark.read + .format("statestore") + .option(StateSourceOptions.PATH, checkpoint) + .option(StateSourceOptions.STATE_VAR_NAME, "valueState") + .load() + + checkAnswer( + stateReaderDf.selectExpr( + "key.value AS groupingKey", + "value.id AS valueId", + "value.name AS valueName"), + Seq(Row("a", 1L, "dummyKey"), Row("b", 1L, "dummyKey"))) + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..7a0664b46a381 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -35,7 +35,6 @@ import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoin import org.apache.spark.sql.execution.python.FlatMapCoGroupsInPandasExec import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.{IntegerType, StructField, StructType} @@ -60,7 +59,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -71,7 +70,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprA, exprB)) @@ -84,7 +83,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprD :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1) EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprC, exprD)) assert(rightKeys === Seq(exprA, exprB)) @@ -126,8 +125,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprC, exprB, exprD)) assert(rightKeys === Seq(exprD, exprA, exprC)) case other => fail(other.toString) @@ -145,7 +144,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec1) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -159,7 +158,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3) EnsureRequirements.apply(smjExec2) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), _) => assert(leftKeys === Seq(exprB, exprA)) assert(rightKeys === Seq(exprB, exprC)) @@ -173,7 +172,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec3) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprB, exprC)) assert(rightKeys === Seq(exprB, exprA)) case other => fail(other.toString) @@ -319,7 +318,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -335,7 +334,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) @@ -353,7 +352,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(p.expressions == Seq(exprC)) @@ -372,7 +371,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprC)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -388,7 +387,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB, exprB)) assert(rightKeys === Seq(exprA, exprC, exprD)) assert(p.expressions == Seq(exprA, exprC, exprA)) @@ -444,8 +443,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(right.numPartitions == 5) case other => fail(other.toString) @@ -462,7 +461,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprC, exprD)) @@ -481,8 +480,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 5) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == 5) @@ -492,7 +491,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 1) assert(right.numPartitions == 1) assert(right.expressions == Seq(exprC)) @@ -510,8 +509,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == conf.numShufflePartitions) assert(left.expressions == Seq(exprA, exprB)) assert(right.numPartitions == conf.numShufflePartitions) @@ -529,7 +528,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: PartitioningCollection, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.numPartitions == 10) assert(right.numPartitions == 10) assert(right.expressions == Seq(exprA)) @@ -545,7 +544,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: PartitioningCollection, _, _), _), _) => assert(left.numPartitions == 20) assert(left.expressions == Seq(exprC)) @@ -587,7 +586,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == 6) @@ -606,7 +605,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -624,7 +623,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -644,7 +643,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) @@ -666,8 +665,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB)) assert(right.expressions === Seq(exprC, exprD)) assert(left.numPartitions == conf.numShufflePartitions) @@ -691,7 +690,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA, exprB)) assert(rightKeys === Seq(exprC, exprD)) assert(left.numPartitions == 9) @@ -714,8 +713,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None, plan1, plan2) EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 20) @@ -733,7 +732,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(right.numPartitions == 10) @@ -765,8 +764,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { } else { EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(leftKeys, rightKeys, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(leftKeys === Seq(exprA)) assert(rightKeys === Seq(exprC)) assert(left.numPartitions == 5) @@ -862,28 +861,9 @@ class EnsureRequirementsSuite extends SharedSparkSession { assert(right.expressions === Seq(bucket(4, exprA), years(exprC))) case other => fail(other.toString) } - - // by default spark.sql.requireAllClusterKeysForCoPartition is true, so when there isn't - // exact match on all partition keys, Spark will fallback to shuffle. - plan1 = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty) - ) - plan2 = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprC) :: Nil, Seq.empty) - ) - smjExec = SortMergeJoinExec( - exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - EnsureRequirements.apply(smjExec) match { - case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => - assert(left.expressions === Seq(exprA, exprB, exprB)) - assert(right.expressions === Seq(exprA, exprC, exprC)) - case other => fail(other.toString) - } } - test(s"KeyedPartitioning with ${REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key} = false") { + test("KeyedPartitioning with subset of join keys") { var plan1 = new DummySparkPlanWithBatchScanChild( outputPartitioning = KeyedPartitioning(bucket(4, exprB) :: years(exprC) :: Nil, Seq.empty) ) @@ -891,9 +871,14 @@ class EnsureRequirementsSuite extends SharedSparkSession { outputPartitioning = KeyedPartitioning(bucket(4, exprC) :: years(exprB) :: Nil, Seq.empty) ) - // simple case + // simple case: join key exprA is not covered by either side's partition keys, so by default + // the coverage check of requireAllClusterKeysForCoPartition falls back to shuffle to avoid + // joining on a partitioning coarser than the join keys var smjExec = SortMergeJoinExec( exprA :: exprB :: exprC :: Nil, exprA :: exprC :: exprB :: Nil, Inner, None, plan1, plan2) + assert(EnsureRequirements.apply(smjExec) + .collect { case s: ShuffleExchangeLike => s }.length == 2) + // with requireAllClusterKeysForCoPartition=false, SPJ is allowed applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), @@ -912,7 +897,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -930,7 +915,20 @@ class EnsureRequirementsSuite extends SharedSparkSession { KeyedPartitioning(years(exprA) :: bucket(4, exprC) :: days(exprA) :: Nil, Seq.empty)) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { + case SortMergeJoinExec(_, _, _, _, + SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), + SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => + assert(left.expressions === Seq(years(exprA), bucket(4, exprB), days(exprA))) + assert(right.expressions === Seq(years(exprA), bucket(4, exprC), days(exprA))) + case other => fail(other.toString) + } + + // a column partitioned by more than one transform: partition expressions outnumber the + // join keys, but every join key is covered, so SPJ is allowed with default configs + smjExec = SortMergeJoinExec( + exprA :: exprB :: Nil, exprA :: exprC :: Nil, Inner, None, plan1, plan2) + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => @@ -951,8 +949,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner, None, plan1, plan2) applyEnsureRequirementsWithSubsetKeys(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprC)) assert(right.expressions === Seq(exprA, exprB, exprC)) case other => fail(other.toString) @@ -967,10 +965,10 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -985,10 +983,10 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) @@ -1006,16 +1004,44 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) smjExec = SortMergeJoinExec( exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) - applyEnsureRequirementsWithSubsetKeys(smjExec) match { + EnsureRequirements.apply(smjExec) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _), _), _) => assert(left.expressions === Seq(exprA, exprB, exprB)) assert(right.expressions === Seq(exprA, exprC, exprC)) case other => fail(other.toString) } } + test("KeyedPartitioning: duplicated join keys in hand-built plans do not block SPJ") { + // Queries produce this key list only in unusual configurations: BooleanSimplification + // normally dedups the conjunction, but it is an excludable rule + // (spark.sql.optimizer.excludedRules), and EnsureRequirements must also stay robust + // for hand-built or rewritten plans. The coverage check treats duplicated cluster + // keys as covered, so SPJ is allowed with either config value. + val plan1 = new DummySparkPlanWithBatchScanChild( + outputPartitioning = + KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprB) :: Nil, Seq.empty)) + val plan2 = new DummySparkPlanWithBatchScanChild( + outputPartitioning = + KeyedPartitioning(bucket(4, exprA) :: bucket(4, exprC) :: Nil, Seq.empty)) + val smjExec = SortMergeJoinExec( + exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner, None, plan1, plan2) + Seq("true", "false").foreach { requireAllKeys => + withSQLConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> requireAllKeys) { + EnsureRequirements.apply(smjExec) match { + case SortMergeJoinExec(_, _, _, _, + SortExec(_, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), _), + SortExec(_, _, DummySparkPlan(_, _, right: KeyedPartitioning, _, _), _), _) => + assert(left.expressions === Seq(bucket(4, exprA), bucket(4, exprB))) + assert(right.expressions === Seq(bucket(4, exprA), bucket(4, exprC))) + case other => fail(s"Expected no shuffle, but got: $other") + } + } + } + } + test("SPARK-41413: check compatibility when partition values mismatch") { withSQLConf(SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true") { val leftPartValues = Seq(Array[Any](1, 1), Array[Any](2, 2)).map(new GenericInternalRow(_)) @@ -1130,7 +1156,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { case ShuffledHashJoinExec(_, _, _, _, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _), - DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) => + DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) => assert(left.expressions == a1 :: Nil) assert(attrs == a1 :: Nil) assert(partitionKeys == pks.map(_.row)) @@ -1234,8 +1260,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled to default partitions assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) @@ -1259,8 +1285,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to key mismatch case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1278,7 +1304,7 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), SortExec(_, _, _: DummySparkPlan, _), _) => // Left side shuffled, right side kept as-is case other => fail(s"Expected shuffle on the left side, but got: $other") @@ -1296,8 +1322,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _), _) => // Both sides shuffled due to canCreatePartitioning = false case other => fail(s"Expected shuffles on both sides, but got: $other") } @@ -1371,8 +1397,8 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(join) match { case SortMergeJoinExec(_, _, _, _, - SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _), _), - SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _), _), _) => + SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _, _), _), + SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _, _), _), _) => // Both sides should be shuffled because partition keys not in join keys assert(p1.numPartitions == 10) assert(p2.numPartitions == 10) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/ExistenceJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/ExistenceJoinSuite.scala index 428d29f2989a2..cabc2a1c34c6d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/ExistenceJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/ExistenceJoinSuite.scala @@ -177,6 +177,43 @@ class ExistenceJoinSuite extends SharedSparkSession { } } + // Condition: a = c (equi-key) AND b < 3.0 (left-only) AND d < 4.0 (right-only) + private lazy val mixedResidualCondition = { + And( + And( + EqualTo(left.col("a").expr, right.col("c").expr), + LessThan(left.col("b").expr, Literal(3.0))), + LessThan(right.col("d").expr, Literal(4.0))) + } + + // Condition: a = c (equi-key) AND d < 4.0 (right-only) + private lazy val rightOnlyResidualCondition = { + And( + EqualTo(left.col("a").expr, right.col("c").expr), + LessThan(right.col("d").expr, Literal(4.0))) + } + + // Condition: a = c (equi-key) AND b < 3.0 (left-only): the residual conjunct references + // only the streamed side, so the split hoists it entirely and restCondition is None. + private lazy val leftOnlyResidualCondition = { + And( + EqualTo(left.col("a").expr, right.col("c").expr), + LessThan(left.col("b").expr, Literal(3.0))) + } + + protected def testWithSplitStreamedSideCondOnAndOff( + testName: String)(f: String => Unit): Unit = { + Seq("false", "true").foreach { configValue => + testWithWholeStageCodegenOnAndOff( + s"$testName (splitStreamedSideJoinCondition=$configValue)") { _ => + withSQLConf( + SQLConf.SPLIT_STREAMED_SIDE_JOIN_CONDITION.key -> configValue) { + f(configValue) + } + } + } + } + // Note: the input dataframes and expression must be evaluated lazily because // the SQLContext should be used only within a test to keep SQL tests stable private def testExistenceJoin( @@ -196,16 +233,22 @@ class ExistenceJoinSuite extends SharedSparkSession { val existsAttr = AttributeReference("exists", BooleanType, false)() val leftSemiPlus = ExistenceJoin(existsAttr) def createLeftSemiPlusJoin(join: SparkPlan): SparkPlan = { - val output = join.output.dropRight(1) - val condition = if (joinType == LeftSemi) { - existsAttr - } else { - Not(existsAttr) + joinType match { + case LeftSemi => + val output = join.output.dropRight(1) + ProjectExec(output, FilterExec(existsAttr, join)) + case LeftAnti => + val output = join.output.dropRight(1) + ProjectExec(output, FilterExec(Not(existsAttr), join)) + case _ => + // Only LeftSemi/LeftAnti can be expressed by filtering an ExistenceJoin result. + join } - ProjectExec(output, FilterExec(condition, join)) } - testWithWholeStageCodegenOnAndOff(s"$testName using ShuffledHashJoin") { _ => + val testSemiPlusWrapper = joinType == LeftSemi || joinType == LeftAnti + + testWithSplitStreamedSideCondOnAndOff(s"$testName using ShuffledHashJoin") { _ => extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => @@ -214,17 +257,19 @@ class ExistenceJoinSuite extends SharedSparkSession { leftKeys, rightKeys, joinType, BuildRight, boundCondition, left, right)), expectedAnswer, sortAnswers = true) - checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => - EnsureRequirements.apply( - createLeftSemiPlusJoin(ShuffledHashJoinExec( - leftKeys, rightKeys, leftSemiPlus, BuildRight, boundCondition, left, right))), - expectedAnswer, - sortAnswers = true) + if (testSemiPlusWrapper) { + checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + createLeftSemiPlusJoin(ShuffledHashJoinExec( + leftKeys, rightKeys, leftSemiPlus, BuildRight, boundCondition, left, right))), + expectedAnswer, + sortAnswers = true) + } } } } - testWithWholeStageCodegenOnAndOff(s"$testName using BroadcastHashJoin") { _ => + testWithSplitStreamedSideCondOnAndOff(s"$testName using BroadcastHashJoin") { _ => extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => @@ -233,17 +278,19 @@ class ExistenceJoinSuite extends SharedSparkSession { leftKeys, rightKeys, joinType, BuildRight, boundCondition, left, right)), expectedAnswer, sortAnswers = true) - checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => - EnsureRequirements.apply( - createLeftSemiPlusJoin(BroadcastHashJoinExec( - leftKeys, rightKeys, leftSemiPlus, BuildRight, boundCondition, left, right))), - expectedAnswer, - sortAnswers = true) + if (testSemiPlusWrapper) { + checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + createLeftSemiPlusJoin(BroadcastHashJoinExec( + leftKeys, rightKeys, leftSemiPlus, BuildRight, boundCondition, left, right))), + expectedAnswer, + sortAnswers = true) + } } } } - testWithWholeStageCodegenOnAndOff(s"$testName using SortMergeJoin") { _ => + testWithSplitStreamedSideCondOnAndOff(s"$testName using SortMergeJoin") { _ => extractJoinParts().foreach { case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => @@ -251,45 +298,56 @@ class ExistenceJoinSuite extends SharedSparkSession { SortMergeJoinExec(leftKeys, rightKeys, joinType, boundCondition, left, right)), expectedAnswer, sortAnswers = true) + if (testSemiPlusWrapper) { + checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + createLeftSemiPlusJoin(SortMergeJoinExec( + leftKeys, rightKeys, leftSemiPlus, boundCondition, left, right))), + expectedAnswer, + sortAnswers = true) + } + } + } + } + + Seq("false", "true").foreach { configValue => + test(s"$testName using BroadcastNestedLoopJoin build left" + + s" (splitStreamedSideJoinCondition=$configValue)") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.SPLIT_STREAMED_SIDE_JOIN_CONDITION.key -> configValue) { checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => EnsureRequirements.apply( - createLeftSemiPlusJoin(SortMergeJoinExec( - leftKeys, rightKeys, leftSemiPlus, boundCondition, left, right))), + BroadcastNestedLoopJoinExec(left, right, BuildLeft, joinType, condition)), expectedAnswer, sortAnswers = true) + if (testSemiPlusWrapper) { + checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + createLeftSemiPlusJoin(BroadcastNestedLoopJoinExec( + left, right, BuildLeft, leftSemiPlus, condition))), + expectedAnswer, + sortAnswers = true) + } } } } - test(s"$testName using BroadcastNestedLoopJoin build left") { - withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { - checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => - EnsureRequirements.apply( - BroadcastNestedLoopJoinExec(left, right, BuildLeft, joinType, condition)), - expectedAnswer, - sortAnswers = true) - checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => - EnsureRequirements.apply( - createLeftSemiPlusJoin(BroadcastNestedLoopJoinExec( - left, right, BuildLeft, leftSemiPlus, condition))), - expectedAnswer, - sortAnswers = true) - } - } - - testWithWholeStageCodegenOnAndOff(s"$testName using BroadcastNestedLoopJoin build right") { _ => + testWithSplitStreamedSideCondOnAndOff( + s"$testName using BroadcastNestedLoopJoin build right") { _ => withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => EnsureRequirements.apply( BroadcastNestedLoopJoinExec(left, right, BuildRight, joinType, condition)), expectedAnswer, sortAnswers = true) - checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => - EnsureRequirements.apply( - createLeftSemiPlusJoin(BroadcastNestedLoopJoinExec( - left, right, BuildRight, leftSemiPlus, condition))), - expectedAnswer, - sortAnswers = true) + if (testSemiPlusWrapper) { + checkAnswer2(leftRows, rightRows, (left: SparkPlan, right: SparkPlan) => + EnsureRequirements.apply( + createLeftSemiPlusJoin(BroadcastNestedLoopJoinExec( + left, right, BuildRight, leftSemiPlus, condition))), + expectedAnswer, + sortAnswers = true) + } } } } @@ -401,4 +459,242 @@ class ExistenceJoinSuite extends SharedSparkSession { Some(And(EqualTo(left.col("a").expr, rightUniqueKey.col("c").expr), LessThan(left.col("b").expr, rightUniqueKey.col("d").expr))), Seq(Row(1, 2.0), Row(1, 2.0), Row(3, 3.0), Row(null, null), Row(null, 5.0), Row(6, null))) + + // ---- Tests for streamed-side-only residual predicate hoisting ---- + + // LeftAnti: rows where b >= 3.0 OR no right match with d < 4.0 + // (1, 2.0): no c=1 match -> emitted + // (2, 1.0): match exists -> dropped + // (3, 3.0): b=3.0 >= 3.0 -> emitted + // (null, null): null key -> emitted + // (null, 5.0): b=5.0 >= 3.0 -> emitted + // (6, null): b=null -> emitted + testExistenceJoin( + "test mixed residual condition for left anti join", + LeftAnti, + left, + right, + Some(mixedResidualCondition), + Seq(Row(1, 2.0), Row(1, 2.0), Row(3, 3.0), Row(null, null), Row(null, 5.0), Row(6, null))) + + // LeftOuter: rows where b < 3.0 and right match with d < 4.0 get matched; others emitted as + // null-padded. Equi-matches: (2,1.0)-(2,3.0), (3,3.0)-(3,2.0), (6,null)-(6,null). + // For (2,1.0): b<3.0 true, right d=3.0<4.0 true -> matched output (2,1.0,2,3.0). + // Both sides contain two rows with the matching key, so the Cartesian product emits 4 rows. + // For (3,3.0): b<3.0 false -> emitted as (3,3.0,null,null) + // For (6,null): b<3.0 null -> emitted as (6,null,null,null) + // For (1,2.0): no equi-match -> emitted as (1,2.0,null,null) + // Null keys are emitted as null-padded. + testExistenceJoin( + "test mixed residual condition for left outer join", + LeftOuter, + left, + right, + Some(mixedResidualCondition), + Seq( + Row(1, 2.0, null, null), + Row(1, 2.0, null, null), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(3, 3.0, null, null), + Row(null, null, null, null), + Row(null, 5.0, null, null), + Row(6, null, null, null))) + + // ExistenceJoin with the same mixed residual condition: exists=true only for (2,1.0). + testExistenceJoin( + "test mixed residual condition for existence join", + ExistenceJoin(AttributeReference("exists", BooleanType, false)()), + left, + right, + Some(mixedResidualCondition), + Seq( + Row(1, 2.0, false), + Row(1, 2.0, false), + Row(2, 1.0, true), + Row(2, 1.0, true), + Row(3, 3.0, false), + Row(null, null, false), + Row(null, 5.0, false), + Row(6, null, false))) + + // Inner join is outside the split whitelist: nothing is hoisted and the results must be + // identical with the split config on and off. + testExistenceJoin( + "test mixed residual condition for inner join", + Inner, + left, + right, + Some(mixedResidualCondition), + Seq( + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0))) + + // Right-only residual: nothing is hoistable (the only residual conjunct references the + // non-streamed side), so d < 4.0 must stay in the join condition and the results must be + // identical with the split config on and off. Compared to the mixed condition above, the + // b < 3.0 conjunct is gone, so (3, 3.0) now matches. + + // LeftAnti: (2, 1.0) and (3, 3.0) match -> dropped; (6, null) has a key match but d is + // null -> no qualifying match -> emitted. + testExistenceJoin( + "test right-only residual condition for left anti join", + LeftAnti, + left, + right, + Some(rightOnlyResidualCondition), + Seq(Row(1, 2.0), Row(1, 2.0), Row(null, null), Row(null, 5.0), Row(6, null))) + + // LeftOuter: (2, 1.0) matches both right (2, 3.0) rows (2 left rows -> 4 output rows); + // (3, 3.0) matches (3, 2.0); everything else is emitted null-padded, including (6, null) + // whose key match has d = null. + testExistenceJoin( + "test right-only residual condition for left outer join", + LeftOuter, + left, + right, + Some(rightOnlyResidualCondition), + Seq( + Row(1, 2.0, null, null), + Row(1, 2.0, null, null), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(3, 3.0, 3, 2.0), + Row(null, null, null, null), + Row(null, 5.0, null, null), + Row(6, null, null, null))) + + // ExistenceJoin with the same right-only residual condition: exists=true for (2, 1.0) and + // (3, 3.0). + testExistenceJoin( + "test right-only residual condition for existence join", + ExistenceJoin(AttributeReference("exists", BooleanType, false)()), + left, + right, + Some(rightOnlyResidualCondition), + Seq( + Row(1, 2.0, false), + Row(1, 2.0, false), + Row(2, 1.0, true), + Row(2, 1.0, true), + Row(3, 3.0, true), + Row(null, null, false), + Row(null, 5.0, false), + Row(6, null, false))) + + // Left-only residual: the only residual conjunct references the streamed side, so the + // split hoists it entirely and restCondition is None. These cases cover the fully-hoisted + // shape (SortMergeJoinExec codegen's conditionForCodegen = None branch and the + // interpreted always-true rest walk). The expected answers happen to match the + // mixed-condition ones above because d < 4.0 is true for every equi-match in the + // fixtures; the point of these cases is the restCondition = None code shape. + + // LeftAnti: (2, 1.0) matches -> dropped; (3, 3.0) fails the hoisted b < 3.0; (6, null) + // evaluates it to null; null keys never match. + testExistenceJoin( + "test left-only residual condition for left anti join", + LeftAnti, + left, + right, + Some(leftOnlyResidualCondition), + Seq(Row(1, 2.0), Row(1, 2.0), Row(3, 3.0), Row(null, null), Row(null, 5.0), Row(6, null))) + + // LeftOuter: (2, 1.0) matches both right (2, 3.0) rows (2 left rows -> 4 output rows); + // everything else is emitted null-padded, including (3, 3.0) which fails the hoisted + // b < 3.0 and (6, null) for which it is null. + testExistenceJoin( + "test left-only residual condition for left outer join", + LeftOuter, + left, + right, + Some(leftOnlyResidualCondition), + Seq( + Row(1, 2.0, null, null), + Row(1, 2.0, null, null), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(3, 3.0, null, null), + Row(null, null, null, null), + Row(null, 5.0, null, null), + Row(6, null, null, null))) + + // ExistenceJoin with the same left-only residual condition: exists=true only for (2, 1.0). + testExistenceJoin( + "test left-only residual condition for existence join", + ExistenceJoin(AttributeReference("exists", BooleanType, false)()), + left, + right, + Some(leftOnlyResidualCondition), + Seq( + Row(1, 2.0, false), + Row(1, 2.0, false), + Row(2, 1.0, true), + Row(2, 1.0, true), + Row(3, 3.0, false), + Row(null, null, false), + Row(null, 5.0, false), + Row(6, null, false))) + + // Inner join is outside the split whitelist: with a fully streamed-side-only residual, + // nothing may be hoisted and the results must be identical with the split config on and + // off. + testExistenceJoin( + "test left-only residual condition for inner join", + Inner, + left, + right, + Some(leftOnlyResidualCondition), + Seq( + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0))) + + // RightOuter streams the right side, so here d < 4.0 is the streamed-side-only conjunct + // and IS hoistable (the mirror of the left-streamed cases above). The hash joins use + // BuildLeft so that the right side is the streamed one; right rows with no qualifying + // left match are emitted null-padded. + testWithSplitStreamedSideCondOnAndOff( + "test right-only residual condition for right outer join") { _ => + val join = Join(left.logicalPlan, right.logicalPlan, + Inner, Some(rightOnlyResidualCondition), JoinHint.NONE) + ExtractEquiJoinKeys.unapply(join).foreach { + case (_, leftKeys, rightKeys, boundCondition, _, _, _, _) => + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val expected = Seq( + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(2, 1.0, 2, 3.0), + Row(3, 3.0, 3, 2.0), + Row(null, null, 4, 1.0), + Row(null, null, null, null), + Row(null, null, null, 5.0), + Row(null, null, 6, null)) + checkAnswer2(left, right, (l: SparkPlan, r: SparkPlan) => + EnsureRequirements.apply( + SortMergeJoinExec(leftKeys, rightKeys, RightOuter, boundCondition, l, r)), + expected, sortAnswers = true) + checkAnswer2(left, right, (l: SparkPlan, r: SparkPlan) => + EnsureRequirements.apply( + BroadcastHashJoinExec( + leftKeys, rightKeys, RightOuter, BuildLeft, boundCondition, l, r)), + expected, sortAnswers = true) + checkAnswer2(left, right, (l: SparkPlan, r: SparkPlan) => + EnsureRequirements.apply( + ShuffledHashJoinExec( + leftKeys, rightKeys, RightOuter, BuildLeft, boundCondition, l, r)), + expected, sortAnswers = true) + } + } + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SplitStreamedSideJoinConditionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SplitStreamedSideJoinConditionSuite.scala new file mode 100644 index 0000000000000..ec3f8d3c8c439 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/joins/SplitStreamedSideJoinConditionSuite.scala @@ -0,0 +1,245 @@ +/* + * 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.execution.joins + +import org.apache.spark.sql.{DataFrame, QueryTest} +import org.apache.spark.sql.catalyst.expressions.{Add, And, LessThan, Literal, Rand} +import org.apache.spark.sql.catalyst.plans.LeftOuter +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests for streamed-side join condition hoisting in whole-stage codegen. The + * generated-code shape under test requires all of the following: + * - spark.sql.join.splitStreamedSideJoinCondition = true + * - a join type that preserves streamed rows (LeftOuter, RightOuter, LeftAnti, ExistenceJoin) + * - a broadcast hash join, so the scan and the join share one whole-stage. Sort-merge and + * shuffled hash joins are out of scope here: their streamed side crosses a Sort or an + * exchange, which advances its cursor before running the inlined consume code, and the + * sort-merge guard is a standalone pre-loop guard that emits and continues before the + * match loop; it is not folded into the match condition. The relevant contrast is that + * it uses continue in its own loop instead of returning from an inlined consumer. + * ExistenceJoinSuite covers those join implementations functionally, with the + * config on and off. + * - a batching streamed producer (vectorized parquet scan -> ColumnarToRowExec) whose loop + * cursor is only written back after the inlined consume code + * - a streamed column (pad) not referenced by the join keys/condition, so + * WholeStageCodegenExec.consume does not wrap the join's doConsume in a function + * - a residual filter that references pad (pad % 2 = 0 is not pushed into the scan), so pad + * is materialized before the join's doConsume and the generated code compiles + * + * Expected results are produced by running the same query with the hoisting disabled. The + * LIMIT below ORDER BY on the hoisting-enabled run is load-bearing: an early-returning guard + * makes the generated code reprocess the same batch forever; the limit is what makes the + * query terminate (with wrong results). + */ +class SplitStreamedSideJoinConditionSuite extends QueryTest with SharedSparkSession { + + // streamed_t: (id, a, pad) with id 0..7, a = id, pad = id * 10 + // build_t: (bid, b) with bid 0..3, b = bid * 100 + private def withStreamedAndBuildTables(f: => Unit): Unit = { + withTempPath { path => + spark.range(8).selectExpr("id", "id AS a", "id * 10 AS pad") + .write.parquet(path.getCanonicalPath + "/streamed") + spark.range(4).selectExpr("id AS bid", "id * 100 AS b") + .write.parquet(path.getCanonicalPath + "/build") + spark.read.parquet(path.getCanonicalPath + "/streamed") + .createOrReplaceTempView("streamed_t") + spark.read.parquet(path.getCanonicalPath + "/build") + .createOrReplaceTempView("build_t") + f + } + } + + // Compares the query with hoisting enabled against the same query with hoisting disabled, + // sorting both sides by all columns. + private def checkDataFrame(query: => DataFrame): Unit = { + val expected = withSQLConf(SQLConf.SPLIT_STREAMED_SIDE_JOIN_CONDITION.key -> "false") { + sortByAllColumns(query).collect().toSeq + } + val df = query + assert(df.queryExecution.executedPlan.toString.contains("BroadcastHashJoin"), + s"expected a broadcast hash join in ${df.queryExecution.executedPlan}") + checkAnswer(sortByAllColumns(df.limit(100)), expected) + } + + private def sortByAllColumns(df: DataFrame): DataFrame = { + val cols = df.columns.toSeq + df.orderBy(cols.head, cols.tail: _*) + } + + private def testWithCodegenOnAndOff(testName: String)( + queryDf: => DataFrame, + codegenOnly: Boolean = false): Unit = { + testWithWholeStageCodegenOnAndOff(testName) { _ => + val confs = Seq( + SQLConf.SPLIT_STREAMED_SIDE_JOIN_CONDITION.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") ++ + (if (codegenOnly) Seq(SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY") else Nil) + withSQLConf(confs: _*) { + withStreamedAndBuildTables { + checkDataFrame(queryDf) + } + } + } + } + + // Streamed-only conjuncts that can throw or are non-deterministic must stay in the + // residual: hoisting evaluates the hoisted part for every streamed row, including rows + // that have no buffered match and would never evaluate the conjunct otherwise. + + test("split does not hoist conjuncts that can throw") { + val streamed = spark.range(8).selectExpr("id", "id AS a").queryExecution.executedPlan + val a = streamed.output.find(_.name == "a").get + val hoistable = LessThan(a, Literal(5L)) + // Add can throw on overflow in ANSI mode but inherits throwable = false from its + // non-throwing children. It must stay in the residual, which is why a !throwable + // gate is not sufficient and the hoisted part is restricted to expression families + // with a reliable non-throwing contract. + val unmarkedThrowable = LessThan(Add(a, Literal(1L)), Literal(5L)) + assert(!unmarkedThrowable.throwable) + val (streamedOnly, rest) = StreamedSideJoinCondition.split( + Some(And(hoistable, unmarkedThrowable)), LeftOuter, streamed, splitEnabled = true) + assert(streamedOnly.contains(hoistable)) + assert(rest.contains(unmarkedThrowable)) + } + + test("split does not hoist non-deterministic conjuncts") { + val streamed = spark.range(8).selectExpr("id", "id AS a").queryExecution.executedPlan + val a = streamed.output.find(_.name == "a").get + val hoistable = LessThan(a, Literal(5L)) + val nonDeterministic = LessThan(Rand(Literal(0L)), Literal(0.5)) + val (streamedOnly, rest) = StreamedSideJoinCondition.split( + Some(And(hoistable, nonDeterministic)), LeftOuter, streamed, splitEnabled = true) + assert(streamedOnly.contains(hoistable)) + assert(rest.contains(nonDeterministic)) + } + + // The streamed-only condition a < 5 is false for ids 5..7, so the hoisted guard fires on + // those rows. pad is selected but not referenced by the join. + + testWithCodegenOnAndOff( + "left outer join, unique build key")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, s.pad, b.bid + |FROM (SELECT * FROM streamed_t WHERE pad % 2 = 0) s + |LEFT OUTER JOIN (SELECT DISTINCT bid FROM build_t) b + |ON s.id = b.bid AND s.a < 5 + """.stripMargin)) + + testWithCodegenOnAndOff( + "left outer join, non-unique build key")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, s.pad, b.bid, b.b + |FROM (SELECT * FROM streamed_t WHERE pad % 2 = 0) s + |LEFT OUTER JOIN (SELECT bid, b FROM build_t UNION ALL SELECT bid, b FROM build_t) b + |ON s.id = b.bid AND s.a < 5 + """.stripMargin)) + + // The "compiles without pre-materialized column" variants drop the residual filter, so pad + // stays lazy through the join and is only materialized by the join's own consume code. They + // exercise how the generated guard / probe code scopes that materialization: if pad's + // declaration lands in a different (nested) scope than its uses, the generated code does + // not compile. CODEGEN_ONLY turns the silent fallback to interpreted execution into an + // error so these tests actually assert compilation succeeds. + testWithCodegenOnAndOff( + "left outer join compiles without pre-materialized column")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, s.pad, b.bid, b.b + |FROM streamed_t s + |LEFT OUTER JOIN build_t b + |ON s.id = b.bid AND s.a < 5 + """.stripMargin), + codegenOnly = true) + + testWithCodegenOnAndOff( + "right outer join")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ b.bid, b.b, s.id, s.a, s.pad + |FROM build_t b + |RIGHT OUTER JOIN (SELECT * FROM streamed_t WHERE pad % 2 = 0) s + |ON b.bid = s.id AND s.a < 5 + """.stripMargin)) + + testWithCodegenOnAndOff( + "left anti join")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, s.pad + |FROM (SELECT * FROM streamed_t WHERE pad % 2 = 0) s + |LEFT ANTI JOIN build_t b + |ON s.id = b.bid AND s.a < 5 + """.stripMargin)) + + testWithCodegenOnAndOff( + "left anti join compiles without pre-materialized column")( + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, s.pad + |FROM streamed_t s + |LEFT ANTI JOIN build_t b + |ON s.id = b.bid AND s.a < 5 + """.stripMargin), + codegenOnly = true) + + // The EXISTS flag is projected (not filtered on) so guard-fired rows (ids 5..7) are + // emitted via the guard path with e = false, making any replay of those rows visible as + // duplicates. Filtering on EXISTS instead would drop them and mask a premature return in + // the guard: the stage buffer stays empty, BufferedRowIterator.hasNext treats that as + // end-of-stream, and since the abandoned rows produce no output anyway, the query would + // return the right answer for the wrong reason. + testWithCodegenOnAndOff( + "existence join")( + spark.sql(""" + |SELECT s.id, s.a, s.pad, + | EXISTS (SELECT /*+ BROADCAST(b) */ 1 FROM build_t b WHERE b.bid = s.id AND s.a < 5) AS e + |FROM (SELECT * FROM streamed_t WHERE pad % 2 = 0) s + """.stripMargin)) + + testWithCodegenOnAndOff( + "existence join compiles without pre-materialized column")( + spark.sql(""" + |SELECT s.id, s.a, s.pad, + | EXISTS (SELECT /*+ BROADCAST(b) */ 1 FROM build_t b WHERE b.bid = s.id AND s.a < 5) AS e + |FROM streamed_t s + """.stripMargin), + codegenOnly = true) + + // Left outer join whose ON clause has two streamed-side-only conjuncts: a hoistable + // a < 5 and a real registered UDF that returns true when its input is true and throws + // when it is false or null. The UDF conjunct holds for the matched rows (ids 0..3) and + // throws for the unmatched ones (ids 4..7), so it throws exactly when it is wrongly + // hoisted and evaluated before probing. A ScalaUDF does not override `throwable`, so + // the old !throwable gate hoisted it; the whitelist keeps it in the residual. + private def throwingUdfJoin: DataFrame = { + spark.udf.register("throw_unless_true", (v: Boolean) => + if (v) true else throw new RuntimeException("throw_unless_true evaluated to false")) + spark.sql(""" + |SELECT /*+ BROADCAST(b) */ s.id, s.a, b.bid, b.b + |FROM streamed_t s + |LEFT OUTER JOIN build_t b + |ON s.id = b.bid AND s.a < 5 AND throw_unless_true(s.a < 4) + """.stripMargin) + } + + // End-to-end: the throwing UDF conjunct must not run for streamed rows 4..7, which have + // no buffered match. The hoistable conjunct a < 5 keeps the hoisted guard on the path + // so the test proves the UDF conjunct is excluded from it, not that hoisting is off. + // Before the fix, rows 4..7 evaluate the hoisted UDF conjunct and throw. + testWithCodegenOnAndOff( + "throwing UDF conjunct is not evaluated for unmatched rows")( + throwingUdfJoin) +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala index 033a0eea5b5ae..7801188d532c1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/SQLLastAttemptMetricPlanShapesSuite.scala @@ -29,7 +29,9 @@ import org.apache.spark.sql.execution.exchange._ import org.apache.spark.sql.functions.udf import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.tags.ExtendedSQLTest +@ExtendedSQLTest class SQLLastAttemptMetricPlanShapesSuite extends SharedSparkSession with SQLMetricsTestUtils diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala new file mode 100644 index 0000000000000..86bbe951ab9b1 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/DSv2PlanMergingSuite.scala @@ -0,0 +1,257 @@ +/* + * 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.execution.planmerging + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.{DataFrame, QueryTest, Row} +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.connector.FakeV2ProviderWithCustomSchema +import org.apache.spark.sql.connector.catalog.{InMemoryScanMergingPartitionFilterCatalog, InMemoryScanMergingReportingCatalog} +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, DataSourceV2ScanRelation} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession + +/** + * End-to-end test for the DSv2 scan-merge gap this change closes: a source whose (equal) filter is + * strict only via the iterative PartitionPredicate second pass. When [[PlanMerger]] rebuilds the + * merged scan it drives the real + * [[org.apache.spark.sql.execution.datasources.v2.V2ScanRelationPushDown]] (Approach 2), so the + * second pass runs and the filter comes back strict -- the merge proceeds. A hand-rolled single + * push-predicates call would only run the first pass, see the filter come back post-scan, + * mis-classify it as non-strict and decline the merge (leaving two scans). + */ +class DSv2PlanMergingSuite extends QueryTest with SharedSparkSession + with BeforeAndAfter { + + private val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + private val tbl = "scanmerge.t" + private val tbl2 = "scanmerge.t2" + + before { + spark.conf.set("spark.sql.catalog.scanmerge", + classOf[InMemoryScanMergingPartitionFilterCatalog].getName) + spark.conf.set("spark.sql.catalog.scanmergereport", + classOf[InMemoryScanMergingReportingCatalog].getName) + } + + after { + spark.sessionState.catalogManager.reset() + spark.conf.unset("spark.sql.catalog.scanmerge") + spark.conf.unset("spark.sql.catalog.scanmergereport") + } + + private def v2Scans(df: DataFrame): Seq[DataSourceV2ScanRelation] = + df.queryExecution.optimizedPlan.collectWithSubqueries { + case s: DataSourceV2ScanRelation => s + } + + // A successful DSv2 merge builds the scan and leaves NO bare DataSourceV2Relation in the plan. + // A leaked deferred scan (e.g. if a future recursion arm forwarded `deferredScan` without + // building it) would surface as an unbuilt placeholder relation the read path cannot plan -- + // cheap guard so a regression fails loudly here, not as a silently-declined merge (see + // MergeContext). + private def assertNoPlaceholderRelation(df: DataFrame): Unit = + assert( + df.queryExecution.optimizedPlan.collectWithSubqueries { + case r: DataSourceV2Relation => r + }.isEmpty, + s"unbuilt placeholder DataSourceV2Relation left in plan:\n${df.queryExecution.optimizedPlan}") + + test("SPARK-40259: merge two DSv2 scans whose filter is strict only via the second pass") { + withTable(tbl) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1, 10), ('a', 2, 20), ('b', 3, 30)") + + // `part_col IN ('a')` is untranslatable/returned in the first pass and only accepted (strict) + // in the iterative PartitionPredicate second pass. The two scalar subqueries differ only in + // their projected data column, so PlanMerger fuses them into a single scan reading {c1, c2}. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE part_col IN ('a')) AS m1, + | (SELECT max(c2) FROM $tbl WHERE part_col IN ('a')) AS m2 + |""".stripMargin) + + // Correctness: the merged scan must still enforce the partition filter. Reading all + // partitions would give (3, 30) instead of the filtered (2, 20). + checkAnswer(df, Row(2, 20)) + + // The merged subquery is referenced once per scalar subquery, so the logical plan duplicates + // it (physical planning reuses it). Dedupe by canonical form: a successful merge leaves a + // single distinct scan reading the union of both columns; declining would leave two distinct + // scans, one per column. + val scans = v2Scans(df) + assert(scans.nonEmpty, s"expected a DSv2 scan:\n${df.queryExecution.optimizedPlan}") + assert(scans.map(_.canonicalized).distinct.length == 1, + s"the two scans should be fused into one:\n${df.queryExecution.optimizedPlan}") + val scan = scans.head + assert(scan.output.map(_.name).toSet == Set("c1", "c2"), + s"the merged scan should read the union of both columns; got ${scan.output}") + // The filter must be re-enforced strictly on the merged scan (present in pushedFilters), + // which only happens because the rebuild runs the iterative second pass. + assert(scan.pushedFilters.exists(_.references.exists(_.name == "part_col")), + s"the part_col filter should be re-pushed strict onto the merged scan; " + + s"got pushedFilters=${scan.pushedFilters.mkString("[", ", ", "]")}") + assertNoPlaceholderRelation(df) + } + } + + test("SPARK-40259: do not merge DSv2 scans with different strict partition filters") { + withTable(tbl) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1, 10), ('a', 2, 20), ('b', 3, 30)") + + // Same table and shape, but different (strict, fully-enforced) partition filters. The strict + // filters are unequal, so the two scans must NOT fuse -- otherwise one side's partition would + // be read for both, giving a wrong answer. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE part_col IN ('a')) AS m1, + | (SELECT max(c2) FROM $tbl WHERE part_col IN ('b')) AS m2 + |""".stripMargin) + + checkAnswer(df, Row(2, 30)) + + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 2, + s"scans with different partition filters must not be fused:\n" + + df.queryExecution.optimizedPlan) + } + } + + test("SPARK-40259: merge two DSv2 scans with differing best-effort filters (OR-widen)") { + withTable(tbl) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1, 10), ('a', 2, 20), ('b', 3, 30)") + + // c1 > 1 / c2 > 1 are data-column filters this source does not enforce (best-effort, + // post-scan), so the two scans carry no strict filter and fuse into one reading {c1, c2}. + // The differing post-scan filters are OR-widened above the merged scan (Phase 2) and each + // aggregate keeps its own filter. This runs the differing-filter path end to end, which the + // plan-shape tests in MergeSubplansSuite cover only structurally. (A rand() on one side is + // not added here: it would make the scalar subquery non-deterministic, so MergeSubplans would + // not extract it at all; the non-deterministic pruning drop is unit-covered in + // MergeSubplansSuite.) + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE c1 > 1) AS m1, + | (SELECT max(c2) FROM $tbl WHERE c2 > 1) AS m2 + |""".stripMargin) + + // c1 > 1 -> {2, 3} -> max 3; c2 > 1 -> {10, 20, 30} -> max 30. + checkAnswer(df, Row(3, 30)) + + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 1, + s"the two scans should be fused into one:\n${df.queryExecution.optimizedPlan}") + assert(scans.head.output.map(_.name).toSet == Set("c1", "c2"), + s"the merged scan should read the union of both columns; got ${scans.head.output}") + assertNoPlaceholderRelation(df) + } + } + } + + test("SPARK-40259: do not merge DSv2 scans from different tables") { + withTable(tbl, tbl2) { + sql(s"CREATE TABLE $tbl (part_col string, c1 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"CREATE TABLE $tbl2 (part_col string, c2 int) USING $v2Source " + + "PARTITIONED BY (part_col)") + sql(s"INSERT INTO $tbl VALUES ('a', 1), ('a', 2)") + sql(s"INSERT INTO $tbl2 VALUES ('a', 10), ('a', 30)") + + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $tbl WHERE part_col IN ('a')) AS m1, + | (SELECT max(c2) FROM $tbl2 WHERE part_col IN ('a')) AS m2 + |""".stripMargin) + + checkAnswer(df, Row(2, 30)) + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 2, + s"scans from different tables must remain separate:\n${df.queryExecution.optimizedPlan}") + } + } + + test("SPARK-58549: a scan merge preserves the sources' reported key-grouped partitioning") { + val t = "scanmergereport.t2" + withTable(t) { + // V2_BUCKETING_ENABLED is what makes the preserved report reach the physical plan (only + // DataSourceV2ScanExecBase.outputPartitioning reads it), which the last assertion checks; AQE + // is off so that `executedPlan` is the plan those assertions can walk. + withSQLConf( + SQLConf.V2_BUCKETING_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + sql(s"CREATE TABLE $t (c1 int, c2 int) USING $v2Source PARTITIONED BY (c1)") + sql(s"INSERT INTO $t VALUES (1, 10), (2, 20), (3, 30)") + + // Both scalar subqueries read the partition column c1, so each scan reports + // KeyGroupedPartitioning on c1; they differ in the extra column read, so PlanMerger fuses + // them into one scan reading {c1, c2}. c1 survives in the union, so the merged scan + // re-derives the same partitioning -- not a degradation, so the merge proceeds by default. + val df = sql( + s""" + |SELECT + | (SELECT max(c1) FROM $t) AS m1, + | (SELECT max(c1 + c2) FROM $t) AS m2 + |""".stripMargin) + checkAnswer(df, Row(3, 33)) + + val scans = v2Scans(df) + assert(scans.map(_.canonicalized).distinct.length == 1, + s"the two scans should be fused into one:\n${df.queryExecution.optimizedPlan}") + val scan = scans.head + assert(scan.output.map(_.name).toSet == Set("c1", "c2"), + s"the merged scan should read the union of both columns; got ${scan.output}") + assert(scan.keyGroupedPartitioning.exists(_.nonEmpty), + s"the merged scan should preserve the reported key-grouped partitioning; " + + s"got ${scan.keyGroupedPartitioning}") + assert(scan.keyGroupedPartitioning.get.flatMap(_.references).exists(_.name == "c1"), + s"the preserved partitioning should be on c1; got ${scan.keyGroupedPartitioning}") + assertNoPlaceholderRelation(df) + + // The preserved report is not just carried on the logical node: the merged scan still + // reports it as its physical output partitioning, which is what lets a storage-partitioned + // join above it skip the shuffle. + val batchScans = df.queryExecution.executedPlan.collectWithSubqueries { + case b: BatchScanExec => b + } + assert(batchScans.nonEmpty, s"expected a BatchScanExec:\n${df.queryExecution.executedPlan}") + batchScans.foreach { b => + b.outputPartitioning match { + case k: KeyedPartitioning => + assert(k.expressions.flatMap(_.references).exists(_.name == "c1"), + s"the merged scan should be key-grouped on c1; got ${k.expressions}") + case other => + fail(s"expected a KeyedPartitioning on the merged scan, got $other") + } + } + } + } + } +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala similarity index 54% rename from sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala rename to sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala index dfa17f926c5a7..670366ee87ed2 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/MergeSubplansSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala @@ -15,15 +15,25 @@ * limitations under the License. */ -package org.apache.spark.sql.catalyst.optimizer +package org.apache.spark.sql.execution.planmerging import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, CreateNamedStruct, GetStructField, If, Literal, Or, ScalarSubquery} +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending, Attribute, AttributeReference, CreateNamedStruct, ExprId, GetStructField, If, Literal, Or, ScalarSubquery, SortOrder, TransformExpression} import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules._ +import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes +import org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier, SupportsRead, Table, TableCapability} +import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, ScalarFunction, UnboundFunction} +import org.apache.spark.sql.connector.expressions.{Expressions, FieldReference, SortDirection => V2SortDirection, SortOrder => V2SortOrder} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownLimit, SupportsPushDownOffset, SupportsPushDownRequiredColumns, SupportsPushDownTableSample, SupportsPushDownTopN, SupportsPushDownV2Filters, SupportsReportOrdering, SupportsReportPartitioning} +import org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning => V2KeyGroupedPartitioning, Partitioning => V2Partitioning, UnknownPartitioning => V2UnknownPartitioning} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown} import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{DataType, IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap class MergeSubplansSuite extends PlanTest { @@ -1966,4 +1976,1282 @@ class MergeSubplansSuite extends PlanTest { comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) } } + + // ---- SPARK-40259: generic DSv2 scan merge ---- + + private val v2Table = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType)))) + + /** Same shape as [[v2Table]] but does NOT declare the `SCAN_MERGING` capability. */ + private val v2TableNoMerge = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType))), + supportsScanMerging = false) + + /** + * Same shape as [[v2Table]], but its scans report `a ASC` as their output ordering, so a scan + * rebuilt over this table re-derives a report instead of none. + */ + private val v2TableReportingA = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType))), + reportedOrderingCols = Seq("a")) + + /** + * Same shape as [[v2Table]], but its scans report `bucket(4, a)` as their key-grouped + * partitioning -- a TRANSFORM report, which resolves to a `TransformExpression` rather than a + * plain attribute (see [[v2ScanReportingOn]]). + */ + private val v2TableBucketedOnA = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType))), + reportedBucketPartition = Some((4, "a"))) + + /** + * Same shape as [[v2Table]], but its scans report `a ASC` ONLY when no filter was pushed into + * them, so the deferred build's two attempts (strict + best-effort, then strict-only) re-derive + * different reports. + */ + private val v2TablePushSensitiveReport = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType))), + reportedOrderingCols = Seq("a"), reportOnlyWhenNothingPushed = true) + + /** A `DataSourceV2ScanRelation` over `table` projecting only the given columns. */ + private def v2ScanReadingOn(table: TestV2Table, cols: Seq[String]): DataSourceV2ScanRelation = { + val fullOutput = toAttributes(table.schema()) + val relation = + DataSourceV2Relation(table, fullOutput, None, None, CaseInsensitiveStringMap.empty()) + val output = cols.map(c => fullOutput.find(_.name == c).get) + val scan = TestV2Scan(StructType(output.map(a => StructField(a.name, a.dataType, a.nullable)))) + // These stand in for a plain scan produced by the column-pruning path, which is mergeable. + DataSourceV2ScanRelation(relation, scan, output, mergeableScan = true) + } + + /** A `DataSourceV2ScanRelation` over [[v2Table]] projecting only the given columns. */ + private def v2ScanReading(cols: String*): DataSourceV2ScanRelation = + v2ScanReadingOn(v2Table, cols) + + /** + * A `DataSourceV2ScanRelation` over `table` whose reported partitioning/ordering is derived the + * way the optimizer derives it: by running `V2ScanPartitioningAndOrdering` over the table's own + * scan, resolving transforms through a function catalog. Every call therefore binds the transform + * function afresh, so the two sides of a merge hold DIFFERENT `BoundFunction` instances -- + * exactly as in production, where each scan relation is annotated by its own derivation. + */ + private def v2ScanReportingOn( + table: TestV2Table, + cols: Seq[String], + funCatalog: FunctionCatalog = TestFreshBindFunctionCatalog): DataSourceV2ScanRelation = { + val fullOutput = toAttributes(table.schema()) + val relation = DataSourceV2Relation( + table, fullOutput, Some(funCatalog), None, CaseInsensitiveStringMap.empty()) + val output = cols.map(c => fullOutput.find(_.name == c).get) + val builder = table.newScanBuilder(CaseInsensitiveStringMap.empty()) + builder.asInstanceOf[SupportsPushDownRequiredColumns].pruneColumns( + StructType(output.map(a => StructField(a.name, a.dataType, a.nullable)))) + V2ScanPartitioningAndOrdering( + DataSourceV2ScanRelation(relation, builder.build(), output, mergeableScan = true)) + .asInstanceOf[DataSourceV2ScanRelation] + } + + /** A `DataSourceV2ScanRelation` over a table without the `SCAN_MERGING` capability. */ + private def v2ScanReadingNoMerge(cols: String*): DataSourceV2ScanRelation = + v2ScanReadingOn(v2TableNoMerge, cols) + + private def v2Scans(plan: LogicalPlan): Seq[DataSourceV2ScanRelation] = + plan.collectWithSubqueries { case s: DataSourceV2ScanRelation => s } + + /** + * Normalizes a merged plan so `comparePlans` can match it: (1) drops each DSv2 scan's + * dynamically-built Scan to a schema-only placeholder (the pushed describe() strings are asserted + * separately), and (2) resets the nested DataSourceV2Relation's output exprIds to a positional + * scheme. The reset is needed because `DataSourceV2ScanRelation` is a `LeafNode`, so its + * `relation` is a constructor arg rather than a child -- PlanTest's `normalizeExprIds` never + * recurses into it, and those exprIds otherwise differ between the expected and actual plans. + */ + private def normalizeScans(plan: LogicalPlan): LogicalPlan = plan.transformWithSubqueries { + case s: DataSourceV2ScanRelation => + val fixedRelation = s.relation.copy( + output = s.relation.output.zipWithIndex.map { case (a, i) => a.withExprId(ExprId(i)) }) + s.copy(relation = fixedRelation, scan = TestV2Scan(s.scan.readSchema())) + } + + // Normalize merged DSv2 scans before the standard comparison (a no-op on plans without them), + // so tests can call plain comparePlans. See normalizeScans for why this is needed. + override protected def comparePlans( + plan1: LogicalPlan, + plan2: LogicalPlan, + checkAnalysis: Boolean = true): Unit = + super.comparePlans(normalizeScans(plan1), normalizeScans(plan2), checkAnalysis) + + test("SPARK-40259: merge DSv2 scans that differ only in projected columns") { + val sub1 = ScalarSubquery(v2ScanReading("a").groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(v2ScanReading("b").groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the two scans fuse into one reading {a, b}, feeding a single aggregate. No filters, + // so no OR-widen propagation. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val mergedSubquery = mergedScan + .groupBy()(sum($"a").as("sum_a"), sum($"b").as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + + test("SPARK-40259: do not merge DSv2 scans when a pushdown is merge-blocking") { + val sub1 = ScalarSubquery(v2ScanReading("a").groupBy()(sum($"a").as("sum_a"))) + val blockingScan = v2ScanReading("b").copy(mergeableScan = false) + val sub2 = ScalarSubquery(blockingScan.groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // A merge-blocking pushdown declines the merge: the plan is left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: do not merge DSv2 scans when a column is read at a nested-pruned type") { + // A struct column read at a narrower (nested-pruned) type carries GetStructField ordinals + // against that narrow layout. The merged scan reads the full struct, so remapping those + // ordinals onto it would read the wrong field -- the merge must decline. np reads s pruned to + // <b>, cp reads a plain column x; without the guard they fuse into a scan of {s:<a,b>, x}. + val schema = StructType(Seq( + StructField("s", StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType)))), + StructField("x", IntegerType))) + val relation = DataSourceV2Relation( + new TestV2Table(schema), toAttributes(schema), None, None, CaseInsensitiveStringMap.empty()) + val sFull = relation.output.find(_.name == "s").get + val xFull = relation.output.find(_.name == "x").get + + // np reads s at struct<b> (nested-pruned); cp reads x. Both opt in and are otherwise mergeable. + val sPruned = AttributeReference( + "s", StructType(Seq(StructField("b", IntegerType))))(sFull.exprId) + val npScan = DataSourceV2ScanRelation(relation, + TestV2Scan(StructType(Seq(StructField("s", sPruned.dataType)))), + Seq(sPruned), mergeableScan = true) + val cpScan = DataSourceV2ScanRelation(relation, + TestV2Scan(StructType(Seq(StructField("x", IntegerType)))), + Seq(xFull), mergeableScan = true) + + val sub1 = ScalarSubquery(npScan.groupBy()(sum(GetStructField(sPruned, 0)).as("sb"))) + val sub2 = ScalarSubquery(cpScan.groupBy()(sum(xFull).as("sx"))) + val originalQuery = testRelation.select(sub1, sub2) + + // The nested-pruned column declines the merge: the plan is left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: merge DSv2 scans with different filters via OR-widen propagation") { + val sub1 = ScalarSubquery( + v2ScanReading("a").where($"a" > 1).groupBy()(sum($"a").as("sum_a"))) // cp + val sub2 = ScalarSubquery( + v2ScanReading("b").where($"b" > 2).groupBy()(sum($"b").as("sum_b"))) // np + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the two scans fuse into one reading {a, b}; the differing filters become + // propagatedFilter aliases OR-widened in a Filter, and each side's aggregate carries its + // filter as a FILTER clause (np = sub2 gets id 0, cp = sub1 gets id 1). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("sum_a"), + sum($"b", Some(npFilter)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the merged scan re-pushes OR(a > 1, b > 2) for row-group pruning, which the + // structural comparison above normalizes out. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed + assert(pushed.nonEmpty && pushed.mkString.contains("a") && pushed.mkString.contains("b"), + s"expected an OR pruning predicate over a and b pushed to the merged scan, got: $pushed") + } + } + + test("SPARK-40259: DSv2 symmetric merge honors its own config when the general one is off") { + // Two DSv2 scans with equal (empty) strict filters but differing best-effort filters. The + // general symmetric propagation is off, so a non-DSv2 pair would not merge; the DSv2-specific + // config allows it, since the equal strict filters mean both sides read the same base set and + // only the best-effort OR pruning is broadened. + val sub1 = ScalarSubquery( + v2ScanReading("a").where($"a" > 1).groupBy()(sum($"a").as("sum_a"))) // cp + val sub2 = ScalarSubquery( + v2ScanReading("b").where($"b" > 2).groupBy()(sum($"b").as("sum_b"))) // np + val originalQuery = testRelation.select(sub1, sub2) + + // Expected when the DSv2 config allows it: the two scans fuse into one reading {a, b}; the + // differing filters become propagatedFilter aliases OR-widened in a Filter, and each side's + // aggregate carries its filter as a FILTER clause (np = sub2 gets id 0, cp = sub1 gets id 1). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("sum_a"), + sum($"b", Some(npFilter)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false", + SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + + // With the DSv2 config also off, the differing-filter pair must NOT merge (plan unchanged). + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false", + SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false") { + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + } + + test("SPARK-40259: DSv2 merge fact propagates to an outer stacked Filter pair") { + // Stacked Filters over each DSv2 scan (cp = sub1: inner a>1, outer a<100; np = sub2: inner b>2, + // outer b<200). The leaf scans merge and the deferral is consumed at the INNER (Filter, Filter) + // pair; the OUTER pair sees dsv2DeferredScan = None. The merge only fuses if the "a DSv2 merge + // happened below" fact (TryMergeResult.dsv2Merged) still reaches the outer pair so it applies + // the DSv2-symmetric exemption -- dsv2DeferredScan alone (consumed at the inner pair) would + // not. Optimize runs only MergeSubplans, so the stacked Filters are not combined before the + // rule sees them. + val sub1 = ScalarSubquery( + v2ScanReading("a").where($"a" > 1).where($"a" < 100).groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery( + v2ScanReading("b").where($"b" > 2).where($"b" < 200).groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected when the DSv2 config is on: the scans fuse into one reading {a, b}. Merge traversal + // (inner-to-outer): + // Inner pair (np: b>2, cp: a>1) OR-widens: + // f0 = Alias(b > 2, "propagatedFilter_0") -- np / sum_b + // f1 = Alias(a > 1, "propagatedFilter_1") -- cp / sum_a + // Outer pair (np: b<200, cp: a<100) OR-widens, AND-combining each side's inner filter: + // f2 = Alias(AND(f0, b < 200), "propagatedFilter_2") -- np + // f3 = Alias(AND(f1, a < 100), "propagatedFilter_3") -- cp + // Aggregate consumes f2/f3 as FILTER clauses: sum_a FILTER f3, sum_b FILTER f2 -- i.e. each + // side's per-side AND(inner, outer). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val f0Alias = Alias($"b" > 2, "propagatedFilter_0")() + val f1Alias = Alias($"a" > 1, "propagatedFilter_1")() + val f0 = f0Alias.toAttribute + val f1 = f1Alias.toAttribute + val innerFilter = mergedScan + .select(mergedScan.output ++ Seq(f0Alias, f1Alias): _*) + .where(Or(f0, f1)) + val f2Alias = Alias(And(f0, $"b" < 200), "propagatedFilter_2")() + val f3Alias = Alias(And(f1, $"a" < 100), "propagatedFilter_3")() + val f2 = f2Alias.toAttribute + val f3 = f3Alias.toAttribute + val mergedSubquery = innerFilter + .select(innerFilter.output ++ Seq(f2Alias, f3Alias): _*) + .where(Or(f2, f3)) + .groupBy()( + sum($"a", Some(f3)).as("sum_a"), + sum($"b", Some(f2)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + // General symmetric off; the DSv2-specific config on. The flag lets BOTH the inner and outer + // Filter pair OR-widen, so the two scans fuse into one and each aggregate FILTER is the + // per-side AND(inner, outer). + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false", + SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + + // With the DSv2 config off the inner pair already declines, so nothing fuses (plan unchanged). + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false", + SQLConf.MERGE_SUBPLANS_DSV2_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "false") { + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + } + + test("SPARK-40259: merge proceeds without pruning when the source rejects the pruning") { + val rejecting = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType))), + acceptsFilters = false) + val s1 = v2ScanReadingOn(rejecting, Seq("a")) + val s2 = v2ScanReadingOn(rejecting, Seq("b")) + val sub1 = ScalarSubquery(s1.where(s1.output.head > 1).groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(s2.where(s2.output.head > 2).groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the same OR-widen merge as the accepting case -- only the pushed pruning differs + // (a rejecting source records none). np = sub2 -> id 0, cp = sub1 -> id 1. + val mergedScan = v2ScanReadingOn(rejecting, Seq("a", "b")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("sum_a"), + sum($"b", Some(npFilter)).as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Phase 2 degrades gracefully: the merge happens; a rejecting source records no pruning. + assert(v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.isEmpty, + "a rejecting source must not record any pushed pruning predicate") + } + } + + test("SPARK-40259: do not merge when a strict pushed filter is not re-enforced on rebuild") { + // The scans claim a strict filter on "a" (in pushedFilters), but the source enforces nothing + // strictly (empty strictColumns -> everything is best-effort). Re-pushing the "strict" filter + // would leave it merely best-effort with nothing above to re-check it, so the merge must abort. + val bestEffort = new TestV2Table(StructType(Seq( + StructField("a", IntegerType), StructField("b", IntegerType), StructField("c", StringType)))) + val s1 = v2ScanReadingOn(bestEffort, Seq("a", "b")) + val s2 = v2ScanReadingOn(bestEffort, Seq("a", "c")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"c").as("max_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // A strict filter the rebuilt scan cannot re-enforce declines the merge: plan left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: do not merge DSv2 scans of a table without the SCAN_MERGING capability") { + val sub1 = ScalarSubquery(v2ScanReadingNoMerge("a").groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(v2ScanReadingNoMerge("b").groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + // A table without the SCAN_MERGING capability declines the merge: plan left unchanged. + comparePlans(Optimize.execute(originalQuery.analyze), originalQuery.analyze) + } + + test("SPARK-40259: merge DSv2 scans with identical filters re-pushes the pruning") { + // Identical post-scan filter (a > 1), differing columns: the two scans fuse and the identical + // condition is re-pushed to the merged scan for row-group pruning (Phase 2, single condition). + val sub1 = ScalarSubquery( + v2ScanReading("a", "b").where($"a" > 1).groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + v2ScanReading("a", "c").where($"a" > 1).groupBy()(sum($"c").as("sum_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: the identical filter a > 1 kept as a single Filter above the merged scan {a, b, c}; + // no OR-widen (the conditions match). The aggregates read b and c. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b", "c")) + val mergedSubquery = mergedScan.where($"a" > 1) + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("sum_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("sum_c"), $"sum_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the identical a > 1 is re-pushed to the merged scan for pruning. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed + assert(pushed.nonEmpty && pushed.mkString.contains("a"), + s"expected the identical filter a > 1 to be re-pushed to the merged scan, got: $pushed") + } + + test("SPARK-40259: merge three DSv2 scans with differing filters re-pushes the full OR") { + // Three-way merge exercises the tagged (MERGED_FILTER_TAG) branch: the leaf re-merge rebuilds + // the scan strict-only each round, so the tagged branch must re-establish the full 3-way OR. + val t = new TestV2Table(StructType(Seq(StructField("a", IntegerType), + StructField("b", IntegerType), StructField("d", IntegerType)))) + val sub1 = ScalarSubquery( + v2ScanReadingOn(t, Seq("a")).where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery( + v2ScanReadingOn(t, Seq("b")).where($"b" > 2).groupBy()(sum($"b").as("s2"))) + val sub3 = ScalarSubquery( + v2ScanReadingOn(t, Seq("d")).where($"d" > 3).groupBy()(sum($"d").as("s3"))) + val originalQuery = testRelation.select(sub1, sub2, sub3) + + // Expected: one merged scan reading {a, b, d}; step 1 merges sub1 (cp) + sub2 (np) -> + // propagatedFilter_0 = b > 2 (np), _1 = a > 1 (cp); step 2 merges sub3 (np) -> _2 = d > 3, + // extending the OR. Each aggregate carries its side's FILTER. + val mergedScan = v2ScanReadingOn(t, Seq("a", "b", "d")) + val npFilter0Alias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilter0Alias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter1Alias = Alias($"d" > 3, "propagatedFilter_2")() + val npFilter0 = npFilter0Alias.toAttribute + val cpFilter0 = cpFilter0Alias.toAttribute + val npFilter1 = npFilter1Alias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilter0Alias, cpFilter0Alias, npFilter1Alias): _*) + .where(Or(Or(npFilter0, cpFilter0), npFilter1)) + .groupBy()( + sum($"a", Some(cpFilter0)).as("s1"), + sum($"b", Some(npFilter0)).as("s2"), + sum($"d", Some(npFilter1)).as("s3")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2", + Literal("s3"), $"s3")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1), + extractorExpression(0, analyzedMergedSubquery.output, 2)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the merged scan carries the full 3-way OR pruning over a, b and d. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a") && pushed.contains("b") && pushed.contains("d"), + s"expected an OR pruning over a, b and d pushed to the merged scan, got: $pushed") + } + } + + test("SPARK-40259: three-way merge reuses an existing propagated filter (reuse branch)") { + // sub3's filter (a > 1) matches sub1's, so the tagged-filter merge takes the reuse branch: + // no new propagatedFilter alias is created and the OR stays two-way; sub3's aggregate reuses + // the existing filter attribute. This exercises the reuse branch against a DSv2 scan, which + // the all-distinct-filters three-way test above does not. + val t = new TestV2Table(StructType(Seq(StructField("a", IntegerType), + StructField("b", IntegerType), StructField("c", IntegerType)))) + val sub1 = ScalarSubquery( + v2ScanReadingOn(t, Seq("a")).where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery( + v2ScanReadingOn(t, Seq("b")).where($"b" > 2).groupBy()(sum($"b").as("s2"))) + val sub3 = ScalarSubquery( + v2ScanReadingOn(t, Seq("a", "c")).where($"a" > 1).groupBy()(sum($"c").as("s3"))) + val originalQuery = testRelation.select(sub1, sub2, sub3) + + // Expected: one merged scan reading {a, b, c}; step 1 merges sub1 (cp) + sub2 (np) -> + // propagatedFilter_0 = b > 2 (np), _1 = a > 1 (cp). Step 2 merges sub3 (np, a > 1): its + // condition matches _1, so no new alias is created and the OR stays two-way; sub3 reuses _1. + val mergedScan = v2ScanReadingOn(t, Seq("a", "b", "c")) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("s1"), + sum($"b", Some(npFilter)).as("s2"), + sum($"c", Some(cpFilter)).as("s3")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2", + Literal("s3"), $"s3")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1), + extractorExpression(0, analyzedMergedSubquery.output, 2)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // The OR references only the two distinct conditions (a > 1, b > 2); c is not a filter. + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a") && pushed.contains("b"), + s"expected an OR pruning over a and b pushed to the merged scan, got: $pushed") + } + } + + test("SPARK-40259: a non-deterministic filter conjunct is not pushed as a pruning predicate") { + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + // Bare non-grouping aggregates (joined) go through the Aggregate extraction path, which -- + // unlike the ScalarSubquery path -- does not gate on determinism, so the OR-widen pruning + // predicate carries a non-deterministic conjunct into the scan-merge Phase 2. + val agg1 = v2ScanReading("a").where(($"a" > 1) && (rand(0) < Literal(0.5))) + .groupBy()(sum($"a").as("s1")) + val agg2 = v2ScanReading("b").where($"b" > 2).groupBy()(sum($"b").as("s2")) + val originalQuery = agg1.join(agg2) + + val scans = v2Scans(Optimize.execute(originalQuery.analyze)) + assert(scans.length == 1, "the scans should still be fused") + val pushed = scans.head.scan.asInstanceOf[TestV2Scan].pushed.mkString + // Pruning predicate `(a > 1 AND rand() < 0.5) OR (b > 2)` is non-deterministic as a whole. + // A source that prunes on it uses its own rand() draw, while the enclosing Filter re-checks + // exactness with a different draw, so pruned rows would be lost. Best-effort pruning degrades + // gracefully: the predicate is dropped rather than weakened, so nothing is pushed and no + // rand() reaches the source. The merge itself is unaffected. + assert(!pushed.contains("RAND"), + s"the non-deterministic rand() conjunct must not be pushed as pruning, got: $pushed") + assert(pushed.isEmpty, + s"a non-deterministic pruning predicate must be dropped wholesale, got: $pushed") + } + } + + test("SPARK-40259: merge DSv2 scans that pushed the same strict filters (re-push path)") { + // A source that fully enforces (strict) predicates on column "a". + val strictOnA = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType), + StructField("c", StringType))), + strictColumns = Set("a")) + val s1 = v2ScanReadingOn(strictOnA, Seq("a", "b")) + val s2 = v2ScanReadingOn(strictOnA, Seq("a", "c")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"b").as("sum_b"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "a").get > 0)) + .groupBy()(sum($"c").as("max_c"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading {a, b, c} that keeps the strict a > 0 in pushedFilters; the + // two aggregates read b and c. No post-scan filter, so no OR-widen. + val mergedScan0 = v2ScanReadingOn(strictOnA, Seq("a", "b", "c")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "a").get > 0)) + val mergedSubquery = mergedScan + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("max_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("max_c"), $"max_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: the strict a > 0 is re-enforced on the merged scan (normalized out above). + assert(v2Scans(optimized).head.pushedFilters.nonEmpty, + "the merged relation must carry the strict pushed filters") + } + + test("SPARK-40259: merge scans with equal strict filters and differing post-scan filters") { + // Mixed source: "p" is fully enforced (strict), "a"/"b" are best-effort. Both sides push + // the same strict filter p > 0 and carry a differing post-scan filter (a > 1 / b > 2). The + // scans fuse on the equal strict p; the differing post-scan filters OR-widen above the + // merged scan, so the rebuild re-pushes p > 0 (strict) and (a > 1 OR b > 2) (pruning) at + // once -- the only path that drives buildMergedScan with both non-empty. + val mixed = new TestV2Table( + StructType(Seq(StructField("p", IntegerType), StructField("a", IntegerType), + StructField("b", IntegerType))), + strictColumns = Set("p")) + val s1 = v2ScanReadingOn(mixed, Seq("p", "a")) + val s2 = v2ScanReadingOn(mixed, Seq("p", "b")) + val sub1 = ScalarSubquery( + s1.copy(pushedFilters = Seq(s1.output.find(_.name == "p").get > 0)) + .where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery( + s2.copy(pushedFilters = Seq(s2.output.find(_.name == "p").get > 0)) + .where($"b" > 2).groupBy()(sum($"b").as("s2"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading {p, a, b} that keeps p > 0 strict (pushedFilters), with + // the differing post-scan filters OR-widened above it (np = sub2 -> id 0, cp = sub1 -> id 1). + val mergedScan0 = v2ScanReadingOn(mixed, Seq("p", "a", "b")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "p").get > 0)) + val npFilterAlias = Alias($"b" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("s1"), + sum($"b", Some(npFilter)).as("s2")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + // Additional: p stays strict (in pushedFilters) and the OR (a, b) is re-pushed as pruning, + // both normalized out of the structural comparison above. + val mergedResult = v2Scans(optimized).head + assert(mergedResult.pushedFilters.exists(_.references.exists(_.name == "p")), + s"the strict filter on p must stay enforced; got ${mergedResult.pushedFilters}") + val pushed = mergedResult.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a") && pushed.contains("b"), + s"the differing post-scan filters must be re-pushed as OR pruning, got: $pushed") + } + } + + test("SPARK-40259: do not merge DSv2 scans when the merge would degrade a reported " + + "key-grouped partitioning or ordering") { + // An input reports key-grouped partitioning or ordering, but the merged scan -- rebuilt over a + // TestV2Scan that reports neither -- re-derives nothing, so merging would degrade what the + // input reported. With the degradation configs off (the default) the merge is declined on both + // the np and cp side and the plan is left unchanged, rather than forcing a shuffle/sort the + // original plan avoided. + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = + Seq( + (withField(v2ScanReading("a")), v2ScanReading("b")), + (v2ScanReading("a"), withField(v2ScanReading("b")))).foreach { case (npScan, cpScan) => + val q = testRelation.select( + ScalarSubquery(npScan.groupBy()(sum($"a").as("sa"))), + ScalarSubquery(cpScan.groupBy()(sum($"b").as("sb")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + } + + test("SPARK-58549: do not merge DSv2 scans reporting incompatible kGP/ordering") { + // Both inputs report a partitioning/ordering, but on the different column each reads, so no + // single rebuilt scan could keep both not-worse. combineRequired* returns None (incompatible), + // so the merge is declined at the leaf -- before any rebuild -- with the degradation configs + // off (the default). This is distinct from the single-side case above (which rebuilds and then + // finds the re-derived report degraded): here the two inputs disagree with each other up front. + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = { + // withField applied to the "a" scan reports on a, applied to the "b" scan reports on b. + val q = testRelation.select( + ScalarSubquery(withField(v2ScanReading("a")).groupBy()(sum($"a").as("sa"))), + ScalarSubquery(withField(v2ScanReading("b")).groupBy()(sum($"b").as("sb")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + } + + test("SPARK-58549: merge DSv2 scans reporting kGP/ordering when the degradation config allows") { + // With the matching degradation config on, a merge that would drop a reported partitioning or + // ordering proceeds anyway (trading it for a single scan). The merged scan re-derives no report + // from the non-reporting TestV2Scan, so the fused plan is the plain column union. + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val mergedSubquery = mergedScan + .groupBy()(sum($"a").as("sum_a"), sum($"b").as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + def assertMerges( + withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation, + confKey: String, + bothSides: Boolean = false): Unit = { + val sub1 = ScalarSubquery(withField(v2ScanReading("a")).groupBy()(sum($"a").as("sum_a"))) + val cpScan = if (bothSides) withField(v2ScanReading("b")) else v2ScanReading("b") + val sub2 = ScalarSubquery(cpScan.groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + withSQLConf(confKey -> "true") { + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + } + + // One side reports: the report combines fine, and the config lets through the degradation the + // post-rebuild check finds (the rebuilt scan re-derives nothing). + assertMerges(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key) + assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key) + // Both sides report, each on the column it reads: INCOMPATIBLE, so here the config is what + // skips the EARLY decline at the leaf, before any rebuild (the mirror of the incompatible + // default-decline test above). + assertMerges(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key, + bothSides = true) + assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key, bothSides = true) + } + + test("SPARK-58549: merge DSv2 scans when the rebuilt scan re-derives the required ordering") { + // The NOT-WORSE side of the check: both inputs report `a ASC` and the table reports the same, + // so the rebuilt merged scan re-derives it and nothing is degraded -- the merge proceeds with + // the configs off. Every other test here rebuilds over a plain TestV2Scan, which reports + // nothing, so any non-empty requirement is degraded by construction and only the declining side + // gets covered. + val withOrdering = (s: DataSourceV2ScanRelation) => + s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))) + val q = testRelation.select( + ScalarSubquery(withOrdering(v2ScanReadingOn(v2TableReportingA, Seq("a", "b"))) + .groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery(withOrdering(v2ScanReadingOn(v2TableReportingA, Seq("a", "c"))) + .groupBy()(sum($"c").as("sum_c")))) + val optimized = Optimize.execute(q.analyze) + + val scans = v2Scans(optimized) + assert(scans.length == 1, s"the two scans should be fused into one:\n$optimized") + assert(scans.head.output.map(_.name).toSet == Set("a", "b", "c"), + s"the merged scan should read the union of both columns; got ${scans.head.output}") + assert(scans.head.ordering.exists(_.map(_.child).collect { + case a: Attribute => a.name + } == Seq("a")), + s"the merged scan should re-derive the reported ordering; got ${scans.head.ordering}") + } + + test("SPARK-58549: do not merge when the rebuilt scan re-derives less than the combined " + + "ordering") { + // One input reports `a ASC`, the other `a ASC, c ASC`, so the combined requirement is the + // STRONGER of the two. The table reports only `a ASC`, so the rebuilt merged scan satisfies the + // weaker input's ordering but not the combined one -- a degradation, declined with the configs + // off. Guards that combineRequiredOrdering keeps the stronger side: keeping the weaker one + // would let this merge through and lose the second sort key. + val scanAB = v2ScanReadingOn(v2TableReportingA, Seq("a", "b")) + val scanAC = v2ScanReadingOn(v2TableReportingA, Seq("a", "c")) + val q = testRelation.select( + ScalarSubquery(scanAB.copy(ordering = Some(Seq(SortOrder(scanAB.output.head, Ascending)))) + .groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery(scanAC.copy(ordering = Some(Seq( + SortOrder(scanAC.output.head, Ascending), SortOrder(scanAC.output.last, Ascending)))) + .groupBy()(sum($"c").as("sum_c")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + test("SPARK-58549: merge DSv2 scans reporting the same bucket transform partitioning") { + // Both inputs report `bucket(4, a)`, each derived independently, so each holds its OWN + // BoundFunction instance -- what production does, since V2ExpressionUtils binds the function + // afresh per derivation. The two reports compare equal only because `TestBucketFunction` + // implements `equals`/`hashCode` the way `BoundFunction` asks a connector to, so this is the + // end-to-end check that Spark honours that contract: the merge proceeds and the rebuilt scan + // re-derives the same report. It is also the only test here whose report is a transform rather + // than a plain attribute, which is what an identity-partitioned source would report. + val q = testRelation.select( + ScalarSubquery(v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "b")) + .groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery(v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "c")) + .groupBy()(sum($"c").as("sum_c")))) + val optimized = Optimize.execute(q.analyze) + + val scans = v2Scans(optimized) + assert(scans.length == 1, s"the two scans should be fused into one:\n$optimized") + val kgp = scans.head.keyGroupedPartitioning + assert(kgp.exists(_.forall(_.isInstanceOf[TransformExpression])), + s"the merged scan should preserve the reported bucket transform; got $kgp") + assert(kgp.get.flatMap(_.references).exists(_.name == "a"), + s"the preserved partitioning should be on a; got $kgp") + } + + test("SPARK-58769: decline the merge when the reported transform is not comparable") { + // The mirror of the test above, and the cost this documents: same query, same reported + // `bucket(4, a)`, but the connector's function does not implement `equals`, so the two + // independently bound instances do not compare equal and Spark cannot tell the two reports + // apart from two different partitionings. It does not derive that identity itself (see + // `BoundFunction#equals`), so the merge is declined rather than done on a report it cannot + // verify, and the plan is left alone. + val q = testRelation.select( + ScalarSubquery( + v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "b"), TestNotComparableFunctionCatalog) + .groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery( + v2ScanReportingOn(v2TableBucketedOnA, Seq("a", "c"), TestNotComparableFunctionCatalog) + .groupBy()(sum($"c").as("sum_c")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + + test("SPARK-58549: enforce the required report on the deferred under-Filter scan build") { + // The above tests fuse scans directly under an Aggregate (no Filter), so they exercise the + // scan build at the leaf. When the scans sit under an (identical) Filter the build is instead + // DEFERRED to the enclosing Filter, and the required report is carried there through + // DSv2DeferredScan. This test drives that deferred path: each scan reads {a, <col>} and reports + // on a, both filter on `a > 1` (so the two fuse into {a, b, c} under one Filter). The merged + // scan is rebuilt over a non-reporting TestV2Scan, so it re-derives no report -- a degradation. + def reportsOnA(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation) + : (ScalarSubquery, ScalarSubquery) = ( + ScalarSubquery( + withField(v2ScanReading("a", "b")).where($"a" > 1).groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery( + withField(v2ScanReading("a", "c")).where($"a" > 1).groupBy()(sum($"c").as("sum_c")))) + + // Default configs: the deferred build declines on the degradation, leaving the plan unchanged. + def assertDeclines(withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = { + val (sub1, sub2) = reportsOnA(withField) + val q = testRelation.select(sub1, sub2) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } + assertDeclines(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head)))) + assertDeclines(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending))))) + + // With the matching config on, the deferred build proceeds: the two scans fuse into {a, b, c} + // with the identical `a > 1` re-pushed for pruning (as in the identical-filter merge test). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b", "c")) + val mergedSubquery = mergedScan.where($"a" > 1) + .groupBy()(sum($"b").as("sum_b"), sum($"c").as("sum_c")) + .select(CreateNamedStruct(Seq( + Literal("sum_b"), $"sum_b", + Literal("sum_c"), $"sum_c")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + def assertMerges( + withField: DataSourceV2ScanRelation => DataSourceV2ScanRelation, confKey: String): Unit = { + val (sub1, sub2) = reportsOnA(withField) + val q = testRelation.select(sub1, sub2) + withSQLConf(confKey -> "true") { + comparePlans(Optimize.execute(q.analyze), correctAnswer.analyze) + } + } + assertMerges(s => s.copy(keyGroupedPartitioning = Some(Seq(s.output.head))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_KEY_GROUPED_PARTITIONING_DEGRADATION.key) + assertMerges(s => s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))), + SQLConf.MERGE_SUBPLANS_DSV2_ALLOW_ORDERING_DEGRADATION.key) + } + + test("SPARK-58549: the deferred build checks the required report per attempt") { + // The deferred build tries strict + best-effort first, then strict-only. The report is checked + // on EACH attempt, not once on whichever came back, so the deferred path stays no weaker than + // the leaf path: the strict-only attempt is exactly what the leaf builds when no Filter sits + // above the scan, so a merge it can satisfy must not be lost just because the first attempt + // could not. Here both inputs report `a ASC` and neither has a strict pushed filter, and the + // source reports only when nothing was pushed: attempt 1 (offered the Filter's condition) + // re-derives no ordering and is rejected, attempt 2 re-derives `a ASC` and the merge lands. + val withOrdering = (s: DataSourceV2ScanRelation) => + s.copy(ordering = Some(Seq(SortOrder(s.output.head, Ascending)))) + val q = testRelation.select( + ScalarSubquery(withOrdering(v2ScanReadingOn(v2TablePushSensitiveReport, Seq("a", "b"))) + .where($"a" > 1).groupBy()(sum($"b").as("sum_b"))), + ScalarSubquery(withOrdering(v2ScanReadingOn(v2TablePushSensitiveReport, Seq("a", "c"))) + .where($"a" > 1).groupBy()(sum($"c").as("sum_c")))) + val optimized = Optimize.execute(q.analyze) + + val scans = v2Scans(optimized) + assert(scans.length == 1, s"the two scans should be fused into one:\n$optimized") + assert(scans.head.ordering.exists(_.map(_.child).collect { + case a: Attribute => a.name + } == Seq("a")), + s"the strict-only attempt should have satisfied the required ordering; " + + s"got ${scans.head.ordering}") + assert(scans.head.scan.asInstanceOf[TestV2ReportingScan].inner.pushed.isEmpty, + "the surviving build should be the strict-only one, with nothing pushed") + } + + test("SPARK-40259: merge DSv2 scans that report empty key-grouped partitioning or ordering") { + // A source implementing SupportsReportPartitioning/SupportsReportOrdering but reporting nothing + // yields Some(Nil), not None (V2ScanPartitioningAndOrdering sets the field unconditionally). An + // empty report carries no partitioning/ordering to drop, so the merge should still proceed -- + // the gate tests the inner Seq, not the Option. The fused plan is the plain column union. + def assertMerges(withEmptyField: DataSourceV2ScanRelation => DataSourceV2ScanRelation): Unit = { + val sub1 = ScalarSubquery(withEmptyField(v2ScanReading("a")).groupBy()(sum($"a").as("sum_a"))) + val sub2 = ScalarSubquery(withEmptyField(v2ScanReading("b")).groupBy()(sum($"b").as("sum_b"))) + val originalQuery = testRelation.select(sub1, sub2) + + val mergedScan = v2ScanReadingOn(v2Table, Seq("a", "b")) + val mergedSubquery = mergedScan + .groupBy()(sum($"a").as("sum_a"), sum($"b").as("sum_b")) + .select(CreateNamedStruct(Seq( + Literal("sum_a"), $"sum_a", + Literal("sum_b"), $"sum_b")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + comparePlans(Optimize.execute(originalQuery.analyze), correctAnswer.analyze) + } + + assertMerges(_.copy(keyGroupedPartitioning = Some(Nil))) + assertMerges(_.copy(ordering = Some(Nil))) + } + + test("SPARK-40259: merge DSv2 scans reading identical columns but differing filters") { + // Both scans read only "a", so the column union adds nothing (the np-only set is empty); the + // merge is driven purely by the differing filters, and the OR is still re-pushed for pruning. + // (sub1 is cp, sub2 is np.) + val sub1 = ScalarSubquery(v2ScanReading("a").where($"a" > 1).groupBy()(sum($"a").as("s1"))) + val sub2 = ScalarSubquery(v2ScanReading("a").where($"a" > 2).groupBy()(sum($"a").as("s2"))) + val originalQuery = testRelation.select(sub1, sub2) + + // Expected: one merged scan reading just {a} (np-only set empty), the differing filters + // OR-widened above it (np = sub2 -> id 0, cp = sub1 -> id 1). + val mergedScan = v2ScanReadingOn(v2Table, Seq("a")) + val npFilterAlias = Alias($"a" > 2, "propagatedFilter_0")() + val cpFilterAlias = Alias($"a" > 1, "propagatedFilter_1")() + val npFilter = npFilterAlias.toAttribute + val cpFilter = cpFilterAlias.toAttribute + val mergedSubquery = mergedScan + .select(mergedScan.output ++ Seq(npFilterAlias, cpFilterAlias): _*) + .where(Or(npFilter, cpFilter)) + .groupBy()( + sum($"a", Some(cpFilter)).as("s1"), + sum($"a", Some(npFilter)).as("s2")) + .select(CreateNamedStruct(Seq( + Literal("s1"), $"s1", + Literal("s2"), $"s2")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + + withSQLConf( + SQLConf.MERGE_SUBPLANS_FILTER_PROPAGATION_ENABLED.key -> "true", + SQLConf.MERGE_SUBPLANS_SYMMETRIC_FILTER_PROPAGATION_ENABLED.key -> "true") { + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + val pushed = v2Scans(optimized).head.scan.asInstanceOf[TestV2Scan].pushed.mkString + assert(pushed.contains("a"), s"the OR over a should be re-pushed as pruning, got: $pushed") + } + } + + test("SPARK-40259: merge three DSv2 scans that pushed the same strict filter") { + // Round-2 merge feeds the already-merged relation back through samePushedFilters. All three + // scans carry the same strict a > 0, so they fuse into a single scan reading a, b, c, d. + val strictOnA = new TestV2Table( + StructType(Seq(StructField("a", IntegerType), StructField("b", IntegerType), + StructField("c", IntegerType), StructField("d", IntegerType))), + strictColumns = Set("a")) + def strictSub(col: String): ScalarSubquery = { + val s = v2ScanReadingOn(strictOnA, Seq("a", col)) + ScalarSubquery( + s.copy(pushedFilters = Seq(s.output.find(_.name == "a").get > 0)) + .groupBy()(sum(s.output.find(_.name == col).get).as(s"agg_$col"))) + } + val originalQuery = testRelation.select(strictSub("b"), strictSub("c"), strictSub("d")) + + // Expected: one merged scan reading {a, b, c, d} keeping the shared strict a > 0; three + // aggregates read b, c and d. No post-scan filter, so no OR-widen. + val mergedScan0 = v2ScanReadingOn(strictOnA, Seq("a", "b", "c", "d")) + val mergedScan = mergedScan0.copy( + pushedFilters = Seq(mergedScan0.output.find(_.name == "a").get > 0)) + val mergedSubquery = mergedScan + .groupBy()(sum($"b").as("agg_b"), sum($"c").as("agg_c"), sum($"d").as("agg_d")) + .select(CreateNamedStruct(Seq( + Literal("agg_b"), $"agg_b", + Literal("agg_c"), $"agg_c", + Literal("agg_d"), $"agg_d")).as("mergedValue")) + val analyzedMergedSubquery = mergedSubquery.analyze + val correctAnswer = WithCTE( + testRelation.select( + extractorExpression(0, analyzedMergedSubquery.output, 0), + extractorExpression(0, analyzedMergedSubquery.output, 1), + extractorExpression(0, analyzedMergedSubquery.output, 2)), + Seq(definitionNode(analyzedMergedSubquery, 0))) + val optimized = Optimize.execute(originalQuery.analyze) + comparePlans(optimized, correctAnswer.analyze) + assert(v2Scans(optimized).head.pushedFilters.exists(_.references.exists(_.name == "a")), + "the merged scan must keep the shared strict filter on a") + } + + test("SPARK-40259: pushed limit/offset/top-N/sample block merging " + + "(hasBlockingPushdown classification)") { + // hasBlockingPushdown flags every pushdown a rebuilt scan cannot reproduce. Drive each term + // through the real pushdown (the stub opts into all of them) so that a term silently dropped + // from the denylist is caught. Limit, offset and sample each set only their own holder field, + // so they are isolated; top-N sets sortOrders AND pushedLimit, so it exercises the top-N path + // but does not isolate the sortOrders term (no plan pushes a sort order without a limit). + def pushDown(op: DataSourceV2Relation => LogicalPlan): DataSourceV2ScanRelation = { + val relation = DataSourceV2Relation( + v2Table, toAttributes(v2Table.schema()), None, None, CaseInsensitiveStringMap.empty()) + V2ScanRelationPushDown(op(relation)) + .collectFirst { case s: DataSourceV2ScanRelation => s }.get + } + + // Default-safe: only the plain column-pruning path grants mergeability, so a plain scan with + // no non-reproducible pushdown must come out mergeable (guards the one granting site). + assert(pushDown(r => Project(Seq(r.output.head), r)).mergeableScan, + "a plain scan with only reproducible pushdowns must be mergeable") + + assert(!pushDown(r => Limit(Literal(1), r)).mergeableScan, + "a pushed limit must block merging") + assert(!pushDown(r => Offset(Literal(1), r)).mergeableScan, + "a pushed offset must block merging") + assert( + !pushDown(r => Limit(Literal(1), + Sort(Seq(SortOrder(r.output.head, Ascending)), global = true, r))) + .mergeableScan, + "a pushed top-N must block merging") + assert(!pushDown(r => Sample(0.0, 0.5, withReplacement = false, Some(0L), r)) + .mergeableScan, + "a pushed sample must block merging") + + // End-to-end: a scan that pushed a sample does not fuse with another (plain) scan of the same + // table -- the classification declines the scan merge, so the plan is left unchanged. The two + // read different columns, so they are not identical subplans and could only fuse via the + // (blocked) scan merge, not the identical-plan reuse short-circuit. + val sampled = pushDown(r => Sample(0.0, 0.5, withReplacement = false, Some(0L), r)) + val plain = v2ScanReading("a") + val q = testRelation.select( + ScalarSubquery(plain.groupBy()(sum($"a").as("s1"))), + ScalarSubquery(sampled.groupBy()(sum(sampled.output(1)).as("s2")))) + comparePlans(Optimize.execute(q.analyze), q.analyze) + } +} + +/** + * Minimal readable DSv2 table whose scan opts into merging, for scan-merge tests. + * + * A pushed predicate that references only columns in `strictColumns` is fully enforced (strict): + * accepted by the source and NOT returned as a post-scan filter. Any other predicate is + * best-effort (stats): accepted for row-group pruning but returned so it is re-checked above the + * scan. When `acceptsFilters` is false the source rejects everything: nothing is accepted and + * every predicate is returned (used to exercise the "no pruning recovery" path). + */ +private class TestV2Table( + tableSchema: StructType, + acceptsFilters: Boolean = true, + strictColumns: Set[String] = Set.empty, + supportsScanMerging: Boolean = true, + reportedOrderingCols: Seq[String] = Nil, + reportedBucketPartition: Option[(Int, String)] = None, + reportOnlyWhenNothingPushed: Boolean = false) + extends Table with SupportsRead { + override def name(): String = "test_v2_table" + override def schema(): StructType = tableSchema + override def capabilities(): java.util.Set[TableCapability] = { + val caps = java.util.EnumSet.of(TableCapability.BATCH_READ) + if (supportsScanMerging) caps.add(TableCapability.SCAN_MERGING) + caps + } + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new TestV2ScanBuilder(tableSchema, acceptsFilters, strictColumns, reportedOrderingCols, + reportedBucketPartition, reportOnlyWhenNothingPushed) +} + +private class TestV2ScanBuilder( + tableSchema: StructType, + acceptsFilters: Boolean, + strictColumns: Set[String], + reportedOrderingCols: Seq[String] = Nil, + reportedBucketPartition: Option[(Int, String)] = None, + reportOnlyWhenNothingPushed: Boolean = false) + extends ScanBuilder with SupportsPushDownRequiredColumns with SupportsPushDownV2Filters + with SupportsPushDownLimit with SupportsPushDownOffset with SupportsPushDownTopN + with SupportsPushDownTableSample { + private var prunedSchema: StructType = tableSchema + private var accepted: Array[Predicate] = Array.empty // strict UNION stats (pushedPredicates) + + override def pruneColumns(requiredSchema: StructType): Unit = prunedSchema = requiredSchema + + // Opts into limit pushdown so V2ScanRelationPushDown records a merge-blocking pushedLimit. + override def pushLimit(limit: Int): Boolean = true + + // Opts into offset/top-N/sample pushdown too, so V2ScanRelationPushDown records the matching + // merge-blocking pushdown for each (used by the hasBlockingPushdown classification test). + override def pushOffset(offset: Int): Boolean = true + override def pushTopN(orders: Array[V2SortOrder], limit: Int): Boolean = true + override def isPartiallyPushed(): Boolean = false + override def pushTableSample( + lowerBound: Double, + upperBound: Double, + withReplacement: Boolean, + seed: Long): Boolean = true + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + if (!acceptsFilters) { + predicates // reject everything: nothing accepted, all returned as post-scan + } else { + accepted = predicates + predicates.filterNot(isStrict) // stats predicates are re-checked above the scan + } + } + + private def isStrict(p: Predicate): Boolean = + p.references().nonEmpty && + p.references().forall(r => strictColumns.contains(r.fieldNames().mkString("."))) + + override def pushedPredicates(): Array[Predicate] = accepted + override def build(): Scan = { + val scan = TestV2Scan(prunedSchema, accepted.map(_.describe()).toSeq) + if (reportedOrderingCols.isEmpty && reportedBucketPartition.isEmpty) { + scan + } else { + TestV2ReportingScan(scan, reportedOrderingCols, reportedBucketPartition, + reportOnlyWhenNothingPushed) + } + } +} + +/** `pushed` records the `describe()` of the predicates pushed to the scan, for test assertions. */ +private case class TestV2Scan(schema: StructType, pushed: Seq[String] = Seq.empty) + extends Scan { + override def readSchema(): StructType = schema +} + +/** + * A [[TestV2Scan]] that also reports an output ordering and/or a key-grouped partitioning, so a + * scan rebuilt by the merge re-derives a report through `V2ScanPartitioningAndOrdering`. The plain + * [[TestV2Scan]] reports nothing, which only ever exercises the degrading side of the not-worse + * check. + */ +private case class TestV2ReportingScan( + inner: TestV2Scan, + orderingCols: Seq[String] = Nil, + bucketPartition: Option[(Int, String)] = None, + onlyWhenNothingPushed: Boolean = false) + extends Scan with SupportsReportOrdering with SupportsReportPartitioning { + override def readSchema(): StructType = inner.readSchema() + + // A source may report per scan: `outputOrdering`/`outputPartitioning` are answered by the Scan + // the connector built AFTER seeing what it accepted, so an ordering that only holds over the + // unpruned file set can legitimately disappear once a filter is pushed. `onlyWhenNothingPushed` + // models that, so a build offered a best-effort filter and one without it report differently. + private def reports: Boolean = !onlyWhenNothingPushed || inner.pushed.isEmpty + + override def outputOrdering(): Array[V2SortOrder] = if (reports) { + orderingCols.map(c => Expressions.sort(FieldReference(c), V2SortDirection.ASCENDING)).toArray + } else { + Array.empty + } + + override def outputPartitioning(): V2Partitioning = bucketPartition.filter(_ => reports) match { + case Some((numBuckets, col)) => + new V2KeyGroupedPartitioning(Array(Expressions.bucket(numBuckets, col)), 0) + case None => new V2UnknownPartitioning(0) + } +} + +/** + * A `FunctionCatalog` that resolves `bucket`, binding it to a FRESH function instance every time -- + * as a real connector does, since `V2ExpressionUtils.loadV2FunctionOpt` binds the function afresh + * for every derivation of a reported transform. Spark's own `UnboundBucketFunction` hands back a + * singleton, which would hide the fact that relating two independently derived reports rests on the + * function's own `equals` (see `BoundFunction#equals`). + */ +private object TestFreshBindFunctionCatalog extends FunctionCatalog { + override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {} + override def name(): String = "test_fresh_bind" + override def listFunctions(namespace: Array[String]): Array[Identifier] = + Array(Identifier.of(Array.empty, "bucket")) + override def loadFunction(ident: Identifier): UnboundFunction = ident.name() match { + case "bucket" => TestUnboundBucketFunction + // loadV2FunctionOpt treats this as "function not found" and reports no partitioning. + case other => throw new UnsupportedOperationException(s"no such function: $other") + } +} + +private object TestUnboundBucketFunction extends UnboundFunction { + override def bind(inputType: StructType): BoundFunction = new TestBucketFunction + override def description(): String = name() + override def name(): String = "bucket" +} + +/** + * A bound `bucket` that implements `equals`/`hashCode` over the state identifying it, as + * `BoundFunction` asks a connector to. That is what lets Spark recognize two separately bound + * instances as the same transform: `V2ExpressionUtils` binds afresh on every derivation, so two + * reported partitionings hold two different instances and nothing but this comparison relates them. + * `canonicalName` is stable too (its default returns a fresh random UUID), which is what a + * storage-partitioned join compares. + */ +private class TestBucketFunction extends ScalarFunction[Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "bucket" + override def canonicalName(): String = "testcat.bucket" + + override def equals(other: Any): Boolean = other match { + case that: TestBucketFunction => + canonicalName() == that.canonicalName() && resultType() == that.resultType() && + // `Array` equality is reference identity, so compare the elements. + inputTypes().sameElements(that.inputTypes()) + case _ => false + } + + override def hashCode(): Int = canonicalName().hashCode +} + +/** + * The same catalog with a bound `bucket` that does NOT implement `equals`/`hashCode` -- a connector + * that meets only the `canonicalName` obligation. Two of its separately bound instances never + * compare equal, so Spark cannot relate two reports of one transform. + */ +private object TestNotComparableFunctionCatalog extends FunctionCatalog { + override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {} + override def name(): String = "test_not_comparable" + override def listFunctions(namespace: Array[String]): Array[Identifier] = + Array(Identifier.of(Array.empty, "bucket")) + override def loadFunction(ident: Identifier): UnboundFunction = ident.name() match { + case "bucket" => TestUnboundNotComparableBucketFunction + case other => throw new UnsupportedOperationException(s"no such function: $other") + } +} + +private object TestUnboundNotComparableBucketFunction extends UnboundFunction { + override def bind(inputType: StructType): BoundFunction = new TestNotComparableBucketFunction + override def description(): String = name() + override def name(): String = "bucket" +} + +/** A plain class, so it inherits the identity comparison from `Object`. */ +private class TestNotComparableBucketFunction extends ScalarFunction[Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "bucket" + override def canonicalName(): String = "testcat.bucket" } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala similarity index 99% rename from sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala rename to sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala index e1109f20e6040..7440c5badd679 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/PlanMergeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/PlanMergingSuite.scala @@ -15,14 +15,15 @@ * limitations under the License. */ -package org.apache.spark.sql +package org.apache.spark.sql.execution.planmerging +import org.apache.spark.sql.Row import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -class PlanMergeSuite extends SharedSparkSession +class PlanMergingSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala index 42da9f060b857..34d63df710cbd 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala @@ -27,7 +27,7 @@ import org.apache.spark.sql.connector.catalog.CatalogManager import org.apache.spark.sql.execution.{FilterExec, InputAdapter, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} -import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.test.{ExamplePointUDT, SharedSparkSession} import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -367,3 +367,31 @@ class MyDummyScalarPandasUDF extends UserDefinedPythonFunction( dataType = BooleanType, pythonEvalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF, udfDeterministic = true) + +class MyDummyScalarArrowUDF extends UserDefinedPythonFunction( + name = "dummyScalarArrowUDF", + func = new DummyUDF, + dataType = BooleanType, + pythonEvalType = PythonEvalType.SQL_SCALAR_ARROW_UDF, + udfDeterministic = true) + +class MyDummyScalarPandasIterUDF extends UserDefinedPythonFunction( + name = "dummyScalarPandasIterUDF", + func = new DummyUDF, + dataType = BooleanType, + pythonEvalType = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, + udfDeterministic = true) + +class MyDummyScalarArrowIterUDF extends UserDefinedPythonFunction( + name = "dummyScalarArrowIterUDF", + func = new DummyUDF, + dataType = BooleanType, + pythonEvalType = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, + udfDeterministic = true) + +class MyDummyUDTPythonUDF extends UserDefinedPythonFunction( + name = "dummyUDTUDF", + func = new DummyUDF, + dataType = new ExamplePointUDT, + pythonEvalType = PythonEvalType.SQL_BATCHED_UDF, + udfDeterministic = true) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala new file mode 100644 index 0000000000000..09d982026006b --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFFromLambdaSuite.scala @@ -0,0 +1,450 @@ +/* + * 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.execution.python + +import org.apache.spark.api.python.PythonEvalType +import org.apache.spark.sql.{AnalysisException, QueryTest} +import org.apache.spark.sql.catalyst.expressions.{LambdaFunction, Literal, NamedArgumentExpression, + PythonUDF} +import org.apache.spark.sql.catalyst.plans.logical.ArrowEvalPython +import org.apache.spark.sql.functions.{array_sort, col, forall, lit, map_filter, map_zip_with, + transform, transform_keys, transform_values, when, zip_with} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.{ExamplePointUDT, SharedSparkSession} +import org.apache.spark.sql.types.{ArrayType, IntegerType, MapType, StringType} + +/** + * Plan-shape tests for [[ExtractPythonUDFFromLambda]]. + * + * These assert the structure the rewrite produces - that no `PythonUDF` is left inside a + * `LambdaFunction`, that the lifted UDF is an element-wise UDF over an array, and that plans + * without a UDF in a lambda are untouched. End-to-end result correctness is covered by + * `pyspark.sql.tests.test_udf_in_higher_order_function`, which needs a real Python worker. + */ +class ExtractPythonUDFFromLambdaSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + + private val pythonUDF = new MyDummyPythonUDF + private val scalarPandasUDF = new MyDummyScalarPandasUDF + private val scalarArrowUDF = new MyDummyScalarArrowUDF + private val scalarPandasIterUDF = new MyDummyScalarPandasIterUDF + private val scalarArrowIterUDF = new MyDummyScalarArrowIterUDF + // Used where a UDF call must receive two arguments, e.g. a pairwise comparator. + private val pythonUDF2 = new MyDummyPythonUDF + private val nondeterministicUDF = new MyDummyNondeterministicPythonUDF + private val udtUDF = new MyDummyUDTPythonUDF + + private def arrayDF = Seq(Seq(1, 2, 3)).toDF("values") + + /** All `PythonUDF`s that remain inside a lambda in the optimized plan. */ + private def udfsInsideLambda(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = { + df.queryExecution.optimizedPlan.expressions.flatMap { e => + e.collect { case l: LambdaFunction => l }.flatMap { l => + l.collect { case u: PythonUDF => u } + } + } + } + + private def liftedUDFs(df: org.apache.spark.sql.DataFrame): Seq[PythonUDF] = { + df.queryExecution.optimizedPlan.collect { + case a: ArrowEvalPython => a.udfs + }.flatten + } + + test("transform: the UDF is lifted out of the lambda as an element-wise array UDF") { + val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r")) + + // The whole point of the rewrite: nothing Python-shaped is left inside a lambda. + assert(udfsInsideLambda(df).isEmpty) + + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF) + // The lifted UDF takes the array and returns an array of the original return type. + assert(lifted.head.dataType == ArrayType(pythonUDF.dataType, containsNull = true)) + assert(lifted.head.children.size == 1) + assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType]) + } + + test("filter/exists/forall: the UDF is lifted out of the lambda") { + val exprs = Seq( + org.apache.spark.sql.functions.filter(col("values"), x => pythonUDF(x)), + org.apache.spark.sql.functions.exists(col("values"), x => pythonUDF(x)), + org.apache.spark.sql.functions.forall(col("values"), x => pythonUDF(x))) + + exprs.foreach { e => + val df = arrayDF.select(e.as("r")) + assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for $e") + assert(liftedUDFs(df).forall(_.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF)) + } + } + + test("aggregate: a UDF anywhere in the fold is rejected") { + // The fold is sequential, so no array can precompute the values the UDF sees. A UDF in `merge` + // or in `finish` therefore has no rewrite and analysis must fail. + val cases = Seq( + "merge" -> + org.apache.spark.sql.functions.aggregate( + col("values"), lit(false), (acc, x) => acc || pythonUDF(x)), + "finish" -> + org.apache.spark.sql.functions.aggregate( + col("values"), lit(0), (acc, x) => acc + x, acc => pythonUDF(acc))) + cases.foreach { case (name, expr) => + val e = intercept[AnalysisException] { + arrayDF.select(expr.as("r")).collect() + } + assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF", + s"expected aggregate with a UDF in $name to be rejected") + } + } + + test("several UDFs and nested UDFs in one lambda are all lifted") { + val several = arrayDF.select( + transform(col("values"), x => pythonUDF(x) || pythonUDF(x + lit(1))).as("r")) + assert(udfsInsideLambda(several).isEmpty) + // Two distinct calls, so two lifted array UDFs. + assert(liftedUDFs(several).size == 2) + + // The same deterministic call twice must be evaluated once. + val duplicated = arrayDF.select( + transform(col("values"), x => pythonUDF(x) || pythonUDF(x)).as("r")) + assert(udfsInsideLambda(duplicated).isEmpty) + assert(liftedUDFs(duplicated).size == 1) + } + + test("a UDF nested inside a composite argument is lifted, not left inside a lambda") { + // SPARK-27052: `f(g(x) + 1)` / `f(-g(x))`. The inner call `g(x)` is already lifted; its + // occurrence buried inside the composite argument must be substituted too, or a raw `g` over a + // lambda variable would be left inside the generated transform for `ExtractPythonUDFs` to + // re-extract (the SPARK-48706 failure mode). Both nesting shapes must leave no UDF in a lambda. + val plusOne = arrayDF.select( + transform(col("values"), x => pythonUDF(pythonUDF(x).cast("int") + lit(1))).as("r")) + assert(udfsInsideLambda(plusOne).isEmpty) + // Two distinct calls (inner `g(x)`, outer `f(g(x) + 1)`), so two lifted array UDFs. + assert(liftedUDFs(plusOne).size == 2) + + val negated = arrayDF.select( + transform(col("values"), x => pythonUDF(-pythonUDF(x).cast("int"))).as("r")) + assert(udfsInsideLambda(negated).isEmpty) + assert(liftedUDFs(negated).size == 2) + } + + test("identical nondeterministic calls are lifted distinctly, not deduplicated") { + // Deduplicating `f(x)` with `f(x)` would collapse two independent draws into one; a + // nondeterministic UDF must keep each call, so both are lifted. + val df = arrayDF.select( + transform(col("values"), x => nondeterministicUDF(x) || nondeterministicUDF(x)).as("r")) + assert(udfsInsideLambda(df).isEmpty) + assert(liftedUDFs(df).size == 2) + } + + test("transform_values whose result type equals the key type replaces values, not keys") { + // SPARK-27052: dispatch is by concrete function, not result type. For map<string, string>, + // transform_values' lambda also returns string, which must not be treated as new keys. + val maps = Seq(Map("a" -> "x")).toDF("m") + val df = maps.select(transform_values(col("m"), (k, v) => pythonUDF(v).cast("string")).as("r")) + assert(udfsInsideLambda(df).isEmpty) + // Keys are preserved: the map still has the original key type and no new-key projection. + val mapType = df.queryExecution.analyzed.schema.head.dataType.asInstanceOf[MapType] + assert(mapType.keyType == StringType) + } + + test("a UDF inside a nested higher-order function's lambda is lifted, deepening the nesting") { + // `transform(matrix, row -> transform(row, x -> f(x)))`: the UDF is lifted onto the inner + // variable and then re-lifted onto the real `array<array<int>>` column, so it becomes a + // depth-2 element-wise UDF (flattening two array levels). Nothing is left inside a lambda. + val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values") + .select(transform(col("values"), inner => + transform(inner, x => pythonUDF(x))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + assert(lifted.head.evalType == PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF) + // Re-lifted once per enclosing lambda: depth 2, over an array<array<...>> argument. + assert(lifted.head.elementwiseNestingDepth == 2) + assert(lifted.head.dataType == + ArrayType(ArrayType(pythonUDF.dataType, containsNull = true), containsNull = true)) + } + + test("a UDF in a nested lambda that captures an enclosing lambda variable is lifted") { + // `transform(m, row -> transform(row, x -> f(x, size(row))))`: the UDF reads the inner element + // and the *enclosing* variable `row`. The captured value is repeated into an aligned array, so + // the UDF still lifts to a depth-2 element-wise UDF over the real column and nothing is left in + // a lambda. Relies on the fixed `HigherOrderFunction.canonicalized`, which no longer leaks the + // captured variable into the lifted UDF's references. + val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values") + .select(transform(col("values"), row => + transform(row, x => pythonUDF2(x, org.apache.spark.sql.functions.size(row)))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + assert(lifted.head.elementwiseNestingDepth == 2) + } + + test("a UDF on the outer element of a nested array is lifted") { + // Here the UDF applies to the outer array's element, which is a real column (the element just + // happens to be an array), so it is lifted onto that column like any other single-level array. + val df = Seq(Seq(Seq(1, 2), Seq(3))).toDF("values") + .select(transform(col("values"), inner => pythonUDF(inner)).as("r")) + assert(udfsInsideLambda(df).isEmpty) + } + + test("a rewritable higher-order function inside another lambda leaves no UDF in any lambda") { + // `transform(arr2, i -> array_max(transform(arr, x -> f(x))) + i)`: the inner `transform` + // iterates the real column `arr` (not the outer lambda variable `i`), so it is rewritable and + // its UDF is lifted out of the inner lambda. The lifted element-wise UDF reads only `arr`, so + // `ExtractPythonUDFs` then hoists it out of the outer lambda too (evaluated once per row). The + // end state must leave no `PythonUDF` inside any lambda. + val df = Seq((Seq(1, 2, 3), Seq(10, 20))).toDF("arr", "arr2") + .select(transform(col("arr2"), i => + org.apache.spark.sql.functions.array_max( + transform(col("arr"), x => pythonUDF(x).cast("int"))) + i).as("r")) + assert(udfsInsideLambda(df).isEmpty) + assert(liftedUDFs(df).size == 1) + } + + test("a lambda with no Python UDF is left unchanged") { + val df = arrayDF.select(transform(col("values"), x => x + lit(1)).as("r")) + val analyzed = df.queryExecution.analyzed + val optimized = df.queryExecution.optimizedPlan + // The rule must be inert: no eval-python node is introduced. + assert(optimized.collect { case a: ArrowEvalPython => a }.isEmpty) + assert(!optimized.toString.contains("pythonUDF")) + assert(analyzed.expressions.flatMap(_.collect { case u: PythonUDF => u }).isEmpty) + } + + test("every mapping higher-order function is rewritten") { + // One assertion per function, so that a shape regressing to "UDF left inside a lambda" is + // caught here rather than only by the end-to-end Python suite. `aggregate` / `reduce` are not + // here: they fold rather than map, so a UDF in them is rejected (see the aggregate test above). + val arrays = Seq((Seq(1, 2), Seq(3, 4))).toDF("l", "r") + val maps = Seq((Map("a" -> 1), Map("a" -> 2))).toDF("l", "r") + + val arrayCases = Seq( + "transform" -> transform(col("l"), x => pythonUDF(x)), + "transform with index" -> transform(col("l"), (x, i) => pythonUDF(x) || i > 0), + "filter" -> org.apache.spark.sql.functions.filter(col("l"), x => pythonUDF(x)), + "exists" -> org.apache.spark.sql.functions.exists(col("l"), x => pythonUDF(x)), + "forall" -> forall(col("l"), x => pythonUDF(x)), + "zip_with" -> zip_with(col("l"), col("r"), (a, b) => pythonUDF(a) || pythonUDF(b)), + "array_sort" -> array_sort(col("l"), + (a, b) => when(pythonUDF(a) === pythonUDF(b), lit(0)).otherwise(lit(1)))) + arrayCases.foreach { case (name, expr) => + val df = arrays.select(expr.as("r")) + assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name") + } + + val mapCases = Seq( + "transform_keys" -> transform_keys(col("l"), (k, v) => pythonUDF(k).cast("string")), + "transform_values" -> transform_values(col("l"), (k, v) => pythonUDF(v)), + "map_filter" -> map_filter(col("l"), (k, v) => pythonUDF(v)), + "map_zip_with" -> map_zip_with(col("l"), col("r"), (k, a, b) => pythonUDF(a) || pythonUDF(b))) + mapCases.foreach { case (name, expr) => + val df = maps.select(expr.as("r")) + assert(udfsInsideLambda(df).isEmpty, s"UDF left inside a lambda for $name") + } + } + + test("a pairwise array_sort comparator is lifted over the cross product") { + // One UDF call receiving both elements has no per-element key, so the UDF is precomputed over + // every ordered pair instead. The call must end up outside every lambda like any other. + val df = arrayDF.select( + array_sort(col("values"), (a, b) => pythonUDF2(a, b).cast("int")).as("r")) + assert(udfsInsideLambda(df).isEmpty) + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + // The UDF takes both pair sides, each a flat array of all n*n pairs. + assert(lifted.head.children.size == 2) + } + + test("a key-form array_sort comparator lifts one UDF per element, not per comparator side") { + // `(a, b) -> udf(a) < udf(b)`: `udf(a)` and `udf(b)` canonicalize differently (distinct + // variable exprIds) but both lift to the same function over the whole array, so they must be + // deduplicated into one lifted UDF - otherwise the Python function runs 2n times instead of n. + val df = arrayDF.select( + array_sort(col("values"), + (a, b) => when(pythonUDF(a) === pythonUDF(b), lit(0)).otherwise(lit(1))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + assert(liftedUDFs(df).size == 1) + } + + test("a vectorized scalar UDF inside a lambda is lifted to its element-wise eval type") { + // Each vectorized scalar flavor lifts to its own element-wise eval type so the worker keeps + // that flavor's batching contract (pandas Series, Arrow Array, or an iterator of batches). + val cases = Seq( + scalarPandasUDF -> PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + scalarArrowUDF -> PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + scalarPandasIterUDF -> PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + scalarArrowIterUDF -> PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF) + cases.foreach { case (udf, expectedEvalType) => + val df = arrayDF.select(transform(col("values"), x => udf(x)).as("r")) + val evalTypeName = PythonEvalType.toString(expectedEvalType) + assert(udfsInsideLambda(df).isEmpty, s"UDF still inside a lambda for $evalTypeName") + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + assert(lifted.head.evalType == expectedEvalType) + assert(lifted.head.dataType == ArrayType(udf.dataType, containsNull = true)) + assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType]) + } + } + + test("a UDF with a UDT type inside a lambda still fails analysis") { + // Lifting forces the Arrow element-wise eval type, which has no UDT fallback, so a UDF whose + // argument or return type involves a UDT is not rewritable and keeps the previous analysis + // error rather than failing at runtime. + val e = intercept[AnalysisException] { + arrayDF.select(transform(col("values"), x => udtUDF(x))).collect() + } + assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF") + } + + test("the rewritable predicate: eval types, zero-argument / iterator-kwarg / UDT UDFs") { + // A zero-arg call has no array to carry the iterated shape and would crash the worker; an + // iterator UDF takes no kwargs; a UDT would hit the Arrow path with no fallback - so the shared + // predicate rejects those, keeping the previous analysis error. A named argument on a + // non-iterator flavor is accepted (the lift keeps the NamedArgumentExpression). + val plain = PythonUDF("f", null, IntegerType, Seq(Literal(1)), + PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true) + assert(PythonUDF.isElementwiseRewritableUDF(plain)) + + // Every row-at-a-time and vectorized scalar eval type is rewritable; each maps to its own + // element-wise lifted eval type. + Seq( + PythonEvalType.SQL_ARROW_BATCHED_UDF -> PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_UDF -> PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF -> + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_UDF -> PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF -> + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF).foreach { + case (base, lifted) => + assert(PythonUDF.isElementwiseRewritableUDF(plain.copy(evalType = base))) + assert(PythonUDF.liftedElementwiseEvalType(base) == lifted) + } + assert(PythonUDF.liftedElementwiseEvalType(PythonEvalType.SQL_BATCHED_UDF) == + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF) + + // An already-lifted element-wise UDF is rewritable again (nested lambdas re-lift it), and its + // lifted eval type is itself - only the nesting depth changes. + Seq( + PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF, + PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF).foreach { ew => + assert(PythonUDF.isElementwiseRewritableUDF(plain.copy(evalType = ew))) + assert(PythonUDF.liftedElementwiseEvalType(ew) == ew) + } + + val zeroArg = plain.copy(children = Seq.empty) + assert(!PythonUDF.isElementwiseRewritableUDF(zeroArg)) + + // A named argument is rewritable on the non-iterator flavors (the lift keeps the + // NamedArgumentExpression as a direct child), but not on an iterator UDF (no kwargs there). + val named = plain.copy(children = Seq(NamedArgumentExpression("k", Literal(1)))) + assert(PythonUDF.isElementwiseRewritableUDF(named)) + assert(PythonUDF.isElementwiseRewritableUDF( + named.copy(evalType = PythonEvalType.SQL_SCALAR_PANDAS_UDF))) + assert(!PythonUDF.isElementwiseRewritableUDF( + named.copy(evalType = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF))) + assert(!PythonUDF.isElementwiseRewritableUDF( + named.copy(evalType = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF))) + + val udtReturn = plain.copy(dataType = new ExamplePointUDT) + assert(!PythonUDF.isElementwiseRewritableUDF(udtReturn)) + + val udtArg = plain.copy(children = Seq(Literal.create(null, new ExamplePointUDT))) + assert(!PythonUDF.isElementwiseRewritableUDF(udtArg)) + } + + test("a zero-argument UDF inside a lambda still fails analysis") { + // `transform(arr, x -> f())` has no argument to carry the iterated shape, so it is not + // rewritable and must keep failing analysis rather than crash the Python worker at runtime. + val e = intercept[AnalysisException] { + arrayDF.select(transform(col("values"), _ => pythonUDF())).collect() + } + assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF") + } + + test("a nondeterministic iterated argument is rejected") { + // The rewrite references the iterated argument several times and nondeterministic expressions + // are not subexpression-eliminated, so a nondeterministic argument like `shuffle(arr)` would be + // evaluated independently per reference and misalign the results. It must fail analysis. + val e = intercept[AnalysisException] { + arrayDF.select( + org.apache.spark.sql.functions.filter( + org.apache.spark.sql.functions.shuffle(col("values")), + x => pythonUDF(x))).collect() + } + assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF") + } + + test("the rewrite can be disabled by conf, restoring the previous error") { + withSQLConf(SQLConf.PYTHON_UDF_IN_HIGHER_ORDER_FUNCTION_ENABLED.key -> "false") { + val e = intercept[AnalysisException] { + arrayDF.select(transform(col("values"), x => pythonUDF(x))).collect() + } + assert(e.getCondition == "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF") + } + } + + test("the rewrite cannot be disabled via excludedRules") { + // The lambda rewrite is driven by `ExtractPythonUDFs` (not a standalone batch rule), and that + // rule is non-excludable, so a plan that only works because of the rewrite cannot be broken by + // excludedRules. Excluding either name must leave the UDF lifted out of the lambda. + withSQLConf( + SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> + Seq(ExtractPythonUDFFromLambda.ruleName, ExtractPythonUDFs.ruleName).mkString(",")) { + val df = arrayDF.select(transform(col("values"), x => pythonUDF(x)).as("r")) + assert(udfsInsideLambda(df).isEmpty) + } + } + + test("a UDF over only constants is still lifted per element") { + // SPARK-27052: `transform(arr, x -> udf(lit(10)))` does not read the element, but it is still + // lifted per element (`overArray` repeats the constant into an aligned array) so it keeps the + // lambda's call domain - once per element, zero times for a null/empty array - rather than + // being left to ExtractPythonUDFs, which would call it once per row. + val df = arrayDF.select(transform(col("values"), _ => pythonUDF(lit(10))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + } + + test("a UDF argument that is an expression over the element is lifted") { + val df = arrayDF.select(transform(col("values"), x => pythonUDF(x * lit(2))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + // The argument became an array-valued expression over the whole array. + assert(lifted.head.children.head.dataType.isInstanceOf[ArrayType]) + } + + test("an outer column argument is repeated into an aligned array for the lifted UDF") { + val df = Seq((Seq(1, 2), 10)).toDF("values", "base") + .select(transform(col("values"), x => pythonUDF(x, col("base"))).as("r")) + assert(udfsInsideLambda(df).isEmpty) + val lifted = liftedUDFs(df) + assert(lifted.size == 1) + // Both arguments are single-level arrays aligned with the iterated array: the element argument + // is the array itself, and the outer column is repeated into an aligned array so the worker + // flattens every argument uniformly. + lifted.head.children.foreach { c => + assert(c.dataType.isInstanceOf[ArrayType]) + assert(c.dataType.asInstanceOf[ArrayType].elementType == IntegerType) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala index 9e47473b42d75..de96f23e519b3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala @@ -22,7 +22,7 @@ import java.io.{File, FileWriter} import org.apache.spark.api.python.PythonException import org.apache.spark.api.python.PythonUtils import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, Row} -import org.apache.spark.sql.execution.FilterExec +import org.apache.spark.sql.execution.{FilterExec, LimitExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.datasources.DataSourceManager import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation} @@ -308,6 +308,280 @@ class PythonDataSourceSuite extends PythonDataSourceSuiteBase { } } + test("data source reader with limit pushdown") { + assume(shouldTestPandasUDFs) + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + | + |class SimpleDataSourceReader(DataSourceReader): + | def __init__(self): + | self.limit = None + | + | def pushLimit(self, limit): + | self.limit = limit + | return True + | + | def partitions(self): + | # A pushed limit lets the reader plan a single partition. + | assert self.limit == 2, self.limit + | return [InputPartition(0)] + | + | def read(self, partition): + | for i in range(self.limit): + | yield (i,) + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load().limit(2) + val plan = df.queryExecution.executedPlan + + collectFirst(plan) { + case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] => + val p = s.scan.asInstanceOf[PythonScan] + assert(p.getMetaData().get("PushedLimit").contains("LIMIT 2")) + }.getOrElse( + fail(s"PythonScan not found in the plan. Actual plan:\n$plan") + ) + + // Spark keeps its own limit: a Python data source is not trusted to honor the pushed limit. + assert(collectFirst(plan) { case l: LimitExec => l }.isDefined, + s"A limit operator should be retained. Actual plan:\n$plan") + + checkAnswer(df, Seq(Row(0), Row(1))) + } + } + + test("data source reader limit pushdown not supported by the reader") { + assume(shouldTestPandasUDFs) + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader + | + |class SimpleDataSourceReader(DataSourceReader): + | def pushLimit(self, limit): + | return False + | + | def read(self, partition): + | yield (0,) + | yield (1,) + | yield (2,) + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load().limit(2) + val plan = df.queryExecution.executedPlan + + collectFirst(plan) { + case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] => + val p = s.scan.asInstanceOf[PythonScan] + assert(!p.getMetaData().contains("PushedLimit")) + }.getOrElse( + fail(s"PythonScan not found in the plan. Actual plan:\n$plan") + ) + + checkAnswer(df, Seq(Row(0), Row(1))) + } + } + + test("limit pushdown is not used for top-N (orderBy then limit)") { + assume(shouldTestPandasUDFs) + // A LIMIT after an ORDER BY is a top-N: the limit cannot be pushed to the scan, because + // pushing it before the sort would drop rows the sort needs. The reader must not see it. + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + | + |class SimpleDataSourceReader(DataSourceReader): + | def pushLimit(self, limit): + | raise AssertionError("pushLimit must not be called for a top-N query") + | + | def partitions(self): + | return [InputPartition(0)] + | + | def read(self, partition): + | yield from [(2,), (1,), (3,)] + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load().orderBy("id").limit(2) + val plan = df.queryExecution.executedPlan + + collectFirst(plan) { + case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] => + val p = s.scan.asInstanceOf[PythonScan] + assert(!p.getMetaData().contains("PushedLimit")) + }.getOrElse( + fail(s"PythonScan not found in the plan. Actual plan:\n$plan") + ) + + checkAnswer(df, Seq(Row(1), Row(2))) + } + } + + test("no limit is pushed for a filter-only query with both pushdowns enabled") { + assume(shouldTestPandasUDFs) + // The filter-pushdown worker also handles limit pushdown, but a query without a LIMIT sends + // the -1 no-limit sentinel, so `pushLimit` must not be called even when the reader implements + // it and limit pushdown is enabled. + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + | + |class SimpleDataSourceReader(DataSourceReader): + | def pushFilters(self, filters): + | # Reject the filter so Spark re-applies it; this test is only about the limit. + | return filters + | + | def pushLimit(self, limit): + | raise AssertionError("pushLimit must not be called when the query has no limit") + | + | def partitions(self): + | return [InputPartition(0)] + | + | def read(self, partition): + | yield from [(1,), (2,)] + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf( + SQLConf.PYTHON_FILTER_PUSHDOWN_ENABLED.key -> "true", + SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load().filter("id = 1") + val plan = df.queryExecution.executedPlan + + collectFirst(plan) { + case s: BatchScanExec if s.scan.isInstanceOf[PythonScan] => + val p = s.scan.asInstanceOf[PythonScan] + assert(!p.getMetaData().contains("PushedLimit")) + }.getOrElse( + fail(s"PythonScan not found in the plan. Actual plan:\n$plan") + ) + + checkAnswer(df, Seq(Row(1))) + } + } + + test("limit pushdown does not leak into a reused base scan") { + assume(shouldTestPandasUDFs) + // A base DataFrame and its `.limit(n)` share the same Python data source instance. The read + // function produced while pushing the limit must stay scoped to the limited scan and not + // leak into the base scan, which would make the base DataFrame read too few rows. + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + | + |class SimpleDataSourceReader(DataSourceReader): + | def pushLimit(self, limit): + | self._limit = limit + | return True + | + | def partitions(self): + | return [InputPartition(0)] + | + | def read(self, partition): + | n = getattr(self, "_limit", 10) + | for i in range(n): + | yield (i,) + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf(SQLConf.PYTHON_LIMIT_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load() + checkAnswer(df.limit(2), Seq(Row(0), Row(1))) + checkAnswer(df, (0 until 10).map(Row(_))) + } + } + + test("filter pushdown does not leak into a reused base scan") { + assume(shouldTestPandasUDFs) + // Same as above for filter pushdown: the read function bound to the pushed filters must not + // leak into the base scan, which would make the base DataFrame return only filtered rows. + val dataSourceScript = + s""" + |from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + | + |class SimpleDataSourceReader(DataSourceReader): + | def pushFilters(self, filters): + | self._filtered = True + | return [] + | + | def partitions(self): + | return [InputPartition(0)] + | + | def read(self, partition): + | if getattr(self, "_filtered", False): + | yield (1,) + | else: + | for i in range(10): + | yield (i,) + | + |class SimpleDataSource(DataSource): + | def schema(self): + | return "id int" + | + | def reader(self, schema): + | return SimpleDataSourceReader() + |""".stripMargin + val schema = StructType.fromDDL("id INT") + val dataSource = + createUserDefinedPythonDataSource(name = dataSourceName, pythonScript = dataSourceScript) + withSQLConf(SQLConf.PYTHON_FILTER_PUSHDOWN_ENABLED.key -> "true") { + spark.dataSource.registerPython(dataSourceName, dataSource) + val df = spark.read.format(dataSourceName).schema(schema).load() + checkAnswer(df.filter("id = 1"), Seq(Row(1))) + checkAnswer(df, (0 until 10).map(Row(_))) + } + } + test("register data source") { assume(shouldTestPandasUDFs) val dataSourceScript = diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala index dfff8ed9e9758..88e1e7a9122b0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.python import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, Row} -import org.apache.spark.sql.functions.{array, avg, col, count, transform} +import org.apache.spark.sql.functions.{aggregate, array, avg, col, count, lit, transform} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types.LongType @@ -117,16 +117,28 @@ class PythonUDFSuite extends SharedSparkSession { assert(df.agg(pandasTestUDF(df("id"))).schema.fieldNames.exists(_.startsWith(udfName))) } - test("SPARK-48706: Negative test case for Python UDF in higher order functions") { + test("SPARK-27052: Python UDF in a higher order function lambda") { assume(shouldTestPythonUDFs) + // SPARK-48706 originally rejected this. `ExtractPythonUDFFromLambda` now rewrites it so the + // UDF is applied to the whole array outside the lambda, so it runs and returns a result. + checkAnswer( + spark.range(1).select(transform(array("id"), x => pythonTestUDF(x))), + Row(Seq(0))) + } + + test("SPARK-27052: Negative test case for Python UDF in higher order functions") { + assume(shouldTestPythonUDFs) + // A UDF reading `aggregate`'s accumulator is sequential, so there is no array to precompute + // over and the rewrite does not apply. This must still fail rather than give a wrong answer. checkError( exception = intercept[AnalysisException] { - spark.range(1).select(transform(array("id"), x => pythonTestUDF(x))).collect() + spark.range(1).select( + aggregate(array("id"), lit(0L), (acc, x) => pythonTestUDF(acc) + x)).collect() }, condition = "UNSUPPORTED_FEATURE.LAMBDA_FUNCTION_WITH_PYTHON_UDF", parameters = Map("funcName" -> "\"pyUDF(namedlambdavariable())\""), context = ExpectedContext( - "transform", s".*${this.getClass.getSimpleName}.*")) + "aggregate", s".*${this.getClass.getSimpleName}.*")) } test("SPARK-48666: Python UDF execution against partitioned column") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala new file mode 100644 index 0000000000000..e56a8a671dd4c --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkPreInitCleanupSuite.scala @@ -0,0 +1,166 @@ +/* + * 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.execution.python.streaming + +import java.io.{DataInputStream, DataOutputStream} +import java.util.{ArrayList, HashMap} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} + +import org.apache.spark.SparkException +import org.apache.spark.api.python.{PythonFunction, SimplePythonFunction} +import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.statefulprocessor.DriverStatefulProcessorHandleImpl +import org.apache.spark.sql.streaming.TimeMode +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.StructType + +class TransformWithStateInPySparkPreInitCleanupSuite extends SharedSparkSession { + + private val groupingKeySchema = StructType(Nil) + + private def newPythonFunction(): PythonFunction = { + new SimplePythonFunction( + command = Seq.empty[Byte], + envVars = new HashMap[String, String](), + pythonIncludes = new ArrayList[String](), + pythonExec = "python3", + pythonVer = "3", + broadcastVars = null, + accumulator = null) + } + + private def newDriverHandle(): DriverStatefulProcessorHandleImpl = { + new DriverStatefulProcessorHandleImpl(TimeMode.None(), null) + } + + private class StubPreInitRunner( + failInitWith: Option[Throwable] = None, + failProcessWith: Option[Throwable] = None, + failStopWith: Option[Throwable] = None) + extends TransformWithStateInPySparkPythonPreInitRunner( + newPythonFunction(), + "pyspark.sql.streaming.transform_with_state_driver_worker", + groupingKeySchema, + newDriverHandle()) { + + val workerAlive = new AtomicBoolean(false) + val initCount = new AtomicInteger(0) + val processCount = new AtomicInteger(0) + val stopCount = new AtomicInteger(0) + + override def init(): (DataOutputStream, DataInputStream) = { + initCount.incrementAndGet() + workerAlive.set(true) + failInitWith.foreach(error => throw error) + (null, null) + } + + override def process(): Unit = { + processCount.incrementAndGet() + failProcessWith.foreach(error => throw error) + } + + override def stop(): Unit = { + stopCount.incrementAndGet() + workerAlive.set(false) + failStopWith.foreach(error => throw error) + } + } + + private def runPreInit(runner: TransformWithStateInPySparkPythonPreInitRunner): Unit = { + TransformWithStateInPySparkExec.runPreInitRunner(runner) + } + + test("init failure after worker creation still stops the runner") { + val initFailure = new RuntimeException("init failed") + val runner = new StubPreInitRunner(failInitWith = Some(initFailure)) + + val thrown = intercept[Throwable] { + runPreInit(runner) + } + + assert(thrown eq initFailure) + assert(runner.initCount.get() === 1) + assert(runner.stopCount.get() === 1) + assert(!runner.workerAlive.get()) + assert(runner.processCount.get() === 0) + } + + test("repeated init failures do not accumulate live workers") { + val runners = (1 to 20).map { _ => + val runner = new StubPreInitRunner(failInitWith = Some(new RuntimeException("init failed"))) + intercept[Throwable] { + runPreInit(runner) + } + runner + } + + assert(runners.forall(_.initCount.get() === 1)) + assert(runners.count(_.workerAlive.get()) === 0) + } + + test("a stop failure does not mask the original init failure") { + val initFailure = new RuntimeException("init failed") + val stopFailure = new IllegalStateException("cleanup failed") + val runner = + new StubPreInitRunner(failInitWith = Some(initFailure), failStopWith = Some(stopFailure)) + + val thrown = intercept[Throwable] { + runPreInit(runner) + } + + assert(thrown eq initFailure) + assert(thrown.getSuppressed.contains(stopFailure)) + assert(runner.stopCount.get() === 1) + assert(!runner.workerAlive.get()) + } + + test("process failure is wrapped but still stops the runner") { + val processFailure = new RuntimeException("worker crashed") + val runner = new StubPreInitRunner(failProcessWith = Some(processFailure)) + + val thrown = intercept[SparkException] { + runPreInit(runner) + } + + assert(thrown.getMessage.contains("exited unexpectedly (crashed)")) + assert(thrown.getCause eq processFailure) + assert(runner.stopCount.get() === 1) + assert(!runner.workerAlive.get()) + } + + test("success path stops the runner exactly once") { + val runner = new StubPreInitRunner() + + runPreInit(runner) + + assert(runner.initCount.get() === 1) + assert(runner.processCount.get() === 1) + assert(runner.stopCount.get() === 1) + assert(!runner.workerAlive.get()) + } + + test("stop before startStateServer ran does not throw NPE") { + val runner = new TransformWithStateInPySparkPythonPreInitRunner( + newPythonFunction(), + "pyspark.sql.streaming.transform_with_state_driver_worker", + groupingKeySchema, + newDriverHandle()) + + runner.stop() + runner.stop() + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala index e253a6aa45c35..ea1c8a0a5d8e0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/streaming/TransformWithStateInPySparkStateServerSuite.scala @@ -16,15 +16,27 @@ */ package org.apache.spark.sql.execution.python.streaming -import java.io.DataOutputStream -import java.nio.channels.ServerSocketChannel +import java.io.{DataOutputStream, InterruptedIOException} +import java.net.{InetSocketAddress, Socket} +import java.nio.ByteBuffer +import java.nio.channels.{ + AsynchronousCloseException, + ClosedByInterruptException, + ClosedChannelException, + ServerSocketChannel, + SocketChannel +} +import java.util.concurrent.atomic.AtomicReference import scala.collection.mutable +import scala.concurrent.duration._ import com.google.protobuf.ByteString import org.mockito.ArgumentMatchers.{any, argThat} import org.mockito.Mockito.{mock, times, verify, when} +import org.mockito.invocation.InvocationOnMock import org.scalatest.BeforeAndAfterEach +import org.scalatest.concurrent.Eventually.{eventually, timeout} import org.apache.spark.SparkFunSuite import org.apache.spark.sql.{Encoder, Row} @@ -110,6 +122,34 @@ class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with Bef .thenReturn(Seq(getIntegerRow(1))) } + test("run closes the accepted socket once the request loop ends") { + val acceptedSocket = mock(classOf[SocketChannel]) + when(serverSocket.accept()).thenReturn(acceptedSocket) + when(acceptedSocket.socket()).thenReturn(mock(classOf[Socket])) + // Ends the request loop right away: this test is about the socket, not the requests. + when(acceptedSocket.isConnected).thenReturn(false) + + stateServer.run() + + verify(acceptedSocket).close() + } + + test("run closes the accepted socket when the client disconnects") { + val acceptedSocket = mock(classOf[SocketChannel]) + when(serverSocket.accept()).thenReturn(acceptedSocket) + when(acceptedSocket.socket()).thenReturn(mock(classOf[Socket])) + when(acceptedSocket.isConnected).thenReturn(true) + // Channels.newInputStream synchronizes on this before reading. + when(acceptedSocket.blockingLock()).thenReturn(new Object) + when(acceptedSocket.isBlocking).thenReturn(true) + // No bytes ever arrive, so the read hits EOF and the loop returns early. + when(acceptedSocket.read(any(classOf[ByteBuffer]))).thenReturn(-1) + + stateServer.run() + + verify(acceptedSocket).close() + } + test("set handle state") { val message = StatefulProcessorCall.newBuilder().setSetHandleState( SetHandleState.newBuilder().setState(HandleState.CREATED).build()).build() @@ -637,6 +677,100 @@ class TransformWithStateInPySparkStateServerSuite extends SparkFunSuite with Bef verify(outputStream).writeInt(argThat((x: Int) => x > 0)) } + Seq( + ("InterruptedException", () => new InterruptedException()), + ("InterruptedIOException", () => new InterruptedIOException()), + ("ClosedByInterruptException", () => new ClosedByInterruptException()) + ).foreach { case (name, newException) => + test(s"run handles $name while waiting for the Python worker") { + Thread.interrupted() + val socket = mock(classOf[ServerSocketChannel]) + when(socket.accept()) + .thenAnswer((_: InvocationOnMock) => throw newException()) + + try { + newStateServer(socket).run() + assert(Thread.currentThread().isInterrupted) + } finally { + Thread.interrupted() + } + + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } + } + + Seq( + ("AsynchronousCloseException", () => new AsynchronousCloseException()), + ("ClosedChannelException", () => new ClosedChannelException()) + ).foreach { case (name, newException) => + test(s"run handles $name while waiting for the Python worker") { + Thread.interrupted() + val socket = mock(classOf[ServerSocketChannel]) + when(socket.accept()) + .thenAnswer((_: InvocationOnMock) => throw newException()) + + newStateServer(socket).run() + + assert(!Thread.currentThread().isInterrupted) + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } + } + + Seq( + ("before accept", true), + ("while blocked in accept", false) + ).foreach { case (name, interruptBeforeRun) => + test(s"run handles real channel shutdown $name") { + val socket = ServerSocketChannel.open() + socket.bind(new InetSocketAddress("127.0.0.1", 0)) + val failure = new AtomicReference[Throwable]() + val listener = new Thread(() => { + if (interruptBeforeRun) { + Thread.currentThread().interrupt() + } + try { + newStateServer(socket).run() + } catch { + case t: Throwable => failure.set(t) + } + }) + + try { + listener.start() + if (!interruptBeforeRun) { + eventually(timeout(10.seconds)) { + assert(listener.getStackTrace.exists(_.getMethodName == "accept")) + } + listener.interrupt() + } + socket.close() + listener.join(10000) + + assert(!listener.isAlive) + assert(failure.get() == null) + assert(!socket.isOpen) + verify(statefulProcessorHandle).setHandleState(StatefulProcessorHandleState.CLOSED) + verify(outputStream, times(0)).writeInt(any[Int]) + } finally { + listener.interrupt() + socket.close() + listener.join(10000) + } + } + } + + private def newStateServer( + socket: ServerSocketChannel): TransformWithStateInPySparkStateServer = { + new TransformWithStateInPySparkStateServer(socket, + statefulProcessorHandle, groupingKeySchema, 2, + batchTimestampMs, eventTimeWatermarkForEviction, + outputStream, valueStateMap, transformWithStateInPySparkDeserializer, + listStateMap, mutable.HashMap[String, Iterator[Row]](), mapStateMap, + mutable.HashMap[String, Iterator[(Row, Row)]](), expiryTimerIter, listTimerMap) + } + private def getIntegerRow(value: Int): Row = { new GenericRowWithSchema(Array(value), stateSchema) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala index df29d94ec32db..eae45e7556775 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/RocksDBSuite.scala @@ -1054,6 +1054,42 @@ class RocksDBSuite extends AlsoTestWithRocksDBFeatures with SharedSparkSession } } + test("RocksDB: split maintenance methods upload snapshots and clean up separately") { + val remoteDir = Utils.createTempDir().toString + new File(remoteDir).delete() + val conf = dbConf.copy(enableChangelogCheckpointing = true, + minVersionsToRetain = 3, minDeltasForSnapshot = 1, minVersionsToDelete = 3) + withDB(remoteDir, conf = conf) { db => + // Commit 5 versions, uploading snapshots after each via doSnapshotMaintenance. + for (version <- 0 to 4) { + db.load(version) + db.put(version.toString, version.toString) + db.commit() + db.doSnapshotMaintenance() + } + assert(snapshotVersionsPresent(remoteDir) == (1 to 5)) + assert(changelogVersionsPresent(remoteDir) == (1 to 5)) + + // Commit 1 more version without maintenance. + // stale versions: (1, 2, 3), keep versions: (4, 5, 6) + db.load(5) + db.put("5", "5") + db.commit() + assert(snapshotVersionsPresent(remoteDir) == (1 to 5)) + assert(changelogVersionsPresent(remoteDir) == (1 to 6)) + + // doSnapshotMaintenance should upload version 6 and not clean up. + db.doSnapshotMaintenance() + assert(snapshotVersionsPresent(remoteDir) == (1 to 6)) + assert(changelogVersionsPresent(remoteDir) == (1 to 6)) + + // doCleanupMaintenance should delete stale versions (1, 2, 3) and not upload a snapshot. + db.doCleanupMaintenance() + assert(snapshotVersionsPresent(remoteDir) == Seq(4, 5, 6)) + assert(changelogVersionsPresent(remoteDir) == Seq(4, 5, 6)) + } + } + testWithStateStoreCheckpointIdsAndColumnFamilies( "RocksDB: minDeltasForSnapshot", TestWithChangelogCheckpointingEnabled) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionAllColumnFamiliesWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionAllColumnFamiliesWriterSuite.scala index 66617a4d0908b..a6d9ad985efda 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionAllColumnFamiliesWriterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionAllColumnFamiliesWriterSuite.scala @@ -32,12 +32,14 @@ import org.apache.spark.sql.streaming.{InputEvent, ListStateTTLProcessor, MapInp import org.apache.spark.sql.streaming.util.{StreamManualClock, TTLProcessorUtils} import org.apache.spark.sql.streaming.util.{EventTimeTimerProcessor, MultiStateVarProcessor, MultiStateVarProcessorTestUtils, TimerTestUtils} import org.apache.spark.sql.types.StructType +import org.apache.spark.tags.SlowSQLTest /** * Test suite for StatePartitionAllColumnFamiliesWriter. * Tests the writer's ability to correctly write raw bytes read from * StatePartitionAllColumnFamiliesReader to a state store without loading previous versions. */ +@SlowSQLTest class StatePartitionAllColumnFamiliesWriterSuite extends StateDataSourceTestBase { import testImplicits._ diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreDecoupledMaintenanceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreDecoupledMaintenanceSuite.scala new file mode 100644 index 0000000000000..cea84e80efd07 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreDecoupledMaintenanceSuite.scala @@ -0,0 +1,1259 @@ +/* + * 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.execution.streaming.state + +import java.util.UUID +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import scala.collection.mutable + +import org.apache.hadoop.conf.Configuration +import org.apache.logging.log4j.Level +import org.scalatest.{BeforeAndAfter, PrivateMethodTester} +import org.scalatest.concurrent.Eventually._ +import org.scalatest.time.SpanSugar._ + +import org.apache.spark._ +import org.apache.spark.internal.Logging +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.SQLConf.STATE_STORE_PROVIDER_CLASS +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest + +/** + * A fake StateStoreProvider that gives tests deterministic control over + * snapshot and cleanup maintenance timing using a latch-based handshake. + * + * Each operation follows a two-phase pattern: + * 1. The op counts down its enteredLatch, telling the test "I'm running." + * 2. The op blocks on its continueSignal until the test counts it down. + * + * The scheduler runs normally and submits tasks to the pools, but those + * tasks block inside the provider's maintenance methods. This holds the + * pool threads mid-execution, keeping downstream logic (source handling, + * queue routing, close) from running until the test releases the latch. + */ +class BlockingMaintenanceProvider extends StateStoreProvider + with Logging { + private var id: StateStoreId = null + + // Per-instance state. No shared static fields, so stale scheduler + // cycles from a previous test use the old instance's latches (already + // counted down) and finish immediately. No cross-test interference. + @volatile var snapshotThreadName: String = "" + @volatile var cleanupThreadName: String = "" + @volatile var closeThreadName: String = "" + @volatile var snapshotShouldThrow: Boolean = false + @volatile var cleanupShouldThrow: Boolean = false + @volatile var closeShouldBlock: Boolean = false + + val snapshotEnteredLatch = new CountDownLatch(1) + val cleanupEnteredLatch = new CountDownLatch(1) + val snapshotContinueSignal = new CountDownLatch(1) + val cleanupContinueSignal = new CountDownLatch(1) + val closeEnteredLatch = new CountDownLatch(1) + val closeContinueSignal = new CountDownLatch(1) + + override def init( + stateStoreId: StateStoreId, + keySchema: StructType, + valueSchema: StructType, + keyStateEncoderSpec: KeyStateEncoderSpec, + useColumnFamilies: Boolean, + storeConfs: StateStoreConf, + hadoopConf: Configuration, + useMultipleValuesPerKey: Boolean = false, + stateSchemaProvider: Option[StateSchemaProvider] = None + ): Unit = { + id = stateStoreId + } + + override def stateStoreId: StateStoreId = id + + override def close(): Unit = { + closeThreadName = Thread.currentThread.getName + if (closeShouldBlock) { + closeEnteredLatch.countDown() + closeContinueSignal.await() + } + } + + /** Returns null because tests using this provider do not need a real + * store. They only exercise the maintenance scheduler and close paths. */ + override def getStore( + version: Long, + uniqueId: Option[String], + forceSnapshotOnCommit: Boolean = false, + loadEmpty: Boolean = false): StateStore = null + + /** Signals entry, then blocks until the test releases the continue latch. */ + override def doSnapshotMaintenance(): Unit = { + snapshotThreadName = Thread.currentThread.getName + logInfo(s"Snapshot maintenance entered on" + + s" ${Thread.currentThread.getName}") + snapshotEnteredLatch.countDown() + snapshotContinueSignal.await() + logInfo(s"Snapshot maintenance continuing on" + + s" ${Thread.currentThread.getName}") + if (snapshotShouldThrow) { + throw new RuntimeException("snapshot error") + } + } + + /** Same handshake as doSnapshotMaintenance but for cleanup. */ + override def doCleanupMaintenance(): Unit = { + cleanupThreadName = Thread.currentThread.getName + logInfo(s"Cleanup maintenance entered on" + + s" ${Thread.currentThread.getName}") + cleanupEnteredLatch.countDown() + cleanupContinueSignal.await() + logInfo(s"Cleanup maintenance continuing on" + + s" ${Thread.currentThread.getName}") + if (cleanupShouldThrow) { + throw new RuntimeException("cleanup error") + } + } +} + +class FakeStateStoreProviderTracksCloseThread extends StateStoreProvider { + import FakeStateStoreProviderTracksCloseThread._ + private var id: StateStoreId = null + + override def init( + stateStoreId: StateStoreId, + keySchema: StructType, + valueSchema: StructType, + keyStateEncoderSpec: KeyStateEncoderSpec, + useColumnFamilies: Boolean, + storeConfs: StateStoreConf, + hadoopConf: Configuration, + useMultipleValuesPerKey: Boolean = false, + stateSchemaProvider: Option[StateSchemaProvider] = None): Unit = { + id = stateStoreId + } + + override def stateStoreId: StateStoreId = id + + override def close(): Unit = { + closeThreadNames = Thread.currentThread.getName :: closeThreadNames + } + + override def getStore( + version: Long, + uniqueId: Option[String], + forceSnapshotOnCommit: Boolean = false, + loadEmpty: Boolean = false): StateStore = null +} + +private object FakeStateStoreProviderTracksCloseThread { + var closeThreadNames: List[String] = Nil +} + +@ExtendedSQLTest +abstract class StateStoreDecoupledMaintenanceSuiteBase[ + ProviderClass <: StateStoreProvider] + extends SparkFunSuite + with BeforeAndAfter + with PrivateMethodTester { + + import StateStoreCoordinatorSuite._ + + before { + StateStore.stop() + require(!StateStore.isMaintenanceRunning) + } + + after { + StateStore.stop() + require(!StateStore.isMaintenanceRunning) + } + + private def getDefaultSQLConf( + minDeltasForSnapshot: Int, + numOfVersToRetainInMemory: Int): SQLConf = { + val sqlConf = new SQLConf() + sqlConf.setConf(SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT, + minDeltasForSnapshot) + sqlConf.setConf(SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY, + numOfVersToRetainInMemory) + sqlConf.setConf(SQLConf.MIN_BATCHES_TO_RETAIN, 2) + sqlConf + } + + private def getUnloadQueue() = { + val f = PrivateMethod[ConcurrentLinkedQueue[(StateStoreProviderId, + StateStoreProvider, MaintenanceOpRequest)]]( + Symbol("unloadedProvidersToClose")) + StateStore invokePrivate f() + } + + private def getSnapshotPartitions() = { + val f = PrivateMethod[mutable.HashSet[StateStoreProviderId]]( + Symbol("snapshotPartitions")) + StateStore invokePrivate f() + } + + private def getCleanupPartitions() = { + val f = PrivateMethod[mutable.HashSet[StateStoreProviderId]]( + Symbol("cleanupPartitions")) + StateStore invokePrivate f() + } + + private def getLoadedProviders() = { + val f = PrivateMethod[ + mutable.HashMap[StateStoreProviderId, StateStoreProvider]]( + Symbol("loadedProviders")) + StateStore invokePrivate f() + } + + private def getMaintenanceTask() = { + val f = PrivateMethod[StateStore.MaintenanceTask]( + Symbol("maintenanceTask")) + StateStore invokePrivate f() + } + + private def getBlockingProvider( + id: StateStoreProviderId): BlockingMaintenanceProvider = { + val loaded = getLoadedProviders() + loaded.synchronized { loaded.get(id).get } + .asInstanceOf[BlockingMaintenanceProvider] + } + + private def maintenanceStoreConf( + providerClass: Class[_], + interval: Long = 100L, + numThreads: Int = 4): StateStoreConf = { + val sqlConf = getDefaultSQLConf( + SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get, + SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get) + sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, interval) + sqlConf.setConf(SQLConf.NUM_STATE_STORE_MAINTENANCE_THREADS, numThreads) + sqlConf.setConf(STATE_STORE_PROVIDER_CLASS, providerClass.getName) + new StateStoreConf(sqlConf) + } + + private def withSparkContext(body: SparkContext => Unit): Unit = { + body(SparkContext.getOrCreate( + new SparkConf().setMaster("local").setAppName("test"))) + } + + private def loadNullProvider( + dir: String, + storeConf: StateStoreConf, + partition: Int = 0): StateStoreProviderId = { + val storeId = StateStoreProviderId( + StateStoreId(dir, 0, partition), UUID.randomUUID) + StateStore.get(storeId, null, null, NoPrefixKeyStateEncoderSpec(null), 0, + stateStoreCkptId = None, stateSchemaBroadcast = None, + useColumnFamilies = false, storeConf, new Configuration()) + storeId + } + + test("SPARK-51596: task thread unload lifecycle " + + "from queue to close") { + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + // Long interval so close can only happen via triggerNow, not the + // periodic scheduler tick. Without triggerNow the test fails at the + // snapshot latch because the first cycle never fires within the 10s + // timeout. + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], interval = 30000L) + val id1 = loadNullProvider("lifecycle", storeConf) + val bp = getBlockingProvider(id1) + + val queue = getUnloadQueue() + assert(StateStore.isLoaded(id1)) + assert(queue.isEmpty, "Queue should start empty") + + // Make stale and load another provider to trigger task thread queueing. + // Use a non-blocking provider for id2 since we don't need to + // observe its maintenance. + coordinatorRef.reportActiveInstance(id1, "otherhost", "otherexec", Seq.empty) + val storeConf2 = maintenanceStoreConf( + classOf[FakeStateStoreProviderTracksCloseThread]) + val id2 = loadNullProvider("lifecycle", storeConf2, partition = 1) + + assert(!StateStore.isLoaded(id1), "Provider1 should be removed") + assert(StateStore.isLoaded(id2), "Provider2 should still be loaded") + + // The task thread queued id1 with All. We can't peek the queue + // here because triggerNow fires immediately after queueing, draining + // it before we can inspect. snapshotEnteredLatch proves the entry + // was consumed and snapshot was submitted. + + // Step 2: Scheduler (via triggerNow) submits first op as + // FromUnloadedProvidersQueue. Snapshot enters latch. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS) && + bp.snapshotEnteredLatch.getCount == 0, "snapshot should have started") + + // Step 3: Release snapshot. Post-work queues remaining op (Cleanup) + // via otherMaintenanceOpRequest. + bp.snapshotContinueSignal.countDown() + + // Ideally we would peek the queue here to verify the entry is Cleanup + // (via otherMaintenanceOpRequest). But triggerNow drains it before we + // can peek. Instead, cleanupEnteredLatch being counted down proves + // Cleanup was queued and submitted. If the entry were Snapshot, + // doSnapshotMaintenance would have been called instead. + + // Step 4: Scheduler (via triggerNow) picks up Cleanup as + // FromUnloadedProvidersQueue. Cleanup enters. + assert(bp.cleanupEnteredLatch.await(10, TimeUnit.SECONDS) && + bp.cleanupEnteredLatch.getCount == 0, "cleanup should have started") + + // Verify intermediate queue state: snapshot's post-work queued Cleanup + // and scheduler drained it. Cleanup is now running, queue is empty. + assert(queue.isEmpty, "queue should be drained while cleanup runs") + + // Step 5: Release cleanup. FromUnloadedProvidersQueue calls closeProvider. + bp.cleanupContinueSignal.countDown() + + // Verify provider was closed on cleanup pool. + eventually(timeout(10.seconds)) { + assert(bp.closeThreadName.contains( + "state-store-maintenance-low-priority"), + "close should happen on cleanup pool thread, but was on: " + + bp.closeThreadName) + assert(!StateStore.isLoaded(id1), + "provider should be removed from loadedProviders") + assert(queue.isEmpty, "Queue should be drained") + } + } + } + } + + test("tryClaimPartition returns true first call, false second, " + + "true for different opType") { + val id = StateStoreProviderId( + StateStoreId("dir", 0, 0), UUID.randomUUID) + val id2 = StateStoreProviderId( + StateStoreId("dir", 0, 1), UUID.randomUUID) + + try { + // First claim for snapshot succeeds + assert(StateStore.tryClaimPartition(id, MaintenanceOpType.Snapshot)) + // Second claim for same id + opType fails + assert(!StateStore.tryClaimPartition(id, MaintenanceOpType.Snapshot)) + // Claim for same id but different opType succeeds + assert(StateStore.tryClaimPartition(id, MaintenanceOpType.Cleanup)) + // That one is also occupied now + assert(!StateStore.tryClaimPartition(id, MaintenanceOpType.Cleanup)) + + // Different id can still claim both + assert(StateStore.tryClaimPartition(id2, MaintenanceOpType.Snapshot)) + assert(StateStore.tryClaimPartition(id2, MaintenanceOpType.Cleanup)) + } finally { + getSnapshotPartitions().clear() + getCleanupPartitions().clear() + } + } + + test("otherMaintenanceOpRequest maps correctly") { + assert(StateStore.otherMaintenanceOpRequest(MaintenanceOpType.Snapshot) + === MaintenanceOpRequest.Cleanup) + assert(StateStore.otherMaintenanceOpRequest(MaintenanceOpType.Cleanup) + === MaintenanceOpRequest.Snapshot) + } + + test("closeProvider sets unloaded even if close() throws") { + val storeId = StateStoreProviderId(StateStoreId("closeTest", 0, 0), UUID.randomUUID) + val callOrder = new mutable.ArrayBuffer[String]() + val provider = new FakeStateStoreProviderTracksCloseThread { + override def close(): Unit = { + callOrder += "close" + throw new RuntimeException("close failed") + } + override def setUnloaded(): Unit = { + callOrder += "setUnloaded" + super.setUnloaded() + } + } + provider.init( + storeId.storeId, null, null, NoPrefixKeyStateEncoderSpec(null), + useColumnFamilies = false, null, null) + + assert(!provider.unloaded) + intercept[RuntimeException] { + StateStore.closeProvider(storeId, provider) + } + assert(provider.unloaded, + "setUnloaded should run even if close() throws") + assert(callOrder === Seq("close", "setUnloaded"), + "setUnloaded should run after close() even if it throws") + } + + test("concurrent snapshot and cleanup on same provider " + + "both succeed") { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val storeId = loadNullProvider("concurrentDir", storeConf) + val bp = getBlockingProvider(storeId) + + // Wait for both snapshot and cleanup to enter + assert(bp.snapshotEnteredLatch + .await(30, TimeUnit.SECONDS), "snapshot should have started") + assert(bp.cleanupEnteredLatch + .await(30, TimeUnit.SECONDS), "cleanup should have started") + + // Scheduler is no longer needed. Stop and wait so no new + // cycles interfere with assertions below. + getMaintenanceTask().stopAndAwait() + + // Both are running on their respective pool threads. + assert(bp.snapshotThreadName + .startsWith("state-store-maintenance-high-priority")) + assert(bp.cleanupThreadName + .startsWith("state-store-maintenance-low-priority")) + + // Partition sets should be claimed while both are running. + assert(!StateStore.tryClaimPartition(storeId, MaintenanceOpType.Snapshot), + "snapshot partition set should be occupied") + assert(!StateStore.tryClaimPartition(storeId, MaintenanceOpType.Cleanup), + "cleanup partition set should be occupied") + + // Read lock should be held while maintenance ops are running. + assert(bp.maintenanceLock.getReadLockCount == 2, + "both pool threads should hold the read lock") + + // Release both to finish + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + + // Verify all ops completed by checking partition sets and read + // lock are released. Use reflection to read the sets without + // claiming (tryClaimPartition has side effects that break + // eventually retries). + eventually(timeout(10.seconds)) { + assert(!getSnapshotPartitions().contains(storeId), + "snapshot partition set should be released") + assert(!getCleanupPartitions().contains(storeId), + "cleanup partition set should be released") + assert(bp.maintenanceLock.getReadLockCount == 0, + "read lock should be released after maintenance completes") + } + } + } + } + + test("partition sets and locks released when maintenance throws, " + + "write lock blocks until read lock is freed") { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val storeId = loadNullProvider("errorDir", storeConf) + val bp = getBlockingProvider(storeId) + bp.snapshotShouldThrow = true + bp.closeShouldBlock = true + + // Wait for all ops to enter. + assert(bp.snapshotEnteredLatch + .await(10, TimeUnit.SECONDS), "snapshot should start") + assert(bp.cleanupEnteredLatch + .await(10, TimeUnit.SECONDS), "cleanup should start") + + // Scheduler is no longer needed. Stop and wait so no new + // cycles interfere with assertions below. + getMaintenanceTask().stopAndAwait() + + // Partition set is claimed while running. + assert(!StateStore.tryClaimPartition(storeId, MaintenanceOpType.Snapshot), + "snapshot set should be occupied") + + assert(bp.maintenanceLock.getReadLockCount == 2, + "both pool threads should hold the read lock") + + // Release snapshot. It will throw. The error handler tries to + // acquire the write lock, but cleanup still holds a read lock. + bp.snapshotContinueSignal.countDown() + + // Wait for the error handler to be blocked on the write lock. + eventually(timeout(5.seconds)) { + assert(bp.maintenanceLock.getQueueLength > 0, + "error handler should be waiting for write lock") + } + // Close has not been entered because the write lock is blocked. + assert(bp.closeEnteredLatch.getCount == 1, + "close should not be entered while write lock is blocked") + + // Release cleanup. Read lock freed. Write lock unblocks. Close called. + bp.cleanupContinueSignal.countDown() + + assert(bp.closeEnteredLatch + .await(10, TimeUnit.SECONDS), "close should be called") + assert(bp.maintenanceLock.isWriteLocked, + "write lock should be held during close") + + // Release close to let error handler finish. + bp.closeContinueSignal.countDown() + + // Wait for the error and finally block to complete. + eventually(timeout(10.seconds)) { + assert(!StateStore.isLoaded(storeId), + "provider should be unloaded after throw") + assert(bp.closeThreadName.nonEmpty, + "provider should be closed after throw") + assert(!getSnapshotPartitions().contains(storeId), + "snapshot set should be released by finally block after throw") + assert(bp.maintenanceLock.getReadLockCount == 0, + "read lock should be released after error handling") + assert(!bp.maintenanceLock.isWriteLocked, + "write lock should be released after error handling") + } + } + } + } + + private def testRequeue(opRequest: MaintenanceOpRequest): Unit = { + val logAppender = new LogAppender("requeue-log", maxEvents = 1000) + logAppender.setThreshold(Level.INFO) + // Scope to the StateStore logger so unrelated INFO logs (e.g. a streaming + // query leaked from a prior suite) cannot flood the appender and trip its cap. + val loggerName = StateStore.getClass.getName.stripSuffix("$") + withLogAppender(logAppender, + loggerNames = Seq(loggerName), level = Some(Level.INFO)) { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val storeId = loadNullProvider("requeueDir", storeConf) + val bp = getBlockingProvider(storeId) + val queue = getUnloadQueue() + + try { + // Wait for both ops to enter, occupying both partition sets. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + assert(bp.cleanupEnteredLatch.await(10, TimeUnit.SECONDS)) + + // Add an entry to the queue. The scheduler will try to drain + // it but the partition set is occupied (held by the blocked + // task above), so it should be requeued. + queue.add((storeId, bp, opRequest)) + + // Wait for the scheduler to attempt draining and verify + // the requeue log. Queue checks are inside eventually to + // handle the momentary gap between the poll removing the + // entry and offer putting the entry back. + // Queue should still have the entry. + eventually(timeout(10.seconds)) { + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("Had to requeue")), + s"scheduler should have logged requeue for $opRequest") + val peeked = queue.peek() + assert(peeked != null, + s"$opRequest entry should have been requeued") + val (requeuedId, _, requeuedOp) = peeked + assert(requeuedId == storeId) + assert(requeuedOp == opRequest) + } + } finally { + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + queue.clear() + } + } + } + } + } + + test("Snapshot entry requeues when snapshot partition set is occupied") { + testRequeue(MaintenanceOpRequest.Snapshot) + } + + test("Cleanup entry requeues when cleanup partition set is occupied") { + testRequeue(MaintenanceOpRequest.Cleanup) + } + + test("All entry requeues when both partition sets are occupied") { + testRequeue(MaintenanceOpRequest.All) + } + + test("When MaintenanceOpRequest is All, cleanup is submitted " + + "if snapshot partition set is occupied") { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + // Long interval so we can set up before the first + // cycle fires (5s initial delay). + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], interval = 5000L) + val id = loadNullProvider("shortCircuit", storeConf) + val bp = getBlockingProvider(id) + + try { + // Claim snapshot partition set. + assert(StateStore.tryClaimPartition(id, MaintenanceOpType.Snapshot)) + + // Remove from loadedProviders so scheduler doesn't submit + // cleanup for this provider by iterating through it. + // Only the queue entry should submit cleanup. + getLoadedProviders().synchronized { getLoadedProviders().remove(id) } + + // Add All entry. Scheduler's first cycle (at 5s) drains it, + // tries snapshot (occupied), falls through to cleanup. + // Cleanup blocks on bp's latch, keeping the partition set. + getUnloadQueue().add((id, bp, MaintenanceOpRequest.All)) + + eventually(timeout(10.seconds)) { + assert(getCleanupPartitions().contains(id), + "cleanup should be claimed") + assert(getUnloadQueue().isEmpty, "queue should be drained") + } + } finally { + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + getSnapshotPartitions().remove(id) + } + } + } + } + + test("canProcess is false and maintenance is skipped " + + "when provider has already been unloaded") { + val logAppender = new LogAppender("canProcess-log", maxEvents = 1000) + logAppender.setThreshold(Level.INFO) + // Scope to the StateStore logger so unrelated INFO logs (e.g. a streaming + // query leaked from a prior suite) cannot flood the appender and trip its cap. + val loggerName = StateStore.getClass.getName.stripSuffix("$") + withLogAppender(logAppender, + loggerNames = Seq(loggerName), level = Some(Level.INFO)) { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], interval = 5000L) + val id = loadNullProvider("canProcess", storeConf) + val bp = getBlockingProvider(id) + + // Mark unloaded before the first maintenance cycle fires (5s + // initial delay). canProcess checks !provider.unloaded and will + // return false, skipping maintenance entirely. + bp.setUnloaded() + + // Wait for the "Skipping maintenance" log proving canProcess + // was false and maintenance was skipped. + eventually(timeout(10.seconds)) { + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("Skipping maintenance")), + "should log skipping maintenance for unloaded provider") + } + assert(bp.snapshotEnteredLatch.getCount == 1, + "snapshot should not have entered") + assert(bp.cleanupEnteredLatch.getCount == 1, + "cleanup should not have entered") + } + } + } + } + + test("canProcess is false and maintenance is skipped " + + "when provider instance differs") { + val logAppender = new LogAppender("canProcess-stale-log", maxEvents = 1000) + logAppender.setThreshold(Level.INFO) + // Scope to the StateStore logger so unrelated INFO logs (e.g. a streaming + // query leaked from a prior suite) cannot flood the appender and trip its cap. + val loggerName = StateStore.getClass.getName.stripSuffix("$") + withLogAppender(logAppender, + loggerNames = Seq(loggerName), level = Some(Level.INFO)) { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + // Long interval so we can block the cleanup pool before the + // first cycle fires. numThreads=2 gives each pool 1 thread. + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], + interval = 5000L, numThreads = 2) + val id = loadNullProvider("canProcessStale", storeConf) + val bp = getBlockingProvider(id) + + // Block the cleanup pool's thread with a dummy task before the + // first cycle fires (5s away). + val cleanupPoolField = PrivateMethod[StateStore.MaintenanceThreadPool]( + Symbol("lowPriorityThreadPool")) + val cleanupPool = StateStore invokePrivate cleanupPoolField() + val blockLatch = new CountDownLatch(1) + cleanupPool.execute(() => blockLatch.await()) + + // Wait for the first cycle. Snapshot runs freely. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + + // Stop the scheduler so the next cycle doesn't run the + // replacement (which lacks thread locals). + getMaintenanceTask().stopAndAwait() + + bp.snapshotContinueSignal.countDown() + + // Replace A with a different instance while cleanup is waiting. + val replacement = new FakeStateStoreProviderTracksCloseThread + replacement.init(id.storeId, null, null, + NoPrefixKeyStateEncoderSpec(null), + useColumnFamilies = false, null, null) + val loaded = getLoadedProviders() + loaded.synchronized { loaded.put(id, replacement) } + + // Release the dummy. Cleanup starts, canProcess sees + // contains(A) is false (replacement is there), skips. + blockLatch.countDown() + + eventually(timeout(10.seconds)) { + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains("Skipping maintenance")), + "should log skipping maintenance for stale instance") + } + assert(bp.cleanupEnteredLatch.getCount == 1, + "cleanup should not have entered") + } + } + } + } + + test("FromLoadedProviders unload: reloaded provider is not removed nor queued") { + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val id = loadNullProvider("staleInstance", storeConf) + val bp = getBlockingProvider(id) + + // Mark as needing to be closed. + coordinatorRef.reportActiveInstance(id, "otherhost", "otherexec", Seq.empty) + + // Wait for snapshot to enter. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + + // Stop only the scheduler (not the pools) so A's threads keep + // running but no new cycles fire after we replace. + getMaintenanceTask().stopAndAwait() + + // Replace provider A with a different instance while A is blocked. + // The scheduler is stopped so no maintenance runs on the + // replacement. When A finishes, loadedProviders.get(id).contains(A) + // is false (replacement is there), removal is skipped. + val replacement = new FakeStateStoreProviderTracksCloseThread + replacement.init(id.storeId, null, null, + NoPrefixKeyStateEncoderSpec(null), useColumnFamilies = false, null, null) + val loaded = getLoadedProviders() + loaded.synchronized { loaded.put(id, replacement) } + + // Release A. + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + + // Wait for A's partition sets to be released. + eventually(timeout(10.seconds)) { + assert(!getSnapshotPartitions().contains(id), + "snapshot partition should be released") + } + + // A should NOT have removed the replacement from loadedProviders. + assert(StateStore.isLoaded(id), + "replacement provider should still be loaded") + // A should NOT have queued anything (instance differs, skip). + assert(getUnloadQueue().isEmpty, + "queue should be empty (stale instance skipped removal)") + } + } + } + + test("FromLoadedProviders unload: with concurrent ops, " + + "only one removes and queues") { + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val id = loadNullProvider("concurrentUnload", storeConf) + val bp = getBlockingProvider(id) + + // Make provider stale so both ops detect inactive in source handling. + coordinatorRef.reportActiveInstance(id, "otherhost", "otherexec", Seq.empty) + + // Wait for both ops to enter. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + assert(bp.cleanupEnteredLatch.await(10, TimeUnit.SECONDS)) + + // Stop the scheduler so no new cycles drain the queue before + // we can check its size. + getMaintenanceTask().stopAndAwait() + + // Release both. Both finish, both see !verifyIfStoreInstanceActive. + // Only one should remove from loadedProviders and queue. + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + + val queue = getUnloadQueue() + eventually(timeout(10.seconds)) { + assert(!StateStore.isLoaded(id), "provider should be removed") + assert(queue.size() == 1, + "only one op should queue, the other should no-op") + } + } + } + } + + test("stale provider from loadedProviders is closed properly " + + "through the full queue routing lifecycle") { + // Load a provider, make it stale via the coordinator, then verify: + // 1. Snapshot and cleanup run on separate pools with partition sets claimed + // 2. Both pool threads hold the read lock (readLockCount == 2) + // 3. Snapshot detects inactive, queues Cleanup via otherMaintenanceOpRequest + // 4. Cleanup runs as FromUnloadedProvidersQueue and closes the provider + // 5. During close: write lock held, no read locks + // 6. After close: all locks released, provider removed, queue drained + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + val storeConf = maintenanceStoreConf(classOf[BlockingMaintenanceProvider]) + val id = loadNullProvider("decoupled", storeConf) + val bp = getBlockingProvider(id) + bp.closeShouldBlock = true + + // Make provider stale so source handling queues it + coordinatorRef.reportActiveInstance(id, "otherhost", "otherexec", Seq.empty) + + // Wait for both tasks to block + assert(bp.snapshotEnteredLatch + .await(5, TimeUnit.SECONDS), "Snapshot task did not start") + assert(bp.cleanupEnteredLatch + .await(5, TimeUnit.SECONDS), "Cleanup task did not start") + + // Snapshot and cleanup should run on separate pools + assert(bp.snapshotThreadName + .contains("state-store-maintenance-high-priority"), + s"Snapshot should run on snapshot pool, was: " + + s"${bp.snapshotThreadName}") + assert(bp.cleanupThreadName + .contains("state-store-maintenance-low-priority"), + s"Cleanup should run on cleanup pool, was: " + + s"${bp.cleanupThreadName}") + + // Both partition sets should be claimed + val snap = getSnapshotPartitions() + val clean = getCleanupPartitions() + assert(snap.contains(id), "Snapshot partition should be claimed") + assert(clean.contains(id), "Cleanup partition should be claimed") + + // Both pool threads hold the read lock. + assert(bp.maintenanceLock.getReadLockCount == 2, + "both pool threads should hold the read lock") + + val queue = getUnloadQueue() + + // Both blocked, queue should be empty, no close + assert(queue.isEmpty, "Queue should be empty while tasks are blocked") + assert(bp.closeThreadName.isEmpty, + "No close while tasks are blocked") + + // Release snapshot. It will detect inactive, remove from + // loadedProviders, and queue the other op. Cleanup stays blocked, + // holding its partition set, so the scheduler cannot process the + // queue entry. + bp.snapshotContinueSignal.countDown() + + // Wait for snapshot source handling to complete. Check inside + // eventually to handle the momentary gap between the poll + // removing the entry and offer putting the entry back. + // Verify the queue entry. Snapshot completed first, so it queued + // otherMaintenanceOpRequest(Snapshot) = Cleanup. + eventually(timeout(5.seconds)) { + val peeked = queue.peek() + assert(peeked != null, "Queue should have one entry") + val (queuedId, _, opRequest) = peeked + assert(queuedId == id) + assert(opRequest == MaintenanceOpRequest.Cleanup, + s"Expected Cleanup, got $opRequest") + } + + // No close should have happened yet + assert(bp.closeThreadName.isEmpty, + "No close before final op runs") + + // Release cleanup. FromUnloadedProvidersQueue will release the + // read lock, acquire the write lock, and call close(). + bp.cleanupContinueSignal.countDown() + + // close() blocks on the latch. Write lock held, no read locks. + assert(bp.closeEnteredLatch + .await(10, TimeUnit.SECONDS), "close should be called") + assert(bp.maintenanceLock.isWriteLocked, + "write lock should be held during close") + assert(bp.maintenanceLock.getReadLockCount == 0, + "no read locks should be held during close") + + // Release close to let the downgrade and cleanup finish. + bp.closeContinueSignal.countDown() + + // Everything released. + eventually(timeout(10.seconds)) { + assert(bp.closeThreadName.nonEmpty, + "Provider should be closed") + assert(!StateStore.isLoaded(id), + "Provider should be removed from loadedProviders") + assert(queue.isEmpty, "Queue should be drained") + assert(snap.isEmpty && clean.isEmpty, + "Both partition sets should be released") + assert(bp.maintenanceLock.getReadLockCount == 0, + "read lock should be released after close") + assert(!bp.maintenanceLock.isWriteLocked, + "write lock should be released after close") + } + } + } + } + + test("scheduler maintenance triggerNow: at-most-one pending, no-op after stop") { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], interval = 60000L) + val id = loadNullProvider("triggerNow", storeConf) + val bp = getBlockingProvider(id) + + // Reflection to extract maintenanceTask, its triggerPending flag, + // the underlying ScheduledThreadPoolExecutor, and loadedProviders. + val task = getMaintenanceTask() + val loaded = getLoadedProviders() + val pendingField = task.getClass.getDeclaredField("triggerPending") + val executorField = task.getClass.getDeclaredField("executor") + pendingField.setAccessible(true) + executorField.setAccessible(true) + val triggerPending = pendingField.get(task).asInstanceOf[AtomicBoolean] + val executor = executorField.get(task) + .asInstanceOf[java.util.concurrent.ScheduledThreadPoolExecutor] + + assert(!triggerPending.get(), "triggerPending should start as false") + + // Hold loadedProviders lock so doMaintenance blocks on Phase 2 + // (loadedProviders.synchronized). This keeps the scheduler executor + // busy, allowing us to test the at-most-one guard. + val lockHeld = new CountDownLatch(1) + val lockRelease = new CountDownLatch(1) + val lockThread = new Thread(() => { + loaded.synchronized { lockHeld.countDown(); lockRelease.await() } + }) + lockThread.start() + lockHeld.await() + + // triggerNow submits to scheduler executor. + // We wait until the triggered task starts executing, which will block on + // the loadedProviders lock. + // processUnloadedOnly=false so the triggered cycle iterates + // loadedProviders and blocks on the lock. + task.triggerNow(processUnloadedOnly = false) + eventually(timeout(5.seconds)) { + assert(!triggerPending.get(), "triggerPending reset when task started") + // Queue has 1 entry: the periodic future. The triggered task is + // currently executing (blocked on the loadedProviders lock). + assert(executor.getQueue.size() == 1, + s"expected 1 (periodic future), got ${executor.getQueue.size()}") + } + + // First call queues one pending run + task.triggerNow() + assert(triggerPending.get(), "one pending run queued") + // Queue has exactly 2: the pending triggered run + the periodic future + assert(executor.getQueue.size() == 2, + s"expected 2 queued tasks, got ${executor.getQueue.size()}") + + // Subsequent calls are no-ops (at-most-one pending). + // Queue size should not increase. + val queueSizeBefore = executor.getQueue.size() + task.triggerNow() + task.triggerNow() + assert(triggerPending.get(), "still one pending") + assert(executor.getQueue.size() == queueSizeBefore, + "queue size should not increase from no-op triggerNow calls") + + // Release the lock so the blocked cycle can proceed. After it + // finishes, the queued pending cycle runs automatically. + lockRelease.countDown() + + // Wait for the first cycle's snapshot to enter + assert(bp.snapshotEnteredLatch + .await(10, TimeUnit.SECONDS), "triggerNow should fire a cycle") + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + + // Verify the queued pending task also ran: after both cycles + // complete, the executor queue should have only the periodic future + // left, and triggerPending should be false. + eventually(timeout(10.seconds)) { + assert(!triggerPending.get(), "pending task should have run") + assert(executor.getQueue.size() == 1, + "only the periodic future should remain in the queue") + } + + // After stop, triggerNow catches RejectedExecutionException and logs. + val logAppender = new LogAppender("triggerNow-warn", maxEvents = 100) + logAppender.setThreshold(Level.WARN) + // Scope to the StateStore logger. Attaching to the root logger is unsafe + // here: concurrent suites in the same JVM mutate the root logger's level + // and appenders (each withLogAppender does setLevel/updateLoggers), which + // can drop the warning this assertion checks even though it was logged. + val loggerName = StateStore.getClass.getName.stripSuffix("$") + withLogAppender(logAppender, + loggerNames = Seq(loggerName), level = Some(Level.WARN)) { + // Use WithoutLock to avoid deadlock from stopMaintenanceTask + // holding loadedProviders lock while awaiting pool termination, but + // pool threads needing to acquire the same lock. + StateStore.stopMaintenanceTaskWithoutLock() + task.triggerNow() + assert(!triggerPending.get(), "reset after rejection") + assert(logAppender.loggingEvents.exists(_.getMessage.getFormattedMessage + .contains("triggerNow called after scheduler maintenance task stopped")), + "should log warning on rejected execution") + } + } + } + } + + test("scheduler maintenance triggerNow: only unloaded queue is processed") { + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + // Long interval so periodic cycles don't interfere. + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], interval = 60000L) + val id1 = loadNullProvider("triggerOnly", storeConf) + val bp1 = getBlockingProvider(id1) + + // Make id1 stale. Loading id2 calls reportActiveInstance which + // detects id1 as stale, queues it, and calls triggerNow. + coordinatorRef.reportActiveInstance( + id1, "otherhost", "otherexec", Seq.empty) + val id2 = loadNullProvider("triggerOnly", storeConf, partition = 1) + val bp2 = getBlockingProvider(id2) + + // triggerNow fires with processUnloadedOnly=true. + // id1 should enter (proves triggerNow processed the queue). + assert(bp1.snapshotEnteredLatch.await(10, TimeUnit.SECONDS), + "id1 should be processed from queue by triggerNow") + + // Give the triggered cycle time to finish. If phase 2 ran, + // id2 would have been submitted to the pool by now. + Thread.sleep(2000) + + // id2 should NOT have entered. The scheduler should NOT have + // iterated loadedProviders and submitted maintenance tasks. + assert(bp2.snapshotEnteredLatch.getCount == 1, + "id2 snapshot should not have been submitted by triggerNow") + assert(bp2.cleanupEnteredLatch.getCount == 1, + "id2 cleanup should not have been submitted by triggerNow") + + bp1.snapshotContinueSignal.countDown() + bp1.cleanupContinueSignal.countDown() + bp2.snapshotContinueSignal.countDown() + bp2.cleanupContinueSignal.countDown() + } + } + } + + private def poolSizeConf( + total: Int, + ratio: Double = 0.5): StateStoreConf = { + val sqlConf = getDefaultSQLConf( + SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get, + SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get) + sqlConf.setConf(SQLConf.NUM_STATE_STORE_MAINTENANCE_THREADS, total) + sqlConf.setConf( + SQLConf.STATE_STORE_MAINTENANCE_SNAPSHOT_THREAD_RATIO, ratio) + new StateStoreConf(sqlConf) + } + + test("getPoolSizes: ratio based split") { + // Default ratio 0.5: even split. + assert(StateStore.getPoolSizes(poolSizeConf(2)) === (1, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(3)) === (2, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(4)) === (2, 2)) + assert(StateStore.getPoolSizes(poolSizeConf(5)) === (3, 2)) + assert(StateStore.getPoolSizes(poolSizeConf(8)) === (4, 4)) + assert(StateStore.getPoolSizes(poolSizeConf(100)) === (50, 50)) + + // Custom ratios. + assert(StateStore.getPoolSizes(poolSizeConf(8, ratio = 0.75)) === (6, 2)) + assert(StateStore.getPoolSizes(poolSizeConf(8, ratio = 0.25)) === (2, 6)) + assert(StateStore.getPoolSizes(poolSizeConf(10, ratio = 0.8)) === (8, 2)) + assert(StateStore.getPoolSizes(poolSizeConf(10, ratio = 0.1)) === (1, 9)) + + // Fractional rounding (math.round rounds 0.5 up). + assert(StateStore.getPoolSizes(poolSizeConf(10, ratio = 1.0/3)) === (3, 7)) + assert(StateStore.getPoolSizes(poolSizeConf(11, ratio = 0.5)) === (6, 5)) + assert(StateStore.getPoolSizes(poolSizeConf(7, ratio = 0.3)) === (2, 5)) + + // Each pool gets at least 1 thread. Total is never exceeded. + assert(StateStore.getPoolSizes(poolSizeConf(2, ratio = 0.99)) === (1, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(2, ratio = 0.01)) === (1, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(4, ratio = 0.99)) === (3, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(4, ratio = 0.01)) === (1, 3)) + assert(StateStore.getPoolSizes(poolSizeConf(10, ratio = 0.99)) === (9, 1)) + assert(StateStore.getPoolSizes(poolSizeConf(10, ratio = 0.01)) === (1, 9)) + + Seq(2, 3, 6, 7, 12, 15, 20, 50).foreach { total => + Seq(0.05, 0.15, 0.33, 0.5, 0.67, 0.85, 0.95).foreach { ratio => + val (s, c) = StateStore.getPoolSizes(poolSizeConf(total, ratio)) + assert(s + c == total, s"total=$total ratio=$ratio: $s + $c != $total") + assert(s >= 1, s"total=$total ratio=$ratio: snapshot=$s < 1") + assert(c >= 1, s"total=$total ratio=$ratio: cleanup=$c < 1") + } + } + } + + /** + * Proves that a full pool does not starve the other. Provider A fills + * the specified pool. Provider B on a different partition can still + * run the other op type. + * @param blockSnapshot if true, fill snapshot pool; if false, fill cleanup pool + */ + private def testPoolIsolation(blockSnapshot: Boolean): Unit = { + withSparkContext { sc => + withCoordinatorRef(sc) { _ => + // numThreads=2: each pool gets exactly 1 thread. + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider], numThreads = 2) + val id1 = loadNullProvider("poolIsolation", storeConf) + val bp1 = getBlockingProvider(id1) + + // Wait for both of A's ops to enter. + assert(bp1.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + assert(bp1.cleanupEnteredLatch.await(10, TimeUnit.SECONDS)) + + // Release the op we DON'T want to fill, keeping the other + // pool's only thread occupied. + if (blockSnapshot) bp1.cleanupContinueSignal.countDown() + else bp1.snapshotContinueSignal.countDown() + + // Load B on a different partition. + val id2 = loadNullProvider("poolIsolation", storeConf, partition = 1) + val bp2 = getBlockingProvider(id2) + + if (blockSnapshot) { + assert(bp2.cleanupEnteredLatch.await(10, TimeUnit.SECONDS), + "B's cleanup should run despite snapshot pool being full") + assert(bp2.snapshotEnteredLatch.getCount == 1, + "B's snapshot should not have entered") + } else { + assert(bp2.snapshotEnteredLatch.await(10, TimeUnit.SECONDS), + "B's snapshot should run despite cleanup pool being full") + assert(bp2.cleanupEnteredLatch.getCount == 1, + "B's cleanup should not have entered") + } + + // Release everything. + bp1.snapshotContinueSignal.countDown() + bp1.cleanupContinueSignal.countDown() + bp2.snapshotContinueSignal.countDown() + bp2.cleanupContinueSignal.countDown() + } + } + } + + test("full snapshot pool does not prevent cleanup from running") { + testPoolIsolation(blockSnapshot = true) + } + + test("full cleanup pool does not prevent snapshot from running") { + testPoolIsolation(blockSnapshot = false) + } + + test("maintenance op skips when write lock is held during close") { + val logAppender = new LogAppender("tryLock-skip", maxEvents = 100) + logAppender.setThreshold(Level.DEBUG) + val loggerName = StateStore.getClass.getName.stripSuffix("$") + withLogAppender(logAppender, + loggerNames = Seq(loggerName), level = Some(Level.DEBUG)) { + withSparkContext { sc => + withCoordinatorRef(sc) { coordinatorRef => + val storeConf = maintenanceStoreConf( + classOf[BlockingMaintenanceProvider]) + val id = loadNullProvider("writeLockBlock", storeConf) + val bp = getBlockingProvider(id) + bp.closeShouldBlock = true + + // Make stale so the close path is triggered. + coordinatorRef.reportActiveInstance( + id, "otherhost", "otherexec", Seq.empty) + + // Wait for both ops to enter, release both. One detects + // inactive, queues remaining op. That op runs as + // FromUnloadedProvidersQueue, acquires write lock, calls close, + // blocks on closeShouldBlock. + assert(bp.snapshotEnteredLatch.await(10, TimeUnit.SECONDS)) + assert(bp.cleanupEnteredLatch.await(10, TimeUnit.SECONDS)) + bp.snapshotContinueSignal.countDown() + bp.cleanupContinueSignal.countDown() + + // Wait for close to hold the write lock. + assert(bp.closeEnteredLatch.await(10, TimeUnit.SECONDS)) + assert(bp.maintenanceLock.isWriteLocked, + "write lock should be held during close") + + // Add provider back to the unload queue. Use All because we + // don't know which op is doing close (and holding that partition + // set). All tries both and submits whichever is free. The + // scheduler's next cycle drains the queue and submits to pool. + getUnloadQueue().add((id, bp, MaintenanceOpRequest.All)) + + // The pool thread calls tryLock(0, SECONDS) which returns + // false because the write lock is held. The op is skipped + // without blocking. The debug log proves the pool thread ran. + eventually(timeout(10.seconds)) { + assert(getUnloadQueue().isEmpty, + "scheduler should have consumed the queue entry") + assert(logAppender.loggingEvents.exists( + _.getMessage.getFormattedMessage.contains( + "could not acquire read lock")), + "pool thread should have logged tryLock failure") + assert(bp.maintenanceLock.getQueueLength == 0, + "no thread should be blocked on the lock") + assert(bp.maintenanceLock.getReadLockCount == 0, + "read lock should not be acquired when write lock is held") + } + assert(bp.maintenanceLock.isWriteLocked, + "write lock should still be held during close") + + // Release close. + bp.closeContinueSignal.countDown() + + eventually(timeout(10.seconds)) { + assert(!bp.maintenanceLock.isWriteLocked, + "write lock should be released after close") + } + } + } + } + } +} + +class StateStoreDecoupledMaintenanceSuite + extends StateStoreDecoupledMaintenanceSuiteBase[HDFSBackedStateStoreProvider] + with SharedSparkSession { + override def beforeEach(): Unit = {} + override def afterEach(): Unit = {} +} + +class StateStoreDecoupledMaintenanceSuiteWithRowChecksum + extends StateStoreDecoupledMaintenanceSuite + with EnableStateStoreRowChecksum + +@ExtendedSQLTest +class RocksDBDecoupledMaintenanceSuite + extends StateStoreDecoupledMaintenanceSuiteBase[RocksDBStateStoreProvider] + with AlsoTestWithEncodingTypes + with AlsoTestWithRocksDBFeatures + with SharedSparkSession { + override def afterEach(): Unit = {} +} + +@ExtendedSQLTest +class RocksDBDecoupledMaintenanceSuiteWithRowChecksum + extends RocksDBDecoupledMaintenanceSuite + with EnableStateStoreRowChecksum diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala index 58d951500c8c5..2015b7eaed1ac 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreInstanceMetricSuite.scala @@ -30,24 +30,28 @@ import org.apache.spark.tags.ExtendedSQLTest // maintenance for partitions 0 and 1 (these are arbitrary choices). This is used to test // snapshot upload lag can be observed through StreamingQueryProgress metrics. class RocksDBSkipMaintenanceOnCertainPartitionsProvider extends RocksDBStateStoreProvider { - override def doMaintenance(): Unit = { - if (stateStoreId.partitionId == 0 || stateStoreId.partitionId == 1) { - return - } - super.doMaintenance() - } + private def shouldSkip: Boolean = + stateStoreId.partitionId == 0 || stateStoreId.partitionId == 1 + + override def doSnapshotMaintenance(): Unit = + if (!shouldSkip) super.doSnapshotMaintenance() + + override def doCleanupMaintenance(): Unit = + if (!shouldSkip) super.doCleanupMaintenance() } // HDFSBackedSkipMaintenanceOnCertainPartitionsProvider is a test-only provider that skips running // maintenance for partitions 0 and 1 (these are arbitrary choices). This is used to test // snapshot upload lag can be observed through StreamingQueryProgress metrics. class HDFSBackedSkipMaintenanceOnCertainPartitionsProvider extends HDFSBackedStateStoreProvider { - override def doMaintenance(): Unit = { - if (stateStoreId.partitionId == 0 || stateStoreId.partitionId == 1) { - return - } - super.doMaintenance() - } + private def shouldSkip: Boolean = + stateStoreId.partitionId == 0 || stateStoreId.partitionId == 1 + + override def doSnapshotMaintenance(): Unit = + if (!shouldSkip) super.doSnapshotMaintenance() + + override def doCleanupMaintenance(): Unit = + if (!shouldSkip) super.doCleanupMaintenance() } @ExtendedSQLTest diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala index b3d85c855e2f3..a34116c494398 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala @@ -21,7 +21,6 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream, File, IOException, import java.net.URI import java.util import java.util.UUID -import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean import scala.collection.mutable @@ -38,7 +37,6 @@ import org.scalatest.time.SpanSugar._ import org.apache.spark._ import org.apache.spark.LocalSparkContext._ -import org.apache.spark.internal.Logging import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection, UnsafeRow} import org.apache.spark.sql.catalyst.util.quietly @@ -54,138 +52,6 @@ import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.Utils -/** - * A test StateStoreProvider implementation that controls maintenance execution - * timing using a CountDownLatch to simulate concurrent maintenance scenarios. - * - * This provider is used to test the scenario where a task thread attempts to - * unload a provider via maintenance while it's already being processed by a - * maintenance thread. This tests the awaitProcessThisPartition functionality - * that ensures proper synchronization in StateStore's maintenance thread pool. - */ -class SignalingStateStoreProvider extends StateStoreProvider with Logging { - import SignalingStateStoreProvider._ - private var id: StateStoreId = null - - override def init( - stateStoreId: StateStoreId, - keySchema: StructType, - valueSchema: StructType, - keyStateEncoderSpec: KeyStateEncoderSpec, - useColumnFamilies: Boolean, - storeConfs: StateStoreConf, - hadoopConf: Configuration, - useMultipleValuesPerKey: Boolean = false, - stateSchemaProvider: Option[StateSchemaProvider] = None): Unit = { - id = stateStoreId - } - - override def stateStoreId: StateStoreId = id - - /** - * Records which thread called close() to verify that only maintenance threads close providers - */ - override def close(): Unit = { - closeThreadName = Thread.currentThread.getName - } - - /** - * This test implementation doesn't need to provide an actual store - */ - override def getStore( - version: Long, - uniqueId: Option[String], - forceSnapshotOnCommit: Boolean = false, - loadEmpty: Boolean = false): StateStore = null - - /** - * Simulates a maintenance operation that blocks until a signal is received. - * This allows testing the scenario where a provider is already under maintenance - * when a task thread tries to trigger another maintenance operation on it. - */ - override def doMaintenance(): Unit = { - maintenanceStarted = true - logInfo(s"Maintenance started on thread: ${Thread.currentThread().getName}") - - // Block until the test signals to continue - continueSignal.await() - - logInfo(s"Maintenance continuing after signal on thread: ${Thread.currentThread().getName}") - } -} - -/** - * Companion object that tracks state and provides synchronization primitives - * for testing concurrent maintenance scenarios - */ -object SignalingStateStoreProvider extends Logging { - // For tracking state across threads - var maintenanceStarted: Boolean = false - var taskSubmittedMaintenance: Boolean = false - var closeThreadName: String = "" - - // Added for queue testing - var providerWasQueued: Boolean = false - - // For coordination between threads - var continueSignal = new CountDownLatch(1) - val maintenanceStartedLatch = new CountDownLatch(1) - val taskAttemptCompletedLatch = new CountDownLatch(1) - - /** - * Resets all test state between test runs - */ - def reset(): Unit = { - maintenanceStarted = false - taskSubmittedMaintenance = false - closeThreadName = "" - - // Reset the latch to ensure maintenance will block again - try { - continueSignal = new CountDownLatch(1) - } catch { - case e: Exception => - logError(s"Error resetting latch: ${e.getMessage}") - } - } -} - -class FakeStateStoreProviderTracksCloseThread extends StateStoreProvider { - import FakeStateStoreProviderTracksCloseThread._ - private var id: StateStoreId = null - - override def init( - stateStoreId: StateStoreId, - keySchema: StructType, - valueSchema: StructType, - keyStateEncoderSpec: KeyStateEncoderSpec, - useColumnFamilies: Boolean, - storeConfs: StateStoreConf, - hadoopConf: Configuration, - useMultipleValuesPerKey: Boolean = false, - stateSchemaProvider: Option[StateSchemaProvider] = None): Unit = { - id = stateStoreId - } - - override def stateStoreId: StateStoreId = id - - override def close(): Unit = { - closeThreadNames = Thread.currentThread.getName :: closeThreadNames - } - - override def getStore( - version: Long, - uniqueId: Option[String], - forceSnapshotOnCommit: Boolean = false, - loadEmpty: Boolean = false): StateStore = null - - override def doMaintenance(): Unit = {} -} - -private object FakeStateStoreProviderTracksCloseThread { - var closeThreadNames: List[String] = Nil -} - // MaintenanceErrorOnCertainPartitionsProvider is a test-only provider that throws an // exception during maintenance for partitions 0 and 1 (these are arbitrary choices). It is // used to test that an exception in a single provider's maintenance does not affect other @@ -211,11 +77,20 @@ class MaintenanceErrorOnCertainPartitionsProvider extends HDFSBackedStateStorePr storeConfs, hadoopConf, useMultipleValuesPerKey) } - override def doMaintenance(): Unit = { + private def maybeThrow(): Unit = { if (id.partitionId == 0 || id.partitionId == 1) { throw new RuntimeException("Intentional maintenance failure") } - super.doMaintenance() + } + + override def doSnapshotMaintenance(): Unit = { + maybeThrow() + super.doSnapshotMaintenance() + } + + override def doCleanupMaintenance(): Unit = { + maybeThrow() + super.doCleanupMaintenance() } } @@ -270,17 +145,24 @@ private object FakeStateStoreProviderWithMaintenanceError { class MaintenanceCountingStateStoreProvider extends HDFSBackedStateStoreProvider { import MaintenanceCountingStateStoreProvider._ - override def doMaintenance(): Unit = { - maintenanceCallCount.incrementAndGet() - super.doMaintenance() + override def doSnapshotMaintenance(): Unit = { + snapshotMaintenanceCallCount.incrementAndGet() + super.doSnapshotMaintenance() + } + + override def doCleanupMaintenance(): Unit = { + cleanupMaintenanceCallCount.incrementAndGet() + super.doCleanupMaintenance() } } private object MaintenanceCountingStateStoreProvider { - val maintenanceCallCount = new java.util.concurrent.atomic.AtomicInteger(0) + val snapshotMaintenanceCallCount = new java.util.concurrent.atomic.AtomicInteger(0) + val cleanupMaintenanceCallCount = new java.util.concurrent.atomic.AtomicInteger(0) def reset(): Unit = { - maintenanceCallCount.set(0) + snapshotMaintenanceCallCount.set(0) + cleanupMaintenanceCallCount.set(0) } } @@ -302,262 +184,6 @@ class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider] require(!StateStore.isMaintenanceRunning) } - test("SPARK-51596: submitMaintenanceWorkForProvider from task thread adds" + - " to queue when timeout occurs") { - // Reset tracking variables for a clean test - SignalingStateStoreProvider.reset() - - val sqlConf = getDefaultSQLConf( - SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get, - SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get - ) - - // Critical: Set a very short timeout to ensure awaitProcessThisPartition fails quickly - sqlConf.setConf(SQLConf.STATE_STORE_MAINTENANCE_PROCESSING_TIMEOUT, 1L) // 1 second - - // Maintenance interval large enough that we control timing manually - sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, 30000L) - sqlConf.setConf(SQLConf.NUM_STATE_STORE_MAINTENANCE_THREADS, 4) - - // Use our test provider - sqlConf.setConf( - SQLConf.STATE_STORE_PROVIDER_CLASS, - classOf[SignalingStateStoreProvider].getName - ) - - val conf = new SparkConf().setMaster("local").setAppName("test") - - withSpark(SparkContext.getOrCreate(conf)) { sc => - withCoordinatorRef(sc) { _ => - val rootLocation = s"${Utils.createTempDir().getAbsolutePath}/spark-51596-timeout-queue" - val providerId = StateStoreProviderId(StateStoreId(rootLocation, 0, 0), UUID.randomUUID) - - // Load the provider to start the maintenance system - StateStore.get( - providerId, - keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema), - 0, None, None, useColumnFamilies = false, - new StateStoreConf(sqlConf), new Configuration() - ) - - // Access the queue via reflection for verification - val queueField = PrivateMethod[ConcurrentLinkedQueue[ - (StateStoreProviderId, StateStoreProvider)]]( - Symbol("unloadedProvidersToClose")) - val queue = StateStore invokePrivate queueField() - assert(queue.isEmpty, "Queue should start empty") - - // Manually trigger maintenance which will block - val maintenanceMethod = PrivateMethod[Unit](Symbol("doMaintenance")) - StateStore invokePrivate maintenanceMethod() - - // Wait for maintenance to start - eventually(timeout(5.seconds)) { - assert(SignalingStateStoreProvider.maintenanceStarted) - assert(StateStore.isLoaded(providerId)) - } - - // Now get access to the provider to simulate a task thread - val loadedProvidersField = PrivateMethod[ - mutable.HashMap[StateStoreProviderId, StateStoreProvider]]( - Symbol("loadedProviders")) - val loadedProviders = StateStore invokePrivate loadedProvidersField() - val provider = loadedProviders.synchronized { loadedProviders.get(providerId).get } - val maintenancePartitionsField = PrivateMethod[ - mutable.HashSet[StateStoreProviderId]]( - Symbol("maintenancePartitions")) - val maintenancePartitions = StateStore invokePrivate maintenancePartitionsField() - - // Create a task thread that will attempt to submit maintenance - val taskThread = new Thread(() => { - try { - // Call submitMaintenanceWorkForProvider directly since that's what we're testing - val submitMaintenanceMethod = PrivateMethod[Unit]( - Symbol("submitMaintenanceWorkForProvider")) - StateStore invokePrivate submitMaintenanceMethod( - providerId, provider, new StateStoreConf(sqlConf), - MaintenanceTaskType.FromTaskThread) - - SignalingStateStoreProvider.taskSubmittedMaintenance = true - SignalingStateStoreProvider.taskAttemptCompletedLatch.countDown() - } catch { - case e: Exception => - logError(s"Error in task thread: ${e.getMessage}", e) - } - }) - - // Start the task thread - it should timeout and add provider to queue - taskThread.start() - - // Wait for task attempt to complete - assert(SignalingStateStoreProvider - .taskAttemptCompletedLatch.await(10, TimeUnit.SECONDS), - "Task thread didn't complete") - - // Critical verification: After timeout, the provider should be in the queue - eventually(timeout(5.seconds)) { - assert(queue.size() == 1, "Provider should be queued after timeout") - } - val (queuedId, _) = queue.peek() - assert(queuedId == providerId, "Queued provider has wrong ID") - - // Now allow the first maintenance to complete - SignalingStateStoreProvider.continueSignal.countDown() - - eventually(timeout(5.seconds)) { - assert(maintenancePartitions.isEmpty, - "Maintenance partitions should be removed from") - } - // Manually trigger another maintenance to process the queue - StateStore invokePrivate maintenanceMethod() - - // Verify the queue eventually gets processed - eventually(timeout(5.seconds)) { - assert(queue.isEmpty, "Queue should be emptied after maintenance") - } - } - } - } - - test("SPARK-51596: queued maintenance tasks get processed when lock is available") { - // Reset tracking variables for a clean test - SignalingStateStoreProvider.reset() - - val sqlConf = getDefaultSQLConf( - SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get, - SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get - ) - // Use a maintenance interval large enough that we control timing explicitly - sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, 30000L) - // Set our special provider class that lets us control maintenance timing - sqlConf.setConf( - SQLConf.STATE_STORE_PROVIDER_CLASS, - classOf[SignalingStateStoreProvider].getName - ) - - val conf = new SparkConf().setMaster("local").setAppName("test") - - withSpark(SparkContext.getOrCreate(conf)) { sc => - withCoordinatorRef(sc) { coordinatorRef => - val rootLocation = s"${Utils.createTempDir().getAbsolutePath}/spark-51596-queue" - - // Create two providers that we'll use for the test - val provider1Id = - StateStoreProviderId(StateStoreId(rootLocation, 0, 0), UUID.randomUUID) - val provider2Id = - StateStoreProviderId(StateStoreId(rootLocation, 0, 1), UUID.randomUUID) - - // Get the first provider to load it - StateStore.get( - provider1Id, - keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema), - 0, None, None, useColumnFamilies = false, - new StateStoreConf(sqlConf), new Configuration() - ) - - // Manually trigger maintenance for provider1, which will block in doMaintenance() - val maintenanceMethod = PrivateMethod[Unit](Symbol("doMaintenance")) - StateStore invokePrivate maintenanceMethod() - - // Wait for maintenance to start before continuing - eventually(timeout(5.seconds)) { - assert(SignalingStateStoreProvider.maintenanceStarted) - assert(StateStore.isLoaded(provider1Id)) - } - - // Now make the first provider "stale" by reporting it active on another executor - coordinatorRef.reportActiveInstance(provider1Id, "otherhost", "otherexec", Seq.empty) - - // Get provider2 which will cause a maintenance task for provider1 to be queued - // (since provider1 is already under maintenance and can't be processed immediately) - StateStore.get( - provider2Id, - keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema), - 0, None, None, useColumnFamilies = false, - new StateStoreConf(sqlConf), new Configuration() - ) - - // Mark that task submitted maintenance - SignalingStateStoreProvider.taskSubmittedMaintenance = true - - // Unblock the first maintenance operation - SignalingStateStoreProvider.continueSignal.countDown() - - // Verify that provider1 is eventually unloaded by the maintenance thread - // after the first maintenance completes and the queued maintenance runs - eventually(timeout(5.seconds)) { - // Provider1 should be unloaded - assert(!StateStore.isLoaded(provider1Id)) - // Provider2 should still be loaded - assert(StateStore.isLoaded(provider2Id)) - // Close should have been called on a maintenance thread - assert(SignalingStateStoreProvider.closeThreadName.contains("maintenance")) - } - - // Get the partitionsForMaintenance field to check the queue is empty - val partitionsField = PrivateMethod[ - ConcurrentLinkedQueue[StateStoreProviderId]](Symbol("unloadedProvidersToClose")) - val queue = StateStore invokePrivate partitionsField() - assert(queue.isEmpty, "Maintenance queue should be empty after processing queued tasks") - } - } - } - - test("SPARK-51596: unloading only occurs on maintenance thread but occurs promptly") { - // Reset closeThreadNames - FakeStateStoreProviderTracksCloseThread.closeThreadNames = Nil - - val sqlConf = getDefaultSQLConf( - SQLConf.STATE_STORE_MIN_DELTAS_FOR_SNAPSHOT.defaultValue.get, - SQLConf.MAX_BATCHES_TO_RETAIN_IN_MEMORY.defaultValue.get - ) - // Make maintenance interval very large (30s) so that task thread runs before maintenance. - sqlConf.setConf(SQLConf.STREAMING_MAINTENANCE_INTERVAL, 30000L) - // Use the `FakeStateStoreProviderTracksCloseThread` to run the test - sqlConf.setConf( - SQLConf.STATE_STORE_PROVIDER_CLASS, - classOf[FakeStateStoreProviderTracksCloseThread].getName - ) - - val conf = new SparkConf().setMaster("local").setAppName("test") - - withSpark(SparkContext.getOrCreate(conf)) { sc => - withCoordinatorRef(sc) { coordinatorRef => - val rootLocation = s"${Utils.createTempDir().getAbsolutePath}/spark-51596" - val providerId = - StateStoreProviderId(StateStoreId(rootLocation, 0, 0), UUID.randomUUID) - val providerId2 = - StateStoreProviderId(StateStoreId(rootLocation, 0, 1), UUID.randomUUID) - - // Create provider to start the maintenance task + pool - StateStore.get( - providerId, - keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema), - 0, None, None, useColumnFamilies = false, new StateStoreConf(sqlConf), new Configuration() - ) - - // Report instance active on another executor - coordinatorRef.reportActiveInstance(providerId, "otherhost", "otherexec", Seq.empty) - - // Load another provider to trigger task unload - StateStore.get( - providerId2, - keySchema, valueSchema, NoPrefixKeyStateEncoderSpec(keySchema), - 0, None, None, useColumnFamilies = false, new StateStoreConf(sqlConf), new Configuration() - ) - - // Wait for close to occur. Timeout is less than maintenance interval, - // so should only close by task triggering. - eventually(timeout(5.seconds)) { - assert(FakeStateStoreProviderTracksCloseThread.closeThreadNames.size == 1) - FakeStateStoreProviderTracksCloseThread.closeThreadNames.foreach { name => - assert(name.contains("state-store-maintenance-thread"))} - } - } - } - } - - test("retaining only two latest versions when MAX_BATCHES_TO_RETAIN_IN_MEMORY set to 2") { tryWithProviderResource( newStoreProvider(minDeltasForSnapshot = 10, numOfVersToRetainInMemory = 2)) { provider => @@ -620,6 +246,31 @@ class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider] } } + test("HDFS: split maintenance methods upload snapshots and clean up old files separately") { + tryWithProviderResource(newStoreProvider(opId = Random.nextInt(), partition = 0, + minDeltasForSnapshot = 5)) { provider => + for (i <- 1 to 21) { + val store = provider.getStore(i - 1) + put(store, "a", 0, i) + store.commit() + // Snapshot and cleanup run as independent operations. + provider.doSnapshotMaintenance() + provider.doCleanupMaintenance() + } + + // Snapshots are uploaded by doSnapshotMaintenance (at versions 6, 12, 18 given + // minDeltasForSnapshot = 5) and doCleanupMaintenance removes old files, retaining only + // the last numVersionsToRetain (default 2) versions anchored on the latest snapshot (18). + val basePath = provider.stateStoreId.storeCheckpointLocation() + val remainingFiles = new File(basePath.toString) + .listFiles().filter(f => f.isFile && !f.getName.startsWith(".")) + .map(_.getName).filterNot(_.endsWith(".crc")).toSet + assert(remainingFiles === + Set("18.snapshot", "18.delta", "19.delta", "20.delta", "21.delta"), + s"Unexpected remaining files: $remainingFiles") + } + } + test("get, put, remove etc operations on non-default col family should fail") { tryWithProviderResource(newStoreProvider(opId = Random.nextInt(), partition = 0, minDeltasForSnapshot = 5)) { provider => @@ -1081,27 +732,30 @@ class StateStoreSuite extends StateStoreSuiteBase[HDFSBackedStateStoreProvider] assert(StateStore.isLoaded(storeProviderId1), "Store is not loaded") } - // Record the current maintenance call count before deactivation - val maintenanceCountBeforeDeactivate = - MaintenanceCountingStateStoreProvider.maintenanceCallCount.get() + // Record the current maintenance call counts before deactivation + val snapshotCountBefore = + MaintenanceCountingStateStoreProvider.snapshotMaintenanceCallCount.get() + val cleanupCountBefore = + MaintenanceCountingStateStoreProvider.cleanupMaintenanceCallCount.get() - // Deactivate the store instance - this should trigger maintenance before unload + // Deactivate the store instance - this should trigger maintenance before close. In the + // decoupled design, the provider is removed from loadedProviders before close completes + // (removal and close are separate events), so we wait for all conditions together. coordinatorRef.deactivateInstances(storeProviderId1.queryRunId) - // Wait for the store to be unloaded eventually(timeout(timeoutDuration)) { assert(!StateStore.isLoaded(storeProviderId1), "Store was not unloaded") + val snapshotCountAfter = + MaintenanceCountingStateStoreProvider.snapshotMaintenanceCallCount.get() + val cleanupCountAfter = + MaintenanceCountingStateStoreProvider.cleanupMaintenanceCallCount.get() + assert(snapshotCountAfter > snapshotCountBefore, + s"Snapshot maintenance should run before close. " + + s"Before: $snapshotCountBefore, After: $snapshotCountAfter") + assert(cleanupCountAfter > cleanupCountBefore, + s"Cleanup maintenance should run before close. " + + s"Before: $cleanupCountBefore, After: $cleanupCountAfter") } - - // Get the maintenance count after unload - val maintenanceCountAfterUnload = - MaintenanceCountingStateStoreProvider.maintenanceCallCount.get() - - // Ensure that maintenance was called at least one more time during unload - assert(maintenanceCountAfterUnload > maintenanceCountBeforeDeactivate, - s"Maintenance should be called before unload. " + - s"Before: $maintenanceCountBeforeDeactivate, " + - s"After: $maintenanceCountAfterUnload") } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala index 8475f283b6628..86e39c668df83 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/TimerSuite.scala @@ -60,6 +60,33 @@ class TimerSuite extends StateVariableSuiteBase { } } + testWithTimeMode("reusable expired timer iterator resumes from prior threshold") { timeMode => + tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => + val store = provider.getStore(0) + assert(store.isInstanceOf[SupportsReusableIterator]) + + ImplicitGroupingKeyTracker.setImplicitKey("test_key") + val timerState = new TimerStateImpl(store, timeMode, stringEncoder) + timerState.registerTimer(1000L) + timerState.registerTimer(3000L) + + assert(timerState.getExpiredTimersReusable(1500L).toSeq === + Seq(("test_key", 1000L))) + + timerState.deleteTimer(1000L) + // The refreshed iterator resumes at the prior 1500 ms threshold, so it does not revisit + // the backdated timer at 500 ms. + timerState.registerTimer(500L) + timerState.registerTimer(1500L) + timerState.registerTimer(2000L) + + assert(timerState.getExpiredTimersReusable(2500L).toSeq === + Seq(("test_key", 1500L), ("test_key", 2000L))) + assert(timerState.getExpiredTimers(2500L).toSeq === + Seq(("test_key", 500L), ("test_key", 1500L), ("test_key", 2000L))) + } + } + testWithTimeMode("multiple instances with single key") { timeMode => tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala index 874c69dee1d61..1df11021a18ee 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/ValueStateSuite.scala @@ -371,6 +371,43 @@ class ValueStateSuite extends StateVariableSuiteBase { } } + test("Value state TTL uses the current processing-time function") { + tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => + val store = provider.getStore(0) + var currentTimestampMs = 1000L + val handle = new StatefulProcessorHandleImpl( + store, + UUID.randomUUID(), + stringEncoder, + TimeMode.ProcessingTime(), + batchTimestampMs = Some(currentTimestampMs), + currentTimestampMs = Some(() => currentTimestampMs)) + val state = handle.getValueState[String]( + "testState", Encoders.STRING, TTLConfig(Duration.ofMillis(100))) + .asInstanceOf[ValueStateImplWithTTL[String]] + + ImplicitGroupingKeyTracker.setImplicitKey("test_key") + try { + state.update("v1") + assert(state.getTTLValue().contains(("v1", 1100L))) + + currentTimestampMs = 1099L + assert(state.get() === "v1") + currentTimestampMs = 1100L + assert(state.get() === null) + + state.update("v2") + assert(state.getTTLValue().contains(("v2", 1200L))) + currentTimestampMs = 1199L + assert(state.get() === "v2") + currentTimestampMs = 1200L + assert(state.get() === null) + } finally { + ImplicitGroupingKeyTracker.removeImplicitKey() + } + } + } + // Guards against the UnsafeRow byte-order bug where a scan boundary row with a // null element-key struct encodes larger than a real entry (null-bitmap bit = 1), // making seek() silently skip boundary entries. Uses a primitive Long grouping @@ -378,7 +415,8 @@ class ValueStateSuite extends StateVariableSuiteBase { // happen to mask the bug via size-based byte differences. test("SPARK-56400: TTL eviction iterator - boundary at prev+1 with fixed-size element key") { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => - val store = provider.getStore(0) + val store = CkptIdCollectingStateStoreWrapper(provider.getStore(0)) + assert(!store.isInstanceOf[SupportsReusableIterator]) val longKeyEncoder = encoderFor(Encoders.scalaLong).asInstanceOf[ExpressionEncoder[Any]] // 1 ms TTL so expiration = batchTimestampMs + 1, hitting the prev+1 boundary @@ -411,12 +449,53 @@ class ValueStateSuite extends StateVariableSuiteBase { Seq(firstBatchTs + 1, firstBatchTs + 1, firstBatchTs + 1)) // The eviction iterator (bounded range scan) should find all three entries. - val evicted = state2.ttlEvictionIterator().toList + val evicted = state2.ttlEvictionIterator(nextBatchTs).toList assert(evicted.size === 3, s"Expected 3 evictable TTL entries at expiration = prevBatch + 1, got ${evicted.size}") } } + test("TTL eviction reusable iterator refreshes and resumes from prior threshold") { + tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => + val store = provider.getStore(0) + assert(store.isInstanceOf[SupportsReusableIterator]) + val longKeyEncoder = encoderFor(Encoders.scalaLong).asInstanceOf[ExpressionEncoder[Any]] + val ttlConfig = TTLConfig(Duration.ofMillis(500)) + var currentTimeMs = 250L + val handle = new StatefulProcessorHandleImpl( + store, + UUID.randomUUID(), + longKeyEncoder, + TimeMode.ProcessingTime(), + batchTimestampMs = Some(5000L), + prevBatchTimestampMs = Some(5000L), + currentTimestampMs = Some(() => currentTimeMs)) + val state = handle.getValueState[Long]( + "testState", Encoders.scalaLong, ttlConfig).asInstanceOf[ValueStateImplWithTTL[Long]] + + def update(key: Long): Unit = { + ImplicitGroupingKeyTracker.setImplicitKey(key) + try { + state.update(key) + } finally { + ImplicitGroupingKeyTracker.removeImplicitKey() + } + } + + // The first scan starts at the column-family beginning and stops on the future entry. + update(1L) // expires at 750 + assert(state.clearExpiredStateForAllKeys(500L) === 0L) + + // Refreshing makes this new entry visible and seeking to the prior threshold revisits the + // entry consumed by the first scan. Seeking to the new threshold would skip it. + currentTimeMs = 500L + update(2L) // expires at 1000 + assert(state.clearExpiredStateForAllKeys(900L) === 1L) + assert(state.clearExpiredStateForAllKeys(1000L) === 1L) + assert(state.getTTLRows().isEmpty) + } + } + test("test null or zero TTL duration throws error") { tryWithProviderResource(newStoreProviderWithStateVariable(true)) { provider => val store = provider.getStore(0) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SparkPlanInfoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SparkPlanInfoSuite.scala index 1ef07bf9ebc15..2fc070e8f763f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SparkPlanInfoSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ui/SparkPlanInfoSuite.scala @@ -17,6 +17,9 @@ package org.apache.spark.sql.execution.ui +import scala.concurrent.duration._ + +import org.apache.spark.sql.DataFrame import org.apache.spark.sql.execution.SparkPlanInfo import org.apache.spark.sql.test.SharedSparkSession @@ -24,6 +27,39 @@ class SparkPlanInfoSuite extends SharedSparkSession { import testImplicits._ + private def collectSparkPlanInfo(sparkPlanInfo: SparkPlanInfo): Seq[SparkPlanInfo] = { + sparkPlanInfo +: sparkPlanInfo.children.flatMap(collectSparkPlanInfo) + } + + private def findSparkPlanInfo(sparkPlanInfo: SparkPlanInfo, nodeName: String): SparkPlanInfo = { + collectSparkPlanInfo(sparkPlanInfo) + .find(_.nodeName == nodeName) + .getOrElse(fail(s"Could not find $nodeName in ${sparkPlanInfo.simpleString}")) + } + + private def collectSparkPlanGraphMetrics( + df: DataFrame): (SparkPlanGraph, Map[Long, String]) = { + val statusStore = spark.sharedState.statusStore + spark.sparkContext.listenerBus.waitUntilEmpty(10000) + val previousExecutionIds = statusStore.executionsList().map(_.executionId).toSet + + df.collect() + spark.sparkContext.listenerBus.waitUntilEmpty(10000) + + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(statusStore.executionsList().map(_.executionId).toSet + .diff(previousExecutionIds).size === 1) + } + val executionIds = statusStore.executionsList().map(_.executionId).toSet + .diff(previousExecutionIds) + val executionId = executionIds.head + eventually(timeout(10.seconds), interval(10.milliseconds)) { + assert(statusStore.execution(executionId).exists(_.metricValues != null)) + } + + (statusStore.planGraph(executionId), statusStore.executionMetrics(executionId)) + } + def validateSparkPlanInfo(sparkPlanInfo: SparkPlanInfo): Unit = { sparkPlanInfo.nodeName match { case "InMemoryTableScan" => assert(sparkPlanInfo.children.length == 1) @@ -41,4 +77,48 @@ class SparkPlanInfoSuite extends SharedSparkSession { validateSparkPlanInfo(planInfoResult) } + + test("SPARK-47017: SparkPlanInfo and SQL UI include SQL plan inside RDDScanExec") { + val source = spark.range(10).where($"id" > 3).select($"id".as("age")) + val recreated = spark.createDataFrame(source.rdd, source.schema) + + val planInfo = SparkPlanInfo.fromSparkPlan(recreated.queryExecution.executedPlan) + val rddScanInfo = findSparkPlanInfo(planInfo, "Scan ExistingRDD") + val internalRDDPlanInfos = rddScanInfo.children.flatMap(collectSparkPlanInfo) + val filterInfo = internalRDDPlanInfos + .find(_.nodeName == "Filter") + .getOrElse(fail(s"Could not find Filter under Scan ExistingRDD in ${planInfo.simpleString}")) + + assert(rddScanInfo.children.nonEmpty) + assert(filterInfo.metrics.exists(_.name == "number of output rows")) + + val unionRDD = spark.sparkContext.union(source.rdd, source.rdd) + val unionPlanInfo = SparkPlanInfo.fromSparkPlan( + spark.createDataFrame(unionRDD, source.schema).queryExecution.executedPlan) + val unionRDDScanInfo = findSparkPlanInfo(unionPlanInfo, "Scan ExistingRDD") + + assert(unionRDDScanInfo.children.size === 1) + + val nested = spark.createDataFrame(recreated.rdd, recreated.schema) + val nestedPlanInfo = SparkPlanInfo.fromSparkPlan(nested.queryExecution.executedPlan) + val nestedRDDScanInfo = findSparkPlanInfo(nestedPlanInfo, "Scan ExistingRDD") + + assert(nestedRDDScanInfo.children.size === 1) + assert(collectSparkPlanInfo(nestedRDDScanInfo.children.head).exists(_.nodeName == "Filter")) + + val (planGraph, metricValues) = collectSparkPlanGraphMetrics(recreated) + val filterNode = planGraph.allNodes + .find(_.name == "Filter") + .getOrElse(fail(s"Could not find Filter in ${recreated.queryExecution.executedPlan}")) + val filterMetric = filterNode.metrics + .find(_.name == "number of output rows") + .getOrElse(fail("Could not find number of output rows metric for Filter")) + val filterMetricValue = metricValues + .getOrElse(filterMetric.accumulatorId, fail("Could not find Filter metric value")) + val outputRows = "\\d+".r.findFirstIn(filterMetricValue.replace(",", "")) + .map(_.toLong) + .getOrElse(fail(s"Could not parse Filter metric value $filterMetricValue")) + + assert(outputRows === 6L) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala index 6f0e39bb9c660..a29a4b2b2f26e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/vectorized/ColumnarBatchSuite.scala @@ -2103,6 +2103,95 @@ class ColumnarBatchSuite extends SparkFunSuite { } } + // Nanosecond-precision timestamp values used to exercise ColumnarRow/ColumnarBatchRow + // dispatch. They cover a null (index 3), the sub-microsecond boundaries (0 and 999), and + // several pre-epoch instants; nanosWithinMicro must survive the round-trip unchanged. + private val nanosTestValues: Seq[Option[TimestampNanosVal]] = Seq( + Some(TimestampNanosVal.fromParts(0L, 0.toShort)), // epoch, no sub-micro + Some(TimestampNanosVal.fromParts(1L, 1.toShort)), + Some(TimestampNanosVal.fromParts(1000000L, 999.toShort)), // max sub-micro + None, // null slot + Some(TimestampNanosVal.fromParts(-1L, 500.toShort)), // pre-epoch + Some(TimestampNanosVal.fromParts(-1000000000000L, 1.toShort)), // pre-epoch, far + Some(TimestampNanosVal.fromParts(1234567890123456L, 789.toShort)), + Some(TimestampNanosVal.fromParts(-42L, 999.toShort)), // pre-epoch, max sub-micro + Some(TimestampNanosVal.fromParts(Long.MaxValue, 0.toShort)), + Some(TimestampNanosVal.fromParts(Long.MinValue, 123.toShort))) + + // Populates a nanos-typed vector (child 0 = epochMicros Long, child 1 = nanosWithinMicro + // Short) from `nanosTestValues`, writing a null for the None slot. + private def putNanosTestValues(column: WritableColumnVector, isLtz: Boolean): Unit = { + nanosTestValues.zipWithIndex.foreach { + case (Some(v), i) => + if (isLtz) column.putTimestampLTZNanos(i, v) else column.putTimestampNTZNanos(i, v) + case (None, i) => + column.putNull(i) + } + } + + // Verifies the typed leaf accessor, get(ordinal, dataType) and the copy() round-trip for a + // single nanosecond-timestamp field at ordinal 0 of `row`. + private def assertNanosRow( + row: InternalRow, + dt: DataType, + isLtz: Boolean, + expected: Option[TimestampNanosVal]): Unit = { + def leaf(r: InternalRow): TimestampNanosVal = + if (isLtz) r.getTimestampLTZNanos(0) else r.getTimestampNTZNanos(0) + expected match { + case Some(v) => + assert(!row.isNullAt(0)) + assert(leaf(row) === v) + assert(leaf(row).nanosWithinMicro == v.nanosWithinMicro) + // get(ordinal, dataType) must return a TimestampNanosVal, not a boxed Long. + val got = row.get(0, dt) + assert(got.isInstanceOf[TimestampNanosVal]) + assert(got === v) + // copy() must yield a GenericInternalRow whose nanos field matches the source. + val copied = row.copy() + assert(copied.isInstanceOf[GenericInternalRow]) + assert(!copied.isNullAt(0)) + assert(leaf(copied) === v) + assert(copied.get(0, dt) === v) + case None => + assert(row.isNullAt(0)) + assert(row.get(0, dt) == null) + assert(row.copy().isNullAt(0)) + } + } + + Seq(7, 8, 9).foreach { p => + Seq( + (TimestampNTZNanosType(p): DataType, false), + (TimestampLTZNanosType(p): DataType, true)).foreach { case (dt, isLtz) => + val family = if (isLtz) "LTZ" else "NTZ" + + testVector( + s"ColumnarBatchRow nanos $family(precision=$p) round-trips copy()/get()", + nanosTestValues.length, + dt) { column => + putNanosTestValues(column, isLtz) + val batchRow = new ColumnarBatchRow(Array(column)) + nanosTestValues.zipWithIndex.foreach { case (expected, i) => + batchRow.rowId = i + assertNanosRow(batchRow, dt, isLtz, expected) + } + } + + // ColumnarRow is produced by getStruct on a struct-typed vector, so nest the nanos + // column inside a struct to exercise ColumnarRow's copy()/get() dispatch. + testVector( + s"ColumnarRow nanos $family(precision=$p) round-trips copy()/get()", + nanosTestValues.length, + new StructType().add("ts", dt)) { column => + putNanosTestValues(column.getChild(0), isLtz) + nanosTestValues.zipWithIndex.foreach { case (expected, i) => + assertNanosRow(column.getStruct(i), dt, isLtz, expected) + } + } + } + } + testVector("WritableColumnVector.reserve(): requested capacity is negative", 1024, ByteType) { column => val ex = intercept[RuntimeException] { column.reserve(-1) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala index 6136ddb0fd536..7b5402119b409 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/expressions/ExpressionInfoSuite.scala @@ -68,8 +68,10 @@ class ExpressionInfoSuite extends SharedSparkSession { new ExpressionInfo( "testClass", null, "testName", null, "", "", "", invalidGroupName, "", "", "") }, - condition = "_LEGACY_ERROR_TEMP_3202", + condition = "MALFORMED_EXPRESSION_INFO.GROUP", + sqlState = Some("22023"), parameters = Map( + "fieldName" -> "group", "exprName" -> "testName", "group" -> invalidGroupName, "validGroups" -> validGroups.mkString("[", ", ", "]"))) @@ -93,8 +95,10 @@ class ExpressionInfoSuite extends SharedSparkSession { new ExpressionInfo( "testClass", null, "testName", null, "", "", "", "", "", "", invalidSource) }, - condition = "_LEGACY_ERROR_TEMP_3203", + condition = "MALFORMED_EXPRESSION_INFO.SOURCE", + sqlState = Some("22023"), parameters = Map( + "fieldName" -> "source", "exprName" -> "testName", "source" -> invalidSource, "validSources" -> validSources.sorted.mkString("[", ", ", "]"))) @@ -106,8 +110,9 @@ class ExpressionInfoSuite extends SharedSparkSession { exception = intercept[SparkIllegalArgumentException] { new ExpressionInfo("testClass", null, "testName", null, "", "", invalidNote, "", "", "", "") }, - condition = "_LEGACY_ERROR_TEMP_3201", - parameters = Map("exprName" -> "testName", "note" -> invalidNote)) + condition = "MALFORMED_EXPRESSION_INFO.NOTE", + sqlState = Some("22023"), + parameters = Map("fieldName" -> "note", "exprName" -> "testName", "note" -> invalidNote)) val invalidSince = "-3.0.0" checkError( @@ -115,8 +120,9 @@ class ExpressionInfoSuite extends SharedSparkSession { new ExpressionInfo( "testClass", null, "testName", null, "", "", "", "", invalidSince, "", "") }, - condition = "_LEGACY_ERROR_TEMP_3204", - parameters = Map("since" -> invalidSince, "exprName" -> "testName")) + condition = "MALFORMED_EXPRESSION_INFO.SINCE", + sqlState = Some("22023"), + parameters = Map("fieldName" -> "since", "since" -> invalidSince, "exprName" -> "testName")) val invalidDeprecated = " invalid deprecated" checkError( @@ -124,8 +130,12 @@ class ExpressionInfoSuite extends SharedSparkSession { new ExpressionInfo( "testClass", null, "testName", null, "", "", "", "", "", invalidDeprecated, "") }, - condition = "_LEGACY_ERROR_TEMP_3205", - parameters = Map("exprName" -> "testName", "deprecated" -> invalidDeprecated)) + condition = "MALFORMED_EXPRESSION_INFO.DEPRECATED", + sqlState = Some("22023"), + parameters = Map( + "fieldName" -> "deprecated", + "exprName" -> "testName", + "deprecated" -> invalidDeprecated)) } test("using _FUNC_ instead of function names in examples") { @@ -164,11 +174,39 @@ class ExpressionInfoSuite extends SharedSparkSession { } test("SPARK-32870: Default expressions in FunctionRegistry should have their " + - "usage, examples, since, and group filled") { + "usage, examples, arguments, since, and group filled") { val ignoreSet = Set( // Cast aliases do not need examples "org.apache.spark.sql.catalyst.expressions.Cast") + // Functions that take no arguments are exempt from the `arguments` documentation + // requirement, since there is nothing to describe. + val noArgumentsSet = Set( + "org.apache.spark.sql.catalyst.expressions.Collations", + "org.apache.spark.sql.catalyst.expressions.CumeDist", + "org.apache.spark.sql.catalyst.expressions.CurDateExpressionBuilder", + "org.apache.spark.sql.catalyst.expressions.CurrentCatalog", + "org.apache.spark.sql.catalyst.expressions.CurrentDatabase", + "org.apache.spark.sql.catalyst.expressions.CurrentDate", + "org.apache.spark.sql.catalyst.expressions.CurrentPath", + "org.apache.spark.sql.catalyst.expressions.CurrentTimeZone", + "org.apache.spark.sql.catalyst.expressions.CurrentTimestamp", + "org.apache.spark.sql.catalyst.expressions.CurrentUser", + "org.apache.spark.sql.catalyst.expressions.EulerNumber", + "org.apache.spark.sql.catalyst.expressions.InputFileBlockLength", + "org.apache.spark.sql.catalyst.expressions.InputFileBlockStart", + "org.apache.spark.sql.catalyst.expressions.InputFileName", + "org.apache.spark.sql.catalyst.expressions.LocalTimestamp", + "org.apache.spark.sql.catalyst.expressions.MonotonicallyIncreasingID", + "org.apache.spark.sql.catalyst.expressions.Now", + "org.apache.spark.sql.catalyst.expressions.Pi", + "org.apache.spark.sql.catalyst.plans.logical.PythonWorkerLogs", + "org.apache.spark.sql.catalyst.expressions.RowNumber", + "org.apache.spark.sql.catalyst.expressions.SQLKeywords", + "org.apache.spark.sql.catalyst.expressions.SparkPartitionID", + "org.apache.spark.sql.catalyst.expressions.SparkVersion", + "org.apache.spark.sql.catalyst.expressions.Uuid") + spark.sessionState.functionRegistry.listFunction().foreach { funcId => val info = spark.sessionState.catalog.lookupFunctionInfo(funcId) if (!ignoreSet.contains(info.getClassName)) { @@ -179,6 +217,9 @@ class ExpressionInfoSuite extends SharedSparkSession { assert(info.getSince.matches("[0-9]+\\.[0-9]+\\.[0-9]+")) assert(info.getGroup.nonEmpty) + if (!noArgumentsSet.contains(info.getClassName)) { + assert(info.getArguments.nonEmpty) + } if (info.getArguments.nonEmpty) { assert(info.getArguments.startsWith("\n Arguments:\n")) assert(info.getArguments.endsWith("\n ")) @@ -207,9 +248,12 @@ class ExpressionInfoSuite extends SharedSparkSession { "org.apache.spark.sql.catalyst.expressions.CurrentDate", "org.apache.spark.sql.catalyst.expressions.CurDateExpressionBuilder", "org.apache.spark.sql.catalyst.expressions.CurrentTimestamp", + "org.apache.spark.sql.catalyst.expressions.CurrentTimestampExpressionBuilder", "org.apache.spark.sql.catalyst.expressions.CurrentTimeZone", "org.apache.spark.sql.catalyst.expressions.Now", + "org.apache.spark.sql.catalyst.expressions.NowExpressionBuilder", "org.apache.spark.sql.catalyst.expressions.LocalTimestamp", + "org.apache.spark.sql.catalyst.expressions.LocalTimestampExpressionBuilder", "org.apache.spark.sql.catalyst.expressions.CurrentTime", // Random output without a seed "org.apache.spark.sql.catalyst.expressions.Rand", @@ -248,6 +292,8 @@ class ExpressionInfoSuite extends SharedSparkSession { val clonedSpark = spark.cloneSession() // Coalescing partitions can change result order, so disable it. clonedSpark.conf.set(SQLConf.COALESCE_PARTITIONS_ENABLED.key, false) + // parse_sql examples require the experimental feature flag. + clonedSpark.conf.set(SQLConf.PARSE_SQL_ENABLED.key, true) val info = clonedSpark.sessionState.catalog.lookupFunctionInfo(funcId) val className = info.getClassName if (!ignoreSet.contains(className)) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala index b654deb12ad4e..6642b527e40e1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCSuite.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser import org.apache.spark.sql.catalyst.plans.logical.ShowCreateTable import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, CharVarcharUtils, DateTimeTestUtils} import org.apache.spark.sql.connector.catalog.Identifier -import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue} +import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue} import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, Predicate} import org.apache.spark.sql.execution.{DataSourceScanExec, ExtendedMode, ProjectExec} import org.apache.spark.sql.execution.command.{ExplainCommand, ShowCreateTableCommand} @@ -93,6 +93,7 @@ class JDBCSuite extends SharedSparkSession { jdbcClientType: String): Metadata = new MetadataBuilder() .putLong("scale", 0) .putBoolean("isTimestampNTZ", false) + .putBoolean("preferTimestampNanos", false) .putBoolean("isSigned", dataType.isInstanceOf[NumericType]) .putString("jdbcClientType", jdbcClientType) .build() @@ -478,6 +479,67 @@ class JDBCSuite extends SharedSparkSession { assert(lastPredicate == """"PartitionColumn" >= '2020-08-02'""") } + test("columnPartition supports TimestampNTZType partition column") { + val schema = StructType(Seq( + StructField("PartitionColumn", TimestampNTZType) + )) + + // (lowerBound, upperBound, numPartitions, expected where clauses in partition order). + val cases = Seq( + ("2018-07-06 10:00:00", "2018-07-06 16:00:00", "3", Seq( + """"PartitionColumn" < '2018-07-06 12:00:00' or "PartitionColumn" is null""", + """"PartitionColumn" >= '2018-07-06 12:00:00' AND """ + + """"PartitionColumn" < '2018-07-06 14:00:00'""", + """"PartitionColumn" >= '2018-07-06 14:00:00'""")), + // Fractional-second bounds parse, and the zoneless midpoint keeps sub-second precision. + ("2018-07-06 10:00:00.100", "2018-07-06 10:00:00.300", "2", Seq( + """"PartitionColumn" < '2018-07-06 10:00:00.2' or "PartitionColumn" is null""", + """"PartitionColumn" >= '2018-07-06 10:00:00.2'""")) + ) + + // NTZ bounds are zoneless, so the generated predicates must be identical regardless of the + // session time zone (unlike TimestampType, which shifts by the zone). + Seq("UTC", "America/Los_Angeles", "Asia/Kolkata").foreach { tz => + cases.foreach { case (lowerBound, upperBound, numPartitions, expected) => + val partitions = JDBCRelation.columnPartition( + schema, + analysis.caseInsensitiveResolution, + tz, + new JDBCOptions(url, "table", Map( + "lowerBound" -> lowerBound, + "upperBound" -> upperBound, + "numPartitions" -> numPartitions, + "partitionColumn" -> "PartitionColumn"))) + + val clauses = partitions.map(_.asInstanceOf[JDBCPartition].whereClause) + assert(clauses === expected.toArray, + s"NTZ partition clauses should be time-zone independent, but differed for tz=$tz " + + s"(bounds $lowerBound..$upperBound)") + } + } + } + + test("columnPartition rejects zoned bounds for a TimestampNTZType partition column") { + val schema = StructType(Seq( + StructField("PartitionColumn", TimestampNTZType) + )) + // allowTimeZone = false: a bound carrying a zone offset is rejected rather than silently + // shifted, so NTZ bounds stay zoneless. + val e = intercept[IllegalArgumentException] { + JDBCRelation.columnPartition( + schema, + analysis.caseInsensitiveResolution, + "America/Los_Angeles", + new JDBCOptions(url, "table", Map( + "lowerBound" -> "2018-07-06 10:00:00+05:00", + "upperBound" -> "2018-07-06 16:00:00+05:00", + "numPartitions" -> "2", + "partitionColumn" -> "PartitionColumn"))) + } + assert(e.getMessage.contains("Cannot parse the bound value")) + assert(e.getMessage.contains("2018-07-06 10:00:00+05:00")) + } + test("overflow of partition bound difference does not give negative stride") { val df = sql("SELECT * FROM partsoverflow") checkNumPartitions(df, expectedNumPartitions = 3) @@ -837,6 +899,72 @@ class JDBCSuite extends SharedSparkSession { } } + test("SPARK-57460: JDBC TIMESTAMP keeps microsecond mapping by default") { + // Without preferTimestampNanos, a driver TIMESTAMP(9) still infers as microsecond + // TimestampType even when the nanos preview feature is enabled. + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + val conn = java.sql.DriverManager.getConnection(urlWithUserAndPass) + try { + conn.createStatement().execute("CREATE TABLE TEST.TS_DEFAULT (t TIMESTAMP(9))") + conn.createStatement().execute( + "INSERT INTO TEST.TS_DEFAULT VALUES (TIMESTAMP '2020-02-02 04:13:14.123456789')") + val df = spark.read.jdbc(urlWithUserAndPass, "TEST.TS_DEFAULT", new Properties()) + assert(df.schema("T").dataType === TimestampType) + } finally { + conn.createStatement().execute("DROP TABLE IF EXISTS TEST.TS_DEFAULT") + conn.close() + } + } + } + + test("SPARK-57460: JDBC TIMESTAMP reads nanosecond LTZ precision when requested") { + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + val conn = java.sql.DriverManager.getConnection(urlWithUserAndPass) + try { + conn.createStatement().execute("CREATE TABLE TEST.TS_NANOS_LTZ (t TIMESTAMP(9))") + conn.createStatement().execute( + "INSERT INTO TEST.TS_NANOS_LTZ VALUES (TIMESTAMP '2020-02-02 04:13:14.123456789')") + val df = spark.read + .option("preferTimestampNanos", "true") + .jdbc(urlWithUserAndPass, "TEST.TS_NANOS_LTZ", new Properties()) + assert(df.schema("T").dataType === TimestampLTZNanosType(9)) + val result = df.collect() + // H2 stores TIMESTAMP as local wall-clock; the LTZ read binds it to the session zone. + val expected = java.time.LocalDateTime.of(2020, 2, 2, 4, 13, 14, 123456789) + .atZone(java.time.ZoneId.systemDefault()).toInstant + assert(result(0).getAs[java.time.Instant](0) === expected) + } finally { + conn.createStatement().execute("DROP TABLE IF EXISTS TEST.TS_NANOS_LTZ") + conn.close() + } + } + } + + test("SPARK-57460: JDBC TIMESTAMP reads nanosecond NTZ precision when requested") { + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + val conn = java.sql.DriverManager.getConnection(urlWithUserAndPass) + try { + conn.createStatement().execute("CREATE TABLE TEST.TS_NANOS_NTZ (t TIMESTAMP(8))") + conn.createStatement().execute( + "INSERT INTO TEST.TS_NANOS_NTZ VALUES (TIMESTAMP '2020-02-02 04:13:14.123456789')") + val df = spark.read + .option("preferTimestampNanos", "true") + .option("preferTimestampNTZ", "true") + .jdbc(urlWithUserAndPass, "TEST.TS_NANOS_NTZ", new Properties()) + // Reported scale 8 -> precision 8. H2 rounds the stored value to 8 fractional digits + // (.123456789 -> .12345679); Spark's read path only floors to precision, so the value is + // preserved at 8 digits. + assert(df.schema("T").dataType === TimestampNTZNanosType(8)) + val result = df.collect() + assert(result(0).getAs[java.time.LocalDateTime](0) === + java.time.LocalDateTime.of(2020, 2, 2, 4, 13, 14, 123456790)) + } finally { + conn.createStatement().execute("DROP TABLE IF EXISTS TEST.TS_NANOS_NTZ") + conn.close() + } + } + } + test("test DATE types") { val rows = spark.read.jdbc( urlWithUserAndPass, "TEST.TIMETYPES", new Properties()).collect() @@ -1454,6 +1582,13 @@ class JDBCSuite extends SharedSparkSession { assert(mySqlDialect.getJDBCType(FloatType).map(_.databaseTypeDefinition).get == "FLOAT") } + test("MySQL blocks casts to double") { + val dialect = MySQLDialect() + val cast = new V2Cast(FieldReference("value"), IntegerType, DoubleType) + + assert(dialect.compileExpression(cast).isEmpty) + } + test("PostgresDialect type mapping") { val Postgres = JdbcDialects.get("jdbc:postgresql://127.0.0.1/db") val md = new MetadataBuilder().putLong("scale", 0).putBoolean("isTimestampNTZ", false) @@ -2413,6 +2548,52 @@ class JDBCSuite extends SharedSparkSession { checkAnswer(df2, expectedResult) } + test("support TimestampNTZType partition column end-to-end") { + val tableName = "timestamp_ntz_partition_table" + // Write a genuine TimestampNTZType column through Spark so it round-trips as a zoneless + // wall-clock value (a raw JDBC TIMESTAMP column would pick up a JVM-time-zone shift on read). + val df = Seq( + "2018-07-06T05:50:00", + "2018-07-06T08:10:08", + "2018-07-08T13:32:01", + "2018-07-12T09:51:15" + ).map(LocalDateTime.parse).toDF("t") + df.write.format("jdbc") + .mode("overwrite") + .option("url", urlWithUserAndPass) + .option("dbtable", tableName) + .save() + + // Bounds are zoneless, so both the generated predicates and the results must be identical + // regardless of the JVM default time zone. + DateTimeTestUtils.outstandingZoneIds.foreach { zoneId => + DateTimeTestUtils.withDefaultTimeZone(zoneId) { + val readDf = spark.read.format("jdbc") + .option("url", urlWithUserAndPass) + .option("dbtable", tableName) + .option("preferTimestampNTZ", true) + .option("partitionColumn", "t") + .option("lowerBound", "2018-07-04 03:30:00") + .option("upperBound", "2018-07-27 14:11:05") + .option("numPartitions", 2) + .load() + + assert(readDf.schema("t").dataType === TimestampNTZType) + + readDf.logicalPlan match { + case LogicalRelationWithTable(JDBCRelation(_, parts, _, _), _) => + val whereClauses = parts.map(_.asInstanceOf[JDBCPartition].whereClause).toSet + assert(whereClauses === Set( + """"t" < '2018-07-15 20:50:32.5' or "t" is null""", + """"t" >= '2018-07-15 20:50:32.5'"""), + s"NTZ partition predicates should be time-zone independent, but differed for " + + s"zone=$zoneId") + } + checkAnswer(readDf, df) + } + } + } + test("throws an exception for unsupported partition column types") { val errMsg = intercept[AnalysisException] { spark.read.format("jdbc") @@ -2801,7 +2982,8 @@ class JDBCSuite extends SharedSparkSession { "hint" -> hint)) }.getMessage assert(e.contains(s"Invalid value `$hint` for option `hint`." + - s" It should start with `/*+ ` and end with ` */`.")) + s" It should start with `/*+ ` and end with ` */`," + + s" for example `/*+ INDEX(t1 id_idx) */`.")) } // dialect supported check diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala index d7e5786ecb9c7..8f456c3e57ea7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala @@ -236,6 +236,12 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { batchStmt.addBatch("INSERT INTO \"test\".\"strings_with_nulls\" VALUES ('a a a')") batchStmt.addBatch("INSERT INTO \"test\".\"strings_with_nulls\" VALUES (null)") + batchStmt.addBatch( + "CREATE TABLE \"test\".\"null_literal\" (s TEXT(32))") + batchStmt.addBatch("INSERT INTO \"test\".\"null_literal\" VALUES ('keep')") + batchStmt.addBatch("INSERT INTO \"test\".\"null_literal\" VALUES ('')") + batchStmt.addBatch("INSERT INTO \"test\".\"null_literal\" VALUES (null)") + batchStmt.executeBatch() conn @@ -1819,7 +1825,8 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { Seq(Row("test", "address", false), Row("test", "people", false), Row("test", "empty_table", false), Row("test", "employee", false), Row("test", "item", false), Row("test", "dept", false), - Row("test", "person", false), Row("test", "view1", false), Row("test", "view2", false), + Row("test", "null_literal", false), Row("test", "person", false), + Row("test", "view1", false), Row("test", "view2", false), Row("test", "datetime", false), Row("test", "binary_tab", false), Row("test", "employee_bonus", false), Row("test", "strings_with_nulls", false))) @@ -2515,7 +2522,7 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { checkAggregateRemoved(df3) checkPushedInfo(df3, """ - |PushedAggregates: [AVG(CASE WHEN BONUS IS NOT NULL THEN BONUS ELSE null END)], + |PushedAggregates: [AVG(CASE WHEN BONUS IS NOT NULL THEN BONUS ELSE NULL END)], |PushedFilters: [DEPT IS NOT NULL, DEPT > 0], |PushedGroupByExpressions: [DEPT], |""".stripMargin.replaceAll("\n", " ")) @@ -2531,7 +2538,7 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { checkAggregateRemoved(df4) checkPushedInfo(df4, """ - |PushedAggregates: [AVG(DISTINCT CASE WHEN BONUS IS NOT NULL THEN BONUS ELSE null END)], + |PushedAggregates: [AVG(DISTINCT CASE WHEN BONUS IS NOT NULL THEN BONUS ELSE NULL END)], |PushedFilters: [DEPT IS NOT NULL, DEPT > 0], |PushedGroupByExpressions: [DEPT], |""".stripMargin.replaceAll("\n", " ")) @@ -3174,4 +3181,17 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { assertResult(expectedMetadata) { jdbcRdd.getDatabaseMetadata } } + + test("SPARK-58782: null literal in aggregate should render as NULL not 'null'") { + val df = sql("SELECT NULLIF(s, '') AS g, COUNT(*) FROM h2.test.null_literal GROUP BY g") + + checkAggregateRemoved(df) + checkPushedInfo(df, + "PushedAggregates: [COUNT(*)]", + "PushedGroupByExpressions: [CASE WHEN S = '' THEN NULL ELSE S END]") + + // The '' row should collapse into the NULL group, not create a separate 'null' string group + checkAnswer(df, Seq(Row("keep", 1), Row(null, 2))) + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala index bce7a883fa2b9..1045a5afa5d7e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCWriteSuite.scala @@ -28,15 +28,12 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.SparkException import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} -import org.apache.spark.sql.{AnalysisException, Column, DataFrame, Row, SaveMode} -import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.parser.ParseException -import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.execution.datasources.jdbc.{JDBCOptions, JdbcUtils} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ -import org.apache.spark.unsafe.types.TimestampNanosVal import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils @@ -701,22 +698,45 @@ class JDBCWriteSuite extends SharedSparkSession with BeforeAndAfter { === java.sql.Timestamp.valueOf("2020-02-02 04:13:14.56789")) } - test("SPARK-57166: nanosecond timestamp types are not supported in JDBC write") { + test("SPARK-57460: JDBC nanosecond timestamp write round-trip") { withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { - Seq(TimestampNTZNanosType(9), TimestampLTZNanosType(9)).foreach { nanosType => - // The nanos literal is built directly from its internal value to avoid relying on - // cast/parser support. - val nanosLiteral = Literal.create(new TimestampNanosVal(0L, 0.toShort), nanosType) - val df = spark.range(1).select(Column(nanosLiteral).as("ts")) - checkErrorMatchPVals( - exception = intercept[AnalysisException] { - df.write.jdbc(url, "TEST.NANOSTYPES", new Properties()) - }, - condition = "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE", - parameters = Map( - "columnName" -> "`ts`", - "columnType" -> java.util.regex.Pattern.quote(s""""${nanosType.sql}""""), - "format" -> ".*")) + Seq(7, 8, 9).foreach { precision => + // TIMESTAMP_NTZ(p): write a LocalDateTime, read it back as a nanos NTZ column. + val ntzTable = "TEST.NANOS_NTZ" + val ldt = java.time.LocalDateTime.of(2020, 2, 2, 4, 13, 14, 123456789) + val ntzSchema = new StructType().add("t", TimestampNTZNanosType(precision)) + val ntzDf = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(ldt))), ntzSchema) + ntzDf.write.mode(SaveMode.Overwrite).jdbc(url, ntzTable, new Properties()) + val ntzReadBack = spark.read + .option("preferTimestampNanos", "true") + .option("preferTimestampNTZ", "true") + .jdbc(url, ntzTable, new Properties()) + assert(ntzReadBack.schema.fields(0).dataType === TimestampNTZNanosType(precision)) + val expectedNtz = precision match { + case 7 => java.time.LocalDateTime.of(2020, 2, 2, 4, 13, 14, 123456700) + case 8 => java.time.LocalDateTime.of(2020, 2, 2, 4, 13, 14, 123456780) + case 9 => ldt + } + assert(ntzReadBack.collect()(0).getAs[java.time.LocalDateTime](0) === expectedNtz) + + // TIMESTAMP_LTZ(p): write an Instant, read it back as a nanos LTZ column. + val ltzTable = "TEST.NANOS_LTZ" + val instant = java.time.Instant.parse("2020-02-02T12:13:14.123456789Z") + val ltzSchema = new StructType().add("t", TimestampLTZNanosType(precision)) + val ltzDf = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(instant))), ltzSchema) + ltzDf.write.mode(SaveMode.Overwrite).jdbc(url, ltzTable, new Properties()) + val ltzReadBack = spark.read + .option("preferTimestampNanos", "true") + .jdbc(url, ltzTable, new Properties()) + assert(ltzReadBack.schema.fields(0).dataType === TimestampLTZNanosType(precision)) + val expectedLtz = precision match { + case 7 => java.time.Instant.parse("2020-02-02T12:13:14.123456700Z") + case 8 => java.time.Instant.parse("2020-02-02T12:13:14.123456780Z") + case 9 => instant + } + assert(ltzReadBack.collect()(0).getAs[java.time.Instant](0) === expectedLtz) } } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala index 95c2fcbd7b5d7..34d87b6f91d7c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/CreateTableAsSelectSuite.scala @@ -30,7 +30,6 @@ import org.apache.spark.util.Utils class CreateTableAsSelectSuite extends DataSourceTest with SharedSparkSession { import testImplicits._ - protected override lazy val sql = spark.sql _ private var path: File = null override def beforeAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala index 786e50eea2e71..cbd92f0af1901 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/FilteredScanSuite.scala @@ -134,7 +134,6 @@ object ColumnsRequired { } class FilteredScanSuite extends DataSourceTest with SharedSparkSession { - protected override lazy val sql = spark.sql _ override def beforeAll(): Unit = { super.beforeAll() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala index f2645e32e8a9b..4668bde453de5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/InsertSuite.scala @@ -34,6 +34,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf.PartitionOverwriteMode import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ +import org.apache.spark.tags.ExtendedSQLTest import org.apache.spark.util.Utils class SimpleInsertSource extends SchemaRelationProvider { @@ -57,10 +58,10 @@ case class SimpleInsert(userSpecifiedSchema: StructType)(@transient val sparkSes } } +@ExtendedSQLTest class InsertSuite extends DataSourceTest with SharedSparkSession { import testImplicits._ - protected override lazy val sql = spark.sql _ private var path: File = null override def beforeAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/PrunedScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/PrunedScanSuite.scala index f242f75f39f20..4e48d0246e57d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/PrunedScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/PrunedScanSuite.scala @@ -54,7 +54,6 @@ case class SimplePrunedScan(from: Int, to: Int)(@transient val sparkSession: Spa } class PrunedScanSuite extends DataSourceTest with SharedSparkSession { - protected override lazy val sql = spark.sql _ override def beforeAll(): Unit = { super.beforeAll() diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/SaveLoadSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/SaveLoadSuite.scala index 1e0d9bd9990e2..571b558c3c762 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/SaveLoadSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/SaveLoadSuite.scala @@ -32,7 +32,6 @@ import org.apache.spark.util.Utils class SaveLoadSuite extends DataSourceTest with SharedSparkSession with BeforeAndAfter { import testImplicits._ - protected override lazy val sql = spark.sql _ private var originalDefaultSource: String = null private var path: File = null private var df: DataFrame = null diff --git a/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala index eac77c2938207..a22fecc237e12 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/sources/TableScanSuite.scala @@ -123,7 +123,6 @@ class LegacyTimestampSource extends RelationProvider { } class TableScanSuite extends DataSourceTest with SharedSparkSession { - protected override lazy val sql = spark.sql _ private lazy val tableWithSchemaExpected = (1 to 10).map { i => Row( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/CommitLogSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/CommitLogSuite.scala index ba785aa830ab5..a733bd6d75719 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/CommitLogSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/CommitLogSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.streaming import java.io.{ByteArrayInputStream, FileInputStream, FileOutputStream} import java.nio.file.Path -import org.apache.spark.sql.execution.streaming.checkpointing.{CommitLog, CommitMetadata, CommitMetadataBase, CommitMetadataV2, CommitMetadataV3, OffsetSeqLog, SinkMetadataInfo} +import org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointVersionManager, CommitLog, CommitLogType, CommitMetadata, CommitMetadataBase, CommitMetadataV2, CommitMetadataV3, OffsetSeqLog, SinkMetadataInfo} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession @@ -250,4 +250,70 @@ class CommitLogSuite extends SharedSparkSession { commitLogFormatVersion = CommitLog.VERSION_1).version === CommitLog.VERSION_1) } } + + /** The commit log version the session config asks for, via the public resolution entry point. */ + private def sessionCommitLogVersion(): Int = { + CheckpointVersionManager.resolveCommitLogVersion(spark, latestCommittedBatch = None) + } + + test("commit log version derives from the state store checkpoint format") { + // Nothing set: defaults to VERSION_1. + assert(sessionCommitLogVersion() === CommitLog.VERSION_1) + + // State store checkpoint format v2 makes each batch write stateUniqueIds, which only a commit + // log at VERSION_2 or above can persist, so it raises the commit log version to v2. + withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + assert(sessionCommitLogVersion() === CommitLog.VERSION_2, + "state store v2 must raise the commit log to v2") + } + + // The resolved version is capped at VERSION_2: VERSION_3 exists only to carry sink-evolution + // metadata and is written exclusively by the sink-evolution path, never derived from a config. + withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "3") { + assert(sessionCommitLogVersion() === CommitLog.VERSION_2, + "a config-derived commit log version must never resolve to v3") + } + } + + test("an existing checkpoint's commit log version wins over the session config") { + // The whole point of resolution: a commit log created at one version keeps being written at + // that version, so a higher state store checkpoint format cannot start writing a format the + // checkpoint lacks. + withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + val existing: CommitMetadataBase = CommitMetadata(nextBatchWatermarkMs = 0) + assert(existing.version === CommitLog.VERSION_1) + val resolved = + CheckpointVersionManager.resolveCommitLogVersion(spark, Some((7L, existing))) + assert(resolved === CommitLog.VERSION_1, + s"an existing V1 commit log must stay V1 even with the state store at v2, got $resolved") + } + } + + test("recording a commit log version sets the implied state store format") { + val sinkMetadataMap = Map("sink" -> SinkMetadataInfo( + sinkName = "sink", + commitOffset = OffsetSeqLog.SERIALIZED_VOID_OFFSET, + providerName = "provider", + apiVersion = "DSv2")) + val session = spark.cloneSession() + val v3WithStateIds = CommitMetadataV3(0, Some(Map.empty), sinkMetadataMap) + CheckpointVersionManager.setFormatVersion( + session, CommitLogType, CommitLog.VERSION_3, Some(v3WithStateIds)) + assert(session.conf.get(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key) === "2", + "state checkpoint ids in a V3 commit imply state store format v2") + + val v3WithoutStateIdsSession = spark.cloneSession() + val v3WithoutStateIds = CommitMetadataV3(0, None, sinkMetadataMap) + CheckpointVersionManager.setFormatVersion( + v3WithoutStateIdsSession, CommitLogType, CommitLog.VERSION_3, Some(v3WithoutStateIds)) + assert(v3WithoutStateIdsSession.conf.get( + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key) === "1", + "a V3 commit without state checkpoint ids must preserve state store format v1") + + val v1Session = spark.cloneSession() + CheckpointVersionManager.setFormatVersion(v1Session, CommitLogType, CommitLog.VERSION_1) + assert(v1Session.conf.get(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key) === "1", + "a V1 commit log cannot carry state store checkpoint ids, so the state store must be v1") + } + } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeTransformWithStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeTransformWithStateSuite.scala new file mode 100644 index 0000000000000..eac33ef717e5f --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/RealTimeTransformWithStateSuite.scala @@ -0,0 +1,2052 @@ +/* + * 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.streaming + +import java.sql.Timestamp +import java.time.Duration +import java.util.concurrent.{CountDownLatch, TimeUnit} + +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.{ + SparkConf, SparkException, SparkRuntimeException, SparkThrowable, TaskContext, TaskContextImpl} +import org.apache.spark.sql.Encoders +import org.apache.spark.sql.execution.SortExec +import org.apache.spark.sql.execution.datasources.v2.{LowLatencyClock, RealTimeStreamScanExec} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.streaming.operators.stateful.transformwithstate.TransformWithStateExec +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.{ + EnableStateStoreRowChecksum, RocksDBConf, RocksDBStateStoreProvider} +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.util.{GlobalSingletonManualClock, StreamManualClock} +import org.apache.spark.tags.SlowSQLTest + +private class RealTimeEagerCountProcessor + extends StatefulProcessor[String, (String, Int), (String, Long)] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState("count", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, Long)] = { + // Eager consumption detects the blocking final group produced by GroupedIterator. + val newCount = Option(countState.get()).getOrElse(0L) + inputRows.size + countState.update(newCount) + // Access state lazily as the output is consumed to verify implicit-key lifecycle handling. + Iterator.single(key).map(currentKey => (currentKey, countState.get())) + } +} + +private class RealTimeRunningCountStatefulProcessor(emitEvery: Long) + extends StatefulProcessor[String, String, (String, Long)] { + + @transient private var countState: MapState[String, Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getMapState( + "countState", Encoders.STRING, Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[String], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.flatMap { row => + val count = countState.getValue(row) + 1L + countState.updateValue(row, count) + if (count % emitEvery == 0L) Iterator.single((row, count)) else Iterator.empty + } + } +} + +private class RealTimeTTLCountProcessor(ttl: Duration = Duration.ofSeconds(10)) + extends StatefulProcessor[String, (String, Int), (String, Long)] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState( + "count", Encoders.scalaLong, TTLConfig(ttl)) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.map { _ => + val newCount = Option(countState.get()).getOrElse(0L) + 1L + countState.update(newCount) + (key, newCount) + } + } +} + +private class RealTimeListTTLProcessor + extends StatefulProcessor[String, (String, Int), (String, Long)] { + + @transient private var listState: ListState[Int] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + listState = getHandle.getListState( + "values", Encoders.scalaInt, TTLConfig(Duration.ofSeconds(10))) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.map { case (_, value) => + listState.appendList(Array(value, value + 1, value + 2)) + (key, listState.get().size.toLong) + } + } +} + +private class RealTimeMapTTLAndTimerProcessor + extends StatefulProcessor[String, (String, Int), (String, String, Long)] { + + @transient private var countState: MapState[String, Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getMapState( + "count", Encoders.STRING, Encoders.scalaLong, TTLConfig(Duration.ofMinutes(10))) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, String, Long)] = { + inputRows.map { case (_, timerDelayMs) => + val count = Option(countState.getValue("count")).getOrElse(0L) + 1L + countState.updateValue("count", count) + getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + timerDelayMs) + (key, "data", count) + } + } + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String, Long)] = { + Iterator.single((key, "timer", expiredTimerInfo.getExpiryTimeInMs())) + } +} + +private class RealTimePartitionProcessor + extends StatefulProcessor[String, (String, Int), (String, Int)] { + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {} + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, Int)] = { + inputRows.map(_ => (key, TaskContext.getPartitionId())) + } +} + +private class RealTimeProcessingTimerProcessor + extends StatefulProcessor[String, (String, Int), (String, String)] { + + @transient private var timerRegistered: ValueState[Boolean] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + timerRegistered = + getHandle.getValueState("timerRegistered", Encoders.scalaBoolean, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, String)] = { + inputRows.map { _ => + if (!Option(timerRegistered.get()).getOrElse(false)) { + getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + 10000L) + timerRegistered.update(true) + } + (key, "data") + } + } + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String)] = { + Iterator.single((key, "timer")) + } +} + +private class RealTimeProcessingTimerValueProcessor + extends StatefulProcessor[String, (String, Int), (String, String, Long)] { + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {} + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, String, Long)] = { + inputRows.map { _ => + val currentTimeMs = timerValues.getCurrentProcessingTimeInMs() + getHandle.registerTimer(currentTimeMs + 10000L) + (key, "data", currentTimeMs) + } + } + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String, Long)] = { + Iterator.single((key, "timer", timerValues.getCurrentProcessingTimeInMs())) + } +} + +private object TaskCompletionListenerCount { + private lazy val listenerStackField = { + val field = classOf[TaskContextImpl].getDeclaredField("onCompleteCallbacks") + field.setAccessible(true) + field + } + + def get(): Int = listenerStackField.get(TaskContext.get()) + .asInstanceOf[java.util.Stack[_]].size() +} + +private class RealTimeTimerIteratorListenerProcessor + extends StatefulProcessor[String, (String, Int), (Int, Boolean)] { + + private var previousListenerCount = -1 + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {} + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(Int, Boolean)] = { + inputRows.map { case (_, value) => + val listenerCount = TaskCompletionListenerCount.get() + val listenerCountIncreased = + previousListenerCount >= 0 && listenerCount > previousListenerCount + previousListenerCount = listenerCount + (value, listenerCountIncreased) + } + } +} + +private class RealTimeTTLIteratorListenerProcessor + extends StatefulProcessor[String, (String, Int), (Int, Int)] { + + @transient private var valueState: ValueState[Int] = _ + private var previousListenerCount = -1 + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + valueState = getHandle.getValueState( + "value", Encoders.scalaInt, TTLConfig(Duration.ofHours(1))) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(Int, Int)] = { + inputRows.map { case (_, value) => + valueState.update(value) + val listenerCount = TaskCompletionListenerCount.get() + val addedListenerCount = if (previousListenerCount >= 0) { + listenerCount - previousListenerCount + } else { + 0 + } + previousListenerCount = listenerCount + (value, addedListenerCount) + } + } +} + +private class RTMStatefulProcessorWithProcTimeTimerWithMultipleTimers(timerExpireTs: Long) + extends RTMStatefulProcessorWithProcTimeTimer(timerExpireTs) { + override def handleInputRows( + key: String, + inputRows: Iterator[String], + timerValues: TimerValues): Iterator[(String, String)] = { + + val currCount = Option(_countState.get()).getOrElse(0L) + if (currCount == 0 && (key == "a" || key == "c")) { + getHandle.registerTimer( + timerValues.getCurrentProcessingTimeInMs() + timerExpireTs + ) + + getHandle.registerTimer( + timerValues.getCurrentProcessingTimeInMs() + (timerExpireTs + 1000) + ) + } + + val count = currCount + 1 + if (count == 3) { + _countState.clear() + Iterator.empty + } else { + _countState.update(count) + Iterator((key, count.toString)) + } + } +} + +private class RTMStatefulProcessorWithProcTimeTimer(timerExpireTs: Long) + extends RunningCountStatefulProcessor { + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String)] = { + _countState.clear() + Iterator((key, "-1")) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[String], + timerValues: TimerValues): Iterator[(String, String)] = { + + val currCount = Option(_countState.get()).getOrElse(0L) + if (currCount == 0 && (key == "a" || key == "c")) { + getHandle.registerTimer( + timerValues.getCurrentProcessingTimeInMs() + timerExpireTs + ) + } + + val count = currCount + 1 + if (count == 3) { + _countState.clear() + Iterator.empty + } else { + _countState.update(count) + Iterator((key, count.toString)) + } + } +} + +private class RTMStatefulProcessorWithProcTimeTimerInputInt(timerExpireTs: Long) + extends StatefulProcessor[Int, Int, (Int, Long)] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState( + "countState", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleExpiredTimer( + key: Int, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(Int, Long)] = { + countState.clear() + Iterator.single((key, -1L)) + } + + override def handleInputRows( + key: Int, + inputRows: Iterator[Int], + timerValues: TimerValues): Iterator[(Int, Long)] = { + val currentCount = Option(countState.get()).getOrElse(0L) + if (currentCount == 0L) { + getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + timerExpireTs) + } + val count = currentCount + 1L + if (count == 3L) { + countState.clear() + Iterator.empty + } else { + countState.update(count) + Iterator.single((key, count)) + } + } +} + +private class RTMStatefulProcessorWithProcTimeTimerInputTuple(timerExpireTs: Long) + extends StatefulProcessor[(Int, String), (Int, String), (Int, String, Long)] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState( + "countState", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleExpiredTimer( + key: (Int, String), + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(Int, String, Long)] = { + countState.clear() + Iterator.single((key._1, key._2, -1L)) + } + + override def handleInputRows( + key: (Int, String), + inputRows: Iterator[(Int, String)], + timerValues: TimerValues): Iterator[(Int, String, Long)] = { + val currentCount = Option(countState.get()).getOrElse(0L) + if (currentCount == 0L) { + getHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs() + timerExpireTs) + } + val count = currentCount + 1L + if (count == 3L) { + countState.clear() + Iterator.empty + } else { + countState.update(count) + Iterator.single((key._1, key._2, count)) + } + } +} + +private class RealTimeEventTimerProcessor + extends StatefulProcessor[String, (Timestamp, String), (String, String, Long)] { + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = {} + + override def handleInputRows( + key: String, + inputRows: Iterator[(Timestamp, String)], + timerValues: TimerValues): Iterator[(String, String, Long)] = { + inputRows.map { case (eventTime, _) => + getHandle.registerTimer(eventTime.getTime) + (key, "data", timerValues.getCurrentWatermarkInMs()) + } + } + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String, Long)] = { + Iterator.single((key, "timer", timerValues.getCurrentWatermarkInMs())) + } +} + +private case class RealTimeEventTimeOutputRow( + key: String, + outputEventTime: Timestamp, + count: Long) + +private class RealTimeEventTimeOutputProcessor( + outputEventTimeOverride: Option[Timestamp] = None) + extends StatefulProcessor[String, (Timestamp, String), RealTimeEventTimeOutputRow] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState("count", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(Timestamp, String)], + timerValues: TimerValues): Iterator[RealTimeEventTimeOutputRow] = { + inputRows.map { case (eventTime, _) => + val newCount = Option(countState.get()).getOrElse(0L) + 1L + countState.update(newCount) + RealTimeEventTimeOutputRow( + key, outputEventTimeOverride.getOrElse(eventTime), newCount) + } + } +} + +private class RealTimeInitialCountProcessor + extends StatefulProcessorWithInitialState[ + String, (String, Int), (String, Long), Long] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState("count", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInitialState( + key: String, + initialState: Long, + timerValues: TimerValues): Unit = { + countState.update(Option(countState.get()).getOrElse(0L) + initialState) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(String, Int)], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.map { _ => + val newCount = Option(countState.get()).getOrElse(0L) + 1L + countState.update(newCount) + (key, newCount) + } + } +} + +private class RealTimeEventInitialCountProcessor + extends StatefulProcessorWithInitialState[ + String, (Timestamp, String), (String, Long), Long] { + + @transient private var countState: ValueState[Long] = _ + + override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { + countState = getHandle.getValueState("count", Encoders.scalaLong, TTLConfig.NONE) + } + + override def handleInitialState( + key: String, + initialState: Long, + timerValues: TimerValues): Unit = { + countState.update(initialState) + } + + override def handleInputRows( + key: String, + inputRows: Iterator[(Timestamp, String)], + timerValues: TimerValues): Iterator[(String, Long)] = { + inputRows.map { _ => + val newCount = Option(countState.get()).getOrElse(0L) + 1L + countState.update(newCount) + (key, newCount) + } + } +} + +private class RealTimeInitialStateProcTimerWithExpiryProcessor + extends StatefulProcessorWithInitialStateProcTimerClass { + + override def handleExpiredTimer( + key: String, + timerValues: TimerValues, + expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String)] = { + super.handleExpiredTimer(key, timerValues, expiredTimerInfo).map { case (expiredKey, _) => + (expiredKey, expiredTimerInfo.getExpiryTimeInMs().toString) + } + } +} + +private object RealTimeInitialStateFailure { + @volatile var enabled: Boolean = false +} + +private object RealTimeInitialStateBootstrapBlock { + @volatile private var enabled = false + @volatile private var taskStarted = new CountDownLatch(0) + @volatile private var releaseTask = new CountDownLatch(0) + + def enable(): Unit = { + taskStarted = new CountDownLatch(1) + releaseTask = new CountDownLatch(1) + enabled = true + } + + def awaitTaskStart(): Boolean = taskStarted.await(1, TimeUnit.MINUTES) + + def awaitReleaseIfEnabled(): Unit = { + if (enabled) { + taskStarted.countDown() + releaseTask.await() + } + } + + def disable(): Unit = { + enabled = false + releaseTask.countDown() + } +} + +@SlowSQLTest +class RealTimeTransformWithStateSuite extends StreamRealTimeModeE2ESuiteBase { + import testImplicits._ + + override protected def sparkConf: SparkConf = super.sparkConf + .set(SQLConf.STATE_STORE_PROVIDER_CLASS.key, classOf[RocksDBStateStoreProvider].getName) + + private def advanceClock(clock: GlobalSingletonManualClock): ExternalAction = { + advanceClock(clock, defaultTrigger.batchDurationMs) + } + + private def advanceClock( + clock: GlobalSingletonManualClock, + advanceMs: Long): ExternalAction = { + new ExternalAction { + override def runAction(): Unit = clock.advance(advanceMs) + } + } + + private def createStringMemoryStream(numPartitions: Int = 2) + : (LowLatencyMemoryStream[String], GlobalSingletonManualClock) = { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + (LowLatencyMemoryStream[String](numPartitions), clock) + } + + private def waitForNextBatchToStart(): Unit = { + eventually(timeout(60.seconds)) { + val tasksRunning = spark.sparkContext.statusTracker + .getExecutorInfos.map(_.numRunningTasks()).sum + assert(tasksRunning >= 1, s"tasksRunning: $tasksRunning") + } + } + + private def initialState(values: Seq[(String, Long)]) = { + values.toDS() + // More initial-state partitions than available task slots exercises finite bootstrap + // scheduling independently from the later pipelined RTM batches. + .repartition(12, $"_1") + .map { value => + RealTimeInitialStateBootstrapBlock.awaitReleaseIfEnabled() + if (RealTimeInitialStateFailure.enabled) { + throw new RuntimeException("injected initial-state bootstrap failure") + } + value + } + .groupByKey(_._1) + .mapValues(_._2) + } + + private def transformWithInitialState( + input: LowLatencyMemoryStream[(String, Int)], + values: Seq[(String, Long)]) = { + input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeInitialCountProcessor, + TimeMode.None(), + OutputMode.Update(), + initialState(values)) + } + + test("processes repeated keys within a long-running batch without a sort") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, _) = createMemoryStream(numPartitions = 2) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeEagerCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1), ("a", 2), ("b", 1), ("a", 3)), + // The RTM batch remains open for five minutes. Seeing the final input here proves that + // TransformWithState processes rows individually rather than waiting for batch end. + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("a", 2L), ("b", 1L), ("a", 3L)), + Execute { q => + val operators = q.lastExecution.executedPlan.collect { + case t: TransformWithStateExec => t + } + assert(operators.size == 1, q.lastExecution.executedPlan) + assert(operators.head.isRealTimeMode) + assert(operators.head.requiredChildOrdering.forall(_.isEmpty)) + assert(!operators.head.child.exists(_.isInstanceOf[SortExec])) + }, + StopStream + ) + } + } + + test("map state can emit an aggregate every other input") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, clock) = createStringMemoryStream(numPartitions = 2) + val result = input.toDS() + .groupByKey(identity) + .transformWithState( + new RealTimeRunningCountStatefulProcessor(2L), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, "a", "b"), + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock), + WaitUntilBatchProcessed(0), + CheckAnswerWithTimeout(60.seconds.toMillis), + Execute { _ => waitForNextBatchToStart() }, + AddData(input, "c", "a", "a", "c"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 2L), ("c", 2L)), + StopStream + ) + } + } + + test("expires TTL state while the RTM batch remains open") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> "1") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + // Processing another key runs the periodic TTL cleanup without closing the RTM batch. + AddData(input, ("b", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L), ("b", 1L)), + AddData(input, ("a", 2)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("b", 1L), ("a", 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.customMetrics + .get("numValuesRemovedDueToTTLExpiry") > 0L) + }, + StopStream + ) + } + } + + test("cleans up expired TTL state at the end of an RTM batch") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> + "86400000") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L)), + advanceClock(clock, defaultTrigger.batchDurationMs + 10001L), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.numRowsTotal == 0L) + assert(batch0.stateOperators.head.customMetrics + .get("numValuesRemovedDueToTTLExpiry") == 1L) + }, + StopStream + ) + } + } + + test("expires value, map, and list state while the RTM batch remains open") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[String](1) + val result = input.toDS() + .groupByKey(identity) + .transformWithState( + new MultiStatefulVariableTTLProcessor(TTLConfig(Duration.ofSeconds(10))), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, "a", "b"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L), ("b", 1L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, "c"), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("b", 1L), ("c", 1L)), + AddData(input, "a"), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("b", 1L), ("c", 1L), ("a", 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.customMetrics + .get("numValuesRemovedDueToTTLExpiry") >= 6L) + }, + StopStream + ) + } + } + + test("hides expired TTL state before the periodic cleanup scan") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> + "86400000") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, ("a", 2)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L), ("a", 1L)), + StopStream + ) + } + } + + test("cleans up every entry in an expired ListState") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> "1") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeListTTLProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1), ("b", 10)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 3L), ("b", 3L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, ("c", 20)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 3L), ("b", 3L), ("c", 3L)), + AddData(input, ("a", 30)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 3L), ("b", 3L), ("c", 3L), ("a", 3L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.customMetrics + .get("numValuesRemovedDueToTTLExpiry") >= 6L) + }, + StopStream + ) + } + } + + test("applies a shorter TTL after an RTM checkpoint restart") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> "1") { + withTempDir { checkpointDir => + val (input, clock) = createMemoryStream(numPartitions = 1) + val checkpoint = checkpointDir.getCanonicalPath + val original = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLCountProcessor(Duration.ofSeconds(400)), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(original, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + StopStream + ) + + val reduced = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLCountProcessor(Duration.ofSeconds(10)), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(reduced, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + AddData(input, ("a", 2)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 2L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, ("b", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 2L), ("b", 1L)), + AddData(input, ("a", 3)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 2L), ("b", 1L), ("a", 1L)), + StopStream + ) + } + } + } + + test("runs transformWithState across multiple state store partitions") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "3") { + val (input, _) = createMemoryStream(numPartitions = 3) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimePartitionProcessor, + TimeMode.None(), + OutputMode.Update()) + val sink = new ContinuousMemorySink() + val rows = (0 until 64).map(i => (s"key-$i", i)) + + testStream(result, OutputMode.Update(), sink = sink)( + StartStream(), + AddData(input, rows: _*), + Execute { _ => + eventually(timeout(60.seconds)) { + assert(sink.allData.size == rows.size) + assert(sink.allData.map(_.getInt(1)).distinct.size > 1) + } + }, + StopStream + ) + } + } + + test("runs transformWithState after a union") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val leftInput = LowLatencyMemoryStream[(String, Int)](1) + val rightInput = LowLatencyMemoryStream[(String, Int)](1) + + val result = leftInput.toDS() + .union(rightInput.toDS()) + .groupByKey(_._1) + .transformWithState( + new RealTimeEagerCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(leftInput, ("a", 1)), + AddData(rightInput, ("a", 2), ("b", 1)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("a", 2L), ("b", 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + AddData(leftInput, ("a", 3)), + AddData(rightInput, ("b", 2)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", 1L), + ("a", 2L), + ("b", 1L), + ("a", 3L), + ("b", 2L)), + Execute { q => + val operators = q.lastExecution.executedPlan.collect { + case transform: TransformWithStateExec => transform + } + assert(operators.size == 1, q.lastExecution.executedPlan) + assert(operators.head.isRealTimeMode) + }, + StopStream + ) + } + } + + test("runs transformWithState after real-time deduplication") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, _) = createMemoryStream(numPartitions = 2) + val result = input.toDS() + .dropDuplicates() + .groupByKey(_._1) + .transformWithState( + new RealTimeEagerCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1), ("b", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L), ("b", 1L)), + AddData(input, ("a", 1), ("a", 1), ("c", 1)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 1L), ("b", 1L), ("c", 1L)), + StopStream + ) + } + } + + test("fires processing-time timers from a later row while the RTM batch remains open") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeProcessingTimerProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "data")), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, ("b", 1)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", "data"), ("b", "data"), ("a", "timer")), + StopStream + ) + } + } + + test("reuses one timer iterator task listener across RTM input rows") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val (input, _) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTimerIteratorListenerProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1), ("a", 2), ("a", 3)), + // The first timer scan adds the reusable iterator's completion listener. Later scans + // refresh that iterator and must not add another listener. + CheckAnswerWithTimeout( + 60.seconds.toMillis, + (1, false), + (2, true), + (3, false)), + StopStream + ) + } + } + + test("reuses one TTL iterator task listener across RTM cleanup scans") { + val changelogKey = + s"${RocksDBConf.ROCKSDB_SQL_CONF_NAME_PREFIX}.changelogCheckpointing.enabled" + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STREAMING_TRANSFORM_WITH_STATE_REAL_TIME_MODE_TTL_EVICTION_INTERVAL_MS.key -> "0", + changelogKey -> "true") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeTTLIteratorListenerProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, (1, 0)), + new ExternalAction { + override def runAction(): Unit = clock.advance(1L) + }, + AddData(input, ("a", 2)), + // The first timer and TTL scans each add one reusable iterator listener. + CheckAnswerWithTimeout(60.seconds.toMillis, (1, 0), (2, 2)), + new ExternalAction { + override def runAction(): Unit = clock.advance(1L) + }, + AddData(input, ("a", 3)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, (1, 0), (2, 2), (3, 0)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + StopStream + ) + } + } + + test("expired timer receives the current RTM processing time") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + GlobalSingletonManualClock.reset() + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeProcessingTimerValueProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "data", 0L)), + new ExternalAction { + override def runAction(): Unit = clock.advance(10001L) + }, + AddData(input, ("b", 1)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 0L), + ("b", "data", 10001L), + ("a", "timer", 10001L)), + StopStream + ) + } + } + + test("fires remaining processing-time timers at the end of an RTM batch") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeProcessingTimerProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, ("a", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "data")), + advanceClock(clock), + WaitUntilBatchProcessed(0), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", "data"), ("a", "timer")), + StopStream + ) + } + } + + test("processing time timers expire after multiple batches") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val (input, clock) = createStringMemoryStream() + + val processor = new RTMStatefulProcessorWithProcTimeTimer( + defaultTrigger.batchDurationMs * 3) + + val result = input + .toDS() + .groupByKey(x => x) + .transformWithState(processor, TimeMode.ProcessingTime(), OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, "a"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "1")), + advanceClock(clock), + WaitUntilBatchProcessed(0), + // In real time mode, batches execute for a fixed amount of time. Wait for the next + // batch's tasks before advancing a manual clock again to avoid skipping its end time. + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock), + WaitUntilBatchProcessed(1), + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock), + WaitUntilBatchProcessed(2), + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock, 1L), + AddData(input, "a"), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", "1"), ("a", "1"), ("a", "-1")), + StopStream + ) + } + } + + test("processing time timers single key multiple registered timers") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, clock) = createStringMemoryStream() + + val processor = new RTMStatefulProcessorWithProcTimeTimerWithMultipleTimers(30000) + + val result = input + .toDS() + .groupByKey(x => x) + .transformWithState(processor, TimeMode.ProcessingTime(), OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, "a"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "1")), + advanceClock(clock, 31001L), + AddData(input, "a"), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "1"), ("a", "2"), ("a", "-1"), ("a", "-1")), + StopStream + ) + } + } + + test("processing time timers with timers from multiple keys") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, clock) = createStringMemoryStream() + + val processor = new RTMStatefulProcessorWithProcTimeTimer(30000) + + val result = input + .toDS() + .groupByKey(x => x) + .transformWithState(processor, TimeMode.ProcessingTime(), OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, "a"), + AddData(input, "b"), + AddData(input, "c"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "1"), ("b", "1"), ("c", "1")), + advanceClock(clock, 30001L), + AddData(input, "a"), + AddData(input, "b"), + AddData(input, "c"), + CheckAnswerRowsContainsWithTimeout( + 60.seconds.toMillis, + ("a", "1"), + ("b", "1"), + ("c", "1"), + ("a", "-1"), + ("b", "2"), + ("c", "-1") + ), + StopStream + ) + } + } + + test("processing time timers with an integer key") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[Int](2) + val result = input.toDS() + .groupByKey(identity) + .transformWithState( + new RTMStatefulProcessorWithProcTimeTimerInputInt(30000L), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, 1), + CheckAnswerWithTimeout(60.seconds.toMillis, (1, 1L)), + advanceClock(clock, 30001L), + AddData(input, 1), + CheckAnswerRowsContainsWithTimeout(60.seconds.toMillis, (1, 1L), (1, -1L)), + StopStream + ) + } + } + + test("processing time timers with a product key") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Int, String)](2) + val result = input.toDS() + .groupByKey(identity) + .transformWithState( + new RTMStatefulProcessorWithProcTimeTimerInputTuple(30000L), + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, (1, "a")), + CheckAnswerWithTimeout(60.seconds.toMillis, (1, "a", 1L)), + advanceClock(clock, 30001L), + AddData(input, (1, "a")), + CheckAnswerRowsContainsWithTimeout(60.seconds.toMillis, (1, "a", -1L)), + StopStream + ) + } + } + + test("processing time timers survive an RTM checkpoint restart") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + withTempDir { checkpointDir => + val (input, clock) = createStringMemoryStream() + val processor = new RTMStatefulProcessorWithProcTimeTimer( + defaultTrigger.batchDurationMs + 1000L) + val result = input + .toDS() + .groupByKey(x => x) + .transformWithState(processor, TimeMode.ProcessingTime(), OutputMode.Update()) + val checkpoint = checkpointDir.getCanonicalPath + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + AddData(input, "a"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "1")), + advanceClock(clock), + WaitUntilBatchProcessed(0), + StopStream + ) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock, 1001L), + AddData(input, "b"), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "-1"), ("b", "1")), + StopStream + ) + } + } + } + + test("uses fixed between-batch watermarks for incremental event-time timers") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Timestamp, String)](1) + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState( + new RealTimeEventTimerProcessor, + TimeMode.EventTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, + (new Timestamp(100L), "a"), + (new Timestamp(200L), "a"), + (new Timestamp(150L), "a"), + (new Timestamp(300L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + // Batch 1 uses batch 0's fixed eviction watermark (300 - 10 = 290). The first + // input row scans timers already below it without waiting for the batch to end. + AddData(input, + (new Timestamp(280L), "a"), + (new Timestamp(400L), "a")), + // The watermark does not move inside batch 1. A newly registered timer at 280 fires + // immediately against the same fixed 290 watermark. + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "data", 290L), + ("a", "timer", 290L), + ("a", "data", 290L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + Execute { _ => waitForNextBatchToStart() }, + // Batch 2 uses eviction watermark 390 and late-events watermark 290. The late 280 row + // still triggers the per-row timer scan even though it is not passed to the processor. + AddData(input, (new Timestamp(280L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "data", 290L), + ("a", "data", 290L), + ("a", "timer", 390L)), + // The watermark stays fixed while the accepted row is processed in the same batch. + AddData(input, (new Timestamp(500L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "data", 0L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "timer", 290L), + ("a", "data", 290L), + ("a", "data", 290L), + ("a", "timer", 390L), + ("a", "data", 390L)), + advanceClock(clock), + WaitUntilBatchProcessed(2), + Execute { q => + val batch2 = q.recentProgress.find(_.batchId == 2).getOrElse { + fail(s"batch 2 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch2.stateOperators.length == 1) + assert(batch2.stateOperators.head.numRowsDroppedByWatermark == 1L) + }, + StopStream + ) + } + } + + test("fires a newly expired event-time timer from its RTM input row") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Timestamp, String)](1) + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState( + new RealTimeEventTimerProcessor, + TimeMode.EventTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, (new Timestamp(300L), "a")), + CheckAnswerWithTimeout(10.seconds.toMillis, ("a", "data", 0L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + // Batch 1's fixed eviction watermark is 290, so the timer registered for this row is + // scanned immediately after the row is processed. + AddData(input, (new Timestamp(280L), "a")), + CheckAnswerWithTimeout( + 10.seconds.toMillis, + ("a", "data", 0L), + ("a", "data", 290L), + ("a", "timer", 290L)), + StopStream + ) + } + } + + test("fires event-time timers in the final scan of an empty RTM batch") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Timestamp, String)](1) + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState( + new RealTimeEventTimerProcessor, + TimeMode.EventTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, + (new Timestamp(100L), "expired"), + (new Timestamp(300L), "watermark")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("expired", "data", 0L), + ("watermark", "data", 0L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + // Batch 1 has no input rows. Its final scan uses batch 0's fixed eviction watermark. + advanceClock(clock), + WaitUntilBatchProcessed(1), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("expired", "data", 0L), + ("watermark", "data", 0L), + ("expired", "timer", 290L)), + Execute { q => + val batch1 = q.recentProgress.find(_.batchId == 1).getOrElse { + fail(s"batch 1 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch1.numInputRows == 0L) + }, + StopStream + ) + } + } + + test("recovers event-time timers and fires them from the next row after restart") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + withTempDir { checkpointDir => + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Timestamp, String)](1) + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState( + new RealTimeEventTimerProcessor, + TimeMode.EventTime(), + OutputMode.Update()) + val checkpoint = checkpointDir.getCanonicalPath + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + AddData(input, + (new Timestamp(100L), "expired"), + (new Timestamp(300L), "watermark")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("expired", "data", 0L), + ("watermark", "data", 0L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + StopStream + ) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpoint), + Execute { _ => waitForNextBatchToStart() }, + // The recovered batch watermark is 300 - 10 = 290. The next input row scans the + // recovered timers against that fixed watermark. + AddData(input, (new Timestamp(400L), "trigger")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("trigger", "data", 290L), + ("expired", "timer", 290L)), + StopStream + ) + } + } + } + + test("supports an output event-time column in RTM") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream.singlePartition[(Timestamp, String)] + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "1 minute") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState[RealTimeEventTimeOutputRow]( + new RealTimeEventTimeOutputProcessor, + "outputEventTime", + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, + (new Timestamp(1000000L), "a"), + (new Timestamp(2000000L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + RealTimeEventTimeOutputRow("a", new Timestamp(1000000L), 1L), + RealTimeEventTimeOutputRow("a", new Timestamp(2000000L), 2L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + AddData(input, (new Timestamp(3000000L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + RealTimeEventTimeOutputRow("a", new Timestamp(1000000L), 1L), + RealTimeEventTimeOutputRow("a", new Timestamp(2000000L), 2L), + RealTimeEventTimeOutputRow("a", new Timestamp(3000000L), 3L)), + StopStream + ) + } + } + + test("uses the fixed prior-batch watermark for output event-time validation") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream.singlePartition[(Timestamp, String)] + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState[RealTimeEventTimeOutputRow]( + new RealTimeEventTimeOutputProcessor(Some(new Timestamp(1L))), + "outputEventTime", + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + AddData(input, (new Timestamp(300L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + RealTimeEventTimeOutputRow("a", new Timestamp(1L), 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + AddData(input, (new Timestamp(400L), "a")), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + RealTimeEventTimeOutputRow("a", new Timestamp(1L), 1L), + RealTimeEventTimeOutputRow("a", new Timestamp(1L), 2L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + Execute { _ => waitForNextBatchToStart() }, + // Batch 2's late-events watermark is fixed at 290. Its accepted input would emit an + // output timestamp of 1, which must be rejected against that same fixed watermark. + AddData(input, (new Timestamp(500L), "a")), + ExpectFailure[SparkRuntimeException] { error => + checkError( + error.asInstanceOf[SparkThrowable], + "EMITTING_ROWS_OLDER_THAN_WATERMARK_NOT_ALLOWED", + parameters = Map( + "currentWatermark" -> "290", + "emittedRowEventTime" -> "1000")) + } + ) + } + } + + test("processing-time timer registered from initial state uses the batch timestamp") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + GlobalSingletonManualClock.reset() + val (input, executorClock) = createStringMemoryStream(numPartitions = 1) + val driverClock = new StreamManualClock(100000L) + val result = input.toDS() + .groupByKey(identity) + .transformWithState( + new RealTimeInitialStateProcTimerWithExpiryProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update(), + Seq("a").toDS().groupByKey(identity)) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(triggerClock = driverClock), + WaitUntilBatchProcessed(0), + Execute { _ => waitForNextBatchToStart() }, + // The initial-state timer is based on batch 0's driver timestamp (100000), not the + // executor clock (0), so advancing the executor by only the timer delay must not fire it. + advanceClock(executorClock, 5001L), + Execute { _ => input.addData(Seq("b")) }, + CheckAnswerWithTimeout(60.seconds.toMillis, ("b", "1")), + advanceClock(executorClock, 100000L), + Execute { _ => input.addData(Seq("c")) }, + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("b", "1"), ("c", "1"), ("a", "105000")), + StopStream + ) + } + } + + test("initial-state bootstrap batch does not use pipelined shuffle") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val (input, _) = createMemoryStream(numPartitions = 2) + val result = transformWithInitialState(input, Seq("a" -> 2L, "b" -> 5L)) + + RealTimeInitialStateBootstrapBlock.enable() + try { + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(), + Execute { q => + try { + assert(RealTimeInitialStateBootstrapBlock.awaitTaskStart()) + val execution = q.lastExecution + assert(execution.currentBatchId == 0L) + val plan = execution.executedPlan + val scans = plan.collect { case scan: RealTimeStreamScanExec => scan } + assert(scans.nonEmpty) + assert(scans.forall(_.batchDurationMs == 0L)) + + val shuffles = plan.collect { case exchange: ShuffleExchangeExec => exchange } + assert(shuffles.nonEmpty) + assert(shuffles.forall(!_.pipelined)) + } finally { + RealTimeInitialStateBootstrapBlock.disable() + } + }, + WaitUntilBatchProcessed(0), + StopStream + ) + } finally { + RealTimeInitialStateBootstrapBlock.disable() + } + } + } + + Seq( + "non-empty" -> Seq("a" -> 2L, "b" -> 5L), + "empty" -> Seq.empty[(String, Long)] + ).foreach { case (description, initialValues) => + test(s"hydrates $description initial state in a finite batch and recovers it") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + withTempDir { checkpointDir => + val (input, clock) = createMemoryStream(numPartitions = 2) + val result = transformWithInitialState(input, initialValues) + val initialCount = initialValues.toMap.getOrElse("a", 0L) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + // Input available when the query starts must wait until initial state is durable. + AddData(input, ("a", 1)), + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.numInputRows == 0) + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.numRowsTotal == initialValues.size) + }, + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", initialCount + 1L)), + Execute { q => + val plan = q.lastExecution.executedPlan + val scans = plan.collect { case scan: RealTimeStreamScanExec => scan } + assert(scans.nonEmpty) + assert(scans.forall(_.batchDurationMs == defaultTrigger.batchDurationMs)) + val streamingShuffles = plan.collect { + case exchange: ShuffleExchangeExec if exchange.exists { + case _: RealTimeStreamScanExec => true + case _ => false + } => exchange + } + assert(streamingShuffles.nonEmpty) + assert(streamingShuffles.forall(_.pipelined)) + }, + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream + ) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + AddData(input, ("a", 2)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", initialCount + 2L)), + advanceClock(clock), + WaitUntilBatchProcessed(2), + StopStream + ) + } + } + } + } + + test("hydrates event-time initial state before the first RTM input batch") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val clock = new GlobalSingletonManualClock() + LowLatencyClock.setClock(clock) + val input = LowLatencyMemoryStream[(Timestamp, String)](1) + val result = input.toDF() + .select(col("_1").as("eventTime"), col("_2").as("key")) + .withWatermark("eventTime", "10 milliseconds") + .as[(Timestamp, String)] + .groupByKey(_._2) + .transformWithState( + new RealTimeEventInitialCountProcessor, + TimeMode.EventTime(), + OutputMode.Update(), + initialState(Seq("a" -> 2L))) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + AddData(input, (new Timestamp(100L), "a")), + StartStream(), + WaitUntilBatchProcessed(0), + Execute { q => + val batch0 = q.recentProgress.find(_.batchId == 0).getOrElse { + fail(s"batch 0 progress was not retained: ${q.recentProgress.toSeq}") + } + assert(batch0.numInputRows == 0) + assert(batch0.stateOperators.length == 1) + assert(batch0.stateOperators.head.numRowsTotal == 1) + }, + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 3L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream + ) + } + } + + test("processes non-contiguous duplicate initial-state keys without sorting") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + val (input, clock) = createMemoryStream(numPartitions = 1) + // A single input partition preserves a, b, a through the one-partition hash exchange. + val duplicateInitialState = Seq("a" -> 1L, "b" -> 5L, "a" -> 2L) + .toDS() + .coalesce(1) + .groupByKey(_._1) + .mapValues(_._2) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeInitialCountProcessor, + TimeMode.None(), + OutputMode.Update(), + duplicateInitialState) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + AddData(input, ("a", 1)), + StartStream(), + WaitUntilBatchProcessed(0), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 4L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream + ) + } + } + + test("retries initial-state bootstrap from batch 0 after failure") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + withTempDir { checkpointDir => + val (input, clock) = createMemoryStream(numPartitions = 2) + val result = transformWithInitialState(input, Seq("a" -> 2L, "b" -> 5L)) + val offsetFile = new java.io.File(checkpointDir, "offsets/0") + + try { + RealTimeInitialStateFailure.enabled = true + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + ExpectFailure[SparkException] { error => + assert(error.getMessage.contains("injected initial-state bootstrap failure")) + } + ) + assert(!offsetFile.exists()) + + RealTimeInitialStateFailure.enabled = false + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + AddData(input, ("a", 1)), + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + WaitUntilBatchProcessed(0), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 3L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream + ) + } finally { + RealTimeInitialStateFailure.enabled = false + } + } + } + } + + Seq(true, false).foreach { changelogCheckpointingEnabled => + test(s"recovers transformWithState state with RocksDB changelog checkpointing " + + s"enabled=$changelogCheckpointingEnabled") { + val changelogKey = + s"${RocksDBConf.ROCKSDB_SQL_CONF_NAME_PREFIX}.changelogCheckpointing.enabled" + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + changelogKey -> changelogCheckpointingEnabled.toString) { + withTempDir { checkpointDir => + val (input, clock) = createMemoryStream(numPartitions = 2) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeEagerCountProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + AddData(input, ("a", 1), ("a", 2)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 1L), ("a", 2L)), + Execute { q => + assert(q.sparkSessionForStream.conf.get(changelogKey) == + changelogCheckpointingEnabled.toString) + }, + advanceClock(clock), + WaitUntilBatchProcessed(0), + StopStream + ) + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getCanonicalPath), + AddData(input, ("a", 3), ("b", 1)), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 3L), ("b", 1L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream + ) + } + } + } + } + + test("keeps transformWithState state across MBM to RTM to MBM restarts") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + withTempDir { checkpointDir => + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeEagerCountProcessor, + TimeMode.None(), + OutputMode.Update()) + val checkpoint = checkpointDir.getCanonicalPath + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + AddData(input, ("a", 1), ("a", 2)), + StartStream( + trigger = Trigger.ProcessingTime(1000), + checkpointLocation = checkpoint), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", 2L)), + WaitUntilBatchProcessed(0), + Execute { q => + val operators = q.lastExecution.executedPlan.collect { + case transform: TransformWithStateExec => transform + } + assert(operators.size == 1) + assert(!operators.head.isRealTimeMode) + }, + StopStream, + StartStream(trigger = defaultTrigger, checkpointLocation = checkpoint), + AddData(input, ("a", 3), ("b", 1)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, ("a", 2L), ("a", 3L), ("b", 1L)), + Execute { q => + val operators = q.lastExecution.executedPlan.collect { + case transform: TransformWithStateExec => transform + } + assert(operators.size == 1) + assert(operators.head.isRealTimeMode) + }, + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream, + AddData(input, ("a", 4), ("b", 2)), + StartStream( + trigger = Trigger.ProcessingTime(1000), + checkpointLocation = checkpoint), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", 2L), ("a", 3L), ("b", 1L), ("a", 4L), ("b", 2L)), + WaitUntilBatchProcessed(2), + Execute { q => + val operators = q.lastExecution.executedPlan.collect { + case transform: TransformWithStateExec => transform + } + assert(operators.size == 1) + assert(!operators.head.isRealTimeMode) + }, + StopStream + ) + } + } + } + + test("keeps MapState TTL and processing-time timers across MBM to RTM to MBM restarts") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + withTempDir { checkpointDir => + GlobalSingletonManualClock.reset() + val (input, clock) = createMemoryStream(numPartitions = 1) + val result = input.toDS() + .groupByKey(_._1) + .transformWithState( + new RealTimeMapTTLAndTimerProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + val checkpoint = checkpointDir.getCanonicalPath + val timerDelayAcrossRtmBoundary = (defaultTrigger.batchDurationMs + 1000L).toInt + + testStream(result, OutputMode.Update(), sink = new ContinuousMemorySink())( + AddData(input, ("a", 10000)), + StartStream( + trigger = Trigger.ProcessingTime(1000), + triggerClock = clock, + checkpointLocation = checkpoint), + CheckAnswerWithTimeout(60.seconds.toMillis, ("a", "data", 1L)), + WaitUntilBatchProcessed(0), + StopStream, + StartStream(trigger = defaultTrigger, checkpointLocation = checkpoint), + Execute { _ => waitForNextBatchToStart() }, + advanceClock(clock, 10001L), + AddData(input, ("a", timerDelayAcrossRtmBoundary)), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 1L), + ("a", "data", 2L), + ("a", "timer", 10000L)), + advanceClock(clock), + WaitUntilBatchProcessed(1), + StopStream, + advanceClock(clock, 1001L), + AddData(input, ("a", 10000)), + StartStream( + trigger = Trigger.ProcessingTime(1000), + triggerClock = clock, + checkpointLocation = checkpoint), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 1L), + ("a", "data", 2L), + ("a", "timer", 10000L), + ("a", "data", 3L), + ("a", "timer", 311001L)), + WaitUntilBatchProcessed(2), + StopStream, + advanceClock(clock, Duration.ofMinutes(10).toMillis + 1L), + AddData(input, ("a", 10000)), + StartStream( + trigger = Trigger.ProcessingTime(1000), + triggerClock = clock, + checkpointLocation = checkpoint), + CheckAnswerWithTimeout( + 60.seconds.toMillis, + ("a", "data", 1L), + ("a", "data", 2L), + ("a", "timer", 10000L), + ("a", "data", 3L), + ("a", "timer", 311001L), + ("a", "data", 1L), + ("a", "timer", 321002L)), + WaitUntilBatchProcessed(3), + StopStream + ) + } + } + } +} + +@SlowSQLTest +class RealTimeTransformWithStateSuiteWithRowChecksum + extends RealTimeTransformWithStateSuite with EnableStateStoreRowChecksum diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala index a9b8d54a6c61b..ee101cbede9d1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala @@ -70,7 +70,6 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { "errorType" -> "operator", "message" -> ( "org.apache.spark.sql.execution.SortExec, " + - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec, " + "org.apache.spark.sql.execution.joins.SortMergeJoinExec are" ) ) @@ -107,41 +106,18 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { } } - // TODO(SPARK-54237) : Remove this test after RTM can shuffle to multiple stages - test("repartition not allowed") { - val inputData = LowLatencyMemoryStream[Int](2) - - val df = inputData.toDF() - .select(col("value").as("key")) - .repartition(4, col("key")) - - val query = runStreamingQuery("repartition_allowlist", df) - - eventually(timeout(60.seconds)) { - checkError( - exception = query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException], - condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", - parameters = Map( - "errorType" -> "operator", - "message" -> ( - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec is" - ) - ) - ) - } - } - - // TODO(SPARK-54236) : Remove this test after RTM supports stateful queries - test("stateful queries not allowed") { + // A repartitionByRange produces a range-partitioned shuffle. Building its RangePartitioner runs a + // separate job that samples the input to compute range bounds, which cannot complete while the + // source keeps producing, so it must be rejected up front rather than stalling the query. + test("range-partitioned shuffle not allowed") { val inputData = LowLatencyMemoryStream[Int](2) val df = inputData.toDF() .select(col("value").as("key")) - .groupBy(col("key")) - .count() - .select(concat(col("key"), lit("-"), col("count"))) + .repartitionByRange(3, col("key")) + .select(col("key")) - val query = runStreamingQuery("repartition_allowlist", df) + val query = runStreamingQuery("range_shuffle_allowlist", df) eventually(timeout(60.seconds)) { checkError( @@ -149,15 +125,16 @@ class StreamRealTimeModeAllowlistSuite extends StreamRealTimeModeE2ESuiteBase { condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", parameters = Map( "errorType" -> "operator", - "message" -> ( - "org.apache.spark.sql.execution.aggregate.HashAggregateExec, " + - "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec, " + - "org.apache.spark.sql.execution.streaming" + - ".operators.stateful.StateStoreRestoreExec, " + - "org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec are" - ) + "message" -> "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec is" ) ) } } + + // The "stateful queries not allowed" test that used to live here asserted that a streaming + // aggregation was rejected, because it planned into the micro-batch aggregation operators and + // HashAggregateExec is not allowlisted. A Real-Time Mode aggregation is now planned as the + // streamline aggregate operator, which is allowlisted, so there is no rejection left to assert. + // StreamlineStreamingAggregationRealTimeSuite covers the aggregation itself end to end, and the + // generic operator-allowlist test above still guards the operators that remain unsupported. } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeCoexistenceSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeCoexistenceSuite.scala new file mode 100644 index 0000000000000..9b1e8b9516d92 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeCoexistenceSuite.scala @@ -0,0 +1,255 @@ +/* + * 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.streaming + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.concurrent.duration._ + +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart, SparkListenerStageCompleted, SparkListenerStageSubmitted} +import org.apache.spark.sql.{ForeachWriter, Row} +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution, + StreamingQueryWrapper} +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, + LowLatencyMemoryStream} +import org.apache.spark.sql.functions.count +import org.apache.spark.sql.internal.SQLConf + +/** + * Tests that Real-Time Mode (RTM) and MicroBatch Mode (MBM) multi-stage queries can run in the + * SAME cluster -- concurrently, in one SparkContext, sharing one set of executors. + * + * This is the coexistence property that makes RTM usable without dedicating a cluster to it. It + * holds because a shuffle's implementation is chosen by DEPENDENCY TYPE, not by a cluster-wide + * setting: `SparkEnv.shuffleManagerFor` routes a `PipelinedShuffleDependency` to the pipelined + * manager (`spark.shuffle.manager.incremental`, the streaming shuffle) and every other + * `ShuffleDependency` to the blocking manager (`spark.shuffle.manager`, sort shuffle). Both + * managers are instantiated in the same JVM and neither query has to know about the other. + * + * Each test uses MULTI-STAGE queries on both sides, since a single-stage query has no shuffle and + * so would not exercise routing at all. The RTM side is verified to be genuinely pipelined (rather + * than merely running) by asserting on the `pipelined` flag of its exchanges, and the MBM side is + * verified to be genuinely NOT pipelined -- a test that only checked both queries produced answers + * would pass even if routing collapsed to a single manager. + */ +class StreamRealTimeModeCoexistenceSuite extends StreamRealTimeModeSuiteBase { + + import testImplicits._ + + /** Every shuffle exchange in the query's last executed plan, with its `pipelined` flag. */ + private def exchangePipelinedFlags(q: StreamExecution): Seq[Boolean] = + q.lastExecution.executedPlan.collect { case s: ShuffleExchangeExec => s.pipelined } + + /** Asserts the query has at least one shuffle and every one of them is pipelined. */ + private def assertAllExchangesPipelined(q: StreamExecution): Unit = { + val flags = exchangePipelinedFlags(q) + assert(flags.nonEmpty, "expected at least one shuffle exchange in the RTM plan") + assert(flags.forall(identity), + s"expected every RTM exchange to be pipelined, got: ${flags.mkString(", ")}") + } + + /** Asserts the query has at least one shuffle and none of them is pipelined. */ + private def assertNoExchangePipelined(q: StreamExecution): Unit = { + val flags = exchangePipelinedFlags(q) + assert(flags.nonEmpty, "expected at least one shuffle exchange in the MBM plan") + assert(!flags.exists(identity), + s"expected no MBM exchange to be pipelined, got: ${flags.mkString(", ")}") + } + + test("an RTM and an MBM multi-stage query run concurrently in the same cluster") { + // Keep both queries' shuffles small: the RTM query's whole pipelined group is gang-admitted, so + // its scan + dedup tasks and the MBM query's tasks must all fit the cluster's slots at once. + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val rtmInput = LowLatencyMemoryStream[(String, Int)] + val mbmInput = MemoryStream[(String, Int)] + + // Both are multi-stage: a shuffle (repartition by key) feeding a stateful dedup. + val rtmQuery = rtmInput.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + val mbmQuery = mbmInput.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + + // Start the MBM query first and leave it running for the whole RTM test. + val mbmHandle = mbmQuery.writeStream + .format("memory") + .queryName("coexistence_mbm") + .outputMode(OutputMode.Update) + .start() + + try { + mbmInput.addData(("a", 1), ("b", 1), ("a", 2)) + mbmHandle.processAllAvailable() + checkAnswer(spark.table("coexistence_mbm"), Seq(Row("a"), Row("b"))) + + val mbmExec = mbmHandle.asInstanceOf[StreamingQueryWrapper].streamingQuery + assertNoExchangePipelined(mbmExec) + + // With the MBM query still active, run an RTM query in the same context. + testStream(rtmQuery, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(rtmInput, ("x", 1), ("y", 1), ("x", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "x", "y"), + Execute { q => + assertAllExchangesPipelined(q) + assert(mbmHandle.isActive, "the MBM query must still be running alongside RTM") + }, + StopStream + ) + + // The MBM query must still make progress AFTER the RTM query has come and gone, proving the + // pipelined shuffle did not disturb the blocking manager's state. + mbmInput.addData(("c", 1), ("a", 3)) + mbmHandle.processAllAvailable() + checkAnswer(spark.table("coexistence_mbm"), Seq(Row("a"), Row("b"), Row("c"))) + assertNoExchangePipelined( + mbmHandle.asInstanceOf[StreamingQueryWrapper].streamingQuery) + } finally { + mbmHandle.stop() + spark.sql("DROP TABLE IF EXISTS coexistence_mbm") + } + } + } + + test("a batch query with a shuffle runs while an RTM query is active") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val rtmInput = LowLatencyMemoryStream[(String, Int)] + val rtmQuery = rtmInput.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + + // A multi-stage BATCH query: groupBy forces a blocking shuffle. Run it mid-RTM-batch. + val batchResult = new AtomicReference[Seq[Row]](null) + + testStream(rtmQuery, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(rtmInput, ("x", 1), ("y", 1), ("x", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "x", "y"), + Execute { q => + assertAllExchangesPipelined(q) + // While the RTM batch is still open, a regular batch job with its own shuffle must run to + // completion on the same executors, using the blocking shuffle manager. + val df = spark.range(0, 100).selectExpr("id % 5 AS k").groupBy("k").agg(count("*")) + batchResult.set(df.orderBy("k").collect().toSeq) + }, + Execute { _ => + val rows = batchResult.get() + assert(rows != null, "the batch query did not run") + assert(rows.length == 5, s"expected 5 groups, got ${rows.length}") + assert(rows.forall(_.getLong(1) == 20L), s"expected 20 rows per group, got $rows") + }, + StopStream + ) + } + } + + test("two RTM queries run concurrently, each with its own pipelined group") { + // Two independent pipelined groups must be admitted and co-scheduled at the same time. Keep the + // partition counts low so both groups' gang demands fit the cluster together. + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val inputA = LowLatencyMemoryStream[(String, Int)] + val inputB = LowLatencyMemoryStream[(String, Int)] + + val queryA = inputA.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + val queryB = inputB.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + + // Track, from the driver, whether a stage of query A and a stage of query B were ever RUNNING + // at the same time. Each query's long-running RTM stages are attributed to its query id (via + // the job property StreamExecution tags), and `sawBothRunning` is set whenever both queries + // have at least one stage running simultaneously -- which only happens if the two pipelined + // groups are genuinely co-scheduled rather than one running after the other. + val idA = new AtomicReference[String](null) + val idB = new AtomicReference[String](null) + val stagesA = ConcurrentHashMap.newKeySet[Int]() + val stagesB = ConcurrentHashMap.newKeySet[Int]() + val runningA = ConcurrentHashMap.newKeySet[Int]() + val runningB = ConcurrentHashMap.newKeySet[Int]() + val sawBothRunning = new AtomicBoolean(false) + val listener = new SparkListener { + override def onJobStart(e: SparkListenerJobStart): Unit = { + val qid = Option(e.properties).map(_.getProperty(StreamExecution.QUERY_ID_KEY)).orNull + if (qid != null && qid == idA.get()) e.stageIds.foreach(stagesA.add(_)) + else if (qid != null && qid == idB.get()) e.stageIds.foreach(stagesB.add(_)) + } + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + val sid = e.stageInfo.stageId + if (stagesA.contains(sid)) runningA.add(sid) + else if (stagesB.contains(sid)) runningB.add(sid) + if (!runningA.isEmpty && !runningB.isEmpty) sawBothRunning.set(true) + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + val sid = e.stageInfo.stageId + runningA.remove(sid) + runningB.remove(sid) + } + } + spark.sparkContext.addSparkListener(listener) + + // ForeachWriter is one of the sinks RTM allows (see RealTimeModeAllowlist.allowedSinks); + // ForeachBatch is not, so it cannot be used to drive a second RTM query here. + val handleB = queryB.writeStream + .foreach(new ForeachWriter[Row] { + override def open(partitionId: Long, epochId: Long): Boolean = true + override def process(value: Row): Unit = () + override def close(errorOrNull: Throwable): Unit = () + }) + .queryName("coexistence_rtm_b") + .outputMode(OutputMode.Update) + .trigger(defaultTrigger) + .start() + + try { + idB.set(handleB.id.toString) + eventually(timeout(60.seconds)) { + assert(handleB.isActive, "second RTM query failed to start") + } + val execB = handleB.asInstanceOf[StreamingQueryWrapper].streamingQuery + + testStream(queryA, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputA, ("x", 1), ("y", 1), ("x", 2)), + StartStream(), + // StartStream only launches the query; its id is not known until it is running. Record it + // here, then drive a SECOND batch below so the listener observes a batch of query A whose + // jobs all start after idA is set (the first batch's jobs may fire before this and be + // dropped from stagesA). This mirrors the mitigation in the sibling co-scheduling test. + Execute(q => idA.set(q.id.toString)), + CheckAnswerWithTimeout(60000, "x", "y"), + AddData(inputA, ("z", 1), ("w", 1)), + CheckAnswerWithTimeout(60000, "x", "y", "z", "w"), + Execute { q => + assert(handleB.isActive, "the second RTM query must still be running") + assert(handleB.exception.isEmpty, + s"second RTM query failed: ${handleB.exception.map(_.getMessage).getOrElse("")}") + // Query A's pipelined group is running now. Feed query B so it keeps scheduling its own + // group (an idle RTM query can run empty batches), and wait until the listener has seen + // a stage of A and a stage of B running at the same time. + eventually(timeout(60.seconds)) { + inputB.addData(("p", 1), ("q", 1)) + assert(sawBothRunning.get(), + "expected a stage of query A and a stage of query B to run concurrently") + } + // Both groups are pipelined shuffles (not materialized). + assertAllExchangesPipelined(q) + assertAllExchangesPipelined(execB) + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + handleB.stop() + } + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeDefaultConfsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeDefaultConfsSuite.scala new file mode 100644 index 0000000000000..65fb2c09a53bc --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeDefaultConfsSuite.scala @@ -0,0 +1,238 @@ +/* + * 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.streaming + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.{SparkIllegalArgumentException, SparkThrowable} +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.execution.streaming.state.{HDFSBackedStateStoreProvider, RocksDBConf, + RocksDBStateStoreProvider} +import org.apache.spark.sql.internal.SQLConf + +/** + * Tests the configuration defaults a Real-Time Mode query applies at query start. + * + * These are not engine-wide defaults because they are only the right choice for a low-latency, + * long-running batch. Each is applied only when the user has not set it, so an explicit choice + * always wins -- both directions are asserted here, since a defaulting block that silently + * overrode a user's setting would be worse than having no default at all. + */ +class StreamRealTimeModeDefaultConfsSuite extends StreamRealTimeModeSuiteBase { + + import testImplicits._ + + private val changelogKey = + s"${RocksDBConf.ROCKSDB_SQL_CONF_NAME_PREFIX}.changelogCheckpointing.enabled" + + /** Runs a trivial RTM query to completion and returns the effective conf values afterwards. */ + private def runRealTimeQueryAndReadConfs(keys: Seq[String]): Map[String, Option[String]] = { + val inputData = LowLatencyMemoryStream[Int] + val mapped = inputData.toDS().map(_ + 1) + var observed: Map[String, Option[String]] = Map.empty + + testStream(mapped, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2, 3), + StartStream(), + CheckAnswerWithTimeout(60000, 2, 3, 4), + Execute { q => + // Read from sparkSessionForStream, the isolated CLONE the batches actually run with + // (StreamExecution.sparkSessionForStream). The defaults are applied there, not to the + // caller's session, so reading `q.sparkSession` would observe the unmodified original. + val conf = q.sparkSessionForStream.conf + observed = keys.map(k => k -> conf.getOption(k)).toMap + }, + StopStream + ) + observed + } + + test("a Real-Time Mode query defaults to checkpoint format v2 with the RocksDB state store") { + val keys = Seq( + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key, + SQLConf.STATE_STORE_PROVIDER_CLASS.key) + val observed = runRealTimeQueryAndReadConfs(keys) + + assert(observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key).contains("2"), + "Real-Time Mode should default the state-store checkpoint format to v2, got " + + observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)) + assert(observed(SQLConf.STATE_STORE_PROVIDER_CLASS.key) + .contains(classOf[RocksDBStateStoreProvider].getName), + "Real-Time Mode should default the state-store provider to RocksDB, got " + + observed(SQLConf.STATE_STORE_PROVIDER_CLASS.key)) + } + + test("the checkpoint-format and provider defaults are applied independently") { + // The two are defaulted by independent guards (matching the runtime): pinning one does not + // suppress the other. Pin only the version (to a compatible v2, since an incompatible explicit + // config is rejected up front by the pre-flight) and confirm the provider is still defaulted. + withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + val observed = runRealTimeQueryAndReadConfs( + Seq(SQLConf.STATE_STORE_PROVIDER_CLASS.key, + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)) + assert(observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key).contains("2"), + "the pinned version must be left alone, got " + + observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)) + assert(observed(SQLConf.STATE_STORE_PROVIDER_CLASS.key) + .contains(classOf[RocksDBStateStoreProvider].getName), + "the provider default is independent of the version, so it must still be RocksDB, got " + + observed(SQLConf.STATE_STORE_PROVIDER_CLASS.key)) + } + } + + test("an explicit checkpoint format version is not overridden by Real-Time Mode") { + // The defaulting guard respects an explicit version rather than forcing 2. v2 is used here + // (not v1) because an explicit v1 is rejected up front by the pre-flight (covered below). + withSQLConf(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "2") { + val observed = runRealTimeQueryAndReadConfs( + Seq(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)) + assert(observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key).contains("2"), + "an explicitly configured checkpoint format version must be left alone, got " + + observed(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key)) + } + } + + test("an explicit incompatible config is rejected up front for a Real-Time Mode query") { + // The pre-flight in StreamingQueryManager rejects explicit RTM-incompatible session configs + // before the query is built: checkpoint format below v2, a non-RocksDB provider, and + // sortBeforeRepartition=true. Each is reported in one SQL_CONFIGURATION_NOT_SUPPORTED error. + def assertRejected(confs: (String, String)*): Unit = { + withSQLConf(confs: _*) { + val inputData = LowLatencyMemoryStream[Int] + val e = intercept[SparkIllegalArgumentException] { + testStream(inputData.toDS(), OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream()) + } + checkError(e, condition = "STREAMING_REAL_TIME_MODE.SQL_CONFIGURATION_NOT_SUPPORTED", + parameters = e.getMessageParameters.asScala.toMap) + } + } + assertRejected(SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") + assertRejected( + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[HDFSBackedStateStoreProvider].getName) + assertRejected(SQLConf.SORT_BEFORE_REPARTITION.key -> "true") + } + + test("the checkpoint-v1 escape hatch bypasses the pre-flight version check") { + // With the escape hatch on, an explicit v1 config is permitted through the pre-flight (the + // existing-v1-checkpoint fail-fast in initializeExecution is covered separately below). + withSQLConf( + SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1.key -> "true", + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1") { + val inputData = LowLatencyMemoryStream[Int] + testStream(inputData.toDS(), OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1), + StartStream(), + CheckAnswerWithTimeout(60000, 1), + StopStream + ) + } + } + + test("a Real-Time Mode query defaults changelog checkpointing on with RocksDB") { + withSQLConf( + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName) { + val observed = runRealTimeQueryAndReadConfs(Seq(changelogKey)) + assert(observed(changelogKey).contains("true"), + s"Real-Time Mode should default $changelogKey to true, got ${observed(changelogKey)}") + } + } + + test("an explicit changelog checkpointing setting is not overridden by Real-Time Mode") { + // The RocksDB provider must be in force, otherwise the `usingRocksDb` gate skips the block and + // the guard under test is never reached -- the assertion would hold vacuously. + withSQLConf( + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName, + changelogKey -> "false") { + val observed = runRealTimeQueryAndReadConfs(Seq(changelogKey)) + assert(observed(changelogKey).contains("false"), + s"an explicitly configured $changelogKey must be left alone, got ${observed(changelogKey)}") + } + } + + test("switching an existing v1 checkpoint to Real-Time Mode fails fast") { + // An existing v1 checkpoint carries no explicit config, so the pre-flight cannot see it; the + // initializeExecution fail-fast catches it instead, from the resolved state-store format. + // Resolution keeps an existing checkpoint at the version it was created with, so rather than + // silently running at v1 (where a re-executed batch can reuse the failed batch's state file + // names and lose data) the query is rejected at start. The rejection is unconditional -- a + // plain (stateless) query is rejected too. + withTempDir { checkpointDir => + val inputData = LowLatencyMemoryStream[Int] + val mapped = inputData.toDS().map(_ + 1) + testStream(mapped, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1), + // Phase 1: a micro-batch trigger, which writes a v1 commit log. + StartStream( + trigger = Trigger.ProcessingTime("1 second"), + checkpointLocation = checkpointDir.getAbsolutePath), + WaitUntilCurrentBatchProcessed, + StopStream, + AddData(inputData, 2), + // Phase 2: the same checkpoint under the Real-Time trigger must be rejected. + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + ExpectFailure[SparkIllegalArgumentException] { e => + checkError( + e.asInstanceOf[SparkThrowable], + condition = "STREAMING_REAL_TIME_MODE.CHECKPOINT_FORMAT_V1_NOT_SUPPORTED", + parameters = Map( + "config" -> SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1.key)) + } + ) + } + } + + test("the escape hatch allows Real-Time Mode on an existing v1 checkpoint") { + withSQLConf( + SQLConf.STREAMING_REAL_TIME_MODE_DANGEROUSLY_ALLOW_CHECKPOINT_V1.key -> "true") { + withTempDir { checkpointDir => + val inputData = LowLatencyMemoryStream[Int] + val mapped = inputData.toDS().map(_ + 1) + testStream(mapped, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1), + StartStream( + trigger = Trigger.ProcessingTime("1 second"), + checkpointLocation = checkpointDir.getAbsolutePath), + WaitUntilCurrentBatchProcessed, + StopStream, + AddData(inputData, 2), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + // The sink accumulates across both phases, so batch 1's row is still present. + CheckAnswerWithTimeout(60000, 2, 3), + StopStream + ) + } + } + } + + test("a fresh Real-Time Mode checkpoint is not rejected") { + // A fresh checkpoint takes state-store format v2 from the Real-Time Mode defaults, so the + // rejection does not apply. Only a resolved state-store format v1 -- from an existing + // incompatible checkpoint or an explicit v1 config without the escape hatch -- is rejected. + withTempDir { checkpointDir => + val inputData = LowLatencyMemoryStream[Int] + testStream(inputData.toDS(), OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, 1, 2), + StopStream + ) + } + } + +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala index 0f584d818c06e..5ca440b244125 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala @@ -17,18 +17,25 @@ package org.apache.spark.sql.streaming -import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.AtomicInteger +import java.io.IOException +import java.util.concurrent.{ConcurrentHashMap, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} import scala.concurrent.duration.Duration +import scala.jdk.CollectionConverters._ import org.scalatest.concurrent.PatienceConfiguration.Timeout -import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} +import org.apache.spark.{SparkException, SparkIllegalArgumentException, SparkIllegalStateException, TaskContext} +import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart, SparkListenerStageCompleted, SparkListenerStageSubmitted} +import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.execution.streaming.RealTimeTrigger -import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamExecution} import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} -import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager, + FailureInjectionFileSystem, RocksDBStateStoreProvider} +import org.apache.spark.sql.functions.{broadcast, concat, lit, udf} import org.apache.spark.sql.internal.SQLConf class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { @@ -209,6 +216,74 @@ class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase { StopStream ) } + + test("pipelined shuffle: a static-side shuffle is not marked pipelined") { + // A broadcast stream-static join can carry a shuffle on its STATIC side, in a subtree with no + // RealTimeStreamScanExec. That shuffle must materialize normally: a static side runs to + // completion rather than streaming, so pulling it into the pipelined group would demand + // concurrent slots for a stage that must instead finish, failing admission. Only shuffles on + // the streaming path are marked, matching what the operator allowlist actually validates. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") { + val staticData = spark.range(0, 3).selectExpr("id as sk", "id * 10 as sv") + .repartition(3, $"sk") + val inputData = LowLatencyMemoryStream[(String, Int)](2) + val streamDf = inputData.toDF().select($"_1".as("key"), $"_2".cast("long").as("sk")) + val result = streamDf + .repartition(2, $"sk") + .join(broadcast(staticData), Seq("sk"), "left") + .select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 2)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + // The streaming-side repartition is pipelined; a static-side shuffle, if it survives + // into this plan, must not be. + val streamingSide = exchanges.filter(_.exists { + case _: RealTimeStreamScanExec => true + case _ => false + }) + assert(streamingSide.nonEmpty, "expected a shuffle on the streaming path") + assert(streamingSide.forall(_.pipelined), + "a streaming-path shuffle must be pipelined") + assert(exchanges.filterNot(streamingSide.contains).forall(!_.pipelined), + "a static-side shuffle must not be marked pipelined") + }, + StopStream + ) + } + } + + test("multiple broadcast joins with the same static table run in Real-Time Mode") { + // Two broadcast joins against the SAME static table -- exchange reuse collapses the second + // broadcast into a ReusedExchangeExec. This is the supported reuse shape in Real-Time Mode: a + // reused BROADCAST exchange (allowlisted), NOT a reused shuffle. The RTM marking rule only + // marks ShuffleExchangeExec, so it correctly leaves the broadcast reuse alone. + val inputData = LowLatencyMemoryStream.singlePartition[Int] + val staticData = Seq((1, "a"), (2, "b"), (3, "c")).toDF("key", "value") + val df = inputData.toDS().toDF("key").join(broadcast(staticData), Seq("key"), "left") + val df2 = df.join(broadcast(staticData), Seq("key"), "left") + testStream(df2, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2, 3), + StartStream(), + CheckAnswerWithTimeout(30000, (1, "a", "a"), (2, "b", "b"), (3, "c", "c")), + Execute { q => + val plan = q.lastExecution.executedPlan + // Verify the shape this test is about, not just the answer: the second broadcast is reused, + // and the marking rule left it alone because it only marks shuffle exchanges. + assert(plan.exists(_.isInstanceOf[ReusedExchangeExec]), + s"expected the second broadcast to be reused, got:\n$plan") + assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "a broadcast stream-static join should introduce no shuffle exchange") + }, + StopStream + ) + } } class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClockSuiteBase { @@ -393,4 +468,393 @@ class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClo StopStream ) } + + test("transformWithState writes batch 0 metadata only after the RTM offset WAL") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName, + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> classOf[RocksDBStateStoreProvider].getName, + SQLConf.SHUFFLE_PARTITIONS.key -> "1") { + withTempDir { checkpointDir => + val injectionState = FailureInjectionFileSystem.registerTempPath(checkpointDir.getPath) + try { + val inputData = LowLatencyMemoryStream[String](1) + val result = inputData.toDS() + .groupByKey(value => value) + .transformWithState( + new RunningCountStatefulProcessor, + TimeMode.ProcessingTime(), + OutputMode.Update()) + val metadataFile = new java.io.File( + checkpointDir, "state/0/_metadata/v2/0") + val stateSchemaDir = new java.io.File( + checkpointDir, "state/0/_stateSchema/default") + + injectionState.failureCreateAtomicRegex = Seq(".*/offsets/0") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, "a"), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, ("a", "1")), + Execute { _ => + assert(Option(stateSchemaDir.listFiles()).exists(_.nonEmpty)) + assert(!metadataFile.exists()) + }, + advanceRealTimeClock, + ExpectFailure[IOException]() + ) + assert(!metadataFile.exists()) + + injectionState.failureCreateAtomicRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, ("a", "1")), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + StopStream + ) + assert(metadataFile.exists()) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(checkpointDir.getPath) + } + } + } + } + + // ======================================================================================== + // Pipelined (streaming) shuffle: a stateful/repartition Real-Time Mode query whose shuffle is a + // PipelinedShuffleDependency, so the producer (source scan) and consumer stages are co-scheduled + // and stream records through a transient shuffle instead of the consumer waiting for the producer + // to fully materialize. + // ======================================================================================== + + override def beforeEach(): Unit = { + super.beforeEach() + StreamRealTimeModeSuite.failTasks = false + } + + /** Assert every shuffle exchange in the query's last executed plan is pipelined. */ + private def assertAllExchangesPipelined(q: StreamExecution): Unit = { + val exchanges = q.lastExecution.executedPlan.collect { case s: ShuffleExchangeExec => s } + assert(exchanges.nonEmpty, "expected at least one shuffle exchange in the plan") + assert(exchanges.forall(_.pipelined), + "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " + + exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", ")) + } + + test("pipelined shuffle: stateful dedup runs in Real-Time Mode and co-schedules its stages") { + // Track, from the driver, whether the producer (source scan) and consumer (dedup) stages of the + // pipelined group were ever RUNNING simultaneously. A sequential producer-then-consumer + // schedule never exceeds one running stage at a time; >= 2 proves genuine co-scheduling. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val queryStageIds = ConcurrentHashMap.newKeySet[Int]() + // Count only stages belonging to the query under test. The suite shares one SparkContext, so a + // stage from any other streaming query would otherwise satisfy the co-scheduling assertion + // below even if this query's producer and consumer actually ran one after the other. The id is + // captured from the query once it is running, and every job is matched against it. + val queryId = new AtomicReference[String](null) + val listener = new SparkListener { + override def onJobStart(e: SparkListenerJobStart): Unit = { + // StreamExecution tags every streaming job with its query id. + val id = queryId.get() + if (id != null && e.properties.getProperty(StreamExecution.QUERY_ID_KEY) == id) { + e.stageIds.foreach(queryStageIds.add(_)) + } + } + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + if (queryStageIds.contains(e.stageInfo.stageId)) { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + // scan --shuffle(repartition by key)--> streaming dropDuplicates --> sink. + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2), ("c", 2)), + StartStream(), + // StartStream only launches the query -- the stream runs on its own thread and can submit + // the first batch's jobs before the id is recorded here, in which case those jobs are not + // attributed. Rather than race, the assertion below relies on the SECOND batch, which is + // driven after this point and so is always observed in full. + Execute(q => queryId.set(q.id.toString)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d"), + Execute { q => + assertAllExchangesPipelined(q) + assert(maxConcurrentStages.get() >= 2, + s"expected >= 2 stages running concurrently, saw max ${maxConcurrentStages.get()}") + }, + StopStream + ) + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("pipelined shuffle: multi-key dedup runs in Real-Time Mode over a pipelined shuffle") { + // dropDuplicates on more than one column: the hash-partitioning is by the composite key. + // Confirms the pipelined path is not specific to a single dedup column. + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("k1"), $"_2".as("k2")) + .dropDuplicates("k1", "k2") + .select(concat($"k1", lit("-"), $"k2").as("out")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + // (a,1) twice -> once; (a,2) is a distinct composite key -> also emitted. + AddData(inputData, ("a", 1), ("a", 1), ("a", 2), ("b", 1)), + StartStream(), + CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("a", 1), ("a", 2), ("b", 2)), + CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1", "b-2"), + Execute { q => assertAllExchangesPipelined(q) }, + StopStream + ) + } + + test("pipelined shuffle: an explicit repartition-by-key runs in Real-Time Mode") { + // A stateless repartition (no dedup): the ShuffleExchangeExec is still marked pipelined in RTM. + val inputData = LowLatencyMemoryStream[Int] + val result = inputData.toDF().repartition(4, $"value").select($"value") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, 1, 2, 3, 4, 5), + StartStream(), + CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5), + Execute { q => assertAllExchangesPipelined(q) }, + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, 6, 7), + CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5, 6, 7), + StopStream + ) + } + + test("pipelined shuffle: a chain of two shuffles is a single co-scheduled pipelined group") { + // A round-robin repartition (RoundRobinPartitioning) feeding a dropDuplicates (HashPartitioning + // on the dedup key) produces TWO distinct shuffle exchanges -- neither distribution satisfies + // the other, so they are not collapsed. BOTH must be marked pipelined and the all-pipelined job + // co-schedules as one group (>= 2 stages running at once). This is the multi-shuffle-per- + // group scenario: more than one pipelined shuffle in a single Real-Time Mode job. + val runningStages = ConcurrentHashMap.newKeySet[Int]() + val maxConcurrentStages = new AtomicInteger(0) + val queryStageIds = ConcurrentHashMap.newKeySet[Int]() + // Count only stages belonging to the query under test. The suite shares one SparkContext, so a + // stage from any other streaming query would otherwise satisfy the co-scheduling assertion + // below even if this query's producer and consumer actually ran one after the other. The id is + // captured from the query once it is running, and every job is matched against it. + val queryId = new AtomicReference[String](null) + val listener = new SparkListener { + override def onJobStart(e: SparkListenerJobStart): Unit = { + // StreamExecution tags every streaming job with its query id. + val id = queryId.get() + if (id != null && e.properties.getProperty(StreamExecution.QUERY_ID_KEY) == id) { + e.stageIds.foreach(queryStageIds.add(_)) + } + } + override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = { + if (queryStageIds.contains(e.stageInfo.stageId)) { + runningStages.add(e.stageInfo.stageId) + maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max) + } + } + override def onStageCompleted(e: SparkListenerStageCompleted): Unit = { + runningStages.remove(e.stageInfo.stageId) + } + } + spark.sparkContext.addSparkListener(listener) + // Keep the shuffle partition count small: the whole group (scan + 2 shuffles) must fit in the + // test cluster's slots at once (gang admission), so a 2-shuffle chain at the default 200 + // partitions would fail with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT. + try { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")) + .repartition(4) // RoundRobinPartitioning -> shuffle #1 + .dropDuplicates("key") // HashPartitioning(key) -> shuffle #2 + .select($"key") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)), + StartStream(), + // StartStream only launches the query -- the stream runs on its own thread and can submit + // the first batch's jobs before the id is recorded here, in which case those jobs are not + // attributed and contribute nothing to maxConcurrentStages. So do not assert on the first + // batch: record the id, let batch 0 finish, then drive a SECOND batch that is guaranteed + // to run entirely after this point, and assert on that one. + Execute(q => queryId.set(q.id.toString)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("d", 1), ("e", 1), ("d", 2)), + CheckAnswerWithTimeout(60000, "a", "b", "c", "d", "e"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got ${exchanges.size}") + assert(exchanges.forall(_.pipelined), + "every exchange in the chain must be pipelined, got: " + + exchanges.map(_.pipelined).mkString(", ")) + assert(maxConcurrentStages.get() >= 2, + s"a 2-shuffle chain must co-schedule >= 2 stages, saw ${maxConcurrentStages.get()}") + }, + StopStream + ) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + + test("pipelined shuffle: a chain of two round-robin repartitions runs in Real-Time Mode") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val inputData = LowLatencyMemoryStream[(String, Int)](2) + // Two round-robin repartitions with a map between them; consecutive repartitions collapse to + // the last one, so the map is what keeps both shuffles in the plan. Both are pipelined, and + // the second one's producer reads the first one's output -- an UNORDERED input. A round-robin + // shuffle is order-sensitive once the deterministic local sort is off, which would escalate + // that producer to INDETERMINATE and get the group rejected; a pipelined shuffle is exempt + // from the order-sensitive marking precisely so this shape runs. + val result = inputData.toDS() + .repartition(3) + .map(row => row) + .repartition(2) + .map(row => row) + .toDF().select($"_1".as("key")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 2), ("c", 3)), + StartStream(), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + Execute { q => + val exchanges = q.lastExecution.executedPlan.collect { + case s: ShuffleExchangeExec => s + } + assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got ${exchanges.size}") + assert(exchanges.forall(_.pipelined), + "every exchange in the chain must be pipelined, got: " + + exchanges.map(_.pipelined).mkString(", ")) + }, + StopStream + ) + } + } + + test("pipelined shuffle: an explicit sortBeforeRepartition=true is rejected") { + // The deterministic local sort before a round-robin repartition never drains an unbounded + // stream, so honouring sortBeforeRepartition=true would hang a Real-Time Mode query forever. + // Rather than silently override it, the pre-flight in StreamingQueryManager rejects the + // explicit value up front. See StreamingQueryManager.throwIfConfsAreRealTimeModeIncompatible. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + SQLConf.SORT_BEFORE_REPARTITION.key -> "true") { + val inputData = LowLatencyMemoryStream[(String, Int)](2) + val result = inputData.toDF().select($"_1".as("key")) + .repartition(4) + .dropDuplicates("key") + .select($"key") + val e = intercept[SparkIllegalArgumentException] { + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)), + StartStream() + ) + } + checkError(e, condition = "STREAMING_REAL_TIME_MODE.SQL_CONFIGURATION_NOT_SUPPORTED", + parameters = e.getMessageParameters.asScala.toMap) + } + } + + test("pipelined shuffle: dedup recovers from a task failure via checkpoint restart") { + withTempDir { checkpointDir => + // A UDF that throws on demand, placed after the dedup so the failure lands in the pipelined + // consumer stage while the query runs. RTM does not retry tasks, so one failure fails it. + val failUDF = udf { (key: String) => + if (StreamRealTimeModeSuite.failTasks) { + throw new RuntimeException(s"forced task failure on $key") + } + key + } + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key") + .select(failUDF($"key").as("key")) + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1), ("a", 2)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + Execute { _ => StreamRealTimeModeSuite.failTasks = true }, + AddData(inputData, ("c", 1)), + advanceRealTimeClock, + ExpectFailure[SparkException] { ex => + val msg = Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage) + assert(msg != null && msg.contains("forced task failure"), + s"expected a forced task failure, got: $msg") + } + ) + // Restart from the same checkpoint: batch-0 dedup state must survive (a, b not re-emitted), + // only the genuinely-new c, d appear -- recovery to the last committed batch. + StreamRealTimeModeSuite.failTasks = false + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } + } + + test("pipelined shuffle: dedup recovers from a commit-log write failure") { + withSQLConf( + SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key -> + classOf[FailureInjectionCheckpointFileManager].getName) { + withTempDir { checkpointDir => + val injectionState = FailureInjectionFileSystem.registerTempPath(checkpointDir.getPath) + try { + val inputData = LowLatencyMemoryStream[(String, Int)] + val result = inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key") + // Batch 0 dedups a, b and commits. Then fail the close() of batch 1's commit-log write so + // batch 1 cannot commit and the query fails after processing. + injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1") + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + AddData(inputData, ("a", 1), ("b", 1)), + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + CheckAnswerWithTimeout(60000, "a", "b"), + advanceRealTimeClock, + WaitUntilBatchProcessed(0), + AddData(inputData, ("c", 1)), + CheckAnswerWithTimeout(60000, "a", "b", "c"), + advanceRealTimeClock, + ExpectFailure[IOException]() + ) + // Clear injection, restart: batch-0 state survives (a, b seen); uncommitted batch 1 + // re-runs so its new key c is still emitted, plus a further new key d. + injectionState.createAtomicDelayCloseRegex = Seq.empty + testStream(result, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(checkpointLocation = checkpointDir.getAbsolutePath), + AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)), + CheckAnswerWithTimeout(60000, "c", "d"), + advanceRealTimeClock, + StopStream + ) + } finally { + FailureInjectionFileSystem.removePathFromTempToInjectionState(checkpointDir.getPath) + } + } + } + } +} + +/** Driver-side switch a UDF reads on executors to fail a task on demand (fault-tolerance tests). */ +object StreamRealTimeModeSuite { + @volatile var failTasks: Boolean = false } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala index b4fd41a6b5501..1bad6cf328770 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamSuite.scala @@ -1236,6 +1236,36 @@ class StreamSuite extends StreamTest { testCurrentTimestampOnStreamingQuery() } + test("SPARK-57837: nanosecond current_timestamp(p)/localtimestamp(p) use the batch timestamp") { + // In micro-batch streaming, current_timestamp()/localtimestamp() are rewritten to the + // replay-stable batch timestamp (CurrentBatchTimestamp), not folded to a wall clock at + // planning time. The nanosecond variants must follow the same rewrite: the batch timestamp is + // millisecond resolution, so the collected value must be millisecond-aligned (a wall-clock + // fold from Instant.now() would carry sub-millisecond digits and fail this check). + val input = MemoryStream[Int] + val df = input.toDS().selectExpr("current_timestamp(9) AS ltz", "localtimestamp(9) AS ntz") + + def assertBatchAligned(rows: Seq[Row]): Unit = { + assert(rows.size === 1) + val row = rows.head + // LTZ collects as java.time.Instant, NTZ as java.time.LocalDateTime. + val ltzNanos = row.getAs[java.time.Instant]("ltz").getNano + val ntzNanos = row.getAs[java.time.LocalDateTime]("ntz").getNano + assert(ltzNanos % 1000000 === 0, + s"current_timestamp(9) sub-second $ltzNanos is not millisecond-aligned") + assert(ntzNanos % 1000000 === 0, + s"localtimestamp(9) sub-second $ntzNanos is not millisecond-aligned") + } + + testStream(df)( + AddData(input, 1), + CheckLastBatch { rows: Seq[Row] => assertBatchAligned(rows) }, + Execute { _ => Thread.sleep(1000) }, + AddData(input, 2), + CheckLastBatch { rows: Seq[Row] => assertBatchAligned(rows) } + ) + } + private def testCurrentTimestampOnStreamingQuery(): Unit = { val input = MemoryStream[Int] val df = input.toDS().withColumn("cur_timestamp", lit(current_timestamp())) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala index f065f1de5cdc4..0055b1f91a7c3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingAggregationSuite.scala @@ -110,6 +110,38 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with Assertions { ) } + testWithAllStateVersions("approximate percentiles preserve existing streaming checkpoints") { + withTempDir { checkpointDir => + val inputData = MemoryStream[(Int, Int)] + val aggregated = inputData.toDF() + .groupBy($"_1") + .agg( + expr("percentile_approx(_2, 0.5)").as("p50"), + expr("percentile_approx(_2, 0.9)").as("p90")) + .as[(Int, Int, Int)] + + // Create the two-digest checkpoint used before percentile fusion was introduced. + testStream(aggregated, Update)( + StartStream( + checkpointLocation = checkpointDir.getAbsolutePath, + additionalConfs = Map( + SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "false")), + AddData(inputData, (0, 1), (0, 2), (0, 3)), + CheckLastBatch((0, 2, 3)), + StopStream) + + // Restart with percentile fusion enabled to verify the existing checkpoint stays valid. + testStream(aggregated, Update)( + StartStream( + checkpointLocation = checkpointDir.getAbsolutePath, + additionalConfs = Map( + SQLConf.COMBINE_APPROXIMATE_PERCENTILES_ENABLED.key -> "true")), + AddData(inputData, (0, 4)), + CheckLastBatch((0, 2, 4)), + StopStream) + } + } + testWithAllStateVersions("count distinct") { val inputData = MemoryStream[(Int, Seq[Int])] @@ -930,6 +962,207 @@ class StreamingAggregationSuite extends StateStoreMetricsTest with Assertions { (inputData2, aggregated2) } + // Streaming aggregation planned as StatefulStreamlineAggregateExec, which merges each input row + // against the state store and emits as it goes, rather than the micro-batch operators that emit + // once the batch ends. Real-Time Mode always plans it; here the config selects it so the operator + // can be covered under a micro-batch trigger too. See StreamlineStreamingAggregationRealTimeSuite + // for the Real-Time Mode coverage. + private val streamlineEnabled = + SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key -> "true" + + testWithAllStateVersions("streamline aggregation: update mode emits per input row", + streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggregated = inputData.toDF() + .groupBy($"value") + .agg(count("*")) + .as[(Int, Long)] + + testStream(aggregated, Update)( + AddData(inputData, 3), + CheckLastBatch((3, 1)), + AddData(inputData, 3, 2), + CheckLastBatch((3, 2), (2, 1)), + StopStream, + StartStream(), + AddData(inputData, 3, 2, 1), + CheckLastBatch((3, 3), (2, 2), (1, 1)), + // The distinguishing behaviour: four rows for the same key produce four outputs, one per + // input. Micro-batch aggregation would emit only the final (4, 4) for this batch. + AddData(inputData, 4, 4, 4, 4), + CheckLastBatch((4, 1), (4, 2), (4, 3), (4, 4)) + ) + } + + testWithAllStateVersions("streamline aggregation: complete mode outputs the whole result table", + streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggregated = inputData.toDF() + .groupBy($"value") + .agg(count("*")) + .as[(Int, Long)] + + testStream(aggregated, Complete)( + AddData(inputData, 3), + CheckLastBatch((3, 1)), + AddData(inputData, 2), + CheckLastBatch((3, 1), (2, 1)), + StopStream, + StartStream(), + AddData(inputData, 3, 2), + CheckLastBatch((3, 2), (2, 2)) + ) + } + + testWithAllStateVersions("streamline aggregation: sum, min, max and avg", streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggregated = inputData.toDF() + .selectExpr("value % 2 AS key", "value") + .groupBy($"key") + .agg(sum("value"), min("value"), max("value"), avg("value")) + .as[(Int, Long, Int, Int, Double)] + + testStream(aggregated, Complete)( + AddData(inputData, 1, 2, 3, 4), + // key 1: values 1, 3 -> sum 4, min 1, max 3, avg 2.0 + // key 0: values 2, 4 -> sum 6, min 2, max 4, avg 3.0 + CheckLastBatch((1, 4L, 1, 3, 2.0), (0, 6L, 2, 4, 3.0)), + AddData(inputData, 5, 6), + // key 1 gains 5 -> sum 9, max 5, avg 3.0; key 0 gains 6 -> sum 12, max 6, avg 4.0 + CheckLastBatch((1, 9L, 1, 5, 3.0), (0, 12L, 2, 6, 4.0)) + ) + } + + testWithAllStateVersions("streamline aggregation: multiple grouping keys", streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggregated = inputData.toDF() + .selectExpr("value", "value % 2 AS k1", "value % 3 AS k2") + .groupBy($"k1", $"k2") + .agg(count("*")) + .as[(Int, Int, Long)] + + testStream(aggregated, Complete)( + AddData(inputData, 1, 2, 3, 4, 5, 6), + CheckLastBatch( + (1, 1, 1), // 1 + (0, 2, 1), // 2 + (1, 0, 1), // 3 + (0, 1, 1), // 4 + (1, 2, 1), // 5 + (0, 0, 1)) // 6 + ) + } + + testWithAllStateVersions("streamline aggregation: recovery from a restart keeps state", + streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggregated = inputData.toDF() + .groupBy($"value") + .agg(count("*")) + .as[(Int, Long)] + + testStream(aggregated, Complete)( + AddData(inputData, 1, 1, 2), + CheckLastBatch((1, 2), (2, 1)), + StopStream, + StartStream(), + // The counts continue from the committed state rather than restarting at 1. + AddData(inputData, 1, 2), + CheckLastBatch((1, 3), (2, 2)) + ) + } + + testWithAllStateVersions("streamline aggregation: computed grouping key", streamlineEnabled) { + val inputData = MemoryStream[Int] + // The grouping key is a COMPUTED expression, not a bare column reference. The final stage is + // planned with the original groupingExpressions rather than the post-shuffle attributes, so if + // that is wrong this is where it shows: the final stage would try to re-evaluate `value % 3` + // against its child's output, which only carries the already-grouped attribute. + val aggregated = inputData.toDF() + .groupBy(($"value" % 3).as("k")) + .agg(count("*"), sum("value")) + .as[(Int, Long, Long)] + + testStream(aggregated, Complete)( + AddData(inputData, 1, 2, 3, 4, 5, 6), + // k=1: 1,4 -> count 2 sum 5 ; k=2: 2,5 -> count 2 sum 7 ; k=0: 3,6 -> count 2 sum 9 + CheckLastBatch((1, 2L, 5L), (2, 2L, 7L), (0, 2L, 9L)) + ) + } + + // Append mode is the one mode that drives the operator's eviction path: a windowed grouping key + // is emitted only once the watermark passes it, via EvictionIterator (which removes in hasNext). + // Mirrors the stateStoreSave Append test above so the streamline operator is held to the same + // watermark/eviction behaviour. + testWithAllStateVersions("streamline aggregation: append mode emits windows past the watermark", + streamlineEnabled) { + val inputData = MemoryStream[Int] + + val aggWithWatermark = inputData.toDF() + .withColumn("eventTime", timestamp_seconds($"value")) + .withWatermark("eventTime", "10 seconds") + .groupBy(window($"eventTime", "5 seconds") as Symbol("window")) + .agg(count("*") as Symbol("count")) + .select($"window".getField("start").cast("long").as[Long], $"count".as[Long]) + + testStream(aggWithWatermark, Append)( + StartStream(additionalConfs = Map(SQLConf.SHUFFLE_PARTITIONS.key -> "3")), + AddData(inputData, 3, 2, 1, 9), + // Nothing is emitted yet: no window has fallen fully below the watermark. + CheckLastBatch(), + AddData(inputData, 25), // Advance watermark to 15s; windows ending <= 15s are now evictable. + // Both closed windows are emitted once, past the watermark, and then evicted: + // [0,5) has 1,2,3 -> count 3 ; [5,10) has 9 -> count 1. + CheckLastBatch((0, 3), (5, 1)) + ) + // Note: state-row *metrics* (e.g. updated rows) are not asserted here because the streamline + // operator writes state per input row rather than once per key per batch, so its counts + // legitimately differ from the micro-batch stateStoreSave operator. Output correctness -- what + // Append actually guarantees -- is covered by the CheckLastBatch assertions above. + } + + // The micro-batch aggregation operator (stateStoreSave) and the streamline operator share the + // same StreamingAggregationStateManager and state format, so a checkpoint written by one can be + // read by the other. The operator-name check in IncrementalExecution must therefore treat the + // switch as a supported transition rather than a changed stateful operator (which would abort the + // query on restart). Both directions are exercised. + test("switching between the micro-batch and streamline aggregation operators keeps state") { + Seq(false -> true, true -> false).foreach { case (firstStreamline, secondStreamline) => + withTempDir { dir => + val inputData = MemoryStream[Int] + val aggregated = inputData.toDF().groupBy($"value").agg(count("*")).as[(Int, Long)] + spark.conf.set( + SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key, firstStreamline.toString) + testStream(aggregated, Complete)( + StartStream(checkpointLocation = dir.getAbsolutePath), + AddData(inputData, 1, 1, 2), + CheckLastBatch((1, 2), (2, 1)), + StopStream, + Execute { _ => + spark.conf.set( + SQLConf.STREAMING_USE_STREAMLINE_AGGREGATOR.key, secondStreamline.toString) + }, + StartStream(checkpointLocation = dir.getAbsolutePath), + // State survives the operator switch: the counts continue rather than resetting. + AddData(inputData, 1, 2), + CheckLastBatch((1, 3), (2, 2)), + StopStream + ) + } + } + } + + // Note: the streamline operator's Append branch asserts a watermark is present, but Append + // aggregation WITHOUT a watermark is already rejected at analysis time by + // UnsupportedOperationChecker (STREAMING_OUTPUT_MODE.UNSUPPORTED_OPERATION), so that assert is a + // defensive internal invariant no normal query reaches. The Append-WITH-watermark path is + // covered by the test above. + @tailrec private def findStateSchemaNotCompatible(exc: Throwable): Option[SparkUnsupportedOperationException] = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationSuite.scala index 003f71d16437a..21261280a2800 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationSuite.scala @@ -684,6 +684,113 @@ class StreamingDeduplicationSuite extends StateStoreMetricsTest sqlConf = spark.sessionState.conf ) } + + // Total incremental removals over the query's retained recent progress (recentProgress is + // retention-bounded, so this is not the full batch history), read from the operator's + // numRowsIncrementallyRemoved custom metric. Where a stateful operator ran (stateOperators is + // non-empty) the metric must be present -- assert rather than default it to 0, so a regression + // that stops emitting the metric fails the test instead of silently reading as zero. + private def numRowsIncrementallyRemoved(q: StreamingQuery): Long = { + q.recentProgress.flatMap(_.stateOperators.headOption).map { op => + assert(op.customMetrics.containsKey("numRowsIncrementallyRemoved"), + s"numRowsIncrementallyRemoved custom metric missing; got ${op.customMetrics}") + op.customMetrics.get("numRowsIncrementallyRemoved").toLong + }.sum + } + + // Total state rows removed over the query's retained recent progress (recentProgress is + // retention-bounded, so this is not the full batch history), read from the operator's first-class + // numRowsRemoved metric (the incremental removals are a subset of these). This asserts that + // eviction actually removed state, which numRowsIncrementallyRemoved alone does not: a batch-end + // drain removes rows without incrementing the incremental counter. + private def totalStateRowsRemoved(q: StreamingQuery): Long = + q.recentProgress.flatMap(_.stateOperators.headOption.map(_.numRowsRemoved)).sum + + test("deduplicate with watermark - incremental cleanup preserves dedup output") { + // With a non-zero incremental cleanup factor, removal of watermark-expired state is spread + // across input-record processing rather than occurring all at once at batch end. The + // deduplicated OUTPUT must be unchanged. (State is evicted against the late-events watermark + // under incremental cleanup, so the timing of state removal lags by a batch compared to the + // factor-0 default -- this test asserts on output and on the late-event safety property, not + // on per-batch state counts.) + withSQLConf(SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR.key -> "10") { + val inputData = MemoryStream[Int] + val result = inputData.toDS() + .withColumn("eventTime", timestamp_seconds($"value")) + .withWatermark("eventTime", "10 seconds") + .dropDuplicates() + .select($"eventTime".cast("long").as[Long]) + + testStream(result, Append)( + // Duplicates within the batch are dropped; each distinct event time is emitted once. + AddData(inputData, (1 to 5).flatMap(_ => (10 to 15)): _*), + CheckAnswer(10 to 15: _*), + + AddData(inputData, 25), // Advances watermark; 25 is new and emitted. + CheckNewAnswer(25), + + // A record at 10 is now below the watermark: it must be dropped, and crucially it must not + // be re-emitted even though incremental cleanup may not yet have removed its key. This is + // the safety property behind evicting against the late-events (not eviction) watermark. + AddData(inputData, 10), + CheckNewAnswer(), + + AddData(inputData, 45), + CheckNewAnswer(45), + + // A duplicate of a surviving recent key (45) is still deduplicated. + AddData(inputData, 45), + CheckNewAnswer() + ) + } + } + + test("deduplicate with watermark - incremental cleanup evicts during record processing") { + // A batch that carries input records AND has state eligible for eviction under the late-events + // watermark should evict incrementally as those records are processed, so + // numRowsIncrementallyRemoved is non-zero. (Eviction uses the late-events watermark under + // incremental cleanup, which lags the eviction watermark by one batch.) + withSQLConf(SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR.key -> "10") { + val inputData = MemoryStream[Int] + val result = inputData.toDS() + .withColumn("eventTime", timestamp_seconds($"value")) + .withWatermark("eventTime", "10 seconds") + .dropDuplicates() + .select($"eventTime".cast("long").as[Long]) + + testStream(result, Append)( + // Batch 0: three distinct keys at 10, 11, 12. Watermark is 0, nothing evictable yet. + AddData(inputData, 10, 11, 12), + CheckAnswer(10, 11, 12), + assertNumStateRows(total = 3, updated = 3), + AssertOnQuery(q => numRowsIncrementallyRemoved(q) == 0, + "no incremental removal before any state is evictable"), + + // Batch 1: a new key at 100 advances the eviction watermark to 90, but the late-events + // watermark used by incremental cleanup still lags (it is the previous batch's eviction + // watermark). The keys [10, 11, 12] are not yet below the late-events watermark, so no + // removal happens in this batch; it is deferred to batch 2. + AddData(inputData, 100), + CheckNewAnswer(100), + + // Batch 2: another new record. By now the late-events watermark has advanced past the + // original keys, so processing this record incrementally evicts them. + AddData(inputData, 101), + CheckNewAnswer(101), + AssertOnQuery(q => numRowsIncrementallyRemoved(q) > 0, + "expired keys should be removed incrementally while processing input records"), + // Assert the state was actually removed, not merely that the incremental counter moved: + // the three original keys [10, 11, 12] must be gone, leaving only the recent keys. Reading + // the first-class numRowsRemoved metric (rather than the output) is what catches a + // regression that leaves state behind while still producing correct output. + AssertOnQuery(q => totalStateRowsRemoved(q) >= 3, + "expired keys must be removed from state, not just counted as incremental"), + // The surviving state is only the most recent keys; correctness is unchanged. + AddData(inputData, 10), // below watermark, dropped + CheckNewAnswer() + ) + } + } } trait StreamingDeduplicationSuiteBase { self: StreamTest => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala index 9645f82ac241b..8a9766c7e8751 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala @@ -22,6 +22,7 @@ import org.apache.spark.sql.catalyst.streaming.InternalOutputModes.Append import org.apache.spark.sql.execution.streaming.operators.stateful.StatefulOperatorsUtils import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.functions.timestamp_seconds +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{LongType, StringType, StructType} import org.apache.spark.tags.SlowSQLTest @@ -164,6 +165,50 @@ class StreamingDeduplicationWithinWatermarkSuite extends StateStoreMetricsTest ) } + test("incremental cleanup factor does not change dropDuplicatesWithinWatermark output") { + // dropDuplicatesWithinWatermark deliberately does not participate in incremental cleanup + // (its dedup key excludes the event time, so mid-batch eviction of a shared key could change + // the output). Setting a non-zero incrementalCleanupFactor must therefore leave the output and + // per-batch state counts identical to the default -- this is the same sequence as + // "deduplicate with subset of columns which event time column is not in subset". + withSQLConf(SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR.key -> "10") { + val inputData = MemoryStream[(String, Int)] + val result = inputData.toDS() + .withColumn("eventTime", timestamp_seconds($"_2")) + .withWatermark("eventTime", "2 seconds") + .dropDuplicatesWithinWatermark("_1") + .select($"_1", $"eventTime".cast("long").as[Long]) + + testStream(result, Append)( + AddData(inputData, "a" -> 17), + CheckNewAnswer("a" -> 17), + assertNumStateRows(total = 1, updated = 1), + + AddData(inputData, "a" -> 16), + CheckNewAnswer(), + assertNumStateRows(total = 1, updated = 0), + + AddData(inputData, "a" -> 13), + CheckNewAnswer(), + assertNumStateRows(total = 1, updated = 0, droppedByWatermark = 1), + + AddData(inputData, "b" -> 22, "c" -> 21), + CheckNewAnswer("b" -> 22, "c" -> 21), + assertNumStateRows(total = 2, updated = 2), + + // "a" -> 21 is emitted as new because the earlier "a" state expired and was evicted at the + // previous batch end -- unchanged from the factor-0 path. + AddData(inputData, "a" -> 21), + CheckNewAnswer("a" -> 21), + assertNumStateRows(total = 3, updated = 1), + + AddData(inputData, "d" -> 25), + CheckNewAnswer("d" -> 25), + assertNumStateRows(total = 2, updated = 1) + ) + } + } + test("SPARK-39650: duplicate with specific keys should allow input to change schema") { withTempDir { checkpoint => val dedupeInputData = MemoryStream[(String, Int)] diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala index c46f0076721b9..c0983f338abe5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala @@ -677,8 +677,8 @@ abstract class StreamingInnerJoinBase extends StreamingJoinSuite { assert(query.lastExecution.executedPlan.collect { case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _, - ShuffleExchangeExec(opA: HashPartitioning, _, _, _), - ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _) + ShuffleExchangeExec(opA: HashPartitioning, _, _, _, _), + ShuffleExchangeExec(opB: HashPartitioning, _, _, _, _), _) if partitionExpressionsColumns(opA.expressions) === Seq("a", "b") && partitionExpressionsColumns(opB.expressions) === Seq("a", "b") && opA.numPartitions == numPartitions && opB.numPartitions == numPartitions => j diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingRealTimeModeSourceCompatSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingRealTimeModeSourceCompatSuite.scala new file mode 100644 index 0000000000000..6b20ddf08dac0 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingRealTimeModeSourceCompatSuite.scala @@ -0,0 +1,271 @@ +/* + * 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.streaming + +import java.util + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability, TableProvider} +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory, Scan, ScanBuilder} +import org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset, PartitionOffset, SupportsRealTimeMode, SupportsRealTimeRead} +import org.apache.spark.sql.connector.read.streaming.SupportsRealTimeRead.RecordStatus +import org.apache.spark.sql.execution.streaming.sources.ContinuousMemorySink +import org.apache.spark.sql.sources.DataSourceRegister +import org.apache.spark.sql.types.{IntegerType, StringType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.unsafe.types.UTF8String + +// scalastyle:off +/** + * ============================================================================================ + * BACKWARD-COMPATIBILITY GUARD -- DO NOT MODIFY THE SOURCE DEFINITIONS BELOW. + * ============================================================================================ + * + * The classes in this file define a self-contained Real-Time Mode (RTM) streaming source that is + * written EXCLUSIVELY against the public, `@Evolving` connector APIs an external connector author + * has access to: + * + * - `org.apache.spark.sql.sources.DataSourceRegister` + * - `org.apache.spark.sql.connector.catalog.{TableProvider, Table, SupportsRead, + * TableCapability}` + * - `org.apache.spark.sql.connector.read.{ScanBuilder, Scan, InputPartition, PartitionReader, + * PartitionReaderFactory}` + * - `org.apache.spark.sql.connector.read.streaming.{MicroBatchStream, Offset, PartitionOffset, + * SupportsRealTimeMode, SupportsRealTimeRead}` + * + * It deliberately does NOT use any `private[spark]` / internal helper (no `LowLatencyMemoryStream`, + * `LongOffset`, `SimpleTableProvider`, RPC endpoints, or the engine-internal low latency clock). + * Its purpose is to pin the source-level backward compatibility of the RTM connector SPI: if a + * future change to `SupportsRealTimeMode` or `SupportsRealTimeRead` breaks external implementors + * the way SPARK-55699 did (it replaced `nextWithTimeout(Long)` rather than adding an overload, see + * SPARK-58386), this file will FAIL TO COMPILE, catching the incompatibility at build time. + * + * `SupportsRealTimeRead` offers two `nextWithTimeout` overloads and a source overrides exactly one. + * Both are covered here so the guard tracks either entry point: + * - `CompatOneArgPartitionReader` overrides only `nextWithTimeout(Long)` -- the method a + * third-party source is expected to implement. The engine invokes the two-arg overload, so this + * reader is exercised through the interface's default delegation. + * - `CompatTwoArgPartitionReader` overrides `nextWithTimeout(Long, Long)` directly -- the + * overload the engine invokes. + * Keep both. + * + * When you add a genuinely new REQUIRED method to one of these interfaces, prefer a `default` + * method so this frozen source keeps compiling. If a required change is truly unavoidable, updating + * this file is a strong signal that external connectors will also break -- treat it accordingly. + * ============================================================================================ + */ +// scalastyle:on + +/** Deterministic, never-changing dataset the guard source serves. */ +private object CompatRealTimeData { + val schema: StructType = + new StructType().add("value", IntegerType).add("name", StringType) + + // A fixed, finite dataset. Frozen on purpose -- do not change. + val records: Array[(Int, String)] = Array((1, "a"), (2, "b"), (3, "c")) +} + +/** A public-API `Offset`: the number of records consumed so far. */ +private case class CompatOffset(consumed: Int) extends Offset { + override def json(): String = consumed.toString +} + +/** A public-API per-partition `PartitionOffset`. */ +private case class CompatPartitionOffset(partitionId: Int, offset: Int) extends PartitionOffset + +/** A serializable `InputPartition` carrying its slice of the frozen dataset. */ +private case class CompatInputPartition( + partitionId: Int, + startOffset: Int, + rows: Array[(Int, String)]) + extends InputPartition + +/** + * Shared reader logic for both `nextWithTimeout` entry points. Concrete subclasses only pick which + * overload of [[SupportsRealTimeRead#nextWithTimeout]] to override; both route to [[pollNext]]. + */ +private abstract class CompatRealTimePartitionReaderBase(partition: CompatInputPartition) + extends SupportsRealTimeRead[InternalRow] { + + private var pos = 0 + private var currentRow: InternalRow = _ + + private def toRow(i: Int): InternalRow = { + val (v, n) = partition.rows(i) + InternalRow(v, UTF8String.fromString(n)) + } + + /** Return the next record, or wait until the timeout elapses and report no record. */ + protected final def pollNext(timeoutMs: java.lang.Long): RecordStatus = { + if (pos < partition.rows.length) { + val (value, _) = partition.rows(pos) + currentRow = toRow(pos) + pos += 1 + // Report the record along with a deterministic synthetic arrival time, exercising the + // arrival-time branch of RecordStatus. + return RecordStatus.newStatusWithArrivalTimeMs(value.toLong) + } + // Exhausted this batch's data: keep waiting until the caller's timeout elapses, then report + // no record -- the same wait-until-timeout behavior a real source has. Measured against the + // wall clock, as a third-party source without the engine's reference clock would do. + val startNs = System.nanoTime() + var elapsedMs = 0L + while (elapsedMs < timeoutMs) { + Thread.sleep(10L) + elapsedMs = (System.nanoTime() - startNs) / 1000000L + } + RecordStatus.newStatusWithoutArrivalTime(false) + } + + override def getOffset: PartitionOffset = + CompatPartitionOffset(partition.partitionId, partition.startOffset + pos) + + override def next(): Boolean = { + if (pos < partition.rows.length) { + currentRow = toRow(pos) + pos += 1 + true + } else { + false + } + } + + override def get(): InternalRow = currentRow + + override def close(): Unit = {} +} + +/** + * Reader that overrides ONLY the single-argument `nextWithTimeout(Long)` -- the method a + * third-party source is expected to implement. Reached through the two-arg default delegation. + */ +private class CompatOneArgPartitionReader(partition: CompatInputPartition) + extends CompatRealTimePartitionReaderBase(partition) { + override def nextWithTimeout(timeoutMs: java.lang.Long): RecordStatus = pollNext(timeoutMs) +} + +/** Reader that overrides the two-argument `nextWithTimeout(Long, Long)` the engine invokes. */ +private class CompatTwoArgPartitionReader(partition: CompatInputPartition) + extends CompatRealTimePartitionReaderBase(partition) { + override def nextWithTimeout( + startTimeMs: java.lang.Long, timeoutMs: java.lang.Long): RecordStatus = pollNext(timeoutMs) +} + +/** A public-API `PartitionReaderFactory`, parameterized by which reader overload to use. */ +private class CompatRealTimeReaderFactory(twoArg: Boolean) extends PartitionReaderFactory { + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = { + val p = partition.asInstanceOf[CompatInputPartition] + if (twoArg) new CompatTwoArgPartitionReader(p) else new CompatOneArgPartitionReader(p) + } +} + +/** The RTM stream: a public `MicroBatchStream` that also implements `SupportsRealTimeMode`. */ +private class CompatRealTimeStream(twoArg: Boolean) + extends MicroBatchStream with SupportsRealTimeMode { + override def initialOffset(): Offset = CompatOffset(0) + override def deserializeOffset(json: String): Offset = CompatOffset(json.toInt) + override def commit(end: Offset): Unit = {} + override def stop(): Unit = {} + + override def latestOffset(): Offset = CompatOffset(CompatRealTimeData.records.length) + override def planInputPartitions(start: Offset, end: Offset): Array[InputPartition] = { + val from = start.asInstanceOf[CompatOffset].consumed + val to = end.asInstanceOf[CompatOffset].consumed + Array(CompatInputPartition(0, from, CompatRealTimeData.records.slice(from, to))) + } + override def createReaderFactory(): PartitionReaderFactory = + new CompatRealTimeReaderFactory(twoArg) + + override def planInputPartitions(start: Offset): Array[InputPartition] = { + val from = start.asInstanceOf[CompatOffset].consumed + Array(CompatInputPartition(0, from, CompatRealTimeData.records.drop(from))) + } + override def mergeOffsets(offsets: Array[PartitionOffset]): Offset = { + val maxOffset = offsets.map(_.asInstanceOf[CompatPartitionOffset].offset).max + CompatOffset(maxOffset) + } +} + +/** Scan + ScanBuilder wired to the RTM stream, using only public APIs. */ +private class CompatRealTimeScan(twoArg: Boolean) extends ScanBuilder with Scan { + override def build(): Scan = this + override def readSchema(): StructType = CompatRealTimeData.schema + override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = + new CompatRealTimeStream(twoArg) +} + +/** + * The top-level source, registered as a `TableProvider` + `DataSourceRegister`. Loaded by fully + * qualified class name via `spark.readStream.format(...)`, so it doesn't need a `META-INF/services` + * registration entry. The `twoArg` option selects which `nextWithTimeout` overload the reader + * implements. + */ +class CompatRealTimeSourceProvider extends TableProvider with DataSourceRegister { + override def shortName(): String = "compat-realtime-source" + + override def inferSchema(options: CaseInsensitiveStringMap): StructType = + CompatRealTimeData.schema + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: util.Map[String, String]): Table = new CompatRealTimeTable +} + +private class CompatRealTimeTable extends Table with SupportsRead { + override def name(): String = "compat-realtime-source" + override def schema(): StructType = CompatRealTimeData.schema + override def capabilities(): util.Set[TableCapability] = + util.EnumSet.of(TableCapability.MICRO_BATCH_READ) + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new CompatRealTimeScan(options.getBoolean("twoArg", false)) +} + +/** + * SPARK-58386: a compile-time and runtime backward-compatibility guard for the public Real-Time + * Mode connector SPI (`SupportsRealTimeMode` / `SupportsRealTimeRead`). The frozen source above + * must keep compiling against these interfaces, and these tests run it end-to-end through a real + * RTM streaming query to prove an external-style source is still driven correctly, for both + * `nextWithTimeout` entry points. + */ +class StreamingRealTimeModeSourceCompatSuite extends StreamRealTimeModeManualClockSuiteBase { + import testImplicits._ + + private def runSourceEndToEnd(twoArg: Boolean): Unit = { + val df = spark.readStream + .format(classOf[CompatRealTimeSourceProvider].getName) + .option("twoArg", twoArg) + .load() + .selectExpr("concat(cast(value as string), '-', name) as output") + + testStream(df, OutputMode.Update, Map.empty, new ContinuousMemorySink())( + StartStream(), + CheckAnswerWithTimeout(10000, "1-a", "2-b", "3-c"), + StopStream + ) + } + + test("RTM source implementing single-arg nextWithTimeout reads end-to-end") { + runSourceEndToEnd(twoArg = false) + } + + test("RTM source implementing two-arg nextWithTimeout reads end-to-end") { + runSourceEndToEnd(twoArg = true) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala new file mode 100644 index 0000000000000..5c5409f7a7d59 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamlineStreamingAggregationRealTimeSuite.scala @@ -0,0 +1,134 @@ +/* + * 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.streaming + +import org.apache.spark.sql.execution.streaming.StatefulStreamlineAggregateExec +import org.apache.spark.sql.execution.streaming.operators.stateful.StreamingAggregationStateManager +import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink, LowLatencyMemoryStream} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.internal.SQLConf + +/** + * Tests for streaming aggregation in Real-Time Mode, which is planned as + * [[StatefulStreamlineAggregateExec]] rather than the micro-batch aggregation operators. + */ +class StreamlineStreamingAggregationRealTimeSuite extends StreamRealTimeModeSuiteBase + with StateStoreMetricsTest { + + import testImplicits._ + + def stateFormatVersions: Seq[Int] = StreamingAggregationStateManager.supportedVersions + + def executeFuncWithStateVersionSQLConf( + stateVersion: Int, + confPairs: Seq[(String, String)], + func: => Any): Unit = { + withSQLConf(confPairs ++ + Seq(SQLConf.STREAMING_AGGREGATION_STATE_FORMAT_VERSION.key -> stateVersion.toString): _*) { + func + } + } + + def testWithAllStateVersions(name: String, confPairs: (String, String)*) + (func: => Any): Unit = { + for (version <- stateFormatVersions) { + test(s"$name - state format version $version") { + executeFuncWithStateVersionSQLConf(version, confPairs, func) + } + } + } + + testWithAllStateVersions("aggregation runs in Real-Time Mode") { + val inputData = LowLatencyMemoryStream[(String, Int)](2) + + val agg = inputData.toDF().select($"_1".as("key"), $"_2".as("value")) + .groupBy($"key") + .agg(sum("value").as("total")) + + testStream(agg, OutputMode.Update, sink = new ContinuousMemorySink())( + StartStream(), + AddData(inputData, ("a", 1), ("b", 2), ("a", 3)), + // Update mode emits an intermediate result per input row, so "a" is seen twice: once with + // its own value and once merged with the earlier one. Micro-batch aggregation would instead + // emit only the final value per key per batch. + CheckAnswerWithTimeout(60000, ("a", 1L), ("b", 2L), ("a", 4L)), + Execute { q => + val aggregates = q.lastExecution.executedPlan.collect { + case a: StatefulStreamlineAggregateExec => a + } + assert(aggregates.size == 1, + s"expected the streamline aggregate operator, got:\n${q.lastExecution.executedPlan}") + }, + StopStream + ) + } + + // The eviction counts below are per state store, so the state has to live in a single partition + // for them to be predictable. Testing incremental cleanup without pinning the partition count + // would need a partition-aware data generator. + testWithAllStateVersions("update mode aggregation with incremental cleanup evicts records " + + "not removed during incremental cleanup", + SQLConf.STREAMING_STATE_INCREMENTAL_CLEANUP_FACTOR.key -> "2", + SQLConf.SHUFFLE_PARTITIONS.key -> "1" + ) { + val inputData = LowLatencyMemoryStream[Int] + val aggWithWatermark = inputData.toDF() + .withColumn("eventTime", timestamp_seconds($"value")) + .withWatermark("eventTime", "50 seconds") + .groupBy(window($"eventTime", "10 seconds") as Symbol("window")) + .agg(count("*") as Symbol("count")) + .select($"window".getField("end").cast("long").as[Long], $"count".as[Long]) + + // With incremental eviction, we evict incrementalCleanupFactor * numInputRows from the + // state store. + testStream(aggWithWatermark, OutputMode.Update, sink = new ContinuousMemorySink())( + StartStream(), + AddData(inputData, 9, 19, 29, 39, 49), + WaitUntilBatchProcessed(0), + // Batch 0: watermark is 0. + CheckAnswerWithTimeout(60000, (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)), + + // Batch 1: watermark starts at 0 and moves past the [40, 50) window end. + AddData(inputData, 101), + // Wait until batch 1 and the no data batch complete + WaitUntilBatchProcessed(2), + CheckAnswerWithTimeout(60000, (110, 1), (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)), + + Execute { q => + val batch1Metrics = q.recentProgress.filter(_.batchId == 1).head + assert( + batch1Metrics.stateOperators.head.customMetrics.get("numRowsIncrementallyRemoved") == 0) + assert(batch1Metrics.stateOperators.head.numRowsRemoved === 0) + }, + + AddData(inputData, 100), + WaitUntilBatchProcessed(3), + CheckAnswerWithTimeout(60000, + (110, 1), (110, 2), (10, 1), (20, 1), (30, 1), (40, 1), (50, 1)), + + Execute { q => + val batch3Metrics = q.recentProgress.filter(_.batchId == 3).head + // 1 record is in the batch, and the incremental cleanup factor is 2. Thus, 2 rows should be + // incrementally cleaned up, but we should still remove 5 total. + assert( + batch3Metrics.stateOperators.head.customMetrics.get("numRowsIncrementallyRemoved") == 2) + assert(batch3Metrics.stateOperators.head.numRowsRemoved === 5) + } + ) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala index 9685f70b86a72..53a02b3f6d209 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateInitialStateSuite.scala @@ -756,6 +756,7 @@ class TransformWithStateInitialStateSuite extends StateStoreMetricsTest } } +@SlowSQLTest class TransformWithStateInitialStateSuiteCheckpointV2 extends TransformWithStateInitialStateSuite { override def beforeAll(): Unit = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/ContinuousSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/ContinuousSuite.scala index c70f21ae144b6..4de3867fd61e5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/ContinuousSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/ContinuousSuite.scala @@ -185,6 +185,24 @@ class ContinuousSuite extends ContinuousSuiteBase { "Continuous processing does not support current time operations.")) } + test("SPARK-57837: nanosecond current-timestamp functions are rejected") { + // current_timestamp(p) -> CurrentTimestampNanos and localtimestamp(p) -> LocalTimestampNanos + // must be rejected in continuous processing just like their microsecond forms. The nanos + // types are enabled by default under Utils.isTesting. + Seq("current_timestamp(9)", "now(9)", "localtimestamp(9)").foreach { fn => + val input = ContinuousMemoryStream[Int] + val df = input.toDF().selectExpr(fn) + + val except = intercept[AnalysisException] { + testStream(df)(StartStream()) + } + + assert(except.message.contains( + "Continuous processing does not support current time operations."), + s"$fn should be rejected for continuous processing") + } + } + test("subquery alias") { withTempView("memory") { val input = ContinuousMemoryStream[Int] diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSinkEvolutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSinkEvolutionSuite.scala index 46de58fad3334..42a2edd016b9b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSinkEvolutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSinkEvolutionSuite.scala @@ -23,8 +23,9 @@ import org.apache.spark.{SparkException, SparkIllegalArgumentException} import org.apache.spark.sql._ import org.apache.spark.sql.execution.streaming.checkpointing.{CommitLog, CommitMetadataV3} import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.streaming.{StreamTest, Trigger} +import org.apache.spark.sql.streaming.{StreamingQuery, StreamTest, Trigger} import org.apache.spark.util.Utils /** @@ -313,6 +314,40 @@ class StreamingSinkEvolutionSuite extends StreamTest with BeforeAndAfterEach { assert(v3.sinkMetadataMap.size === 1) } + testWithSinkEvolution("restart V3 commit log with state checkpoint format V1") { + val checkpointDir = newMetadataDir + val input = MemoryStream[Int] + + withSQLConf( + SQLConf.STATE_STORE_CHECKPOINT_FORMAT_VERSION.key -> "1", + SQLConf.STATE_STORE_PROVIDER_CLASS.key -> + classOf[HDFSBackedStateStoreProvider].getName) { + def startQuery(): StreamingQuery = input.toDF() + .groupBy("value") + .count() + .writeStream + .format("noop") + .outputMode("update") + .name("stateful_sink") + .option("checkpointLocation", checkpointDir) + .start() + + input.addData(1) + val firstRun = startQuery() + firstRun.processAllAvailable() + firstRun.stop() + + input.addData(2) + val restarted = startQuery() + restarted.processAllAvailable() + restarted.stop() + } + + val commitLog = new CommitLog(spark, s"$checkpointDir/commits", readOnly = true) + val v3 = commitLog.getLatest().get._2.asInstanceOf[CommitMetadataV3] + assert(v3.stateUniqueIds.isEmpty) + } + // ============== // Helper Methods // ============== diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSourceEvolutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSourceEvolutionSuite.scala index cdf8cb76d8e38..ee52768ceb5f2 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSourceEvolutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/test/StreamingSourceEvolutionSuite.scala @@ -29,12 +29,14 @@ import org.apache.spark.sql.execution.streaming.checkpointing.{OffsetMap, Offset import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.streaming.StreamTest import org.apache.spark.sql.streaming.Trigger._ +import org.apache.spark.tags.SlowSQLTest import org.apache.spark.util.Utils /** * Test suite for streaming source naming and validation. * Tests cover the naming API, validation rules, and resolution pipeline. */ +@SlowSQLTest class StreamingSourceEvolutionSuite extends StreamTest { private def newMetadataDir = diff --git a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala index 19e499942e310..17aa7de7f86d5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/util/PartitionKeyedAccumulatorSuite.scala @@ -104,4 +104,26 @@ class PartitionKeyedAccumulatorSuite extends SparkFunSuite { assert(acc.accumulatedNumPartitions == 2) assert(acc.foldValues("")((s, v) => s + v).length == 2) // "c" + "b" (each partition once) } + + test("SPARK-58272: fold returns an atomic snapshot only after every partition completes") { + val accumulator = new PartitionKeyedAccumulator[Stats] + accumulator.add((0, (10L, 100L))) + + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.isEmpty) + + accumulator.add((1, (5L, 50L))) + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.contains((15L, 150L))) + + accumulator.add((1, (7L, 70L))) + assert(accumulator.foldValuesIfComplete(2, (0L, 0L)) { + case ((rows, bytes), (partitionRows, partitionBytes)) => + (rows + partitionRows, bytes + partitionBytes) + }.contains((17L, 170L))) + } } diff --git a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceSuite.scala b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceSuite.scala index f2a3812b59307..be6dca482fa17 100644 --- a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceSuite.scala @@ -25,7 +25,9 @@ import org.scalatest.PrivateMethodTester import org.apache.spark.{JobExecutionStatus, SparkFunSuite} import org.apache.spark.sql.execution.ui.{SparkPlanGraph, SparkPlanGraphCluster, SparkPlanGraphEdge, SparkPlanGraphNode, SQLExecutionUIData, SQLPlanMetric} -import org.apache.spark.status.api.v1.JacksonMessageWriter +import org.apache.spark.status.{AppStatusStore, StageDataWrapper} +import org.apache.spark.status.api.v1.{JacksonMessageWriter, StageData, StageStatus} +import org.apache.spark.util.kvstore.InMemoryStore object SqlResourceSuite { @@ -157,6 +159,89 @@ object SqlResourceSuite { assert(executionData.errorMessage == null) assert(executionData.rootExecutionId == 1) assert(executionData.modifiedConfigs == MODIFIED_CONFIGS) + // The fixture execution has no stages to aggregate, so the task time is + // unknown and reported as -1 rather than a misleading zero. + assert(executionData.totalTaskTime == -1L) + } + + private def newAppStore(stageDatas: Seq[StageData]): AppStatusStore = { + val kvStore = new InMemoryStore() + val store = new AppStatusStore(kvStore) + stageDatas.foreach { s => + kvStore.write(new StageDataWrapper(s, Set.empty, Map.empty)) + } + store + } + + private def stageData( + stageId: Int, + attemptId: Int, + executorRunTime: Long): StageData = { + new StageData( + status = StageStatus.COMPLETE, + stageId = stageId, + attemptId = attemptId, + numTasks = 1, + numActiveTasks = 0, + numCompleteTasks = 1, + numFailedTasks = 0, + numKilledTasks = 0, + numCompletedIndices = 1, + submissionTime = Some(new Date(0)), + firstTaskLaunchedTime = Some(new Date(0)), + completionTime = Some(new Date(1)), + failureReason = None, + executorDeserializeTime = 0, + executorDeserializeCpuTime = 0, + executorRunTime = executorRunTime, + executorCpuTime = 0, + resultSize = 0, + jvmGcTime = 0, + resultSerializationTime = 0, + memoryBytesSpilled = 0, + diskBytesSpilled = 0, + peakExecutionMemory = 0, + inputBytes = 0, + inputRecords = 0, + outputBytes = 0, + outputRecords = 0, + shuffleRemoteBlocksFetched = 0, + shuffleLocalBlocksFetched = 0, + shuffleFetchWaitTime = 0, + shuffleRemoteBytesRead = 0, + shuffleRemoteBytesReadToDisk = 0, + shuffleLocalBytesRead = 0, + shuffleReadBytes = 0, + shuffleReadRecords = 0, + shuffleCorruptMergedBlockChunks = 0, + shuffleMergedFetchFallbackCount = 0, + shuffleMergedRemoteBlocksFetched = 0, + shuffleMergedLocalBlocksFetched = 0, + shuffleMergedRemoteChunksFetched = 0, + shuffleMergedLocalChunksFetched = 0, + shuffleMergedRemoteBytesRead = 0, + shuffleMergedLocalBytesRead = 0, + shuffleRemoteReqsDuration = 0, + shuffleMergedRemoteReqsDuration = 0, + shuffleWriteBytes = 0, + shuffleWriteTime = 0, + shuffleWriteRecords = 0, + name = null, + description = None, + details = "", + schedulingPool = "", + rddIds = Seq.empty, + accumulatorUpdates = Seq.empty, + tasks = None, + executorSummary = None, + speculationSummary = None, + killedTasksSummary = Map.empty, + resourceProfileId = 0, + peakExecutorMetrics = None, + taskMetricsDistributions = None, + executorMetricsDistributions = None, + isShufflePushEnabled = false, + shuffleMergersCount = 0) } } @@ -174,7 +259,8 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { test("Prepare ExecutionData when details = false and planDescription = false") { val executionData = sqlResource invokePrivate prepareExecutionData( - sqlExecutionUIData, SparkPlanGraph(Seq.empty, Seq.empty), false, false) + sqlExecutionUIData, SparkPlanGraph(Seq.empty, Seq.empty), false, false, + newAppStore(Seq.empty)) verifyExpectedExecutionData(executionData, edges = Seq.empty, nodes = Seq.empty, planDescription = "") } @@ -182,7 +268,8 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { test("Prepare ExecutionData when details = true and planDescription = false") { val executionData = sqlResource invokePrivate prepareExecutionData( - sqlExecutionUIData, SparkPlanGraph(nodes, edges), true, false) + sqlExecutionUIData, SparkPlanGraph(nodes, edges), true, false, + newAppStore(Seq.empty)) verifyExpectedExecutionData( executionData, nodes = getNodes(), @@ -193,7 +280,8 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { test("Prepare ExecutionData when details = true and planDescription = true") { val executionData = sqlResource invokePrivate prepareExecutionData( - sqlExecutionUIData, SparkPlanGraph(nodes, edges), true, true) + sqlExecutionUIData, SparkPlanGraph(nodes, edges), true, true, + newAppStore(Seq.empty)) verifyExpectedExecutionData( executionData, nodes = getNodes(), @@ -204,7 +292,8 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { test("Prepare ExecutionData when details = true and planDescription = false and WSCG = off") { val executionData = sqlResource invokePrivate prepareExecutionData( - sqlExecutionUIData, SparkPlanGraph(nodesWhenCodegenIsOff, edges), true, false) + sqlExecutionUIData, SparkPlanGraph(nodesWhenCodegenIsOff, edges), true, false, + newAppStore(Seq.empty)) verifyExpectedExecutionData( executionData, nodes = getExpectedNodesWhenWholeStageCodegenIsOff(), @@ -237,7 +326,7 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { val executionData = sqlResource invokePrivate prepareExecutionData( d, - SparkPlanGraph(nodes, edges), true, true) + SparkPlanGraph(nodes, edges), true, true, newAppStore(Seq.empty)) assert(executionData.status == "FAILED") assert(executionData.errorMessage == "now you see me, now you don't") assert(executionData.rootExecutionId == 1) @@ -254,12 +343,41 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { errorMessage = None, queryId = null) val executionData = sqlResource invokePrivate prepareExecutionData( - d, SparkPlanGraph(Seq.empty, Seq.empty), false, false) + d, SparkPlanGraph(Seq.empty, Seq.empty), false, false, newAppStore(Seq.empty)) assert(executionData.queryId == null) assert(executionData.errorMessage == null) assert(executionData.rootExecutionId == -1) } + test("SPARK-58552: totalTaskTime aggregates executorRunTime across all attempts " + + "of all stages") { + // Stage 0 has two attempts (10 and 30 ms) - the retried attempt still + // consumed task time, so both count. Stage 1 has a single 20 ms attempt. + val store = newAppStore(Seq( + stageData(stageId = 0, attemptId = 0, executorRunTime = 10L), + stageData(stageId = 0, attemptId = 1, executorRunTime = 30L), + stageData(stageId = 1, attemptId = 0, executorRunTime = 20L))) + val exec = new SQLExecutionUIData( + executionId = 0, + rootExecutionId = 0, + description = "agg", + details = "", + physicalPlanDescription = "", + modifiedConfigs = Map.empty, + metrics = Seq.empty, + submissionTime = 0L, + completionTime = Some(new Date(1L)), + jobs = Map.empty[Int, JobExecutionStatus], + stages = Set(0, 1), + metricValues = Map.empty, + errorMessage = None, + queryId = null) + val executionData = + sqlResource invokePrivate prepareExecutionData( + exec, SparkPlanGraph(Seq.empty, Seq.empty), false, false, store) + assert(executionData.totalTaskTime == 60L) + } + test("SPARK-57987: JSON serialization of default modifiedConfigs and node desc") { val mapper = new JacksonMessageWriter().mapper val nodeWithEmptyDesc = Node(0, SCAN_TEXT, metrics = Seq.empty) @@ -279,5 +397,7 @@ class SqlResourceSuite extends SparkFunSuite with PrivateMethodTester { assert(executionJson.contains("\"modifiedConfigs\":{}")) assert(executionJson.contains( "\"nodes\":[{\"nodeId\":0,\"nodeName\":\"Scantext\",\"metrics\":[]}]")) + // totalTaskTime defaults to -1 (unknown) when not provided. + assert(executionData.totalTaskTime == -1L) } } diff --git a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala index e7072a1c0d28a..336d28f87c454 100644 --- a/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/status/api/v1/sql/SqlResourceWithActualMetricsSuite.scala @@ -197,6 +197,7 @@ class SqlResourceWithActualMetricsSuite assert((firstRow \ "status").extract[String].nonEmpty) assert((firstRow \ "description").extract[String] != null) assert((firstRow \ "duration").extract[Long] >= 0) + assert((firstRow \ "totalTaskTime").extract[Long] >= 0) // Test search filter val searchUrl = new URI( diff --git a/sql/core/src/test/scala/test/org/apache/spark/sql/ExpressionToColumnSuite.scala b/sql/core/src/test/scala/test/org/apache/spark/sql/ExpressionToColumnSuite.scala new file mode 100644 index 0000000000000..c03453b4856a4 --- /dev/null +++ b/sql/core/src/test/scala/test/org/apache/spark/sql/ExpressionToColumnSuite.scala @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package test.org.apache.spark.sql + +import org.apache.spark.sql.{Column, QueryTest, Row} +import org.apache.spark.sql.catalyst.expressions.{Expression, Literal} +import org.apache.spark.sql.classic.{ClassicConversions, ColumnConversions} +import org.apache.spark.sql.classic.ClassicConversions._ +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests the public Expression <-> Column conversions from a package outside of + * `org.apache.spark`, which is the only way to catch that a step of the path is package-private. + * Compiling this suite is as much a part of the test as running it. + */ +class ExpressionToColumnSuite extends QueryTest with SharedSparkSession { + + test("SPARK-49828: build a Column from an Expression via the Column companion") { + val e: Expression = Literal(1) + val c: Column = Column(e) + assert(ColumnConversions.expression(c) == e) + } + + test("SPARK-49828: build a Column from an Expression via ClassicConversions.column") { + val e: Expression = Literal(1) + val c: Column = ClassicConversions.column(e) + assert(ColumnConversions.expression(c) == e) + } + + test("SPARK-49828: a Column built from an Expression is usable in a query") { + val df = spark.range(2).select(Column(Literal(1)).as("one")) + checkAnswer(df, Seq(Row(1), Row(1))) + } +} diff --git a/sql/gen-sql-api-docs.py b/sql/gen-sql-api-docs.py index f5ed47f6ec7e4..1b53285fc1b45 100644 --- a/sql/gen-sql-api-docs.py +++ b/sql/gen-sql-api-docs.py @@ -22,7 +22,6 @@ from pyspark.java_gateway import launch_gateway - ExpressionInfo = namedtuple( "ExpressionInfo", "className name usage arguments examples note since deprecated group") diff --git a/sql/gen-sql-config-docs.py b/sql/gen-sql-config-docs.py index 4db22ff3b8e46..473d747c79567 100644 --- a/sql/gen-sql-config-docs.py +++ b/sql/gen-sql-config-docs.py @@ -17,16 +17,13 @@ import os import re - from collections import namedtuple from textwrap import dedent # To avoid adding a new direct dependency, we import markdown from within mkdocs. from mkdocs.structure.pages import markdown - from pyspark.java_gateway import launch_gateway - SQLConfEntry = namedtuple( "SQLConfEntry", ["name", "default", "description", "version"]) diff --git a/sql/gen-sql-functions-docs.py b/sql/gen-sql-functions-docs.py index 2ae00f6db8221..d39b5fcf79eb3 100644 --- a/sql/gen-sql-functions-docs.py +++ b/sql/gen-sql-functions-docs.py @@ -22,10 +22,8 @@ # To avoid adding a new direct dependency, we import markdown from within mkdocs. from mkdocs.structure.pages import markdown - from pyspark.java_gateway import launch_gateway - ExpressionInfo = namedtuple("ExpressionInfo", "name usage examples group") groups = { @@ -243,6 +241,8 @@ def generate_functions_examples_html(jvm, jspark, html_output_dir): """ print("Enabling TIME data type") jspark.sql("SET spark.sql.timeType.enabled = true") + print("Enabling parse_sql function") + jspark.sql("SET spark.sql.function.parseSql.enabled = true") print("Running SQL examples to generate formatted output.") for key, infos in _list_grouped_function_infos(jvm): examples = _make_pretty_examples(jspark, infos) diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtils.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtils.scala index df31ca311d46b..c4a7d1727bdcc 100644 --- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtils.scala +++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtils.scala @@ -147,8 +147,12 @@ object RowSetUtils { // types that reach this branch do not use the `nested` flag in `toHiveString`. Now, // Geospatial types use it for wrapping EWKT in quotes when nested = true, so we need // to set `nested` here to false to avoid spurious quotes for standalone geo values. + // String types need the same treatment: the fast path above matches only the + // default-collation StringType singleton, so CHAR/VARCHAR and collated strings land + // here and would otherwise be rendered as "value" instead of value. val nested = typ match { case _: GeometryType | _: GeographyType => false + case _: StringType => false case _ => true } toHiveString((row.get(ordinal), typ), nested, timeFormatters, binaryFormatter) diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetColumnsOperation.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetColumnsOperation.scala index b429086db22f2..9dff17b803764 100644 --- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetColumnsOperation.scala +++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetColumnsOperation.scala @@ -126,16 +126,18 @@ private[hive] class SparkGetColumnsOperation( } /** - * For boolean, numeric and datetime types, it returns the default size of its catalyst type + * For boolean, numeric and datetime types, this method returns the input type's default size. + * For CHAR(n) and VARCHAR(n), it returns the declared character length n. * For struct type, when its elements are fixed-size, the summation of all element sizes will be * returned. - * For array, map, string, and binaries, the column size is variable, return null as unknown. + * For array, map, unbounded string, and binaries, the column size is variable; return null. */ private def getColumnSize(typ: DataType): Option[Int] = typ match { case dt @ (BooleanType | _: NumericType | DateType | TimestampType | TimestampNTZType | CalendarIntervalType | NullType | _: AnsiIntervalType) => Some(dt.defaultSize) case c: CharType => Some(c.length) + case v: VarcharType => Some(v.length) case StructType(fields) => val sizeArr = fields.map(f => getColumnSize(f.dataType)) if (sizeArr.contains(None)) { @@ -146,6 +148,22 @@ private[hive] class SparkGetColumnsOperation( case other => None } + /** + * JDBC CHAR_OCTET_LENGTH is a byte capacity. Spark CHAR/VARCHAR lengths are in + * characters, so report 4 * n (UTF-8 maximum bytes per character), saturating at + * Int.MaxValue. Unbounded STRING and non-character types stay null (not applicable). + */ + private def getCharOctetLength(typ: DataType): Option[Int] = typ match { + case c: CharType => Some(maxUtf8OctetLength(c.length)) + case v: VarcharType => Some(maxUtf8OctetLength(v.length)) + case _ => None + } + + private def maxUtf8OctetLength(numChars: Int): Int = { + val maxChars = Int.MaxValue / 4 + if (numChars > maxChars) Int.MaxValue else numChars * 4 + } + /** * The number of fractional digits for this type. * Null is returned for data types where this is not applicable. @@ -222,7 +240,7 @@ private[hive] class SparkGetColumnsOperation( null, // COLUMN_DEF null, // SQL_DATA_TYPE null, // SQL_DATETIME_SUB - null, // CHAR_OCTET_LENGTH + getCharOctetLength(column.dataType).map(_.asInstanceOf[AnyRef]).orNull, ordinal.asInstanceOf[AnyRef], // ORDINAL_POSITION, 1-based (if (column.nullable) "YES" else "NO"), // IS_NULLABLE null, // SCOPE_CATALOG diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetFunctionsOperation.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetFunctionsOperation.scala index d79f2821f95e9..088416fab6aac 100644 --- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetFunctionsOperation.scala +++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkGetFunctionsOperation.scala @@ -88,7 +88,7 @@ private[hive] class SparkGetFunctionsOperation( s"Usage: ${info.getUsage}\nExtended Usage:${info.getExtended}", // REMARKS DatabaseMetaData.functionResultUnknown.asInstanceOf[AnyRef], // FUNCTION_TYPE info.getClassName) // SPECIFIC_NAME - rowSet.addRow(rowData); + rowSet.addRow(rowData) } } setState(OperationState.FINISHED) diff --git a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala index e1553ce9b368b..ab322d2cb4078 100644 --- a/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala +++ b/sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLCLIDriver.scala @@ -72,7 +72,8 @@ private[hive] object SparkSQLCLIDriver extends Logging { def installSignalHandler(): Unit = { HiveInterruptUtils.add(() => { if (SparkSQLEnv.sparkContext != null) { - SparkSQLEnv.sparkContext.cancelAllJobs() + SparkSQLEnv.sparkContext.cancelAllJobs( + "because the user interrupted the Spark SQL CLI with Ctrl+C") } }) } diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtilsSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtilsSuite.scala new file mode 100644 index 0000000000000..b2f0d462a3a7b --- /dev/null +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/RowSetUtilsSuite.scala @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.hive.thriftserver + +import scala.jdk.CollectionConverters._ + +import org.apache.hive.service.rpc.thrift.TProtocolVersion + +import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.{CharType, DataType, StringType, VarcharType} + +class RowSetUtilsSuite extends SparkFunSuite { + + private def stringValues(value: String, dataType: DataType): Seq[String] = { + val rowSet = RowSetUtils.toTRowSet( + 0, + Seq(Row(value)), + Array(dataType), + TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V8) + rowSet.getColumns.asScala.head.getStringVal.getValues.asScala.toSeq + } + + // Only the default-collation StringType singleton takes the fast path in toTColumn. CHAR, + // VARCHAR and collated strings fall through to the generic branch, which renders values with + // toHiveString and used to quote them there. + test("SPARK-58794: CHAR, VARCHAR and collated string values are not quoted") { + Seq( + CharType(4), + VarcharType(4), + StringType("UTF8_LCASE"), + StringType).foreach { dataType => + assert(stringValues("ab", dataType) === Seq("ab"), s"$dataType was rendered with quotes") + } + } +} diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SparkMetadataOperationSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SparkMetadataOperationSuite.scala index b3194cfdef3f6..af0e115b24929 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SparkMetadataOperationSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/SparkMetadataOperationSuite.scala @@ -316,6 +316,7 @@ class SparkMetadataOperationSuite extends HiveThriftServer2TestBase { |using parquet""".stripMargin withJdbcStatement(tableName) { statement => + statement.execute(s"SET ${SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key}=true") statement.execute(ddl) val databaseMetaData = statement.getConnection.getMetaData @@ -338,10 +339,21 @@ class SparkMetadataOperationSuite extends HiveThriftServer2TestBase { val colSize = rowSet.getInt("COLUMN_SIZE") schema(pos).dataType match { - case StringType | BinaryType | _: ArrayType | _: MapType | _: VarcharType => + case StringType | BinaryType | _: ArrayType | _: MapType => assert(colSize === 0) + case c: CharType => assert(colSize === c.length) + case v: VarcharType => assert(colSize === v.length) case o => assert(colSize === o.defaultSize) } + if (schema(pos).name == "c17") assert(colSize === 255) + if (schema(pos).name == "c18") assert(colSize === 1024) + + val octetLength = rowSet.getInt("CHAR_OCTET_LENGTH") + schema(pos).dataType match { + case c: CharType => assert(octetLength === c.length * 4) + case v: VarcharType => assert(octetLength === v.length * 4) + case _ => assert(octetLength === 0) // JDBC getInt on SQL NULL + } assert(rowSet.getInt("BUFFER_LENGTH") === 0) // not used val decimalDigits = rowSet.getInt("DECIMAL_DIGITS") @@ -372,6 +384,23 @@ class SparkMetadataOperationSuite extends HiveThriftServer2TestBase { } } + test("SPARK-58794: result metadata preserves CHAR and VARCHAR") { + withJdbcStatement() { statement => + statement.execute(s"SET ${SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key}=true") + val resultSet = statement.executeQuery( + "SELECT CAST('ab' AS CHAR(4)) AS c, CAST('cd' AS VARCHAR(6)) AS v") + assert(resultSet.next()) + + val metadata = resultSet.getMetaData + assert(metadata.getColumnType(1) === java.sql.Types.CHAR) + assert(metadata.getColumnTypeName(1) === "char") + assert(metadata.getPrecision(1) === 4) + assert(metadata.getColumnType(2) === java.sql.Types.VARCHAR) + assert(metadata.getColumnTypeName(2) === "varchar") + assert(metadata.getPrecision(2) === 6) + } + } + test("get columns operation should handle interval column properly") { val viewName = "view_interval" val ddl = s"CREATE GLOBAL TEMP VIEW $viewName as select interval 1 day as i" diff --git a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala index fbc150f79ab43..f8fd426881e55 100644 --- a/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala +++ b/sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerWithSparkContextSuite.scala @@ -214,7 +214,7 @@ trait ThriftServerWithSparkContextSuite extends SharedThriftServer { val sessionHandle = client.openSession(user, "") val infoValue = client.getInfo(sessionHandle, GetInfoType.CLI_ODBC_KEYWORDS) // scalastyle:off line.size.limit - assert(infoValue.getStringValue == "ADD,AFTER,AGGREGATE,ALIGN,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,APPLY,APPROX,ARCHIVE,ARRAY,AS,ASC,ASENSITIVE,ASOF,AT,ATOMIC,AUTHORIZATION,AUTO,BEGIN,BERNOULLI,BETWEEN,BIGINT,BIN,BINARY,BINDING,BIN_DISTRIBUTE_RATIO,BIN_END,BIN_START,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALL,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CDC,CHANGE,CHANGES,CHAR,CHARACTER,CHECK,CLEAR,CLOSE,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONDITION,CONSTRAINT,CONTAINS,CONTINUE,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATABASE,CURRENT_DATE,CURRENT_PATH,CURRENT_SCHEMA,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,CURSOR,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFAULT_PATH,DEFINED,DEFINER,DELAY,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTANCE,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,ELSEIF,END,ENFORCED,ESCAPE,ESCAPED,EVOLUTION,EXACT,EXCEPT,EXCHANGE,EXCLUDE,EXCLUSIVE,EXECUTE,EXISTS,EXIT,EXPLAIN,EXPORT,EXTEND,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FLOW,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FOUND,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GEOGRAPHY,GEOMETRY,GLOBAL,GRANT,GROUP,GROUPING,HANDLER,HAVING,HISTORY,HOUR,HOURS,IDENTIFIED,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCLUSIVE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSENSITIVE,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,JSON,KEY,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LEVEL,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCALTIME,LOCATION,LOCK,LOCKS,LOGICAL,LONG,LOOP,MACRO,MAP,MATCHED,MATCH_CONDITION,MATERIALIZED,MAX,MEASURE,MERGE,METRICS,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NEAREST,NEXT,NO,NONE,NORELY,NOT,NULL,NULLS,NUMERIC,OF,OFFSET,ON,ONLY,OPEN,OPTION,OPTIONS,OR,ORDER,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PATH,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROCEDURE,PROCEDURES,PROPERTIES,PURGE,QUALIFY,QUARTER,QUERY,RANGE,READ,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,RECURSION,RECURSIVE,REDUCE,REFERENCES,REFRESH,RELY,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCD,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SEQUENCE,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SIMILARITY,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,SQLEXCEPTION,SQLSTATE,START,STATISTICS,STORED,STRATIFY,STREAM,STREAMING,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM,SYSTEM_PATH,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRACK,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNIFORM,UNION,UNIQUE,UNKNOWN,UNLOCK,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUE,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WATERMARK,WEEK,WEEKS,WHEN,WHERE,WHILE,WIDTH,WINDOW,WITH,WITHIN,WITHOUT,X,YEAR,YEARS,ZONE") + assert(infoValue.getStringValue == "ADD,AFTER,AGGREGATE,ALIGN,ALL,ALTER,ALWAYS,ANALYZE,AND,ANTI,ANY,ANY_VALUE,APPLY,APPROX,ARCHIVE,ARRAY,AS,ASC,ASENSITIVE,ASOF,AT,ATOMIC,AUTHORIZATION,AUTO,BEGIN,BERNOULLI,BETWEEN,BIGINT,BIN,BINARY,BINDING,BIN_DISTRIBUTE_RATIO,BIN_END,BIN_START,BOOLEAN,BOTH,BUCKET,BUCKETS,BY,BYTE,CACHE,CALL,CALLED,CASCADE,CASE,CAST,CATALOG,CATALOGS,CDC,CHANGE,CHANGES,CHAR,CHARACTER,CHECK,CLEAR,CLOSE,CLUSTER,CLUSTERED,CODEGEN,COLLATE,COLLATION,COLLATIONS,COLLECTION,COLUMN,COLUMNS,COMMENT,COMMIT,COMPACT,COMPACTIONS,COMPENSATION,COMPUTE,CONCATENATE,CONDITION,CONDITIONAL,CONSTRAINT,CONTAINS,CONTINUE,COST,CREATE,CROSS,CUBE,CURRENT,CURRENT_DATABASE,CURRENT_DATE,CURRENT_PATH,CURRENT_SCHEMA,CURRENT_TIME,CURRENT_TIMESTAMP,CURRENT_USER,CURSOR,DATA,DATABASE,DATABASES,DATE,DATEADD,DATEDIFF,DATE_ADD,DATE_DIFF,DAY,DAYOFYEAR,DAYS,DBPROPERTIES,DEC,DECIMAL,DECLARE,DEFAULT,DEFAULT_PATH,DEFINED,DEFINER,DELAY,DELETE,DELIMITED,DESC,DESCRIBE,DETERMINISTIC,DFS,DIRECTORIES,DIRECTORY,DISTANCE,DISTINCT,DISTRIBUTE,DIV,DO,DOUBLE,DROP,ELSE,ELSEIF,EMPTY,END,ENFORCED,ERROR,ESCAPE,ESCAPED,EVOLUTION,EXACT,EXCEPT,EXCHANGE,EXCLUDE,EXCLUSIVE,EXECUTE,EXISTS,EXIT,EXPLAIN,EXPORT,EXTEND,EXTENDED,EXTERNAL,EXTRACT,FALSE,FETCH,FIELDS,FILEFORMAT,FILTER,FIRST,FLOAT,FLOW,FOLLOWING,FOR,FOREIGN,FORMAT,FORMATTED,FOUND,FROM,FULL,FUNCTION,FUNCTIONS,GENERATED,GEOGRAPHY,GEOMETRY,GLOBAL,GRANT,GROUP,GROUPING,HANDLER,HAVING,HISTORY,HOUR,HOURS,IDENTIFIED,IDENTIFIER,IDENTITY,IF,IGNORE,ILIKE,IMMEDIATE,IMPORT,IN,INCLUDE,INCLUSIVE,INCREMENT,INDEX,INDEXES,INNER,INPATH,INPUT,INPUTFORMAT,INSENSITIVE,INSERT,INT,INTEGER,INTERSECT,INTERVAL,INTO,INVOKER,IS,ITEMS,ITERATE,JOIN,JSON,JSON_EXISTS,JSON_QUERY,JSON_TABLE,JSON_VALUE,KEEP,KEY,KEYS,LANGUAGE,LAST,LATERAL,LAZY,LEADING,LEAVE,LEFT,LEVEL,LIKE,LIMIT,LINES,LIST,LOAD,LOCAL,LOCALTIME,LOCATION,LOCK,LOCKS,LOGICAL,LONG,LOOP,MACRO,MAP,MATCHED,MATCH_CONDITION,MATERIALIZED,MAX,MEASURE,MERGE,METRICS,MICROSECOND,MICROSECONDS,MILLISECOND,MILLISECONDS,MINUS,MINUTE,MINUTES,MODIFIES,MONTH,MONTHS,MSCK,NAME,NAMESPACE,NAMESPACES,NANOSECOND,NANOSECONDS,NATURAL,NEAREST,NEXT,NO,NONE,NORELY,NOT,NULL,NULLS,NUMERIC,OBJECT,OF,OFFSET,OMIT,ON,ONLY,OPEN,OPTION,OPTIONS,OR,ORDER,ORDINALITY,OUT,OUTER,OUTPUTFORMAT,OVER,OVERLAPS,OVERLAY,OVERWRITE,PARTITION,PARTITIONED,PARTITIONS,PATH,PERCENT,PIVOT,PLACING,POSITION,PRECEDING,PRIMARY,PRINCIPALS,PROCEDURE,PROCEDURES,PROPERTIES,PURGE,QUALIFY,QUARTER,QUERY,QUOTES,RANGE,READ,READS,REAL,RECORDREADER,RECORDWRITER,RECOVER,RECURSION,RECURSIVE,REDUCE,REFERENCES,REFRESH,RELY,RENAME,REPAIR,REPEAT,REPEATABLE,REPLACE,RESET,RESPECT,RESTRICT,RETURN,RETURNING,RETURNS,REVOKE,RIGHT,ROLE,ROLES,ROLLBACK,ROLLUP,ROW,ROWS,SCD,SCHEMA,SCHEMAS,SECOND,SECONDS,SECURITY,SELECT,SEMI,SEPARATED,SEQUENCE,SERDE,SERDEPROPERTIES,SESSION_USER,SET,SETS,SHORT,SHOW,SIMILARITY,SINGLE,SKEWED,SMALLINT,SOME,SORT,SORTED,SOURCE,SPECIFIC,SQL,SQLEXCEPTION,SQLSTATE,START,STATISTICS,STORED,STRATIFY,STREAM,STREAMING,STRING,STRUCT,SUBSTR,SUBSTRING,SYNC,SYSTEM,SYSTEM_PATH,SYSTEM_TIME,SYSTEM_VERSION,TABLE,TABLES,TABLESAMPLE,TARGET,TBLPROPERTIES,TERMINATED,THEN,TIME,TIMEDIFF,TIMESTAMP,TIMESTAMPADD,TIMESTAMPDIFF,TIMESTAMP_LTZ,TIMESTAMP_NTZ,TINYINT,TO,TOUCH,TRACK,TRAILING,TRANSACTION,TRANSACTIONS,TRANSFORM,TRIM,TRUE,TRUNCATE,TRY_CAST,TYPE,UNARCHIVE,UNBOUNDED,UNCACHE,UNCONDITIONAL,UNIFORM,UNION,UNIQUE,UNKNOWN,UNLOCK,UNNEST,UNPIVOT,UNSET,UNTIL,UPDATE,USE,USER,USING,VALUE,VALUES,VAR,VARCHAR,VARIABLE,VARIANT,VERSION,VIEW,VIEWS,VOID,WATERMARK,WEEK,WEEKS,WHEN,WHERE,WHILE,WIDTH,WINDOW,WITH,WITHIN,WITHOUT,WRAPPER,X,YEAR,YEARS,ZONE") // scalastyle:on line.size.limit } } diff --git a/sql/hive/src/main/java/org/apache/hadoop/hive/ql/exec/HiveFunctionRegistryUtils.java b/sql/hive/src/main/java/org/apache/hadoop/hive/ql/exec/HiveFunctionRegistryUtils.java index 333ae0151a5a6..1cb527d92d05c 100644 --- a/sql/hive/src/main/java/org/apache/hadoop/hive/ql/exec/HiveFunctionRegistryUtils.java +++ b/sql/hive/src/main/java/org/apache/hadoop/hive/ql/exec/HiveFunctionRegistryUtils.java @@ -18,6 +18,10 @@ package org.apache.hadoop.hive.ql.exec; import org.apache.hadoop.hive.ql.metadata.HiveException; +import org.apache.hadoop.hive.ql.udf.SettableUDF; +import org.apache.hadoop.hive.ql.udf.generic.GenericUDF; +import org.apache.hadoop.hive.ql.udf.generic.GenericUDFBridge; +import org.apache.hadoop.hive.ql.udf.generic.GenericUDFMacro; import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector.PrimitiveCategory; import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils.PrimitiveGrouping; import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector.Category; @@ -29,6 +33,8 @@ import java.util.Iterator; import java.util.List; +import org.apache.hadoop.util.ReflectionUtils; + import org.apache.spark.internal.SparkLogger; import org.apache.spark.internal.SparkLoggerFactory; @@ -339,4 +345,53 @@ static void filterMethodsByTypeAffinity(List<Method> udfMethods, List<TypeInfo> } } } + + /** + * Create a copy of an existing GenericUDF. + */ + public static GenericUDF cloneGenericUDF(GenericUDF genericUDF) { + if (null == genericUDF) { + return null; + } + + GenericUDF clonedUDF; + if (genericUDF instanceof GenericUDFBridge) { + GenericUDFBridge bridge = (GenericUDFBridge) genericUDF; + clonedUDF = new GenericUDFBridge(bridge.getUdfName(), bridge.isOperator(), + bridge.getUdfClassName()); + } else if (genericUDF instanceof GenericUDFMacro) { + GenericUDFMacro bridge = (GenericUDFMacro) genericUDF; + clonedUDF = new GenericUDFMacro(bridge.getMacroName(), bridge.getBody().clone(), + bridge.getColNames(), bridge.getColTypes()); + } else { + clonedUDF = ReflectionUtils.newInstance(genericUDF.getClass(), null); + } + + if (clonedUDF != null) { + // Copy info that may be required in the new copy. + // The SettableUDF calls below could be replaced using this mechanism as well. + try { + genericUDF.copyToNewInstance(clonedUDF); + } catch (UDFArgumentException err) { + throw new IllegalArgumentException(err); + } + + // The original may have settable info that needs to be added to the new copy. + if (genericUDF instanceof SettableUDF) { + try { + TypeInfo typeInfo = ((SettableUDF)genericUDF).getTypeInfo(); + if (typeInfo != null) { + ((SettableUDF)clonedUDF).setTypeInfo(typeInfo); + } + } catch (UDFArgumentException err) { + // In theory this should not happen - if the original copy of the UDF had this + // data, we should be able to set the UDF copy with this same settableData. + LOG.error("Unable to add settable data to UDF " + genericUDF.getClass()); + throw new IllegalArgumentException(err); + } + } + } + + return clonedUDF; + } } diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala index 7f76132e25727..76d9a08f603d0 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveInspectors.scala @@ -967,6 +967,9 @@ private[hive] trait HiveInspectors { toInspector(sqlType) // Hive has no TIME type, so it cannot be represented by any Hive object inspector. case _: TimeType => throw unsupportedHiveType(dataType) + // Hive has no nanosecond-precision timestamp type, so it cannot be represented by any Hive + // object inspector. Reject it instead of silently downgrading to microsecond precision. + case _: AnyTimestampNanoType => throw unsupportedHiveType(dataType) } private def unsupportedHiveType(dataType: DataType): AnalysisException = { @@ -1044,6 +1047,10 @@ private[hive] trait HiveInspectors { // Hive has no TIME type, so a TIME constant cannot be mapped to a Hive object inspector. case Literal(_, dt: TimeType) => throw unsupportedHiveType(dt) + // Hive has no nanosecond-precision timestamp type, so such a constant cannot be mapped to a + // Hive object inspector. + case Literal(_, dt: AnyTimestampNanoType) => + throw unsupportedHiveType(dt) // We will enumerate all of the possible constant expressions, throw exception if we missed case Literal(_, dt) => throw SparkException.internalError(s"Hive doesn't support the constant type [$dt].") @@ -1064,6 +1071,7 @@ private[hive] trait HiveInspectors { case _: CurrentTime => false case _: CurrentTimestampLike => false case _: LocalTimestamp => false + case _: LocalTimestampNanos => false case _ => e.children.forall(canEarlyEval) } @@ -1297,6 +1305,8 @@ private[hive] trait HiveInspectors { case _: YearMonthIntervalType => intervalYearMonthTypeInfo // Hive has no TIME type, so there is no Hive TypeInfo to map it to. case _: TimeType => throw unsupportedHiveType(dt) + // Hive has no nanosecond-precision timestamp type, so there is no Hive TypeInfo to map it to. + case _: AnyTimestampNanoType => throw unsupportedHiveType(dt) case dt => throw unsupportedHiveType(dt) } } diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala index 32d8928836976..445cc80e625c4 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/client/HiveShim.scala @@ -351,7 +351,12 @@ private[client] class Shim_v2_0 extends Shim with Logging { val filter = convertFilters(table, predicates) val partitions = - if (filter.isEmpty) { + if (referencesCharVarcharPartitionKey(catalogTable, predicates)) { + // convertFilters skips CHAR/VARCHAR keys. If another conjunct is supported, its + // non-empty filter would otherwise fetch a superset before residual pruning. + prunePartitionsFastFallback( + hive, table, catalogTable, predicates, forceClientSide = true) + } else if (filter.isEmpty) { prunePartitionsFastFallback(hive, table, catalogTable, predicates) } else { logDebug(s"Hive metastore filter is '$filter'.") @@ -391,11 +396,29 @@ private[client] class Shim_v2_0 extends Shim with Logging { partitions.asScala.toSeq } + private def referencesCharVarcharPartitionKey( + catalogTable: CatalogTable, + predicates: Seq[Expression]): Boolean = { + SQLConf.get.charVarcharStandardSemantics && { + val charVarcharPartNames = catalogTable.partitionSchema.fields.collect { + case f if CharVarcharUtils.hasCharVarchar(f.dataType) => f.name + } + charVarcharPartNames.nonEmpty && { + val resolver = SQLConf.get.resolver + predicates.exists(_.exists { + case a: Attribute => charVarcharPartNames.exists(n => resolver(n, a.name)) + case _ => false + }) + } + } + } + private def prunePartitionsFastFallback( hive: Hive, table: Table, catalogTable: CatalogTable, - predicates: Seq[Expression]): java.util.Collection[Partition] = { + predicates: Seq[Expression], + forceClientSide: Boolean = false): java.util.Collection[Partition] = { val timeZoneId = SQLConf.get.sessionLocalTimeZone // Because there is no way to know whether the partition properties has timeZone, @@ -408,9 +431,9 @@ private[client] class Shim_v2_0 extends Shim with Logging { } } - if (!SQLConf.get.metastorePartitionPruningFastFallback || - predicates.isEmpty || - predicates.exists(hasTimeZoneAwareExpression)) { + if ((!forceClientSide && !SQLConf.get.metastorePartitionPruningFastFallback) || + predicates.isEmpty || + predicates.exists(hasTimeZoneAwareExpression)) { recordHiveCall() hive.getAllPartitionsOf(table) } else { diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveFileFormat.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveFileFormat.scala index acbc72fbf7e0f..4545acc7f2035 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveFileFormat.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/HiveFileFormat.scala @@ -41,7 +41,7 @@ import org.apache.spark.sql.execution.datasources.{FileFormat, OutputWriter, Out import org.apache.spark.sql.hive.{HiveInspectors, HiveTableUtil} import org.apache.spark.sql.internal.SessionStateHelper import org.apache.spark.sql.sources.DataSourceRegister -import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType, TimeType, UserDefinedType} +import org.apache.spark.sql.types.{AnyTimestampNanoType, ArrayType, DataType, MapType, StructType, TimeType, UserDefinedType} import org.apache.spark.util.SerializableJobConf /** @@ -120,6 +120,10 @@ case class HiveFileFormat(fileSinkConf: FileSinkDesc) // (recursing into nested types) while preserving the default behavior for all other types. case _: TimeType => false + // Hive has no nanosecond-precision timestamp type. Reject it explicitly rather than silently + // downgrading to microsecond precision (which the Hive serde would otherwise do). + case _: AnyTimestampNanoType => false + case st: StructType => st.forall { f => supportDataType(f.dataType) } case ArrayType(elementType, _) => supportDataType(elementType) diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFEvaluators.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFEvaluators.scala index 866d88dee8783..f09ee7eed93ef 100644 --- a/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFEvaluators.scala +++ b/sql/hive/src/main/scala/org/apache/spark/sql/hive/hiveUDFEvaluators.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.hive.HiveShim.HiveFunctionWrapper import org.apache.spark.sql.types.DataType abstract class HiveUDFEvaluatorBase[UDFType <: AnyRef]( - funcWrapper: HiveFunctionWrapper, children: Seq[Expression]) + protected val funcWrapper: HiveFunctionWrapper, children: Seq[Expression]) extends HiveInspectors with Serializable { @transient @@ -115,6 +115,15 @@ class HiveGenericUDFEvaluator( funcWrapper: HiveFunctionWrapper, children: Seq[Expression]) extends HiveUDFEvaluatorBase[GenericUDF](funcWrapper, children) { + // SPARK-58792: copied expression nodes (e.g. via withNewChildrenInternal) share one + // HiveFunctionWrapper, whose cached GenericUDF instance is mutable: initialize() + // rewrites its converters and output holders based on the arguments of whichever + // copy initialized it last. Give every evaluator its own clone so copied nodes + // cannot corrupt each other. + @transient + override lazy val function: GenericUDF = + HiveFunctionRegistryUtils.cloneGenericUDF(funcWrapper.createFunction[GenericUDF]()) + @transient private lazy val argumentInspectors = children.map(toInspector).toArray diff --git a/sql/hive/src/test/resources/conf/binding-policy-exceptions/configs-without-binding-policy-exceptions b/sql/hive/src/test/resources/conf/binding-policy-exceptions/configs-without-binding-policy-exceptions index 37598f184c021..69f1adf5f62a6 100644 --- a/sql/hive/src/test/resources/conf/binding-policy-exceptions/configs-without-binding-policy-exceptions +++ b/sql/hive/src/test/resources/conf/binding-policy-exceptions/configs-without-binding-policy-exceptions @@ -453,7 +453,6 @@ spark.sql.ansi.enabled spark.sql.ansi.enforceReservedKeywords spark.sql.ansi.relationPrecedence spark.sql.artifact.cacheStorageLevel -spark.sql.artifact.copyFromLocalToFs.allowDestLocal spark.sql.artifact.isolation.alwaysApplyClassloader spark.sql.artifact.isolation.enabled spark.sql.assumeAnsiFalseIfNotPersisted.enabled diff --git a/sql/hive/src/test/resources/data/scripts/cat.py b/sql/hive/src/test/resources/data/scripts/cat.py index 420d9f832a184..392a2f02fec31 100644 --- a/sql/hive/src/test/resources/data/scripts/cat.py +++ b/sql/hive/src/test/resources/data/scripts/cat.py @@ -16,8 +16,8 @@ # specific language governing permissions and limitations # under the License. # -import sys import os +import sys table_name = None if os.environ in 'hive_streaming_tablename': diff --git a/sql/hive/src/test/resources/data/scripts/input20_script.py b/sql/hive/src/test/resources/data/scripts/input20_script.py index 1dc5567d12ae9..97729c4602e75 100644 --- a/sql/hive/src/test/resources/data/scripts/input20_script.py +++ b/sql/hive/src/test/resources/data/scripts/input20_script.py @@ -16,8 +16,9 @@ # specific language governing permissions and limitations # under the License. # -import sys import re +import sys + line = sys.stdin.readline() x = 1 while line: diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveCharVarcharTestSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveCharVarcharTestSuite.scala index 90cb5501ee6f6..73a9b312d5ec5 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveCharVarcharTestSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveCharVarcharTestSuite.scala @@ -17,9 +17,11 @@ package org.apache.spark.sql.hive -import org.apache.spark.sql.{CharVarcharTestSuite, Row} +import org.apache.spark.metrics.source.HiveCatalogMetrics +import org.apache.spark.sql.{CharVarcharTestSuite, QueryTest, Row} import org.apache.spark.sql.execution.command.CharVarcharDDLTestBase import org.apache.spark.sql.hive.test.TestHiveSingleton +import org.apache.spark.sql.internal.SQLConf class HiveCharVarcharTestSuite extends CharVarcharTestSuite with TestHiveSingleton { @@ -91,6 +93,72 @@ class HiveCharVarcharTestSuite extends CharVarcharTestSuite with TestHiveSinglet } } } + + test("SPARK-59001: CHAR/VARCHAR partition filters prune client-side") { + // Keep the relation a HiveTableRelation, otherwise the scan is converted to a file index + // and never reaches HiveShim's metastore filter conversion. Hive convertFilters still + // skips CHAR/VARCHAR keys; pruning is client-side under standardSemantics. + withSQLConf( + SQLConf.CHAR_VARCHAR_STANDARD_SEMANTICS.key -> "true", + SQLConf.HIVE_METASTORE_PARTITION_PRUNING.key -> "true", + HiveUtils.CONVERT_METASTORE_PARQUET.key -> "false") { + val partitionValues = Seq("a", "b", "c", "d", "e") + + def withHivePartTable(partitionType: String)(body: => Unit): Unit = { + withTable("std_hive_part") { + sql( + s"""CREATE TABLE std_hive_part (i INT, p $partitionType) + |USING $format PARTITIONED BY (p)""".stripMargin) + partitionValues.foreach { v => + sql(s"INSERT INTO std_hive_part PARTITION (p='$v') VALUES (1)") + } + body + } + } + + withHivePartTable("CHAR(5)") { + HiveCatalogMetrics.reset() + // Store assignment pads CHAR(5); compare is not PAD SPACE, so the literal must match. + // checkToRDD = false: checkAnswer would otherwise scan twice and double the metric. + QueryTest.checkAnswer( + sql("SELECT i FROM std_hive_part WHERE p = 'a '"), + Seq(Row(1)), + checkToRDD = false) + assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount === 1) + HiveCatalogMetrics.reset() + QueryTest.checkAnswer( + sql("SELECT i FROM std_hive_part WHERE p = 'a'"), + Nil, + checkToRDD = false) + assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount === 0) + } + + withHivePartTable("VARCHAR(5)") { + HiveCatalogMetrics.reset() + QueryTest.checkAnswer( + sql("SELECT i FROM std_hive_part WHERE p = 'a'"), + Seq(Row(1)), + checkToRDD = false) + assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount === 1) + } + + // A supported conjunct must not bypass client-side pruning of the CHAR predicate. + withTable("std_hive_part") { + sql( + s"""CREATE TABLE std_hive_part (i INT, ds INT, p CHAR(5)) + |USING $format PARTITIONED BY (ds, p)""".stripMargin) + partitionValues.foreach { v => + sql(s"INSERT INTO std_hive_part PARTITION (ds=1, p='$v') VALUES (1)") + } + HiveCatalogMetrics.reset() + QueryTest.checkAnswer( + sql("SELECT i FROM std_hive_part WHERE ds = 1 AND p = 'a '"), + Seq(Row(1)), + checkToRDD = false) + assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount === 1) + } + } + } } class HiveCharVarcharDDLTestSuite extends CharVarcharDDLTestBase with TestHiveSingleton { diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveExternalCatalogSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveExternalCatalogSuite.scala index db522b72e4cca..fb74d83c04d9d 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveExternalCatalogSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveExternalCatalogSuite.scala @@ -23,15 +23,17 @@ import org.apache.logging.log4j.Level import org.apache.spark.SparkConf import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.catalog._ +import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.execution.QueryExecutionException import org.apache.spark.sql.execution.command.DDLUtils -import org.apache.spark.sql.types.{StringType, StructField, StructType} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampLTZNanosType, TimestampNTZNanosType} /** * Test suite for the [[HiveExternalCatalog]]. */ -class HiveExternalCatalogSuite extends ExternalCatalogSuite { +class HiveExternalCatalogSuite extends ExternalCatalogSuite with SQLHelper { private val externalCatalog: HiveExternalCatalog = { val catalog = new HiveExternalCatalog(new SparkConf, new Configuration) @@ -243,6 +245,50 @@ class HiveExternalCatalogSuite extends ExternalCatalogSuite { assert(DataTypeUtils.sameType(alteredTable.schema, newSchema)) } + test("SPARK-57835: restore a persisted nanos-typed table when the preview flag is off") { + val catalog = newBasicCatalog() + val tableName = "nanos_tbl" + + // Nanos types are not Hive-compatible (SPARK-57831), so a datasource table persists an empty + // schema in the metastore and the real schema as JSON in the table properties. The table is + // created with the preview flag on (the test default via Utils.isTesting). + val nanosSchema = StructType(Seq( + StructField("id", IntegerType), + StructField("ntz", TimestampNTZNanosType(9)), + StructField("ltz", TimestampLTZNanosType(7)))) + + val tableDDL = CatalogTable( + identifier = TableIdentifier(tableName, Some("db1")), + tableType = CatalogTableType.MANAGED, + storage = storageFormat, + schema = nanosSchema, + provider = Some("parquet")) + + catalog.createTable(tableDDL, ignoreIfExists = false) + + // Because nanos types are not Hive-compatible, the metastore-visible (raw) schema is the + // placeholder EMPTY_DATA_SCHEMA; the true schema lives in the Spark-specific table properties. + val rawTable = externalCatalog.getRawTable("db1", tableName) + assert(rawTable.schema == HiveExternalCatalog.EMPTY_DATA_SCHEMA) + + // Read-through policy (SPARK-57835): with the preview flag off, catalog restoration must still + // reconstruct the persisted nanos schema so the table remains describable and droppable. + // Before this change getTable threw FEATURE_NOT_ENABLED here. + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "false") { + val restored = externalCatalog.getTable("db1", tableName) + assert(DataTypeUtils.sameType(restored.schema, nanosSchema)) + // The restored nanos columns render with their precision (used by DESCRIBE / SHOW CREATE). + assert(restored.schema("ntz").dataType === TimestampNTZNanosType(9)) + assert(restored.schema("ltz").dataType === TimestampLTZNanosType(7)) + } + + // And it still round-trips with the flag on. + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + assert(DataTypeUtils.sameType( + externalCatalog.getTable("db1", tableName).schema, nanosSchema)) + } + } + test("SPARK-50137: Avoid fallback to Hive-incompatible ways on thrift exception") { val hadoopConf = new Configuration() // Use an unavailable uri to mock client connection timeout. diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveInspectorSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveInspectorSuite.scala index b7fb506f07b53..a88dcd0ef56aa 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveInspectorSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/HiveInspectorSuite.scala @@ -308,4 +308,24 @@ class HiveInspectorSuite extends SparkFunSuite with HiveInspectors { condition = "UNSUPPORTED_DATATYPE", parameters = expectedParams) } + + test("SPARK-57815: nanosecond timestamp types are unsupported in Hive object inspectors") { + Seq( + TimestampNTZNanosType(9), TimestampNTZNanosType(7), + TimestampLTZNanosType(9), TimestampLTZNanosType(8)).foreach { nanosType => + val expectedParams = Map("typeName" -> s"\"${nanosType.sql}\"") + checkError( + exception = intercept[AnalysisException](toInspector(nanosType)), + condition = "UNSUPPORTED_DATATYPE", + parameters = expectedParams) + checkError( + exception = intercept[AnalysisException](toInspector(Literal.create(null, nanosType))), + condition = "UNSUPPORTED_DATATYPE", + parameters = expectedParams) + checkError( + exception = intercept[AnalysisException](nanosType.toTypeInfo), + condition = "UNSUPPORTED_DATATYPE", + parameters = expectedParams) + } + } } diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala index 7385410e1a712..db1743862e907 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/InsertSuite.scala @@ -708,6 +708,40 @@ class InsertSuite extends QueryTest with TestHiveSingleton with BeforeAndAfter { } } + test("SPARK-57815: nanosecond timestamp is unsupported when writing to a Hive serde directory") { + // Disable native data source conversion so that the write goes through the Hive serde path + // (HiveFileFormat) instead of a native data source that supports nanosecond timestamps. + withSQLConf( + HiveUtils.CONVERT_METASTORE_INSERT_DIR.key -> "false", + SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "true") { + Seq("ORC", "PARQUET").foreach { fileFormat => + Seq( + "TIMESTAMP_NTZ(9)" -> TimestampNTZNanosType(9), + "TIMESTAMP_LTZ(9)" -> TimestampLTZNanosType(9)).foreach { case (typeStr, dt) => + withTempDir { dir => + // InsertIntoHiveDirCommand wraps the failure in a SparkException, so assert on the + // cause. Rejecting here avoids a silent downgrade to microsecond precision. + val e = intercept[SparkException] { + sql( + s""" + |INSERT OVERWRITE LOCAL DIRECTORY '${dir.toURI.getPath}' + |STORED AS $fileFormat + |SELECT CAST('2025-01-06 12:30:45.123456789' AS $typeStr) AS c + """.stripMargin) + } + checkError( + exception = e.getCause.asInstanceOf[AnalysisException], + condition = "UNSUPPORTED_DATA_TYPE_FOR_DATASOURCE", + parameters = Map( + "columnName" -> "`c`", + "columnType" -> s"\"${dt.sql}\"", + "format" -> "Hive")) + } + } + } + } + } + test("insert overwrite to dir from temp table") { withTempView("test_insert_table") { spark.range(10).selectExpr("id", "id AS str").createOrReplaceTempView("test_insert_table") diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveDDLSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveDDLSuite.scala index 6ce4e8702c32f..6b79870979167 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveDDLSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveDDLSuite.scala @@ -3417,4 +3417,94 @@ class HiveDDLSuite any[String], any[String], any[StructType]) } } + + test("SPARK-57835: read persisted nanos-typed tables when the preview flag is off") { + withTable("nanos_ddl_tbl") { + // Create the table with the preview flag on (the test default via Utils.isTesting). + sql( + """CREATE TABLE nanos_ddl_tbl (id INT, ntz TIMESTAMP_NTZ(9), ltz TIMESTAMP_LTZ(7)) + |USING parquet""".stripMargin) + + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "false") { + // Read-through policy (SPARK-57835): metadata reads succeed and render the nanos columns + // with their precision, even though the preview flag is off. Before this change these + // commands failed at getTable time with FEATURE_NOT_ENABLED. + val describeRows = sql("DESCRIBE TABLE nanos_ddl_tbl").collect() + .map(r => r.getString(0) -> r.getString(1)).toMap + assert(describeRows("ntz") === "timestamp_ntz(9)") + assert(describeRows("ltz") === "timestamp_ltz(7)") + + val showCreate = sql("SHOW CREATE TABLE nanos_ddl_tbl").head().getString(0) + assert(showCreate.contains("TIMESTAMP_NTZ(9)")) + assert(showCreate.contains("TIMESTAMP_LTZ(7)")) + + // But actually reading the data still fails with an actionable, feature-flag error. + checkError( + exception = intercept[SparkException] { + sql("SELECT * FROM nanos_ddl_tbl").collect() + }, + condition = "FEATURE_NOT_ENABLED", + parameters = Map( + "featureName" -> "Nanosecond-precision timestamp types", + "configKey" -> "spark.sql.timestampNanosTypes.enabled", + "configValue" -> "true")) + + // The table remains droppable with the flag off (a key reason for read-through: a table + // written with the flag on must never become un-manageable once it is off). + sql("DROP TABLE nanos_ddl_tbl") + assert(!spark.sessionState.catalog.tableExists(TableIdentifier("nanos_ddl_tbl"))) + } + } + } + + test("SPARK-57835: read a persisted view over a nanos column when the preview flag is off") { + withTable("nanos_view_base") { + withView("nanos_view") { + sql("CREATE TABLE nanos_view_base (id INT, ntz TIMESTAMP_NTZ(9)) USING parquet") + // The view persists its analyzed output schema (which includes the nanos column) into + // table properties; this is created with the flag on. + sql("CREATE VIEW nanos_view AS SELECT id, ntz FROM nanos_view_base") + + withSQLConf(SQLConf.TIMESTAMP_NANOS_TYPES_ENABLED.key -> "false") { + // Restoring the view schema from properties (DataType.fromJson) succeeds, so DESCRIBE + // renders the nanos column. + val describeRows = sql("DESCRIBE TABLE nanos_view").collect() + .map(r => r.getString(0) -> r.getString(1)).toMap + assert(describeRows("ntz") === "timestamp_ntz(9)") + + // The view can still be dropped with the flag off. + sql("DROP VIEW nanos_view") + assert(!spark.sessionState.catalog.tableExists(TableIdentifier("nanos_view"))) + } + } + } + } + + test("SPARK-56822: DESCRIBE TABLE and SHOW CREATE TABLE render nanos columns with the flag on") { + withTable("nanos_basic_render", "nanos_basic_render_rt") { + // The preview flag is on by default in tests (via Utils.isTesting), so this exercises the + // normal happy path: create a table with nanos columns and introspect it with the flag on. + // SPARK-57835 covers the flag-off read-through path; this covers the basic flag-on case. + sql( + """CREATE TABLE nanos_basic_render (id INT, ntz TIMESTAMP_NTZ(9), ltz TIMESTAMP_LTZ(7)) + |USING parquet""".stripMargin) + + // DESCRIBE TABLE renders each column's type name (lowercase, with precision). + val describeRows = sql("DESCRIBE TABLE nanos_basic_render").collect() + .map(r => r.getString(0) -> r.getString(1)).toMap + assert(describeRows("ntz") === "timestamp_ntz(9)") + assert(describeRows("ltz") === "timestamp_ltz(7)") + + // SHOW CREATE TABLE renders the parseable, uppercased DDL type for each column. + val showCreate = sql("SHOW CREATE TABLE nanos_basic_render").head().getString(0) + assert(showCreate.contains("TIMESTAMP_NTZ(9)")) + assert(showCreate.contains("TIMESTAMP_LTZ(7)")) + + // Round-trip: the emitted DDL re-parses and re-creates an identical nanos schema. + sql(showCreate.replace("nanos_basic_render", "nanos_basic_render_rt")) + val rtSchema = spark.table("nanos_basic_render_rt").schema + assert(rtSchema("ntz").dataType === TimestampNTZNanosType(9)) + assert(rtSchema("ltz").dataType === TimestampLTZNanosType(7)) + } + } } diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala index 01af6b3cc570d..942172d1411c3 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/HiveUDFSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.hive.execution import java.io.{DataInput, DataOutput, File, PrintWriter} +import java.sql.{Date, Timestamp} import java.util.{ArrayList, Arrays, Properties} import scala.jdk.CollectionConverters._ @@ -35,13 +36,17 @@ import org.apache.hadoop.io.{LongWritable, Writable} import org.apache.spark.{SparkException, SparkFiles, TestUtils} import org.apache.spark.sql.{AnalysisException, QueryTest, Row} -import org.apache.spark.sql.catalyst.expressions.CodegenObjectFactoryMode -import org.apache.spark.sql.catalyst.plans.logical.Project +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, CodegenObjectFactoryMode, Literal} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, Project} +import org.apache.spark.sql.catalyst.util.DateTimeUtils import org.apache.spark.sql.execution.WholeStageCodegenExec import org.apache.spark.sql.functions.{call_function, max} +import org.apache.spark.sql.hive.HiveGenericUDF +import org.apache.spark.sql.hive.HiveShim.HiveFunctionWrapper import org.apache.spark.sql.hive.test.{TestHiveSingleton, TestUDTFJar} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.TimeType +import org.apache.spark.sql.types.{TimestampType, TimeType} import org.apache.spark.tags.SlowHiveTest import org.apache.spark.util.Utils @@ -59,7 +64,6 @@ case class ListStringCaseClass(l: Seq[String]) @SlowHiveTest class HiveUDFSuite extends QueryTest with TestHiveSingleton { import spark.implicits._ - import testImplicits.castToImpl import spark.udf @@ -887,6 +891,75 @@ class HiveUDFSuite extends QueryTest with TestHiveSingleton { } hiveContext.reset() } + + test("SPARK-58792: copied HiveGenericUDF nodes must not share a mutable GenericUDF") { + val tsAttr = AttributeReference("ts", TimestampType, nullable = false)() + val constTs = Literal( + DateTimeUtils.fromJavaTimestamp(Timestamp.valueOf("2024-09-10 01:02:03")), TimestampType) + val original = HiveGenericUDF( + "default.date_add", + HiveFunctionWrapper(classOf[GenericUDFDateAdd].getName), + Seq(tsAttr, Literal(1))) + // Optimizer-created copies of a HiveGenericUDF (via withNewChildrenInternal) share + // one HiveFunctionWrapper, and used to share the single mutable GenericUDF + // instance cached inside it. + val constCopy = original.copy(children = Seq(constTs, Literal(1))) + assert(original.funcWrapper eq constCopy.funcWrapper) + + val boundOriginal = BindReferences.bindReference(original, Seq(tsAttr)) + val input = InternalRow(DateTimeUtils.fromJavaTimestamp( + Timestamp.valueOf("2023-12-25 10:00:00"))) + // Interleave evaluation of the two copies the way a Project/Filter would. With a + // shared instance, the second evaluation of constCopy reuses the converters that + // the attribute copy's initialize() installed on the shared instance, and throws + // ClassCastException: TimestampWritable cannot be cast to java.sql.Timestamp. + (1 to 2).foreach { _ => + assert(constCopy.eval(input) == DateTimeUtils.fromJavaDate(Date.valueOf("2024-09-11"))) + assert(boundOriginal.eval(input) == DateTimeUtils.fromJavaDate(Date.valueOf("2023-12-26"))) + } + } + + test("SPARK-58792: inferred literal-binding conjunct must not corrupt a copied Hive UDF") { + // InferFiltersFromConstraints substitutes the pt = literal binding into the UDF + // conjunct and ANDs the substituted copy into the same Filter, so the Filter holds + // two copies of one UDF expression whose argument constness differs (literal vs. + // attribute). The conjunct is a bare boolean UDF call rather than a + // BinaryComparison, so ConstantPropagation (which substitutes into BinaryComparisons + // only) never rewrites the original conjunct into the same constant form, and the + // divergent pair reaches execution with stock rules - no rules excluded. A real + // table is used because a LocalRelation source lets the optimizer evaluate the + // Filter on the driver per-conjunct, which does not interleave the two copies. + // With a shared GenericUDF instance, the second row threw ClassCastException: + // TimestampWritable cannot be cast to java.sql.Timestamp. + withUserDefinedFunction("hive_gt" -> true) { + sql(s"CREATE TEMPORARY FUNCTION hive_gt AS '${classOf[GenericUDFOPGreaterThan].getName}'") + withTable("gt_table") { + sql("CREATE TABLE gt_table (id INT, pt DATE, created_at TIMESTAMP) STORED AS PARQUET") + sql(""" + |INSERT INTO gt_table VALUES + | (1, DATE '2024-09-10', TIMESTAMP '2024-09-01 00:00:00'), + | (2, DATE '2024-09-10', TIMESTAMP '2024-09-05 00:00:00'), + | (3, DATE '2024-09-10', TIMESTAMP '2024-09-15 00:00:00'), + | (4, DATE '2024-09-11', TIMESTAMP '2024-09-01 00:00:00') + |""".stripMargin) + val df = sql(""" + |SELECT id FROM gt_table + |WHERE pt = DATE '2024-09-10' + | AND hive_gt(CAST(pt AS TIMESTAMP), created_at) + |ORDER BY id + |""".stripMargin) + // Guard against the test going silently vacuous: the optimized Filter must + // hold both copies, exactly one of them with a literal first argument. + val udfs = df.queryExecution.optimizedPlan.collect { + case f: Filter => f.condition.collect { case u: HiveGenericUDF => u } + }.flatten + assert(udfs.size == 2) + assert(udfs.count(_.children.head.isInstanceOf[Literal]) == 1) + checkAnswer(df, Row(1) :: Row(2) :: Nil) + } + } + hiveContext.reset() + } } class TestPair(x: Int, y: Int) extends Writable with Serializable { diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala index 68f46cbcb1126..74d317ba1e585 100644 --- a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala +++ b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala @@ -594,6 +594,29 @@ abstract class SQLQuerySuiteBase extends QueryTest with TestHiveSingleton { } } + test("SPARK-56558: CTAS IF NOT EXISTS Hive Table should be with non-existent " + + "or empty location") { + withSQLConf(SQLConf.ALLOW_NON_EMPTY_LOCATION_IN_CTAS.key -> "false") { + withTempDir { dir => + val tempLocation = dir.toURI.toString + withTable("ctas1", "ctas_with_existing_location") { + sql(s"CREATE TABLE ctas1(id string) stored as rcfile LOCATION '$tempLocation/ctas1'") + sql("INSERT INTO TABLE ctas1 SELECT 'A' ") + // The target table does not exist in the catalog, so IF NOT EXISTS must not skip the + // non-empty location check and overwrite the data of table ctas1. + val m = intercept[AnalysisException] { + sql(s"""CREATE TABLE IF NOT EXISTS ctas_with_existing_location stored as rcfile + |LOCATION '$tempLocation' + |AS SELECT key k, value FROM src ORDER BY k, value""".stripMargin) + }.getMessage + assert(m.contains("CREATE-TABLE-AS-SELECT cannot create " + + "table with location to a non-empty directory")) + checkAnswer(spark.table("ctas1"), Row("A")) + } + } + } + } + test("CTAS with serde") { withTable("ctas1", "ctas2", "ctas3", "ctas4", "ctas5") { sql("CREATE TABLE ctas1 AS SELECT key k, value FROM src ORDER BY k, value") diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/ChangeArgs.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/ChangeArgs.scala index a8c3485da31ce..fe28decab357c 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/ChangeArgs.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/ChangeArgs.scala @@ -74,9 +74,8 @@ object ColumnSelection { * @param columnSelection The user-provided selection. `None` is a no-op and returns `schema` * unchanged. * @param resolver Determines whether two column names are considered equal. Callers - * should pass the session resolver, e.g. - * `session.sessionState.conf.resolver`, so column matching stays - * consistent with `spark.sql.caseSensitive`. + * should pass the resolver for the effective `spark.sql.caseSensitive` + * of the operation being validated. */ def applyToSchema( schemaName: String, diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala index c93460c2ad812..3f2a2e7415323 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessor.scala @@ -94,17 +94,27 @@ case class Scd2BatchProcessor( // tail can detect its own redundancy via LEAD(1): if the next row is a non-tail at // the same instant, the synthetic close the tail encodes is already represented by // that event and the tail is dropped downstream. - // - // Any tiebreaking beyond this rule only meaningfully fires when the user's source - // has emitted two or more events at the same sequence, violating the uniqueness - // contract above. Behavior in that case is publicly undefined and the remaining - // tiebreaker clauses exist as a best-effort to keep retries and replays deterministic. orderDecompositionTailsFirst, // Upsert-representing rows sort before tombstones because rows detect if they are being // bisected by LEAD(1). This allows upserts to match against same-sequence deletes, an // arbitrary but deterministic convention. When this happens, the delete event will survive // and persist as a tombstone in the auxiliary table. - orderUpsertRepresentingRowsFirst + orderUpsertRepresentingRowsFirst, + // Amongst upsert-representing rows, there's one valid case where rows are still tied, even + // if the user's change feed source did not emit duplicate sequences: the auxiliary merge + // commits before the target merge, which re-reads that table, so a row this batch wrote to + // the auxiliary table re-enters the window beside the copy the microbatch or the target + // table still holds. The copies differ only in the boundaries each one recorded, and the + // two keys below break that tie deterministically - dropRedundantRowsPostDecomposition + // drops a tie's leading row, so the copy sorting last is the one that survives. + // + // A batch that moves a run's start leaves the copies disagreeing on where their interval + // begins. Nulls are inert on this key: only decomposition tails carry a null startAt, and + // orderDecompositionTailsFirst has already separated those by the time it is consulted. + startAtCol.desc_nulls_first, + // Copies agreeing on their run start may still disagree on its closure, so nulls sort + // first and are dropped: a copy that recorded no boundary never displaces one that did. + endAtCol.desc_nulls_first ) } @@ -254,33 +264,15 @@ case class Scd2BatchProcessor( } /** - * Find the auxiliary-table rows whose state matters for reconciling the microbatch. - * - * @param rawAuxiliaryTableDf - * the auxiliary table in its native schema, which is expected to contain - * [[deletedByBatchIdColName]] in addition to all of the columns in the target table. - * @param perKeyMinimumSequenceInMicrobatchDf - * one row per distinct key as produced by [[computeMinimumSequencePerKey]], representing - * the minimum sequence for that key in the microbatch. - * @param batchId - * the underlying Spark streaming query's batchId, used to scope aux-row visibility for - * replay-stability across retries of the same microbatch. - * @return - * a dataframe containing all the affected aux rows, with the aux-only - * [[deletedByBatchIdColName]] column dropped so the result is union-compatible with - * preprocessed microbatch rows and target-table rows downstream. + * Restrict the auxiliary table to the rows that are live for this microbatch, and drop the + * aux-only [[deletedByBatchIdColName]] column so the result shares the canonical SCD2 row + * schema with target-table rows and preprocessed-microbatch rows. */ - private[autocdc] def findAffectedRowsFromAuxiliaryTable( + private def filterLiveAuxiliaryRows( rawAuxiliaryTableDf: DataFrame, - perKeyMinimumSequenceInMicrobatchDf: DataFrame, - batchId: Long - ): DataFrame = { - val auxTableRecordStartAtField = Scd2BatchProcessor.recordStartAtOf( - F.col(AutoCdcReservedNames.cdcMetadataColName) - ) + batchId: Long): DataFrame = { val auxTableDeletedByBatchIdCol = F.col(Scd2BatchProcessor.deletedByBatchIdColName) - - val reducedAuxiliaryTableDf = rawAuxiliaryTableDf + rawAuxiliaryTableDf .filter( // [[deletedByBatchIdColName]] carries the batchId whose MERGE logically deleted the // row, or null on live aux rows. Rows deleted by other batches are excluded - those @@ -290,115 +282,136 @@ case class Scd2BatchProcessor( auxTableDeletedByBatchIdCol.isNull || auxTableDeletedByBatchIdCol === F.lit(batchId) ) - // Drop the aux-only idempotency column so the output schema matches target-table rows - // and preprocessed-microbatch rows (which share the same canonical SCD2 row schema). .drop(Scd2BatchProcessor.deletedByBatchIdColName) + } - val perKeyMinimumSequenceInMicrobatchCol = F.col(Scd2BatchProcessor.minSequenceColName) + /** + * Project a table of canonical SCD2 rows down to `[key1, ... keyN, effectiveRecordStartAt]`. + */ + private def projectEffectiveRecordStartAtPerRow(rowsDf: DataFrame): DataFrame = + rowsDf.select( + keysQuoted.map(F.col) :+ + Scd2BatchProcessor.canonicalRowIntervalColumns.effectiveRecordStartAt + .as(Scd2BatchProcessor.effectiveRecordStartAtColName): _* + ) - // Per key, identify the sequence value associated with the anchor row in the aux table. - // - // The anchor row is the aux row with the largest [[recordStartAtFieldName]] strictly less - // than the min sequence in the incoming microbatch for that key. The reconciler needs this - // "left context" in two cases: - // (1) Incoming no-op upsert: without the anchor, it would look like a new run head, when in - // reality it's a part of an existing no-op run/head. - // (2) Incoming state-changing upsert that bisects two aux no-ops: the anchor surfaces - // the before-half so both halves can be promoted to target. (The after-half is - // picked up by the >= minSeq branch.) - // - // Because no-op upserts are stored only in the aux table, the anchor concept only exists when - // pulling in rows from the aux table, and is not relevant for the target table. - // - // Keys with no aux row strictly before the min sequence have no anchor; their affected set - // reduces to "all aux rows at or after the min sequence." - // - // The shape of this DataFrame is: [key1, key2, ... keyN, anchorSequence] - val perKeyAnchorSequenceDf = reducedAuxiliaryTableDf + /** + * Per key; calculate the earliest point in time (sequence) at or after which all existing rows + * across the auxiliary and target tables may be affected by the microbatch, and therefore should + * be pulled in for reconciliation. The row sitting exactly at the cutoff is itself included. + * + * Returns a dataframe with one row per distinct key in [[perKeyMinimumSequenceInMicrobatchDf]], + * with the key columns and the calculated [[affectedSequenceCutoffColName]] column. + */ + private[autocdc] def computePerKeyAffectedSequenceCutoff( + rawAuxiliaryTableDf: DataFrame, + targetTableDf: DataFrame, + perKeyMinimumSequenceInMicrobatchDf: DataFrame, + batchId: Long + ): DataFrame = { + val effectiveRecordStartAtCol = F.col(Scd2BatchProcessor.effectiveRecordStartAtColName) + val perKeyMinimumSequenceInMicrobatchCol = F.col(Scd2BatchProcessor.minSequenceColName) + val latestSequenceBeforeMicrobatchCol = + F.col(Scd2BatchProcessor.latestSequenceBeforeMicrobatchColName) + + // In order to determine the affected sequence cutoff per key, we first need a "global" (across + // both auxiliary and target tables, hence unioned) timeline of existing effective sequences per + // key. + val allRowsByEffectiveRecordStartAt = + projectEffectiveRecordStartAtPerRow(filterLiveAuxiliaryRows(rawAuxiliaryTableDf, batchId)) + .unionByName(projectEffectiveRecordStartAtPerRow(targetTableDf)) + + // Across all existing rows per key, we need to find the latest one that starts + // (i.e effectiveRecordStartAt) before the first event for that key in the microbatch. This is + // an "anchor", where any existing row that predates this one on the timeline will definitely + // not be affected by the microbatch. + val perKeyLatestSequenceBeforeMicrobatchDf = allRowsByEffectiveRecordStartAt // The number of rows in [[perKeyMinimumSequenceInMicrobatchDf]] is bounded by the - // number of unique keys in the microbatch, which should typically be small. The - // auxiliary table should generally also be small, containing only no-op upsert runs - // and tombstones per key. Therefore this join should be cheap, and broadcast joinable. + // number of unique keys in the microbatch, which should typically be small, so this + // join should be cheap and broadcast joinable. .join(perKeyMinimumSequenceInMicrobatchDf, keysRaw) - .filter(auxTableRecordStartAtField < perKeyMinimumSequenceInMicrobatchCol) + // We only care about rows that definitely start before the first event for the same key in + // the microbatch. + .filter(effectiveRecordStartAtCol < perKeyMinimumSequenceInMicrobatchCol) .groupBy(keysQuoted.map(F.col): _*) + // Of all the existing rows that start before the earliest event in the microbatch, we want + // the latest, hence max-by. .agg( - F.max(auxTableRecordStartAtField).as(Scd2BatchProcessor.anchorSequenceColName) + F.max(effectiveRecordStartAtCol) + .as(Scd2BatchProcessor.latestSequenceBeforeMicrobatchColName) ) - val anchorSequenceCol = F.col(Scd2BatchProcessor.anchorSequenceColName) - val auxRowIsAnchorRow = auxTableRecordStartAtField === anchorSequenceCol - - // Now that we have the minimum sequence in the microbatch and the sequence of the anchor row, - // we have enough information to compute the full set of auxiliary rows that may affect or - // be affected by the microbatch. Membership here is a conservative superset: every row that - // could possibly participate in reconciliation is included, but downstream reconciliation - // determines the actual outcome per row. - val auxRowIsAtOrAfterMinSequenceInMicrobatch = - auxTableRecordStartAtField >= perKeyMinimumSequenceInMicrobatchCol - - val auxRowAffectsMicrobatch = auxRowIsAtOrAfterMinSequenceInMicrobatch || auxRowIsAnchorRow - - val affectedRowsFromAuxiliaryTable = reducedAuxiliaryTableDf - // Per row, project the minimum microbatch sequence and anchor sequence for that row's key - // set onto the row, so the affected-row predicate can be evaluated in a single filter. - .join(perKeyMinimumSequenceInMicrobatchDf, keysRaw) - .join( - perKeyAnchorSequenceDf, - keysRaw, - joinType = "left" - ) - .filter(auxRowAffectsMicrobatch) - .drop(perKeyMinimumSequenceInMicrobatchCol, anchorSequenceCol) - affectedRowsFromAuxiliaryTable + // Now for all unique keys in the microbatch, we need to calculate its affected sequence + // cutoff. If a key has existing rows and at least one of them precedes the earliest event + // for that same key in the microbatch, use its sequence as the cutoff. In all other cases + // (key does not yet exist in aux/target, or microbatch contains event for key that precedes + // all existing rows), the min sequence in the microbatch will be the cutoff point for the key; + // all existing rows will necessarily be considered affected. This is a very cheap join to + // make, where both dataframes cardinality is limited by the number of unique keys in the + // microbatch. + perKeyMinimumSequenceInMicrobatchDf + .join(perKeyLatestSequenceBeforeMicrobatchDf, keysRaw, joinType = "left") + .select( + keysQuoted.map(F.col) :+ + F.coalesce(latestSequenceBeforeMicrobatchCol, perKeyMinimumSequenceInMicrobatchCol) + .as(Scd2BatchProcessor.affectedSequenceCutoffColName): _* + ) } /** - * Find the target-table rows whose state matters for reconciling the microbatch. + * Per key, keep only rows in [[rowsDf]] that are included by the sequence cutoff. * - * @param targetTableDf - * the target table in its native schema. - * @param perKeyMinimumSequenceInMicrobatchDf - * one row per distinct key as produced by [[computeMinimumSequencePerKey]], representing - * the minimum sequence for that key in the microbatch. - * @return - * a dataframe containing the affected target rows, with all columns passed-through. + * A plain `effectiveRecordStartAt >= cutoff` threshold suffices, in both directions: + * + * 1. Nothing needed for reconciliation is missed. Once a row is pulled in, every later row for + * that key must come with it, so that boundaries reconcile and rows are promoted or demoted + * correctly - even when the two live in different tables, as when a run's visible tail sits + * in the target table after the run's hidden head in the auxiliary table. The cutoff is a + * single threshold per key over the unified ordering across both tables, and so selects + * exactly a complete suffix of the global, unified view. This includes the live rows (open + * upserts in the target) too. + * 2. Nothing dropped was needed for reconciliation. The cutoff is the position of the last row + * preceding the microbatch per key, so every row below it is separated from the microbatch + * by at least one intervening row, and cannot be affected by it. An open row is the one case + * that does not follow from separation alone, since it is affected by anything after it - + * but no row can intervene above an open row, as that row would have closed it. */ - private[autocdc] def findAffectedRowsFromTargetTable( - targetTableDf: DataFrame, - perKeyMinimumSequenceInMicrobatchDf: DataFrame - ): DataFrame = { - val targetEndAtCol = F.col(Scd2BatchProcessor.endAtColName) - val perKeyMinimumSequenceInMicrobatchCol = F.col(Scd2BatchProcessor.minSequenceColName) + private def selectRowsAtOrAfterCutoff( + rowsDf: DataFrame, + perKeyAffectedSequenceCutoffDf: DataFrame): DataFrame = { + val affectedSequenceCutoffCol = F.col(Scd2BatchProcessor.affectedSequenceCutoffColName) + rowsDf + .join(perKeyAffectedSequenceCutoffDf, keysRaw) + .filter( + Scd2BatchProcessor.canonicalRowIntervalColumns.effectiveRecordStartAt >= + affectedSequenceCutoffCol) + .drop(affectedSequenceCutoffCol) + } - // Per key, identify all the rows in the target table that may be affected by the - // incoming microbatch. - // - // Unlike the auxiliary table, the target table holds visible rows only: no hidden open - // no-op upsert rows, no tombstones. Visible rows for a given key form a non-overlapping - // interval partition over the sequencing axis, and at most one row has a null [[endAtColName]] - // (the currently active row per key). - // - // Hence we can simply grab all rows that were active at some point after the min sequencing - // per key, which can be determined entirely by the row's [[endAtColName]]. - val isCurrentlyActiveRow = targetEndAtCol.isNull - - // `>=` (rather than strict `>`) additionally pulls in the row that closes exactly at the - // smallest incoming sequence: the consecutive left neighbor of that incoming event. This - // provides "left context" for the smallest event, analogous to the anchor row in - // [[findAffectedRowsFromAuxiliaryTable]]. It may need to be demoted from a target run - // boundary to an aux no-op continuation if the incoming event at minSeq turns out to - // extend an earlier run. - val rowEndsAfterMinimumSequence = targetEndAtCol >= perKeyMinimumSequenceInMicrobatchCol - val rowMayBeAffected = isCurrentlyActiveRow || rowEndsAfterMinimumSequence - - val affectedRowsFromTargetTable = targetTableDf - .join(perKeyMinimumSequenceInMicrobatchDf, keysRaw) - .filter(rowMayBeAffected) - .drop(perKeyMinimumSequenceInMicrobatchCol) + /** + * Retrieve all rows from the auxiliary table that are possibly affected by this microbatch, and + * need to participate in reconciliation. + */ + private[autocdc] def findAffectedRowsFromAuxiliaryTable( + rawAuxiliaryTableDf: DataFrame, + perKeyAffectedSequenceCutoffDf: DataFrame, + batchId: Long + ): DataFrame = selectRowsAtOrAfterCutoff( + rowsDf = filterLiveAuxiliaryRows(rawAuxiliaryTableDf, batchId), + perKeyAffectedSequenceCutoffDf = perKeyAffectedSequenceCutoffDf + ) - affectedRowsFromTargetTable - } + /** + * Retrieve all rows from the target table that are possibly affected by this microbatch, and + * need to participate in reconciliation. + */ + private[autocdc] def findAffectedRowsFromTargetTable( + targetTableDf: DataFrame, + perKeyAffectedSequenceCutoffDf: DataFrame + ): DataFrame = selectRowsAtOrAfterCutoff( + rowsDf = targetTableDf, + perKeyAffectedSequenceCutoffDf = perKeyAffectedSequenceCutoffDf + ) /** * For every closed non-tombstone row in the input dataframe whose immediate window-order @@ -611,6 +624,12 @@ case class Scd2BatchProcessor( * upsert drops in favor of a same-sequence tombstone (delete wins over upsert at the * same instant). * + * Two copies of one row that differ only in the boundaries each recorded - which the + * auxiliary merge can produce for the target merge to read - collide the same way, and + * are ordered by [[startAtColName]] first and [[endAtColName]] second, each descending + * with nulls first. A copy that recorded no boundary therefore never outlives one that + * did, so of an open and a closed copy at the same recordStartAt the closed one survives. + * * @param decomposedRowsPerKey * the output of [[decomposeOutOfOrderRows]]: a dataframe conforming to the canonical * SCD2 row schema `[user_cols..., [[startAtColName]], [[endAtColName]], @@ -738,7 +757,8 @@ case class Scd2BatchProcessor( F.when( isWindowLocalUpsertRunHead, // The first row in the window may be a window-local run head but not a global run - // head (e.g., an aux anchor row pulled in for left context). In that case, `startAt` + // head (e.g., the row at the affected sequence cutoff, pulled in from either the + // auxiliary or the target table for left context). In that case, `startAt` // may be strictly less than `recordStartAt`, encoding the true global run start, and // we propagate it forward to later in-window continuations of the same run. // For every later window-local upsert run head, `recordStartAt` is the run start. @@ -1492,12 +1512,35 @@ object Scd2BatchProcessor { s"${AutoCdcReservedNames.prefix}is_redundant_delete_encoding" /** - * Name of the temporary column used to identify the sequence associated with the anchor - * row found in the auxiliary table for the incoming microbatch. Since sequences must be unique - * amongst all rows for a key (or risk undefined behavior), this sequence value uniquely - * identifies an exact row in the aux. + * Name of the temporary column carrying a row's [[Scd2IntervalColumns.effectiveRecordStartAt]] + * in the narrow, key-plus-ordering-metadata union of the auxiliary and target tables that + * [[Scd2BatchProcessor.computePerKeyAffectedSequenceCutoff]] aggregates over. + * + * Temporary in that the column has no observable side effect or persistence across microbatches. */ - private val anchorSequenceColName: String = s"${AutoCdcReservedNames.prefix}anchor_sequence" + private val effectiveRecordStartAtColName: String = + s"${AutoCdcReservedNames.prefix}effective_record_start_at" + + /** + * Name of the temporary column holding, per key, the largest effective ordering position + * strictly below the key's minimum microbatch sequence, across the auxiliary AND target tables + * jointly. Null for a key with no such row, in which case the affected sequence cutoff falls + * back to the microbatch minimum. + * + * Temporary in that the column has no observable side effect or persistence across microbatches. + */ + private val latestSequenceBeforeMicrobatchColName: String = + s"${AutoCdcReservedNames.prefix}latest_sequence_before_microbatch" + + /** + * Name of the temporary column holding the single per-key cutoff on effective ordering + * position that gates affected-row selection from both the auxiliary and the target table, as + * computed by [[Scd2BatchProcessor.computePerKeyAffectedSequenceCutoff]]. + * + * Temporary in that the column has no observable side effect or persistence across microbatches. + */ + private val affectedSequenceCutoffColName: String = + s"${AutoCdcReservedNames.prefix}affected_sequence_cutoff" /** * Name of the temporary column projected by [[Scd2BatchProcessor.identifyAndTagAuxRows]] to @@ -1524,6 +1567,17 @@ object Scd2BatchProcessor { private def recordStartAtOf(cdcMetadataCol: Column): Column = cdcMetadataCol.getField(recordStartAtFieldName) + /** + * The [[Scd2IntervalColumns]] of a row read from either the auxiliary or the target table, in + * the canonical SCD2 row schema. The columns are unresolved name references, so they read from + * whichever dataframe the expressions are applied to. + */ + private def canonicalRowIntervalColumns: Scd2IntervalColumns = Scd2IntervalColumns( + recordStartAt = recordStartAtOf(F.col(AutoCdcReservedNames.cdcMetadataColName)), + startAt = F.col(startAtColName), + endAt = F.col(endAtColName) + ) + /** * Schema of the CDC metadata struct column for SCD2 rows. */ diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandler.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandler.scala index 3e59fa09d25a5..e4adfcb516be7 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandler.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandler.scala @@ -75,23 +75,35 @@ case class Scd2ForeachBatchHandler( ) val auxTableDf = batchDf.sparkSession.read.table(auxiliaryTableIdentifier.quotedString) - val affectedRowsFromAuxiliaryTable = batchProcessor.findAffectedRowsFromAuxiliaryTable( + val targetTableDf = batchDf.sparkSession.read.table(targetTableIdentifier.quotedString) + + val perKeyAffectedSequenceCutoffDf = batchProcessor.computePerKeyAffectedSequenceCutoff( rawAuxiliaryTableDf = auxTableDf, + targetTableDf = targetTableDf, perKeyMinimumSequenceInMicrobatchDf = perKeyMinimumSequenceInMicrobatchDf, batchId = batchId ) - val targetTableDf = batchDf.sparkSession.read.table(targetTableIdentifier.quotedString) + val affectedRowsFromAuxiliaryTable = batchProcessor.findAffectedRowsFromAuxiliaryTable( + rawAuxiliaryTableDf = auxTableDf, + perKeyAffectedSequenceCutoffDf = perKeyAffectedSequenceCutoffDf, + batchId = batchId + ) + val affectedRowsFromTargetTable = batchProcessor.findAffectedRowsFromTargetTable( targetTableDf = targetTableDf, - perKeyMinimumSequenceInMicrobatchDf = perKeyMinimumSequenceInMicrobatchDf + perKeyAffectedSequenceCutoffDf = perKeyAffectedSequenceCutoffDf ) - // All three share the canonical schema; findAffectedRowsFromAuxiliaryTable drops the aux-only - // deletedByBatchId column. + // The three inputs share the canonical SCD2 row schema by name, but not necessarily by column + // set: after cross-run schema evolution the target (and the aux table, which mirrors it) can + // carry user columns that the current microbatch no longer emits. `allowMissingColumns` pads + // such columns with null on the side that lacks them (recursing into structs and arrays; map + // types are not supported) instead of failing the union. (findAffectedRowsFromAuxiliaryTable + // drops the aux-only deletedByBatchId column.) val microbatchAndAffectedRows = preprocessedBatchDf - .unionByName(affectedRowsFromAuxiliaryTable) - .unionByName(affectedRowsFromTargetTable) + .unionByName(affectedRowsFromAuxiliaryTable, allowMissingColumns = true) + .unionByName(affectedRowsFromTargetTable, allowMissingColumns = true) val decomposedDf = microbatchAndAffectedRows .transform(batchProcessor.decomposeOutOfOrderRows) diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTable.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTable.scala index 9f3dd2dba0d16..da5b7eb8ad125 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTable.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTable.scala @@ -27,8 +27,9 @@ import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.catalyst.analysis.Resolver import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Table => CatalogTable, TableCatalog} -import org.apache.spark.sql.pipelines.autocdc.{AutoCdcReservedNames, Scd2BatchProcessor, ScdType} -import org.apache.spark.sql.types.{LongType, StructField, StructType} +import org.apache.spark.sql.pipelines.autocdc.{AutoCdcReservedNames, Scd1BatchProcessor, + Scd2BatchProcessor, ScdType} +import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType} /** * Helpers to construct and validate an AutoCDC flow's auxiliary table within the context of a @@ -61,20 +62,29 @@ object AutoCdcAuxiliaryTable { s"${PipelinesTableProperties.pipelinesPrefix}autocdc.keyColumnNames" /** - * Serialize key column names to the JSON form stored at [[keyColumnNamesProperty]]. - * Round-trips an empty list as `[]`; callers are expected to enforce a non-empty key set - * upstream. + * Table property recording the resolved SCD2 track-history column names as a JSON string array. + * These columns define an SCD2 run (a change in any of them opens a new historical record), so + * changing the set would reinterpret already-reconciled history. SCD2-only; absent for SCD1. + * Full-refresh is the only way to change it. */ - private[graph] def serializeKeyColumnNames(names: Seq[String]): String = { + val trackHistoryColumnNamesProperty: String = + s"${PipelinesTableProperties.pipelinesPrefix}autocdc.trackHistoryColumnNames" + + /** + * A JSON string-array codec for column-name lists persisted as auxiliary-table properties (used + * for both [[keyColumnNamesProperty]] and [[trackHistoryColumnNamesProperty]]). Round-trips an + * empty list as `[]` -- both an empty key set and an empty track-history set are meaningful to + * some caller, so no non-empty invariant is assumed here. + */ + private[graph] def serializeColumnNames(names: Seq[String]): String = { compact(JArray(names.map(JString(_)).toList)) } /** - * Parse a [[keyColumnNamesProperty]] value. `None` if it is not a JSON array of strings. - * Round-trips an empty list as `[]`; callers are expected to enforce a non-empty key set - * upstream. + * Parse a value written by [[serializeColumnNames]]. `None` if it is not a JSON array of strings. + * Round-trips an empty list as `[]` (see [[serializeColumnNames]]). */ - private[graph] def parseKeyColumnNames(raw: String): Option[Seq[String]] = { + private[graph] def parseColumnNames(raw: String): Option[Seq[String]] = { val parsed = try Some(parse(raw)) catch { case NonFatal(_) => None } parsed.flatMap { case JArray(elems) => @@ -133,7 +143,7 @@ object AutoCdcAuxiliaryTable { ): AuxiliaryTableSpec = { val scd1AuxiliaryTableIdentifier = identifier(targetTable.identifier) - val resolver = inputAutoCdcFlow.df.sparkSession.sessionState.conf.resolver + val resolver = inputAutoCdcFlow.effectiveResolver val autoCdcKeyColumnNames = inputAutoCdcFlow.changeArgs.keys.map(_.name) // The auxiliary table should derive its schema from the exact same key/CDC metadata column @@ -163,7 +173,7 @@ object AutoCdcAuxiliaryTable { Map(scdTypePropertyKey -> ScdType.Type1.label) ++ // Persist the AutoCDC key column names as a JSON list; immutable post-creation (full-refresh // is the only way to change it). - Map(keyColumnNamesProperty -> serializeKeyColumnNames(keyFields.map(_.name))) ++ + Map(keyColumnNamesProperty -> serializeColumnNames(keyFields.map(_.name))) ++ // Inherit the target's format so MERGE semantics line up. When unspecified, omit the provider // so the catalog falls back to its default. targetTable.format.map(TableCatalog.PROP_PROVIDER -> _) @@ -174,7 +184,9 @@ object AutoCdcAuxiliaryTable { properties = scd1AuxiliaryTableProperties, targetTableIdentifier = targetTable.identifier, expectedKeyFields = keyFields, - expectedScdType = ScdType.Type1 + expectedScdType = ScdType.Type1, + expectedSequencingType = inputAutoCdcFlow.sequencingType, + expectedTrackHistoryColumnNames = None ) } @@ -205,7 +217,7 @@ object AutoCdcAuxiliaryTable { ): AuxiliaryTableSpec = { val scd2AuxiliaryTableIdentifier = identifier(targetTable.identifier) - val resolver = inputAutoCdcFlow.df.sparkSession.sessionState.conf.resolver + val resolver = inputAutoCdcFlow.effectiveResolver val autoCdcKeyColumnNames = inputAutoCdcFlow.changeArgs.keys.map(_.name) // Resolve the key fields from the (evolved) target schema, exactly as SCD1 does, so the @@ -227,13 +239,27 @@ object AutoCdcAuxiliaryTable { StructField(Scd2BatchProcessor.deletedByBatchIdColName, LongType, nullable = true) val scd2AuxiliaryTableSchema = StructType(targetTableSchema.fields :+ deletedByBatchIdField) + // The effective track-history column set, resolved by the flow from its user-selected source + // schema (see [[AutoCdcMergeFlow.trackHistoryColumnNames]]) -- NOT recomputed here from the + // evolved target schema, which would keep tracking columns the flow no longer selects and would + // miss implicit (default / `* EXCEPT`) tracked-set changes. A change in this set reinterprets + // which transitions open a new historical record, so it is drift-checked. + val trackHistoryColumnNames = inputAutoCdcFlow.trackHistoryColumnNames.getOrElse( + throw SparkException.internalError( + "SCD2 AutoCDC flow is missing its resolved track-history column set." + ) + ) + val scd2AuxiliaryTableProperties = // Record which SCD strategy this auxiliary table serves so downstream readers can identify it // without inspecting the schema. Map(scdTypePropertyKey -> ScdType.Type2.label) ++ // Persist the AutoCDC key column names as a JSON list; immutable post-creation (full-refresh // is the only way to change it). - Map(keyColumnNamesProperty -> serializeKeyColumnNames(keyFields.map(_.name))) ++ + Map(keyColumnNamesProperty -> serializeColumnNames(keyFields.map(_.name))) ++ + // Persist the resolved track-history column names; a change reinterprets already-reconciled + // history, so it is immutable post-creation (full-refresh is the only way to change it). + Map(trackHistoryColumnNamesProperty -> serializeColumnNames(trackHistoryColumnNames)) ++ // Inherit the target's format so MERGE semantics line up. When unspecified, omit the provider // so the catalog falls back to its default. targetTable.format.map(TableCatalog.PROP_PROVIDER -> _) @@ -244,7 +270,9 @@ object AutoCdcAuxiliaryTable { properties = scd2AuxiliaryTableProperties, targetTableIdentifier = targetTable.identifier, expectedKeyFields = keyFields, - expectedScdType = ScdType.Type2 + expectedScdType = ScdType.Type2, + expectedSequencingType = inputAutoCdcFlow.sequencingType, + expectedTrackHistoryColumnNames = Some(trackHistoryColumnNames) ) } @@ -255,7 +283,7 @@ object AutoCdcAuxiliaryTable { * * @param targetTableSchema the AutoCDC target's evolved schema to resolve against * @param fieldName the column name to resolve - * @param resolver the session resolver used for case-sensitivity-aware field lookups + * @param resolver the effective resolver used for case-sensitivity-aware field lookups * @param targetTableIdentifier the AutoCDC target's identifier, named in the error message * @param autoCdcFlowIdentifier the AutoCDC flow writing to the target, named in the error message * @return the matching field @@ -373,6 +401,114 @@ object AutoCdcAuxiliaryTable { } } + /** + * Reject an incremental update to an existing AutoCDC target table whose sequencing type has + * drifted. The AutoCDC sequencing *expression* may legitimately change across runs (e.g. a new + * timestamp parse format), but its resolved result type must not: the target persists the + * sequencing type inside its `_cdc_metadata` struct (and, for SCD2, in the interval columns), so + * a changed type would make new events incomparable with the persisted history and would + * otherwise surface only as a generic CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE during schema + * evolution. Runs against the target table (before its schema is evolved), not the auxiliary + * table. The remedy is a full refresh. + * + * @param existingTargetSchema the schema of the already-materialized target table. + * @param expectedScdType the SCD type of the incoming AutoCDC flow, which determines which inner + * `_cdc_metadata` field carries the recorded sequencing type. + * @param expectedSequencingType the resolved sequencing type of the incoming AutoCDC flow. + * @param resolver the effective resolver, used to match the reserved column and inner field names + * the same case-aware way as every other schema lookup in this file. + */ + private[graph] def validateNoTargetSequencingTypeDrift( + existingTargetSchema: StructType, + targetTableIdentifier: TableIdentifier, + expectedScdType: ScdType, + expectedSequencingType: DataType, + resolver: Resolver): Unit = { + // The sequencing type is embedded as an inner field of the reserved _cdc_metadata struct: for + // SCD1 the delete/upsert sequence fields, for SCD2 the recordStartAt field. Look the field up + // by name (not by position) so a future metadata field added at position 0 cannot silently + // shift this to an unrelated type, and via the resolver so a case-differing hand-written target + // DDL resolves the same way it does everywhere else. If the metadata column is absent, not a + // struct, or lacks the expected inner field, this is not a recognizable AutoCDC target state; + // skip rather than misreport (schema evolution will surface any genuine incompatibility). + val sequencingFieldName = expectedScdType match { + case ScdType.Type1 => Scd1BatchProcessor.cdcUpsertSequenceFieldName + case ScdType.Type2 => Scd2BatchProcessor.recordStartAtFieldName + } + val recordedSequencingType: Option[DataType] = existingTargetSchema.fields + .find(f => resolver(f.name, AutoCdcReservedNames.cdcMetadataColName)) + .map(_.dataType) + .collect { case s: StructType => s } + .flatMap(_.fields.find(f => resolver(f.name, sequencingFieldName))) + .map(_.dataType) + + recordedSequencingType.foreach { recordedType => + if (!recordedType.sameType(expectedSequencingType)) { + throw new AnalysisException( + errorClass = "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT", + messageParameters = Map( + "tableName" -> targetTableIdentifier.unquotedString, + "expectedSequencingType" -> expectedSequencingType.sql, + "recordedSequencingType" -> recordedType.sql + ) + ) + } + } + } + + /** + * Reject an existing SCD2 auxiliary table whose recorded track-history column set differs from + * `expected` (order-insensitive, resolver-aware). These columns define an SCD2 run - a change in + * any of them opens a new historical record - so changing the set would reinterpret already + * reconciled history. The remedy is a full refresh. + * + * `expected` is `None` for SCD1 (no track-history concept); the check is then a no-op. + */ + private[graph] def validateNoTrackHistoryDrift( + existingAuxiliaryTable: CatalogTable, + targetTableIdentifier: TableIdentifier, + expectedTrackHistoryColumnNames: Option[Seq[String]], + resolver: Resolver): Unit = { + expectedTrackHistoryColumnNames.foreach { expectedNames => + val rawRecorded = Option( + existingAuxiliaryTable.properties().get(trackHistoryColumnNamesProperty) + ).getOrElse { + throw new AnalysisException( + errorClass = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MISSING", + messageParameters = Map( + "tableName" -> targetTableIdentifier.unquotedString, + "propertyName" -> trackHistoryColumnNamesProperty + ) + ) + } + val recordedNames = parseColumnNames(rawRecorded).getOrElse { + throw new AnalysisException( + errorClass = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MALFORMED", + messageParameters = Map( + "tableName" -> targetTableIdentifier.unquotedString, + "propertyName" -> trackHistoryColumnNamesProperty, + "rawValue" -> rawRecorded + ) + ) + } + // Set equality, resolver-aware: same arity and every expected name has a recorded + // counterpart. Order is irrelevant to run semantics. + val drifted = + recordedNames.length != expectedNames.length || + expectedNames.exists(e => !recordedNames.exists(r => resolver(r, e))) + if (drifted) { + throw new AnalysisException( + errorClass = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + messageParameters = Map( + "tableName" -> targetTableIdentifier.unquotedString, + "expectedTrackHistoryColumns" -> expectedNames.mkString(", "), + "recordedTrackHistoryColumns" -> recordedNames.mkString(", ") + ) + ) + } + } + } + /** * Read [[keyColumnNamesProperty]] off an existing auxiliary table and parse it into the ordered * list of recorded AutoCDC key column names. @@ -391,7 +527,7 @@ object AutoCdcAuxiliaryTable { ) ) } - parseKeyColumnNames(rawKeyColumnNamesStr).getOrElse { + parseColumnNames(rawKeyColumnNamesStr).getOrElse { throw new AnalysisException( errorClass = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MALFORMED", messageParameters = Map( diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AuxiliaryTableSpec.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AuxiliaryTableSpec.scala index 55cd19a1f8c43..71ad220c4cea5 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AuxiliaryTableSpec.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/AuxiliaryTableSpec.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.pipelines.graph import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.pipelines.autocdc.ScdType -import org.apache.spark.sql.types.{StructField, StructType} +import org.apache.spark.sql.types.{DataType, StructField, StructType} /** * A specification for an internal auxiliary table whose lifecycle follows a materialized @@ -70,6 +70,13 @@ sealed trait AuxiliaryTableSpec { * it already exists, in order (names and types). * @param expectedScdType the SCD type the auxiliary table is expected to have recorded, if it * already exists. + * @param expectedSequencingType the resolved [[org.apache.spark.sql.types.DataType]] of the + * AutoCDC sequencing expression. The expression may change across runs + * but its result type must not, so the persisted interval/recordStartAt + * values stay comparable. + * @param expectedTrackHistoryColumnNames the resolved SCD2 track-history column names the auxiliary + * table is expected to have recorded, in order. `None` for SCD1 (which + * has no track-history concept), so the check is skipped there. */ final case class AutoCdcAuxiliaryTableSpec( identifier: TableIdentifier, @@ -77,4 +84,6 @@ final case class AutoCdcAuxiliaryTableSpec( properties: Map[String, String], targetTableIdentifier: TableIdentifier, expectedKeyFields: Seq[StructField], - expectedScdType: ScdType) extends AuxiliaryTableSpec + expectedScdType: ScdType, + expectedSequencingType: DataType, + expectedTrackHistoryColumnNames: Option[Seq[String]]) extends AuxiliaryTableSpec diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/CoreDataflowNodeProcessor.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/CoreDataflowNodeProcessor.scala index 66f2995ee02d9..affa4725a2bfa 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/CoreDataflowNodeProcessor.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/CoreDataflowNodeProcessor.scala @@ -32,9 +32,9 @@ import org.apache.spark.sql.pipelines.graph.DataflowGraphTransformer.{ * Processor that is responsible for analyzing each flow and sort the nodes in * topological order */ -class CoreDataflowNodeProcessor(rawGraph: DataflowGraph) { +class CoreDataflowNodeProcessor(rawGraph: DataflowGraph, sessionCaseSensitive: Boolean) { - private val flowResolver = new FlowResolver(rawGraph) + private val flowResolver = new FlowResolver(rawGraph, sessionCaseSensitive) // Map of input identifier to resolved [[Input]]. private val resolvedInputs = new ConcurrentHashMap[TableIdentifier, Input]() @@ -86,7 +86,8 @@ class CoreDataflowNodeProcessor(rawGraph: DataflowGraph) { identifier = table.identifier, specifiedSchema = table.specifiedSchema, incomingFlowIdentifiers = flowsToTable.map(_.identifier).toSet, - availableFlows = resolvedFlowsToTable + availableFlows = resolvedFlowsToTable, + sessionCaseSensitive = sessionCaseSensitive ) resolvedInputs.put(table.identifier, virtualTableInput) Seq(table) @@ -110,7 +111,7 @@ class CoreDataflowNodeProcessor(rawGraph: DataflowGraph) { } } -private class FlowResolver(rawGraph: DataflowGraph) { +private class FlowResolver(rawGraph: DataflowGraph, sessionCaseSensitive: Boolean) { /** Helper used to track which confs were set by which flows. */ private case class FlowConf(key: String, value: String, flowIdentifier: TableIdentifier) @@ -203,7 +204,7 @@ private class FlowResolver(rawGraph: DataflowGraph) { flow: UnresolvedFlow, funcResult: FlowFunctionResult): ResolvedFlow = { flow match { - case acf: AutoCdcFlow => new AutoCdcMergeFlow(acf, funcResult) + case acf: AutoCdcFlow => new AutoCdcMergeFlow(acf, funcResult, sessionCaseSensitive) case utf: UntypedFlow => transformUntypedFlowToResolvedFlow(utf, funcResult) } } diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraph.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraph.scala index c5210976d3f98..cd1dbf63ccdf1 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraph.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraph.scala @@ -22,7 +22,7 @@ import scala.util.Try import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.pipelines.graph.DataflowGraph.mapUnique -import org.apache.spark.sql.pipelines.util.SchemaMergingUtils +import org.apache.spark.sql.pipelines.util.SchemaInferenceUtils import org.apache.spark.sql.types.StructType /** @@ -146,7 +146,9 @@ case class DataflowGraph( * upstream flows * @return The reanalyzed flow */ - protected[graph] def reanalyzeFlow(srcFlow: Flow): ResolvedFlow = { + protected[graph] def reanalyzeFlow( + srcFlow: Flow, + sessionCaseSensitive: Boolean): ResolvedFlow = { val upstreamDatasetIdentifiers = dfsInternal( flowNodes(srcFlow.identifier).output, downstream = false, @@ -164,35 +166,46 @@ case class DataflowGraph( tables = table.get(srcFlow.destinationIdentifier).toSeq, sinks = sink.get(srcFlow.destinationIdentifier).toSeq ) - subgraph.resolve().resolvedFlow(srcFlow.identifier) + subgraph.resolve(sessionCaseSensitive).resolvedFlow(srcFlow.identifier) } /** * A map of the inferred schema of each table, computed by merging the analyzed schemas * of all flows writing to that table. + * + * The merge honors the effective `spark.sql.caseSensitive` of the flows writing to each table: + * under case-insensitive analysis two flows emitting column names that differ only in case + * contribute a single column rather than both, which would otherwise produce a target schema the + * engine's own resolver cannot disambiguate. Which spelling survives such a fold is fixed by + * [[SchemaInferenceUtils.inferSchemaFromFlows]], which merges in sorted flow identifier order. */ - lazy val inferredSchema: Map[TableIdentifier, StructType] = { - flowsTo.view.mapValues { flows => - flows - .map { flow => - resolvedFlow(flow.identifier).schema - } - .reduce(SchemaMergingUtils.mergeSchemas) - }.toMap + def inferSchemas(sessionCaseSensitive: Boolean): Map[TableIdentifier, StructType] = { + flowsTo.map { case (destinationIdentifier, flows) => + val resolvedFlows = flows.map { flow => + resolvedFlow(flow.identifier) + } + destinationIdentifier -> SchemaInferenceUtils.inferSchemaFromFlows( + tableIdentifier = destinationIdentifier, + flows = resolvedFlows, + userSpecifiedSchema = None, + sessionCaseSensitive = sessionCaseSensitive) + } } /** * The internal auxiliary tables owned by each destination [[Table]], derived from the resolved - * flows writing to it and that destination's [[inferredSchema]]. Keyed by the destination's - * identifier; only destinations that actually require auxiliary tables appear. Today only - * AutoCDC flow destination tables have an auxiliary table, and exactly one. + * flows writing to it and the destination schemas inferred from them. Keyed by the destination's + * identifier; only destinations that actually require auxiliary tables appear. Today only AutoCDC + * flow destination tables have an auxiliary table, and exactly one. * * Auxiliary tables are deliberately NOT part of the logical graph (they are never resolved, * connected, or exposed as [[Input]]s); this is purely a derived view used during dataset * materialization to create/evolve them alongside their owning table. The derivation is pure and * performs no catalog access. */ - lazy val auxiliaryTableSpecs: Map[TableIdentifier, AuxiliaryTableSpec] = { + def auxiliaryTableSpecs( + inferredSchemas: Map[TableIdentifier, StructType] + ): Map[TableIdentifier, AuxiliaryTableSpec] = { resolvedFlowsTo.flatMap { case (destinationTableIdentifier, flowsToDestinationTable) => table.get(destinationTableIdentifier).flatMap { destinationTable => flowsToDestinationTable @@ -203,7 +216,7 @@ case class DataflowGraph( .map { autoCdcFlow => val spec = AutoCdcAuxiliaryTable.buildAuxiliaryTableSpecFor( targetTable = destinationTable, - targetTableSchema = inferredSchema(destinationTableIdentifier), + targetTableSchema = inferredSchemas(destinationTableIdentifier), inputAutoCdcFlow = autoCdcFlow ) destinationTableIdentifier -> spec @@ -213,22 +226,22 @@ case class DataflowGraph( } /** Ensure that the [[DataflowGraph]] is valid and throws errors if not. */ - def validate(): DataflowGraph = { - validationFailure.toOption match { + def validate(sessionCaseSensitive: Boolean): DataflowGraph = { + validationFailure(sessionCaseSensitive).toOption match { case Some(exception) => throw exception case None => this } } /** - * Validate the current [[DataflowGraph]] and cache the validation failure. + * Validate the current [[DataflowGraph]] and return the validation failure, if one exists. * * To add more validations, add them in a helper function that throws an exception if the * validation fails, and invoke the helper function here. */ - private lazy val validationFailure: Try[Throwable] = Try { + private def validationFailure(sessionCaseSensitive: Boolean): Try[Throwable] = Try { validateSuccessfulFlowAnalysis() - validateUserSpecifiedSchemas() + validateUserSpecifiedSchemas(sessionCaseSensitive) // Connecting the graph sorts it topologically validateGraphIsTopologicallySorted() validateMultiQueryTables() @@ -236,7 +249,6 @@ case class DataflowGraph( validateEveryDatasetHasFlow() validateTablesAreResettable() validateFlowStreamingness() - inferredSchema }.failed /** @@ -259,10 +271,12 @@ case class DataflowGraph( def resolved: Boolean = flows.forall(f => resolvedFlow.contains(f.identifier)) - def resolve(): DataflowGraph = + def resolve(sessionCaseSensitive: Boolean): DataflowGraph = DataflowGraphTransformer.withDataflowGraphTransformer(this) { transformer => val coreDataflowNodeProcessor = - new CoreDataflowNodeProcessor(rawGraph = this) + new CoreDataflowNodeProcessor( + rawGraph = this, + sessionCaseSensitive = sessionCaseSensitive) transformer .transformDownNodes(coreDataflowNodeProcessor.processNode) .getDataflowGraph diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DatasetManager.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DatasetManager.scala index 4f96f2f709908..61d188f302e13 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DatasetManager.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/DatasetManager.scala @@ -24,7 +24,7 @@ import org.apache.spark.SparkException import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.TableIdentifier -import org.apache.spark.sql.catalyst.analysis.PersistedView +import org.apache.spark.sql.catalyst.analysis.{NoSuchTableException, PersistedView, Resolver} import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connector.catalog.{ CatalogV2Util, @@ -39,9 +39,12 @@ import org.apache.spark.sql.connector.catalog.CatalogV2Util.v2ColumnsToStructTyp import org.apache.spark.sql.connector.expressions.{ClusterByTransform, Expressions, Transform} import org.apache.spark.sql.execution.command.CreateViewCommand import org.apache.spark.sql.pipelines.graph.QueryOrigin.ExceptionHelpers -import org.apache.spark.sql.pipelines.util.PipelinesCatalogUtils +import org.apache.spark.sql.pipelines.util.{ + PipelinesCatalogUtils, + SchemaInferenceUtils, + SchemaMergingUtils +} import org.apache.spark.sql.pipelines.util.SchemaInferenceUtils.diffSchemas -import org.apache.spark.sql.pipelines.util.SchemaMergingUtils import org.apache.spark.sql.types.StructType /** @@ -97,6 +100,9 @@ object DatasetManager extends Logging { val tablesToMaterialize = { tablesToMatz(resolvedDataflowGraph).map(t => t.table.identifier -> t).toMap } + val sessionCaseSensitive = context.spark.sessionState.conf.caseSensitiveAnalysis + val inferredSchemas = resolvedDataflowGraph.inferSchemas(sessionCaseSensitive) + val auxiliaryTableSpecs = resolvedDataflowGraph.auxiliaryTableSpecs(inferredSchemas) // materialized [[DataflowGraph]] where each table has been materialized and each table // has metadata (e.g., normalized table storage path) populated @@ -107,16 +113,30 @@ object DatasetManager extends Logging { if (tablesToMaterialize.keySet.contains(table.identifier)) { try { val isFullRefresh = tablesToMaterialize(table.identifier).isFullRefresh + // Load the existing auxiliary table (if any) once here and thread the snapshot into + // both materializeTable (which uses it for AutoCDC config-drift validation) and + // materializeAuxiliaryTable (which uses it to decide evolve-vs-create). Nothing + // between them mutates the auxiliary table, so a single load is safe and avoids a + // redundant catalog round trip. + val auxiliaryTableSpecOpt = auxiliaryTableSpecs.get(table.identifier) + val existingAuxiliaryTable = auxiliaryTableSpecOpt.flatMap { spec => + val (auxCatalog, auxId) = + PipelinesCatalogUtils.resolveTableCatalog(context.spark, spec.identifier) + loadTableIfExists(auxCatalog, auxId) + } val (tableWithMaterializationMetadata, catalogTableEntity) = materializeTable( resolvedDataflowGraph = resolvedDataflowGraph, table = table, + inferredSchemas = inferredSchemas, isFullRefresh = isFullRefresh, + auxiliaryTableSpecOpt = auxiliaryTableSpecOpt, + existingAuxiliaryTable = existingAuxiliaryTable, context = context ) // Auxiliary tables' lifecycle should follow the table that it is complementary to. // If this table has any auxiliary tables, validate the target can host them and // materialize/full-refresh them accordingly. - resolvedDataflowGraph.auxiliaryTableSpecs.get(table.identifier).foreach { + auxiliaryTableSpecOpt.foreach { auxiliaryTableSpec => // If this table is an AutoCDC target table, as identified by being // accompanied by an AutoCDC auxiliary table, additionally validate that the @@ -134,6 +154,11 @@ object DatasetManager extends Logging { materializeAuxiliaryTable( auxiliaryTableSpec = auxiliaryTableSpec, isFullRefresh = isFullRefresh, + existingAuxiliaryTable = existingAuxiliaryTable, + // The auxiliary schema is derived from its target's, so it evolves under the + // target's effective case sensitivity. + caseSensitive = effectiveCaseSensitivityFor( + resolvedDataflowGraph, table.identifier, context), context = context ) } @@ -267,9 +292,14 @@ object DatasetManager extends Logging { /** * Materializes a table in the catalog. This method will create or update the table in the * catalog based on the given table and context. - * @param resolvedDataflowGraph The resolved [[DataflowGraph]] used to infer the table schema. + * @param resolvedDataflowGraph The resolved [[DataflowGraph]] used for table metadata. * @param table The table to be materialized. + * @param inferredSchemas The schemas inferred from the resolved graph, keyed by table. * @param isFullRefresh Whether this table should be full refreshed or not. + * @param auxiliaryTableSpecOpt The spec for the auxiliary table (if this table has one) + * @param existingAuxiliaryTable The already-loaded auxiliary table for this target (if it has one + * and it exists), used for AutoCDC config-drift validation. Loaded + * once by the caller and shared with [[materializeAuxiliaryTable]]. * @param context The context for the pipeline update. * @return The materialized graph [[Table]] (with additional metadata set) paired with the loaded * DSv2 handle of the just created/evolved table. @@ -277,7 +307,10 @@ object DatasetManager extends Logging { private def materializeTable( resolvedDataflowGraph: DataflowGraph, table: Table, + inferredSchemas: Map[TableIdentifier, StructType], isFullRefresh: Boolean, + auxiliaryTableSpecOpt: Option[AuxiliaryTableSpec], + existingAuxiliaryTable: Option[V2Table], context: PipelineUpdateContext): (Table, V2Table) = { logInfo(log"Materializing metadata for table ${MDC(LogKeys.TABLE_NAME, table.identifier)}.") // Get the DSv2 catalog handler and identifier for the table. @@ -285,7 +318,7 @@ object DatasetManager extends Logging { PipelinesCatalogUtils.resolveTableCatalog(context.spark, table.identifier) val outputSchema = table.specifiedSchema.getOrElse( - resolvedDataflowGraph.inferredSchema(table.identifier).asNullable + inferredSchemas(table.identifier).asNullable ) val mergedProperties = resolveTableProperties(table, identifier) val partitioning = table.partitionCols.toSeq.flatten.map(Expressions.identity) @@ -330,16 +363,57 @@ object DatasetManager extends Logging { context.spark.sql(s"TRUNCATE TABLE ${table.identifier.quotedString}") } + val autoCdcAuxTableSpecOpt = auxiliaryTableSpecOpt.collect { + case autoCdcSpec: AutoCdcAuxiliaryTableSpec => autoCdcSpec + } + val effectiveCaseSensitive = effectiveCaseSensitivityFor( + resolvedDataflowGraph, table.identifier, context) + val effectiveResolver = SchemaInferenceUtils.resolverFor(effectiveCaseSensitive) + + // For an incrementally-updated AutoCDC target, validate that the AutoCDC configuration recorded + // on the auxiliary table has not drifted, BEFORE anything is created or evolved this run. These + // checks read the auxiliary table, so they run whenever IT exists -- independent of whether the + // target exists. That matters when a user drops and recreates the target without dropping the + // internal auxiliary table: the target is then absent (so it is re-created below) but the stale + // auxiliary table survives, and `materializeAuxiliaryTable`'s additive evolve would otherwise + // silently overwrite the recorded key/SCD-type/track-history properties with this run's values. + // Running here turns that into one clear drift error (remedy: full refresh). + if (isTableIncrementallyUpdated) { + autoCdcAuxTableSpecOpt.foreach { + validateNoAutoCdcAuxConfigDrift(_, existingAuxiliaryTable, effectiveResolver) + } + } + // Create the table if absent, otherwise evolve it (schema + properties). existingTableOpt match { case Some(existingTable) => + // The sequencing-type check needs the existing target schema (the type is embedded in the + // target's `_cdc_metadata`), so it runs here, and BEFORE `evolveTable`: `evolveTable` + // ALTERs the target (additively) in place, so a check that ran afterwards would leave the + // target already mutated by a run it then rejects -- and the drift remedy could not undo + // that. Running first means a rejected run leaves the target untouched, and surfaces the + // change as an actionable SEQUENCING_TYPE_DRIFT rather than a generic + // CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE from the schema merge. + if (isTableIncrementallyUpdated) { + autoCdcAuxTableSpecOpt.foreach { autoCdcSpec => + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = + CatalogV2Util.v2ColumnsToStructType(existingTable.columns()), + targetTableIdentifier = autoCdcSpec.targetTableIdentifier, + expectedScdType = autoCdcSpec.expectedScdType, + expectedSequencingType = autoCdcSpec.expectedSequencingType, + resolver = effectiveResolver + ) + } + } evolveTable( catalog = catalog, tableIdentifier = identifier, existingTable = existingTable, desiredSchema = outputSchema, properties = mergedProperties, - mergeWithExistingSchema = isTableIncrementallyUpdated + mergeWithExistingSchema = isTableIncrementallyUpdated, + caseSensitive = effectiveCaseSensitive ) case None => createTable( @@ -394,11 +468,18 @@ object DatasetManager extends Logging { * * @param auxiliaryTableSpec the spec describing the auxiliary table to create/evolve. * @param isFullRefresh whether the owning table is being fully refreshed. + * @param existingAuxiliaryTable the already-loaded auxiliary table (if it exists), loaded once by + * the caller and shared with the config-drift validation in + * [[materializeTable]] rather than re-loaded here. + * @param caseSensitive the effective case sensitivity of the flows writing to the auxiliary + * table's TARGET, whose schema the auxiliary schema is derived from. * @param context the context for the pipeline update. */ private def materializeAuxiliaryTable( auxiliaryTableSpec: AuxiliaryTableSpec, isFullRefresh: Boolean, + existingAuxiliaryTable: Option[V2Table], + caseSensitive: Boolean, context: PipelineUpdateContext): Unit = { // Get the DSv2 catalog handler and identifier for the aux table. val (catalog, auxiliaryTableIdentifier) = @@ -434,33 +515,23 @@ object DatasetManager extends Logging { transforms = Seq.empty ) } else { - loadTableIfExists(catalog, auxiliaryTableIdentifier) match { - case Some(existingAuxiliaryTable) => - auxiliaryTableSpec match { - case autoCdcSpec: AutoCdcAuxiliaryTableSpec => - // For AutoCDC auxiliary tables specifically, we persist metadata about the AutoCDC - // configuration that should be invariant for the flow's lifetime; i.e until it is - // full-refreshed. Validate these configurations remain invariant before attempting - // to evolve the auxiliary table's schema, to prevent corrupting the table. - AutoCdcAuxiliaryTable.validateNoKeyColumnDrift( - existingAuxiliaryTable = existingAuxiliaryTable, - targetTableIdentifier = autoCdcSpec.targetTableIdentifier, - expectedKeyFields = autoCdcSpec.expectedKeyFields, - resolver = context.spark.sessionState.conf.resolver - ) - AutoCdcAuxiliaryTable.validateNoScdTypeDrift( - existingAuxiliaryTable = existingAuxiliaryTable, - targetTableIdentifier = autoCdcSpec.targetTableIdentifier, - expectedScdType = autoCdcSpec.expectedScdType - ) - } + // Uses the auxiliary-table snapshot loaded by the caller (see materializeDatasets), rather + // than re-loading it here. + existingAuxiliaryTable match { + case Some(existingAuxTable) => + // NOTE: AutoCDC configuration-drift validation (key columns, SCD type, sequencing type, + // track-history columns) intentionally runs in [[materializeTable]] BEFORE the target's + // schema is evolved, not here -- see `validateNoAutoCdcAuxConfigDrift`. Validating here + // would be too late: the target has already been ALTERed by the time an aux-owned check + // could reject the run. evolveTable( catalog = catalog, tableIdentifier = auxiliaryTableIdentifier, - existingTable = existingAuxiliaryTable, + existingTable = existingAuxTable, desiredSchema = auxiliaryTableSpec.schema, properties = auxiliaryTableSpec.properties, - mergeWithExistingSchema = true + mergeWithExistingSchema = true, + caseSensitive = caseSensitive ) case None => createTable( @@ -474,11 +545,78 @@ object DatasetManager extends Logging { } } - /** Loads the table at `identifier` from `catalog`, or `None` if it does not exist. */ + /** + * Validate that an incrementally-updated AutoCDC flow's configuration has not drifted from what + * its auxiliary table recorded. Called from [[materializeTable]] before the target and auxiliary + * tables are created/evolved this run, so a rejected run leaves both untouched. + * + * Covers the checks that read the recorded configuration off the existing auxiliary table: key + * columns, SCD type, and track-history columns. Runs whenever the auxiliary table exists, + * independent of the target's existence (see the call site). If the auxiliary table does not + * exist yet (first AutoCDC run), all three are skipped -- there is no recorded configuration to + * drift from. The sequencing-type check is separate ([[AutoCdcAuxiliaryTable]] + * `.validateNoTargetSequencingTypeDrift`) because it reads the target schema, not the aux table. + * + * @param autoCdcSpec the auxiliary-table spec carrying this run's expected AutoCDC configuration. + * @param existingAuxiliaryTableOpt the already-loaded auxiliary table (if it exists), shared with + * the caller and [[materializeAuxiliaryTable]] to avoid a + * redundant load. + * @param resolver the effective resolver of the flows writing to the AutoCDC target. + */ + private def validateNoAutoCdcAuxConfigDrift( + autoCdcSpec: AutoCdcAuxiliaryTableSpec, + existingAuxiliaryTableOpt: Option[V2Table], + resolver: Resolver): Unit = { + existingAuxiliaryTableOpt.foreach { existingAuxiliaryTable => + AutoCdcAuxiliaryTable.validateNoKeyColumnDrift( + existingAuxiliaryTable = existingAuxiliaryTable, + targetTableIdentifier = autoCdcSpec.targetTableIdentifier, + expectedKeyFields = autoCdcSpec.expectedKeyFields, + resolver = resolver + ) + AutoCdcAuxiliaryTable.validateNoScdTypeDrift( + existingAuxiliaryTable = existingAuxiliaryTable, + targetTableIdentifier = autoCdcSpec.targetTableIdentifier, + expectedScdType = autoCdcSpec.expectedScdType + ) + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existingAuxiliaryTable, + targetTableIdentifier = autoCdcSpec.targetTableIdentifier, + expectedTrackHistoryColumnNames = autoCdcSpec.expectedTrackHistoryColumnNames, + resolver = resolver + ) + } + } + + /** + * The effective `spark.sql.caseSensitive` for schema evolution of `tableIdentifier`, read from + * the flows writing to it rather than from the session, so evolution stays consistent with the + * flows whose schemas it is evolving (a pipeline-level `SET` never reaches the session). Fails if + * those flows disagree; see [[SchemaInferenceUtils.effectiveCaseSensitivity]]. + */ + private def effectiveCaseSensitivityFor( + resolvedDataflowGraph: DataflowGraph, + tableIdentifier: TableIdentifier, + context: PipelineUpdateContext): Boolean = { + SchemaInferenceUtils.effectiveCaseSensitivity( + tableIdentifier = tableIdentifier, + flows = resolvedDataflowGraph.flowsTo.getOrElse(tableIdentifier, Seq.empty), + sessionCaseSensitive = context.spark.sessionState.conf.caseSensitiveAnalysis + ) + } + + /** + * Loads the table at `identifier` from `catalog`, or `None` if it does not exist. A single + * `loadTable` guarded by a `NoSuchTableException` catch, rather than a `tableExists` + + * `loadTable` pair: one catalog round trip instead of two, and no window where the table can + * disappear between the existence check and the load. A missing *namespace* is not treated as + * "table absent" -- it surfaces as its own `NoSuchDatabaseException` rather than being swallowed. + */ private def loadTableIfExists( catalog: TableCatalog, identifier: Identifier): Option[V2Table] = { - Option.when(catalog.tableExists(identifier))(catalog.loadTable(identifier)) + try Some(catalog.loadTable(identifier)) + catch { case _: NoSuchTableException => None } } /** @@ -524,6 +662,15 @@ object DatasetManager extends Logging { * @param mergeWithExistingSchema whether the effective schema is the merge of the existing and * desired schemas (additive evolution) rather than the desired * schema as-is. + * @param caseSensitive whether the additive schema merge treats field names differing + * only in case as distinct columns. Callers should pass the + * effective `spark.sql.caseSensitive` used to resolve the schema + * being evolved. When `false`, an incoming column differing from + * an existing one only in case is folded onto it rather than + * added as a duplicate. Only affects the merge (i.e. + * `mergeWithExistingSchema = true`); the subsequent diff always + * keys columns on their exact names, so a case-only rename on a + * non-merging path stays an explicit drop-then-add. */ private def evolveTable( catalog: TableCatalog, @@ -531,13 +678,18 @@ object DatasetManager extends Logging { existingTable: V2Table, desiredSchema: StructType, properties: Map[String, String], - mergeWithExistingSchema: Boolean): Unit = { + mergeWithExistingSchema: Boolean, + caseSensitive: Boolean): Unit = { val currentSchema = v2ColumnsToStructType(existingTable.columns()) val targetSchema = if (mergeWithExistingSchema) { - SchemaMergingUtils.mergeSchemas(currentSchema, desiredSchema) + SchemaMergingUtils.mergeSchemas(currentSchema, desiredSchema, caseSensitive) } else { desiredSchema } + // `diffSchemas` keys column identity on exact field names. On the incremental path the merge + // above has already folded a case-only-differing incoming field onto the persisted one. On the + // non-merging paths (materialized views, full refresh), `targetSchema` is the declared schema + // as-is, where exact-name matching keeps a case-only rename visible as a schema change. val columnChanges = diffSchemas(currentSchema, targetSchema) val existingProperties = existingTable.properties() diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/Flow.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/Flow.scala index 54525aab87126..a8be0a4c2ec33 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/Flow.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/Flow.scala @@ -23,6 +23,7 @@ import org.apache.spark.SparkException import org.apache.spark.internal.Logging import org.apache.spark.sql.{functions => F, AnalysisException, Column} import org.apache.spark.sql.catalyst.{AliasIdentifier, TableIdentifier} +import org.apache.spark.sql.catalyst.analysis.Resolver import org.apache.spark.sql.classic.DataFrame import org.apache.spark.sql.pipelines.autocdc.{ AutoCdcReservedNames, @@ -33,6 +34,7 @@ import org.apache.spark.sql.pipelines.autocdc.{ Scd2BatchProcessor, ScdType } +import org.apache.spark.sql.pipelines.util.SchemaInferenceUtils import org.apache.spark.sql.types.{DataType, StructField, StructType} /** @@ -250,8 +252,15 @@ class AppendOnceFlow( */ class AutoCdcMergeFlow( val flow: AutoCdcFlow, - val funcResult: FlowFunctionResult + val funcResult: FlowFunctionResult, + sessionCaseSensitive: Boolean ) extends ResolvedFlow { + private[graph] val effectiveResolver: Resolver = SchemaInferenceUtils.resolverFor( + SchemaInferenceUtils.effectiveCaseSensitivity( + tableIdentifier = destinationIdentifier, + flows = Seq(this), + sessionCaseSensitive = sessionCaseSensitive)) + requireReservedPrefixAbsentInSourceColumns() requireReservedFrameworkColumnsAbsentInSourceColumns() @@ -263,14 +272,11 @@ class AutoCdcMergeFlow( schemaName = "changeDataFeed", schema = df.schema, columnSelection = changeArgs.columnSelection, - resolver = spark.sessionState.conf.resolver + resolver = effectiveResolver ) // AutoCDC flows require all key columns to be present in the user-selected source schema, // so that they survive into the target table where SCD reconciliation needs them. requireKeysPresentInSelectedSchema(selectedSchema) - // SCD2 flows may specify history-tracking columns; validate they resolve to eligible columns - // of the selected schema at construction time, rather than failing mid-stream on first batch. - requireTrackHistoryColumnsResolvableInSelectedSchema(selectedSchema) selectedSchema } @@ -278,6 +284,38 @@ class AutoCdcMergeFlow( private[graph] val sequencingType: DataType = df.select(changeArgs.sequencing).schema.head.dataType + /** + * SCD2 only: the effective set of history-tracking column names for this flow, resolved from the + * [[userSelectedSchema]] (the user-selected source columns), not from the persisted/evolved + * target schema. This is the single source of truth for the tracked set: + * + * - Resolving against [[userSelectedSchema]] means the tracked set follows the flow's own + * selection. In particular, under default or `* EXCEPT` tracking, dropping a column from the + * source (or from the `COLUMNS` selection) removes it from the tracked set -- rather than the + * column lingering as tracked because it still exists in the (sticky) target schema. + * - It is recorded on the auxiliary table and drift-checked across runs: a change to this set + * reinterprets which transitions open a new SCD2 record, which cannot be applied to + * already-reconciled history, so any change requires a full refresh (see + * [[AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift]]). Because the set is + * selection-derived, adding or dropping a source column under default / `* EXCEPT` tracking + * is such a change; this is an intended divergence from SCD1, where non-key schema evolution + * needs no full refresh. + * + * Computing it here (rather than in the aux-table spec builder from the target schema) also + * validates the selection at construction time: an unresolvable or ineligible explicit + * `TRACK HISTORY ON` selection throws `AUTOCDC_COLUMNS_NOT_FOUND_IN_SCHEMA` here, before the + * first microbatch, rather than failing deep inside the SCD2 batch processor. `None` for SCD1. + */ + private[graph] val trackHistoryColumnNames: Option[Seq[String]] = + changeArgs.storedAsScdType match { + case ScdType.Type2 => + Some(Scd2BatchProcessor.computeTrackedHistoryColumns( + schema = userSelectedSchema, + changeArgs = changeArgs, + resolver = effectiveResolver)) + case ScdType.Type1 => None + } + /** * Returns the augmented output schema of this flow, which can differ from the schema of the * source change-data-feed dataframe. @@ -376,7 +414,7 @@ class AutoCdcMergeFlow( * names that use the reserved Spark AutoCDC prefix. */ private def requireReservedPrefixAbsentInSourceColumns(): Unit = { - val resolver = spark.sessionState.conf.resolver + val resolver = effectiveResolver val reservedPrefix = AutoCdcReservedNames.prefix def nameContainsReservedPrefix(name: String): Boolean = { @@ -408,7 +446,7 @@ class AutoCdcMergeFlow( * during preprocessing. No-op for SCD1, which has no such columns. */ private def requireReservedFrameworkColumnsAbsentInSourceColumns(): Unit = { - val resolver = spark.sessionState.conf.resolver + val resolver = effectiveResolver val reservedPrefix = AutoCdcReservedNames.prefix // Only the non-prefixed reserved names need checking here; prefixed ones are already rejected @@ -440,7 +478,7 @@ class AutoCdcMergeFlow( * Validate all keys specified in changeArgs are actually present in the user-selected schema. */ private def requireKeysPresentInSelectedSchema(selectedSchema: StructType): Unit = { - val resolver = spark.sessionState.conf.resolver + val resolver = effectiveResolver changeArgs.keys .find(key => !selectedSchema.fieldNames.exists(name => resolver(name, key.name))) @@ -455,25 +493,4 @@ class AutoCdcMergeFlow( } } - /** - * Validate that this flow's [[ChangeArgs.trackHistorySelection]] (SCD2 `TRACK HISTORY ON ...`) - * resolves against the user-selected source schema at construction time. Without this, an - * unresolvable or ineligible (key/framework) tracking column would only surface when the first - * microbatch runs reconciliation, deep inside the SCD2 batch processor. - * - * Delegates to [[Scd2BatchProcessor.computeTrackedHistoryColumns]] -- the same resolution used at - * runtime -- so the two can never diverge; it throws `AUTOCDC_COLUMNS_NOT_FOUND_IN_SCHEMA` on an - * unresolvable selection. `trackHistorySelection` is `None` for SCD1 (enforced by [[ChangeArgs]]) - * and for SCD2 flows that do not restrict tracking, in which case resolution is a no-op. - */ - private def requireTrackHistoryColumnsResolvableInSelectedSchema( - selectedSchema: StructType): Unit = { - if (changeArgs.trackHistorySelection.isDefined) { - Scd2BatchProcessor.computeTrackedHistoryColumns( - schema = selectedSchema, - changeArgs = changeArgs, - resolver = spark.sessionState.conf.resolver - ) - } - } } diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala index c36d824bdf4a6..c665c344c53b4 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala @@ -105,6 +105,10 @@ trait FlowExecution { /** Context about this pipeline update. */ def updateContext: PipelineUpdateContext + /** The session's `spark.sql.caseSensitive` fallback for resolving this flow. */ + protected def sessionCaseSensitive: Boolean = + spark.sessionState.conf.caseSensitiveAnalysis + /** The thread execution context for the current `FlowExecution`. */ implicit val executionContext: ExecutionContext = { ExecutionContext.fromExecutor(FlowExecution.threadPool) @@ -232,7 +236,7 @@ class StreamingTableWrite( override def getOrigin: QueryOrigin = flow.origin def startStream(): StreamingQuery = { - val data = graph.reanalyzeFlow(flow).df + val data = graph.reanalyzeFlow(flow, sessionCaseSensitive).df val dataStreamWriter = data .writeStream .queryName(displayName) @@ -260,7 +264,7 @@ class BatchTableWrite( def executeInternal(): Future[Unit] = { SparkSessionUtils.withSqlConf(spark, sqlConf.toList: _*) { updateContext.flowProgressEventLogger.recordRunning(flow = flow) - val data = graph.reanalyzeFlow(flow).df + val data = graph.reanalyzeFlow(flow, sessionCaseSensitive).df Future { val dataFrameWriter = data.write destination.format.foreach(dataFrameWriter.format) @@ -298,7 +302,7 @@ class SinkWrite( override def getOrigin: QueryOrigin = flow.origin def startStream(): StreamingQuery = { - val data = graph.reanalyzeFlow(flow).df + val data = graph.reanalyzeFlow(flow, sessionCaseSensitive).df data.writeStream .queryName(displayName) .option("checkpointLocation", checkpointPath) @@ -328,7 +332,7 @@ class Scd1MergeStreamingWrite( override def getOrigin: QueryOrigin = flow.origin override def startStream(): StreamingQuery = { - val sourceChangeDataFeed = graph.reanalyzeFlow(flow).df + val sourceChangeDataFeed = graph.reanalyzeFlow(flow, sessionCaseSensitive).df // The auxiliary table is created and evolved during dataset materialization (see // [[DatasetManager]]), so it already exists by the time this flow executes; resolve its @@ -373,7 +377,7 @@ class Scd2MergeStreamingWrite( override def getOrigin: QueryOrigin = flow.origin override def startStream(): StreamingQuery = { - val sourceChangeDataFeed = graph.reanalyzeFlow(flow).df + val sourceChangeDataFeed = graph.reanalyzeFlow(flow, sessionCaseSensitive).df // The auxiliary table is created and evolved during dataset materialization (see // [[DatasetManager]]), so it already exists by the time this flow executes; resolve its diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphErrors.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphErrors.scala index c835665a0f380..7bc26e6288c14 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphErrors.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphErrors.scala @@ -109,4 +109,33 @@ object GraphErrors { cause = Option(cause.orNull) ) } + + /** + * Throws if the flows writing to one table disagree on a configuration whose value determines + * how the table's schema is derived, so that the resulting schema would otherwise depend on the + * order the flows happen to be evaluated in. + * + * @param tableIdentifier the destination table the conflicting flows write to + * @param configKey the configuration the flows disagree on + * @param valuesByFlow the distinct values, each with the flows that declared it + */ + def conflictingFlowConfigurationError( + tableIdentifier: TableIdentifier, + configKey: String, + valuesByFlow: Map[String, Seq[TableIdentifier]]): AnalysisException = { + val rendered = valuesByFlow.toSeq + .sortBy(_._1) + .map { case (value, flows) => + s"$value (${flows.map(_.unquotedString).sorted.mkString(", ")})" + } + .mkString("; ") + new AnalysisException( + errorClass = "CONFLICTING_PIPELINE_FLOW_CASE_SENSITIVITY", + messageParameters = Map( + "tableName" -> tableIdentifier.unquotedString, + "configKey" -> configKey, + "flowConfigurations" -> rendered + ) + ) + } } diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphExecution.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphExecution.scala index c687c7f01ed7a..db1e6c546da12 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphExecution.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphExecution.scala @@ -261,12 +261,32 @@ object GraphExecution extends Logging { } /** - * Analyze the exception thrown by flow execution and figure out if we should retry the execution, - * or we need to reanalyze the flow entirely to resolve issues like schema changes. + * Represents that the `FlowExecution` should be stopped because a streaming flow's set of + * sources changed since the last run. This is unrecoverable without a full refresh, so the flow + * must not be retried regardless of the remaining retry budget. + */ + private case class StreamingSourcesChanged( + cause: Throwable, + flowDisplayName: String + ) extends FlowExecutionStopReason { + override lazy val runTerminationReason: RunTerminationReason = { + StreamingSourcesChangedFailure(flowDisplayName, Option(cause)) + } + override lazy val failureMessage: String = { + s"Flow '$flowDisplayName' had streaming sources added or removed. It will not be " + + s"retried. Please perform a full refresh to rebuild it against the current sources." + } + } + + /** + * Analyze the exception thrown by flow execution and decide whether to retry the execution or + * stop it. The result is either RetryFlowExecution or StopFlowExecution; this function does not + * reanalyze the flow itself. * This should be the narrow waist for all exception analysis in flow execution. - * TODO: currently it only handles schema change and max retries, we should aim to extend this to - * include other non-retryable exception as well so we can have a single SoT for all these error - * matching logic. + * Currently it handles max retries and streaming source changes; other non-retryable errors are + * still routed through the retry path. + * TODO: extend this to include other non-retryable exceptions as well so we can have a single + * SoT for all these error matching logic. * @param ex Exception to analyze. * @param flowDisplayName The user facing flow name with the error. * @param currentNumTries Number of times the flow has been tried. @@ -278,8 +298,12 @@ object GraphExecution extends Logging { currentNumTries: => Int, maxAllowedRetries: => Int ): FlowExecutionAction = { - val flowExecutionNonRetryableReasonOpt = if (currentNumTries > maxAllowedRetries) { - Some(MaxRetryExceeded(ex, flowDisplayName, maxAllowedRetries)) + val error = ex + val flowExecutionNonRetryableReasonOpt = if (PipelinesErrors.streamingSourcesChanged(error)) { + // Source-set changes need a full refresh, so they are never retried. + Some(StreamingSourcesChanged(error, flowDisplayName)) + } else if (currentNumTries > maxAllowedRetries) { + Some(MaxRetryExceeded(error, flowDisplayName, maxAllowedRetries)) } else { None } diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphValidations.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphValidations.scala index ebe90b677aa21..5b4eaba4e7095 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphValidations.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/GraphValidations.scala @@ -252,7 +252,7 @@ trait GraphValidations extends Logging { } } - protected def validateUserSpecifiedSchemas(): Unit = { + protected def validateUserSpecifiedSchemas(sessionCaseSensitive: Boolean): Unit = { // Look up tables by their destination identifier, not by the flow's own identifier. The two // coincide only for an implicit/default flow (whose identifier equals its destination // table's); for a named flow (e.g. `CREATE FLOW <name> AS AUTO CDC INTO <target>`) they @@ -262,8 +262,10 @@ trait GraphValidations extends Logging { // schema of all incoming flows. This must be equivalent to the declared schema. val inferredSchema = SchemaInferenceUtils .inferSchemaFromFlows( + tableIdentifier = t.identifier, flowsTo(t.identifier).map(f => resolvedFlow(f.identifier)), - userSpecifiedSchema = t.specifiedSchema + userSpecifiedSchema = t.specifiedSchema, + sessionCaseSensitive = sessionCaseSensitive ) t.specifiedSchema.foreach { ss => diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelineExecution.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelineExecution.scala index d35d701d44e57..724ccc56e35c7 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelineExecution.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelineExecution.scala @@ -110,7 +110,8 @@ class PipelineExecution(context: PipelineUpdateContext) { private def resolveGraph(): DataflowGraph = { try { - context.unresolvedGraph.resolve().validate() + val sessionCaseSensitive = context.spark.sessionState.conf.caseSensitiveAnalysis + context.unresolvedGraph.resolve(sessionCaseSensitive).validate(sessionCaseSensitive) } catch { case e: UnresolvedPipelineException => handleInvalidPipeline(e) diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelinesErrors.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelinesErrors.scala index b194e9c235fba..ef2ac9732c995 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelinesErrors.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/PipelinesErrors.scala @@ -63,6 +63,30 @@ object PipelinesErrors extends Logging { getExceptionChain(throwable).exists(check) } + /** + * Returns true if `ex` (or any of its causes) indicates that a streaming flow's set of sources + * changed since the last run. This is unrecoverable without a full refresh, so a flow that fails + * with this error must not be retried. + * + * Structured Streaming reports this as a bare `AssertionError`, so the only signal available + * here is the message text, produced by the source-count assertion in + * `org.apache.spark.sql.execution.streaming.checkpointing.OffsetSeq.toStreamProgress`. Since + * this predicate drives the retry decision, a change to that message silently turns these + * failures back into retried ones; keep the two in sync. Giving that assertion an error + * condition and matching on it here would remove the coupling. + */ + private[graph] def streamingSourcesChanged(ex: Throwable): Boolean = { + checkCauses( + throwable = ex, + check = cause => { + cause.isInstanceOf[AssertionError] && + cause.getMessage != null && + cause.getMessage.contains("sources in the checkpoint offsets and now there are") && + cause.getMessage.contains("sources requested by the query. Cannot continue.") + } + ) + } + /** * Checks an error for streaming specific handling. This is a pretty messy signature as a result * of unifying some divergences between the triggered caller in TriggeredGraphExecution and the @@ -87,28 +111,7 @@ object PipelinesErrors extends Logging { maxRetries: Int, onRetry: => Unit ): Unit = { - if (PipelinesErrors.checkCauses( - throwable = ex, - check = ex => { - ex.isInstanceOf[AssertionError] && - ex.getMessage != null && - ex.getMessage.contains("sources in the checkpoint offsets and now there are") && - ex.getMessage.contains("sources requested by the query. Cannot continue.") - } - )) { - val message = s""" - |Flow '${flow.displayName}' had streaming sources added or removed. Please perform a - |full refresh in order to rebuild '${flow.displayName}' against the current set of - |sources. - |""".stripMargin - - env.flowProgressEventLogger.recordFailed( - flow = flow, - exception = ex, - logAsWarn = false, - messageOpt = Option(message) - ) - } else if (flow.once && ex == null) { + if (flow.once && ex == null) { // No need to do anything if this is a ONCE flow with no exception. That just means it's done. } else { val actionFromError = GraphExecution.determineFlowExecutionActionFromError( @@ -120,7 +123,8 @@ object PipelinesErrors extends Logging { actionFromError match { // Simply retry case GraphExecution.RetryFlowExecution => onRetry - // Schema change exception + // Non-retryable stop reason (max retries exceeded, streaming sources changed, ...). + // When shouldRethrow is true, this rethrows so the run stops eagerly on these reasons. case GraphExecution.StopFlowExecution(reason) => val msg = reason.failureMessage if (reason.warnInsteadOfError) { diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/RunTerminationReason.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/RunTerminationReason.scala index c95ce6a197eeb..195042286ed9d 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/RunTerminationReason.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/RunTerminationReason.scala @@ -85,6 +85,21 @@ case class QueryExecutionFailure( } } +/** + * Indicates that a run has failed because a streaming flow's set of sources changed since the last + * run. The flow is not retried, since only a full refresh can recover it. + */ +case class StreamingSourcesChangedFailure( + flowName: String, + override val cause: Option[Throwable]) + extends RunFailure { + override def isFatal: Boolean = false + + override def message: String = + s"Run is $terminalState since flow '$flowName' had streaming sources added or removed. " + + s"Perform a full refresh to rebuild it against the current sources." +} + /** Abstract class used to identify failures related to failures stopping an operation/timeouts. */ abstract class FailureStoppingOperation extends RunFailure { diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/elements.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/elements.scala index 885755fd78ece..baef25176e9a5 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/elements.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/elements.scala @@ -160,12 +160,14 @@ case class Table( * @param identifier The identifier of the parent table. * @param specifiedSchema The user-specified schema for the parent table. * @param incomingFlowIdentifiers The identifiers of all flows that write to the parent table. + * @param sessionCaseSensitive The session's `spark.sql.caseSensitive` fallback. * @param availableFlows All resolved flows that write to the parent table. */ case class VirtualTableInput( identifier: TableIdentifier, specifiedSchema: Option[StructType], incomingFlowIdentifiers: Set[TableIdentifier], + sessionCaseSensitive: Boolean, availableFlows: Seq[ResolvedFlow] = Nil ) extends TableElement with Input with Logging { @@ -186,7 +188,11 @@ case class VirtualTableInput( // Otherwise infer the schema from a combination of the incoming flows and the // user-specified schema, if provided. case _ => - SchemaInferenceUtils.inferSchemaFromFlows(availableFlows, specifiedSchema) + SchemaInferenceUtils.inferSchemaFromFlows( + tableIdentifier = identifier, + flows = availableFlows, + userSpecifiedSchema = specifiedSchema, + sessionCaseSensitive = sessionCaseSensitive) } // Produce either a streaming or batch dataframe, depending on whether this is a virtual diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtils.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtils.scala index 4777772342d7d..fd60684150742 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtils.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtils.scala @@ -20,14 +20,89 @@ package org.apache.spark.sql.pipelines.util import scala.util.control.NonFatal import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.analysis.{ + caseInsensitiveResolution, + caseSensitiveResolution, + Resolver +} import org.apache.spark.sql.connector.catalog.TableChange +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.pipelines.common.DatasetType -import org.apache.spark.sql.pipelines.graph.{GraphElementTypeUtils, GraphErrors, ResolvedFlow} +import org.apache.spark.sql.pipelines.graph.{ + Flow, + GraphElementTypeUtils, + GraphErrors, + ResolvedFlow +} import org.apache.spark.sql.types.{StructField, StructType} object SchemaInferenceUtils { + def resolverFor(caseSensitive: Boolean): Resolver = { + if (caseSensitive) { + caseSensitiveResolution + } else { + caseInsensitiveResolution + } + } + + /** + * The effective `spark.sql.caseSensitive` for schema derivation on `tableIdentifier`, taken from + * the flows writing to it rather than from the session. + * + * A pipeline can set `spark.sql.caseSensitive` for itself, and a `SET` in pipeline source does + * not touch the session: [[org.apache.spark.sql.pipelines.graph.GraphRegistrationContext]] folds + * it into each flow's `sqlConf`, and it is applied when the flow is analyzed and executed. Schema + * derivation therefore has to read it from the same place, or evolution can disagree with the + * flows whose schemas it is deriving from -- e.g. folding an incoming `Value` onto a persisted + * `value` while the flow, resolving case-sensitively, expects `Value` to be its own column. + * + * All flows writing to a table must agree: the value decides whether names differing only in case + * identify the same column, so a disagreement would make the resulting schema depend on the order + * the flows are evaluated in. Throws + * [[org.apache.spark.sql.pipelines.graph.GraphErrors.conflictingFlowConfigurationError]] if they + * disagree. Flows that do not set it at all inherit the session's value. + */ + def effectiveCaseSensitivity( + tableIdentifier: TableIdentifier, + flows: Seq[Flow], + sessionCaseSensitive: Boolean): Boolean = { + val declaredByFlow = flows.flatMap { flow => + flow.sqlConf.get(SQLConf.CASE_SENSITIVE.key).map(value => value -> flow.identifier) + } + if (declaredByFlow.isEmpty) { + return sessionCaseSensitive + } + + // Compare the parsed booleans, so that e.g. "TRUE" and "true" are not reported as a conflict, + // but report the values as written to keep the error recognizable to the user. + val byParsedValue = declaredByFlow.groupBy { case (value, _) => value.trim.toBoolean } + // A flow that leaves the conf unset inherits the session value, which conflicts just as much as + // an explicitly opposite value. + val flowsWithoutDeclaration = flows.filterNot { flow => + flow.sqlConf.contains(SQLConf.CASE_SENSITIVE.key) + } + val effectiveValues = byParsedValue.keySet ++ + Option.when(flowsWithoutDeclaration.nonEmpty)(sessionCaseSensitive) + if (effectiveValues.sizeIs > 1) { + val valuesByFlow = declaredByFlow + .groupBy { case (value, _) => value } + .map { case (value, entries) => value -> entries.map { case (_, id) => id } } ++ + Option + .when(flowsWithoutDeclaration.nonEmpty)( + s"$sessionCaseSensitive (session default)" -> flowsWithoutDeclaration.map(_.identifier) + ) + .toMap + throw GraphErrors.conflictingFlowConfigurationError( + tableIdentifier = tableIdentifier, + configKey = SQLConf.CASE_SENSITIVE.key, + valuesByFlow = valuesByFlow + ) + } + effectiveValues.head + } + /** * Given a set of flows that write to the same destination and possibly a user-specified schema, * we infer the schema of the destination dataset. The logic is as follows: @@ -39,42 +114,70 @@ object SchemaInferenceUtils { * The user-specified schema will take precedence over the inferred schema. * Returns an error if encountered during schema inference or merging the inferred schema with * the user-specified one. + * + * All merges honor the effective `spark.sql.caseSensitive` of the flows writing to + * `tableIdentifier`, falling back to `sessionCaseSensitive` for flows that do not set it. Under + * case-insensitive analysis, flows emitting column names that differ only in case contribute a + * single column, and a declared column matches a flow column differing only in case -- consistent + * with how the rest of the engine resolves those names. + * + * When flows differ only in column casing, the surviving spelling is the one from the flow with + * the lowest identifier: `flows` is merged in sorted identifier order, not in the order given. We + * sort on the identifier's parts (catalog, database, table) to avoid collisions for identifiers + * whose parts contain dots. + * Sorting here rather than at the call sites keeps every caller agreeing on the result, since the + * schemas they derive are compared against each other -- the graph's inferred schema materializes + * the table, while [[org.apache.spark.sql.pipelines.graph.VirtualTableInput]] produces the schema + * downstream flows resolve against, and `diffSchemas` keys column identity on the exact name. Two + * callers ordering the same flows differently would spell one column two ways, leaving a + * downstream view disagreeing with its source and turning the next refresh into a drop-then-add. */ def inferSchemaFromFlows( + tableIdentifier: TableIdentifier, flows: Seq[ResolvedFlow], - userSpecifiedSchema: Option[StructType]): StructType = { + userSpecifiedSchema: Option[StructType], + sessionCaseSensitive: Boolean): StructType = { if (flows.isEmpty) { return userSpecifiedSchema.getOrElse(new StructType()) } require( - flows.forall(_.destinationIdentifier == flows.head.destinationIdentifier), + flows.forall(_.destinationIdentifier == tableIdentifier), "Expected all flows to have the same destination" ) - val inferredSchema = flows.map(_.schema).fold(new StructType()) { (schemaSoFar, schema) => - try { - SchemaMergingUtils.mergeSchemas(schemaSoFar, schema) - } catch { - case NonFatal(e) => - throw GraphErrors.unableToInferSchemaError( - flows.head.destinationIdentifier, - schemaSoFar, - schema, - cause = Option(e) - ) + val caseSensitive = effectiveCaseSensitivity( + tableIdentifier = tableIdentifier, + flows = flows, + sessionCaseSensitive = sessionCaseSensitive + ) + + val inferredSchema = flows + .sortBy(f => (f.identifier.catalog, f.identifier.database, f.identifier.table)) + .map(_.schema) + .fold(new StructType()) { (schemaSoFar, schema) => + try { + SchemaMergingUtils.mergeSchemas(schemaSoFar, schema, caseSensitive) + } catch { + case NonFatal(e) => + throw GraphErrors.unableToInferSchemaError( + tableIdentifier, + schemaSoFar, + schema, + cause = Option(e) + ) + } } - } - val identifier = flows.head.destinationIdentifier val datasetType = GraphElementTypeUtils.getDatasetTypeForMaterializedViewOrStreamingTable(flows) // We merge the inferred schema with the user-specified schema to pick up any schema metadata // that is provided by the user, e.g., comments or column masks. mergeInferredAndUserSchemasIfNeeded( - identifier, + tableIdentifier, datasetType, inferredSchema, - userSpecifiedSchema + userSpecifiedSchema, + caseSensitive ) } @@ -82,12 +185,13 @@ object SchemaInferenceUtils { tableIdentifier: TableIdentifier, datasetType: DatasetType, inferredSchema: StructType, - userSpecifiedSchema: Option[StructType]): StructType = { + userSpecifiedSchema: Option[StructType], + caseSensitive: Boolean): StructType = { userSpecifiedSchema match { case Some(userSpecifiedSchema) => try { // Merge the inferred schema with the user-provided schema hint - SchemaMergingUtils.mergeSchemas(userSpecifiedSchema, inferredSchema) + SchemaMergingUtils.mergeSchemas(userSpecifiedSchema, inferredSchema, caseSensitive) } catch { case NonFatal(e) => throw GraphErrors.incompatibleUserSpecifiedAndInferredSchemasError( @@ -110,6 +214,15 @@ object SchemaInferenceUtils { * 1. New columns that need to be added * 2. Existing columns that need type updates * + * Column identity is keyed on the exact field name, not on a case-normalized one. On the + * incremental streaming-table path, `targetSchema` is the merge of the current and desired + * schemas, and [[SchemaMergingUtils.mergeSchemas]] has already folded an incoming + * case-only-differing field onto the persisted one. On the non-merging paths (materialized views + * and any full refresh), `targetSchema` is the run's declared schema as-is, so exact-name keying + * keeps a case-only rename visible as an explicit drop-then-add. + * Exact keying also avoids silently collapsing two genuinely distinct declared columns that + * differ only in case (`value` and `Value`) into an arbitrary one of the two. + * * @param currentSchema The current schema of the table * @param targetSchema The target schema that we want the table to have * @return A sequence of TableChange objects representing the necessary changes diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaMergingUtils.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaMergingUtils.scala index d15e7ac6425cc..aeab5e623431f 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaMergingUtils.scala +++ b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/util/SchemaMergingUtils.scala @@ -20,7 +20,26 @@ package org.apache.spark.sql.pipelines.util import org.apache.spark.sql.types.StructType object SchemaMergingUtils { - def mergeSchemas(tableSchema: StructType, dataSchema: StructType): StructType = { - StructType.merge(tableSchema, dataSchema).asInstanceOf[StructType] + + /** + * Additively merges `dataSchema` into `tableSchema`, returning a schema that is the union of the + * two (recursing into nested structs/arrays). On a field present in both, `tableSchema`'s name + * and position win; `dataSchema` only contributes fields absent from `tableSchema`. + * + * @param caseSensitive whether two field names that differ only in case are considered distinct. + * When `false` (mirroring a case-insensitive session), `dataSchema`'s field + * is folded onto the matching `tableSchema` field rather than added as a + * separate, case-differing column. Deliberately has no default: every caller + * merges schemas that some pipeline will later resolve names against, so the + * choice belongs to the caller and should be visible at the call site rather + * than silently inherited. Callers should pass the effective + * `spark.sql.caseSensitive` of the flows involved (see + * [[SchemaInferenceUtils.effectiveCaseSensitivity]]). + */ + def mergeSchemas( + tableSchema: StructType, + dataSchema: StructType, + caseSensitive: Boolean): StructType = { + StructType.merge(tableSchema, dataSchema, caseSensitive).asInstanceOf[StructType] } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/AutoCdcFlowSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/AutoCdcFlowSuite.scala index a3d26e6f47e55..afe1915a22ab5 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/AutoCdcFlowSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/AutoCdcFlowSuite.scala @@ -44,6 +44,8 @@ import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StringType, */ class AutoCdcFlowSuite extends QueryTest with SharedSparkSession { + import testImplicits._ + private val testIdentifier = TableIdentifier("cdc_target", Some("db")) /** A no-op [[FlowFunction]] that throws if invoked; AutoCdcFlow tests should never call it. */ @@ -176,13 +178,14 @@ class AutoCdcFlowSuite extends QueryTest with SharedSparkSession { trackHistorySelection = trackHistorySelection ) ) - new AutoCdcMergeFlow(flow, successfulFuncResult(sourceDf)) + new AutoCdcMergeFlow( + flow, + successfulFuncResult(sourceDf), + spark.sessionState.conf.caseSensitiveAnalysis) } /** A stable 3-column source streaming dataframe used across most schema tests. */ private def threeColumnSourceDf(): DataFrame = { - val session = spark - import session.implicits._ MemoryStream[(Int, String, Option[Long])].toDS().toDF("id", "name", "seq") } @@ -413,8 +416,6 @@ class AutoCdcFlowSuite extends QueryTest with SharedSparkSession { test("AutoCdcMergeFlow.schema's SCD2 framework columns use the resolved sequencing type") { // A non-Long sequencing expression must flow through to __START_AT / __END_AT and the // metadata struct's record-start-at field. - val session = spark - import session.implicits._ val sourceDf = MemoryStream[(Int, String, Int)].toDS().toDF("id", "name", "seq") val resolvedFlow = newAutoCdcMergeFlow(sourceDf, storedAsScdType = ScdType.Type2) @@ -431,8 +432,6 @@ class AutoCdcFlowSuite extends QueryTest with SharedSparkSession { // A multi-column sequence is expressed as a struct over the ordering columns; its resolved // type is the corresponding StructType, which must flow into __START_AT / __END_AT and the // metadata struct's record-start-at field unchanged, including each field's own nullability. - val session = spark - import session.implicits._ // seq1 is a non-null Int; seq2 is a nullable Long (Option[Long]), so the struct carries // mixed per-field nullability. val sourceDf = @@ -547,8 +546,6 @@ class AutoCdcFlowSuite extends QueryTest with SharedSparkSession { /** Builds an empty source df with `id` + `seq` + the supplied extra columns. */ private def sourceDfWithExtraColumns(extraColumns: (String, DataType)*): DataFrame = { - val session = spark - import session.implicits._ val baseStream = MemoryStream[(Int, Option[Long])].toDS().toDF("id", "seq") extraColumns.foldLeft(baseStream) { case (acc, (name, dt)) => acc.withColumn(name, F.lit(null).cast(dt)) diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessorSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessorSuite.scala index 42b294248413d..de82f05047cd5 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessorSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2BatchProcessorSuite.scala @@ -89,6 +89,44 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { spark.createDataFrame(spark.sparkContext.parallelize(rows), schema) } + /** + * Select the affected aux rows the way [[Scd2ForeachBatchHandler]] does: compute the shared + * per-key affected-sequence cutoff across both tables first, then apply it to the aux table. + */ + private def findAffectedAuxRows( + processor: Scd2BatchProcessor, + aux: DataFrame, + target: DataFrame, + minSeq: DataFrame, + batchId: Long = 100L): DataFrame = + processor.findAffectedRowsFromAuxiliaryTable( + rawAuxiliaryTableDf = aux, + perKeyAffectedSequenceCutoffDf = processor.computePerKeyAffectedSequenceCutoff( + rawAuxiliaryTableDf = aux, + targetTableDf = target, + perKeyMinimumSequenceInMicrobatchDf = minSeq, + batchId = batchId + ), + batchId = batchId + ) + + /** Target-side counterpart of [[findAffectedAuxRows]]. */ + private def findAffectedTargetRows( + processor: Scd2BatchProcessor, + target: DataFrame, + aux: DataFrame, + minSeq: DataFrame, + batchId: Long = 100L): DataFrame = + processor.findAffectedRowsFromTargetTable( + targetTableDf = target, + perKeyAffectedSequenceCutoffDf = processor.computePerKeyAffectedSequenceCutoff( + rawAuxiliaryTableDf = aux, + targetTableDf = target, + perKeyMinimumSequenceInMicrobatchDf = minSeq, + batchId = batchId + ) + ) + /** * Build a [[Scd2BatchProcessor]] suitable for `findAffected*` and * `computeMinimumSequencePerKey` tests. The `sequencing` is fixed to `F.col("seq")`, @@ -822,17 +860,21 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val keySchema = new StructType().add("id", IntegerType) val userSchema = keySchema.add("value", StringType) - // Two keys to demonstrate per-key anchor isolation. + // Two keys to demonstrate per-key cutoff isolation. // // Input row shape per `auxTableOf`: // (id, value, __START_AT, __END_AT, Row(recordStartAt), deletedByBatchId) // - // Key 1: aux rows at recordStartAt 3, 5, 10. minSeq = 10. - // - 3 -> older than the anchor; dropped. - // - 5 -> anchor (max < 10); included. - // - 10 -> at minSeq; included via the >= branch (NOT as anchor; selection is strict <). - // Key 2: only one aux row at 7, minSeq = 7. - // - 7 -> at minSeq; included via >= branch. No anchor (no rows < 7 for this key). + // The target table is empty throughout, so every cutoff below comes from the aux table. + // + // Key 1: aux rows at recordStartAt 3, 5, 10. minSeq = 10, so the cutoff is 5 (the largest + // recordStartAt strictly below minSeq). + // - 3 -> below the cutoff; dropped. + // - 5 -> sits at the cutoff; included. + // - 10 -> after the cutoff; included. + // Key 2: only one aux row at 7, minSeq = 7. Nothing precedes minSeq, so the cutoff falls + // back to minSeq itself. + // - 7 -> sits at the cutoff; included. val aux = auxTableOf(userSchema)( Row(1, "v1.3", 3L, null, Row(3L), null), Row(1, "v1.5", 5L, null, Row(5L), null), @@ -843,19 +885,16 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(1, 10L), Row(2, 7L) ) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, expectedAnswer = Seq( - Row(1, "v1.5", 5L, null, Row(5L)), // anchor for key=1 - Row(1, "v1.10", 10L, null, Row(10L)), // >= minSeq for key=1 - Row(2, "v2.7", 7L, null, Row(7L)) // >= minSeq for key=2 (no anchor) + Row(1, "v1.5", 5L, null, Row(5L)), // sits at key=1's cutoff + Row(1, "v1.10", 10L, null, Row(10L)), // after key=1's cutoff + Row(2, "v2.7", 7L, null, Row(7L)) // sits at key=2's cutoff, which fell back to minSeq ) ) } @@ -867,27 +906,23 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { // Aux carries a mix of row kinds for one key. The find function does NOT distinguish // between them - it filters purely on `recordStartAt` - so a tombstone, a no-op upsert - // run head, and a continuation are all eligible anchor candidates and all eligible for - // the >= minSeq inclusion branch. + // run head, and a continuation can all set the cutoff, and can all be selected by it. val aux = auxTableOf(userSchema)( // Tombstone at recordStartAt = 3 (deleted at sequence 3): startAt = endAt = 3. - // Older than the anchor; dropped. + // Below the cutoff; dropped. Row(1, null, 3L, 3L, Row(3L), null), // No-op upsert continuation at recordStartAt = 7: startAt inherits its run head's - // recordStartAt, endAt is null. Anchor for minSeq=10 (max < 10). + // recordStartAt, endAt is null. Sets the cutoff for minSeq=10 (nearest below it). Row(1, "alice", 5L, null, Row(7L), null), - // Tombstone at recordStartAt = 12: at-or-after minSeq, included via >= branch. + // Tombstone at recordStartAt = 12: after the cutoff; included. Row(1, null, 12L, 12L, Row(12L), null), - // No-op upsert continuation at recordStartAt = 15: included via >= branch. + // No-op upsert continuation at recordStartAt = 15: after the cutoff; included. Row(1, "bob", 13L, null, Row(15L), null) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, @@ -910,19 +945,16 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(1, "alice", 2L, null, Row(12L), null) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, expectedAnswer = Seq( - // Row with record start at of 8 gets pulled in as an anchor, + // Row with record start at of 8 sets the cutoff, Row(1, "alice", 2L, null, Row(8L)), - // Row with record start at of 12 gets pulled in as a regular affected row. + // Row with record start at of 12 sits after the cutoff. Row(1, "alice", 2L, null, Row(12L)) ) ) @@ -933,31 +965,28 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val keySchema = new StructType().add("id", IntegerType) val userSchema = keySchema.add("value", StringType) - // Tombstone-as-anchor is incidental: the find function selects the anchor purely on - // `max recordStartAt < minSeq`, so a tombstone qualifies just like any other row kind. - // Downstream reconciliation does not actually rely on the anchor when it is a - // tombstone (a delete already closed the prior run, so any subsequent incoming event - // is necessarily a fresh run head regardless of whether the anchor is surfaced). We - // still pull it in as a harmless side effect of the range filter, and this behavior is + // A tombstone setting the cutoff is incidental: the cutoff is the largest effective + // record start below minSeq across both tables, so a tombstone qualifies just like any + // other row kind. Downstream reconciliation does not actually rely on that row when it + // is a tombstone (a delete already closed the prior run, so any subsequent incoming + // event is necessarily a fresh run head regardless of whether it is surfaced). We still + // pull it in as a harmless side effect of the range filter, and this behavior is // documented via test. val aux = auxTableOf(userSchema)( Row(1, null, 7L, 7L, Row(7L), null), Row(1, null, 12L, 12L, Row(12L), null) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, expectedAnswer = Seq( - // Pulled in as anchor. + // Sets the cutoff. Row(1, null, 7L, 7L, Row(7L)), - // Pulled in as regular affected row. + // Sits after the cutoff. Row(1, null, 12L, 12L, Row(12L)) ) ) @@ -973,9 +1002,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { // The idempotency filter retains rows deleted by `currentBatchId` (so a mid-flight // retry sees its own prior writes) and drops rows deleted by any other batch. This - // applies uniformly to both the anchor and non-anchor affected rows. + // applies uniformly to the row at the cutoff and to the rows after it. val aux = auxTableOf(userSchema)( - // Anchor candidate (recordStartAt < minSeq): + // Cutoff candidate (recordStartAt < minSeq): Row(1, "anchor", 5L, null, Row(5L), currentBatchId), // deleted by current -> kept // At-or-after minSeq: Row(1, "live", 10L, null, Row(10L), null), // not deleted -> kept @@ -983,10 +1012,13 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(1, "ignored", 12L, null, Row(12L), differentBatchId) // deleted by another -> dropped ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, + val result = findAffectedAuxRows( + processor, + aux = aux, + target = target, + minSeq = minSeq, batchId = currentBatchId ) @@ -1009,22 +1041,25 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val currentBatchId = 100L val differentBatchId = 99L - // Codifies the step-ordering invariant inside `findAffectedRowsFromAuxiliaryTable`: the - // idempotency filter MUST run before the anchor `max(...)` aggregation. Here the closest + // Codifies the step-ordering invariant inside `computePerKeyAffectedSequenceCutoff`: the + // idempotency filter MUST run before the cutoff `max(...)` aggregation. Here the closest // pre-minSeq candidate (recordStartAt=7) was logically deleted by a different batch, so - // it is filtered out and the anchor falls back to recordStartAt=3. If a future refactor + // it is filtered out and the cutoff falls back to recordStartAt=3. If a future refactor // were to flip these two steps (e.g. as a "perf optimization"), this test would catch it - // because the natural-anchor row (7) would otherwise be selected and then dropped, leaving - // no anchor at all. + // because row 7 would set the cutoff and then be dropped, leaving nothing at the cutoff + // at all. val aux = auxTableOf(userSchema)( Row(1, "live3", 3L, null, Row(3L), null), Row(1, "stale7", 7L, null, Row(7L), differentBatchId) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, + val result = findAffectedAuxRows( + processor, + aux = aux, + target = target, + minSeq = minSeq, batchId = currentBatchId ) @@ -1046,12 +1081,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { // `_cdc_metadata` struct schema untouched. val aux = auxTableOf(userSchema)(Row(1, "v", 5L, null, Row(5L), null)) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) assert(!result.columns.contains(Scd2BatchProcessor.deletedByBatchIdColName)) val cdcMetadataField = result.schema(AutoCdcReservedNames.cdcMetadataColName) @@ -1065,14 +1097,11 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val aux = auxTableOf(userSchema)(Row(1, "v", 5L, null, Row(5L), null)) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) - // The lone aux row is the anchor (recordStartAt=5 < minSeq=10, no other candidates). + // The lone aux row sets the cutoff (recordStartAt=5 < minSeq=10, no other candidates). checkAnswer( df = result, expectedAnswer = Seq(Row(1, "v", 5L, null, Row(5L))) @@ -1087,9 +1116,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val processor = processorWithKeys(Seq("region", "customer_id")) - // Three composite keys: (US, 1), (EU, 1), (US, 2). Each is independent. - // (US, 1): anchor at 3; row at 10 included via >=. - // (EU, 1): anchor at 4; no rows at or after 12 -> only the anchor. + // Three composite keys: (US, 1), (EU, 1), (US, 2). Each gets its own cutoff. + // (US, 1): cutoff at 3; the row at 10 follows it. + // (EU, 1): cutoff at 4; nothing follows it, so only the cutoff row is selected. // (US, 2): no aux rows -> contributes nothing. val aux = auxTableOf(userSchema)( Row("US", 1, "us1.3", 3L, null, Row(3L), null), @@ -1101,12 +1130,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row("EU", 1, 12L), Row("US", 2, 100L) ) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, @@ -1125,12 +1151,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val aux = auxTableOf(userSchema)() val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) assert(result.collect().isEmpty) } @@ -1144,12 +1167,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { // Aux only has rows for key=1. Microbatch only sees key=2. val aux = auxTableOf(userSchema)(Row(1, "v", 5L, null, Row(5L), null)) val minSeq = minSeqOf(keySchema)(Row(2, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) assert(result.collect().isEmpty) } @@ -1166,12 +1186,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(2, "v2", 7L, null, Row(7L), null) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val target = targetTableOf(userSchema)() - val result = processor.findAffectedRowsFromAuxiliaryTable( - rawAuxiliaryTableDf = aux, - perKeyMinimumSequenceInMicrobatchDf = minSeq, - batchId = 100L - ) + val result = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq) checkAnswer( df = result, @@ -1186,11 +1203,13 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val keySchema = new StructType().add("id", IntegerType) val userSchema = keySchema.add("value", StringType) - // Single key with four target rows: - // - row closed at endAt=5 -> < minSeq=10 -> excluded - // - row closed at endAt=10 -> = minSeq=10 -> included (>=) - // - row closed at endAt=15 -> > minSeq=10 -> included - // - row active (endAt=null) -> always included + // Single key with four target rows and an empty aux table, so the cutoff is the target + // row nearest below minSeq=10: recordStartAt=5. Selection is on recordStartAt, not on + // endAt, so a row's interval width is irrelevant - only where it starts matters. + // - recordStartAt=1 -> below the cutoff -> excluded + // - recordStartAt=5 -> sits at the cutoff -> included + // - recordStartAt=10 -> after the cutoff -> included + // - recordStartAt=15 -> after the cutoff -> included (and is the active row) val target = targetTableOf(userSchema)( Row(1, "old", 1L, 5L, Row(1L)), Row(1, "edge", 5L, 10L, Row(5L)), @@ -1198,11 +1217,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(1, "active", 15L, null, Row(15L)) ) val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val aux = auxTableOf(userSchema)() - val result = processor.findAffectedRowsFromTargetTable( - targetTableDf = target, - perKeyMinimumSequenceInMicrobatchDf = minSeq - ) + val result = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq) checkAnswer( df = result, @@ -1214,19 +1231,66 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { ) } + test("affected sequence cutoff derives from the target table") { + val processor = processorWithKeys(Seq("id")) + val keySchema = new StructType().add("id", IntegerType) + val userSchema = keySchema.add("value", StringType) + + // A standalone delete at 40 found nothing live to close, so it survives in the auxiliary + // table as a tombstone. The upsert at 42 then opened a run in the gap after it. + val aux = auxTableOf(userSchema)(Row(1, null, 40L, 40L, Row(40L), null)) + val target = targetTableOf(userSchema)(Row(1, "target", 42L, null, Row(42L))) + val minSeq = minSeqOf(keySchema)(Row(1, 50L)) + + // The target's row at 42 is the cutoff, so the auxiliary tombstone at 40 falls below it. + checkAnswer( + df = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq), + expectedAnswer = Seq(Row(1, "target", 42L, null, Row(42L))) + ) + checkAnswer( + df = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq), + expectedAnswer = Seq.empty[Row] + ) + } + + test("affected sequence cutoff derives from the auxiliary table") { + val processor = processorWithKeys(Seq("id")) + val keySchema = new StructType().add("id", IntegerType) + val userSchema = keySchema.add("value", StringType) + + // A delete at 41 closed the target's run, leaving no tombstone of its own since the closed + // row already carries that boundary. A later standalone delete at 42 landed in the gap after + // it with nothing to close, and so survives in the auxiliary table. + val aux = auxTableOf(userSchema)(Row(1, null, 42L, 42L, Row(42L), null)) + val target = targetTableOf(userSchema)(Row(1, "target", 40L, 41L, Row(40L))) + val minSeq = minSeqOf(keySchema)(Row(1, 50L)) + + // The tombstone at 42 is the cutoff, so the target's row at 40 falls below it - correctly, + // since that interval already closed at 41, before anything in the microbatch. + checkAnswer( + df = findAffectedAuxRows(processor, aux = aux, target = target, minSeq = minSeq), + expectedAnswer = Seq(Row(1, null, 42L, 42L, Row(42L))) + ) + checkAnswer( + df = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq), + expectedAnswer = Seq.empty[Row] + ) + } + test("findAffectedRowsFromTargetTable computes inclusion independently per key") { val processor = processorWithKeys(Seq("id")) val keySchema = new StructType().add("id", IntegerType) val userSchema = keySchema.add("value", StringType) - // Two keys with overlapping endAt ranges but different per-key minSeqs. Each key is - // reconciled independently against its own minSeq. + // Two keys with overlapping intervals but different per-key minSeqs. Each key gets its + // own cutoff, computed independently. val target = targetTableOf(userSchema)( - // Key 1: minSeq=10. "active" (null) and "recent" (15) are at/after 10. + // Key 1: minSeq=10, so the cutoff is recordStartAt=5. "k1.recent" sits at it and + // "k1.active" follows it; "k1.old" at recordStartAt=1 falls below it. Row(1, "k1.old", 1L, 5L, Row(1L)), Row(1, "k1.recent", 5L, 15L, Row(5L)), Row(1, "k1.active", 15L, null, Row(15L)), - // Key 2: minSeq=20. Only "active" (null) is at/after 20. + // Key 2: minSeq=20, so the cutoff is recordStartAt=18. Only "k2.active" survives. Row(2, "k2.old", 1L, 10L, Row(1L)), Row(2, "k2.recent", 10L, 18L, Row(10L)), Row(2, "k2.active", 18L, null, Row(18L)) @@ -1235,11 +1299,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row(1, 10L), Row(2, 20L) ) + val aux = auxTableOf(userSchema)() - val result = processor.findAffectedRowsFromTargetTable( - targetTableDf = target, - perKeyMinimumSequenceInMicrobatchDf = minSeq - ) + val result = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq) checkAnswer( df = result, @@ -1259,9 +1321,10 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val processor = processorWithKeys(Seq("region", "customer_id")) - // (US, 1) and (EU, 1) are distinct composite keys. (US, 1)'s active row is included - // for minSeq=10; (EU, 1)'s active row is included for minSeq=12; (EU, 1)'s old closed - // row at endAt=5 is excluded (5 < 12). (US, 2) has no target rows. + // (US, 1) and (EU, 1) are distinct composite keys, each getting its own cutoff. + // (US, 1)'s only row is at recordStartAt=1, which becomes its cutoff for minSeq=10. + // (EU, 1)'s cutoff for minSeq=12 is recordStartAt=5, so its older row at + // recordStartAt=1 falls below it. (US, 2) has no target rows. val target = targetTableOf(userSchema)( Row("US", 1, "us1", 1L, null, Row(1L)), Row("EU", 1, "eu1.old", 1L, 5L, Row(1L)), @@ -1272,11 +1335,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { Row("EU", 1, 12L), Row("US", 2, 100L) ) + val aux = auxTableOf(userSchema)() - val result = processor.findAffectedRowsFromTargetTable( - targetTableDf = target, - perKeyMinimumSequenceInMicrobatchDf = minSeq - ) + val result = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq) checkAnswer( df = result, @@ -1294,11 +1355,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val target = targetTableOf(userSchema)() val minSeq = minSeqOf(keySchema)(Row(1, 10L)) + val aux = auxTableOf(userSchema)() - val result = processor.findAffectedRowsFromTargetTable( - targetTableDf = target, - perKeyMinimumSequenceInMicrobatchDf = minSeq - ) + val result = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq) assert(result.collect().isEmpty) } @@ -1312,11 +1371,9 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { // Target only has rows for key=1. Microbatch only sees key=2. val target = targetTableOf(userSchema)(Row(1, "v", 1L, null, Row(1L))) val minSeq = minSeqOf(keySchema)(Row(2, 10L)) + val aux = auxTableOf(userSchema)() - val result = processor.findAffectedRowsFromTargetTable( - targetTableDf = target, - perKeyMinimumSequenceInMicrobatchDf = minSeq - ) + val result = findAffectedTargetRows(processor, target = target, aux = aux, minSeq = minSeq) assert(result.collect().isEmpty) } @@ -1946,7 +2003,8 @@ class Scd2BatchProcessorSuite extends QueryTest with SharedSparkSession { val processor = processorWithKeys(Seq("id")) val userSchema = new StructType().add("id", IntegerType).add("value", StringType) - // The first row is an aux anchor (startAt < recordStartAt), pulled in as left context + // The first row sits at the affected sequence cutoff (startAt < recordStartAt), pulled in + // as left context // for a run that began at startAt=2. Because the row sits at the front of the window, // its existing startAt encodes the true global run start and must be preserved - // and propagated to the in-window continuation. diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandlerSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandlerSuite.scala index 287447cc06642..1384bcd078aaf 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandlerSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/autocdc/Scd2ForeachBatchHandlerSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.pipelines.autocdc import org.scalatest.BeforeAndAfter import org.apache.spark.sql.{functions => F, AnalysisException, QueryTest, Row} -import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.classic.{DataFrame, Dataset} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ @@ -715,6 +715,132 @@ class Scd2ForeachBatchHandlerSuite checkAnswer(auxTable, auxRow(1, null, 20L, 20L, 20L, 2L)) } + test("the auxiliary merge's own writes do not duplicate the rows the target merge reads") { + // [[Scd2ForeachBatchHandler.execute]] builds ONE reconciliation source and hands it to two + // MERGEs in turn. The auxiliary merge commits first, and the source reads the auxiliary table, + // so the target merge re-evaluates that source over a table this very batch just rewrote. A + // batch that DEMOTES a visible target row into a hidden auxiliary no-op stresses that: the + // same (id, recordStartAt) - the pair both MERGEs match on - ends up in both tables at once. + // Reconciliation has to collapse the two copies, keeping the one that knows the run's closure. + // + // Seed: the "a" run covers [5, 20), its head hidden at 5 and its visible tail at 10, closed by + // "b" at 20. + createAuxTable(auxRow(1, "a", 5L, null, 5L, null)) + createTargetTable( + targetRow(1, "a", 5L, 20L, 10L), + targetRow(1, "b", 20L, null, 20L) + ) + + // A no-op "a" at 15 lands inside the run and becomes its visible tail, demoting the tail at 10 + // into the auxiliary table - where the target table still holds its own copy of it. + val batchId = 11L + val batchDf = microbatchOf(sourceSchema)(upsert(1, "a", 15L)) + val reconciled = exec.reconcileMicrobatch(batchDf, batchId) + val sourceTheAuxMergeReads = reconciled.reconciledAndRoutedDf.collect().toSeq + processor.mergeRowsIntoAuxiliaryTable( + reconciledDfWithAuxRowsTagged = reconciled.reconciledAndRoutedDf, + originalAffectedRowsFromAuxiliaryTable = reconciled.affectedRowsFromAuxiliaryTable, + auxiliaryTableIdentifier = defaultAuxTableIdentifier, + batchId = batchId + ) + + // The demotion landed: the auxiliary table now holds a copy of the event at 10 alongside the + // run head at 5, and the target table has not been touched yet, so it holds one too. + checkAnswer( + auxTable, + Seq( + auxRow(1, "a", 5L, null, 5L, null), + auxRow(1, "a", 5L, null, 10L, null) + ) + ) + + // Evaluate the source the way the target merge does: a fresh execution over the same plan, so + // its scans are rebuilt against the auxiliary table as the merge above left it. It must yield + // exactly what the auxiliary merge consumed - the duplicated event at 10 collapsing back to + // the one copy that was there before, the copy that knows the run's closure at 20. + checkAnswer( + Dataset.ofRows(spark, reconciled.reconciledAndRoutedDf.logicalPlan), + sourceTheAuxMergeReads + ) + + processor.mergeRowsIntoTargetTable( + reconciledDfWithAuxRowsTagged = reconciled.reconciledAndRoutedDf, + affectedRowsFromTargetTable = reconciled.affectedRowsFromTargetTable, + targetTableIdentifier = defaultTargetTableIdentifier + ) + + // The merge lands exactly those rows: the run still covers [5, 20) with the event at 15 as its + // visible tail, and "b" is untouched. + checkAnswer( + targetTable, + Seq( + targetRow(1, "a", 5L, 20L, 15L), + targetRow(1, "b", 20L, null, 20L) + ) + ) + } + + test("a demoted row's two copies reconcile even when they disagree on where the run starts") { + // Same cross-merge duplication as the test above, but here the batch also moves the run's + // start, so the two copies of the demoted event disagree on startAt rather than only on + // endAt - the case the startAt sort key exists to decide. + // + // Seed: the "a" run covers [10, 20) as a single visible event, closed by "b" at 20. Nothing + // is hidden yet. + createAuxTable() + createTargetTable( + targetRow(1, "a", 10L, 20L, 10L), + targetRow(1, "b", 20L, null, 20L) + ) + + // Two no-op "a" events. The one at 5 predates the run and pulls its start back to 5; the one + // at 15 becomes the run's visible tail, demoting the event at 10 into the auxiliary table. + val batchId = 11L + val batchDf = microbatchOf(sourceSchema)(upsert(1, "a", 5L), upsert(1, "a", 15L)) + val reconciled = exec.reconcileMicrobatch(batchDf, batchId) + val sourceTheAuxMergeReads = reconciled.reconciledAndRoutedDf.collect().toSeq + processor.mergeRowsIntoAuxiliaryTable( + reconciledDfWithAuxRowsTagged = reconciled.reconciledAndRoutedDf, + originalAffectedRowsFromAuxiliaryTable = reconciled.affectedRowsFromAuxiliaryTable, + auxiliaryTableIdentifier = defaultAuxTableIdentifier, + batchId = batchId + ) + + // Both hidden members of the run now record the run's new start at 5, while the target table + // still holds its own copy of the event at 10 recording the old start at 10. + checkAnswer( + auxTable, + Seq( + auxRow(1, "a", 5L, null, 5L, null), + auxRow(1, "a", 5L, null, 10L, null) + ) + ) + checkAnswer(targetTable.filter("value = 'a'"), targetRow(1, "a", 10L, 20L, 10L)) + + // The target merge re-evaluates the same plan against the auxiliary table as the merge above + // left it, so the two disagreeing copies of the event at 10 meet in the window. They must + // collapse back to what the auxiliary merge already consumed. + checkAnswer( + Dataset.ofRows(spark, reconciled.reconciledAndRoutedDf.logicalPlan), + sourceTheAuxMergeReads + ) + + processor.mergeRowsIntoTargetTable( + reconciledDfWithAuxRowsTagged = reconciled.reconciledAndRoutedDf, + affectedRowsFromTargetTable = reconciled.affectedRowsFromTargetTable, + targetTableIdentifier = defaultTargetTableIdentifier + ) + + // The run now covers [5, 20) with the event at 15 as its visible tail, and "b" is untouched. + checkAnswer( + targetTable, + Seq( + targetRow(1, "a", 5L, 20L, 15L), + targetRow(1, "b", 20L, null, 20L) + ) + ) + } + test("duplicate events at the same key and sequence in one microbatch collapse to one record") { createAuxTable() createTargetTable() @@ -908,6 +1034,75 @@ class Scd2ForeachBatchHandlerSuite resolvedSequencingType = LongType ) + test("a closed run of untracked-only changes is one record holding the last event's values") { + // SPARK-58937 Regression test; A row hidden inside a no-op run should not be prematurely + // considered affected by the microbatch and promoted into the target table. + createTable(defaultAuxIdent, defaultAuxTableIdentifier, trackedAuxSchema) + createTable(defaultTargetIdent, defaultTargetTableIdentifier, trackedCanonicalSchema) + val handler = execWith(trackedProcessor) + + Seq( + // Start a run whose tracked `name` is "alice". + Seq(Row(1, "alice", 39, 39L, false)), + // Continue that run with an untracked-only score change. + Seq(Row(1, "alice", 41, 41L, false)), + // Close the run, then process a later event that needs left context from the same key. + Seq(Row(1, "bob", 42, 42L, false)), + Seq(Row(1, "carol", 43, 43L, false)) + ).zipWithIndex.foreach { case (batch, batchId) => + handler.execute(microbatchOf(trackedSourceSchema)(batch: _*), batchId) + } + + checkAnswer( + targetTable, + Seq( + Row(1, "alice", 41, 39L, 42L, meta(41L)), + Row(1, "bob", 42, 42L, 43L, meta(42L)), + Row(1, "carol", 43, 43L, null, meta(43L)) + ) + ) + + // The run head stays hidden in the auxiliary table, holding the values it had before the + // untracked-only change continued the run. + checkAnswer(auxTable, Row(1, "alice", 39, 39L, null, meta(39L), null)) + } + + test("an event after a deletion leaves the run the deletion closed exactly as it was") { + // SPARK-58937 Regression test; A row hidden behind a run that a deletion closed should not be + // considered affected by a microbatch that only touches instants after that deletion. + createAuxTable() + createTargetTable() + + // Batch 1: a two-event no-op run for "a". The head at 10 is hidden in the aux table and the + // tail at 11 is the single visible row covering the run. + runBatch(1L)(upsert(1, "a", 10L), upsert(1, "a", 11L)) + checkAnswer(targetTable, targetRow(1, "a", 10L, null, 11L)) + checkAnswer(auxTable, auxRow(1, "a", 10L, null, 10L, null)) + + // Batch 2: a delete closes that run at 12. The key is now absent from 12 onwards, so no + // target row covers any instant at or after 12 - the history has a gap there. + runBatch(2L)(del(1, 12L)) + checkAnswer(targetTable, targetRow(1, "a", 10L, 12L, 11L)) + checkAnswer(auxTable, auxRow(1, "a", 10L, null, 10L, null)) + + // Batch 3: a new value at 20, strictly after the gap. It cannot interact with the run that + // already closed at 12, so that run's single visible row must be left exactly as it is - in + // particular the head hidden at 10 must not be pulled back into the target table beside it. + runBatch(3L)(upsert(1, "b", 20L)) + + checkAnswer( + targetTable, + Seq( + targetRow(1, "a", 10L, 12L, 11L), + targetRow(1, "b", 20L, null, 20L) + ) + ) + + // The head hidden at 10 is also left untouched - the new event neither closed it nor + // duplicated it. + checkAnswer(auxTable, auxRow(1, "a", 10L, null, 10L, null)) + } + test("changing only an untracked column updates the current record without adding history") { createTable(defaultAuxIdent, defaultAuxTableIdentifier, trackedAuxSchema) createTable(defaultTargetIdent, defaultTargetTableIdentifier, trackedCanonicalSchema) diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTableSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTableSuite.scala index fb368198f72d3..27ae39d9bc99c 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTableSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcAuxiliaryTableSuite.scala @@ -22,8 +22,15 @@ import scala.jdk.CollectionConverters._ import org.apache.spark.SparkFunSuite import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.TableIdentifier +import org.apache.spark.sql.catalyst.analysis.{caseInsensitiveResolution, caseSensitiveResolution} import org.apache.spark.sql.connector.catalog.{Table, TableCapability} -import org.apache.spark.sql.pipelines.autocdc.ScdType +import org.apache.spark.sql.pipelines.autocdc.{ + AutoCdcReservedNames, + Scd1BatchProcessor, + Scd2BatchProcessor, + ScdType +} +import org.apache.spark.sql.types.{IntegerType, LongType, StringType, StructType} /** * Unit tests for the [[AutoCdcAuxiliaryTable]] companion object. @@ -44,9 +51,9 @@ class AutoCdcAuxiliaryTableSuite extends SparkFunSuite { // concern -- SQL identifier quoting (backticks) is never part of the stored bytes. private def assertKeyColumnNamesRoundTrip(names: Seq[String]): Unit = { - val json = AutoCdcAuxiliaryTable.serializeKeyColumnNames(names) + val json = AutoCdcAuxiliaryTable.serializeColumnNames(names) assert( - AutoCdcAuxiliaryTable.parseKeyColumnNames(json).contains(names), + AutoCdcAuxiliaryTable.parseColumnNames(json).contains(names), s"round-trip failed: input=${names}, serialized=${json}" ) } @@ -59,19 +66,19 @@ class AutoCdcAuxiliaryTableSuite extends SparkFunSuite { override def properties(): java.util.Map[String, String] = props.asJava } - test("serializeKeyColumnNames/parseKeyColumnNames round-trip preserves plain ASCII names") { + test("serializeColumnNames/parseColumnNames round-trip preserves plain ASCII names") { assertKeyColumnNamesRoundTrip(Seq("id")) assertKeyColumnNamesRoundTrip(Seq("id", "region")) assertKeyColumnNamesRoundTrip(Seq("id", "region", "country")) } - test("serializeKeyColumnNames/parseKeyColumnNames round-trip preserves the empty list") { + test("serializeColumnNames/parseColumnNames round-trip preserves the empty list") { // Empty key sets are not user-reachable (AutoCdcMergeFlow rejects them upstream), but the // helpers themselves must round-trip a `[]` JSON array faithfully. assertKeyColumnNamesRoundTrip(Seq.empty) } - test("serializeKeyColumnNames/parseKeyColumnNames preserves names containing JSON-escaped " + + test("serializeColumnNames/parseColumnNames preserves names containing JSON-escaped " + "characters (quote, backslash, control chars)") { // JSON serializer must escape `"` -> `\"`, `\` -> `\\`, and control chars; the parser // must invert those escapes and yield the original literal bytes. @@ -83,7 +90,7 @@ class AutoCdcAuxiliaryTableSuite extends SparkFunSuite { assertKeyColumnNamesRoundTrip(Seq("a\"b\\c\nd")) } - test("serializeKeyColumnNames/parseKeyColumnNames preserves names containing characters " + + test("serializeColumnNames/parseColumnNames preserves names containing characters " + "that JSON does not escape (single quote, dot, space, backtick)") { // JSON does not escape these, but they are common in real-world identifiers (especially // when users backtick-quote at the API boundary). They must flow through verbatim. @@ -96,18 +103,18 @@ class AutoCdcAuxiliaryTableSuite extends SparkFunSuite { assertKeyColumnNamesRoundTrip(Seq("it's", "name with spaces", "a.b.c", "back`tick")) } - test("parseKeyColumnNames returns None for inputs that are not a JSON array of strings") { + test("parseColumnNames returns None for inputs that are not a JSON array of strings") { // None of these are a top-level JSON array of strings; the parser must reject every shape // with `None` so callers can surface a structured INTERNAL_ERROR with consistent wording. - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("not-json").isEmpty) - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("").isEmpty) - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("\"id\"").isEmpty) // bare string - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("null").isEmpty) - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("{\"id\": 1}").isEmpty) // object - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("[1, 2, 3]").isEmpty) // numbers - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("[\"id\", 1]").isEmpty) // mixed types - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("[\"id\", null]").isEmpty) - assert(AutoCdcAuxiliaryTable.parseKeyColumnNames("[[\"id\"]]").isEmpty) // nested array + assert(AutoCdcAuxiliaryTable.parseColumnNames("not-json").isEmpty) + assert(AutoCdcAuxiliaryTable.parseColumnNames("").isEmpty) + assert(AutoCdcAuxiliaryTable.parseColumnNames("\"id\"").isEmpty) // bare string + assert(AutoCdcAuxiliaryTable.parseColumnNames("null").isEmpty) + assert(AutoCdcAuxiliaryTable.parseColumnNames("{\"id\": 1}").isEmpty) // object + assert(AutoCdcAuxiliaryTable.parseColumnNames("[1, 2, 3]").isEmpty) // numbers + assert(AutoCdcAuxiliaryTable.parseColumnNames("[\"id\", 1]").isEmpty) // mixed types + assert(AutoCdcAuxiliaryTable.parseColumnNames("[\"id\", null]").isEmpty) + assert(AutoCdcAuxiliaryTable.parseColumnNames("[[\"id\"]]").isEmpty) // nested array } test("validateNoScdTypeDrift accepts an auxiliary table whose recorded SCD type matches") { @@ -156,4 +163,263 @@ class AutoCdcAuxiliaryTableSuite extends SparkFunSuite { "tableName" -> TableIdentifier("target", Some("ns"), Some("cat")).unquotedString, "propertyName" -> AutoCdcAuxiliaryTable.scdTypePropertyKey)) } + + private val targetIdent = TableIdentifier("target", Some("ns"), Some("cat")) + + /** An auxiliary table stub recording the given track-history column names as JSON. */ + private def auxTableWithTrackHistory(names: Seq[String]): Table = + auxTableWithProperties(Map( + AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty -> + AutoCdcAuxiliaryTable.serializeColumnNames(names))) + + test("validateNoTrackHistoryDrift is a no-op when the expected column set is None") { + // A None expected set means the flow does not constrain track-history (SCD1, or an SCD2 flow + // whose default resolution has not been computed here); the validator must not even read the + // property. Passing an empty-properties table proves nothing is dereferenced. + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = auxTableWithProperties(Map.empty), + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = None, + resolver = caseInsensitiveResolution) + } + + test("validateNoTrackHistoryDrift accepts a recorded set that matches regardless of order") { + val existing = auxTableWithTrackHistory(Seq("name", "amount", "seq")) + // Same set, different order: must not throw. + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existing, + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("seq", "name", "amount")), + resolver = caseInsensitiveResolution) + } + + test("validateNoTrackHistoryDrift compares case-insensitively under the default resolver, " + + "even when the stored property names differ only in case") { + // Isolates the resolver-aware comparison: the stored property holds `Name`/`AMOUNT` while the + // expected set holds `name`/`amount`. In the end-to-end path both sides are normalized to + // actual schema field names before comparison, so only a direct unit test can exercise a + // genuine case difference reaching the resolver. Under the default resolver, no drift. + val existing = auxTableWithTrackHistory(Seq("Name", "AMOUNT")) + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existing, + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("name", "amount")), + resolver = caseInsensitiveResolution) + } + + test("validateNoTrackHistoryDrift throws TRACK_HISTORY_DRIFT under the case-sensitive resolver " + + "when the stored property names differ only in case") { + // The mirror of the case-insensitive test: with the case-sensitive resolver, `Name` and + // `name` are distinct, so the same-cardinality sets do not match and the validator drifts. + val existing = auxTableWithTrackHistory(Seq("Name", "amount")) + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existing, + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("name", "amount")), + resolver = caseSensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "expectedTrackHistoryColumns" -> "name, amount", + "recordedTrackHistoryColumns" -> "Name, amount")) + } + + test("validateNoTrackHistoryDrift throws TRACK_HISTORY_DRIFT when the recorded set differs") { + val existing = auxTableWithTrackHistory(Seq("name")) + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existing, + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("amount")), + resolver = caseInsensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "expectedTrackHistoryColumns" -> "amount", + "recordedTrackHistoryColumns" -> "name")) + } + + test("validateNoTrackHistoryDrift throws AUXILIARY_TABLE_PROPERTY_MISSING when the " + + "track-history property is absent") { + // An SCD2 aux table created before this property existed: the validator must surface a + // structured error (remedy: full refresh) rather than skipping the check. + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = auxTableWithProperties(Map.empty), + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("name")), + resolver = caseInsensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MISSING", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "propertyName" -> AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty)) + } + + test("validateNoTrackHistoryDrift throws AUXILIARY_TABLE_PROPERTY_MALFORMED when the " + + "track-history property is not a JSON array of strings") { + val existing = auxTableWithProperties(Map( + AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty -> "not-a-json-array")) + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTrackHistoryDrift( + existingAuxiliaryTable = existing, + targetTableIdentifier = targetIdent, + expectedTrackHistoryColumnNames = Some(Seq("name")), + resolver = caseInsensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MALFORMED", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "propertyName" -> AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty, + "rawValue" -> "not-a-json-array")) + } + + // =========================================================================================== + // validateNoTargetSequencingTypeDrift + // =========================================================================================== + + private val meta = AutoCdcReservedNames.cdcMetadataColName + + /** An SCD1 target schema whose `_cdc_metadata` carries sequence fields of `seqType`. */ + private def scd1TargetSchema(seqType: org.apache.spark.sql.types.DataType): StructType = + new StructType() + .add("id", IntegerType, nullable = false) + .add(meta, new StructType() + .add(Scd1BatchProcessor.cdcDeleteSequenceFieldName, seqType) + .add(Scd1BatchProcessor.cdcUpsertSequenceFieldName, seqType)) + + /** An SCD2 target schema whose `_cdc_metadata` carries a recordStartAt field of `seqType`. */ + private def scd2TargetSchema(seqType: org.apache.spark.sql.types.DataType): StructType = + new StructType() + .add("id", IntegerType, nullable = false) + .add(meta, new StructType() + .add(Scd2BatchProcessor.recordStartAtFieldName, seqType)) + + test("validateNoTargetSequencingTypeDrift accepts a matching SCD1 sequencing type") { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = scd1TargetSchema(LongType), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type1, + expectedSequencingType = LongType, + resolver = caseInsensitiveResolution) + } + + test("validateNoTargetSequencingTypeDrift accepts a matching SCD2 sequencing type") { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = scd2TargetSchema(LongType), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = LongType, + resolver = caseInsensitiveResolution) + } + + test("validateNoTargetSequencingTypeDrift throws SEQUENCING_TYPE_DRIFT when the recorded " + + "type differs (SCD1)") { + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = scd1TargetSchema(LongType), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type1, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "expectedSequencingType" -> IntegerType.sql, + "recordedSequencingType" -> LongType.sql)) + } + + test("validateNoTargetSequencingTypeDrift throws SEQUENCING_TYPE_DRIFT when the recorded " + + "type differs (SCD2)") { + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = scd2TargetSchema(LongType), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + checkError( + exception = ex, + condition = "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT", + sqlState = "42000", + parameters = Map( + "tableName" -> targetIdent.unquotedString, + "expectedSequencingType" -> IntegerType.sql, + "recordedSequencingType" -> LongType.sql)) + } + + test("validateNoTargetSequencingTypeDrift is a silent no-op when _cdc_metadata is absent") { + // Not a recognizable AutoCDC target state: skip rather than misreport. A genuine + // incompatibility would surface later during schema evolution. + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = new StructType().add("id", IntegerType, nullable = false), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + + test("validateNoTargetSequencingTypeDrift is a silent no-op when _cdc_metadata is not a " + + "struct") { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = new StructType() + .add("id", IntegerType, nullable = false) + .add(meta, StringType), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + + test("validateNoTargetSequencingTypeDrift is a silent no-op when _cdc_metadata is an empty " + + "struct") { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = new StructType() + .add("id", IntegerType, nullable = false) + .add(meta, new StructType()), + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + + test("validateNoTargetSequencingTypeDrift resolves _cdc_metadata and the inner field via the " + + "resolver (case-insensitive)") { + // A hand-written target DDL may differ in case; under the default resolver the check still + // finds the metadata column and the recordStartAt field, so a real type drift is caught. + val upperCasedSchema = new StructType() + .add("id", IntegerType, nullable = false) + .add(meta.toUpperCase(java.util.Locale.ROOT), new StructType() + .add(Scd2BatchProcessor.recordStartAtFieldName.toLowerCase(java.util.Locale.ROOT), + LongType)) + val ex = intercept[AnalysisException] { + AutoCdcAuxiliaryTable.validateNoTargetSequencingTypeDrift( + existingTargetSchema = upperCasedSchema, + targetTableIdentifier = targetIdent, + expectedScdType = ScdType.Type2, + expectedSequencingType = IntegerType, + resolver = caseInsensitiveResolution) + } + assert(ex.getCondition == "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT") + } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcConfigDriftSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcConfigDriftSuite.scala new file mode 100644 index 0000000000000..c26a5248488b9 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcConfigDriftSuite.scala @@ -0,0 +1,616 @@ +/* + * 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.pipelines.graph + +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * End-to-end tests covering AutoCDC configuration-drift validation for the sequencing result type + * (SCD1 and SCD2) and the SCD2 track-history column set, validated at flow execution-init time + * against the auxiliary table's recorded configuration (mirroring + * [[AutoCdcScd1KeyDriftSuite]] for keys). + * + * Guiding principle: guard the invariants that keep already-persisted state coherent, not the + * expressions themselves. The sequencing expression and delete condition may change across runs; + * the sequencing result *type* and the SCD2 track-history column *set* may not. + */ +class AutoCdcConfigDriftSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + private def targetName: String = + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString + + /** SCD2 target DDL: user columns + the SCD2 framework columns (sequencing type long). */ + private def createScd2Target(userCols: String): Unit = { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target ($userCols, $scd2MetadataDdl)" + ) + } + + // =========================================================================================== + // Sequencing type drift + // =========================================================================================== + + test("AutoCDC source validation uses pipeline case sensitivity, not session default") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val stream = MemoryStream[(Int, Long, Long)] + stream.addData((1, 1L, 1L)) + + val ctx = new TestGraphRegistrationContext( + spark, + Map(SQLConf.CASE_SENSITIVE.key -> "true")) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "version", "__start_at")), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + } + + ctx.resolveToDataflowGraph() + } + } + + test("an SCD1 flow whose sequencing type differs from the recorded type triggers " + + "SEQUENCING_TYPE_DRIFT") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, seq_long BIGINT, seq_int INT, $scd1MetadataDdl)" + ) + + // Pipeline #1 sequences by a BIGINT column; aux records sequencingType = long. + val stream1 = MemoryStream[(Int, Long, Int)] + stream1.addData((1, 1L, 1)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "seq_long", "seq_int"), + keys = Seq("id"), + sequencing = $"seq_long")) + + // Pipeline #2 sequences by an INT column - type drift (int vs long), even though the + // expression (a different column) is otherwise a legal change. + val stream2 = MemoryStream[(Int, Long, Int)] + stream2.addData((1, 2L, 2)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "seq_long", "seq_int"), + keys = Seq("id"), + sequencing = $"seq_int") + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedSequencingType" -> "INT", + "recordedSequencingType" -> "BIGINT" + ) + ) + } + + test("an SCD1 flow that changes the sequencing expression but keeps the same type does NOT " + + "trigger drift") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, seq BIGINT, $scd1MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 10L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "seq"), + keys = Seq("id"), + sequencing = $"seq")) + + // A different expression over the same column, still yielding BIGINT: legal, no drift. + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((1, 20L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "seq"), + keys = Seq("id"), + sequencing = $"seq" + 1L)) + } + + test("an SCD2 flow whose sequencing type differs from the recorded type triggers " + + "SEQUENCING_TYPE_DRIFT") { + createScd2Target("id INT NOT NULL, seq_long BIGINT, seq_int INT") + + val stream1 = MemoryStream[(Int, Long, Int)] + stream1.addData((1, 1L, 1)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "seq_long", "seq_int"), + keys = Seq("id"), + sequencing = $"seq_long", + scdType = ScdType.Type2)) + + val stream2 = MemoryStream[(Int, Long, Int)] + stream2.addData((1, 2L, 2)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "seq_long", "seq_int"), + keys = Seq("id"), + sequencing = $"seq_int", + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.SEQUENCING_TYPE_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedSequencingType" -> "INT", + "recordedSequencingType" -> "BIGINT" + ) + ) + } + + // =========================================================================================== + // Track-history drift (SCD2 only) + // =========================================================================================== + + test("an SCD2 flow that changes its explicit TRACK HISTORY column set triggers " + + "TRACK_HISTORY_DRIFT") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Pipeline #1 tracks history on `name` only. + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name")))))) + + // Pipeline #2 tracks history on `amount` - a different set. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("amount"))))) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedTrackHistoryColumns" -> "amount", + "recordedTrackHistoryColumns" -> "name" + ) + ) + } + + test("an SCD2 flow that reorders the same TRACK HISTORY columns does NOT trigger drift") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some(ColumnSelection.IncludeColumns( + Seq(UnqualifiedColumnName("name"), UnqualifiedColumnName("amount")))))) + + // Same set, reversed order: run semantics are order-insensitive, so no drift. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some(ColumnSelection.IncludeColumns( + Seq(UnqualifiedColumnName("amount"), UnqualifiedColumnName("name")))))) + } + + test("an SCD2 flow with no TRACK HISTORY (default = all eligible columns) followed by an " + + "explicit selection of that same set does NOT trigger drift") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Pipeline #1 omits trackHistorySelection: the recorded set is the default, every eligible + // (non-key, non-framework) column, i.e. name, amount, seq. + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2)) + + // Pipeline #2 explicitly lists that same default set: the drift check compares resolved sets, + // not user syntax, so an explicit restatement of the default must NOT drift. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some(ColumnSelection.IncludeColumns(Seq( + UnqualifiedColumnName("name"), + UnqualifiedColumnName("amount"), + UnqualifiedColumnName("seq")))))) + } + + test("an SCD2 flow with no TRACK HISTORY (default = all eligible columns) followed by an " + + "explicit subset triggers TRACK_HISTORY_DRIFT") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Pipeline #1 records the default set (name, amount, seq). + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2)) + + // Pipeline #2 narrows to an explicit subset (name only): a real change to which transitions + // open a new record, so it must drift against the recorded default set. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name"))))) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedTrackHistoryColumns" -> "name", + "recordedTrackHistoryColumns" -> "name, amount, seq" + ) + ) + } + + test("an SCD2 flow's EXCEPT-based TRACK HISTORY followed by an equivalent explicit include " + + "set does NOT trigger drift") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Pipeline #1 uses TRACK HISTORY ON * EXCEPT (amount): the resolved set is the eligible + // columns minus `amount`, i.e. name, seq. + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.ExcludeColumns(Seq(UnqualifiedColumnName("amount")))))) + + // Pipeline #2 states the same set as an explicit include list: EXCLUDE and INCLUDE that + // resolve to the same set must not drift, since only the resolved set is recorded. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some(ColumnSelection.IncludeColumns( + Seq(UnqualifiedColumnName("name"), UnqualifiedColumnName("seq")))))) + } + + test("SCD2 track-history drift validation is resolver-aware: a case-only difference does NOT " + + "trigger drift under the default (case-insensitive) resolver") { + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Pipeline #1 records track-history on `name`. + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name")))))) + + // Pipeline #2 selects `NAME` (different case). The source DF column is still lowercase `name` + // so it resolves against the schema; only the tracking-column casing differs. Under the + // default case-insensitive resolver the two sets are equal, so there must be no drift. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = Some( + ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("NAME")))))) + } + + test("an existing SCD2 aux table missing the trackHistoryColumnNames property requires a " + + "full refresh (AUXILIARY_TABLE_PROPERTY_MISSING)") { + // Back-compat guard: an SCD2 auxiliary table created before this change carries no + // trackHistoryColumnNames property. Track-history drift validation surfaces this as a + // structured AUXILIARY_TABLE_PROPERTY_MISSING (remedy: full refresh) rather than silently + // skipping the check. Simulate the pre-existing table by unsetting the property after the + // first run, then run again. + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + val stream = MemoryStream[(Int, String, Int, Long)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2) + + stream.addData((1, "a", 10, 1L)) + runPipeline(buildCtx()) + + // Drop the property to mimic an aux table materialized before this change. + spark.sql( + s"ALTER TABLE ${auxTableNameFor("target")} " + + s"UNSET TBLPROPERTIES ('${AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty}')" + ) + + stream.addData((1, "a", 20, 2L)) + val ex = intercept[RuntimeException] { runPipeline(buildCtx()) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MISSING", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "propertyName" -> AutoCdcAuxiliaryTable.trackHistoryColumnNamesProperty + ) + ) + } + + // =========================================================================================== + // Sequencing type: SCD2 expression-change symmetry with the SCD1 case above + // =========================================================================================== + + test("an SCD2 flow that changes the sequencing expression but keeps the same type does NOT " + + "trigger drift") { + createScd2Target("id INT NOT NULL, seq BIGINT") + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 10L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2)) + + // A different expression over the same column, still yielding BIGINT: legal, no drift. + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((1, 20L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "seq"), + keys = Seq("id"), + sequencing = $"seq" + 1L, + scdType = ScdType.Type2)) + } + + // =========================================================================================== + // Intended divergence from SCD1: additive source-schema evolution under default / EXCEPT + // tracking changes the effective tracked set, and so requires a full refresh. + // =========================================================================================== + + test("adding a source column under default (all-column) tracking triggers TRACK_HISTORY_DRIFT") { + // The effective tracked set is derived from the flow's selected source schema, so under default + // tracking every selected non-key column is tracked. Adding a source column therefore changes + // the tracked set, which reinterprets which transitions open a new SCD2 record and cannot be + // applied to already-reconciled history. Unlike SCD1 (where a new nullable column is absorbed + // by schema evolution), SCD2 requires a full refresh. Pins that intended divergence. + createScd2Target("id INT NOT NULL, name STRING, seq BIGINT") + + // Run #1: source (id, name, seq); recorded tracked set = {name, seq}. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "a", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2)) + + // Run #2: source gains a nullable `city`; default tracking now resolves to {name, city, seq}. + val stream2 = MemoryStream[(Int, String, String, Long)] + stream2.addData((1, "a", "nyc", 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "city", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedTrackHistoryColumns" -> "name, city, seq", + "recordedTrackHistoryColumns" -> "name, seq")) + + // The drift check runs before the target's schema is evolved, so the rejected run must leave + // the target untouched: `city` must NOT have been added. (Were it added, the "correct the + // flow" remedy would then wedge the pipeline on a column-count mismatch during reconciliation.) + assert( + !spark.table(s"$catalog.$namespace.target").schema.fieldNames.contains("city"), + "rejected run must not have evolved the target schema to add `city`") + } + + test("dropping a source column under default (all-column) tracking triggers " + + "TRACK_HISTORY_DRIFT") { + // The mirror of the additive case: removing a selected column shrinks the default tracked set, + // which is likewise a tracked-set change requiring a full refresh. + createScd2Target("id INT NOT NULL, name STRING, city STRING, seq BIGINT") + + // Run #1: source (id, name, city, seq); recorded tracked set = {name, city, seq}. + val stream1 = MemoryStream[(Int, String, String, Long)] + stream1.addData((1, "a", "nyc", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "city", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2)) + + // Run #2: `city` dropped from the source; default tracking now resolves to {name, seq}. + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((1, "a", 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedTrackHistoryColumns" -> "name, seq", + "recordedTrackHistoryColumns" -> "name, city, seq")) + } + + test("dropping the target (but not the auxiliary table) between runs still detects drift") { + // A user drops and recreates the target to reset it, but does not know to also drop the + // internal auxiliary table. On the next run the target is absent (so it is re-created), but the + // stale auxiliary table survives with the old recorded configuration. Drift validation reads + // the auxiliary table, so it must still fire regardless of the target's existence -- otherwise + // the aux table's additive evolve would silently overwrite the recorded track-history property + // with the new run's value. Mirrors AutoCdcScd1AuxiliaryTableDurabilitySuite's + // "auxiliary table is dropped between runs" case, with the two tables swapped. + createScd2Target("id INT NOT NULL, name STRING, amount INT, seq BIGINT") + + // Run #1: track history on `name`; records trackHistoryColumnNames = [name] on the aux table. + val stream1 = MemoryStream[(Int, String, Int, Long)] + stream1.addData((1, "a", 10, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "target", + sourceDf = stream1.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = + Some(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name")))))) + + // Drop ONLY the target; the auxiliary table survives with its recorded config. + spark.sql(s"DROP TABLE $catalog.$namespace.target") + assert(spark.catalog.tableExists(auxTableNameFor("target")), + "auxiliary table should survive dropping the target") + + // Run #2: track history on `amount` instead -- a changed tracked set. + // This will recreate the target. + val stream2 = MemoryStream[(Int, String, Int, Long)] + stream2.addData((1, "a", 20, 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "target", + sourceDf = stream2.toDF().toDF("id", "name", "amount", "seq"), + keys = Seq("id"), + sequencing = $"seq", + scdType = ScdType.Type2, + trackHistorySelection = + Some(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("amount"))))) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> targetName, + "expectedTrackHistoryColumns" -> "amount", + "recordedTrackHistoryColumns" -> "name")) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcCrossScdConvergenceSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcCrossScdConvergenceSuite.scala new file mode 100644 index 0000000000000..64170f48c0b7a --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcCrossScdConvergenceSuite.scala @@ -0,0 +1,108 @@ +/* + * 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.pipelines.graph + +import scala.util.Random + +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{Scd2BatchProcessor, ScdType} +import org.apache.spark.sql.pipelines.utils.ExecutionTest +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Differential test for cross-SCD current-state agreement: given the same randomly-generated + * CDF, every live key's SCD Type 1 target row must equal (in user data columns) the current + * open SCD Type 2 row for that key - this is by definitions of an SCD1 and SCD2 transformation. + * + * By asserting the final outcome of SCD1 equals the final live rows of SCD2, each implementation + * is an effective verifier of the other, and catches regressions or behavior changes made to one + * implementation but not the other. + * + * CDC metadata and SCD2 interval bounds are not compared. + */ +class AutoCdcCrossScdConvergenceSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin + with AutoCdcRandomCdcTestMixin { + + /** + * Assert SCD1 live rows equal SCD2 current open rows (`__END_AT IS NULL`) on user data + * columns only. + */ + private def assertCrossScdAgreement( + scd1Table: String, + scd2Table: String, + expectedLiveKeyCount: Int): Unit = { + val scd1Data = spark.table(s"$catalog.$namespace.$scd1Table").select( + dataColumnNames.map(functions.col): _* + ) + val scd2CurrentData = spark.table(s"$catalog.$namespace.$scd2Table") + .where(functions.col(Scd2BatchProcessor.endAtColName).isNull) + .select(dataColumnNames.map(functions.col): _*) + + // Verify the number of live keys (i.e rows that haven't been fully deleted) are the same in + // both SCD1 and SCD2, after all events are applied. + val scd1LiveKeyCount = scd1Data.count() + val scd2LiveKeyCount = scd2CurrentData.count() + assert( + scd1LiveKeyCount == expectedLiveKeyCount, + s"Expected $expectedLiveKeyCount live SCD1 keys, found $scd1LiveKeyCount") + assert( + scd2LiveKeyCount == expectedLiveKeyCount, + s"Expected $expectedLiveKeyCount live SCD2 keys, found $scd2LiveKeyCount") + + checkAnswer(scd1Data, scd2CurrentData) + } + + private val crossScdConvergenceTestName = + "SCD1 current rows match SCD2 open rows for the same shuffled CDC stream" + + test(crossScdConvergenceTestName) { + val numDistinctKeys = resolveNumDistinctKeys() + val maxUniqueEventsPerKey = resolveMaxUniqueEventsPerKey() + val numBatches = resolveNumBatches() + + forEachConvergenceSeed(crossScdConvergenceTestName) { (seed, seedIndex) => + val rand = new Random(seed) + val sortedEventStream = generateRandomCdcEventStream(rand) + val shuffledEventStream = rand.shuffle(sortedEventStream) + val expectedLiveKeyCount = sortedEventStream + .groupBy(_.key) + .values + .count(events => !events.maxBy(_.sequence).isDelete) + + // Avoid dumping thousands of events into every clue string (ScalaTest evaluates clues + // eagerly). + withClue( + s"\ncross-SCD convergence testName=$crossScdConvergenceTestName " + + s"seedIndex=$seedIndex seed=$seed " + + s"(rerun this test with -D$convergenceReproSeedSystemProperty=$seed to reproduce)\n" + + s"keys=$numDistinctKeys maxEventsPerKey=$maxUniqueEventsPerKey " + + s"numBatches=$numBatches expectedLiveKeys=$expectedLiveKeyCount " + + s"events=${sortedEventStream.size}\n" + ) { + val scd1Table = s"cross_scd1_$seedIndex" + val scd2Table = s"cross_scd2_$seedIndex" + runRandomCdcPipeline(scd1Table, ScdType.Type1, shuffledEventStream, numBatches) + runRandomCdcPipeline(scd2Table, ScdType.Type2, shuffledEventStream, numBatches) + assertCrossScdAgreement(scd1Table, scd2Table, expectedLiveKeyCount) + } + } + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcGraphExecutionTestMixin.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcGraphExecutionTestMixin.scala index 305b03e0bdcfb..c789627005b6b 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcGraphExecutionTestMixin.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcGraphExecutionTestMixin.scala @@ -208,7 +208,8 @@ trait AutoCdcGraphExecutionTestMixin extends BeforeAndAfterEach { sequencing: Column, columnSelection: Option[ColumnSelection] = None, deleteCondition: Option[Column] = None, - scdType: ScdType = ScdType.Type1 + scdType: ScdType = ScdType.Type1, + trackHistorySelection: Option[ColumnSelection] = None ): AutoCdcFlow = AutoCdcFlow( identifier = fullyQualifiedIdentifier(name, Some(catalog), Some(namespace)), destinationIdentifier = fullyQualifiedIdentifier(target, Some(catalog), Some(namespace)), @@ -223,7 +224,8 @@ trait AutoCdcGraphExecutionTestMixin extends BeforeAndAfterEach { sequencing = sequencing, columnSelection = columnSelection, deleteCondition = deleteCondition, - storedAsScdType = scdType + storedAsScdType = scdType, + trackHistorySelection = trackHistorySelection ) ) @@ -241,7 +243,8 @@ trait AutoCdcGraphExecutionTestMixin extends BeforeAndAfterEach { sequencing: Column, columnSelection: Option[ColumnSelection] = None, deleteCondition: Option[Column] = None, - scdType: ScdType = ScdType.Type1): TestGraphRegistrationContext = + scdType: ScdType = ScdType.Type1, + trackHistorySelection: Option[ColumnSelection] = None): TestGraphRegistrationContext = new TestGraphRegistrationContext(spark) { registerTable(target, catalog = Some(catalog), database = Some(namespace)) registerFlow(autoCdcFlow( @@ -252,7 +255,8 @@ trait AutoCdcGraphExecutionTestMixin extends BeforeAndAfterEach { sequencing = sequencing, columnSelection = columnSelection, deleteCondition = deleteCondition, - scdType = scdType + scdType = scdType, + trackHistorySelection = trackHistorySelection )) } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcOutOfOrderConvergenceSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcOutOfOrderConvergenceSuite.scala index f164b2e70f8ac..e8e4eb5fdfce0 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcOutOfOrderConvergenceSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcOutOfOrderConvergenceSuite.scala @@ -17,37 +17,12 @@ package org.apache.spark.sql.pipelines.graph -import scala.collection.mutable.ArrayBuffer import scala.util.Random -import org.apache.spark.sql.execution.streaming.runtime.MemoryStream -import org.apache.spark.sql.functions -import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} -import org.apache.spark.sql.pipelines.graph.AutoCdcOutOfOrderConvergenceSuite.SourceRow -import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.pipelines.autocdc.ScdType +import org.apache.spark.sql.pipelines.utils.ExecutionTest import org.apache.spark.sql.test.SharedSparkSession -object AutoCdcOutOfOrderConvergenceSuite { - /** - * A single CDC event in the source stream. - * - * @param key Identity column (the AutoCDC `keys`). - * @param name Data column (nullable string). - * @param amount Data column (nullable int). - * @param active Data column (nullable boolean). - * @param sequence Sequencing value (the AutoCDC `sequencing` expression). - * @param isDelete Drives the AutoCDC `deleteCondition`; `true` marks the event as a delete, - * `false` as an upsert. Excluded from the target via `columnSelection`. - */ - case class SourceRow( - key: Int, - name: Option[String], - amount: Option[Int], - active: Option[Boolean], - sequence: Long, - isDelete: Boolean) -} - /** * Differential test for the AutoCDC merge's order-invariance property, for both SCD Type 1 and * SCD Type 2: feeding the same randomly-generated CDC event stream as a single sorted micro-batch @@ -56,118 +31,8 @@ object AutoCdcOutOfOrderConvergenceSuite { class AutoCdcOutOfOrderConvergenceSuite extends ExecutionTest with SharedSparkSession - with AutoCdcGraphExecutionTestMixin { - - // Distinct keys in the generated event stream. - private val numDistinctKeys: Int = 5 - // Upper bound on unique events (one per sequence) generated per key, before intentionally - // duplicating some events. - private val maxUniqueEventsPerKey: Int = 80 - // Probability an event is a delete; (1 - this) is the upsert probability. - private val deleteEventProbability: Double = 0.20 - // Probability an event is immediately re-emitted with the same sequence and payload. - private val duplicateEventProbability: Double = 0.15 - // Probability an optional payload column is non-null; (1 - this) is the null probability. - private val nonNullProbability: Double = 0.75 - // Number of microbatches the out-of-order pipeline splits the shuffled events across. - private val numOutOfOrderBatches: Int = 8 - - // System property used to pin the test seed for reproduction. If unset, the suite generates a - // fresh seed on each run and reports it in the failure message so a failing seed can be replayed - // by setting this property. Mirrors the convention used by `RandomDataGenerator` and other Spark - // suites that expose tunables via `spark.sql.test.<feature>` system properties. - private val seedSystemProperty: String = - "spark.sql.test.autocdc.outOfOrderConvergenceSeed" - - private def resolveTestSeed(): Long = { - Option(System.getProperty(seedSystemProperty)).map(_.toLong).getOrElse(Random.nextLong()) - } - - private val keyColumn: String = "key" - private val nameColumn: String = "name" - private val amountColumn: String = "amount" - private val activeColumn: String = "active" - private val sequenceColumn: String = "sequence" - private val isDeleteColumn: String = "is_delete" - - private val sourceColumnNames: Seq[String] = - Seq(keyColumn, nameColumn, amountColumn, activeColumn, sequenceColumn, isDeleteColumn) - - private def randomUpsertOrDelete( - rand: Random, key: Int, sequence: Long, isDelete: Boolean): SourceRow = { - val colorPalette = Seq("red", "blue", "green", "yellow") - SourceRow( - key = key, - name = Option.when(rand.nextDouble() < nonNullProbability)( - colorPalette(rand.nextInt(colorPalette.length))), - amount = Option.when(rand.nextDouble() < nonNullProbability)(rand.nextInt(100)), - active = Option.when(rand.nextDouble() < nonNullProbability)(rand.nextBoolean()), - sequence = sequence, - isDelete = isDelete - ) - } - - private def generateRandomCdcEventStream(rand: Random): Seq[SourceRow] = { - var nextSequence: Long = 0L - val events = ArrayBuffer.empty[SourceRow] - (0 until numDistinctKeys).foreach { key => - val numUniqueEventsForKey = rand.between(1, maxUniqueEventsPerKey + 1) - (0 until numUniqueEventsForKey).foreach { _ => - val isDelete = rand.nextDouble() < deleteEventProbability - val event = randomUpsertOrDelete(rand, key, nextSequence, isDelete) - nextSequence += 1 - events += event - if (rand.nextDouble() < duplicateEventProbability) { - events += event - } - } - } - events.sortBy(_.sequence).toSeq - } - - /** Build a pipeline context with a single AutoCDC flow of `scdType` reading from `stream`. */ - private def buildPipelineContext( - targetTable: String, - stream: MemoryStream[SourceRow], - scdType: ScdType): TestGraphRegistrationContext = { - new TestGraphRegistrationContext(spark) { - registerTable(targetTable, catalog = Some(catalog), database = Some(namespace)) - registerFlow(autoCdcFlow( - name = s"${targetTable}_flow", - target = targetTable, - query = dfFlowFunc(stream.toDF().toDF(sourceColumnNames: _*)), - keys = Seq(keyColumn), - sequencing = functions.col(sequenceColumn), - deleteCondition = Some(functions.col(isDeleteColumn) === true), - columnSelection = Some(ColumnSelection.ExcludeColumns( - Seq(UnqualifiedColumnName(isDeleteColumn)) - )), - scdType = scdType - )) - } - } - - /** - * DDL fragment for the SCD-type-specific reserved columns a target table carries after the - * user-selected data columns: the CDC metadata column for SCD1, and the interval bounds plus - * metadata column for SCD2. The sequencing type is BIGINT here. - */ - private def reservedColumnsDdl(scdType: ScdType): String = scdType match { - case ScdType.Type1 => scd1MetadataDdl - case ScdType.Type2 => scd2MetadataDdl - } - - private def createTargetTable(targetTable: String, scdType: ScdType): Unit = { - spark.sql( - s"CREATE TABLE $catalog.$namespace.$targetTable (" + - s"`$keyColumn` INT NOT NULL, " + - s"`$nameColumn` STRING, " + - s"`$amountColumn` INT, " + - s"`$activeColumn` BOOLEAN, " + - s"`$sequenceColumn` BIGINT NOT NULL, " + - s"${reservedColumnsDdl(scdType)})" - ) - } + with AutoCdcGraphExecutionTestMixin + with AutoCdcRandomCdcTestMixin { private def assertTargetsConverge(inOrderTable: String, outOfOrderTable: String): Unit = { checkAnswer( @@ -176,57 +41,49 @@ class AutoCdcOutOfOrderConvergenceSuite ) } - private def runConvergenceTest(seed: Long, scdType: ScdType): Unit = { - val session = spark - import session.implicits._ - - val rand = new Random(seed) - val sortedEventStream = generateRandomCdcEventStream(rand) - val shuffledEventStream = rand.shuffle(sortedEventStream) - - withClue( - s"\nscdType=${scdType.label} seed=$seed " + - s"(rerun with -D$seedSystemProperty=$seed to reproduce)\n" + - s"events (${sortedEventStream.size} total, sorted by sequence):\n" + - sortedEventStream.map(r => s" $r").mkString("\n") + "\n" - ) { - // Table names are scd-type-suffixed purely for readability: the SCD1 and SCD2 tests run as - // separate test cases and the mixin's afterEach resets the catalog between them, so they - // could not collide even with identical names; the suffix just makes a failing run's tables - // self-identifying. - val suffix = scdType.label.toLowerCase(java.util.Locale.ROOT) - val inOrderTable = s"inorder_target_$suffix" - val outOfOrderTable = s"outoforder_target_$suffix" - createTargetTable(inOrderTable, scdType) - createTargetTable(outOfOrderTable, scdType) - - val inOrderStream = MemoryStream[SourceRow] - val inOrderCtx = buildPipelineContext(inOrderTable, inOrderStream, scdType) - inOrderStream.addData(sortedEventStream: _*) - runPipeline(inOrderCtx) - - val outOfOrderStream = MemoryStream[SourceRow] - val outOfOrderCtx = buildPipelineContext(outOfOrderTable, outOfOrderStream, scdType) - val totalEvents = shuffledEventStream.size - (0 until numOutOfOrderBatches).foreach { batchIndex => - val batchStart = batchIndex * totalEvents / numOutOfOrderBatches - val batchEnd = (batchIndex + 1) * totalEvents / numOutOfOrderBatches - outOfOrderStream.addData(shuffledEventStream.slice(batchStart, batchEnd): _*) - runPipeline(outOfOrderCtx) + private def runConvergenceTest(scdType: ScdType, testName: String): Unit = { + val numDistinctKeys = resolveNumDistinctKeys() + val maxUniqueEventsPerKey = resolveMaxUniqueEventsPerKey() + val numBatches = resolveNumBatches() + + forEachConvergenceSeed(testName) { (seed, seedIndex) => + val rand = new Random(seed) + val sortedEventStream = generateRandomCdcEventStream(rand) + val shuffledEventStream = rand.shuffle(sortedEventStream) + + withClue( + s"\nout-of-order convergence scdType=${scdType.label} testName=$testName " + + s"seedIndex=$seedIndex seed=$seed " + + s"(rerun this test with -D$convergenceReproSeedSystemProperty=$seed to reproduce)\n" + + s"keys=$numDistinctKeys maxEventsPerKey=$maxUniqueEventsPerKey " + + s"numBatches=$numBatches events=${sortedEventStream.size}\n" + ) { + val inOrderTable = s"inorder_target_$seedIndex" + val outOfOrderTable = s"outoforder_target_$seedIndex" + + // In-order baseline: one microbatch with the sequence-sorted stream. + runRandomCdcPipeline(inOrderTable, scdType, sortedEventStream, numBatches = 1) + // Out-of-order: same events shuffled across the configured number of microbatches. + runRandomCdcPipeline(outOfOrderTable, scdType, shuffledEventStream, numBatches) + + // Only the user-visible target must converge. The auxiliary tables legitimately differ by + // arrival order (e.g. deletedByBatchId stamps and cross-batch GC depend on how events are + // batched), so they are not compared. + assertTargetsConverge(inOrderTable, outOfOrderTable) } - - // Only the user-visible target must converge. The auxiliary tables legitimately differ by - // arrival order (e.g. deletedByBatchId stamps and cross-batch GC depend on how events are - // batched), so they are not compared. - assertTargetsConverge(inOrderTable, outOfOrderTable) } } - test("SCD1 merge converges across micro-batch shuffling for randomly generated CDC events") { - runConvergenceTest(resolveTestSeed(), ScdType.Type1) + private val scd1OutOfOrderTestName = + "SCD1 merge converges across micro-batch shuffling for randomly generated CDC events" + private val scd2OutOfOrderTestName = + "SCD2 merge converges across micro-batch shuffling for randomly generated CDC events" + + test(scd1OutOfOrderTestName) { + runConvergenceTest(ScdType.Type1, scd1OutOfOrderTestName) } - test("SCD2 merge converges across micro-batch shuffling for randomly generated CDC events") { - runConvergenceTest(resolveTestSeed(), ScdType.Type2) + test(scd2OutOfOrderTestName) { + runConvergenceTest(ScdType.Type2, scd2OutOfOrderTestName) } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcRandomCdcTestMixin.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcRandomCdcTestMixin.scala new file mode 100644 index 0000000000000..df9746e8c0934 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcRandomCdcTestMixin.scala @@ -0,0 +1,295 @@ +/* + * 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.pipelines.graph + +import scala.collection.mutable.ArrayBuffer +import scala.util.Random + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} +import org.apache.spark.sql.pipelines.graph.AutoCdcRandomCdcTestMixin.SourceRow +import org.apache.spark.sql.pipelines.utils.ExecutionTest +import org.apache.spark.sql.test.SharedSparkSession + +object AutoCdcRandomCdcTestMixin { + /** + * A single CDC event in a randomly-generated AutoCDC source stream. + * + * @param key Identity column (the AutoCDC `keys`). + * @param name Data column (nullable string). + * @param amount Data column (nullable int). + * @param active Data column (nullable boolean). + * @param sequence Sequencing value (the AutoCDC `sequencing` expression). + * @param isDelete Drives the AutoCDC `deleteCondition`; `true` marks the event as a delete, + * `false` as an upsert. Excluded from the target via `columnSelection`. + */ + case class SourceRow( + key: Int, + name: Option[String], + amount: Option[Int], + active: Option[Boolean], + sequence: Long, + isDelete: Boolean) +} + +/** + * Shared random-CDC fixture helpers for AutoCDC differential convergence suites + * ([[AutoCdcOutOfOrderConvergenceSuite]], [[AutoCdcCrossScdConvergenceSuite]]). + * + * Owns the common event schema, stream generator, and microbatch feeding for random-data AutoCDC + * suites. + * + * Exposed random data generation knobs (optional; defaults are CI-sized): + * - `spark.sql.test.autocdc.convergenceMultiRunBaseSeed` (multi-run mode) + * - `spark.sql.test.autocdc.convergenceMultiRunNumSeeds` (multi-run mode) + * - `spark.sql.test.autocdc.convergenceReproSeed` (repro mode; overrides multi-run) + * - `spark.sql.test.autocdc.convergenceNumKeys` + * - `spark.sql.test.autocdc.convergenceMaxEventsPerKey` + * - `spark.sql.test.autocdc.convergenceNumBatches` + * + * Suites may override these defaults when they genuinely need a different baseline. For local + * stress testing, for example: + * {{{ + * build/sbt \ + * -Dspark.sql.test.autocdc.convergenceMultiRunNumSeeds=10 \ + * -Dspark.sql.test.autocdc.convergenceNumKeys=100 \ + * 'pipelines/testOnly *AutoCdcCrossScdConvergenceSuite' + * }}} + * + * Two execution modes: + * - Multi-run (default): each test generates `convergenceMultiRunNumSeeds` random CDC streams + * from `convergenceMultiRunBaseSeed`. Used for CI and local stress testing. + * - Repro: when `convergenceReproSeed` is set, each test runs a single stream with that seed + * and ignores `convergenceMultiRunBaseSeed` and `convergenceMultiRunNumSeeds`. Used to replay + * one failing case from the seed printed in a failure message. + */ +trait AutoCdcRandomCdcTestMixin extends Logging { + self: ExecutionTest with SharedSparkSession with AutoCdcGraphExecutionTestMixin => + + // Probability an event is a delete; (1 - this) is the upsert probability. + protected val deleteEventProbability: Double = 0.20 + // Probability an event is immediately re-emitted with the same sequence and payload. + protected val duplicateEventProbability: Double = 0.15 + // Probability an upsert repeats the previous upsert's payload at a new sequence. In SCD2 this + // will produce a no-op upsert row, provided that sequence is excluded from track-history column + // selection. + protected val noOpContinuationProbability: Double = 0.15 + // Probability an optional payload column is non-null; (1 - this) is the null probability. + protected val nonNullProbability: Double = 0.75 + + // CI-sized defaults shared by every convergence suite. Override in a suite only when that + // suite genuinely needs a different baseline; prefer the shared system properties for + // local stress scaling so both suites stay aligned under normal CI. + protected val defaultBaseSeed: Long = 0x5EEDL + protected val defaultNumDistinctKeys: Int = 5 + protected val defaultMaxUniqueEventsPerKey: Int = 80 + protected val defaultNumBatches: Int = 8 + protected val defaultNumSeedsPerRun: Int = 1 + + // Exposed so suite failure clues can tell callers how to force a deterministic replay. + protected val multiRunBaseSeedSystemProperty: String = + "spark.sql.test.autocdc.convergenceMultiRunBaseSeed" + protected val multiRunNumSeedsSystemProperty: String = + "spark.sql.test.autocdc.convergenceMultiRunNumSeeds" + protected val convergenceReproSeedSystemProperty: String = + "spark.sql.test.autocdc.convergenceReproSeed" + private val numKeysSystemProperty: String = + "spark.sql.test.autocdc.convergenceNumKeys" + private val maxEventsPerKeySystemProperty: String = + "spark.sql.test.autocdc.convergenceMaxEventsPerKey" + private val numBatchesSystemProperty: String = + "spark.sql.test.autocdc.convergenceNumBatches" + + private def positiveIntProp(name: String, default: Int): Int = { + val value = Option(System.getProperty(name)).map(_.toInt).getOrElse(default) + require(value > 0, s"$name must be positive, but got $value") + value + } + + protected def configuredBaseSeed: Long = + Option(System.getProperty(multiRunBaseSeedSystemProperty)) + .map(_.toLong) + .getOrElse(defaultBaseSeed) + + private def configuredReproSeed: Option[Long] = + Option(System.getProperty(convergenceReproSeedSystemProperty)).map(_.toLong) + + private def resolveNumSeeds(): Int = + positiveIntProp(multiRunNumSeedsSystemProperty, defaultNumSeedsPerRun) + + protected def resolveNumDistinctKeys(): Int = + positiveIntProp(numKeysSystemProperty, defaultNumDistinctKeys) + + protected def resolveMaxUniqueEventsPerKey(): Int = + positiveIntProp(maxEventsPerKeySystemProperty, defaultMaxUniqueEventsPerKey) + + protected def resolveNumBatches(): Int = + positiveIntProp(numBatchesSystemProperty, defaultNumBatches) + + /** + * Invoke `callback(seed, seedIndex)` for each convergence iteration. + * + * When [[configuredReproSeed]] is set, invokes the callback once with that seed (repro mode). + * Otherwise derives `numSeeds` iterations from [[configuredBaseSeed]] and [[testName]] + * (multi-run mode). + */ + protected def forEachConvergenceSeed(testName: String)(callback: (Long, Int) => Unit): Unit = { + configuredReproSeed match { + case Some(reproSeed) => + logInfo( + s"AutoCDC convergence repro mode for test '$testName': " + + s"-D$convergenceReproSeedSystemProperty=$reproSeed " + + s"(-D$multiRunBaseSeedSystemProperty and " + + s"-D$multiRunNumSeedsSystemProperty are ignored)") + callback(reproSeed, 0) + case None => + val perTestBaseSeed = configuredBaseSeed ^ testName.hashCode.toLong + val numSeeds = resolveNumSeeds() + val masterRand = new Random(perTestBaseSeed) + val seeds = perTestBaseSeed +: Seq.fill(numSeeds - 1)(masterRand.nextLong()) + seeds.zipWithIndex.foreach { case (seed, seedIndex) => + callback(seed, seedIndex) + } + } + } + + // Forward declare key, sequence, and data columns, so that inheriting suites can reference them. + protected val keyColumn: String = "key" + protected val nameColumn: String = "name" + protected val amountColumn: String = "amount" + protected val activeColumn: String = "active" + protected val sequenceColumn: String = "sequence" + protected val isDeleteColumn: String = "is_delete" + + protected val sourceColumnNames: Seq[String] = + Seq(keyColumn, nameColumn, amountColumn, activeColumn, sequenceColumn, isDeleteColumn) + + /** User data columns on the target; excludes CDC metadata and SCD2 interval bounds. */ + protected val dataColumnNames: Seq[String] = + Seq(keyColumn, nameColumn, amountColumn, activeColumn, sequenceColumn) + + private def randomUpsertOrDelete( + rand: Random, key: Int, sequence: Long, isDelete: Boolean): SourceRow = { + val colorPalette = Seq("red", "blue", "green", "yellow") + SourceRow( + key = key, + name = Option.when(rand.nextDouble() < nonNullProbability)( + colorPalette(rand.nextInt(colorPalette.length))), + amount = Option.when(rand.nextDouble() < nonNullProbability)(rand.nextInt(100)), + active = Option.when(rand.nextDouble() < nonNullProbability)(rand.nextBoolean()), + sequence = sequence, + isDelete = isDelete + ) + } + + /** + * Generate a sequence-sorted CDC event stream. + */ + protected def generateRandomCdcEventStream(rand: Random): Seq[SourceRow] = { + val numDistinctKeys = resolveNumDistinctKeys() + val maxUniqueEventsPerKey = resolveMaxUniqueEventsPerKey() + + var nextSequence: Long = 0L + val allEvents = ArrayBuffer.empty[SourceRow] + (0 until numDistinctKeys).foreach { key => + val numUniqueEventsForKey = rand.between(1, maxUniqueEventsPerKey + 1) + val eventsForKey = ArrayBuffer.empty[SourceRow] + + (0 until numUniqueEventsForKey).foreach { _ => + val isDelete = rand.nextDouble() < deleteEventProbability + val event = if (isDelete) { + randomUpsertOrDelete(rand, key, nextSequence, isDelete = true) + } else { + val previousEventIfUpsertOpt = eventsForKey.lastOption.filterNot(_.isDelete) + val upsertToNoOpContinueOpt = previousEventIfUpsertOpt.filter( + _ => rand.nextDouble() < noOpContinuationProbability) + + upsertToNoOpContinueOpt match { + case Some(upsertToNoOpContinue) => + // If we're no-op continuing a previous upsert, reuse the same [tracked history] + // columns, incrementing only the sequence. This relies on sequence being the single + // non-track-history column in the AutoCDC configuration. + upsertToNoOpContinue.copy(sequence = nextSequence) + case _ => + // If we're not no-op continuing a previous upsert, create a new upsert event. + randomUpsertOrDelete(rand, key, nextSequence, isDelete = false) + } + } + + // By AutoCDC contract, only exact duplicate re-emissions (handled separately below) may + // reuse sequences for a particular key. Otherwise, the behavior for two unique events for + // the same key with the same sequence leads to undefined behavior. Each distinct event + // creation for this key should increment `nextSequence`. + nextSequence += 1 + eventsForKey += event + + if (rand.nextDouble() < duplicateEventProbability) { + // Full duplicate events are intentionally not counted against `numUniqueEventsForKey`. + // These differ from no-op upsert continuation events, as they share the same sequence as + // their preceding event too, in addition to all other columns. + eventsForKey += event + } + } + + allEvents.addAll(eventsForKey) + } + allEvents.sortBy(_.sequence).toSeq + } + + /** + * Feed `events` through an AutoCDC pipeline of `scdType` across `numBatches` microbatches + * (one pipeline run per microbatch). The target and auxiliary tables are created by pipeline + * materialization from the flow's inferred schema. + */ + protected def runRandomCdcPipeline( + targetTable: String, + scdType: ScdType, + events: Seq[SourceRow], + numBatches: Int): Unit = { + val session = spark + import session.implicits._ + + val stream = MemoryStream[SourceRow] + val ctx = singleAutoCdcFlowPipeline( + flowName = s"${targetTable}_flow", + target = targetTable, + sourceDf = stream.toDF().toDF(sourceColumnNames: _*), + keys = Seq(keyColumn), + sequencing = functions.col(sequenceColumn), + columnSelection = Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName(isDeleteColumn)) + )), + deleteCondition = Some(functions.col(isDeleteColumn) === true), + scdType = scdType, + trackHistorySelection = scdType match { + case ScdType.Type1 => None + case ScdType.Type2 => Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName(sequenceColumn)))) + } + ) + val totalEvents = events.size + (0 until numBatches).foreach { batchIndex => + val batchStart = batchIndex * totalEvents / numBatches + val batchEnd = (batchIndex + 1) * totalEvents / numBatches + stream.addData(events.slice(batchStart, batchEnd): _*) + runPipeline(ctx) + } + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableDurabilitySuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableDurabilitySuite.scala index 422b95b3c2678..73ef85c50189e 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableDurabilitySuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableDurabilitySuite.scala @@ -39,10 +39,9 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { - test("a higher-sequence event in a later pipeline run correctly upserts the row") { - val session = spark - import session.implicits._ + import testImplicits._ + test("a higher-sequence event in a later pipeline run correctly upserts the row") { spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -82,9 +81,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("an event with a sequence lower than what was applied in a prior pipeline run " + "is suppressed") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -120,9 +116,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("the auxiliary table places the AutoCDC key column first, ahead of any non-key " + "source columns") { - val session = spark - import session.implicits._ - // Source DF column order is (name, id, version): the AutoCDC key column `id` does NOT // appear first in the source DF. The auxiliary table must still write `id` as its // leading column. @@ -150,9 +143,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("the auxiliary table preserves the user's declared key order, independent of the " + "source DataFrame and target table column orders") { - val session = spark - import session.implicits._ - // Source DF: (value, id, region, version). Target table: (value, id, region, version, // _cdc_metadata) -- same ordering as the source. The user, however, declares // `keys = Seq("region", "id")` -- the OPPOSITE order from how those columns appear in @@ -183,9 +173,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("a dry run resolves and validates the graph without provisioning the auxiliary " + "table") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -208,9 +195,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("if the AutoCDC auxiliary table is dropped between runs, it is transparently " + "recreated") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -247,9 +231,6 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite test("auxiliary key-column-names property survives identifiers containing special " + "characters that exercise both JSON and SQL string-literal escaping") { - val session = spark - import session.implicits._ - // This test exercises the full identifier-text persistence path with composite keys whose // names collectively cover every escape class: // - `it's` -- single quote: not escaped by JSON; the writer must double it @@ -317,7 +298,7 @@ class AutoCdcScd1AuxiliaryTableDurabilitySuite s"auxiliary table $auxName is missing the " + s"${AutoCdcAuxiliaryTable.keyColumnNamesProperty} property; got: ${rows.toSeq}" )) - AutoCdcAuxiliaryTable.parseKeyColumnNames(prop.getString(1)) + AutoCdcAuxiliaryTable.parseColumnNames(prop.getString(1)) .getOrElse(fail( s"auxiliary table $auxName has a malformed " + s"${AutoCdcAuxiliaryTable.keyColumnNamesProperty} property: '${prop.getString(1)}'" diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableSpecSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableSpecSuite.scala index 530cae7cfa585..8bd2f2a0d1e84 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableSpecSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1AuxiliaryTableSpecSuite.scala @@ -43,12 +43,12 @@ import org.apache.spark.sql.types.LongType */ class AutoCdcScd1AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSession { + import testImplicits._ + private def targetIdentifier = fullyQualifiedIdentifier("target") /** Source change feed with data columns `(id, name, version)`. */ private def sourceDf = { - val session = spark - import session.implicits._ val stream = MemoryStream[(Int, String, Long)] stream.addData((1, "alice", 1L)) stream.toDF().toDF("id", "name", "version") @@ -75,7 +75,9 @@ class AutoCdcScd1AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe deleteCondition = None, storedAsScdType = ScdType.Type1))) val graph = ctx.resolveToDataflowGraph() - graph.auxiliaryTableSpecs(targetIdentifier).asInstanceOf[AutoCdcAuxiliaryTableSpec] + val inferredSchemas = graph.inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis) + graph.auxiliaryTableSpecs(inferredSchemas)(targetIdentifier) + .asInstanceOf[AutoCdcAuxiliaryTableSpec] } test("SCD1 aux schema is exactly the key columns plus the CDC metadata column") { @@ -106,7 +108,7 @@ class AutoCdcScd1AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe assert(spec.properties(AutoCdcAuxiliaryTable.scdTypePropertyKey) == ScdType.Type1.label) assert(spec.expectedKeyFields.map(_.name) == Seq("id")) assert( - AutoCdcAuxiliaryTable.parseKeyColumnNames( + AutoCdcAuxiliaryTable.parseColumnNames( spec.properties(AutoCdcAuxiliaryTable.keyColumnNamesProperty)).contains(Seq("id"))) assert(spec.identifier == AutoCdcAuxiliaryTable.identifier(targetIdentifier)) assert(spec.targetTableIdentifier == targetIdentifier) @@ -116,7 +118,7 @@ class AutoCdcScd1AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe val spec = scd1AuxSpec(keys = Seq("id", "name")) assert(spec.expectedKeyFields.map(_.name) == Seq("id", "name")) assert( - AutoCdcAuxiliaryTable.parseKeyColumnNames( + AutoCdcAuxiliaryTable.parseColumnNames( spec.properties(AutoCdcAuxiliaryTable.keyColumnNamesProperty)).contains(Seq("id", "name"))) // The aux schema is the two keys followed by the metadata column, and nothing else. assert( diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1FullRefreshSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1FullRefreshSuite.scala index 549c79116d361..d7c676a211179 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1FullRefreshSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1FullRefreshSuite.scala @@ -37,10 +37,9 @@ class AutoCdcScd1FullRefreshSuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { - test("full refresh wipes target rows and the auxiliary table for the refreshed flow") { - val session = spark - import session.implicits._ + import testImplicits._ + test("full refresh wipes target rows and the auxiliary table for the refreshed flow") { spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -96,9 +95,6 @@ class AutoCdcScd1FullRefreshSuite test("after a full refresh, an event with a sequence below the previous run's " + "watermark now lands") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -157,9 +153,6 @@ class AutoCdcScd1FullRefreshSuite } test("selective full refresh wipes only the requested target's auxiliary state") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.t_a " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1KeyDriftSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1KeyDriftSuite.scala index 1d27b7d69dab2..e180056761f20 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1KeyDriftSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1KeyDriftSuite.scala @@ -307,6 +307,46 @@ class AutoCdcScd1KeyDriftSuite } } + test("AutoCDC key drift validation uses pipeline case sensitivity, not session default") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("id"))) + + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((1, 2L)) + val ctx2 = new TestGraphRegistrationContext( + spark, + Map(SQLConf.CASE_SENSITIVE.key -> "true")) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "flow_v2", + target = "target", + query = dfFlowFunc(stream2.toDF().toDF("Id", "version")), + keys = Seq("Id"), + sequencing = $"version")) + } + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "expectedKeySchema" -> "Id INT NOT NULL", + "recordedKeySchema" -> "id INT NOT NULL" + ) + ) + } + } + test("under the default (case-insensitive) resolver, an AutoCDC flow whose key differs only " + "in case from the recorded key does NOT trigger drift") { // Pairs with the case-sensitive test above: same recorded key, but under the default @@ -315,11 +355,11 @@ class AutoCdcScd1KeyDriftSuite // case-sensitive resolver in the validator is caught. // // Note that only the *key declaration* (`Seq("Id")`) has different casing here -- the - // source DF column name still matches the target's `id` exactly. Differing the source DF - // column casing as well would not exercise drift: [[SchemaMergingUtils.mergeSchemas]] is - // case-sensitive on column names and would add `Id` as a new column to the target, - // producing AMBIGUOUS_REFERENCE during the streaming write rather than letting drift - // validation make the call. + // source DF column name still matches the target's `id` exactly. This keeps the test focused + // on the drift validator: whether the source DF column were `id` or `Id`, under the default + // (case-insensitive) resolver schema evolution folds it onto the existing `id` (SPARK-58517), + // so the streaming write itself would not fail either way and drift validation remains the + // sole decision-maker. spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1MultiPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1MultiPipelineSuite.scala index 7191f2e60cd08..c24f7283af2af 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1MultiPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1MultiPipelineSuite.scala @@ -38,11 +38,10 @@ class AutoCdcScd1MultiPipelineSuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { + import testImplicits._ + test("two AutoCDC pipelines targeting separate tables maintain independent target and " + "auxiliary tables") { - val session = spark - import session.implicits._ - // Two distinct target tables created up-front. spark.sql( s"CREATE TABLE $catalog.$namespace.t_a " + @@ -92,9 +91,6 @@ class AutoCdcScd1MultiPipelineSuite test("a downstream pipeline can read an AutoCDC target written by a different pipeline " + "without observing the CDC metadata column") { - val session = spark - import session.implicits._ - // Pipeline #1 writes into target `src` via AutoCDC. spark.sql( s"CREATE TABLE $catalog.$namespace.src " + @@ -130,9 +126,6 @@ class AutoCdcScd1MultiPipelineSuite test("two AutoCDC pipelines targeting the same table with identical key and data " + "schemas merge into a shared target table") { - val session = spark - import session.implicits._ - // Target table is created once up-front; both pipelines target it with the same // AutoCDC `keys` and the same source-DF data schema. The two pipelines have distinct // flow names ("flow_v1" / "flow_v2") so they own independent streaming checkpoints, @@ -189,9 +182,6 @@ class AutoCdcScd1MultiPipelineSuite test("two AutoCDC pipelines targeting the same table with the same key but different " + "data columns evolve the shared target schema") { - val session = spark - import session.implicits._ - // Target is created up-front with pipeline #1's schema only; pipeline #2 brings a new // top-level nullable `age` column that the dataset materialization layer is expected // to schema-merge into the target. @@ -260,9 +250,6 @@ class AutoCdcScd1MultiPipelineSuite test("a second pipeline targeting an existing AutoCDC table with different keys " + "fails with KEY_SCHEMA_DRIFT") { - val session = spark - import session.implicits._ - // Target table with both candidate keys present so the second pipeline would otherwise // be schema-compatible with the first; only the AutoCDC `keys` differ between flows. spark.sql( diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SchemaEvolutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SchemaEvolutionSuite.scala index 635e3d93de56d..448820443d13d 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SchemaEvolutionSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SchemaEvolutionSuite.scala @@ -34,8 +34,9 @@ import org.apache.spark.sql.test.SharedSparkSession /** * Tests covering AutoCDC's interaction with non-key schema evolution across pipeline runs. The * suite documents the supported additive cases (new top-level columns, new nested fields in - * array-of-struct, broadening / narrowing column selection) and the cases that fail loudly - * today (subtractive nested evolution, type-incompatible changes, case-only renames). + * array-of-struct, broadening / narrowing column selection, and -- under case-insensitive + * resolution -- a source column differing from an existing one only in case) and the cases that + * fail loudly today (subtractive nested evolution, type-incompatible changes). * * These behaviors are largely inherited from the lower layers (`SchemaMergingUtils` for * schema merge, the v2 writer's column-resolution layer for nested-field handling) rather @@ -47,10 +48,9 @@ class AutoCdcScd1SchemaEvolutionSuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { - test("a nullable non-key column merges correctly with mixed NULL and non-NULL values") { - val session = spark - import session.implicits._ + import testImplicits._ + test("a nullable non-key column merges correctly with mixed NULL and non-NULL values") { // Single MemoryStream with `email` as nullable from the start. Run #1 emits a row with // a NULL email; run #2 emits an upsert with a non-NULL email. spark.sql( @@ -86,9 +86,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("widening a non-key column's type between runs fails with " + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { - val session = spark - import session.implicits._ - // Changing a non-key column's type between pipeline runs is rejected by // `SchemaMergingUtils` with CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE even when the new type // is strictly wider. Users must full-refresh the target to change column types. @@ -130,9 +127,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("narrowing a non-key column's type between runs fails with " + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { - val session = spark - import session.implicits._ - // Mirror image of the widening test above: changing a non-key column's type between // pipeline runs is rejected by SchemaMergingUtils with CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE // even when the new type is strictly narrower. @@ -175,9 +169,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("a new top-level nullable column appearing in the source DF between runs is " + "added to the target") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -223,9 +214,6 @@ class AutoCdcScd1SchemaEvolutionSuite } test("additive target-column evolution leaves the SCD1 auxiliary table schema unchanged") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -274,9 +262,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("broadening the column selection between runs adds the newly-included column to " + "the target") { - val session = spark - import session.implicits._ - // Source DF schema is fixed at (id, name, email, version) across both runs. Only the // `columnSelection` knob differs: run #1 includes (id, name, version); run #2 selects // None (= all source columns). mergeSchemas adds `email` to the target via the same @@ -322,9 +307,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("narrowing the column selection between runs preserves the dropped column on " + "existing rows and leaves it NULL on new rows") { - val session = spark - import session.implicits._ - // Validates the additive-only column-selection contract on the narrowing side: // tightening `columnSelection` between runs leaves the dropped column in place at the // schema level (SDP's `SchemaMergingUtils.mergeSchemas` is a union, never a subtraction). @@ -369,9 +351,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("a top-level column dropped from the source DF between runs is preserved on " + "existing rows and left NULL on new rows") { - val session = spark - import session.implicits._ - // Symmetric to the new-source-column case (which exercises the source DF *gaining* a // column). Validates that the additive-only column-selection contract holds when the // narrowing is driven by the source DF's own schema shrinking, rather than by a @@ -420,9 +399,6 @@ class AutoCdcScd1SchemaEvolutionSuite } test("dropping a nested struct field between runs fails with INCOMPATIBLE_DATA_FOR_TABLE") { - val session = spark - import session.implicits._ - // The v2 writer's column-resolution layer requires every nested target field to be // present in the microbatch DF. When run #2's source projection drops `b.c`, the merge // fails with INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_FIND_DATA. Users who want to drop a @@ -479,9 +455,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("a new field added inside an array<struct> element between runs is added to the " + "target") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(key INT NOT NULL, version BIGINT NOT NULL, " + @@ -533,9 +506,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("dropping a field inside an array<struct> element between runs fails with " + "INCOMPATIBLE_DATA_FOR_TABLE") { - val session = spark - import session.implicits._ - // Symmetric to the nested-struct case, but for `array<struct>`. The v2 writer rejects // the merge because it cannot find data for the target's `vals.element.b.d` column // when run #2's projection drops `d` from the element struct. Users must full-refresh @@ -585,18 +555,15 @@ class AutoCdcScd1SchemaEvolutionSuite ) } - test("a source DF column whose name differs from the target only by case fails with " + - "AMBIGUOUS_REFERENCE under case-insensitive resolution") { - val session = spark - import session.implicits._ - - // `DatasetManager`'s schema-merge compares the existing target schema and the flow's - // output schema *case-sensitively*: `SchemaMergingUtils.mergeSchemas` calls - // `StructType.merge` without forwarding the session-level case-sensitivity. When the - // target has `value` and the source DF emits `Value`, the merged schema ends up with - // both as separate columns. Reference resolution downstream is case-insensitive - // (Spark's default), so the MERGE plan trips on the duplicate and reports - // AMBIGUOUS_REFERENCE. + test("a source DF column whose name differs from the target only by case is folded onto the " + + "existing column under case-insensitive resolution") { + // Under case-insensitive resolution (Spark's default), a target `value` and a source `Value` + // are the same column. Schema evolution honors that by threading case-sensitivity into + // `SchemaMergingUtils.mergeSchemas`: the merge maps `Value` onto the existing `value`, so + // `diffSchemas` has no case-only difference left to process. No second column is added, and + // the write succeeds. (Before SPARK-58517 the merge ran case-sensitively regardless of the + // session and added a duplicate `Value` column, after which the case-insensitive MERGE plan + // tripped on the ambiguous reference.) withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { spark.sql( s"CREATE TABLE $catalog.$namespace.target " + @@ -605,8 +572,7 @@ class AutoCdcScd1SchemaEvolutionSuite val stream = MemoryStream[(Int, Long, String)] stream.addData((1, 1L, "alice")) - // Source DF emits `Value` (capital), differing only in case from the target's - // `value` column. + // Source DF emits `Value` (capital), differing only in case from the target's `value` column. val df = stream.toDF().toDF("key", "version", "Value") val ctx = singleAutoCdcFlowPipeline( flowName = "auto_cdc_flow", @@ -615,33 +581,24 @@ class AutoCdcScd1SchemaEvolutionSuite keys = Seq("key"), sequencing = functions.col("version")) - val ex = intercept[RuntimeException] { runPipeline(ctx) } - // The exact `name` and `referenceNames` parameters depend on internal merge-plan - // synthesis; the condition match is the meaningful invariant for this test. - checkErrorInPipelineFailure( - failure = ex, - condition = "AMBIGUOUS_REFERENCE", - parameters = Map( - "name" -> ".*", - "referenceNames" -> ".*" - ), - matchPVals = true, - queryContext = Array( - ExpectedContext( - fragment = s"`$catalog`.`$namespace`.`target`.`Value`", - start = 0, - stop = 27 - ) - ) + runPipeline(ctx) + + // The target schema is unchanged (still a single `value` column, original case), and the + // row lands with the emitted value folded into it. + assert( + spark.table(s"$catalog.$namespace.target").schema.fieldNames.toSeq === + Seq("key", "version", "value", AutoCdcReservedNames.cdcMetadataColName), + "the target should keep its single `value` column, not gain a `Value` column" + ) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, 1L, "alice", cdcMeta(None, Some(1L)))) ) } } test("extra columns on the target that the AutoCDC flow does not emit are preserved " + "across the merge") { - val session = spark - import session.implicits._ - // The target is wider than the AutoCDC flow's source DF: column `extra` is present on // the target but never produced by the flow. AutoCDC must tolerate the extra target // column -- pre-existing rows keep their `extra` value, and newly-inserted rows @@ -676,9 +633,6 @@ class AutoCdcScd1SchemaEvolutionSuite test("changing a non-key column type from TIMESTAMP to STRING between runs fails with " + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { - val session = spark - import session.implicits._ - // `mergeSchemas` rejects an incompatible type change between TIMESTAMP and STRING. // Captured alongside the type-widening / type-narrowing tests; users must full-refresh // the target to change a column's type. diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SinglePipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SinglePipelineSuite.scala index 9ea7a41d0372e..f37b86497505c 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SinglePipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1SinglePipelineSuite.scala @@ -40,10 +40,9 @@ class AutoCdcScd1SinglePipelineSuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { - test("an upsert event lands a new row in an empty target table") { - val session = spark - import session.implicits._ + import testImplicits._ + test("an upsert event lands a new row in an empty target table") { spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -73,9 +72,6 @@ class AutoCdcScd1SinglePipelineSuite test("consecutive upsert, delete, and re-upsert events for the same key in one run " + "converge to the latest event") { - val session = spark - import session.implicits._ - // Target schema deliberately omits `is_delete`: the source carries it as a control // column, drives the deleteCondition, and is excluded from the target projection. spark.sql( @@ -117,9 +113,6 @@ class AutoCdcScd1SinglePipelineSuite test("two AutoCDC flows targeting separate tables in one pipeline produce independent " + "results") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.t_a " + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -168,9 +161,6 @@ class AutoCdcScd1SinglePipelineSuite test("an AutoCDC flow targeting a table whose format does not support row-level " + "operations fails with AUTOCDC_TARGET_DOES_NOT_SUPPORT_MERGE") { - val session = spark - import session.implicits._ - // Intentionally use a non-merge-compatible catalog, whose default table format is parquet. val catalog = TestGraphRegistrationContext.DEFAULT_CATALOG val database = TestGraphRegistrationContext.DEFAULT_DATABASE diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1TargetTableDurabilitySuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1TargetTableDurabilitySuite.scala index 0f777708398b6..e3ac46db18916 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1TargetTableDurabilitySuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd1TargetTableDurabilitySuite.scala @@ -35,11 +35,10 @@ class AutoCdcScd1TargetTableDurabilitySuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { + import testImplicits._ + test("pre-loaded rows: an event with a lower sequence is suppressed and a higher one " + "wins") { - val session = spark - import session.implicits._ - spark.sql( s"CREATE TABLE $catalog.$namespace.target " + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd1MetadataDdl)" @@ -75,9 +74,6 @@ class AutoCdcScd1TargetTableDurabilitySuite test("pre-loaded target rows merge correctly on the first AutoCDC run, and the " + "auxiliary table is created lazily") { - val session = spark - import session.implicits._ - // Target was populated by some external process; this is the first AutoCDC run. spark.sql( s"CREATE TABLE $catalog.$namespace.target " + @@ -118,9 +114,6 @@ class AutoCdcScd1TargetTableDurabilitySuite test("a target table created without the CDC metadata column gets the column " + "auto-added on the first AutoCDC run") { - val session = spark - import session.implicits._ - // User creates the target without the AutoCDC metadata column. DatasetManager evolves // the existing table schema by merging it with the AutoCdcMergeFlow's output schema, // which includes the metadata column. The first run therefore proceeds normally, and diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableDurabilitySuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableDurabilitySuite.scala new file mode 100644 index 0000000000000..e6b1119383ced --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableDurabilitySuite.scala @@ -0,0 +1,376 @@ +/* + * 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.pipelines.graph + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{ + ColumnSelection, + Scd2BatchProcessor, + ScdType, + UnqualifiedColumnName +} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests covering the durability of the SCD Type 2 AutoCDC auxiliary table across pipeline runs: + * the per-key history recorded in the auxiliary table must persist between incremental runs, and + * the auxiliary table must be transparently recreated if it is deleted out-of-band. The SCD2 + * analog of [[AutoCdcScd1AuxiliaryTableDurabilitySuite]]. Unlike SCD1, the SCD2 auxiliary table's + * schema is the full target row schema plus the aux-only deleted-by-batch-id marker, so the + * schema-layout assertions differ accordingly. + */ +class AutoCdcScd2AuxiliaryTableDurabilitySuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + test("a higher-sequence event in a later pipeline run correctly closes and opens records") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Single MemoryStream reused across both pipeline runs so the streaming checkpoint can + // resume cleanly. + val changeDataFeedStream = MemoryStream[(Int, String, Long)] + def buildGraphRegistrationContext(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = changeDataFeedStream.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + // Run #1: insert id=1 at seq=1. + changeDataFeedStream.addData((1, "alice", 1L)) + runPipeline(buildGraphRegistrationContext()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2: upsert id=1 at seq=2 (closes the seq=1 record, opens a new one) and insert id=2 at + // seq=1 (new key). The auxiliary table from run #1 persists and supplies the prior history. + changeDataFeedStream.addData((1, "alice2", 2L), (2, "bob", 1L)) + runPipeline(buildGraphRegistrationContext()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 1L, 1L, 2L, scd2Meta(1L)), + Row(1, "alice2", 2L, 2L, null, scd2Meta(2L)), + Row(2, "bob", 1L, 1L, null, scd2Meta(1L)) + ) + ) + } + + test("an event with a sequence lower than what was applied in a prior pipeline run " + + "is woven in as a closed prior record") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Single MemoryStream reused across both runs so the streaming checkpoint can resume. + val stream = MemoryStream[(Int, String, Long)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + // Run #1: upsert id=1 at seq=10. Auxiliary table records the open record at seq=10. + stream.addData((1, "alice", 10L)) + runPipeline(buildCtx()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", 10L, 10L, null, scd2Meta(10L))) + ) + + // Run #2: late upsert at seq=5 (< the persisted seq=10). Unlike SCD1, SCD2 does not suppress + // it: the aux history lets reconciliation weave it in as a closed prior record ending at 10, + // while the seq=10 record stays open. + stream.addData((1, "early", 5L)) + runPipeline(buildCtx()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "early", 5L, 5L, 10L, scd2Meta(5L)), + Row(1, "alice", 10L, 10L, null, scd2Meta(10L)) + ) + ) + } + + test("the SCD2 auxiliary table schema is the full target row schema plus the " + + "deleted-by-batch-id marker, and records the key columns property") { + // Source DF column order is (name, id, version): the AutoCDC key column `id` does NOT appear + // first in the source DF. The SCD2 auxiliary table mirrors the full target row schema (all + // user + framework columns) with the aux-only deleted-by-batch-id marker appended, and records + // the key columns in the key-column-names property. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(name STRING, id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(String, Int, Long)] + stream.addData(("alice", 1, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("name", "id", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2)) + + val targetSchema = spark.table(s"$catalog.$namespace.target").schema.fieldNames.toSeq + val auxSchema = spark.table(auxTableNameFor("target")).schema.fieldNames.toSeq + // The aux schema is the full target row schema with the marker appended. + assert(auxSchema == targetSchema :+ Scd2BatchProcessor.deletedByBatchIdColName) + assert(getAuxTableKeyColumnNames(target = "target") == Seq("id")) + } + + test("the auxiliary table preserves the user's declared key order in the key-columns " + + "property, independent of the source DataFrame and target table column orders") { + // The user declares `keys = Seq("region", "id")` -- the OPPOSITE order from how those columns + // appear in both the source DF and the target. The recorded key-column-names property should + // honor the user's declared key order so subsequent runs compare keys against the same layout. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(value STRING, id INT NOT NULL, region STRING NOT NULL, " + + s"version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(String, Int, String, Long)] + stream.addData(("v", 1, "us", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("value", "id", "region", "version"), + keys = Seq("region", "id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2)) + + assert(getAuxTableKeyColumnNames(target = "target") == Seq("region", "id")) + } + + test("a dry run resolves and validates the graph without provisioning the auxiliary " + + "table") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + val updateCtx = TestPipelineUpdateContext(spark, ctx.toDataflowGraph, storageRoot) + updateCtx.pipelineExecution.dryRunPipeline() + + assert(!spark.catalog.tableExists(auxTableNameFor("target"))) + } + + test("if the SCD2 AutoCDC auxiliary table is dropped between runs, it is transparently " + + "recreated") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Single MemoryStream reused across both runs so the streaming checkpoint can resume. + val stream = MemoryStream[(Int, Long)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + stream.addData((1, 1L)) + runPipeline(buildCtx()) + assert(spark.catalog.tableExists(auxTableNameFor("target"))) + + // Manually drop the auxiliary table. + spark.sql(s"DROP TABLE ${auxTableNameFor("target")}") + assert(!spark.catalog.tableExists(auxTableNameFor("target"))) + + stream.addData((1, 2L)) + runPipeline(buildCtx()) + + // The dropped auxiliary table must be transparently recreated. Here the seq=1 record also + // lives in the target as a visible row, and SCD2 reconciliation reads affected rows from the + // target as well as the aux table, so this particular history survives the drop: the seq=2 + // event still closes the seq=1 record and opens a new one. (This is NOT a general guarantee + // that the aux table is disposable -- state the aux holds that is NOT mirrored in the target, + // e.g. a tombstone from a delete-only run, is lost on a drop; see the aux-sole-holder test + // below.) + assert(spark.catalog.tableExists(auxTableNameFor("target"))) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, 1L, 1L, 2L, scd2Meta(1L)), + Row(1, 2L, 2L, null, scd2Meta(2L)) + ) + ) + } + + test("the auxiliary table durably holds state absent from the target: a tombstone from a " + + "delete-only run closes a later lower-sequence upsert") { + // Unlike the transparently-recreated test above (where the surviving state also lived in the + // target as a visible row), here the auxiliary table is the SOLE holder of the state. A + // delete-only first run leaves the target empty but records a tombstone at seq=10 in the aux; + // the durability of THAT aux-only row is what lets a later, lower-sequence upsert land as a + // closed prior record. Drop the aux and this history is gone (the upsert would instead open a + // current record) -- which is exactly why the aux is not disposable. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Single MemoryStream reused across both runs so the streaming checkpoint can resume. + val stream = MemoryStream[(Int, String, Long, Boolean)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "version", "is_delete"), + keys = Seq("id"), + sequencing = functions.col("version"), + deleteCondition = Some(functions.col("is_delete") === true), + columnSelection = Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName("is_delete")) + )), + scdType = ScdType.Type2) + + // Run #1: a delete at seq=10. The target stays empty; the aux records a tombstone at seq=10. + stream.addData((1, "alice", 10L, true)) + runPipeline(buildCtx()) + checkAnswer(spark.table(s"$catalog.$namespace.target"), Seq.empty) + // The tombstone lives only in the aux, as a live row (its deleted-by-batch-id marker is null; + // a non-null marker is what flags a row logically deleted), with no matching visible target + // row. + assert(spark.table(auxTableNameFor("target")).count() == 1, + "the delete-only run should record exactly one aux tombstone row") + + // Run #2 (aux retained): a later upsert at seq=5, BELOW the recorded seq=10. Because the aux + // still holds the seq=10 tombstone, reconciliation weaves seq=5 in as a closed prior record + // ending at 10 rather than an open current record -- state the target alone could not supply. + stream.addData((1, "early", 5L, false)) + runPipeline(buildCtx()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "early", 5L, 5L, 10L, scd2Meta(5L))) + ) + } + + test("auxiliary key-column-names property survives identifiers containing special " + + "characters that exercise both JSON and SQL string-literal escaping") { + // This test exercises the full identifier-text persistence path with composite keys whose + // names collectively cover every escape class: + // - `it's` -- single quote: not escaped by JSON; the writer must double it + // to `''` to keep the SQL TBLPROPERTIES literal well-formed. + // - `name with spaces` -- whitespace identifier: backtick-quoted in DDL, no escaping + // needed in the JSON or the property value. + // - `a"b` -- literal double quote: JSON escapes as `\"`. + // - `c\d` -- literal backslash: JSON escapes as `\\`. + // If any layer drops, splits, or misescapes a name, the post-run lookup of the + // [[AutoCdcAuxiliaryTable.keyColumnNamesProperty]] property either fails to read or + // returns a value that is no longer a parseable JSON array of strings. + val keyNames = Seq("it's", "name with spaces", "a\"b", "c\\d") + + // SQL DDL identifier rendering: backticks delimit each identifier; an embedded backtick + // would have to be escaped by doubling, but none of these names contain one. + val targetTableDdl = keyNames + .map(name => s"`$name` STRING NOT NULL") + .mkString(", ") + s", version BIGINT NOT NULL, $scd2MetadataDdl" + spark.sql(s"CREATE TABLE $catalog.$namespace.target ($targetTableDdl)") + + // The AutoCDC API runs every key through `UnqualifiedColumnName.apply`, which calls + // `CatalystSqlParser.parseMultipartIdentifier`. To get a single-part identifier whose + // text includes special characters, the API caller has to backtick-quote at the boundary; + // we mirror that here by wrapping each name in backticks (and doubling any embedded + // backtick -- not needed for these names but kept for parity with how a user would call + // the API). + val backtickQuotedKeys = keyNames.map(name => s"`${name.replace("`", "``")}`") + + // Single MemoryStream reused across both runs so the streaming checkpoint can resume. + val stream = MemoryStream[(String, String, String, String, Long)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF((keyNames :+ "version"): _*), + keys = backtickQuotedKeys, + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + // Run #1: a single insert with arbitrary non-empty key values. + stream.addData(("v1", "v2", "v3", "v4", 1L)) + runPipeline(buildCtx()) + + // The persisted property must round-trip every name byte-for-byte. + assert(getAuxTableKeyColumnNames(target = "target") == keyNames) + + // Run #2: same keys, a higher sequence -- drift validation reads the property back, parses + // the JSON, and looks up each recorded name in the aux schema. If any layer mangled the + // identifier text (lost an escape, dropped a `'`, split on a `.`, ...), validation would + // either throw KEY_SCHEMA_DRIFT (name lookup miss) or INTERNAL_ERROR (recorded name absent + // from aux schema). Reaching the second run successfully proves the round-trip works. + stream.addData(("v1", "v2", "v3", "v4", 2L)) + runPipeline(buildCtx()) + + // The persisted property is immutable across non-full-refresh runs, so it must still be + // intact after run #2. + assert(getAuxTableKeyColumnNames(target = "target") == keyNames) + } + + private def getAuxTableKeyColumnNames(target: String): Seq[String] = { + val auxName = auxTableNameFor(target) + val rows = spark.sql(s"SHOW TBLPROPERTIES $auxName").collect() + val prop = rows + .find(_.getString(0) == AutoCdcAuxiliaryTable.keyColumnNamesProperty) + .getOrElse(fail( + s"auxiliary table $auxName is missing the " + + s"${AutoCdcAuxiliaryTable.keyColumnNamesProperty} property; got: ${rows.toSeq}" + )) + AutoCdcAuxiliaryTable.parseColumnNames(prop.getString(1)) + .getOrElse(fail( + s"auxiliary table $auxName has a malformed " + + s"${AutoCdcAuxiliaryTable.keyColumnNamesProperty} property: '${prop.getString(1)}'" + )) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableSpecSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableSpecSuite.scala index d86f3ba89bddf..46fbea24c6ae3 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableSpecSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2AuxiliaryTableSpecSuite.scala @@ -44,12 +44,12 @@ import org.apache.spark.sql.types.{LongType, StructField, StructType} */ class AutoCdcScd2AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSession { + import testImplicits._ + private def targetIdentifier = fullyQualifiedIdentifier("target") /** Source change feed with data columns `(id, name, version)`. */ private def sourceDf = { - val session = spark - import session.implicits._ val stream = MemoryStream[(Int, String, Long)] stream.addData((1, "alice", 1L)) stream.toDF().toDF("id", "name", "version") @@ -76,7 +76,9 @@ class AutoCdcScd2AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe deleteCondition = None, storedAsScdType = ScdType.Type2))) val graph = ctx.resolveToDataflowGraph() - graph.auxiliaryTableSpecs(targetIdentifier).asInstanceOf[AutoCdcAuxiliaryTableSpec] + val inferredSchemas = graph.inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis) + graph.auxiliaryTableSpecs(inferredSchemas)(targetIdentifier) + .asInstanceOf[AutoCdcAuxiliaryTableSpec] } /** The SCD2 target (inferred) schema for the default single-flow graph. */ @@ -97,7 +99,8 @@ class AutoCdcScd2AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe columnSelection = None, deleteCondition = None, storedAsScdType = ScdType.Type2))) - ctx.resolveToDataflowGraph().inferredSchema(targetIdentifier) + ctx.resolveToDataflowGraph() + .inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis)(targetIdentifier) } test("SCD2 aux schema is the full target schema plus the deleted-by-batch-id marker") { @@ -139,7 +142,7 @@ class AutoCdcScd2AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe assert(spec.properties(AutoCdcAuxiliaryTable.scdTypePropertyKey) == ScdType.Type2.label) assert(spec.expectedKeyFields.map(_.name) == Seq("id")) assert( - AutoCdcAuxiliaryTable.parseKeyColumnNames( + AutoCdcAuxiliaryTable.parseColumnNames( spec.properties(AutoCdcAuxiliaryTable.keyColumnNamesProperty)).contains(Seq("id"))) assert(spec.identifier == AutoCdcAuxiliaryTable.identifier(targetIdentifier)) assert(spec.targetTableIdentifier == targetIdentifier) @@ -149,7 +152,7 @@ class AutoCdcScd2AuxiliaryTableSpecSuite extends PipelineTest with SharedSparkSe val spec = scd2AuxSpec(keys = Seq("id", "name")) assert(spec.expectedKeyFields.map(_.name) == Seq("id", "name")) assert( - AutoCdcAuxiliaryTable.parseKeyColumnNames( + AutoCdcAuxiliaryTable.parseColumnNames( spec.properties(AutoCdcAuxiliaryTable.keyColumnNamesProperty)).contains(Seq("id", "name"))) // Both keys survive into the aux schema. assert(spec.schema.fieldNames.contains("id")) diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2ColumnEvolutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2ColumnEvolutionSuite.scala new file mode 100644 index 0000000000000..25cb185a7ff49 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2ColumnEvolutionSuite.scala @@ -0,0 +1,326 @@ +/* + * 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.pipelines.graph + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * End-to-end tests for SCD Type 2 AutoCDC column-schema evolution across runs: a microbatch that + * is narrower than the already-evolved target (a source column dropped, a nested struct/array field + * dropped, or the `COLUMNS` selection narrowed) must reconcile correctly instead of failing the + * internal union. + * + * These exercise the fix for SPARK-58418. Before it, `Scd2ForeachBatchHandler.reconcileMicrobatch` + * unioned the microbatch with the affected target/aux rows without `allowMissingColumns`, so a + * narrower microbatch failed with NUM_COLUMNS_MISMATCH (top-level) or INCOMPATIBLE_COLUMN_TYPE + * (nested). The contract asserted here is additive-tolerant: records already written keep their + * values for the no-longer-emitted column, and only records opened by the narrower microbatch carry + * null for it. + * + * This matches SCD1's behavior for a dropped top-level column + * ([[AutoCdcScd1SchemaEvolutionSuite]]). SCD2 applies the same behavior to a dropped *nested* + * struct/array field, where SCD1 instead fails with INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_FIND_DATA: + * SCD1's MERGE source is missing the nested field and the v2 writer's resolver rejects it, whereas + * SCD2's `allowMissingColumns` pads the field before the union/MERGE. So SCD2 handles nested + * subtractive evolution consistently with the top-level case (and with SCD1's top-level case), + * rather than reproducing SCD1's nested-drop limitation. + * + * Changing the effective *tracked-history* column set is a distinct, separately-scoped concern + * (SPARK-58452 / SPARK-58391) and is deliberately not exercised here: every scenario keeps the + * effective tracked set unchanged across runs, so the only thing evolving is the set of user + * columns the flow emits. Each scenario tracks history explicitly on `name` and drops the + * non-tracked `email`, so the tracked set ({name}) is unchanged and these stay valid once the + * track-history drift guard (SPARK-58391) lands. + */ +class AutoCdcScd2ColumnEvolutionSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + /** An explicit SCD2 `TRACK HISTORY ON (name)` selection, shared across the scenarios below. */ + private val trackName: Option[ColumnSelection] = + Some(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name")))) + + test("a source column dropped between runs is preserved on existing records and null on new " + + "ones") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, email STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Shared stream; run #2 projects `email` away so the microbatch is narrower than the target. + // Track history on `name` only, so the dropped `email` is a non-tracked column and dropping it + // is column-schema narrowing rather than a tracked-set change (robust once SPARK-58391 lands). + val stream = MemoryStream[(Int, String, String, Long)] + def buildCtx(includeEmail: Boolean): TestGraphRegistrationContext = { + val df = stream.toDF().toDF("id", "name", "email", "version") + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = if (includeEmail) df else df.drop("email"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2, + trackHistorySelection = trackName) + } + + // Run #1 (wide): key=1 opens a record carrying email=a@x. + stream.addData((1, "alice", "a@x", 1L)) + runPipeline(buildCtx(includeEmail = true)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", "a@x", 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2 (narrow): update key=1 (closes its record) + insert key=2. The dropped `email` is + // preserved on key=1's now-closed record and is null on the newly-opened records. + stream.addData((1, "alice2", "ignored", 2L), (2, "bob", "ignored", 1L)) + runPipeline(buildCtx(includeEmail = false)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", "a@x", 1L, 1L, 2L, scd2Meta(1L)), + Row(1, "alice2", null, 2L, 2L, null, scd2Meta(2L)), + Row(2, "bob", null, 1L, 1L, null, scd2Meta(1L)) + ) + ) + } + + test("narrowing the COLUMNS selection to drop a non-tracked column preserves it on existing " + + "records and leaves it null on new ones") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, email STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // The source DF is fixed at (id, name, email, version) across both runs; only the flow's + // `columnSelection` narrows. Tracking history on `name` only makes `email` a + // selected-but-not-tracked column, so dropping it from the selection is pure column-schema + // narrowing -- the effective tracked set ({name}) is unchanged, so this stays column evolution + // rather than a tracked-set change even once the track-history drift guard (SPARK-58391) lands. + val stream = MemoryStream[(Int, String, String, Long)] + def buildCtx(selection: Option[ColumnSelection]): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "email", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + columnSelection = selection, + scdType = ScdType.Type2, + trackHistorySelection = trackName) + + // Run #1: no selection (all columns); key=1 carries email=a@x. + stream.addData((1, "alice", "a@x", 1L)) + runPipeline(buildCtx(selection = None)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", "a@x", 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2: narrow the selection to (id, name, version), dropping the non-tracked `email`. + // Because `name` (the sole tracked column) changes, this opens a new record; key=1's closed + // record keeps a@x, and the new records carry null. + stream.addData((1, "alice2", "ignored", 2L), (2, "bob", "ignored", 1L)) + runPipeline(buildCtx(selection = Some(ColumnSelection.IncludeColumns( + Seq("id", "name", "version").map(UnqualifiedColumnName(_)) + )))) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", "a@x", 1L, 1L, 2L, scd2Meta(1L)), + Row(1, "alice2", null, 2L, 2L, null, scd2Meta(2L)), + Row(2, "bob", null, 1L, 1L, null, scd2Meta(1L)) + ) + ) + } + + test("a late narrower event weaves into history without rewriting existing records") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, email STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Track history on `name` only, so dropping `email` is column narrowing with an unchanged + // tracked set (robust once SPARK-58391 lands). + val stream = MemoryStream[(Int, String, String, Long)] + def buildCtx(includeEmail: Boolean): TestGraphRegistrationContext = { + val df = stream.toDF().toDF("id", "name", "email", "version") + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = if (includeEmail) df else df.drop("email"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2, + trackHistorySelection = trackName) + } + + // Run #1 (wide): two distinct-name records for key=1 at seq 10 and 30. + stream.addData((1, "alice", "a@x", 10L), (1, "alicia", "b@x", 30L)) + runPipeline(buildCtx(includeEmail = true)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", "a@x", 10L, 10L, 30L, scd2Meta(10L)), + Row(1, "alicia", "b@x", 30L, 30L, null, scd2Meta(30L)) + ) + ) + + // Run #2 (narrow): a late event at seq=20 with a new name bisects the seq=10 record. The + // pre-existing records keep their email values; the newly-inserted seq=20 record has null. + stream.addData((1, "annie", "ignored", 20L)) + runPipeline(buildCtx(includeEmail = false)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", "a@x", 10L, 10L, 20L, scd2Meta(10L)), + Row(1, "annie", null, 20L, 20L, 30L, scd2Meta(20L)), + Row(1, "alicia", "b@x", 30L, 30L, null, scd2Meta(30L)) + ) + ) + } + + test("a nested struct field dropped between runs is preserved on existing records and null on " + + "new ones (SCD2 is more permissive than SCD1 here)") { + // Contrast with SCD1: AutoCdcScd1SchemaEvolutionSuite rejects this exact shape with + // INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_FIND_DATA, because its MERGE source is missing `value.b.c` + // and the v2 writer's resolver cannot find data for the target's nested field. SCD2's + // `allowMissingColumns` pads the missing field before the union/MERGE, so the nested drop is + // handled the same additive-tolerant way as a top-level drop -- preserved on existing records, + // null on new ones. This is a deliberate, consistency-improving divergence from SCD1, not + // parity: SCD1's nested-drop failure is a writer limitation, not an intended policy. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, " + + s"value STRUCT<a:INT,b:STRUCT<c:INT,d:INT>>, $scd2MetadataDdl)" + ) + + // Default tracking is fine here: the tracked set is a set of top-level column *names*, and only + // the nested shape of `value` changes across runs -- the top-level name `value` is retained -- + // so the effective tracked set ({value}) is unchanged and this stays column evolution, not a + // tracked-set change, even once SPARK-58391 lands. + val stream = MemoryStream[(Int, Long, Int, Int, Int)] + def buildCtx(includeC: Boolean): TestGraphRegistrationContext = { + val src = stream.toDF().toDF("id", "version", "a", "b_c", "b_d") + val inner = if (includeC) { + functions.struct(functions.col("b_c").as("c"), functions.col("b_d").as("d")) + } else { + functions.struct(functions.col("b_d").as("d")) + } + val projected = src.select( + functions.col("id"), + functions.col("version"), + functions.struct(functions.col("a"), inner.as("b")).as("value") + ) + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projected, + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + } + + // Run #1 (wide): value.b carries both c and d for key=1. + stream.addData((1, 1L, 1, 10, 100)) + runPipeline(buildCtx(includeC = true)) + + // Run #2 (narrow): value.b drops `c`. The reconciliation union no longer fails; key=1's closed + // record keeps c=10, and the newly-opened record has c=null (d flows through). + stream.addData((1, 2L, 2, 99, 200)) + runPipeline(buildCtx(includeC = false)) + + checkAnswer( + spark.table(s"$catalog.$namespace.target").selectExpr( + "id", "version", "value.a", "value.b.c", "value.b.d", "__START_AT", "__END_AT"), + Seq( + Row(1, 1L, 1, 10, 100, 1L, 2L), + Row(1, 2L, 2, null, 200, 2L, null) + ) + ) + } + + test("a field dropped inside an array<struct> element between runs is preserved on existing " + + "records and null on new ones (SCD2 is more permissive than SCD1 here)") { + // The array<struct> analog of the nested-struct-drop test above, and the counterpart to + // AutoCdcScd1SchemaEvolutionSuite's array<struct> case, which fails with + // INCOMPATIBLE_DATA_FOR_TABLE.CANNOT_FIND_DATA on `vals.element.b.d`. SCD2's + // allowMissingColumns recurses into arrays as well as structs, so it pads the dropped element + // field before the union/MERGE and reconciles additively. Default tracking is fine: the tracked + // set is the top-level name `vals`, unchanged when only a nested element field is dropped. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, " + + s"vals ARRAY<STRUCT<a:INT,b:STRUCT<c:INT,d:INT>>>, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, Long, Int, Int, Int)] + def buildCtx(includeD: Boolean): TestGraphRegistrationContext = { + val src = stream.toDF().toDF("id", "version", "a", "b_c", "b_d") + val inner = if (includeD) { + functions.struct(functions.col("b_c").as("c"), functions.col("b_d").as("d")) + } else { + functions.struct(functions.col("b_c").as("c")) + } + val projected = src.select( + functions.col("id"), + functions.col("version"), + functions.array( + functions.struct(functions.col("a"), inner.as("b")) + ).as("vals") + ) + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projected, + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + } + + // Run #1 (wide): vals[0].b carries both c and d for key=1. + stream.addData((1, 1L, 1, 10, 100)) + runPipeline(buildCtx(includeD = true)) + + // Run #2 (narrow): drop `d` from the element struct. The union pads the missing nested element + // field; key=1's closed record keeps d=100, and the newly-opened record has d=null. + stream.addData((1, 2L, 2, 200, 99)) + runPipeline(buildCtx(includeD = false)) + + checkAnswer( + spark.table(s"$catalog.$namespace.target") + .selectExpr("id", "version", "inline(vals) as (a, b)", "__START_AT", "__END_AT") + .selectExpr("id", "version", "a", "b.c", "b.d", "__START_AT", "__END_AT"), + Seq( + Row(1, 1L, 1, 10, 100, 1L, 2L), + Row(1, 2L, 2, 200, null, 2L, null) + ) + ) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2FullRefreshSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2FullRefreshSuite.scala new file mode 100644 index 0000000000000..3b6151d6c929f --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2FullRefreshSuite.scala @@ -0,0 +1,263 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.pipelines.graph + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{ + ColumnSelection, + ScdType, + UnqualifiedColumnName +} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests covering AutoCDC's full-refresh semantics for SCD Type 2 targets: full refresh must wipe + * both the target rows and the (richer) SCD2 auxiliary table for the refreshed targets, and must + * leave non-refreshed targets untouched in selective-refresh mode. The SCD2 analog of + * [[AutoCdcScd1FullRefreshSuite]]. + */ +class AutoCdcScd2FullRefreshSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + /** Create an SCD2 target with user columns `(id, name, version)` plus the framework columns. */ + private def createScd2Target(table: String): Unit = { + spark.sql( + s"CREATE TABLE $table (" + + s"id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + } + + test("full refresh wipes target rows and the auxiliary table for the refreshed flow") { + createScd2Target(s"$catalog.$namespace.target") + + // Run #1: populate target + auxiliary table. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "alice", 5L)) + val ctx1 = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream1.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx1) + assert( + spark.catalog.tableExists(auxTableNameFor("target")), + "Auxiliary table should exist after first run" + ) + + // Run #2 (full refresh): auxiliary table should be dropped by DatasetManager, target + // truncated. The new run brings only id=2 at seq=1. + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((2, "bob", 1L)) + val ctx2 = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream2.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + val updateCtx = TestPipelineUpdateContext( + spark, + ctx2.toDataflowGraph, + storageRoot, + fullRefreshTables = AllTables + ) + updateCtx.pipelineExecution.runPipeline() + updateCtx.pipelineExecution.awaitCompletion() + + // Only id=2 remains, as a single open current record; id=1 from run #1 is wiped. + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(2, "bob", 1L, 1L, null, scd2Meta(1L))) + ) + } + + test("after a full refresh, an event with a sequence below the previous run's " + + "watermark now lands") { + createScd2Target(s"$catalog.$namespace.target") + + // Run #1: delete at seq=10 sets a high watermark in the auxiliary table. + val stream1 = MemoryStream[(Int, String, Long, Boolean)] + stream1.addData((1, "alice", 10L, true)) + val ctx1 = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream1.toDF().toDF("id", "name", "version", "is_delete")), + keys = Seq("id"), + sequencing = functions.col("version"), + deleteCondition = Some(functions.col("is_delete") === true), + columnSelection = Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName("is_delete")) + )), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx1) + + // Run #2 (full refresh): auxiliary table is dropped, watermark reset. seq=5 should + // now land as an open current record. + val stream2 = MemoryStream[(Int, String, Long, Boolean)] + stream2.addData((1, "fresh", 5L, false)) + val ctx2 = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream2.toDF().toDF("id", "name", "version", "is_delete")), + keys = Seq("id"), + sequencing = functions.col("version"), + deleteCondition = Some(functions.col("is_delete") === true), + columnSelection = Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName("is_delete")) + )), + scdType = ScdType.Type2 + )) + } + val updateCtx = TestPipelineUpdateContext( + spark, + ctx2.toDataflowGraph, + storageRoot, + fullRefreshTables = AllTables + ) + updateCtx.pipelineExecution.runPipeline() + updateCtx.pipelineExecution.awaitCompletion() + + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "fresh", 5L, 5L, null, scd2Meta(5L))) + ) + } + + test("selective full refresh wipes only the requested target's auxiliary state") { + createScd2Target(s"$catalog.$namespace.t_a") + createScd2Target(s"$catalog.$namespace.t_b") + + // t_b's run #1 is delete-only so that the state proving the aux was spared lives ONLY in the + // aux: the target stays empty and the aux holds a seq=10 tombstone. In run #2, t_b's seq=5 + // upsert landing closed at 10 is possible only because the selective refresh left t_b's aux + // intact -- target state alone could not supply the seq=10 closure. (An open upsert in run #1 + // would instead route to the target and leave the aux empty, so the assertion would hold even + // if the aux were wiped, attributing the outcome to the wrong state.) + // + // streamA is replaced across runs because t_a is full-refreshed in run #2 (its streaming + // checkpoint is reset by full-refresh, so a fresh source is fine and matches the user-visible + // semantics). streamB is reused across runs because t_b is NOT full-refreshed -- its + // streaming checkpoint must resume against the same MemoryStream instance, otherwise the + // seq=5 assertion below could pass for the wrong reason (the source never produced seq=5 + // in run #2 instead of the aux tombstone shaping it). + val streamA1 = MemoryStream[(Int, String, Long)] + val streamB = MemoryStream[(Int, String, Long, Boolean)] + streamA1.addData((1, "a", 10L)) + streamB.addData((1, "b", 10L, true)) // delete at seq=10: target empty, aux tombstone at 10 + // dfFlowFunc is a TestGraphRegistrationContext method, so it can only be called inside the + // context blocks below; flowB takes the already-built query and adds t_b's delete knobs. + def flowB(query: FlowFunction): AutoCdcFlow = autoCdcFlow( + name = "flow_b", + target = "t_b", + query = query, + keys = Seq("id"), + sequencing = functions.col("version"), + deleteCondition = Some(functions.col("is_delete") === true), + columnSelection = Some(ColumnSelection.ExcludeColumns( + Seq(UnqualifiedColumnName("is_delete")) + )), + scdType = ScdType.Type2 + ) + val ctx1 = new TestGraphRegistrationContext(spark) { + registerTable("t_a", catalog = Some(catalog), database = Some(namespace)) + registerTable("t_b", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "flow_a", + target = "t_a", + query = dfFlowFunc(streamA1.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + registerFlow(flowB(dfFlowFunc(streamB.toDF().toDF("id", "name", "version", "is_delete")))) + } + runPipeline(ctx1) + // Precondition: t_b's run #1 left the target empty with the seq=10 tombstone only in the aux. + checkAnswer(spark.table(s"$catalog.$namespace.t_b"), Seq.empty) + + // Run #2: full refresh ONLY on t_a; t_b's auxiliary state must persist. + val streamA2 = MemoryStream[(Int, String, Long)] + // t_a's aux is wiped, so seq=5 is the only record it has ever seen: a fresh open record. + streamA2.addData((1, "a2", 5L)) + // t_b keeps its aux (seq=10 tombstone). The late seq=5 upsert is woven into history as a + // closed prior record ending at seq=10 -- a closure only the retained aux can supply. + streamB.addData((1, "b2", 5L, false)) + val ctx2 = new TestGraphRegistrationContext(spark) { + registerTable("t_a", catalog = Some(catalog), database = Some(namespace)) + registerTable("t_b", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "flow_a", + target = "t_a", + query = dfFlowFunc(streamA2.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + registerFlow(flowB(dfFlowFunc(streamB.toDF().toDF("id", "name", "version", "is_delete")))) + } + val updateCtx = TestPipelineUpdateContext( + spark, + ctx2.toDataflowGraph, + storageRoot, + fullRefreshTables = SomeTables(Set( + fullyQualifiedIdentifier("t_a", Some(catalog), Some(namespace)) + )) + ) + updateCtx.pipelineExecution.runPipeline() + updateCtx.pipelineExecution.awaitCompletion() + + // t_a: refreshed, so the seq=5 event lands as a fresh open current record. + checkAnswer( + spark.table(s"$catalog.$namespace.t_a"), + Seq(Row(1, "a2", 5L, 5L, null, scd2Meta(5L))) + ) + // t_b: aux retained, so the late seq=5 event is woven in as a closed prior record ending at + // the tombstoned seq=10. With no open successor (the seq=10 event was a delete), the closed + // [5, 10) record is the only visible row. + checkAnswer( + spark.table(s"$catalog.$namespace.t_b"), + Seq(Row(1, "b2", 5L, 5L, 10L, scd2Meta(5L))) + ) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2KeyDriftSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2KeyDriftSuite.scala new file mode 100644 index 0000000000000..43fc60c23bdde --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2KeyDriftSuite.scala @@ -0,0 +1,487 @@ +/* + * 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.pipelines.graph + +import org.apache.spark.sql.classic.DataFrame +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.pipelines.autocdc.{Scd2BatchProcessor, ScdType} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.MetadataBuilder + +/** + * End-to-end tests covering AutoCDC SCD2 key-drift validation: the AutoCDC flow's declared keys + * are validated against the auxiliary table's recorded keys at flow execution-init time. A change + * in keys across runs without a full refresh corrupts the merge semantics; validation detects this + * and fails fast with a structured [[AUTOCDC_INVALID_STATE]] error. The SCD2 analog of + * [[AutoCdcScd1KeyDriftSuite]]. + * + * Key-drift validation itself is SCD-type-agnostic (it compares recorded vs declared key + * (name, dataType) sets), so these mirror the SCD1 cases with SCD2 flows. The tests that + * pre-create a tampered auxiliary table build an SCD2-shaped aux table (full target row schema + * plus the deleted-by-batch-id marker) carrying the SCD2 scd-type property. + * + * Two deliberate departures from a strict one-to-one mirror keep both suites at fourteen tests: + * [[AutoCdcScd1KeyDriftSuite]]'s "AutoCDC key drift validation uses pipeline case sensitivity, not + * session default" has no analog here -- that the validator reads the pipeline's conf rather than + * the ambient session conf is SCD-type-agnostic and already pinned by the SCD1 suite -- and in its + * place this suite adds the SCD_TYPE_DRIFT case, which the SCD1 suite structurally cannot cover. + */ +class AutoCdcScd2KeyDriftSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** + * Properties clause seeding an SCD2 auxiliary table with the scd-type property and the given + * JSON key-column-names array, so drift validation classifies it as SCD2 and reads its keys. + */ + private def scd2AuxProps(keyColumnNamesJson: String): String = + s"TBLPROPERTIES (" + + s"'${AutoCdcAuxiliaryTable.scdTypePropertyKey}' = '${ScdType.Type2.label}', " + + s"'${AutoCdcAuxiliaryTable.keyColumnNamesProperty}' = '$keyColumnNamesJson')" + + test("a pipeline execution that adds a key column to an existing AutoCDC flow triggers " + + "KEY_SCHEMA_DRIFT") { + // Target table carries both candidate key columns up-front so only the AutoCDC `keys` + // declaration differs between the two pipelines. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, region STRING NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1 declares one key (`id`). + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "us", 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "region", "version"), Seq("id"))) + + // Pipeline #2 declares two keys (`region` + `id`) - arity drift. + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((1, "us", 2L)) + val ctx2 = buildPipeline( + "flow_v2", stream2.toDF().toDF("id", "region", "version"), Seq("region", "id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + // `region` is nullable here because Scala `String` is a reference type and the + // [[MemoryStream]] tuple encoder treats reference types as nullable. Only Scala + // primitives (`Int`, `Long`, ...) yield `NOT NULL` columns. + "expectedKeySchema" -> "region STRING,id INT NOT NULL", + "recordedKeySchema" -> "id INT NOT NULL" + ) + ) + } + + test("a pipeline execution that drops a key column from an existing AutoCDC flow triggers " + + "KEY_SCHEMA_DRIFT") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(region STRING NOT NULL, id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1 declares two keys [region, id]. + val stream1 = MemoryStream[(String, Int, Long)] + stream1.addData(("us", 1, 1L)) + runPipeline(buildPipeline( + "flow_v1", stream1.toDF().toDF("region", "id", "version"), Seq("region", "id"))) + + // Pipeline #2 declares only [id] - arity drift. + val stream2 = MemoryStream[(String, Int, Long)] + stream2.addData(("us", 1, 2L)) + val ctx2 = buildPipeline("flow_v2", stream2.toDF().toDF("region", "id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "expectedKeySchema" -> "id INT NOT NULL", + // `region` is nullable here because Scala `String` is a reference type; see the + // analogous comment in the "adds a key column" test above. + "recordedKeySchema" -> "region STRING,id INT NOT NULL" + ) + ) + } + + test("a pipeline execution that swaps a key in an existing AutoCDC flow for a different name " + + "(same arity) triggers KEY_SCHEMA_DRIFT") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, region STRING NOT NULL, country STRING NOT NULL, " + + s"version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1 declares [id, region]. + val stream1 = MemoryStream[(Int, String, String, Long)] + stream1.addData((1, "us", "USA", 1L)) + runPipeline(buildPipeline( + "flow_v1", stream1.toDF().toDF("id", "region", "country", "version"), Seq("id", "region"))) + + // Pipeline #2 declares [id, country] - same arity, different key set. + val stream2 = MemoryStream[(Int, String, String, Long)] + stream2.addData((1, "us", "USA", 2L)) + val ctx2 = buildPipeline( + "flow_v2", stream2.toDF().toDF("id", "region", "country", "version"), Seq("id", "country")) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + // `country` and `region` are nullable here because Scala `String` is a reference type; + // see the analogous comment in the "adds a key column" test above. + "expectedKeySchema" -> "id INT NOT NULL,country STRING", + "recordedKeySchema" -> "id INT NOT NULL,region STRING" + ) + ) + } + + test("a pipeline whose recorded aux key dataType differs from the flow's source dataType " + + "triggers KEY_SCHEMA_DRIFT") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + // Pre-seed an SCD2-shaped aux table whose recorded key `id` is BIGINT, differing from the + // flow's INT source key. + spark.sql( + s"""CREATE TABLE ${auxTableNameFor("target")} """ + + s"""(id BIGINT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl, """ + + s"""${Scd2BatchProcessor.deletedByBatchIdColName} BIGINT) ${scd2AuxProps("[\"id\"]")}""" + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = buildPipeline("flow", stream.toDF().toDF("id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "expectedKeySchema" -> "id INT NOT NULL", + "recordedKeySchema" -> "id BIGINT NOT NULL" + ) + ) + } + + test("a composite key reorder ([a,b] -> [b,a]) does NOT trigger drift validation") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(a INT NOT NULL, b STRING NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1 declares keys [a, b]; pipeline #2 declares the same set reversed [b, a]. Drift + // validation is order-independent, so pipeline #2 must NOT throw. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "x", 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("a", "b", "version"), Seq("a", "b"))) + + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((2, "y", 1L)) + runPipeline(buildPipeline("flow_v2", stream2.toDF().toDF("a", "b", "version"), Seq("b", "a"))) + } + + test("a pipeline execution that changes a key column's nullability or metadata in an " + + "existing AutoCDC flow does NOT trigger drift") { + // Drift validation compares (name, dataType) pairs as a set; nullability and column metadata + // are not part of [[DataType]], so they do not gate semantic equivalence. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1: source carries `id INT NOT NULL` (Scala primitive `Int`), no metadata. + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("id"))) + + // Pipeline #2: source carries `id INT` (nullable, via `Option[Int]`) AND attaches non-empty + // column metadata. Same name and `dataType` as the recorded key, but every [[StructField]] + // aspect outside `dataType` differs. + val stream2 = MemoryStream[(Option[Int], Long)] + stream2.addData((Some(2), 2L)) + val baseDf = stream2.toDF().toDF("id", "version") + val md = new MetadataBuilder() + .putString("description", "primary key") + .build() + val sourceDfWithMetadata = baseDf.select(baseDf("id").as("id", md), baseDf("version")) + runPipeline(buildPipeline("flow_v2", sourceDfWithMetadata, Seq("id"))) + } + + test("a pipeline execution that wraps an existing AutoCDC flow's key in backticks does NOT " + + "trigger drift") { + // Backticks are a SQL-parse syntactic device, not part of the identifier itself. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("id"))) + + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((2, 1L)) + runPipeline(buildPipeline("flow_v2", stream2.toDF().toDF("id", "version"), Seq("`id`"))) + } + + test("a pipeline execution that drops backticks around an existing AutoCDC flow's " + + "previously-backtick-quoted key does NOT trigger drift") { + // The reverse direction: drift validation must be backtick-invariant on both the write side + // (recorded property strips backticks) and the read side (resolver-aware lookup). + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("`id`"))) + + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((2, 1L)) + runPipeline(buildPipeline("flow_v2", stream2.toDF().toDF("id", "version"), Seq("id"))) + } + + test("under spark.sql.caseSensitive = true, an AutoCDC flow whose key differs only in case " + + "from the recorded key triggers KEY_SCHEMA_DRIFT") { + // validateNoKeyColumnDrift uses spark.sessionState.conf.resolver, so its behavior on + // `Id` vs `id` flips with the session conf. Pipeline #1 seeds the aux under the default + // resolver with recorded key `["id"]`; pipeline #2 runs under the case-sensitive resolver + // with key `["Id"]`, which is a distinct identifier there, so drift must fire. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("id"))) + + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((1, 2L)) + val ctx2 = buildPipeline("flow_v2", stream2.toDF().toDF("Id", "version"), Seq("Id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "expectedKeySchema" -> "Id INT NOT NULL", + "recordedKeySchema" -> "id INT NOT NULL" + ) + ) + } + } + + test("under the default (case-insensitive) resolver, an AutoCDC flow whose key differs only " + + "in case from the recorded key does NOT trigger drift") { + // Pairs with the case-sensitive test above: under the default resolver the two identifiers + // are equivalent, so drift validation must accept pipeline #2. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long)] + stream1.addData((1, 1L)) + runPipeline(buildPipeline("flow_v1", stream1.toDF().toDF("id", "version"), Seq("id"))) + + val stream2 = MemoryStream[(Int, Long)] + stream2.addData((1, 2L)) + runPipeline(buildPipeline("flow_v2", stream2.toDF().toDF("id", "version"), Seq("Id"))) + } + + test("a pipeline whose aux table is missing the keyColumnNames property fails with " + + "AUXILIARY_TABLE_PROPERTY_MISSING") { + // Pre-create the aux table directly without the [[keyColumnNamesProperty]] to simulate + // corrupt metadata. Validation must surface a structured AUTOCDC_INVALID_STATE error. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + spark.sql( + s"""CREATE TABLE ${auxTableNameFor("target")} """ + + s"""(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl, """ + + s"""${Scd2BatchProcessor.deletedByBatchIdColName} BIGINT) """ + + s"""TBLPROPERTIES ('${AutoCdcAuxiliaryTable.scdTypePropertyKey}' = """ + + s"""'${ScdType.Type2.label}')""" + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = buildPipeline("flow", stream.toDF().toDF("id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MISSING", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "propertyName" -> AutoCdcAuxiliaryTable.keyColumnNamesProperty + ) + ) + } + + test("a pipeline whose aux table has a malformed keyColumnNames property fails with " + + "AUXILIARY_TABLE_PROPERTY_MALFORMED") { + // Pre-create the aux table directly with a non-JSON-array property value to simulate corrupt + // metadata. Validation must surface a structured AUTOCDC_INVALID_STATE error. + val malformedKeysArray = "not-a-json-array" + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + spark.sql( + s"""CREATE TABLE ${auxTableNameFor("target")} """ + + s"""(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl, """ + + s"""${Scd2BatchProcessor.deletedByBatchIdColName} BIGINT) """ + + scd2AuxProps(malformedKeysArray) + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = buildPipeline("flow", stream.toDF().toDF("id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_PROPERTY_MALFORMED", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "propertyName" -> AutoCdcAuxiliaryTable.keyColumnNamesProperty, + "rawValue" -> malformedKeysArray + ) + ) + } + + test("a pipeline whose aux table records a key absent from its schema fails with " + + "AUXILIARY_TABLE_KEY_COLUMN_MISSING") { + // Pre-create the aux table with the [[keyColumnNamesProperty]] pointing at a column that does + // not exist in the aux schema. Validation must surface a structured AUTOCDC_INVALID_STATE + // error rather than KEY_SCHEMA_DRIFT. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + spark.sql( + s"""CREATE TABLE ${auxTableNameFor("target")} """ + + s"""(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl, """ + + s"""${Scd2BatchProcessor.deletedByBatchIdColName} BIGINT) ${scd2AuxProps("[\"region\"]")}""" + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = buildPipeline("flow", stream.toDF().toDF("id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.AUXILIARY_TABLE_KEY_COLUMN_MISSING", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "keyColumnName" -> "region", + "propertyName" -> AutoCdcAuxiliaryTable.keyColumnNamesProperty + ) + ) + } + + test("a pipeline execution whose recorded SCD type differs from the flow's SCD type triggers " + + "SCD_TYPE_DRIFT") { + // SCD_TYPE_DRIFT is reachable end-to-end from DatasetManager (validateNoScdTypeDrift runs + // right after validateNoKeyColumnDrift when evolving the aux table), but the SCD1 suite + // structurally cannot cover it, so this SCD2 suite is its home. Pre-create an SCD2-shaped aux + // table whose recorded scd-type property is Type1, with matching keys so key-drift validation + // passes and the scd-type check is the one that fires, then run a Type2 flow. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + spark.sql( + s"""CREATE TABLE ${auxTableNameFor("target")} """ + + s"""(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl, """ + + s"""${Scd2BatchProcessor.deletedByBatchIdColName} BIGINT) """ + + s"""TBLPROPERTIES (""" + + s"""'${AutoCdcAuxiliaryTable.scdTypePropertyKey}' = '${ScdType.Type1.label}', """ + + s"""'${AutoCdcAuxiliaryTable.keyColumnNamesProperty}' = '[\"id\"]')""" + ) + + val stream = MemoryStream[(Int, Long)] + stream.addData((1, 1L)) + val ctx = buildPipeline("flow", stream.toDF().toDF("id", "version"), Seq("id")) + + val ex = intercept[RuntimeException] { runPipeline(ctx) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.SCD_TYPE_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("target", Some(catalog), Some(namespace)).unquotedString, + "expectedScdType" -> ScdType.Type2.label, + "recordedScdType" -> ScdType.Type1.label + ) + ) + } + + /** + * Build a single-flow SCD2 pipeline targeting `cat.ns1.target` with the given source DF and key + * column list. Thin wrapper over [[singleAutoCdcFlowPipeline]] since every drift test targets + * the same `target` table. + */ + private def buildPipeline( + flowName: String, + sourceDf: DataFrame, + keys: Seq[String]): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = flowName, + target = "target", + sourceDf = sourceDf, + keys = keys, + sequencing = $"version", + scdType = ScdType.Type2) +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2MultiPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2MultiPipelineSuite.scala new file mode 100644 index 0000000000000..cd161282e7294 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2MultiPipelineSuite.scala @@ -0,0 +1,302 @@ +/* + * 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.pipelines.graph + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * End-to-end tests that exercise interactions between separate SCD Type 2 AutoCDC pipelines (i.e. + * distinct [[DataflowGraph]] / [[TestPipelineUpdateContext]] invocations) sharing the same v2 + * catalog. The SCD2 analog of [[AutoCdcScd1MultiPipelineSuite]]: independent target/auxiliary + * tables per target, downstream reads that ignore the framework columns, a shared target written + * by two pipelines, schema evolution across pipelines, and key-drift rejection. + */ +class AutoCdcScd2MultiPipelineSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + test("two AutoCDC pipelines targeting separate tables maintain independent target and " + + "auxiliary tables") { + // Two distinct target tables created up-front. + spark.sql( + s"CREATE TABLE $catalog.$namespace.t_a " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + spark.sql( + s"CREATE TABLE $catalog.$namespace.t_b " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1 only knows about `t_a`. Its auxiliary table must not affect pipeline #2's `t_b`. + val streamA = MemoryStream[(Int, String, Long)] + streamA.addData((1, "alice", 100L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_a", + target = "t_a", + sourceDf = streamA.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + // Pipeline #2 only knows about `t_b`. Uses a deliberately *lower* sequence to verify the + // watermark from pipeline #1's auxiliary table (seq=100) does not leak into pipeline #2. + val streamB = MemoryStream[(Int, String, Long)] + streamB.addData((9, "bob", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_b", + target = "t_b", + sourceDf = streamB.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + checkAnswer( + spark.table(s"$catalog.$namespace.t_a"), + Seq(Row(1, "alice", 100L, 100L, null, scd2Meta(100L))) + ) + checkAnswer( + spark.table(s"$catalog.$namespace.t_b"), + Seq(Row(9, "bob", 1L, 1L, null, scd2Meta(1L))) + ) + + // Each target has its own auxiliary table; no cross-contamination. + assert(spark.catalog.tableExists(auxTableNameFor("t_a"))) + assert(spark.catalog.tableExists(auxTableNameFor("t_b"))) + } + + test("a downstream pipeline can read an AutoCDC target written by a different pipeline " + + "without observing the framework columns") { + // Pipeline #1 writes into target `src` via AutoCDC. + spark.sql( + s"CREATE TABLE $catalog.$namespace.src " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "alice", 1L), (2, "bob", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "writer", + target = "src", + sourceDf = stream.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + // Pipeline #2 is a regular materialized view that selects the user-data columns from `src` + // (a different graph entirely). It must observe the merged AutoCDC rows and be able to ignore + // the framework columns without them polluting downstream consumers. + val ctxReader = new TestGraphRegistrationContext(spark) { + registerMaterializedView( + "downstream_mv", + query = dfFlowFunc( + spark.read.table(s"$catalog.$namespace.src").select("id", "name", "version") + ) + ) + } + runPipeline(ctxReader) + + checkAnswer( + spark.table(fullyQualifiedIdentifier("downstream_mv").toString), + Seq(Row(1, "alice", 1L), Row(2, "bob", 1L)) + ) + } + + test("two AutoCDC pipelines targeting the same table with identical key and data " + + "schemas merge into a shared target table") { + // Target table is created once up-front; both pipelines target it with the same AutoCDC + // `keys` and the same source-DF data schema. The two pipelines have distinct flow names so + // they own independent streaming checkpoints, but share the target and its auxiliary table. + spark.sql( + s"CREATE TABLE $catalog.$namespace.shared_target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1: inserts rows with id=1 and id=2 at version=1. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "alice", 1L), (2, "bob", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "shared_target", + sourceDf = stream1.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + // Sanity-check pipeline #1's effect before pipeline #2 runs. + checkAnswer( + spark.table(s"$catalog.$namespace.shared_target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L)), + Row(2, "bob", 1L, 1L, null, scd2Meta(1L)) + ) + ) + + // Pipeline #2: updates id=2 (existing key) to a higher sequence and inserts id=3 (new key). + // id=1 is untouched and must survive into the final target unchanged. + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((2, "bob-v2", 2L), (3, "carol", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "shared_target", + sourceDf = stream2.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + // Final target: id=1 untouched; id=2's original record closed at seq=2 with a new open record; + // id=3 freshly inserted by pipeline #2. + checkAnswer( + spark.table(s"$catalog.$namespace.shared_target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L)), + Row(2, "bob", 1L, 1L, 2L, scd2Meta(1L)), + Row(2, "bob-v2", 2L, 2L, null, scd2Meta(2L)), + Row(3, "carol", 1L, 1L, null, scd2Meta(1L)) + ) + ) + + // The auxiliary table for the shared target is itself shared across both pipelines. + assert(spark.catalog.tableExists(auxTableNameFor("shared_target"))) + } + + test("two AutoCDC pipelines targeting the same table with the same key but different " + + "data columns evolve the shared target schema") { + // Target is created up-front with pipeline #1's schema only; pipeline #2 brings a new + // top-level nullable `age` column that the dataset materialization layer is expected to + // schema-merge into the target. + spark.sql( + s"CREATE TABLE $catalog.$namespace.shared_target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1: source DF schema is (id, name, version); inserts id=1 and id=2. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "alice", 1L), (2, "bob", 1L)) + val ctx1 = singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "shared_target", + sourceDf = stream1.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2, + // Both pipelines pin the tracked set to `name`, the one non-key column they share. Left to + // the default (selection-derived) tracking, pipeline #2's extra `age` would also widen the + // tracked set, which is rejected as TRACK_HISTORY_DRIFT (SPARK-58391); the axis under test + // here is the shared target's data-column evolution. + trackHistorySelection = + Option(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name"))))) + runPipeline(ctx1) + + // Sanity-check pipeline #1's state before schema evolution kicks in. + checkAnswer( + spark.table(s"$catalog.$namespace.shared_target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L)), + Row(2, "bob", 1L, 1L, null, scd2Meta(1L)) + ) + ) + + // Pipeline #2: source DF schema is (id, name, age, version). The new nullable `age` column + // should be added to the target by dataset materialization; pipeline #1's untouched id=1 row + // is backfilled to NULL. The `age` column lands after the framework columns in the target. + val stream2 = MemoryStream[(Int, String, Option[Int], Long)] + stream2.addData((2, "bob-v2", Some(25), 2L), (3, "carol", Some(30), 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "shared_target", + sourceDf = stream2.toDF().toDF("id", "name", "age", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2, + trackHistorySelection = + Option(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name")))))) + + checkAnswer( + spark.table(s"$catalog.$namespace.shared_target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L), null), + Row(2, "bob", 1L, 1L, 2L, scd2Meta(1L), null), + Row(2, "bob-v2", 2L, 2L, null, scd2Meta(2L), 25), + Row(3, "carol", 1L, 1L, null, scd2Meta(1L), 30) + ) + ) + + // NOTE: the SCD1 analog of this test additionally re-runs the narrower pipeline #1 against the + // now-wider evolved target. For SCD2 that microbatch-narrower-than-target path is covered + // separately by AutoCdcScd2ColumnEvolutionSuite (SPARK-58418), so it is not duplicated here. + } + + test("a second pipeline targeting an existing AutoCDC table with different keys " + + "fails with KEY_SCHEMA_DRIFT") { + // Target table with both candidate keys present so the second pipeline would otherwise be + // schema-compatible with the first; only the AutoCDC `keys` differ between flows. + spark.sql( + s"CREATE TABLE $catalog.$namespace.shared_target " + + s"(id INT NOT NULL, name STRING NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Pipeline #1: AutoCDC flow keyed on `id`. + val stream1 = MemoryStream[(Int, String, Long)] + stream1.addData((1, "alice", 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "flow_v1", + target = "shared_target", + sourceDf = stream1.toDF().toDF("id", "name", "version"), + keys = Seq("id"), + sequencing = $"version", + scdType = ScdType.Type2)) + + // Pipeline #2: completely separate graph, but targets the same physical `shared_target` + // table with `keys = Seq("name")`. + val stream2 = MemoryStream[(Int, String, Long)] + stream2.addData((2, "alice", 1L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "flow_v2", + target = "shared_target", + sourceDf = stream2.toDF().toDF("id", "name", "version"), + keys = Seq("name"), + sequencing = $"version", + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.KEY_SCHEMA_DRIFT", + sqlState = Some("42000"), + parameters = Map( + "tableName" -> + fullyQualifiedIdentifier("shared_target", Some(catalog), Some(namespace)).unquotedString, + // Pipeline #2's AutoCDC key resolves from the source DF, where `MemoryStream[(Int, String, + // Long)]` produces a nullable StringType for `name`. + "expectedKeySchema" -> "name STRING", + // Pipeline #1 persisted the aux table from a source DF whose `id` was a non-null Scala + // primitive (`Int`), so the recorded key carries `NOT NULL`. + "recordedKeySchema" -> "id INT NOT NULL" + ) + ) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SchemaEvolutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SchemaEvolutionSuite.scala new file mode 100644 index 0000000000000..838a5f8ee2f5b --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SchemaEvolutionSuite.scala @@ -0,0 +1,526 @@ +/* + * 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.pipelines.graph + +import java.sql.Timestamp + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.pipelines.autocdc.{ColumnSelection, ScdType, UnqualifiedColumnName} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests covering SCD Type 2 AutoCDC's interaction with non-key schema evolution across pipeline + * runs. The SCD2 analog of [[AutoCdcScd1SchemaEvolutionSuite]]; documents the supported additive + * cases (new top-level columns, a new field inside an array<struct> element, broadening column + * selection) and the cases that fail loudly (incompatible type changes). + * + * Unlike SCD1, an SCD2 upsert to an existing key does not overwrite the row: it closes the prior + * record and opens a new one, so evolution assertions carry the full interval history. + * + * The additive cases pin `trackHistorySelection` to the columns present in every run. Under default + * tracking the tracked set is derived from the selected non-key columns, so adding a column widens + * it -- rejected as `AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT` (SPARK-58391) because it would + * reinterpret already-reconciled history. Pinning keeps each additive test on the axis it names; + * the drift rejection itself has its own test here. This is an intended SCD2-vs-SCD1 divergence: + * the same additive evolution needs no full refresh under SCD1. + * + * Scope notes -- cases intentionally covered elsewhere rather than duplicated here: + * - The *narrowing* / dropped-column cases (a microbatch narrower than the already-evolved + * target, incl. dropped nested struct/array fields) live in [[AutoCdcScd2ColumnEvolutionSuite]] + * under SPARK-58418, which makes them reconcile correctly. + * - Changing `trackHistorySelection` between runs -- the SCD2-only evolution axis, which decides + * whether an upsert opens a new record -- is also exercised end-to-end in + * [[AutoCdcScd2ColumnEvolutionSuite]], so it is not repeated here. + * + * One SCD1 evolution case has no SCD2 analog and so is absent here by design: "extra columns on + * the target that the AutoCDC flow does not emit are preserved" relies on SCD1's in-place overwrite + * (an SCD2 upsert instead reads unemitted target columns as NULL onto the newly-opened record), and + * there is no SCD2-specific preservation invariant to assert. + */ +class AutoCdcScd2SchemaEvolutionSuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + test("a nullable non-key column merges correctly with mixed NULL and non-NULL values") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, email STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, String, Option[String], Long)] + def buildCtx(): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "email", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + // Run #1: insert with NULL email opens a current record. + stream.addData((1, "alice", None, 1L)) + runPipeline(buildCtx()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", null, 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2: upsert with non-NULL email at higher seq closes the prior record and opens a new one. + stream.addData((1, "alice2", Some("a@x.com"), 2L)) + runPipeline(buildCtx()) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", null, 1L, 1L, 2L, scd2Meta(1L)), + Row(1, "alice2", "a@x.com", 2L, 2L, null, scd2Meta(2L)) + ) + ) + } + + test("widening a non-key column's type between runs fails with " + + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, age INT, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Int, Long)] + stream1.addData((1, 30, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream1.toDF().toDF("id", "age", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2)) + + // Run #2: widen `age` from Int to Long. + val stream2 = MemoryStream[(Int, Long, Long)] + stream2.addData((1, 31L, 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream2.toDF().toDF("id", "age", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE", + sqlState = Some("42825"), + parameters = Map( + "left" -> "\"INT\"", + "right" -> "\"BIGINT\"" + ) + ) + } + + test("narrowing a non-key column's type between runs fails with " + + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, payload BIGINT, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long, Long)] + stream1.addData((1, 100L, 1L)) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream1.toDF().toDF("id", "payload", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2)) + + // Run #2: narrow `payload` from Long (BIGINT) to Int (INT). + val stream2 = MemoryStream[(Int, Int, Long)] + stream2.addData((1, 5, 2L)) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream2.toDF().toDF("id", "payload", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE", + sqlState = Some("42825"), + parameters = Map( + "left" -> "\"BIGINT\"", + "right" -> "\"INT\"" + ) + ) + } + + test("a new top-level nullable column appearing in the source DF between runs is " + + "added to the target") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, String, Option[String], Long)] + def buildCtx(includeEmail: Boolean): TestGraphRegistrationContext = { + val sourceDf = stream.toDF().toDF("id", "name", "email", "version") + val projectedDf = if (includeEmail) sourceDf else sourceDf.drop("email") + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projectedDf, + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2, + // Pin the tracked set to the columns present in both runs. Under default tracking the set + // is selection-derived, so adding `email` would also widen it -- a change that reinterprets + // already-reconciled history and is rejected as TRACK_HISTORY_DRIFT (SPARK-58391). Pinning + // isolates the axis under test here: additive evolution of the target's data columns. + trackHistorySelection = + Option(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name"))))) + } + + // Run #1: source projects (id, name, version). Target schema is unchanged. + stream.addData((1, "alice", None, 1L)) + runPipeline(buildCtx(includeEmail = false)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2: source projects (id, name, email, version) for a new key id=2. mergeSchemas appends + // `email` after the framework columns; the existing id=1 row gets NULL for the new column. + stream.addData((2, "bob", Some("b@x.com"), 2L)) + runPipeline(buildCtx(includeEmail = true)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L), null), + Row(2, "bob", 2L, 2L, null, scd2Meta(2L), "b@x.com") + ) + ) + } + + test("adding a source column under default tracking is rejected as track-history drift") { + // The counterpart to the additive tests above, which pin `trackHistorySelection` precisely to + // avoid this: with default tracking the tracked set is derived from the selected non-key + // columns (see `Scd2BatchProcessor.computeTrackedHistoryColumns`), so adding `email` silently + // widens it from [name, version] to [name, email, version]. That reinterprets which transitions + // open a new historical record and cannot be applied to already-reconciled history, so + // SPARK-58391 rejects it rather than letting the second run write history under different rules + // than the first. This is an intended SCD2-vs-SCD1 divergence: for SCD1 the same additive + // evolution needs no full refresh. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, String, Option[String], Long)] + def buildCtx(includeEmail: Boolean): TestGraphRegistrationContext = { + val sourceDf = stream.toDF().toDF("id", "name", "email", "version") + val projectedDf = if (includeEmail) sourceDf else sourceDf.drop("email") + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projectedDf, + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + } + + // Run #1 records the tracked set as [name] on the auxiliary table. + stream.addData((1, "alice", None, 1L)) + runPipeline(buildCtx(includeEmail = false)) + + // Run #2 would track [email, name]; the recorded set no longer matches. + stream.addData((2, "bob", Some("b@x.com"), 2L)) + val ex = intercept[RuntimeException] { runPipeline(buildCtx(includeEmail = true)) } + checkErrorInPipelineFailure( + failure = ex, + condition = "AUTOCDC_INVALID_STATE.TRACK_HISTORY_DRIFT", + parameters = Map( + "tableName" -> s"$catalog.$namespace.target", + // `version` is tracked too: only the keys and the reserved framework columns are excluded + // from the eligible set, and the sequencing column is neither. + "expectedTrackHistoryColumns" -> "name, email, version", + "recordedTrackHistoryColumns" -> "name, version" + ) + ) + + // The rejection is a pre-write validation: the target still holds only run #1's row, and has + // not gained the `email` column. + val fieldNames = spark.table(s"$catalog.$namespace.target").schema.fieldNames.toSeq + assert(!fieldNames.contains("email"), s"target should not have gained `email`; got $fieldNames") + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", 1L, 1L, null, scd2Meta(1L))) + ) + } + + test("additive target-column evolution extends the SCD2 auxiliary table schema") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + // Shared (id, name, version) stream; run #1 projects away `name`, run #2 keeps it so the + // target (and, unlike SCD1, the aux table -- which mirrors the full target row) gain `name`. + val stream = MemoryStream[(Int, String, Long)] + def buildCtx(includeName: Boolean): TestGraphRegistrationContext = { + val sourceDf = stream.toDF().toDF("id", "name", "version") + val projectedDf = if (includeName) sourceDf else sourceDf.drop("name") + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projectedDf, + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2, + // Pin the tracked set to `version` -- an eligible non-key column (the sequencing column is + // neither a key nor a reserved framework column) present in both runs, so the set stays + // stable while `name` is added. Tracking `name` itself would widen the set when it appears + // in run #2 and trip TRACK_HISTORY_DRIFT (SPARK-58391); see the note in the preceding test. + trackHistorySelection = + Option(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("version"))))) + } + + // Run #1: target is (id, version, framework); aux mirrors it plus the marker. + stream.addData((1, "ignored", 1L)) + runPipeline(buildCtx(includeName = false)) + val auxAfterRun1 = spark.table(auxTableNameFor("target")).schema.fieldNames.toSeq + assert(!auxAfterRun1.contains("name"), + s"aux schema after run #1 should not yet contain `name`; got $auxAfterRun1") + + // Run #2: `name` is added to the target for a new key id=2. The SCD2 aux table mirrors the + // full target row schema, so it gains `name` too (unlike SCD1, whose aux holds only keys + + // metadata and is unaffected by non-key evolution). + stream.addData((2, "bob", 2L)) + runPipeline(buildCtx(includeName = true)) + + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, 1L, 1L, null, scd2Meta(1L), null), + Row(2, 2L, 2L, null, scd2Meta(2L), "bob") + ) + ) + assert(spark.table(auxTableNameFor("target")).schema.fieldNames.contains("name")) + } + + test("a new field added inside an array<struct> element between runs is added to the " + + "target") { + // SCD2 analog of AutoCdcScd1SchemaEvolutionSuite's array<struct> additive case: unlike the + // top-level scalar additions above, this exercises unionByName / mergeSchemas recursing into + // an array element struct. Unlike SCD1's overwrite-in-place, the SCD2 upsert closes the prior + // record (which never saw `vals.element.b.d`, so it reads NULL there) and opens a new one + // carrying the widened value. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(key INT NOT NULL, version BIGINT NOT NULL, " + + s"vals ARRAY<STRUCT<a:INT,b:STRUCT<c:INT>>>, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, Long, Int, Int, Int)] + def buildCtx(includeD: Boolean): TestGraphRegistrationContext = { + val src = stream.toDF().toDF("key", "version", "a", "b_c", "b_d") + val inner = if (includeD) { + functions.struct(functions.col("b_c").as("c"), functions.col("b_d").as("d")) + } else { + functions.struct(functions.col("b_c").as("c")) + } + val projected = src.select( + functions.col("key"), + functions.col("version"), + functions.array( + functions.struct(functions.col("a"), inner.as("b")) + ).as("vals") + ) + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = projected, + keys = Seq("key"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + } + + // Run #1: element struct is (a, b.c); no b.d yet. Opens key=1's current record at version=1. + stream.addData((1, 1L, 1, 1, 99)) + runPipeline(buildCtx(includeD = false)) + + // Run #2 widens the element struct with b.d. The version=2 upsert to key=1 closes its + // version=1 record (which predates b.d, so reads NULL) and opens a new one with b.d=2; the + // new key=3 lands as an open record with the full widened struct. + stream.addData((1, 2L, 1, 1, 2), (3, 1L, 3, 3, 3)) + runPipeline(buildCtx(includeD = true)) + + // Inline-explode flattens the array<struct>; carry the interval bounds to prove the closed + // prior record reads NULL for the newly-added nested field. + checkAnswer( + spark.table(s"$catalog.$namespace.target") + .selectExpr("key", "__START_AT", "__END_AT", "inline(vals) as (a, b)") + .select("key", "__START_AT", "__END_AT", "a", "b.c", "b.d"), + Seq( + Row(1, 1L, 2L, 1, 1, null), + Row(1, 2L, null, 1, 1, 2), + Row(3, 1L, null, 3, 3, 3) + ) + ) + } + + test("broadening the column selection between runs adds the newly-included column to " + + "the target") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, String, String, Long)] + def buildCtx(selection: Option[ColumnSelection]): TestGraphRegistrationContext = + singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream.toDF().toDF("id", "name", "email", "version"), + keys = Seq("id"), + sequencing = functions.col("version"), + columnSelection = selection, + scdType = ScdType.Type2, + // Broadening `columnSelection` would also widen the default (selection-derived) tracked + // set, which is rejected as TRACK_HISTORY_DRIFT (SPARK-58391). Pin it to `name` -- selected + // in both runs -- so this test covers only the column-selection axis. + trackHistorySelection = + Option(ColumnSelection.IncludeColumns(Seq(UnqualifiedColumnName("name"))))) + + // Run #1: only (id, name, version) selected; `email` is dropped before the MERGE. + stream.addData((1, "alice", "ignored", 1L)) + runPipeline(buildCtx(selection = Some(ColumnSelection.IncludeColumns( + Seq("id", "name", "version").map(UnqualifiedColumnName(_)) + )))) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, "alice", 1L, 1L, null, scd2Meta(1L))) + ) + + // Run #2: broaden to no selection for a new key id=2. mergeSchemas adds `email`; the existing + // id=1 row gets NULL, the new row gets the actual value. + stream.addData((2, "bob", "b@x.com", 2L)) + runPipeline(buildCtx(selection = None)) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 1L, 1L, null, scd2Meta(1L), null), + Row(2, "bob", 2L, 2L, null, scd2Meta(2L), "b@x.com") + ) + ) + } + + test("a source DF column whose name differs from the target only by case folds onto the " + + "target column under case-insensitive resolution") { + // SPARK-58517: schema evolution honors `spark.sql.caseSensitive`, so a source `Value` maps onto + // the target's existing `value` instead of evolving the target to carry both spellings. The + // target keeps its own spelling (the merge is left-biased) and no new column appears. + // + // Before that fix the merge ran case-sensitively and the target gained a second, case-differing + // column -- a schema self-inconsistent under the case-insensitive resolver. The breakage + // surfaced later, during microbatch reconciliation rather than at table creation: + // `Scd2ForeachBatchHandler.reconcileMicrobatch` read the two-column target back into its + // affected-rows `unionByName`, where `ResolveUnion`'s case-insensitive duplicate check reported + // COLUMN_ALREADY_EXISTS. + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(key INT NOT NULL, version BIGINT NOT NULL, value STRING, $scd2MetadataDdl)" + ) + + val stream = MemoryStream[(Int, Long, String)] + stream.addData((1, 1L, "alice")) + val df = stream.toDF().toDF("key", "version", "Value") + val ctx = singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = df, + keys = Seq("key"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + runPipeline(ctx) + + // A single `value` column, spelled as the target declared it -- not a second `Value`. + assert( + spark.table(s"$catalog.$namespace.target").schema.fieldNames.count( + _.equalsIgnoreCase("value")) === 1) + assert(spark.table(s"$catalog.$namespace.target").schema.fieldNames.contains("value")) + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq(Row(1, 1L, "alice", 1L, null, scd2Meta(1L))) + ) + } + } + + test("changing a non-key column type from TIMESTAMP to STRING between runs fails with " + + "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") { + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(key INT NOT NULL, version BIGINT NOT NULL, value TIMESTAMP, $scd2MetadataDdl)" + ) + + val stream1 = MemoryStream[(Int, Long, Timestamp)] + stream1.addData((1, 1L, Timestamp.valueOf("2024-01-01 10:00:00"))) + runPipeline(singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream1.toDF().toDF("key", "version", "value"), + keys = Seq("key"), + sequencing = functions.col("version"), + scdType = ScdType.Type2)) + + // Run #2 emits `value` as STRING. mergeSchemas rejects the type change. + val stream2 = MemoryStream[(Int, Long, String)] + stream2.addData((1, 2L, "2024-01-02 11:00:00")) + val ctx2 = singleAutoCdcFlowPipeline( + flowName = "auto_cdc_flow", + target = "target", + sourceDf = stream2.toDF().toDF("key", "version", "value"), + keys = Seq("key"), + sequencing = functions.col("version"), + scdType = ScdType.Type2) + + val ex = intercept[RuntimeException] { runPipeline(ctx2) } + checkErrorInPipelineFailure( + failure = ex, + condition = "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE", + sqlState = Some("42825"), + parameters = Map( + "left" -> "\"TIMESTAMP\"", + "right" -> "\"STRING\"" + ) + ) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SinglePipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SinglePipelineSuite.scala index e413399c4033e..595a3df45e28b 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SinglePipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2SinglePipelineSuite.scala @@ -45,6 +45,8 @@ class AutoCdcScd2SinglePipelineSuite with SharedSparkSession with AutoCdcGraphExecutionTestMixin { + import testImplicits._ + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) @@ -62,8 +64,6 @@ class AutoCdcScd2SinglePipelineSuite } test("SCD2: an upsert lands an open current record in an empty target table") { - val session = spark - import session.implicits._ createScd2Target(s"$catalog.$namespace.target") val stream = MemoryStream[(Int, String, Long)] @@ -91,8 +91,6 @@ class AutoCdcScd2SinglePipelineSuite } test("SCD2: an update to a key closes the prior record and opens a new one") { - val session = spark - import session.implicits._ createScd2Target(s"$catalog.$namespace.target") val stream = MemoryStream[(Int, String, Long)] @@ -123,8 +121,6 @@ class AutoCdcScd2SinglePipelineSuite } test("SCD2: a delete closes the current record with no open record remaining") { - val session = spark - import session.implicits._ // Target omits `is_delete`: the source carries it as a control column driving the delete // condition, and it is excluded from the target projection. createScd2Target(s"$catalog.$namespace.target") @@ -158,8 +154,6 @@ class AutoCdcScd2SinglePipelineSuite } test("SCD2: the auxiliary table is materialized for the target") { - val session = spark - import session.implicits._ createScd2Target(s"$catalog.$namespace.target") val stream = MemoryStream[(Int, String, Long)] diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2TargetTableDurabilitySuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2TargetTableDurabilitySuite.scala new file mode 100644 index 0000000000000..ff4e903117406 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/AutoCdcScd2TargetTableDurabilitySuite.scala @@ -0,0 +1,285 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.pipelines.graph + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.functions +import org.apache.spark.sql.pipelines.autocdc.{AutoCdcReservedNames, Scd2BatchProcessor, ScdType} +import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext} +import org.apache.spark.sql.test.SharedSparkSession + +/** + * Tests covering SCD Type 2 AutoCDC's behavior when the target table is pre-populated by something + * other than a prior AutoCDC run: hand-loaded open ("current") records and a target created + * without the framework columns. These verify AutoCDC interoperates gracefully with users who + * hand-populate the target. The SCD2 analog of [[AutoCdcScd1TargetTableDurabilitySuite]]. + */ +class AutoCdcScd2TargetTableDurabilitySuite + extends ExecutionTest + with SharedSparkSession + with AutoCdcGraphExecutionTestMixin { + + import testImplicits._ + + /** The SCD2 target's `_cdc_metadata` struct value for a given recordStartAt. */ + private def scd2Meta(recordStartAt: Long): Row = Row(recordStartAt) + + /** Create an SCD2 target with user columns `(id, name, version)` plus the framework columns. */ + private def createScd2Target(table: String): Unit = { + spark.sql( + s"CREATE TABLE $table (" + + s"id INT NOT NULL, name STRING, version BIGINT NOT NULL, $scd2MetadataDdl)" + ) + } + + /** + * Insert a pre-existing open ("current") SCD2 record into a target table, as if a previous + * AutoCDC run had opened it at sequencing version [[sequence]]: `__START_AT` = `sequence`, + * `__END_AT` = NULL (still active), and `_cdc_metadata.__RECORD_START_AT` = `sequence`. + * + * @param table Fully-qualified table name (catalog.schema.table). + * @param colValues Comma-separated SQL literals for the user-defined columns, in declared + * order, excluding the trailing framework columns. + * @param sequence Value to seed the interval start and the record-start-at with. + */ + private def insertPreloadedCurrentRecord( + table: String, colValues: String, sequence: Long): Unit = { + val recordStartAt = Scd2BatchProcessor.recordStartAtFieldName + spark.sql( + s"INSERT INTO $table SELECT $colValues, " + + s"CAST($sequence AS BIGINT), CAST(NULL AS BIGINT), " + + s"named_struct('$recordStartAt', CAST($sequence AS BIGINT))" + ) + } + + /** + * Insert a pre-existing closed (historical) SCD2 record into a target table, as if a previous + * AutoCDC run had opened it at [[startAt]] and later closed it at [[endAt]]: `__START_AT` = + * `startAt`, `__END_AT` = `endAt` (no longer active), and `_cdc_metadata.__RECORD_START_AT` = + * `startAt`. + * + * @param table Fully-qualified table name (catalog.schema.table). + * @param colValues Comma-separated SQL literals for the user-defined columns, in declared + * order, excluding the trailing framework columns. + * @param startAt Interval start (and record-start-at) of the pre-existing record. + * @param endAt Interval end (exclusive) at which the pre-existing record was closed. + */ + private def insertPreloadedClosedRecord( + table: String, colValues: String, startAt: Long, endAt: Long): Unit = { + val recordStartAt = Scd2BatchProcessor.recordStartAtFieldName + spark.sql( + s"INSERT INTO $table SELECT $colValues, " + + s"CAST($startAt AS BIGINT), CAST($endAt AS BIGINT), " + + s"named_struct('$recordStartAt', CAST($startAt AS BIGINT))" + ) + } + + test("pre-loaded current record: a higher-sequence upsert closes it and opens a new record") { + createScd2Target(s"$catalog.$namespace.target") + insertPreloadedCurrentRecord(s"$catalog.$namespace.target", "1, 'alice', 5", 5L) + + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "alicia", 10L)) // > pre-existing seq=5 -> closes it, opens a new record + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx) + + // The pre-existing record is closed at the incoming event's sequence; the new value is open. + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 5L, 5L, 10L, scd2Meta(5L)), + Row(1, "alicia", 10L, 10L, null, scd2Meta(10L)) + ) + ) + } + + test("pre-loaded current record: a lower-sequence upsert is woven in as a closed prior record") { + createScd2Target(s"$catalog.$namespace.target") + insertPreloadedCurrentRecord(s"$catalog.$namespace.target", "1, 'alice', 10", 10L) + + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "early", 5L)) // < pre-existing seq=10 -> closed prior record ending at 10 + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx) + + // Unlike SCD1, the late lower-sequence event is not suppressed: it becomes a closed prior + // record ending where the pre-existing record starts, which stays the open current record. + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "early", 5L, 5L, 10L, scd2Meta(5L)), + Row(1, "alice", 10L, 10L, null, scd2Meta(10L)) + ) + ) + } + + test("pre-loaded closed record: an event landing inside its interval bisects it") { + // The interop shape unique to SCD2: a hand-loaded *closed* record -- a target row with no aux + // counterpart -- split by an event landing inside its interval. Pre-load [5, 20) and feed an + // event at seq=10. The pre-existing record is closed early, at 10, and the incoming event + // takes over the remainder of the span, [10, 20), so the two records partition the original + // interval with no row left open. + createScd2Target(s"$catalog.$namespace.target") + insertPreloadedClosedRecord(s"$catalog.$namespace.target", "1, 'alice', 5", startAt = 5L, + endAt = 20L) + + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "mid", 10L)) // 5 < 10 < 20 -> bisects the pre-existing [5, 20) record + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx) + + // The pre-existing record is split at the incoming event's sequence: its value carries into + // [5, 10), the incoming "mid" opens [10, 20), and no row remains open (endAt=20 was the + // pre-existing closure). + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 5L, 5L, 10L, scd2Meta(5L)), + Row(1, "mid", 10L, 10L, 20L, scd2Meta(10L)) + ) + ) + } + + test("pre-loaded target rows merge correctly on the first AutoCDC run, and the " + + "auxiliary table is created lazily") { + // Target was populated by some external process; this is the first AutoCDC run. + createScd2Target(s"$catalog.$namespace.target") + insertPreloadedCurrentRecord(s"$catalog.$namespace.target", "1, 'alice', 1", 1L) + + assert( + !spark.catalog.tableExists(auxTableNameFor("target")), + "Auxiliary table should not exist before the first AutoCDC run" + ) + + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "bob", 2L)) + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx) + + // seq=2 > pre-existing seq=1, so the pre-existing record closes at 2 and "bob" opens. + checkAnswer( + spark.table(s"$catalog.$namespace.target"), + Seq( + Row(1, "alice", 1L, 1L, 2L, scd2Meta(1L)), + Row(1, "bob", 2L, 2L, null, scd2Meta(2L)) + ) + ) + assert( + spark.catalog.tableExists(auxTableNameFor("target")), + "Auxiliary table should be created lazily on the first AutoCDC run" + ) + } + + test("a target table created without the framework columns gets them " + + "auto-added on the first AutoCDC run") { + // User creates the target without the AutoCDC framework columns. DatasetManager evolves the + // existing table schema by merging it with the AutoCdcMergeFlow's output schema, which + // includes __START_AT / __END_AT and the metadata column. The first run therefore proceeds + // normally, and subsequent reads see the framework columns alongside the user's data columns. + spark.sql( + s"CREATE TABLE $catalog.$namespace.target " + + s"(id INT NOT NULL, name STRING, version BIGINT NOT NULL)" + ) + + val stream = MemoryStream[(Int, String, Long)] + stream.addData((1, "alice", 1L)) + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable("target", catalog = Some(catalog), database = Some(namespace)) + registerFlow(autoCdcFlow( + name = "auto_cdc_flow", + target = "target", + query = dfFlowFunc(stream.toDF().toDF("id", "name", "version")), + keys = Seq("id"), + sequencing = functions.col("version"), + scdType = ScdType.Type2 + )) + } + runPipeline(ctx) + + val schema = spark.table(s"$catalog.$namespace.target").schema + Seq( + Scd2BatchProcessor.startAtColName, + Scd2BatchProcessor.endAtColName, + AutoCdcReservedNames.cdcMetadataColName + ).foreach { col => + assert( + schema.fieldNames.contains(col), + s"Target must have $col after first AutoCDC run; got ${schema.fieldNames.toSeq}" + ) + } + // Schema evolution appends the framework columns after the user columns in the flow's output + // order (__START_AT, __END_AT, then the metadata column -- the same order scd2MetadataDdl + // declares for a pre-created target). Assert by name so the row matches regardless of the + // physical column order. + checkAnswer( + spark.table(s"$catalog.$namespace.target").select( + "id", "name", "version", + Scd2BatchProcessor.startAtColName, + Scd2BatchProcessor.endAtColName, + AutoCdcReservedNames.cdcMetadataColName + ), + Seq(Row(1, "alice", 1L, 1L, null, scd2Meta(1L))) + ) + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectInvalidPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectInvalidPipelineSuite.scala index 62680ebd3835f..172d8354932b8 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectInvalidPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectInvalidPipelineSuite.scala @@ -31,6 +31,9 @@ import org.apache.spark.sql.types.{IntegerType, StructType} */ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { + private def validateGraph(graph: DataflowGraph): DataflowGraph = + graph.validate(spark.sessionState.conf.caseSensitiveAnalysis) + test("Missing source") { class P extends TestGraphRegistrationContext(spark) { registerPersistedView("b", query = readFlowFunc("a")) @@ -39,7 +42,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { val dfg = new P().resolveToDataflowGraph() assert(!dfg.resolved, "Pipeline should not have resolved properly") val ex = intercept[UnresolvedPipelineException] { - dfg.validate() + validateGraph(dfg) } assert(ex.getMessage.contains("Failed to resolve flows in the pipeline")) assertAnalysisException( @@ -64,7 +67,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { val dfg = new P().resolveToDataflowGraph() assert(!dfg.resolved, "Pipeline should not have resolved properly") val ex = intercept[UnresolvedPipelineException] { - dfg.validate() + validateGraph(dfg) } assert(ex.getMessage.contains("Failed to resolve flows in the pipeline")) assert( @@ -141,7 +144,9 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerFlow("a", "a_2", sqlFlowFunc(spark, "SELECT non_existent_col FROM RANGE(5)")) registerTable("b", query = Option(readFlowFunc("a"))) } - val ex = intercept[UnresolvedPipelineException] { new P().resolveToDataflowGraph().validate() } + val ex = intercept[UnresolvedPipelineException] { + validateGraph(new P().resolveToDataflowGraph()) + } assert(ex.directFailures.keySet == Set(fullyQualifiedIdentifier("a_2"))) assert(ex.downstreamFailures.keySet == Set(fullyQualifiedIdentifier("b"))) @@ -158,7 +163,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { val dfg = new P().resolveToDataflowGraph() val ex = intercept[UnresolvedPipelineException] { - dfg.validate() + validateGraph(dfg) }.directFailures(fullyQualifiedIdentifier("b")).getMessage verifyUnresolveColumnError(ex, "x", Seq("z")) } @@ -175,7 +180,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { val dfg = new P().resolveToDataflowGraph() val ex = intercept[UnresolvedPipelineException] { - dfg.validate() + validateGraph(dfg) } assert( ex.directFailures(fullyQualifiedIdentifier("c")) @@ -200,7 +205,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { val dfg = new P().resolveToDataflowGraph() assert(!dfg.resolved) val ex = intercept[UnresolvedPipelineException] { - dfg.validate() + validateGraph(dfg) } assert( ex.directFailures(fullyQualifiedIdentifier("c")) @@ -217,7 +222,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerPersistedView("a", query = readFlowFunc("a")) } val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } assert(e.upstreamDataset == fullyQualifiedIdentifier("a")) assert(e.downstreamTable == fullyQualifiedIdentifier("a")) @@ -229,7 +234,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerPersistedView("b", query = readFlowFunc("a")) } val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } val cycle = Set( fullyQualifiedIdentifier("a"), @@ -260,7 +265,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { fullyQualifiedIdentifier("d") ) val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } assert(e.upstreamDataset != e.downstreamTable) assert(cycle.contains(e.upstreamDataset)) @@ -287,7 +292,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { fullyQualifiedIdentifier("d") ) val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } assert(e.upstreamDataset != e.downstreamTable) assert(cycle.contains(e.upstreamDataset)) @@ -313,7 +318,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { fullyQualifiedIdentifier("d") ) val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } assert(e.upstreamDataset != e.downstreamTable) assert(cycle.contains(e.upstreamDataset)) @@ -340,7 +345,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { fullyQualifiedIdentifier("d") ) val e = intercept[CircularDependencyException] { - new P().resolveToDataflowGraph().validate() + validateGraph(new P().resolveToDataflowGraph()) } assert(e.upstreamDataset != e.downstreamTable) assert(cycle.contains(e.upstreamDataset)) @@ -408,7 +413,9 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerPersistedView("a", query = dfFlowFunc(Seq(1).toDF())) registerTable("b", query = Option(readStreamFlowFunc("a"))) } - val ex = intercept[UnresolvedPipelineException] { p.resolveToDataflowGraph().validate() } + val ex = intercept[UnresolvedPipelineException] { + validateGraph(p.resolveToDataflowGraph()) + } assert( ex.directFailures(fullyQualifiedIdentifier("b")) .getMessage @@ -429,7 +436,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerPersistedView("a", query = dfFlowFunc(mem.toDF())) registerTable("b", query = Option(readFlowFunc("a"))) } - val ex = intercept[UnresolvedPipelineException] { p.resolveToDataflowGraph().validate() } + val ex = intercept[UnresolvedPipelineException] { validateGraph(p.resolveToDataflowGraph()) } assert( ex.directFailures(fullyQualifiedIdentifier("b")) .getMessage @@ -449,7 +456,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -471,7 +478,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -499,7 +506,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -522,7 +529,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { specifiedSchema = Option(new StructType().add("x", IntegerType)) ) }.resolveToDataflowGraph() - val ex1 = intercept[AnalysisException] { graph1.validate() } + val ex1 = intercept[AnalysisException] { validateGraph(graph1) } assert( ex1.getMessage.contains( s"'${fullyQualifiedIdentifier("a").unquotedString}' " + @@ -535,7 +542,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { registerTable("a", specifiedSchema = Option(new StructType().add("x", IntegerType))) registerFlow("a", "a", query = dfFlowFunc(Seq(true, false).toDF("x")), once = true) }.resolveToDataflowGraph() - val ex2 = intercept[AnalysisException] { graph2.validate() } + val ex2 = intercept[AnalysisException] { validateGraph(graph2) } assert( ex2.getMessage.contains( s"'${fullyQualifiedIdentifier("a").unquotedString}' " + @@ -592,7 +599,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -642,7 +649,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -704,7 +711,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( @@ -749,7 +756,7 @@ class ConnectInvalidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() val ex = intercept[AnalysisException] { - graph.validate() + validateGraph(graph) } checkError( diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala index ea6d3202ba868..b90ccbbde7ce5 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala @@ -504,7 +504,7 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession { registerFlow("sink_a", "sink_flow", query = readStreamFlowFunc("a")) } val g = P.resolveToDataflowGraph() - g.validate() + g.validate(spark.sessionState.conf.caseSensitiveAnalysis) assert(g.resolved) assert(g.sink(TableIdentifier("sink_a")).isInstanceOf[Sink]) val sink = g.sink(TableIdentifier("sink_a")) @@ -751,7 +751,7 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession { }.resolveToDataflowGraph() assert(!unresolved.resolved, "case-sensitive consumer flow should fail to resolve") val ex = intercept[UnresolvedPipelineException] { - unresolved.validate() + unresolved.validate(spark.sessionState.conf.caseSensitiveAnalysis) } assertAnalysisException( ex.directFailures(fullyQualifiedIdentifier("consumer")), diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/GraphExecutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/GraphExecutionSuite.scala new file mode 100644 index 0000000000000..f3d7108e8c225 --- /dev/null +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/GraphExecutionSuite.scala @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.pipelines.graph + +import org.apache.spark.SparkFunSuite + +/** + * Unit tests for the flow-execution decision logic: `PipelinesErrors.streamingSourcesChanged` and + * `GraphExecution.determineFlowExecutionActionFromError`, the single source of truth for whether a + * failed flow is retried. + * + * These tests are intentionally session-less - both entry points are pure functions of a + * `Throwable` and two counters. End-to-end behavior (event levels, retry counts, and the terminal + * run state a source change produces) is covered by `TriggeredGraphExecutionSuite`. + */ +class GraphExecutionSuite extends SparkFunSuite { + + /** The message Structured Streaming asserts with when a stream's source set changes. */ + private def sourceSetChangeMessage: String = + "There are [2] sources in the checkpoint offsets and now there are [3] sources " + + "requested by the query. Cannot continue." + + test("streamingSourcesChanged matches only a streaming source-set change error") { + // Structured Streaming raises a bare AssertionError; it is usually wrapped in another + // exception by the time the pipeline sees it, so the whole cause chain has to be checked. + val sourceChange = new AssertionError(sourceSetChangeMessage) + assert(PipelinesErrors.streamingSourcesChanged(sourceChange)) + assert(PipelinesErrors.streamingSourcesChanged(new RuntimeException("wrapper", sourceChange))) + // Unrelated errors - including an AssertionError with a different message - must not match. + assert(!PipelinesErrors.streamingSourcesChanged(new RuntimeException("boom"))) + assert(!PipelinesErrors.streamingSourcesChanged(new AssertionError("a different assertion"))) + } + + test("determineFlowExecutionActionFromError stops on a source change before checking retries") { + val sourceChange = + new RuntimeException("stream failed", new AssertionError(sourceSetChangeMessage)) + + // The source-change check comes before the retry budget, so the flow stops even with retries + // left, and the run-level reason names the source change rather than an exhausted budget. + GraphExecution.determineFlowExecutionActionFromError( + ex = sourceChange, + flowDisplayName = "flow_a", + currentNumTries = 1, + maxAllowedRetries = 3) match { + case GraphExecution.StopFlowExecution(reason) => + assert(reason.failureMessage.contains("streaming sources added or removed")) + assert( + reason.runTerminationReason == + StreamingSourcesChangedFailure("flow_a", Some(sourceChange))) + case other => fail(s"expected StopFlowExecution, got $other") + } + } + + test("determineFlowExecutionActionFromError retries other errors until budget is exhausted") { + val transient = new RuntimeException("transient") + // Retries remaining -> retry. + assert( + GraphExecution.determineFlowExecutionActionFromError( + ex = transient, flowDisplayName = "flow_a", currentNumTries = 1, maxAllowedRetries = 3) == + GraphExecution.RetryFlowExecution) + // Budget exhausted -> stop as max-retries-exceeded, not as a source change. + GraphExecution.determineFlowExecutionActionFromError( + ex = transient, + flowDisplayName = "flow_a", + currentNumTries = 4, + maxAllowedRetries = 3) match { + case GraphExecution.StopFlowExecution(reason) => + assert(!reason.failureMessage.contains("streaming sources added or removed")) + assert( + reason.runTerminationReason == + QueryExecutionFailure("flow_a", maxRetries = 3, Some(transient))) + case other => fail(s"expected StopFlowExecution, got $other") + } + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/MaterializeTablesSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/MaterializeTablesSuite.scala index 46e2d6d9ae631..916169c4de2d9 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/MaterializeTablesSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/MaterializeTablesSuite.scala @@ -32,6 +32,7 @@ import org.apache.spark.sql.connector.catalog.{ } import org.apache.spark.sql.connector.expressions.{ClusterByTransform, Expressions, FieldReference} import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.pipelines.graph.DatasetManager.TableMaterializationException import org.apache.spark.sql.pipelines.utils.{BaseCoreExecutionTest, TestGraphRegistrationContext} import org.apache.spark.sql.test.SharedSparkSession @@ -565,7 +566,7 @@ abstract class MaterializeTablesSuite extends BaseCoreExecutionTest { val graph1 = new TestGraphRegistrationContext(spark) { registerTable("a", query = Option(dfFlowFunc(spark.readStream.format("rate").load()))) - }.resolveToDataflowGraph().validate() + }.resolveToDataflowGraph().validate(spark.sessionState.conf.caseSensitiveAnalysis) materializeGraph(graph1, storageRoot = storageRoot) } @@ -638,7 +639,7 @@ abstract class MaterializeTablesSuite extends BaseCoreExecutionTest { new TestGraphRegistrationContext(spark) { registerView("a", query = dfFlowFunc(streamInts.toDF())) registerTable("b", query = Option(sqlFlowFunc(spark, "SELECT value AS x FROM STREAM a"))) - }.resolveToDataflowGraph().validate() + }.resolveToDataflowGraph().validate(spark.sessionState.conf.caseSensitiveAnalysis) val (refreshSelection, fullRefreshSelection) = if (isFullRefresh) { (NoTables, AllTables) @@ -668,7 +669,7 @@ abstract class MaterializeTablesSuite extends BaseCoreExecutionTest { new TestGraphRegistrationContext(spark) { registerView("a", query = dfFlowFunc(streamInts.toDF())) registerTable("b", query = Option(sqlFlowFunc(spark, "SELECT value AS y FROM STREAM a"))) - }.resolveToDataflowGraph().validate(), + }.resolveToDataflowGraph().validate(spark.sessionState.conf.caseSensitiveAnalysis), contextOpt = updateContextOpt, storageRoot = storageRoot ) @@ -1214,6 +1215,401 @@ abstract class MaterializeTablesSuite extends BaseCoreExecutionTest { } } + test("SPARK-58517: re-materializing with a case-only column difference is a no-op under " + + "case-insensitive resolution") { + withRecordingCatalog { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + // Create the table with `value`, then re-materialize with the same column cased as `Value`. + // Under case-insensitive resolution these are the same column, so schema evolution must + // fold `Value` onto the existing `value`: no alterTable, and the persisted column keeps its + // original name/case. Before SPARK-58517 the case-sensitive merge instead added a second + // `Value` column, corrupting the table. + materializeStreamingTable( + "t", new StructType().add("id", IntegerType).add("value", StringType), Map.empty) + assert(recordingCatalog.recordedAlters.isEmpty) + + materializeStreamingTable( + "t", new StructType().add("id", IntegerType).add("Value", StringType), Map.empty) + assert(recordingCatalog.recordedAlters.isEmpty, + s"expected no alter, got: ${recordingCatalog.recordedAlters}") + + assert( + loadTableFromRecordingCatalog("t").columns() sameElements + CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("value", StringType) + ), + "the persisted schema should keep the original `value` column, not gain a `Value` column" + ) + } + } + } + + test("SPARK-58517: multi-flow schema inference folds case-only column differences under " + + "case-insensitive resolution") { + withRecordingCatalog { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + // Two append flows write to the same streaming table, one emitting `value` and the other + // `Value`. The target schema is INFERRED by merging the flows' schemas, which happens + // before the evolveTable path runs -- so inference must honor case-insensitivity too, + // otherwise the table is created with both columns and the engine's own resolver cannot + // disambiguate them. + val df1 = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(1, "a"))), + new StructType().add("id", IntegerType).add("value", StringType)) + val df2 = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(2, "b"))), + new StructType().add("id", IntegerType).add("Value", StringType)) + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable( + "t", + catalog = Option(recordingCatalogName), + database = Option(recordingNamespace)) + registerFlow( + "t", "f1", dfFlowFunc(df1), + catalog = Option(recordingCatalogName), database = Option(recordingNamespace)) + registerFlow( + "t", "f2", dfFlowFunc(df2), + catalog = Option(recordingCatalogName), database = Option(recordingNamespace)) + } + + val graph = ctx.resolveToDataflowGraph() + val inferredSchemas = graph.inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis) + val (_, inferred) = inferredSchemas.head + // The two spellings fold into a single column, and the lowest flow identifier supplies the + // surviving spelling: `f1` sorts before `f2`, so `value` wins. + assert(inferred.fieldNames.toSeq === Seq("id", "value")) + } + } + } + + test("SPARK-58517: a case-only fold picks the same spelling for the materialized table and for " + + "downstream resolution, in either flow declaration order") { + // The table's schema is derived twice from the same flows, by two different callers: the graph + // materializes the table from `inferSchemas`, while downstream flows resolve against the + // `VirtualTableInput` schema, whose `availableFlows` is in declaration order, not sorted. Both + // go through `inferSchemaFromFlows`, which merges in sorted flow identifier order, so the + // surviving spelling is the same on both paths and does not depend on declaration order. Were + // the two to disagree, the downstream view would persist a column spelled differently from the + // source column it selects, and -- because `diffSchemas` keys on exact names -- reordering the + // flow definitions would turn the next refresh into a drop-then-add of that column. + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + // `f_eu` sorts before `f_us`, so `value` is the surviving spelling in both orders. + def graphWithFlowsDeclared(usFirst: Boolean): DataflowGraph = { + val usDf = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(1, "a"))), + new StructType().add("id", IntegerType).add("Value", StringType)) + val euDf = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(2, "b"))), + new StructType().add("id", IntegerType).add("value", StringType)) + + new TestGraphRegistrationContext(spark) { + registerTable("events") + val registerUs = () => registerFlow("events", "f_us", dfFlowFunc(usDf)) + val registerEu = () => registerFlow("events", "f_eu", dfFlowFunc(euDf)) + if (usFirst) { + registerUs() + registerEu() + } else { + registerEu() + registerUs() + } + // Reads the table, so it resolves against the VirtualTableInput schema. + registerMaterializedView( + "events_summary", + query = sqlFlowFunc(spark, "SELECT id, value FROM events")) + }.resolveToDataflowGraph() + } + + Seq(true, false).foreach { usFirst => + val graph = graphWithFlowsDeclared(usFirst) + val sessionCaseSensitive = spark.sessionState.conf.caseSensitiveAnalysis + val eventsIdentifier = fullyQualifiedIdentifier("events") + + val materializedSchema = graph.inferSchemas(sessionCaseSensitive)(eventsIdentifier) + assert( + materializedSchema.fieldNames.toSeq === Seq("id", "value"), + s"materialized schema for usFirst=$usFirst") + + // The downstream view's own schema reflects what it resolved against upstream: Spark takes + // the resolved attribute's name, so a `Value` upstream would surface here as `Value`. + val summarySchema = graph + .inferSchemas(sessionCaseSensitive)(fullyQualifiedIdentifier("events_summary")) + assert( + summarySchema.fieldNames.toSeq === Seq("id", "value"), + s"downstream schema for usFirst=$usFirst") + } + } + } + + test("multi-flow schema inference keeps case-only column differences distinct under " + + "case-sensitive resolution") { + withRecordingCatalog { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + // The case-sensitive control: `value` and `Value` are distinct columns, so inference + // contributes both. + val df1 = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(1, "a"))), + new StructType().add("id", IntegerType).add("value", StringType)) + val df2 = spark.createDataFrame( + spark.sparkContext.parallelize(Seq(Row(2, "b"))), + new StructType().add("id", IntegerType).add("Value", StringType)) + + val ctx = new TestGraphRegistrationContext(spark) { + registerTable( + "t", + catalog = Option(recordingCatalogName), + database = Option(recordingNamespace)) + registerFlow( + "t", "f1", dfFlowFunc(df1), + catalog = Option(recordingCatalogName), database = Option(recordingNamespace)) + registerFlow( + "t", "f2", dfFlowFunc(df2), + catalog = Option(recordingCatalogName), database = Option(recordingNamespace)) + } + + val graph = ctx.resolveToDataflowGraph() + val inferred = graph.inferSchemas( + spark.sessionState.conf.caseSensitiveAnalysis).values.head + // Both spellings survive as distinct columns in sorted flow identifier order. + assert(inferred.fieldNames.toSeq === Seq("id", "value", "Value")) + } + } + } + + test("SPARK-58517: a materialized view's case-only column rename is applied under " + + "case-insensitive resolution") { + // The non-merging path: for a materialized view `targetSchema` is the run's declared schema + // as-is (no merge with the persisted schema), so a case-only rename must remain visible to + // `diffSchemas` as a drop-then-add. Case-insensitive matching here would emit no change at all + // and freeze the persisted spelling forever -- the table would permanently disagree with its + // definition, with no error pointing at the discrepancy, and a colleague materializing the same + // definition against a fresh table would get the declared casing instead. + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + materializeGraph( + new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, 2L)).toDF("id", "total"))) + registerMaterializedView("mv", query = sqlFlowFunc(spark, "SELECT id, total FROM src")) + }.resolveToDataflowGraph(), + storageRoot = storageRoot + ) + + val catalog = spark.sessionState.catalogManager.currentCatalog.asInstanceOf[TableCatalog] + val identifier = Identifier.of(Array(TestGraphRegistrationContext.DEFAULT_DATABASE), "mv") + assert( + catalog.loadTable(identifier).columns() sameElements CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("total", LongType)) + ) + + // Re-materialize with the column cased as `Total`. The table must follow the definition. + materializeGraph( + new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, 2L)).toDF("id", "total"))) + registerMaterializedView( + "mv", query = sqlFlowFunc(spark, "SELECT id, total AS Total FROM src")) + }.resolveToDataflowGraph(), + storageRoot = storageRoot + ) + assert( + catalog.loadTable(identifier).columns() sameElements CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("Total", LongType)), + "the materialized view should adopt the declared `Total` casing, not keep `total`" + ) + } + } + + test("SPARK-58517: a full-refreshed streaming table's case-only column rename is applied " + + "under case-insensitive resolution") { + // The streaming-table analog of the materialized-view case above: a full refresh also takes + // `targetSchema` as the declared schema without merging, so the same case-only rename must be + // applied rather than silently ignored. + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val graph = materializeGraph( + new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, 2L)).toDF("id", "total"))) + registerTable("st", query = Option(sqlFlowFunc(spark, "SELECT id, total FROM src"))) + }.resolveToDataflowGraph(), + storageRoot = storageRoot + ) + + val catalog = spark.sessionState.catalogManager.currentCatalog.asInstanceOf[TableCatalog] + val identifier = Identifier.of(Array(TestGraphRegistrationContext.DEFAULT_DATABASE), "st") + assert( + catalog.loadTable(identifier).columns() sameElements CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("total", LongType)) + ) + + val renamedGraph = + new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, 2L)).toDF("id", "total"))) + registerTable( + "st", query = Option(sqlFlowFunc(spark, "SELECT id, total AS Total FROM src"))) + }.resolveToDataflowGraph() + + materializeGraph( + renamedGraph, + contextOpt = Option( + TestPipelineUpdateContext( + spark = spark, + unresolvedGraph = graph, + refreshTables = NoTables, + fullRefreshTables = AllTables, + storageRoot = storageRoot + ) + ), + storageRoot = storageRoot + ) + assert( + catalog.loadTable(identifier).columns() sameElements CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("Total", LongType)), + "a full-refreshed streaming table should adopt the declared `Total` casing" + ) + } + } + + test("re-materializing with a case-only column difference adds a column under case-sensitive " + + "resolution") { + withRecordingCatalog { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + // The case-sensitive counterpart: `value` and `Value` are distinct, so `Value` is added. + materializeStreamingTable( + "t", new StructType().add("id", IntegerType).add("value", StringType), Map.empty) + assert(recordingCatalog.recordedAlters.isEmpty) + + materializeStreamingTable( + "t", new StructType().add("id", IntegerType).add("Value", StringType), Map.empty) + assert(recordingCatalog.recordedAlters.size == 1) + val changes = recordingCatalog.recordedAlters.flatten + assert(changes.collect { case ac: TableChange.AddColumn => ac.fieldNames()(0) } == + Seq("Value")) + } + } + } + + test("SPARK-58517: schema evolution uses the pipeline's case sensitivity, not the session's") { + // A pipeline-level `SET spark.sql.caseSensitive` never reaches the session, so evolution must + // read it from the flows. Here the session default is case-INsensitive while the pipeline asks + // for case-SENSITIVE, so `Value` must become its own column alongside the persisted `value` -- + // matching how the flow itself resolves the name. + withRecordingCatalog { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + materializeGraph( + new TestGraphRegistrationContext( + spark, Map(SQLConf.CASE_SENSITIVE.key -> "true")) { + registerView("src", query = dfFlowFunc(Seq((1, "a")).toDF("id", "value"))) + registerTable( + "t", + query = Option(sqlFlowFunc(spark, "SELECT id, value FROM src")), + catalog = Option(recordingCatalogName), + database = Option(recordingNamespace)) + }.resolveToDataflowGraph(), + storageRoot = storageRoot + ) + assert( + loadTableFromRecordingCatalog("t").columns() sameElements + CatalogV2Util.structTypeToV2Columns( + new StructType().add("id", IntegerType).add("value", StringType))) + + materializeGraph( + new TestGraphRegistrationContext( + spark, Map(SQLConf.CASE_SENSITIVE.key -> "true")) { + registerView("src", query = dfFlowFunc(Seq((1, "a")).toDF("id", "value"))) + registerTable( + "t", + query = Option(sqlFlowFunc(spark, "SELECT id, value AS Value FROM src")), + catalog = Option(recordingCatalogName), + database = Option(recordingNamespace)) + }.resolveToDataflowGraph(), + storageRoot = storageRoot + ) + assert( + loadTableFromRecordingCatalog("t").columns() sameElements + CatalogV2Util.structTypeToV2Columns( + new StructType() + .add("id", IntegerType) + .add("value", StringType) + .add("Value", StringType)), + "the pipeline asked for case-sensitive resolution, so `Value` must be its own column") + } + } + } + + test("SPARK-58517: flows writing to one table that disagree on case sensitivity are rejected") { + // The effective value decides whether names differing only in case identify the same column, so + // if the flows disagree the resulting schema would depend on the order they are evaluated in. + // Fail with a clear error instead of picking one arbitrarily. + val ctx = new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, "a")).toDF("id", "value"))) + registerTable("t") + registerFlow( + "t", "f1", sqlFlowFunc(spark, "SELECT id, value FROM src"), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "true")) + registerFlow( + "t", "f2", sqlFlowFunc(spark, "SELECT id, value AS Value FROM src"), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "false")) + } + + val ex = intercept[AnalysisException] { + ctx.resolveToDataflowGraph().inferSchemas( + spark.sessionState.conf.caseSensitiveAnalysis) + } + checkError( + exception = ex, + condition = "CONFLICTING_PIPELINE_FLOW_CASE_SENSITIVITY", + parameters = Map( + "tableName" -> "spark_catalog.test_db.t", + "configKey" -> SQLConf.CASE_SENSITIVE.key, + "flowConfigurations" -> + ("false (spark_catalog.test_db.f2); true (spark_catalog.test_db.f1)") + ) + ) + } + + test("SPARK-58517: a flow inheriting the session value conflicts with one that overrides it") { + // f2 leaves the conf unset, so it inherits the session's case-INsensitive default, which + // conflicts with f1's explicit case-sensitive request just as much as an opposite explicit + // value would. + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val ctx = new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, "a")).toDF("id", "value"))) + registerTable("t") + registerFlow( + "t", "f1", sqlFlowFunc(spark, "SELECT id, value FROM src"), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "true")) + registerFlow("t", "f2", sqlFlowFunc(spark, "SELECT id, value FROM src")) + } + + val ex = intercept[AnalysisException] { + ctx.resolveToDataflowGraph().inferSchemas( + spark.sessionState.conf.caseSensitiveAnalysis) + } + assert(ex.getCondition === "CONFLICTING_PIPELINE_FLOW_CASE_SENSITIVITY") + assert(ex.getMessage.contains("session default")) + } + } + + test("SPARK-58517: flows that agree on case sensitivity are accepted") { + // The negative control: identical explicit values are not a conflict, and neither is a value + // that merely differs in spelling from the session's ("TRUE" vs "true"). + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val ctx = new TestGraphRegistrationContext(spark) { + registerView("src", query = dfFlowFunc(Seq((1, "a")).toDF("id", "value"))) + registerTable("t") + registerFlow( + "t", "f1", sqlFlowFunc(spark, "SELECT id, value FROM src"), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "true")) + registerFlow( + "t", "f2", sqlFlowFunc(spark, "SELECT id, value AS Value FROM src"), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "TRUE")) + } + val inferred = ctx.resolveToDataflowGraph() + .inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis)( + fullyQualifiedIdentifier("t")) + // Case-sensitive, so both spellings survive. + assert(inferred.fieldNames.toSeq === Seq("id", "value", "Value")) + } + } + test("re-materializing with a dropped property neither removes it nor issues an alterTable") { withRecordingCatalog { val schema = new StructType().add("id", IntegerType) diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala index cec8db6ec5288..8ce7971b405fa 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SqlPipelineSuite.scala @@ -30,6 +30,12 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { private val externalTable1Ident = fullyQualifiedIdentifier("external_t1") private val externalTable2Ident = fullyQualifiedIdentifier("external_t2") + private def resolveGraph(graph: DataflowGraph): DataflowGraph = + graph.resolve(spark.sessionState.conf.caseSensitiveAnalysis) + + private def validateGraph(graph: DataflowGraph): DataflowGraph = + graph.validate(spark.sessionState.conf.caseSensitiveAnalysis) + override def beforeEach(): Unit = { super.beforeEach() // Create mock external tables that tests can reference, ex. to stream from. @@ -53,7 +59,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |SELECT * FROM STREAM $externalTable2Ident; |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert(resolvedDataflowGraph.flows.size == 4) assert(resolvedDataflowGraph.tables.size == 2) @@ -127,7 +133,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { "CREATE MATERIALIZED VIEW a COMMENT 'this is a comment' AS SELECT * FROM range(1, 4)" ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) val flowA = resolvedDataflowGraph.resolvedFlows @@ -144,7 +150,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert( resolvedDataflowGraph.resolvedFlows @@ -168,7 +174,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) Seq("a", "b", "c", "d").foreach { datasetName => val backingFlow = resolvedDataflowGraph.resolvedFlows @@ -258,7 +264,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |PARTITIONED BY (id1, id2) |AS SELECT id as id1, id as id2 FROM range(1,2) """.stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert( resolvedDataflowGraph.tables @@ -363,7 +369,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { val unresolvedDataflowGraph = unresolvedDataflowGraphFromSql( sqlText = "CREATE STREAMING TABLE st TBLPROPERTIES ('prop1'='foo', 'prop2'='bar') AS SELECT 1" ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert( resolvedDataflowGraph.tables .find(_.identifier == fullyQualifiedIdentifier("st")) @@ -387,7 +393,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert( resolvedDataflowGraph.flows @@ -518,7 +524,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) assert( resolvedDataflowGraph.resolutionFailedFlows @@ -577,10 +583,10 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { |""".stripMargin ) - val resolvedDataflowGraph = unresolvedDataflowGraph.resolve() + val resolvedDataflowGraph = resolveGraph(unresolvedDataflowGraph) // Let inferred/declared schema mismatch detection execute - resolvedDataflowGraph.validate() + validateGraph(resolvedDataflowGraph) val expectedSchema = new StructType().add(name = "id", dataType = LongType, nullable = false) @@ -651,7 +657,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { val unresolvedDataflowGraph = unresolvedDataflowGraphFromSql( sqlText = s"CREATE VIEW b COMMENT 'my persisted comment' AS SELECT * FROM range(1, 4);" ) - val graph = unresolvedDataflowGraph.resolve().validate() + val graph = validateGraph(resolveGraph(unresolvedDataflowGraph)) val view = graph.views.last @@ -897,9 +903,7 @@ class SqlPipelineSuite extends PipelineTest with SharedSparkSession { checkError( exception = intercept[AnalysisException] { - unresolvedDataflowGraph - .resolve() - .validate() + validateGraph(resolveGraph(unresolvedDataflowGraph)) }, condition = "PIPELINE_DATASET_WITHOUT_FLOW", sqlState = Option("0A000"), diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala index db8c368ca89df..9525a61dedd53 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala @@ -60,7 +60,7 @@ class TriggeredGraphExecutionSuite extends ExecutionTest with SharedSparkSession registerMaterializedView("b", query = readFlowFunc("a")) } val unresolvedGraph = pipelineDef.toDataflowGraph - val resolvedGraph = unresolvedGraph.resolve() + val resolvedGraph = unresolvedGraph.resolve(spark.sessionState.conf.caseSensitiveAnalysis) assert(resolvedGraph.flows.size == 2) assert(unresolvedGraph.flows.size == 2) assert(unresolvedGraph.tables.size == 2) @@ -109,7 +109,7 @@ class TriggeredGraphExecutionSuite extends ExecutionTest with SharedSparkSession } val unresolvedGraph = pipelineDef.toDataflowGraph - val resolvedGraph = unresolvedGraph.resolve() + val resolvedGraph = unresolvedGraph.resolve(spark.sessionState.conf.caseSensitiveAnalysis) assert(resolvedGraph.flows.size == 4) assert(resolvedGraph.tables.size == 3) assert(resolvedGraph.views.size == 1) @@ -462,6 +462,8 @@ class TriggeredGraphExecutionSuite extends ExecutionTest with SharedSparkSession updateContext2.pipelineExecution.runPipeline() updateContext2.pipelineExecution.awaitCompletion() + // A streaming source change is unrecoverable without a full refresh, so the flow must not be + // retried: we should see exactly one failure rather than maxFlowRetryAttempts + 1 of them. assertFlowProgressEvent( eventBuffer = updateContext2.eventBuffer, identifier = fullyQualifiedIdentifier("input_table"), @@ -469,6 +471,17 @@ class TriggeredGraphExecutionSuite extends ExecutionTest with SharedSparkSession expectedEventLevel = EventLevel.ERROR, msgChecker = _.contains( s"Flow '${eventLogName("input_table")}' had streaming sources added or removed." + ), + expectedNumOfEvents = Option(1) + ) + + // The run should fail because of the source change, not because the flow exhausted its retries. + assertRunProgressEvent( + eventBuffer = updateContext2.eventBuffer, + state = RunState.FAILED, + expectedEventLevel = EventLevel.ERROR, + msgChecker = _.contains( + s"flow '${eventLogName("input_table")}' had streaming sources added or removed." ) ) } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/UserSpecifiedSchemaValidationSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/UserSpecifiedSchemaValidationSuite.scala index d216539c93bc3..fee3493d01b97 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/UserSpecifiedSchemaValidationSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/UserSpecifiedSchemaValidationSuite.scala @@ -20,10 +20,11 @@ package org.apache.spark.sql.pipelines.graph import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.execution.streaming.runtime.MemoryStream import org.apache.spark.sql.functions +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.pipelines.autocdc.{ChangeArgs, ScdType, UnqualifiedColumnName} import org.apache.spark.sql.pipelines.utils.{PipelineTest, TestGraphRegistrationContext} import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.types.{IntegerType, StringType, StructType} /** * Tests for `GraphValidations.validateUserSpecifiedSchemas`, which requires a table's @@ -100,10 +101,14 @@ class UserSpecifiedSchemaValidationSuite extends PipelineTest with SharedSparkSe /** The full inferred AUTO CDC output schema (data columns plus the reserved metadata column). */ private def autoCdcInferredSchema(flowName: String): StructType = - autoCdcGraph(flowName, declaredSchema = None).inferredSchema(targetIdentifier) + autoCdcGraph(flowName, declaredSchema = None) + .inferSchemas(spark.sessionState.conf.caseSensitiveAnalysis)(targetIdentifier) + + private def validateGraph(graph: DataflowGraph): DataflowGraph = + graph.validate(spark.sessionState.conf.caseSensitiveAnalysis) private def assertSchemaIncompatible(graph: DataflowGraph): Unit = { - val ex = intercept[AnalysisException](graph.validate()) + val ex = intercept[AnalysisException](validateGraph(graph)) assert(ex.getCondition == "USER_SPECIFIED_AND_INFERRED_SCHEMA_NOT_COMPATIBLE") assert(ex.getMessage.contains(targetIdentifier.unquotedString)) } @@ -111,7 +116,7 @@ class UserSpecifiedSchemaValidationSuite extends PipelineTest with SharedSparkSe // Plain flows: the inferred schema is exactly the source's data columns. test("compatible user-specified schema is accepted for an implicit plain flow") { - plainGraph(flowName = "target", declaredSchema = Some(dataSchema)).validate() + validateGraph(plainGraph(flowName = "target", declaredSchema = Some(dataSchema))) } test("incompatible user-specified schema is rejected for an implicit plain flow") { @@ -120,7 +125,7 @@ class UserSpecifiedSchemaValidationSuite extends PipelineTest with SharedSparkSe } test("compatible user-specified schema is accepted for a named plain flow") { - plainGraph(flowName = "plain_flow", declaredSchema = Some(dataSchema)).validate() + validateGraph(plainGraph(flowName = "plain_flow", declaredSchema = Some(dataSchema))) } test("incompatible user-specified schema is rejected for a named plain flow") { @@ -128,6 +133,53 @@ class UserSpecifiedSchemaValidationSuite extends PipelineTest with SharedSparkSe dataSchemaMissingColumn))) } + test("user-specified schema validation uses pipeline case sensitivity, not session default") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val ctx = new TestGraphRegistrationContext( + spark, + Map(SQLConf.CASE_SENSITIVE.key -> "true")) { + val session = spark + import session.implicits._ + + registerView("src", query = dfFlowFunc(Seq((1, "alice")).toDF("id", "value"))) + registerTable( + "target", + specifiedSchema = Some( + new StructType().add("id", IntegerType).add("value", StringType))) + registerFlow( + destinationName = "target", + name = "case_sensitive_flow", + query = sqlFlowFunc(spark, "SELECT id, value AS Value FROM src")) + } + + assertSchemaIncompatible(ctx.resolveToDataflowGraph()) + } + } + + test("user-specified schema validation uses case sensitivity inherited from upstream view") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "false") { + val ctx = new TestGraphRegistrationContext(spark) { + val session = spark + import session.implicits._ + + registerPersistedView( + "src", + query = dfFlowFunc(Seq((1, "alice")).toDF("id", "value")), + sqlConf = Map(SQLConf.CASE_SENSITIVE.key -> "true")) + registerTable( + "target", + specifiedSchema = Some( + new StructType().add("id", IntegerType).add("value", StringType))) + registerFlow( + destinationName = "target", + name = "case_sensitive_flow", + query = sqlFlowFunc(spark, "SELECT id, value AS Value FROM src")) + } + + assertSchemaIncompatible(ctx.resolveToDataflowGraph()) + } + } + // AUTO CDC flows: the inferred schema appends a reserved metadata column to the data columns. test("data-only user-specified schema is rejected for an implicit AUTO CDC flow") { @@ -144,12 +196,14 @@ class UserSpecifiedSchemaValidationSuite extends PipelineTest with SharedSparkSe // Schema includes the appended metadata column, matching the inferred schema exactly. autoCdcGraph( flowName = "target", - declaredSchema = Some(autoCdcInferredSchema("target"))).validate() + declaredSchema = Some(autoCdcInferredSchema("target"))).validate( + spark.sessionState.conf.caseSensitiveAnalysis) } test("full user-specified schema is accepted for a named AUTO CDC flow") { autoCdcGraph( flowName = "auto_cdc_flow", - declaredSchema = Some(autoCdcInferredSchema("auto_cdc_flow"))).validate() + declaredSchema = Some(autoCdcInferredSchema("auto_cdc_flow"))).validate( + spark.sessionState.conf.caseSensitiveAnalysis) } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtilsSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtilsSuite.scala index 41d5bbe14a6b1..92e1c700b99ba 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtilsSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/util/SchemaInferenceUtilsSuite.scala @@ -17,11 +17,66 @@ package org.apache.spark.sql.pipelines.util -import org.apache.spark.SparkFunSuite +import scala.util.Success + +import org.apache.spark.SparkException +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.connector.catalog.TableChange +import org.apache.spark.sql.pipelines.graph.{ + FlowFunction, + FlowFunctionResult, + Input, + QueryContext, + QueryOrigin, + ResolvedFlow, + StreamingFlow, + UntypedFlow +} +import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.types._ -class SchemaInferenceUtilsSuite extends SparkFunSuite { +class SchemaInferenceUtilsSuite extends QueryTest with SharedSparkSession { + + /** A [[FlowFunction]] that throws if invoked; the inferSchemaFromFlows test builds resolved + * flows directly. */ + private val noOpFlowFunction: FlowFunction = new FlowFunction { + override def call( + allInputs: Set[TableIdentifier], + availableInputs: Seq[Input], + configuration: Map[String, String], + queryContext: QueryContext, + queryOrigin: QueryOrigin): FlowFunctionResult = + throw new UnsupportedOperationException( + "noOpFlowFunction.call should not be invoked from SchemaInferenceUtilsSuite tests") + } + + private val queryContext = QueryContext(currentCatalog = Some("c"), currentDatabase = Some("d")) + + /** A resolved flow with the given identifier and output schema, writing to `destination`. */ + private def resolvedFlow( + identifier: TableIdentifier, + destination: TableIdentifier, + schema: StructType): ResolvedFlow = { + val df = spark.createDataFrame(spark.sparkContext.emptyRDD[Row], schema) + val flow = UntypedFlow( + identifier = identifier, + destinationIdentifier = destination, + func = noOpFlowFunction, + queryContext = queryContext, + sqlConf = Map.empty, + once = false, + origin = QueryOrigin.empty) + new StreamingFlow( + flow, + FlowFunctionResult( + requestedInputs = Set.empty, + batchInputs = Set.empty, + streamingInputs = Set.empty, + usedExternalInputs = Set.empty, + dataFrame = Success(df), + sqlConf = Map.empty)) + } test("determineColumnChanges - adding new columns") { val currentSchema = new StructType() @@ -270,4 +325,161 @@ class SchemaInferenceUtilsSuite extends SparkFunSuite { assert(addedColumnNames === Set("full_name", "email")) assert(deletedColumnNames === Set("first_name", "last_name")) } + + test("determineColumnChanges - a case-only difference is a drop-then-add, not a match") { + // diffSchemas keys column identity on the EXACT field name, with no case normalization. So a + // target `Value` against a persisted `value` is a distinct column: `value` is dropped and + // `Value` added. This is what makes a case-only rename visible on the non-merging paths + // (materialized views, full refresh), where targetSchema is the declared schema as-is. + val currentSchema = new StructType().add("id", IntegerType).add("value", StringType) + val targetSchema = new StructType().add("id", IntegerType).add("Value", StringType) + + val changes = SchemaInferenceUtils.diffSchemas(currentSchema, targetSchema) + + val addChanges = changes.collect { case ac: TableChange.AddColumn => ac.fieldNames()(0) } + val deleteChanges = changes.collect { case dc: TableChange.DeleteColumn => dc.fieldNames()(0) } + assert(addChanges === Seq("Value")) + assert(deleteChanges === Seq("value")) + } + + test("determineColumnChanges - two declared columns differing only in case are both kept") { + // A declared schema carrying both `value` and `Value` reaches diffSchemas verbatim (nothing on + // the create path rejects duplicate-cased columns). Exact-name keying must surface BOTH as + // additions; normalizing the lookup key would collapse them and silently keep an arbitrary one + // (whichever came last), losing a column the user declared. + val currentSchema = new StructType().add("id", IntegerType) + val targetSchema = new StructType() + .add("id", IntegerType) + .add("value", StringType) + .add("Value", IntegerType) + + val changes = SchemaInferenceUtils.diffSchemas(currentSchema, targetSchema) + + val added = changes.collect { case ac: TableChange.AddColumn => + ac.fieldNames()(0) -> ac.dataType() + }.toMap + assert(added === Map("value" -> StringType, "Value" -> IntegerType)) + } + + test("mergeSchemas - a nested field differing only in case folds onto the existing field when " + + "case-insensitive") { + // The nested analog of the top-level case-only fold. `StructType.merge` propagates the + // case-sensitivity flag into nested struct merges (SPARK-58525), so the incoming `s.Value` is + // matched to the existing `s.value` and the struct keeps a single field with the persisted + // (left) spelling -- rather than growing a second, case-differing nested field. + val currentSchema = new StructType() + .add("id", IntegerType) + .add("s", new StructType().add("value", StringType)) + val dataSchema = new StructType() + .add("id", IntegerType) + .add("s", new StructType().add("Value", StringType)) + + val merged = + SchemaMergingUtils.mergeSchemas(currentSchema, dataSchema, caseSensitive = false) + assert(merged === currentSchema) + + // Because the merge is a no-op, evolution derives no table changes at all: in particular the + // nested struct is NOT rewritten (which would be an UpdateColumnType on `s`). + assert( + SchemaInferenceUtils.diffSchemas(currentSchema, merged).isEmpty) + } + + test("mergeSchemas - a nested field differing only in case stays distinct when case-sensitive") { + // The case-sensitive control: `s.value` and `s.Value` are different fields, so the merged + // struct carries both. + val currentSchema = new StructType().add("s", new StructType().add("value", StringType)) + val dataSchema = new StructType().add("s", new StructType().add("Value", StringType)) + + val merged = SchemaMergingUtils.mergeSchemas(currentSchema, dataSchema, caseSensitive = true) + val expectedStruct = new StructType().add("value", StringType).add("Value", StringType) + assert(merged === new StructType().add("s", expectedStruct)) + + // Unlike the case-insensitive test above (where the merge is a no-op and no changes are + // derived), evolution here must rewrite the top-level `s` column. `diffSchemas` compares nested + // types wholesale, so the growth of a nested field surfaces as a single UpdateColumnType on `s` + // carrying the full new struct -- not as an add of `s.Value`. + val changes = SchemaInferenceUtils.diffSchemas(currentSchema, merged) + assert(changes.length === 1) + val typeChange = changes.collect { case tc: TableChange.UpdateColumnType => tc } + assert(typeChange.length === 1) + assert(typeChange.head.fieldNames() === Array("s")) + assert(typeChange.head.newDataType() === expectedStruct) + } + + test("mergeSchemas - a nested case-only field whose type also changes fails to merge, and " + + "diffSchemas reports it as a type change") { + // A nested field that differs only in case AND changes type is rejected rather than silently + // resolved. Note this is a *type* incompatibility, not a case one: `StructType.merge` never + // widens numeric types, so `int` -> `long` fails identically for a same-cased field and at the + // top level. The value of pinning it here is that case-insensitive matching does not turn an + // incompatible type change into a silent merge -- the run still fails loudly, and the user's + // remedy is a full refresh. + val currentSchema = new StructType().add("s", new StructType().add("value", IntegerType)) + val dataSchema = new StructType().add("s", new StructType().add("Value", LongType)) + + val ex = intercept[SparkException] { + SchemaMergingUtils.mergeSchemas(currentSchema, dataSchema, caseSensitive = false) + } + assert(ex.getCondition === "CANNOT_MERGE_INCOMPATIBLE_DATA_TYPE") + + // Same-cased and top-level widening fail the same way, confirming the rejection is about the + // type change rather than the case difference. + intercept[SparkException] { + SchemaMergingUtils.mergeSchemas( + currentSchema, + new StructType().add("s", new StructType().add("value", LongType)), + caseSensitive = false) + } + intercept[SparkException] { + SchemaMergingUtils.mergeSchemas( + new StructType().add("v", IntegerType), + new StructType().add("v", LongType), + caseSensitive = false) + } + + // Diffing the two schemas directly (rather than diffing against their merge, which fails + // above) reports a TYPE change on the enclosing `s` column -- not a field-name mismatch, i.e. + // not an add of `s.Value` plus a delete of `s.value`. `diffSchemas` keys column identity only + // at the top level and compares nested types wholesale, so the case difference inside the + // struct never surfaces as an add/delete pair. + { + val changes = SchemaInferenceUtils.diffSchemas(currentSchema, dataSchema) + assert(changes.length === 1, s"changes=$changes") + val typeChange = changes.collect { case tc: TableChange.UpdateColumnType => tc } + assert(typeChange.length === 1, s"changes=$changes") + assert(typeChange.head.fieldNames() === Array("s")) + assert(typeChange.head.newDataType() === new StructType().add("Value", LongType)) + assert(!changes.exists(_.isInstanceOf[TableChange.AddColumn])) + assert(!changes.exists(_.isInstanceOf[TableChange.DeleteColumn])) + } + } + + test("inferSchemaFromFlows folds a case-only column to the same spelling regardless of flow " + + "order, even when identifier names contain dots") { + // The merge order decides which spelling of a case-only-differing column survives, so it must + // not depend on the incoming flow order (the nondeterministic flow-resolution completion + // order). The two identifiers below differ only in where the dot falls, so a dot-joined sort + // key would render them identical; sorting on the identifier parts keeps them distinct. + val destination = TableIdentifier("t", Some("d"), Some("c")) + val flowA = resolvedFlow( + identifier = TableIdentifier("x", Some("a.b"), Some("c")), + destination = destination, + schema = new StructType().add("id", IntegerType).add("value", StringType)) + val flowB = resolvedFlow( + identifier = TableIdentifier("b.x", Some("a"), Some("c")), + destination = destination, + schema = new StructType().add("id", IntegerType).add("Value", StringType)) + + // The lower identifier (flowB: database "a" precedes "a.b") supplies the surviving spelling, in + // either input order. + val expected = new StructType().add("id", IntegerType).add("Value", StringType) + Seq(Seq(flowA, flowB), Seq(flowB, flowA)).foreach { flows => + val inferred = SchemaInferenceUtils.inferSchemaFromFlows( + tableIdentifier = destination, + flows = flows, + userSpecifiedSchema = None, + sessionCaseSensitive = false) + assert(inferred === expected, s"unexpected schema for input order $flows") + } + } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/utils/TestGraphRegistrationContext.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/utils/TestGraphRegistrationContext.scala index 068171a46aa16..9d6067cd51cd3 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/utils/TestGraphRegistrationContext.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/utils/TestGraphRegistrationContext.scala @@ -305,7 +305,8 @@ class TestGraphRegistrationContext( query: FlowFunction, once: Boolean = false, catalog: Option[String] = None, - database: Option[String] = None + database: Option[String] = None, + sqlConf: Map[String, String] = Map.empty ): Unit = { val rawFlowIdentifier = GraphIdentifierManager.parseTableIdentifier(name, spark) val rawDestinationIdentifier = @@ -345,7 +346,7 @@ class TestGraphRegistrationContext( currentCatalog = catalog.orElse(Some(defaultCatalog)), currentDatabase = database.orElse(Some(defaultDatabase)) ), - sqlConf = Map.empty, + sqlConf = sqlConf, once = once, origin = QueryOrigin( objectName = Option(flowIdentifier.unquotedString), @@ -407,7 +408,8 @@ class TestGraphRegistrationContext( * Generates a dataflow graph from this pipeline definition and resolves it. * @return */ - def resolveToDataflowGraph(): DataflowGraph = toDataflowGraph.resolve() + def resolveToDataflowGraph(): DataflowGraph = + toDataflowGraph.resolve(spark.sessionState.conf.caseSensitiveAnalysis) } object TestGraphRegistrationContext { diff --git a/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala b/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala index a47b08c7b949b..1fd7f5ac624ca 100644 --- a/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala +++ b/streaming/src/main/scala/org/apache/spark/streaming/ui/BatchPage.scala @@ -270,19 +270,6 @@ private[ui] class BatchPage(parent: StreamingTab) extends WebUIPage("batch") { } } - private def generateOutputOperationStatusForUI(failure: String): String = { - if (failure.startsWith("org.apache.spark.SparkException")) { - "Failed due to Spark job error\n" + failure - } else { - var nextLineIndex = failure.indexOf("\n") - if (nextLineIndex < 0) { - nextLineIndex = failure.length - } - val firstLine = failure.substring(0, nextLineIndex) - s"Failed due to error: $firstLine\n$failure" - } - } - /** * Generate the job table for the batch. */ diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala index 4278313fdf394..f6f934bc55184 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala @@ -22,9 +22,9 @@ import org.apache.spark.udf.worker.{CancelResponse, ExecutionError, FinishRespon /** * :: Experimental :: * The terminal outcome a [[WorkerSession]] settles on, returned by - * [[WorkerSession#close]]. Mirrors the four terminal `WorkerSession.SessionState`s, - * so close() reports the outcome faithfully rather than collapsing failures into - * a clean cancel. + * [[WorkerSession#close]]. Enumerates the terminal outcomes, carried by the + * single `WorkerSession.SessionState.Terminal`, so close() reports the outcome + * faithfully rather than collapsing failures into a clean cancel. * * '''Clean outcomes''' ([[Finished]] / [[Cancelled]]) wrap the worker's * `FinishResponse` / `CancelResponse` -- per-execution metrics, an optional @@ -35,11 +35,12 @@ import org.apache.spark.udf.worker.{CancelResponse, ExecutionError, FinishRespon * `FinishResponse` was already produced when a `Cancel` arrives the engine still * receives [[Finished]], otherwise [[Cancelled]]. * - * '''Failure outcomes''' ([[Failed]] / [[TransportFailed]]) carry the cause - * instead of a proto terminator (none arrived, so they have no metrics). They - * exist so a failure is not reported as a benign cancel -- in particular an - * error raised during finish/close, '''after all data has been drained''', - * reaches the caller only through this value, never through the result iterator. + * '''Failure outcomes''' ([[Failed]] / [[TransportFailed]] / [[Interrupted]]) + * carry the cause instead of a proto terminator (none arrived, so they have no + * metrics). They prevent a failure from being reported as a benign cancel. In + * particular, an error raised during finish/close, '''after all data has been + * drained''', reaches the caller only through this value, never through the + * result iterator. */ @Experimental sealed trait Termination @@ -59,8 +60,25 @@ object Termination { final case class Failed(error: ExecutionError) extends Termination /** - * A transport failure, timeout, or interrupt tore the stream down before any - * terminator arrived. Carries the underlying cause. + * A transport failure or timeout tore the stream down before any terminator + * arrived. Carries the underlying cause. Leaves the worker in an unknown state, + * so it is not salvageable (see [[WorkerSession.isWorkerSalvageable]]). */ final case class TransportFailed(cause: Throwable) extends Termination + + /** + * The session was interrupted (an [[InterruptedException]] on the engine + * thread driving it, e.g. a Spark task kill) before a terminator arrived. A + * best-effort `Cancel` is sent and its `CancelResponse` awaited only briefly + * (the interrupted thread must unwind promptly rather than block); this + * terminal is settled only when that brief wait expires without an ack -- if + * the worker acks in time the session settles a cooperative [[Cancelled]] + * instead. So it is distinct from a [[Cancelled]] (which carries the drained + * `CancelResponse`) and from a [[TransportFailed]] (a genuine transport + * fault). Although the interrupt is an engine-side event, the missing + * acknowledgement leaves the worker's state unknown, so it is not salvageable + * without a separate liveness proof (see [[WorkerSession.isWorkerSalvageable]]). + * Carries the interrupt cause. + */ + final case class Interrupted(cause: Throwable) extends Termination } diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UDFDispatcherManager.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UDFDispatcherManager.scala index 7fd624b9d0d8e..5b66df68c619c 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UDFDispatcherManager.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/UDFDispatcherManager.scala @@ -44,9 +44,10 @@ class UDFDispatcherManager( workerLogger: WorkerLogger = WorkerLogger.NoOp ) { - // Guarded by `rwLock`. The read lock is used by getDispatcher - // (with upgrade when a new dispatcher must be added) and the - // write lock is used by close. + // Guarded by `rwLock`. getDispatcher takes the read lock, releasing it + // and re-acquiring the write lock when a new dispatcher must be added + // (ReentrantReadWriteLock does not support upgrading a held read lock); + // close takes the write lock. private val rwLock = new ReentrantReadWriteLock() private val dispatchers = new HashMap[UDFWorkerSpecification, WorkerDispatcher]() diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerDispatcher.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerDispatcher.scala index e938c3e04be5b..39ab55b75ed4f 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerDispatcher.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerDispatcher.scala @@ -32,8 +32,8 @@ import org.apache.spark.udf.worker.UDFWorkerSpecification * worker that backed it MUST NOT be returned to any reuse pool. A transport * error leaves the worker in an unknown state; only workers that complete * sessions cleanly are eligible for reuse. Implementations are responsible for - * tracking this condition -- typically [[WorkerSession.doProcess]] flags the - * worker as invalid before [[WorkerSession.doClose]] releases it, so the + * tracking this condition -- [[WorkerSession.close]] marks the worker invalid + * (via [[WorkerSession.isWorkerSalvageable]]) before releasing it, so the * dispatcher can distinguish a clean release from a failed one. */ @Experimental diff --git a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala index 6ee9b129ac945..0afde972bc29a 100644 --- a/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala +++ b/udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala @@ -89,10 +89,10 @@ import org.apache.spark.udf.worker.{Cancel, DataRequest, DataResponse, Finish, I * | (from any non-terminal state) * }}} * The clean path (via Finishing) and the cancel path (via Cancelling) settle the - * same `(terminal)`. The four terminals are: + * same `(terminal)`. The terminals are: * {{{ * Finished(FinishResponse) | Cancelled(CancelResponse) - * Failed(ExecutionError) | TransportFailed(Throwable) + * Failed(ExecutionError) | TransportFailed(Throwable) | Interrupted(Throwable) * }}} * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse` * (metrics + finish/cancel callback `data`/`error`); the failure terminals carry @@ -156,6 +156,7 @@ abstract class WorkerSession( * the worker needs to start processing. */ final def init(message: Init): InitResponse = { + require(message != null, "message is required") if (!state.compareAndSet(SessionState.Created, SessionState.Initializing)) { throw new IllegalStateException( s"init must be called exactly once before process (current state: ${state.get()})") @@ -348,18 +349,20 @@ abstract class WorkerSession( /** * Whether the underlying worker is in a state safe to reuse after this * session ends. The default treats only a dead or unknown transport as - * unsafe: a [[Termination.TransportFailed]] outcome (transport failure, - * timeout, or interrupt) -- or a session that never settled -- leaves the - * worker in an unknown state and is not salvageable. Every other terminal is - * salvageable: a clean [[Termination.Finished]], a cooperative - * [[Termination.Cancelled]], and also an execution [[Termination.Failed]], - * which is typically a user-code (UDF) error reported by a still-healthy - * worker rather than a worker fault. A `false` result tells [[close]] to mark - * `workerHandle` invalid so no reuse pool recycles the worker. Subclasses may - * override for protocol-specific nuances. + * unsafe: a [[Termination.TransportFailed]] outcome (transport failure or + * timeout), a [[Termination.Interrupted]] without a worker acknowledgement, or + * a session that never settled leaves the worker in an unknown state and is + * not salvageable. The salvageable terminals are a clean + * [[Termination.Finished]], a cooperative [[Termination.Cancelled]], and an + * execution [[Termination.Failed]] (typically a user-code (UDF) error reported + * by a still-healthy worker rather than a worker fault). + * A `false` result tells [[close]] to mark `workerHandle` invalid so no reuse + * pool recycles the worker. Subclasses may override for protocol-specific + * nuances. */ protected def isWorkerSalvageable: Boolean = state.get() match { case SessionState.Terminal(_: Termination.TransportFailed) => false + case SessionState.Terminal(_: Termination.Interrupted) => false case t if t.isTerminal => true case _ => false } @@ -461,8 +464,9 @@ object WorkerSession { /** * The session is over; no further writes are valid. The single terminal * state carries the public [[Termination]] outcome (`Finished` / - * `Cancelled` / `Failed` / `TransportFailed`), so those four outcome cases - * live once on [[Termination]] rather than being mirrored on the state. + * `Cancelled` / `Failed` / `TransportFailed` / `Interrupted`), so those + * outcome cases live once on [[Termination]] rather than being mirrored on + * the state. */ final case class Terminal(termination: Termination) extends SessionState { override def isTerminal: Boolean = true diff --git a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala index 37380d0d77655..c8bbcfffd7bb2 100644 --- a/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala +++ b/udf/worker/core/src/test/scala/org/apache/spark/udf/worker/core/WorkerSessionSuite.scala @@ -95,6 +95,18 @@ class WorkerSessionSuite extends AnyFunSuite { assert(ex.getMessage.contains("exactly once")) } + test("init rejects a null message before changing state") { + var initCalled = false + val s = new FakeWorkerSession(onInit = _ => { + initCalled = true + InitResponse.getDefaultInstance + }) + val ex = intercept[IllegalArgumentException](s.init(null)) + assert(ex.getMessage.contains("message is required")) + assert(!initCalled) + assert(s.state == SessionState.Created) + } + test("process before init is rejected") { val s = new FakeWorkerSession() val ex = intercept[IllegalStateException](s.process(Iterator.empty)) @@ -166,6 +178,21 @@ class WorkerSessionSuite extends AnyFunSuite { assert(h.released == 1) } + test("close marks a worker invalid after an unacknowledged Interrupted termination") { + val h = new RecordingHandle + val cause = new InterruptedException("task killed") + // The interrupt itself is an engine-side event, but Interrupted means no + // CancelResponse was received. Without that acknowledgement (or a separate + // liveness proof), the worker's state is unknown and it must not be reused. + val s = new FakeWorkerSession(handle = h, onCloseHook = (self, _) => { + self.settle(Termination.Interrupted(cause)) + self.term + }) + assert(s.close() == Termination.Interrupted(cause)) + assert(h.invalidated == 1) + assert(h.released == 1) + } + test("close enforces the doClose terminal post-condition") { val h = new RecordingHandle // doClose returns a Termination without settling any terminal -- a subclass @@ -261,6 +288,7 @@ class WorkerSessionSuite extends AnyFunSuite { assert(termFor(Termination.Cancelled(can)) == Termination.Cancelled(can)) assert(termFor(Termination.Failed(err)) == Termination.Failed(err)) assert(termFor(Termination.TransportFailed(cause)) == Termination.TransportFailed(cause)) + assert(termFor(Termination.Interrupted(cause)) == Termination.Interrupted(cause)) } test("settledTermination throws before a terminal is settled") { diff --git a/udf/worker/grpc/pom.xml b/udf/worker/grpc/pom.xml index a2ce417b5ea76..aa563d157d9c0 100644 --- a/udf/worker/grpc/pom.xml +++ b/udf/worker/grpc/pom.xml @@ -65,6 +65,11 @@ <artifactId>spark-udf-worker-proto_${scala.binary.version}</artifactId> <version>${project.version}</version> </dependency> + <dependency> + <groupId>org.apache.spark</groupId> + <artifactId>spark-udf-worker-core_${scala.binary.version}</artifactId> + <version>${project.version}</version> + </dependency> <dependency> <groupId>org.scala-lang</groupId> <artifactId>scala-library</artifactId> diff --git a/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala new file mode 100644 index 0000000000000..c0c3511889a2f --- /dev/null +++ b/udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSession.scala @@ -0,0 +1,1182 @@ +/* + * 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.udf.worker.grpc + +import java.util.Objects +import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue, TimeoutException, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference} + +import scala.util.control.NonFatal + +import io.grpc.{ConnectivityState, ManagedChannel} +import io.grpc.stub.StreamObserver + +import org.apache.spark.annotation.Experimental +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlRequest, + UdfControlResponse, UdfRequest, UdfResponse, UdfWorkerGrpc} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger, WorkerSession} +import org.apache.spark.udf.worker.core.WorkerSession.SessionState +import org.apache.spark.udf.worker.grpc.GrpcWorkerSession._ + +/** + * :: Experimental :: + * gRPC implementation of [[WorkerSession]] for the `UdfWorker.Execute` + * bidirectional RPC. + * + * Drives one bidirectional `Execute` stream against the worker per the + * ordering invariants documented in `udf_message.proto` (`PayloadChunk*` + * between Init and InitResponse omitted here; chunking is not yet + * implemented -- see the TODO below): + * {{{ + * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)? + * | Cancel + * Worker -> Engine: InitResponse -> (DataResponse)* -> + * (ErrorResponse)? -> (FinishResponse | CancelResponse) + * }}} + * + * Knows nothing about how the worker was provisioned -- the dispatcher + * constructs this with a [[WorkerHandle]] and channel; the base + * [[WorkerSession]] handles dispatcher-side cleanup on close. + * + * '''Driving model.''' Consumption-driven (Volcano / pull): the thread that + * consumes the [[doProcess]] result iterator is the one that pulls input and + * sends each `DataRequest`; the gRPC callback thread only receives output. It is + * pull-driven but not one-input-per-output -- `advance` sends the next input + * whenever the output queue is momentarily empty, so under async delivery it may + * push several input batches before any output is read. HTTP/2 flow control + * bounds wire traffic, but application-level buffering is intentionally left to + * a follow-up (see the TODO on [[outputQueue]]). + * + * '''State machine.''' This class does not keep its own state machine: it + * drives the single [[WorkerSession.SessionState]] owned by the base. The base + * advances `Created -> Initializing` (in `init`) and `Initialized -> Streaming` + * (in `process`); this class advances the protocol-event edges through + * [[compareAndSetState]] / [[completeTerminal]] as it exchanges messages: + * {{{ + * Initializing --(InitResponse ok)--> Initialized [handleControl] + * Streaming ----(input exhausted)-------> Finishing [ProcessIterator] + * <any non-terminal> --(Cancel send starts)--> Cancelling [sendCancelInternal] + * <any non-terminal> --(terminator/error)--> terminal [completeTerminal] + * }}} + * The two clean terminals carry the worker's `FinishResponse` / `CancelResponse` + * (metrics + finish/cancel callback `data`/`error`) so [[close]] can return + * them. Cancellation intent is deliberately tracked outside the machine in + * [[cancelRequested]], which makes the wire write idempotent and suppresses any + * in-flight Data/Finish once cancellation begins. + * + * Threading: + * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`, + * returning it. + * - [[doProcess]] returns an iterator. Input batches are forwarded inline + * (the iterator's `next()` thread also sends `DataRequest`). Output + * batches arrive via the response observer (gRPC callback thread) and + * are consumed by the same iterator. A terminator (`FinishResponse`, + * `CancelResponse`, `ErrorResponse`, gRPC stream error) is published + * once. + * - [[doClose]] is thread-safe and idempotent: it settles + returns the + * terminator (cancelling in-flight work if the stream had not finished) + * and terminates the request side. + * + * TODO [SPARK-55278]: this class does not yet implement payload chunking; + * the entire [[org.apache.spark.udf.worker.UdfPayload]] is sent inline. Chunking will be added + * when a UDF payload large enough to exceed gRPC's default message size + * limit is introduced. + * + * @param workerHandle dispatcher-side handle for releasing the worker on + * [[close]] (see [[WorkerSession]]). + * @param channel a gRPC channel built and owned by the caller (the + * dispatcher). Not closed here -- the dispatcher tears it + * down via [[WorkerHandle]]. + * @param logger diagnostics. Defaults to [[WorkerLogger.NoOp]]. + * @param initResponseTimeoutMs upper bound on the wait for `InitResponse` + * after [[doInit]] sends `Init`. + * @param terminalTimeoutMs upper bound on the wait for a stream + * terminator (`FinishResponse`, + * `CancelResponse`, or `ErrorResponse`). + * Each output-queue poll resets this wait; + * see [[doProcess]] / `ProcessIterator`. + * @param interruptCancelTimeoutMs upper bound on the wait for a `CancelResponse` + * after an interrupt (cancelled query / killed + * task) sends `Cancel`. Short by design; on expiry + * the session settles `Interrupted`. See + * `handleInterrupt`. + */ +@Experimental +class GrpcWorkerSession( + workerHandle: WorkerHandle, + channel: ManagedChannel, + logger: WorkerLogger = WorkerLogger.NoOp, + initResponseTimeoutMs: Long = DEFAULT_INIT_RESPONSE_TIMEOUT_MS, + terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS, + interruptCancelTimeoutMs: Long = DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS) + extends WorkerSession(workerHandle, logger) { + + require(channel != null, "channel is required") + + private val asyncStub = UdfWorkerGrpc.newStub(channel) + + // Output batches from the worker, drained by the process() iterator. + // Intentionally unbounded in this first implementation: a bounded queue would + // block the gRPC callback (Netty event-loop) thread when full, stalling + // terminator/control delivery on the whole channel. The consumer normally + // drains promptly, but HTTP/2 flow control alone does not bound this application + // queue or gRPC's asynchronous send buffers. + // + // TODO [SPARK-55278]: add application-level gRPC flow control in a follow-up, + // using ClientCallStreamObserver readiness for requests and manual inbound + // demand for responses, before wiring this transport into a production path. + // TODO [SPARK-57324]: expose queue depth as a metric (early warning for a + // stalled consumer). + private val outputQueue = new LinkedBlockingQueue[QueueItem]() + + // This value couples the worker's `InitResponse` (success or error) with the + // latch on which init() blocks. The latch fires when the InitResponse arrives + // (`complete`), when a pre-init ErrorResponse arrives, or when a terminal + // settles first without an InitResponse (`signalWithoutValue`). Until it fires + // we have no proof the worker accepted the session. + // + // Settle-before-release rule (referenced from every callback that both settles + // a terminal and fires this latch): settle the terminal FIRST, then complete / + // signal. init() blocks on the latch, and only a latch await/release pair + // establishes a happens-before edge, so a woken init() is guaranteed to + // observe the terminal rather than a transient state. OneShotValue keeps that + // publish-then-release in one place instead of every caller remembering to + // count down after setting the reference. + private val initValue = new OneShotValue[InitResponse] + + // Fired when the session reaches a terminal [[SessionState]]. doClose() and + // the init-error path block on this to drain the terminator. + private val terminalLatch = new CountDownLatch(1) + + // Captures an ErrorResponse encountered during the data phase so that + // the CancelResponse terminator can attribute the failure to the original + // user / worker / protocol error rather than reporting a bare "Cancelled". + private val executionError = new AtomicReference[Option[ExecutionError]](None) + + // True immediately before Finish is handed to the request observer. It is set + // before onNext so a reentrant callback may deliver FinishResponse while the + // Finish request is still being written. + private val finishSendStarted = new AtomicBoolean(false) + + // Cancellation INTENT -- distinct from the `Cancelling` state, which is reached + // once a Cancel send is attempted under [[requestLock]] ([[sendCancelInternal]]). + // Intent is set first and can outrun (or never reach) that attempt, so + // `cancelRequested` does NOT imply state `Cancelling`. It is kept outside the + // [[SessionState]] machine to (a) make cancellation idempotent across all call + // sites and (b) suppress any Data/Finish that would otherwise race a Cancel + // onto the wire (re-read inside [[requestLock]]). + private val cancelRequested = new AtomicBoolean(false) + + // gRPC requires serialized writes to a request StreamObserver. + private val requestLock = new Object + + // Initialized in init() -- before that, close() is a no-op on the request + // side, which is exactly the contract the wrapping WorkerSession expects. + @volatile private var requestObserver: StreamObserver[UdfRequest] = _ + + // True after the request side has been half-closed, aborted, or observed the + // transport closing. It prevents this session from issuing duplicate terminal + // calls, including following onError with onCompleted from close(). + private val requestSideTerminated = new AtomicBoolean(false) + + private val responseObserver: StreamObserver[UdfResponse] = new StreamObserver[UdfResponse] { + override def onNext(response: UdfResponse): Unit = { + response.getResponseCase match { + case UdfResponse.ResponseCase.DATA => + // A DataResponse before InitResponse violates the protocol (InitResponse + // must precede any DataResponse): fast-fail rather than enqueue it and let + // init() block to its timeout. Settling the terminal wakes a blocked + // init() via onTerminalSettled. + if (!initResolved) { + Transitions.transportFailed(new IllegalStateException( + "worker sent a DataResponse before InitResponse")) + } else { + outputQueue.put(QueueItem.Batch(response.getData)) + } + + case UdfResponse.ResponseCase.CONTROL => + handleControl(response.getControl) + + case other => + // A malformed response (empty / unknown oneof) is a terminal transport + // failure; fast-fail init the same way. + Transitions.transportFailed(new IllegalStateException( + s"unexpected response oneof: $other")) + } + } + + override def onError(t: Throwable): Unit = { + // Transport-level failure: the stream is dead, no further writes possible. + // Settling the terminal wakes a blocked init() (via onTerminalSettled) so it + // surfaces the transport cause instead of the initResponseTimeoutMs error. + requestSideTerminated.set(true) + Transitions.transportFailed(t) + } + + override def onCompleted(): Unit = { + // Worker half-closed its side without sending a terminator (FinishResponse + // / CancelResponse). Treat as transport error so the engine sees a + // failure, not a silent end-of-stream. Settling the terminal wakes a blocked + // init() via onTerminalSettled. + requestSideTerminated.set(true) + if (!currentState.isTerminal) { + Transitions.transportFailed(new IllegalStateException( + "worker response stream closed without a terminator")) + } + } + } + + /** + * Wakes everything that can be blocked when the base settles a terminal: + * the result iterator (on [[outputQueue]]), a thread on [[terminalLatch]] + * (close), and a thread still blocked in [[doInit]] on [[initValue]]. Invoked + * once, by the caller that wins [[completeTerminal]], so this is the '''single''' + * place a settled terminal wakes a blocked [[doInit]] -- callers that settle a + * terminal (any response-callback failure/terminator branch, the timeout paths, + * or close()) do NOT signal [[initValue]] themselves; settling is enough. + * Settle-before-release (see [[initValue]]) holds because this runs after the + * terminal CAS in [[completeTerminal]]. The only direct [[initValue]] signals + * left are the non-terminal init paths ([[handleControl]]'s INIT-ok / INIT-error + * / pre-init-ERROR branches), which must wake [[doInit]] '''without''' settling a + * terminal. INIT-ok first advances the session to `Initialized`; the error + * branches remain `Initializing` only until [[doInit]] sends Cancel and throws. + */ + override protected def onTerminalSettled(termination: Termination): Unit = { + outputQueue.put(QueueItem.EndOfStream) + initValue.signalWithoutValue() + terminalLatch.countDown() + } + + private def handleControl(ctrl: UdfControlResponse): Unit = ctrl.getControlCase match { + case UdfControlResponse.ControlCase.INIT => + val resp = ctrl.getInit + if (resp.hasError) { + // Record the error and publish the InitResponse; do NOT settle a terminal + // or send Cancel here. The proto requires the engine to Cancel after an + // init error (udf_message.proto); doInit owns that synchronous cleanup + // after the initial onNext returns, then drains the CancelResponse before + // throwing. Settle-before-release (see initValue): there is no terminal to + // settle first here, just publish. + executionError.compareAndSet(None, Some(resp.getError)) + initValue.complete(resp) + } else { + // InitResponse OK. Only advance from Initializing so a terminal that raced + // in (e.g. a transport error) still wins; process() then opens the data + // phase. Settle-before-release (see initValue): publish after the CAS. + Transitions.initAccepted() + initValue.complete(resp) + } + + case UdfControlResponse.ControlCase.ERROR => + val err = ctrl.getError.getError + executionError.compareAndSet(None, Some(err)) + if (!initResolved) { + // Pre-init ErrorResponse. Leave the session in Initializing and just wake + // init(): the recorded executionError tells doInit the worker failed, and + // doInit sends the Cancel and drains the CancelResponse after the initial + // onNext returns, just like the INIT-error branch. + initValue.signalWithoutValue() + } else { + // Data-phase ErrorResponse: requestObserver is published, so cancel here. + // Cancel -> CancelResponse settles the terminal; the iterator surfaces the + // recorded executionError. The init latch already fired in init(), so no + // signalWithoutValue is needed. + sendCancelInternal(() => cancelWithReason("aborting after ErrorResponse")) + } + + case UdfControlResponse.ControlCase.FINISH => + // The FinishResponse carries metrics + the finish-callback data/error. + // Keep it on the terminal so close() can return it; the iterator inspects + // its error field to decide whether to throw. A FinishResponse is valid only + // after the engine has started sending Finish; accepting one earlier could + // silently truncate input and report a clean result. + if (finishSendStarted.get()) { + Transitions.finished(ctrl.getFinish) + } else { + Transitions.transportFailed(new IllegalStateException( + "worker sent FinishResponse before the engine sent Finish")) + } + + case UdfControlResponse.ControlCase.CANCEL => + // The CancelResponse carries metrics + the cancel-callback error. Keep it + // on the terminal so close() can return it; any prior ErrorResponse is + // tracked in executionError and surfaced by the iterator. Settling the + // terminal wakes a blocked init() (a CANCEL before InitResponse) via + // onTerminalSettled. + Transitions.cancelled(ctrl.getCancel) + + case UdfControlResponse.ControlCase.CONTROL_NOT_SET => + // Settling the terminal wakes a blocked init() via onTerminalSettled. + Transitions.transportFailed(new IllegalStateException( + "empty UdfControlResponse oneof")) + } + + /** + * True once init is no longer pending -- i.e. the stream is past `Initializing`. + * Not "init succeeded": a terminal (including a failure) also counts as resolved. + */ + private def initResolved: Boolean = currentState match { + case SessionState.Created | SessionState.Initializing => false + case _ => true + } + + private def cancelWithReason(reason: String): Cancel = + Cancel.newBuilder().setReason(reason).build() + + /** + * The protocol transition graph in one place: names for the edges, not a + * second source of truth. Every edge acts on the single + * [[WorkerSession.SessionState]] owned by the base via [[compareAndSetState]] + * (non-terminal) or [[completeTerminal]] (terminal), so the base's CAS is the + * only synchronization and a terminal that arrived first always wins (the + * non-terminal CASes fail against it; [[completeTerminal]] is first-wins). + * + * Edges this class drives -- edge, method, then driver site(s) / thread (the + * base drives the API-call edges: `Created -> Initializing` in `init`, + * `Initialized -> Streaming` in `process`): + * {{{ + * Initializing -> Initialized initAccepted handleControl INIT-ok [gRPC cb] + * Streaming -> Finishing beginFinish advance branch 3 [engine] + * non-terminal -> Cancelling beginCancelFrom sendCancelInternal, immediately + * before the Cancel send [gRPC cb | engine | init | close] + * non-terminal -> Terminal finished/cancelled/transportFailed/interrupted + * handleControl / doInit / doClose / advance / onError + * }}} + * `Cancelling` is reached only when a Cancel send starts; a pre-stream or + * raced cancel goes straight to a `Cancelled`/`TransportFailed` terminal (see + * [[sendCancelInternal]], [[doClose]], `ProcessIterator`) or nowhere -- so + * `cancelRequested` (intent) does not imply state `Cancelling`. + */ + private object Transitions { + /** `InitResponse` OK: `Initializing -> Initialized`. */ + def initAccepted(): Boolean = + compareAndSetState(SessionState.Initializing, SessionState.Initialized) + + /** Input exhausted: `Streaming -> Finishing` (once). */ + def beginFinish(): Boolean = + compareAndSetState(SessionState.Streaming, SessionState.Finishing) + + /** `Cancel` send starts: `cur -> Cancelling`, from any non-terminal `cur`. */ + def beginCancelFrom(cur: SessionState): Boolean = + !cur.isTerminal && compareAndSetState(cur, SessionState.Cancelling) + + /** Clean terminal carrying the worker's `FinishResponse`. */ + def finished(response: FinishResponse): Boolean = + completeTerminal(Termination.Finished(response)) + + /** Clean terminal carrying the worker's `CancelResponse`. */ + def cancelled(response: CancelResponse): Boolean = + completeTerminal(Termination.Cancelled(response)) + + /** Failure terminal carrying a transport-level cause. */ + def transportFailed(cause: Throwable): Boolean = + completeTerminal(Termination.TransportFailed(cause)) + + /** + * Terminal for an engine-thread interrupt (e.g. a Spark task kill). Distinct + * from [[transportFailed]] because the cause is an engine-side interrupt, but + * still unsalvageable without a worker acknowledgement. Settled by + * [[handleInterrupt]] only when the brief post-Cancel drain expires without a + * `CancelResponse`; if the worker acks in time the session settles a + * cooperative [[cancelled]] instead. + */ + def interrupted(cause: Throwable): Boolean = + completeTerminal(Termination.Interrupted(cause)) + } + + // ---- WorkerSession hooks ------------------------------------------------ + + override protected def doInit(message: Init): InitResponse = { + // Construct the request before opening the RPC. A malformed Init must not + // leave an otherwise-unused Execute stream waiting for its first request. + val initRequest = UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setInit(message).build()) + .build() + + // Fail fast if the channel is already shut down. Without this check, + // asyncStub.execute(...) would still succeed and the failure would + // surface ~initResponseTimeoutMs later as a misleading "InitResponse + // timed out" error. + if (channel.getState(false) == ConnectivityState.SHUTDOWN) { + val ex = new IllegalStateException("gRPC channel is shut down") + Transitions.transportFailed(ex) + throw new GrpcWorkerSessionException("UDF worker channel is closed", ex) + } + try { + requestLock.synchronized { + // Serialize stream creation, publication, and Init with close/cancel. A + // close that wins the lock first prevents the RPC from opening; one that + // loses cannot release the worker until the initial write returns. + if (!currentState.isTerminal) { + val stream = asyncStub.execute(responseObserver) + requestObserver = stream + // Stream creation may invoke a response callback synchronously. Do not + // send Init if that callback has already settled a terminal. + if (!currentState.isTerminal) { + // With directExecutor, gRPC callbacks run synchronously on the caller's + // thread, so InitResponse can reach responseObserver/handleControl from + // *inside* this stream.onNext. + // requestObserver is already published, so an immediate ErrorResponse + // can write its required Cancel reentrantly without losing the request. + sendOnNext(initRequest) + } + } + } + } catch { + case NonFatal(e) => + Transitions.transportFailed(e) + // Surface as GrpcWorkerSessionException so the engine integration layer + // (which catches that type and wraps it) sees a uniform init-failure + // exception rather than the raw transport error. + throw new GrpcWorkerSessionException("UDF worker stream failed during init", e) + } + + try { + initValue.await(initResponseTimeoutMs) + } catch { + case _: InterruptedException => + // Interrupt (cancelled query / killed task) while awaiting InitResponse: + // cooperatively Cancel with a bounded drain, settling Cancelled if the + // worker acks in time, else Interrupted. See handleInterrupt. + handleInterrupt("waiting for InitResponse") + throw new InterruptedException("interrupted while waiting for InitResponse") + case e: TimeoutException => + sendCancelInternal(() => cancelWithReason("InitResponse timed out")) + // Settle the terminal so a subsequent close() does not stall for a + // second `terminalTimeoutMs` waiting for a worker that already missed + // its init deadline. + Transitions.transportFailed(e) + // Surface as GrpcWorkerSessionException (carrying the timeout cause) so + // the engine integration layer that catches that type can wrap it. + throw new GrpcWorkerSessionException( + s"timed out waiting for InitResponse after ${initResponseTimeoutMs}ms", e) + } + + initValue.get match { + case Some(resp) if resp.hasError => + failInitWithError(resp.getError, + s"UDF worker init failed: ${describeError(resp.getError)}") + case Some(resp) => + executionError.get() match { + case Some(err) => + // A reentrant callback can report an ErrorResponse immediately after + // InitResponse and before the initial onNext returns. The callback + // has already sent Cancel; drain its response + // and surface the original error instead of returning init success + // for a session that is already cancelling or terminal. + failInitWithError(err, + s"UDF worker reported an error as init completed: ${describeError(err)}") + case None if currentState != SessionState.Initialized => + // Likewise, an immediate protocol failure or concurrent close may + // have moved the session out of Initialized after publishing the + // InitResponse. A normal return would violate init()'s contract and + // leave process() to fail later with only an ordering error. + failInitFromCurrentState() + case None => + resp + } + case None => + // No InitResponse arrived but the latch fired. + // + // A pre-init ErrorResponse leaves the session in Initializing with the + // error recorded in executionError (see handleControl ERROR branch): the + // engine must now Cancel and drain the CancelResponse, which failInit does. + executionError.get() match { + case Some(err) if !currentState.isTerminal => + failInitWithError(err, + s"UDF worker reported an error before init completed: ${describeError(err)}") + case _ => + // Otherwise a terminal already settled -- the worker terminated the + // stream before sending InitResponse (transport error, half-close, or + // a premature FinishResponse/CancelResponse). Surface it as an init + // failure rather than letting the caller proceed as if init succeeded. + failInitFromCurrentState() + } + } + } + + /** Surfaces the terminal that prevented init from completing normally. */ + private def failInitFromCurrentState(): Nothing = { + executionError.get() match { + case Some(err) => + failInitWithError(err, + s"UDF worker reported an error as init completed: ${describeError(err)}") + case None => () + } + // Cancelling can be observed after a concurrent close or a reentrant error + // before its CancelResponse arrives. Drain it so the exception reflects the + // stable terminal rather than a transient state. + if (currentState == SessionState.Cancelling) { + awaitTerminal() + } + currentState match { + case SessionState.Terminal(Termination.TransportFailed(cause)) => + throw new GrpcWorkerSessionException("UDF worker stream failed during init", cause) + case SessionState.Terminal(Termination.Failed(err)) => + throw new GrpcWorkerSessionException( + s"UDF worker reported an error before init completed: ${describeError(err)}", err) + case SessionState.Terminal(Termination.Cancelled(_)) => + throw new GrpcWorkerSessionException( + "UDF worker stream was cancelled before init completed") + case SessionState.Terminal(Termination.Interrupted(cause)) => + throw new GrpcWorkerSessionException("UDF worker init was interrupted", cause) + case SessionState.Terminal(Termination.Finished(_)) => + throw new GrpcWorkerSessionException("UDF worker finished before init completed") + case other => + throw new IllegalStateException( + s"init completed without an accepted session or terminal: $other") + } + } + + /** + * Fails init when the worker reports an error before init returns: an + * `InitResponse` carrying an error, a pre-init `ErrorResponse`, or an immediate + * post-init `ErrorResponse`. The protocol requires the engine to send `Cancel` + * and the worker to reply with `CancelResponse` (udf_message.proto); we send it, + * drain the terminator so no stream is left dangling, and throw the structured + * error. A responsive worker settles `Cancelled`; a failed drain settles + * `TransportFailed`. + */ + private def failInitWithError(err: ExecutionError, message: String): Nothing = { + sendCancelInternal(() => cancelWithReason("init failed")) + awaitTerminal() + throw new GrpcWorkerSessionException(message, err) + } + + override protected def doProcess( + input: Iterator[DataRequest], + finish: () => Finish): Iterator[DataResponse] = { + // Init success is guaranteed by the base [[WorkerSession]] lifecycle: if + // doInit had failed it would have thrown and process() would never run. + new ProcessIterator(input, finish) + } + + override protected def doClose(cancel: () => Cancel): Termination = { + // Coordinate the no-stream decision with doInit's publication + Init write. + // If close wins this lock, doInit observes the terminal and never sends Init; + // if doInit wins, close observes the published observer and sends Cancel only + // after the Init send has returned. + val noPublishedStream = requestLock.synchronized { + if (requestObserver == null) { + // init() never put a stream on the wire (closed before/around init, or + // init threw before publishing). There is no protocol terminator; if no + // terminal has settled yet, treat the session as cancelled-before-start. + // A terminal may already be settled here (e.g. the channel-shutdown + // TransportFailed in doInit also leaves requestObserver null); return the + // settled terminal as-is rather than a bare Cancelled that disagrees with + // the state. The base WorkerSession still releases the worker handle. + if (!currentState.isTerminal) { + Transitions.cancelled(CancelResponse.getDefaultInstance) + } + true + } else { + false + } + } + if (noPublishedStream) { + return settledTermination + } + // If the stream has not finished on its own, cancel anything in flight so + // the worker can clean up, then wait for the terminator. sendCancelInternal + // evaluates the cancel thunk only when it attempts a Cancel, and at + // most once across all callers. + if (!currentState.isTerminal) { + sendCancelInternal(cancel) + try { + terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => + // Interrupted mid-close: settle Interrupted via the + // shared handler rather than falling through to the TransportFailed + // guard below. The Cancel was already attempted, so handleInterrupt's + // sendCancelInternal is a no-op and it just runs the bounded drain. + // + // TODO [SPARK-57640]: close() is a finalizer often already on the task + // kill path, yet this makes an interrupted close block up to another + // interruptCancelTimeoutMs (chasing a clean, salvageable Cancelled + // rather than immediately unwinding to Interrupted). Current choice: + // spend the bounded wait to obtain proof the worker is recyclable; it + // is small next to terminalTimeoutMs. Revisit if killed-task teardown + // latency (especially many sessions torn down at once) makes an + // immediate unwind preferable -- e.g. skip this second drain in the + // close() path, since close already gave the worker its terminalTimeoutMs window. + handleInterrupt("closing the session") + } + } + // If the worker still has not settled (timeout above, or an interrupt whose + // bounded drain did not settle a terminal), record a terminal so + // isWorkerSalvageable returns a definite answer and any other thread reading + // the state sees a stable value. + if (!currentState.isTerminal) { + Transitions.transportFailed(new TimeoutException( + s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")) + } + // Close the request side according to the settled outcome. Clean protocol + // terminators are half-closed; failures abort the RPC so close() never follows + // an onError with onCompleted. + currentState match { + case SessionState.Terminal(Termination.Finished(_)) | + SessionState.Terminal(Termination.Cancelled(_)) => + completeRequestSide() + case SessionState.Terminal(Termination.Failed(error)) => + abortRequestSide(new GrpcWorkerSessionException( + s"UDF execution failed: ${describeError(error)}", error)) + case SessionState.Terminal(Termination.TransportFailed(cause)) => + abortRequestSide(cause) + case SessionState.Terminal(Termination.Interrupted(cause)) => + abortRequestSide(cause) + case other => + logger.debug(s"UDF Execute stream closed without a terminal outcome: $other") + } + // Derive the Termination from the settled terminal. close() is the + // finalizer/cleanup path and must NOT re-throw: a UDF / data-phase error is + // already surfaced while the result iterator is consumed, and an init error + // is surfaced from init(). settledTermination returns the settled terminal + // as-is: the response proto for the clean terminators, and the failure + // terminals (Failed / TransportFailed / Interrupted) carrying their cause -- + // not a bare Cancelled. + settledTermination + } + + // ---- Internal request helpers --------------------------------------------- + + /** + * Writes one request and aborts the request side if the observer rejects it. + * Callers hold [[requestLock]], so this also serializes the compensating + * `onError` with every other request-side operation. + */ + private def sendOnNext(req: UdfRequest): Unit = { + try { + requestObserver.onNext(req) + } catch { + case NonFatal(e) => + abortRequestSide(e) + throw e + } + } + + /** Aborts the request side at most once. Safe to call while holding [[requestLock]]. */ + private def abortRequestSide(cause: Throwable): Unit = requestLock.synchronized { + if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) { + try { + requestObserver.onError(cause) + } catch { + case NonFatal(e) => logger.debug("Error aborting UDF Execute stream", e) + } + } + } + + /** Half-closes the request side at most once. */ + private def completeRequestSide(): Unit = requestLock.synchronized { + if (requestObserver != null && requestSideTerminated.compareAndSet(false, true)) { + try { + requestObserver.onCompleted() + } catch { + case NonFatal(e) => logger.debug("Error half-closing UDF Execute stream", e) + } + } + } + + /** + * Sends a Data or Finish request to the worker. Three invariants are + * checked inside [[requestLock]]: + * - terminal state: writes are unsafe (transport dead or terminal + * received). Throws, so the caller's terminal/exception path runs. + * - request-side termination: no write may follow `onError` / `onCompleted`. + * Throws so the caller settles or surfaces the failure. + * - [[cancelRequested]]: a Cancel has been (or is about to be) sent. + * Silently no-ops so a Data/Finish never appears on the wire after + * Cancel; the caller's iterator falls through to the terminator wait. + * + * Init is NOT sent through this helper -- [[doInit]] publishes + * [[requestObserver]] and writes Init together under [[requestLock]], so a + * concurrent cancel observes the observer but cannot acquire the lock and send + * Cancel until after Init. + */ + private def sendRequest(req: UdfRequest): Unit = + requestLock.synchronized { + if (currentState.isTerminal) { + throw new IllegalStateException( + "cannot send request: UDF Execute stream is already closed") + } + if (requestSideTerminated.get()) { + throw new IllegalStateException( + "cannot send request: UDF Execute request stream is already closed") + } + // Suppress the write when a cancel has been (or is about to be) flushed + // through this lock; that preserves the proto ordering invariant (no + // Data/Finish after Cancel). Test the flag rather than `return`-ing from + // this by-name `synchronized` body: a non-local return compiles to a thrown + // NonLocalReturnControl, which is brittle (also relied on in sendCancelInternal). + if (!cancelRequested.get()) { + if (req.hasControl && req.getControl.hasFinish) { + // Set before onNext: directExecutor delivery may return FinishResponse + // reentrantly from inside this call. + finishSendStarted.set(true) + } + sendOnNext(req) + } + } + + /** + * Sends a `Cancel` control message, returning `true` iff the request observer + * accepts the `Cancel`. Idempotent across ALL call sites (close()'s + * cancel, in-band cancels from `handleControl`'s INIT error / ERROR paths, + * and engine-internal cancels from the iterator) via [[cancelRequested]]: + * only the first caller attempts a `Cancel`, and the `cancel` thunk is evaluated + * only by that caller and only when the send is about to start -- + * so a side-effecting thunk (e.g. one carrying a client cancel callback) runs + * at most once. Setting [[cancelRequested]] also blocks subsequent Data/Finish + * writes from [[ProcessIterator]] (see [[sendRequest]]), preserving the proto + * invariant that nothing follows `Cancel` on the engine-to-worker side. + */ + private def sendCancelInternal(cancel: () => Cancel): Boolean = { + // Do not consume the one-shot cancellation intent until there is an observer + // that can carry it. doInit publishes the observer before sending Init while + // holding requestLock, so every response to Init sees a non-null observer. + if (requestObserver == null) return false + if (!cancelRequested.compareAndSet(false, true)) return false + if (currentState.isTerminal) return false + try { + requestLock.synchronized { + // A terminal that arrived first must win; bail without writing. Compute + // the block's value rather than `return` (see sendRequest for why non-local + // return is avoided -- here it would also escape the NonFatal catch below). + val cur = currentState + if (cur.isTerminal || requestSideTerminated.get()) { + false + } else { + val request = UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setCancel(cancel()).build()) + .build() + // Record that a Cancel send has started by advancing to Cancelling. + Transitions.beginCancelFrom(cur) + sendOnNext(request) + true + } + } + } catch { + case NonFatal(e) => + logger.debug(s"Cancel send failed (stream may already be torn down): ${e.getMessage}") + Transitions.transportFailed(e) + false + } + } + + private def awaitTerminal(): Unit = { + if (currentState.isTerminal) return + try { + if (!terminalLatch.await(terminalTimeoutMs, TimeUnit.MILLISECONDS)) { + Transitions.transportFailed(new TimeoutException( + s"timed out waiting for stream terminator after ${terminalTimeoutMs}ms")) + } + } catch { + case _: InterruptedException => handleInterrupt("waiting for stream terminator") + } + } + + /** + * Handles an [[InterruptedException]] observed while blocked on a worker event + * (init, terminal drain, close, or the result-iterator poll). An interrupt is + * an engine-side event -- typically a cancelled query / killed task -- not a + * worker fault, so we cancel cooperatively and keep the worker salvageable only + * when it acks: + * + * 1. Send a best-effort `Cancel` (idempotent across call sites). + * 2. Wait up to [[interruptCancelTimeoutMs]] -- short, so the interrupted + * thread unwinds promptly -- for the worker's `CancelResponse`. A healthy + * worker acks in that window and the callback settles a clean `Cancelled` + * terminal (worker salvageable, response drained). + * 3. If the ack does not arrive in time, settle `Interrupted`. Without a + * terminator or liveness proof the worker is not salvageable and must not + * be returned to a reuse pool. + * + * The bounded wait deliberately runs '''before''' the thread's interrupt flag + * is restored: [[InterruptedException]] clears the flag on throw, so re-setting + * it first would make the wait below throw immediately and defeat the drain. + * The flag is restored at the end so the caller's unwind still observes the + * interrupt. + * + * TODO [SPARK-57640]: add a worker liveness/heartbeat check so a future version + * may prove that an interrupted worker which missed the short ack window is + * nevertheless safe to recycle. + */ + private def handleInterrupt(waitContext: String): Unit = { + // NB: flag is currently clear (InterruptedException cleared it); do not + // restore it until after the bounded drain below. + if (!currentState.isTerminal) { + sendCancelInternal(() => cancelWithReason(s"interrupted while $waitContext")) + val acked = try { + terminalLatch.await(interruptCancelTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => + // A second interrupt during the drain: give up the wait immediately. + false + } + if (!acked && !currentState.isTerminal) { + // No CancelResponse in time: settle Interrupted rather than + // TransportFailed, and stop a later close() from blocking on a terminator + // that will not be drained on this thread's behalf. Interrupted remains + // unsalvageable until a separate liveness check can prove otherwise. + Transitions.interrupted( + new InterruptedException(s"interrupted while $waitContext")) + } + } + Thread.currentThread().interrupt() + } + + // ---- ProcessIterator ------------------------------------------------------ + + /** + * Iterator returned by [[doProcess]]. Drives the data phase of the + * stream end-to-end: each call to `hasNext` / `next` may send a + * `DataRequest` or `Finish` to the worker, and reads result batches + * out of [[outputQueue]] as the worker emits them. + * + * The iterator is single-threaded with respect to the engine, but + * coexists with the gRPC callback thread (which enqueues responses + * and may settle the state) and with any thread that finalizes via + * [[close]]. It drives the `Streaming -> Finishing` transition (sending + * `Finish` once input is exhausted, built lazily from the `finish` thunk). + */ + private class ProcessIterator(input: Iterator[DataRequest], finish: () => Finish) + extends Iterator[DataResponse] { + + // Latched once the terminator sentinel has been observed. Without this, + // a second hasNext() after the iterator naturally exhausts would re-enter + // advance(), fall through all branches, and block branch 4 for the + // full terminalTimeoutMs before returning. Callers that probe hasNext + // an extra time (iterator.size, instrumentation wrappers) would hang. + // Iterator-local and only touched by the single engine thread. + private val exhausted = new AtomicBoolean(false) + @volatile private var prefetched: DataResponse = _ + + /** + * Runs a caller-supplied thunk (`input.hasNext`, `input.next()`, or the + * `finish` builder) that may throw, and on failure Cancels the in-flight + * stream before rethrowing. Without the Cancel a failure would leave the + * worker awaiting more input with no terminator owed on the wire; with it the + * worker tears down and the exception still propagates to the engine. Setting + * `cancelRequested` also suppresses any further Data/Finish this loop might + * otherwise attempt (see [[sendRequest]]). An [[InterruptedException]] uses + * [[handleInterrupt]] so its CancelResponse drain remains bounded by + * [[interruptCancelTimeoutMs]]. + */ + private def cancelOnThrow[T](reason: String)(op: => T): T = + try op catch { + case _: InterruptedException => + handleInterrupt(reason) + throw new InterruptedException(s"interrupted while $reason") + case NonFatal(e) => + sendCancelInternal(() => cancelWithReason(reason)) + throw e + } + + override def hasNext: Boolean = { + if (prefetched ne null) return true + advance() + prefetched ne null + } + + override def next(): DataResponse = { + if (prefetched eq null) advance() + val out = prefetched + if (out eq null) { + throw new NoSuchElementException("ProcessIterator exhausted") + } + prefetched = null + out + } + + /** + * Fills [[prefetched]] with the next output batch, or leaves it null at the + * terminator (which also sets [[exhausted]], so `hasNext` reads null as "done"). + * Each loop iteration tries in order: (1) drain queued output -- a batch (fill + * and return) or the terminator sentinel (return); (2) while `Streaming`, send + * the next input batch and loop; (3) once input is exhausted, send `Finish` + * (CAS `Streaming -> Finishing`, once) and loop; (4) otherwise block for late + * output or the terminator, then return. Branches 2-3 loop; 1 and 4 return. + * + * '''Per-poll timeout.''' Branch 4 waits up to [[terminalTimeoutMs]] per + * poll, reset by every worker event -- the contract is "emit at least one + * event every [[terminalTimeoutMs]] after Finish", not "finish the UDF within + * [[terminalTimeoutMs]]". A worker expecting a long post-Finish silence MAY + * emit an empty `DataResponse` as a heartbeat to reset the wait; it is + * surfaced to the caller, so it should be a batch the caller recognises as + * empty (e.g. a zero-row Arrow batch). + */ + private def advance(): Unit = { + if (exhausted.get()) return // terminator already seen; never re-block + while (prefetched eq null) { + // (1) Drain anything the worker has already produced. + outputQueue.poll() match { + case null => // queue empty, fall through to send/wait branches below + case QueueItem.EndOfStream => + exhausted.set(true) + throwIfTerminalError() + return + case QueueItem.Batch(b) => + prefetched = b + return + } + + // (2) Send next input batch while the stream is open for data. + // `input.hasNext` may itself fetch/compute the next element (many Spark + // iterators prefetch), so it can throw just like `input.next()`; both go + // through cancelOnThrow so a failing upstream Cancels the stream instead + // of stranding the worker waiting for input that will never arrive. + if (!cancelRequested.get() && currentState == SessionState.Streaming && + cancelOnThrow("input iterator failed")(input.hasNext)) { + val request = cancelOnThrow("input iterator failed") { + val batch = Objects.requireNonNull( + input.next(), "input iterator returned null") + UdfRequest.newBuilder().setData(batch).build() + } + if (!sendOrEndOnRacedTerminal(request)) return + } else if (!cancelRequested.get() && Transitions.beginFinish()) { + // (3) No more input; send Finish exactly once (unless cancelled). The + // `finish` thunk is caller-supplied (it may run a finish callback), so + // it can throw; cancelOnThrow Cancels the stream before rethrowing. The + // Streaming -> Finishing CAS has already run, but we have not written + // Finish yet, so the Cancel is the only engine-to-worker message that + // reaches the wire -- the proto "nothing after Cancel" invariant holds. + val request = cancelOnThrow("finish callback failed") { + val finishMsg = Objects.requireNonNull( + finish(), "finish callback returned null") + UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder().setFinish(finishMsg).build()) + .build() + } + if (!sendOrEndOnRacedTerminal(request)) { + return + } + } else { + // (4) Block for late output or the terminator. See class doc + // above for the per-poll-vs-total-session timeout semantics. + val item = try { + outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS) + } catch { + case _: InterruptedException => + // Interrupt (cancelled query / killed task) while reading results: + // cooperatively Cancel with a bounded drain, settling Cancelled if + // the worker acks in time, else Interrupted. See handleInterrupt. + handleInterrupt("reading UDF result") + exhausted.set(true) + throw new InterruptedException("interrupted while reading UDF result") + } + item match { + case null => + Transitions.transportFailed(new IllegalStateException( + s"timed out waiting for UDF output after ${terminalTimeoutMs}ms")) + exhausted.set(true) + throwIfTerminalError() + return + case QueueItem.EndOfStream => + exhausted.set(true) + throwIfTerminalError() + return + case QueueItem.Batch(b) => + prefetched = b + return + } + } + } + } + + /** + * Sends one data-phase request (`DataRequest` / `Finish`), recovering from a + * terminator that raced the write. Returns `true` to keep looping (sent, or + * suppressed by a pending cancel), `false` if a terminator settled and the + * iterator should end (caller `return`s). + * + * On a write failure a terminator has usually already settled (the worker + * finished/failed early) and the write only failed because the stream is + * closed: record a transport terminal if none is set ([[completeTerminal]] is + * a no-op once settled), then [[throwIfTerminalError]] -- which throws for an + * error/transport terminal, or returns for a clean `Finished` that raced the + * send, in which case the benign "stream closed" error is dropped. + */ + private def sendOrEndOnRacedTerminal(req: UdfRequest): Boolean = { + try { + sendRequest(req) + true + } catch { + case NonFatal(e) => + if (!currentState.isTerminal) { + Transitions.transportFailed(e) + } + exhausted.set(true) + throwIfTerminalError() + false + } + } + + /** + * Surfaces a failed terminator as an exception when the result iterator is + * drained. A data-phase [[ExecutionError]] (captured in [[executionError]]) + * takes precedence over the terminator's own error, then the finish/cancel + * callback error carried on the terminator, then a bare cancellation. + */ + private def throwIfTerminalError(): Unit = currentState match { + case SessionState.Terminal(Termination.Finished(response)) => + responseError(response.hasError, response.getError, executionError.get()) + .foreach(err => throw new GrpcWorkerSessionException( + s"UDF execution failed: ${describeError(err)}", err)) + case SessionState.Terminal(Termination.Cancelled(response)) => + responseError(response.hasError, response.getError, executionError.get()) match { + case Some(err) => + throw new GrpcWorkerSessionException( + s"UDF execution failed: ${describeError(err)}", err) + case None => + throw new GrpcWorkerSessionException("UDF execution was cancelled") + } + case SessionState.Terminal(Termination.Failed(err)) => + throw new GrpcWorkerSessionException( + s"UDF execution failed: ${describeError(err)}", err) + case SessionState.Terminal(Termination.TransportFailed(t)) => + throw new GrpcWorkerSessionException("UDF worker stream failed", t) + case SessionState.Terminal(Termination.Interrupted(t)) => + throw new GrpcWorkerSessionException("UDF execution was interrupted", t) + case other => + throw new IllegalStateException(s"terminator sentinel without terminal: $other") + } + + /** Picks the error to surface: prior data-phase error, else the terminator's. */ + private def responseError( + hasError: Boolean, + error: ExecutionError, + priorError: Option[ExecutionError]): Option[ExecutionError] = + priorError.orElse(if (hasError) Some(error) else None) + } +} + +object GrpcWorkerSession { + /** Upper bound on the wait for `InitResponse`. */ + val DEFAULT_INIT_RESPONSE_TIMEOUT_MS: Long = 30000L + + /** Upper bound on the wait for `FinishResponse` / `CancelResponse`. */ + val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L + + /** + * Upper bound on the wait for a `CancelResponse` after an interrupt (e.g. a + * cancelled query / killed task) sends `Cancel`. Much shorter than + * [[DEFAULT_TERMINAL_TIMEOUT_MS]]: the interrupted thread must unwind + * promptly, so a healthy worker gets a brief window to ack the Cancel (clean + * `Cancelled` terminal, worker salvageable) before the session falls back to + * an unsalvageable `Interrupted` terminal. See `handleInterrupt`. + */ + val DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS: Long = 2000L + + // Distinguishes a (possibly empty) data batch from end-of-stream, and makes + // the iterator's match exhaustive. + private sealed trait QueueItem + private object QueueItem { + final case class Batch(response: DataResponse) extends QueueItem + case object EndOfStream extends QueueItem + } + + /** + * A write-once value paired with the latch a waiter blocks on, so the "publish + * the value, then release the waiter" ordering lives in one place instead of + * every call site having to remember to count the latch down after setting the + * reference. The value is set at most once ([[complete]]); the latch can also + * be released without a value ([[signalWithoutValue]]) when init fails through + * a pre-init error or terminal. All reads/writes go through the latch, so a + * waiter released by [[await]] has a happens-before edge to the [[complete]] + * that set the value. + */ + private final class OneShotValue[A] { + private val latch = new CountDownLatch(1) + private val ref = new AtomicReference[Option[A]](None) + private val completed = new AtomicBoolean(false) + + /** Publishes the value (first writer wins) and releases any waiter. */ + def complete(value: A): Unit = { + if (completed.compareAndSet(false, true)) { + ref.set(Some(value)) + latch.countDown() + } + } + + /** Completes without a value (init failed through another event) and releases the waiter. */ + def signalWithoutValue(): Unit = { + if (completed.compareAndSet(false, true)) { + latch.countDown() + } + } + + /** + * Blocks up to `timeoutMs` for a release, throwing [[TimeoutException]] if + * none arrives so the caller handles the timeout on the exception path rather + * than by inspecting a return value. + */ + def await(timeoutMs: Long): Unit = + if (!latch.await(timeoutMs, TimeUnit.MILLISECONDS)) { + throw new TimeoutException( + s"timed out waiting for value after ${timeoutMs}ms") + } + + /** The published value, or None if none was ever set. */ + def get: Option[A] = ref.get() + } + + private[grpc] def describeError(err: ExecutionError): String = err.getKindCase match { + case ExecutionError.KindCase.USER => + val u = err.getUser + val cls = if (u.hasErrorClass) s"[${u.getErrorClass}] " else "" + s"$cls${u.getMessage}" + case ExecutionError.KindCase.WORKER => + s"WorkerError: ${err.getWorker.getMessage}" + case ExecutionError.KindCase.PROTOCOL => + s"ProtocolError: ${err.getProtocol.getMessage}" + case ExecutionError.KindCase.KIND_NOT_SET => + "ExecutionError without kind" + } +} + +/** + * :: Experimental :: + * Exception thrown by [[GrpcWorkerSession]] when the UDF execution fails + * at the engine-protocol layer: init failure, ErrorResponse from the + * worker, a failure terminator with no response, transport failure, or + * (from the result iterator) a cancellation. + * + * This extends plain [[RuntimeException]] rather than Spark's + * `SparkRuntimeException` so the udf-worker modules stay free of a spark-core + * dependency. The engine integration layer (which already depends on + * spark-core) is expected to catch this and wrap it in a + * `SparkRuntimeException` with an appropriate error class when surfacing UDF + * failures to users. [[executionError]] is preserved to carry the structured + * cause across that boundary. + * + * @param executionError the structured protocol error, when present: a worker + * `ErrorResponse`, an init error, or a failure terminator + * with an error. `null` when there is no structured cause: + * a transport failure, a timeout, or a cancellation without + * an error. Callers must null-check before use. + */ +@Experimental +class GrpcWorkerSessionException( + message: String, + cause: Throwable = null, + @javax.annotation.Nullable val executionError: ExecutionError = null) + extends RuntimeException(message, cause) { + + def this(message: String, error: ExecutionError) = + this(message, null, error) +} diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala index eae6ef1a639a9..4de8342208399 100644 --- a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/EchoProtocolSuite.scala @@ -469,6 +469,10 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { @volatile var executionError: Option[ExecutionError] = None @volatile var streamError: Option[Throwable] = None private val requestCompleted = new AtomicBoolean(false) + // gRPC forbids concurrent calls to the request StreamObserver too: the response + // observer half-closes it (completeRequestStream) on a callback thread while the + // test thread may still be sending a trailing Cancel. Serialize both through this. + private val requestLock = new Object // Counted down on InitResponse (success or failure) or on terminal error. // The engine MUST wait for this before sending any DataRequest or Finish. private val initResponseLatch = new CountDownLatch(1) @@ -637,22 +641,23 @@ class EchoProtocolSuite extends AnyFunSuite with BeforeAndAfterEach { } } - def sendCancel(reason: String = ""): Unit = { - // If a terminator already arrived (FinishResponse / CancelResponse), - // the request stream has been half-closed and Cancel arrives too - // late -- silently ignore, matching the proto's Cancel-after-Finish - // contract. - if (requestCompleted.get()) return - requestObserver.onNext(UdfRequest.newBuilder() - .setControl(UdfControlRequest.newBuilder() - .setCancel(Cancel.newBuilder().setReason(reason).build()) + def sendCancel(reason: String = ""): Unit = requestLock.synchronized { + // If a terminator already arrived (FinishResponse / CancelResponse), the request + // stream has been half-closed and Cancel arrives too late -- silently ignore, + // matching the proto's Cancel-after-Finish contract. requestLock makes this check + // and the send atomic against completeRequestStream's half-close. + if (!requestCompleted.get()) { + requestObserver.onNext(UdfRequest.newBuilder() + .setControl(UdfControlRequest.newBuilder() + .setCancel(Cancel.newBuilder().setReason(reason).build()) + .build()) .build()) - .build()) + } // Request stream stays open until the response terminator arrives; // completeRequestStream() is called by the response observer. } - def completeRequestStream(): Unit = { + def completeRequestStream(): Unit = requestLock.synchronized { if (requestCompleted.compareAndSet(false, true)) { requestObserver.onCompleted() } diff --git a/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala new file mode 100644 index 0000000000000..fb0eeae9ed376 --- /dev/null +++ b/udf/worker/grpc/src/test/scala/org/apache/spark/udf/worker/grpc/GrpcWorkerSessionConcurrencySuite.scala @@ -0,0 +1,1581 @@ +/* + * 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.udf.worker.grpc + +import java.util.Locale +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicReference} + +import scala.jdk.CollectionConverters._ + +import com.google.protobuf.ByteString +import io.grpc.{CallOptions, ClientCall, ConnectivityState, ForwardingClientCall, ManagedChannel, + Metadata, MethodDescriptor, Server} +import io.grpc.inprocess.{InProcessChannelBuilder, InProcessServerBuilder} +import io.grpc.stub.StreamObserver +import org.scalatest.BeforeAndAfterEach +// scalastyle:off funsuite +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.udf.worker.{Cancel, CancelResponse, DataRequest, DataResponse, + ErrorResponse, ExecutionError, Finish, FinishResponse, Init, InitResponse, UdfControlResponse, + UdfPayload, UdfRequest, UdfResponse, UDFWorkerDataFormat, UdfWorkerGrpc, UserError, + WorkerRequest, WorkerResponse} +import org.apache.spark.udf.worker.core.{Termination, WorkerHandle, WorkerLogger} + +/** + * Concurrency tests for [[GrpcWorkerSession]] that pin the wire-ordering and + * fast-fail invariants under concurrent and worker-misbehavior scenarios: + * - Cancel must never appear on the wire before Init. + * - A worker terminator (ERROR / FINISH / CANCEL / onCompleted) arriving + * before InitResponse must fail [[GrpcWorkerSession#init]] fast, not hang + * for `initResponseTimeoutMs`. + * - Repeated [[Iterator#hasNext]] after natural iterator exhaustion must + * return immediately, not block for `terminalTimeoutMs`. + * - Close racing the initial write must either prevent Init or send Cancel + * after Init without releasing the worker underneath an active stream. + * - An immediate ErrorResponse after InitResponse must fail init and send the + * protocol-required Cancel; a premature FinishResponse must not truncate input. + * - Cancel and close concurrent with an in-progress data phase terminate + * cleanly without leaks or unbounded hangs. + * + * Runs entirely in-process: no subprocess, no UDS. Server services are + * custom-built per test so we can drive specific worker misbehavior. + */ +class GrpcWorkerSessionConcurrencySuite + extends AnyFunSuite with BeforeAndAfterEach { +// scalastyle:on funsuite + + /** Used by tests to keep stale in-flight infra reachable for teardown. */ + private val openServers = new ConcurrentLinkedQueue[Server]() + private val openChannels = new ConcurrentLinkedQueue[ManagedChannel]() + private val openSessions = new ConcurrentLinkedQueue[GrpcWorkerSession]() + + override def afterEach(): Unit = { + // Shut channels down first. This fires onError on any still-live stream, + // which settles the session terminal and counts down the init/terminal + // latches. That unblocks both the session.close() below and any worker + // thread a failing test left parked on a (deliberately large) timeout, so + // teardown never hangs even when a test asserts via assertFinishesWithin. + openChannels.asScala.foreach { c => + try c.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () } + } + openChannels.clear() + openSessions.asScala.foreach { s => try s.close(emptyCancel) catch { case _: Throwable => () } } + openSessions.clear() + openServers.asScala.foreach { s => + try s.shutdownNow().awaitTermination(2, TimeUnit.SECONDS) catch { case _: Throwable => () } + } + openServers.clear() + super.afterEach() + } + + // A session timeout large enough that a correct test never reaches it; a + // regression that fails to short-circuit blocks here for minutes and is caught + // by assertFinishesWithin (below) instead of a flaky `elapsed < timeout` bound. + private val NeverReachedTimeoutMs = TimeUnit.MINUTES.toMillis(10) + + /** + * Runs `body` on a daemon thread and asserts it finishes within `withinMs`, + * rethrowing whatever `body` threw (so an `intercept` inside `body` still + * works). Pair with [[NeverReachedTimeoutMs]]: the correct fast path returns + * in milliseconds, so `withinMs` (seconds) has an enormous safety margin and + * does not flake, while a regression that parks on the timeout never returns + * within `withinMs` and fails the assertion. + */ + private def assertFinishesWithin(withinMs: Long, name: String)(body: => Unit): Unit = { + val thrown = new AtomicReference[Throwable]() + val done = new CountDownLatch(1) + val worker = new Thread(() => { + try body catch { case t: Throwable => thrown.set(t) } finally done.countDown() + }, name) + worker.setDaemon(true) + worker.start() + assert(done.await(withinMs, TimeUnit.MILLISECONDS), + s"$name did not finish within ${withinMs}ms; it parked on a session timeout " + + "that the fast path should have short-circuited") + Option(thrown.get()).foreach(t => throw t) + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private class TestWorkerHandle extends WorkerHandle { + val invalidated = new AtomicBoolean(false) + val released = new AtomicBoolean(false) + override def id: String = "test-worker" + override def markInvalid(): Unit = invalidated.set(true) + override def releaseSession(): Unit = released.set(true) + } + + /** + * Builds an in-process server/channel pair. + * + * @param directExecutor when `true` (default), responses are delivered + * reentrantly on the calling thread -- the worst case for the session's + * reentrancy handling. When `false`, both server and channel use their + * default executors, so responses arrive on a *separate* thread. That + * cross-thread delivery mirrors the production Netty transport (the gRPC + * callback runs on an event-loop thread, never reentrantly), which the + * directExecutor tests deliberately do not exercise. + */ + private def startServer( + service: UdfWorkerGrpc.UdfWorkerImplBase, + directExecutor: Boolean = true): (Server, ManagedChannel) = { + val name = InProcessServerBuilder.generateName() + val serverBuilder = InProcessServerBuilder.forName(name).addService(service) + val channelBuilder = InProcessChannelBuilder.forName(name) + if (directExecutor) { + serverBuilder.directExecutor() + channelBuilder.directExecutor() + } + val server = serverBuilder.build().start() + val channel = channelBuilder.build() + openServers.add(server) + openChannels.add(channel) + (server, channel) + } + + /** A [[ManagedChannel]] wrapper whose operations delegate to `underlying`. */ + private class DelegatingManagedChannel(underlying: ManagedChannel) extends ManagedChannel { + override def shutdown(): ManagedChannel = { + underlying.shutdown() + this + } + + override def isShutdown: Boolean = underlying.isShutdown + + override def isTerminated: Boolean = underlying.isTerminated + + override def shutdownNow(): ManagedChannel = { + underlying.shutdownNow() + this + } + + override def awaitTermination(timeout: Long, unit: TimeUnit): Boolean = + underlying.awaitTermination(timeout, unit) + + override def newCall[ReqT, RespT]( + methodDescriptor: MethodDescriptor[ReqT, RespT], + callOptions: CallOptions): ClientCall[ReqT, RespT] = + underlying.newCall(methodDescriptor, callOptions) + + override def authority(): String = underlying.authority() + + override def getState(requestConnection: Boolean): ConnectivityState = + underlying.getState(requestConnection) + + override def notifyWhenStateChanged( + source: ConnectivityState, + callback: Runnable): Unit = underlying.notifyWhenStateChanged(source, callback) + + override def resetConnectBackoff(): Unit = underlying.resetConnectBackoff() + + override def enterIdle(): Unit = underlying.enterIdle() + } + + private def newSession( + channel: ManagedChannel, + initResponseTimeoutMs: Long = 5000L, + terminalTimeoutMs: Long = 5000L, + interruptCancelTimeoutMs: Long = 5000L): GrpcWorkerSession = { + val session = new GrpcWorkerSession( + new TestWorkerHandle, channel, WorkerLogger.NoOp, + initResponseTimeoutMs = initResponseTimeoutMs, + terminalTimeoutMs = terminalTimeoutMs, + interruptCancelTimeoutMs = interruptCancelTimeoutMs) + openSessions.add(session) + session + } + + /** + * The [[TestWorkerHandle]] backing a session, for asserting the handle + * lifecycle. `workerHandle` is `private[worker]`, and this suite lives under + * `org.apache.spark.udf.worker`, so the access is in scope. + */ + private def handleOf(session: GrpcWorkerSession): TestWorkerHandle = + session.workerHandle.asInstanceOf[TestWorkerHandle] + + // Protocol version carried on Init. The in-process fake workers in this suite + // do not validate it; any sane value is fine. + private val SupportedVersion = 1 + + private def basicInit(payload: String = "echo"): Init = Init.newBuilder() + .setProtocolVersion(SupportedVersion) + .setDataFormat(UDFWorkerDataFormat.ARROW) + .setUdf(UdfPayload.newBuilder() + .setPayload(ByteString.copyFromUtf8(payload)) + .setFormat("echo") + .build()) + .build() + + // Default lifecycle messages for the data phase and finalization. + private val emptyFinish: () => Finish = () => Finish.getDefaultInstance + private val emptyCancel: () => Cancel = () => Cancel.getDefaultInstance + + /** Wraps strings as input [[DataRequest]] batches. */ + private def echoIn(batches: String*): Iterator[DataRequest] = + batches.iterator.map(s => DataRequest.newBuilder() + .setData(ByteString.copyFromUtf8(s)).build()) + + /** + * Captures every incoming request in order. Replies follow a user-supplied + * function, so tests can drive arbitrary worker misbehavior. The default + * `onRequest` matches an Echo worker: InitResponse on Init, echo on Data, + * FinishResponse on Finish, CancelResponse on Cancel. + */ + private class CapturingService( + val captured: ConcurrentLinkedQueue[UdfRequest] = new ConcurrentLinkedQueue(), + onRequest: (UdfRequest, StreamObserver[UdfResponse]) => Unit = null) + extends UdfWorkerGrpc.UdfWorkerImplBase { + + private val handler: (UdfRequest, StreamObserver[UdfResponse]) => Unit = + if (onRequest != null) onRequest else defaultEcho + + override def execute(resp: StreamObserver[UdfResponse]): StreamObserver[UdfRequest] = + new StreamObserver[UdfRequest] { + // gRPC requires serialized writes to a request StreamObserver; the + // capturing service may reply from multiple control paths so we + // synchronize on `resp` rather than relying on directExecutor. + override def onNext(req: UdfRequest): Unit = { + captured.add(req) + resp.synchronized { handler(req, resp) } + } + override def onError(t: Throwable): Unit = () + override def onCompleted(): Unit = resp.synchronized { resp.onCompleted() } + } + + override def manage( + request: WorkerRequest, + responseObserver: StreamObserver[WorkerResponse]): Unit = () + + private def defaultEcho(req: UdfRequest, resp: StreamObserver[UdfResponse]): Unit = { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + c.getControlCase match { + case _ if c.hasInit => + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build() + ).build()) + case _ if c.hasFinish => + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setFinish( + FinishResponse.getDefaultInstance).build() + ).build()) + case _ if c.hasCancel => + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build() + ).build()) + case _ => () + } + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + case _ => () + } + } + } + + // --------------------------------------------------------------------------- + // Cancel-never-precedes-Init invariant. Publication of requestObserver and the + // Init write are serialized with close/cancel under requestLock: close either + // wins first and prevents Init entirely, or waits until Init is on the wire + // before writing Cancel. + // --------------------------------------------------------------------------- + + // --------------------------------------------------------------------------- + // Terminator-before-InitResponse: must fail init fast, not hang for + // initResponseTimeoutMs (onTerminalSettled completes initValue without a value). + // --------------------------------------------------------------------------- + + test("worker emits ErrorResponse before InitResponse: init fails fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("simulated pre-init error") + .setErrorClass("PreInitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + // Huge init timeout: a regression that failed to complete initValue would + // park init here until the timeout, so finishing quickly is proof the fast + // path ran -- no wall-clock threshold needed. + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("error") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected init/error in message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("worker emits InitResponse with error: init fails fast (no terminal-timeout stall)") { + // Regression for directExecutor reentrancy where InitResponse(error) is + // delivered inside stream.onNext. requestObserver is already published, so + // doInit can send the required Cancel after onNext returns and drain the + // CancelResponse without stalling on terminalTimeoutMs. + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("simulated init failure") + .setErrorClass("InitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + // Huge terminalTimeoutMs: a regression (awaiting a CancelResponse that never + // arrives) would park init in awaitTerminal for the full timeout; the fix + // settles the terminal at once so init returns in milliseconds. Finishing + // well within assertFinishesWithin is the signal the stall did not happen. + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected an init-failure message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("terminal signal wins over a late InitResponse") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + // DATA before InitResponse is a protocol failure and completes the + // init one-shot without a value. A later InitResponse must not overwrite + // that completion and make init() return successfully. + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder() + .setData(ByteString.copyFromUtf8("premature")).build()) + .build()) + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("init") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("stream"), + s"expected the earlier protocol failure to win, got: ${ex.getMessage}") + assert(session.close(emptyCancel).isInstanceOf[Termination.TransportFailed]) + } + + test("ErrorResponse immediately after InitResponse fails init and sends Cancel") { + val cancelSeen = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + // Both responses are delivered reentrantly from inside the Init write. + // The second response is a valid generator-style data-phase failure. + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("generator failed") + .setErrorClass("GeneratorError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + cancelSeen.countDown() + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, terminalTimeoutMs = NeverReachedTimeoutMs) + + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.contains("generator failed"), + s"the immediate ErrorResponse must fail init with the worker error, got: ${ex.getMessage}") + assert(cancelSeen.await(5, TimeUnit.SECONDS), + "an immediate post-init ErrorResponse must send the protocol-required Cancel") + assert(session.close(emptyCancel).isInstanceOf[Termination.Cancelled]) + } + + test("FinishResponse before Finish is rejected instead of truncating input") { + val replied = new AtomicBoolean(false) + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (req.hasData && replied.compareAndSet(false, true)) { + // A buggy worker claims success after the first batch even though the + // engine has neither exhausted its input nor sent Finish. + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setFinish( + FinishResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, terminalTimeoutMs = NeverReachedTimeoutMs) + session.init(basicInit()) + + val ex = intercept[GrpcWorkerSessionException] { + session.process(echoIn("first", "must-not-be-consumed"), emptyFinish).foreach(_ => ()) + } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("stream failed"), + s"a premature FinishResponse must be a protocol failure, got: ${ex.getMessage}") + val dataRequests = service.captured.asScala.count(_.hasData) + assert(dataRequests == 1, + s"the premature terminator should stop input after one batch, got: $dataRequests") + assert(session.close(emptyCancel).isInstanceOf[Termination.TransportFailed]) + } + + test("worker emits FinishResponse before InitResponse: init fails fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setFinish( + FinishResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("finish") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected init/finish in message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("worker emits CancelResponse before InitResponse: init fails fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("cancel") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected init/cancel in message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("worker half-closes response stream before InitResponse: init fails fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onCompleted() + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + } + session.close(emptyCancel) + } + + test("worker emits malformed UdfResponse before InitResponse: init fails fast") { + // A UdfResponse with no oneof set (RESPONSE_NOT_SET) is malformed: the + // session's responseObserver settles a transport-failure terminal in its + // catch-all `case other` branch. Regression guard for failing to wake the + // initValue waiter when that terminal settles: without it, init() blocks + // until initResponseTimeoutMs and reports a misleading "timed out" error + // instead of failing fast with the malformed-response cause. + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.getDefaultInstance) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected init in message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + // --------------------------------------------------------------------------- + // Exhaustion guard: repeated hasNext after the iterator drains must not + // block for terminalTimeoutMs (exhausted flag fix). + // --------------------------------------------------------------------------- + + test("repeated hasNext after natural exhaustion returns immediately") { + val service = new CapturingService() + val (_, channel) = startServer(service) + // Huge terminalTimeoutMs: without the exhausted-flag short-circuit, a second + // hasNext() would re-enter the output-queue poll and park for the full + // timeout. Probing within assertFinishesWithin therefore proves the probes + // are non-blocking, with no dependence on a wall-clock threshold. + val session = newSession(channel, terminalTimeoutMs = NeverReachedTimeoutMs) + session.init(basicInit()) + val it = session.process(echoIn("hello"), emptyFinish) + assert(new String(it.next().getData.toByteArray) == "hello") + // Drain to terminator. + assert(!it.hasNext, "iterator should be exhausted after the single echo batch") + // Now probe many times; each call must return false without blocking. + assertFinishesWithin(10000, "repeated-hasNext") { + (1 to 5).foreach { _ => + assert(!it.hasNext, "exhausted iterator should keep returning false") + } + } + session.close(emptyCancel) + } + + // --------------------------------------------------------------------------- + // close() concurrent with process(): clean termination. + // --------------------------------------------------------------------------- + + test("close concurrent with process: iterator surfaces cancellation cleanly") { + // Worker echoes data but never sends FinishResponse (only after Cancel + // arrives, it sends CancelResponse). This pins the timing so the + // engine-side iterator is genuinely waiting when close() intervenes. + val readyToCancel = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + // Finish ignored: the worker never replies, so the iterator is + // blocked waiting on output until cancel intervenes. + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + readyToCancel.countDown() + case _ => () + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, terminalTimeoutMs = 30000L) + session.init(basicInit()) + + val handle = new AtomicReference[Throwable]() + val processThread = new Thread(() => { + try session.process(echoIn("hello"), emptyFinish).foreach(_ => ()) + catch { case t: Throwable => handle.set(t) } + }, "process-cancel") + processThread.start() + assert(readyToCancel.await(5, TimeUnit.SECONDS), + "worker never received the data batch") + // close() from another thread is the cancellation trigger: it sends Cancel, + // the worker replies CancelResponse, and the in-flight iterator surfaces it. + session.close(emptyCancel) + processThread.join(10000) + assert(!processThread.isAlive, "process thread should terminate after close") + val t = handle.get() + assert(t != null, "expected the iterator to surface a cancellation exception") + assert(t.isInstanceOf[GrpcWorkerSessionException], + s"expected GrpcWorkerSessionException, got ${t.getClass.getName}") + } + + // --------------------------------------------------------------------------- + // Close concurrent with in-progress process(): bounded, no hang. + // --------------------------------------------------------------------------- + + test("close concurrent with process: bounded shutdown, no leak") { + val readyToClose = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + readyToClose.countDown() + case _ => () + } + }) + val (_, channel) = startServer(service) + // Huge terminalTimeoutMs: the worker replies CancelResponse, so a correct + // close() returns on that terminator in milliseconds. A regression that + // failed to terminate on the terminator would instead park on the timeout, + // which assertFinishesWithin catches without a wall-clock threshold. + val session = newSession(channel, terminalTimeoutMs = NeverReachedTimeoutMs) + session.init(basicInit()) + + val processThread = new Thread(() => { + try session.process(echoIn("hello"), emptyFinish).foreach(_ => ()) + catch { case _: Throwable => () } + }, "process-close") + processThread.start() + assert(readyToClose.await(5, TimeUnit.SECONDS), + "worker never received the data batch") + + assertFinishesWithin(10000, "close")(session.close(emptyCancel)) + processThread.join(10000) + assert(!processThread.isAlive, "process thread should terminate after close") + } + + // --------------------------------------------------------------------------- + // Cross-thread delivery (no directExecutor). + // + // The tests above deliver worker responses reentrantly on the caller's thread + // (directExecutor), which is the harness for the reentrancy-hardening paths + // but is NOT how the production Netty transport behaves -- there the gRPC + // callback always runs on a separate event-loop thread. These tests re-run the + // key fast-fail and data-phase paths with cross-thread delivery so the + // session's correctness does not silently depend on reentrant delivery. + // --------------------------------------------------------------------------- + + test("cross-thread delivery: ErrorResponse before InitResponse fails init fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("simulated pre-init error") + .setErrorClass("PreInitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service, directExecutor = false) + // requestObserver is published before the cross-thread ErrorResponse can be + // delivered, so unlike the directExecutor case a Cancel does reach the wire; + // either way init must fail fast rather than park on a session timeout. + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("error") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected init/error in message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("cross-thread delivery: InitResponse with error fails init fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("simulated init failure") + .setErrorClass("InitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected an init-failure message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + test("cross-thread delivery: data-phase ErrorResponse surfaces through the iterator") { + // Worker accepts init, then replies to the first DataRequest with an + // ErrorResponse instead of an echo. Per the protocol the engine follows with + // Cancel and the worker replies CancelResponse. Exercises the data-phase + // race-recovery (sendOrEndOnRacedTerminal) under cross-thread delivery: the + // error terminal can settle while the iterator is mid-send. + val erroredOnce = new AtomicBoolean(false) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + case UdfRequest.RequestCase.DATA => + if (erroredOnce.compareAndSet(false, true)) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("boom in UDF") + .setErrorClass("UdfError").build()).build()).build()).build()) + .build()) + } + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + val it = session.process(echoIn("a", "b", "c"), emptyFinish) + val ex = intercept[GrpcWorkerSessionException] { it.foreach(_ => ()) } + assert(ex.getMessage.contains("boom in UDF"), + s"expected the worker's UDF error to surface, got: ${ex.getMessage}") + assert(ex.executionError != null, "structured executionError should be preserved") + assert(ex.executionError.getUser.getErrorClass == "UdfError") + session.close(emptyCancel) + } + + test("cross-thread delivery: multi-batch echo round-trips in order then finishes") { + val service = new CapturingService() + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + val it = session.process(echoIn("a", "b", "c"), emptyFinish) + val out = it.map(r => new String(r.getData.toByteArray)).toList + assert(out == List("a", "b", "c"), + s"echo worker should return inputs in order over cross-thread delivery, got: $out") + assert(!it.hasNext, "iterator should be exhausted after the FinishResponse terminator") + val handle = handleOf(session) + val termination = session.close(emptyCancel) + // Handle lifecycle on a clean finish: the session is released back to the + // dispatcher exactly once and, because a Finished terminal is salvageable, + // the worker is NOT marked invalid (it stays eligible for reuse). + assert(termination.isInstanceOf[Termination.Finished], + s"a fully drained echo session should settle Finished, got: $termination") + assert(handle.released.get(), "close() must release the worker handle") + assert(!handle.invalidated.get(), + "a clean Finished termination is salvageable; the worker must not be marked invalid") + // Idempotent close(): still released exactly once, still not invalidated. + session.close(emptyCancel) + assert(handle.released.get() && !handle.invalidated.get(), + "a repeat close() must not change the handle lifecycle outcome") + } + + // --------------------------------------------------------------------------- + // An init error surfaces through init()'s exception (carrying the structured + // ExecutionError); per the protocol the engine then sends Cancel and the worker + // replies CancelResponse, so the terminal -- and close()'s return -- is the + // proto terminator Cancelled, not Failed. A genuine transport failure (worker + // never replies to Cancel) still reports TransportFailed faithfully rather than + // collapsing into a clean Cancelled. + // --------------------------------------------------------------------------- + + test("pre-init error: init throws the structured error; close returns Cancelled") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("pre-init boom") + .setErrorClass("PreInitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + // Proto: the engine must Cancel after an init error and the worker + // replies CancelResponse, settling the terminal as Cancelled. + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, terminalTimeoutMs = NeverReachedTimeoutMs) + // The structured error is surfaced through init()'s exception, not the terminal. + val ex = intercept[GrpcWorkerSessionException](session.init(basicInit())) + assert(ex.executionError != null && ex.executionError.getUser.getErrorClass == "PreInitError", + s"init() must throw the structured error, got: ${ex.executionError}") + + // close() returns the proto terminator: Cancelled, not a generic Failed outcome. + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.Cancelled], + s"expected the proto terminator Termination.Cancelled, got: $termination") + } + + test("close that times out without a terminator returns a TransportFailed termination") { + // Worker accepts Init but never replies to Cancel, so close() must give up + // after terminalTimeoutMs and report the failure rather than masquerade as a + // clean Cancelled. This is the "error only during close" case: nothing is + // draining the iterator, so the Termination is the only channel for it. + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } + // Deliberately ignore Cancel: no terminator ever arrives. + }) + val (_, channel) = startServer(service) + val session = newSession(channel, terminalTimeoutMs = 1000L) + session.init(basicInit()) + + val handle = handleOf(session) + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.TransportFailed], + s"a close that times out without a terminator must report TransportFailed, got: $termination") + // Handle lifecycle on an unsalvageable termination: the session is released + // AND the worker is marked invalid so the dispatcher will not recycle a + // worker left in an unknown state by the timed-out stream. + assert(handle.released.get(), "close() must release the worker handle") + assert(handle.invalidated.get(), + "a TransportFailed termination is unsalvageable; the worker must be marked invalid") + } + + // --------------------------------------------------------------------------- + // Worker sends a DataResponse before InitResponse: a protocol violation that + // must fail init fast (fast-fail in responseObserver.onNext DATA branch), + // not enqueue the stray batch and let init() hang for initResponseTimeoutMs. + // --------------------------------------------------------------------------- + + test("worker emits DataResponse before InitResponse: init fails fast") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + // Reply with a DATA response instead of the required InitResponse. + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder() + .setData(ByteString.copyFromUtf8("premature")).build()) + .build()) + } + }) + val (_, channel) = startServer(service) + // Huge init timeout: a regression that enqueued the stray batch instead of + // fast-failing would park init here until the timeout, so finishing quickly + // is the proof the fast path ran -- no wall-clock threshold needed. + val session = newSession(channel, initResponseTimeoutMs = NeverReachedTimeoutMs) + + assertFinishesWithin(10000, "init") { + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("init"), + s"expected an init-failure message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + // --------------------------------------------------------------------------- + // Proto compliance on an init error: the engine MUST send Cancel after an init + // error and the worker replies CancelResponse (udf_message.proto). init() + // sends the Cancel -- from where requestObserver is published, not the response + // callback that may still see it null under reentrant delivery -- and drains + // the CancelResponse, so the terminal is Cancelled. The structured error is + // surfaced through init()'s exception, not the terminal. + // --------------------------------------------------------------------------- + + test("cross-thread InitResponse error: engine sends Cancel, close returns Cancelled") { + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("init failure") + .setErrorClass("InitError").build()).build()).build()).build()) + .build()) + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + // Cross-thread delivery: requestObserver is published before the InitResponse + // error arrives, so the Cancel reaches the wire and its CancelResponse settles + // the terminal Cancelled -- the proto terminator. + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = NeverReachedTimeoutMs) + val ex = intercept[GrpcWorkerSessionException](session.init(basicInit())) + assert(ex.executionError != null && ex.executionError.getUser.getErrorClass == "InitError", + s"init() must throw the structured error, got: ${ex.executionError}") + + // Proto invariant: a Cancel was actually written to the worker on the init error. + assert(service.captured.asScala.exists(r => r.hasControl && r.getControl.hasCancel), + "the engine must send Cancel after an init error (udf_message.proto)") + + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.Cancelled], + s"expected the proto terminator Termination.Cancelled, got: $termination") + } + + // --------------------------------------------------------------------------- + // A throwing input iterator (hasNext or next) must Cancel the stream so the + // worker is not stranded awaiting input, and rethrow to the engine. + // --------------------------------------------------------------------------- + + test("input iterator throwing in hasNext: cancels the stream and surfaces the error") { + val cancelSeen = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + cancelSeen.countDown() + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + // An input iterator whose hasNext throws on first probe (mimics a Spark + // upstream that computes the next element eagerly in hasNext). + val boom = new Iterator[DataRequest] { + override def hasNext: Boolean = throw new RuntimeException("hasNext boom") + override def next(): DataRequest = throw new NoSuchElementException() + } + val it = session.process(boom, emptyFinish) + val ex = intercept[RuntimeException] { it.foreach(_ => ()) } + assert(ex.getMessage.contains("hasNext boom"), + s"the upstream failure must propagate, got: ${ex.getMessage}") + assert(cancelSeen.await(5, TimeUnit.SECONDS), + "a throwing input iterator must Cancel the stream so the worker is not stranded") + session.close(emptyCancel) + } + + test("input iterator returning null: cancels the stream and surfaces the error") { + val service = new CapturingService() + val (_, channel) = startServer(service) + val session = newSession(channel) + session.init(basicInit()) + + val input = Iterator.single(null.asInstanceOf[DataRequest]) + val it = session.process(input, emptyFinish) + val ex = intercept[NullPointerException] { it.hasNext } + assert(ex.getMessage.contains("input iterator returned null"), + s"the invalid input must propagate, got: ${ex.getMessage}") + val requests = service.captured.asScala.toSeq + assert(requests.size == 2 && requests.head.getControl.hasInit && + requests(1).getControl.hasCancel, + s"a null input must produce exactly Init then Cancel, got: $requests") + assert(session.close(emptyCancel).isInstanceOf[Termination.Cancelled]) + } + + // --------------------------------------------------------------------------- + // A throwing finish() thunk (caller-supplied, may run a finish callback) must + // Cancel the stream before rethrowing the callback failure. + // --------------------------------------------------------------------------- + + test("finish thunk throwing: cancels the stream and surfaces the error") { + val cancelSeen = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + cancelSeen.countDown() + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + // finish() throws when the input is exhausted and the iterator tries to + // build the Finish message. Drive the whole iterator inside intercept: + // under cross-thread delivery input flows ahead of output, so finish() may + // fire before the echo of "only" is read back -- the throw can surface on + // any probe, so we must not read a batch outside the intercept. + val throwingFinish: () => Finish = () => throw new RuntimeException("finish boom") + val it = session.process(echoIn("only"), throwingFinish) + val ex = intercept[RuntimeException] { it.foreach(_ => ()) } + assert(ex.getMessage.contains("finish boom"), + s"the finish-thunk failure must propagate, got: ${ex.getMessage}") + assert(cancelSeen.await(5, TimeUnit.SECONDS), + "a throwing finish thunk must Cancel the stream so the worker is not stranded") + session.close(emptyCancel) + } + + test("finish thunk returning null: cancels the stream and surfaces the error") { + val service = new CapturingService() + val (_, channel) = startServer(service) + val session = newSession(channel) + session.init(basicInit()) + + val it = session.process(Iterator.empty, () => null) + val ex = intercept[NullPointerException] { it.hasNext } + assert(ex.getMessage.contains("finish callback returned null"), + s"the invalid finish result must propagate, got: ${ex.getMessage}") + val requests = service.captured.asScala.toSeq + assert(requests.size == 2 && requests.head.getControl.hasInit && + requests(1).getControl.hasCancel, + s"a null finish result must produce exactly Init then Cancel, got: $requests") + assert(session.close(emptyCancel).isInstanceOf[Termination.Cancelled]) + } + + // --------------------------------------------------------------------------- + // close() concurrent with init(): close must serialize with stream opening, and + // a close-settled terminal must wake an init blocked on its response. + // + // Regression guard for onTerminalSettled waking initValue. The terminal here is + // settled by close() ITSELF (the terminalTimeoutMs path in doClose, because the + // worker ignores Cancel), NOT by the response callback -- so unlike every other + // terminator this path does not run handleControl's signalWithoutValue(). Only + // onTerminalSettled can wake the init() blocked on initValue; without it, init() + // stays parked until initResponseTimeoutMs. + // --------------------------------------------------------------------------- + + test("close concurrent with stream opening waits and sends Cancel only after Init") { + val callStartEntered = new CountDownLatch(1) + val releaseCallStart = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, underlying) = startServer(service) + val channel = new DelegatingManagedChannel(underlying) { + override def newCall[ReqT, RespT]( + methodDescriptor: MethodDescriptor[ReqT, RespT], + callOptions: CallOptions): ClientCall[ReqT, RespT] = + new ForwardingClientCall.SimpleForwardingClientCall[ReqT, RespT]( + super.newCall(methodDescriptor, callOptions)) { + override def start( + responseListener: ClientCall.Listener[RespT], + headers: Metadata): Unit = { + callStartEntered.countDown() + assert(releaseCallStart.await(10, TimeUnit.SECONDS), + "timed out waiting to release ClientCall.start") + super.start(responseListener, headers) + } + } + } + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = NeverReachedTimeoutMs) + + val initThrown = new AtomicReference[Throwable]() + val initThread = new Thread(() => { + try session.init(basicInit()) catch { case t: Throwable => initThrown.set(t) } + }, "init-open-close-race") + initThread.start() + assert(callStartEntered.await(5, TimeUnit.SECONDS), "client call never started opening") + + val closeResult = new AtomicReference[Termination]() + val closeThrown = new AtomicReference[Throwable]() + val closeDone = new CountDownLatch(1) + val closeThread = new Thread(() => { + try closeResult.set(session.close(emptyCancel)) + catch { case t: Throwable => closeThrown.set(t) } + finally closeDone.countDown() + }, "close-during-stream-open") + closeThread.start() + + // The close thread must block on requestLock while ClientCall.start is still + // opening the RPC. Before the fix it returned through the requestObserver-null + // path and released the worker while stream creation was still in progress. + val blockedDeadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(5) + while (closeThread.getState != Thread.State.BLOCKED && closeDone.getCount != 0 && + System.nanoTime() < blockedDeadlineNanos) { + Thread.sleep(10L) + } + assert(closeThread.getState == Thread.State.BLOCKED, + s"close must wait for stream creation, state=${closeThread.getState}") + assert(!handleOf(session).released.get(), + "close must not release the worker while the RPC is still opening") + + releaseCallStart.countDown() + closeThread.join(10000) + initThread.join(10000) + assert(!closeThread.isAlive && !initThread.isAlive, + "close and init must both finish after the worker receives Cancel") + assert(closeThrown.get() == null, s"close failed: ${closeThrown.get()}") + assert(initThrown.get().isInstanceOf[GrpcWorkerSessionException], + s"init must fail after concurrent close, got: ${initThrown.get()}") + assert(closeResult.get().isInstanceOf[Termination.Cancelled], + s"worker acknowledgement should settle Cancelled, got: ${closeResult.get()}") + val requests = service.captured.asScala.toSeq + assert(requests.size == 2 && requests.head.getControl.hasInit && + requests(1).getControl.hasCancel, + s"expected exactly Init then Cancel, got: $requests") + } + + test("outgoing send failure aborts the request side and close does not half-close") { + val cancelCalls = new AtomicInteger(0) + val halfCloseCalls = new AtomicInteger(0) + val failNextSend = new AtomicBoolean(true) + val service = new CapturingService() + val (_, underlying) = startServer(service) + val channel = new DelegatingManagedChannel(underlying) { + override def newCall[ReqT, RespT]( + methodDescriptor: MethodDescriptor[ReqT, RespT], + callOptions: CallOptions): ClientCall[ReqT, RespT] = + new ForwardingClientCall.SimpleForwardingClientCall[ReqT, RespT]( + super.newCall(methodDescriptor, callOptions)) { + override def sendMessage(message: ReqT): Unit = { + if (failNextSend.compareAndSet(true, false)) { + throw new RuntimeException("send boom") + } + super.sendMessage(message) + } + + override def cancel(message: String, cause: Throwable): Unit = { + cancelCalls.incrementAndGet() + super.cancel(message, cause) + } + + override def halfClose(): Unit = { + halfCloseCalls.incrementAndGet() + super.halfClose() + } + } + } + val session = newSession(channel) + + val ex = intercept[GrpcWorkerSessionException] { session.init(basicInit()) } + assert(ex.getCause != null && ex.getCause.getMessage.contains("send boom"), + s"the outgoing failure must surface from init, got: $ex") + assert(cancelCalls.get() == 1, + "an outgoing onNext failure must terminate the request side with onError") + + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.TransportFailed], + s"the outgoing failure must settle TransportFailed, got: $termination") + assert(halfCloseCalls.get() == 0, + "close must not invoke onCompleted after the request side was aborted") + assert(cancelCalls.get() == 1, + "close must not terminate an already-aborted request side again") + } + + test("close concurrent with blocked init: init is woken and fails fast") { + val initReceived = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + // Accept the Init stream but never reply -- init() blocks awaiting + // InitResponse. Ignore Cancel too, so the terminator never arrives from + // the worker and close() must settle the terminal on its own timeout. + if (req.hasControl && req.getControl.hasInit) { + initReceived.countDown() + } + }) + val (_, channel) = startServer(service) + // Huge initResponseTimeoutMs: a regression that fails to wake init() on the + // close-settled terminal would park it here for 10 minutes, so finishing + // within the join below is the proof it was woken -- no wall-clock threshold. + // Small terminalTimeoutMs: close() gives up waiting for the (never-arriving) + // CancelResponse quickly and settles TransportFailed itself. + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, terminalTimeoutMs = 1000L) + + val thrown = new AtomicReference[Throwable]() + val initThread = new Thread(() => { + try session.init(basicInit()) catch { case t: Throwable => thrown.set(t) } + }, "blocked-init") + initThread.start() + assert(initReceived.await(5, TimeUnit.SECONDS), + "worker never received the Init request") + + // close() from the test thread settles the terminal (TransportFailed, after + // terminalTimeoutMs of no CancelResponse), which must wake the parked init(). + session.close(emptyCancel) + initThread.join(10000) + assert(!initThread.isAlive, + "init() parked on initResponseTimeoutMs; the close-settled terminal did not " + + "wake it (onTerminalSettled must signal initValue)") + val t = thrown.get() + assert(t != null, "init() must surface an init-failure exception") + assert(t.isInstanceOf[GrpcWorkerSessionException], + s"expected GrpcWorkerSessionException, got ${if (t == null) "null" else t.getClass.getName}") + } + + // --------------------------------------------------------------------------- + // Terminator-carried callback errors: a FinishResponse / CancelResponse may + // carry an error raised by the finish / cancel callback (udf_message.proto). + // The result iterator must surface it, and a prior data-phase ErrorResponse + // must take precedence over the terminator's own callback error + // (throwIfTerminalError / responseError). + // --------------------------------------------------------------------------- + + test("FinishResponse carrying a callback error: iterator throws it") { + // Worker echoes the single batch, then finishes with a FinishResponse whose + // error field is set (a finish-callback failure). No data-phase ErrorResponse + // precedes it, so the terminator's own error is the one surfaced. + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasFinish) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setFinish( + FinishResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("finish callback boom") + .setErrorClass("FinishCallbackError").build()).build()) + .build()).build()).build()) + } + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + val it = session.process(echoIn("hello"), emptyFinish) + val ex = intercept[GrpcWorkerSessionException] { it.foreach(_ => ()) } + assert(ex.getMessage.contains("finish callback boom"), + s"the FinishResponse callback error must surface, got: ${ex.getMessage}") + assert(ex.executionError != null && + ex.executionError.getUser.getErrorClass == "FinishCallbackError", + "the structured finish-callback error should be preserved") + session.close(emptyCancel) + } + + test("data-phase ErrorResponse takes precedence over a CancelResponse callback error") { + // Worker replies to the first DataRequest with an ErrorResponse (data-phase + // failure), then -- per the protocol, the engine sends Cancel -- answers with + // a CancelResponse that ALSO carries a (cancel-callback) error. The iterator + // must surface the original data-phase error, not the terminator's, per + // responseError's precedence rule. + val erroredOnce = new AtomicBoolean(false) + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } else if (c.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("cancel callback boom") + .setErrorClass("CancelCallbackError").build()).build()) + .build()).build()).build()) + } + case UdfRequest.RequestCase.DATA => + if (erroredOnce.compareAndSet(false, true)) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setError( + ErrorResponse.newBuilder().setError( + ExecutionError.newBuilder().setUser( + UserError.newBuilder().setMessage("data-phase boom") + .setErrorClass("DataPhaseError").build()).build()).build()) + .build()).build()) + } + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + val session = newSession(channel, terminalTimeoutMs = 5000L) + session.init(basicInit()) + + val it = session.process(echoIn("a", "b", "c"), emptyFinish) + val ex = intercept[GrpcWorkerSessionException] { it.foreach(_ => ()) } + assert(ex.getMessage.contains("data-phase boom"), + s"the prior data-phase error must take precedence, got: ${ex.getMessage}") + assert(ex.executionError != null && + ex.executionError.getUser.getErrorClass == "DataPhaseError", + s"expected the data-phase error to be preserved, got: ${ex.executionError}") + session.close(emptyCancel) + } + + // --------------------------------------------------------------------------- + // Data-phase output timeout: after Finish, if the worker goes silent without a + // terminator, advance()'s per-poll wait (terminalTimeoutMs) must expire and the + // iterator must surface a TransportFailed rather than block forever. + // --------------------------------------------------------------------------- + + test("data-phase output timeout: silent worker after Finish surfaces a timeout") { + // Worker accepts Init and echoes data, but never sends a terminator and + // ignores Finish/Cancel -- so once input is exhausted the iterator's output + // poll has nothing to drain and must time out on terminalTimeoutMs. + val service = new CapturingService( + onRequest = (req, resp) => { + req.getRequestCase match { + case UdfRequest.RequestCase.CONTROL => + val c = req.getControl + if (c.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } + // Finish and Cancel ignored: no terminator ever arrives. + case UdfRequest.RequestCase.DATA => + resp.onNext(UdfResponse.newBuilder() + .setData(DataResponse.newBuilder().setData(req.getData.getData).build()) + .build()) + case _ => () + } + }) + val (_, channel) = startServer(service, directExecutor = false) + // Small terminalTimeoutMs: the poll gives up quickly; assertFinishesWithin + // guards against a regression that blocks the iterator indefinitely. + val session = newSession(channel, terminalTimeoutMs = 1000L) + session.init(basicInit()) + + val it = session.process(echoIn("hello"), emptyFinish) + assertFinishesWithin(10000, "output-timeout") { + val ex = intercept[GrpcWorkerSessionException] { it.foreach(_ => ()) } + assert(ex.getMessage.toLowerCase(Locale.ROOT).contains("timed out") || + ex.getMessage.toLowerCase(Locale.ROOT).contains("stream failed"), + s"expected a timeout/transport-failure message, got: ${ex.getMessage}") + } + session.close(emptyCancel) + } + + // --------------------------------------------------------------------------- + // Interrupt handling (cancelled query / killed task). An interrupt while + // blocked on init or pulling input sends a best-effort Cancel and waits + // interruptCancelTimeoutMs for the CancelResponse: a responsive worker settles + // a clean, salvageable Cancelled; an unresponsive one falls back to Interrupted + // and is invalidated because no acknowledgement or liveness proof makes it safe + // to reuse. + // --------------------------------------------------------------------------- + + test("interrupt during init: responsive worker settles Cancelled, worker salvageable") { + // Worker never sends InitResponse (so init() blocks), but DOES reply to + // Cancel -- so the interrupt's bounded drain observes the CancelResponse and + // settles a clean Cancelled. + val initReceived = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + initReceived.countDown() // block: no InitResponse + } else if (req.hasControl && req.getControl.hasCancel) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setCancel( + CancelResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service, directExecutor = false) + // Huge initResponseTimeoutMs so only the interrupt can end the init wait; + // generous interruptCancelTimeoutMs so the ack is observed. + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, interruptCancelTimeoutMs = 5000L) + val handle = handleOf(session) + + val thrown = new AtomicReference[Throwable]() + val initThread = new Thread(() => { + try session.init(basicInit()) catch { case t: Throwable => thrown.set(t) } + }, "interrupt-init-responsive") + initThread.start() + assert(initReceived.await(5, TimeUnit.SECONDS), "worker never received Init") + initThread.interrupt() + initThread.join(10000) + assert(!initThread.isAlive, "init() must return after the interrupt") + assert(thrown.get().isInstanceOf[InterruptedException], + s"init() must rethrow InterruptedException, got: ${thrown.get()}") + + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.Cancelled], + s"a responsive worker's ack should settle Cancelled, got: $termination") + assert(!handle.invalidated.get(), + "an acknowledged Cancelled outcome is salvageable; the worker must not be invalidated") + assert(handle.released.get(), "close() must release the worker handle") + } + + test("interrupt during init: unresponsive worker settles Interrupted and is invalidated") { + // Worker never sends InitResponse and ignores Cancel -- so the interrupt's + // bounded drain times out and the session falls back to Interrupted. + val initReceived = new CountDownLatch(1) + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + initReceived.countDown() // block: no InitResponse, and Cancel ignored + } + }) + val (_, channel) = startServer(service, directExecutor = false) + // Short interruptCancelTimeoutMs so the fallback fires quickly; huge + // terminalTimeoutMs so a regression that used the wrong timeout would stall. + val session = newSession(channel, + initResponseTimeoutMs = NeverReachedTimeoutMs, + terminalTimeoutMs = NeverReachedTimeoutMs, + interruptCancelTimeoutMs = 500L) + val handle = handleOf(session) + + val thrown = new AtomicReference[Throwable]() + val initThread = new Thread(() => { + try session.init(basicInit()) catch { case t: Throwable => thrown.set(t) } + }, "interrupt-init-unresponsive") + initThread.start() + assert(initReceived.await(5, TimeUnit.SECONDS), "worker never received Init") + initThread.interrupt() + // Bounded by interruptCancelTimeoutMs (500ms), not terminalTimeoutMs (10min): + // if init parked on the wrong timeout this join would fail. + initThread.join(10000) + assert(!initThread.isAlive, + "init() must return within interruptCancelTimeoutMs of the interrupt") + assert(thrown.get().isInstanceOf[InterruptedException], + s"init() must rethrow InterruptedException, got: ${thrown.get()}") + + val termination = session.close(emptyCancel) + assert(termination.isInstanceOf[Termination.Interrupted], + s"an unresponsive worker should settle Interrupted, got: $termination") + assert(handle.invalidated.get(), + "an unacknowledged Interrupted outcome must invalidate the worker") + assert(handle.released.get(), "close() must release the worker handle") + } + + test("interrupt while pulling input: bounded Cancel drain settles Interrupted") { + // Worker accepts Init but ignores Cancel. The interrupted input pull must use + // interruptCancelTimeoutMs rather than leaving close() to wait terminalTimeoutMs. + val service = new CapturingService( + onRequest = (req, resp) => { + if (req.hasControl && req.getControl.hasInit) { + resp.onNext(UdfResponse.newBuilder().setControl( + UdfControlResponse.newBuilder().setInit( + InitResponse.getDefaultInstance).build()).build()) + } + }) + val (_, channel) = startServer(service) + val session = newSession(channel, + terminalTimeoutMs = NeverReachedTimeoutMs, + interruptCancelTimeoutMs = 500L) + val handle = handleOf(session) + session.init(basicInit()) + + val inputNextEntered = new CountDownLatch(1) + val neverReleaseInput = new CountDownLatch(1) + val input = new Iterator[DataRequest] { + override def hasNext: Boolean = true + override def next(): DataRequest = { + inputNextEntered.countDown() + neverReleaseInput.await() + throw new IllegalStateException("blocked input unexpectedly resumed") + } + } + val thrown = new AtomicReference[Throwable]() + val termination = new AtomicReference[Termination]() + val processThread = new Thread(() => { + try { + session.process(input, emptyFinish).hasNext + } catch { + case t: Throwable => thrown.set(t) + } finally { + termination.set(session.close(emptyCancel)) + } + }, "interrupt-input-pull") + processThread.setDaemon(true) + processThread.start() + assert(inputNextEntered.await(5, TimeUnit.SECONDS), "input.next() was never entered") + + processThread.interrupt() + // Bounded by interruptCancelTimeoutMs (500ms), not terminalTimeoutMs (10min). + processThread.join(10000) + assert(!processThread.isAlive, + "input interruption must unwind within interruptCancelTimeoutMs") + assert(thrown.get().isInstanceOf[InterruptedException], + s"processing must rethrow InterruptedException, got: ${thrown.get()}") + assert(termination.get().isInstanceOf[Termination.Interrupted], + s"an unresponsive worker should settle Interrupted, got: ${termination.get()}") + assert(service.captured.asScala.exists(r => r.hasControl && r.getControl.hasCancel), + "interrupting an input pull must send Cancel") + assert(handle.invalidated.get(), + "an unacknowledged Interrupted outcome must invalidate the worker") + assert(handle.released.get(), "close() must release the worker handle") + } +} diff --git a/ui-test/package-lock.json b/ui-test/package-lock.json index 72e93e20a57f0..447a8c5144b79 100644 --- a/ui-test/package-lock.json +++ b/ui-test/package-lock.json @@ -1058,9 +1058,9 @@ } }, "node_modules/@jest/reporters/node_modules/brace-expansion": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.2.tgz", - "integrity": "sha512-w5JZcKgdhDOgOwm8H+KgbosopHMuGcl6qbulwjtz3SM7I7P3yW1eAjzMPLrIE+NQ9vjgANKHWeMHnrT0OXW1oA==", + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.4.tgz", + "integrity": "sha512-hGfVzPxthbf3+2yjg/RBs60cB0FhqBS/zvdV/4wn4/BmN0bNMMHPc4V/BbFieqf1TKAGGAHnY4eSjajCl0f2Xg==", "dev": true, "license": "MIT", "dependencies": { @@ -2049,16 +2049,16 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.7", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.7.tgz", - "integrity": "sha512-7oFy703dxfY3/NLxC1fh2SUCQ0H9rmAY+5EpDVfXjUTTs+HEwR2nYaqLv+GWcTsumwxPfiz6CzCNkwXwBUwqCA==", + "version": "5.0.9", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz", + "integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==", "dev": true, "license": "MIT", "dependencies": { "balanced-match": "^4.0.2" }, "engines": { - "node": "18 || 20 || >=22" + "node": "20 || >=22" } }, "node_modules/brace-expansion/node_modules/@isaacs/cliui": { @@ -3131,9 +3131,9 @@ } }, "node_modules/jest-config/node_modules/brace-expansion": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.2.tgz", - "integrity": "sha512-w5JZcKgdhDOgOwm8H+KgbosopHMuGcl6qbulwjtz3SM7I7P3yW1eAjzMPLrIE+NQ9vjgANKHWeMHnrT0OXW1oA==", + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.4.tgz", + "integrity": "sha512-hGfVzPxthbf3+2yjg/RBs60cB0FhqBS/zvdV/4wn4/BmN0bNMMHPc4V/BbFieqf1TKAGGAHnY4eSjajCl0f2Xg==", "dev": true, "license": "MIT", "dependencies": { @@ -3539,9 +3539,9 @@ } }, "node_modules/jest-runtime/node_modules/brace-expansion": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.2.tgz", - "integrity": "sha512-w5JZcKgdhDOgOwm8H+KgbosopHMuGcl6qbulwjtz3SM7I7P3yW1eAjzMPLrIE+NQ9vjgANKHWeMHnrT0OXW1oA==", + "version": "2.1.4", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.1.4.tgz", + "integrity": "sha512-hGfVzPxthbf3+2yjg/RBs60cB0FhqBS/zvdV/4wn4/BmN0bNMMHPc4V/BbFieqf1TKAGGAHnY4eSjajCl0f2Xg==", "dev": true, "license": "MIT", "dependencies": { @@ -3785,9 +3785,9 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "funding": [ { @@ -4669,9 +4669,9 @@ } }, "node_modules/test-exclude/node_modules/brace-expansion": { - "version": "1.1.16", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.16.tgz", - "integrity": "sha512-IDw48K2/2kRkg9LdJxurvq3lV3aBgq0REY89duEqFRthjlPdXHKMj7EnQOXVckxzgisinf3nHfrcE2FufFLXMw==", + "version": "1.1.18", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.18.tgz", + "integrity": "sha512-Edep/X9fGqVNmzKBVsDYIOtD+z1tuezV70LBjdCst9Tqu76lsnvRiZ6oTic1n+/BIwX6QDGAO94PN4N2SADvtw==", "dev": true, "license": "MIT", "dependencies": {